[diffusion] Run LTX-2 VAE decode in channels_last_3d (faster decode, lower peak memory) (#27431)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py

@@ -6,6 +6,11 @@ import torch.nn as nn
 from safetensors.torch import load_file as safetensors_load_file
 
 from sglang.multimodal_gen.configs.models import ModelConfig
+from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig
+from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
+    QwenImagePipelineConfig,
+)
+from sglang.multimodal_gen.configs.pipeline_configs.wan import WanT2V480PConfig
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
     ComponentLoader,
 )
@@ -76,10 +81,13 @@ def _should_use_channels_last_3d(
     if server_args is None:
         return False
 
-    pipeline_name = server_args.pipeline_config.__class__.__name__
+    pipeline_config = server_args.pipeline_config
-    if pipeline_name.startswith("QwenImage"):
+    if isinstance(pipeline_config, QwenImagePipelineConfig):
         return True
-    if "Wan" in pipeline_name and server_args.num_gpus == 1:
+    if (
+        isinstance(pipeline_config, (WanT2V480PConfig, LTX2PipelineConfig))
+        and server_args.num_gpus == 1
+    ):
         return True
     return False
 

python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py

@@ -1,3 +1,4 @@
+from functools import lru_cache
 from typing import Optional, Tuple, Union
 
 import torch
@@ -14,6 +15,23 @@ from sglang.multimodal_gen.configs.models.vaes.ltx_video import LTXVideoVAEConfi
 from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE
 
 
+@lru_cache(maxsize=128)
+def _is_channels_last_3d_stride(size: tuple[int, ...], stride: tuple[int, ...]) -> bool:
+    if len(size) != 5:
+        return False
+
+    expected_stride = 1
+    for dim in (1, 4, 3, 2, 0):
+        if size[dim] == 0:
+            return True
+        if size[dim] == 1:
+            continue
+        if stride[dim] != expected_stride:
+            return False
+        expected_stride *= size[dim]
+    return True
+
+
 class PerChannelRMSNorm(nn.Module):
     """
     Per-pixel (per-location) RMS normalization layer.
@@ -87,6 +105,42 @@ class LTX2VideoCausalConv3d(nn.Module):
             padding_mode=spatial_padding_mode,
         )
 
+    def _weight_is_channels_last_3d(self) -> bool:
+        w = self.conv.weight
+        return hasattr(torch, "channels_last_3d") and _is_channels_last_3d_stride(
+            tuple(w.size()), tuple(w.stride())
+        )
+
+    def _causal_temporal_pad_channels_last(
+        self,
+        x: torch.Tensor,
+        left: int,
+        right: int,
+        left_pad: Optional[torch.Tensor] = None,
+        right_pad: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # Build the temporally-padded tensor directly in channels_last_3d so the
+        # cuDNN NDHWC conv path is preserved. The single allocate-and-copy_ does
+        # double duty (replication pad + layout fix) and avoids the expensive
+        # contiguous->channels_last reconversion that repeat()+concatenate()
+        # would otherwise force on the decode hot path. Numerically identical to
+        # the repeat/concatenate path (replicate-edge == repeat-edge frame).
+        b, c, t, h, w = x.shape
+        out = torch.empty(
+            (b, c, t + left + right, h, w),
+            dtype=x.dtype,
+            device=x.device,
+            memory_format=torch.channels_last_3d,
+        )
+        out[:, :, left : left + t].copy_(x)
+        if left:
+            out[:, :, :left].copy_(x[:, :, :1] if left_pad is None else left_pad)
+        if right:
+            out[:, :, left + t :].copy_(
+                x[:, :, -1:] if right_pad is None else right_pad
+            )
+        return out
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -95,13 +149,13 @@ class LTX2VideoCausalConv3d(nn.Module):
         cache_key: Optional[str] = None,
     ) -> torch.Tensor:
         time_kernel_size = self.kernel_size[0]
+        use_channels_last_pad = (
+            hidden_states.dim() == 5 and self._weight_is_channels_last_3d()
+        )
 
         if causal:
-            if (
+            left, right = time_kernel_size - 1, 0
-                conv_cache is not None
+            if conv_cache is not None and cache_key is not None and left:
-                and cache_key is not None
-                and time_kernel_size > 1
-            ):
                 # Streaming: prepend the previous chunk's last (k-1) frames of the
                 # PADDED conv input (first chunk: replicate frame 0, the "sink"),
                 # then store the last (k-1) frames of THIS padded input for the next
@@ -110,34 +164,55 @@ class LTX2VideoCausalConv3d(nn.Module):
                 # stride 1 keeps the output T invariant.
                 prev = conv_cache.get(cache_key)
                 if prev is None:
-                    pad_left = hidden_states[:, :, :1, :, :].repeat(
+                    if use_channels_last_pad:
-                        (1, 1, time_kernel_size - 1, 1, 1)
+                        hidden_states = self._causal_temporal_pad_channels_last(
-                    )
+                            hidden_states, left, right
+                        )
+                    else:
+                        pad_left = hidden_states[:, :, :1, :, :].repeat(
+                            (1, 1, left, 1, 1)
+                        )
+                        hidden_states = torch.concatenate(
+                            [pad_left, hidden_states], dim=2
+                        )
                 else:
                     pad_left = prev.to(
                         device=hidden_states.device, dtype=hidden_states.dtype
                     )
-                hidden_states = torch.concatenate([pad_left, hidden_states], dim=2)
+                    if use_channels_last_pad:
+                        hidden_states = self._causal_temporal_pad_channels_last(
+                            hidden_states, left, right, left_pad=pad_left
+                        )
+                    else:
+                        hidden_states = torch.concatenate(
+                            [pad_left, hidden_states], dim=2
+                        )
                 conv_cache[cache_key] = (
-                    hidden_states[:, :, -(time_kernel_size - 1) :, :, :]
+                    hidden_states[:, :, -left:, :, :].detach().clone()
-                    .detach()
-                    .clone()
                 )
             else:
-                pad_left = hidden_states[:, :, :1, :, :].repeat(
+                if use_channels_last_pad:
-                    (1, 1, time_kernel_size - 1, 1, 1)
+                    hidden_states = self._causal_temporal_pad_channels_last(
-                )
+                        hidden_states, left, right
-                hidden_states = torch.concatenate([pad_left, hidden_states], dim=2)
+                    )
+                else:
+                    pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, left, 1, 1))
+                    hidden_states = torch.concatenate([pad_left, hidden_states], dim=2)
         else:
-            pad_left = hidden_states[:, :, :1, :, :].repeat(
+            left = right = (time_kernel_size - 1) // 2
-                (1, 1, (time_kernel_size - 1) // 2, 1, 1)
+
-            )
+            if use_channels_last_pad:
-            pad_right = hidden_states[:, :, -1:, :, :].repeat(
+                hidden_states = self._causal_temporal_pad_channels_last(
-                (1, 1, (time_kernel_size - 1) // 2, 1, 1)
+                    hidden_states, left, right
-            )
+                )
-            hidden_states = torch.concatenate(
+            else:
-                [pad_left, hidden_states, pad_right], dim=2
+                pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, left, 1, 1))
-            )
+                parts = [pad_left, hidden_states]
+                if right:
+                    parts.append(
+                        hidden_states[:, :, -1:, :, :].repeat((1, 1, right, 1, 1))
+                    )
+                hidden_states = torch.concatenate(parts, dim=2)
 
         hidden_states = self.conv(hidden_states)
         return hidden_states

python/sglang/multimodal_gen/test/unit/test_ltx2_vae_channels_last.py

@@ -0,0 +1,93 @@
+import unittest
+
+import torch
+
+from sglang.multimodal_gen.runtime.models.vaes.ltx_2_vae import LTX2VideoCausalConv3d
+
+
+@unittest.skipUnless(
+    hasattr(torch, "channels_last_3d"), "channels_last_3d is unavailable"
+)
+class TestLTX2CausalConvChannelsLast(unittest.TestCase):
+    """The channels_last_3d causal-conv path must stay numerically identical to
+    the original repeat()+concatenate() temporal padding (only the conv kernel's
+    floating-point accumulation order may differ, which fp32 keeps negligible)."""
+
+    def _check(self, causal: bool, kernel_size):
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        torch.manual_seed(0)
+        conv = LTX2VideoCausalConv3d(8, 8, kernel_size).to(device, torch.float32).eval()
+        x = torch.randn(1, 8, 5, 6, 7, dtype=torch.float32, device=device)
+
+        # Default (contiguous) weight -> original repeat/concat padding branch.
+        self.assertFalse(conv._weight_is_channels_last_3d())
+        with torch.no_grad():
+            y_ref = conv(x.clone(), causal=causal)
+
+        # channels_last_3d weight -> layout-preserving pad branch.
+        conv.conv.weight.data = conv.conv.weight.data.to(
+            memory_format=torch.channels_last_3d
+        )
+        self.assertTrue(conv._weight_is_channels_last_3d())
+        with torch.no_grad():
+            y_cl = conv(x.clone(), causal=causal)
+
+        self.assertEqual(y_ref.shape, y_cl.shape)
+        # On CUDA, cuDNN preserves channels_last_3d through the conv; CPU conv3d
+        # returns a contiguous tensor regardless. The pad output layout itself is
+        # asserted device-independently in test_pad_replicates_edge_frames_exactly.
+        if device == "cuda":
+            self.assertTrue(y_cl.is_contiguous(memory_format=torch.channels_last_3d))
+        # Equal weights in either memory format -> equal math; only the conv
+        # kernel's float accumulation order may differ (negligible in fp32).
+        torch.testing.assert_close(y_ref, y_cl, rtol=1e-4, atol=1e-4)
+
+    def test_causal_matches_reference(self):
+        self._check(causal=True, kernel_size=3)
+
+    def test_non_causal_matches_reference(self):
+        self._check(causal=False, kernel_size=3)
+
+    def test_temporal_only_kernel_matches_reference(self):
+        self._check(causal=True, kernel_size=(3, 1, 1))
+
+    def test_causal_cache_matches_monolithic_reference(self):
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        torch.manual_seed(0)
+        conv = LTX2VideoCausalConv3d(4, 5, 3).to(device, torch.float32).eval()
+        conv.conv.weight.data = conv.conv.weight.data.to(
+            memory_format=torch.channels_last_3d
+        )
+        self.assertTrue(conv._weight_is_channels_last_3d())
+        x = torch.randn(1, 4, 6, 3, 3, dtype=torch.float32, device=device)
+
+        with torch.no_grad():
+            whole = conv(x, causal=True)
+            cache = {}
+            parts = [
+                conv(chunk, causal=True, conv_cache=cache, cache_key="conv")
+                for chunk in (x[:, :, :2], x[:, :, 2:])
+            ]
+            chunked = torch.cat(parts, dim=2)
+
+        self.assertEqual(chunked.shape, whole.shape)
+        torch.testing.assert_close(whole, chunked, rtol=1e-4, atol=1e-4)
+
+    def test_pad_replicates_edge_frames_exactly(self):
+        # The temporal pad must replicate the first (and, when non-causal, last)
+        # frame exactly -- assert against an explicit reference construction.
+        conv = LTX2VideoCausalConv3d(4, 4, 3).to(torch.float32)
+        conv.conv.weight.data = conv.conv.weight.data.to(
+            memory_format=torch.channels_last_3d
+        )
+        x = torch.randn(1, 4, 3, 2, 2, dtype=torch.float32)
+        padded = conv._causal_temporal_pad_channels_last(x, left=2, right=1)
+        expected = torch.cat(
+            [x[:, :, :1].repeat(1, 1, 2, 1, 1), x, x[:, :, -1:]], dim=2
+        )
+        self.assertTrue(padded.is_contiguous(memory_format=torch.channels_last_3d))
+        torch.testing.assert_close(padded, expected, rtol=0, atol=0)
+
+
+if __name__ == "__main__":
+    unittest.main()

python/sglang/multimodal_gen/test/unit/test_vae_loader.py

@@ -3,6 +3,15 @@ from unittest.mock import patch
 
 import torch
 
+from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig
+from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
+    QwenImagePipelineConfig,
+)
+from sglang.multimodal_gen.configs.pipeline_configs.wan import (
+    FastWan2_2_TI2V_5B_Config,
+    Wan2_2_I2V_A14B_Config,
+    WanT2V480PConfig,
+)
 from sglang.multimodal_gen.runtime.loader.component_loaders import vae_loader
 from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import (
     _backfill_ltx2_audio_vae_latent_stats,
@@ -17,26 +26,6 @@ class _FakeServerArgs:
         self.num_gpus = num_gpus
 
 
-class QwenImagePipelineConfig:
-    pass
-
-
-class WanT2V480PConfig:
-    pass
-
-
-class FastWan2_2_TI2V_5B_Config:
-    pass
-
-
-class Wan2_2_I2V_A14B_Config:
-    pass
-
-
-class LTX2PipelineConfig:
-    pass
-
-
 class TestVAELoader(unittest.TestCase):
     def test_backfill_ltx2_audio_vae_latent_stats_maps_official_keys(self):
         loaded = {
@@ -109,7 +98,16 @@ class TestVAELoader(unittest.TestCase):
             server_args = _FakeServerArgs(Wan2_2_I2V_A14B_Config(), num_gpus=2)
             self.assertFalse(_should_use_channels_last_3d(server_args, "video_vae"))
 
-    def test_channels_last_3d_defaults_false_for_ltx_on_cuda(self):
+    def test_channels_last_3d_defaults_true_for_single_gpu_ltx_on_cuda(self):
+        with (
+            patch.dict("os.environ", {}, clear=True),
+            patch.object(vae_loader.current_platform, "is_cuda", return_value=True),
+            patch.object(vae_loader.current_platform, "is_rocm", return_value=False),
+        ):
+            server_args = _FakeServerArgs(LTX2PipelineConfig(), num_gpus=1)
+            self.assertTrue(_should_use_channels_last_3d(server_args, "video_vae"))
+
+    def test_channels_last_3d_defaults_false_for_multi_gpu_ltx_on_cuda(self):
         with (
             patch.dict("os.environ", {}, clear=True),
             patch.object(vae_loader.current_platform, "is_cuda", return_value=True),