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

[diffusion] fix: further align ltx2.3 accuracy with tp (#24660)

Browse Source
This commit is contained in:
Mick
2026-05-11 13:42:08 +08:00
committed by GitHub
parent ed70226ec1
commit 6e5b4de01a
16 changed files with 435 additions and 311 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py
+19 -2
View File
@@ -2,7 +2,10 @@
# SPDX-License-Identifier: Apache-2.0
from contextlib import nullcontext
import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( # FlashAttentionMetadata,
AttentionBackend,
@@ -14,6 +17,13 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS = [
SDPBackend.CUDNN_ATTENTION,
SDPBackend.FLASH_ATTENTION,
SDPBackend.EFFICIENT_ATTENTION,
SDPBackend.MATH,
]
class SDPABackend(AttentionBackend):
@@ -51,6 +61,7 @@ class SDPAImpl(AttentionImpl):
self.causal = causal
self.softmax_scale = softmax_scale
self.dropout = extra_impl_args.get("dropout_p", 0.0)
self.allow_cudnn_sdp = bool(extra_impl_args.get("allow_cudnn_sdp", False))
def forward(
self,
@@ -71,8 +82,14 @@ class SDPAImpl(AttentionImpl):
}
if query.shape[1] != key.shape[1]:
attn_kwargs["enable_gqa"] = True
output = torch.nn.functional.scaled_dot_product_attention(
query, key, value, **attn_kwargs
sdpa_context = (
sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
if self.allow_cudnn_sdp and query.device.type == "cuda"
else nullcontext()
)
with sdpa_context:
output = torch.nn.functional.scaled_dot_product_attention(
query, key, value, **attn_kwargs
)
output = output.transpose(1, 2)
return output
python/sglang/multimodal_gen/runtime/layers/attention/layer.py
+56 -27
View File
@@ -1,10 +1,12 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from contextlib import nullcontext
from typing import Type
import torch
import torch.nn as nn
from torch.nn.attention import SDPBackend, sdpa_kernel
from sglang.multimodal_gen.runtime.distributed.communication_op import (
sequence_model_parallel_all_gather,
@@ -35,6 +37,13 @@ from sglang.multimodal_gen.runtime.managers.forward_context import (
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
from sglang.multimodal_gen.utils import get_compute_dtype
_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS = [
SDPBackend.CUDNN_ATTENTION,
SDPBackend.FLASH_ATTENTION,
SDPBackend.EFFICIENT_ATTENTION,
SDPBackend.MATH,
]
class UlyssesAttention(nn.Module):
"""Ulysses-style SequenceParallelism attention layer."""
@@ -246,6 +255,7 @@ class LocalAttention(nn.Module):
head_size, dtype, supported_attention_backends=supported_attention_backends
)
impl_cls = attn_backend.get_impl_cls()
self.allow_cudnn_sdp = bool(extra_impl_args.get("allow_cudnn_sdp", False))
self.attn_impl = impl_cls(
num_heads=num_heads,
head_size=head_size,
@@ -304,15 +314,21 @@ class LocalAttention(nn.Module):
mask = mask[:, None, :, :]
mask = (mask - 1.0) * torch.finfo(q_.dtype).max
return torch.nn.functional.scaled_dot_product_attention(
q_,
k_,
v_,
attn_mask=mask,
dropout_p=0.0,
is_causal=False,
scale=self.softmax_scale,
).transpose(1, 2)
sdpa_context = (
sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
if self.allow_cudnn_sdp and q_.device.type == "cuda"
else nullcontext()
)
with sdpa_context:
return torch.nn.functional.scaled_dot_product_attention(
q_,
k_,
v_,
attn_mask=mask,
dropout_p=0.0,
is_causal=False,
scale=self.softmax_scale,
).transpose(1, 2)
output = self.attn_impl.forward(q, k, v, attn_metadata=ctx_attn_metadata)
return output
@@ -373,6 +389,7 @@ class USPAttention(nn.Module):
f"Please ensure your platform supports these backends."
)
impl_cls: Type["AttentionImpl"] = attn_backend.get_impl_cls()
self.allow_cudnn_sdp = bool(extra_impl_args.get("allow_cudnn_sdp", False))
self.attn_impl = impl_cls(
num_heads=num_heads,
head_size=head_size,
@@ -453,15 +470,21 @@ class USPAttention(nn.Module):
k_ = k.transpose(1, 2)
v_ = v.transpose(1, 2)
mask = _prepare_sdpa_mask(attn_mask, dtype=q_.dtype, device=q_.device)
return torch.nn.functional.scaled_dot_product_attention(
q_,
k_,
v_,
attn_mask=mask,
dropout_p=0.0,
is_causal=False,
scale=self.softmax_scale,
).transpose(1, 2)
sdpa_context = (
sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
if self.allow_cudnn_sdp and q_.device.type == "cuda"
else nullcontext()
)
with sdpa_context:
return torch.nn.functional.scaled_dot_product_attention(
q_,
k_,
v_,
attn_mask=mask,
dropout_p=0.0,
is_causal=False,
scale=self.softmax_scale,
).transpose(1, 2)
if get_ring_parallel_world_size() > 1:
raise NotImplementedError(
@@ -489,15 +512,21 @@ class USPAttention(nn.Module):
k_ = k.transpose(1, 2)
v_ = v.transpose(1, 2)
mask = _prepare_sdpa_mask(gathered_mask, dtype=q_.dtype, device=q_.device)
out = torch.nn.functional.scaled_dot_product_attention(
q_,
k_,
v_,
attn_mask=mask,
dropout_p=0.0,
is_causal=False,
scale=self.softmax_scale,
).transpose(1, 2)
sdpa_context = (
sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
if self.allow_cudnn_sdp and q_.device.type == "cuda"
else nullcontext()
)
with sdpa_context:
out = torch.nn.functional.scaled_dot_product_attention(
q_,
k_,
v_,
attn_mask=mask,
dropout_p=0.0,
is_causal=False,
scale=self.softmax_scale,
).transpose(1, 2)
if sp_size > 1:
out = _usp_output_all_to_all(out, head_dim=2)
return out
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
+29 -32
View File
@@ -4,11 +4,8 @@
from __future__ import annotations
import functools
import math
from typing import Any, Optional, Tuple, Union
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
@@ -98,17 +95,6 @@ def _ltx2_build_batched_perturbation_states(
return states
@functools.lru_cache(maxsize=5)
def _ltx2_rope_freq_grid_np(theta: float, num_pos_dims: int, dim: int) -> torch.Tensor:
# Official LTX uses NumPy float64 for double-precision RoPE frequencies.
n_elem = 2 * num_pos_dims
pow_indices = np.power(
theta,
np.linspace(0.0, 1.0, dim // n_elem, dtype=np.float64),
)
return torch.tensor(pow_indices * math.pi / 2.0, dtype=torch.float32)
def apply_interleaved_rotary_emb(
x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor]
) -> torch.Tensor:
@@ -351,22 +337,20 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module):
).to(device)
num_rope_elems = num_pos_dims * 2
if self.double_precision:
freqs = _ltx2_rope_freq_grid_np(self.theta, num_pos_dims, self.dim).to(
device=device
)
else:
pow_indices = torch.pow(
self.theta,
torch.linspace(
start=0.0,
end=1.0,
steps=self.dim // num_rope_elems,
dtype=torch.float32,
device=device,
),
)
freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32)
# LTX-2.3 HQ is sensitive to RoPE rounding; keep frequency generation on
# the target device instead of caching a CPU/NumPy tensor.
freqs_dtype = torch.float64 if self.double_precision else torch.float32
pow_indices = torch.pow(
self.theta,
torch.linspace(
start=0.0,
end=1.0,
steps=self.dim // num_rope_elems,
dtype=freqs_dtype,
device=device,
),
)
freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32)
freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs
freqs = freqs.transpose(-1, -2).flatten(2)
@@ -647,6 +631,8 @@ class LTX2Attention(nn.Module):
causal=False,
supported_attention_backends=supported_attention_backends,
prefix=f"{prefix}.attn",
# official LTX2 torch_sdpa uses cuDNN; cuda setup disables it
allow_cudnn_sdp=True,
)
else:
self.attn = USPAttention(
@@ -658,6 +644,8 @@ class LTX2Attention(nn.Module):
causal=False,
supported_attention_backends=supported_attention_backends,
prefix=f"{prefix}.attn",
# official LTX2 torch_sdpa uses cuDNN; cuda setup disables it
allow_cudnn_sdp=True,
)
def forward(
@@ -1422,8 +1410,17 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
if hasattr(arch.rope_type, "value")
else str(arch.rope_type)
)
rope_double_precision = bool(
hf_config.get("rope_double_precision", arch.double_precision_rope)
frequencies_precision = hf_config.get("frequencies_precision")
if frequencies_precision is None:
frequencies_precision = getattr(arch, "frequencies_precision", None)
# diffusers/LTX configs use `frequencies_precision` for this RoPE switch
rope_double_precision = (
str(frequencies_precision) == "float64"
if frequencies_precision is not None
else bool(
hf_config.get("rope_double_precision", arch.double_precision_rope)
)
)
self.quantize_video_rope_coords_to_hidden_dtype = bool(
hf_config.get("quantize_video_rope_coords_to_hidden_dtype", False)
python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py
+29 -35
View File
@@ -146,14 +146,21 @@ class Gemma3Attention(nn.Module):
prefix=f"{prefix}.o_proj",
)
self.layer_type = (
config.text_config.layer_types[layer_id]
if hasattr(config.text_config, "layer_types")
else None
)
self.is_sliding = (
config.text_config.layer_types[layer_id] == "sliding_attention"
)
layer_types = getattr(config.text_config, "layer_types", None)
if layer_types:
self.layer_type = layer_types[layer_id]
self.is_sliding = self.layer_type == "sliding_attention"
else:
# official Gemma3 uses sliding_window_pattern when layer_types is absent
sliding_window_pattern = getattr(
config.text_config, "sliding_window_pattern", None
)
self.is_sliding = (
bool((layer_id + 1) % sliding_window_pattern)
if sliding_window_pattern
else False
)
self.layer_type = "sliding_attention" if self.is_sliding else None
rope_parameters = getattr(config.text_config, "rope_parameters", None) or {}
layer_rope_params = {}
@@ -204,7 +211,7 @@ class Gemma3Attention(nn.Module):
self.sliding_window = None
self.window_size = (-1, -1)
self.rotary_emb = get_rope(
self.rotary_pos_emb = get_rope(
self.head_dim,
rotary_dim=self.head_dim,
max_position=config.text_config.max_position_embeddings,
@@ -213,12 +220,6 @@ class Gemma3Attention(nn.Module):
is_neox_style=True,
)
self.rope_scaling_factor = (
float(rope_scaling["factor"])
if rope_scaling and rope_scaling.get("factor")
else None
)
# Local Attention not support attention mask, we use global attention instead.
# self.attn = LocalAttention(
# self.num_heads,
@@ -238,26 +239,19 @@ class Gemma3Attention(nn.Module):
dim=self.head_dim, eps=config.text_config.rms_norm_eps
)
def rotary_emb(self, positions, q, k):
"""Apply RoPE using the same device-side inv_freq materialization as LTX."""
positions_flat = positions.flatten().float()
def _apply_rotary_pos_emb(self, positions, q, k):
positions_flat = positions.flatten().to(
device=self.rotary_pos_emb.cos_sin_cache.device, dtype=torch.long
)
cos_sin = self.rotary_pos_emb.cos_sin_cache.index_select(0, positions_flat)
cos, sin = cos_sin.chunk(2, dim=-1)
# match HF Gemma3: expand half-dim freqs to full head dim before rotate_half
cos = torch.cat((cos, cos), dim=-1).to(device=q.device, dtype=q.dtype)
sin = torch.cat((sin, sin), dim=-1).to(device=q.device, dtype=q.dtype)
cos = cos.unsqueeze(1)
sin = sin.unsqueeze(1)
num_tokens = positions_flat.shape[0]
with torch.autocast(device_type=q.device.type, enabled=False):
freq_indices = (
torch.arange(
0, self.head_dim, 2, dtype=torch.int64, device=q.device
).float()
/ self.head_dim
)
inv_freq = 1.0 / (self.rope_theta**freq_indices)
if self.rope_scaling_factor is not None:
inv_freq = inv_freq / self.rope_scaling_factor
freqs = torch.outer(positions_flat, inv_freq)
emb = freqs.repeat(1, 2)
cos = emb.cos().to(q.dtype).unsqueeze(1)
sin = emb.sin().to(q.dtype).unsqueeze(1)
q = q.reshape(num_tokens, -1, self.head_dim)
k = k.reshape(num_tokens, -1, self.head_dim)
q = q * cos + _rotate_half(q) * sin
@@ -283,7 +277,7 @@ class Gemma3Attention(nn.Module):
k = self.k_norm(k)
# Apply RoPE
q, k = self.rotary_emb(positions, q, k)
q, k = self._apply_rotary_pos_emb(positions, q, k)
q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
@@ -306,7 +300,7 @@ class Gemma3Attention(nn.Module):
attn_mask = attn_mask.masked_fill(causal, False)
if self.is_sliding and self.sliding_window is not None:
idx = torch.arange(seq_len, device=hidden_states.device)
dist = idx[None, :] - idx[:, None]
dist = idx[:, None] - idx[None, :]
too_far = dist > self.sliding_window
attn_mask = attn_mask.masked_fill(too_far, False)
python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py
-1
View File
@@ -193,7 +193,6 @@ class LTX2SigmaPreparationStage(PipelineStage):
int(batch.num_inference_steps),
number_of_tokens=latent_num_frames * latent_height * latent_width,
)
batch.sigmas.append(0.0011)
else:
batch.sigmas = build_official_ltx2_sigmas(
int(batch.num_inference_steps)
python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py
+5 -1
View File
@@ -82,8 +82,12 @@ class ParallelExecutor(PipelineExecutor):
elif paradigm == StageParallelismType.CFG_PARALLEL:
obj_list = [batch] if rank == 0 else []
# `dist.broadcast(src=...)` expects a global rank for process groups.
broadcasted_list = broadcast_pyobj(
obj_list, rank=rank, dist_group=cfg_group.cpu_group, src=0
obj_list,
rank=get_world_rank(),
dist_group=cfg_group.cpu_group,
src=cfg_group.ranks[0],
)
if rank != 0:
batch = broadcasted_list[0]
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py
+6 -6
View File
@@ -354,12 +354,11 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
scheduler = clone_scheduler_runtime(original_batch_scheduler or self.scheduler)
distilled_device = scheduler.sigmas.device
num_steps = len(self.distilled_sigmas) - 1
# Inject `0.0011` before the terminal `0.0` to avoid the
# `sigma_next==0` singularity in res2s' `(sample - denoised) /
# (sigma - sigma_next)`. Official `res2s_denoising_loop` does this
# exact injection (samplers.py:262); official `euler_denoising_loop`
# does NOT — it uses `sigma_next` directly. So gate on the active
# sampler, not on the model variant.
# (sigma - sigma_next)`. This changes the final sigma pair only; it
# must not add an extra denoising timestep.
if self.sampler_name == "res2s" and self.distilled_sigmas[-1].item() == 0.0:
scheduler_sigmas = torch.cat(
[
@@ -372,9 +371,10 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
scheduler_sigmas = self.distilled_sigmas
scheduler.sigmas = scheduler_sigmas
num_steps = len(scheduler_sigmas) - 1
scheduler.num_inference_steps = num_steps
scheduler.timesteps = (scheduler_sigmas[:num_steps] * 1000).to(distilled_device)
scheduler.timesteps = (self.distilled_sigmas[:num_steps] * 1000).to(
distilled_device
)
scheduler._step_index = None
scheduler._begin_index = None
python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py
+141 -121
View File
@@ -422,45 +422,76 @@ class LTX2DenoisingStage(DenoisingStage):
return pred * factor
@classmethod
def _ltx2_combine_guided_x0_parallel(
def _ltx2_combine_guided_x0_parallel_av(
cls,
*,
latents: torch.Tensor,
local_velocities: dict[str, torch.Tensor],
sigma: float | torch.Tensor,
cfg_scale: float,
stg_scale: float,
rescale_scale: float,
modality_scale: float,
) -> torch.Tensor:
"""Combine stage-1 guidance passes that were split across CFG ranks.
video_latents: torch.Tensor,
audio_latents: torch.Tensor,
local_video_velocities: dict[str, torch.Tensor],
local_audio_velocities: dict[str, torch.Tensor],
video_sigma: float | torch.Tensor,
audio_sigma: float | torch.Tensor,
video_cfg_scale: float,
video_stg_scale: float,
video_rescale_scale: float,
video_modality_scale: float,
audio_cfg_scale: float,
audio_stg_scale: float,
audio_rescale_scale: float,
audio_modality_scale: float,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Reconstruct CFG branches once for both modalities before guider math."""
first_video_velocity = next(iter(local_video_velocities.values()))
first_audio_velocity = next(iter(local_audio_velocities.values()))
video_template = cls._ltx2_velocity_to_x0(
video_latents, first_video_velocity, video_sigma
)
audio_template = cls._ltx2_velocity_to_x0(
audio_latents, first_audio_velocity, audio_sigma
)
video_numel = video_template.numel()
branches: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
Each pass is one model forward with a different conditioning setup:
positive prompt, negative prompt, attention-disabled perturbation, or
audio/video cross-attention disabled. A rank only owns some passes, so
it contributes weighted x0 terms for those passes and all-reduce
reconstructs the full guided x0 on every rank.
"""
coefficients = {
"cond": cfg_scale + stg_scale + modality_scale - 1.0,
"neg": 1.0 - cfg_scale,
"perturbed": -stg_scale,
"modality": 1.0 - modality_scale,
}
first_velocity = next(iter(local_velocities.values()))
template = cls._ltx2_velocity_to_x0(latents, first_velocity, sigma)
cond_partial = torch.zeros_like(template)
pred_partial = torch.zeros_like(template)
for name in ("cond", "neg", "perturbed", "modality"):
if name in local_video_velocities:
local_video = cls._ltx2_velocity_to_x0(
video_latents, local_video_velocities[name], video_sigma
)
local_audio = cls._ltx2_velocity_to_x0(
audio_latents, local_audio_velocities[name], audio_sigma
)
else:
local_video = torch.zeros_like(video_template)
local_audio = torch.zeros_like(audio_template)
flat = torch.cat((local_video.reshape(-1), local_audio.reshape(-1)))
flat = cfg_model_parallel_all_reduce(flat)
branches[name] = (
flat[:video_numel].reshape_as(video_template),
flat[video_numel:].reshape_as(audio_template),
)
for name, velocity in local_velocities.items():
denoised = cls._ltx2_velocity_to_x0(latents, velocity, sigma)
if name == "cond":
cond_partial = cond_partial + denoised
pred_partial = pred_partial + denoised * coefficients[name]
cond = cfg_model_parallel_all_reduce(cond_partial)
pred = cfg_model_parallel_all_reduce(pred_partial)
return cls._ltx2_apply_rescale(cond, pred, rescale_scale)
# folding the coefficients changes bf16 rounding and drifts from single-GPU
guided_video = cls._ltx2_calculate_guided_x0(
cond=branches["cond"][0],
uncond_text=branches["neg"][0],
uncond_perturbed=branches["perturbed"][0],
uncond_modality=branches["modality"][0],
cfg_scale=video_cfg_scale,
stg_scale=video_stg_scale,
rescale_scale=video_rescale_scale,
modality_scale=video_modality_scale,
)
guided_audio = cls._ltx2_calculate_guided_x0(
cond=branches["cond"][1],
uncond_text=branches["neg"][1],
uncond_perturbed=branches["perturbed"][1],
uncond_modality=branches["modality"][1],
cfg_scale=audio_cfg_scale,
stg_scale=audio_stg_scale,
rescale_scale=audio_rescale_scale,
modality_scale=audio_modality_scale,
)
return guided_video, guided_audio
@staticmethod
def _ltx2_channelwise_normalize(noise: torch.Tensor) -> torch.Tensor:
@@ -556,7 +587,9 @@ class LTX2DenoisingStage(DenoisingStage):
dtype=sliced.dtype,
)
sliced = torch.cat([sliced, pad], dim=1)
return sliced.to(dtype=reference_tensor.dtype)
return sliced
# The native HQ SDE path consumes this through `.float()`; keep the
# original downcast boundary before that upcast to preserve the trajectory.
return cls._ltx2_res2s_new_noise(reference_tensor, generator).to(
dtype=reference_tensor.dtype
)
@@ -703,6 +736,8 @@ class LTX2DenoisingStage(DenoisingStage):
update (midpoint SDE, bongmath anchor refinement, midpoint re-eval,
final RK2 combination with SDE noise). Mirrors the guided stage-1 res2s
math but without CFG/STG (stage-2 HQ uses the simple CFG path).
Raw HQ model timesteps are only inputs to the DiT call; stage-2 res2s
math stays in scheduler sigma space.
"""
sigma_val = float(sigma.item())
sigma_next_val = float(sigma_next.item())
@@ -711,24 +746,10 @@ class LTX2DenoisingStage(DenoisingStage):
denoised_video = ctx.latents.float()
denoised_audio = ctx.audio_latents.float()
else:
video_sigma_for_x0 = (
model_video_timestep
if ctx.use_ltx23_hq_timestep_semantics
and model_video_timestep is not None
else sigma
denoised_video = ctx.latents.float() - sigma * model_video_velocity.float()
denoised_audio = (
ctx.audio_latents.float() - sigma * model_audio_velocity.float()
)
audio_sigma_for_x0 = (
model_audio_timestep
if ctx.use_ltx23_hq_timestep_semantics
and model_audio_timestep is not None
else sigma
)
denoised_video = self._ltx2_velocity_to_x0(
ctx.latents, model_video_velocity, video_sigma_for_x0
).float()
denoised_audio = self._ltx2_velocity_to_x0(
ctx.audio_latents, model_audio_velocity, audio_sigma_for_x0
).float()
if sigma_val == 0.0 or sigma_next_val == 0.0:
next_video = denoised_video.to(dtype=ctx.latents.dtype)
@@ -738,14 +759,9 @@ class LTX2DenoisingStage(DenoisingStage):
sigma_d = sigma.double()
sigma_next_d = sigma_next.double()
if ctx.use_ltx23_hq_timestep_semantics:
h = self._ltx2_res2s_step_size_scalar(sigma_d, sigma_next_d)
a21, b1, b2 = self._ltx2_get_res2s_coefficients_scalar(h)
h_value = h
else:
h = -torch.log(torch.clamp(sigma_next_d / sigma_d, min=1e-12))
a21, b1, b2 = self._ltx2_get_res2s_coefficients(h)
h_value = float(h.item())
h = -torch.log(torch.clamp(sigma_next_d / sigma_d, min=1e-12))
a21, b1, b2 = self._ltx2_get_res2s_coefficients(h)
h_value = float(h.item())
sub_sigma = torch.sqrt(torch.clamp(sigma_d * sigma_next_d, min=0.0))
anchor_video = ctx.latents.double()
@@ -809,22 +825,8 @@ class LTX2DenoisingStage(DenoisingStage):
midpoint_video_model_latents, midpoint_audio_model_latents, sub_sigma
)
mid_video_sigma_for_x0 = (
mid_video_timestep
if ctx.use_ltx23_hq_timestep_semantics and mid_video_timestep is not None
else sub_sigma
)
mid_audio_sigma_for_x0 = (
mid_audio_timestep
if ctx.use_ltx23_hq_timestep_semantics and mid_audio_timestep is not None
else sub_sigma
)
midpoint_denoised_video = self._ltx2_velocity_to_x0(
midpoint_video_latents, mid_v, mid_video_sigma_for_x0
).float()
midpoint_denoised_audio = self._ltx2_velocity_to_x0(
midpoint_audio_latents, mid_a, mid_audio_sigma_for_x0
).float()
midpoint_denoised_video = midpoint_video_latents.float() - sub_sigma * mid_v
midpoint_denoised_audio = midpoint_audio_latents.float() - sub_sigma * mid_a
eps2_video = midpoint_denoised_video.double() - anchor_video
eps2_audio = midpoint_denoised_audio.double() - anchor_audio
@@ -848,23 +850,19 @@ class LTX2DenoisingStage(DenoisingStage):
ctx.audio_latents, batch
).float()
)
sde_sigma = sigma if ctx.use_ltx23_hq_timestep_semantics else sigma_d
sde_sigma_next = (
sigma_next if ctx.use_ltx23_hq_timestep_semantics else sigma_next_d
)
next_video = self._ltx2_res2s_sde_step(
sample=anchor_video,
denoised_sample=next_video_det,
sigma=sde_sigma,
sigma_next=sde_sigma_next,
sigma=sigma_d,
sigma_next=sigma_next_d,
noise=step_noise_video,
terminal=False,
)
next_audio = self._ltx2_res2s_sde_step(
sample=anchor_audio,
denoised_sample=next_audio_det,
sigma=sde_sigma,
sigma_next=sde_sigma_next,
sigma=sigma_d,
sigma_next=sigma_next_d,
noise=step_noise_audio,
terminal=False,
)
@@ -1574,9 +1572,19 @@ class LTX2DenoisingStage(DenoisingStage):
ctx.denoise_mask = ctx.denoise_mask.to(device)
if ctx.clean_latent is not None:
ctx.clean_latent = ctx.clean_latent.to(device)
self._move_ltx2_scheduler_tensors_to_device(ctx.scheduler, device)
self._move_ltx2_scheduler_tensors_to_device(ctx.audio_scheduler, device)
return ctx
@staticmethod
def _move_ltx2_scheduler_tensors_to_device(scheduler: object, device) -> None:
# cfg-parallel batches are broadcast from rank 0; scheduler state must be local
for name in ("sigmas", "timesteps"):
value = getattr(scheduler, name, None)
if isinstance(value, torch.Tensor):
setattr(scheduler, name, value.to(device))
def _before_denoising_loop(
self, ctx: LTX2DenoisingContext, batch: Req, server_args: ServerArgs
) -> None:
@@ -1652,9 +1660,13 @@ class LTX2DenoisingStage(DenoisingStage):
)
use_official_cfg_path = stage1_guider_params is None
if use_official_cfg_path:
cfg_parallel = (
server_args.enable_cfg_parallel and batch.do_classifier_free_guidance
)
do_two_branch_cfg = batch.do_classifier_free_guidance
if ctx.stage == "stage2" and is_ltx2_two_stage_pipeline_name(
server_args.pipeline_class_name
):
# official two-stage stage 2 is a distilled positive-only denoiser
do_two_branch_cfg = False
cfg_parallel = server_args.enable_cfg_parallel and do_two_branch_cfg
cfg_rank = get_classifier_free_guidance_rank() if cfg_parallel else 0
if cfg_parallel:
@@ -1691,7 +1703,7 @@ class LTX2DenoisingStage(DenoisingStage):
audio_encoder_hidden_states=batch.audio_prompt_embeds[0],
encoder_attention_mask=prompt_attention_mask,
)
if batch.do_classifier_free_guidance:
if do_two_branch_cfg:
cfg_batch_size = batch_size * 2
model_kwargs = self._repeat_ltx2_model_kwargs_batch(
model_kwargs, cfg_batch_size
@@ -1735,7 +1747,7 @@ class LTX2DenoisingStage(DenoisingStage):
model_video, model_audio = self._combine_cfg_parallel_av(
model_video, model_audio, float(batch.guidance_scale), cfg_rank
)
elif batch.do_classifier_free_guidance:
elif do_two_branch_cfg:
model_video_uncond, model_video_text = model_video.chunk(2)
model_audio_uncond, model_audio_text = model_audio.chunk(2)
model_video = model_video_uncond + (
@@ -1780,7 +1792,7 @@ class LTX2DenoisingStage(DenoisingStage):
audio_encoder_hidden_states=batch.audio_prompt_embeds[0],
encoder_attention_mask=prompt_attention_mask,
)
if batch.do_classifier_free_guidance:
if do_two_branch_cfg:
cfg_batch_size = batch_size_local * 2
model_kwargs_local = self._repeat_ltx2_model_kwargs_batch(
model_kwargs_local, cfg_batch_size
@@ -1827,7 +1839,7 @@ class LTX2DenoisingStage(DenoisingStage):
mid_v = mid_v.float()
mid_a = mid_a.float()
if batch.do_classifier_free_guidance:
if do_two_branch_cfg:
mid_v_u, mid_v_t = mid_v.chunk(2)
mid_a_u, mid_a_t = mid_a.chunk(2)
mid_v = mid_v_u + batch.guidance_scale * (mid_v_t - mid_v_u)
@@ -2202,20 +2214,43 @@ class LTX2DenoisingStage(DenoisingStage):
) * ctx.denoise_mask.squeeze(-1)
if stage1_cfg_parallel:
guided_video = self._ltx2_combine_guided_x0_parallel(
latents=video_latents,
local_velocities={
name: output[0] for name, output in pass_outputs.items()
},
sigma=video_sigma_for_x0,
cfg_scale=float(stage1_guider_params["video_cfg_scale"]),
stg_scale=float(stage1_guider_params["video_stg_scale"]),
rescale_scale=float(
stage1_guider_params["video_rescale_scale"]
),
modality_scale=float(
stage1_guider_params["video_modality_scale"]
),
guided_video, guided_audio = (
self._ltx2_combine_guided_x0_parallel_av(
video_latents=video_latents,
audio_latents=audio_latents,
local_video_velocities={
name: output[0] for name, output in pass_outputs.items()
},
local_audio_velocities={
name: output[1] for name, output in pass_outputs.items()
},
video_sigma=video_sigma_for_x0,
audio_sigma=audio_sigma_for_x0,
video_cfg_scale=float(
stage1_guider_params["video_cfg_scale"]
),
video_stg_scale=float(
stage1_guider_params["video_stg_scale"]
),
video_rescale_scale=float(
stage1_guider_params["video_rescale_scale"]
),
video_modality_scale=float(
stage1_guider_params["video_modality_scale"]
),
audio_cfg_scale=float(
stage1_guider_params["audio_cfg_scale"]
),
audio_stg_scale=float(
stage1_guider_params["audio_stg_scale"]
),
audio_rescale_scale=float(
stage1_guider_params["audio_rescale_scale"]
),
audio_modality_scale=float(
stage1_guider_params["audio_modality_scale"]
),
)
)
if video_skip and ctx.last_denoised_video is not None:
denoised_video_local = ctx.last_denoised_video
@@ -2224,21 +2259,6 @@ class LTX2DenoisingStage(DenoisingStage):
if update_skip_cache:
ctx.last_denoised_video = guided_video
guided_audio = self._ltx2_combine_guided_x0_parallel(
latents=audio_latents,
local_velocities={
name: output[1] for name, output in pass_outputs.items()
},
sigma=audio_sigma_for_x0,
cfg_scale=float(stage1_guider_params["audio_cfg_scale"]),
stg_scale=float(stage1_guider_params["audio_stg_scale"]),
rescale_scale=float(
stage1_guider_params["audio_rescale_scale"]
),
modality_scale=float(
stage1_guider_params["audio_modality_scale"]
),
)
if audio_skip and ctx.last_denoised_audio is not None:
denoised_audio_local = ctx.last_denoised_audio
else:
python/sglang/multimodal_gen/runtime/server_args.py
+9 -3
View File
@@ -614,9 +614,15 @@ class ServerArgs(DisaggArgsMixin):
if component_name is None:
continue
key = component_name.replace("-", "_")
backend = self.component_attention_backends.get(key)
if backend is not None:
return AttentionBackendEnum[backend.upper()], key
fallback_keys = [key]
if key.endswith("_2"):
# Secondary two-stage components inherit the base component
# backend unless explicitly overridden.
fallback_keys.append(key[:-2])
for backend_key in fallback_keys:
backend = self.component_attention_backends.get(backend_key)
if backend is not None:
return AttentionBackendEnum[backend.upper()], backend_key
return None, None
def _adjust_warmup(self):
python/sglang/multimodal_gen/test/server/consistency_threshold.json
+9 -9
View File
@@ -110,15 +110,15 @@
"mean_abs_diff_threshold": 10.0
},
"ltx_2.3_one_stage_ti2v": {
"clip_threshold": 0.64,
"ssim_threshold": 0.42,
"psnr_threshold": 8.8,
"mean_abs_diff_threshold": 59.0
"clip_threshold": 0.84,
"ssim_threshold": 0.78,
"psnr_threshold": 20.5,
"mean_abs_diff_threshold": 11.0
},
"ltx_2.3_two_stage_t2v_2gpus": {
"clip_threshold": 0.79,
"ssim_threshold": 0.12,
"psnr_threshold": 12.1,
"psnr_threshold": 16.0,
"mean_abs_diff_threshold": 51.0
},
"wan2_1_t2v_1.3b_teacache_enabled": {
@@ -254,10 +254,10 @@
"mean_abs_diff_threshold": 45.0
},
"ltx_2_3_two_stage_ti2v_2gpus": {
"clip_threshold": 0.92,
"ssim_threshold": 0.58,
"psnr_threshold": 17.5,
"mean_abs_diff_threshold": 20.0
"clip_threshold": 0.55,
"ssim_threshold": 0.29,
"psnr_threshold": 11.7,
"mean_abs_diff_threshold": 47.0
}
},
"default_clip_threshold_image": 0.92,
python/sglang/multimodal_gen/test/server/gpu_cases.py
+5 -6
View File
@@ -601,8 +601,7 @@ TWO_GPU_CASES = [
"ltx_2_two_stage_t2v",
DiffusionServerArgs(
model_path="Lightricks/LTX-2",
ulysses_degree=2,
dit_layerwise_offload=True,
cfg_parallel=True,
extras=["--pipeline-class-name LTX2TwoStagePipeline"],
),
T2V_sampling_params,
@@ -613,7 +612,7 @@ TWO_GPU_CASES = [
model_path="Lightricks/LTX-2.3",
cfg_parallel=True,
extras=[
"--pipeline-class-name LTX2TwoStagePipeline --ltx2-two-stage-device-mode original"
"--pipeline-class-name LTX2TwoStagePipeline --ltx2-two-stage-device-mode original",
],
),
TI2V_sampling_params,
@@ -634,10 +633,10 @@ TWO_GPU_CASES = [
cfg_parallel=True,
extras=[
"--pipeline-class-name LTX2TwoStagePipeline",
"--ltx2-two-stage-device-mode original",
"--component-attention-backends transformer=fa",
],
),
T2V_sampling_params,
DiffusionSamplingParams(prompt=T2V_PROMPT, extras={"seed": 42}),
run_component_accuracy_check=False,
),
# I2V LoRA test case
@@ -708,7 +707,7 @@ TWO_GPU_CASES = [
"ltx_2.3_one_stage_ti2v",
DiffusionServerArgs(
model_path="Lightricks/LTX-2.3",
ulysses_degree=2,
cfg_parallel=True,
),
TI2V_sampling_params,
run_component_accuracy_check=False,
python/sglang/multimodal_gen/test/server/perf_baselines.json
+61 -59
View File
@@ -1096,68 +1096,70 @@
},
"ltx_2_two_stage_t2v": {
"stages_ms": {
"InputValidationStage": 0.1,
"TextEncodingStage": 1830.33,
"LTX2TextConnectorStage": 9.61,
"LTX2HalveResolutionStage": 0.06,
"LTX2LoRASwitchStage": 11578.48,
"LTX2SigmaPreparationStage": 0.18,
"TimestepPreparationStage": 19.37,
"LTX2AVLatentPreparationStage": 0.23,
"LTX2AVDenoisingStage": 53177.49,
"LTX2UpsampleStage": 2104.17,
"LTX2RefinementStage": 3947.13,
"LTX2AVDecodingStage": 332.07
"InputValidationStage": 0.03,
"TextEncodingStage": 1463.27,
"LTX2TextConnectorStage": 12.85,
"LTX2HalveResolutionStage": 0.11,
"LTX2LoRASwitchStage": 325.46,
"LTX2SigmaPreparationStage": 0.25,
"TimestepPreparationStage": 7.85,
"LTX2AVLatentPreparationStage": 0.34,
"LTX2ImageEncodingStage": 0.02,
"LTX2AVDenoisingStage": 7744.44,
"LTX2UpsampleStage": 2.98,
"LTX2RefinementStage": 666.08,
"LTX2AVDecodingStage": 338.87,
"per_frame_generation": null
},
"denoise_step_ms": {
"0": 1186.27,
"1": 1331.86,
"2": 1330.41,
"3": 1331.28,
"4": 1331.5,
"5": 1331.45,
"6": 1331.79,
"7": 1331.59,
"8": 1331.55,
"9": 1331.55,
"10": 1331.51,
"11": 1331.34,
"12": 1331.48,
"13": 1331.38,
"14": 1331.74,
"15": 1331.84,
"16": 1331.09,
"17": 1331.79,
"18": 1332.34,
"19": 1337.7,
"20": 1337.53,
"21": 1334.48,
"22": 1334.87,
"23": 1333.28,
"24": 1333.15,
"25": 1333.82,
"26": 1333.55,
"27": 1339.35,
"28": 1336.96,
"29": 1335.25,
"30": 1331.8,
"31": 1339.52,
"32": 1334.1,
"33": 1331.96,
"34": 1331.78,
"35": 1332.5,
"36": 1331.3,
"37": 1331.75,
"38": 1331.94,
"39": 1331.84,
"40": 1278.82,
"41": 1330.68,
"42": 1331.7
"0": 165.1,
"1": 309.19,
"2": 166.63,
"3": 175.53,
"4": 158.77,
"5": 191.1,
"6": 203.59,
"7": 202.98,
"8": 205.09,
"9": 195.92,
"10": 226.96,
"11": 203.85,
"12": 190.75,
"13": 192.07,
"14": 193.85,
"15": 191.34,
"16": 193.56,
"17": 192.05,
"18": 189.53,
"19": 191.61,
"20": 187.52,
"21": 192.57,
"22": 190.67,
"23": 189.47,
"24": 187.38,
"25": 190.22,
"26": 196.7,
"27": 185.05,
"28": 189.59,
"29": 209.85,
"30": 194.47,
"31": 189.43,
"32": 189.58,
"33": 188.41,
"34": 198.11,
"35": 188.45,
"36": 187.06,
"37": 188.65,
"38": 200.22,
"39": 156.46,
"40": 220.06,
"41": 223.51,
"42": 221.07
},
"expected_e2e_ms": 73463.94,
"expected_avg_denoise_ms": 1328.22,
"expected_median_denoise_ms": 1331.79,
"estimated_full_test_time_s": 133.1
"expected_e2e_ms": 10601.1,
"expected_avg_denoise_ms": 195.04,
"expected_median_denoise_ms": 191.34,
"estimated_full_test_time_s": 345.4
},
"wan2_2_ti2v_5b": {
"stages_ms": {
python/sglang/multimodal_gen/test/test_consistency_metrics.py
+3
View File
@@ -169,3 +169,6 @@ def test_save_consistency_failure_artifact(tmp_path, monkeypatch):
assert artifact_path.suffix == ".png"
assert (tmp_path / "consistency_failures" / "summary.json").exists()
assert (tmp_path / "consistency_failures" / "index.html").exists()
assert (
tmp_path / "consistency_failures" / "generated" / "unit_image_fail_1gpu.png"
).exists()
python/sglang/multimodal_gen/test/test_utils.py
+55 -2
View File
@@ -50,12 +50,12 @@ SGL_TEST_FILES_CONSISTENCY_GT_BASES = (
SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_BASE,
SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE,
)
# Keep non-comparable LTX CI scenarios on sglang_generated rather than hiding
# remaining semantic gaps behind very loose official thresholds.
# LTX cases listed here compare against official-generated GT.
SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_CASES = frozenset(
{
"ltx_2.3_one_stage_ti2v",
"ltx_2.3_two_stage_t2v_2gpus",
"ltx_2_3_two_stage_ti2v_2gpus",
}
)
CONSISTENCY_THRESHOLD_JSON_PATH = (
@@ -1458,6 +1458,7 @@ def _consistency_failure_record(
is_video: bool,
output_format: str | None,
image_name: str,
generated_files: list[str],
gt_remote_files: list[tuple[str, str]] | None,
) -> dict[str, Any]:
return {
@@ -1466,6 +1467,7 @@ def _consistency_failure_record(
"is_video": is_video,
"output_format": output_format,
"comparison_png": image_name,
"generated_files": generated_files,
"metrics": {
"min_clip_similarity": _json_metric_value(result.min_similarity),
"min_ssim": _json_metric_value(result.min_ssim),
@@ -1499,6 +1501,36 @@ def _consistency_failure_record(
}
def _save_generated_artifact_images(
out_dir: Path,
case_id: str,
num_gpus: int,
output_frames: list[np.ndarray],
is_video: bool,
output_format: str | None,
) -> list[str]:
generated_dir = out_dir / "generated"
generated_dir.mkdir(parents=True, exist_ok=True)
safe_case_id = _safe_artifact_name(case_id)
if is_video:
suffixes = ("frame_0", "frame_mid", "frame_last")
filenames = [
f"{safe_case_id}_{num_gpus}gpu_{suffix}.png"
for suffix in suffixes[: len(output_frames)]
]
else:
ext = output_format_to_ext(output_format)
filenames = [f"{safe_case_id}_{num_gpus}gpu.{ext}"]
generated_files = []
for frame, filename in zip(output_frames, filenames):
path = generated_dir / filename
Image.fromarray(_ensure_rgb_uint8_image(frame)).save(path)
generated_files.append(str(path.relative_to(out_dir)))
return generated_files
def _write_consistency_failure_index(
out_dir: Path,
records: list[dict[str, Any]],
@@ -1508,6 +1540,15 @@ def _write_consistency_failure_index(
case_id = html.escape(record["case_id"])
png = html.escape(record["comparison_png"])
metrics = record["metrics"]
generated_links = "".join(
f'<li><a href="{html.escape(path)}">{html.escape(path)}</a></li>'
for path in record.get("generated_files", [])
)
generated_html = (
f"<p>Generated images:</p><ul>{generated_links}</ul>"
if generated_links
else ""
)
sections.append(
"<section>"
f"<h2>{case_id} ({record['num_gpus']} GPU)</h2>"
@@ -1518,6 +1559,7 @@ def _write_consistency_failure_index(
f"mean_abs_diff={metrics['max_mean_abs_diff']}"
"</p>"
f'<img src="{png}" alt="{case_id} comparison">'
f"{generated_html}"
"</section>"
)
@@ -1529,6 +1571,7 @@ def _write_consistency_failure_index(
"section{margin:0 0 28px;padding:16px;background:white;border:1px solid #ddd;border-radius:6px}"
"h2{font-size:18px;margin:0 0 8px}"
"p{margin:0 0 12px;color:#444}"
"ul{margin:0 0 12px;padding-left:20px}"
"img{max-width:100%;height:auto;border:1px solid #ddd}"
"</style></head><body>"
"<h1>Diffusion consistency failures</h1>" + "".join(sections) + "</body></html>"
@@ -1566,6 +1609,15 @@ def save_consistency_failure_artifact(
)
comparison.save(image_path)
generated_files = _save_generated_artifact_images(
out_dir=out_dir,
case_id=case_id,
num_gpus=num_gpus,
output_frames=output_frames,
is_video=is_video,
output_format=output_format,
)
record = _consistency_failure_record(
case_id=case_id,
num_gpus=num_gpus,
@@ -1573,6 +1625,7 @@ def save_consistency_failure_artifact(
is_video=is_video,
output_format=output_format,
image_name=image_name,
generated_files=generated_files,
gt_remote_files=gt_remote_files,
)
case_json_path = out_dir / f"{safe_case_id}.json"