From 6e5b4de01a68985f532702abe82f056218ff30c7 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 11 May 2026 13:42:08 +0800 Subject: [PATCH] [diffusion] fix: further align ltx2.3 accuracy with tp (#24660) --- .../cookbook/diffusion/LTX/LTX2 & LTX2.3.mdx | 13 +- .../src/snippets/diffusion/ltx-deployment.jsx | 2 +- .../runtime/layers/attention/backends/sdpa.py | 21 +- .../runtime/layers/attention/layer.py | 83 ++++-- .../runtime/models/dits/ltx_2.py | 61 ++-- .../runtime/models/encoders/gemma_3.py | 64 ++--- .../runtime/pipelines/ltx_2_pipeline.py | 1 - .../executors/parallel_executor.py | 6 +- .../pipelines_core/stages/denoising_av.py | 12 +- .../pipelines_core/stages/ltx_2_denoising.py | 262 ++++++++++-------- .../multimodal_gen/runtime/server_args.py | 12 +- .../test/server/consistency_threshold.json | 18 +- .../multimodal_gen/test/server/gpu_cases.py | 11 +- .../test/server/perf_baselines.json | 120 ++++---- .../test/test_consistency_metrics.py | 3 + .../sglang/multimodal_gen/test/test_utils.py | 57 +++- 16 files changed, 435 insertions(+), 311 deletions(-) diff --git a/docs_new/cookbook/diffusion/LTX/LTX2 & LTX2.3.mdx b/docs_new/cookbook/diffusion/LTX/LTX2 & LTX2.3.mdx index 2eaf29650..c73169bc1 100644 --- a/docs_new/cookbook/diffusion/LTX/LTX2 & LTX2.3.mdx +++ b/docs_new/cookbook/diffusion/LTX/LTX2 & LTX2.3.mdx @@ -71,7 +71,7 @@ Other deployment flags: - `--lora-weight-name`: Select the exact safetensors file when the LoRA repository contains multiple weight files. -For native LTX-2.3 two-stage serving without a user LoRA, `resident` is the fastest high-VRAM path. When you pass `--lora-path`, SGLang still applies the user LoRA during the two-stage switch, so use `resident` on H200-class GPUs for enough VRAM, but do not expect the same premerged-stage2 benefit as the no-user-LoRA path. +For native LTX-2.3 two-stage serving without a user LoRA, `resident` is the fastest high-VRAM path. LTX-2 still applies the distilled LoRA during the stage switch, so `--ltx2-two-stage-device-mode` is mainly an LTX-2.3 optimization. When you pass `--lora-path`, SGLang still applies the user LoRA during the two-stage switch, so use `resident` on H200-class GPUs for enough VRAM, but do not expect the same premerged-stage2 benefit as the no-user-LoRA path. ### 3.3 Fast multi-GPU presets @@ -80,11 +80,12 @@ For latency-oriented LTX serving, prefer CFG parallel over sequence parallelism. | Target | Recommended server flags | Notes | | --- | --- | --- | -| 1 high-VRAM GPU | `--ltx2-two-stage-device-mode resident` | Fastest two-stage setup when both DiTs fit. | -| 1 standard GPU | `--ltx2-two-stage-device-mode snapshot` | Lower VRAM than `resident`; use this when H100-class memory is tight. | -| 2 GPUs | `--num-gpus 2 --enable-cfg-parallel --ltx2-two-stage-device-mode resident` | Fastest common 2-GPU setup. | -| 4 GPUs | `--num-gpus 4 --tp-size 2 --enable-cfg-parallel --ltx2-two-stage-device-mode resident` | Fastest common 4-GPU layout: TP2 inside each CFG branch. | -| Official comparison | `--ltx2-two-stage-device-mode original` | Use this only when matching the original stage-switch semantics matters. | +| LTX-2.3, 1 high-VRAM GPU | `--ltx2-two-stage-device-mode resident` | Fastest two-stage setup when both DiTs fit. | +| LTX-2.3, 1 standard GPU | `--ltx2-two-stage-device-mode snapshot` | Lower VRAM than `resident`; use this when H100-class memory is tight. | +| LTX-2, 2 GPUs | `--num-gpus 2 --enable-cfg-parallel` | Fastest verified 2-GPU setup; keep `--dit-layerwise-offload` disabled unless memory is tight. | +| LTX-2.3, 2 GPUs | `--num-gpus 2 --enable-cfg-parallel --ltx2-two-stage-device-mode resident` | Fastest common 2-GPU setup. | +| LTX-2.3, 4 GPUs | `--num-gpus 4 --tp-size 2 --enable-cfg-parallel --ltx2-two-stage-device-mode resident` | Fastest common 4-GPU layout: TP2 inside each CFG branch. | +| Official comparison | `--ltx2-two-stage-device-mode original` | Use this only when matching the original LTX-2.3 stage-switch semantics matters. | Use `--enable-cfg-parallel` for degree-2 CFG parallel. Use `--cfg-parallel-size` only when you explicitly need a different CFG branch count. If `resident` exceeds available VRAM, keep the same parallelism preset and switch only the device mode to `snapshot`. diff --git a/docs_new/src/snippets/diffusion/ltx-deployment.jsx b/docs_new/src/snippets/diffusion/ltx-deployment.jsx index e4d7ca77a..2f15a7e3c 100644 --- a/docs_new/src/snippets/diffusion/ltx-deployment.jsx +++ b/docs_new/src/snippets/diffusion/ltx-deployment.jsx @@ -141,7 +141,7 @@ export const LTXDeployment = () => { let command = `sglang serve \\\n --model-path ${config.repoId} \\\n --pipeline-class-name ${pipelineClass}`; command += getParallelFlags(); - if (values.pipeline !== 'one-stage') { + if (values.model === 'ltx23' && values.pipeline !== 'one-stage') { command += ` \\\n --ltx2-two-stage-device-mode ${getDeviceMode()}`; } diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py index 5523f6559..9345785ed 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index ee0127c22..71aa1c4f1 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index 96f513a92..178f362ab 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py b/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py index c5e6ae354..16ed42134 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py index 3259ab363..39e2872b4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py index 3a4cf3419..07f765cbb 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py @@ -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] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py index 6fb7a0280..aeac9cd73 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py index e21620f51..577b86aa2 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py @@ -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: diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 7514e0de9..d98cfddd3 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -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): diff --git a/python/sglang/multimodal_gen/test/server/consistency_threshold.json b/python/sglang/multimodal_gen/test/server/consistency_threshold.json index f3e85286f..e67ae720c 100644 --- a/python/sglang/multimodal_gen/test/server/consistency_threshold.json +++ b/python/sglang/multimodal_gen/test/server/consistency_threshold.json @@ -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, diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index 2f5c000c3..39fb54b53 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -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, diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 08685b75c..393a73137 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -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": { diff --git a/python/sglang/multimodal_gen/test/test_consistency_metrics.py b/python/sglang/multimodal_gen/test/test_consistency_metrics.py index b77b42d3a..7cb78928b 100644 --- a/python/sglang/multimodal_gen/test/test_consistency_metrics.py +++ b/python/sglang/multimodal_gen/test/test_consistency_metrics.py @@ -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() diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index ccda47743..6c64ce709 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -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'
  • {html.escape(path)}
    • ' + for path in record.get("generated_files", []) + ) + generated_html = ( + f"

    Generated images:

    " + if generated_links + else "" + ) sections.append( "
    " f"

    {case_id} ({record['num_gpus']} GPU)

    " @@ -1518,6 +1559,7 @@ def _write_consistency_failure_index( f"mean_abs_diff={metrics['max_mean_abs_diff']}" "

    " f'{case_id} comparison' + f"{generated_html}" "
    " ) @@ -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}" "" "

    Diffusion consistency failures

    " + "".join(sections) + "" @@ -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"