[diffusion][cache-dit] add dual-transformer Cache-DiT adapter specs (#30150)

Co-authored-by: Yihao Wang <42559837+AgainstEntropy@users.noreply.github.com>

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

@@ -222,6 +222,43 @@ class CacheDitConfig:
     steps_computation_policy: str = "dynamic"
 
 
+@dataclass(frozen=True)
+class DualTransformerBlockAdapterSpec:
+    """BlockAdapter metadata for dual-transformer DiT pipelines.
+
+    This describes the cache-dit-facing structure of a pair of transformers.
+    The denoising loop semantics live in DenoisingStage; this spec only covers
+    how cache-dit should find blocks and interpret each block's forward.
+    """
+
+    blocks_attr: tuple[str, str]
+    blocks_name: Optional[List[str]]
+    forward_pattern: List[ForwardPattern]
+    check_forward_pattern: bool
+    check_num_outputs: bool
+    has_separate_cfg: bool
+
+
+DUAL_TRANSFORMER_BLOCK_ADAPTER_SPECS: dict[str, DualTransformerBlockAdapterSpec] = {
+    "wan2.2": DualTransformerBlockAdapterSpec(
+        blocks_attr=("blocks", "blocks"),
+        blocks_name=None,
+        forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2],
+        check_forward_pattern=True,
+        check_num_outputs=False,
+        has_separate_cfg=True,
+    ),
+    "ideogram4": DualTransformerBlockAdapterSpec(
+        blocks_attr=("layers", "layers"),
+        blocks_name=["layers", "layers"],
+        forward_pattern=[ForwardPattern.Pattern_3, ForwardPattern.Pattern_3],
+        check_forward_pattern=False,
+        check_num_outputs=False,
+        has_separate_cfg=False,
+    ),
+}
+
+
 # Custom BlockAdapter for DiT models absent from cache-dit's BlockAdapterRegister.
 # Value: (blocks attr, forward_pattern). forward_pattern must
 # match the block's forward signature (see cache_dit.ForwardPattern; e.g., ERNIE
@@ -404,9 +441,10 @@ def enable_cache_on_dual_transformer(
 ) -> tuple[torch.nn.Module, torch.nn.Module]:
     """Enable cache-dit on dual transformers using BlockAdapter.
 
-    For models with two transformers (high-noise expert and low-noise expert),
+    For models with two transformers, cache-dit requires enabling cache on both
-    cache-dit requires enabling cache on both simultaneously via BlockAdapter.
+    simultaneously via BlockAdapter. The two transformers may be split by denoising
-    This cannot be done by calling enable_cache separately on each transformer.
+    range, or run as paired conditional/unconditional branches. This cannot be done
+    by calling enable_cache separately on each transformer.
 
     Args:
         primary_config: CacheDitConfig for primary transformer.
@@ -414,14 +452,11 @@ def enable_cache_on_dual_transformer(
         sp_group: Sequence parallel process group (for Ulysses/Ring).
         tp_group: Tensor parallel process group.
     """
-    _supported_dual_transformer_models = [
+    adapter_spec = DUAL_TRANSFORMER_BLOCK_ADAPTER_SPECS.get(model_name)
-        "wan2.2",
+    if adapter_spec is None:
-        "ideogram4",
-    ]
-    if model_name not in _supported_dual_transformer_models:
         raise ValueError(
             f"Dual-transformer cache-dit is only supported for "
-            f"{_supported_dual_transformer_models}, got {model_name}."
+            f"{sorted(DUAL_TRANSFORMER_BLOCK_ADAPTER_SPECS)}, got {model_name}."
         )
 
     if not primary_config.enabled:
@@ -533,46 +568,28 @@ def enable_cache_on_dual_transformer(
         transformer_2, parallelism_config, sp_group, tp_group
     )
 
-    if model_name == "wan2.2":
+    transformer_blocks_attr, transformer_2_blocks_attr = adapter_spec.blocks_attr
-        transformer_blocks = getattr(transformer, "blocks", None)
+    transformer_blocks = getattr(transformer, transformer_blocks_attr, None)
-        transformer_2_blocks = getattr(transformer_2, "blocks", None)
+    transformer_2_blocks = getattr(transformer_2, transformer_2_blocks_attr, None)
-        blocks_name = None
-        forward_pattern = [ForwardPattern.Pattern_2, ForwardPattern.Pattern_2]
-        check_forward_pattern = True
-        check_num_outputs = False
-        has_separate_cfg = True
-    elif model_name == "ideogram4":
-        transformer_blocks = getattr(transformer, "layers", None)
-        transformer_2_blocks = getattr(transformer_2, "layers", None)
-        blocks_name = ["layers", "layers"]
-        forward_pattern = [ForwardPattern.Pattern_3, ForwardPattern.Pattern_3]
-        check_forward_pattern = False
-        check_num_outputs = False
-        has_separate_cfg = False
-    else:
-        raise ValueError(
-            f"Dual-transformer is not implemented for model {model_name} yet."
-        )
-
     if transformer_blocks is None or transformer_2_blocks is None:
-        expected_attr = "layers" if model_name == "ideogram4" else "blocks"
         raise ValueError(
-            f"Dual transformers for {model_name} must have '{expected_attr}' "
+            f"Dual transformers for {model_name} must expose cache-dit block "
-            "attribute for cache-dit. "
+            f"attributes {adapter_spec.blocks_attr}. "
-            f"transformer has {expected_attr}: {transformer_blocks is not None}, "
+            f"transformer has {transformer_blocks_attr}: "
-            f"secondary transformer has {expected_attr}: {transformer_2_blocks is not None}"
+            f"{transformer_blocks is not None}, secondary transformer has "
+            f"{transformer_2_blocks_attr}: {transformer_2_blocks is not None}"
         )
 
     cache_dit.enable_cache(
         BlockAdapter(
             transformer=[transformer, transformer_2],
             blocks=[transformer_blocks, transformer_2_blocks],
-            blocks_name=blocks_name,
+            blocks_name=adapter_spec.blocks_name,
-            forward_pattern=forward_pattern,
+            forward_pattern=adapter_spec.forward_pattern,
             params_modifiers=[primary_modifier, secondary_modifier],
-            check_forward_pattern=check_forward_pattern,
+            check_forward_pattern=adapter_spec.check_forward_pattern,
-            check_num_outputs=check_num_outputs,
+            check_num_outputs=adapter_spec.check_num_outputs,
-            has_separate_cfg=has_separate_cfg,
+            has_separate_cfg=adapter_spec.has_separate_cfg,
         ),
         parallelism_config=None,
     )

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

@@ -12,6 +12,13 @@ class _FakeDBCacheConfig:
         return kwargs
 
 
+class _FakeForwardPattern:
+    # A class (not a SimpleNamespace instance) so it is a valid type in
+    # annotations like List[ForwardPattern], matching the real Enum.
+    Pattern_2 = "Pattern_2"
+    Pattern_3 = "Pattern_3"
+
+
 def _install_cache_dit_stub():
     cache_dit = types.ModuleType("cache_dit")
     cache_dit.refresh_calls = []
@@ -36,7 +43,7 @@ def _install_cache_dit_stub():
     cache_dit.steps_mask = steps_mask
     cache_dit.BlockAdapter = types.SimpleNamespace
     cache_dit.DBCacheConfig = _FakeDBCacheConfig
-    cache_dit.ForwardPattern = types.SimpleNamespace(Pattern_3="Pattern_3")
+    cache_dit.ForwardPattern = _FakeForwardPattern
     cache_dit.ParamsModifier = object
     cache_dit.TaylorSeerCalibratorConfig = object