[Cache Refactor 1/N] Simplify CacheDiT Integration (#2527)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/diffusion/cache/test_cache_backends.py

@@ -13,10 +13,12 @@ This module tests the cache backend implementations:
 
 from unittest.mock import Mock, patch
 
+import cache_dit
 import pytest
+from cache_dit import ForwardPattern
 
 from vllm_omni.diffusion.cache.cache_dit_backend import (
-    CUSTOM_DIT_ENABLERS,
+    CacheDiTAdapterConfig,
     CacheDiTBackend,
 )
 from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
@@ -45,13 +47,20 @@ class TestCacheDiTBackend:
         assert backend.config.Fn_compute_blocks == 4
         assert backend.enabled is False
 
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
-    def test_enable_single_transformer(self, mock_cache_dit):
+    def test_enable_single_transformer(self, mock_cache_dit, mock_block_adapter):
         """Test enabling cache-dit on single-transformer pipeline."""
         # Mock pipeline
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "DiTPipeline"
         mock_transformer = Mock()
+        mock_transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={
+                "layers": ForwardPattern.Pattern_0,
+            },
+        )
+
         mock_pipeline.transformer = mock_transformer
 
         # Mock cache_dit functions
@@ -66,14 +75,20 @@ class TestCacheDiTBackend:
         assert backend._refresh_func is not None
         mock_cache_dit.enable_cache.assert_called_once()
 
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
-    def test_refresh(self, mock_cache_dit):
+    def test_refresh(self, mock_cache_dit, mock_block_adapter):
         """Test refreshing cache context with SCM mask policy updates when num_inference_steps changes."""
         # Mock pipeline
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "DiTPipeline"
         mock_transformer = Mock()
         mock_pipeline.transformer = mock_transformer
+        mock_transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={
+                "layers": ForwardPattern.Pattern_0,
+            },
+        )
 
         # Mock cache_dit functions
         mock_cache_dit.enable_cache = Mock()
@@ -124,23 +139,24 @@ class TestCacheDiTBackend:
         assert "cache_config" in call_args[1]
         assert backend._last_num_inference_steps == 100
 
-    def test_hunyuan_custom_enabler_registered(self):
-        """Test HunyuanImage3 custom cache-dit enabler is registered."""
-        assert "HunyuanImage3Pipeline" in CUSTOM_DIT_ENABLERS
-
-    def test_dreamid_custom_enabler_registered(self):
-        """Test DreamIDOmni custom cache-dit enabler is registered."""
-        assert "DreamIDOmniPipeline" in CUSTOM_DIT_ENABLERS
-
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
     def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
-        """Test HunyuanImage3 custom enabler uses pipeline.model for cache enable/refresh."""
+        """Test HunyuanImage3 uses pipeline.transformer for cache enable/refresh.
+
+        NOTE: HunyuanImage3 no longer has a custom enabler, so this tests against the generic path.
+        """
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline"
         mock_pipeline.model = Mock()
         mock_pipeline.model.layers = Mock()
-
+        mock_pipeline.model._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={
+                "layers": ForwardPattern.Pattern_4,
+            },
+        )
+        # NOTE: pipe.transformer is pipe.model in HunyuanImage3Pipelines
+        mock_pipeline.transformer = mock_pipeline.model
         mock_cache_dit.enable_cache = Mock()
         mock_cache_dit.refresh_context = Mock()
 
@@ -152,9 +168,9 @@ class TestCacheDiTBackend:
         mock_block_adapter.assert_called_once()
         adapter_kwargs = mock_block_adapter.call_args.kwargs
         assert adapter_kwargs["transformer"] is mock_pipeline.model
-        assert adapter_kwargs["blocks"] is mock_pipeline.model.layers
-        assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_4
-        assert len(adapter_kwargs["params_modifiers"]) == 1
+        assert len(adapter_kwargs["blocks"]) == 1
+        assert adapter_kwargs["blocks"][0] == mock_pipeline.model.layers
+        assert adapter_kwargs["forward_pattern"][0] == ForwardPattern.Pattern_4
         mock_cache_dit.enable_cache.assert_called_once()
 
         backend.refresh(mock_pipeline, num_inference_steps=12)
@@ -163,24 +179,23 @@ class TestCacheDiTBackend:
         assert call_args[0][0] is mock_pipeline.model
         assert call_args[1]["num_inference_steps"] == 12
 
-    def test_enable_hunyuan_pipeline_requires_model_layers(self):
-        """Test HunyuanImage3 enabler fails with a formatted pipeline class name."""
-        mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline"
-        mock_pipeline.model = Mock(spec=[])
-
-        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
-
-        with pytest.raises(ValueError, match="HunyuanImage3Pipeline"):
-            backend.enable(mock_pipeline)
-
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
     @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
     def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
+        """Test DreamID uses pipeline.transformer for cache enable/refresh.
+
+        NOTE: DreamID no longer has a custom enabler, so this tests against the generic path.
+        """
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
         mock_pipeline.transformer = Mock()
         mock_pipeline.transformer.fused_blocks = Mock()
+        mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={
+                "fused_blocks": ForwardPattern.Pattern_0,
+            },
+            has_separate_cfg=True,
+        )
 
         mock_cache_dit.enable_cache = Mock()
         mock_cache_dit.refresh_context = Mock()
@@ -193,10 +208,10 @@ class TestCacheDiTBackend:
         mock_block_adapter.assert_called_once()
         adapter_kwargs = mock_block_adapter.call_args.kwargs
         assert adapter_kwargs["transformer"] is mock_pipeline.transformer
-        assert adapter_kwargs["blocks"] is mock_pipeline.transformer.fused_blocks
-        assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_0
+        assert len(adapter_kwargs["blocks"]) == 1
+        assert adapter_kwargs["blocks"][0] == mock_pipeline.transformer.fused_blocks
+        assert adapter_kwargs["forward_pattern"][0] == ForwardPattern.Pattern_0
         assert adapter_kwargs["has_separate_cfg"] is True
-        assert len(adapter_kwargs["params_modifiers"]) == 1
         mock_cache_dit.enable_cache.assert_called_once()
 
         backend.refresh(mock_pipeline, num_inference_steps=12)
@@ -205,15 +220,43 @@ class TestCacheDiTBackend:
         assert call_args[0][0] is mock_pipeline.transformer
         assert call_args[1]["num_inference_steps"] == 12
 
-    def test_enable_dreamid_pipeline_requires_fused_blocks(self):
+    @pytest.mark.parametrize("num_inference_steps", [1, 7])
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+    def test_refresh_scm_bypassed_for_unsupported_step_counts(
+        self, mock_cache_dit, mock_block_adapter, num_inference_steps
+    ):
+        """Ensure SCM is bypassed when num_inference_steps < 8 and not in (4, 6),
+        because cache_dit.steps_mask() raises for unsupported step count.
+        For these cases, we fall back to the non-SCM path to avoid crashing.
+        """
         mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
-        mock_pipeline.transformer = Mock(spec=[])
+        mock_pipeline.__class__.__name__ = "DiTPipeline"
+        mock_transformer = Mock()
+        mock_pipeline.transformer = mock_transformer
+        mock_transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={
+                "layers": ForwardPattern.Pattern_0,
+            },
+        )
 
-        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+        mock_cache_dit.enable_cache = Mock()
+        mock_cache_dit.refresh_context = Mock()
+        mock_cache_dit.steps_mask = cache_dit.steps_mask
 
-        with pytest.raises(ValueError, match="DreamIDOmniPipeline"):
-            backend.enable(mock_pipeline)
+        # Create a cache config with an scm policy & enable it
+        config = DiffusionCacheConfig(scm_steps_mask_policy="fast")
+        backend = CacheDiTBackend(config)
+        backend.enable(mock_pipeline)
+
+        backend.refresh(mock_pipeline, num_inference_steps=num_inference_steps)
+
+        mock_cache_dit.refresh_context.assert_called_once()
+        call_args = mock_cache_dit.refresh_context.call_args
+        # Ensure that we properly guard, i.e., cache config is filtered
+        assert call_args[0][0] == mock_transformer
+        assert call_args[1]["num_inference_steps"] == num_inference_steps
+        assert "cache_config" not in call_args[1]
 
 
 class TestTeaCacheBackend:

tests/diffusion/cache/test_cache_dit.py

@@ -5,6 +5,7 @@
 Model specific tests for CacheDiT enablement.
 """
 
+import sys
 from unittest.mock import Mock, patch
 
 import pytest
@@ -12,36 +13,49 @@ import torch
 from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
 
 import vllm_omni.diffusion.cache.cache_dit_backend as cd_backend
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, CacheDiTBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
+from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
+from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel
 from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
 from vllm_omni.platforms import current_omni_platform
 
+# NOTE: We patch DreamID Omni's modules here with mocks so that we can import and inspect
+# the class even though the dependency may not be set up correctly; this is ok for these
+# tests because we just inspect it and never initialize the model.
+for mod in ("dreamid_omni", "dreamid_omni.modules", "dreamid_omni.modules.model"):
+    sys.modules.setdefault(mod, Mock())
+# isort: split
+from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel as DreamIdOmniModel  # noqa: E402
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
-SEPARATE_CFG_ENABLERS = [
-    cd_backend.enable_cache_for_dreamid_omni,
-    cd_backend.enable_cache_for_ltx2,
-    cd_backend.enable_cache_for_helios,
-    cd_backend.enable_cache_for_wan22,
-    cd_backend.enable_cache_for_longcat_image,
-    cd_backend.enable_cache_for_cosmos3,
+SEPARATE_CFG_TRANSFORMERS = [
+    DreamIdOmniModel,
+    LTX2VideoTransformer3DModel,
+    HeliosTransformer3DModel,
+    LongCatImageTransformer2DModel,
+    Cosmos3VFMTransformer,
 ]
 
 SAMPLE_CACHE_CONFIG = DiffusionCacheConfig()
 
 
-@pytest.mark.parametrize("enabler", SEPARATE_CFG_ENABLERS)
+@pytest.mark.parametrize("transformer_model", SEPARATE_CFG_TRANSFORMERS)
+def test_cache_dit_configs_have_separate_cfg(transformer_model):
+    """Check models with separate CFG set has_separate_cfg=True in their configs."""
+    assert hasattr(transformer_model, "_cache_dit_adapter_config")
+    assert isinstance(transformer_model._cache_dit_adapter_config, CacheDiTAdapterConfig)
+    assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True
+
+
 @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
 @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
-def test_separate_cfg(mock_cache_dit, mock_block_adapter, enabler):
-    """Ensure that custom enablers for models with separate CFG pass
-    the param through to cache_dit correctly.
-
-    Regression test for: https://github.com/vllm-project/vllm-omni/pull/2860
-    """
+def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter):
+    """Ensure that Wan22, which has a custom enabler, setts custom CFG correctly."""
     mock_pipeline = Mock()
-    enabler(mock_pipeline, SAMPLE_CACHE_CONFIG)
+    cd_backend.enable_cache_for_wan22(mock_pipeline, SAMPLE_CACHE_CONFIG)
 
     mock_cache_dit.enable_cache.assert_called_once()
     adapter_kwargs = mock_block_adapter.call_args.kwargs
@@ -55,6 +69,7 @@ def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter):
     mock_pipeline = Mock()
     gen_layers = object()
     mock_pipeline.transformer.gen_layers = gen_layers
+    mock_pipeline.transformer._cache_dit_adapter_config = Cosmos3VFMTransformer._cache_dit_adapter_config
 
     cd_backend.enable_cache_for_cosmos3(mock_pipeline, SAMPLE_CACHE_CONFIG)