From 46b92b22e2505be598a0804318942f7e92140302 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Mon, 24 Aug 2026 18:02:00 +0800 Subject: [PATCH] [diffusion] Accelerate LingBot Video RMSNorm in quality=high (#35969) --- docs/docs/sglang-diffusion/fused_kernels.mdx | 2 + .../sglang/kernels/ops/diffusion/__init__.py | 5 + .../sites/lingbot_video_rmsnorm_site.py | 101 ++++++++++++++++++ .../runtime/models/dits/lingbot_video_moe.py | 11 ++ .../pipelines_core/stages/denoising.py | 7 ++ .../kernels/ops/diffusion/test_sites.py | 48 +++++++++ 6 files changed, 174 insertions(+) create mode 100644 python/sglang/kernels/ops/diffusion/sites/lingbot_video_rmsnorm_site.py diff --git a/docs/docs/sglang-diffusion/fused_kernels.mdx b/docs/docs/sglang-diffusion/fused_kernels.mdx index af3201324..25ac51cbd 100644 --- a/docs/docs/sglang-diffusion/fused_kernels.mdx +++ b/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 | diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index efeb37173..2cba1cbb2 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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", diff --git a/python/sglang/kernels/ops/diffusion/sites/lingbot_video_rmsnorm_site.py b/python/sglang/kernels/ops/diffusion/sites/lingbot_video_rmsnorm_site.py new file mode 100644 index 000000000..578d4f3d1 --- /dev/null +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_video_moe.py b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_video_moe.py index c7d2e353a..b01a48ab7 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_video_moe.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 6b787333f..836ed2a57 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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, diff --git a/test/registered/kernels/ops/diffusion/test_sites.py b/test/registered/kernels/ops/diffusion/test_sites.py index 866185a49..ed71093e1 100644 --- a/test/registered/kernels/ops/diffusion/test_sites.py +++ b/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"]))