[Diffusion] enable cache-dit for ERNIE-Image model (#28266)

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

@@ -222,6 +222,40 @@ class CacheDitConfig:
     steps_computation_policy: str = "dynamic"
 
 
+# Custom BlockAdapter for DiT models absent from cache-dit's BlockAdapterRegister.
+# Value: (blocks attr, forward_pattern, has_separate_cfg). 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
+# fields.
+_CUSTOM_BLOCK_ADAPTER_SPECS: dict[str, tuple[str, ForwardPattern, bool]] = {
+    "ErnieImageTransformer2DModel": ("layers", ForwardPattern.Pattern_3, True),
+}
+
+
+def _build_custom_block_adapter(
+    transformer: torch.nn.Module,
+) -> 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 = getattr(transformer, blocks_attr, None)
+    if blocks is None:
+        raise ValueError(
+            f"Transformer {transformer.__class__.__name__} has no attribute "
+            f"{blocks_attr!r} for cache-dit blocks."
+        )
+    return BlockAdapter(
+        transformer=transformer,
+        blocks=blocks,
+        forward_pattern=forward_pattern,
+        has_separate_cfg=has_separate_cfg,
+    )
+
+
 def enable_cache_on_transformer(
     transformer: torch.nn.Module,
     config: CacheDitConfig,
@@ -249,16 +283,21 @@ def enable_cache_on_transformer(
             "Please provide it in CacheDitConfig."
         )
 
-    # Check if the transformer is pre-registered in cache-dit
+    # Prefer the standard path (transformer pre-registered in cache-dit). For
+    # models absent from the registry, fall back to a manual BlockAdapter (see
+    # _build_custom_block_adapter).
+    custom_adapter = None
     if not BlockAdapterRegister.is_supported(transformer):
-        transformer_cls_name = transformer.__class__.__name__
-        raise ValueError(
-            f"{transformer_cls_name} is not officially supported by cache-dit. "
-            "Supported cache-dit DiT families include Flux, QwenImage, HunyuanDiT, "
-            "HunyuanVideo, Wan, CogVideoX, Mochi, and others. "
-            "Please ensure your transformer belongs to one of these families or "
-            "define a custom BlockAdapter."
-        )
+        custom_adapter = _build_custom_block_adapter(transformer)
+        if custom_adapter is None:
+            transformer_cls_name = transformer.__class__.__name__
+            raise ValueError(
+                f"{transformer_cls_name} is not officially supported by cache-dit. "
+                "Supported cache-dit DiT families include Flux, QwenImage, HunyuanDiT, "
+                "HunyuanVideo, Wan, CogVideoX, Mochi, and others. "
+                "Please ensure your transformer belongs to one of these families or "
+                "define a custom BlockAdapter."
+            )
 
     # Build cache config (including SCM fields if provided)
     cache_config = DBCacheConfig(
@@ -312,8 +351,18 @@ def enable_cache_on_transformer(
 
     _mark_transformer_parallelized(transformer, parallelism_config, sp_group, tp_group)
 
+    # Custom path: pass a pre-built BlockAdapter, bypassing the registry.
+    # Standard path: let enable_cache discover the registered adapter.
+    target = transformer
+    if custom_adapter is not None:
+        target = custom_adapter
+        logger.info(
+            "Enabling cache-dit on %s via custom BlockAdapter (%s).",
+            model_name,
+            custom_adapter.forward_pattern,
+        )
     cache_dit.enable_cache(
-        transformer,
+        target,
         cache_config=cache_config,
         calibrator_config=calibrator_config,
         parallelism_config=None,

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

@@ -34,9 +34,9 @@ def _install_cache_dit_stub():
 
     cache_dit.refresh_context = refresh_context
     cache_dit.steps_mask = steps_mask
-    cache_dit.BlockAdapter = object
+    cache_dit.BlockAdapter = types.SimpleNamespace
     cache_dit.DBCacheConfig = _FakeDBCacheConfig
-    cache_dit.ForwardPattern = object
+    cache_dit.ForwardPattern = types.SimpleNamespace(Pattern_3="Pattern_3")
     cache_dit.ParamsModifier = object
     cache_dit.TaylorSeerCalibratorConfig = object
 
@@ -125,28 +125,29 @@ def _install_torch_stub():
     }
 
 
-class TestCacheDitRefreshContext(unittest.TestCase):
-    def _import_module_with_stub(self):
-        stub_modules = _install_cache_dit_stub()
-        stub_modules.update(_install_sglang_dependency_stubs())
-        stub_modules.update(_install_torch_stub())
-        module_path = (
-            Path(__file__).resolve().parents[2]
-            / "runtime"
-            / "cache"
-            / "cache_dit_integration.py"
+def _import_module_with_stub():
+    stub_modules = _install_cache_dit_stub()
+    stub_modules.update(_install_sglang_dependency_stubs())
+    stub_modules.update(_install_torch_stub())
+    module_path = (
+        Path(__file__).resolve().parents[2]
+        / "runtime"
+        / "cache"
+        / "cache_dit_integration.py"
+    )
+    with patch.dict(sys.modules, stub_modules):
+        spec = importlib.util.spec_from_file_location(
+            "test_cache_dit_integration_target", module_path
         )
-        with patch.dict(sys.modules, stub_modules):
-            spec = importlib.util.spec_from_file_location(
-                "test_cache_dit_integration_target", module_path
-            )
-            module = importlib.util.module_from_spec(spec)
-            assert spec.loader is not None
-            spec.loader.exec_module(module)
-        return module
+        module = importlib.util.module_from_spec(spec)
+        assert spec.loader is not None
+        spec.loader.exec_module(module)
+    return module
 
+
+class TestCacheDitRefreshContext(unittest.TestCase):
     def test_refresh_context_without_scm_preset_skips_steps_mask(self):
-        module = self._import_module_with_stub()
+        module = _import_module_with_stub()
         module.refresh_context_on_transformer(
             transformer="transformer",
             num_inference_steps=50,
@@ -166,7 +167,7 @@ class TestCacheDitRefreshContext(unittest.TestCase):
         )
 
     def test_refresh_context_with_scm_preset_uses_steps_mask(self):
-        module = self._import_module_with_stub()
+        module = _import_module_with_stub()
         module.refresh_context_on_transformer(
             transformer="transformer",
             num_inference_steps=8,
@@ -187,7 +188,7 @@ class TestCacheDitRefreshContext(unittest.TestCase):
         )
 
     def test_dual_refresh_without_scm_preset_skips_steps_mask(self):
-        module = self._import_module_with_stub()
+        module = _import_module_with_stub()
         module.refresh_context_on_dual_transformer(
             transformer="transformer",
             transformer_2="transformer_2",
@@ -216,5 +217,39 @@ class TestCacheDitRefreshContext(unittest.TestCase):
         )
 
 
+def _make_transformer(class_name, layers=None):
+    transformer = type(class_name, (), {})()
+    if layers is not None:
+        transformer.layers = layers
+    return transformer
+
+
+class TestBuildCustomBlockAdapter(unittest.TestCase):
+    def test_builds_adapter_for_registered_class(self):
+        module = _import_module_with_stub()
+        blocks = ["block_0", "block_1"]
+        transformer = _make_transformer("ErnieImageTransformer2DModel", blocks)
+
+        adapter = module._build_custom_block_adapter(transformer)
+
+        self.assertIsNotNone(adapter)
+        self.assertEqual(adapter.blocks, blocks)
+        self.assertEqual(adapter.forward_pattern, "Pattern_3")
+        self.assertTrue(adapter.has_separate_cfg)
+
+    def test_returns_none_for_unknown_class(self):
+        module = _import_module_with_stub()
+        transformer = _make_transformer("SomeUnregisteredTransformer", ["b0"])
+
+        self.assertIsNone(module._build_custom_block_adapter(transformer))
+
+    def test_raises_when_blocks_attr_missing(self):
+        module = _import_module_with_stub()
+        transformer = _make_transformer("ErnieImageTransformer2DModel")
+
+        with self.assertRaises(ValueError):
+            module._build_custom_block_adapter(transformer)
+
+
 if __name__ == "__main__":
     unittest.main()