# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for LTX video VAE tiling and distributed decode behavior.""" from collections.abc import Callable from dataclasses import dataclass from types import SimpleNamespace from typing import Any import pytest import torch pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @pytest.mark.parametrize("initial_deterministic", [False, True]) @pytest.mark.parametrize("fail", [False, True]) def test_ltx_vocoder_deterministic_context(monkeypatch, initial_deterministic, fail): from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _deterministic_ltx_vocoder monkeypatch.setattr(torch.backends.cudnn, "deterministic", initial_deterministic) if torch.backends.cudnn.is_available(): monkeypatch.setattr(torch.backends.cudnn, "benchmark_limit", 17) settings = { name: getattr(torch.backends.cudnn, name) for name in ("enabled", "benchmark", "benchmark_limit", "allow_tf32") } try: with _deterministic_ltx_vocoder(): assert torch.backends.cudnn.deterministic assert {name: getattr(torch.backends.cudnn, name) for name in settings} == settings if fail: raise RuntimeError("injected failure") except RuntimeError as exc: assert fail and str(exc) == "injected failure" assert torch.backends.cudnn.deterministic == initial_deterministic assert {name: getattr(torch.backends.cudnn, name) for name in settings} == settings @dataclass(frozen=True) class _OperatorSetProbe: fna: Callable[..., torch.Tensor] | None @dataclass class _AttentionProbe: kernel_size: tuple[int, int, int] head_dim: int to_out: tuple[torch.nn.Module, ...] projected: torch.Tensor def project_qkv(self, _hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: return self.projected, self.projected, self.projected def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch): import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime class FakeBaseVocoder(torch.nn.Module): def __init__(self): super().__init__() self.input_dtype = None def forward(self, value): self.input_dtype = value.dtype return value monkeypatch.setattr( ltx_runtime.torch, "autocast", lambda **_kwargs: (_ for _ in ()).throw(AssertionError("base vocoder must not use autocast")), ) vocoder = FakeBaseVocoder() output = ltx_runtime._run_ltx_vocoder(vocoder, torch.ones(1, dtype=torch.bfloat16)) assert vocoder.input_dtype == torch.bfloat16 assert output.dtype == torch.bfloat16 class TestLTXDiffusionDecoder: def test_diffusion_decoder_reuses_diffusers_with_scoped_overrides(self, monkeypatch): from diffusers.models.autoencoders.ltx2_diffusion_decoder import ( LTX2VideoDiffusionDecoderModel as DiffusersModel, ) from vllm_omni.diffusion.models.ltx2.vae.decoder import ( LTX2VideoDiffusionDecoderModel, LTX2VideoVaeDiffusionNABlock, LTX2VideoVaeNeighborhoodAttention, LTX2VideoVaeSwiGLU, ) original_init = DiffusersModel.__init__ original_parameters = {} def capture_parameters(self, *args, **kwargs): original_init(self, *args, **kwargs) original_parameters.update(self.named_parameters()) with torch.device("meta"): reference = DiffusersModel() monkeypatch.setattr(DiffusersModel, "__init__", capture_parameters) model = LTX2VideoDiffusionDecoderModel() assert isinstance(model, DiffusersModel) assert all(parameter is original_parameters[name] for name, parameter in model.named_parameters()) assert {name: tensor.shape for name, tensor in model.state_dict().items()} == { name: tensor.shape for name, tensor in reference.state_dict().items() } modules = tuple(model.decoder.modules()) expected_blocks = sum(model.config.decoder_stage_depths) assert sum(isinstance(module, LTX2VideoVaeNeighborhoodAttention) for module in modules) == expected_blocks assert sum(isinstance(module, LTX2VideoVaeSwiGLU) for module in modules) == expected_blocks assert ( sum(isinstance(module, LTX2VideoVaeDiffusionNABlock) for module in modules) == model.config.decoder_stage_depths[-1] ) assert type(model.decoder).__module__.startswith("vllm_omni.") assert model.tiled_decode.__func__ is DiffusersModel.tiled_decode assert type(reference.decoder) is type(model.decoder).__bases__[0] def test_natten_processor_routes_stage5_through_tilelang_fna(self, monkeypatch): from vllm_omni.diffusion.models.ltx2.vae import decoder as diffvae_modules processor = object.__new__(diffvae_modules.LTX2VideoVaeNeighborhoodNattenProcessor) processor.backend = None fallback_calls = [] processor._na3d = lambda *args, **kwargs: fallback_calls.append((args, kwargs)) hidden_states = torch.randn(1, 11, 11, 11, 256, dtype=torch.bfloat16) projected = hidden_states.view(1, 11, 11, 11, 4, 64) class FakeAttention: kernel_size = (11, 11, 11) head_dim = 64 to_out = (torch.nn.Identity(),) @staticmethod def project_qkv(_hidden_states): return projected, projected, projected def token_permute(value, _tile, *, flip_tiled_dims): assert flip_tiled_dims return value, value.shape, None def token_unpermute(value, *, token_layout_shape, tile_shape, flip_tiled_dims): assert token_layout_shape == projected.shape assert tile_shape == (4, 4, 4) assert flip_tiled_dims return value fast_path_calls = [] def fast_path(query, key, value, *, shape, window): fast_path_calls.append((query, key, value, shape, window)) return query monkeypatch.setattr(diffvae_modules, "is_ltx2_fna_eligible", lambda _tensor: True) monkeypatch.setattr( diffvae_modules, "resolve_ltx2_vae_operators", lambda _device: _OperatorSetProbe(fna=fast_path) ) monkeypatch.setattr( diffvae_modules, "_load_natten_token_permutation", lambda: (token_permute, token_unpermute), ) output = processor(FakeAttention(), hidden_states) assert torch.equal(output, hidden_states) assert len(fast_path_calls) == 1 assert fast_path_calls[0][3:] == ((11, 11, 11), (11, 11, 11)) assert not fallback_calls @pytest.mark.parametrize("reason", ["device", "head_dim", "channels", "dtype", "tilelang_missing", "kernel"]) def test_natten_processor_keeps_natten_when_fast_path_does_not_apply(self, monkeypatch, reason): from vllm_omni.diffusion.models.ltx2.vae import decoder as ops processor = object.__new__(ops.LTX2VideoVaeNeighborhoodNattenProcessor) processor.backend = None channels = 128 if reason == "channels" else 256 head_dim = 32 if reason == "head_dim" else 64 x = torch.zeros(1, 11, 11, 11, channels, dtype=torch.float32 if reason == "dtype" else torch.bfloat16) projected = x.view(1, 11, 11, 11, channels // head_dim, head_dim) attn = _AttentionProbe( kernel_size=(3, 5, 5) if reason == "kernel" else (11, 11, 11), head_dim=head_dim, to_out=(torch.nn.Identity(),), projected=projected, ) calls = [] def natten(q, k, v, **kwargs): calls.append(kwargs) return q def unexpected_native_load(*_args, **_kwargs): raise AssertionError("Native FNA must not be loaded for an unsupported configuration") processor._na3d = natten monkeypatch.setattr(ops, "is_ltx2_fna_eligible", lambda _tensor: reason != "device") monkeypatch.setattr( ops, "resolve_ltx2_vae_operators", lambda _device: _OperatorSetProbe(fna=None if reason == "tilelang_missing" else unexpected_native_load), ) assert torch.equal(processor(attn, x), x) assert calls == [{"kernel_size": attn.kernel_size, "scale": 1.0, "backend": None}] def test_diffusion_decoder_noneligible_block_uses_diffusers_forward(self, monkeypatch): from vllm_omni.diffusion.models.ltx2.vae import decoder as diffvae_modules block = object.__new__(diffvae_modules.LTX2VideoVaeDiffusionNABlock) torch.nn.Module.__init__(block) expected = torch.randn(1, 1, 1, 1, 1) calls = [] def upstream_forward(_self, hidden_states, latent_context, modulation, block_mask): assert block_mask is None calls.append((hidden_states, latent_context, modulation)) return expected monkeypatch.setattr( diffvae_modules.diffusers_decoder.LTX2VideoVaeDiffusionNABlock, "forward", upstream_forward, ) hidden_states = torch.randn(1, 1, 1, 1, 1) latent_context = torch.randn_like(hidden_states) modulation = (torch.randn_like(hidden_states),) assert block.forward(hidden_states, latent_context, modulation) is expected assert calls == [(hidden_states, latent_context, modulation)] @pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width"]) def test_distributed_diffusion_decoder_rejects_non_tile_parallel_modes(self, mode): from vllm_omni.diffusion.models.ltx2.vae.distributed import ( DistributedLTX2VideoDiffusionDecoderModel, ) model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel) torch.nn.Module.__init__(model) model.distributed_executor = SimpleNamespace(set_parallel_size=lambda *_args, **_kwargs: None) with pytest.raises(ValueError, match="only supports vae_parallel_mode='tile'"): model.set_parallel_size(2, mode=mode) def test_distributed_diffusion_decoder_accepts_tile_parallel_mode(self): from vllm_omni.diffusion.models.ltx2.vae.distributed import ( DistributedLTX2VideoDiffusionDecoderModel, ) calls = [] model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel) torch.nn.Module.__init__(model) model.distributed_executor = SimpleNamespace( set_parallel_size=lambda parallel_size, mode: calls.append((parallel_size, mode)) ) model.set_parallel_size(4, mode="tile") assert calls == [(4, "tile")] def test_short_clip_keeps_stage5_temporal_context_then_crops_output(self): from vllm_omni.diffusion.models.ltx2.vae.decoder import ( LTX2VideoDiffusionDecoder3d, LTX2VideoDiffusionDecoderModel, ) class IdentityAttention: def build_block_mask(self, _hidden_states): return None class IdentityBlock(torch.nn.Module): def __init__(self): super().__init__() self.attn = IdentityAttention() def forward(self, hidden_states, _block_mask): return hidden_states class IdentityUpsample(torch.nn.Module): def forward(self, hidden_states, *, drop_leading_frame): assert drop_leading_frame return hidden_states decoder = object.__new__(LTX2VideoDiffusionDecoder3d) torch.nn.Module.__init__(decoder) decoder.det_stages = torch.nn.ModuleList([torch.nn.ModuleList([IdentityBlock()])]) decoder.upsamples = torch.nn.ModuleList([IdentityUpsample()]) decoder.trailing_pad_latent_frames = 2 decoder.temporal_compression_ratio = 8 decoder.stage5_kernel = (11, 3, 3) # 17 frames contain one real frame followed by 16 propagated ghost # frames. The reference keeps 11 for stage-5 NATTEN instead of # cropping the tensor immediately back to one. context = decoder.forward_stage_4(torch.zeros(1, 17, 3, 3, 4)) assert context.shape == (1, 11, 3, 3, 4) class FakeDecoder(torch.nn.Module): def forward(self, z, *, generator, num_inference_steps): assert generator is None assert num_inference_steps is None return torch.zeros(z.shape[0], 3, 11, z.shape[3] * 2, z.shape[4] * 2) model = object.__new__(LTX2VideoDiffusionDecoderModel) torch.nn.Module.__init__(model) model.decoder = FakeDecoder() model.use_tiling = False model.tile_sample_min_num_frames = 80 model.tile_sample_min_height = 768 model.tile_sample_min_width = 768 model.temporal_compression_ratio = 8 model.spatial_compression_ratio = 2 output = model.decode(torch.zeros(1, 4, 1, 2, 3), generator=None, return_dict=False)[0] assert output.shape == (1, 3, 1, 4, 6) @pytest.mark.parametrize( ("model_version", "extras", "expected"), [ ("2.5", {}, True), ("2.5", {"ltx2_use_conv_vae": True}, False), ("2.3", {}, False), ], ) def test_ltx25_diffusion_decoder_is_default_with_conv_vae_opt_in(self, model_version, extras, expected): from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder assert _ltx2_use_diffusion_decoder(SimpleNamespace(extras=extras), model_version) is expected def test_conv_vae_opt_in_rejects_non_boolean_model_extra(self): from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder with pytest.raises(TypeError, match="ltx2_use_conv_vae"): _ltx2_use_diffusion_decoder(SimpleNamespace(extras={"ltx2_use_conv_vae": "true"}), "2.5") def test_native_diffusion_decoder_conversion_splits_qkv_and_folds_gates(self): from vllm_omni.diffusion.models.ltx2.vae.decoder import ( convert_ltx25_native_diffusion_decoder_state_dict, ) qkv_weight = torch.arange(24, dtype=torch.float32).reshape(6, 4) qkv_bias = torch.arange(6, dtype=torch.float32) projection_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2) projection_bias = torch.ones(2) gate = torch.tensor([2.0, 3.0]) mean = torch.tensor([0.1, 0.2]) std = torch.tensor([0.3, 0.4]) native = { "encoder.ignored": torch.ones(1), "per_channel_statistics.mean-of-means": mean, "vae.per_channel_statistics.std-of-means": std, "decoder.type_emb": torch.ones(2), "decoder.coarse_head.weight": torch.ones(1), "decoder.det_stages.0.0.attn.qkv.weight": qkv_weight, "decoder.det_stages.0.0.attn.qkv.bias": qkv_bias, "decoder.det_stages.0.0.attn.q_norm.weight": torch.tensor([5.0, 6.0]), "decoder.det_stages.0.0.attn.k_norm.weight": torch.tensor([7.0, 8.0]), "decoder.diff_blocks.0.attn.proj.weight": projection_weight, "decoder.diff_blocks.0.attn.proj.bias": projection_bias, "decoder.diff_blocks.0.gate_msa": gate, "vae.decoder.t_embedder.mlp.0.weight": torch.ones(2, 2), } converted = convert_ltx25_native_diffusion_decoder_state_dict(native) torch.testing.assert_close(converted["latents_mean"], mean) torch.testing.assert_close(converted["latents_std"], std) torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.weight"], qkv_weight[:2]) torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.weight"], qkv_weight[2:4]) torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.weight"], qkv_weight[4:]) torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.bias"], qkv_bias[:2]) torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.bias"], qkv_bias[2:4]) torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.bias"], qkv_bias[4:]) torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_q.weight"], torch.tensor([5.0, 6.0])) torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_k.weight"], torch.tensor([7.0, 8.0])) torch.testing.assert_close( converted["decoder.diff_blocks.0.attn.to_out.0.weight"], gate.unsqueeze(1) * projection_weight ) torch.testing.assert_close(converted["decoder.diff_blocks.0.attn.to_out.0.bias"], gate * projection_bias) assert "decoder.t_embedder.timestep_embedder.linear_1.weight" in converted assert not any("type_emb" in key or "coarse" in key or "gate_msa" in key for key in converted) def test_native_diffusion_decoder_conversion_rejects_invalid_fused_qkv(self): from vllm_omni.diffusion.models.ltx2.vae.decoder import ( convert_ltx25_native_diffusion_decoder_state_dict, ) with pytest.raises(ValueError, match="not divisible by 3"): convert_ltx25_native_diffusion_decoder_state_dict( {"decoder.det_stages.0.0.attn.qkv.weight": torch.ones(5, 2)} ) def test_decode_uses_diffusion_decoder_without_conv_vae_conditioning(self): from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline seen = {} class ConvVae: dtype = torch.float32 config = SimpleNamespace(timestep_conditioning=True) def decode(self, *_args, **_kwargs): raise AssertionError("the convolutional VAE decoder must not run") class DiffusionDecoder: dtype = torch.float32 def decode(self, latents, *, generator, return_dict): seen["latents"] = latents seen["generator"] = generator seen["return_dict"] = return_dict return (latents + 1,) class AudioVae: dtype = torch.float32 def decode(self, audio_latents, *, return_dict): assert return_dict is False return (audio_latents + 2,) class VideoProcessor: def postprocess_video(self, video, *, output_type): seen["output_type"] = output_type return video pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) pipe.use_diffusion_decoder = True pipe.vae = ConvVae() pipe.diffusion_decoder = DiffusionDecoder() pipe.audio_vae = AudioVae() pipe.vocoder = torch.nn.Identity() pipe.video_processor = VideoProcessor() generator = torch.Generator().manual_seed(123) latents = torch.ones(1, 1) output = pipe._decode_output( latents=latents, audio_latents=torch.ones(1, 1), output_type="pt", connector_prompt_embeds=torch.ones(1, 1), generator=generator, device=torch.device("cpu"), decode_timestep=0.7, decode_noise_scale=0.5, prompt_batch_size=1, ) torch.testing.assert_close(seen["latents"], latents) assert seen["generator"] is generator assert seen["return_dict"] is False assert seen["output_type"] == "pt" torch.testing.assert_close(output.output[0], latents + 1) torch.testing.assert_close(output.output[1], torch.full((1, 1), 3.0)) def test_diffusion_decode_denormalizes_with_decoder_statistics(self): from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) pipe.use_diffusion_decoder = True pipe.transformer_spatial_patch_size = 1 pipe.transformer_temporal_patch_size = 1 pipe.vae = SimpleNamespace( latents_mean=torch.full((2,), 100.0), latents_std=torch.full((2,), 10.0), config=SimpleNamespace(scaling_factor=1.0), ) pipe.diffusion_decoder = SimpleNamespace( latents_mean=torch.tensor([10.0, 20.0]), latents_std=torch.tensor([2.0, 4.0]), config=SimpleNamespace(scaling_factor=2.0), ) pipe.audio_vae = SimpleNamespace( latents_mean=torch.tensor(0.0), latents_std=torch.tensor(1.0), ) forward_ctx = SimpleNamespace( latent_num_frames=1, latent_height=1, latent_width=1, original_audio_num_frames=1, latent_mel_bins=2, ) video, audio = pipe._unpack_and_denormalize_stage( forward_ctx, torch.tensor([[[1.0, 2.0]]]), torch.zeros(1, 1, 2), ) torch.testing.assert_close(video.flatten(), torch.tensor([11.0, 24.0])) torch.testing.assert_close(audio, torch.zeros(1, 1, 1, 2)) def test_non_output_rank_does_not_enter_diffusion_decoder(self, monkeypatch): import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline class ConvVae: config = SimpleNamespace(timestep_conditioning=False) def is_distributed_enabled(self): raise AssertionError("DiffVAE decode must not inspect ConvVAE collectives") class DiffusionDecoder: dtype = torch.float32 def decode(self, *_args, **_kwargs): raise AssertionError("only the output rank may run DiffVAE decode") monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True) monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1) pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) pipe.use_diffusion_decoder = True pipe.vae = ConvVae() pipe.diffusion_decoder = DiffusionDecoder() output = pipe._decode_output( latents=torch.ones(1, 1), audio_latents=torch.ones(1, 1), output_type="pt", connector_prompt_embeds=torch.ones(1, 1), generator=None, device=torch.device("cpu"), decode_timestep=0.0, decode_noise_scale=None, prompt_batch_size=1, ) assert output.output[0].numel() == 0 assert output.output[1].numel() == 0 def test_non_output_rank_enters_distributed_diffusion_decoder(self, monkeypatch): import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline class ConvVae: config = SimpleNamespace(timestep_conditioning=False) def is_distributed_enabled(self): raise AssertionError("DiffVAE decode must query the active decoder") class DiffusionDecoder: dtype = torch.float32 def __init__(self): self.decode_calls = 0 def is_distributed_enabled(self): return True def decode(self, *_args, **_kwargs): self.decode_calls += 1 return (torch.empty(0, 3, 0, 0, 0),) monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True) monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1) pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) pipe.use_diffusion_decoder = True pipe.vae = ConvVae() pipe.diffusion_decoder = DiffusionDecoder() output = pipe._decode_output( latents=torch.ones(1, 1), audio_latents=torch.ones(1, 1), output_type="pt", connector_prompt_embeds=torch.ones(1, 1), generator=None, device=torch.device("cpu"), decode_timestep=0.0, decode_noise_scale=None, prompt_batch_size=1, ) assert pipe.diffusion_decoder.decode_calls == 1 assert output.output[0].numel() == 0 assert output.output[1].numel() == 0 def test_diffusion_decoder_patch_parallel_size_one_uses_native_tiling(self, monkeypatch): from vllm_omni.diffusion.models.ltx2.vae.decoder import LTX2VideoDiffusionDecoderModel from vllm_omni.diffusion.models.ltx2.vae.distributed import ( DistributedLTX2VideoDiffusionDecoderModel, ) expected = torch.ones(1, 3, 1, 2, 2) seen = {} def native_tiled_decode(self, z, generator=None, num_inference_steps=None): seen["z"] = z seen["generator"] = generator seen["num_inference_steps"] = num_inference_steps return expected monkeypatch.setattr(LTX2VideoDiffusionDecoderModel, "tiled_decode", native_tiled_decode) model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel) torch.nn.Module.__init__(model) model.is_distributed_enabled = lambda: False generator = torch.Generator().manual_seed(17) z = torch.zeros(1, 1, 1, 1, 1) result = DistributedLTX2VideoDiffusionDecoderModel.tiled_decode( model, z, generator=generator, num_inference_steps=1, ) assert result is expected assert seen["z"] is z assert seen["generator"] is generator assert seen["num_inference_steps"] == 1 def test_distributed_diffusion_tiles_preserve_serial_noise_order(self, monkeypatch): from diffusers.utils.torch_utils import randn_tensor from vllm_omni.diffusion.models.ltx2.vae.distributed import ( DistributedLTX2VideoDiffusionDecoderModel, ) model = DistributedLTX2VideoDiffusionDecoderModel( out_channels=1, latent_channels=1, patch_size=1, decoder_head_dim=8, decoder_stage_channels=(8, 8, 8, 8, 8), decoder_stage_depths=(1, 1, 1, 1, 1), decoder_stage_kernels=((1, 1, 1),) * 4, decoder_upsample_strides=((1, 1, 1),) * 4, decoder_upsample_channel_reductions=(1, 1, 1, 1), decoder_stage5_kernel=(1, 1, 1), spatial_compression_ratio=1, temporal_compression_ratio=1, ) model.tile_sample_min_num_frames = 2 model.tile_sample_stride_num_frames = 1 model.tile_sample_min_height = 2 model.tile_sample_stride_height = 1 model.tile_sample_min_width = 2 model.tile_sample_stride_width = 1 model.distributed_executor = SimpleNamespace(group=None) model.decoder.forward_stages_1_to_3 = lambda z: z.permute(0, 2, 3, 4, 1) monkeypatch.setattr(torch.distributed, "broadcast", lambda *_args, **_kwargs: None) z = torch.zeros(1, 1, 2, 3, 3) distributed_generator = torch.Generator().manual_seed(123) reference_generator = torch.Generator().manual_seed(123) tasks, grid_spec = model._distributed_tile_split(z, distributed_generator, num_inference_steps=1) assert grid_spec.grid_shape == (1, 3, 3) assert grid_spec.split_dims == (2, 3, 4) for task in tasks: tile_shape = model._tiled_pixel_shape_from_features( task.tensor, drop_leading_frame=task.drop_leading_frame, crop_trailing_ghost=task.crop_trailing_ghost, ) expected = randn_tensor(tile_shape, generator=reference_generator, device=z.device, dtype=z.dtype) actual = randn_tensor(tile_shape, generator=task.noise_generator, device=z.device, dtype=z.dtype) torch.testing.assert_close(actual, expected) assert torch.equal(distributed_generator.get_state(), reference_generator.get_state()) class TestLTXOutputRank: @pytest.mark.parametrize( ("distributed_vae_state", "expected_decode_calls"), [(False, 0), (True, 1), (RuntimeError("unavailable"), None)], ) def test_non_output_rank_only_enters_collective_vae_decode( self, monkeypatch, distributed_vae_state, expected_decode_calls, ): import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline class FakeVae: dtype = torch.float32 config = SimpleNamespace(timestep_conditioning=False) def __init__(self): self.decode_calls = 0 def is_distributed_enabled(self): if isinstance(distributed_vae_state, Exception): raise distributed_vae_state return distributed_vae_state def decode(self, *_args, **_kwargs): self.decode_calls += 1 return (torch.ones(1, 1),) monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True) monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1) pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) pipe.vae = FakeVae() decode_kwargs = { "latents": torch.ones(1, 1), "audio_latents": torch.ones(1, 1), "output_type": "np", "connector_prompt_embeds": torch.ones(1, 1), "generator": None, "device": torch.device("cpu"), "decode_timestep": 0.0, "decode_noise_scale": None, "prompt_batch_size": 1, } if isinstance(distributed_vae_state, Exception): with pytest.raises(RuntimeError, match="unavailable"): pipe._decode_output(**decode_kwargs) assert pipe.vae.decode_calls == 0 return output = pipe._decode_output(**decode_kwargs) assert pipe.vae.decode_calls == expected_decode_calls assert output.output[0].numel() == 0 assert output.output[1].numel() == 0 class TestLTX23VaeDistributedDecode: """Test LTX-2.3 distributed VAE helpers without loading weights.""" def test_ltx23_video_vae_is_distributed_tile_only_class(self): from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( DistributedAutoencoderKLLTX2Video, ) from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin assert issubclass(DistributedAutoencoderKLLTX2Video, DistributedVaeMixin) assert not hasattr(DistributedAutoencoderKLLTX2Video, "patch_split") def test_ltx23_vae_executor_gathers_known_tile_shapes_and_returns_empty_on_non_rank0(self): from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import LTX2VaeExecutor from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( DistributedOperator, GridSpec, TileTask, ) z = torch.zeros(1, 1, 1, 1, 1) tile_output_shapes = { 0: (1, 1, 1, 2, 2), 1: (1, 1, 1, 2, 1), 2: (1, 1, 1, 1, 2), 3: (1, 1, 1, 1, 1), } tasks = [ TileTask(0, (0, 0), z, workload=4), TileTask(1, (0, 1), z, workload=2), TileTask(2, (1, 0), z, workload=2), TileTask(3, (1, 1), z, workload=1), ] grid_spec = GridSpec( split_dims=(3, 4), grid_shape=(2, 2), tile_spec={ "max_tile_output_shape": (1, 1, 1, 2, 2), "tile_output_shapes": tile_output_shapes, }, output_dtype=torch.float32, ) seen: dict[str, Any] = {} def exec_tile(task): return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1)) def merge_tiles(coord_tensor_map, passed_grid_spec): seen["merged_shapes"] = {coord: tuple(tile.shape) for coord, tile in coord_tensor_map.items()} assert passed_grid_spec is grid_spec return torch.stack( [ coord_tensor_map[(0, 0)].flatten()[0], coord_tensor_map[(0, 1)].flatten()[0], coord_tensor_map[(1, 0)].flatten()[0], coord_tensor_map[(1, 1)].flatten()[0], ] ) operator = DistributedOperator(split=lambda _z: (tasks, grid_spec), exec=exec_tile, merge=merge_tiles) rank0_executor = object.__new__(LTX2VaeExecutor) rank0_executor.parallel_size = 2 rank0_executor.world_size = 2 rank0_executor.rank = 0 def gather_rank0(local_tile_tensor): assigned = rank0_executor._balance_tasks(tasks, 2) rank1_results = [(task.tile_id, exec_tile(task)) for task in assigned[1]] rank1_tile_tensor = rank0_executor._pack_local_tiles_without_meta( rank1_results, list(local_tile_tensor.shape), z.device, torch.float32, ) seen["rank0_gather_shape"] = tuple(local_tile_tensor.shape) return [local_tile_tensor, rank1_tile_tensor] def fail_final_sync(*_args, **_kwargs): raise AssertionError("broadcast_result=False should not sync the final result") rank0_executor.gather_tensors = gather_rank0 rank0_executor._sync_final_result = fail_final_sync rank0_result = rank0_executor.execute(z, operator, broadcast_result=False) torch.testing.assert_close(rank0_result, torch.tensor([1.0, 2.0, 3.0, 4.0])) assert seen["rank0_gather_shape"] == (2, 1, 1, 1, 2, 2) assert seen["merged_shapes"] == { (0, 0): (1, 1, 1, 2, 2), (0, 1): (1, 1, 1, 2, 1), (1, 0): (1, 1, 1, 1, 2), (1, 1): (1, 1, 1, 1, 1), } non_rank0_executor = object.__new__(LTX2VaeExecutor) non_rank0_executor.parallel_size = 2 non_rank0_executor.world_size = 2 non_rank0_executor.rank = 1 def gather_rank1(local_tile_tensor): seen["rank1_gather_shape"] = tuple(local_tile_tensor.shape) return None def fail_non_rank0_merge(*_args, **_kwargs): raise AssertionError("non-rank0 should not merge gathered tiles") non_rank0_executor.gather_tensors = gather_rank1 non_rank0_executor._sync_final_result = fail_final_sync empty_result = non_rank0_executor.execute( z, DistributedOperator( split=lambda _z: (tasks, grid_spec), exec=exec_tile, merge=fail_non_rank0_merge, ), broadcast_result=False, ) assert tuple(empty_result.shape) == (0,) assert seen["rank1_gather_shape"] == (2, 1, 1, 1, 2, 2) class TestLTX23VaeTiling: """Test LTX-2.3 video VAE tile helpers without loading weights.""" def test_ltx23_video_vae_tile_split_uses_native_ltx23_tile_geometry(self): from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( DistributedAutoencoderKLLTX2Video, ) vae = SimpleNamespace( spatial_compression_ratio=32, tile_sample_min_height=512, tile_sample_min_width=512, tile_sample_stride_height=448, tile_sample_stride_width=448, temporal_compression_ratio=8, dtype=torch.float32, ) z = torch.zeros(1, 2, 5, 16, 24) tasks, grid_spec = DistributedAutoencoderKLLTX2Video.tile_split(vae, z) assert grid_spec.grid_shape == (2, 2) assert grid_spec.split_dims == (3, 4) assert grid_spec.tile_spec["sample_height"] == 512 assert grid_spec.tile_spec["sample_width"] == 768 assert grid_spec.tile_spec["blend_height"] == 64 assert grid_spec.tile_spec["blend_width"] == 64 assert grid_spec.tile_spec["max_tile_output_shape"] == (1, 3, 33, 512, 512) assert grid_spec.tile_spec["tile_output_shapes"] == { 0: (1, 3, 33, 512, 512), 1: (1, 3, 33, 512, 320), 2: (1, 3, 33, 64, 512), 3: (1, 3, 33, 64, 320), } assert [task.grid_coord for task in tasks] == [(0, 0), (0, 1), (1, 0), (1, 1)] assert [tuple(task.tensor.shape) for task in tasks] == [ (1, 2, 5, 16, 16), (1, 2, 5, 16, 10), (1, 2, 5, 2, 16), (1, 2, 5, 2, 10), ] assert [task.workload for task in tasks] == [5 * 16 * 16, 5 * 16 * 10, 5 * 2 * 16, 5 * 2 * 10] def test_ltx23_video_vae_tile_merge_blends_and_crops_like_tiled_decode(self): from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( DistributedAutoencoderKLLTX2Video, ) from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import GridSpec class FakeVae: def __init__(self): self.blend_calls = [] def clear_cache(self): pass def blend_v(self, _previous, current, blend_height): self.blend_calls.append(("v", blend_height)) return current def blend_h(self, _previous, current, blend_width): self.blend_calls.append(("h", blend_width)) return current fake_vae = FakeVae() grid_spec = GridSpec( split_dims=(3, 4), grid_shape=(2, 2), tile_spec={ "sample_height": 10, "sample_width": 10, "blend_height": 1, "blend_width": 2, "tile_sample_stride_height": 5, "tile_sample_stride_width": 5, }, ) tiles = { (0, 0): torch.full((1, 3, 2, 6, 6), 1.0), (0, 1): torch.full((1, 3, 2, 6, 6), 2.0), (1, 0): torch.full((1, 3, 2, 6, 6), 3.0), (1, 1): torch.full((1, 3, 2, 6, 6), 4.0), } merged = DistributedAutoencoderKLLTX2Video.tile_merge(fake_vae, tiles, grid_spec) assert merged.shape == (1, 3, 2, 10, 10) assert fake_vae.blend_calls == [("h", 2), ("v", 1), ("v", 1), ("h", 2)] torch.testing.assert_close(merged[:, :, :, :5, :5], torch.ones(1, 3, 2, 5, 5)) torch.testing.assert_close(merged[:, :, :, :5, 5:], torch.full((1, 3, 2, 5, 5), 2.0)) torch.testing.assert_close(merged[:, :, :, 5:, :5], torch.full((1, 3, 2, 5, 5), 3.0)) torch.testing.assert_close(merged[:, :, :, 5:, 5:], torch.full((1, 3, 2, 5, 5), 4.0)) def test_ltx23_video_vae_tiled_decode_dispatches_to_tile_operator(self): from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_ltx2 z = torch.zeros(1, 2, 1, 16, 24) expected = torch.ones(1, 3, 1, 512, 768) seen = {} class FakeExecutor: def execute(self, tensor, operator, broadcast_result=True): seen["tensor"] = tensor seen["operator"] = operator seen["broadcast_result"] = broadcast_result return expected vae = SimpleNamespace(distributed_executor=FakeExecutor(), is_distributed_enabled=lambda: True) vae.tile_split = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_split.__get__(vae) vae.tile_exec = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_exec.__get__(vae) vae.tile_merge = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_merge.__get__(vae) output = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tiled_decode( vae, z, temb=torch.tensor(0.5), return_dict=False, ) assert len(output) == 1 assert output[0] is expected assert seen["tensor"] is z assert seen["broadcast_result"] is False assert seen["operator"].split.__name__ == "tile_split" assert seen["operator"].merge.__name__ == "tile_merge"