From 546965fc72a73b147323fc839ad0aa955ef64aff Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Tue, 11 Aug 2026 18:23:08 +0800 Subject: [PATCH] [diffusion] LTX-2: mount the bit-exact fused modulate at the 8 bare adaLN sites (ltx23-one-stage denoise -2.8% H100 / -2.6% H200) (#34315) Co-authored-by: Claude Fable 5 --- .../runtime/models/dits/ltx_2.py | 83 +++++++++++++++---- .../test/unit/test_ltx2_modulate_mount.py | 48 +++++++++++ 2 files changed, 117 insertions(+), 14 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_ltx2_modulate_mount.py 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 dda58fefc..0c31ac6af 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -10,6 +10,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from sglang.kernels.ops.diffusion.bitexact_gate import BitExactFusionGate from sglang.kernels.ops.diffusion.fused_linear_gelu import ( can_fuse_linear_gelu, fused_gelu_active, @@ -26,6 +27,10 @@ from sglang.kernels.ops.diffusion.ltx2_rmsnorm_modulate import ( ltx2_rms_norm_modulate_active, mark_ltx2_rms_norm_modulate_site, ) +from sglang.kernels.ops.diffusion.modulate_scale_shift import ( + can_use_modulate_scale_shift_cuda, + modulate_scale_shift_cuda, +) from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2ArchConfig, LTX2Config from sglang.multimodal_gen.configs.models.fsdp import ( @@ -140,6 +145,54 @@ def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int: ) +_LTX2_MODULATE = BitExactFusionGate("LTX-2 fused modulate") + + +def _ltx2_modulate( + x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor +) -> torch.Tensor: + """``x * (1 + scale) + shift`` via the bit-exact fused CUDA modulate. + + The adaLN scale/shift rows here are ``(B, 1, D)`` views (from + ``unbind``/``squeeze`` of the combined tables), so they are made dense and + squeezed to the ``(B, D)`` layout the kernel expects — a bit-exact copy of + tiny tensors. Per-token modulation rows (``dim 1 > 1``) and any other + unsupported layout keep the eager chain; the first fused call + self-verifies ``torch.equal`` against the eager chain and falls back + permanently on any mismatch. + """ + verified = _LTX2_MODULATE.verified + if ( + not _LTX2_MODULATE.disabled + and x.dim() == 3 + and x.is_contiguous() + and scale.dim() == 3 + and scale.shape == (x.shape[0], 1, x.shape[-1]) + and shift.shape == scale.shape + and (verified or _LTX2_MODULATE.can_attempt_once()) + ): + scale_rows = scale.squeeze(1).contiguous() + shift_rows = shift.squeeze(1).contiguous() + if can_use_modulate_scale_shift_cuda(x, scale_rows, shift_rows): + try: + out = modulate_scale_shift_cuda(x, scale_rows, shift_rows) + except Exception as exc: + _LTX2_MODULATE.on_exception(exc, logger=logger) + else: + if verified: + return out + return _LTX2_MODULATE.accept_or_fallback( + out, + x * (1 + scale) + shift, + logger=logger, + mismatch_msg=( + "LTX-2 fused modulate is not bit-exact on this " + "platform; falling back to eager" + ), + ) + return x * (1 + scale) + shift + + def _ltx2_rms_norm_modulate( block: nn.Module, rms_norm: nn.Module, @@ -1301,8 +1354,8 @@ class LTX2TransformerBlock(nn.Module): norm_hidden_states = _ltx2_rms_norm_modulate( self, self.rms_norm, hidden_states, vscale_q, vshift_q, self.norm_eps ) - mod_encoder_hidden_states = ( - encoder_hidden_states * (1 + v_prompt_scale) + v_prompt_shift + mod_encoder_hidden_states = _ltx2_modulate( + encoder_hidden_states, v_prompt_scale, v_prompt_shift ) attn_hidden_states = self.attn2( norm_hidden_states, @@ -1333,8 +1386,8 @@ class LTX2TransformerBlock(nn.Module): ashift_q, self.norm_eps, ) - mod_audio_encoder_hidden_states = ( - audio_encoder_hidden_states * (1 + a_prompt_scale) + a_prompt_shift + mod_audio_encoder_hidden_states = _ltx2_modulate( + audio_encoder_hidden_states, a_prompt_scale, a_prompt_shift ) attn_audio_hidden_states = self.audio_attn2( norm_audio_hidden_states, @@ -1423,11 +1476,11 @@ class LTX2TransformerBlock(nn.Module): v2a_gate = audio_ca_gate[0].squeeze(2) # A2V - mod_norm_hidden_states = ( - norm_hidden_states * (1 + video_a2v_ca_scale) + video_a2v_ca_shift + mod_norm_hidden_states = _ltx2_modulate( + norm_hidden_states, video_a2v_ca_scale, video_a2v_ca_shift ) - mod_norm_audio_hidden_states = ( - norm_audio_hidden_states * (1 + audio_a2v_ca_scale) + audio_a2v_ca_shift + mod_norm_audio_hidden_states = _ltx2_modulate( + norm_audio_hidden_states, audio_a2v_ca_scale, audio_a2v_ca_shift ) if not skip_a2v_cross_attn: @@ -1448,11 +1501,11 @@ class LTX2TransformerBlock(nn.Module): ) # V2A - mod_norm_hidden_states = ( - norm_hidden_states * (1 + video_v2a_ca_scale) + video_v2a_ca_shift + mod_norm_hidden_states = _ltx2_modulate( + norm_hidden_states, video_v2a_ca_scale, video_v2a_ca_shift ) - mod_norm_audio_hidden_states = ( - norm_audio_hidden_states * (1 + audio_v2a_ca_scale) + audio_v2a_ca_shift + mod_norm_audio_hidden_states = _ltx2_modulate( + norm_audio_hidden_states, audio_v2a_ca_scale, audio_v2a_ca_shift ) if not skip_v2a_cross_attn: @@ -2220,7 +2273,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] with torch.autocast(device_type=hidden_states.device.type, enabled=False): hidden_states = self.norm_out(hidden_states) - hidden_states = hidden_states * (1 + scale) + shift + hidden_states = _ltx2_modulate(hidden_states, scale, shift) hidden_states, _ = self.proj_out(hidden_states) # Audio @@ -2233,7 +2286,9 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): ) with torch.autocast(device_type=audio_hidden_states.device.type, enabled=False): audio_hidden_states = self.audio_norm_out(audio_hidden_states) - audio_hidden_states = audio_hidden_states * (1 + audio_scale) + audio_shift + audio_hidden_states = _ltx2_modulate( + audio_hidden_states, audio_scale, audio_shift + ) audio_hidden_states, _ = self.audio_proj_out(audio_hidden_states) # Unpatchify if requested (default True for pipeline compatibility) return_latents = kwargs.get("return_latents", True) diff --git a/python/sglang/multimodal_gen/test/unit/test_ltx2_modulate_mount.py b/python/sglang/multimodal_gen/test/unit/test_ltx2_modulate_mount.py new file mode 100644 index 000000000..33b71673f --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_ltx2_modulate_mount.py @@ -0,0 +1,48 @@ +import unittest + +import torch + +from sglang.multimodal_gen.runtime.models.dits.ltx_2 import _ltx2_modulate + + +class TestLtx2ModulateMount(unittest.TestCase): + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") + def test_row_broadcast_matches_eager(self): + torch.manual_seed(0) + x = torch.randn(2, 517, 4096, device="cuda", dtype=torch.bfloat16) + scale = torch.randn(2, 1, 4096, device="cuda", dtype=torch.bfloat16) + shift = torch.randn(2, 1, 4096, device="cuda", dtype=torch.bfloat16) + reference = x * (1 + scale) + shift + self.assertTrue(torch.equal(reference, _ltx2_modulate(x, scale, shift))) + + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") + def test_non_contiguous_rows_match_eager(self): + # unbind()/squeeze() of the combined adaLN tables produces strided + # (B, 1, D) views; the helper densifies them before the kernel. + torch.manual_seed(1) + x = torch.randn(2, 33, 512, device="cuda", dtype=torch.bfloat16) + table = torch.randn(2, 1, 4, 512, device="cuda", dtype=torch.bfloat16) + scale, shift = table.unbind(dim=2)[:2] + reference = x * (1 + scale) + shift + self.assertTrue(torch.equal(reference, _ltx2_modulate(x, scale, shift))) + + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") + def test_per_token_rows_fall_back(self): + torch.manual_seed(2) + x = torch.randn(2, 16, 128, device="cuda", dtype=torch.bfloat16) + scale = torch.randn(2, 16, 128, device="cuda", dtype=torch.bfloat16) + shift = torch.randn(2, 16, 128, device="cuda", dtype=torch.bfloat16) + reference = x * (1 + scale) + shift + self.assertTrue(torch.equal(reference, _ltx2_modulate(x, scale, shift))) + + def test_cpu_falls_back(self): + torch.manual_seed(3) + x = torch.randn(1, 9, 64, dtype=torch.float32) + scale = torch.randn(1, 1, 64, dtype=torch.float32) + shift = torch.randn(1, 1, 64, dtype=torch.float32) + reference = x * (1 + scale) + shift + self.assertTrue(torch.equal(reference, _ltx2_modulate(x, scale, shift))) + + +if __name__ == "__main__": + unittest.main()