minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-06-09 23:26:40 +08:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 17d8c5801d
commit aa18a68ac5
4 changed files with 223 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
+11 -3
View File
@@ -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
python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py
+100 -25
View File
@@ -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
python/sglang/multimodal_gen/test/unit/test_ltx2_vae_channels_last.py
+93
View File
@@ -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
+19 -21
View File
@@ -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),