From dd6f32cb47b196ec29b29a4a01d8304e3e46b0d2 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Tue, 11 Aug 2026 07:38:08 +0800 Subject: [PATCH] [Perf][Diffusion] Add LTX2 channels-last VAE path Signed-off-by: Hongsheng Liu --- tests/config/test_omni_config.py | 24 +++++ tests/diffusion/models/ltx2/test_ltx2_vae.py | 79 +++++++++++++++ tests/diffusion/test_diffusion_engine.py | 12 +++ .../test_async_omni_diffusion_config.py | 21 ++++ vllm_omni/config/omni_config.py | 1 + vllm_omni/config/stage_config.py | 1 + vllm_omni/diffusion/data.py | 4 + .../autoencoders/autoencoder_kl_ltx2.py | 96 +++++++++++++++++++ .../diffusion/models/ltx2/ltx2_components.py | 2 + vllm_omni/engine/arg_utils.py | 1 + vllm_omni/engine/async_omni_engine.py | 1 + vllm_omni/entrypoints/cli/serve.py | 8 ++ 12 files changed, 250 insertions(+) diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 3d2caeabf..2cad1261d 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -776,6 +776,30 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic(): omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false") +def test_structured_diffusion_config_rejects_non_boolean_vae_channels_last_3d(): + with pytest.raises(ValidationError, match="vae_use_channels_last_3d"): + omni_config_module._DiffusionConfigProjection(vae_use_channels_last_3d="true") + + +def test_from_pipeline_config_routes_vae_channels_last_3d(tmp_path): + deploy_path = tmp_path / "dreamzero_vae_layout.yaml" + deploy_path.write_text( + "\n".join( + [ + "pipeline: dreamzero", + "async_chunk: false", + "stages:", + " - stage_id: 0", + " vae_use_channels_last_3d: true", + ] + ) + ) + + stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0) + + assert stage.diffusion_config.vae_use_channels_last_3d is True + + def test_structured_diffusion_config_rejects_invalid_compile_granularity(): with pytest.raises(ValidationError, match="diffusion_compile_granularity"): omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block") diff --git a/tests/diffusion/models/ltx2/test_ltx2_vae.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py index 108bb960d..4d69712db 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_vae.py +++ b/tests/diffusion/models/ltx2/test_ltx2_vae.py @@ -11,6 +11,85 @@ import torch pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +class TestLTXChannelsLast3D: + @pytest.mark.parametrize( + ("causal", "kernel_size"), + [ + (True, 3), + (False, 3), + (True, (3, 1, 1)), + (True, 1), + ], + ) + def test_layout_preserving_causal_conv_matches_original(self, causal, kernel_size): + from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d + + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + _enable_channels_last_3d, + _weight_is_channels_last_3d, + ) + + torch.manual_seed(0) + conv = LTX2VideoCausalConv3d(4, 5, kernel_size).to(dtype=torch.float32).eval() + hidden_states = torch.randn(1, 4, 5, 6, 7, dtype=torch.float32) + + # Pointwise Conv3d storage is valid in both contiguous formats because + # all spatial/temporal kernel dimensions are singleton. + if kernel_size != 1: + assert not _weight_is_channels_last_3d(conv) + with torch.no_grad(): + expected = conv(hidden_states.clone(), causal=causal) + + converted, patched = _enable_channels_last_3d(conv) + + assert (converted, patched) == (1, 1) + assert _weight_is_channels_last_3d(conv) + with torch.no_grad(): + actual = conv(hidden_states.clone(), causal=causal) + + torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4) + + def test_temporal_pad_replicates_edges_in_channels_last_3d(self): + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + _temporal_pad_channels_last_3d, + ) + + torch.manual_seed(0) + hidden_states = torch.randn(1, 4, 3, 2, 2, dtype=torch.float32) + + actual = _temporal_pad_channels_last_3d(hidden_states, left=2, right=1) + expected = torch.cat( + [ + hidden_states[:, :, :1].repeat(1, 1, 2, 1, 1), + hidden_states, + hidden_states[:, :, -1:], + ], + dim=2, + ) + + assert actual.is_contiguous(memory_format=torch.channels_last_3d) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_contiguous_weight_keeps_original_padding_path(self, monkeypatch): + from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d + + import vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 as ltx2_vae + + conv = LTX2VideoCausalConv3d(4, 5, 3).to(dtype=torch.float32).eval() + hidden_states = torch.randn(1, 4, 3, 2, 2, dtype=torch.float32) + + def fail_channels_last_pad(*_args, **_kwargs): + raise AssertionError("contiguous weights must retain the original padding path") + + monkeypatch.setattr(ltx2_vae, "_temporal_pad_channels_last_3d", fail_channels_last_pad) + + with torch.no_grad(): + expected = conv(hidden_states, causal=True) + actual = ltx2_vae._layout_preserving_causal_conv_forward(conv, hidden_states, causal=True) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + class TestLTXOutputRank: @pytest.mark.parametrize( ("distributed_vae_state", "expected_decode_calls"), diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index a0c6acb02..5edb3c311 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -534,6 +534,18 @@ class TestDiffusionCompileConfig: with pytest.raises(TypeError, match="diffusion_compile_dynamic"): OmniDiffusionConfig(model="test", diffusion_compile_dynamic="false") + def test_channels_last_3d_is_opt_in_and_strict(self) -> None: + assert OmniDiffusionConfig(model="test").vae_use_channels_last_3d is False + assert ( + OmniDiffusionConfig.from_kwargs( + model="test", + vae_use_channels_last_3d=True, + ).vae_use_channels_last_3d + is True + ) + with pytest.raises(TypeError, match="vae_use_channels_last_3d"): + OmniDiffusionConfig(model="test", vae_use_channels_last_3d="true") + @pytest.mark.parametrize( "kwargs, feature", [ diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 64b586bf2..a5bf40010 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -20,6 +20,7 @@ def test_default_stage_config_includes_cache_backend(): "cache_backend": "cache_dit", "cache_config": '{"Fn_compute_blocks": 2}', "vae_use_slicing": True, + "vae_use_channels_last_3d": True, "ulysses_degree": 2, } )[0] @@ -29,6 +30,7 @@ def test_default_stage_config_includes_cache_backend(): assert engine_args["cache_backend"] == "cache_dit" assert engine_args["cache_config"]["Fn_compute_blocks"] == 2 assert engine_args["vae_use_slicing"] is True + assert engine_args["vae_use_channels_last_3d"] is True assert engine_args["parallel_config"].ulysses_degree == 2 assert engine_args["model_stage"] == "diffusion" @@ -340,6 +342,25 @@ def test_serve_cli_accepts_diffusion_compile_controls(): assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is False +def test_serve_cli_accepts_vae_channels_last_3d_opt_in(): + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="command") + OmniServeCommand().subparser_init(subparsers) + args = parser.parse_args( + [ + "serve", + "Lightricks/LTX-2", + "--omni", + "--vae-use-channels-last-3d", + ] + ) + + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(args.get_explicit_kwargs_dict())[0] + + assert args.vae_use_channels_last_3d is True + assert stage_cfg["engine_args"]["vae_use_channels_last_3d"] is True + + def test_serve_cli_accepts_diffusion_attention_backend(): """Ensure diffusion serve CLI exposes the shorthand backend flag.""" parser = TrackingArgumentParser() diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 8997b12e1..ead9eded3 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -573,6 +573,7 @@ class _DiffusionConfigProjection: fa_deterministic: bool = False vae_use_slicing: bool = False vae_use_tiling: bool = False + vae_use_channels_last_3d: bool = Field(default=False, strict=True) mask_strategy_file_path: str | None = None skip_time_steps: int = 15 VSA_sparsity: float = 0.0 diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 5b6b9a01a..492add581 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -401,6 +401,7 @@ class StageDeployConfig: step_execution: bool | None = None vae_use_slicing: bool | None = None vae_use_tiling: bool | None = None + vae_use_channels_last_3d: bool | None = None boundary_ratio: float | None = None flow_shift: float | None = None diffusion_kv_cache_dtype: str | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 73051692d..4ec1dd6af 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -711,6 +711,8 @@ class OmniDiffusionConfig: # VAE memory optimization parameters vae_use_slicing: bool = False vae_use_tiling: bool = False + # Opt in because CL3D performance varies materially by model shape and GPU. + vae_use_channels_last_3d: bool = False # STA (Sliding Tile Attention) parameters mask_strategy_file_path: str | None = None @@ -930,6 +932,8 @@ class OmniDiffusionConfig: ) if not isinstance(self.diffusion_compile_dynamic, bool): raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}") + if not isinstance(self.vae_use_channels_last_3d, bool): + raise TypeError(f"vae_use_channels_last_3d must be a bool, got {type(self.vae_use_channels_last_3d)!r}") self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode) if self.omni_kv_config is None: diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_ltx2.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_ltx2.py index 0e7410808..f56305852 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_ltx2.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_ltx2.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import MethodType from typing import Any import torch from diffusers import AutoencoderKLLTX2Video +from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d from diffusers.models.autoencoders.vae import DecoderOutput +from torch import nn from vllm.logger import init_logger from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( @@ -19,6 +22,89 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor logger = init_logger(__name__) +# Adapted from https://github.com/sgl-project/sglang/pull/27431. +def _weight_is_channels_last_3d(module: LTX2VideoCausalConv3d) -> bool: + weight = module.conv.weight + return ( + weight.dim() == 5 + and hasattr(torch, "channels_last_3d") + and weight.is_contiguous(memory_format=torch.channels_last_3d) + ) + + +def _temporal_pad_channels_last_3d( + hidden_states: torch.Tensor, + *, + left: int, + right: int, +) -> torch.Tensor: + """Replicate temporal edge frames directly into CL3D storage.""" + batch, channels, frames, height, width = hidden_states.shape + padded = torch.empty( + (batch, channels, frames + left + right, height, width), + dtype=hidden_states.dtype, + device=hidden_states.device, + memory_format=torch.channels_last_3d, + ) + padded[:, :, left : left + frames].copy_(hidden_states) + if left: + padded[:, :, :left].copy_(hidden_states[:, :, :1]) + if right: + padded[:, :, left + frames :].copy_(hidden_states[:, :, -1:]) + return padded + + +def _layout_preserving_causal_conv_forward( + module: LTX2VideoCausalConv3d, + hidden_states: torch.Tensor, + causal: bool = True, +) -> torch.Tensor: + """Diffusers LTX causal-conv math with a CL3D-preserving pad path.""" + time_kernel_size = module.kernel_size[0] + use_channels_last_pad = hidden_states.dim() == 5 and _weight_is_channels_last_3d(module) + + if use_channels_last_pad: + if causal: + left, right = time_kernel_size - 1, 0 + else: + left = right = (time_kernel_size - 1) // 2 + if left == 0 and right == 0: + hidden_states = hidden_states.contiguous(memory_format=torch.channels_last_3d) + else: + hidden_states = _temporal_pad_channels_last_3d( + hidden_states, + left=left, + right=right, + ) + elif causal: + pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, time_kernel_size - 1, 1, 1)) + hidden_states = torch.concatenate([pad_left, hidden_states], dim=2) + else: + pad_size = (time_kernel_size - 1) // 2 + pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, pad_size, 1, 1)) + pad_right = hidden_states[:, :, -1:, :, :].repeat((1, 1, pad_size, 1, 1)) + hidden_states = torch.concatenate([pad_left, hidden_states, pad_right], dim=2) + + return module.conv(hidden_states) + + +def _enable_channels_last_3d(model: nn.Module) -> tuple[int, int]: + """Convert LTX Conv3d weights and install layout-preserving causal pads.""" + if not hasattr(torch, "channels_last_3d"): + return 0, 0 + + converted_convs = 0 + patched_causal_convs = 0 + for module in model.modules(): + if isinstance(module, LTX2VideoCausalConv3d): + module.forward = MethodType(_layout_preserving_causal_conv_forward, module) + patched_causal_convs += 1 + if isinstance(module, nn.Conv3d): + module.weight.data = module.weight.data.contiguous(memory_format=torch.channels_last_3d) + converted_convs += 1 + return converted_convs, patched_causal_convs + + class LTX2VaeExecutor(DistributedVaeExecutor): def _pack_local_tiles_without_meta(self, local_results, global_padding_shape, device, dtype): tile_tensor = torch.zeros(global_padding_shape, device=device, dtype=dtype) @@ -80,6 +166,16 @@ class DistributedAutoencoderKLLTX2Video(AutoencoderKLLTX2Video, DistributedVaeMi model.init_distributed() return model + def enable_channels_last_3d(self) -> None: + """Opt into the layout-preserving CL3D path for this video VAE.""" + converted_convs, patched_causal_convs = _enable_channels_last_3d(self) + logger.info( + "LTX2 video VAE: converted %d Conv3d weights to channels_last_3d and " + "installed %d layout-preserving causal conv paths", + converted_convs, + patched_causal_convs, + ) + def init_distributed(self): self.distributed_executor = LTX2VaeExecutor() diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index 54c866eb7..24ddef958 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -365,6 +365,8 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: local_files_only=local_files_only, dtype=dtype, ) + if getattr(od_config, "vae_use_channels_last_3d", False): + pipeline.vae.enable_channels_last_3d() pipeline.audio_vae = _load_component( AutoencoderKLLTX2Audio, model, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 02ec3587b..bc4249085 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -491,6 +491,7 @@ class OrchestratorArgs: step_execution: bool = False vae_use_slicing: bool = False vae_use_tiling: bool = False + vae_use_channels_last_3d: bool = False enable_multithread_weight_load: bool = True num_weight_load_threads: int = 4 enable_cpu_offload: bool = False diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index fc80cfa93..14a87ce4a 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1026,6 +1026,7 @@ class AsyncOmniEngine: "request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0), "vae_use_slicing": kwargs.get("vae_use_slicing", False), "vae_use_tiling": kwargs.get("vae_use_tiling", False), + "vae_use_channels_last_3d": kwargs.get("vae_use_channels_last_3d", False), "cache_backend": cache_backend, "cache_config": cache_config, "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 7fb9cb8cd..0bd1b1ff5 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -629,6 +629,14 @@ class OmniServeCommand(CLISubcommand): action="store_true", help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).", ) + omni_config_group.add_argument( + "--vae-use-channels-last-3d", + action="store_true", + help=( + "Use channels_last_3d Conv3d weights for supported VAEs. This is opt-in because performance " + "depends on model shape and GPU architecture." + ), + ) # Parallel weight loading (faster diffusion startup) omni_config_group.add_argument(