[Refactor][1/N]diffusion/cache: cache-dit (#5226)

Signed-off-by: yangjianjuan <66548561+yangjianjuan@users.noreply.github.com>

tests/diffusion/cache/test_cache_backends.py

@@ -17,13 +17,14 @@ import cache_dit
 import pytest
 from cache_dit import ForwardPattern
 
-from vllm_omni.diffusion.cache.cache_dit_backend import (
+from vllm_omni.diffusion.cache.cachedit import (
     CacheDiTAdapterConfig,
     CacheDiTBackend,
+    CacheDiTConfig,
 )
-from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
+from vllm_omni.diffusion.cache.magcache import MagCacheBackend
 from vllm_omni.diffusion.cache.selector import get_cache_backend
-from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
+from vllm_omni.diffusion.cache.teacache import TeaCacheBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -32,6 +33,19 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 class TestCacheDiTBackend:
     """Test CacheDiTBackend implementation."""
 
+    def test_config_projection_contains_only_cache_dit_fields(self):
+        shared_config = DiffusionCacheConfig(
+            Fn_compute_blocks=4,
+            max_warmup_steps=8,
+            rel_l1_thresh=0.3,
+        )
+
+        config = CacheDiTConfig.from_diffusion_config(shared_config)
+
+        assert config.Fn_compute_blocks == 4
+        assert config.max_warmup_steps == 8
+        assert not hasattr(config, "rel_l1_thresh")
+
     def test_init_with_dict(self):
         """Test initialization with dictionary config."""
         config_dict = {"Fn_compute_blocks": 4, "max_warmup_steps": 8}
@@ -47,8 +61,8 @@ 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")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
     def test_enable_single_transformer(self, mock_cache_dit, mock_block_adapter):
         """Test enabling cache-dit on single-transformer pipeline."""
         # Mock pipeline
@@ -75,10 +89,61 @@ 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")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.logger")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
+    def test_enable_without_model_adapter_uses_cache_dit_registry(self, mock_cache_dit, mock_logger):
+        """A missing model-declared adapter falls back to Cache-DiT's registry."""
+
+        class BuiltinTransformer:
+            pass
+
+        class BuiltinPipeline:
+            transformer = BuiltinTransformer()
+
+        pipeline = BuiltinPipeline()
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+
+        backend.enable(pipeline)
+
+        assert backend.enabled is True
+        assert backend._refresh_func is not None
+        assert mock_cache_dit.enable_cache.call_args.args[0] is pipeline.transformer
+        mock_logger.info.assert_any_call(
+            "Transformer %s does not declare _cache_dit_adapter_config; "
+            "falling back to Cache-DiT's built-in adapter registry.",
+            "BuiltinTransformer",
+        )
+
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
+    def test_enable_without_compatible_adapter_has_contextual_error(self, mock_cache_dit):
+        """An unsupported fallback reports both pipeline and transformer names."""
+
+        class UnsupportedTransformer:
+            pass
+
+        class UnsupportedPipeline:
+            transformer = UnsupportedTransformer()
+
+        pipeline = UnsupportedPipeline()
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+        mock_cache_dit.enable_cache.side_effect = ValueError("unsupported")
+
+        with pytest.raises(
+            ValueError,
+            match=(
+                "Failed to enable Cache-DiT for pipeline UnsupportedPipeline with transformer UnsupportedTransformer"
+            ),
+        ) as exc_info:
+            backend.enable(pipeline)
+
+        assert isinstance(exc_info.value.__cause__, ValueError)
+        assert backend.enabled is False
+        assert backend._refresh_func is None
+
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.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."""
+        """Test refreshing cache context with SCM mask policy updates."""
         # Mock pipeline
         mock_pipeline = Mock()
         mock_pipeline.__class__.__name__ = "DiTPipeline"
@@ -107,8 +172,6 @@ class TestCacheDiTBackend:
 
         # First refresh with 50 steps
         backend.refresh(mock_pipeline, num_inference_steps=50)
-        assert backend._last_num_inference_steps == 50
-
         # Verify steps_mask was called with mask policy (not direct steps mask)
         mock_cache_dit.steps_mask.assert_called_with(mask_policy="fast", total_steps=50)
         assert mock_cache_dit.steps_mask.call_count == 1
@@ -137,10 +200,28 @@ class TestCacheDiTBackend:
         call_args = mock_cache_dit.refresh_context.call_args
         assert call_args[0][0] == mock_transformer
         assert "cache_config" in call_args[1]
-        assert backend._last_num_inference_steps == 100
+        assert mock_cache_dit.refresh_context.call_count == 1
 
-    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
-    @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
+    def test_refresh_resets_same_step_count(self, mock_cache_dit, mock_block_adapter):
+        """Every generation must reset context, even when the step count is unchanged."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "DiTPipeline"
+        mock_pipeline.transformer = Mock()
+        mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
+            block_forward_patterns={"layers": ForwardPattern.Pattern_0}
+        )
+
+        backend = CacheDiTBackend(DiffusionCacheConfig())
+        backend.enable(mock_pipeline)
+        backend.refresh(mock_pipeline, num_inference_steps=20)
+        backend.refresh(mock_pipeline, num_inference_steps=20)
+
+        assert mock_cache_dit.refresh_context.call_count == 2
+
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
     def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
         """Test HunyuanImage3 uses pipeline.transformer for cache enable/refresh.
 
@@ -179,8 +260,8 @@ class TestCacheDiTBackend:
         assert call_args[0][0] is mock_pipeline.model
         assert call_args[1]["num_inference_steps"] == 12
 
-    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
-    @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cachedit.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.
 
@@ -221,8 +302,8 @@ class TestCacheDiTBackend:
         assert call_args[1]["num_inference_steps"] == 12
 
     @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")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
     def test_refresh_scm_bypassed_for_unsupported_step_counts(
         self, mock_cache_dit, mock_block_adapter, num_inference_steps
     ):
@@ -362,8 +443,6 @@ class TestCacheSelector:
 class TestMagCacheBackend:
     """Test MagCacheBackend implementation."""
 
-    from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
-
     def test_init(self):
         """Test initialization."""
         config = DiffusionCacheConfig(mag_threshold=0.1, mag_max_skip_steps=2, mag_calibrate=True)

tests/diffusion/cache/test_cache_dit.py

@@ -5,15 +5,18 @@
 Model specific tests for CacheDiT enablement.
 """
 
+import ast
 import sys
+from pathlib import Path
 from unittest.mock import Mock, patch
 
 import pytest
 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 CacheDiTAdapterConfig, CacheDiTBackend, cache_summary
+import vllm_omni.diffusion.cache.cachedit as cd_backend
+import vllm_omni.diffusion.cache.cachedit.model_specific as cd_model_specific
+from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig, CacheDiTBackend, cache_summary
 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
@@ -42,8 +45,73 @@ SEPARATE_CFG_TRANSFORMERS = [
 SAMPLE_CACHE_CONFIG = DiffusionCacheConfig()
 
 
+def test_custom_cache_dit_enablers_are_registered_explicitly():
+    expected_enablers = {
+        "Wan22Pipeline": cd_model_specific.enable_cache_for_wan22,
+        "Wan22I2VPipeline": cd_model_specific.enable_cache_for_wan22,
+        "Wan22TI2VPipeline": cd_model_specific.enable_cache_for_wan22,
+        "Wan22VACEPipeline": cd_model_specific.enable_cache_for_wan22,
+        "Wan22S2VPipeline": cd_model_specific.enable_cache_for_wan22_s2v,
+        "Cosmos3OmniDiffusersPipeline": cd_model_specific.enable_cache_for_cosmos3,
+        "Cosmos3OmniPipeline": cd_model_specific.enable_cache_for_cosmos3,
+        "Krea2Pipeline": cd_model_specific.enable_cache_for_krea2,
+    }
+
+    with patch.dict(cd_backend.CUSTOM_DIT_ENABLERS, {}, clear=True):
+        cd_model_specific.register_custom_dit_enablers()
+        assert cd_backend.CUSTOM_DIT_ENABLERS == expected_enablers
+
+
 def test_wan22_vace_uses_wan22_custom_cache_dit_enabler():
-    assert cd_backend.CUSTOM_DIT_ENABLERS["Wan22VACEPipeline"] is cd_backend.enable_cache_for_wan22
+    assert cd_backend.CUSTOM_DIT_ENABLERS["Wan22VACEPipeline"] is cd_model_specific.enable_cache_for_wan22
+
+
+@pytest.mark.parametrize(
+    "pipeline_name",
+    ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"],
+)
+def test_cosmos3_aliases_use_cosmos3_custom_cache_dit_enabler(pipeline_name: str):
+    assert cd_backend.CUSTOM_DIT_ENABLERS[pipeline_name] is cd_model_specific.enable_cache_for_cosmos3
+
+
+def test_cachedit_public_api_is_explicit():
+    assert set(cd_backend.__all__) == {
+        "BagelCachedAdapter",
+        "CUSTOM_DIT_ENABLERS",
+        "CacheDiTAdapterConfig",
+        "CacheDiTBackend",
+        "CacheDiTConfig",
+        "SensenovaCachedAdapter",
+        "cache_summary",
+        "enable_cache_for_dit",
+    }
+    assert not hasattr(cd_backend, "enable_cache_for_wan22")
+    assert not hasattr(cd_backend, "enable_cache_for_wan22_s2v")
+
+
+def test_cachedit_consumers_use_package_api():
+    cache_dir = Path(cd_backend.__file__).resolve().parents[1]
+    package_root = cache_dir.parents[1]
+    legacy_module = "vllm_omni.diffusion.cache.cache_dit_backend"
+    internal_prefix = "vllm_omni.diffusion.cache.cachedit."
+
+    invalid_imports = []
+    for source_path in package_root.rglob("*.py"):
+        if source_path.is_relative_to(cache_dir):
+            continue
+        tree = ast.parse(source_path.read_text(encoding="utf-8"))
+        for node in ast.walk(tree):
+            modules = []
+            if isinstance(node, ast.ImportFrom) and node.module is not None:
+                modules.append(node.module)
+            elif isinstance(node, ast.Import):
+                modules.extend(alias.name for alias in node.names)
+
+            for module in modules:
+                if module == legacy_module or module.startswith(internal_prefix):
+                    invalid_imports.append(f"{source_path.relative_to(package_root)}:{node.lineno}: {module}")
+
+    assert not invalid_imports, "Cache-DiT consumers must use the package API:\n" + "\n".join(invalid_imports)
 
 
 @pytest.mark.parametrize("transformer_model", SEPARATE_CFG_TRANSFORMERS)
@@ -54,12 +122,12 @@ def test_cache_dit_configs_have_separate_cfg(transformer_model):
     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")
+@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
+@patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit")
 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()
-    cd_backend.enable_cache_for_wan22(mock_pipeline, SAMPLE_CACHE_CONFIG)
+    cd_model_specific.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
@@ -67,8 +135,8 @@ def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_blo
 
 
 @pytest.mark.parametrize("has_transformer_2", [False, True])
-@patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
-@patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
+@patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit")
 def test_wan22_custom_enabler_passes_taylorseer_calibrator(
     mock_cache_dit,
     mock_block_adapter,
@@ -82,7 +150,7 @@ def test_wan22_custom_enabler_passes_taylorseer_calibrator(
         mock_pipeline.transformer_2 = None
     cache_config = DiffusionCacheConfig(enable_taylorseer=True, taylorseer_order=1)
 
-    cd_backend.enable_cache_for_wan22(mock_pipeline, cache_config)
+    cd_model_specific.enable_cache_for_wan22(mock_pipeline, cache_config)
 
     enable_cache_kwargs = mock_cache_dit.enable_cache.call_args.kwargs
     calibrator_config = enable_cache_kwargs["calibrator_config"]
@@ -94,8 +162,8 @@ def test_wan22_custom_enabler_passes_taylorseer_calibrator(
         assert modifier._context_kwargs["calibrator_config"] is calibrator_config
 
 
-@patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
-@patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
+@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
 def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter):
     """Cosmos3 should cache only the repeated GEN pathway blocks."""
     mock_pipeline = Mock()
@@ -103,7 +171,7 @@ def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter):
     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)
+    cd_model_specific.enable_cache_for_cosmos3(mock_pipeline, SAMPLE_CACHE_CONFIG)
 
     mock_cache_dit.enable_cache.assert_called_once()
     adapter_kwargs = mock_block_adapter.call_args.kwargs