[mickqian]

Stacked on #46245 (FFN residual chain), #46244 and #46228; review the commits after #46245.

Motivation

Each Wan block projects the 512 text tokens to the cross-attention K and V (and RMS-normalizes K) in every denoising step: on Wan2.2 A14B 720p (4xB300) two small GEMMs and a norm, ~0.07 ms of every block, 40 blocks x 40 steps. The text and the DiT weights do not change inside a denoising loop, and the denoising stage hands the transformer the same prompt-embedding tensor at every step of the loop (the CFG branch kwargs are built once per loop). LTX-2 does the same per loop in #46105 through its own denoising stage. Wan runs the shared stage, so the transformer keys the projection by the prompt tensor, and the shared stage frees it when the loop ends.

Modifications

managers/forward_context.py: DENOISING_LOOP_STATE, a Req.extra key for what models keep for one denoising loop. Its leading _ keeps it out of disaggregated request transfers, like _skip_softmax_params.

pipelines_core/stages/denoising.py: _post_denoising_loop drops it. The standard, progressive, DMD and LTX-2 loops all end there, before decode. The stage knows nothing about what is inside, so there is no per-model hook and no getattr/hasattr fallback.

models/dits/wanvideo.py:

• WanT2VCrossAttention.project_context_kv(context) returns the normed K and the V that forward projects; forward(..., context_kv=None) uses them when given.
• WanTransformerBlock.forward(..., text_kv=None) passes them to attn2.
• WanTransformer3DModel._text_kvs(source, context): on the first step with a prompt tensor source, it projects every block's K/V from the embedded context. They are stored in the forward context's request, under req.extra[DENOISING_LOOP_STATE], keyed by (id(model), id(source)). The entry holds source so its id cannot be reused while cached. The module holds none of them: they go when the loop ends and never outlive the request, through VAE decode, across requests, or while component residency has the DiT offloaded. Later steps of the loop read them. A new request (whose LoRA or weights may differ) projects again; inference tensors carry no version counter, so the request is the validity scope. An expert's first step drops the other expert's entries, and at most one other CFG branch is kept, so with serial CFG at most two entries (~840 MB for A14B) live during the loop. K/V cached on a device other than the activation's are projected again. It stays off, and every step projects as before, when a cross-attention projection is LoRA-wrapped, weights are layerwise-offloaded, under compile or CUDA-graph capture, for I2V cross-attention, or when a cache adapter or hook replaced the blocks (_plain_blocks).

The same GEMMs and norm run on the same inputs, so outputs are bit-identical.

Accuracy Tests

• test_model_fast_paths.py::test_wan_projected_text_kv_keeps_bits (Wan section): a Wan block given project_context_kv(context) gives the same bits as the block projecting it itself.
• test_model_fast_paths.py::test_wan_text_kv_lives_for_one_loop_of_one_request, run with two models standing in for the two experts:
  • Every step of a loop reads the K/V projected on its first step, for both CFG branches, with the bits of a direct projection.
  • Another request projects its own, and K/V left on another device are projected again.
  • The second expert's first step frees the first expert's K/V.
  • The real DenoisingStage._post_denoising_loop frees the rest.
  • A request dropped mid-loop takes its K/V with it.
  • Negative checks: a model that keeps a reference fails the expert-switch check, and a stage that keeps the state fails the loop-end check.
• Wan2.2 T2V A14B 720p, 4xB300, 3-step lossless (crf 0) mp4: md5 identical to the build without the Wan stack for 3 prompts, before and after the fix.

Speed Tests and Profiling

Wan2.2 T2V A14B 720p, 4xB300, harness single_e2e, fresh server per run, interleaved: base (#46245 plus the spatial-parallel VAE changes #46171 and #46174) 89.800 / 89.791 / 89.691 / 89.771 s (mean 89.763) -> 89.604 / 89.652 / 89.665 / 89.616 / 89.691 s (mean 89.646), -0.13%. After the memory-lifetime fix, the full stack measured 89.517 / 89.540 s against 89.556 s for the pre-fix stack in the same session: no regression.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ✅ Run #40439978673
Latest PR Test (Extra): ❌ Run #40439978258
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.

[mickqian]

/tag-and-rerun-ci