[Perf][Diffusion] Add LTX2 channels-last VAE path

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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")

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"),

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",
         [

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()

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

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

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:

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()
 

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,

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

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),

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(