[diffusion] optimize LTX2.3 HQ denoising split passes (#24298)

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

@@ -364,13 +364,17 @@ class LTX2DenoisingStage(DenoisingStage):
         sigma_next: torch.Tensor,
         noise: torch.Tensor,
         eta: float = 0.5,
+        terminal: bool = False,
     ) -> torch.Tensor:
+        # The caller decides terminal steps from Python scalars before entering
+        # this helper. Keep that branch on host to avoid a CUDA bool sync in
+        # every res2s SDE update.
+        if terminal:
+            return denoised_sample.to(dtype=sample.dtype)
         alpha_ratio, sigma_down, sigma_up = cls._ltx2_get_sde_coeff(
             sigma_next,
             sigma_up=sigma_next * eta,
         )
-        if bool((sigma_up == 0).any()) or bool((sigma_next == 0).any()):
-            return denoised_sample.to(dtype=sample.dtype)
         eps_next = (sample - denoised_sample) / (sigma - sigma_next)
         denoised_next = sample - sigma * eps_next
         x_noised = (
@@ -446,6 +450,7 @@ class LTX2DenoisingStage(DenoisingStage):
             sigma=sigma_d,
             sigma_next=sub_sigma,
             noise=sub_noise_video,
+            terminal=False,
         )
         midpoint_audio_latents = self._ltx2_res2s_sde_step(
             sample=anchor_audio,
@@ -453,6 +458,7 @@ class LTX2DenoisingStage(DenoisingStage):
             sigma=sigma_d,
             sigma_next=sub_sigma,
             noise=sub_noise_audio,
+            terminal=False,
         )
         midpoint_video_latents = self._ltx2_apply_clean_latent_mask(
             midpoint_video_latents.to(dtype=ctx.latents.dtype), ctx
@@ -502,6 +508,7 @@ class LTX2DenoisingStage(DenoisingStage):
             sigma=sigma_d,
             sigma_next=sigma_next_d,
             noise=step_noise_video,
+            terminal=False,
         )
         next_audio = self._ltx2_res2s_sde_step(
             sample=anchor_audio,
@@ -509,6 +516,7 @@ class LTX2DenoisingStage(DenoisingStage):
             sigma=sigma_d,
             sigma_next=sigma_next_d,
             noise=step_noise_audio,
+            terminal=False,
         )
 
         next_video = self._ltx2_apply_clean_latent_mask(
@@ -977,6 +985,48 @@ class LTX2DenoisingStage(DenoisingStage):
             kwargs["perturbation_configs"] = perturbation_configs
         return kwargs
 
+    @staticmethod
+    def _ltx2_guidance_perturbation_config(
+        pass_spec: LTX2GuidancePassSpec,
+    ) -> dict[str, object]:
+        return {
+            "skip_video_self_attn_blocks": pass_spec.skip_video_self_attn_blocks,
+            "skip_audio_self_attn_blocks": pass_spec.skip_audio_self_attn_blocks,
+            "skip_a2v_cross_attn": pass_spec.disable_a2v_cross_attn,
+            "skip_v2a_cross_attn": pass_spec.disable_v2a_cross_attn,
+        }
+
+    @classmethod
+    def _build_ltx2_guidance_perturbation_configs(
+        cls,
+        pass_specs: list[LTX2GuidancePassSpec],
+        batch_size: int,
+    ) -> tuple[dict[str, object], ...]:
+        return tuple(
+            cls._ltx2_guidance_perturbation_config(pass_spec)
+            for pass_spec in pass_specs
+            for _ in range(batch_size)
+        )
+
+    @staticmethod
+    def _apply_ltx2_guidance_pass_kwargs(
+        model_kwargs: dict[str, object],
+        pass_spec: LTX2GuidancePassSpec,
+    ) -> None:
+        """Copy disable-attention options from pass_spec into model_kwargs."""
+        if pass_spec.skip_video_self_attn_blocks:
+            model_kwargs["skip_video_self_attn_blocks"] = (
+                pass_spec.skip_video_self_attn_blocks
+            )
+        if pass_spec.skip_audio_self_attn_blocks:
+            model_kwargs["skip_audio_self_attn_blocks"] = (
+                pass_spec.skip_audio_self_attn_blocks
+            )
+        if pass_spec.disable_a2v_cross_attn:
+            model_kwargs["disable_a2v_cross_attn"] = True
+        if pass_spec.disable_v2a_cross_attn:
+            model_kwargs["disable_v2a_cross_attn"] = True
+
     @classmethod
     def _repeat_ltx2_model_kwargs_batch(
         cls,
@@ -1210,6 +1260,8 @@ class LTX2DenoisingStage(DenoisingStage):
             device=ctx.latents.device, dtype=torch.float32
         )
         dt = sigma_next - sigma
+        sigma_val = float(sigma.item())
+        sigma_next_val = float(sigma_next.item())
 
         stage1_guider_params = self._get_ltx2_stage1_guider_params(
             batch, server_args, ctx.stage
@@ -1435,6 +1487,21 @@ class LTX2DenoisingStage(DenoisingStage):
                 and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0
             )
         )
+        # "Perturbation" means disabling selected attention paths
+        # for that item (self-attention blocks or audio/video cross-attention)
+        # to compute STG/modality guidance.
+        #
+
+        # Decide whether to use different pass kwargs for split model calls
+        # 1. HQ splits the expanded batch into one-item model calls. Since each
+        # call has only one perturbation setting, pass the disable options
+        # directly as model arguments.
+        # 2. TI2V/non-HQ may keep several expanded
+        # items with different settings in one model call, so it needs
+        # perturbation_configs: one config dict per expanded item.
+        use_split_pass_kwargs = (
+            server_args.pipeline_class_name == "LTX2TwoStageHQPipeline"
+        )
         skip_v2a_cross_attn_for_video_gt = bool(
             batch.extra.get("ltx2_skip_v2a_cross_attn_for_video_gt", False)
         )
@@ -1580,16 +1647,6 @@ class LTX2DenoisingStage(DenoisingStage):
 
                     num_passes = len(pass_specs)
                     expanded_batch_size = batch_size_local * num_passes
-                    perturbation_configs = tuple(
-                        {
-                            "skip_video_self_attn_blocks": pass_spec.skip_video_self_attn_blocks,
-                            "skip_audio_self_attn_blocks": pass_spec.skip_audio_self_attn_blocks,
-                            "skip_a2v_cross_attn": pass_spec.disable_a2v_cross_attn,
-                            "skip_v2a_cross_attn": pass_spec.disable_v2a_cross_attn,
-                        }
-                        for pass_spec in pass_specs
-                        for _ in range(batch_size_local)
-                    )
                     batched_model_kwargs = self._repeat_ltx2_model_kwargs_batch(
                         base_model_kwargs_local, expanded_batch_size
                     )
@@ -1619,18 +1676,37 @@ class LTX2DenoisingStage(DenoisingStage):
                     )
                     if use_split_stage1_guided_passes:
                         split_sizes = [1] * expanded_batch_size
+                        split_pass_specs = tuple(
+                            pass_spec
+                            for pass_spec in pass_specs
+                            for _ in range(batch_size_local)
+                        )
+                        split_perturbation_configs = (
+                            ()
+                            if use_split_pass_kwargs
+                            else self._build_ltx2_guidance_perturbation_configs(
+                                pass_specs, batch_size_local
+                            )
+                        )
                         batched_video_chunks = []
                         batched_audio_chunks = []
                         with self._ltx2_model_forward_context(ctx, step):
-                            for model_kwargs_chunk, perturbation_config in zip(
+                            for index, (model_kwargs_chunk, pass_spec) in enumerate(
+                                zip(
                                     self._split_ltx2_model_kwargs(
                                         batched_model_kwargs, split_sizes
                                     ),
-                                perturbation_configs,
+                                    split_pass_specs,
                                     strict=True,
+                                )
                             ):
+                                if use_split_pass_kwargs:
+                                    self._apply_ltx2_guidance_pass_kwargs(
+                                        model_kwargs_chunk, pass_spec
+                                    )
+                                else:
                                     model_kwargs_chunk["perturbation_configs"] = (
-                                    perturbation_config,
+                                        split_perturbation_configs[index],
                                     )
                                 video_chunk, audio_chunk = step.current_model(
                                     **model_kwargs_chunk
@@ -1641,6 +1717,11 @@ class LTX2DenoisingStage(DenoisingStage):
                         batched_video = torch.cat(batched_video_chunks, dim=0)
                         batched_audio = torch.cat(batched_audio_chunks, dim=0)
                     else:
+                        perturbation_configs = (
+                            self._build_ltx2_guidance_perturbation_configs(
+                                pass_specs, batch_size_local
+                            )
+                        )
                         with self._ltx2_model_forward_context(ctx, step):
                             batched_video, batched_audio = step.current_model(
                                 **batched_model_kwargs,
@@ -1772,7 +1853,6 @@ class LTX2DenoisingStage(DenoisingStage):
                 ctx.latents = original_video_latents
                 ctx.audio_latents = original_audio_latents
 
-        sigma_val = float(sigma.item())
         denoised_video, denoised_audio = evaluate_stage1_guided_x0(
             video_latents=ctx.latents,
             audio_latents=ctx.audio_latents,
@@ -1781,7 +1861,7 @@ class LTX2DenoisingStage(DenoisingStage):
         )
 
         if self.sampler_name == "res2s":
-            if sigma_val == 0.0 or float(sigma_next.item()) == 0.0:
+            if sigma_val == 0.0 or sigma_next_val == 0.0:
                 next_video_latents = denoised_video.to(dtype=ctx.latents.dtype)
                 next_audio_latents = denoised_audio.to(dtype=ctx.audio_latents.dtype)
             else:
@@ -1822,6 +1902,7 @@ class LTX2DenoisingStage(DenoisingStage):
                     sigma=sigma_d,
                     sigma_next=sub_sigma,
                     noise=substep_video_noise,
+                    terminal=False,
                 )
                 midpoint_audio_latents = self._ltx2_res2s_sde_step(
                     sample=anchor_audio,
@@ -1829,6 +1910,7 @@ class LTX2DenoisingStage(DenoisingStage):
                     sigma=sigma_d,
                     sigma_next=sub_sigma,
                     noise=substep_audio_noise,
+                    terminal=False,
                 )
 
                 midpoint_video_latents = self._ltx2_apply_clean_latent_mask(
@@ -1888,6 +1970,7 @@ class LTX2DenoisingStage(DenoisingStage):
                     sigma=sigma_d,
                     sigma_next=sigma_next_d,
                     noise=step_video_noise,
+                    terminal=False,
                 )
                 next_audio_latents = self._ltx2_res2s_sde_step(
                     sample=anchor_audio,
@@ -1895,6 +1978,7 @@ class LTX2DenoisingStage(DenoisingStage):
                     sigma=sigma_d,
                     sigma_next=sigma_next_d,
                     noise=step_audio_noise,
+                    terminal=False,
                 )
 
                 next_video_latents = self._ltx2_apply_clean_latent_mask(