olivame/unifolm-world-model-action
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
64 Commits 12 Branches 0 Tags
second
Commit Graph

64 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
olivame
57ba85d147 KV 融合实现完成。改动总结： 速度微弱提升psnr略微上升 
attention.py — 3处改动：
  1. __init__ 添加 _kv_fused = False 标志
  2.新增 fuse_kv() 方法：将 to_k + to_v → to_kv，同时处理 _ip/_as/_aa 辅助 KV 对
  2. bmm_forward 两个分支加_kv_fused 判断，用to_kv().chunk(2, dim=-1) 替代分别调用
 2026-02-10 18:15:52 +00:00
olivame
2cef3e9e45 ├─────┼─────────────────────────────────┼───────────────────────┼───────────────────┤ 
│ 1   │ CUDA Stream 预创建              │ wma_model.py          │ 50次 → 0次        │
  ├─────┼─────────────────────────────────┼───────────────────────┼───────────────────┤
  │ 2   │ noise buffer 预分配             │ ddim.py               │ 50次 alloc → 0次  │
  ├─────┼─────────────────────────────────┼───────────────────────┼───────────────────┤
  │ 3   │ global_feature expand提到循环外 │ conditional_unet1d.py │ ~700次 → ~100次   │
  ├─────┼─────────────────────────────────┼───────────────────────┼───────────────────┤
  │ 4   │ alpha/sigma dtype 预转换        │ ddim.py               │ 200次 .to() → 0次 │
效果不算特别明显
 2026-02-10 13:40:52 +00:00
olivame
a09d35ae5b - state_unet 放到一个独立的 CUDA stream 上执行 
- action_unet 在默认 stream 上同时执行
  - 用 wait_stream 确保两者都完成后再返回
两个 1D UNet 输入完全独立，共享的 hs_a 和 context_action 都是只读的。GPU 利用率只有 ~31%，小张量 kernel 不会打满 GPU，两个 stream 可以真正并行。
 2026-02-10 10:47:10 +00:00
olivame
db848bca01 profile 结果 2026-02-10 07:02:20 +00:00
olivame
b0ebb7006e 添加三层迭代级性能分析工具 profile_iteration.py 
Layer1: CUDA Events 精确测量每个itr内10个阶段耗时
Layer2: torch.profiler GPU timeline trace
Layer3: CSV输出支持A/B对比

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
 2026-02-10 05:42:11 +00:00
olivame
125b85ce68 实现fs_embed 缓存，收益不明显，精度不降低 2026-02-09 18:49:44 +00:00
olivame
0b3b0e534a 复用 DDIMSampler + make_schedule微弱提升 2026-02-09 18:26:39 +00:00
olivame
6dca3696d8 实现了Context 预计算和缓存功能，提升了采样效率。 psnr不下降 2026-02-09 17:42:47 +00:00
olivame
f192c8aca9 添加CrossAttention kv缓存，减少重复计算，提升性能，psnr=31.8022 dB 2026-02-09 17:04:23 +00:00
olivame
4288c9d8c9 减少了一路视频vae解码 2026-02-09 16:48:16 +00:00
olivame
a2cd34dd51 1. einsum('b i d, b j d -> b i j') → torch.bmm(q, k.transpose(-1,-2)) — 直接映射 rocBLAS batched GEMM 
2. baddbmm 把 scale 融合进 GEMM，少一次 kernel launch
3. 第二个 einsum 同理换torch.bm
每一轮加速1到两秒
 2026-02-08 18:54:48 +00:00
olivame
7338cc384a ddim.py — torch.float16 → torch.bfloat16，修复 dtype 不匹配 
attention.py — 4 处 softmax 都包裹了 torch.amp.autocast('cuda', enabled=False)，阻止 autocast 将 bf16 提升到 fp32
 2026-02-08 17:02:05 +00:00
olivame
f86ab51a04 全链路 bf16 混合精度修正与 UNet FLOPS profiling 
- GroupNorm/LayerNorm bypass autocast，消除 bf16→fp32→bf16 转换开销
  - DDIM 调度系数 cast 到输入 dtype，attention mask 直接用 bf16 分配
  - alphas_cumprod 提升到 float64 保证数值精度
  - SinusoidalPosEmb 输出 dtype跟随模型精度
  - 新增 profile_unet.py 脚本及FLOPS 分析结果
  - 启用 TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL
  - case1 PSNR: 30.45 → 30.24（bf16 精度预期内波动）
 2026-02-08 16:01:30 +00:00
olivame
75c798ded0 DDIM loop 内小张量分配优化，attention mask 缓存到 GPU 2026-02-08 14:20:48 +00:00
olivame
e588182642 修复混合精度vae相关的配置错误，确保在推理阶段正确使用了混合精度模型，并且导出了正确精度的检查点文件。 2026-02-08 12:35:59 +00:00
olivame
e6c55a648c 所有case的baseline，amd版本的ground truth都上传了 2026-02-08 09:42:14 +00:00
olivame
fffc5a9956 init 2026-02-08 03:29:15 +00:00
UniGen-X
cbaebc016f Update README.md 2025-10-01 10:13:04 +08:00
hengguo
05d2d82236 update readme 2025-09-23 16:57:54 +08:00
yuchen-x
50e8c3ed55 update readme_cn 2025-09-23 16:15:59 +08:00
yuchen-x
ddb5848d86 Merge branch 'main' of github.com:unitreerobotics/unifolm-world-model-action 2025-09-23 16:11:47 +08:00
yuchen-x
cb0cf4a353 updata readme_cn 2025-09-23 16:11:29 +08:00
hengguo
54f61a4336 update readme_cn 2025-09-23 16:00:53 +08:00
yuchen-x
118ada7c35 update readme 2025-09-23 15:22:27 +08:00
yuchen-x
8d5546d322 update readme 2025-09-23 15:21:12 +08:00
yuchen-x
eccd1680c1 update readme 2025-09-23 15:19:02 +08:00
yuchen-x
f12b478265 upload real-robot deployment code 2025-09-23 15:13:22 +08:00
yuchen-x
5dcd1ca503 fix a typo on COLUMNS definition 2025-09-22 17:33:28 +08:00
UniGen-X
7b4d383611 Update README.md 2025-09-21 17:34:40 +08:00
UniGen-X
733e228bb8 Update README.md 2025-09-21 17:19:28 +08:00
UniGen-X
e9c60f6e62 Update README_cn.md 2025-09-19 10:18:38 +08:00
UniGen-X
2d4d79ab3a Update README.md 2025-09-19 10:18:07 +08:00
UniGen-X
1f21fe7fd8 Update README.md 2025-09-19 10:16:58 +08:00
UniGen-X
a57037ab03 Update README_cn.md 2025-09-17 11:45:09 +08:00
UniGen-X
be43dfef9d Update README.md 2025-09-17 11:44:32 +08:00
UniGen-X
884dcce130 Update README_cn.md 2025-09-17 10:49:42 +08:00
UniGen-X
29f3101a1f Update README.md 2025-09-17 10:49:20 +08:00
UniGen-X
712a289c28 Update README_cn.md 2025-09-17 10:44:52 +08:00
UniGen-X
c0fad43420 Update README.md 2025-09-17 10:44:30 +08:00
UniGen-X
3b83374922 Merge pull request #8 from hu-po/patch-1 
Fix typo in Acknowledgement section of README
 2025-09-17 10:31:31 +08:00
UniGen-X
6a33fddb99 Update README_cn.md 2025-09-17 10:29:47 +08:00
UniGen-X
a0010c4036 Update README.md 2025-09-17 10:27:02 +08:00
UniGen-X
a1f6430c38 Update README.md 2025-09-17 10:26:51 +08:00
hu-po
ee45fe75e2 Fix typo in Acknowledgement section of README 2025-09-16 13:22:33 -05:00
yuchen-x
9c9942f1d4 upload eval_utils.py file 2025-09-16 21:48:15 +08:00
UniGen-X
a457f2f049 Update README.md 2025-09-16 20:55:49 +08:00
UniGen-X
ea7260a7e0 Update README_cn.md 2025-09-15 18:28:52 +08:00
UniGen-X
a06384d865 Update README.md 2025-09-15 18:28:19 +08:00
yuchen-x
43a7750825 update python version 2025-09-15 18:25:26 +08:00
UniGen-X
f44065f889 Update README_cn.md 2025-09-15 18:24:38 +08:00