[diffusion] model: support Ltx 2.3 two stage ti2v (#22667)

python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py

@@ -118,6 +118,24 @@ def is_ltx23_native_variant(arch_config: object) -> bool:
     return str(getattr(arch_config, "ltx_variant", "ltx_2")) == "ltx_2_3"
 
 
+def sync_ltx23_runtime_vae_markers(
+    arch_config: object,
+    loaded_vae_config: object | None,
+) -> None:
+    if loaded_vae_config is None:
+        return
+    source = getattr(loaded_vae_config, "arch_config", loaded_vae_config)
+    for key in (
+        "ltx_variant",
+        "condition_encoder_subdir",
+        "video_decoder_variant",
+        "video_decoder_config",
+    ):
+        value = getattr(source, key, None)
+        if value is not None:
+            setattr(arch_config, key, value)
+
+
 def _gemma_postprocess_func(
     outputs: BaseEncoderOutput,
     text_inputs: dict,

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -609,6 +609,7 @@ class SamplingParams:
 
         user_kwargs = dict(kwargs)
         user_kwargs.pop("diffusers_kwargs", None)
+
         user_sampling_params = SamplingParams(*args, **user_kwargs)
         # TODO: refactor
         sampling_params._merge_with_user_params(
@@ -793,7 +794,7 @@ class SamplingParams:
             "--cfg-normalization",
             type=float,
             dest="cfg_normalization",
-            help=("CFG renormalization factor (for Z-Image). "),
+            help="CFG renormalization factor (for Z-Image). ",
         )
         add_argument(
             "--boundary-ratio",

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

@@ -27,6 +27,19 @@ from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 logger = init_logger(__name__)
 
 
+def _backfill_ltx2_audio_vae_latent_stats(
+    loaded: dict[str, torch.Tensor], component_name: str
+) -> None:
+    if component_name != "audio_vae":
+        return
+    mean_key = "per_channel_statistics.mean-of-means"
+    std_key = "per_channel_statistics.std-of-means"
+    if "latents_mean" not in loaded and mean_key in loaded:
+        loaded["latents_mean"] = loaded[mean_key]
+    if "latents_std" not in loaded and std_key in loaded:
+        loaded["latents_std"] = loaded[std_key]
+
+
 def _convert_conv3d_weights_to_channels_last_3d(module: nn.Module) -> int:
     """
     Convert Conv3d weights to channels_last_3d (NDHWC) memory format.
@@ -142,6 +155,7 @@ class VAELoader(ComponentLoader):
         loaded = {}
         for sf_path in safetensors_list:
             loaded.update(safetensors_load_file(sf_path))
+        _backfill_ltx2_audio_vae_latent_stats(loaded, component_name)
         vae.load_state_dict(loaded, strict=False)
 
         state_keys = set(vae.state_dict().keys())

python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py

@@ -517,6 +517,7 @@ class LTX2TextConnectors(nn.Module):
     ):
         super().__init__()
         caption_channels = config.caption_channels
+        self.caption_channels = caption_channels
         text_proj_in_factor = config.text_proj_in_factor
         video_connector_num_attention_heads = config.video_connector_num_attention_heads
         video_connector_attention_head_dim = config.video_connector_attention_head_dim
@@ -641,7 +642,7 @@ class LTX2TextConnectors(nn.Module):
                 audio_hidden_states = audio_hidden_states.to(
                     self.audio_aggregate_embed.weight.dtype
                 )
-            source_dim = self.video_aggregate_embed.out_features
+            source_dim = self.caption_channels
             video_hidden_states = self._rescale_v2_features(
                 video_hidden_states,
                 self.video_aggregate_embed.out_features,

python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py

@@ -48,13 +48,54 @@ def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int:
     )
 
 
+def _ltx2_is_perturbed(
+    perturbation_config: dict[str, object],
+    key: str,
+    block_idx: int,
+) -> bool:
+    value = perturbation_config.get(key)
+    if value is None:
+        return False
+    if key.endswith("_blocks"):
+        return block_idx in value
+    return bool(value)
+
+
+def _ltx2_batched_perturbation_mask(
+    perturbation_configs: tuple[dict[str, object], ...] | None,
+    key: str,
+    block_idx: int,
+    values: torch.Tensor,
+) -> tuple[torch.Tensor | None, bool]:
+    if not perturbation_configs:
+        return None, False
+
+    mask = torch.ones(
+        (len(perturbation_configs),), device=values.device, dtype=values.dtype
+    )
+    any_perturbed = False
+    all_perturbed = True
+    for batch_idx, config in enumerate(perturbation_configs):
+        perturbed = _ltx2_is_perturbed(config, key, block_idx)
+        any_perturbed = any_perturbed or perturbed
+        all_perturbed = all_perturbed and perturbed
+        if perturbed:
+            mask[batch_idx] = 0
+
+    if not any_perturbed:
+        return None, False
+    if all_perturbed:
+        return None, True
+    return mask.view(mask.numel(), *([1] * (values.ndim - 1))), False
+
+
 def apply_interleaved_rotary_emb(
     x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor]
 ) -> torch.Tensor:
     cos, sin = freqs
     x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1)
     x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
-    return (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+    return x * cos + x_rotated * sin
 
 
 def apply_split_rotary_emb(
@@ -76,7 +117,7 @@ def apply_split_rotary_emb(
         )
     r = last // 2
 
-    split_x = x.reshape(*x.shape[:-1], 2, r).float()
+    split_x = x.reshape(*x.shape[:-1], 2, r)
     first_x = split_x[..., :1, :]
     second_x = split_x[..., 1:, :]
 
@@ -248,9 +289,13 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module):
         return self.prepare_audio_coords(*args, **kwargs)
 
     def forward(
-        self, coords: torch.Tensor, device: Optional[Union[str, torch.device]] = None
+        self,
+        coords: torch.Tensor,
+        device: Optional[Union[str, torch.device]] = None,
+        out_dtype: Optional[torch.dtype] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         device = device or coords.device
+        out_dtype = out_dtype or coords.dtype
         num_pos_dims = coords.shape[1]
 
         if coords.ndim == 4:
@@ -317,7 +362,7 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module):
             cos_freqs = torch.swapaxes(cos_freq, 1, 2)
             sin_freqs = torch.swapaxes(sin_freq, 1, 2)
 
-        return cos_freqs, sin_freqs
+        return cos_freqs.to(dtype=out_dtype), sin_freqs.to(dtype=out_dtype)
 
 
 def rms_norm(x: torch.Tensor, eps: float) -> torch.Tensor:
@@ -653,6 +698,8 @@ class LTX2Attention(nn.Module):
                 )
 
             if perturbation_mask is not None:
+                if perturbation_mask.ndim == out.ndim - 1:
+                    perturbation_mask = perturbation_mask.unsqueeze(-1)
                 out = out * perturbation_mask + v * (1 - perturbation_mask)
 
         if not use_attention:
@@ -918,6 +965,10 @@ class LTX2TransformerBlock(nn.Module):
         skip_audio_self_attn: bool = False,
         skip_a2v_cross_attn: bool = False,
         skip_v2a_cross_attn: bool = False,
+        video_self_attn_perturbation_mask: Optional[torch.Tensor] = None,
+        audio_self_attn_perturbation_mask: Optional[torch.Tensor] = None,
+        a2v_cross_attn_perturbation_mask: Optional[torch.Tensor] = None,
+        v2a_cross_attn_perturbation_mask: Optional[torch.Tensor] = None,
         audio_replicated_for_sp: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
 
@@ -934,6 +985,7 @@ class LTX2TransformerBlock(nn.Module):
             norm_hidden_states,
             mask=video_self_attention_mask,
             pe=video_rotary_emb,
+            perturbation_mask=video_self_attn_perturbation_mask,
             all_perturbed=skip_video_self_attn,
             gather_context_kv_for_sp=audio_replicated_for_sp,
         )
@@ -949,6 +1001,7 @@ class LTX2TransformerBlock(nn.Module):
             norm_audio_hidden_states,
             mask=audio_self_attention_mask,
             pe=audio_rotary_emb,
+            perturbation_mask=audio_self_attn_perturbation_mask,
             all_perturbed=skip_audio_self_attn,
             skip_sequence_parallel_override=audio_replicated_for_sp,
         )
@@ -1097,6 +1150,10 @@ class LTX2TransformerBlock(nn.Module):
                 mask=a2v_cross_attention_mask,
                 skip_sequence_parallel_override=audio_replicated_for_sp,
             )
+            if a2v_cross_attn_perturbation_mask is not None:
+                a2v_attn_hidden_states = (
+                    a2v_attn_hidden_states * a2v_cross_attn_perturbation_mask
+                )
             hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states
 
         # V2A
@@ -1116,6 +1173,10 @@ class LTX2TransformerBlock(nn.Module):
                 mask=v2a_cross_attention_mask,
                 gather_context_kv_for_sp=audio_replicated_for_sp,
             )
+            if v2a_cross_attn_perturbation_mask is not None:
+                v2a_attn_hidden_states = (
+                    v2a_attn_hidden_states * v2a_cross_attn_perturbation_mask
+                )
             audio_hidden_states = (
                 audio_hidden_states + v2a_gate * v2a_attn_hidden_states
             )
@@ -1532,6 +1593,12 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
             raise ValueError(
                 "audio_num_frames must be provided for RoPE coordinate generation."
             )
+        perturbation_configs = kwargs.get("perturbation_configs")
+        if perturbation_configs is not None and len(perturbation_configs) != batch_size:
+            raise ValueError(
+                "perturbation_configs length must match batch size, got "
+                f"{len(perturbation_configs)=} {batch_size=}."
+            )
 
         if video_coords is None:
             # Wan-style SP-RoPE: when SP is enabled, each rank runs on its local
@@ -1566,15 +1633,25 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
             video_coords, hidden_states.device, hidden_states.dtype
         )
         audio_coords = audio_coords.to(device=audio_hidden_states.device)
-        video_rotary_emb = self.rope(video_coords, device=hidden_states.device)
+        video_rotary_emb = self.rope(
+            video_coords,
+            device=hidden_states.device,
+            out_dtype=hidden_states.dtype,
+        )
         audio_rotary_emb = self.audio_rope(
-            audio_coords, device=audio_hidden_states.device
+            audio_coords,
+            device=audio_hidden_states.device,
+            out_dtype=audio_hidden_states.dtype,
         )
         ca_video_rotary_emb = self.cross_attn_rope(
-            video_coords[:, 0:1, :], device=hidden_states.device
+            video_coords[:, 0:1, :],
+            device=hidden_states.device,
+            out_dtype=hidden_states.dtype,
         )
         ca_audio_rotary_emb = self.cross_attn_audio_rope(
-            audio_coords[:, 0:1, :], device=audio_hidden_states.device
+            audio_coords[:, 0:1, :],
+            device=audio_hidden_states.device,
+            out_dtype=audio_hidden_states.dtype,
         )
 
         # 2. Patchify input projections
@@ -1673,6 +1750,55 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
         skip_video_self_attn_blocks = set(skip_video_self_attn_blocks or ())
         skip_audio_self_attn_blocks = set(skip_audio_self_attn_blocks or ())
         for block in self.transformer_blocks:
+            video_self_attn_perturbation_mask = None
+            audio_self_attn_perturbation_mask = None
+            a2v_cross_attn_perturbation_mask = None
+            v2a_cross_attn_perturbation_mask = None
+            skip_video_self_attn = block.idx in skip_video_self_attn_blocks
+            skip_audio_self_attn = block.idx in skip_audio_self_attn_blocks
+            skip_a2v_cross_attn = disable_a2v_cross_attn
+            skip_v2a_cross_attn = disable_v2a_cross_attn
+            if perturbation_configs is not None:
+                if not skip_video_self_attn:
+                    (
+                        video_self_attn_perturbation_mask,
+                        skip_video_self_attn,
+                    ) = _ltx2_batched_perturbation_mask(
+                        perturbation_configs,
+                        "skip_video_self_attn_blocks",
+                        block.idx,
+                        hidden_states,
+                    )
+                if not skip_audio_self_attn:
+                    (
+                        audio_self_attn_perturbation_mask,
+                        skip_audio_self_attn,
+                    ) = _ltx2_batched_perturbation_mask(
+                        perturbation_configs,
+                        "skip_audio_self_attn_blocks",
+                        block.idx,
+                        audio_hidden_states,
+                    )
+                if not skip_a2v_cross_attn:
+                    (
+                        a2v_cross_attn_perturbation_mask,
+                        skip_a2v_cross_attn,
+                    ) = _ltx2_batched_perturbation_mask(
+                        perturbation_configs,
+                        "skip_a2v_cross_attn",
+                        block.idx,
+                        hidden_states,
+                    )
+                if not skip_v2a_cross_attn:
+                    (
+                        v2a_cross_attn_perturbation_mask,
+                        skip_v2a_cross_attn,
+                    ) = _ltx2_batched_perturbation_mask(
+                        perturbation_configs,
+                        "skip_v2a_cross_attn",
+                        block.idx,
+                        audio_hidden_states,
+                    )
             hidden_states, audio_hidden_states = block(
                 hidden_states,
                 audio_hidden_states,
@@ -1699,10 +1825,14 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
                 audio_self_attention_mask=audio_self_attention_mask,
                 a2v_cross_attention_mask=a2v_cross_attention_mask,
                 v2a_cross_attention_mask=v2a_cross_attention_mask,
-                skip_video_self_attn=block.idx in skip_video_self_attn_blocks,
+                skip_video_self_attn=skip_video_self_attn,
-                skip_audio_self_attn=block.idx in skip_audio_self_attn_blocks,
+                skip_audio_self_attn=skip_audio_self_attn,
-                skip_a2v_cross_attn=disable_a2v_cross_attn,
+                skip_a2v_cross_attn=skip_a2v_cross_attn,
-                skip_v2a_cross_attn=disable_v2a_cross_attn,
+                skip_v2a_cross_attn=skip_v2a_cross_attn,
+                video_self_attn_perturbation_mask=video_self_attn_perturbation_mask,
+                audio_self_attn_perturbation_mask=audio_self_attn_perturbation_mask,
+                a2v_cross_attn_perturbation_mask=a2v_cross_attn_perturbation_mask,
+                v2a_cross_attn_perturbation_mask=v2a_cross_attn_perturbation_mask,
                 audio_replicated_for_sp=audio_replicated_for_sp,
             )
 

python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py

@@ -7,6 +7,7 @@ from diffusers import FlowMatchEulerDiscreteScheduler
 
 from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
     is_ltx23_native_variant,
+    sync_ltx23_runtime_vae_markers,
 )
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
     PipelineComponentLoader,
@@ -249,6 +250,10 @@ class _BaseLTX2Pipeline(LoRAPipeline):
     def initialize_pipeline(self, server_args: ServerArgs):
         orig = self.get_module("scheduler")
         self.modules["scheduler"] = LTX2FlowMatchScheduler.from_config(orig.config)
+        sync_ltx23_runtime_vae_markers(
+            server_args.pipeline_config.vae_config.arch_config,
+            getattr(self.get_module("vae"), "config", None),
+        )
 
 
 class LTX2Pipeline(_BaseLTX2Pipeline):

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py

@@ -3,9 +3,7 @@ import copy
 import torch
 from diffusers.utils.torch_utils import randn_tensor
 
-from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
+from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import is_ltx23_native_variant
-    is_ltx23_native_variant,
-)
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import (
     LTX2DenoisingStage,
@@ -156,23 +154,59 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         """Run the distilled refinement schedule on top of the shared AV denoiser."""
         batch.extra["ltx2_phase"] = "stage2"
+        original_clean_latent_background = getattr(
+            batch, "ltx2_ti2v_clean_latent_background", None
+        )
+        is_native_ti2v = (
+            is_ltx23_native_variant(server_args.pipeline_config.vae_config.arch_config)
+            and batch.image_path is not None
+            and isinstance(batch.latents, torch.Tensor)
+        )
+        if is_native_ti2v:
+            # Official two-stage TI2V keeps the upsampled stage-2 latent as the
+            # clean background and only overwrites the conditioned frame tokens.
+            batch.ltx2_ti2v_clean_latent_background = batch.latents.detach().clone()
+        else:
+            batch.ltx2_ti2v_clean_latent_background = None
         if self._should_reset_stage2_generators(server_args):
             self._reset_stage2_generators(batch)
-        noise_scale = self.distilled_sigmas[0].to(batch.latents.device)
+        noise_scale = float(self.distilled_sigmas[0].item())
+        if is_native_ti2v:
+            prepared_latents, denoise_mask, _ = self._prepare_ltx2_ti2v_clean_state(
+                latents=batch.latents,
+                image_latent=batch.image_latent,
+                num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)),
+                zero_clean_latent=True,
+                clean_latent_background=batch.ltx2_ti2v_clean_latent_background,
+            )
+            video_noise = self._randn_like_with_batch_generators(
+                prepared_latents, batch
+            )
+            scaled_mask = (
+                denoise_mask.to(device=prepared_latents.device, dtype=torch.float32)
+                * noise_scale
+            )
+            batch.latents = (
+                video_noise * scaled_mask + prepared_latents * (1 - scaled_mask)
+            ).to(prepared_latents.dtype)
+        else:
             video_noise = self._randn_like_with_batch_generators(batch.latents, batch)
-        batch.latents = video_noise * noise_scale + batch.latents * (1 - noise_scale)
+            batch.latents = (
+                video_noise * noise_scale + batch.latents * (1 - noise_scale)
+            ).to(batch.latents.dtype)
 
         if isinstance(batch.audio_latents, torch.Tensor):
             audio_noise = self._randn_like_with_batch_generators(
                 batch.audio_latents, batch
             )
-            audio_noise_scale = noise_scale.to(
+            audio_scaled_mask = (
-                batch.audio_latents.device, batch.audio_latents.dtype
+                torch.ones_like(batch.audio_latents[..., :1], dtype=torch.float32)
+                * noise_scale
             )
             batch.audio_latents = (
-                audio_noise * audio_noise_scale
+                audio_noise * audio_scaled_mask
-                + batch.audio_latents * (1 - audio_noise_scale)
+                + batch.audio_latents * (1 - audio_scaled_mask)
-            )
+            ).to(batch.audio_latents.dtype)
         if not is_ltx23_native_variant(
             server_args.pipeline_config.vae_config.arch_config
         ):
@@ -214,5 +248,6 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
             batch.timesteps = original_batch_timesteps
             batch.num_inference_steps = original_batch_num_inference_steps
             batch.do_classifier_free_guidance = original_do_cfg
+            batch.ltx2_ti2v_clean_latent_background = original_clean_latent_background
 
         return batch

python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py

@@ -59,6 +59,8 @@ class LTX2DenoisingContext(DenoisingContext):
     latent_width: int = 0
     denoise_mask: torch.Tensor | None = None
     clean_latent: torch.Tensor | None = None
+    last_denoised_video: torch.Tensor | None = None
+    last_denoised_audio: torch.Tensor | None = None
     trajectory_audio_latents: list[torch.Tensor] = field(default_factory=list)
 
 
@@ -167,6 +169,7 @@ class LTX2DenoisingStage(DenoisingStage):
         image_latent: torch.Tensor,
         num_img_tokens: int,
         zero_clean_latent: bool,
+        clean_latent_background: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         latents = latents.clone()
         conditioned = image_latent[:, :num_img_tokens, :].to(
@@ -179,7 +182,13 @@ class LTX2DenoisingStage(DenoisingStage):
             dtype=torch.float32,
         )
         denoise_mask[:, :num_img_tokens, :] = 0.0
-        if zero_clean_latent:
+        if clean_latent_background is not None:
+            clean_latent = (
+                clean_latent_background.detach()
+                .clone()
+                .to(device=latents.device, dtype=latents.dtype)
+            )
+        elif zero_clean_latent:
             clean_latent = torch.zeros_like(latents)
         else:
             clean_latent = latents.detach().clone()
@@ -642,6 +651,14 @@ class LTX2DenoisingStage(DenoisingStage):
         if do_ti2v:
             if not (isinstance(ctx.latents, torch.Tensor) and ctx.latents.ndim == 3):
                 raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].")
+            clean_latent_background = getattr(
+                batch, "ltx2_ti2v_clean_latent_background", None
+            )
+            if not (
+                isinstance(clean_latent_background, torch.Tensor)
+                and clean_latent_background.shape == ctx.latents.shape
+            ):
+                clean_latent_background = None
             # Keep conditioned tokens clean and reuse the mask during every step update.
             ctx.latents, ctx.denoise_mask, ctx.clean_latent = (
                 self._prepare_ltx2_ti2v_clean_state(
@@ -649,6 +666,7 @@ class LTX2DenoisingStage(DenoisingStage):
                     image_latent=batch.image_latent,
                     num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)),
                     zero_clean_latent=ctx.is_ltx23_variant,
+                    clean_latent_background=clean_latent_background,
                 )
             )
         return ctx
@@ -948,6 +966,11 @@ class LTX2DenoisingStage(DenoisingStage):
             ctx.audio_latents = ctx.audio_scheduler.step(
                 model_audio, step.t_device, ctx.audio_latents, return_dict=False
             )[0]
+            if ctx.denoise_mask is not None and ctx.clean_latent is not None:
+                ctx.latents = (
+                    ctx.latents.float() * ctx.denoise_mask
+                    + ctx.clean_latent.float() * (1.0 - ctx.denoise_mask)
+                ).to(dtype=ctx.latents.dtype)
             ctx.latents = self.post_forward_for_ti2v_task(
                 batch, server_args, ctx.reserved_frames_mask, ctx.latents, ctx.z
             )
@@ -956,6 +979,32 @@ class LTX2DenoisingStage(DenoisingStage):
         encoder_hidden_states = batch.prompt_embeds[0]
         audio_encoder_hidden_states = batch.audio_prompt_embeds[0]
         encoder_attention_mask = prompt_attention_mask
+        negative_encoder_hidden_states = batch.negative_prompt_embeds[0]
+        negative_audio_encoder_hidden_states = batch.negative_audio_prompt_embeds[0]
+        negative_encoder_attention_mask = self._get_ltx_prompt_attention_mask(
+            batch,
+            is_ltx23_variant=(
+                ctx.is_ltx23_variant and not ctx.use_ltx23_legacy_one_stage
+            ),
+            negative=True,
+        )
+
+        video_skip = self._ltx2_should_skip_step(
+            step.step_index, int(stage1_guider_params["video_skip_step"])
+        )
+        audio_skip = self._ltx2_should_skip_step(
+            step.step_index, int(stage1_guider_params["audio_skip_step"])
+        )
+        need_perturbed = (
+            float(stage1_guider_params["video_stg_scale"]) != 0.0
+            or float(stage1_guider_params["audio_stg_scale"]) != 0.0
+        )
+        need_modality = (
+            float(stage1_guider_params["video_modality_scale"]) != 1.0
+            or float(stage1_guider_params["audio_modality_scale"]) != 1.0
+        )
+
+        if ctx.use_ltx23_legacy_one_stage:
             with set_forward_context(
                 current_timestep=step.step_index, attn_metadata=step.attn_metadata
             ):
@@ -966,78 +1015,21 @@ class LTX2DenoisingStage(DenoisingStage):
                         encoder_attention_mask=encoder_attention_mask,
                     )
                 )
-
-            if stage1_guider_params is not None or batch.do_classifier_free_guidance:
                 v_neg, a_v_neg = step.current_model(
                     **build_model_kwargs(
-                        encoder_hidden_states=batch.negative_prompt_embeds[0],
+                        encoder_hidden_states=negative_encoder_hidden_states,
-                        audio_encoder_hidden_states=batch.negative_audio_prompt_embeds[
+                        audio_encoder_hidden_states=negative_audio_encoder_hidden_states,
-                            0
+                        encoder_attention_mask=negative_encoder_attention_mask,
-                        ],
-                        encoder_attention_mask=self._get_ltx_prompt_attention_mask(
-                            batch,
-                            is_ltx23_variant=(
-                                ctx.is_ltx23_variant
-                                and not ctx.use_ltx23_legacy_one_stage
-                            ),
-                            negative=True,
-                        ),
                     )
                 )
-            else:
-                v_neg = None
-                a_v_neg = None
 
             v_pos = v_pos.float()
             a_v_pos = a_v_pos.float()
-        if v_neg is not None:
             v_neg = v_neg.float()
-        if a_v_neg is not None:
             a_v_neg = a_v_neg.float()
 
-        sigma_val = float(sigma.item())
+            v_ptb = None
-        video_sigma_for_x0: float | torch.Tensor = sigma_val
+            a_v_ptb = None
-        if ctx.denoise_mask is not None:
-            video_sigma_for_x0 = sigma.to(
-                device=ctx.latents.device, dtype=torch.float32
-            ) * ctx.denoise_mask.squeeze(-1)
-
-        denoised_video = self._ltx2_velocity_to_x0(
-            ctx.latents, v_pos, video_sigma_for_x0
-        )
-        denoised_audio = self._ltx2_velocity_to_x0(
-            ctx.audio_latents, a_v_pos, sigma_val
-        )
-        denoised_video_neg = None
-        denoised_audio_neg = None
-        denoised_video_perturbed = None
-        denoised_audio_perturbed = None
-        denoised_video_modality = None
-        denoised_audio_modality = None
-
-        if (
-            (stage1_guider_params is not None or batch.do_classifier_free_guidance)
-            and v_neg is not None
-            and a_v_neg is not None
-        ):
-            denoised_video_neg = self._ltx2_velocity_to_x0(
-                ctx.latents, v_neg, video_sigma_for_x0
-            )
-            denoised_audio_neg = self._ltx2_velocity_to_x0(
-                ctx.audio_latents, a_v_neg, sigma_val
-            )
-        if stage1_guider_params is not None:
-            video_skip = self._ltx2_should_skip_step(
-                step.step_index, int(stage1_guider_params["video_skip_step"])
-            )
-            audio_skip = self._ltx2_should_skip_step(
-                step.step_index, int(stage1_guider_params["audio_skip_step"])
-            )
-
-            need_perturbed = (
-                float(stage1_guider_params["video_stg_scale"]) != 0.0
-                or float(stage1_guider_params["audio_stg_scale"]) != 0.0
-            )
             if need_perturbed:
                 with set_forward_context(
                     current_timestep=step.step_index, attn_metadata=step.attn_metadata
@@ -1055,17 +1047,11 @@ class LTX2DenoisingStage(DenoisingStage):
                             ),
                         )
                     )
-                denoised_video_perturbed = self._ltx2_velocity_to_x0(
+                v_ptb = v_ptb.float()
-                    ctx.latents, v_ptb.float(), video_sigma_for_x0
+                a_v_ptb = a_v_ptb.float()
-                )
-                denoised_audio_perturbed = self._ltx2_velocity_to_x0(
-                    ctx.audio_latents, a_v_ptb.float(), sigma_val
-                )
 
-            need_modality = (
+            v_mod = None
-                float(stage1_guider_params["video_modality_scale"]) != 1.0
+            a_v_mod = None
-                or float(stage1_guider_params["audio_modality_scale"]) != 1.0
-            )
             if need_modality:
                 with set_forward_context(
                     current_timestep=step.step_index, attn_metadata=step.attn_metadata
@@ -1079,21 +1065,354 @@ class LTX2DenoisingStage(DenoisingStage):
                             disable_v2a_cross_attn=True,
                         )
                     )
-                denoised_video_modality = self._ltx2_velocity_to_x0(
+                v_mod = v_mod.float()
-                    ctx.latents, v_mod.float(), video_sigma_for_x0
+                a_v_mod = a_v_mod.float()
+        else:
+            # NOTE: this flag must be identical across all SP ranks so that
+            # every rank executes the same number of model-forward calls (each
+            # of which contains NCCL collectives).
+            # _should_apply_ltx2_ti2v() is SP-rank-dependent (only the rank owning the first latent
+            # frame returns True), so we must NOT use it here.
+            # Instead we check the rank-invariant attribute that is always set on every
+            # rank when the request is a TI2V request.
+            use_split_two_stage_ti2v_guider = (
+                server_args.pipeline_class_name == "LTX2TwoStagePipeline"
+                and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0
             )
-                denoised_audio_modality = self._ltx2_velocity_to_x0(
+
-                    ctx.audio_latents, a_v_mod.float(), sigma_val
+            def cat_or_none(items: list[torch.Tensor | None]) -> torch.Tensor | None:
+                if items[0] is None:
+                    return None
+                return torch.cat(items, dim=0)
+
+            pass_specs: list[
+                tuple[
+                    str,
+                    torch.Tensor,
+                    torch.Tensor,
+                    torch.Tensor | None,
+                    dict[str, object],
+                ]
+            ] = [
+                (
+                    "cond",
+                    encoder_hidden_states,
+                    audio_encoder_hidden_states,
+                    encoder_attention_mask,
+                    {
+                        "skip_video_self_attn_blocks": (),
+                        "skip_audio_self_attn_blocks": (),
+                        "skip_a2v_cross_attn": False,
+                        "skip_v2a_cross_attn": False,
+                    },
+                ),
+                (
+                    "neg",
+                    negative_encoder_hidden_states,
+                    negative_audio_encoder_hidden_states,
+                    negative_encoder_attention_mask,
+                    {
+                        "skip_video_self_attn_blocks": (),
+                        "skip_audio_self_attn_blocks": (),
+                        "skip_a2v_cross_attn": False,
+                        "skip_v2a_cross_attn": False,
+                    },
+                ),
+            ]
+            if need_perturbed:
+                pass_specs.append(
+                    (
+                        "perturbed",
+                        encoder_hidden_states,
+                        audio_encoder_hidden_states,
+                        encoder_attention_mask,
+                        {
+                            "skip_video_self_attn_blocks": tuple(
+                                stage1_guider_params["video_stg_blocks"]
+                            ),
+                            "skip_audio_self_attn_blocks": tuple(
+                                stage1_guider_params["audio_stg_blocks"]
+                            ),
+                            "skip_a2v_cross_attn": False,
+                            "skip_v2a_cross_attn": False,
+                        },
+                    )
+                )
+            if need_modality:
+                pass_specs.append(
+                    (
+                        "modality",
+                        encoder_hidden_states,
+                        audio_encoder_hidden_states,
+                        encoder_attention_mask,
+                        {
+                            "skip_video_self_attn_blocks": (),
+                            "skip_audio_self_attn_blocks": (),
+                            "skip_a2v_cross_attn": True,
+                            "skip_v2a_cross_attn": True,
+                        },
+                    )
+                )
+
+            num_passes = len(pass_specs)
+            expanded_batch_size = batch_size * num_passes
+            perturbation_configs = tuple(
+                perturbation_config
+                for _, _, _, _, perturbation_config in pass_specs
+                for _ in range(batch_size)
+            )
+            batched_hidden_states = self._repeat_batch_dim(
+                latent_model_input, expanded_batch_size
+            )
+            batched_audio_hidden_states = self._repeat_batch_dim(
+                audio_latent_model_input, expanded_batch_size
+            )
+            batched_encoder_hidden_states = torch.cat(
+                [item[1] for item in pass_specs], dim=0
+            )
+            batched_audio_encoder_hidden_states = torch.cat(
+                [item[2] for item in pass_specs], dim=0
+            )
+            batched_timestep_video = self._repeat_batch_dim(
+                timestep_video, expanded_batch_size
+            )
+            batched_timestep_audio = self._repeat_batch_dim(
+                timestep_audio, expanded_batch_size
+            )
+            batched_prompt_timestep_video = (
+                None
+                if prompt_timestep_video is None
+                else self._repeat_batch_dim(prompt_timestep_video, expanded_batch_size)
+            )
+            batched_prompt_timestep_audio = (
+                None
+                if prompt_timestep_audio is None
+                else self._repeat_batch_dim(prompt_timestep_audio, expanded_batch_size)
+            )
+            batched_encoder_attention_mask = cat_or_none(
+                [item[3] for item in pass_specs]
+            )
+            batched_audio_encoder_attention_mask = cat_or_none(
+                [item[3] for item in pass_specs]
+            )
+            batched_video_coords = (
+                None
+                if video_coords is None
+                else self._repeat_batch_dim(video_coords, expanded_batch_size)
+            )
+            batched_audio_coords = (
+                None
+                if audio_coords is None
+                else self._repeat_batch_dim(audio_coords, expanded_batch_size)
+            )
+            batched_video_self_attention_mask = (
+                None
+                if video_self_attention_mask is None
+                else self._repeat_batch_dim(
+                    video_self_attention_mask, expanded_batch_size
+                )
+            )
+            batched_audio_self_attention_mask = (
+                None
+                if audio_self_attention_mask is None
+                else self._repeat_batch_dim(
+                    audio_self_attention_mask, expanded_batch_size
+                )
+            )
+            batched_a2v_cross_attention_mask = (
+                None
+                if a2v_cross_attention_mask is None
+                else self._repeat_batch_dim(
+                    a2v_cross_attention_mask, expanded_batch_size
+                )
+            )
+            batched_v2a_cross_attention_mask = (
+                None
+                if v2a_cross_attention_mask is None
+                else self._repeat_batch_dim(
+                    v2a_cross_attention_mask, expanded_batch_size
+                )
+            )
+            if use_split_two_stage_ti2v_guider:
+                split_sizes = [1] * expanded_batch_size
+
+                def split_or_none(
+                    tensor: torch.Tensor | None,
+                ) -> list[torch.Tensor | None]:
+                    if tensor is None:
+                        return [None] * len(split_sizes)
+                    return list(tensor.split(split_sizes, dim=0))
+
+                batched_video_chunks = []
+                batched_audio_chunks = []
+                with set_forward_context(
+                    current_timestep=step.step_index, attn_metadata=step.attn_metadata
+                ):
+                    for (
+                        hidden_states_chunk,
+                        audio_hidden_states_chunk,
+                        encoder_hidden_states_chunk,
+                        audio_encoder_hidden_states_chunk,
+                        timestep_video_chunk,
+                        timestep_audio_chunk,
+                        prompt_timestep_video_chunk,
+                        prompt_timestep_audio_chunk,
+                        encoder_attention_mask_chunk,
+                        audio_encoder_attention_mask_chunk,
+                        video_coords_chunk,
+                        audio_coords_chunk,
+                        video_self_attention_mask_chunk,
+                        audio_self_attention_mask_chunk,
+                        a2v_cross_attention_mask_chunk,
+                        v2a_cross_attention_mask_chunk,
+                        perturbation_config_chunk,
+                    ) in zip(
+                        batched_hidden_states.split(split_sizes, dim=0),
+                        batched_audio_hidden_states.split(split_sizes, dim=0),
+                        batched_encoder_hidden_states.split(split_sizes, dim=0),
+                        batched_audio_encoder_hidden_states.split(split_sizes, dim=0),
+                        batched_timestep_video.split(split_sizes, dim=0),
+                        batched_timestep_audio.split(split_sizes, dim=0),
+                        split_or_none(batched_prompt_timestep_video),
+                        split_or_none(batched_prompt_timestep_audio),
+                        split_or_none(batched_encoder_attention_mask),
+                        split_or_none(batched_audio_encoder_attention_mask),
+                        split_or_none(batched_video_coords),
+                        split_or_none(batched_audio_coords),
+                        split_or_none(batched_video_self_attention_mask),
+                        split_or_none(batched_audio_self_attention_mask),
+                        split_or_none(batched_a2v_cross_attention_mask),
+                        split_or_none(batched_v2a_cross_attention_mask),
+                        ((cfg,) for cfg in perturbation_configs),
+                        strict=True,
+                    ):
+                        video_chunk, audio_chunk = step.current_model(
+                            hidden_states=hidden_states_chunk,
+                            audio_hidden_states=audio_hidden_states_chunk,
+                            encoder_hidden_states=encoder_hidden_states_chunk,
+                            audio_encoder_hidden_states=audio_encoder_hidden_states_chunk,
+                            timestep=timestep_video_chunk,
+                            audio_timestep=timestep_audio_chunk,
+                            prompt_timestep=prompt_timestep_video_chunk,
+                            audio_prompt_timestep=prompt_timestep_audio_chunk,
+                            encoder_attention_mask=encoder_attention_mask_chunk,
+                            audio_encoder_attention_mask=audio_encoder_attention_mask_chunk,
+                            num_frames=ctx.latent_num_frames_for_model,
+                            height=ctx.latent_height,
+                            width=ctx.latent_width,
+                            fps=batch.fps,
+                            audio_num_frames=audio_num_frames_latent,
+                            video_coords=video_coords_chunk,
+                            audio_coords=audio_coords_chunk,
+                            video_self_attention_mask=video_self_attention_mask_chunk,
+                            audio_self_attention_mask=audio_self_attention_mask_chunk,
+                            a2v_cross_attention_mask=a2v_cross_attention_mask_chunk,
+                            v2a_cross_attention_mask=v2a_cross_attention_mask_chunk,
+                            audio_replicated_for_sp=ctx.replicate_audio_for_sp,
+                            perturbation_configs=perturbation_config_chunk,
+                            return_latents=False,
+                            return_dict=False,
+                        )
+                        batched_video_chunks.append(video_chunk)
+                        batched_audio_chunks.append(audio_chunk)
+
+                batched_video = torch.cat(batched_video_chunks, dim=0)
+                batched_audio = torch.cat(batched_audio_chunks, dim=0)
+            else:
+                with set_forward_context(
+                    current_timestep=step.step_index, attn_metadata=step.attn_metadata
+                ):
+                    batched_video, batched_audio = step.current_model(
+                        hidden_states=batched_hidden_states,
+                        audio_hidden_states=batched_audio_hidden_states,
+                        encoder_hidden_states=batched_encoder_hidden_states,
+                        audio_encoder_hidden_states=batched_audio_encoder_hidden_states,
+                        timestep=batched_timestep_video,
+                        audio_timestep=batched_timestep_audio,
+                        prompt_timestep=batched_prompt_timestep_video,
+                        audio_prompt_timestep=batched_prompt_timestep_audio,
+                        encoder_attention_mask=batched_encoder_attention_mask,
+                        audio_encoder_attention_mask=batched_audio_encoder_attention_mask,
+                        num_frames=ctx.latent_num_frames_for_model,
+                        height=ctx.latent_height,
+                        width=ctx.latent_width,
+                        fps=batch.fps,
+                        audio_num_frames=audio_num_frames_latent,
+                        video_coords=batched_video_coords,
+                        audio_coords=batched_audio_coords,
+                        video_self_attention_mask=batched_video_self_attention_mask,
+                        audio_self_attention_mask=batched_audio_self_attention_mask,
+                        a2v_cross_attention_mask=batched_a2v_cross_attention_mask,
+                        v2a_cross_attention_mask=batched_v2a_cross_attention_mask,
+                        audio_replicated_for_sp=ctx.replicate_audio_for_sp,
+                        perturbation_configs=perturbation_configs,
+                        return_latents=False,
+                        return_dict=False,
+                    )
+
+            batched_video = batched_video.float()
+            batched_audio = batched_audio.float()
+            pass_outputs = {
+                pass_name: (
+                    video_chunk,
+                    audio_chunk,
+                )
+                for (pass_name, _, _, _, _), video_chunk, audio_chunk in zip(
+                    pass_specs,
+                    batched_video.chunk(num_passes, dim=0),
+                    batched_audio.chunk(num_passes, dim=0),
+                    strict=True,
+                )
+            }
+            v_pos, a_v_pos = pass_outputs["cond"]
+            v_neg, a_v_neg = pass_outputs["neg"]
+            v_ptb, a_v_ptb = pass_outputs.get("perturbed", (None, None))
+            v_mod, a_v_mod = pass_outputs.get("modality", (None, None))
+
+        sigma_val = float(sigma.item())
+        video_sigma_for_x0: float | torch.Tensor = sigma_val
+        if ctx.denoise_mask is not None:
+            video_sigma_for_x0 = sigma.to(
+                device=ctx.latents.device, dtype=torch.float32
+            ) * ctx.denoise_mask.squeeze(-1)
+
+        denoised_video = self._ltx2_velocity_to_x0(
+            ctx.latents, v_pos, video_sigma_for_x0
+        )
+        denoised_audio = self._ltx2_velocity_to_x0(
+            ctx.audio_latents, a_v_pos, sigma_val
+        )
+        denoised_video_neg = self._ltx2_velocity_to_x0(
+            ctx.latents, v_neg, video_sigma_for_x0
+        )
+        denoised_audio_neg = self._ltx2_velocity_to_x0(
+            ctx.audio_latents, a_v_neg, sigma_val
+        )
+        denoised_video_perturbed = (
+            None
+            if v_ptb is None
+            else self._ltx2_velocity_to_x0(ctx.latents, v_ptb, video_sigma_for_x0)
+        )
+        denoised_audio_perturbed = (
+            None
+            if a_v_ptb is None
+            else self._ltx2_velocity_to_x0(ctx.audio_latents, a_v_ptb, sigma_val)
+        )
+        denoised_video_modality = (
+            None
+            if v_mod is None
+            else self._ltx2_velocity_to_x0(ctx.latents, v_mod, video_sigma_for_x0)
+        )
+        denoised_audio_modality = (
+            None
+            if a_v_mod is None
+            else self._ltx2_velocity_to_x0(ctx.audio_latents, a_v_mod, sigma_val)
         )
 
         if not video_skip:
             denoised_video = self._ltx2_calculate_guided_x0(
                 cond=denoised_video,
-                    uncond_text=(
+                uncond_text=denoised_video_neg,
-                        denoised_video_neg
-                        if denoised_video_neg is not None
-                        else denoised_video
-                    ),
                 uncond_perturbed=(
                     denoised_video_perturbed
                     if denoised_video_perturbed is not None
@@ -1109,14 +1428,14 @@ class LTX2DenoisingStage(DenoisingStage):
                 rescale_scale=float(stage1_guider_params["video_rescale_scale"]),
                 modality_scale=float(stage1_guider_params["video_modality_scale"]),
             )
+            ctx.last_denoised_video = denoised_video
+        elif ctx.last_denoised_video is not None:
+            denoised_video = ctx.last_denoised_video
+
         if not audio_skip:
             denoised_audio = self._ltx2_calculate_guided_x0(
                 cond=denoised_audio,
-                    uncond_text=(
+                uncond_text=denoised_audio_neg,
-                        denoised_audio_neg
-                        if denoised_audio_neg is not None
-                        else denoised_audio
-                    ),
                 uncond_perturbed=(
                     denoised_audio_perturbed
                     if denoised_audio_perturbed is not None
@@ -1132,23 +1451,15 @@ class LTX2DenoisingStage(DenoisingStage):
                 rescale_scale=float(stage1_guider_params["audio_rescale_scale"]),
                 modality_scale=float(stage1_guider_params["audio_modality_scale"]),
             )
-        elif (
+            ctx.last_denoised_audio = denoised_audio
-            batch.do_classifier_free_guidance
+        elif ctx.last_denoised_audio is not None:
-            and denoised_video_neg is not None
+            denoised_audio = ctx.last_denoised_audio
-            and denoised_audio_neg is not None
-        ):
-            denoised_video = denoised_video + (batch.guidance_scale - 1.0) * (
-                denoised_video - denoised_video_neg
-            )
-            denoised_audio = denoised_audio + (batch.guidance_scale - 1.0) * (
-                denoised_audio - denoised_audio_neg
-            )
 
         if ctx.denoise_mask is not None and ctx.clean_latent is not None:
             denoised_video = (
                 denoised_video * ctx.denoise_mask
                 + ctx.clean_latent.float() * (1.0 - ctx.denoise_mask)
-            )
+            ).to(denoised_video.dtype)
 
         # 6. Convert x0 predictions back to velocity and update both latent streams.
         if sigma_val == 0.0:

python/sglang/multimodal_gen/test/server/accuracy_config.py

@@ -68,11 +68,6 @@ SKIP_COMPONENTS: Dict[str, Dict[ComponentType, ComponentSkip]] = {
             "HF AutoencoderDC checkpoint leaves required to_qkv_multiscale weights missing, so VAE transfer would compare against partially initialized reference weights"
         )
     },
-    "mova_360p_1gpu": {
-        ComponentType.TRANSFORMER: ComponentSkip(
-            "HF reference transformer cannot be materialized from the video_dit repo layout"
-        )
-    },
     "qwen_image_t2i_cache_dit_enabled": {
         ComponentType.VAE: ComponentSkip(
             "Representative VAE accuracy is already covered by qwen_image_t2i for the same source component and topology"
@@ -324,14 +319,6 @@ SKIP_COMPONENTS: Dict[str, Dict[ComponentType, ComponentSkip]] = {
             "Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.31) after 100% matched weight transfer"
         ),
     },
-    "mova_360p_ring2_uly1": {
-        ComponentType.TRANSFORMER: ComponentSkip(
-            "HF reference transformer cannot be materialized from the MOVA video_dit repo layout"
-        ),
-        ComponentType.TEXT_ENCODER: ComponentSkip(
-            "Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.31) after 100% matched weight transfer"
-        ),
-    },
     "flux_image_t2i_2_gpus": {
         ComponentType.TEXT_ENCODER: ComponentSkip(
             "Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.47) after 100% matched weight transfer"

python/sglang/multimodal_gen/test/server/consistency_threshold.json

@@ -79,12 +79,6 @@
             "psnr_threshold": 28.0,
             "mean_abs_diff_threshold": 8.0
         },
-        "mova_360p_1gpu": {
-            "clip_threshold": 0.90,
-            "ssim_threshold": 0.87,
-            "psnr_threshold": 24.0,
-            "mean_abs_diff_threshold": 10.0
-        },
         "wan2_1_t2v_1_3b_lora_1gpu": {
             "clip_threshold": 0.88,
             "ssim_threshold": 0.75,
@@ -223,12 +217,6 @@
             "psnr_threshold": 22.2,
             "mean_abs_diff_threshold": 10.0
         },
-        "mova_360p_ring2_uly1": {
-            "clip_threshold": 0.90,
-            "ssim_threshold": 0.91,
-            "psnr_threshold": 24.0,
-            "mean_abs_diff_threshold": 10.0
-        },
         "wan2_1_i2v_14b_480P_2gpu": {
             "clip_threshold": 0.76,
             "ssim_threshold": 0.51,

python/sglang/multimodal_gen/test/server/perf_baselines.json

@@ -2616,7 +2616,63 @@
             },
             "expected_e2e_ms": 34384.39,
             "expected_avg_denoise_ms": 782.76,
-            "expected_median_denoise_ms": 806.95
+            "expected_median_denoise_ms": 806.95,
+            "estimated_full_test_time_s": 216.7
+        },
+        "ltx_2_3_two_stage_ti2v_2gpus": {
+            "stages_ms": {
+                "InputValidationStage": 3.61,
+                "TextEncodingStage": 1775.95,
+                "LTX2TextConnectorStage": 27.44,
+                "LTX2HalveResolutionStage": 0.06,
+                "LTX2LoRASwitchStage": 107.23,
+                "LTX2SigmaPreparationStage": 0.41,
+                "TimestepPreparationStage": 24.7,
+                "LTX2AVLatentPreparationStage": 0.2,
+                "LTX2AVDenoisingStage": 26851.54,
+                "LTX2UpsampleStage": 408.81,
+                "LTX2RefinementStage": 1561.91,
+                "LTX2AVDecodingStage": 862.29,
+                "per_frame_generation": null
+            },
+            "denoise_step_ms": {
+                "0": 812.4,
+                "1": 817.22,
+                "2": 1159.79,
+                "3": 852.65,
+                "4": 857.69,
+                "5": 864.02,
+                "6": 868.72,
+                "7": 859.11,
+                "8": 836.39,
+                "9": 841.12,
+                "10": 1093.34,
+                "11": 902.56,
+                "12": 830.56,
+                "13": 818.54,
+                "14": 837.61,
+                "15": 837.03,
+                "16": 853.87,
+                "17": 830.49,
+                "18": 833.02,
+                "19": 1179.55,
+                "20": 846.55,
+                "21": 856.85,
+                "22": 887.05,
+                "23": 836.35,
+                "24": 843.8,
+                "25": 836.38,
+                "26": 821.39,
+                "27": 826.6,
+                "28": 1165.21,
+                "29": 833.04,
+                "30": 540.96,
+                "31": 334.98,
+                "32": 312.19
+            },
+            "expected_e2e_ms": 40552.46,
+            "expected_avg_denoise_ms": 840.21,
+            "expected_median_denoise_ms": 837.61
         }
     }
 }

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -549,17 +549,6 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
             extras={"enable_upscaling": True, "upscaling_scale": 4},
         ),
     ),
-    DiffusionTestCase(
-        "mova_360p_1gpu",
-        DiffusionServerArgs(
-            model_path=DEFAULT_MOVA_360P_MODEL_NAME_FOR_TEST,
-            modality="video",
-            num_gpus=1,
-            dit_layerwise_offload=True,
-        ),
-        TI2V_sampling_params,
-        run_perf_check=False,
-    ),
 ]
 
 HUNYUAN3D_SHAPE_sampling_params = DiffusionSamplingParams(
@@ -917,19 +906,6 @@ TWO_GPU_CASES_A = [
         TI2V_sampling_params,
         run_perf_check=False,
     ),
-    DiffusionTestCase(
-        "mova_360p_ring2_uly1",
-        DiffusionServerArgs(
-            model_path=DEFAULT_MOVA_360P_MODEL_NAME_FOR_TEST,
-            modality="video",
-            num_gpus=2,
-            ring_degree=2,
-            ulysses_degree=1,
-            dit_layerwise_offload=True,
-        ),
-        TI2V_sampling_params,
-        run_perf_check=False,
-    ),
     DiffusionTestCase(
         "ltx_2_two_stage_t2v",
         DiffusionServerArgs(
@@ -940,6 +916,16 @@ TWO_GPU_CASES_A = [
         ),
         T2V_sampling_params,
     ),
+    DiffusionTestCase(
+        "ltx_2_3_two_stage_ti2v_2gpus",
+        DiffusionServerArgs(
+            model_path="Lightricks/LTX-2.3",
+            modality="video",
+            num_gpus=2,
+            extras=["--pipeline-class-name LTX2TwoStagePipeline"],
+        ),
+        TI2V_sampling_params,
+    ),
 ]
 
 TWO_GPU_CASES_B = [
@@ -1098,7 +1084,6 @@ ACCURACY_ONE_GPU_CASES_A_IDS = (
     "qwen_image_edit_2511_ti2i",
     "qwen_image_layered_i2i",
     "flux_2_image_t2i_upscaling_4x",
-    "mova_360p_1gpu",
 )
 
 ACCURACY_ONE_GPU_CASES_B_IDS = (
@@ -1130,7 +1115,6 @@ ACCURACY_TWO_GPU_CASES_A_IDS = (
     "fsdp-inference",
     "mova_360p_tp2",
     "mova_360p_ring1_uly2",
-    "mova_360p_ring2_uly1",
     "ltx_2_two_stage_t2v",
 )
 

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

@@ -1,8 +1,14 @@
 import argparse
 import math
 import unittest
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
+from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
+    LTX2PipelineConfig,
+    is_ltx23_native_variant,
+    sync_ltx23_runtime_vae_markers,
+)
 from sglang.multimodal_gen.configs.sample.diffusers_generic import (
     DiffusersGenericSamplingParams,
 )
@@ -141,6 +147,33 @@ class TestSamplingParamsSubclass(unittest.TestCase):
                     expected,
                 )
 
+    def test_ltx23_runtime_vae_markers_sync_variant_and_decoder_metadata(self):
+        arch_config = LTX2PipelineConfig().vae_config.arch_config
+
+        self.assertFalse(is_ltx23_native_variant(arch_config))
+        self.assertEqual(arch_config.video_decoder_variant, "ltx_2")
+        self.assertEqual(arch_config.condition_encoder_subdir, "")
+
+        sync_ltx23_runtime_vae_markers(
+            arch_config,
+            SimpleNamespace(
+                arch_config=SimpleNamespace(
+                    ltx_variant="ltx_2_3",
+                    condition_encoder_subdir="ltx23_image_encoder",
+                    video_decoder_variant="ltx_2_3",
+                    video_decoder_config={"_class_name": "AutoencoderKLLTX2Video"},
+                )
+            ),
+        )
+
+        self.assertTrue(is_ltx23_native_variant(arch_config))
+        self.assertEqual(arch_config.condition_encoder_subdir, "ltx23_image_encoder")
+        self.assertEqual(arch_config.video_decoder_variant, "ltx_2_3")
+        self.assertEqual(
+            arch_config.video_decoder_config,
+            {"_class_name": "AutoencoderKLLTX2Video"},
+        )
+
 
 class TestSamplingParamsCliArgs(unittest.TestCase):
     def _parse_cli_kwargs(self, argv: list[str]) -> dict:

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

@@ -0,0 +1,48 @@
+import unittest
+
+import torch
+
+from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import (
+    _backfill_ltx2_audio_vae_latent_stats,
+)
+
+
+class TestVAELoader(unittest.TestCase):
+    def test_backfill_ltx2_audio_vae_latent_stats_maps_official_keys(self):
+        loaded = {
+            "per_channel_statistics.mean-of-means": torch.tensor([1.0, 2.0]),
+            "per_channel_statistics.std-of-means": torch.tensor([3.0, 4.0]),
+        }
+
+        _backfill_ltx2_audio_vae_latent_stats(loaded, "audio_vae")
+
+        self.assertTrue(torch.equal(loaded["latents_mean"], torch.tensor([1.0, 2.0])))
+        self.assertTrue(torch.equal(loaded["latents_std"], torch.tensor([3.0, 4.0])))
+
+    def test_backfill_ltx2_audio_vae_latent_stats_does_not_override_existing(self):
+        loaded = {
+            "per_channel_statistics.mean-of-means": torch.tensor([1.0, 2.0]),
+            "per_channel_statistics.std-of-means": torch.tensor([3.0, 4.0]),
+            "latents_mean": torch.tensor([5.0, 6.0]),
+            "latents_std": torch.tensor([7.0, 8.0]),
+        }
+
+        _backfill_ltx2_audio_vae_latent_stats(loaded, "audio_vae")
+
+        self.assertTrue(torch.equal(loaded["latents_mean"], torch.tensor([5.0, 6.0])))
+        self.assertTrue(torch.equal(loaded["latents_std"], torch.tensor([7.0, 8.0])))
+
+    def test_backfill_ltx2_audio_vae_latent_stats_skips_non_audio_vae(self):
+        loaded = {
+            "per_channel_statistics.mean-of-means": torch.tensor([1.0]),
+            "per_channel_statistics.std-of-means": torch.tensor([2.0]),
+        }
+
+        _backfill_ltx2_audio_vae_latent_stats(loaded, "vae")
+
+        self.assertNotIn("latents_mean", loaded)
+        self.assertNotIn("latents_std", loaded)
+
+
+if __name__ == "__main__":
+    unittest.main()

scripts/ci/utils/diffusion/comparison_configs.json

@@ -136,18 +136,16 @@
             }
         },
         {
-            "id": "ltx2.3_twostage_t2v_2gpus",
+            "id": "ltx2.3_twostage_ti2v_2gpus",
             "model": "Lightricks/LTX-2.3",
-            "task": "text-to-video",
+            "task": "image-to-video",
-            "prompt": "A cat and a dog baking a cake together in a kitchen.",
+            "prompt": "The cat starts walking slowly towards the camera.",
-            "width": 768,
-            "height": 512,
             "num_frames": 121,
             "seed": 42,
             "num_gpus": 2,
             "frameworks": {
                 "sglang": {
-                    "serve_args": "--enable-torch-compile --warmup --enable-cfg-parallel --pipeline-class-name LTX2TwoStagePipeline",
+                    "serve_args": "--enable-torch-compile --warmup --pipeline-class-name LTX2TwoStagePipeline",
                     "extra_env": {}
                 }
             }