From 6ff246bf3cac978b34c8404c17fd8125ad8201e9 Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Fri, 19 Jun 2026 18:30:44 -0600 Subject: [PATCH] [Cache Refactor 1/N] Simplify CacheDiT Integration (#2527) Signed-off-by: Alex Brooks --- tests/diffusion/cache/test_cache_backends.py | 115 +- tests/diffusion/cache/test_cache_dit.py | 47 +- .../models/ltx2/test_ltx2_3_pipeline.py | 17 +- .../diffusion/cache/cache_dit_backend.py | 1511 +++-------------- .../models/bagel/bagel_transformer.py | 9 + .../models/cosmos3/transformer_cosmos3.py | 12 + .../diffusion/models/dreamid_omni/fusion.py | 9 + .../ernie_image/ernie_image_transformer.py | 8 + .../diffusion/models/flux/flux_transformer.py | 9 + .../models/flux2/flux2_transformer.py | 9 + .../flux2_klein/flux2_klein_transformer.py | 9 + .../models/glm_image/glm_image_transformer.py | 10 + .../models/helios/helios_transformer.py | 9 + .../hunyuan_image3_transformer.py | 9 + .../hunyuan_video_15_transformer.py | 9 + .../longcat_image_transformer.py | 10 + .../diffusion/models/ltx2/ltx2_transformer.py | 10 + .../diffusion/models/sd3/sd3_transformer.py | 8 + .../sensenova_u1/pipeline_sensenova_u1.py | 2 + .../sensenova_u1/sensenova_u1_transformer.py | 10 + .../stable_audio/stable_audio_transformer.py | 8 + 21 files changed, 508 insertions(+), 1332 deletions(-) diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py index 021473ca5..f5981b927 100644 --- a/tests/diffusion/cache/test_cache_backends.py +++ b/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: diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index 146084a50..0a05aeb76 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/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) diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py index bf6a75636..9f38f097a 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py +++ b/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: diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index ff1eea6ae..5add71daa 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -10,11 +10,19 @@ pipelines in vllm-omni, supporting both single and dual-transformer architecture import functools from collections.abc import Callable from contextlib import ExitStack -from typing import Any, Optional +from dataclasses import dataclass +from typing import Any, Optional, TypeAlias import cache_dit import torch -from cache_dit import BlockAdapter, DBCacheConfig, ForwardPattern, ParamsModifier, TaylorSeerCalibratorConfig +from cache_dit import ( + BlockAdapter, + CalibratorConfig, + DBCacheConfig, + ForwardPattern, + ParamsModifier, + TaylorSeerCalibratorConfig, +) from cache_dit.caching.block_adapters import FakeDiffusionPipeline from cache_dit.caching.cache_adapters.cache_adapter import CachedAdapter from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2 @@ -25,10 +33,29 @@ from vllm.logger import init_logger from vllm_omni.diffusion.cache.base import CacheBackend from vllm_omni.diffusion.data import DiffusionCacheConfig, OmniDiffusionConfig -from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline logger = init_logger(__name__) +RefreshCacheContextFunc: TypeAlias = Callable[[Any, int, bool], None] + + +@dataclass +class CacheDiTAdapterConfig: + """Config for creating a Cache DiT's block adapter; to enable CacheDiT, + most models just need to define an instance of this class as a class + var in the DiT. + """ + + block_forward_patterns: dict[str, ForwardPattern] + has_separate_cfg: bool = False + cached_adapter_cls: type[CachedAdapter] | None = None + check_forward_pattern: bool = True + + +# Registry of custom cache-dit enablers for specific models +# Maps pipeline names to their cache-dit enablement functions +CUSTOM_DIT_ENABLERS: dict[str, Callable] = {} + # Small helper to centralize cache-dit summaries. def cache_summary(pipeline: Any, details: bool = True) -> None: @@ -37,14 +64,60 @@ def cache_summary(pipeline: Any, details: bool = True) -> None: cache_dit.summary(pipeline.transformer_2, details=details) -# Registry of custom cache-dit enablers for specific models -# Maps pipeline names to their cache-dit enablement functions -# Models in this registry require custom handling (e.g., dual-transformer architectures) -# Will be populated after function definitions -CUSTOM_DIT_ENABLERS: dict[str, Callable] = {} +def default_get_pipeline_transformer(pipeline: Any) -> Any: + return pipeline.transformer -def _build_db_cache_config(cache_config: Any) -> DBCacheConfig: +def build_cache_context_refresh( + cache_config: DiffusionCacheConfig, + get_pipeline_transformer: Callable[[Any], Any] = default_get_pipeline_transformer, +) -> RefreshCacheContextFunc: + """Build the cache context refresh func for a single Transformer.""" + + def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + """Refresh cache context for the transformer with new num_inference_steps. + + Args: + pipeline: The pipeline instance. + num_inference_steps: New number of inference steps. + """ + transformer = get_pipeline_transformer(pipeline) + + # Bypass SCM for step counts that don't support predefined masks (e.g., vLLM's 1-step dummy run) + scm_supported_steps = num_inference_steps >= 8 or num_inference_steps in (4, 6) + + if cache_config.scm_steps_mask_policy is None or not scm_supported_steps: + cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) + else: + cache_dit.refresh_context( + transformer, + cache_config=DBCacheConfig().reset( + num_inference_steps=num_inference_steps, + steps_computation_mask=cache_dit.steps_mask( + mask_policy=cache_config.scm_steps_mask_policy, + total_steps=num_inference_steps, + ), + steps_computation_policy=cache_config.scm_steps_policy, + ), + verbose=verbose, + ) + + return refresh_cache_context + + +def _resolve_calibrator_config(cache_config: DiffusionCacheConfig) -> CalibratorConfig | None: + """Resolves the Calibrator subconfig For DiT cache.""" + calibrator = None + if cache_config.enable_taylorseer: + taylorseer_order = cache_config.taylorseer_order + calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) + logger.info(f"TaylorSeer enabled with order={taylorseer_order}") + # In the future, more calibrators will likely be added to DiT Cache, + # e.g., focal; handle them generically here. + return calibrator + + +def _build_db_cache_config(cache_config: DiffusionCacheConfig) -> DBCacheConfig: """Build DBCacheConfig with optional SCM (Step Computation Masking) support. Args: @@ -53,9 +126,8 @@ def _build_db_cache_config(cache_config: Any) -> DBCacheConfig: Returns: DBCacheConfig instance with SCM support if configured. """ - return DBCacheConfig( - # we will refresh the context when gets num_inference_steps in the first inference request + # we will refresh the context when we get num_inference_steps in the first inference request num_inference_steps=None, Fn_compute_blocks=cache_config.Fn_compute_blocks, Bn_compute_blocks=cache_config.Bn_compute_blocks, @@ -68,7 +140,100 @@ def _build_db_cache_config(cache_config: Any) -> DBCacheConfig: ) -def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int], None]: +def enable_cache_for_dit( + pipeline: Any, + cache_config: Any, + block_adapter: BlockAdapter | None = None, + adapter_cls: type[CachedAdapter] | None = None, +) -> RefreshCacheContextFunc: + """Enable cache-dit for regular single-transformer DiT models. + + Args: + pipeline: The diffusion pipeline instance. + cache_config: DiffusionCacheConfig instance with cache configuration. + block_adapter: Custom block adapters for specific model architectures. + adapter_cls: Custom cached adapter class for specific model architectures. + + Returns: + A refresh function that can be called to update cache context with new num_inference_steps. + """ + # Build DBCacheConfig with optional SCM support + db_cache_config = _build_db_cache_config(cache_config) + + # Build calibrator config if TaylorSeer is enabled + calibrator_config = _resolve_calibrator_config(cache_config) + + logger.info( + f"Enabling cache-dit on transformer: " + f"Fn={db_cache_config.Fn_compute_blocks}, " + f"Bn={db_cache_config.Bn_compute_blocks}, " + f"W={db_cache_config.max_warmup_steps}, " + ) + + # Enable cache-dit on the transformer + transformer = default_get_pipeline_transformer(pipeline) + + # If we have a custom cached adapter subclass, call apply directly + if adapter_cls is not None: + adapter_cls.apply( + transformer if block_adapter is None else block_adapter, + cache_config=db_cache_config, + calibrator_config=calibrator_config, + ) + else: + # Otherwise, call enable cache, which will call CachedAdapter.apply for us + cache_dit.enable_cache( + transformer if block_adapter is None else block_adapter, + cache_config=db_cache_config, + calibrator_config=calibrator_config, + ) + + return build_cache_context_refresh(cache_config) + + +### Complex / custom enablers for DiT cache +# NOTE (Alex): This case is rare; you should only really need to do this if you have a dual transformer +# architecture, since it hasn't been handled generically yet, or if the model class has unique attributes +# that it sets during Cache DiT enablement. +# +# For the vast majority of models, you should only need to add a _cache_dit_adapter_config attribute +# to the Transformer class, which controls the forward pattern, whether or not we have separate CFG, +# and so on. + + +# from https://github.com/vipshop/cache-dit/pull/542 +def _split_wan22_inference_steps(pipeline, num_inference_steps: int) -> tuple[int, int]: + """Split inference steps into high-noise and low-noise steps for Wan2.2. + + This is an internal helper function specific to Wan2.2's dual-transformer + architecture that uses boundary_ratio to determine the split point. + + Args: + num_inference_steps: Total number of inference steps. + + Returns: + A tuple of (num_high_noise_steps, num_low_noise_steps). + """ + if pipeline.boundary_ratio is not None: + boundary_timestep = pipeline.boundary_ratio * pipeline.scheduler.config.num_train_timesteps + else: + boundary_timestep = None + + # Set timesteps to calculate the split + device = next(pipeline.transformer.parameters()).device + pipeline.scheduler.set_timesteps(num_inference_steps, device=device) + + timesteps = pipeline.scheduler.timesteps + num_high_noise_steps = 0 # high-noise steps for transformer + for t in timesteps: + if boundary_timestep is None or t >= boundary_timestep: + num_high_noise_steps += 1 + # low-noise steps for transformer_2 + num_low_noise_steps = num_inference_steps - num_high_noise_steps + return num_high_noise_steps, num_low_noise_steps + + +def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc: """Enable cache-dit for Wan2.2 single or dual-transformer architecture. Wan2.2 can use single or dual transformers (transformer and transformer_2) that need @@ -81,10 +246,11 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int], Returns: A refresh function that can be called to update cache context with new num_inference_steps. """ + # Build DBCacheConfig with optional SCM support + db_cache_config = _build_db_cache_config(cache_config) if getattr(pipeline, "transformer_2", None) is None: logger.info("transformer_2 not found, enabling cache-dit for single transformer mode") - db_cache_config = _build_db_cache_config(cache_config) cache_dit.enable_cache( BlockAdapter( transformer=pipeline.transformer, @@ -97,27 +263,7 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int], ), cache_config=db_cache_config, ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for single transformer.""" - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context( - pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose - ) - else: - cache_dit.refresh_context( - pipeline.transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_inference_steps - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context + return build_cache_context_refresh(cache_config) cache_dit.enable_cache( BlockAdapter( @@ -150,47 +296,11 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int], ], has_separate_cfg=True, ), - cache_config=DBCacheConfig( - Fn_compute_blocks=cache_config.Fn_compute_blocks, - Bn_compute_blocks=cache_config.Bn_compute_blocks, - max_warmup_steps=cache_config.max_warmup_steps, - max_cached_steps=cache_config.max_cached_steps, - max_continuous_cached_steps=cache_config.max_continuous_cached_steps, - residual_diff_threshold=cache_config.residual_diff_threshold, - num_inference_steps=None, - ), + cache_config=db_cache_config, ) - # from https://github.com/vipshop/cache-dit/pull/542 - def _split_inference_steps(num_inference_steps: int) -> tuple[int, int]: - """Split inference steps into high-noise and low-noise steps for Wan2.2. - - This is an internal helper function specific to Wan2.2's dual-transformer - architecture that uses boundary_ratio to determine the split point. - - Args: - num_inference_steps: Total number of inference steps. - - Returns: - A tuple of (num_high_noise_steps, num_low_noise_steps). - """ - if pipeline.boundary_ratio is not None: - boundary_timestep = pipeline.boundary_ratio * pipeline.scheduler.config.num_train_timesteps - else: - boundary_timestep = None - - # Set timesteps to calculate the split - device = next(pipeline.transformer.parameters()).device - pipeline.scheduler.set_timesteps(num_inference_steps, device=device) - - timesteps = pipeline.scheduler.timesteps - num_high_noise_steps = 0 # high-noise steps for transformer - for t in timesteps: - if boundary_timestep is None or t >= boundary_timestep: - num_high_noise_steps += 1 - # low-noise steps for transformer_2 - num_low_noise_steps = num_inference_steps - num_high_noise_steps - return num_high_noise_steps, num_low_noise_steps + refresh_trans_one = build_cache_context_refresh(cache_config) + refresh_trans_two = build_cache_context_refresh(cache_config, lambda pipeline: pipeline.transformer_2) def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: """Refresh cache context for both transformers with new num_inference_steps. @@ -199,50 +309,14 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int], pipeline: The Wan2.2 pipeline instance. num_inference_steps: New number of inference steps. """ - - num_high_noise_steps, num_low_noise_steps = _split_inference_steps(num_inference_steps) - # Refresh context for high-noise transformer - if cache_config.scm_steps_mask_policy is None: - # cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_high_noise_steps, verbose=verbose) - cache_dit.refresh_context( - pipeline.transformer, - num_inference_steps=num_high_noise_steps, - verbose=verbose, - ) - cache_dit.refresh_context( - pipeline.transformer_2, - num_inference_steps=num_low_noise_steps, - verbose=verbose, - ) - else: - cache_dit.refresh_context( - pipeline.transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_high_noise_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_high_noise_steps - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - cache_dit.refresh_context( - pipeline.transformer_2, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_low_noise_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_low_noise_steps - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) + num_high_noise_steps, num_low_noise_steps = _split_wan22_inference_steps(pipeline, num_inference_steps) + refresh_trans_one(pipeline, num_high_noise_steps, verbose) + refresh_trans_two(pipeline, num_low_noise_steps, verbose) return refresh_cache_context -def enable_cache_for_wan22_s2v(pipeline: Any, cache_config: Any) -> Callable[[int], None]: +def enable_cache_for_wan22_s2v(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc: """Enable cache-dit for Wan2.2 S2V. S2V uses a single transformer, but unlike the other Wan2.2 variants its @@ -316,557 +390,6 @@ def enable_cache_for_wan22_s2v(pipeline: Any, cache_config: Any) -> Callable[[in return refresh_cache_context -def enable_cache_for_longcat_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for LongCatImage pipeline. - - Args: - pipeline: The LongCatImage pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - """ - # Build DBCacheConfig for transformer - db_cache_config = _build_db_cache_config(cache_config) - - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - # Build ParamsModifier for transformer - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) - - logger.info( - f"Enabling cache-dit on LongCatImage transformer with BlockAdapter: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - # Enable cache-dit using BlockAdapter for transformer - cache_dit.enable_cache( - ( - BlockAdapter( - transformer=pipeline.transformer, - blocks=[ - pipeline.transformer.transformer_blocks, - pipeline.transformer.single_transformer_blocks, - ], - forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1], - params_modifiers=[modifier], - has_separate_cfg=True, - ) - ), - cache_config=db_cache_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. - - Args: - pipeline: The LongCatImage pipeline instance. - num_inference_steps: New number of inference steps. - """ - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - pipeline.transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_flux(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for Flux.1-dev pipeline. - - Args: - pipeline: The Flux pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - Returns: - A refresh function that can be called with a new ``num_inference_steps`` - to update the cache context for the pipeline. - """ - # Build DBCacheConfig for transformer - db_cache_config = _build_db_cache_config(cache_config) - - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - # Build ParamsModifier for transformer - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) - - logger.info( - f"Enabling cache-dit on Flux transformer with BlockAdapter: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - # Enable cache-dit using BlockAdapter for transformer - cache_dit.enable_cache( - ( - BlockAdapter( - transformer=pipeline.transformer, - blocks=[ - pipeline.transformer.transformer_blocks, - pipeline.transformer.single_transformer_blocks, - ], - forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1], - params_modifiers=[modifier], - ) - ), - cache_config=db_cache_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. - - Args: - pipeline: The Flux pipeline instance. - num_inference_steps: New number of inference steps. - """ - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - pipeline.transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_flux2_klein(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for FLUX.2-klein-4B pipeline. - - Args: - pipeline: The FLUX.2-klein-4B pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - Returns: - A refresh function that can be called with a new ``num_inference_steps`` - to update the cache context for the pipeline. - """ - # Build DBCacheConfig for transformer - db_cache_config = _build_db_cache_config(cache_config) - # The Fn_compute_blocks = 2 override is the most important decision here, - # and the rationale (quality degradation at Fn=1) only lives in flux2_klein. - db_cache_config.Fn_compute_blocks = 2 - - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - # Build ParamsModifier for transformer - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) - - logger.info( - f"Enabling cache-dit on Flux2-Klein transformer with BlockAdapter: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - # Enable cache-dit using BlockAdapter for transformer - cache_dit.enable_cache( - ( - BlockAdapter( - transformer=pipeline.transformer, - blocks=[ - pipeline.transformer.transformer_blocks, - pipeline.transformer.single_transformer_blocks, - ], - forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_2], - params_modifiers=[modifier], - ) - ), - cache_config=db_cache_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. - - Args: - pipeline: The FLUX.2-klein-4B pipeline instance. - num_inference_steps: New number of inference steps. - """ - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - pipeline.transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - Fn_compute_blocks=db_cache_config.Fn_compute_blocks, - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_stable_audio_open(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for Stable Audio Open pipeline. - - Args: - pipeline: The StableAudioPipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - - Returns: - A refresh function that can be called to update cache context with new num_inference_steps. - """ - db_cache_config = _build_db_cache_config(cache_config) - - calibrator_config = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - # StableAudio is officially registered in CacheDiT as Pattern_3: - # https://github.com/vipshop/cache-dit/blob/69e82bd1/src/cache_dit/caching/block_adapters/__init__.py#L562 - # - # Pattern_3 is required because StableAudioDiT uses cross-attention - # with static encoder_hidden_states that do not change inside the - # transformer block loop. - cache_dit.enable_cache( - BlockAdapter( - transformer=pipeline.transformer, - blocks=pipeline.transformer.transformer_blocks, - forward_pattern=ForwardPattern.Pattern_3, - params_modifiers=[ - ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - ], - ), - cache_config=db_cache_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. - - Args: - pipeline: The StableAudioPipeline instance. - num_inference_steps: New number of inference steps. - verbose: Whether to log refresh operations. - """ - # Bypass SCM for step counts that don't support predefined masks (e.g., vLLM's 1-step dummy run) - scm_supported_steps = num_inference_steps >= 8 or num_inference_steps in (4, 6) - - if cache_config.scm_steps_mask_policy is None or not scm_supported_steps: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - updated_scm_config = DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ) - - cache_dit.refresh_context( - pipeline.transformer, - cache_config=updated_scm_config, - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_sd3(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for StableDiffusion3Pipeline. - - Args: - pipeline: The StableDiffusion3 pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - """ - # Build DBCacheConfig for transformer - db_cache_config = _build_db_cache_config(cache_config) - - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - # Build ParamsModifier for transformer - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) - - logger.info( - f"Enabling cache-dit on StableDiffusion3 transformer with BlockAdapter: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - # Enable cache-dit using BlockAdapter for transformer - cache_dit.enable_cache( - ( - BlockAdapter( - transformer=pipeline.transformer, - blocks=pipeline.transformer.transformer_blocks, - forward_pattern=ForwardPattern.Pattern_1, - params_modifiers=[modifier], - ) - ), - cache_config=db_cache_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. - - Args: - pipeline: The LongCatImage pipeline instance. - num_inference_steps: New number of inference steps. - """ - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - pipeline.transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_ltx2(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for LTX2 pipelines (audio-video transformer blocks).""" - transformer = get_transformer_from_pipeline(pipeline) - - db_cache_config = _build_db_cache_config(cache_config) - - calibrator_config = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - blocks = transformer.transformer_blocks - - logger.info( - f"Enabling cache-dit on LTX2 transformer: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - cache_dit.enable_cache( - BlockAdapter( - transformer=transformer, - blocks=blocks, - # LTX2 blocks return (hidden_states, audio_hidden_states) - forward_pattern=ForwardPattern.Pattern_0, - # Treat audio_hidden_states as encoder_hidden_states in Pattern_0 - check_forward_pattern=False, - has_separate_cfg=True, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - transformer = get_transformer_from_pipeline(pipeline) - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for regular single-transformer DiT models. - - Args: - pipeline: The diffusion pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - - Returns: - A refresh function that can be called to update cache context with new num_inference_steps. - """ - # Build DBCacheConfig with optional SCM support - db_cache_config = _build_db_cache_config(cache_config) - - # Build calibrator config if TaylorSeer is enabled - calibrator_config = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - logger.info( - f"Enabling cache-dit on transformer: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - # Enable cache-dit on the transformer - transformer = get_transformer_from_pipeline(pipeline) - cache_dit.enable_cache( - transformer, - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. - - Args: - pipeline: The diffusion pipeline instance. - num_inference_steps: New number of inference steps. - """ - transformer = get_transformer_from_pipeline(pipeline) - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_hunyuan_image3(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for HunyuanImage3 pipeline. - - HunyuanImage3 stores its main transformer stack at ``pipeline.model`` with - decoder blocks in ``pipeline.model.layers``. - """ - db_cache_config = _build_db_cache_config(cache_config) - - calibrator_config = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - transformer = getattr(pipeline, "model", None) - if transformer is None or not hasattr(transformer, "layers"): - raise ValueError( - f"HunyuanImage3 cache-dit enabler expects pipeline.model.layers, got pipeline={pipeline.__class__.__name__}" - ) - - logger.info( - f"Enabling cache-dit on HunyuanImage3 model: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - # HunyuanImage3 decoder layers are single-stream `hidden_states` blocks - # (`HunyuanImage3DecoderLayer.forward(hidden_states, ...)`) while caller - # still reads outputs as tuple (`layer_outputs[0]`). Pattern_4 matches this: - # hidden-only input with hidden-first tuple output. - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - cache_dit.enable_cache( - BlockAdapter( - transformer=transformer, - blocks=transformer.layers, - forward_pattern=ForwardPattern.Pattern_4, - params_modifiers=[modifier], - check_forward_pattern=False, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - transformer = getattr(pipeline, "model", None) - if transformer is None: - raise ValueError("HunyuanImage3 cache-dit refresh expects pipeline.model to exist.") - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - class BagelCachedContextManager(CachedContextManager): """ Custom CachedContextManager for Bagel that safely handles NaiveCache objects @@ -1323,73 +846,6 @@ class BagelCachedAdapter(CachedAdapter): return total_cached_blocks -def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for Bagel model (via OmniDiffusion pipeline). - - Args: - pipeline: The OmniDiffusion pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - - Returns: - A refresh function that can be called to update cache context with new num_inference_steps. - """ - # Build DBCacheConfig - db_cache_config = _build_db_cache_config(cache_config) - - # Build calibrator config if TaylorSeer is enabled - calibrator_config = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - # Access the transformer: BagelPipeline -> Qwen2MoTForCausalLM -> Qwen2MoTModel - # BagelPipeline has self.language_model which is Qwen2MoTForCausalLM - # Qwen2MoTForCausalLM has self.model which is Qwen2MoTModel - transformer = pipeline.language_model.model - - logger.info( - f"Enabling cache-dit on Bagel transformer: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - # Enable cache-dit on the transformer - # Pattern_0 corresponds to (hidden_states, encoder_hidden_states) input, output - # Custom adapter for Bagel to handle NaiveCache correctly - # from vllm_omni.diffusion.cache.bagel_cache_adapter import BagelCachedAdapter # No longer needed - BagelCachedAdapter.apply( - BlockAdapter( - transformer=transformer, - blocks=transformer.layers, - forward_pattern=ForwardPattern.Pattern_0, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - transformer = pipeline.language_model.model - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - class SensenovaCachedBlocks(CachedBlocks_Pattern_3_4_5): """ Custom CachedBlocks for SenseNova-U1 that only caches image-token hidden @@ -1477,474 +933,7 @@ class SensenovaCachedAdapter(CachedAdapter): return total_cached_blocks -def enable_cache_for_sensenova_u1(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for SenseNova-U1 model. - - Args: - pipeline: The SenseNova-U1 pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - - Returns: - A refresh function that can be called to update cache context with new num_inference_steps. - """ - transformer = getattr(getattr(pipeline, "language_model", None), "model", None) - if transformer is None or not hasattr(transformer, "layers"): - raise ValueError("SenseNovaU1Pipeline cache-dit expects pipeline.language_model.model.layers.") - - # Build DBCacheConfig - db_cache_config = _build_db_cache_config(cache_config) - - calibrator_config = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - logger.info( - f"Enabling cache-dit on SenseNova-U1 decoder layers: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - SensenovaCachedAdapter.apply( - BlockAdapter( - transformer=transformer, - blocks=transformer.layers, - forward_pattern=ForwardPattern.Pattern_3, - params_modifiers=[modifier], - check_forward_pattern=True, - has_separate_cfg=True, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - transformer = pipeline.language_model.model - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_flux2(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for Flux.2-dev pipeline. - - Args: - pipeline: The Flux2 pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - Returns: - A refresh function that can be called with a new ``num_inference_steps`` - to update the cache context for the pipeline. - """ - # Build DBCacheConfig for transformer - db_cache_config = _build_db_cache_config(cache_config) - - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - # Build ParamsModifier for transformer - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) - - logger.info( - f"Enabling cache-dit on Flux transformer with BlockAdapter: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - # Enable cache-dit using BlockAdapter for transformer - cache_dit.enable_cache( - ( - BlockAdapter( - transformer=pipeline.transformer, - blocks=[ - pipeline.transformer.transformer_blocks, - pipeline.transformer.single_transformer_blocks, - ], - forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_2], - params_modifiers=[modifier], - ) - ), - cache_config=db_cache_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. - - Args: - pipeline: The Flux2 pipeline instance. - num_inference_steps: New number of inference steps. - """ - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - pipeline.transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_glm_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for GlmImage pipeline. - - Args: - pipeline: The GlmImage pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - Returns: - A refresh function that can be called with a new ``num_inference_steps`` - to update the cache context for the pipeline. - """ - # Build DBCacheConfig for transformer - db_cache_config = _build_db_cache_config(cache_config) - - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - # Build ParamsModifier for transformer - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) - - logger.info( - f"Enabling cache-dit on GlmImage transformer with BlockAdapter: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - # Enable cache-dit using BlockAdapter for transformer - # Note: We don't use patch_functor here because it's designed for diffusers' GlmImage, - # and our vllm-omni implementation has a different forward signature. - # We use ForwardPattern.Pattern_0 because our block returns (hidden_states, encoder_hidden_states) - cache_dit.enable_cache( - ( - BlockAdapter( - transformer=pipeline.transformer, - blocks=pipeline.transformer.transformer_blocks, - forward_pattern=ForwardPattern.Pattern_0, - params_modifiers=[modifier], - patch_functor=None, - has_separate_cfg=True, - ) - ), - cache_config=db_cache_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. - - Args: - pipeline: The GlmImage pipeline instance. - num_inference_steps: New number of inference steps. - """ - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - pipeline.transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_dreamid_omni(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for DreamID-Omni fused pipeline. - - Args: - pipeline: The DreamIDOmni pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - - Returns: - A refresh function that can be called with a new ``num_inference_steps`` - to update the cache context for the pipeline. - """ - transformer = getattr(pipeline, "transformer", None) - fused_blocks = getattr(transformer, "fused_blocks", None) - if transformer is None or fused_blocks is None: - raise ValueError("DreamIDOmniPipeline cache-dit expects pipeline.transformer.fused_blocks.") - - db_cache_config = _build_db_cache_config(cache_config) - - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) - - logger.info( - f"Enabling cache-dit on DreamID-Omni fused blocks: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - cache_dit.enable_cache( - BlockAdapter( - transformer=transformer, - blocks=fused_blocks, - forward_pattern=ForwardPattern.Pattern_0, - params_modifiers=[modifier], - has_separate_cfg=True, - ), - cache_config=db_cache_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - transformer = pipeline.transformer - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_ernie_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for ERNIE-Image pipeline. - - ERNIE-Image blocks have signature: - forward(x, rotary_pos_emb, temb, attention_mask) -> x - - Where x is hidden_states (concatenated image + text tokens). - This matches Pattern_3 which expects: - - Input: hidden_states only - - Output: hidden_states only - - Args: - pipeline: The ERNIE-Image pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - Returns: - A refresh function that can be called to update cache context with new num_inference_steps. - """ - db_cache_config = _build_db_cache_config(cache_config) - - calibrator_config = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - logger.info( - f"Enabling cache-dit on ERNIE-Image transformer: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - cache_dit.enable_cache( - BlockAdapter( - transformer=pipeline.transformer, - blocks=pipeline.transformer.layers, - forward_pattern=ForwardPattern.Pattern_3, - check_forward_pattern=False, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - pipeline.transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_helios(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for Helios pipeline. - - Helios extends Wan2.2 with multi-term memory patches and guidance cross-attention. - Its transformer blocks have the same single-output signature as Wan2.2 blocks - (returns hidden_states only), so we use ForwardPattern.Pattern_2. - - Args: - pipeline: The HeliosPipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - - Returns: - A refresh function that can be called to update cache context with new num_inference_steps. - """ - db_cache_config = _build_db_cache_config(cache_config) - - calibrator_config = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - cache_dit.enable_cache( - BlockAdapter( - transformer=pipeline.transformer, - blocks=[pipeline.transformer.blocks], - forward_pattern=[ForwardPattern.Pattern_2], - params_modifiers=[ - ParamsModifier(cache_config=db_cache_config, calibrator_config=calibrator_config), - ], - has_separate_cfg=True, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - logger.info( - f"Enabling cache-dit on Helios transformer: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - pipeline.transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_inference_steps - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_hunyuan_video_15(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for HunyuanVideo 1.5 pipeline. - - HunyuanVideo 1.5 uses a single transformer with has_separate_cfg=True - (separate conditional/unconditional forward passes). The _sp_plan scatter - is applied at the transformer input boundary (empty-string key), so CacheDiT - sees already-sharded hidden_states throughout transformer_blocks. - """ - db_cache_config = _build_db_cache_config(cache_config) - - calibrator_config = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - logger.info( - f"Enabling cache-dit on HunyuanVideo15 transformer: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - cache_dit.enable_cache( - BlockAdapter( - transformer=pipeline.transformer, - blocks=pipeline.transformer.transformer_blocks, - forward_pattern=ForwardPattern.Pattern_0, - check_forward_pattern=True, - has_separate_cfg=True, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - pipeline.transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> Callable[[int], None]: +def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc: """Enable cache-dit for Cosmos3. Cosmos3 has a dual-pathway architecture (UND + GEN) but only the GEN @@ -1959,41 +948,6 @@ def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> Callable[[int] Returns: A refresh function that can be called to update cache context with new num_inference_steps. """ - db_cache_config = _build_db_cache_config(cache_config) - - calibrator_config = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - logger.info( - f"Enabling cache-dit on Cosmos3 gen_layers: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - cache_dit.enable_cache( - BlockAdapter( - transformer=pipeline.transformer, - blocks=[pipeline.transformer.gen_layers], - # Cosmos3 GEN blocks return only hidden_states. Per-layer UND K/V - # conditioning uses the transformer's cache-dit fallback path. - forward_pattern=[ForwardPattern.Pattern_3], - params_modifiers=[ - ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ), - ], - check_forward_pattern=False, - has_separate_cfg=True, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - # The T2I denoising loop skips the unconditional forward outside the # guidance interval as a speed optimization. cache-dit distinguishes the # conditional vs unconditional passes purely by transformer-forward parity @@ -2001,25 +955,8 @@ def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> Callable[[int] # step accounting. Still do both cond/uncond CFG steps when cache-dit is active. # CFG is instead neutralized via scale=1.0 outside the interval. pipeline._cache_dit_requires_paired_cfg = True - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - pipeline.transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_inference_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, - total_steps=num_inference_steps, - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - return refresh_cache_context + block_adapter = CacheDiTBackend.maybe_build_block_adapter(pipeline) + return enable_cache_for_dit(pipeline, cache_config, block_adapter) # Register custom cache-dit enablers after function definitions @@ -2027,29 +964,8 @@ CUSTOM_DIT_ENABLERS.update( { "Wan22Pipeline": enable_cache_for_wan22, "Wan22I2VPipeline": enable_cache_for_wan22, + "Wan22TI2VPipeline": enable_cache_for_wan22, "Wan22S2VPipeline": enable_cache_for_wan22_s2v, - "HunyuanImage3Pipeline": enable_cache_for_hunyuan_image3, - "FluxPipeline": enable_cache_for_flux, - "Flux2KleinPipeline": enable_cache_for_flux2_klein, - "LongCatImagePipeline": enable_cache_for_longcat_image, - "LongCatImageEditPipeline": enable_cache_for_longcat_image, - "StableAudioPipeline": enable_cache_for_stable_audio_open, - "StableDiffusion3Pipeline": enable_cache_for_sd3, - "LTX2Pipeline": enable_cache_for_ltx2, - "LTX2ImageToVideoPipeline": enable_cache_for_ltx2, - "LTX2TwoStagesPipeline": enable_cache_for_ltx2, - "LTX2ImageToVideoTwoStagesPipeline": enable_cache_for_ltx2, - "LTX23Pipeline": enable_cache_for_ltx2, - "LTX23ImageToVideoPipeline": enable_cache_for_ltx2, - "BagelPipeline": enable_cache_for_bagel, - "SenseNovaU1Pipeline": enable_cache_for_sensenova_u1, - "GlmImagePipeline": enable_cache_for_glm_image, - "Flux2Pipeline": enable_cache_for_flux2, - "DreamIDOmniPipeline": enable_cache_for_dreamid_omni, - "ErnieImagePipeline": enable_cache_for_ernie_image, - "HunyuanVideo15Pipeline": enable_cache_for_hunyuan_video_15, - "HunyuanVideo15I2VPipeline": enable_cache_for_hunyuan_video_15, - "HeliosPipeline": enable_cache_for_helios, "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3, } ) @@ -2092,6 +1008,43 @@ class CacheDiTBackend(CacheBackend): self._refresh_func: Callable[[Any, int, bool], None] | None = None self._last_num_inference_steps: int | None = None + @staticmethod + def maybe_build_block_adapter(pipeline) -> BlockAdapter | None: + """If a module defines `_cache_dit_adapter_config`, build the corresponding + block adapter. + """ + transformer = default_get_pipeline_transformer(pipeline) + + adapter_cfg: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None) + if adapter_cfg is None: + return None + + block_attrs, forward_pattern = zip(*(adapter_cfg.block_forward_patterns).items()) + missing_attrs = [block_attr for block_attr in block_attrs if not hasattr(transformer, block_attr)] + + if missing_attrs: + logger.warning("Missing Cache DiT block attributes: %s", missing_attrs) + + block_adapter = BlockAdapter( + transformer=transformer, + blocks=[getattr(transformer, block_attr) for block_attr in block_attrs], + forward_pattern=list(forward_pattern), + has_separate_cfg=adapter_cfg.has_separate_cfg, + check_forward_pattern=adapter_cfg.check_forward_pattern, + ) + return block_adapter + + @staticmethod + def maybe_get_cached_adapter_cls(pipeline) -> type[CachedAdapter] | None: + """If a module has a custom cached adapter type registered, e.g., SenseNova, retrieve it + from the transformer's CacheDiTAdapterConfig.""" + transformer = default_get_pipeline_transformer(pipeline) + + adapter_cfg: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None) + if adapter_cfg is None: + return None + return adapter_cfg.cached_adapter_cls + def enable(self, pipeline: Any) -> None: """Enable cache-dit on the pipeline if configured. @@ -2110,8 +1063,12 @@ class CacheDiTBackend(CacheBackend): logger.info(f"Using custom cache-dit enabler for model: {pipeline_name}") self._refresh_func = CUSTOM_DIT_ENABLERS[pipeline_name](pipeline, self.config) else: - # For regular single-transformer models - self._refresh_func = enable_cache_for_dit(pipeline, self.config) + # Common case; either the model doesn't explicitly support dit cache yet, + # Or it defines its _cache_dit_adapter_config, which describes how we should + # create its block adapter. + block_adapter = self.maybe_build_block_adapter(pipeline) + adapter_cls = self.maybe_get_cached_adapter_cls(pipeline) + self._refresh_func = enable_cache_for_dit(pipeline, self.config, block_adapter, adapter_cls) self.enabled = True logger.info(f"Cache-dit enabled successfully on {pipeline_name}") diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index bc3eef49d..e68025967 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/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 diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 590866846..4543d6922 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/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"] diff --git a/vllm_omni/diffusion/models/dreamid_omni/fusion.py b/vllm_omni/diffusion/models/dreamid_omni/fusion.py index 423935427..acf24dd2f 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/fusion.py +++ b/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 = { diff --git a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py index aa9b0d8b7..8398cb7ab 100644 --- a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py +++ b/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"] diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py index 88d8126c8..9d89368be 100644 --- a/vllm_omni/diffusion/models/flux/flux_transformer.py +++ b/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 diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py index 0c6896c13..5a65db257 100644 --- a/vllm_omni/diffusion/models/flux2/flux2_transformer.py +++ b/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 = { "": { diff --git a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py index a708a8ece..cf2c9f78a 100644 --- a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py +++ b/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"] diff --git a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py index 56a7066e3..c933d4e7f 100644 --- a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py +++ b/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 diff --git a/vllm_omni/diffusion/models/helios/helios_transformer.py b/vllm_omni/diffusion/models/helios/helios_transformer.py index ed0a50fee..730c10f1c 100644 --- a/vllm_omni/diffusion/models/helios/helios_transformer.py +++ b/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 = { diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index 5ce3098ad..e4c9e0b2a 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/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": { diff --git a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py index 9af9e8488..9224973dd 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py +++ b/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 = { diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py index 8f0ff446a..744c1806a 100644 --- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py +++ b/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"] diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index dcf848390..6072c0456 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/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"] diff --git a/vllm_omni/diffusion/models/sd3/sd3_transformer.py b/vllm_omni/diffusion/models/sd3/sd3_transformer.py index 89f061577..d8e037051 100644 --- a/vllm_omni/diffusion/models/sd3/sd3_transformer.py +++ b/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"] diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py index 44446ae83..c1ba258b2 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/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) diff --git a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py index f9ff8aaf4..53e1a47c2 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py +++ b/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 diff --git a/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py b/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py index 83a669d92..c7bc9b2d4 100644 --- a/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py +++ b/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]