Repository navigation
[Feat] DreamZero fix tp & cross attn cache - #4154
Conversation
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
…edRMSNorm DreamZero per-step inference is decode-like (small query, long KV cache), so the many tiny qk-norm all-reduces dominate TP latency and cause the known TP scaling gap. This first cut switches DistributedRMSNorm from a raw torch.distributed.all_reduce on the TP process group to vLLM's tensor_model_parallel_all_reduce, which routes small tensors through the custom all-reduce / symmetric-memory fast path. It also drops the redundant .clone() by taking the collective return value, matching the convention already used in dreamid_omni/wan2_2 and wan2_2_transformer. Numerically identical: the rsqrt/weight formula is unchanged; only the collective mechanism differs. Signed-off-by: dph97 <nenbaying@163.com>
…duce Self-attention applies DistributedRMSNorm to both q and k every step, each issuing its own tiny TP all-reduce of the per-token sum-of-squares. In DreamZero's decode-like per-step forward these latency-bound micro-collectives dominate (2 per block x 40 layers x every denoise step). Add fused_qk_rms_norm(), which packs both sum-of-squares into one tensor and reduces once (2 collectives -> 1), then applies each norm's own RMS and weight shard. DistributedRMSNorm is refactored to expose _local_sum_sq()/_scale() so the fused path reuses the exact same numerics. Numerically identical to norm_q(q), norm_k(k): all-reduce is elementwise, so packing along the last dim reduces each slice independently with the same fp32 accumulation. Only self-attention is fused (q and k share shape there); cross- attention is left as-is because its k is computed once and cached, so it already issues a single all-reduce in steady state. Falls back to independent application when qk_norm is disabled (nn.Identity). Signed-off-by: dph97 <nenbaying@163.com>
…reduce Lock in that fused_qk_rms_norm is numerically identical to applying norm_q/ norm_k separately, and that it issues exactly one (fused) all-reduce over the packed q+k pair. Also covers the nn.Identity fallback (qk_norm=False). CPU-only (core_model + cpu): the TP world size and the collective are mocked so a single process simulates one rank (all-reduce = identity), which still exercises the cat/split and count-scaling logic. Distinct random norm weights ensure a swapped q/k slice would fail. Signed-off-by: dph97 <nenbaying@163.com>
…lelLinear
CausalWanSelfAttention projected q/k/v with three separate ColumnParallelLinear.
Since all three read the same hidden states, replace them with a single
QKVParallelLinear and split the fused output, cutting three GEMMs/launches per
self-attn down to one — a win for DreamZero's decode-like per-step forward. The
existing fused_qk_rms_norm still applies to the split q/k. cross_attn is left
untouched (q comes from x, k/v from context — not fusible).
The checkpoint still ships separate q/k/v weights, so load_weights routes each
.self_attn.{q,k,v}.{weight,bias} into the packed .self_attn.qkv param via
param.weight_loader(param, tensor, shard_id), matching the stacked_params_mapping
convention used by dreamid_omni and wan2_2_s2v. cross_attn q/k/v are not matched
(needle is ".self_attn.q." etc.).
Adds a CPU equivalence test (world_size=1 TP group) proving that loading separate
q/k/v via shard id then splitting reproduces the three independent projections,
plus a guard that the shard ids are not transposed.
Signed-off-by: dph97 <nenbaying@163.com>
Document that WanI2VCrossAttention recomputes k_img/v_img every step while the text k/v are cached. context_img is session-invariant, so k_img/v_img are cacheable the same way — a follow-up W5/W6 optimization. Comment only, no behavior change. Signed-off-by: dph97 <nenbaying@163.com>
9559fc3 to
396afbf
Compare
context_img (= img_emb(clip_feature)) is session-invariant: state.clip_feas is set once at current_start_frame==0 and only cleared on reset (which also rebuilds crossattn_cache). So the image cross-attention k_img/v_img can be cached on the first forward and reused, exactly like the text k/v already are. Only q depends on the per-step x and is still recomputed every step. This removes the per-step norm_k_img all-reduce and the k_img/v_img projection GEMMs in steady state, with no change to numerical results. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: dph97 <nenbaying@163.com>
crossattn_cache reached _forward_inference but was never passed down to _forward_blocks or the per-layer block call, so cross-attention always hit the None branch and recomputed text k/v and image k_img/v_img every step. This made both the text k/v cache and the k_img/v_img cache dead code. Thread crossattn_cache (indexed per layer, like kv_cache) from _forward_inference -> _forward_blocks -> block -> cross_attn so the caches are actually used after the first forward of a session. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: dph97 <nenbaying@163.com>
Lock in the cross-attention cache behaviour and the crossattn_cache threading fix: - WanI2VCrossAttention.forward initialises the cache (text k/v + image k_img/v_img) on the first call and reuses it on later calls; recomputes every call when no cache is provided. - CausalWanModel._forward_blocks threads crossattn_cache down to each per-layer block (regression guard: the cache was previously dead code because the dict never reached the block). CPU-only: __init__ (which needs a TP group) is bypassed and submodules are stubbed to count calls and exercise only the caching/threading logic. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: dph97 <nenbaying@163.com>
…test Remove test_forward_blocks_threads_crossattn_cache_to_each_block and its helpers (_RecordingBlock, _make_model_with_recording_block) and the now-unused CausalWanModel import. Keep the two cross-attention cache behaviour tests (init-and-reuse, and recompute-when-no-cache). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: dph97 <nenbaying@163.com>
…-attn fusion Two silent-failure guards for the self-attn QKV fusion path: - After building the fused QKVParallelLinear, assert total_num_kv_heads == total_num_heads. The forward splits qkv into three equal q/k/v shards, which is only correct without GQA; a future GQA checkpoint would otherwise misalign k/v silently. - In fused_qk_rms_norm, assert q.shape == k.shape. The fused path reduces one packed sum-of-squares and reuses q's token width as the RMS count for both q and k, so a shape mismatch (same batch/seq, different feature width) would silently normalize k with the wrong count. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: dph97 <nenbaying@163.com>
|
@RuixiangMa Hi, could you please help review and merge the code? |
|
LGTM |
In short, TP + CFG parallelism is sufficient for both configurations; no additional parallelism settings are required. |
|
@TKONIY PTAL for a final check |
TP Performance Reproduction on PCIe + Pro6000All durations are milliseconds. PR-body seconds are converted to ms.
Local TP2 PCIe improvement from this PR:
The absolute latencies are much lower than the PR body on this host, so I would |
|
LGTM. Implementation is clean. The TP performance on PCIe node performs much better then before. |
Thanks for helping review the code. |
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> Signed-off-by: z00806815 <zhuhonghan@huawei.com>
|
CI failure is not related |
Bring PR vllm-project#4534 up to date with upstream main (DreamZero CUDA graph / torch.compile / DiT step-cache vllm-project#4213, tp & cross-attn cache fix vllm-project#4154, GR00T, stepcache, etc.). Conflict resolution: - pipeline_dreamzero.py: * diffuse loop: adopt upstream's run_dit (DiT step-cache) wrapper, keep BDE's _kv_get / _kv_get_cross proxies for all KV reads. * reset: adopt upstream's reset_reason (session/inference) model; route both session and inference (window-boundary) resets through _kv_reset so the BDE pool resets too (parity). _kv_reset gains clear_video_latents. * take upstream's state.append_video_latents(video_out). * cross-attn: BDE now uses the model-local cross-attn cache (upstream vllm-project#4154 added k_img/v_img image-token caching that BDE's text-only pool lacked); _kv_get_cross returns model-local, _kv_populate_cross is a no-op. BDE's self-attn paged pool is unchanged. - causal_wan_model.py: silent auto-merge duplicated the crossattn_cache arg in _forward_blocks (upstream vllm-project#4154 and our wiring-gap fix are the same change) — de-duplicated to one param. - video_export_worker.py: keep both gpu_mem_stats (ours) and _latents_to_uint8_frames (upstream); switch gpu_mem_stats off the now-banned torch.cuda.* to torch.accelerator.* (TID251). - bde_perf_compare.py: update to upstream's new _build_observations signature. Validated on H200 (DreamZero-DROID): tests/bde 77 passing; whole package compiles; KV-off (base engine) and BDE_KV_ENABLE=1 both run on the merged code and stay bit-exact (max_abs_diff=0, PSNR=inf). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Restore BDE's engine-managed cross-attn (§5.2) after the main rebase. Upstream vllm-project#4154 added image-token k_img/v_img caching to the I2V cross-attn; BDE's pool only held text k/v, so it had been falling back to the model-local cache. Now BDE caches both: - BDEKVCache: add per-layer _cross_k_img/_cross_v_img pools (2, 257, heads, dim) for I2V; write_cross_kv/read_cross_kv handle the image variants (read returns k_img/v_img when the image pool is allocated). - runner: derive cross_attn_img_length (257 for i2v, 0 for t2v) and thread it through build_kv_cache. - pipeline: _kv_get_cross routes back to the pool; _kv_populate_cross projects text (text_embedding) and image (img_emb(clip_feature)) K/V, and is moved into the image-encode block so clip_feas is available (runs on the first forward and after each window-boundary reset). Validated on H200 (DreamZero-DROID): tests/bde 77 passing; bit-exact parity (KV-off vs BDE_KV_ENABLE=1, max_abs_diff=0, PSNR=inf) with the cross-attn k_img/v_img pool active. Image pool ~0.5 GiB. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Bring PR vllm-project#4534 up to date with upstream main (DreamZero CUDA graph / torch.compile / DiT step-cache vllm-project#4213, tp & cross-attn cache fix vllm-project#4154, GR00T, stepcache, etc.). Conflict resolution: - pipeline_dreamzero.py: * diffuse loop: adopt upstream's run_dit (DiT step-cache) wrapper, keep BDE's _kv_get / _kv_get_cross proxies for all KV reads. * reset: adopt upstream's reset_reason (session/inference) model; route both session and inference (window-boundary) resets through _kv_reset so the BDE pool resets too (parity). _kv_reset gains clear_video_latents. * take upstream's state.append_video_latents(video_out). * cross-attn: BDE now uses the model-local cross-attn cache (upstream vllm-project#4154 added k_img/v_img image-token caching that BDE's text-only pool lacked); _kv_get_cross returns model-local, _kv_populate_cross is a no-op. BDE's self-attn paged pool is unchanged. - causal_wan_model.py: silent auto-merge duplicated the crossattn_cache arg in _forward_blocks (upstream vllm-project#4154 and our wiring-gap fix are the same change) — de-duplicated to one param. - video_export_worker.py: keep both gpu_mem_stats (ours) and _latents_to_uint8_frames (upstream); switch gpu_mem_stats off the now-banned torch.cuda.* to torch.accelerator.* (TID251). - bde_perf_compare.py: update to upstream's new _build_observations signature. Validated on H200 (DreamZero-DROID): tests/bde 77 passing; whole package compiles; KV-off (base engine) and BDE_KV_ENABLE=1 both run on the merged code and stay bit-exact (max_abs_diff=0, PSNR=inf). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Restore BDE's engine-managed cross-attn (§5.2) after the main rebase. Upstream vllm-project#4154 added image-token k_img/v_img caching to the I2V cross-attn; BDE's pool only held text k/v, so it had been falling back to the model-local cache. Now BDE caches both: - BDEKVCache: add per-layer _cross_k_img/_cross_v_img pools (2, 257, heads, dim) for I2V; write_cross_kv/read_cross_kv handle the image variants (read returns k_img/v_img when the image pool is allocated). - runner: derive cross_attn_img_length (257 for i2v, 0 for t2v) and thread it through build_kv_cache. - pipeline: _kv_get_cross routes back to the pool; _kv_populate_cross projects text (text_embedding) and image (img_emb(clip_feature)) K/V, and is moved into the image-encode block so clip_feas is available (runs on the first forward and after each window-boundary reset). Validated on H200 (DreamZero-DROID): tests/bde 77 passing; bit-exact parity (KV-off vs BDE_KV_ENABLE=1, max_abs_diff=0, PSNR=inf) with the cross-attn k_img/v_img pool active. Image pool ~0.5 GiB. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
Purpose
#4127 cc @TKONIY
DreamZero's per-step inference is decode-like: each control step runs a small query over a long, persistent KV cache, 40 layers × N denoise steps, in eager mode. Per-step latency bounds the control frequency, so it's the product metric.
The TP baseline is correct but doesn't scale — the per-step forward is dominated by many small, latency-bound collectives and per-projection kernel launches in the self-attention path. This PR is W6 (TP performance fix) from the DreamZero performance RFC. It removes the highest-ROI, lowest-risk TP overheads in the self-attn projection path, without changing model outputs (the QKV GEMM is mathematically identical, not bit-identical).
Change
DistributedRMSNorm— replace rawtorch.distributed.all_reduceon the TP group with vLLM'stensor_model_parallel_all_reduce(custom-AR / symmetric-memory fast path for tiny tensors).ColumnParallelLinear(q/k/v) iCausalWanSelfAttentionwith oneQKVParallelLinear+ split (3 GEMMs/launches → 1).k_img/v_imgevery step althoughcontext_imgis session-invariant; they are cacheable like the text k/v.Test Plan
Profiled the DreamZero-DROID multi-GPU TP path with a one-chunk Torch Profiler setup to separate communication and computation overhead, identify the causes of the original slow TP computation, and evaluate the optimization results.
Test Result
TP performance
Single GPU performance(fix cross attn kvc)
profile
pipeline_forwardtotal CUDAaten::addmmtimeaten::addmmshare of total latencyaten::addmmcallsCommand Buffer FullCommand Buffer Fullshare of total latencyaten::mulaten::copy_aten::cataten::addnative_layer_normcudnn_convolutionaten::mmEssential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model. Please runmkdocs serveto sync the documentation editions to./docs.BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)