diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index 40a564485..0fd07393d 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/tests/diffusion/cache/test_cache_dit.py @@ -8,9 +8,14 @@ Model specific tests for CacheDiT enablement. from unittest.mock import Mock, patch import pytest +import torch +from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2 import vllm_omni.diffusion.cache.cache_dit_backend as cd_backend +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend from vllm_omni.diffusion.data import DiffusionCacheConfig +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] @@ -40,3 +45,85 @@ def test_separate_cfg(mock_cache_dit, mock_block_adapter, enabler): mock_cache_dit.enable_cache.assert_called_once() adapter_kwargs = mock_block_adapter.call_args.kwargs assert adapter_kwargs["has_separate_cfg"] is True + + +# 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) diff --git a/tests/helpers/fixtures/env.py b/tests/helpers/fixtures/env.py index 452ae2d2d..2d38e4907 100644 --- a/tests/helpers/fixtures/env.py +++ b/tests/helpers/fixtures/env.py @@ -2,6 +2,7 @@ import os import pytest import torch +from vllm.distributed import parallel_state @pytest.fixture(scope="session", autouse=True) @@ -59,3 +60,15 @@ def default_vllm_config(): with set_current_vllm_config(VllmConfig(device_config=device_config)): yield + + +@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 diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index 2a65ad6b0..64bb67e22 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py @@ -1939,10 +1939,6 @@ class LTX2VideoTransformer3DModel(nn.Module): # 5. Run transformer blocks for block in self.transformer_blocks: block_kwargs = { - "hidden_states": hidden_states, - "audio_hidden_states": audio_hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "audio_encoder_hidden_states": audio_encoder_hidden_states, "temb": temb, "temb_audio": temb_audio, "temb_ca_scale_shift": video_cross_attn_scale_shift, @@ -1958,11 +1954,27 @@ class LTX2VideoTransformer3DModel(nn.Module): "encoder_attention_mask": encoder_attention_mask, "audio_encoder_attention_mask": audio_encoder_attention_mask, } + # NOTE: We intentionally pass the first 4 args here positionally, + # because CacheDiT expects the first 2 args to be hidden_states + # and encoder_hidden_states, so passing them as kwargs will cause + # positional / keywords arg collisions. if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, audio_hidden_states = self._gradient_checkpointing_func(block, **block_kwargs) + hidden_states, audio_hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + audio_hidden_states, + encoder_hidden_states, + audio_encoder_hidden_states, + **block_kwargs, + ) else: - hidden_states, audio_hidden_states = block(**block_kwargs) - + hidden_states, audio_hidden_states = block( + hidden_states, + audio_hidden_states, + encoder_hidden_states, + audio_encoder_hidden_states, + **block_kwargs, + ) # 6. Output layers (including unpatchification) scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None] shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]