[diffusion] LTX-2 quality=high fused RMSNorm+modulate + FFN GELU epilogue (H200 ltx23-one-stage denoise 45.85->43.24 s, ~matches torch.compile) (#34172)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

python/sglang/kernels/ops/diffusion/ltx2_rmsnorm_modulate.py

@@ -0,0 +1,80 @@
+"""Weightless RMSNorm + adaLN modulate folded into one kernel for LTX-2.
+
+``rms_norm(x) * (1 + scale) + shift`` at the LTX-2 transformer-block adaLN
+sites is otherwise an aten ``F.rms_norm`` plus a separate ``mul``/``add``
+modulate (one reduction kernel plus several pointwise passes per site). This
+folds the whole chain into a single ``fused_rmsnorm_scale_shift_bitexact``
+launch.
+
+The fused kernel reproduces the RMSNorm math via ``rsqrt.approx`` (the
+flashinfer CuTe form), which differs from aten's refined ``rsqrtf`` by at
+most one bf16 ULP on a small fraction of elements. It is therefore *not*
+bit-exact vs the eager reference, so the fold is opt-in per batch: model code
+marks its adaLN sites with :func:`mark_ltx2_rms_norm_modulate_site` (default
+off, reference path) and the denoising stage calls
+:func:`mount_ltx2_rms_norm_modulate` / :func:`unmount_ltx2_rms_norm_modulate`
+at batch boundaries for ``quality="high"`` requests.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from sglang.kernels.ops.diffusion.quality_gate import QualityGatedFusion
+from sglang.kernels.ops.diffusion.triton.rmsnorm_scale_shift_bitexact import (
+    can_use_fused_rmsnorm_scale_shift,
+    fused_rmsnorm_scale_shift_bitexact,
+)
+
+_SITE_MARKER_ATTR = "_sgl_ltx2_rms_norm_modulate_site"
+_SITE_ENABLED_ATTR = "_sgl_ltx2_rms_norm_modulate_enabled"
+_FUSION = QualityGatedFusion(
+    name="LTX-2 RMSNorm+modulate",
+    marker_attr=_SITE_MARKER_ATTR,
+    enabled_attr=_SITE_ENABLED_ATTR,
+)
+
+# ``RMSNormNoWeight`` applies no scale, so a ones weight reproduces it exactly.
+_ONES_WEIGHT_CACHE: dict[tuple[torch.device, int], torch.Tensor] = {}
+
+
+def mark_ltx2_rms_norm_modulate_site(module: nn.Module) -> None:
+    """Mark ``module`` as an LTX-2 RMSNorm+modulate fusion site (mounted off)."""
+    _FUSION.mark(module)
+
+
+def ltx2_rms_norm_modulate_active(module: nn.Module) -> bool:
+    return _FUSION.is_enabled(module)
+
+
+def mount_ltx2_rms_norm_modulate(root: nn.Module) -> bool:
+    return _FUSION.mount(root)
+
+
+def unmount_ltx2_rms_norm_modulate(root: nn.Module) -> None:
+    _FUSION.unmount(root)
+
+
+def _ones_weight(x: torch.Tensor) -> torch.Tensor:
+    key = (x.device, int(x.shape[-1]))
+    w = _ONES_WEIGHT_CACHE.get(key)
+    if w is None:
+        w = torch.ones(x.shape[-1], device=x.device, dtype=torch.bfloat16)
+        _ONES_WEIGHT_CACHE[key] = w
+    return w
+
+
+def can_fuse_ltx2_rms_norm_modulate(
+    x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor
+) -> bool:
+    if x.dtype is not torch.bfloat16 or not x.is_cuda:
+        return False
+    return can_use_fused_rmsnorm_scale_shift(x, _ones_weight(x), scale, shift)
+
+
+def fused_ltx2_rms_norm_modulate(
+    x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, eps: float
+) -> torch.Tensor:
+    """``rms_norm(x) * (1 + scale) + shift`` as one kernel (weightless RMSNorm)."""
+    return fused_rmsnorm_scale_shift_bitexact(x, _ones_weight(x), scale, shift, eps)

python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py

@@ -10,10 +10,22 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from sglang.kernels.ops.diffusion.fused_linear_gelu import (
+    can_fuse_linear_gelu,
+    fused_gelu_active,
+    fused_linear_gelu_tanh,
+    mark_fused_gelu_site,
+)
 from sglang.kernels.ops.diffusion.ltx2_qknorm_split_rope import (
     can_use_ltx2_qknorm_split_rope_cuda,
     ltx2_qknorm_split_rope_cuda,
 )
+from sglang.kernels.ops.diffusion.ltx2_rmsnorm_modulate import (
+    can_fuse_ltx2_rms_norm_modulate,
+    fused_ltx2_rms_norm_modulate,
+    ltx2_rms_norm_modulate_active,
+    mark_ltx2_rms_norm_modulate_site,
+)
 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.runtime.distributed import (
@@ -125,6 +137,29 @@ def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int:
     )
 
 
+def _ltx2_rms_norm_modulate(
+    block: nn.Module,
+    rms_norm: nn.Module,
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """``rms_norm(x) * (1 + scale) + shift`` for the LTX-2 adaLN sites.
+
+    Folds the weightless RMSNorm and the modulate into one kernel when the
+    ``quality="high"`` fusion is mounted on ``block`` and the per-call guard
+    passes; otherwise the verbatim eager reference chain (the ``lossless``
+    default). The fused kernel is not bit-exact (<=1 bf16 ULP) so it is gated
+    on the request-scoped mount rather than a runtime self-check.
+    """
+    if ltx2_rms_norm_modulate_active(block) and can_fuse_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
+
+
 def _ltx2_disable_fused_ada_values(exc: Exception) -> None:
     global _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED
     _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED = True
@@ -977,8 +1012,12 @@ class LTX2FeedForward(nn.Module):
             input_is_parallel=True,
             quant_config=quant_config,
         )
+        mark_fused_gelu_site(self, "proj_in")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if fused_gelu_active(self) and can_fuse_linear_gelu(self.proj_in, x):
+            x = fused_linear_gelu_tanh(x, self.proj_in.weight, self.proj_in.bias)
+        else:
             x, _ = self.proj_in(x)
             x = self.act(x)
         x, _ = self.proj_out(x)
@@ -1108,6 +1147,7 @@ class LTX2TransformerBlock(nn.Module):
 
         # 4. Feedforward layers
         self.ff = LTX2FeedForward(dim, dim_out=dim, quant_config=quant_config)
+        mark_ltx2_rms_norm_modulate_site(self)
         self.audio_ff = LTX2FeedForward(
             audio_dim, dim_out=audio_dim, quant_config=quant_config
         )
@@ -1200,8 +1240,8 @@ class LTX2TransformerBlock(nn.Module):
             )
         else:
             vshift_msa, vscale_msa, vgate_msa = video_ada_values[0:3]
-        norm_hidden_states = (
+        norm_hidden_states = _ltx2_rms_norm_modulate(
-            self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_msa) + vshift_msa
+            self, self.rms_norm, hidden_states, vscale_msa, vshift_msa, self.norm_eps
         )
         attn_hidden_states = self.attn1(
             norm_hidden_states,
@@ -1220,9 +1260,13 @@ class LTX2TransformerBlock(nn.Module):
             )
         else:
             ashift_msa, ascale_msa, agate_msa = audio_ada_values[0:3]
-        norm_audio_hidden_states = (
+        norm_audio_hidden_states = _ltx2_rms_norm_modulate(
-            self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_msa)
+            self,
-            + ashift_msa
+            self.rms_norm,
+            audio_hidden_states,
+            ascale_msa,
+            ashift_msa,
+            self.norm_eps,
         )
         attn_audio_hidden_states = self.audio_attn1(
             norm_audio_hidden_states,
@@ -1251,8 +1295,8 @@ class LTX2TransformerBlock(nn.Module):
             v_prompt_shift, v_prompt_scale = self.get_ada_values(
                 self.prompt_scale_shift_table, batch_size, temb_prompt, slice(None)
             )
-            norm_hidden_states = (
+            norm_hidden_states = _ltx2_rms_norm_modulate(
-                self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_q) + vshift_q
+                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
@@ -1278,9 +1322,13 @@ class LTX2TransformerBlock(nn.Module):
                 temb_audio_prompt,
                 slice(None),
             )
-            norm_audio_hidden_states = (
+            norm_audio_hidden_states = _ltx2_rms_norm_modulate(
-                self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_q)
+                self,
-                + ashift_q
+                self.rms_norm,
+                audio_hidden_states,
+                ascale_q,
+                ashift_q,
+                self.norm_eps,
             )
             mod_audio_encoder_hidden_states = (
                 audio_encoder_hidden_states * (1 + a_prompt_scale) + a_prompt_shift
@@ -1428,8 +1476,8 @@ class LTX2TransformerBlock(nn.Module):
             )
         else:
             vshift_mlp, vscale_mlp, vgate_mlp = video_ada_values[3:6]
-        norm_hidden_states = (
+        norm_hidden_states = _ltx2_rms_norm_modulate(
-            self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
+            self, self.rms_norm, hidden_states, vscale_mlp, vshift_mlp, self.norm_eps
         )
         ff_output = self.ff(norm_hidden_states)
         hidden_states = residual_gate_add(hidden_states, ff_output, vgate_mlp)
@@ -1440,9 +1488,13 @@ class LTX2TransformerBlock(nn.Module):
             )
         else:
             ashift_mlp, ascale_mlp, agate_mlp = audio_ada_values[3:6]
-        norm_audio_hidden_states = (
+        norm_audio_hidden_states = _ltx2_rms_norm_modulate(
-            self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_mlp)
+            self,
-            + ashift_mlp
+            self.rms_norm,
+            audio_hidden_states,
+            ascale_mlp,
+            ashift_mlp,
+            self.norm_eps,
         )
         audio_ff_output = self.audio_ff(norm_audio_hidden_states)
         audio_hidden_states = residual_gate_add(

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -32,6 +32,10 @@ from sglang.kernels.ops.diffusion.fused_ln_modulate import (
     mount_fused_ln_modulate,
     unmount_fused_ln_modulate,
 )
+from sglang.kernels.ops.diffusion.ltx2_rmsnorm_modulate import (
+    mount_ltx2_rms_norm_modulate,
+    unmount_ltx2_rms_norm_modulate,
+)
 from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType, STA_Mode
 from sglang.multimodal_gen.configs.pipeline_configs.flux import (
@@ -156,6 +160,11 @@ _QUALITY_FUSION_HANDLERS: tuple[
         mount_fused_ln_modulate,
         unmount_fused_ln_modulate,
     ),
+    (
+        "LTX-2 fused RMSNorm+modulate",
+        mount_ltx2_rms_norm_modulate,
+        unmount_ltx2_rms_norm_modulate,
+    ),
     (
         "fused gate RMSNorm (BF16-native Triton)",
         mount_fused_gate_rmsnorm,

test/registered/kernels/ops/diffusion/test_ltx2_rms_norm_modulate.py

@@ -0,0 +1,71 @@
+"""LTX-2 quality=high RMSNorm+modulate fusion: gated, close to eager."""
+
+import sys
+
+import pytest
+import torch
+from torch import nn
+
+from sglang.kernels.ops.diffusion.ltx2_rmsnorm_modulate import (
+    fused_ltx2_rms_norm_modulate,
+    mark_ltx2_rms_norm_modulate_site,
+    mount_ltx2_rms_norm_modulate,
+    unmount_ltx2_rms_norm_modulate,
+)
+from sglang.multimodal_gen.runtime.layers.layernorm import RMSNormNoWeight
+from sglang.multimodal_gen.runtime.models.dits.ltx_2 import _ltx2_rms_norm_modulate
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+
+register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=8, suite="nightly-amd-kernel-1-gpu", nightly=True)
+
+
+@pytest.fixture(autouse=True)
+def _setup():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    torch.cuda.manual_seed(0)
+
+
+def _eager(rms, x, scale, shift, eps):
+    return rms(x, eps) * (1 + scale) + shift
+
+
+def _inputs(hidden, batch=1, seq=4096):
+    rms = RMSNormNoWeight()
+    x = torch.randn(batch, seq, hidden, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
+    shift = torch.randn(batch, 1, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
+    return rms, x, scale, shift
+
+
+# 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.
+    block = nn.Module()
+    mark_ltx2_rms_norm_modulate_site(block)
+    rms, x, scale, shift = _inputs(hidden)
+    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()
+    mark_ltx2_rms_norm_modulate_site(block)
+    assert mount_ltx2_rms_norm_modulate(block)
+    try:
+        rms, x, scale, shift = _inputs(hidden)
+        out = _ltx2_rms_norm_modulate(block, rms, x, scale, shift, 1e-6)
+        # The mounted path routes through the fused kernel exactly.
+        assert torch.equal(out, fused_ltx2_rms_norm_modulate(x, scale, shift, 1e-6))
+        # And stays within half-precision rounding of the eager reference.
+        ref = _eager(rms, x, scale, shift, 1e-6)
+        assert torch.allclose(out.float(), ref.float(), atol=3e-2, rtol=1e-2)
+    finally:
+        unmount_ltx2_rms_norm_modulate(block)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))