[diffusion] chore: align LTX-2 with official (#24313)

.github/workflows/diffusion-ci-gt-gen.yml

@@ -39,10 +39,15 @@ on:
         default: ''
         type: string
       official_source_group:
-        description: 'Official GT source group filter: all, diffusers, wan21, or ltx23. Used only when run_official_cases is true.'
+        description: 'Official GT source group filter: all, diffusers, wan21, ltx, or ltx23. Used only when run_official_cases is true.'
         required: false
         default: ''
         type: string
+      ci_data_ref:
+        description: 'ci-data ref to use for repro scripts when running official GT cases.'
+        required: false
+        default: 'main'
+        type: string
 
 concurrency:
   group: diffusion-ci-gt-gen-${{ github.ref }}-${{ inputs.output_name || inputs.case_ids || inputs.official_case_ids || inputs.official_source_group || inputs.run_official_cases || 'default' }}
@@ -90,6 +95,13 @@ jobs:
                   "zimage_image_t2i",
               ],
               "wan21": ["wan2_1_t2v_1.3b"],
+              "ltx": [
+                  "ltx_2_two_stage_t2v",
+                  "ltx_2.3_two_stage_t2v_2gpus",
+                  "ltx_2_3_two_stage_ti2v_2gpus",
+                  "ltx_2.3_one_stage_ti2v",
+                  "ltx_2_3_hq_pipeline",
+              ],
               "ltx23": [
                   "ltx_2.3_two_stage_t2v_2gpus",
                   "ltx_2.3_one_stage_ti2v",
@@ -97,7 +109,7 @@ jobs:
           }
           source_group = os.environ["OFFICIAL_SOURCE_GROUP"].strip() or "all"
           if source_group == "all":
-              selected_groups = list(groups)
+              selected_groups = ["diffusers", "wan21", "ltx"]
           elif source_group in groups:
               selected_groups = [source_group]
           else:
@@ -109,7 +121,6 @@ jobs:
               "flux_2_image_t2i",
               "flux_2_klein_image_t2i",
               "flux_2_ti2i",
-              "ltx_2.3_two_stage_t2v_2gpus",
           }
           include = []
           for group in selected_groups:
@@ -154,9 +165,13 @@ jobs:
         uses: actions/checkout@v4
         with:
           repository: sgl-project/ci-data
-          ref: main
+          ref: ${{ inputs.ci_data_ref || 'main' }}
           path: ci-data
           token: ${{ secrets.GH_PAT_FOR_NIGHTLY_CI_DATA }}
+          sparse-checkout: |
+            diffusion-ci/repro_scripts
+            diffusion-ci/consistency_gt/official_generated/case_map.json
+          sparse-checkout-cone-mode: false
 
       - name: Prepare sgl-kernel/dist for prebuilt wheel
         if: inputs.kernel_artifact_run_id != ''
@@ -179,8 +194,8 @@ jobs:
           CUSTOM_BUILD_SGL_KERNEL="${{ inputs.kernel_artifact_run_id != '' && 'true' || 'false' }}" \
             bash scripts/ci/cuda/ci_install_dependency.sh diffusion
 
-      - name: Install official LTX-2.3 repro dependencies
+      - name: Install official LTX repro dependencies
-        if: matrix.source_group == 'ltx23'
+        if: matrix.source_group == 'ltx' || matrix.source_group == 'ltx23'
         run: |
           UV_SYSTEM_PYTHON=1 uv pip install \
             "transformers==4.52.4" \
@@ -201,7 +216,7 @@ jobs:
           git -C /tmp/mmgen-official-code/Wan2.1 rev-parse HEAD
 
       - name: Checkout official LTX-2 repo
-        if: matrix.source_group == 'ltx23'
+        if: matrix.source_group == 'ltx' || matrix.source_group == 'ltx23'
         run: |
           mkdir -p /tmp/mmgen-official-code
           if [ ! -d /tmp/mmgen-official-code/LTX-2/.git ]; then
@@ -211,65 +226,6 @@ jobs:
           git -C /tmp/mmgen-official-code/LTX-2 checkout 41d924371612b692c0fd1e4d9d94c3dfb3c02cb3
           git -C /tmp/mmgen-official-code/LTX-2 rev-parse HEAD
 
-      - name: Prepare official LTX-2.3 assets
-        if: matrix.source_group == 'ltx23'
-        env:
-          HF_TOKEN: ${{ secrets.SGLANG_DIFFUSION_CI_HF_TOKEN || secrets.HF_TOKEN }}
-          HUGGING_FACE_HUB_TOKEN: ${{ secrets.SGLANG_DIFFUSION_CI_HF_TOKEN || secrets.HF_TOKEN }}
-        run: |
-          python - <<'PY'
-          import os
-          from pathlib import Path
-
-          from huggingface_hub import hf_hub_download, snapshot_download
-
-          root = Path.home() / ".cache/sgl_diffusion/materialized_models/Lightricks__LTX-2.3-official-assets"
-          root.mkdir(parents=True, exist_ok=True)
-          text_encoder = root / "text_encoder"
-          tokenizer = root / "tokenizer"
-          text_encoder.mkdir(exist_ok=True)
-          tokenizer.mkdir(exist_ok=True)
-
-          gemma_root = Path(
-              snapshot_download(
-                  repo_id="google/gemma-3-12b-it-qat-q4_0-unquantized",
-                  ignore_patterns=["*.onnx", "*.msgpack"],
-                  max_workers=8,
-              )
-          )
-          tokenizer_names = {
-              "tokenizer.json",
-              "tokenizer.model",
-              "tokenizer_config.json",
-              "special_tokens_map.json",
-          }
-
-          def link(src: Path, dst: Path) -> None:
-              if dst.exists() or dst.is_symlink():
-                  return
-              try:
-                  dst.symlink_to(src)
-              except OSError:
-                  os.link(src, dst)
-
-          for src in gemma_root.iterdir():
-              if not src.is_file():
-                  continue
-              link(src, text_encoder / src.name)
-              if src.name in tokenizer_names:
-                  link(src, tokenizer / src.name)
-
-          for filename in (
-              "ltx-2.3-22b-distilled-lora-384.safetensors",
-              "ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
-          ):
-              src = Path(hf_hub_download("Lightricks/LTX-2.3", filename))
-              link(src, root / filename)
-
-          print(root)
-          print(sorted(p.name for p in root.iterdir()))
-          PY
-
       - name: Generate official output
         env:
           HF_TOKEN: ${{ secrets.SGLANG_DIFFUSION_CI_HF_TOKEN || secrets.HF_TOKEN }}
@@ -301,17 +257,22 @@ jobs:
               --dtype bf16 \
               --device-map none \
               --generator-device cuda
-          elif [ "$SOURCE_GROUP" = "ltx23" ]; then
+          elif [ "$SOURCE_GROUP" = "ltx" ] || [ "$SOURCE_GROUP" = "ltx23" ]; then
             cd python
+            extra_ltx_args=()
+            if [ "$CASE_ID" = "ltx_2_3_hq_pipeline" ]; then
+              extra_ltx_args+=(--num-frames 24)
+            fi
             set +e
             PYTHONPATH=/tmp/mmgen-official-code/LTX-2/packages/ltx-core/src:/tmp/mmgen-official-code/LTX-2/packages/ltx-pipelines/src:$PWD:$PYTHONPATH \
               python ../ci-data/diffusion-ci/repro_scripts/gen_official_ltx23.py \
                 --out-dir ./${{ env.OUTPUT_NAME }} \
-                --case-ids "$CASE_ID"
+                --case-ids "$CASE_ID" \
+                "${extra_ltx_args[@]}"
             status=$?
             set -e
             if [ "$status" -ne 0 ]; then
-              find ./${{ env.OUTPUT_NAME }} -name 'official_ltx23_manifest.json' -print -exec cat {} \;
+              find ./${{ env.OUTPUT_NAME }} \( -name 'official_ltx_manifest.json' -o -name 'official_ltx23_manifest.json' \) -print -exec cat {} \;
               exit "$status"
             fi
           else
@@ -327,6 +288,7 @@ jobs:
             python/${{ env.OUTPUT_NAME }}/*.jpg
             python/${{ env.OUTPUT_NAME }}/*.png
             python/${{ env.OUTPUT_NAME }}/official_gt_manifest_*.json
+            python/${{ env.OUTPUT_NAME }}/official_ltx_manifest.json
             python/${{ env.OUTPUT_NAME }}/official_ltx23_manifest.json
           retention-days: 7
 

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

@@ -47,6 +47,7 @@ class LTX2SamplingParams(SamplingParams):
 class LTX23SamplingParams(LTX2SamplingParams):
     """Sampling parameters matching official LTX-2.3 one-stage defaults."""
 
+    seed: int = 42
     generator_device: str = "cuda"
     guidance_scale: float = 3.0
     num_inference_steps: int = 30

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

@@ -1,6 +1,8 @@
+import functools
 import math
 from typing import Optional, Tuple, Union
 
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -19,8 +21,7 @@ def apply_interleaved_rotary_emb(
     cos, sin = freqs
     x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1)  # [B, S, C // 2]
     x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
-    out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+    return x * cos + x_rotated * sin
-    return out
 
 
 def apply_split_rotary_emb(
@@ -47,7 +48,7 @@ def apply_split_rotary_emb(
     r = last // 2
 
     # (..., 2, r)
-    split_x = x.reshape(*x.shape[:-1], 2, r).float()
+    split_x = x.reshape(*x.shape[:-1], 2, r)
     first_x = split_x[..., :1, :]  # (..., 1, r)
     second_x = split_x[..., 1:, :]  # (..., 1, r)
 
@@ -70,6 +71,19 @@ def apply_split_rotary_emb(
     return out
 
 
+@functools.lru_cache(maxsize=5)
+def _ltx2_connector_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)
+
+
 class LTX2Attention(torch.nn.Module):
     r"""
     Attention class for all LTX-2.0 attention layers. Compared to LTX-1.0, this supports specifying the query and key
@@ -261,14 +275,18 @@ class LTX2RotaryPosEmbed1d(nn.Module):
 
         # 2. Calculate 1D RoPE frequencies
         num_rope_elems = 2  # 1 (because 1D) * 2 (for cos, sin) = 2
-        freqs_dtype = torch.float64 if self.double_precision else torch.float32
+        if self.double_precision:
+            freqs = _ltx2_connector_rope_freq_grid_np(self.theta, 1, 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=freqs_dtype,
+                    dtype=torch.float32,
                     device=device,
                 ),
             )
@@ -487,7 +505,12 @@ class LTX2ConnectorTransformer1d(nn.Module):
             attention_mask = torch.zeros_like(attention_mask)
 
         # 2. Calculate 1D RoPE positional embeddings
-        rotary_emb = self.rope(batch_size, seq_len, device=hidden_states.device)
+        rotary_emb = self.rope(
+            batch_size,
+            seq_len,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
 
         # 3. Run 1D transformer blocks
         for block in self.transformer_blocks:

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

@@ -4,8 +4,11 @@
 
 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
@@ -89,6 +92,17 @@ def _ltx2_batched_perturbation_mask(
     return mask.view(mask.numel(), *([1] * (values.ndim - 1))), False
 
 
+@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:
@@ -331,14 +345,18 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module):
         ).to(device)
 
         num_rope_elems = num_pos_dims * 2
-        freqs_dtype = torch.float64 if self.double_precision else torch.float32
+        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=freqs_dtype,
+                    dtype=torch.float32,
                     device=device,
                 ),
             )

python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py

@@ -213,18 +213,11 @@ class Gemma3Attention(nn.Module):
             is_neox_style=True,
         )
 
-        # NOTE(gmixiaojin): The shared RotaryEmbedding above computes inv_freq on
+        self.rope_scaling_factor = (
-        # GPU and uses the x1*cos - x2*sin formula, which causes slight
+            float(rope_scaling["factor"])
-        # numerical differences vs HuggingFace (see the NOTE in
+            if rope_scaling and rope_scaling.get("factor")
-        # rotary_embedding.py:_compute_inv_freq).  For HF-exact alignment we
+            else None
-        # precompute inv_freq on CPU and use rotate_half in self.rotary_emb().
-        freq_indices = (
-            torch.arange(0, self.head_dim, 2, dtype=torch.int64).float() / self.head_dim
         )
-        inv_freq = 1.0 / (self.rope_theta**freq_indices)
-        if rope_scaling and rope_scaling.get("factor"):
-            inv_freq = inv_freq / float(rope_scaling["factor"])
-        self.register_buffer("_hf_inv_freq", inv_freq, persistent=False)
 
         # Local Attention not support attention mask, we use global attention instead.
         # self.attn = LocalAttention(
@@ -246,12 +239,21 @@ class Gemma3Attention(nn.Module):
         )
 
     def rotary_emb(self, positions, q, k):
-        """Apply RoPE using HF-exact formula with precomputed inv_freq."""
+        """Apply RoPE using the same device-side inv_freq materialization as LTX."""
         positions_flat = positions.flatten().float()
         num_tokens = positions_flat.shape[0]
 
         with torch.autocast(device_type=q.device.type, enabled=False):
-            freqs = torch.outer(positions_flat, self._hf_inv_freq.float())
+            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)
@@ -290,11 +292,10 @@ class Gemma3Attention(nn.Module):
         key = k.transpose(1, 2)
         value = v.transpose(1, 2)
 
-        min_val = torch.finfo(query.dtype).min
+        attn_mask = torch.ones(
-        attn_mask = torch.zeros(
             (seq_len, seq_len),
             device=hidden_states.device,
-            dtype=query.dtype,
+            dtype=torch.bool,
         )
         causal = torch.triu(
             torch.ones(
@@ -302,30 +303,39 @@ class Gemma3Attention(nn.Module):
             ),
             diagonal=1,
         )
-        attn_mask = attn_mask.masked_fill(causal, min_val)
+        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]
             too_far = dist > self.sliding_window
-            attn_mask = attn_mask.masked_fill(too_far, min_val)
+            attn_mask = attn_mask.masked_fill(too_far, False)
 
-        key_pad = ~attention_mask.to(torch.bool)
         attn_mask = attn_mask[None, None, :, :].expand(batch_size, 1, seq_len, seq_len)
-        attn_mask = attn_mask.masked_fill(
+        attn_mask = attn_mask & attention_mask.to(torch.bool)[:, None, None, :]
-            key_pad[:, None, None, :].expand(batch_size, 1, seq_len, seq_len),
-            min_val,
-        )
 
-        attn_kwargs = {
-            "attn_mask": attn_mask,
-            "dropout_p": 0.0,
-            "is_causal": False,
-            "scale": self.scaling,
-        }
         if query.shape[1] != key.shape[1]:
-            attn_kwargs["enable_gqa"] = True
+            num_key_value_groups = query.shape[1] // key.shape[1]
+            key = key[:, :, None, :, :].expand(
+                batch_size, key.shape[1], num_key_value_groups, seq_len, self.head_dim
+            )
+            value = value[:, :, None, :, :].expand(
+                batch_size,
+                value.shape[1],
+                num_key_value_groups,
+                seq_len,
+                self.head_dim,
+            )
+            key = key.reshape(batch_size, query.shape[1], seq_len, self.head_dim)
+            value = value.reshape(batch_size, query.shape[1], seq_len, self.head_dim)
+
         attn_output = torch.nn.functional.scaled_dot_product_attention(
-            query, key, value, **attn_kwargs
+            query,
+            key,
+            value,
+            attn_mask=attn_mask,
+            dropout_p=0.0,
+            is_causal=False,
+            scale=self.scaling,
         )
         attn_output = attn_output.transpose(1, 2)
 

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

@@ -193,6 +193,7 @@ 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)
@@ -631,13 +632,14 @@ class LTX2SnapshotResidencyStrategy(LTX2TwoStageResidencyStrategy):
     def _pin_stage1_transformer_if_beneficial(self) -> None:
         """Optionally pin stage-1 DiT on GPU to remove first-stage cold H2D stall.
 
-        We only do this on high-VRAM CUDA machines with CPU offload enabled and
+        We only do this outside low-VRAM mode on high-VRAM CUDA machines with
-        without FSDP inference. It trades extra steady-state VRAM for lower
+        CPU offload enabled and without FSDP inference. It trades extra
-        request latency before the first denoise step.
+        steady-state VRAM for lower request latency before the first denoise step.
         """
         if (
             not self.server_args.dit_cpu_offload
             or self.server_args.use_fsdp_inference
+            or self._snapshot_low_vram_mode
             or not current_platform.is_cuda()
             or current_platform.get_device_total_memory() / BYTES_PER_GB < 70
         ):
@@ -934,6 +936,11 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
                 set_lora_kwargs["merge_weights"] = (
                     self._should_merge_stage2_distilled_lora(self.server_args)
                 )
+            elif phase == "stage1" and self.pipeline_name == "LTX2TwoStageHQPipeline":
+                # Official HQ also builds stage 1 with distilled LoRA fused.
+                set_lora_kwargs["merge_weights"] = (
+                    self._should_merge_stage2_distilled_lora(self.server_args)
+                )
             self.set_lora(
                 **set_lora_kwargs,
             )

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

@@ -361,11 +361,9 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
             scheduler_sigmas = self.distilled_sigmas
 
         scheduler.sigmas = scheduler_sigmas
-        num_steps = len(self.distilled_sigmas) - 1
+        num_steps = len(scheduler_sigmas) - 1
         scheduler.num_inference_steps = num_steps
-        scheduler.timesteps = (self.distilled_sigmas[:num_steps] * 1000).to(
+        scheduler.timesteps = (scheduler_sigmas[:num_steps] * 1000).to(distilled_device)
-            distilled_device
-        )
         scheduler._step_index = None
         scheduler._begin_index = None
 

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

@@ -1,3 +1,4 @@
+import math
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 
@@ -281,9 +282,7 @@ class LTX2DenoisingStage(DenoisingStage):
         )
         if generator is None:
             raise ValueError("LTX-2 res2s noise generator was not initialized.")
-        return cls._ltx2_res2s_new_noise(reference_tensor, generator).to(
+        return cls._ltx2_res2s_new_noise(reference_tensor, generator)
-            dtype=reference_tensor.dtype
-        )
 
     @staticmethod
     def _ltx2_apply_clean_latent_mask(
@@ -319,6 +318,34 @@ class LTX2DenoisingStage(DenoisingStage):
         b1 = cls._ltx2_phi_1(-h) - b2
         return a21, b1, b2
 
+    @staticmethod
+    def _ltx2_phi_scalar(j: int, neg_h: float) -> float:
+        if abs(neg_h) < 1e-10:
+            return 1.0 / math.factorial(j)
+        remainder = sum(neg_h**k / math.factorial(k) for k in range(j))
+        return (math.exp(neg_h) - remainder) / (neg_h**j)
+
+    @classmethod
+    def _ltx2_get_res2s_coefficients_scalar(
+        cls, h: float, c2: float = 0.5
+    ) -> tuple[float, float, float]:
+        a21 = c2 * cls._ltx2_phi_scalar(1, -h * c2)
+        b2 = cls._ltx2_phi_scalar(2, -h) / c2
+        b1 = cls._ltx2_phi_scalar(1, -h) - b2
+        return a21, b1, b2
+
+    @staticmethod
+    def _ltx2_res2s_step_size_scalar(
+        sigma: torch.Tensor, sigma_next: torch.Tensor
+    ) -> float:
+        return float(
+            (
+                -torch.log(
+                    sigma_next.detach().double().cpu() / sigma.detach().double().cpu()
+                )
+            ).item()
+        )
+
     @staticmethod
     def _ltx2_get_sde_coeff(
         sigma_next: torch.Tensor,
@@ -328,31 +355,29 @@ class LTX2DenoisingStage(DenoisingStage):
         sigma_max: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         if sigma_down is not None:
-            alpha_ratio = (1.0 - sigma_next) / (1.0 - sigma_down)
+            alpha_ratio = (1 - sigma_next) / (1 - sigma_down)
-            sigma_up = torch.sqrt(
+            sigma_up = (sigma_next**2 - sigma_down**2 * alpha_ratio**2).clamp(
-                torch.clamp(
+                min=0
-                    sigma_next.square() - sigma_down.square() * alpha_ratio.square(),
+            ) ** 0.5
-                    min=0.0,
-                )
-            )
         elif sigma_up is not None:
-            sigma_up = torch.minimum(sigma_up, sigma_next * 0.9999)
+            sigma_up.clamp_(max=sigma_next * 0.9999)
             sigmax = sigma_max if sigma_max is not None else torch.ones_like(sigma_next)
             sigma_signal = sigmax - sigma_next
-            sigma_residual = torch.sqrt(
+            sigma_residual = (sigma_next**2 - sigma_up**2).clamp(min=0) ** 0.5
-                torch.clamp(sigma_next.square() - sigma_up.square(), min=0.0)
-            )
             alpha_ratio = sigma_signal + sigma_residual
             sigma_down = sigma_residual / alpha_ratio
         else:
             alpha_ratio = torch.ones_like(sigma_next)
             sigma_down = sigma_next
             sigma_up = torch.zeros_like(sigma_next)
-        return (
+
-            torch.nan_to_num(alpha_ratio),
+        sigma_up = torch.nan_to_num(
-            torch.nan_to_num(sigma_down),
+            sigma_up if sigma_up is not None else torch.zeros_like(sigma_next), 0.0
-            torch.nan_to_num(sigma_up),
         )
+        nan_mask = torch.isnan(sigma_down)
+        sigma_down[nan_mask] = sigma_next[nan_mask].to(sigma_down.dtype)
+        alpha_ratio = torch.nan_to_num(alpha_ratio, 1.0)
+        return alpha_ratio, sigma_down, sigma_up
 
     @classmethod
     def _ltx2_res2s_sde_step(
@@ -391,6 +416,8 @@ class LTX2DenoisingStage(DenoisingStage):
         sigma_next: torch.Tensor,
         model_video_velocity: torch.Tensor,
         model_audio_velocity: torch.Tensor,
+        model_video_timestep: torch.Tensor | None,
+        model_audio_timestep: torch.Tensor | None,
         midpoint_model_call,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """res2s RK2 step for unguided stage-2 refinement (HQ pipeline).
@@ -407,10 +434,24 @@ class LTX2DenoisingStage(DenoisingStage):
             denoised_video = ctx.latents.float()
             denoised_audio = ctx.audio_latents.float()
         else:
-            denoised_video = ctx.latents.float() - sigma * model_video_velocity.float()
+            video_sigma_for_x0 = (
-            denoised_audio = (
+                model_video_timestep
-                ctx.audio_latents.float() - sigma * model_audio_velocity.float()
+                if ctx.use_ltx23_hq_timestep_semantics
+                and model_video_timestep is not None
+                else sigma
             )
+            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)
@@ -420,8 +461,14 @@ 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())
         sub_sigma = torch.sqrt(torch.clamp(sigma_d * sigma_next_d, min=0.0))
 
         anchor_video = ctx.latents.double()
@@ -433,12 +480,12 @@ class LTX2DenoisingStage(DenoisingStage):
         midpoint_audio_det = anchor_audio + h * a21 * eps1_audio
 
         sub_noise_video = (
-            self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=True).float()
+            self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=True)
             if ctx.use_native_hq_res2s_sde_noise
             else self._randn_like_with_batch_generators(ctx.latents, batch).float()
         )
         sub_noise_audio = (
-            self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=True).float()
+            self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=True)
             if ctx.use_native_hq_res2s_sde_noise
             else self._randn_like_with_batch_generators(
                 ctx.audio_latents, batch
@@ -460,15 +507,15 @@ class LTX2DenoisingStage(DenoisingStage):
             noise=sub_noise_audio,
             terminal=False,
         )
-        midpoint_video_latents = self._ltx2_apply_clean_latent_mask(
+        midpoint_video_model_latents = self._ltx2_apply_clean_latent_mask(
             midpoint_video_latents.to(dtype=ctx.latents.dtype), ctx
         )
-        midpoint_audio_latents = midpoint_audio_latents.to(
+        midpoint_audio_model_latents = midpoint_audio_latents.to(
             dtype=ctx.audio_latents.dtype
         )
 
         # Bongmath anchor refinement for the first stage-2 step.
-        if float(h.item()) < 0.5 and sigma_val > 0.03:
+        if h_value < 0.5 and sigma_val > 0.03:
             x_mid_v = midpoint_video_latents.double()
             x_mid_a = midpoint_audio_latents.double()
             for _ in range(100):
@@ -477,12 +524,26 @@ class LTX2DenoisingStage(DenoisingStage):
                 anchor_audio = x_mid_a - h * a21 * eps1_audio
                 eps1_audio = denoised_audio.double() - anchor_audio
 
-        mid_v, mid_a = midpoint_model_call(
+        mid_v, mid_a, mid_video_timestep, mid_audio_timestep = midpoint_model_call(
-            midpoint_video_latents, midpoint_audio_latents, sub_sigma
+            midpoint_video_model_latents, midpoint_audio_model_latents, sub_sigma
         )
 
-        midpoint_denoised_video = midpoint_video_latents.float() - sub_sigma * mid_v
+        mid_video_sigma_for_x0 = (
-        midpoint_denoised_audio = midpoint_audio_latents.float() - sub_sigma * mid_a
+            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()
 
         eps2_video = midpoint_denoised_video.double() - anchor_video
         eps2_audio = midpoint_denoised_audio.double() - anchor_audio
@@ -491,30 +552,34 @@ class LTX2DenoisingStage(DenoisingStage):
         next_audio_det = anchor_audio + h * (b1 * eps1_audio + b2 * eps2_audio)
 
         step_noise_video = (
-            self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=False).float()
+            self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=False)
             if ctx.use_native_hq_res2s_sde_noise
             else self._randn_like_with_batch_generators(ctx.latents, batch).float()
         )
         step_noise_audio = (
-            self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=False).float()
+            self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=False)
             if ctx.use_native_hq_res2s_sde_noise
             else self._randn_like_with_batch_generators(
                 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=sigma_d,
+            sigma=sde_sigma,
-            sigma_next=sigma_next_d,
+            sigma_next=sde_sigma_next,
             noise=step_noise_video,
             terminal=False,
         )
         next_audio = self._ltx2_res2s_sde_step(
             sample=anchor_audio,
             denoised_sample=next_audio_det,
-            sigma=sigma_d,
+            sigma=sde_sigma,
-            sigma_next=sigma_next_d,
+            sigma_next=sde_sigma_next,
             noise=step_noise_audio,
             terminal=False,
         )
@@ -1345,7 +1410,12 @@ class LTX2DenoisingStage(DenoisingStage):
                     video_latents: torch.Tensor,
                     audio_latents: torch.Tensor,
                     sigma_value: torch.Tensor,
-                ) -> tuple[torch.Tensor, torch.Tensor]:
+                ) -> tuple[
+                    torch.Tensor,
+                    torch.Tensor,
+                    torch.Tensor | None,
+                    torch.Tensor | None,
+                ]:
                     original_video_latents = ctx.latents
                     original_audio_latents = ctx.audio_latents
                     ctx.latents = video_latents
@@ -1419,7 +1489,12 @@ class LTX2DenoisingStage(DenoisingStage):
                             mid_a_u, mid_a_t = mid_a.chunk(2)
                             mid_v = mid_v_u + batch.guidance_scale * (mid_v_t - mid_v_u)
                             mid_a = mid_a_u + batch.guidance_scale * (mid_a_t - mid_a_u)
-                        return mid_v, mid_a
+                        return (
+                            mid_v,
+                            mid_a,
+                            model_inputs_local.timestep_video,
+                            model_inputs_local.timestep_audio,
+                        )
                     finally:
                         ctx.latents = original_video_latents
                         ctx.audio_latents = original_audio_latents
@@ -1431,6 +1506,8 @@ class LTX2DenoisingStage(DenoisingStage):
                     sigma_next=sigma_next,
                     model_video_velocity=model_video,
                     model_audio_velocity=model_audio,
+                    model_video_timestep=model_inputs.timestep_video,
+                    model_audio_timestep=model_inputs.timestep_audio,
                     midpoint_model_call=_stage2_midpoint_model_call,
                 )
             else:
@@ -1748,23 +1825,28 @@ class LTX2DenoisingStage(DenoisingStage):
                     v_mod, a_v_mod = pass_outputs.get("modality", (None, None))
 
                 sigma_value_float = float(sigma_value.item())
-                video_sigma_for_x0: float | torch.Tensor = sigma_value_float
+                audio_sigma_for_x0: float | torch.Tensor = sigma_value_float
-                if ctx.denoise_mask is not None:
+                if ctx.use_ltx23_hq_timestep_semantics:
+                    video_sigma_for_x0 = model_inputs_local.timestep_video
+                    audio_sigma_for_x0 = model_inputs_local.timestep_audio
+                elif ctx.denoise_mask is not None:
                     video_sigma_for_x0 = sigma_value.to(
                         device=video_latents.device, dtype=torch.float32
                     ) * ctx.denoise_mask.squeeze(-1)
+                else:
+                    video_sigma_for_x0 = sigma_value_float
 
                 denoised_video_local = self._ltx2_velocity_to_x0(
                     video_latents, v_pos, video_sigma_for_x0
                 )
                 denoised_audio_local = self._ltx2_velocity_to_x0(
-                    audio_latents, a_v_pos, sigma_value_float
+                    audio_latents, a_v_pos, audio_sigma_for_x0
                 )
                 denoised_video_neg = self._ltx2_velocity_to_x0(
                     video_latents, v_neg, video_sigma_for_x0
                 )
                 denoised_audio_neg = self._ltx2_velocity_to_x0(
-                    audio_latents, a_v_neg, sigma_value_float
+                    audio_latents, a_v_neg, audio_sigma_for_x0
                 )
                 denoised_video_perturbed = (
                     None
@@ -1777,7 +1859,7 @@ class LTX2DenoisingStage(DenoisingStage):
                     None
                     if a_v_ptb is None
                     else self._ltx2_velocity_to_x0(
-                        audio_latents, a_v_ptb, sigma_value_float
+                        audio_latents, a_v_ptb, audio_sigma_for_x0
                     )
                 )
                 denoised_video_modality = (
@@ -1791,7 +1873,7 @@ class LTX2DenoisingStage(DenoisingStage):
                     None
                     if a_v_mod is None
                     else self._ltx2_velocity_to_x0(
-                        audio_latents, a_v_mod, sigma_value_float
+                        audio_latents, a_v_mod, audio_sigma_for_x0
                     )
                 )
 
@@ -1867,8 +1949,14 @@ class LTX2DenoisingStage(DenoisingStage):
             else:
                 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())
                 sub_sigma = torch.sqrt(torch.clamp(sigma_d * sigma_next_d, min=0.0))
 
                 anchor_video = ctx.latents.double()
@@ -1880,16 +1968,14 @@ class LTX2DenoisingStage(DenoisingStage):
                 midpoint_audio_deterministic = anchor_audio + h * a21 * eps1_audio
 
                 substep_video_noise = (
-                    self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=True).float()
+                    self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=True)
                     if ctx.use_native_hq_res2s_sde_noise
                     else self._randn_like_with_batch_generators(
                         ctx.latents, batch
                     ).float()
                 )
                 substep_audio_noise = (
-                    self._ltx2_res2s_noise_like(
+                    self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=True)
-                        ctx.audio_latents, ctx, substep=True
-                    ).float()
                     if ctx.use_native_hq_res2s_sde_noise
                     else self._randn_like_with_batch_generators(
                         ctx.audio_latents, batch
@@ -1913,15 +1999,15 @@ class LTX2DenoisingStage(DenoisingStage):
                     terminal=False,
                 )
 
-                midpoint_video_latents = self._ltx2_apply_clean_latent_mask(
+                midpoint_video_model_latents = self._ltx2_apply_clean_latent_mask(
                     midpoint_video_latents.to(dtype=ctx.latents.dtype),
                     ctx,
                 )
-                midpoint_audio_latents = midpoint_audio_latents.to(
+                midpoint_audio_model_latents = midpoint_audio_latents.to(
                     dtype=ctx.audio_latents.dtype
                 )
 
-                if float(h.item()) < 0.5 and sigma_val > 0.03:
+                if h_value < 0.5 and sigma_val > 0.03:
                     x_mid_v = midpoint_video_latents.double()
                     x_mid_a = midpoint_audio_latents.double()
                     for _ in range(100):
@@ -1932,8 +2018,8 @@ class LTX2DenoisingStage(DenoisingStage):
 
                 midpoint_denoised_video, midpoint_denoised_audio = (
                     evaluate_stage1_guided_x0(
-                        video_latents=midpoint_video_latents,
+                        video_latents=midpoint_video_model_latents,
-                        audio_latents=midpoint_audio_latents,
+                        audio_latents=midpoint_audio_model_latents,
                         sigma_value=sub_sigma,
                         update_skip_cache=False,
                     )
@@ -1949,34 +2035,36 @@ class LTX2DenoisingStage(DenoisingStage):
                 )
 
                 step_video_noise = (
-                    self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=False).float()
+                    self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=False)
                     if ctx.use_native_hq_res2s_sde_noise
                     else self._randn_like_with_batch_generators(
                         ctx.latents, batch
                     ).float()
                 )
                 step_audio_noise = (
-                    self._ltx2_res2s_noise_like(
+                    self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=False)
-                        ctx.audio_latents, ctx, substep=False
-                    ).float()
                     if ctx.use_native_hq_res2s_sde_noise
                     else self._randn_like_with_batch_generators(
                         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_latents = self._ltx2_res2s_sde_step(
                     sample=anchor_video,
                     denoised_sample=next_video_deterministic,
-                    sigma=sigma_d,
+                    sigma=sde_sigma,
-                    sigma_next=sigma_next_d,
+                    sigma_next=sde_sigma_next,
                     noise=step_video_noise,
                     terminal=False,
                 )
                 next_audio_latents = self._ltx2_res2s_sde_step(
                     sample=anchor_audio,
                     denoised_sample=next_audio_deterministic,
-                    sigma=sigma_d,
+                    sigma=sde_sigma,
-                    sigma_next=sigma_next_d,
+                    sigma_next=sde_sigma_next,
                     noise=step_audio_noise,
                     terminal=False,
                 )

python/sglang/multimodal_gen/runtime/server_args.py

@@ -23,6 +23,7 @@ from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.configs.models.encoders import T5Config
 from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
 from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
+    LTX2PipelineConfig,
     is_ltx23_native_variant,
 )
 from sglang.multimodal_gen.configs.quantization.nunchaku import NunchakuSVDQuantArgs
@@ -487,6 +488,22 @@ class ServerArgs(DisaggArgsMixin):
                 self._parse_attention_backend_config(self.attention_backend_config)
             )
 
+        if self.backend != Backend.DIFFUSERS and isinstance(
+            self.pipeline_config, LTX2PipelineConfig
+        ):
+            text_backend = self.component_attention_backends.get("text_encoder")
+            if text_backend != "torch_sdpa":
+                if text_backend is None:
+                    logger.info(
+                        "Automatically set torch_sdpa backend for component text_encoder to preserve LTX2 official attention semantics"
+                    )
+                else:
+                    logger.warning(
+                        "Overriding %s backend with torch_sdpa for component text_encoder to preserve LTX2 official attention semantics",
+                        text_backend,
+                    )
+                self.component_attention_backends["text_encoder"] = "torch_sdpa"
+
         if self.ring_degree > 1:
             if self.attention_backend is not None and self.attention_backend not in (
                 "fa",

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

@@ -103,6 +103,12 @@
             "psnr_threshold": 25.0,
             "mean_abs_diff_threshold": 8.0
         },
+        "ltx_2_two_stage_t2v": {
+            "clip_threshold": 0.90,
+            "ssim_threshold": 0.89,
+            "psnr_threshold": 24.0,
+            "mean_abs_diff_threshold": 10.0
+        },
         "ltx_2.3_one_stage_ti2v": {
             "clip_threshold": 0.64,
             "ssim_threshold": 0.42,
@@ -244,8 +250,8 @@
         "ltx_2_3_hq_pipeline": {
             "clip_threshold": 0.78,
             "ssim_threshold": 0.48,
-            "psnr_threshold": 12.0,
+            "psnr_threshold": 13.0,
-            "mean_abs_diff_threshold": 48.0
+            "mean_abs_diff_threshold": 45.0
         },
         "ltx_2_3_two_stage_ti2v_2gpus": {
             "clip_threshold": 0.92,

python/sglang/multimodal_gen/test/test_utils.py

@@ -33,7 +33,7 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
-SGL_TEST_FILES_CI_DATA_REVISION = "4d9eff3b05b0ffe1d3529e8bb148b63af11a4b92"
+SGL_TEST_FILES_CI_DATA_REVISION = "6e7b99e16b857c98285277fe3b4ffef30559bde9"
 SGL_TEST_FILES_CONSISTENCY_GT_ROOT = (
     "https://raw.githubusercontent.com/"
     f"sgl-project/ci-data/{SGL_TEST_FILES_CI_DATA_REVISION}/"
@@ -50,6 +50,14 @@ 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.
+SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_CASES = frozenset(
+    {
+        "ltx_2.3_one_stage_ti2v",
+        "ltx_2.3_two_stage_t2v_2gpus",
+    }
+)
 CONSISTENCY_THRESHOLD_JSON_PATH = (
     Path(__file__).resolve().parent / "server" / "consistency_threshold.json"
 )
@@ -955,9 +963,27 @@ def _remote_consistency_gt_candidates(
 
 
 def _remote_file_exists(url: str) -> bool:
+    for method in ("head", "get"):
         try:
-        return requests.head(url, timeout=10, allow_redirects=True).status_code == 200
+            if method == "head":
+                resp = requests.head(url, timeout=10, allow_redirects=True)
+            else:
+                resp = requests.get(
+                    url,
+                    timeout=10,
+                    allow_redirects=True,
+                    headers={"Range": "bytes=0-0"},
+                    stream=True,
+                )
+            try:
+                if resp.status_code in (200, 206):
+                    return True
+                if resp.status_code not in (403, 405, 429) and resp.status_code < 500:
+                    return False
+            finally:
+                resp.close()
         except requests.RequestException:
+            pass
     return False
 
 
@@ -967,7 +993,12 @@ def _find_remote_consistency_gt_files(
     is_video: bool,
     output_format: str | None = None,
 ) -> list[tuple[str, str]]:
-    for base_url in SGL_TEST_FILES_CONSISTENCY_GT_BASES:
+    if case_id in SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_CASES:
+        bases = SGL_TEST_FILES_CONSISTENCY_GT_BASES
+    else:
+        # Avoid accidentally comparing non-comparable CI cases against official GT.
+        bases = (SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE,)
+    for base_url in bases:
         candidates = _remote_consistency_gt_candidates(
             base_url, case_id, num_gpus, is_video, output_format
         )