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 819269c98..1efab045c 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -213,7 +213,15 @@ def _ltx2_rms_norm_modulate( x, scale, shift ): return fused_ltx2_rms_norm_modulate(x, scale, shift, eps) - return rms_norm(x, eps) * (1 + scale) + shift + normed = rms_norm(x, eps) + if torch.compiler.is_compiling(): + # Let Inductor fuse this chain into its surrounding graph. Routing a + # compiled call through the opaque custom op would be a regression. + return normed * (1 + scale) + shift + # Reuse the bit-exact first-sight-verified eager modulate kernel. This + # removes two large broadcast pointwise launches without changing the + # reference rounding. + return _ltx2_modulate(normed, scale, shift) def _ltx2_disable_fused_ada_values(exc: Exception) -> None: diff --git a/test/registered/kernels/ops/diffusion/test_ltx2_rms_norm_modulate.py b/test/registered/kernels/ops/diffusion/test_ltx2_rms_norm_modulate.py index 292452469..37b6ea37b 100644 --- a/test/registered/kernels/ops/diffusion/test_ltx2_rms_norm_modulate.py +++ b/test/registered/kernels/ops/diffusion/test_ltx2_rms_norm_modulate.py @@ -6,6 +6,7 @@ import pytest import torch from torch import nn +import sglang.multimodal_gen.runtime.models.dits.ltx_2 as ltx2_module from sglang.kernels.ops.diffusion.ltx2_rmsnorm_modulate import ( fused_ltx2_rms_norm_modulate, mark_ltx2_rms_norm_modulate_site, @@ -42,7 +43,8 @@ def _inputs(hidden, batch=1, seq=4096): # hidden 4096 = LTX-2 video stream, 2048 = audio stream. @pytest.mark.parametrize("hidden", [4096, 2048]) def test_lossless_default_is_bitexact(hidden): - # A marked-but-unmounted site (the lossless default) runs verbatim eager. + # A marked-but-unmounted site uses only the self-verified bit-exact + # modulate fast path after the reference aten RMSNorm. block = nn.Module() mark_ltx2_rms_norm_modulate_site(block) rms, x, scale, shift = _inputs(hidden) @@ -50,6 +52,21 @@ def test_lossless_default_is_bitexact(hidden): assert torch.equal(out, _eager(rms, x, scale, shift, 1e-6)) +def test_lossless_compile_keeps_expression_visible_to_inductor(monkeypatch): + block = nn.Module() + mark_ltx2_rms_norm_modulate_site(block) + rms, x, scale, shift = _inputs(2048, seq=126) + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + monkeypatch.setattr( + ltx2_module, + "_ltx2_modulate", + lambda *_args: pytest.fail("compiled path must not call the opaque custom op"), + ) + + out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6) + assert torch.equal(out, _eager(rms, x, scale, shift, 1e-6)) + + @pytest.mark.parametrize("hidden", [4096, 2048]) def test_mounted_high_uses_fused_kernel(hidden): block = nn.Module()