Fix LTX2 channels-last review scope

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
(cherry picked from commit dc072a12bd)

tests/config/test_omni_config.py

@@ -751,7 +751,6 @@ def test_from_pipeline_config_routes_regional_compile_dynamic(tmp_path):
                 "  - stage_id: 0",
                 "    diffusion_compile_granularity: regional",
                 "    diffusion_compile_dynamic: false",
-                "    diffusion_compile_vae: false",
             ]
         )
     )
@@ -763,16 +762,13 @@ def test_from_pipeline_config_routes_regional_compile_dynamic(tmp_path):
         cli_overrides={
             "diffusion_compile_granularity": "full",
             "diffusion_compile_dynamic": True,
-            "diffusion_compile_vae": True,
         },
     ).stage_by_id(0)
 
     assert configured_stage.diffusion_config.diffusion_compile_granularity == "regional"
     assert configured_stage.diffusion_config.diffusion_compile_dynamic is False
-    assert configured_stage.diffusion_config.diffusion_compile_vae is False
     assert overridden_stage.diffusion_config.diffusion_compile_granularity == "full"
     assert overridden_stage.diffusion_config.diffusion_compile_dynamic is True
-    assert overridden_stage.diffusion_config.diffusion_compile_vae is True
 
 
 def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
@@ -780,9 +776,24 @@ 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_compile_vae():
-    with pytest.raises(ValidationError, match="diffusion_compile_vae"):
-        omni_config_module._DiffusionConfigProjection(diffusion_compile_vae="false")
+def test_from_pipeline_config_routes_ltx2_model_extra(tmp_path):
+    deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    extras:",
+                "      ltx2_use_channels_last_3d: true",
+            ]
+        )
+    )
+
+    stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
+
+    assert stage.diffusion_config.extras["ltx2_use_channels_last_3d"] is True
 
 
 def test_structured_diffusion_config_rejects_invalid_compile_granularity():

tests/diffusion/models/ltx2/test_ltx2_vae.py

@@ -11,6 +11,100 @@ import torch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+class TestLTXChannelsLast3D:
+    @pytest.mark.parametrize(
+        ("extras", "expected"),
+        [({}, False), ({"ltx2_use_channels_last_3d": True}, True)],
+    )
+    def test_channels_last_opt_in_is_ltx2_model_extra(self, extras, expected):
+        from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_channels_last_3d
+
+        assert _ltx2_use_channels_last_3d(SimpleNamespace(extras=extras)) is expected
+
+    def test_channels_last_opt_in_rejects_non_boolean_model_extra(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_channels_last_3d
+
+        with pytest.raises(TypeError, match="ltx2_use_channels_last_3d"):
+            _ltx2_use_channels_last_3d(SimpleNamespace(extras={"ltx2_use_channels_last_3d": "true"}))
+
+    @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_compile.py

@@ -5,7 +5,7 @@ import pytest
 import torch.nn as nn
 
 import vllm_omni.diffusion.compile as compile_module
-from vllm_omni.diffusion.compile import compile_callable_with_fallback, regionally_compile
+from vllm_omni.diffusion.compile import regionally_compile
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -74,50 +74,3 @@ def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypat
         regionally_compile(model, dynamic=True)
 
     assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards
-
-
-def test_compile_callable_falls_back_once_after_lazy_failure(monkeypatch, caplog):
-    compiled_calls = 0
-    eager_calls = 0
-
-    def eager(value):
-        nonlocal eager_calls
-        eager_calls += 1
-        return value + 1
-
-    def compile_fn(fn, **kwargs):
-        assert fn is eager
-        assert kwargs == {"mode": "default"}
-
-        def compiled(value):
-            nonlocal compiled_calls
-            compiled_calls += 1
-            raise RuntimeError("lazy failure")
-
-        return compiled
-
-    monkeypatch.setattr(compile_module.torch, "compile", compile_fn)
-    callable_fn, activated = compile_callable_with_fallback(eager, label="test callable", mode="default")
-
-    assert activated
-    assert callable_fn(1) == 2
-    assert callable_fn(2) == 3
-    assert compiled_calls == 1
-    assert eager_calls == 2
-    assert "disabling compilation and retrying eagerly" in caplog.text
-
-
-def test_compile_callable_keeps_eager_on_setup_failure(monkeypatch, caplog):
-    def eager(value):
-        return value + 1
-
-    def compile_fn(*args, **kwargs):
-        raise RuntimeError("setup failure")
-
-    monkeypatch.setattr(compile_module.torch, "compile", compile_fn)
-    callable_fn, activated = compile_callable_with_fallback(eager, label="test callable")
-
-    assert not activated
-    assert callable_fn is eager
-    assert callable_fn(1) == 2
-    assert "continuing eagerly" in caplog.text

tests/diffusion/test_diffusion_engine.py

@@ -515,19 +515,16 @@ class TestDiffusionCompileConfig:
 
         assert config.diffusion_compile_granularity == "regional"
         assert config.diffusion_compile_dynamic is True
-        assert config.diffusion_compile_vae is True
 
     def test_from_kwargs_preserves_compile_controls(self) -> None:
         config = OmniDiffusionConfig.from_kwargs(
             model="test",
             diffusion_compile_granularity="full",
             diffusion_compile_dynamic=False,
-            diffusion_compile_vae=False,
         )
 
         assert config.diffusion_compile_granularity == "full"
         assert config.diffusion_compile_dynamic is False
-        assert config.diffusion_compile_vae is False
 
     def test_config_rejects_invalid_compile_granularity(self) -> None:
         with pytest.raises(ValueError, match="diffusion_compile_granularity"):
@@ -537,10 +534,6 @@ class TestDiffusionCompileConfig:
         with pytest.raises(TypeError, match="diffusion_compile_dynamic"):
             OmniDiffusionConfig(model="test", diffusion_compile_dynamic="false")
 
-    def test_config_rejects_non_boolean_compile_vae(self) -> None:
-        with pytest.raises(TypeError, match="diffusion_compile_vae"):
-            OmniDiffusionConfig(model="test", diffusion_compile_vae="false")
-
     @pytest.mark.parametrize(
         "kwargs, feature",
         [

tests/diffusion/test_diffusion_model_runner.py

@@ -130,23 +130,6 @@ class _CompileTrackingModel:
         return self
 
 
-class _DecodeVAE(torch.nn.Module):
-    def decode(self, value):
-        return value + 1
-
-
-class _PipelineWithVAE(torch.nn.Module):
-    _dit_modules = []
-    _encoder_modules = []
-    _vae_modules = ["vae", "vae_alias"]
-    _resident_modules = []
-
-    def __init__(self):
-        super().__init__()
-        self.vae = _DecodeVAE()
-        self.vae_alias = self.vae
-
-
 def _make_request():
     sampling_params = SimpleNamespace(
         generator=None,
@@ -337,31 +320,6 @@ def test_compile_transformer_falls_back_after_synchronous_setup_failure(monkeypa
     assert "lazy compilation errors" in caplog.text
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_compile_vae_decodes_uses_discovery_and_deduplicates_aliases(monkeypatch):
-    runner = object.__new__(DiffusionModelRunner)
-    runner.pipeline = _PipelineWithVAE()
-    compile_calls = []
-
-    def compile_fn(fn, **kwargs):
-        compile_calls.append((fn, kwargs))
-        return lambda *args, **call_kwargs: fn(*args, **call_kwargs) + 10
-
-    monkeypatch.setattr(model_runner_module.torch, "compile", compile_fn)
-
-    DiffusionModelRunner._compile_vae_decodes(runner)
-
-    assert len(compile_calls) == 1
-    assert compile_calls[0][1] == {
-        "mode": "default",
-        "fullgraph": False,
-        "dynamic": None,
-    }
-    assert runner.pipeline.vae.decode(1) == 12
-    assert runner.pipeline.vae_alias.decode is runner.pipeline.vae.decode
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):

tests/diffusion/test_diffusion_scheduler.py

@@ -971,44 +971,6 @@ class TestDiffusionEngine:
         admitted_request = engine.add_req_and_wait_for_response.call_args.args[0]
         assert admitted_request.request_id == "dummy_req_id"
 
-    def test_compile_warmup_repeats_through_normal_request_path(self, mocker: MockerFixture) -> None:
-        engine = DiffusionEngine.__new__(DiffusionEngine)
-        engine.od_config = SimpleNamespace(
-            model_class_name="mock_model",
-            diffusion_load_format="default",
-            diffusion_compile_vae=True,
-            enforce_eager=False,
-        )
-        engine.add_req_and_wait_for_response = mocker.Mock(
-            side_effect=[DiffusionOutput(output=None), DiffusionOutput(output=None)]
-        )
-
-        engine._dummy_run()
-
-        requests = [call.args[0] for call in engine.add_req_and_wait_for_response.call_args_list]
-        assert len(requests) == 2
-        assert requests[0].is_dummy_run()
-        assert not requests[1].is_dummy_run()
-        assert requests[1].prompt == {"prompt": "A detailed image."}
-        assert requests[1].sampling_params.need_kv_receive is False
-        assert requests[1].sampling_params.num_inference_steps == 1
-
-    def test_compile_real_path_warmup_failure_is_non_fatal(self, mocker: MockerFixture, caplog) -> None:
-        engine = DiffusionEngine.__new__(DiffusionEngine)
-        engine.od_config = SimpleNamespace(
-            model_class_name="mock_model",
-            diffusion_load_format="default",
-            diffusion_compile_vae=True,
-            enforce_eager=False,
-        )
-        engine.add_req_and_wait_for_response = mocker.Mock(
-            side_effect=[DiffusionOutput(output=None), RuntimeError("prewarm failure")]
-        )
-
-        engine._dummy_run()
-
-        assert "serving will continue" in caplog.text
-
 
 class TestStepScheduler:
     def setup_method(self) -> None:

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,
+            "extras": {"ltx2_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["extras"]["ltx2_use_channels_last_3d"] is True
     assert engine_args["parallel_config"].ulysses_degree == 2
     assert engine_args["model_stage"] == "diffusion"
 
@@ -87,7 +89,6 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
             "enforce_eager": None,
             "diffusion_compile_granularity": None,
             "diffusion_compile_dynamic": None,
-            "diffusion_compile_vae": None,
         }
     )[0]
 
@@ -99,7 +100,6 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
     assert stage_cfg["engine_args"]["enforce_eager"] is False
     assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "regional"
     assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is True
-    assert stage_cfg["engine_args"]["diffusion_compile_vae"] is True
 
 
 def test_default_stage_config_propagates_ulysses_mode():
@@ -317,7 +317,7 @@ def test_serve_cli_forwards_distributed_offload_residency():
 
 
 def test_serve_cli_accepts_diffusion_compile_controls():
-    """Ensure all compile controls reach the diffusion stage."""
+    """Ensure both compile controls reach the diffusion stage."""
     parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
     OmniServeCommand().subparser_init(subparsers)
@@ -330,7 +330,6 @@ def test_serve_cli_accepts_diffusion_compile_controls():
             "--diffusion-compile-granularity",
             "full",
             "--no-diffusion-compile-dynamic",
-            "--no-diffusion-compile-vae",
         ]
     )
 
@@ -339,10 +338,8 @@ def test_serve_cli_accepts_diffusion_compile_controls():
 
     assert args.diffusion_compile_granularity == "full"
     assert args.diffusion_compile_dynamic is False
-    assert args.diffusion_compile_vae is False
     assert stage_cfg["engine_args"]["diffusion_compile_granularity"] == "full"
     assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is False
-    assert stage_cfg["engine_args"]["diffusion_compile_vae"] is False
 
 
 def test_serve_cli_accepts_diffusion_attention_backend():

vllm_omni/config/omni_config.py

@@ -570,7 +570,6 @@ class _DiffusionConfigProjection:
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
     diffusion_compile_dynamic: bool = Field(default=True, strict=True)
-    diffusion_compile_vae: bool = Field(default=True, strict=True)
     fa_deterministic: bool = False
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False

vllm_omni/config/stage_config.py

@@ -394,7 +394,6 @@ class StageDeployConfig:
     # Diffusion execution, cache, and VAE behavior.
     diffusion_compile_granularity: str | None = None
     diffusion_compile_dynamic: bool | None = None
-    diffusion_compile_vae: bool | None = None
     fa_deterministic: bool | None = None
     cache_backend: str | None = None
     cache_config: dict[str, Any] | None = None

vllm_omni/diffusion/compile.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Callable
 from typing import Any
 
 import torch
@@ -10,59 +9,6 @@ from vllm.logger import init_logger
 logger = init_logger(__name__)
 
 
-class _FallbackCompiledCallable:
-    """Run a compiled callable until it fails, then permanently use eager."""
-
-    def __init__(
-        self,
-        eager: Callable[..., Any],
-        compiled: Callable[..., Any],
-        *,
-        label: str,
-    ) -> None:
-        self._eager = eager
-        self._compiled: Callable[..., Any] | None = compiled
-        self._label = label
-
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        compiled = self._compiled
-        if compiled is None:
-            return self._eager(*args, **kwargs)
-        try:
-            return compiled(*args, **kwargs)
-        except Exception as exc:
-            self._compiled = None
-            logger.warning(
-                "%s failed during lazy torch.compile execution (%s); disabling compilation and retrying eagerly.",
-                self._label,
-                exc,
-            )
-            return self._eager(*args, **kwargs)
-
-
-def compile_callable_with_fallback(
-    fn: Callable[..., Any],
-    *,
-    label: str,
-    **compile_kwargs: Any,
-) -> tuple[Callable[..., Any], bool]:
-    """Compile ``fn`` with a fail-closed eager fallback.
-
-    The boolean return value reports whether lazy compilation was activated.
-    Synchronous setup failures return the original callable unchanged.
-    """
-    try:
-        compiled = torch.compile(fn, **compile_kwargs)
-    except Exception as exc:
-        logger.warning(
-            "%s torch.compile setup failed (%s); continuing eagerly.",
-            label,
-            exc,
-        )
-        return fn, False
-    return _FallbackCompiledCallable(fn, compiled, label=label), True
-
-
 def _matches_repeated_block(
     name: str,
     module: nn.Module,

vllm_omni/diffusion/data.py

@@ -726,9 +726,6 @@ class OmniDiffusionConfig:
     # provide its own setup_compile() implementation.
     diffusion_compile_granularity: str = "regional"
     diffusion_compile_dynamic: bool = True
-    # Compile discovered VAE ``decode`` callables whenever generic diffusion
-    # torch.compile is enabled. Runtime failures permanently fall back to eager.
-    diffusion_compile_vae: bool = True
 
     # Parallel weight loading (for faster diffusion model startup)
     enable_multithread_weight_load: bool = True
@@ -933,8 +930,6 @@ 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.diffusion_compile_vae, bool):
-            raise TypeError(f"diffusion_compile_vae must be a bool, got {type(self.diffusion_compile_vae)!r}")
         self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
 
         if self.omni_kv_config is None:

vllm_omni/diffusion/diffusion_engine.py

@@ -56,7 +56,6 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 _ASYNC_OUTPUT_TIMEOUT = 30.0  # seconds
-_VAE_COMPILE_WARMUP_REQUEST_ID = "vae_compile_warmup_req_id"
 
 __all__ = [
     "DiffusionEngine",
@@ -852,26 +851,28 @@ class DiffusionEngine:
         num_inference_steps = 1
         height = 512
         width = 512
+        prompt: OmniTextPrompt = {"prompt": "dummy run"}
 
         supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config)
+        if supports_image_input:
+            # Provide a dummy image input if the model supports it
+            color_format = image_color_format(self.od_config.model_class_name)
+            dummy_image = PIL.Image.new(color_format, (width, height))
+            prompt.setdefault("multi_modal_data", {})["image"] = dummy_image
+
+        if supports_audio_input:
+            audio_sr = 16000
+            dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
+            prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio
+
         num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input)
         if num_frames <= 0:
             logger.info("Skipping dummy warmup run (num_frames=0)")
             return
-
-        def build_request(request_id: str, prompt_text: str, *, real_path: bool) -> OmniDiffusionRequest:
-            prompt: OmniTextPrompt = {"prompt": prompt_text}
-            if supports_image_input:
-                color_format = image_color_format(self.od_config.model_class_name)
-                prompt.setdefault("multi_modal_data", {})["image"] = PIL.Image.new(
-                    color_format,
-                    (width, height),
-                )
-            if supports_audio_input:
-                audio_sr = 16000
-                prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32)
-
-            sampling_params = OmniDiffusionSamplingParams(
+        req = OmniDiffusionRequest(
+            prompt=prompt,
+            request_id=DUMMY_DIFFUSION_REQUEST_ID,
+            sampling_params=OmniDiffusionSamplingParams(
                 height=height,
                 width=width,
                 num_inference_steps=num_inference_steps,
@@ -884,47 +885,13 @@ class DiffusionEngine:
                 # Disable CFG for warmup to avoid triggering CFG parallel
                 # validation when cfg_parallel_size > 1.
                 extra_args={"cfg_text_scale": 1.0, "cfg_img_scale": 1.0},
-            )
-            if real_path:
-                # Internal pretraffic must exercise ordinary pipeline branches
-                # without waiting for KV produced by an upstream stage.
-                sampling_params.need_kv_receive = False
-            return OmniDiffusionRequest(
-                prompt=prompt,
-                request_id=request_id,
-                sampling_params=sampling_params,
-            )
-
-        req = build_request(DUMMY_DIFFUSION_REQUEST_ID, "dummy run", real_path=False)
+            ),
+        )
         logger.info("dummy run to warm up the model")
         output = self.add_req_and_wait_for_response(req)
         if output.error:
             raise RuntimeError(f"Dummy run failed: {output.error}")
 
-        compile_real_path = bool(
-            getattr(self.od_config, "diffusion_compile_vae", False)
-            and not getattr(self.od_config, "enforce_eager", True)
-        )
-        if not compile_real_path:
-            return
-
-        logger.info("real request-path warmup for compiled VAE decode")
-        real_path_req = build_request(
-            _VAE_COMPILE_WARMUP_REQUEST_ID,
-            "A detailed image.",
-            real_path=True,
-        )
-        try:
-            real_path_output = self.add_req_and_wait_for_response(real_path_req)
-        except Exception as exc:
-            logger.warning("Real request-path compile warmup failed (%s); serving will continue.", exc)
-            return
-        if real_path_output.error:
-            logger.warning(
-                "Real request-path compile warmup returned an error (%s); serving will continue.",
-                real_path_output.error,
-            )
-
     def _submit_rpc(
         self,
         method: str,

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

@@ -57,6 +57,17 @@ _LTX_COMPONENT_SUBFOLDERS = (
 )
 logger = logging.getLogger(__name__)
 
+_LTX2_CHANNELS_LAST_3D_EXTRA = "ltx2_use_channels_last_3d"
+
+
+def _ltx2_use_channels_last_3d(od_config: Any) -> bool:
+    """Read the LTX2-only channels-last opt-in from model extras."""
+    extras = getattr(od_config, "extras", {}) or {}
+    enabled = extras.get(_LTX2_CHANNELS_LAST_3D_EXTRA, False)
+    if not isinstance(enabled, bool):
+        raise TypeError(f"{_LTX2_CHANNELS_LAST_3D_EXTRA} must be a bool, got {type(enabled)!r}")
+    return enabled
+
 
 @dataclass(frozen=True)
 class LTXComponentProfile:
@@ -365,6 +376,8 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         local_files_only=local_files_only,
         dtype=dtype,
     )
+    if _ltx2_use_channels_last_3d(od_config):
+        pipeline.vae.enable_channels_last_3d()
     pipeline.audio_vae = _load_component(
         AutoencoderKLLTX2Audio,
         model,

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -28,7 +28,7 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import (
     resolve_prompt_embed_cache_config,
 )
 from vllm_omni.diffusion.cache.selector import get_cache_backend
-from vllm_omni.diffusion.compile import compile_callable_with_fallback, regionally_compile
+from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -40,7 +40,6 @@ from vllm_omni.diffusion.models.interface import (
     supports_step_execution,
 )
 from vllm_omni.diffusion.offloader import get_offload_backend
-from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob
@@ -189,39 +188,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             compile_dynamic,
         )
 
-    def _compile_vae_decodes(self) -> None:
-        """Compile declared VAE decode callables with eager fallback."""
-        assert self.pipeline is not None
-        modules = ModuleDiscovery.discover(self.pipeline)
-        compiled_count = 0
-        for name, vae in zip(modules.vae_names, modules.vaes, strict=True):
-            decode = getattr(vae, "decode", None)
-            if not callable(decode):
-                logger.debug("Model runner: VAE component %s has no decode callable; skipping compile.", name)
-                continue
-
-            label = f"VAE component {name}.decode"
-            compiled_decode, activated = compile_callable_with_fallback(
-                decode,
-                label=label,
-                mode="default",
-                fullgraph=False,
-                dynamic=None,
-            )
-            if not activated:
-                continue
-            setattr(vae, "decode", compiled_decode)
-            compiled_count += 1
-
-        if compiled_count:
-            logger.info(
-                "Model runner: configured lazy torch.compile for %d VAE decode callable(s); "
-                "startup warmup will compile the real request path.",
-                compiled_count,
-            )
-        else:
-            logger.info("Model runner: no compile-compatible VAE decode callable was discovered.")
-
     def load_model(
         self,
         memory_pool_context_fn: Callable[[str], AbstractContextManager[Any]] | None = None,
@@ -311,8 +277,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         # Apply torch.compile if not in eager mode
         if not self.od_config.enforce_eager:
             if current_omni_platform.supports_torch_inductor():
-                pipeline_owns_compile = hasattr(self.pipeline, "setup_compile")
-                if pipeline_owns_compile:
+                if hasattr(self.pipeline, "setup_compile"):
                     try:
                         self.pipeline.setup_compile()
                     except Exception as exc:
@@ -326,9 +291,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                         transformer_attrs = ("transformer", "transformer_2")
                     for attr_name in transformer_attrs:
                         self._compile_transformer(attr_name)
-
-                if self.od_config.diffusion_compile_vae and not pipeline_owns_compile:
-                    self._compile_vae_decodes()
             else:
                 logger.warning(
                     "Model runner: Platform %s does not support torch inductor, skipping torch.compile.",

vllm_omni/engine/arg_utils.py

@@ -485,7 +485,6 @@ class OrchestratorArgs:
     diffusion_attention_config: str | None = None
     diffusion_compile_granularity: str | None = None
     diffusion_compile_dynamic: bool | None = None
-    diffusion_compile_vae: bool | None = None
     cache_backend: str = "none"
     cache_config: str | None = None
     enable_cache_dit_summary: bool = False

vllm_omni/engine/async_omni_engine.py

@@ -1043,9 +1043,6 @@ class AsyncOmniEngine:
             "diffusion_compile_dynamic": (
                 True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"]
             ),
-            "diffusion_compile_vae": (
-                True if kwargs.get("diffusion_compile_vae") is None else kwargs["diffusion_compile_vae"]
-            ),
             "fa_deterministic": bool(kwargs.get("fa_deterministic", False)),
             "boundary_ratio": kwargs.get("boundary_ratio", None),
             "flow_shift": kwargs.get("flow_shift", None),
@@ -1072,6 +1069,7 @@ class AsyncOmniEngine:
             "streaming_output": kwargs.get("diffusion_streaming_output", False),
             "enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
             "extras": {
+                **(kwargs.get("extras") or {}),
                 "auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None),
                 "default_llama_model_id": kwargs.get("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
             },

vllm_omni/entrypoints/cli/serve.py

@@ -446,15 +446,6 @@ class OmniServeCommand(CLISubcommand):
                 "Disable for fixed-shape workloads with --no-diffusion-compile-dynamic."
             ),
         )
-        omni_config_group.add_argument(
-            "--diffusion-compile-vae",
-            action=argparse.BooleanOptionalAction,
-            default=None,
-            help=(
-                "Compile discovered VAE decode callables when diffusion torch.compile is enabled. "
-                "Use --no-diffusion-compile-vae to keep VAE decode eager."
-            ),
-        )
         omni_config_group.add_argument(
             "--fa-deterministic",
             dest="fa_deterministic",
@@ -638,7 +629,6 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).",
         )
-
         # Parallel weight loading (faster diffusion startup)
         omni_config_group.add_argument(
             "--disable-multithread-weight-load",