Repository navigation
[Perf] diffusion: Write Wan's normed q/k straight into the Ulysses exchange - #46247
Merged
Merged
Conversation
Wan normalizes q and k across all heads (rms_norm_across_heads) with FlashInfer's RMSNorm, then rotates them with FlashInfer's interleaved cos/sin-cache RoPE: three full passes over q and k per block (at 720p on 4xB300 37800 x 5120 tokens per rank, ~0.54 ms of every block). fused_rmsnorm_interleaved_rope_ is FlashInfer's CuTe-DSL RMSNormKernel (same tiling and reduction) whose store rotates each (even, odd) pair of the rounded bf16 output in fp32 with the same expression as the RoPE kernel, in place. Its output is bit-identical to the two kernels it replaces. The Wan block uses it when the q/k norm runs FlashInfer's kernel (CUDA forward, no TP across heads) and the transformer supplies the cos/sin cache. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Each Wan block projects the same normalized hidden states three times. After loading, to_q / to_k / to_v now keep views of one [3 * dim, dim] weight and [3 * dim] bias, and the block runs a single GEMM over it while every layer still reads its view; a LoRA wrapper, offload or a weight swap that replaces a tensor sends the block back to the three projections. cuBLAS picks the same kernel for the fused shape and the output is bit-identical (B300, 37800 x 5120 x 15360: 3.976 ms against 3 x 1.354 ms at the power cap). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…'s norm1 A Wan block ends with MulAdd (hidden + ff_output * gate) and the next block starts with norm1 over that sum: two passes over the residual stream. When every block is a plain WanTransformerBlock in a plain module list and both ops run their CUDA kernels, a block now hands (ff_output, gate) to the next one, whose norm1 runs the CuTe residual norm (residual + gate * x, then LayerNorm and scale/shift) in one kernel. The residual it writes equals MulAdd's output bit for bit (both fuse the multiply-add), and the norm reads that same bf16 sum, so the blocks' outputs keep their bits; the last block still adds its own residual. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Every Wan block projected the text to the cross-attention K and V (and normed K) in every denoising step, although the prompt embeddings the stage passes are the same tensor for the whole loop and the weights do not change inside it. The transformer now projects all blocks' K/V on the first step it sees a prompt tensor, keeps them keyed by that tensor (held, so its id is not reused) and its version, and hands each block its pair; the same GEMMs and norm run on the same inputs, so outputs are bit-identical. One entry per CFG branch is kept. LoRA-wrapped projections, layerwise offload, compile, graph capture and block hooks (cache-dit) keep the per-step projection. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…s exchange A same-device copy (cuMemcpy2DAsync D2D or copy_) runs on the SMs, so it waits out the attention kernel that is running. The head-group pipeline copied this rank's own rows into its receive slots on the copy stream, between the peer sends: each group's own copy waited for the previous group's attention and held the peer sends queued behind it. The output merges (own rows and the peers' blocks into the merged output) were same-device copies too, and all ran after the last group. - This rank's rows of every group go into its slots in one Triton pass (copy_head_groups) on the main stream before the first group attends; the copy stream carries only peer sends. - When the backend writes through a strided output (forward_into; FA4), each group attends straight into a token-major buffer of every row and head. This rank's rows are the result, the peers' rows are sent from there, and the peers' blocks come back in one pass per peer. Outputs are bit-identical: the same kernels see the same values. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…change Under pipelined Ulysses, each Wan block normed and rotated q and k in place and the exchange then copied this rank's own rows of every head group into its receive slots (one same-device copy per tensor, before the first group attends). The pipelined exchange now takes ``fill_all``: one call that gets every destination block at once (this rank's own slot rows and each peer's send block), which an across-heads norm needs since it reads whole rows. Wan's fused RMSNorm + RoPE gains an "into blocks" mode that writes each 512-column run (one head group) to its block from a per-run address table, so q and k are read once and written once, straight where attention and the copy engine read them. USPAttention.forward_qk_filled runs that path and returns None, with q/k untouched, whenever the call would not pipeline; the block then takes its old path. Same kernels on the same values, so outputs are bit-identical. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The block kept the fused [3 * dim, dim] weight and bias as plain attributes. Moving the DiT off the GPU (component residency swapping an A14B expert, --dit-cpu-offload) or layerwise offload gives to_q / to_k / to_v new data, but those attributes kept the fused buffers on the GPU: ~157 MB per block, ~6.3 GB per A14B expert, which defeated the offload. The forward already fell back to three GEMMs, but the memory stayed pinned. The block now holds no reference to them. _fused_qkv reads the fused weight and bias back from the three layers as views, while the layers are still consecutive rows of one buffer on the input's device, so the layers alone keep the buffers alive. The check costs ~10 us per block call on the host (4.4 us before) and is skipped under compile. The test checks that moving the block frees every byte it held on the GPU, and that re-fusing on CPU and GPU gives the same weights and the same output bits. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…-per-prompt Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…quest The model kept up to two entries of every block's cross-attention K/V in a module dict (~420 MB per entry for an A14B expert; ~1.7 GB with serial CFG over both experts) until the next request evicted them: through VAE decode, across the expert switch, and while component residency had the DiT offloaded, since a plain dict does not move with the module. They now live in the request, under the new Req.extra key DENOISING_LOOP_STATE, which the shared DenoisingStage drops in _post_denoising_loop (standard, progressive, DMD and LTX-2 loops all end there), before decode. The module holds none of them, so a request that ends early takes its K/V with it. An expert's first step drops the other expert's entries (at most one other CFG branch is kept), and K/V cached on a device other than the current activation's are projected again. The key starts with "_" so disaggregated transfers skip it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…-rope-into-exchange Conflict in wanvideo.py: _project_qkv (split out here) now reads the fused q/k/v weight through _fused_qkv, and the block keeps _qk_tables next to qkv_fusable instead of the removed qkv_weight / qkv_bias. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…paths.py The fused q/k/v test moves from its own file into a new "Wan" section of test_model_fast_paths.py, the per-model fast-path file. The section's _wan_block_env / _wan_block / _wan_block_inputs helpers hold the single-rank model-parallel setup, the attention patches and the block inputs, so the next Wan tests reuse them instead of repeating them. Like the file's other bit-exact fusion tests it is CUDA-only (requires_inline_ptx), which keeps it out of the AMD nightly run. est_time 110 -> 125. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…l_fast_paths.py The chained-blocks test moves from its own file into the Wan section of test_model_fast_paths.py, on the section's shared block setup. est_time 125 -> 140. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…-per-prompt Conflict (modify/delete) in test_wan_ffn_residual_chain.py: the chain branch moved its test into test_model_fast_paths.py and deleted the file, which this branch extended with the text K/V tests. The merge keeps this branch's file; the next commit moves the text K/V tests to the Wan section and deletes it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ths.py The text K/V tests move into the Wan section of test_model_fast_paths.py, on its shared block setup, and test_wan_ffn_residual_chain.py goes away. The check that a block given projected K/V keeps its bits was the tail of the chain test and is now its own test (test_wan_projected_text_kv_keeps_bits), next to the lifetime test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…-rope-into-exchange Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…n_backend.py The forward_into tests (FA4 writing through a strided out, and the default declining) move from their own file into the existing CUDA attention backend tests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
mickqian
requested review from
AgainstEntropy,
HaiShaw,
ping1jing2 and
yichiche
as code owners
October 24, 2026 04:38
mickqian
requested review from
BBuf,
kevin-mii,
niehen6174 and
yingluosanqian
as code owners
October 24, 2026 04:38
Collaborator
Author
|
/tag-and-rerun-ci |
github-actions
Bot
added
jit-kernel
run-ci
This was referenced Oct 24, 2026
arathi-hlab
added a commit
to arathi-hlab/sglang
that referenced
this pull request
Oct 26, 2026
… in _text_kvs sgl-project#46247 added torch.cuda.is_current_stream_capturing() to the text K/V cache guard; on XPU torch it raises "Tried to instantiate dummy base class _cuda_isCurrentStreamCapturing", failing every Wan denoising step. Use is_device_stream_capturing(context.device) instead.
1 of 2 tasks
1 of 2 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #46246 (text K/V once per loop), which sits on #46245, #46244 and #46228 (the fused q/k norm-RoPE this PR extends). Also depends on #46145 (pipelined Ulysses same-device copies), carried here as cherry-picks of its two commits; review the commits after #46246 other than those.
Motivation
Under pipelined Ulysses (Wan2.2 A14B 720p, 4xB300, CFG 2 x Ulysses 2) every block normed and rotated q and k in place, and the exchange then copied this rank's own rows of every head group into its receive slots: one same-device copy per tensor (
_copy_head_groups_kernel, ~0.073 ms each, 0.27 ms per block for q, k, v and the output gather), on the critical path before the first head group attends. The existingfillhook lets a per-head QK-norm write each block itself, but Wan normalizes across all heads, so a block-at-a-time fill would reread whole rows once per block.Modifications
distributed/device_communicators/ipc_a2a_multi.py:pipelined_attention(..., fill_all=None):fillgiven every destination block at once,[(head_start, head_count, q_dst, k_dst)]over this rank's own slot rows and each peer's send block; the rest of the fill path (v's own rows copied, the peers' q/k blocks sent from the send buffers) is shared withfill.layers/attention/layer.py:USPAttention.forward_qk_filled(q, k, v, fill_all)runs the dense pipelined path with that hook (no first-sight reference: the pipelined backends attend a head group exactly as inside the full call) and returns None, having written nothing, whenever the call would not pipeline.kernels/ops/diffusion/rope/rmsnorm_interleaved_rope_cutedsl.py:fused_rmsnorm_interleaved_rope_into_blocks(x, ..., dst_table): the same kernel, but each thread's 512-column run goes to row m of the contiguous block atdst_table[run]instead of back intox.models/dits/wanvideo.py: the block projects q/k/v, then_attend_qk_filledwrites q and k into the exchange blocks through the address table (built once per set of exchange buffers) and attends; when that returns None (no Ulysses pipeline, capture, compile, offline q/k rotation, another attention op) the block runs its previous path on the untouched projections.Same kernels on the same values, so outputs are bit-identical.
Accuracy Tests
test_rmsnorm_interleaved_rope.py::test_into_blocks_matches_in_place: every run lands in its block with the in-place result's bits, from a strided (fused-qkv) input, which stays untouched.test_usp_pipelined_attention_multi_gpu.py(2 and 4 ranks) passes on the branch.Speed Tests and Profiling
Wan2.2 T2V A14B 720p, 4xB300, harness single_e2e, fresh server per run, interleaved in one round with TRT-LLM: base (#46246) 89.616 / 89.691 s -> 89.454 / 89.482 / 89.488 s (mean 89.475), -0.19%. After the memory-lifetime fixes in #46244 and #46246, the full stack measured 89.517 / 89.540 s against 89.556 s for this commit's stack in the same session: no regression.
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #40220916482
Latest PR Test (Extra): ❌ Run #40220916459
Latest PR Test (AMD ROCm 10): ❌ Run #40220916429