Skip to content

[Feat] DreamZero fix tp & cross attn cache - #4154

Merged
Gaohan123 merged 16 commits into
vllm-project:mainfrom
AbelSara:zhh/tp-fix-phase1
Jun 18, 2026
Merged

Gaohan123 merged 16 commits into
vllm-project:mainfrom
AbelSara:zhh/tp-fix-phase1

Conversation

@AbelSara

@AbelSara AbelSara commented Jun 4, 2026 •

Copy link
Copy Markdown
Contributor

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

  1. Custom all-reduce in DistributedRMSNorm — replace raw torch.distributed.all_reduce on the TP group with vLLM's tensor_model_parallel_all_reduce (custom-AR / symmetric-memory fast path for tiny tensors).
  2. Fused q/k qk-norm all-reduce — self-attn norms q and k every step, each issuing its own tiny all-reduce of the per-token sum-of-squares. Pack both into one tensor and reduce once (2 collectives → 1). Cross-attn is left as-is (its k is cached, so steady state already issues a single all-reduce).
  3. Fused QKV projection — replace the three ColumnParallelLinear (q/k/v) i CausalWanSelfAttention with one QKVParallelLinear + split (3 GEMMs/launches → 1).
  4. k_img/v_img caching — : i2v cross-attn recomputes k_img/v_img every step although context_img is 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

Run / Folder Interconnect first (s) steady (s) gen_total (s) video_fps action_hz
base_tp1_cfg1 (baseline) Single 12.681 12.970 25.651 0.663 1.871
pcie_base_tp1_cfg2 PCIe 7.056 6.607 13.664 1.244 3.513
pcie_base_tp2_cfg1 PCIe 14.439 14.252 28.691 0.593 1.673
pcie_pr_tp2_cfg1 (current) PCIe 14.329 14.008 28.338 0.600 1.694
nvlink_base_tp2_cfg1 NVLink 8.826 8.398 17.224 0.987 2.787
nvlink_pr_tp2_cfg1 (current) NVLink 8.600 8.071 16.671 1.020 2.879

Single GPU performance(fix cross attn kvc)

Metric base_single qkv_single kvc_single
Mean latency (s) 12.825 12.159 11.412
First request (s) 12.681 12.040 11.340
Steady-state request (s) 12.970 12.277 11.483
model_generate_total (s) 25.651 24.317 22.823
action_hz 1.871 1.974 2.103
model_video_fps 0.663 0.699 0.745
request_hz 0.0780 0.0822 0.0876

profile

Metric TP1 CFG1 TP2 CFG1 TP1 CFG2
pipeline_forward total CUDA 11.519s 7.639s 5.940s
Communication time 0s 0.785s all-reduce 0.005s all-gather
Non-communication / compute time 11.519s 6.854s 5.935s
Total speedup vs TP1 CFG1 1.00x 1.51x 1.94x
Compute speedup vs TP1 CFG1 1.00x 1.68x 1.94x
aten::addmm time 5.704s 3.089s 2.842s
aten::addmm share of total latency 49.5% 40.4% 47.8%
aten::addmm calls 8432 8432 4216
FlashAttention time 1.576s 0.828s 0.784s
FlashAttention share of total latency 13.7% 10.8% 13.2%
FlashAttention calls 4080 4080 2040
Command Buffer Full 1.658s 0.564s 0.494s
Command Buffer Full share of total latency 14.4% 7.4% 8.3%
aten::mul 0.847s 0.536s 0.431s
aten::copy_ 0.806s 0.423s 0.419s
aten::cat 0.791s 0.419s 0.398s
aten::add 0.783s 0.756s 0.397s
native_layer_norm 0.137s 0.135s 0.068s
cudnn_convolution 0.079s 0.079s 0.078s
aten::mm 0.046s 0.046s 0.046s
Accuracy error baseline max abs error: 0.050 0.0
Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
  • The test results. Please paste the results comparison before and after, or the e2e results.
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
  • (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

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)

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

@AbelSara AbelSara changed the title [WIP] tp fix phase1 DreamZero tp fix phase1 Jun 5, 2026
@AbelSara AbelSara changed the title DreamZero tp fix phase1 [Feat] DreamZero tp fix phase1 Jun 6, 2026
AbelSara added 5 commits June 6, 2026 20:12
…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>
@AbelSara
AbelSara force-pushed the zhh/tp-fix-phase1 branch from 9559fc3 to 396afbf Compare June 6, 2026 12:12
AbelSara and others added 2 commits June 6, 2026 23:17
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>
@AbelSara AbelSara changed the title [Feat] DreamZero tp fix phase1 [Feat] DreamZero fix tp & cross attn cache Jun 7, 2026
AbelSara and others added 2 commits June 7, 2026 22:49
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>
@mergify mergify Bot mentioned this pull request Jun 7, 2026
5 tasks
Comment thread vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
Comment thread vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
…-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>
@AbelSara
AbelSara requested a review from RuixiangMa June 9, 2026 14:00
@AbelSara

AbelSara commented Jun 9, 2026

Copy link
Copy Markdown
Contributor Author

@RuixiangMa Hi, could you please help review and merge the code?

@amy-why-3459

Copy link
Copy Markdown
Collaborator

LGTM

@hsliuustc0106 hsliuustc0106 added the ready label to trigger buildkite CI label Jun 12, 2026
Comment thread vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
Comment thread vllm_omni/diffusion/models/dreamzero/causal_wan_model.py Outdated
Comment thread tests/dreamzero/test_qkv_fusion.py
@AbelSara

Copy link
Copy Markdown
Contributor Author

based on these findings, what's the recommended deployment settings for users with specific hardwares, e.g., H20/5090

In short, TP + CFG parallelism is sufficient for both configurations; no additional parallelism settings are required.

@hsliuustc0106

Copy link
Copy Markdown
Collaborator

@TKONIY PTAL for a final check

@Gaohan123 Gaohan123 added this to the v0.24.0 milestone Jun 17, 2026
@TKONIY

TKONIY commented Jun 18, 2026 •

Copy link
Copy Markdown
Contributor

TP Performance Reproduction on PCIe + Pro6000

All durations are milliseconds. PR-body seconds are converted to ms.

PR row PR body gen_total (ms) Local status Local ref Local GPUs Local first (ms) Local steady (ms) Local gen_total (ms) Local video_fps Local action_hz Note
base_tp1_cfg1 25651 ok 7b83794 1 7458.0 7474.7 14932.7 1.138 3.214 Single GPU baseline
pcie_base_tp1_cfg2 13664 ok 7b83794 1,2 4451.1 3960.2 8411.3 2.021 5.707 CFG parallel baseline
pcie_base_tp2_cfg1 28691 ok 7b83794 1,2 7372.1 6767.7 14139.8 1.202 3.395 TP baseline
pcie_pr_tp2_cfg1 28338 ok f819f81 1,2 6797.9 6387.3 13185.2 1.289 3.640 PR TP path
nvlink_base_tp2_cfg1 17224 skipped - - - - - - - No NVLink topology on host
nvlink_pr_tp2_cfg1 16671 skipped - - - - - - - No NVLink topology on host

Local TP2 PCIe improvement from this PR:

  • gen_total: 14139.8 ms -> 13185.2 ms, 6.75% faster.
  • steady: 6767.7 ms -> 6387.3 ms, 5.62% faster.
  • action_hz: 3.395 -> 3.640.

The absolute latencies are much lower than the PR body on this host, so I would
not compare raw latency across machines. The local one-variable TP2 comparison
does show the PR path faster than the merge-base path.

@TKONIY

TKONIY commented Jun 18, 2026

Copy link
Copy Markdown
Contributor

LGTM. Implementation is clean. The TP performance on PCIe node performs much better then before.

@AbelSara

Copy link
Copy Markdown
Contributor Author

LGTM. Implementation is clean. The TP performance on PCIe node performs much better then before.

Thanks for helping review the code.

@Gaohan123 Gaohan123 removed the ready label to trigger buildkite CI label Jun 18, 2026

@Gaohan123 Gaohan123 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Thanks

@Gaohan123 Gaohan123 added the ready label to trigger buildkite CI label Jun 18, 2026
@Gaohan123
Gaohan123 enabled auto-merge (squash) June 18, 2026 06:34
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Signed-off-by: z00806815 <zhuhonghan@huawei.com>
@Gaohan123 Gaohan123 added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Jun 18, 2026
@hsliuustc0106
hsliuustc0106 disabled auto-merge June 18, 2026 09:07
@hsliuustc0106

Copy link
Copy Markdown
Collaborator

CI failure is not related

@Gaohan123 Gaohan123 added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Jun 18, 2026
@Gaohan123
Gaohan123 enabled auto-merge (squash) June 18, 2026 10:18
@Gaohan123
Gaohan123 merged commit 8f28b96 into vllm-project:main Jun 18, 2026
6 of 8 checks passed
tzhouam added a commit to tzhouam/vllm-omni that referenced this pull request Jun 20, 2026
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>
tzhouam added a commit to tzhouam/vllm-omni that referenced this pull request Jun 20, 2026
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>
alex-jw-brooks pushed a commit to alex-jw-brooks/vllm-omni that referenced this pull request Jun 23, 2026
@AbelSara
AbelSara deleted the zhh/tp-fix-phase1 branch June 28, 2026 04:53
tzhouam added a commit to tzhouam/vllm-omni that referenced this pull request Jul 4, 2026
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>
tzhouam added a commit to tzhouam/vllm-omni that referenced this pull request Jul 4, 2026
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>
khairulkabir1661 pushed a commit to khairulkabir1661/vllm-omni that referenced this pull request Sep 25, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready label to trigger buildkite CI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants