[diffusion] model: support LTX2.3 high quality pipeline (#23366)

docs/diffusion/compatibility_matrix.md

@@ -33,23 +33,29 @@ default parameters when initializing and generating videos.
 | TurboWan2.1 T2V 14B          | `IPostYellow/TurboWan2.1-T2V-14B-Diffusers`       | 480p                 |    ✅     |         ❌         |     ❌     |              ❌               |               ✅               |                   ✅                    |             ⭕             |
 | TurboWan2.1 T2V 14B 720P     | `IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers`  | 720p                 |    ✅     |         ❌         |     ❌     |              ❌               |               ✅               |                   ✅                    |             ⭕             |
 | TurboWan2.2 I2V A14B         | `IPostYellow/TurboWan2.2-I2V-A14B-Diffusers`      | 720p                 |    ✅     |         ❌         |     ❌     |              ❌               |               ✅               |                   ✅                    |             ⭕             |
-| Wan2.1 Fun 1.3B InP           | `weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers`          | 480p                 |    ✅     |         ✅         |     ✅     |              ⭕               |               ❌               |                   ❌                    |             ✅             |
+| Wan2.1 Fun 1.3B InP          | `weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers`          | 480p                 |    ✅     |         ✅         |     ✅     |              ⭕               |               ❌               |                   ❌                    |             ✅             |
-| Helios Base                   | `BestWishYsh/Helios-Base`                          | 720p                 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
+| Helios Base                  | `BestWishYsh/Helios-Base`                          | 720p                 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
-| Helios Mid                    | `BestWishYsh/Helios-Mid`                           | 720p                 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
+| Helios Mid                   | `BestWishYsh/Helios-Mid`                           | 720p                 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
-| Helios Distilled              | `BestWishYsh/Helios-Distilled`                     | 720p                 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
+| Helios Distilled             | `BestWishYsh/Helios-Distilled`                     | 720p                 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
-| LTX-2 (one and two stages)   | `Lightricks/LTX-2`                                | 768×512<br>1536×1024 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
+| LTX-2 (one/two-stage/TI2V)   | `Lightricks/LTX-2`                                | 768×512<br>1536×1024 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
-| LTX-2.3 (one and two stages) | `Lightricks/LTX-2.3`                              | 768×512<br>1536×1024 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
+| LTX-2.3 (one/two-stage/TI2V/HQ) | `Lightricks/LTX-2.3`                           | 768×512<br>1536×1024<br>1920×1088 (HQ default) |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
 
 **Note**:
 
 1. Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
 2. SageSLA is based on SpargeAttn. Install it first with `pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation`
-3. LTX-2 and LTX-2.3 two-stage generation uses `--pipeline-class-name LTX2TwoStagePipeline`. The spatial upsampler and distilled LoRA are auto-resolved from the model snapshot by default, and can still be overridden with `--spatial-upsampler-path` and `--distilled-lora-path`.
+3. LTX pipeline selection:
+   - One-stage: `--pipeline-class-name LTX2Pipeline`
+   - Two-stage: `--pipeline-class-name LTX2TwoStagePipeline`
+   - Two-stage HQ: `--pipeline-class-name LTX2TwoStageHQPipeline` (HQ defaults to 1920×1088; you can still override `--width/--height`)
+   - LTX-2 and LTX-2.3 support both T2V and TI2V (`--image-path`) on one-stage and two-stage pipelines (including HQ).
+   - The spatial upsampler and distilled LoRA are auto-resolved from the model snapshot by default, and can still be overridden with `--spatial-upsampler-path` and `--distilled-lora-path`.
    - For LTX models, the `Resolutions` column uses output video `width×height` semantics, matching `sglang generate --width ... --height ...`.
-4. LTX-2.3 two-stage also supports `--ltx2-two-stage-device-mode {legacy,snapshot,resident}`:
+4. LTX-2 / LTX-2.3 two-stage also supports `--ltx2-two-stage-device-mode {original,snapshot,resident}`:
    - `snapshot` is the default and recommended mode.
    - `resident` usually provides the best latency/throughput but uses much more VRAM.
-   - `legacy` preserves the historical switching path for fallback/debug.
+   - `original` keeps official two-stage semantics without the premerged stage-2 transformer path.
+   - Example (one prior run): `original` `154.67s`, `snapshot` `114.05s`, `resident` `75.71s`; peak VRAM trend is `original < snapshot < resident`.
 
 ### Image Generation Models
 

docs_new/docs/sglang-diffusion/compatibility_matrix.mdx

@@ -285,7 +285,7 @@ Optimization columns are abbreviated to keep the matrix readable:
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
     </tr>
 <tr>
-      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>LTX-2 (one and two stages)</td>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>LTX-2 (one/two-stage/TI2V)</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>Lightricks/LTX-2</code></td>
       <td style={{padding: "9px 8px", backgroundColor: "rgba(255,255,255,0.02)"}}>768×512<br />1536×1024</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
@@ -297,9 +297,9 @@ Optimization columns are abbreviated to keep the matrix readable:
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
     </tr>
 <tr>
-      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>LTX-2.3 (one and two stages)</td>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>LTX-2.3 (one/two-stage/TI2V/HQ)</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>Lightricks/LTX-2.3</code></td>
-      <td style={{padding: "9px 8px", backgroundColor: "rgba(255,255,255,0.02)"}}>768×512<br />1536×1024</td>
+      <td style={{padding: "9px 8px", backgroundColor: "rgba(255,255,255,0.02)"}}>768×512<br />1536×1024<br />1920×1088 (HQ default)</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
@@ -315,12 +315,18 @@ Optimization columns are abbreviated to keep the matrix readable:
 
 1. Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
 2. SageSLA is based on SpargeAttn. Install it first with `pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation`
-3. LTX-2 and LTX-2.3 two-stage generation uses `--pipeline-class-name LTX2TwoStagePipeline`. The spatial upsampler and distilled LoRA are auto-resolved from the model snapshot by default, and can still be overridden with `--spatial-upsampler-path` and `--distilled-lora-path`.
+3. LTX pipeline selection:
-   - For LTX models, the `Resolution` column uses output video `width×height` semantics, matching `sglang generate --width ... --height ...`.
+   - One-stage: `--pipeline-class-name LTX2Pipeline`
-4. LTX-2.3 two-stage also supports `--ltx2-two-stage-device-mode &#123;legacy,snapshot,resident&#125;`:
+   - Two-stage: `--pipeline-class-name LTX2TwoStagePipeline`
+   - Two-stage HQ: `--pipeline-class-name LTX2TwoStageHQPipeline` (HQ defaults to 1920×1088; you can still override `--width/--height`)
+   - LTX-2 and LTX-2.3 support both T2V and TI2V (`--image-path`) on one-stage and two-stage pipelines (including HQ).
+   - The spatial upsampler and distilled LoRA are auto-resolved from the model snapshot by default, and can still be overridden with `--spatial-upsampler-path` and `--distilled-lora-path`.
+   - For LTX models, the `Resolutions` column uses output video `width×height` semantics, matching `sglang generate --width ... --height ...`.
+4. LTX-2 / LTX-2.3 two-stage also supports `--ltx2-two-stage-device-mode {original,snapshot,resident}`:
    - `snapshot` is the default and recommended mode.
    - `resident` usually provides the best latency/throughput but uses much more VRAM.
-   - `legacy` preserves the historical switching path for fallback/debug.
+   - `original` keeps official two-stage semantics without the premerged stage-2 transformer path.
+   - Example (one prior run): `original` `154.67s`, `snapshot` `114.05s`, `resident` `75.71s`; peak VRAM trend is `original < snapshot < resident`.
 
 ### Image Generation Models
 

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

@@ -244,6 +244,9 @@ class LTX2PipelineConfig(PipelineConfig):
     def tokenize_prompt(self, prompt: list[str], tokenizer, tok_kwargs) -> dict:
         # Adapted from diffusers_pipeline.py _get_gemma_prompt_embeds
         # But we only need tokenization here, the embedding happens in TextEncodingStage
+        # Official LTX Gemma tokenizer trims surrounding whitespace before
+        # tokenization.
+        prompt = [text.strip() for text in prompt]
 
         # Gemma expects left padding for chat-style prompts
         tokenizer.padding_side = "left"

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

@@ -82,3 +82,38 @@ class LTX23SamplingParams(LTX2SamplingParams):
             "audio_stg_blocks": self.audio_stg_blocks,
         }
         return extra
+
+
+@dataclasses.dataclass
+class LTX23HQSamplingParams(LTX23SamplingParams):
+    """Sampling parameters matching official LTX-2.3 HQ two-stage defaults."""
+
+    height: int = 1088
+    width: int = 1920
+    num_inference_steps: int = 15
+    distilled_lora_strength_stage_1: float = 0.25
+    distilled_lora_strength_stage_2: float = 0.5
+
+    video_cfg_scale: float = 3.0
+    video_stg_scale: float = 0.0
+    video_rescale_scale: float = 0.45
+    video_modality_scale: float = 3.0
+    video_skip_step: int = 0
+    video_stg_blocks: list[int] = field(default_factory=list)
+
+    audio_cfg_scale: float = 7.0
+    audio_stg_scale: float = 0.0
+    audio_rescale_scale: float = 1.0
+    audio_modality_scale: float = 3.0
+    audio_skip_step: int = 0
+    audio_stg_blocks: list[int] = field(default_factory=list)
+
+    def build_request_extra(self) -> dict[str, Any]:
+        extra = super().build_request_extra()
+        extra["ltx2_distilled_lora_strength_stage_1"] = float(
+            self.distilled_lora_strength_stage_1
+        )
+        extra["ltx2_distilled_lora_strength_stage_2"] = float(
+            self.distilled_lora_strength_stage_2
+        )
+        return extra

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

@@ -571,18 +571,28 @@ class SamplingParams:
     def from_user_sampling_params_args(
         model_path: str, server_args: "ServerArgs", *args, **kwargs
     ):
+        pipeline_class_name = getattr(server_args, "pipeline_class_name", None)
         try:
-            sampling_params = SamplingParams.from_pretrained(
+            sampling_params = None
-                model_path, backend=server_args.backend, model_id=server_args.model_id
+            if pipeline_class_name:
-            )
+                from sglang.multimodal_gen.registry import get_pipeline_config_classes
-        except (AttributeError, ValueError) as e:
+
+                config_classes = get_pipeline_config_classes(pipeline_class_name)
+                if config_classes is not None:
+                    _, sampling_params_cls = config_classes
+                    sampling_params = sampling_params_cls()
+
+            if sampling_params is None:
+                sampling_params = SamplingParams.from_pretrained(
+                    model_path,
+                    backend=server_args.backend,
+                    model_id=server_args.model_id,
+                )
+        except (AttributeError, ValueError):
             # Handle safetensors files or other cases where model_index.json is not available
             # Use appropriate SamplingParams based on pipeline_class_name from registry
             if os.path.isfile(model_path) and model_path.endswith(".safetensors"):
                 # Determine which sampling params to use based on pipeline_class_name
-                pipeline_class_name = getattr(server_args, "pipeline_class_name", None)
-
-                # Try to get SamplingParams from registry
                 from sglang.multimodal_gen.registry import get_pipeline_config_classes
 
                 config_classes = (
@@ -616,7 +626,7 @@ class SamplingParams:
         user_kwargs = dict(kwargs)
         user_kwargs.pop("diffusers_kwargs", None)
 
-        user_sampling_params = SamplingParams(*args, **user_kwargs)
+        user_sampling_params = type(sampling_params)(*args, **user_kwargs)
         # TODO: refactor
         sampling_params._merge_with_user_params(
             user_sampling_params, explicit_fields=set(user_kwargs.keys())
@@ -982,7 +992,14 @@ class SamplingParams:
         for field in dataclasses.fields(user_params):
             field_name = field.name
             user_value = getattr(user_params, field_name)
-            default_class_value = getattr(SamplingParams, field_name)
+            if hasattr(SamplingParams, field_name):
+                default_class_value = getattr(SamplingParams, field_name)
+            elif field.default is not dataclasses.MISSING:
+                default_class_value = field.default
+            elif field.default_factory is not dataclasses.MISSING:
+                default_class_value = field.default_factory()
+            else:
+                default_class_value = dataclasses.MISSING
 
             is_user_modified = user_value != default_class_value or (
                 explicit_fields is not None and field_name in explicit_fields

python/sglang/multimodal_gen/registry.py

@@ -99,6 +99,7 @@ from sglang.multimodal_gen.configs.sample.hunyuan import (
 from sglang.multimodal_gen.configs.sample.hunyuan3d import Hunyuan3DSamplingParams
 from sglang.multimodal_gen.configs.sample.ltx_2 import (
     LTX2SamplingParams,
+    LTX23HQSamplingParams,
     LTX23SamplingParams,
 )
 from sglang.multimodal_gen.configs.sample.mova import (
@@ -632,6 +633,11 @@ def _register_configs():
             lambda path: "ltx-2.3" in path.lower(),
         ],
     )
+    # register dedicated sampling params for LTX2TwoStageHQPipeline
+    _PIPELINE_CONFIG_REGISTRY.setdefault(
+        "LTX2TwoStageHQPipeline",
+        (LTX2PipelineConfig, LTX23HQSamplingParams),
+    )
 
     # Hunyuan
     register_configs(

python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py

@@ -31,6 +31,32 @@ from sglang.multimodal_gen.utils import FlexibleArgumentParser
 logger = init_logger(__name__)
 
 
+def _resolve_cli_sampling_params_cls(server_args: ServerArgs) -> type[SamplingParams]:
+    pipeline_class_name = getattr(server_args, "pipeline_class_name", None)
+    if pipeline_class_name:
+        from sglang.multimodal_gen.registry import get_pipeline_config_classes
+
+        config_classes = get_pipeline_config_classes(pipeline_class_name)
+        if config_classes is not None:
+            _, sampling_params_cls = config_classes
+            return sampling_params_cls
+
+    try:
+        from sglang.multimodal_gen.registry import get_model_info
+
+        model_info = get_model_info(
+            server_args.model_path,
+            backend=server_args.backend,
+            model_id=server_args.model_id,
+        )
+        if model_info is not None:
+            return model_info.sampling_param_cls
+    except Exception as exc:
+        logger.debug("Falling back to base SamplingParams for CLI parsing: %s", exc)
+
+    return SamplingParams
+
+
 def add_multimodal_gen_generate_args(parser: argparse.ArgumentParser):
     """Add the arguments for the generate command."""
     parser.add_argument(
@@ -130,6 +156,7 @@ def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None
     args.request_id = "mocked_fake_id_for_offline_generate"
 
     server_args = ServerArgs.from_cli_args(args, unknown_args)
+    sampling_params_cls = _resolve_cli_sampling_params_cls(server_args)
 
     sampling_params_kwargs = {}
     config_file = getattr(args, "config", None)
@@ -137,7 +164,7 @@ def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None
     if config_file:
         config_args = ServerArgs.load_config_file(config_file) or {}
         sampling_param_fields = {
-            field.name for field in dataclasses.fields(SamplingParams)
+            field.name for field in dataclasses.fields(sampling_params_cls)
         }
         sampling_params_kwargs.update(
             {
@@ -147,7 +174,7 @@ def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None
             }
         )
 
-    sampling_params_kwargs.update(SamplingParams.get_cli_args(args))
+    sampling_params_kwargs.update(sampling_params_cls.get_cli_args(args))
     _apply_output_file_path_override(args, sampling_params_kwargs)
     sampling_params_kwargs["request_id"] = generate_request_id()
 

python/sglang/multimodal_gen/runtime/layers/lora/linear.py

@@ -36,6 +36,9 @@ from sglang.multimodal_gen.utils import get_mixed_precision_state
 torch._dynamo.config.recompile_limit = 64
 
 
+LORA_MERGE_CHUNK_BYTES = 32 * 1024 * 1024
+
+
 class BaseLayerWithLoRA(nn.Module):
 
     def __init__(
@@ -176,16 +179,48 @@ class BaseLayerWithLoRA(nn.Module):
         """
         # Merge all LoRA adapters in order
         for lora_A, lora_B, _, lora_strength in lora_list:
-            lora_delta = self.slice_lora_b_weights(
+            lora_A_sliced = self.slice_lora_a_weights(lora_A.to(data))
-                lora_B.to(data)
+            lora_B_sliced = self.slice_lora_b_weights(lora_B.to(data))
-            ) @ self.slice_lora_a_weights(lora_A.to(data))
+
-            # Apply lora_alpha / lora_rank scaling for consistency with forward()
+            scale = lora_strength
-            if self.lora_alpha is not None and self.lora_rank is not None:
+            if (
-                if self.lora_alpha != self.lora_rank:
+                self.lora_alpha is not None
-                    lora_delta = lora_delta * (self.lora_alpha / self.lora_rank)
+                and self.lora_rank is not None
-            if lora_delta.dim() > 2:
+                and self.lora_alpha != self.lora_rank
-                lora_delta = lora_delta.reshape(-1, lora_delta.shape[-1])
+            ):
-            data += lora_strength * lora_delta
+                scale *= self.lora_alpha / self.lora_rank
+
+            if not isinstance(lora_B_sliced, torch.Tensor):
+                lora_delta = lora_B_sliced @ lora_A_sliced
+                if isinstance(lora_delta, torch.Tensor) and lora_delta.dim() > 2:
+                    lora_delta = lora_delta.reshape(-1, lora_delta.shape[-1])
+                data.add_(lora_delta, alpha=scale)
+                continue
+
+            if lora_A_sliced.dim() > 2 or lora_B_sliced.dim() > 2:
+                lora_delta = lora_B_sliced @ lora_A_sliced
+                if lora_delta.dim() > 2:
+                    lora_delta = lora_delta.reshape(-1, lora_delta.shape[-1])
+                data_2d = data.reshape(-1, data.shape[-1]) if data.dim() > 2 else data
+                data_2d.add_(lora_delta, alpha=scale)
+                continue
+
+            data_2d = data.reshape(-1, data.shape[-1]) if data.dim() > 2 else data
+            lora_B_2d = (
+                lora_B_sliced.reshape(-1, lora_B_sliced.shape[-1])
+                if lora_B_sliced.dim() > 2
+                else lora_B_sliced
+            )
+
+            chunk_rows = max(
+                1,
+                LORA_MERGE_CHUNK_BYTES
+                // (data_2d.shape[-1] * max(1, data_2d.element_size())),
+            )
+            for start in range(0, lora_B_2d.shape[0], chunk_rows):
+                end = min(start + chunk_rows, lora_B_2d.shape[0])
+                chunk_delta = lora_B_2d[start:end] @ lora_A_sliced
+                data_2d[start:end].add_(chunk_delta, alpha=scale)
 
     @torch.no_grad()
     def merge_lora_weights(self, strength: float | None = None) -> None:

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

@@ -1219,6 +1219,14 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
             return timestep
         return timestep.amax(dim=tuple(range(1, timestep.ndim)))
 
+    def _scale_timestep_for_adaln(self, timestep: torch.Tensor) -> torch.Tensor:
+        ltx_variant = str(getattr(self.config.arch_config, "ltx_variant", "ltx_2"))
+        if ltx_variant == "ltx_2_3" and bool(
+            getattr(self, "_sglang_use_ltx23_hq_timestep_semantics", False)
+        ):
+            return timestep * float(self.timestep_scale_multiplier)
+        return timestep
+
     def _validate_tp_config(self, *, arch: LTX2ArchConfig, tp_size: int) -> None:
         """Validate TP-related dimension constraints (fail-fast)."""
         if tp_size < 1:
@@ -1663,8 +1671,10 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
         audio_hidden_states, _ = self.audio_patchify_proj(audio_hidden_states)
         # 3. Prepare timestep embeddings
         # 3.1. Prepare global modality (video and audio) timestep embedding and modulation parameters
+        timestep_for_adaln = self._scale_timestep_for_adaln(timestep)
+        audio_timestep_for_adaln = self._scale_timestep_for_adaln(audio_timestep)
         temb, embedded_timestep = self.adaln_single(
-            timestep.flatten(),
+            timestep_for_adaln.flatten(),
             hidden_dtype=hidden_states.dtype,
         )
         temb = temb.view(batch_size, -1, temb.size(-1))
@@ -1673,7 +1683,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
         )
 
         temb_audio, audio_embedded_timestep = self.audio_adaln_single(
-            audio_timestep.flatten(),
+            audio_timestep_for_adaln.flatten(),
             hidden_dtype=audio_hidden_states.dtype,
         )
         temb_audio = temb_audio.view(batch_size, -1, temb_audio.size(-1))
@@ -1688,8 +1698,9 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
                 if prompt_timestep is None
                 else prompt_timestep
             )
+            prompt_timestep_for_adaln = self._scale_timestep_for_adaln(prompt_timestep)
             temb_prompt, _ = self.prompt_adaln_single(
-                prompt_timestep.flatten(), hidden_dtype=hidden_states.dtype
+                prompt_timestep_for_adaln.flatten(), hidden_dtype=hidden_states.dtype
             )
             temb_prompt = temb_prompt.view(batch_size, -1, temb_prompt.size(-1))
         if self.audio_prompt_adaln_single is not None:
@@ -1698,8 +1709,11 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
                 if audio_prompt_timestep is None
                 else audio_prompt_timestep
             )
+            audio_prompt_timestep_for_adaln = self._scale_timestep_for_adaln(
+                audio_prompt_timestep
+            )
             temb_audio_prompt, _ = self.audio_prompt_adaln_single(
-                audio_prompt_timestep.flatten(),
+                audio_prompt_timestep_for_adaln.flatten(),
                 hidden_dtype=audio_hidden_states.dtype,
             )
             temb_audio_prompt = temb_audio_prompt.view(
@@ -1714,8 +1728,14 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
             prompt_timestep,
             audio_prompt_timestep,
         )
+        av_ca_video_timestep_for_adaln = self._scale_timestep_for_adaln(
+            av_ca_video_timestep
+        )
+        av_ca_audio_timestep_for_adaln = self._scale_timestep_for_adaln(
+            av_ca_audio_timestep
+        )
         temb_ca_scale_shift, _ = self.av_ca_video_scale_shift_adaln_single(
-            av_ca_video_timestep.flatten(), hidden_dtype=hidden_dtype
+            av_ca_video_timestep_for_adaln.flatten(), hidden_dtype=hidden_dtype
         )
         temb_ca_scale_shift = temb_ca_scale_shift.view(
             batch_size, -1, temb_ca_scale_shift.shape[-1]
@@ -1723,20 +1743,21 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
 
         av_ca_gate_factor = self._get_av_ca_gate_timestep_factor()
         temb_ca_gate, _ = self.av_ca_a2v_gate_adaln_single(
-            av_ca_video_timestep.flatten() * av_ca_gate_factor,
+            av_ca_video_timestep_for_adaln.flatten() * av_ca_gate_factor,
             hidden_dtype=hidden_dtype,
         )
         temb_ca_gate = temb_ca_gate.view(batch_size, -1, temb_ca_gate.shape[-1])
 
         temb_ca_audio_scale_shift, _ = self.av_ca_audio_scale_shift_adaln_single(
-            av_ca_audio_timestep.flatten(), hidden_dtype=audio_hidden_states.dtype
+            av_ca_audio_timestep_for_adaln.flatten(),
+            hidden_dtype=audio_hidden_states.dtype,
         )
         temb_ca_audio_scale_shift = temb_ca_audio_scale_shift.view(
             batch_size, -1, temb_ca_audio_scale_shift.shape[-1]
         )
 
         temb_ca_audio_gate, _ = self.av_ca_v2a_gate_adaln_single(
-            av_ca_audio_timestep.flatten() * av_ca_gate_factor,
+            av_ca_audio_timestep_for_adaln.flatten() * av_ca_gate_factor,
             hidden_dtype=audio_hidden_states.dtype,
         )
         temb_ca_audio_gate = temb_ca_audio_gate.view(

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

@@ -6,9 +6,11 @@ import torch
 from diffusers import FlowMatchEulerDiscreteScheduler
 
 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.ltx_2 import LTX23HQSamplingParams
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
     PipelineComponentLoader,
@@ -120,12 +122,18 @@ def build_official_ltx2_sigmas(
     stretch: bool = True,
     terminal: float = 0.1,
     default_number_of_tokens: int = MAX_SHIFT_ANCHOR,
+    number_of_tokens: int | None = None,
 ) -> list[float]:
     sigmas = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float32)
 
     mm = (max_shift - base_shift) / (MAX_SHIFT_ANCHOR - BASE_SHIFT_ANCHOR)
     b = base_shift - mm * BASE_SHIFT_ANCHOR
-    sigma_shift = float(default_number_of_tokens) * mm + b
+    tokens = (
+        int(number_of_tokens)
+        if number_of_tokens is not None
+        else int(default_number_of_tokens)
+    )
+    sigma_shift = float(tokens) * mm + b
 
     non_zero_mask = sigmas != 0
     shifted = torch.where(
@@ -136,8 +144,9 @@ def build_official_ltx2_sigmas(
 
     if stretch:
         one_minus_z = 1.0 - shifted[non_zero_mask]
-        scale_factor = one_minus_z[-1] / (1.0 - terminal)
+        if bool(torch.any(one_minus_z != 0)):
-        shifted[non_zero_mask] = 1.0 - (one_minus_z / scale_factor)
+            scale_factor = one_minus_z[-1] / (1.0 - terminal)
+            shifted[non_zero_mask] = 1.0 - (one_minus_z / scale_factor)
 
     return shifted[:-1].tolist()
 
@@ -148,7 +157,28 @@ class LTX2SigmaPreparationStage(PipelineStage):
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         batch.extra["ltx2_phase"] = "stage1"
         if is_ltx23_native_variant(server_args.pipeline_config.vae_config.arch_config):
-            batch.sigmas = build_official_ltx2_sigmas(int(batch.num_inference_steps))
+            # Resolution-aware sigma shift is only required for the HQ pipeline
+            # (which targets 1080p+ resolutions and was aligned against official
+            # LTX-2.3 HQ sigmas). Legacy one-stage and two-stage LTX-2.3 paths
+            # were baselined against the constant-anchor schedule.
+            if server_args.pipeline_class_name == "LTX2TwoStageHQPipeline":
+                latent_num_frames = (int(batch.num_frames) - 1) // int(
+                    server_args.pipeline_config.vae_temporal_compression
+                ) + 1
+                latent_height = int(batch.height) // int(
+                    server_args.pipeline_config.vae_scale_factor
+                )
+                latent_width = int(batch.width) // int(
+                    server_args.pipeline_config.vae_scale_factor
+                )
+                batch.sigmas = build_official_ltx2_sigmas(
+                    int(batch.num_inference_steps),
+                    number_of_tokens=latent_num_frames * latent_height * latent_width,
+                )
+            else:
+                batch.sigmas = build_official_ltx2_sigmas(
+                    int(batch.num_inference_steps)
+                )
         else:
             batch.sigmas = np.linspace(
                 1.0,
@@ -171,7 +201,11 @@ def _add_ltx2_front_stages(pipeline: ComposedPipelineBase):
     )
 
 
-def _add_ltx2_stage1_generation_stages(pipeline: ComposedPipelineBase):
+def _add_ltx2_stage1_generation_stages(
+    pipeline: ComposedPipelineBase,
+    *,
+    denoising_sampler_name: str = "euler",
+):
     pipeline.add_stage(LTX2SigmaPreparationStage())
     pipeline.add_standard_timestep_preparation_stage(
         prepare_extra_kwargs=[prepare_ltx2_mu]
@@ -191,6 +225,7 @@ def _add_ltx2_stage1_generation_stages(pipeline: ComposedPipelineBase):
                 scheduler=pipeline.get_module("scheduler"),
                 vae=pipeline.get_module("vae"),
                 audio_vae=pipeline.get_module("audio_vae"),
+                sampler_name=denoising_sampler_name,
                 pipeline=pipeline,
             ),
         ]
@@ -542,18 +577,25 @@ class LTX2TwoStageDeviceManager:
             module.to("cpu")
             return
 
+        pin_memory = bool(
+            self.server_args.pin_cpu_memory and torch.get_device_module().is_available()
+        )
         for name, param in module.named_parameters():
             snapshot = param_snapshots.get(name)
             if snapshot is None:
-                raise KeyError(
+                snapshot = self._clone_cpu_tensor_snapshot(
-                    f"Missing CPU parameter snapshot for {module_name}.{name}"
+                    param.data, pin_memory=pin_memory
                 )
+                param_snapshots[name] = snapshot
             param.data = snapshot
 
         for name, buffer in module.named_buffers():
             snapshot = buffer_snapshots.get(name)
             if snapshot is None:
-                raise KeyError(f"Missing CPU buffer snapshot for {module_name}.{name}")
+                snapshot = self._clone_cpu_tensor_snapshot(
+                    buffer.data, pin_memory=pin_memory
+                )
+                buffer_snapshots[name] = snapshot
             # Preserve runtime-updated buffers (e.g., lazily built caches) when
             # releasing back to CPU snapshots.
             if buffer.device.type == "cuda":
@@ -666,6 +708,10 @@ class LTX2TwoStageDeviceManager:
 class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
     pipeline_name = "LTX2TwoStagePipeline"
     STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
+    STAGE_1_DISTILLED_LORA_STRENGTH = 0.0
+    STAGE_2_DISTILLED_LORA_STRENGTH = 1.0
+    STAGE_1_DENOISING_SAMPLER_NAME = "euler"
+    STAGE_2_DENOISING_SAMPLER_NAME = "euler"
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -712,6 +758,7 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
         self._stage1_lora_path = server_args.lora_path
         self._stage1_lora_scale = float(server_args.lora_scale)
         self._active_lora_phase = None
+        self._active_lora_signature = None
         self._use_premerged_stage2_transformer = False
 
     def _initialize_premerged_stage2_transformer(self, server_args: ServerArgs) -> None:
@@ -733,7 +780,7 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
             lora_nickname="ltx2_stage2_distilled",
             lora_path=self._distilled_lora_path,
             target="transformer_2",
-            strength=1.0,
+            strength=self.STAGE_2_DISTILLED_LORA_STRENGTH,
             merge_weights=True,
         )
 
@@ -757,16 +804,48 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
             "resident",
         )
 
-    def _can_short_circuit_lora_switch(self, phase: str) -> bool:
+    def _get_stage_distilled_lora_strength(
-        return (
+        self, phase: str, batch: Req | None
-            phase in ("stage1", "stage2")
+    ) -> float:
-            and self._use_premerged_stage2_transformer
+        if phase == "stage1":
-            and self._stage1_lora_path is None
+            default_strength = self.STAGE_1_DISTILLED_LORA_STRENGTH
-        )
+            extra_key = "ltx2_distilled_lora_strength_stage_1"
+        elif phase == "stage2":
+            default_strength = self.STAGE_2_DISTILLED_LORA_STRENGTH
+            extra_key = "ltx2_distilled_lora_strength_stage_2"
+        else:
+            raise ValueError(f"Unknown LTX2 two-stage LoRA phase: {phase}")
+
+        if batch is None:
+            return float(default_strength)
+
+        request_strength = batch.extra.get(extra_key)
+        if request_strength is None:
+            return float(default_strength)
+        return float(request_strength)
+
+    def _can_short_circuit_lora_switch(
+        self, phase: str, batch: Req | None = None
+    ) -> bool:
+        distilled_lora_strength = self._get_stage_distilled_lora_strength(phase, batch)
+        if phase == "stage1":
+            return (
+                self._use_premerged_stage2_transformer
+                and self._stage1_lora_path is None
+                and distilled_lora_strength == 0.0
+            )
+        if phase == "stage2":
+            return (
+                self._use_premerged_stage2_transformer
+                and self._stage1_lora_path is None
+                and distilled_lora_strength == self.STAGE_2_DISTILLED_LORA_STRENGTH
+            )
+        return False
 
     def _build_lora_switch_spec(
-        self, phase: str
+        self, phase: str, batch: Req | None = None
     ) -> tuple[list[str], list[str], list[float], list[str]]:
+        distilled_lora_strength = self._get_stage_distilled_lora_strength(phase, batch)
         lora_nicknames: list[str] = []
         lora_paths: list[str] = []
         lora_strengths: list[float] = []
@@ -778,33 +857,42 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
                 lora_paths.append(self._stage1_lora_path)
                 lora_strengths.append(self._stage1_lora_scale)
                 lora_targets.append("transformer")
+            if distilled_lora_strength != 0.0:
+                lora_nicknames.append("ltx2_stage1_distilled")
+                lora_paths.append(self._distilled_lora_path)
+                lora_strengths.append(distilled_lora_strength)
+                lora_targets.append("transformer")
         elif phase == "stage2":
             if self._stage1_lora_path:
                 lora_nicknames.append("ltx2_stage1_base")
                 lora_paths.append(self._stage1_lora_path)
                 lora_strengths.append(self._stage1_lora_scale)
                 lora_targets.append("transformer")
-            lora_nicknames.append("ltx2_stage2_distilled")
+            if distilled_lora_strength != 0.0:
-            lora_paths.append(self._distilled_lora_path)
+                lora_nicknames.append("ltx2_stage2_distilled")
-            lora_strengths.append(1.0)
+                lora_paths.append(self._distilled_lora_path)
-            lora_targets.append("transformer")
+                lora_strengths.append(distilled_lora_strength)
+                lora_targets.append("transformer")
         else:
             raise ValueError(f"Unknown LTX2 two-stage LoRA phase: {phase}")
 
         return lora_nicknames, lora_paths, lora_strengths, lora_targets
 
-    def switch_lora_phase(self, phase: str) -> None:
+    def switch_lora_phase(self, phase: str, batch: Req | None = None) -> None:
-        if phase == self._active_lora_phase:
+        distilled_lora_strength = self._get_stage_distilled_lora_strength(phase, batch)
+        phase_signature = (phase, distilled_lora_strength)
+        if phase_signature == self._active_lora_signature:
             return
 
         if self._device_manager.switch_phase(
             phase
-        ) and self._can_short_circuit_lora_switch(phase):
+        ) and self._can_short_circuit_lora_switch(phase, batch):
             self._active_lora_phase = phase
+            self._active_lora_signature = phase_signature
             return
 
         lora_nicknames, lora_paths, lora_strengths, lora_targets = (
-            self._build_lora_switch_spec(phase)
+            self._build_lora_switch_spec(phase, batch)
         )
         if lora_nicknames:
             set_lora_kwargs = dict(
@@ -830,6 +918,7 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
             self.deactivate_lora_weights(target="transformer")
 
         self._active_lora_phase = phase
+        self._active_lora_signature = phase_signature
 
     def create_pipeline_stages(self, server_args: ServerArgs):
         _add_ltx2_front_stages(self)
@@ -837,7 +926,10 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
         self.add_stage(
             LTX2LoRASwitchStage(pipeline=self, phase="stage1"),
         )
-        _add_ltx2_stage1_generation_stages(self)
+        _add_ltx2_stage1_generation_stages(
+            self,
+            denoising_sampler_name=self.STAGE_1_DENOISING_SAMPLER_NAME,
+        )
         self.add_stages(
             [
                 LTX2UpsampleStage(
@@ -863,10 +955,21 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
                     vae=self.get_module("vae"),
                     audio_vae=self.get_module("audio_vae"),
                     pipeline=self,
+                    sampler_name=self.STAGE_2_DENOISING_SAMPLER_NAME,
                 ),
             ]
         )
         _add_ltx2_decoding_stage(self)
 
 
-EntryClass = [LTX2Pipeline, LTX2TwoStagePipeline]
+class LTX2TwoStageHQPipeline(LTX2TwoStagePipeline):
+    pipeline_name = "LTX2TwoStageHQPipeline"
+    pipeline_config_cls = LTX2PipelineConfig
+    sampling_params_cls = LTX23HQSamplingParams
+    STAGE_1_DISTILLED_LORA_STRENGTH = 0.25
+    STAGE_2_DISTILLED_LORA_STRENGTH = 0.5
+    STAGE_1_DENOISING_SAMPLER_NAME = "res2s"
+    STAGE_2_DENOISING_SAMPLER_NAME = "res2s"
+
+
+EntryClass = [LTX2Pipeline, LTX2TwoStagePipeline, LTX2TwoStageHQPipeline]

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

@@ -112,8 +112,16 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
         vae=None,
         audio_vae=None,
         pipeline=None,
+        sampler_name: str = "euler",
     ):
-        super().__init__(transformer, scheduler, vae, audio_vae, pipeline=pipeline)
+        super().__init__(
+            transformer,
+            scheduler,
+            vae,
+            audio_vae,
+            pipeline=pipeline,
+            sampler_name=sampler_name,
+        )
         self.distilled_sigmas = torch.tensor(distilled_sigmas)
 
     @staticmethod
@@ -171,6 +179,41 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
             return False
         return "LTX-2.3" not in str(getattr(server_args, "model_path", ""))
 
+    @staticmethod
+    def _build_stage2_renoise_generator(
+        batch: Req, reference_tensor: torch.Tensor
+    ) -> torch.Generator:
+        seeds = getattr(batch, "seeds", None)
+        if seeds:
+            seed = int(seeds[0])
+        else:
+            seed = int(getattr(batch, "seed", 10))
+        device = reference_tensor.device
+        dtype = reference_tensor.dtype
+        generator = torch.Generator(device=device).manual_seed(seed)
+        video_shape = batch.extra.get("ltx2_stage1_packed_video_shape")
+        audio_shape = batch.extra.get("ltx2_stage1_packed_audio_shape")
+        if video_shape is not None:
+            _ = torch.randn(
+                tuple(video_shape), device=device, dtype=dtype, generator=generator
+            )
+        if audio_shape is not None:
+            _ = torch.randn(
+                tuple(audio_shape), device=device, dtype=dtype, generator=generator
+            )
+        return generator
+
+    @staticmethod
+    def _ltx2_renoise_like(
+        reference_tensor: torch.Tensor, generator: torch.Generator
+    ) -> torch.Tensor:
+        return torch.randn(
+            reference_tensor.shape,
+            device=reference_tensor.device,
+            dtype=reference_tensor.dtype,
+            generator=generator,
+        )
+
     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"
@@ -191,6 +234,23 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
         if self._should_reset_stage2_generators(server_args):
             self._reset_stage2_generators(batch)
         noise_scale = float(self.distilled_sigmas[0].item())
+        # HQ pipeline uses a dedicated, deterministic renoise generator seeded
+        # from the request seed and advanced by stage-1 packed shapes to match
+        # official LTX-2.3 HQ output. Legacy two-stage paths were baselined
+        # against `batch.generator`'s natural advance through stage-1, so keep
+        # them on the original `_randn_like_with_batch_generators` sampling.
+        is_hq_pipeline = server_args.pipeline_class_name == "LTX2TwoStageHQPipeline"
+        if is_hq_pipeline:
+            video_reference_for_gen = (
+                batch.latents if isinstance(batch.latents, torch.Tensor) else None
+            )
+            if video_reference_for_gen is None:
+                video_reference_for_gen = batch.audio_latents
+            renoise_generator = self._build_stage2_renoise_generator(
+                batch, video_reference_for_gen
+            )
+        else:
+            renoise_generator = None
         if is_native_ti2v:
             prepared_latents, denoise_mask, _ = self._prepare_ltx2_ti2v_clean_state(
                 latents=batch.latents,
@@ -199,34 +259,63 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
                 zero_clean_latent=True,
                 clean_latent_background=batch.ltx2_ti2v_clean_latent_background,
             )
-            video_noise = self._randn_like_with_batch_generators(
+            if is_hq_pipeline:
-                prepared_latents, batch
+                video_noise = self._ltx2_renoise_like(
-            )
+                    prepared_latents, renoise_generator
+                )
+            else:
+                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 = (
+            if is_hq_pipeline:
-                video_noise * scaled_mask + prepared_latents * (1 - scaled_mask)
+                batch.latents = (
-            ).to(prepared_latents.dtype)
+                    video_noise.float() * scaled_mask
+                    + prepared_latents.float() * (1.0 - scaled_mask)
+                ).to(prepared_latents.dtype)
+            else:
+                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)
+            if is_hq_pipeline:
-            batch.latents = (
+                video_noise = self._ltx2_renoise_like(batch.latents, renoise_generator)
-                video_noise * noise_scale + batch.latents * (1 - noise_scale)
+                batch.latents = (
-            ).to(batch.latents.dtype)
+                    video_noise.float() * noise_scale
+                    + batch.latents.float() * (1.0 - noise_scale)
+                ).to(batch.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)
+                ).to(batch.latents.dtype)
 
         if isinstance(batch.audio_latents, torch.Tensor):
-            audio_noise = self._randn_like_with_batch_generators(
+            if is_hq_pipeline:
-                batch.audio_latents, batch
+                audio_noise = self._ltx2_renoise_like(
-            )
+                    batch.audio_latents, renoise_generator
-            audio_scaled_mask = (
+                )
-                torch.ones_like(batch.audio_latents[..., :1], dtype=torch.float32)
+                batch.audio_latents = (
-                * noise_scale
+                    audio_noise.float() * noise_scale
-            )
+                    + batch.audio_latents.float() * (1.0 - noise_scale)
-            batch.audio_latents = (
+                ).to(batch.audio_latents.dtype)
-                audio_noise * audio_scaled_mask
+            else:
-                + batch.audio_latents * (1 - audio_scaled_mask)
+                audio_noise = self._randn_like_with_batch_generators(
-            ).to(batch.audio_latents.dtype)
+                    batch.audio_latents, batch
+                )
+                audio_scaled_mask = (
+                    torch.ones_like(batch.audio_latents[..., :1], dtype=torch.float32)
+                    * noise_scale
+                )
+                batch.audio_latents = (
+                    audio_noise * audio_scaled_mask
+                    + 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
         ):
@@ -244,8 +333,25 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
 
         self.scheduler = copy.deepcopy(original_scheduler)
         distilled_device = self.scheduler.sigmas.device
-        self.scheduler.sigmas = self.distilled_sigmas.to(distilled_device)
         num_steps = len(self.distilled_sigmas) - 1
+        # HQ pipeline extends the sigma schedule so the final step targets a
+        # small non-zero sigma (0.0011) instead of 0.0, matching official
+        # LTX-2.3 HQ's last-step behavior. Legacy two-stage baselines used the
+        # un-extended schedule (final step goes to 0.0).
+        if (
+            server_args.pipeline_class_name == "LTX2TwoStageHQPipeline"
+            and self.distilled_sigmas[-1].item() == 0.0
+        ):
+            scheduler_sigmas = torch.cat(
+                [
+                    self.distilled_sigmas[:-1],
+                    torch.tensor([0.0011, 0.0], dtype=self.distilled_sigmas.dtype),
+                ],
+                dim=0,
+            )
+        else:
+            scheduler_sigmas = self.distilled_sigmas
+        self.scheduler.sigmas = scheduler_sigmas.to(distilled_device)
         self.scheduler.num_inference_steps = num_steps
         self.scheduler.timesteps = (self.distilled_sigmas[:num_steps] * 1000).to(
             distilled_device

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

@@ -155,14 +155,16 @@ class LTX2AVLatentPreparationStage(LatentPreparationStage):
             latent_shape = server_args.pipeline_config.prepare_latent_shape(
                 batch, batch_size, num_frames
             )
+            packed_video_shape = self._packed_video_latent_shape(
+                latent_shape, server_args.pipeline_config
+            )
             latents = randn_tensor(
-                self._packed_video_latent_shape(
+                packed_video_shape,
-                    latent_shape, server_args.pipeline_config
-                ),
                 generator=generator,
                 device=device,
                 dtype=dtype,
             )
+            batch.extra["ltx2_stage1_packed_video_shape"] = tuple(packed_video_shape)
 
             latent_ids = server_args.pipeline_config.maybe_prepare_latent_ids(latents)
             if latent_ids is not None:
@@ -196,13 +198,14 @@ class LTX2AVLatentPreparationStage(LatentPreparationStage):
             latent_shape = server_args.pipeline_config.prepare_audio_latent_shape(
                 batch, batch_size, batch.num_frames
             )
-
+            packed_audio_shape = self._packed_audio_latent_shape(latent_shape)
             audio_latents = randn_tensor(
-                self._packed_audio_latent_shape(latent_shape),
+                packed_audio_shape,
                 generator=generator,
                 device=device,
                 dtype=dtype,
             )
+            batch.extra["ltx2_stage1_packed_audio_shape"] = tuple(packed_audio_shape)
         else:
             audio_latents = audio_latents.to(device)
             audio_latents = server_args.pipeline_config.maybe_pack_audio_latents(

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

@@ -47,7 +47,6 @@ class LTX2TextConnectorStage(PipelineStage):
 
         # Handle CFG: Concatenate negative and positive inputs
         if batch.do_classifier_free_guidance:
-
             # Concatenate: [Negative, Positive]
             prompt_embeds = torch.cat([neg_prompt_embeds, prompt_embeds], dim=0)
             prompt_attention_mask = torch.cat(
@@ -57,7 +56,9 @@ class LTX2TextConnectorStage(PipelineStage):
         # Prepare additive mask for connectors (as per Diffusers implementation)
         dtype = prompt_embeds.dtype
 
-        additive_attention_mask = (1 - prompt_attention_mask.to(dtype)) * -1000000.0
+        additive_attention_mask = (prompt_attention_mask.to(torch.int64) - 1).to(
+            dtype
+        ) * torch.finfo(dtype).max
 
         # Call connectors
         # Expects: prompt_embeds, attention_mask, additive_mask=True

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

@@ -55,7 +55,7 @@ class LTX2LoRASwitchStage(PipelineStage):
             raise ValueError(
                 "LTX2LoRASwitchStage requires pipeline.switch_lora_phase()"
             )
-        switch_fn(self.phase)
+        switch_fn(self.phase, batch=batch)
         batch.extra["ltx2_phase"] = self.phase
         return batch
 

python/sglang/multimodal_gen/runtime/server_args.py

@@ -73,7 +73,7 @@ logger = init_logger(__name__)
 # GPUs on the faster no-offload default while preserving some headroom.
 WAN_LAYERWISE_OFFLOAD_AUTO_DISABLE_MEM_GB = 130
 LTX2_TWO_STAGE_DEVICE_MODES = ("original", "snapshot", "resident")
-LTX2_TWO_STAGE_PIPELINE_NAMES = ("LTX2TwoStagePipeline",)
+LTX2_TWO_STAGE_PIPELINE_NAMES = ("LTX2TwoStagePipeline", "LTX2TwoStageHQPipeline")
 # H200-class GPUs (>=130 GiB total) can usually keep both LTX2 DiTs resident.
 LTX2_RESIDENT_AUTO_ENABLE_MEM_GB = 130
 

python/sglang/multimodal_gen/runtime/utils/model_overlay.py

@@ -26,10 +26,8 @@ logger = init_logger(__name__)
 # Built-in diffusion model overlay registry.
 BUILTIN_MODEL_OVERLAY_REGISTRY: dict[str, dict[str, Any]] = {
     "Lightricks/LTX-2.3": {
-        # TODO: consider move to lmsys hf repo
         "overlay_repo_id": "MickJ/LTX-2.3-overlay",
-        "overlay_revision": "main",
+        "overlay_revision": "e0cc94f279ec16bb87c230134d40319f6ce40c5e",
-        "bundled_overlay_subdir": "ltx_2_3",
     },
 }
 

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

@@ -332,6 +332,16 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [
     #         num_frames=33,
     #     ),
     # ),
+    DiffusionTestCase(
+        "ltx_2_3_hq_pipeline",
+        DiffusionServerArgs(
+            model_path="Lightricks/LTX-2.3",
+            extras=[
+                "--pipeline-class-name LTX2TwoStageHQPipeline --ltx2-two-stage-device-mode snapshot"
+            ],
+        ),
+        T2I_sampling_params,
+    ),
 ]
 
 # Skip hunyuan3d on AMD: marching_cubes surface extraction produces invalid SDF on ROCm.
@@ -519,6 +529,7 @@ TWO_GPU_CASES = [
         "ltx_2_3_two_stage_ti2v_2gpus",
         DiffusionServerArgs(
             model_path="Lightricks/LTX-2.3",
+            ulysses_degree=2,
             extras=[
                 "--pipeline-class-name LTX2TwoStagePipeline --ltx2-two-stage-device-mode original"
             ],
@@ -537,6 +548,7 @@ TWO_GPU_CASES = [
         "ltx_2.3_two_stage_t2v_2gpus",
         DiffusionServerArgs(
             model_path="Lightricks/LTX-2.3",
+            ulysses_degree=2,
             extras=[
                 "--pipeline-class-name LTX2TwoStagePipeline",
                 "--ltx2-two-stage-device-mode original",
@@ -619,6 +631,7 @@ TWO_GPU_CASES = [
         "ltx_2.3_one_stage_ti2v",
         DiffusionServerArgs(
             model_path="Lightricks/LTX-2.3",
+            ulysses_degree=2,
         ),
         TI2V_sampling_params,
     ),

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

@@ -2629,6 +2629,47 @@
             "expected_avg_denoise_ms": 890.17,
             "expected_median_denoise_ms": 896.09,
             "estimated_full_test_time_s": 155.3
+        },
+        "ltx_2_3_hq_pipeline": {
+            "stages_ms": {
+                "InputValidationStage": 0.11,
+                "TextEncodingStage": 984.78,
+                "LTX2TextConnectorStage": 30.42,
+                "LTX2HalveResolutionStage": 0.1,
+                "LTX2LoRASwitchStage": 0.01,
+                "LTX2SigmaPreparationStage": 0.36,
+                "TimestepPreparationStage": 21.28,
+                "LTX2AVLatentPreparationStage": 0.13,
+                "LTX2ImageEncodingStage": 0.03,
+                "LTX2AVDenoisingStage": 20227.05,
+                "LTX2UpsampleStage": 157.73,
+                "LTX2RefinementStage": 1676.07,
+                "LTX2AVDecodingStage": 521.04,
+                "per_frame_generation": null
+            },
+            "denoise_step_ms": {
+                "0": 1406.0,
+                "1": 1362.75,
+                "2": 1306.18,
+                "3": 1299.65,
+                "4": 1282.16,
+                "5": 1290.32,
+                "6": 1284.64,
+                "7": 1265.01,
+                "8": 1304.06,
+                "9": 1246.21,
+                "10": 1102.19,
+                "11": 1379.4,
+                "12": 1467.28,
+                "13": 1469.49,
+                "14": 734.96,
+                "15": 547.19,
+                "16": 543.27,
+                "17": 539.0
+            },
+            "expected_e2e_ms": 24150.97,
+            "expected_avg_denoise_ms": 1157.21,
+            "expected_median_denoise_ms": 1287.48
         }
     }
 }