[Perf] diffusion: Cut LTX-2 denoising-step host overhead

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"

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 ())

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

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):

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()