[diffusion] model: support LTX2.3 two stage

docs/diffusion/compatibility_matrix.md

@@ -33,16 +33,16 @@ 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                 |    ✅     |         ❌         |     ❌     |              ❌               |               ✅               |                   ✅                    |             ⭕             |
-| LTX-2                        | `Lightricks/LTX-2`                                | 768×512<br>1536×1024 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
-| LTX-2.3                      | `Lightricks/LTX-2.3`                              | 768×512<br>1536×1024 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
+| LTX-2 (one and two stages)   | `Lightricks/LTX-2`                                | 768×512<br>1536×1024 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
+| LTX-2.3 (one and two stages) | `Lightricks/LTX-2.3`                              | 768×512<br>1536×1024 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |
 
 **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 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`.
-4. `Lightricks/LTX-2.3` is supported through the bundled native overlay materialization path. One-stage generation uses the default `LTX2Pipeline`; two-stage generation uses `--pipeline-class-name LTX2TwoStagePipeline`.
-5. For LTX models, the `Resolutions` column uses output video `width×height` semantics, matching `sglang generate --width ... --height ...`. One-stage generation is validated at `768×512`; two-stage generation is validated at `1536×1024`.
+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`.
+  - For LTX models, the `Resolutions` column uses output video `width×height` semantics, matching `sglang generate --width ... --height ...`.
 
 ### Image Generation Models
 

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

@@ -365,6 +365,42 @@ def gather_latents_for_sp(self, latents, batch=None):
         latents = sequence_model_parallel_all_gather(latents, dim=2)
         return latents
 
+    def can_shard_audio_latents_for_sp(self, audio_latents) -> bool:
+        """Return whether this pipeline uses packed audio latents that can be SP-sharded."""
+        return False
+
+    def shard_audio_latents_for_sp(self, batch, audio_latents):
+        """Shard packed audio latents for SP. Pipelines without packed audio latents should return the input unchanged."""
+        return audio_latents, False
+
+    def gather_audio_latents_for_sp(self, audio_latents, batch):
+        """Gather SP-sharded audio latents back to full sequence length."""
+        return audio_latents
+
+    def prepare_video_rope_coords_for_sp(
+        self,
+        model,
+        batch,
+        latent_model_input,
+        *,
+        num_frames,
+        height,
+        width,
+    ):
+        """Prepare model-side video RoPE coordinates for the local SP shard when the pipeline requires them."""
+        return None
+
+    def prepare_audio_rope_coords_for_sp(
+        self,
+        model,
+        batch,
+        audio_latent_model_input,
+        *,
+        num_frames,
+    ):
+        """Prepare model-side audio RoPE coordinates for the local SP shard when the pipeline requires them."""
+        return None
+
     def gather_noise_pred_for_sp(self, batch, noise_pred):
         noise_pred = self.gather_latents_for_sp(noise_pred)
         raw_latent_shape = getattr(batch, "raw_latent_shape", None)

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

@@ -345,6 +345,7 @@ def shard_latents_for_sp(self, batch, latents):
         latent_frames, tokens_per_frame = (
             self._infer_video_latent_frames_and_tokens_per_frame(batch, seq_len)
         )
+        orig_latent_frames = int(latent_frames)
 
         # Pad whole frames so `latent_frames` is divisible by `sp_world_size`.
         pad_frames = (sp_world_size - (latent_frames % sp_world_size)) % sp_world_size
@@ -360,6 +361,9 @@ def shard_latents_for_sp(self, batch, latents):
 
         local_frames = int(latent_frames) // int(sp_world_size)
         start_frame = int(sp_rank) * int(local_frames)
+        valid_local_frames = max(
+            min(int(orig_latent_frames) - int(start_frame), int(local_frames)), 0
+        )
         start = int(start_frame) * int(tokens_per_frame)
         end = int(start) + int(local_frames) * int(tokens_per_frame)
         latents = latents[:, start:end, :]
@@ -368,6 +372,9 @@ def shard_latents_for_sp(self, batch, latents):
         batch.sp_video_latent_num_frames = int(local_frames)
         batch.sp_video_start_frame = int(start_frame)
         batch.sp_video_tokens_per_frame = int(tokens_per_frame)
+        batch.sp_video_valid_token_count = int(valid_local_frames) * int(
+            tokens_per_frame
+        )
 
         return latents, True
 
@@ -379,6 +386,104 @@ def gather_latents_for_sp(self, latents, batch=None):
             return sequence_model_parallel_all_gather(latents.contiguous(), dim=1)
         return super().gather_latents_for_sp(latents, batch=batch)
 
+    def shard_audio_latents_for_sp(self, batch, audio_latents):
+        sp_world_size = get_sp_world_size()
+        if sp_world_size <= 1:
+            return audio_latents, False
+        if not (isinstance(audio_latents, torch.Tensor) and audio_latents.ndim == 3):
+            return audio_latents, False
+
+        sp_rank = get_sp_parallel_rank()
+        seq_len = int(audio_latents.shape[1])
+        batch.sp_audio_orig_num_frames = int(seq_len)
+
+        pad_frames = (sp_world_size - (seq_len % sp_world_size)) % sp_world_size
+        if pad_frames:
+            pad = torch.zeros(
+                (audio_latents.shape[0], pad_frames, audio_latents.shape[2]),
+                device=audio_latents.device,
+                dtype=audio_latents.dtype,
+            )
+            audio_latents = torch.cat([audio_latents, pad], dim=1)
+            seq_len += int(pad_frames)
+
+        local_frames = seq_len // sp_world_size
+        start_frame = sp_rank * local_frames
+        end_frame = start_frame + local_frames
+        valid_local_frames = max(
+            min(
+                int(batch.sp_audio_orig_num_frames) - int(start_frame),
+                int(local_frames),
+            ),
+            0,
+        )
+        audio_latents = audio_latents[:, start_frame:end_frame, :]
+
+        batch.sp_audio_latent_num_frames = int(local_frames)
+        batch.sp_audio_start_frame = int(start_frame)
+        batch.sp_audio_valid_token_count = int(valid_local_frames)
+        return audio_latents, True
+
+    def can_shard_audio_latents_for_sp(self, audio_latents) -> bool:
+        return (
+            get_sp_world_size() > 1
+            and isinstance(audio_latents, torch.Tensor)
+            and audio_latents.ndim == 3
+        )
+
+    def gather_audio_latents_for_sp(self, audio_latents, batch):
+        if get_sp_world_size() <= 1:
+            return audio_latents
+        if not (isinstance(audio_latents, torch.Tensor) and audio_latents.ndim == 3):
+            return audio_latents
+
+        audio_latents = sequence_model_parallel_all_gather(
+            audio_latents.contiguous(), dim=1
+        )
+        orig_num_frames = int(batch.sp_audio_orig_num_frames)
+        if orig_num_frames > 0:
+            audio_latents = audio_latents[:, :orig_num_frames, :]
+        return audio_latents
+
+    def prepare_video_rope_coords_for_sp(
+        self,
+        model,
+        batch,
+        latent_model_input,
+        *,
+        num_frames,
+        height,
+        width,
+    ):
+        if not batch.did_sp_shard_latents:
+            return None
+        return model.rope.prepare_video_coords(
+            batch_size=int(latent_model_input.shape[0]),
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            device=latent_model_input.device,
+            fps=batch.fps,
+            start_frame=int(batch.sp_video_start_frame),
+        )
+
+    def prepare_audio_rope_coords_for_sp(
+        self,
+        model,
+        batch,
+        audio_latent_model_input,
+        *,
+        num_frames,
+    ):
+        if not batch.did_sp_shard_audio_latents:
+            return None
+        return model.audio_rope.prepare_audio_coords(
+            batch_size=int(audio_latent_model_input.shape[0]),
+            num_frames=num_frames,
+            device=audio_latent_model_input.device,
+            start_frame=int(batch.sp_audio_start_frame),
+        )
+
     def maybe_pack_audio_latents(self, latents, batch_size, batch):
         # If already packed (3D shape [B, T, C*F]), skip packing
         if latents.dim() == 3: