- state_unet 放到一个独立的 CUDA stream 上执行

- action_unet 在默认 stream 上同时执行
  - 用 wait_stream 确保两者都完成后再返回
两个 1D UNet 输入完全独立,共享的 hs_a 和 context_action 都是只读的。GPU 利用率只有 ~31%,小张量 kernel 不会打满 GPU,两个 stream 可以真正并行。
This commit is contained in:
qhy
2026-02-10 21:41:48 +08:00
parent ff43432ef9
commit dcbcb2c377
4 changed files with 28 additions and 19 deletions

2
.envrc Normal file
View File

@@ -0,0 +1,2 @@
eval "$(conda shell.bash hook 2>/dev/null)"
conda activate unifolm-wma