[diffusion] keep Qwen-Image 2.1 prefix KV per layer under Cache-DiT

[niehen6174]

Motivation

#39983 registered Cache-DiT on 2.1 (Pattern_3 / transformer_blocks). The B200 check ran, but UnifiedBlocks still forwards the same extras to every layer.

2.1 prefix KV is per layer. The native loop sliced at the call site ([cache[i] for cache in prefix_caches]). After the wrap that loop runs once with i == 0. Layer 0 fills the empty caches[0] dict; layers 1–31 treat it as a hit and reuse layer 0's prefix. Same knobs, same denoise wall-clock: teapot at ~28–30 dB vs color noise at 8.02 dB. Wiring bug, not an aggressive residual.

How to fix

Each block keeps a stable _layer_id and slices prefix_caches itself. The outer loop always passes the full list, so Cache-DiT wrap and the native path share one contract (same idea as Hunyuan's _layer_id). Allocate fallback caches from self.num_layers, not len(transformer_blocks) after the wrap.

Do not land the earlier hasattr / per-call _prefix_caches bind: it fixed the noise but is a side-channel, and a missed bind silently shares layer 0 again. Sage / kitchen stay flag-only.

Lossy stack on 1× RTX 4090 24GB (Ada)

T2I 1024×1024, 40-step, CFG=1, seed 42, A red ceramic teapot on a wooden table. DiT layerwise (prefetch=1). Base = BF16 + SDPA, no Cache / Sage / quant. PSNR is vs that same-regime base PNG. Numbers are denoise only — Cache-DiT does not change encode.

stack	Denoise	per-step	vs base	PSNR	needs code?
base BF16+SDPA	26.69 s	651 ms	1.00×	—	—
Sage 2.2	26.89 s	662 ms	0.99×	29.00 dB	no (--attention-backend sage_attn)
kitchen INT8	14.20 s	338 ms	1.88×	31.15 dB	no (--quantization kitchen_int8)
Sage + kitchen	13.71 s	334 ms	1.95×	29.27 dB	no
Cache-DiT Fn=1,Bn=0,W=4,R=0.24,MC=3	10.31 s	195 ms	2.59×	29.57 dB	this PR (--enable-cache-dit)
Cache + kitchen	5.59 s	99 ms	4.77×	28.87 dB	this PR + flags
Cache + kitchen + Sage	4.74 s	89 ms	5.63×	27.81 dB	this PR + flags
Cache-DiT without this PR	10.29 s	195 ms	2.59×	8.02 dB	broken; do not cite

All successful rows are a clean red teapot. Hit count stays 13 full / 27 cached. Kitchen halves DiT H2D on full steps (~670 → ~343 ms). Sage is invisible alone (FA hidden behind copy); on Cache+kitchen it shaves another 0.85 s.

Layerwise 24GB Ada only. Do not cross-compare with resident H200/B200. Sage3 / SVDQ / NVFP4 / Spectrum / TeaCache / Distill stay paused here.

Recommend (this card): quality → base, or kitchen INT8 if lossy is OK (31.15 dB). Denoise speed → Cache + kitchen + Sage. Do not enable Sage by itself.

Cookbook

docs/cookbook/diffusion/Qwen-Image/Qwen-Image-2.1.mdx documents --enable-cache-dit and that each block slices by _layer_id (without that, color noise). It does not publish the 4090 speedup table — that is one-card layerwise evidence, not a picker recipe. cache_dit.mdx lists Qwen-Image 2.1.

Test plan

• python -m pytest sglang/multimodal_gen/test/unit/test_qwen_image21.py -k "prefix_kv or first_slot"
• --enable-cache-dit true 40-step T2I on 4090: teapot, not color noise (28.24 dB vs same-run base)
• no Cache-DiT: prefix KV path unchanged (same-run base teapot)
• Sage / kitchen still work as flags on top of this patch

---

CI States

Latest PR Test (Base): ✅ Run #35520933028
Latest PR Test (Extra): ❌ Run #35520932843
Latest PR Test (AMD ROCm 10): ❌ Run #35520933042

[niehen6174]

/tag-and-rerun-ci

[niehen6174]

/tag-and-rerun-ci