Fix LTX2 channels-last review scope

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

.pre-commit-config.yaml

@@ -66,7 +66,10 @@ repos:
       entry: python tools/pre_commit/check_tts_adapter.py
       language: python
       pass_filenames: false
-      files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)$
+      files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
+      # The budget-is-now-slack reminder is printed on a PASSING run, and
+      # pre-commit hides output from passing hooks unless the hook is verbose.
+      verbose: true
 
     # Keep `suggestion` last
     - id: suggestion

docs/user_guide/diffusion/attention_backends.md

@@ -283,23 +283,20 @@ blocks — use Ulysses SP (`ring_degree=1`).
 
 ### Which geometries run sparse
 
-Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count
-is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own
-tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows
-that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense
-attention and logs `RAINFUSION_ATTN staying dense` with the row count it computed — but it gets no
-speedup, so pick an aligned geometry when you want one.
+MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
+necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
+does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
+promoted rows remain visible to every query and key.
 
-Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8
-spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact
-patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns
-are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially
-distant positions and selection can no longer rank them meaningfully. This is silent, and invisible
-to a `sparsity=0` check, because a fully populated mask does not care how blocks are grouped.
+For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
+128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
+MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
+and kernel tiling aligned, at the cost of retaining more blocks.
 
-"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned
-resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`,
-which you can buy back with `start_step`.
+8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
+creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
+grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
+that implements protected video tails for rf_v2.
 
 ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
 

docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md

@@ -99,7 +99,7 @@ a native vLLM-Omni pipeline for those cases.
 ### Attention Backends
 
 The diffusers backend converts
-[vLLM-Omni standard of attention backend setting](../../../docs/user_guide/diffusion/attention_backends.md)
+[vLLM-Omni standard of attention backend setting](../../diffusion/attention_backends.md)
 to [diffusers standard](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends).
 
 Specifically for `FLASH_ATTN`, it will first attempt to use FlashAttention-3 and then FlashAttention-2.

tests/config/test_omni_config.py

@@ -776,6 +776,26 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
         omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="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():
     with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
         omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")

tests/diffusion/attention/test_rainfusion_plan.py

@@ -5,10 +5,9 @@
 
 The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip:
 1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row
-kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the
-first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
-mask from the video and prefix segments separately, so the two only agree on the
-block count when the video segment ends on a boundary.
+kernel block, and 1344x768 gives 62x24x42 whose 62496 rows do not. Both are
+handed to MindIE-SD: its rf_v2 preprocessing promotes an irregular video tail
+to the always-kept prefix segment before generating the block mask.
 """
 
 import dataclasses
@@ -57,10 +56,13 @@ def test_block_aligned_video_segment_runs_sparse():
     assert plan.latent_shape == list(ALIGNED_GRID)
 
 
-def test_misaligned_video_segment_stays_dense():
-    # Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask:
-    # pad keys would take a share of every real query's softmax denominator.
-    assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None
+def test_misaligned_video_segment_runs_sparse_with_updated_mindiesd():
+    plan = make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID))
+
+    assert plan is not None
+    assert plan.prefix_len == PREFIX_ROWS
+    assert plan.used_len == PREFIX_ROWS + 62496
+    assert plan.latent_shape == list(MISALIGNED_GRID)
 
 
 @pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
@@ -155,3 +157,28 @@ def test_fully_populated_mask_reproduces_dense_attention():
     ).transpose(1, 2)
     error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
     assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.")
+def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd():
+    """vLLM forwards the full irregular grid; MindIE-SD owns its tail handling."""
+    torch.manual_seed(1)
+    prefix_len, grid = 384, (4, 16, 10)
+    video_len = grid[0] * grid[1] * grid[2]
+    used = prefix_len + video_len
+    heads, head_dim = 4, 128
+    q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
+
+    impl = make_impl()
+    impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
+    plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
+    out = impl._forward_sparse_npu(q, k, v, plan)
+
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        q.transpose(1, 2),
+        k.transpose(1, 2),
+        v.transpose(1, 2),
+        scale=impl.softmax_scale,
+    ).transpose(1, 2)
+    error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
+    assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"

tests/diffusion/cache/test_teacache_extractors.py

@@ -488,6 +488,14 @@ class _MiniMaxH3FakeAttention(nn.Module):
     class FakeBackend:
         supports_prefix_kv_slicing = False
 
+        # Mirror AttentionBackend.supports_packed_mask_free: H3's packed
+        # attention path calls attn_backend.supports_packed_mask_free() to decide
+        # whether to build the padding mask. Return False so the fake exercises
+        # the masked path (the abstract backend default).
+        @classmethod
+        def supports_packed_mask_free(cls) -> bool:
+            return False
+
     def __init__(self, **kwargs):
         del kwargs
         super().__init__()

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

tools/pre_commit/check_tts_adapter.py

@@ -123,18 +123,29 @@ def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]:
     return found
 
 
-def _check(label: str, path: Path, actual: int, budget: int, sample: list[str], budget_name: str) -> list[str]:
+def _check(
+    label: str,
+    path: Path,
+    actual: int,
+    budget: int,
+    sample: list[str],
+    budget_name: str,
+    notices: list[str],
+) -> list[str]:
     if actual > budget:
         listing = "\n".join(f"      {line}" for line in sample)
         return [
             f"{path}: {actual} {label} exceeds the budget of {budget}.\n{listing}\n\n{_GUIDANCE}\n",
         ]
     if actual < budget:
-        return [
-            f"{path}: {actual} {label}, down from a budget of {budget}. "
-            f"Lower {budget_name} to {actual} in {Path(__file__).relative_to(REPO_ROOT)} "
-            f"so the ratchet holds the new ground.",
-        ]
+        # Progress is never a build failure. Removing branches is the goal, so a
+        # PR that does it must not hand-edit this file to stay green -- nor break
+        # `main` for everyone else when it merges from a stale base.
+        notices.append(
+            f"{path}: {actual} {label}, below the budget of {budget}. "
+            f"Lower {budget_name} to {actual} in "
+            f"{Path(__file__).relative_to(REPO_ROOT)} to hold the new ground."
+        )
     return []
 
 
@@ -152,6 +163,7 @@ def main(argv: list[str] | None = None) -> int:
             print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr)
             return 1
 
+    notices: list[str] = []
     branches = _model_type_branches(ast.parse(serving_path.read_text()))
     detectors = _legacy_detectors(ast.parse(adapters_path.read_text()))
 
@@ -162,6 +174,7 @@ def main(argv: list[str] | None = None) -> int:
         MAX_MODEL_TYPE_BRANCHES,
         [f"line {lineno}: {src}" for lineno, src in branches],
         "MAX_MODEL_TYPE_BRANCHES",
+        notices,
     )
     errors += _check(
         "unmigrated model types in LEGACY_TTS_DETECTORS",
@@ -170,8 +183,11 @@ def main(argv: list[str] | None = None) -> int:
         MAX_LEGACY_DETECTORS,
         [f"line {lineno}: {name}" for lineno, name in detectors],
         "MAX_LEGACY_DETECTORS",
+        notices,
     )
 
+    for notice in notices:
+        print(f"check_tts_adapter: {notice}", file=sys.stderr)
     for error in errors:
         print(f"check_tts_adapter: {error}", file=sys.stderr)
     return 1 if errors else 0

vllm_omni/diffusion/attention/backends/rainfusion_attn.py

@@ -128,9 +128,9 @@ class RainFusionAttentionImpl(AttentionImpl):
     whose extent the model publishes as ``AttentionMetadata.video_layout``. Every
     other case — warmup denoise steps, exempt layers, a layer that does not declare
     ``qkv_layout="BSND"``, sequences without a published video segment, video
-    segments too short to pay for block selection or not aligned to the kernel's
-    block size — delegates to FlashAttention, so a model can select this backend
-    unconditionally.
+    segments too short to pay for block selection — delegates to FlashAttention,
+    so a model can select this backend unconditionally. MindIE-SD handles an
+    irregular video tail internally, retaining it outside the sparse blocks.
     """
 
     def __init__(
@@ -292,25 +292,6 @@ class RainFusionAttentionImpl(AttentionImpl):
                 _MIN_VIDEO_BLOCKS,
             )
             return None
-        if video_len % _BLOCK_SIZE != 0:
-            # rf_v2 tiles the sequence into ceil((video + prefix) / 128) blocks but
-            # builds its block mask from the two segments pooled apart. The counts
-            # only agree when the video segment ends on a block boundary; otherwise
-            # a prefix block is forced on out of range, and the block straddling the
-            # seam holds real prefix rows that selection may drop. Padding the
-            # sequence to realign it is not an option because rf_v2 ignores
-            # attn_mask, so pad keys would take a share of every softmax denominator.
-            logger.warning_once(
-                "RAINFUSION_ATTN staying dense: latent grid %s gives %d video rows, which is not a "
-                "multiple of the %d-row kernel block. Choose a geometry whose latent_t * h * w is a "
-                "multiple of %d to enable sparsity.",
-                tuple(latent_shape),
-                video_len,
-                _BLOCK_SIZE,
-                _BLOCK_SIZE,
-            )
-            return None
-
         logger.info_once(
             "RAINFUSION_ATTN active: sparsity=%.2f, start_step=%d, exempt_layers=%d, "
             "latent_grid=%s, prefix_rows=%d, video_rows=%d. Realized sparsity is lower than nominal "

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/engine/async_omni_engine.py

@@ -1069,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

@@ -629,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",