[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>
2026-08-10 09:46:17 +08:00
parent c20e99bd22
commit 441910f926
4 changed files with 229 additions and 17 deletions
@@ -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)
@@ -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,10 +1012,14 @@ 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:
-        x, _ = self.proj_in(x)
-        x = self.act(x)
+        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)
        return 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 = (
-            self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_msa) + vshift_msa
+        norm_hidden_states = _ltx2_rms_norm_modulate(
+            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 = (
-            self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_msa)
-            + ashift_msa
+        norm_audio_hidden_states = _ltx2_rms_norm_modulate(
+            self,
+            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 = (
-                self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_q) + vshift_q
+            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
@@ -1278,9 +1322,13 @@ class LTX2TransformerBlock(nn.Module):
                temb_audio_prompt,
                slice(None),
            )
-            norm_audio_hidden_states = (
-                self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_q)
-                + ashift_q
+            norm_audio_hidden_states = _ltx2_rms_norm_modulate(
+                self,
+                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 = (
-            self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
+        norm_hidden_states = _ltx2_rms_norm_modulate(
+            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 = (
-            self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_mlp)
-            + ashift_mlp
+        norm_audio_hidden_states = _ltx2_rms_norm_modulate(
+            self,
+            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(
@@ -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,