[diffusion] Accelerate LingBot Video RMSNorm in quality=high (#35969)

docs/docs/sglang-diffusion/fused_kernels.mdx

@@ -66,6 +66,7 @@ These fusion families mount under `quality="high"`:
 | LTX-2 RMSNorm + modulate | `rms_norm(x) * (1 + scale) + shift` in one launch |
 | Gate RMSNorm (BF16-native) | `RMSNorm + tanh + mul + add` in one pass |
 | HunyuanVideo strided QK RMSNorm | Per-head QK RMSNorm over the packed QKV layout |
+| LingBot Video fused RMSNorm | Replaces the handwritten cast, square, mean, rsqrt, and multiply chain with existing Triton RMSNorm kernels |
 | SANA-Video BF16-input linear attention | Keeps the first linear-attention GEMM's inputs in BF16 with FP32 accumulation/output; the second GEMM remains FP32 |
 
 ## Kernel inventory
@@ -156,6 +157,7 @@ Kernels are written against a specific eager chain in a specific model, so cover
 | LTX-2 | QK-norm + split RoPE, ada-values split, RMSNorm+modulate, modulate, residual-gate add, linear+GELU |
 | LTX-2.5 decoder | paired 3D RoPE with shared axis-table cache |
 | HunyuanVideo | QKV+RoPE pack, strided QK RMSNorm, linear+GELU |
+| LingBot Video MoE | Fused RMSNorm at `quality=high` |
 | Sana | LN+modulate, GLUMB bias+SiLU / bias+GLU, residual-gate add |
 | SANA-Video | Packed QKV/KV; paired fp64 interleaved RoPE; LN+modulate, GLUMB bias+SiLU / bias+GLU, and residual-gate add during BCG; BF16-input linear attention at `quality=high` |
 | Sana-WM | bidirectional gated delta-net, fused QK inverse-RMS |

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

@@ -470,6 +470,11 @@ _EXPORTS: dict[str, str] = {
     "mark_ltx2_rms_norm_modulate_site": "sites.ltx2_rmsnorm_modulate_site",
     "mount_ltx2_rms_norm_modulate": "sites.ltx2_rmsnorm_modulate_site",
     "unmount_ltx2_rms_norm_modulate": "sites.ltx2_rmsnorm_modulate_site",
+    "lingbot_video_rmsnorm_active": "sites.lingbot_video_rmsnorm_site",
+    "mark_lingbot_video_rmsnorm_site": "sites.lingbot_video_rmsnorm_site",
+    "mount_lingbot_video_rmsnorm": "sites.lingbot_video_rmsnorm_site",
+    "try_lingbot_video_rmsnorm": "sites.lingbot_video_rmsnorm_site",
+    "unmount_lingbot_video_rmsnorm": "sites.lingbot_video_rmsnorm_site",
     "mark_sana_video_linear_attention_site": "sites.sana_video_linear_attention_site",
     "mount_sana_video_linear_attention": "sites.sana_video_linear_attention_site",
     "sana_video_linear_attention_active": "sites.sana_video_linear_attention_site",

python/sglang/kernels/ops/diffusion/sites/lingbot_video_rmsnorm_site.py

@@ -0,0 +1,101 @@
+"""LingBot Video fused RMSNorm, gated by request quality.
+
+LingBot Video keeps norm weights and statistics in FP32. Its reference module
+spells RMSNorm as separate cast, square, mean, rsqrt, multiply, weight, and
+output-cast operations. Some checkpoint norm weights remain BF16 after load;
+the reference formula still promotes their multiplication through the FP32
+hidden states. For ``quality="high"``, existing diffusion Triton RMSNorm
+kernels replace that chain. Wide rows with FP32 weights use the one-row
+``norm_infer`` kernel; the remaining sites use the tiled one-pass kernel. Their
+reduction order is not bit-exact, so the default ``quality="lossless"`` path
+remains unchanged.
+"""
+
+from __future__ import annotations
+
+import logging
+from importlib import import_module
+
+import torch
+import torch.nn as nn
+
+from sglang.kernels.ops.diffusion.sites.quality_gate import QualityGatedFusion
+
+logger = logging.getLogger(__name__)
+
+_FUSION = QualityGatedFusion(
+    name="LingBot Video fused RMSNorm",
+    marker_attr="_sgl_lingbot_video_rmsnorm_site",
+    enabled_attr="_sgl_lingbot_video_rmsnorm_enabled",
+)
+
+
+def mark_lingbot_video_rmsnorm_site(module: nn.Module) -> None:
+    """Mark a LingBot RMSNorm module; it starts on the reference path."""
+    _FUSION.mark(module)
+
+
+def lingbot_video_rmsnorm_active(module: nn.Module) -> bool:
+    return _FUSION.is_enabled(module)
+
+
+def _site_reject_reason(site: nn.Module) -> str | None:
+    try:
+        import_module("triton")
+    except ImportError:
+        return "triton unavailable"
+    weight = getattr(site, "weight", None)
+    if not torch.is_tensor(weight) or weight.dim() != 1:
+        return "missing or non-1D norm weight"
+    if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        return f"unsupported norm weight dtype {weight.dtype}"
+    if weight.stride(0) != 1:
+        return "non-contiguous norm weight"
+    return None
+
+
+def mount_lingbot_video_rmsnorm(root: nn.Module) -> bool:
+    return _FUSION.mount(root, reject_reason=_site_reject_reason, logger=logger)
+
+
+def unmount_lingbot_video_rmsnorm(root: nn.Module) -> None:
+    _FUSION.unmount(root)
+
+
+def try_lingbot_video_rmsnorm(
+    site: nn.Module,
+    hidden_states: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor | None:
+    """Return the quality-gated RMSNorm result, or ``None`` to fall back."""
+    if not (
+        _FUSION.is_enabled(site)
+        and hidden_states.is_cuda
+        and hidden_states.dtype in (torch.float16, torch.bfloat16)
+        and hidden_states.is_contiguous()
+        and hidden_states.shape[-1] == weight.numel()
+        and weight.is_cuda
+        and weight.device == hidden_states.device
+        and weight.dtype in (hidden_states.dtype, torch.float32)
+    ):
+        return None
+
+    hidden_size = hidden_states.shape[-1]
+    if weight.dtype == torch.float32 and hidden_size > 128:
+        from sglang.kernels.ops.diffusion.norm.norm_triton import norm_infer
+
+        shape = hidden_states.shape
+        return norm_infer(
+            hidden_states.view(-1, hidden_size),
+            weight,
+            bias=None,
+            eps=eps,
+            is_rms_norm=True,
+        ).view(shape)
+
+    from sglang.kernels.ops.diffusion.norm.rmsnorm_onepass_triton import (
+        triton_one_pass_rms_norm,
+    )
+
+    return triton_one_pass_rms_norm(hidden_states, weight, eps)

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

@@ -10,6 +10,10 @@ import torch.nn.functional as F
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from torch import nn
 
+from sglang.kernels.ops.diffusion import (
+    mark_lingbot_video_rmsnorm_site,
+    try_lingbot_video_rmsnorm,
+)
 from sglang.multimodal_gen.configs.models.dits.lingbot_video_moe import (
     LingBotVideoMoEConfig,
 )
@@ -75,8 +79,15 @@ class LingBotVideoRMSNorm(nn.Module):
         super().__init__()
         self.weight = nn.Parameter(torch.ones(dim))
         self.variance_epsilon = eps
+        mark_lingbot_video_rmsnorm_site(self)
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        fused = try_lingbot_video_rmsnorm(
+            self, hidden_states, self.weight, self.variance_epsilon
+        )
+        if fused is not None:
+            return fused
+
         input_dtype = hidden_states.dtype
         hidden_states = hidden_states.to(torch.float32)
         variance = hidden_states.pow(2).mean(-1, keepdim=True)

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

@@ -25,12 +25,14 @@ from sglang.kernels.ops.diffusion import (
     mount_fused_linear_gelu,
     mount_fused_ln_modulate,
     mount_hunyuan_qknorm,
+    mount_lingbot_video_rmsnorm,
     mount_ltx2_rms_norm_modulate,
     mount_sana_video_linear_attention,
     unmount_fused_gate_rmsnorm,
     unmount_fused_linear_gelu,
     unmount_fused_ln_modulate,
     unmount_hunyuan_qknorm,
+    unmount_lingbot_video_rmsnorm,
     unmount_ltx2_rms_norm_modulate,
     unmount_sana_video_linear_attention,
 )
@@ -184,6 +186,11 @@ _QUALITY_FUSION_HANDLERS: tuple[
         mount_hunyuan_qknorm,
         unmount_hunyuan_qknorm,
     ),
+    (
+        "LingBot Video fused RMSNorm",
+        mount_lingbot_video_rmsnorm,
+        unmount_lingbot_video_rmsnorm,
+    ),
     (
         "SANA-Video BF16-input linear attention",
         mount_sana_video_linear_attention,

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

@@ -27,6 +27,7 @@ import torch.nn.functional as F
 
 import sglang.kernels.ops.diffusion.sites.fused_gate_rmsnorm_site as gate_rmsnorm
 import sglang.kernels.ops.diffusion.sites.fused_linear_gelu_site as linear_gelu
+import sglang.kernels.ops.diffusion.sites.lingbot_video_rmsnorm_site as lingbot_video_rmsnorm
 import sglang.kernels.ops.diffusion.sites.sana_video_linear_attention_site as sana_video_linear_attention
 from sglang.kernels.ops.diffusion import (
     BitExactFusionGate,
@@ -456,5 +457,52 @@ def test_sana_video_linear_attention_quality_path_and_guards():
     assert not sana_video_linear_attention.sana_video_linear_attention_active(site)
 
 
+@requires_cuda
+@torch.no_grad()
+def test_lingbot_video_rmsnorm_quality_path_and_guards():
+    torch.manual_seed(1)
+    site = nn.Module()
+    site.weight = nn.Parameter(torch.randn(2048, device="cuda", dtype=torch.float32))
+    lingbot_video_rmsnorm.mark_lingbot_video_rmsnorm_site(site)
+    hidden_states = torch.randn(1, 128, 2048, device="cuda", dtype=torch.bfloat16)
+
+    assert (
+        lingbot_video_rmsnorm.try_lingbot_video_rmsnorm(
+            site, hidden_states, site.weight, 1e-6
+        )
+        is None
+    )
+    assert lingbot_video_rmsnorm.mount_lingbot_video_rmsnorm(site)
+    output = lingbot_video_rmsnorm.try_lingbot_video_rmsnorm(
+        site, hidden_states, site.weight, 1e-6
+    )
+    states_fp32 = hidden_states.float()
+    variance = states_fp32.pow(2).mean(-1, keepdim=True)
+    normalized = states_fp32 * torch.rsqrt(variance + 1e-6)
+    reference = (site.weight * normalized).bfloat16()
+    torch.testing.assert_close(output, reference, atol=2e-2, rtol=2e-2)
+
+    bf16_site = nn.Module()
+    bf16_site.weight = nn.Parameter(
+        torch.randn(2048, device="cuda", dtype=torch.bfloat16)
+    )
+    lingbot_video_rmsnorm.mark_lingbot_video_rmsnorm_site(bf16_site)
+    assert lingbot_video_rmsnorm.mount_lingbot_video_rmsnorm(bf16_site)
+    bf16_output = lingbot_video_rmsnorm.try_lingbot_video_rmsnorm(
+        bf16_site, hidden_states, bf16_site.weight, 1e-6
+    )
+    bf16_reference = (bf16_site.weight * normalized).bfloat16()
+    torch.testing.assert_close(bf16_output, bf16_reference, atol=3e-2, rtol=3e-2)
+
+    assert (
+        lingbot_video_rmsnorm.try_lingbot_video_rmsnorm(
+            site, hidden_states.float(), site.weight, 1e-6
+        )
+        is None
+    )
+    lingbot_video_rmsnorm.unmount_lingbot_video_rmsnorm(site)
+    assert not lingbot_video_rmsnorm.lingbot_video_rmsnorm_active(site)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v"]))