From 6d6071aff27460484416294a703f6b7f2452e0bf Mon Sep 17 00:00:00 2001
From: niehen6174 <1639206518@qq.com>
Date: Sun, 20 Sep 2026 09:30:53 +0000
Subject: [PATCH 1/3] [diffusion] slice Qwen-Image 2.1 prefix KV by a stable
block _layer_id
UnifiedBlocks forwards the same extras to every layer, so call-site
[cache[i]] after wrap reused caches[0] and wrote color noise.
---
.../runtime/models/dits/qwen_image21.py | 21 +++++++++++++------
1 file changed, 15 insertions(+), 6 deletions(-)
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 297bcee9f42f..d83fa868a45f 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 @@ def forward(self, x, ropes, prefixes, layouts, caches):
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 @@ def forward(
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 @@ def __init__(self, config, hf_config, quant_config=None, **kwargs):
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 @@ def forward(
)
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 @@ def forward(
)
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 @@ def forward(
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:
From 5cb7737aafa64380ee41f115fc38a511684fcb5a Mon Sep 17 00:00:00 2001
From: niehen6174 <1639206518@qq.com>
Date: Sun, 20 Sep 2026 09:30:53 +0000
Subject: [PATCH 2/3] [diffusion] test Qwen-Image 2.1 prefix KV under a
Cache-DiT wrap
Cover the native loop, the UnifiedBlocks same-args path, and the
old first-slot broadcast that shared caches[0] across layers.
---
.../test/unit/test_qwen_image21.py | 69 +++++++++++++++++++
1 file changed, 69 insertions(+)
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 d387376123d2..155218e3aa26 100644
--- a/python/sglang/multimodal_gen/test/unit/test_qwen_image21.py
+++ b/python/sglang/multimodal_gen/test/unit/test_qwen_image21.py
@@ -155,6 +155,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"
From a8ed51032f2b0f138d7eecead58b0f702c999636 Mon Sep 17 00:00:00 2001
From: niehen6174 <1639206518@qq.com>
Date: Sun, 20 Sep 2026 09:30:53 +0000
Subject: [PATCH 3/3] [diffusion] document Qwen-Image 2.1 Cache-DiT prefix KV
Cookbook the per-layer KV wrap and list 2.1 on the Cache-DiT
supported-model table.
---
docs/cookbook/diffusion/Qwen-Image/Qwen-Image-2.1.mdx | 8 ++++++++
docs/docs/sglang-diffusion/cache_dit.mdx | 2 +-
2 files changed, 9 insertions(+), 1 deletion(-)
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 33d1be64587c..8fa584c205b4 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 d26ac88b9ffa..9ca452450db1 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 |