From 54c2c99feb452f9a64ece1f70ba4392e59e67b0a Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Fri, 4 Sep 2026 14:35:13 +0800 Subject: [PATCH] [Diffusion] Fuse LingBot per-token gated residual and RMSNorm modulate (#37910) --- docs/docs/sglang-diffusion/fused_kernels.mdx | 3 +- .../csrc/diffusion/residual_gate_add.cuh | 40 ++++- .../kda_kernels/residual_gate_add_jit.py | 28 +++- .../sglang/kernels/ops/diffusion/__init__.py | 7 + .../norm/rmsnorm_scale_shift_triton.py | 149 ++++++++++++++++++ .../lingbot_video_gated_residual_site.py | 72 +++++++++ .../runtime/models/dits/lingbot_video_moe.py | 68 +++++++- .../pipelines_core/stages/denoising.py | 7 + .../kernels/ops/diffusion/test_modulate.py | 89 +++++++++++ 9 files changed, 448 insertions(+), 15 deletions(-) create mode 100644 python/sglang/kernels/ops/diffusion/norm/rmsnorm_scale_shift_triton.py create mode 100644 python/sglang/kernels/ops/diffusion/sites/lingbot_video_gated_residual_site.py diff --git a/docs/docs/sglang-diffusion/fused_kernels.mdx b/docs/docs/sglang-diffusion/fused_kernels.mdx index 1aa20412f..0bd3c8a47 100644 --- a/docs/docs/sglang-diffusion/fused_kernels.mdx +++ b/docs/docs/sglang-diffusion/fused_kernels.mdx @@ -86,6 +86,7 @@ These fusion families mount under both `quality="extra-high"` and | 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 | +| LingBot Video per-token gated residual + RMSNorm modulate | Folds `residual + gate * update` (per-token `[B, S, 1]` gate) and the `rmsnorm(x) * (1 + scale) + shift` adaLN chain (strided `[B, S, 6D]` chunk views) into single 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 | | FLUX-family VAE fast paths | Channels-last decode, GroupNorm(+SiLU), upsample, and attention replacements for FLUX.2 and AutoencoderKL-based FLUX.1, Z-Image, and SD3 pipelines | | Wan VAE RMSNorm + SiLU | Replaces the channel-first RMSNorm/SiLU chain while keeping the decode in `channels_last_3d` | @@ -192,7 +193,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 / Helios | QKV+RoPE pack, strided QK RMSNorm, linear+GELU; Helios also has paired in-place Q/K RoPE | -| LingBot Video MoE | Fused RMSNorm at `quality=extra-high` or `quality=high` | +| LingBot Video MoE | Fused RMSNorm, per-token gated residual, and fused RMSNorm+modulate at `quality=extra-high` or `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=extra-high` or `quality=high` | | Sana-WM | bidirectional gated delta-net, fused QK inverse-RMS | diff --git a/python/sglang/kernels/kda_kernels/csrc/diffusion/residual_gate_add.cuh b/python/sglang/kernels/kda_kernels/csrc/diffusion/residual_gate_add.cuh index f70e464b5..31ad61ea8 100644 --- a/python/sglang/kernels/kda_kernels/csrc/diffusion/residual_gate_add.cuh +++ b/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 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(gridDim.x) * blockDim.x; for (int64_t index = static_cast(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(std::min(div_ceil(num_vectors, static_cast(kBlockSize)), kMaxGrid)); LaunchKernel(blocks, kBlockSize, device.unwrap())( residual_gate_add_vec_kernel, out_ptr, residual_ptr, update_ptr, gate_ptr, num_vectors); return; } + if (gate_mode == 2) { + const auto blocks = + static_cast(std::min(div_ceil(numel, static_cast(kBlockSize)), kMaxGrid)); + LaunchKernel(blocks, kBlockSize, device.unwrap())( + residual_gate_add_scalar_kernel, + 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(std::min(div_ceil(numel, static_cast(kBlockSize)), kMaxGrid)); - if (broadcast_gate) { + if (gate_mode == 2) { + LaunchKernel(blocks, kBlockSize, device.unwrap())( + residual_gate_add_scalar_kernel, + 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, out_ptr, diff --git a/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py b/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py index 594b0cad8..a3741b85c 100644 --- a/python/sglang/kernels/kda_kernels/residual_gate_add_jit.py +++ b/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) diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 5e2a31262..1c8cea5fe 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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", diff --git a/python/sglang/kernels/ops/diffusion/norm/rmsnorm_scale_shift_triton.py b/python/sglang/kernels/ops/diffusion/norm/rmsnorm_scale_shift_triton.py new file mode 100644 index 000000000..cd1c745ad --- /dev/null +++ b/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", +] diff --git a/python/sglang/kernels/ops/diffusion/sites/lingbot_video_gated_residual_site.py b/python/sglang/kernels/ops/diffusion/sites/lingbot_video_gated_residual_site.py new file mode 100644 index 000000000..c16b07c05 --- /dev/null +++ b/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) 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 012cd4781..c475e1206 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 @@ -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 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 23e319508..f5ac2cd03 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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, diff --git a/test/registered/kernels/ops/diffusion/test_modulate.py b/test/registered/kernels/ops/diffusion/test_modulate.py index b6f4b11db..436e060e4 100644 --- a/test/registered/kernels/ops/diffusion/test_modulate.py +++ b/test/registered/kernels/ops/diffusion/test_modulate.py @@ -16,6 +16,7 @@ from sglang.kernels.jit.utils import get_ci_test_range from sglang.kernels.ops.diffusion import ( can_use_modulate_scale_shift_cuda, can_use_residual_gate_add_cuda, + can_use_rmsnorm_scale_shift_per_token, fuse_layernorm_scale_shift_gate_select01_kernel, fuse_residual_layernorm_scale_shift_gate_select01_kernel, fuse_scale_shift_kernel, @@ -25,6 +26,7 @@ from sglang.kernels.ops.diffusion import ( norm_infer, residual_gate_add, residual_gate_add_cuda, + rmsnorm_scale_shift_per_token, timestep_embedding, try_fused_scaled_residual_add_exact, ) @@ -164,6 +166,37 @@ def test_residual_gate_add_matches_torch(residual_shape, gate_shape): assert torch.equal(residual_gate_add(residual, update, gate), ref) +# LingBot per-token gates are [B, S, 1]: one scalar per token, broadcast +# along the hidden dimension. +PER_TOKEN_GATE_CASES = [ + ((1, 2560, 512), (1, 2560, 1)), + ((1, 17, 65), (1, 17, 1)), + ((2, 33, 128), (2, 33, 1)), +] + + +@pytest.mark.parametrize("residual_shape,gate_shape", PER_TOKEN_GATE_CASES) +def test_residual_gate_add_per_token_matches_torch(residual_shape, gate_shape): + residual = torch.randn(residual_shape, device=DEVICE, dtype=torch.bfloat16) + update = torch.randn_like(residual) + gate = torch.randn(gate_shape, device=DEVICE, dtype=torch.bfloat16) + + assert can_use_residual_gate_add_cuda(residual, update, gate) + ref = residual + update * gate + _assert_gate_add(residual_gate_add_cuda(residual, update, gate), ref) + assert torch.equal(residual_gate_add(residual, update, gate), ref) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_residual_gate_add_per_token_dtypes(dtype): + residual = torch.randn((1, 2560, 512), device=DEVICE, dtype=dtype) + update = torch.randn_like(residual) + gate = torch.randn((1, 2560, 1), device=DEVICE, dtype=dtype) + _assert_gate_add( + residual_gate_add_cuda(residual, update, gate), residual + update * gate + ) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @pytest.mark.parametrize("gate_shape", [(1, 1, 64), (1, 9, 64)]) def test_residual_gate_add_dtypes(dtype, gate_shape): @@ -524,3 +557,59 @@ def test_timestep_embedding_matches_diffusers( if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) + +# --------------------------------------------------------------------------- +# fused RMSNorm + per-token adaLN scale/shift (quality-gated, LingBot) +# --------------------------------------------------------------------------- + + +def _eager_lingbot_norm_modulate(x, weight, scale, shift, eps): + xf = x.to(torch.float32) + var = xf.pow(2).mean(-1, keepdim=True) + xf = xf * torch.rsqrt(var + eps) + normed = (weight.to(torch.float32) * xf).to(x.dtype) + return (normed * (1.0 + scale.to(torch.float32)) + shift.to(torch.float32)).to( + x.dtype + ) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("shape", [(1, 4813, 2048), (1, 2560, 512), (2, 33, 128)]) +def test_rmsnorm_scale_shift_per_token_matches_eager(shape, dtype): + B, S, H = shape + x = torch.randn(shape, device=DEVICE, dtype=dtype) + weight = torch.randn(H, device=DEVICE, dtype=torch.float32) + # scale/shift are non-contiguous chunk views of the [B, S, 6D] modulation, + # matching the LingBot adaLN layout the kernel is built for. + mod = torch.randn((B, S, 6 * H), device=DEVICE, dtype=torch.float32) + shift, scale = mod.chunk(6, dim=-1)[0], mod.chunk(6, dim=-1)[1] + eps = 1e-6 + + assert can_use_rmsnorm_scale_shift_per_token(x, weight, scale, shift) + ref = _eager_lingbot_norm_modulate(x, weight, scale, shift, eps) + out = rmsnorm_scale_shift_per_token(x, weight, scale, shift, eps) + assert out.dtype == x.dtype and out.shape == x.shape + # Not bit-exact (single fp32 pass); assert bf16/fp16 rounding tolerance. + torch.testing.assert_close(out, ref, atol=0.13, rtol=0.02) + + +def test_rmsnorm_scale_shift_per_token_guards(): + B, S, H = 1, 64, 128 + x = torch.randn((B, S, H), device=DEVICE, dtype=torch.bfloat16) + weight = torch.randn(H, device=DEVICE, dtype=torch.float32) + scale = torch.randn((B, S, H), device=DEVICE, dtype=torch.float32) + shift = torch.randn((B, S, H), device=DEVICE, dtype=torch.float32) + assert can_use_rmsnorm_scale_shift_per_token(x, weight, scale, shift) + + assert not can_use_rmsnorm_scale_shift_per_token( + x.cpu(), weight, scale, shift + ) # not on device + assert not can_use_rmsnorm_scale_shift_per_token( + x, weight, scale, shift[:, :, ::2] + ) # strided rows (stride(2) != 1) + assert not can_use_rmsnorm_scale_shift_per_token( + x, weight, scale.float(), shift.double() + ) # mismatched scale/shift dtype + assert not can_use_rmsnorm_scale_shift_per_token( + x, weight[:-1], scale, shift + ) # weight size mismatch