[diffusion] fix: keep Qwen-Image 2.1 prefix KV per layer under Cache-DiT (#40472)

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:

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"