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