From aa18a68ac52cd3d4cb56bac551be4e43fe0d7516 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Tue, 9 Jun 2026 23:26:40 +0800 Subject: [PATCH] [diffusion] Run LTX-2 VAE decode in channels_last_3d (faster decode, lower peak memory) (#27431) Co-authored-by: Claude Opus 4.8 --- .../loader/component_loaders/vae_loader.py | 14 +- .../runtime/models/vaes/ltx_2_vae.py | 125 ++++++++++++++---- .../test/unit/test_ltx2_vae_channels_last.py | 93 +++++++++++++ .../test/unit/test_vae_loader.py | 40 +++--- 4 files changed, 223 insertions(+), 49 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_ltx2_vae_channels_last.py diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 9c21d44d1..0c1f3dec7 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/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__ - if pipeline_name.startswith("QwenImage"): + pipeline_config = server_args.pipeline_config + 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 diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py index 207362e86..cc652a66f 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py +++ b/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 ( - conv_cache is not None - and cache_key is not None - and time_kernel_size > 1 - ): + left, right = time_kernel_size - 1, 0 + if conv_cache is not None and cache_key is not None and left: # 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( - (1, 1, time_kernel_size - 1, 1, 1) - ) + if use_channels_last_pad: + 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) :, :, :] - .detach() - .clone() + hidden_states[:, :, -left:, :, :].detach().clone() ) else: - pad_left = hidden_states[:, :, :1, :, :].repeat( - (1, 1, time_kernel_size - 1, 1, 1) - ) - 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 + ) + 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( - (1, 1, (time_kernel_size - 1) // 2, 1, 1) - ) - pad_right = hidden_states[:, :, -1:, :, :].repeat( - (1, 1, (time_kernel_size - 1) // 2, 1, 1) - ) - hidden_states = torch.concatenate( - [pad_left, hidden_states, pad_right], dim=2 - ) + left = right = (time_kernel_size - 1) // 2 + + if use_channels_last_pad: + hidden_states = self._causal_temporal_pad_channels_last( + hidden_states, left, right + ) + else: + 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 diff --git a/python/sglang/multimodal_gen/test/unit/test_ltx2_vae_channels_last.py b/python/sglang/multimodal_gen/test/unit/test_ltx2_vae_channels_last.py new file mode 100644 index 000000000..1eea62767 --- /dev/null +++ b/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() diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py index df0e9761a..71dda79e5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py +++ b/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),