# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Model specific tests for CacheDiT enablement. """ import ast from pathlib import Path from unittest.mock import Mock, patch import pytest import torch from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2 from vllm.distributed import parallel_state import vllm_omni.diffusion.cache.cachedit as cd_backend import vllm_omni.diffusion.cache.cachedit.model_specific as cd_model_specific from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig, CacheDiTBackend, cache_summary from vllm_omni.diffusion.data import DiffusionCacheConfig from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel 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 pytestmark = [pytest.mark.core_model, pytest.mark.cpu] SEPARATE_CFG_TRANSFORMERS = [ HeliosTransformer3DModel, LongCatImageTransformer2DModel, Cosmos3VFMTransformer, ] SAMPLE_CACHE_CONFIG = DiffusionCacheConfig() def test_custom_cache_dit_enablers_are_registered_explicitly(): expected_enablers = { "Wan22Pipeline": cd_model_specific.enable_cache_for_wan22, "Wan22I2VPipeline": cd_model_specific.enable_cache_for_wan22, "Wan22TI2VPipeline": cd_model_specific.enable_cache_for_wan22, "Wan22VACEPipeline": cd_model_specific.enable_cache_for_wan22, "Wan22S2VPipeline": cd_model_specific.enable_cache_for_wan22_s2v, "Cosmos3OmniDiffusersPipeline": cd_model_specific.enable_cache_for_cosmos3, "Cosmos3OmniPipeline": cd_model_specific.enable_cache_for_cosmos3, "Krea2Pipeline": cd_model_specific.enable_cache_for_krea2, } with patch.dict(cd_backend.CUSTOM_DIT_ENABLERS, {}, clear=True): cd_model_specific.register_custom_dit_enablers() assert cd_backend.CUSTOM_DIT_ENABLERS == expected_enablers @pytest.fixture() def init_fake_tp_group(mocker): """Provide a fake TP group so vLLM linear layers can be instantiated.""" mock_tp = mocker.MagicMock() mock_tp.world_size = 1 mock_tp.rank_in_group = 0 old = parallel_state._TP parallel_state._TP = mock_tp yield parallel_state._TP = old def test_wan22_vace_uses_wan22_custom_cache_dit_enabler(): assert cd_backend.CUSTOM_DIT_ENABLERS["Wan22VACEPipeline"] is cd_model_specific.enable_cache_for_wan22 @pytest.mark.parametrize( "pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"], ) def test_cosmos3_aliases_use_cosmos3_custom_cache_dit_enabler(pipeline_name: str): assert cd_backend.CUSTOM_DIT_ENABLERS[pipeline_name] is cd_model_specific.enable_cache_for_cosmos3 def test_cachedit_public_api_is_explicit(): assert set(cd_backend.__all__) == { "BagelCachedAdapter", "CUSTOM_DIT_ENABLERS", "CacheDiTAdapterConfig", "CacheDiTBackend", "CacheDiTConfig", "CacheDiTRequestSpec", "RequestScopedCacheDiTRuntime", "SensenovaCachedAdapter", "cache_summary", "enable_cache_for_dit", } assert not hasattr(cd_backend, "enable_cache_for_wan22") assert not hasattr(cd_backend, "enable_cache_for_wan22_s2v") def test_cachedit_consumers_use_package_api(): cache_dir = Path(cd_backend.__file__).resolve().parents[1] package_root = cache_dir.parents[1] legacy_module = "vllm_omni.diffusion.cache.cache_dit_backend" internal_prefix = "vllm_omni.diffusion.cache.cachedit." invalid_imports = [] for source_path in package_root.rglob("*.py"): if source_path.is_relative_to(cache_dir): continue tree = ast.parse(source_path.read_text(encoding="utf-8")) for node in ast.walk(tree): modules = [] if isinstance(node, ast.ImportFrom) and node.module is not None: modules.append(node.module) elif isinstance(node, ast.Import): modules.extend(alias.name for alias in node.names) for module in modules: if module == legacy_module or module.startswith(internal_prefix): invalid_imports.append( f"{source_path.relative_to(package_root)}:{getattr(node, 'lineno', '?')}: {module}" ) assert not invalid_imports, "Cache-DiT consumers must use the package API:\n" + "\n".join(invalid_imports) @pytest.mark.parametrize("transformer_model", SEPARATE_CFG_TRANSFORMERS) 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 def test_ltx2_cache_dit_uses_one_forward_per_denoise_step(): """LTX batches all guidance passes into one Transformer invocation.""" adapter_config = LTX2VideoTransformer3DModel._cache_dit_adapter_config assert adapter_config.has_separate_cfg is False @patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter") @patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit") def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter): """Ensure that Wan22, which has a custom enabler, setts custom CFG correctly.""" mock_pipeline = Mock() cd_model_specific.enable_cache_for_wan22(mock_pipeline, SAMPLE_CACHE_CONFIG) mock_cache_dit.enable_cache.assert_called_once() adapter_kwargs = mock_block_adapter.call_args.kwargs assert adapter_kwargs["has_separate_cfg"] is True @pytest.mark.parametrize("has_transformer_2", [False, True]) @patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter") @patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit") def test_wan22_custom_enabler_passes_taylorseer_calibrator( mock_cache_dit, mock_block_adapter, has_transformer_2, ): mock_pipeline = Mock() mock_pipeline.transformer.blocks = [Mock()] if has_transformer_2: mock_pipeline.transformer_2.blocks = [Mock()] else: mock_pipeline.transformer_2 = None cache_config = DiffusionCacheConfig(enable_taylorseer=True, taylorseer_order=1) cd_model_specific.enable_cache_for_wan22(mock_pipeline, cache_config) enable_cache_kwargs = mock_cache_dit.enable_cache.call_args.kwargs calibrator_config = enable_cache_kwargs["calibrator_config"] assert calibrator_config is not None assert calibrator_config.taylorseer_order == 1 adapter_kwargs = mock_block_adapter.call_args.kwargs for modifier in adapter_kwargs["params_modifiers"]: assert modifier._context_kwargs["calibrator_config"] is calibrator_config @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter") @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit") def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter): """Cosmos3 should cache only the repeated GEN pathway blocks.""" mock_pipeline = Mock() gen_layers = object() mock_pipeline.transformer.gen_layers = gen_layers mock_pipeline.transformer._cache_dit_adapter_config = Cosmos3VFMTransformer._cache_dit_adapter_config cd_model_specific.enable_cache_for_cosmos3(mock_pipeline, SAMPLE_CACHE_CONFIG) mock_cache_dit.enable_cache.assert_called_once() adapter_kwargs = mock_block_adapter.call_args.kwargs assert adapter_kwargs["transformer"] is mock_pipeline.transformer assert adapter_kwargs["blocks"] == [gen_layers] assert adapter_kwargs["has_separate_cfg"] is True assert adapter_kwargs["check_forward_pattern"] is False # This test is skipped on ROCm since rocm_unquantized_gemm doesn't support CPU backend @pytest.mark.skipif( current_omni_platform.is_rocm(), reason="vLLM ROCm custom ops lack CPU fallback", ) def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group): """CacheDiT Pattern_0 treats the second positional arg as encoder_hidden_states, which is a collision for one of the kwargs in LTX2 since we treat the audio hidden states as encoder_hidden_states. This test ensures that a tiny LTX2 transformer can be initialized and run through the cache DiT backend without hitting a collision on the kwargs. """ seq_len = 4 video_in = torch.full((1, seq_len, 16), 1.0) audio_in = torch.full((1, seq_len, 16), 2.0) text_in = torch.full((1, seq_len, 16), 3.0) audio_text_in = torch.full((1, seq_len, 16), 4.0) model = LTX2VideoTransformer3DModel( in_channels=16, out_channels=16, patch_size=1, patch_size_t=1, num_attention_heads=2, attention_head_dim=8, cross_attention_dim=16, audio_in_channels=16, audio_out_channels=16, audio_num_attention_heads=2, audio_attention_head_dim=8, audio_cross_attention_dim=16, num_layers=2, caption_channels=16, ) # NOTE: This is currently using the LTX2 custom enabler, but the custom # enablers will be consolidated after # https://github.com/vllm-project/vllm-omni/pull/2527 lands. LTX2Pipeline = type("LTX2Pipeline", (), {}) pipeline = LTX2Pipeline() pipeline.transformer = model backend = CacheDiTBackend(DiffusionCacheConfig()) backend.enable(pipeline) backend.refresh(pipeline, num_inference_steps=5) # Wrap call_Fn_blocks in CacheDiT so that we can verify the # hidden/encoder states are what we expect them to be captured = {} original = CachedBlocks_Pattern_0_1_2.call_Fn_blocks def call_Fn_blocks_and_capture(self, hidden_states, encoder_hidden_states, *a, **kw): captured["hidden_states"] = hidden_states captured["encoder_hidden_states"] = encoder_hidden_states return original(self, hidden_states, encoder_hidden_states, *a, **kw) # Also, map projections to identity so that we can just check # the captured tensors directly instead of having to reproject identity = torch.nn.Identity() with ( patch.object(model, "proj_in", identity), patch.object(model, "audio_proj_in", identity), patch.object(CachedBlocks_Pattern_0_1_2, "call_Fn_blocks", call_Fn_blocks_and_capture), torch.no_grad(), ): model( hidden_states=video_in, audio_hidden_states=audio_in, encoder_hidden_states=text_in, audio_encoder_hidden_states=audio_text_in, timestep=torch.tensor([[1000.0] * seq_len]), num_frames=1, height=2, width=2, audio_num_frames=seq_len, ) # Pattern_0 maps (hidden_states, encoder_hidden_states) to (video, audio) assert torch.equal(captured["hidden_states"], video_in) assert torch.equal(captured["encoder_hidden_states"], audio_in) def test_summary_with_no_transformer_is_nonfatal(): """Regression test for https://github.com/vllm-project/vllm-omni/issues/4325.""" class FakePipeline: pass cache_summary(pipeline=FakePipeline())