From c67d338637f2574fcf66f64b233557c15f8b310c Mon Sep 17 00:00:00 2001 From: Thomas Date: Tue, 23 Jun 2026 21:08:03 +0800 Subject: [PATCH] [Diffusion] enable cache-dit for ERNIE-Image model (#28266) --- .../runtime/cache/cache_dit_integration.py | 69 +++++++++++++--- .../test/unit/test_cache_dit_integration.py | 81 +++++++++++++------ 2 files changed, 117 insertions(+), 33 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py b/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py index fb5f0f712..e204404ab 100644 --- a/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py +++ b/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, diff --git a/python/sglang/multimodal_gen/test/unit/test_cache_dit_integration.py b/python/sglang/multimodal_gen/test/unit/test_cache_dit_integration.py index a4834f0ca..4075259e3 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cache_dit_integration.py +++ b/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()