[Diffusion] Fuse LingBot per-token gated residual and RMSNorm modulate (#37910)

python/sglang/kernels/kda_kernels/csrc/diffusion/residual_gate_add.cuh

@@ -29,7 +29,7 @@ constexpr uintptr_t kAlignment = 16;
 constexpr uint32_t kTransposeTile = 32;
 constexpr uint32_t kTransposeBlockSize = 256;
 
-enum class GateMode : int { kFull, kBroadcastRow };
+enum class GateMode : int { kFull, kBroadcastRow, kPerToken };
 
 template <typename T>
 SGL_DEVICE T residual_gate_value(T residual, T update, T gate) {
@@ -108,7 +108,10 @@ __global__ void residual_gate_add_scalar_kernel(
     int64_t hidden_size) {
   const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
   for (int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; index < numel; index += stride) {
-    const T gate_value = kGateMode == GateMode::kFull ? gate[index] : SGLANG_LDG(gate + index % hidden_size);
+    const T gate_value = kGateMode == GateMode::kFull
+                             ? gate[index]
+                             : (kGateMode == GateMode::kPerToken ? SGLANG_LDG(gate + index / hidden_size)
+                                                                 : SGLANG_LDG(gate + index % hidden_size));
     out[index] = residual_gate_value(residual[index], update[index], gate_value);
   }
 }
@@ -180,7 +183,7 @@ struct ResidualGateAddKernel {
       tvm::ffi::TensorView update,
       tvm::ffi::TensorView gate,
       int64_t hidden_size,
-      bool broadcast_gate) {
+      int64_t gate_mode) {
     using namespace host;
 
     auto N = SymbolicSize{"numel"};
@@ -192,7 +195,10 @@ struct ResidualGateAddKernel {
 
     const int64_t numel = N.unwrap();
     CHECK_HOST(hidden_size > 0 && numel % hidden_size == 0) << "hidden size must be positive and divide the input size";
-    CHECK_HOST(G.unwrap() == (broadcast_gate ? hidden_size : numel)) << "gate size does not match its mode";
+    // gate_mode: 0 = full (numel), 1 = broadcast row (hidden_size), 2 = per-token (numel/hidden_size)
+    const int64_t rows = numel / hidden_size;
+    const int64_t expected_gate = gate_mode == 0 ? numel : (gate_mode == 1 ? hidden_size : rows);
+    CHECK_HOST(G.unwrap() == expected_gate) << "gate size does not match its mode";
     if (numel == 0) {
       return;
     }
@@ -212,13 +218,26 @@ struct ResidualGateAddKernel {
     const bool vectorized = aligned && hidden_size % kVec == 0;
     if (vectorized) {
       const int64_t num_vectors = numel / kVec;
-      if (!broadcast_gate) {
+      if (gate_mode == 0) {
         const auto blocks =
             static_cast<uint32_t>(std::min<int64_t>(div_ceil(num_vectors, static_cast<int64_t>(kBlockSize)), kMaxGrid));
         LaunchKernel(blocks, kBlockSize, device.unwrap())(
             residual_gate_add_vec_kernel<T, kVec>, out_ptr, residual_ptr, update_ptr, gate_ptr, num_vectors);
         return;
       }
+      if (gate_mode == 2) {
+        const auto blocks =
+            static_cast<uint32_t>(std::min<int64_t>(div_ceil(numel, static_cast<int64_t>(kBlockSize)), kMaxGrid));
+        LaunchKernel(blocks, kBlockSize, device.unwrap())(
+            residual_gate_add_scalar_kernel<T, GateMode::kPerToken>,
+            out_ptr,
+            residual_ptr,
+            update_ptr,
+            gate_ptr,
+            numel,
+            hidden_size);
+        return;
+      }
 
       const int64_t rows = numel / hidden_size;
       const int64_t row_vectors = hidden_size / kVec;
@@ -233,7 +252,16 @@ struct ResidualGateAddKernel {
 
     const auto blocks =
         static_cast<uint32_t>(std::min<int64_t>(div_ceil(numel, static_cast<int64_t>(kBlockSize)), kMaxGrid));
-    if (broadcast_gate) {
+    if (gate_mode == 2) {
+      LaunchKernel(blocks, kBlockSize, device.unwrap())(
+          residual_gate_add_scalar_kernel<T, GateMode::kPerToken>,
+          out_ptr,
+          residual_ptr,
+          update_ptr,
+          gate_ptr,
+          numel,
+          hidden_size);
+    } else if (gate_mode == 1) {
       LaunchKernel(blocks, kBlockSize, device.unwrap())(
           residual_gate_add_scalar_kernel<T, GateMode::kBroadcastRow>,
           out_ptr,

python/sglang/kernels/kda_kernels/residual_gate_add_jit.py

@@ -73,24 +73,42 @@ def _residual_gate_add_custom_op(
     if _is_transposed_dense_residual(residual, update, gate):
         module.residual_gate_add_transposed(out, residual, update, gate)
         return out
-    broadcast_gate = gate.shape != residual.shape
+    gate_mode = _gate_mode(residual, gate)
     module.residual_gate_add(
         out.view(-1),
         residual.view(-1),
         update.view(-1),
         gate.view(-1),
         residual.shape[-1],
-        broadcast_gate,
+        gate_mode,
     )
     return out
 
 
+def _gate_mode(residual: torch.Tensor, gate: torch.Tensor) -> int:
+    """0 = full, 1 = broadcast row (hidden_size), 2 = per-token (rows)."""
+    if gate.shape == residual.shape:
+        return 0
+    if _is_row_broadcast_gate(residual, gate):
+        return 1
+    return 2
+
+
 def _is_row_broadcast_gate(residual: torch.Tensor, gate: torch.Tensor) -> bool:
     if gate.dim() != residual.dim() or gate.shape[-1] != residual.shape[-1]:
         return False
     return all(size == 1 for size in gate.shape[:-1])
 
 
+def _is_per_token_gate(residual: torch.Tensor, gate: torch.Tensor) -> bool:
+    """Gate holds one scalar per token (row), broadcast along the hidden dim."""
+    return (
+        gate.dim() == residual.dim()
+        and gate.shape[-1] == 1
+        and gate.shape[:-1] == residual.shape[:-1]
+    )
+
+
 def _is_transposed_dense_residual(
     residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
 ) -> bool:
@@ -121,7 +139,11 @@ def can_use_residual_gate_add_cuda(
         and residual.dim() >= 2
         and residual.numel() > 0
         and update.shape == residual.shape
-        and (gate.shape == residual.shape or _is_row_broadcast_gate(residual, gate))
+        and (
+            gate.shape == residual.shape
+            or _is_row_broadcast_gate(residual, gate)
+            or _is_per_token_gate(residual, gate)
+        )
         and (
             (residual.is_contiguous() and update.is_contiguous())
             or _is_transposed_dense_residual(residual, update, gate)

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

@@ -581,6 +581,13 @@ _EXPORTS: dict[str, str] = {
     "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",
+    "lingbot_video_gated_residual_active": "sites.lingbot_video_gated_residual_site",
+    "mark_lingbot_video_gated_residual_site": "sites.lingbot_video_gated_residual_site",
+    "mount_lingbot_video_gated_residual": "sites.lingbot_video_gated_residual_site",
+    "try_lingbot_video_gated_residual": "sites.lingbot_video_gated_residual_site",
+    "unmount_lingbot_video_gated_residual": "sites.lingbot_video_gated_residual_site",
+    "can_use_rmsnorm_scale_shift_per_token": "norm.rmsnorm_scale_shift_triton",
+    "rmsnorm_scale_shift_per_token": "norm.rmsnorm_scale_shift_triton",
     "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/norm/rmsnorm_scale_shift_triton.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused RMSNorm + adaLN scale/shift with per-token modulation (Triton).
+
+Computes, for each token row,
+
+    y = rmsnorm(x) * (1 + scale) + shift
+
+where ``x`` is ``[B, S, D]`` (bf16/fp16/fp32), the RMSNorm weight is ``[D]``, and
+the per-token ``scale``/``shift`` are ``[B, S, D]`` views (cast to fp32 in
+kernel). ``scale``/``shift`` are typically non-contiguous ``chunk`` views of the
+``[B, S, 6D]`` modulation tensor, so the kernel takes their row stride and reads
+them strided instead of materializing contiguous copies. The whole chain runs in
+one kernel in fp32 with a single rounding to the output dtype, replacing the
+eager LingBot chain ``(norm(x) * (1 + scale) + shift).to(dtype)``.
+
+Unlike the bit-exact ``rmsnorm_scale_shift_bitexact`` (ERNIE), this kernel does
+*not* reproduce PyTorch's parallel variance reduction order bit-for-bit, so it
+is intended for the request-gated (``quality="extra-high"``/``"high"``) fusion
+path, matching the existing quality-gated LingBot RMSNorm fusion.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+from sglang.srt.utils.custom_op import register_custom_op
+
+
+@triton.jit
+def _rmsnorm_scale_shift_kernel(
+    y_ptr,
+    x_ptr,
+    w_ptr,
+    scale_ptr,
+    shift_ptr,
+    mod_row_stride,
+    SEQ,
+    DIM: tl.constexpr,
+    EPS: tl.constexpr,
+    BLOCK_SIZE_SEQ: tl.constexpr,
+    BLOCK_SIZE_DIM: tl.constexpr,
+):
+    seq_blk_id = tl.program_id(0)
+    seq_id = seq_blk_id * BLOCK_SIZE_SEQ
+
+    seq_offset = seq_id + tl.arange(0, BLOCK_SIZE_SEQ)[:, None]
+    s_mask = seq_offset < SEQ
+    d_offset = tl.arange(0, BLOCK_SIZE_DIM)[None, :]
+    d_mask = d_offset < DIM
+    mask = s_mask & d_mask
+
+    xy_ptr = seq_offset * DIM + d_offset
+    mod_ptr = seq_offset * mod_row_stride + d_offset
+
+    x = tl.load(x_ptr + xy_ptr, mask=mask, other=0.0).to(tl.float32)
+    mean_square = tl.sum(x * x, axis=1, keep_dims=True) / DIM
+    rstd = tl.math.rsqrt(mean_square + EPS)
+    w = tl.load(w_ptr + d_offset, mask=d_mask).to(tl.float32)
+    scale = tl.load(scale_ptr + mod_ptr, mask=mask, other=0.0).to(tl.float32)
+    shift = tl.load(shift_ptr + mod_ptr, mask=mask, other=0.0).to(tl.float32)
+
+    y = (x * rstd * w) * (1.0 + scale) + shift
+    tl.store(y_ptr + xy_ptr, y, mask=mask)
+
+
+@register_custom_op(op_name="rmsnorm_scale_shift_per_token_cuda", out_shape="x")
+def _rmsnorm_scale_shift_per_token_cuda(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    shape = x.shape
+    out = torch.empty_like(x)
+    x2 = x.reshape(-1, shape[-1])
+    out2 = out.reshape(-1, shape[-1])
+    scale2 = scale.reshape(-1, shape[-1])
+    shift2 = shift.reshape(-1, shape[-1])
+    S, D = x2.shape
+
+    # scale/shift share the row stride of their parent modulation tensor
+    # (contiguous rows, possibly strided across the 6*D row).
+    assert scale2.stride(0) == shift2.stride(0)
+    mod_row_stride = scale2.stride(0)
+
+    block_size_seq = min(16, triton.next_power_of_2(max(1, S // 512)))
+    grid = (triton.cdiv(S, block_size_seq),)
+    with torch.get_device_module().device(x.device):
+        _rmsnorm_scale_shift_kernel[grid](
+            out2,
+            x2,
+            weight,
+            scale2,
+            shift2,
+            mod_row_stride,
+            S,
+            DIM=D,
+            EPS=eps,
+            BLOCK_SIZE_DIM=triton.next_power_of_2(D),
+            BLOCK_SIZE_SEQ=block_size_seq,
+        )
+    return out
+
+
+def can_use_rmsnorm_scale_shift_per_token(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+) -> bool:
+    return (
+        x.is_cuda
+        and x.dim() == 3
+        and x.is_contiguous()
+        and x.dtype in (torch.float16, torch.bfloat16, torch.float32)
+        and weight.is_cuda
+        and weight.dim() == 1
+        and weight.numel() == x.shape[-1]
+        and scale.shape == x.shape
+        and shift.shape == x.shape
+        and scale.dtype == shift.dtype
+        and scale.dtype in (torch.float16, torch.bfloat16, torch.float32)
+        # rows must be contiguous (stride(1) == 1); row stride may differ
+        # (non-contiguous chunk views of the [B, S, 6D] modulation tensor).
+        and scale.stride(0) == shift.stride(0)
+        and scale.stride(2) == 1
+        and shift.stride(2) == 1
+        and x.numel() > 0
+    )
+
+
+def rmsnorm_scale_shift_per_token(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """Fused ``rmsnorm(x) * (1 + scale) + shift`` with per-token scale/shift."""
+    return _rmsnorm_scale_shift_per_token_cuda(x, weight, scale, shift, eps)
+
+
+__all__ = [
+    "can_use_rmsnorm_scale_shift_per_token",
+    "rmsnorm_scale_shift_per_token",
+]

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

@@ -0,0 +1,72 @@
+"""LingBot Video per-token gated-residual fusion, gated by request quality.
+
+Each LingBot block applies ``residual + (gate * update).to(residual.dtype)`` at
+the attention and FFN updates, where ``update`` (post-norm) and the per-token
+``gate`` (``[B, S, 1]``) stay in FP32 while ``residual`` is BF16. The shared
+``residual_gate_add`` kernel computes ``residual + update * gate`` in a single
+pass but requires one dtype, so the gate and update are first cast to BF16.
+That reordering of the FP32 multiply is numerically equivalent only at
+half-precision rounding level (not bit-exact), so the fusion is opt-in:
+``quality="extra-high"`` and ``quality="high"`` mount it, while the default
+``quality="lossless"`` keeps the reference FP32-multiply form bit-for-bit.
+"""
+
+from __future__ import annotations
+
+import logging
+
+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 per-token gated residual",
+    marker_attr="_sgl_lingbot_video_gated_residual_site",
+    enabled_attr="_sgl_lingbot_video_gated_residual_enabled",
+)
+
+
+def mark_lingbot_video_gated_residual_site(module: nn.Module) -> None:
+    """Mark a LingBot block; it starts on the reference path."""
+    _FUSION.mark(module)
+
+
+def lingbot_video_gated_residual_active(module: nn.Module) -> bool:
+    return _FUSION.is_enabled(module)
+
+
+def mount_lingbot_video_gated_residual(root: nn.Module) -> bool:
+    return _FUSION.mount(root, logger=logger)
+
+
+def unmount_lingbot_video_gated_residual(root: nn.Module) -> None:
+    _FUSION.unmount(root)
+
+
+def try_lingbot_video_gated_residual(
+    site: nn.Module,
+    residual: torch.Tensor,
+    update: torch.Tensor,
+    gate: torch.Tensor,
+) -> torch.Tensor | None:
+    """Return the fused ``residual + update * gate`` when the site is enabled.
+
+    Returns ``None`` (caller runs the reference path) when the site is off or
+    the tensors are not eligible for the per-token fast path.
+    """
+    if not _FUSION.is_enabled(site):
+        return None
+    from sglang.kernels.ops.diffusion import (
+        can_use_residual_gate_add_cuda,
+        residual_gate_add,
+    )
+
+    if residual.dtype != update.dtype or residual.dtype != gate.dtype:
+        update = update.to(residual.dtype)
+        gate = gate.to(residual.dtype)
+    if not can_use_residual_gate_add_cuda(residual, update, gate):
+        return None
+    return residual_gate_add(residual, update, gate)

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

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import logging
 import math
 from typing import Any, Iterable, Iterator, Optional
 
@@ -11,7 +12,12 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from torch import nn
 
 from sglang.kernels.ops.diffusion import (
+    can_use_rmsnorm_scale_shift_per_token,
+    lingbot_video_gated_residual_active,
+    mark_lingbot_video_gated_residual_site,
     mark_lingbot_video_rmsnorm_site,
+    rmsnorm_scale_shift_per_token,
+    try_lingbot_video_gated_residual,
     try_lingbot_video_rmsnorm,
 )
 from sglang.multimodal_gen.configs.models.dits.lingbot_video_moe import (
@@ -64,6 +70,9 @@ LINGBOT_VIDEO_FP32_MODULES = (
 )
 
 
+logger = logging.getLogger(__name__)
+
+
 def is_lingbot_block(name: str, _module: object) -> bool:
     return "blocks" in name and name.split(".")[-1].isdigit()
 
@@ -95,6 +104,51 @@ class LingBotVideoRMSNorm(nn.Module):
         return (self.weight * hidden_states).to(input_dtype)
 
 
+def _lingbot_gated_residual(
+    block: LingBotVideoBlock,
+    residual: torch.Tensor,
+    update: torch.Tensor,
+    gate: torch.Tensor,
+) -> torch.Tensor:
+    """``residual + (gate * update).to(residual.dtype)`` at a block update.
+
+    Uses the request-gated per-token ``residual_gate_add`` fast path when the
+    block's quality-gated site is enabled; otherwise the reference FP32
+    multiply form (bit-exact for ``quality="lossless"``).
+    """
+    fused = try_lingbot_video_gated_residual(block, residual, update, gate)
+    if fused is not None:
+        return fused
+    return residual + (gate * update).to(residual.dtype)
+
+
+def _lingbot_norm_modulate(
+    block: LingBotVideoBlock,
+    norm: LingBotVideoRMSNorm,
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    """``(norm(x) * scale + shift).to(out_dtype)`` at a block modulate site.
+
+    Uses the request-gated fused RMSNorm+scale+shift kernel when the block's
+    quality-gated site is enabled; otherwise the reference chain. ``scale`` is
+    the raw ``scale_msa/mlp``; the kernel applies ``x * (1 + scale)``.
+    """
+    if lingbot_video_gated_residual_active(
+        block
+    ) and can_use_rmsnorm_scale_shift_per_token(x, norm.weight, scale, shift):
+        try:
+            out = rmsnorm_scale_shift_per_token(
+                x, norm.weight, scale, shift, norm.variance_epsilon
+            )
+            return out if out.dtype == out_dtype else out.to(out_dtype)
+        except Exception:
+            pass
+    return (norm(x) * (1.0 + scale) + shift).to(out_dtype)
+
+
 def make_joint_position_ids(
     text_len: int, grid_t: int, grid_h: int, grid_w: int, device: torch.device
 ) -> torch.Tensor:
@@ -283,6 +337,7 @@ class LingBotVideoBlock(nn.Module):
         self.layer_idx = layer_idx
         h = hidden_size
         self.scale_shift_table = nn.Parameter(torch.zeros(1, 6 * h))
+        mark_lingbot_video_gated_residual_site(self)
         self.norm1 = LingBotVideoRMSNorm(h, norm_eps)
         self.attn = LingBotVideoAttention(
             h,
@@ -336,22 +391,25 @@ class LingBotVideoBlock(nn.Module):
             6, dim=-1
         )
         gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
-        scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
 
         bulk_dtype = self.attn.to_q.weight.dtype
-        attn_in = (self.norm1(x) * scale_msa + shift_msa).to(bulk_dtype)
+        attn_in = _lingbot_norm_modulate(
+            self, self.norm1, x, scale_msa, shift_msa, bulk_dtype
+        )
         attn_out = self.attn(
             attn_in,
             freqs_cis,
             attention_mask=attention_mask,
             attn_mask_meta=attn_mask_meta,
         )
-        x = x + (gate_msa * self.norm_post_attn(attn_out)).to(x.dtype)
+        x = _lingbot_gated_residual(self, x, self.norm_post_attn(attn_out), gate_msa)
 
-        ffn_in = (self.norm2(x) * scale_mlp + shift_mlp).to(bulk_dtype)
+        ffn_in = _lingbot_norm_modulate(
+            self, self.norm2, x, scale_mlp, shift_mlp, bulk_dtype
+        )
         ffn_out = self.ffn(ffn_in)
         ffn_normed = self.norm_post_ffn(ffn_out)
-        x = x + (gate_mlp * ffn_normed).to(x.dtype)
+        x = _lingbot_gated_residual(self, x, ffn_normed, gate_mlp)
         return x
 
 

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

@@ -26,6 +26,7 @@ from sglang.kernels.ops.diffusion import (
     mount_fused_linear_gelu,
     mount_fused_ln_modulate,
     mount_hunyuan_qknorm,
+    mount_lingbot_video_gated_residual,
     mount_lingbot_video_rmsnorm,
     mount_ltx2_rms_norm_modulate,
     mount_nvfp4_bias_gelu,
@@ -36,6 +37,7 @@ from sglang.kernels.ops.diffusion import (
     unmount_fused_linear_gelu,
     unmount_fused_ln_modulate,
     unmount_hunyuan_qknorm,
+    unmount_lingbot_video_gated_residual,
     unmount_lingbot_video_rmsnorm,
     unmount_ltx2_rms_norm_modulate,
     unmount_nvfp4_bias_gelu,
@@ -215,6 +217,11 @@ _QUALITY_FUSION_HANDLERS: tuple[
         mount_lingbot_video_rmsnorm,
         unmount_lingbot_video_rmsnorm,
     ),
+    (
+        "LingBot Video per-token gated residual",
+        mount_lingbot_video_gated_residual,
+        unmount_lingbot_video_gated_residual,
+    ),
     (
         "SANA-Video BF16-input linear attention",
         mount_sana_video_linear_attention,