From 4d1bb64f30178867bd4d5dd2279a4ea829985fd7 Mon Sep 17 00:00:00 2001 From: Mick Qian Date: Thu, 22 Oct 2026 22:59:11 +0800 Subject: [PATCH 1/2] [Perf] diffusion/ltx2: read step sigmas without a host sync Every LTX-2 denoising step moved scheduler.sigmas[i] and [i+1] to the device and called .item() on both, and the stage-1 guider path called .item() on sigma again after the forward. Each read blocks the host until the stream drains, so the host can never enqueue the next step while the GPU runs this one -- the bubble that remains once the forward itself is graph-replayed. Stage 2 also paid a pageable H2D copy per sigma, since its distilled schedule lives on the CPU. Take fp32 host and device copies of the schedule once per loop: steps index the device copy for tensor math and the host copy for Python scalars. The guider's x0 conversion takes the device sigma instead of its float; both scale in fp32, so outputs are bit-identical. Co-Authored-By: Claude Opus 5.5 --- .../model_specific_stages/ltx_2/denoising.py | 31 ++++---- .../test/unit/test_ltx2_bcg_coords.py | 74 +++++++++++++++++-- 2 files changed, 85 insertions(+), 20 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py index 358d035ed63c..b9aac3e9a02f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py @@ -74,6 +74,9 @@ class LTX2DenoisingContext(DenoisingContext): use_ltx23_hq_timestep_semantics: bool = False res2s_step_noise_generator: torch.Generator | None = None res2s_substep_noise_generator: torch.Generator | None = None + # fp32 copies of scheduler.sigmas taken once per loop, so steps never sync on them + sigmas_host: torch.Tensor | None = None + sigmas_device: torch.Tensor | None = None @dataclass(slots=True) @@ -1603,6 +1606,11 @@ def _prepare_denoising_loop( self._move_ltx2_scheduler_tensors_to_device(ctx.scheduler, device) self._move_ltx2_scheduler_tensors_to_device(ctx.audio_scheduler, device) + sigmas = ctx.scheduler.sigmas + if not isinstance(sigmas, torch.Tensor): + raise ValueError("Expected scheduler.sigmas to be a tensor for LTX-2.") + ctx.sigmas_host = sigmas.to(device="cpu", dtype=torch.float32) + ctx.sigmas_device = sigmas.to(device=ctx.latents.device, dtype=torch.float32) return ctx @staticmethod @@ -1656,19 +1664,12 @@ def _run_denoising_step( if ctx.audio_scheduler is None: raise ValueError("LTX-2 audio scheduler was not prepared.") - # 1. Read the scheduler sigma pair and derive the Euler delta. - sigmas = getattr(ctx.scheduler, "sigmas", None) - if sigmas is None or not isinstance(sigmas, torch.Tensor): - raise ValueError("Expected scheduler.sigmas to be a tensor for LTX-2.") - sigma = sigmas[step.step_index].to( - device=ctx.latents.device, dtype=torch.float32 - ) - sigma_next = sigmas[step.step_index + 1].to( - device=ctx.latents.device, dtype=torch.float32 - ) + # 1. Read the sigma pair from the per-loop copies and derive the Euler delta. + sigma = ctx.sigmas_device[step.step_index] + sigma_next = ctx.sigmas_device[step.step_index + 1] dt = sigma_next - sigma - sigma_val = float(sigma.item()) - sigma_next_val = float(sigma_next.item()) + sigma_val = float(ctx.sigmas_host[step.step_index]) + sigma_next_val = float(ctx.sigmas_host[step.step_index + 1]) stage1_guider_params = self._get_ltx2_stage1_guider_params( batch, server_args, ctx.stage @@ -2172,9 +2173,9 @@ def evaluate_stage1_guided_x0( v_ptb, a_v_ptb = pass_outputs.get("perturbed", (None, None)) 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 + # x0 scales by sigma in fp32 either way; the tensor form avoids a sync + video_sigma_for_x0: torch.Tensor = sigma_value + audio_sigma_for_x0: torch.Tensor = sigma_value 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 diff --git a/python/sglang/multimodal_gen/test/unit/test_ltx2_bcg_coords.py b/python/sglang/multimodal_gen/test/unit/test_ltx2_bcg_coords.py index 28f1fde7b0e8..887c9f5714a8 100644 --- a/python/sglang/multimodal_gen/test/unit/test_ltx2_bcg_coords.py +++ b/python/sglang/multimodal_gen/test/unit/test_ltx2_bcg_coords.py @@ -201,7 +201,7 @@ def _model(**kwargs): kwargs["audio_hidden_states"] * 0.5 + prompt, ) - def _run_step(self, *, mode, warmup, runner, value): + def _run_step(self, *, mode, warmup, runner, value, sigmas=None): stage = object.__new__(LTX2DenoisingStage) stage.sampler_name = "euler" stage._attention_backend_active_override = None @@ -218,11 +218,19 @@ def _run_step(self, *, mode, warmup, runner, value): stage.server_args = args video = torch.arange(8, dtype=torch.float32).view(1, 2, 4) + value audio = torch.arange(4, dtype=torch.float32).view(1, 1, 4) + value + if sigmas is None: + sigmas = torch.tensor([1.0, 0.5]) + scheduler = SimpleNamespace( + sigmas=sigmas, + step=lambda output, t, sample, return_dict: (sample - 0.5 * output,), + ) ctx = SimpleNamespace( latents=video, audio_latents=audio, - scheduler=SimpleNamespace(sigmas=torch.tensor([1.0, 0.5])), - audio_scheduler=object(), + scheduler=scheduler, + sigmas_host=torch.tensor([1.0, 0.5]), + sigmas_device=sigmas, + audio_scheduler=scheduler, stage="stage1", is_warmup=warmup, is_ltx23_variant=False, @@ -235,8 +243,14 @@ def _run_step(self, *, mode, warmup, runner, value): reserved_frames_mask=None, z=None, ) - step = SimpleNamespace(step_index=0, current_model=self._model) - params = LTX23SamplingParams(video_rescale_scale=0.0, audio_rescale_scale=0.0) + step = SimpleNamespace( + step_index=0, current_model=self._model, t_device=torch.tensor(1000.0) + ) + params = ( + LTX2SamplingParams() + if mode == "official" + else LTX23SamplingParams(video_rescale_scale=0.0, audio_rescale_scale=0.0) + ) embeds = torch.full((1, 2, 4), value) batch = SimpleNamespace( quality="exact" if mode == "legacy" else "lossless", @@ -248,6 +262,9 @@ def _run_step(self, *, mode, warmup, runner, value): prompt_attention_mask=torch.ones(1, 2), negative_attention_mask=torch.ones(1, 2), ltx2_num_image_tokens=int(mode == "ti2v"), + do_classifier_free_guidance=True, + guidance_scale=4.0, + rollout_controllers={}, ) stage._prepare_ltx2_model_inputs = lambda *_: SimpleNamespace( latent_model_input=ctx.latents, @@ -291,6 +308,53 @@ def test_guidance_paths_capture_only_at_warmup_and_preserve_predictions(self): kwargs["perturbation_configs"][3]["skip_a2v_cross_attn"] ) + def test_steps_never_read_device_sigmas_on_host(self): + modes = ("official", "batched", "ti2v", "hq", "legacy", "legacy_batched") + for mode in modes: + with self.subTest(mode=mode): + expected = self._run_step( + mode=mode, warmup=False, runner=None, value=1.0 + ) + actual = self._run_step( + mode=mode, + warmup=False, + runner=None, + value=1.0, + sigmas=torch.tensor([1.0, 0.5]).as_subclass(_DeviceSigmas), + ) + for result, reference in zip(actual, expected): + torch.testing.assert_close( + result.as_subclass(torch.Tensor), reference, rtol=0, atol=0 + ) + + def test_tensor_sigma_x0_matches_python_float(self): + torch.manual_seed(0) + sample = torch.randn(2, 6, 8, dtype=torch.bfloat16) + velocity = torch.randn(2, 6, 8, dtype=torch.bfloat16) + for sigma in ( + torch.tensor(0.7182, dtype=torch.float32), + torch.tensor(0.3141592653589793, dtype=torch.float64), + ): + with self.subTest(dtype=sigma.dtype): + self.assertTrue( + torch.equal( + LTX2DenoisingStage._ltx2_velocity_to_x0( + sample, velocity, sigma + ), + LTX2DenoisingStage._ltx2_velocity_to_x0( + sample, velocity, float(sigma) + ), + ) + ) + + +class _DeviceSigmas(torch.Tensor): + # a CUDA schedule in a CPU test: any host read of a value would sync + def item(self): + raise AssertionError("denoising step read a device sigma on the host") + + __float__ = __bool__ = item + class _AttentionBlock(torch.nn.Module): def __init__(self, attention): From a1b333d1fdfe33730123eb62379b473991817875 Mon Sep 17 00:00:00 2001 From: Mick Qian Date: Thu, 22 Oct 2026 23:03:23 +0800 Subject: [PATCH 2/2] [Perf] diffusion/ltx2: build the additive prompt mask once per forward LTX-2 passes the 0/1 prompt key mask down to every block, and each block's two prompt cross-attentions (LocalAttention) turned it into the additive SDPA mask again: a cast, a subtract and a multiply per call, 288 tiny launches per forward over 48 blocks. Convert it once in the model forward with the helper LocalAttention now shares (prepare_sdpa_mask, previously private to the USP paths), and let the blocks pass the float mask through. Same ops on the same inputs, so outputs are bit-identical. Co-Authored-By: Claude Opus 5.5 --- .../runtime/layers/attention/layer.py | 23 +--- .../runtime/models/dits/ltx_2.py | 15 +++ .../test/unit/test_ltx2_prompt_mask.py | 113 ++++++++++++++++++ 3 files changed, 133 insertions(+), 18 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_ltx2_prompt_mask.py diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 2541e2ef06ab..4fc482ab980f 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -172,7 +172,7 @@ def _count_active_replicated_modes( ) -def _prepare_sdpa_mask( +def prepare_sdpa_mask( mask: torch.Tensor, *, dtype: torch.dtype, device: torch.device ) -> torch.Tensor: mask = mask.to(device=device) @@ -757,20 +757,7 @@ def forward( q_ = q.transpose(1, 2) k_ = k.transpose(1, 2) v_ = v.transpose(1, 2) - - if torch.is_floating_point(attn_mask): - mask = attn_mask.to(dtype=q_.dtype, device=q_.device) - if mask.dim() == 2: - mask = mask[:, None, None, :] - elif mask.dim() == 3: - mask = mask[:, None, :, :] - else: - mask = attn_mask.to(dtype=q_.dtype, device=q_.device) - if mask.dim() == 2: - mask = mask[:, None, None, :] - elif mask.dim() == 3: - mask = mask[:, None, :, :] - mask = (mask - 1.0) * torch.finfo(q_.dtype).max + mask = prepare_sdpa_mask(attn_mask, dtype=q_.dtype, device=q_.device) if q_.shape[1] != k_.shape[1]: repeat_factor = q_.shape[1] // k_.shape[1] @@ -1269,7 +1256,7 @@ def forward( q_ = q.transpose(1, 2) k_ = k.transpose(1, 2) v_ = v.transpose(1, 2) - mask = _prepare_sdpa_mask(attn_mask, dtype=q_.dtype, device=q_.device) + mask = prepare_sdpa_mask(attn_mask, dtype=q_.dtype, device=q_.device) sdpa_context = ( sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS, set_priority=True) if self.allow_cudnn_sdp and q_.device.type == "cuda" @@ -1448,7 +1435,7 @@ def masked_attention(q, k, v): q_ = q.transpose(1, 2) k_ = k.transpose(1, 2) v_ = v.transpose(1, 2) - mask = _prepare_sdpa_mask( + mask = prepare_sdpa_mask( gathered_mask, dtype=q_.dtype, device=q_.device ) sdpa_context = ( @@ -2114,7 +2101,7 @@ def _masked_local_attention( q_ = q.transpose(1, 2) k_ = k.transpose(1, 2) v_ = v.transpose(1, 2) - mask = _prepare_sdpa_mask(attn_mask, dtype=q_.dtype, device=q_.device) + mask = prepare_sdpa_mask(attn_mask, dtype=q_.dtype, device=q_.device) sdpa_context = ( sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS, set_priority=True) if self.allow_cudnn_sdp and q_.device.type == "cuda" diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index 0374d65f93a3..f5d5b4c21678 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -41,6 +41,7 @@ tensor_model_parallel_all_reduce, ) from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention +from sglang.multimodal_gen.runtime.layers.attention.layer import prepare_sdpa_mask from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, RMSNormNoWeight from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, @@ -2183,6 +2184,20 @@ def forward( ): encoder_attention_mask = None + # prompt cross-attention reads the additive mask; build it once, not per block + if encoder_attention_mask is not None: + encoder_attention_mask = prepare_sdpa_mask( + encoder_attention_mask, + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + if audio_encoder_attention_mask is not None: + audio_encoder_attention_mask = prepare_sdpa_mask( + audio_encoder_attention_mask, + dtype=audio_hidden_states.dtype, + device=audio_hidden_states.device, + ) + # 5. Run blocks skip_video_self_attn_blocks = set(skip_video_self_attn_blocks or ()) skip_audio_self_attn_blocks = set(skip_audio_self_attn_blocks or ()) diff --git a/python/sglang/multimodal_gen/test/unit/test_ltx2_prompt_mask.py b/python/sglang/multimodal_gen/test/unit/test_ltx2_prompt_mask.py new file mode 100644 index 000000000000..2d929f171bc0 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_ltx2_prompt_mask.py @@ -0,0 +1,113 @@ +import unittest +from contextlib import ExitStack +from unittest.mock import patch + +import torch + +import sglang.multimodal_gen.runtime.models.dits.ltx_2 as ltx_2 +from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2ArchConfig, LTX2Config +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + maybe_init_distributed_environment_and_model_parallel, + model_parallel_is_initialized, +) +from sglang.multimodal_gen.runtime.layers.attention import layer +from sglang.multimodal_gen.runtime.layers.attention.backends.sdpa import SDPABackend +from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import ( + ensure_distributed_env_defaults, +) +from sglang.test.test_utils import CustomTestCase + + +def _patch_attention_layers(stack: ExitStack) -> None: + for name, value in ( + ("get_attn_backend", SDPABackend), + ("get_compute_dtype", torch.bfloat16), + ("_resolve_sp_attention_mode", ("ulysses", False)), + ("_maybe_install_backend_autotune", None), + ): + stack.enter_context(patch.object(layer, name, return_value=value)) + stack.enter_context(set_forward_context(0, None)) + + +class TestLTX2PromptMask(CustomTestCase): + @torch.no_grad() + def test_local_attention_takes_prebuilt_mask_unchanged(self): + with ExitStack() as stack: + _patch_attention_layers(stack) + attention = layer.LocalAttention( + num_heads=4, + head_size=16, + is_cross_attention=True, + allow_cudnn_sdp=True, + capture_dense_attention=True, + ) + q = torch.randn(2, 8, 4, 16, dtype=torch.bfloat16) + k = torch.randn(2, 12, 4, 16, dtype=torch.bfloat16) + v = torch.randn(2, 12, 4, 16, dtype=torch.bfloat16) + mask = torch.ones(2, 12, dtype=torch.int64) + mask[1, 7:] = 0 + prebuilt = layer.prepare_sdpa_mask(mask, dtype=q.dtype, device=q.device) + self.assertTrue( + torch.equal( + attention(q, k, v, attn_mask=prebuilt), + attention(q, k, v, attn_mask=mask), + ) + ) + + @torch.no_grad() + def test_model_mask_built_once_matches_per_block_conversion(self): + if not model_parallel_is_initialized(): + ensure_distributed_env_defaults() + maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1) + with ExitStack() as stack: + _patch_attention_layers(stack) + torch.manual_seed(0) + arch = LTX2ArchConfig( + num_layers=2, + num_attention_heads=2, + attention_head_dim=32, + in_channels=16, + out_channels=16, + cross_attention_dim=64, + caption_channels=64, + audio_num_attention_heads=2, + audio_attention_head_dim=16, + audio_in_channels=16, + audio_out_channels=16, + audio_cross_attention_dim=32, + ) + model = ltx_2.LTX2VideoTransformer3DModel( + LTX2Config(arch_config=arch), hf_config={} + ).to(torch.bfloat16) + for parameter in model.parameters(): + parameter.normal_(std=0.05) + mask = torch.ones(1, 6, dtype=torch.int64) + mask[:, 4:] = 0 + kwargs = dict( + hidden_states=torch.randn(1, 32, 16, dtype=torch.bfloat16), + audio_hidden_states=torch.randn(1, 5, 16, dtype=torch.bfloat16), + encoder_hidden_states=torch.randn(1, 6, 64, dtype=torch.bfloat16), + audio_encoder_hidden_states=torch.randn(1, 6, 64, dtype=torch.bfloat16), + timestep=torch.tensor([500.0]), + encoder_attention_mask=mask, + audio_encoder_attention_mask=mask, + num_frames=2, + height=4, + width=4, + audio_num_frames=5, + return_latents=False, + ) + built_once = model(**kwargs) + # the blocks' LocalAttention converts the raw mask on every call instead + with patch.object( + ltx_2, "prepare_sdpa_mask", side_effect=lambda mask, **_: mask + ) as convert: + per_block = model(**kwargs) + self.assertEqual(convert.call_count, 2) + for actual, expected in zip(built_once, per_block, strict=True): + self.assertTrue(torch.equal(actual, expected)) + + +if __name__ == "__main__": + unittest.main()