[diffusion][cache-dit] support Krea-2 + run-driven has_separate_cfg (#29688)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py

@@ -223,25 +223,27 @@ class CacheDitConfig:
 
 
 # Custom BlockAdapter for DiT models absent from cache-dit's BlockAdapterRegister.
-# Value: (blocks attr, forward_pattern, has_separate_cfg). forward_pattern must
+# Value: (blocks attr, forward_pattern). forward_pattern must
 # match the block's forward signature (see cache_dit.ForwardPattern; e.g., ERNIE
-# uses Pattern_3). has_separate_cfg=True aligns cache-dit's step counter for
-# sequential CFG (two forwards per step); cache-dit auto-resolves the remaining
+# uses Pattern_3). has_separate_cfg follows the run (passed by
+# enable_cache_on_transformer); cache-dit auto-resolves the remaining
 # fields.
-_CUSTOM_BLOCK_ADAPTER_SPECS: dict[str, tuple[str, ForwardPattern, bool]] = {
-    "ErnieImageTransformer2DModel": ("layers", ForwardPattern.Pattern_3, True),
+_CUSTOM_BLOCK_ADAPTER_SPECS: dict[str, tuple[str, ForwardPattern]] = {
+    "ErnieImageTransformer2DModel": ("layers", ForwardPattern.Pattern_3),
+    "Krea2Transformer2DModel": ("transformer_blocks", ForwardPattern.Pattern_3),
 }
 
 
 def _build_custom_block_adapter(
     transformer: torch.nn.Module,
+    has_separate_cfg: bool = False,
 ) -> Optional[BlockAdapter]:
     """Build a manual BlockAdapter for a model absent from cache-dit's registry,
     or None if the class is unknown."""
     spec = _CUSTOM_BLOCK_ADAPTER_SPECS.get(transformer.__class__.__name__)
     if spec is None:
         return None
-    blocks_attr, forward_pattern, has_separate_cfg = spec
+    blocks_attr, forward_pattern = spec
     blocks = getattr(transformer, blocks_attr, None)
     if blocks is None:
         raise ValueError(
@@ -262,6 +264,7 @@ def enable_cache_on_transformer(
     model_name: str = "transformer",
     sp_group: Optional[torch.distributed.ProcessGroup] = None,
     tp_group: Optional[torch.distributed.ProcessGroup] = None,
+    has_separate_cfg: bool = False,
 ) -> torch.nn.Module:
     """Enable cache-dit on a transformer module, by wrapping the module with cache-dit
 
@@ -272,6 +275,9 @@ def enable_cache_on_transformer(
         model_name: Name of the model for logging purposes.
         sp_group: Sequence parallel process group (for Ulysses/Ring).
         tp_group: Tensor parallel process group.
+        has_separate_cfg: Whether the run issues separate conditional/unconditional
+            passes per step (CFG). Used by custom adapters (ERNIE, Krea-2); a
+            mismatch only disables caching, never corrupts output.
 
     """
     if not config.enabled:
@@ -288,7 +294,9 @@ def enable_cache_on_transformer(
     # _build_custom_block_adapter).
     custom_adapter = None
     if not BlockAdapterRegister.is_supported(transformer):
-        custom_adapter = _build_custom_block_adapter(transformer)
+        custom_adapter = _build_custom_block_adapter(
+            transformer, has_separate_cfg=has_separate_cfg
+        )
         if custom_adapter is None:
             transformer_cls_name = transformer.__class__.__name__
             raise ValueError(

python/sglang/multimodal_gen/runtime/models/dits/krea2.py

@@ -445,7 +445,7 @@ class SingleStreamBlock(nn.Module):
 
     def forward(
         self,
-        x: Tensor,
+        hidden_states: Tensor,
         vec: Tensor,
         freqs: Tensor,
         key_mask: Tensor | None = None,
@@ -455,20 +455,28 @@ class SingleStreamBlock(nn.Module):
         prescale, preshift, pregate, postscale, postshift, postgate = mod.chunk(
             6, dim=-1
         )
-        x = x + pregate * self.attn(
+        hidden_states = hidden_states + pregate * self.attn(
             norm_scale_shift(
-                x, self.norm1.weight + 1, prescale, preshift, self.norm1.eps
+                hidden_states,
+                self.norm1.weight + 1,
+                prescale,
+                preshift,
+                self.norm1.eps,
             ),
             freqs,
             key_mask,
             mask_meta,
         )
-        x = x + postgate * self.ff(
+        hidden_states = hidden_states + postgate * self.ff(
             norm_scale_shift(
-                x, self.norm2.weight + 1, postscale, postshift, self.norm2.eps
+                hidden_states,
+                self.norm2.weight + 1,
+                postscale,
+                postshift,
+                self.norm2.eps,
             )
         )
-        return x
+        return hidden_states
 
 
 # --------------------------------------------------------------------------- #

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -545,6 +545,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
                 model_name="transformer",
                 sp_group=sp_group,
                 tp_group=tp_group,
+                has_separate_cfg=batch.do_classifier_free_guidance,
             )
             logger.info(
                 "cache-dit enabled on transformer (steps=%d, Fn=%d, Bn=%d, rdt=%.3f)",

python/sglang/multimodal_gen/test/unit/test_cache_dit_integration.py

@@ -230,7 +230,7 @@ class TestBuildCustomBlockAdapter(unittest.TestCase):
         blocks = ["block_0", "block_1"]
         transformer = _make_transformer("ErnieImageTransformer2DModel", blocks)
 
-        adapter = module._build_custom_block_adapter(transformer)
+        adapter = module._build_custom_block_adapter(transformer, has_separate_cfg=True)
 
         self.assertIsNotNone(adapter)
         self.assertEqual(adapter.blocks, blocks)
@@ -250,6 +250,28 @@ class TestBuildCustomBlockAdapter(unittest.TestCase):
         with self.assertRaises(ValueError):
             module._build_custom_block_adapter(transformer)
 
+    def test_has_separate_cfg_follows_runtime(self):
+        # No model pins the mode; has_separate_cfg always follows the run's CFG mode
+        # (Krea-2 Raw -> True, Krea-2 Turbo -> False).
+        module = _import_module_with_stub()
+        blocks = ["block_0", "block_1"]
+
+        transformer_raw = _make_transformer("Krea2Transformer2DModel")
+        transformer_raw.transformer_blocks = blocks
+        adapter_raw = module._build_custom_block_adapter(
+            transformer_raw, has_separate_cfg=True
+        )
+        self.assertEqual(adapter_raw.blocks, blocks)
+        self.assertEqual(adapter_raw.forward_pattern, "Pattern_3")
+        self.assertTrue(adapter_raw.has_separate_cfg)
+
+        transformer_turbo = _make_transformer("Krea2Transformer2DModel")
+        transformer_turbo.transformer_blocks = blocks
+        adapter_turbo = module._build_custom_block_adapter(
+            transformer_turbo, has_separate_cfg=False
+        )
+        self.assertFalse(adapter_turbo.has_separate_cfg)
+
 
 if __name__ == "__main__":
     unittest.main()