[diffusion] Reuse bit-exact modulation fast path for LTX-2.3 (#34930)

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:

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