From b912db67ea51c19874476d88477fe35afa248265 Mon Sep 17 00:00:00 2001
From: WenhaoZhang <42087078+niehen6174@users.noreply.github.com>
Date: Mon, 21 Sep 2026 08:48:31 +0800
Subject: [PATCH] [diffusion] fix: keep Qwen-Image 2.1 prefix KV per layer
under Cache-DiT (#40472)
---
.../diffusion/Qwen-Image/Qwen-Image-2.1.mdx | 8 +++
docs/docs/sglang-diffusion/cache_dit.mdx | 2 +-
.../runtime/models/dits/qwen_image21.py | 21 ++++--
.../test/unit/test_qwen_image21.py | 69 +++++++++++++++++++
4 files changed, 93 insertions(+), 7 deletions(-)
diff --git a/docs/cookbook/diffusion/Qwen-Image/Qwen-Image-2.1.mdx b/docs/cookbook/diffusion/Qwen-Image/Qwen-Image-2.1.mdx
index 33d1be645..8fa584c20 100644
--- a/docs/cookbook/diffusion/Qwen-Image/Qwen-Image-2.1.mdx
+++ b/docs/cookbook/diffusion/Qwen-Image/Qwen-Image-2.1.mdx
@@ -216,3 +216,11 @@ Keep eager execution as the default. Breakable CUDA Graph replay requires
matching resolution and condition-prefix length; unseen shapes run eagerly.
Text buckets alone do not guarantee replay. SageAttention and Cache-DiT can
change numerical results and require quality checks for your workload.
+
+### Cache-DiT
+
+Enable `--enable-cache-dit true` or `SGLANG_CACHE_DIT_ENABLED=true`. 2.1 prefix
+KV is per layer: each block slices caches by `_layer_id`. Cache-DiT wraps
+`transformer_blocks` and forwards the same extras to every layer; without that
+slice, later layers reuse layer 0 and the image collapses to color noise.
+See the [Cache-DiT guide](/docs/sglang-diffusion/cache_dit).
diff --git a/docs/docs/sglang-diffusion/cache_dit.mdx b/docs/docs/sglang-diffusion/cache_dit.mdx
index d26ac88b9..9ca452450 100644
--- a/docs/docs/sglang-diffusion/cache_dit.mdx
+++ b/docs/docs/sglang-diffusion/cache_dit.mdx
@@ -783,7 +783,7 @@ SGLang Diffusion x Cache-DiT supports almost all models originally supported in
| Qwen |
- Qwen-Image, Qwen-Image-Edit |
+ Qwen-Image, Qwen-Image-Edit, Qwen-Image 2.1 |
| Hunyuan |
diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image21.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image21.py
index 297bcee9f..d83fa868a 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image21.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image21.py
@@ -381,8 +381,9 @@ class QwenImage21Attention(nn.Module):
class QwenImage21TransformerBlock(nn.Module):
- def __init__(self, ac, quant_config, prefix):
+ def __init__(self, ac, quant_config, prefix, layer_id):
super().__init__()
+ self._layer_id = layer_id
self.img_norm1 = nn.LayerNorm(
ac.hidden_size, eps=ac.eps, elementwise_affine=False
)
@@ -404,6 +405,9 @@ class QwenImage21TransformerBlock(nn.Module):
ropes,
caches,
):
+ # Cache-DiT's UnifiedBlocks forwards the same args to every layer.
+ # Slice here so prefix KV stays per-layer after that wrap.
+ caches = [cache[self._layer_id] for cache in caches]
scale1, gate1, scale2, gate2 = modulation
prefixes = [
apply_modulation(
@@ -482,9 +486,12 @@ class QwenImage21Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin
self.modulation = nn.Sequential(
nn.SiLU(), nn.Linear(ac.hidden_size, ac.hidden_size * 4, bias=False)
)
+ self.num_layers = ac.num_layers
self.transformer_blocks = nn.ModuleList(
[
- QwenImage21TransformerBlock(ac, quant_config, f"transformer_blocks.{i}")
+ QwenImage21TransformerBlock(
+ ac, quant_config, f"transformer_blocks.{i}", i
+ )
for i in range(ac.num_layers)
]
)
@@ -528,7 +535,7 @@ class QwenImage21Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin
)
prefix_modulation = self.prepare_modulation(zero_temb)
if prefix_caches is None:
- prefix_caches = [[None] * len(self.transformer_blocks) for _ in layouts]
+ prefix_caches = [[None] * self.num_layers for _ in layouts]
prefix_states, ropes = [], []
for sample, layout in enumerate(layouts):
prefix = None
@@ -544,8 +551,10 @@ class QwenImage21Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin
)
prefix_states.append({"hidden_states": prefix})
ropes.append(layout["target_rope"][start:end])
- # visit each block once so layerwise offload transfers weights once per batch
- for i, block in enumerate(self.transformer_blocks):
+ # Same extras for every block so Cache-DiT's UnifiedBlocks wrap is valid.
+ # Each block slices prefix_caches by _layer_id. Visit once per layer for
+ # layerwise offload.
+ for block in self.transformer_blocks:
images = block(
images,
modulation,
@@ -553,7 +562,7 @@ class QwenImage21Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin
prefix_modulation,
layouts,
ropes,
- [cache[i] for cache in prefix_caches],
+ prefix_caches,
)
output = self.proj_out(self.norm_out(images, temb))
if sp > 1:
diff --git a/python/sglang/multimodal_gen/test/unit/test_qwen_image21.py b/python/sglang/multimodal_gen/test/unit/test_qwen_image21.py
index 198a15c0f..5056f4944 100644
--- a/python/sglang/multimodal_gen/test/unit/test_qwen_image21.py
+++ b/python/sglang/multimodal_gen/test/unit/test_qwen_image21.py
@@ -178,6 +178,75 @@ def test_condition_slots_expand_to_actual_latent_grid():
torch.testing.assert_close(collapsed[slots][0], hidden[2])
+class _RecordingBlock(torch.nn.Module):
+ def __init__(self, layer_id):
+ super().__init__()
+ self._layer_id = layer_id
+ self.seen = None
+
+ def forward(self, hidden_states, *args):
+ caches = [cache[self._layer_id] for cache in args[-1]]
+ self.seen = caches[0]
+ return hidden_states
+
+
+class _UnifiedBlocks(torch.nn.Module):
+ def __init__(self, blocks):
+ super().__init__()
+ self.transformer_blocks = torch.nn.ModuleList(blocks)
+
+ def forward(self, hidden_states, *args):
+ x = hidden_states
+ for block in self.transformer_blocks:
+ x = block(x, *args)
+ return x
+
+
+def _run_blocks(blocks, prefix_caches):
+ x = torch.zeros(1)
+ for block in blocks:
+ x = block(x, prefix_caches)
+ return x
+
+
+def test_cache_dit_wrapper_keeps_per_layer_prefix_kv():
+ inner = [_RecordingBlock(0), _RecordingBlock(1), _RecordingBlock(2)]
+ wrapped = torch.nn.ModuleList([_UnifiedBlocks(inner)])
+ prefix_caches = [[{"layer": 0}, {"layer": 1}, {"layer": 2}]]
+
+ _run_blocks(wrapped, prefix_caches)
+ assert [block.seen for block in inner] == prefix_caches[0]
+
+
+def test_plain_blocks_still_get_per_layer_prefix_kv():
+ blocks = torch.nn.ModuleList([_RecordingBlock(0), _RecordingBlock(1)])
+ prefix_caches = [[{"layer": 0}, {"layer": 1}]]
+
+ _run_blocks(blocks, prefix_caches)
+ assert [block.seen for block in blocks] == prefix_caches[0]
+
+
+class _FirstSlotBlock(torch.nn.Module):
+ """Old loop body: take caches[0] and broadcast it to every layer."""
+
+ def __init__(self):
+ super().__init__()
+ self.seen = None
+
+ def forward(self, hidden_states, *args):
+ self.seen = args[-1][0]
+ return hidden_states
+
+
+def test_first_slot_only_caches_are_shared_across_layers():
+ inner = [_FirstSlotBlock(), _FirstSlotBlock()]
+ unified = _UnifiedBlocks(inner)
+ prefix_caches = [[{"layer": 0}, {"layer": 1}]]
+
+ unified(torch.zeros(1), [cache[0] for cache in prefix_caches])
+ assert [block.seen for block in inner] == [prefix_caches[0][0], prefix_caches[0][0]]
+
+
def test_adjacent_image_slots_stay_distinct():
layout = build_layout(
[False, True, True, False], [(1, 2, 2), (1, 4, 2), (1, 2, 2)], (4, 6, 6), "cpu"