From 2bfc5d3bb158c63de2b0acd3a6995c62040ae14a Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 3 May 2026 16:37:46 +0800 Subject: [PATCH] [diffusion] optimize LTX2.3 HQ denoising split passes (#24298) --- .../pipelines_core/stages/ltx_2_denoising.py | 130 ++++++++++++++---- 1 file changed, 107 insertions(+), 23 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py index b20996e5c..54dba7f21 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py +++ b/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,19 +1676,38 @@ 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( - self._split_ltx2_model_kwargs( - batched_model_kwargs, split_sizes - ), - perturbation_configs, - strict=True, - ): - model_kwargs_chunk["perturbation_configs"] = ( - perturbation_config, + for index, (model_kwargs_chunk, pass_spec) in enumerate( + zip( + self._split_ltx2_model_kwargs( + batched_model_kwargs, split_sizes + ), + 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"] = ( + 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(