[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)
 

tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py

@@ -692,16 +692,27 @@ class TestRegistryIntegration:
             assert name in _DIFFUSION_POST_PROCESS_FUNCS, f"{name} not in _DIFFUSION_POST_PROCESS_FUNCS"
             assert _DIFFUSION_POST_PROCESS_FUNCS[name] == "get_ltx2_post_process_func"
 
-    def test_cache_dit_enablers_registered(self):
-        """Pipeline variants must be registered in CUSTOM_DIT_ENABLERS."""
+    def test_cache_dit_for_ltx2_does_not_have_custom_enablers_registered(self):
+        """Pipeline variants are *not* registered in CUSTOM_DIT_ENABLERS."""
         from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS
 
+        # NOTE: We used to have custom enablers for this model, but refactored to handle
+        # it more generically. Now we only need to ensure it has git cache adapter config.
         expected = [
             "LTX23Pipeline",
             "LTX23ImageToVideoPipeline",
         ]
         for name in expected:
-            assert name in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS"
+            assert name not in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS"
+
+    def test_ltx2_transformer_has_dit_cache_config(self):
+        """Ensure LTX2 has a Cache DiT adapter config and that it uses separate CFG."""
+        from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
+        from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
+
+        adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
+        assert isinstance(adapter_config, CacheDiTAdapterConfig)
+        assert adapter_config.has_separate_cfg
 
 
 class TestVocoderSampleRateDetection:

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -13,6 +13,7 @@ from dataclasses import dataclass
 import numpy as np
 import torch
 import torch.distributed as dist
+from cache_dit import ForwardPattern
 from torch import nn
 from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
 from transformers.models.qwen2.modeling_qwen2 import (
@@ -36,6 +37,7 @@ from vllm.transformers_utils.configs.bagel import BagelConfig
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata as DiffusionAttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention
+from vllm_omni.diffusion.cache.cache_dit_backend import BagelCachedAdapter, CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_cfg_group,
@@ -772,6 +774,13 @@ class Qwen2MoTDecoderLayer(nn.Module):
 
 
 class Qwen2MoTModel(Qwen2PreTrainedModel):
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "layers": ForwardPattern.Pattern_0,
+        },
+        cached_adapter_cls=BagelCachedAdapter,
+    )
+
     _layerwise_offload_blocks_attrs = ["layers"]
 
     @staticmethod

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -18,6 +18,7 @@ import torch
 import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
+from cache_dit import ForwardPattern
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import (
@@ -30,6 +31,7 @@ from vllm.model_executor.layers.quantization.base_config import (
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
@@ -964,6 +966,16 @@ class Cosmos3VFMTransformer(nn.Module):
     ``Attention`` layer (with Ulysses all-to-all) or plain SDPA accordingly.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        # Cosmos3 GEN blocks return only hidden_states.  Per-layer UND K/V
+        # conditioning uses the transformer's cache-dit fallback path.
+        block_forward_patterns={
+            "gen_layers": ForwardPattern.Pattern_3,
+        },
+        has_separate_cfg=True,
+        check_forward_pattern=False,
+    )
+
     _repeated_blocks = ["Cosmos3GenDecoderLayer"]
 
     _layerwise_offload_blocks_attrs = ["gen_layers"]

vllm_omni/diffusion/models/dreamid_omni/fusion.py

@@ -3,11 +3,13 @@ from typing import TYPE_CHECKING
 
 import torch
 import torch.nn as nn
+from cache_dit import ForwardPattern
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
@@ -251,6 +253,13 @@ class FusedBlock(nn.Module):
 
 
 class FusionModel(nn.Module):
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "fused_blocks": ForwardPattern.Pattern_0,
+        },
+        has_separate_cfg=True,
+    )
+
     _layerwise_offload_blocks_attrs = ["fused_blocks"]
 
     packed_modules_mapping = {

vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING
 
 import torch
 import torch.nn as nn
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from vllm.distributed import get_tensor_model_parallel_world_size
@@ -21,6 +22,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
@@ -459,6 +461,12 @@ class ErnieImageAdaLNContinuous(nn.Module):
 
 
 class ErnieImageTransformer2DModel(nn.Module):
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "layers": ForwardPattern.Pattern_3,
+        },
+        check_forward_pattern=False,
+    )
     _repeated_blocks = ["ErnieImageSharedAdaLNBlock"]
     _layerwise_offload_blocks_attrs = ["layers"]
 

vllm_omni/diffusion/models/flux/flux_transformer.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import (
     CombinedTimestepGuidanceTextProjEmbeddings,
     CombinedTimestepTextProjEmbeddings,
@@ -25,6 +26,7 @@ from vllm.model_executor.layers.linear import (
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
@@ -524,6 +526,13 @@ class FluxTransformer2DModel(nn.Module):
             The dimensions to use for the rotary positional embeddings.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_1,
+            "single_transformer_blocks": ForwardPattern.Pattern_1,
+        }
+    )
+
     # the small and frequently-repeated block(s) of a model
     # -- typically a transformer layer
     # used for torch compile optimizations

vllm_omni/diffusion/models/flux2/flux2_transformer.py

@@ -5,6 +5,7 @@ from types import SimpleNamespace
 from typing import TYPE_CHECKING, Any
 
 import torch
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import (
     TimestepEmbedding,
     Timesteps,
@@ -25,6 +26,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
@@ -755,6 +757,13 @@ class Flux2Transformer2DModel(nn.Module):
     Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_1,
+            "single_transformer_blocks": ForwardPattern.Pattern_2,
+        }
+    )
+
     _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
     _sp_plan = {
         "": {

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

@@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn as nn
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import (
     TimestepEmbedding,
     Timesteps,
@@ -40,6 +41,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
@@ -768,6 +770,13 @@ class Flux2Transformer2DModel(nn.Module):
     Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_1,
+            "single_transformer_blocks": ForwardPattern.Pattern_2,
+        }
+    )
+
     _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 

vllm_omni/diffusion/models/glm_image/glm_image_transformer.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from cache_dit import ForwardPattern
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.transformers.transformer_glm_image import GlmImageCombinedTimestepSizeEmbeddings
 from vllm.logger import init_logger
@@ -20,6 +21,8 @@ from vllm.model_executor.layers.linear import (
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
+
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
@@ -935,6 +938,13 @@ class GlmImageTransformer2DModel(CachedTransformer):
             `od_config.tf_model_config`.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_0,
+        },
+        has_separate_cfg=True,
+    )
+
     _repeated_blocks = ["GlmImageTransformerBlock"]
     # SP plan using GlmImagePrepare module for sharding hidden_states and RoPE together.
     # Similar to Qwen-Image's ImageRopePrepare, this creates a module boundary where

vllm_omni/diffusion/models/helios/helios_transformer.py

@@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import FP32LayerNorm
@@ -25,6 +26,7 @@ from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelL
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
     SequenceParallelOutput,
@@ -602,6 +604,13 @@ class HeliosTransformer3DModel(nn.Module):
     guidance cross-attention, and chunked video generation support.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "blocks": ForwardPattern.Pattern_2,
+        },
+        has_separate_cfg=True,
+    )
+
     _repeated_blocks = ["HeliosTransformerBlock"]
     _layerwise_offload_blocks_attrs = ["blocks"]
     packed_modules_mapping = {

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -11,6 +11,7 @@ from typing import Any, cast
 import numpy as np
 import regex as re
 import torch
+from cache_dit import ForwardPattern
 from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
 from diffusers.image_processor import VaeImageProcessor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
@@ -59,6 +60,7 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionMetadata,
 )
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_cfg_group,
     get_classifier_free_guidance_rank,
@@ -2018,6 +2020,13 @@ class HunyuanImagePostprocessor(nn.Module):
 
 
 class HunyuanImage3Model(nn.Module):
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "layers": ForwardPattern.Pattern_4,
+        },
+        check_forward_pattern=False,
+    )
+
     _sp_plan = {
         # Split custom_pos_emb tuple elements (cos, sin) at model forward input
         "pre_processor": {

vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn.functional as F
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import (
     CombinedTimestepTextProjEmbeddings,
     TimestepEmbedding,
@@ -22,6 +23,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
 from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size
@@ -591,6 +593,13 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
     tensor-parallel layers for the 54 main transformer blocks.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_0,
+        },
+        has_separate_cfg=True,
+    )
+
     _repeated_blocks = ["HunyuanVideo15TransformerBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     packed_modules_mapping = {

vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py

@@ -6,6 +6,7 @@ from typing import Any
 
 import torch
 import torch.nn as nn
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
@@ -18,6 +19,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
@@ -581,6 +583,14 @@ class LongCatImageTransformer2DModel(nn.Module):
     Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_1,
+            "single_transformer_blocks": ForwardPattern.Pattern_1,
+        },
+        has_separate_cfg=True,
+    )
+
     _repeated_blocks = ["LongCatImageTransformerBlock", "LongCatImageSingleTransformerBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 

vllm_omni/diffusion/models/ltx2/ltx2_transformer.py

@@ -23,6 +23,7 @@ import torch
 import torch.distributed
 import torch.nn as nn
 import torch.nn.functional as F
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings, PixArtAlphaTextProjection
 from diffusers.utils import (
     BaseOutput,
@@ -44,6 +45,7 @@ if TYPE_CHECKING:
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
@@ -1449,6 +1451,14 @@ class LTX2VideoTransformer3DModel(nn.Module):
             The normalization layer to use.
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_0,
+        },
+        has_separate_cfg=True,
+        check_forward_pattern=False,
+    )
+
     _supports_gradient_checkpointing = True
     _skip_layerwise_casting_patterns = ["norm"]
     _repeated_blocks = ["LTX2VideoTransformerBlock"]

vllm_omni/diffusion/models/sd3/sd3_transformer.py

@@ -5,6 +5,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 # TODO replace this with vLLM implementation
+from cache_dit import ForwardPattern
 from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, SD35AdaLayerNormZeroX
@@ -19,6 +20,7 @@ from vllm.model_executor.layers.linear import (
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 
 logger = init_logger(__name__)
@@ -386,6 +388,12 @@ class SD3Transformer2DModel(nn.Module):
     The Transformer model introduced in [Stable Diffusion 3](https://huggingface.co/papers/2403.03206).
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_1,
+        }
+    )
+
     _repeated_blocks = ["SD3TransformerBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
 

vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py

@@ -534,6 +534,8 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
             self.llm_cfg,
             prefix="language_model",
         )
+        # We define this for Cache DiT compatibility
+        self.transformer = self.language_model.model
 
         # Vision model (understanding branch)
         self.vision_model = NEOVisionModel(self.vis_cfg)

vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py

@@ -16,6 +16,7 @@ from dataclasses import dataclass
 
 import torch
 import torch.nn as nn
+from cache_dit import ForwardPattern
 from transformers.cache_utils import DynamicCache
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -30,6 +31,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, SensenovaCachedAdapter
 
 
 @dataclass
@@ -606,6 +608,14 @@ class SenseNovaU1DecoderLayer(nn.Module):
 
 
 class SenseNovaU1Model(nn.Module):
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "layers": ForwardPattern.Pattern_3,
+        },
+        has_separate_cfg=True,
+        cached_adapter_cls=SensenovaCachedAdapter,
+    )
+
     def __init__(self, config, quant_config=None, prefix: str = ""):
         super().__init__()
         self.config = config

vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py

@@ -9,12 +9,14 @@ from collections.abc import Iterable
 
 import torch
 import torch.nn as nn
+from cache_dit import ForwardPattern
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import ReplicatedLinear
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
 from vllm_omni.diffusion.layers.fourier import GaussianFourierProjection
@@ -464,6 +466,12 @@ class StableAudioDiTModel(nn.Module):
     - Output: [B, out_channels, L]
     """
 
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={
+            "transformer_blocks": ForwardPattern.Pattern_3,
+        },
+    )
+
     _repeated_blocks = ["StableAudioDiTBlock"]
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     _hsdp_shard_conditions = [is_transformer_block_module]