[diffusion] fuse LingBot MoE group-limited top-k index selection (#38044)
Co-authored-by: BBuf <bbuf@users.noreply.github.com> Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
This commit is contained in:
Xiaoyu Zhang
co-authored by
BBuf
Mick Qian
BBuf
Mick Qian
parent
50c1bf0db0
commit
bd16c22a04
@@ -4,7 +4,7 @@ description: "The fused CUDA/Triton kernels SGLang Diffusion ships, what each on
|
||||
tag: "preserve"
|
||||
---
|
||||
|
||||
Diffusion transformers and VAEs spend a large share of their non-GEMM time on short elementwise chains — adaLN modulate, residual gating, QK-norm, RoPE, norm epilogues — each of which is a separate kernel launch and a separate HBM round trip in eager PyTorch. SGLang Diffusion replaces these chains with fused kernels under [`sglang/kernels/ops/diffusion`](https://github.com/sgl-project/sglang/tree/main/python/sglang/kernels/ops/diffusion).
|
||||
Diffusion transformers and VAEs spend a large share of their non-GEMM time on short elementwise and selection chains — adaLN modulate, residual gating, QK-norm, RoPE, norm epilogues, and MoE routing — each of which can require several kernel launches and HBM round trips in eager PyTorch. SGLang Diffusion replaces these chains with fused kernels under [`sglang/kernels/ops/diffusion`](https://github.com/sgl-project/sglang/tree/main/python/sglang/kernels/ops/diffusion).
|
||||
|
||||
This page is an inventory: what each kernel fuses, what its numerical contract is, and which models use it. It is not a lever you tune — most of these kernels are on by default and require no flag. The one switch is `--quality`, described below.
|
||||
|
||||
@@ -18,6 +18,8 @@ Multi-step denoising amplifies a per-step rounding difference into visible quali
|
||||
|
||||
**Model/checkpoint-native.** Generic close-contract kernels, sparse operators, and FP8/NVFP4 producers can be part of a model implementation or a separately selected deployment path. They are documented in the inventory, but `quality` does not select or undo those choices.
|
||||
|
||||
**Selection-equivalent routing — enabled unconditionally.** LingBot Video's fused group-limited top-k returns the same selected expert-id set as its guarded CUDA `torch.topk(..., sorted=False)` reference chain. The order of those ids is not part of either path's contract. Because the selected experts are unchanged, this path does not depend on the request `quality` tier.
|
||||
|
||||
<Note>
|
||||
A plain fp32 single-pass norm fusion looks harmless and is not. On ERNIE-Image it moved the 50-step trajectory to 18.83 dB PSNR, which is what motivated the bit-exact rewrite of that path.
|
||||
</Note>
|
||||
@@ -93,19 +95,19 @@ These fusion families mount under both `quality="extra-high"` and
|
||||
|
||||
## Kernel inventory
|
||||
|
||||
43 operators are registered in the kernel registry across 47 implementations (some operators carry several backends). Backends are named by provenance, not device: `JIT` compiles under nvcc *and* hipcc, `TRITON` runs on CUDA and ROCm, `CUTE_DSL` needs CUTLASS, `FLYDSL` is ROCm gfx950 only, `AOT` comes from the `sgl_kernel` wheel.
|
||||
45 operators are registered in the kernel registry across 51 implementations (some operators carry several backends). Backends are named by provenance, not device: `KDA` identifies Kernel Design Agents implementations, `JIT` compiles under nvcc *and* hipcc, `TRITON` identifies Triton sources, `CUTE_DSL` needs CUTLASS, `FLYDSL` is ROCm gfx950 only, and `AOT` comes from the `sgl_kernel` wheel. Per-operator capability metadata determines which devices can load each implementation.
|
||||
|
||||
### Normalization
|
||||
|
||||
| Operator | Backend | Contract | Replaces |
|
||||
| --- | --- | --- | --- |
|
||||
| `rmsnorm_scale_shift` | Triton | bit-exact | RMSNorm + `* (1 + scale) + shift` (4 kernels) |
|
||||
| `scale_residual_norm_scale_shift` | Triton / CuTe-DSL / FlyDSL | bit-exact (Triton) | the above plus the preceding `residual + gate * update` |
|
||||
| `scale_residual_norm_scale_shift` | KDA / Triton / CuTe-DSL / FlyDSL | bit-exact (Triton) or backend-specific close contract | the above plus the preceding `residual + gate * update` |
|
||||
| `scale_residual_norm_scale_shift_nvfp4` | JIT CUDA | matches the selected NVFP4 producer contract | Qwen residual LayerNorm/modulation + FC1 NVFP4 quantization |
|
||||
| `layernorm_modulate` | Triton | bit-exact | affine-free LayerNorm + adaLN modulate |
|
||||
| `qk_head_layernorm` | Triton | bit-exact | per-head LayerNorm on q/k |
|
||||
| `qk_rmsnorm_native` | Triton | bit-exact | Z-Image per-head QK RMSNorm |
|
||||
| `norm_scale_shift` | CuTe-DSL / FlyDSL | fp32 statistics | LN-or-RMS + scale/shift, many broadcast modes |
|
||||
| `norm_scale_shift` | KDA / CuTe-DSL / FlyDSL | backend-specific close contract | LN-or-RMS + scale/shift, many broadcast modes |
|
||||
| `rmsnorm_scale`, `rmsnorm_tanh_residual` | Triton | bf16-native statistics | `RMSNorm(x) * scale`, `x + tanh(gate) * RMSNorm(y)` |
|
||||
| `apply_group_norm_silu` | Triton | close | `GroupNorm + SiLU`, NCHW-contiguous |
|
||||
| `group_norm_silu_4d`, `group_norm_silu_rows` | Triton | close | channels-last GroupNorm(+SiLU); what lets a VAE decoder run channels_last end to end with no `nchwToNhwc` transposes |
|
||||
@@ -116,7 +118,7 @@ These fusion families mount under both `quality="extra-high"` and
|
||||
| Operator | Backend | Contract | Replaces |
|
||||
| --- | --- | --- | --- |
|
||||
| `modulate_scale_shift` | JIT CUDA | bit-exact | `x * (1 + scale) + shift` |
|
||||
| `residual_gate_add` | JIT CUDA | bit-exact | `residual + gate * update` |
|
||||
| `residual_gate_add` | KDA (JIT CUDA) | bit-exact | `residual + gate * update` |
|
||||
| `timestep_embedding` | JIT CUDA | close | sinusoidal timestep embedding |
|
||||
| `temb_table_slices` | Triton | bit-exact | see note below |
|
||||
| `ltx2_ada_values` | Triton | bit-exact | LTX-2 nine-way adaLN value split, slices come out contiguous |
|
||||
@@ -130,12 +132,12 @@ These fusion families mount under both `quality="extra-high"` and
|
||||
| Operator | Backend | Contract | Replaces |
|
||||
| --- | --- | --- | --- |
|
||||
| `fused_inplace_qknorm_rope` | JIT CUDA | one bf16 rounding step vs the split baseline; exact with `round_norm_before_rope=True` | separate QK-norm kernel + RoPE |
|
||||
| `flux2_qkv_epilogue` | JIT CUDA | bit-exact against its selected BF16 reference chain | FLUX.2 QK RMSNorm + RoPE + joint text/image QKV packing |
|
||||
| `flux2_qkv_epilogue` | KDA (JIT CUDA) | bit-exact against its selected BF16 reference chain | FLUX.2 QK RMSNorm + RoPE + joint text/image QKV packing |
|
||||
| `qwen_qkv_epilogue` | JIT CUDA | bit-exact against its selected BF16 reference chain | Qwen-Image QK RMSNorm + RoPE + joint QKV writes on SM100+ |
|
||||
| `rope_rotate_half` | Triton | bit-exact | `chunk` → `cat(-x2, x1)` → two muls + add → `cat(tail)`, about 7 kernels per projection |
|
||||
| `interleaved_rope_fp64` | JIT CUDA | bit-exact | paired SANA-Video Q/K RoPE with fp64 tables, about 14 eager kernels |
|
||||
| `helios_qk_rope` | JIT CUDA | bit-exact | paired in-place Helios Q/K RoPE with transposed frequency layout |
|
||||
| `ltx2_qknorm_split_rope` | JIT CUDA | close (validated on B200) | LTX-2 QK-norm + split RoPE |
|
||||
| `ltx2_qknorm_split_rope` | KDA (JIT CUDA) | close (validated on B200) | LTX-2 QK-norm + split RoPE |
|
||||
| `ltx25_decoder_rope` | JIT CUDA | bit-exact | paired LTX-2.5 decoder 3D RoPE from cached compact axis tables |
|
||||
| `hunyuan_qkv_rope_pack` | Triton | bit-exact | QKV pack and RoPE in one pass |
|
||||
|
||||
@@ -154,6 +156,12 @@ These fusion families mount under both `quality="extra-high"` and
|
||||
| `sparse_linear_attn_fwd` | Triton | block-map, compression, and forward for sparse linear attention |
|
||||
| `bigdn` | Triton | Sana-WM bidirectional gated delta-net; the chunkwise form splits phase A along the KV and Z streams so two blocks stay resident per SM, and stores `(I - P)` so phase B's MMA folds the identity-add in |
|
||||
|
||||
### MoE routing
|
||||
|
||||
| Operator | Backend | Contract | Replaces |
|
||||
| --- | --- | --- | --- |
|
||||
| `group_limited_topk` | Triton | selected expert-id set matches the guarded CUDA reference; output order is unspecified | LingBot Video's per-group top-2 reduction, group top-k, mask construction, masked expert scores, and final expert top-k |
|
||||
|
||||
### Data movement
|
||||
|
||||
Every kernel here only moves values (plus zero fill, plus at most one same-order add), so each is bitwise identical to the aten chain it replaces.
|
||||
@@ -164,7 +172,7 @@ Every kernel here only moves values (plus zero fill, plus at most one same-order
|
||||
| `pack_qkv_destination_major` | Triton | Ulysses destination-major QKV pack |
|
||||
| `varlen_pack_qkv`, `varlen_scatter_to_padded` | Triton | varlen gather/scatter around the masked attention path |
|
||||
| `varlen_pack_segmented_qkv` | Triton | varlen gather from a virtual prefix/main Q/K/V sequence |
|
||||
| `causal_conv3d_cat_pad` | JIT CUDA / Triton | causal Conv3d `cat` + `pad` |
|
||||
| `causal_conv3d_cat_pad` | KDA (JIT CUDA) / Triton | causal Conv3d `cat` + `pad` |
|
||||
| `cat_pad_channels_last_3d` | Triton | Wan causal VAE `cat + F.pad + contiguous` (three passes plus cache bookkeeping) in one pass |
|
||||
| `dup_up3d_add` | Triton | `repeat_interleave + permute().contiguous() + add` |
|
||||
|
||||
@@ -174,7 +182,8 @@ These kernels preserve the quantized checkpoint path's selected reference operat
|
||||
|
||||
| Operator | Backend | Replaces |
|
||||
| --- | --- | --- |
|
||||
| `flux2_token_cat_fp8` | Triton | FLUX.2 single-block attention/MLP concatenation plus static FP8 quantization |
|
||||
| `flux2_layernorm_modulate_fp8_quant` | KDA (Triton) | FLUX.2 LayerNorm plus adaLN modulation directly into static FP8 output |
|
||||
| `flux2_token_cat_fp8` | KDA (Triton) | FLUX.2 single-block attention/MLP concatenation plus static FP8 quantization |
|
||||
| `flux2_token_cat_nvfp4` | JIT CUDA | FLUX.2 single-block attention/MLP concatenation plus NVFP4 quantization |
|
||||
|
||||
## Coverage by model
|
||||
@@ -184,7 +193,7 @@ Kernels are written against a specific eager chain in a specific model, so cover
|
||||
| Model | Fused paths |
|
||||
| --- | --- |
|
||||
| FLUX.1 | LN+modulate, modulate, residual-gate add, linear+GELU |
|
||||
| FLUX.2 | LN+modulate, packed SwiGLU, gated residual/norm, residual-gate add, QK RMSNorm+RoPE+joint QKV packing, FP8/NVFP4 token-cat producers |
|
||||
| FLUX.2 | LN+modulate, fused LN+modulate-to-FP8, packed SwiGLU, gated residual/norm, residual-gate add, QK RMSNorm+RoPE+joint QKV packing, FP8/NVFP4 token-cat producers |
|
||||
| Qwen-Image | linear+GELU, select-0/1 LN modulation, added-QKV fusion, QK RMSNorm+RoPE+joint QKV writes, residual norm/modulate+NVFP4 producer |
|
||||
| GLM-Image | LN+modulate, per-head qk LN, residual-gate add, linear+GELU |
|
||||
| ERNIE-Image | RMSNorm+scale/shift, residual-gated variant, rotate-half RoPE, residual-gate add |
|
||||
@@ -193,7 +202,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, per-token gated residual, and fused RMSNorm+modulate at `quality=extra-high` or `quality=high` |
|
||||
| LingBot Video MoE | Default-on group-limited top-k expert selection; 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 |
|
||||
|
||||
@@ -38,6 +38,7 @@ modulate/ adaLN modulate, gating, timestep conditioning
|
||||
rope/ rotary embeddings and the QK-norm chains fused into them
|
||||
activation/ SiLU / GLU / GELU fusions
|
||||
attention/ sparse linear attention, gated delta-net
|
||||
routing/ diffusion-model MoE routing and expert selection
|
||||
layout/ pure data movement: USP/Ulysses relayout, varlen pack, causal pad
|
||||
common/ numerics primitives, platform predicates, non-Triton fallbacks
|
||||
sites/ request-scoped mount policy — NOT kernels (see below)
|
||||
@@ -49,7 +50,7 @@ ext/ JIT C++/CUDA extensions (Hunyuan3D raster/inpaint) — NOT kernels
|
||||
|
||||
**Bit-exact (`torch.equal` vs the eager chain) → mounted unconditionally.**
|
||||
These kernels reproduce every aten rounding boundary, sometimes down to the
|
||||
reduction tree: `norm/layernorm_modulate_triton.py` replicates torch 2.11's
|
||||
reduction tree: `../../kda_kernels/layernorm_modulate_triton.py` replicates torch 2.11's
|
||||
`vectorized_layer_norm_kernel` (128-thread Welford, `_rcp4` guarded
|
||||
reciprocal, `shfl.down` fold order, `div.rn` + `MUFU.RSQ`), and
|
||||
`norm/rmsnorm_scale_shift_bitexact.py` replicates flashinfer's CuTe-DSL
|
||||
@@ -139,7 +140,7 @@ tensor copy per residual site.
|
||||
|---|---|---|
|
||||
| `fused_inplace_qknorm_rope` | JIT CUDA | one bf16 rounding step vs split baseline; `round_norm_before_rope=True` makes it exact; supports compact and full-width NeoX/interleaved caches |
|
||||
| `fused_qknorm_rope_pack_kv` | JIT CUDA | as above, also packs prefix K/V |
|
||||
| `try_fused_flux2_qkv_epilogue` | JIT CUDA | bit-exact vs the selected BF16 chain | FLUX.2 QK RMSNorm + RoPE + joint QKV packing |
|
||||
| `try_fused_flux2_qkv_epilogue` | KDA (JIT CUDA) | bit-exact vs the selected BF16 chain | FLUX.2 QK RMSNorm + RoPE + joint QKV packing |
|
||||
| `try_fused_qwen_qkv_epilogue` | JIT CUDA | bit-exact vs the selected BF16 chain | Qwen-Image QK RMSNorm + RoPE + joint QKV writes; SM100+ |
|
||||
| `fused_rope_rotate_half_bitexact` | Triton | bit-exact (elementwise only) |
|
||||
| `fused_interleaved_rope_fp64` | JIT CUDA | bit-exact vs paired SANA-Video fp64 RoPE |
|
||||
@@ -149,6 +150,12 @@ tensor copy per residual site.
|
||||
| `apply_rotary_embedding` | Triton (+fallbacks) | close; the generic entry point |
|
||||
| `hunyuan_qkv_rope_pack` | Triton | bit-exact; packs QKV and applies RoPE in one pass |
|
||||
|
||||
### MoE routing
|
||||
|
||||
| Entry point | Backend | Contract | Applies to |
|
||||
|---|---|---|---|
|
||||
| `group_limited_topk` | Triton | selected expert-id set matches the guarded CUDA `torch.topk(..., sorted=False)` chain; output order is unspecified | LingBot Video's default-on sigmoid+bias group-limited routing; contiguous fp32 `[tokens, experts]`, at least two power-of-two experts per group |
|
||||
|
||||
### Data movement and quantized layout producers
|
||||
|
||||
`usp_merge_heads`, `pack_qkv_destination_major`, `fused_pack_qkv`,
|
||||
@@ -156,9 +163,10 @@ tensor copy per residual site.
|
||||
`fused_causal_conv3d_cat_pad_cuda`,
|
||||
`cat_pad_channels_last_3d`, `dup_up3d_add`, `fused_temb_table_slices`,
|
||||
and `ltx2_ada_values9` are bit-exact data movement or same-order arithmetic.
|
||||
`try_flux2_token_cat_fp8` and `try_flux2_token_cat_nvfp4` fuse branch
|
||||
concatenation directly into the quantized representation selected by the
|
||||
FLUX.2 checkpoint path.
|
||||
`fused_layernorm_modulate_fp8_quant_raw` folds FLUX.2 LayerNorm, adaLN
|
||||
modulation, and static FP8 quantization. `try_flux2_token_cat_fp8` and
|
||||
`try_flux2_token_cat_nvfp4` fuse branch concatenation directly into the
|
||||
quantized representation selected by the FLUX.2 checkpoint path.
|
||||
|
||||
`fused_temb_table_slices` is worth knowing about: the eager
|
||||
`(table + temb.float()).chunk(6, dim=2)` materializes ~8 GB of fp32 at
|
||||
|
||||
@@ -9,7 +9,8 @@ caller to the file layout; ``test_import_surface.py`` guards against it. The
|
||||
one exception is a test that deliberately exercises a single backend.
|
||||
|
||||
Layout -- ordinary implementations use one subpackage per **operator domain**
|
||||
(``norm``, ``modulate``, ``rope``, ``activation``, ``attention``, ``layout``).
|
||||
(``norm``, ``modulate``, ``rope``, ``activation``, ``attention``, ``routing``,
|
||||
``layout``).
|
||||
Implementations generated by kernel-design agents live in
|
||||
``sglang.kernels.kda_kernels`` and are still exported through this facade.
|
||||
``common`` holds shared numerics and platform plumbing, ``sites`` the
|
||||
@@ -327,6 +328,13 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
|
||||
_CUDA,
|
||||
"Sana-WM bidirectional gated delta-net.",
|
||||
),
|
||||
(
|
||||
"diffusion.group_limited_topk",
|
||||
KernelBackend.TRITON,
|
||||
"routing.group_limited_topk_triton:group_limited_topk",
|
||||
_CUDA,
|
||||
"LingBot Video group-limited MoE top-k expert selection.",
|
||||
),
|
||||
(
|
||||
"diffusion.usp_merge_heads",
|
||||
KernelBackend.JIT,
|
||||
@@ -518,6 +526,9 @@ _EXPORTS: dict[str, str] = {
|
||||
"prepare_rope_tables": "attention.sana_wm_gdn_triton",
|
||||
"_attn_fwd": "attention.sparse_linear_attn_triton",
|
||||
"get_block_map": "attention.sparse_linear_attn_triton",
|
||||
# MoE routing
|
||||
"can_use_group_limited_topk": "routing.group_limited_topk_triton",
|
||||
"group_limited_topk": "routing.group_limited_topk_triton",
|
||||
# Data movement: bitwise identical to the aten chains they replace
|
||||
"can_use_fused_causal_conv3d_cat_pad_cuda": "sglang.kernels.kda_kernels.causal_conv3d_cat_pad_jit",
|
||||
"fused_causal_conv3d_cat_pad_cuda": "sglang.kernels.kda_kernels.causal_conv3d_cat_pad_jit",
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Diffusion-specific routing kernels."""
|
||||
@@ -0,0 +1,187 @@
|
||||
"""Fused group-limited MoE top-k index selection for diffusion routers.
|
||||
|
||||
The reference LingBot Video router builds the group-limited top-k with a chain
|
||||
of small kernels: per-group top-2 and sum, group top-k, a ``scatter_`` into a
|
||||
zero mask, an ``expand``/``reshape`` broadcast, a ``masked_fill`` with
|
||||
``-inf``, and the final expert top-k. On a launch-bound single GPU that chain
|
||||
is pure overhead: every intermediate tensor is tiny and the whole computation
|
||||
is bandwidth- and launch-bound. The later score gather remains in the caller.
|
||||
|
||||
This module fuses the entire selection into a single Triton kernel: one
|
||||
program per token loads its score row once, reduces the per-group sums in
|
||||
registers, masks non-selected groups with ``-inf``, and writes the top-k
|
||||
expert ids. The selected expert-id set matches the reference CUDA
|
||||
``torch.topk`` chain for the guarded layouts, including the production
|
||||
128-expert / 4-group / 2-selected-group / top-8 configuration. The output
|
||||
order is intentionally unspecified, matching the reference's ``sorted=False``
|
||||
contract.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _group_limited_topk_kernel(
|
||||
scores_ptr, # [T, E] f32, rows are scores_for_choice (scores + bias)
|
||||
out_idx_ptr, # [T, TOP_K] i64
|
||||
stride_st,
|
||||
E: tl.constexpr,
|
||||
N_GROUP: tl.constexpr,
|
||||
TOPK_GROUP: tl.constexpr,
|
||||
TOP_K: tl.constexpr,
|
||||
EPG: tl.constexpr, # experts per group = E // N_GROUP
|
||||
BLOCK_E: tl.constexpr, # padded E
|
||||
BLOCK_EPG: tl.constexpr, # padded experts-per-group
|
||||
BLOCK_G: tl.constexpr, # padded N_GROUP
|
||||
):
|
||||
t = tl.program_id(0)
|
||||
offs_e = tl.arange(0, BLOCK_E)
|
||||
e_mask = offs_e < E
|
||||
scores = tl.load(
|
||||
scores_ptr + t * stride_st + offs_e, mask=e_mask, other=float("-inf")
|
||||
)
|
||||
|
||||
# Per-group scores -> [BLOCK_G, BLOCK_EPG], pad with -inf so padded lanes
|
||||
# never win the per-group top-2 reduction.
|
||||
g = tl.reshape(scores, (BLOCK_G, BLOCK_EPG), can_reorder=False)
|
||||
epg_mask = tl.arange(0, BLOCK_EPG)[None, :] < EPG
|
||||
g = tl.where(epg_mask, g, float("-inf"))
|
||||
|
||||
# group score = sum of top-2 experts within each group.
|
||||
group_e = tl.arange(0, BLOCK_EPG)[None, :]
|
||||
m1 = tl.max(g, axis=1)
|
||||
# Remove exactly one copy of the first maximum. Masking every value equal
|
||||
# to m1 would lose the second top-k entry when a group contains duplicate
|
||||
# maxima, which changes both the group score and the selected experts.
|
||||
m1_idx = tl.min(
|
||||
tl.where(g == m1[:, None], group_e, BLOCK_EPG),
|
||||
axis=1,
|
||||
)
|
||||
g2 = tl.where(group_e == m1_idx[:, None], float("-inf"), g)
|
||||
m2 = tl.max(g2, axis=1)
|
||||
group_scores = m1 + m2
|
||||
|
||||
# Select TOPK_GROUP groups by descending group score with an explicit
|
||||
# left tie-break. Correctness tests compare the selected set because the
|
||||
# reference uses torch.topk(..., sorted=False).
|
||||
group_idx = tl.arange(0, BLOCK_G)
|
||||
gs_valid = tl.where(group_idx < N_GROUP, group_scores, float("-inf"))
|
||||
selected_group = tl.zeros((BLOCK_G,), dtype=tl.int1)
|
||||
for _ in tl.static_range(TOPK_GROUP):
|
||||
picked_idx = tl.argmax(
|
||||
gs_valid,
|
||||
axis=0,
|
||||
tie_break_left=True,
|
||||
)
|
||||
is_pick = group_idx == picked_idx
|
||||
selected_group = selected_group | is_pick
|
||||
gs_valid = tl.where(is_pick, float("-inf"), gs_valid)
|
||||
|
||||
# Mask experts in non-selected groups, then flat top-k (same tie-break).
|
||||
masked = tl.where(selected_group[:, None], g, float("-inf"))
|
||||
flat = tl.reshape(masked, (BLOCK_E,), can_reorder=False)
|
||||
flat = tl.where(e_mask, flat, float("-inf"))
|
||||
for kk in tl.static_range(TOP_K):
|
||||
idx = tl.argmax(flat, axis=0, tie_break_left=True)
|
||||
tl.store(out_idx_ptr + t * TOP_K + kk, idx.to(tl.int64))
|
||||
flat = tl.where(offs_e == idx, float("-inf"), flat)
|
||||
|
||||
|
||||
def _next_pow2(n: int) -> int:
|
||||
return 1 << (n - 1).bit_length()
|
||||
|
||||
|
||||
def can_use_group_limited_topk(
|
||||
scores_for_choice: torch.Tensor,
|
||||
n_group: int,
|
||||
topk_group: int,
|
||||
top_k: int,
|
||||
) -> bool:
|
||||
"""Return whether the fused CUDA path supports this routing problem."""
|
||||
if not scores_for_choice.is_cuda or torch.version.hip is not None:
|
||||
return False
|
||||
if scores_for_choice.ndim != 2 or scores_for_choice.dtype != torch.float32:
|
||||
return False
|
||||
if not scores_for_choice.is_contiguous() or scores_for_choice.shape[0] == 0:
|
||||
return False
|
||||
|
||||
num_experts = scores_for_choice.shape[1]
|
||||
if n_group <= 1 or num_experts == 0 or num_experts % n_group != 0:
|
||||
return False
|
||||
experts_per_group = num_experts // n_group
|
||||
if experts_per_group < 2 or experts_per_group & (experts_per_group - 1):
|
||||
return False
|
||||
return 0 < topk_group <= n_group and 0 < top_k <= topk_group * experts_per_group
|
||||
|
||||
|
||||
def _fake_group_limited_topk(
|
||||
scores_for_choice: torch.Tensor,
|
||||
n_group: int,
|
||||
topk_group: int,
|
||||
top_k: int,
|
||||
) -> torch.Tensor:
|
||||
del n_group, topk_group
|
||||
return scores_for_choice.new_empty(
|
||||
(scores_for_choice.shape[0], top_k), dtype=torch.int64
|
||||
)
|
||||
|
||||
|
||||
@register_custom_op(
|
||||
op_name="diffusion_group_limited_topk",
|
||||
mutates_args=[],
|
||||
fake_impl=_fake_group_limited_topk,
|
||||
)
|
||||
def _group_limited_topk_cuda(
|
||||
scores_for_choice: torch.Tensor,
|
||||
n_group: int,
|
||||
topk_group: int,
|
||||
top_k: int,
|
||||
) -> torch.Tensor:
|
||||
t, e = scores_for_choice.shape
|
||||
epg = e // n_group
|
||||
out = torch.empty((t, top_k), dtype=torch.int64, device=scores_for_choice.device)
|
||||
_group_limited_topk_kernel[(t,)](
|
||||
scores_for_choice,
|
||||
out,
|
||||
scores_for_choice.stride(0),
|
||||
E=e,
|
||||
N_GROUP=n_group,
|
||||
TOPK_GROUP=topk_group,
|
||||
TOP_K=top_k,
|
||||
EPG=epg,
|
||||
BLOCK_E=_next_pow2(e),
|
||||
BLOCK_EPG=_next_pow2(epg),
|
||||
BLOCK_G=_next_pow2(n_group),
|
||||
num_warps=4,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def group_limited_topk(
|
||||
scores_for_choice: torch.Tensor,
|
||||
n_group: int,
|
||||
topk_group: int,
|
||||
top_k: int,
|
||||
) -> torch.Tensor:
|
||||
"""Fused group-limited top-k expert ids.
|
||||
|
||||
``scores_for_choice`` is the per-token expert score used for selection
|
||||
(already includes the correction bias), shape ``[T, E]`` float32. Returns
|
||||
the selected expert ids as ``[T, top_k]`` int64. The selected set matches
|
||||
the reference two-stage group-limited selection; output order is not part
|
||||
of the contract.
|
||||
"""
|
||||
if not can_use_group_limited_topk(scores_for_choice, n_group, topk_group, top_k):
|
||||
raise ValueError(
|
||||
"group_limited_topk requires a nonempty contiguous CUDA float32 "
|
||||
"[tokens, experts] tensor, at least two power-of-two experts per "
|
||||
"group, 1 < n_group, 0 < topk_group <= n_group, and top_k no "
|
||||
"larger than the selected-group capacity"
|
||||
)
|
||||
return _group_limited_topk_cuda(scores_for_choice, n_group, topk_group, top_k)
|
||||
+3
-1
@@ -51,7 +51,7 @@ If any benchmark, perf-dump, or `torch.profiler` command prints one of those sig
|
||||
## Main Reference
|
||||
|
||||
- [benchmark-and-profile.md](benchmark-and-profile.md) — canonical denoise benchmark, perf dump, and `torch.profiler` workflow; uses checked-in nightly-aligned presets plus current-source extras such as LongCat image/edit, Qwen base edit/layered, SD3.5, SANA-Video/SANA-WM, LingBot Video/World, Cosmos3 Edge/Super I2V/distilled and the explicit Super TP2 x CFG2 comparator, LTX-2.5 and its diffusion decoder, MiniMax-H3, FLUX.2 Klein, Ideogram4, ERNIE/GLM/SANA image models, FastWan2.1/2.2, the Blackwell-only Wan2.2 NVFP4 comparator, `LTX-2.3`, HunyuanVideo, MOVA, Helios, image edit, Hunyuan3D shape, and a separate Pi0.5 action-policy lane
|
||||
- [existing-fast-paths.md](existing-fast-paths.md) — map bottlenecks to existing fused kernels, packed QKV paths, fused `QK norm + RoPE`, distributed overlap patterns, and open optimization PRs before proposing new code
|
||||
- [existing-fast-paths.md](existing-fast-paths.md) — map bottlenecks to existing fused kernels, MoE routing, packed QKV paths, fused `QK norm + RoPE`, distributed overlap patterns, and open optimization PRs before proposing new code
|
||||
- [scripts/diffusion_skill_env.py](scripts/diffusion_skill_env.py) — preflight helper: repo root discovery from the skill's owning checkout before falling back to `sglang.__file__`, write-access probe, benchmark/profile output directories, idle GPU selection
|
||||
- [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark preset runner via `sglang generate`; defaults to eager/lossless, supports explicit quality and BCG comparators plus a same-GPU applicability matrix, rejects invalid BCG capture/fallback logs and late high-quality DiT fusion mounts, forces H3 to its eager consistency mode, enables synchronized stage attribution, validates nightly preset drift, and can clean one isolated model cache after the full matrix in a `finally` block with a JSONL ledger
|
||||
|
||||
@@ -74,6 +74,8 @@ Always rule out these existing families first:
|
||||
- bit-exact diffusion adaLN modulation and fused LayerNorm + modulation for
|
||||
FLUX.1, GLM-Image, and SANA
|
||||
- request-scoped DiT and VAE fast paths at `quality=extra-high` or `quality=high`
|
||||
- LingBot Video's default-on fused group-limited top-k expert selection before
|
||||
treating its router's `topk`/mask/gather chain as a new hotspot
|
||||
- Wan causal-VAE cache/padding and DupUp3D data-movement fusions
|
||||
- fused diffusion `QK norm + RoPE`
|
||||
- LTX2 split RoPE
|
||||
|
||||
+81
-58
@@ -20,13 +20,14 @@ framework-specific optimization workflow.
|
||||
- `python/sglang/kernels/ops/diffusion/norm/group_norm_silu_twopass_triton.py`
|
||||
- `python/sglang/kernels/ops/diffusion/norm/norm_triton.py`
|
||||
- `python/sglang/kernels/ops/diffusion/norm/rmsnorm_onepass_triton.py`
|
||||
- `python/sglang/kernels/ops/diffusion/norm/layernorm_modulate_triton.py`
|
||||
- `python/sglang/kernels/kda_kernels/layernorm_modulate_triton.py`
|
||||
- `python/sglang/kernels/ops/diffusion/norm/native_bf16_rmsnorm_triton.py`
|
||||
- `python/sglang/kernels/ops/diffusion/norm/zimage_qk_rmsnorm_triton.py`
|
||||
- `python/sglang/kernels/ops/diffusion/rope/rotary_triton.py`
|
||||
- `python/sglang/kernels/ops/diffusion/rope/helios_qk_rope_jit.py`
|
||||
- `python/sglang/kernels/ops/diffusion/rope/ltx2_rotary_triton.py`
|
||||
- `python/sglang/kernels/ops/diffusion/rope/ltx2_qknorm_split_rope_jit.py`
|
||||
- `python/sglang/kernels/kda_kernels/ltx2_qknorm_split_rope_jit.py`
|
||||
- `python/sglang/kernels/ops/diffusion/routing/group_limited_topk_triton.py`
|
||||
- `python/sglang/kernels/ops/diffusion/sites/ltx2_rmsnorm_modulate_site.py`
|
||||
- `python/sglang/kernels/ops/diffusion/modulate/indexed_modulation_triton.py`
|
||||
- `python/sglang/kernels/ops/diffusion/layout/ulysses_qkv_triton.py`
|
||||
@@ -39,8 +40,8 @@ framework-specific optimization workflow.
|
||||
- `python/sglang/multimodal_gen/runtime/models/decoders/ltx_2_5_diffusion_decoder.py`
|
||||
- `python/sglang/multimodal_gen/runtime/layers/moe.py`
|
||||
- `python/sglang/srt/layers/moe/topk.py`
|
||||
- `python/sglang/kernels/ops/diffusion/modulate/residual_gate_add_jit.py`
|
||||
- `python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh`
|
||||
- `python/sglang/kernels/kda_kernels/residual_gate_add_jit.py`
|
||||
- `python/sglang/kernels/kda_kernels/csrc/diffusion/residual_gate_add.cuh`
|
||||
- `python/sglang/kernels/ops/diffusion/layout/varlen_pack_pad_triton.py`
|
||||
- `python/sglang/kernels/ops/diffusion/layout/wan_causal_cache_triton.py`
|
||||
- `python/sglang/kernels/ops/diffusion/norm/scale_residual_norm_cutedsl.py`
|
||||
@@ -49,20 +50,13 @@ framework-specific optimization workflow.
|
||||
- `python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py`
|
||||
- `python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py`
|
||||
- `python/sglang/multimodal_gen/runtime/breakable_cuda_graph/runner.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_qwen_image_modulation.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_group_norm_silu.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_residual_gate_add.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_varlen_pack_pad.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_varlen_uspattn_equivalence.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_native_bf16_rmsnorm.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_flux_ln_modulate.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_glm_image_ln_modulate.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_sana_ln_modulate.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_quality_gate.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_ltx2_rms_norm_modulate.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_bitexact_gate.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_wan_causal_cache.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_stage_profiler_sync.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_modulate.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_norm.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_layout.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_routing.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_sites.py`
|
||||
- `test/registered/kernels/ops/diffusion/test_model_fast_paths.py`
|
||||
- `test/registered/profiling/test_diffusion_stage_profiler_sync.py`
|
||||
- `test/registered/kernels/benchmark/diffusion/bench_qwen_image_modulation.py`
|
||||
- `test/registered/kernels/benchmark/diffusion/bench_group_norm_silu.py`
|
||||
- `test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py`
|
||||
@@ -75,7 +69,7 @@ framework-specific optimization workflow.
|
||||
|
||||
1. Scale/Shift elementwise and gate fusion (AdaLN modulation)
|
||||
- Kernels: `fuse_scale_shift_kernel`, `fuse_layernorm_scale_shift_gate_select01_kernel`, `fuse_residual_layernorm_scale_shift_gate_select01_kernel`
|
||||
- Locations: `elementwise.py`, `layernorm.py`, `fused_scale_shift_gate.py`, `qwen_image.py`, `triton/scale_shift.py`
|
||||
- Locations: `elementwise.py`, `layernorm.py`, `fused_scale_shift_gate.py`, `qwen_image.py`, and `kernels/ops/diffusion/modulate/scale_shift_triton.py`
|
||||
- Use cases: `x * (1 + scale) + shift`, `a * (k + b) + c`, and Qwen-style `(layernorm/residual layernorm) + scale/shift + gate select`.
|
||||
- Constraints: `x` must be CUDA and contiguous. `scale/shift` support 0D/1D/2D/3D/4D broadcast. 4D `[B, F, 1, C]` requires `L % F == 0`.
|
||||
- Causal-video cold start: the 4D path uses a static capped power-of-two
|
||||
@@ -83,11 +77,11 @@ framework-specific optimization workflow.
|
||||
autotuning here: LingBot-World calls this path once per transformer block,
|
||||
and tuning overhead can dominate its first denoise step.
|
||||
- NPU fallback: `scale_shift.py` swaps to `npu_fallback` native path.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_qwen_image_modulation.py`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_modulate.py` and `test_model_fast_paths.py`.
|
||||
|
||||
2. Norm + Scale/Shift fusion (CuTe DSL)
|
||||
- Kernels: `fused_norm_scale_shift`, `fused_scale_residual_norm_scale_shift`
|
||||
- Locations: `layernorm.py`, `cutedsl/scale_residual_norm_scale_shift.py`
|
||||
- Locations: `layernorm.py` and `kernels/ops/diffusion/norm/scale_residual_norm_cutedsl.py`
|
||||
- Use cases:
|
||||
- `y = norm(x) * (1 + scale) + shift`
|
||||
- `y = norm(residual + gate * x) * (1 + scale) + shift`
|
||||
@@ -97,7 +91,7 @@ framework-specific optimization workflow.
|
||||
3. Bit-exact adaLN modulation and LayerNorm + modulation
|
||||
- Kernels: `modulate_scale_shift`, `fused_layernorm_modulate`, and
|
||||
`fused_qk_head_layernorm`.
|
||||
- Locations: `modulate_scale_shift.py`, `triton/layernorm_modulate.py`,
|
||||
- Locations: `kernels/ops/diffusion/modulate/modulate_scale_shift_jit.py`, `kernels/kda_kernels/layernorm_modulate_triton.py`,
|
||||
`runtime/models/dits/flux.py`, `glm_image.py`, and `sana.py`.
|
||||
- Use cases:
|
||||
- `x * (1 + scale[:, None]) + shift[:, None]` as one JIT CUDA launch.
|
||||
@@ -108,9 +102,8 @@ framework-specific optimization workflow.
|
||||
fp16/bf16 BLC inputs with `[B, D]` scale/shift. The Triton LayerNorm path is
|
||||
BF16-specific and only claims bit-exactness for its guarded aten dispatch;
|
||||
FLUX/GLM/SANA run a live eager equality check and fail closed on mismatch.
|
||||
- Validation: `test_flux_ln_modulate.py`, `test_glm_image_ln_modulate.py`,
|
||||
`test_sana_ln_modulate.py`, `test_modulate_scale_shift.py`, and
|
||||
`test_fused_ln_modulate.py`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_norm.py`,
|
||||
`test_modulate.py`, `test_sites.py`, and `test_model_fast_paths.py`.
|
||||
- Workflow rule: if these models show separate norm and modulation kernels,
|
||||
check dtype, alignment, shape, BCG/compile context, and the one-time equality
|
||||
self-test before proposing another fusion.
|
||||
@@ -135,40 +128,40 @@ framework-specific optimization workflow.
|
||||
the 2D conv runs channels_last end-to-end).
|
||||
- Do not confuse request `--quality` with `--output-quality`, which controls
|
||||
output-file compression rather than model math.
|
||||
- Validation: `test_quality_gate.py`, `test_fused_ln_modulate.py`,
|
||||
`test_flux2_vae_fastpath.py`, `test_wan_vae_fastpath.py`, and
|
||||
`test_vae_fast_path_gate.py`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_sites.py`,
|
||||
`test_norm.py`, and `test_model_fast_paths.py`.
|
||||
|
||||
5. Z-Image bf16-native RMSNorm modulation (Triton)
|
||||
- Kernels: `rmsnorm_scale`, `rmsnorm_tanh_residual`
|
||||
- Locations: `triton/native_bf16_rmsnorm.py`, with wrappers in `zimage.py` and
|
||||
`fused_gate_rmsnorm.py`. Note: `triton/zimage_native_norm.py` is QK-only.
|
||||
- Locations: `kernels/ops/diffusion/norm/native_bf16_rmsnorm_triton.py`, with
|
||||
callers in `runtime/models/dits/zimage.py` and
|
||||
`kernels/ops/diffusion/sites/fused_gate_rmsnorm_site.py`.
|
||||
- Use cases:
|
||||
- `y = rmsnorm(x) * scale`
|
||||
- `y = residual + tanh(gate) * rmsnorm(x)`
|
||||
- Constraints: CUDA bf16 tensors, contiguous weights, flattenable row strides,
|
||||
compatible modulation row counts, and `D <= 8192`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_native_bf16_rmsnorm.py`
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_norm.py`
|
||||
- Behavior: the kernels preserve Z-Image's native bf16 arithmetic. They return
|
||||
`None` when an eligibility guard fails, and the runtime wrapper executes the
|
||||
native PyTorch formula.
|
||||
|
||||
6. Triton LayerNorm/RMSNorm fusion
|
||||
- Kernels: `rms_norm_fn`, `layer_norm_fn`, `norm_infer`
|
||||
- Locations: `triton/norm.py`, `layernorm.py`
|
||||
- Locations: `kernels/ops/diffusion/norm/norm_triton.py`, `layernorm.py`
|
||||
- Use cases: fp32 RMSNorm with residual/dropout/rowscale/x1 branches, and inference-friendly `norm_infer`.
|
||||
- Constraints: last dim must be contiguous, and `N * element_size < 64KB`.
|
||||
- Validation: `test/registered/kernels/ops/layernorm/test_rmsnorm.py`.
|
||||
|
||||
7. Triton one-pass RMSNorm (small hidden size fast path)
|
||||
- Kernel: `triton_one_pass_rms_norm`
|
||||
- Locations: `triton/rmsnorm_onepass.py`, `layernorm.py`
|
||||
- Locations: `kernels/ops/diffusion/norm/rmsnorm_onepass_triton.py`, `layernorm.py`
|
||||
- Use case: `hidden_size <= 128` in `RMSNorm.forward_cuda`.
|
||||
- `torch.compile` note: keep this path behind the custom-op wrapper in `rmsnorm_onepass.py`; direct `wrap_triton` can recompile on dynamic row counts.
|
||||
|
||||
8. Triton RoPE fusion
|
||||
- Kernel: `apply_rotary_embedding`
|
||||
- Locations: `triton/rotary.py`, `rotary_embedding/utils.py`
|
||||
- Locations: `kernels/ops/diffusion/rope/rotary_triton.py`, `rotary_embedding/utils.py`
|
||||
- Use case: GPT-J style RoPE when not Neox.
|
||||
- Constraints: `head_size` must be even.
|
||||
- NPU fallback: `npu_fallback.apply_rotary_embedding_native`.
|
||||
@@ -176,24 +169,24 @@ framework-specific optimization workflow.
|
||||
|
||||
9. LTX2 split RoPE fusion
|
||||
- Kernel: `apply_ltx2_split_rotary_emb`
|
||||
- Locations: `triton/ltx2_rotary.py`, `runtime/models/dits/ltx_2.py`
|
||||
- Locations: `kernels/ops/diffusion/rope/ltx2_rotary_triton.py`, `runtime/models/dits/ltx_2.py`
|
||||
- Use case: LTX-2 split rotary embedding over `[B, S, num_heads * head_dim]` with separate `cos` and `sin` tensors.
|
||||
- Constraints: `cos` and `sin` shapes must match `[B, H, S, head_dim / 2]`, and `inner_dim == H * head_dim`.
|
||||
- Workflow rule: if LTX-2 traces show a large split-RoPE PyTorch chain, check whether the LTX2-specific Triton path was disabled by shape or dtype before proposing a new RoPE kernel.
|
||||
|
||||
10. Shared residual-gate add fusion (LTX2, LongCat-Image, SANA, and SANA-Video)
|
||||
- Kernel: `diffusion_residual_gate_add`
|
||||
- Locations: `kernels/ops/diffusion/modulate/residual_gate_add_jit.py`, `kernels/jit/csrc/diffusion/residual_gate_add.cuh`, `runtime/models/dits/ltx_2.py`, `runtime/models/dits/longcat_image.py`, `runtime/models/dits/sana.py`, and `runtime/models/dits/sana_video.py`.
|
||||
- Locations: `kernels/kda_kernels/residual_gate_add_jit.py`, `kernels/kda_kernels/csrc/diffusion/residual_gate_add.cuh`, `runtime/models/dits/ltx_2.py`, `runtime/models/dits/longcat_image.py`, `runtime/models/dits/sana.py`, and `runtime/models/dits/sana_video.py`.
|
||||
- Use case: `residual + update * gate` in LTX2 attention/MLP residuals, LongCat-Image joint- and single-stream transformer residuals, and SANA/SANA-Video transformer blocks.
|
||||
- Constraints: inputs must be same-device CUDA tensors with one dtype (`fp16`, `bf16`, or `fp32`) and `update.shape == residual.shape`. The ordinary path accepts contiguous inputs and a full or row-broadcast gate. The SANA-Video path also accepts a transposed-dense 3D residual (`stride == (tokens * hidden, 1, tokens)`), a contiguous update, and a contiguous `[1, 1, hidden]` gate; it preserves the residual stride in its output.
|
||||
- Behavior: model code calls `residual_gate_add(...)` directly. The CUDA custom op is used while guards pass. On a runtime exception outside `torch.compile`, it logs once, disables the fast path for that device/dtype, and falls back to `residual + update * gate`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_modulate.py`, `python/sglang/multimodal_gen/test/unit/test_longcat_image_residual_gate.py`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_modulate.py` and `test_model_fast_paths.py`.
|
||||
- Microbench: `test/registered/kernels/benchmark/diffusion/bench_residual_gate_add.py`.
|
||||
- Workflow rule: if LTX2, LongCat-Image, or SANA traces show repeated elementwise `mul` + `add` ladders around attention or MLP residuals, inspect input strides and check whether this existing CUDA path was disabled by shape, dtype, layout, or a prior runtime failure before proposing another elementwise fusion. For a transposed residual plus contiguous update, do not force `.contiguous()`; the tiled path is designed to fuse the mixed-layout access.
|
||||
|
||||
11. MiniMax-H3 indexed AdaLN modulation and gated residual fusion
|
||||
- Kernels: `indexed_scale_shift_bf16_`, `indexed_gate_bf16_`
|
||||
- Locations: `triton/indexed_modulation.py`, `runtime/models/dits/minimax_h3.py`
|
||||
- Locations: `kernels/ops/diffusion/modulate/indexed_modulation_triton.py`, `runtime/models/dits/minimax_h3.py`
|
||||
- Use cases: H3's packed video/audio/text rows select per-token modulation with `combined_indices`; the Triton paths replace `index_select` plus scale/shift or gated residual chains in place.
|
||||
- Constraints: CUDA BF16 H3 tensors, BF16 modulation tensors, contiguous disposable inputs; the gated path also requires contiguous `other`. Unsupported shapes/dtypes retain the eager formula.
|
||||
- Numerical contract: the kernels explicitly reproduce H3's eager BF16 rounding boundaries. Do not replace them with a mathematically equivalent contraction without the H3 consistency check.
|
||||
@@ -201,7 +194,7 @@ framework-specific optimization workflow.
|
||||
|
||||
12. MiniMax-H3 packed Ulysses QKV and output relayout
|
||||
- Kernels: `pack_qkv_destination_major`, `usp_merge_heads`
|
||||
- Locations: `triton/ulysses_qkv.py`, `usp_relayout.py`, `runtime/layers/usp.py`, `runtime/models/dits/minimax_h3.py`
|
||||
- Locations: `kernels/ops/diffusion/layout/ulysses_qkv_triton.py`, `kernels/ops/diffusion/layout/usp_relayout_jit.py`, `runtime/layers/usp.py`, and `runtime/models/dits/minimax_h3.py`
|
||||
- Use cases: one destination-major QKV pack plus one collective replaces three separately prepared Ulysses input exchanges; the output JIT kernel replaces `permute(...).contiguous()` when merging gathered heads.
|
||||
- Constraints: packed QKV fast packing requires CUDA fp16/bf16 Q/K/V with matching dtypes, contiguous head dimension, and eager execution. `usp_merge_heads` requires a nonempty contiguous 5D CUDA fp16/bf16/fp32 tensor and is disabled inside `torch.compile`.
|
||||
- Related transport: 2-rank, peer-accessible CUDA groups can use the existing IPC A2A transport; larger or unsupported groups fall back to the normal collective path.
|
||||
@@ -209,16 +202,16 @@ framework-specific optimization workflow.
|
||||
|
||||
13. HunyuanVideo / LTX upsampler GroupNorm + SiLU fusion
|
||||
- Kernel: `triton_group_norm_silu`
|
||||
- Locations: `diffusion/group_norm_silu.py`, `triton/group_norm_silu.py`, `runtime/models/vaes/hunyuanvae.py`, `runtime/models/upsampler/latent_upsampler.py`
|
||||
- Locations: `kernels/ops/diffusion/norm/group_norm_silu.py`, `kernels/ops/diffusion/norm/group_norm_silu_triton.py`, `runtime/models/vaes/hunyuanvae.py`, and `runtime/models/upsampler/latent_upsampler.py`
|
||||
- Use case: `activation(group_norm(x))` when the activation is non-inplace `nn.SiLU` and the GroupNorm is affine.
|
||||
- Enablement: mainline uses `apply_group_norm_silu(...)` in HunyuanVideo VAE paths and LTX latent upsampler paths by default; there is no env toggle. The wrapper dispatches to Triton only when guards pass.
|
||||
- Constraints: CUDA inference path only; no grad, `x.requires_grad == False`, `nn.GroupNorm`, `nn.SiLU(inplace=False)`, affine norm with weight and bias. Unsupported cases fall back to native `activation(norm(x))`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_group_norm_silu.py`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_norm.py` and `python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py`.
|
||||
- Microbench: `test/registered/kernels/benchmark/diffusion/bench_group_norm_silu.py`.
|
||||
|
||||
14. Wan causal-VAE data-movement fusion
|
||||
- Kernels: `cat_pad_channels_last_3d` and `dup_up3d_add`.
|
||||
- Locations: `triton/wan_causal_cache.py` and
|
||||
- Locations: `kernels/ops/diffusion/layout/wan_causal_cache_triton.py` and
|
||||
`runtime/models/vaes/wanvae.py`.
|
||||
- Use cases: build causal Conv3d input plus the next compact feature cache in
|
||||
one channels-last-3D pass, and fuse `main + DupUp3D(src)` without
|
||||
@@ -230,8 +223,8 @@ framework-specific optimization workflow.
|
||||
`cat_pad_channels_last_3d` + compact-cache helper for every causal conv
|
||||
slot; single-frame image decodes keep the compact cache at the reference
|
||||
size (one frame) so peak memory does not grow.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_wan_causal_cache.py`
|
||||
and the `test_qwen_vae_*` cases in `test_model_fast_paths.py`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_layout.py` and the
|
||||
`test_qwen_vae_*` cases in `test_model_fast_paths.py`.
|
||||
|
||||
15. Helios paired transposed RoPE
|
||||
- Kernel: `fused_inplace_helios_qk_rope`.
|
||||
@@ -257,6 +250,30 @@ framework-specific optimization workflow.
|
||||
ladders per block, check TP mode, dtype, shape, contiguity, and pointer
|
||||
alignment before proposing another RoPE kernel.
|
||||
|
||||
16. LingBot Video group-limited MoE routing
|
||||
- Kernel: `group_limited_topk`.
|
||||
- Locations: `kernels/ops/diffusion/routing/group_limited_topk_triton.py` and
|
||||
`runtime/layers/moe.py`.
|
||||
- Use case: fuse the group top-2 reduction, selected-group mask construction,
|
||||
masked expert scores, and final expert top-k into one Triton program per
|
||||
token. The production configuration is 128 experts, 4 groups, 2 selected
|
||||
groups, and 8 selected experts.
|
||||
- Constraints: NVIDIA CUDA, nonempty contiguous fp32 `[tokens, experts]`
|
||||
scores, at least two power-of-two experts per group, valid group counts, and
|
||||
`top_k` no larger than the selected-group capacity. Unsupported inputs keep
|
||||
the eager `torch.topk`/mask path.
|
||||
- Numerical contract: the selected expert-id set matches the guarded CUDA
|
||||
reference; order is unspecified because the reference uses
|
||||
`torch.topk(..., sorted=False)`. This path is selection-equivalent and
|
||||
default-on, independent of request `quality`. The launch is registered as a
|
||||
custom op with a fake implementation, so the guarded path also works under
|
||||
`torch.compile(fullgraph=True)`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_routing.py`.
|
||||
- Workflow rule: if a LingBot Video trace still shows two group-score `topk`
|
||||
calls plus scatter/mask/final-topk launches, check the score dtype,
|
||||
contiguity, group layout, and platform before proposing another router
|
||||
kernel.
|
||||
|
||||
**Faster CUDA Kernel Usage Points**
|
||||
|
||||
1. sgl-kernel RMSNorm and fused add RMSNorm
|
||||
@@ -277,9 +294,9 @@ framework-specific optimization workflow.
|
||||
- Behavior: `flashinfer.rope.apply_rope_with_cos_sin_cache_inplace` when available, otherwise Triton RoPE fallback.
|
||||
|
||||
4. Varlen USP attention pack/scatter
|
||||
- Locations: `runtime/layers/attention/layer.py`, `triton/varlen_pack_pad.py`
|
||||
- Locations: `runtime/layers/attention/layer.py`, `kernels/ops/diffusion/layout/varlen_pack_pad_triton.py`
|
||||
- Behavior: masked `USPAttention.forward` can gather dense Q/K/V into packed `[total_valid, H, D]` rows with `fused_pack_qkv`, run varlen attention, then scatter back with `fused_scatter_to_padded`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_varlen_pack_pad.py` and `test/registered/kernels/ops/diffusion/test_varlen_uspattn_equivalence.py`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_layout.py` and `test_model_fast_paths.py`.
|
||||
- Workflow rule: if a masked attention trace spends time in Python/advanced indexing pack or scatter, first check whether this fused varlen path should have engaged.
|
||||
|
||||
**QK Norm Optimization**
|
||||
@@ -307,7 +324,7 @@ framework-specific optimization workflow.
|
||||
- `can_use_fused_inplace_qknorm_rope(head_dim, rope_dim, is_neox, dtype)` returns true.
|
||||
- Supported head dims: `64, 128, 256`.
|
||||
- Behavior: `apply_qk_norm_rope` prefers the fused JIT kernel when all guards pass; otherwise it falls back to `apply_qk_norm(...)` plus `apply_flashinfer_rope_qk_inplace(...)`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_qknorm_rope.py`.
|
||||
- Validation: `test/registered/kernels/ops/diffusion/test_rope.py`.
|
||||
- MiniMax-H3: the H3 DiT calls `fused_inplace_qknorm_rope` directly for BF16
|
||||
head dim 128 with 96 rotary dims, NeoX layout, and
|
||||
`round_norm_before_rope=True`. This flag is part of H3's eager numerical
|
||||
@@ -372,12 +389,12 @@ framework-specific optimization workflow.
|
||||
one-time contiguous-layout helpers. Reuse or extract those helpers before
|
||||
authoring a video-only kernel.
|
||||
- LingBot Video MoE's router implements sigmoid+bias grouped top-k in
|
||||
`multimodal_gen/runtime/layers/moe.py`. Check parameter and output-order
|
||||
compatibility with `srt/layers/moe/topk.py::biased_grouped_topk` before
|
||||
writing a new router kernel. Current main mounts fused Triton RMSNorm row
|
||||
kernels by weight dtype and hidden size for `quality=extra-high` and
|
||||
`quality=high`; check the quality-site guards before treating an expanded
|
||||
`pow/mean/rsqrt` chain as a new opportunity.
|
||||
`multimodal_gen/runtime/layers/moe.py` and uses the default-on fused Triton
|
||||
selector when its CUDA layout guard passes. Check that guard before treating
|
||||
the `topk`/mask chain as a new opportunity. Its fused Triton RMSNorm row
|
||||
kernels remain request-gated by weight dtype and hidden size at
|
||||
`quality=extra-high` and `quality=high`; check those separate quality-site
|
||||
guards before treating an expanded `pow/mean/rsqrt` chain as new work.
|
||||
- LTX-2.5 reuses the mature LTX-2 DiT paths. Treat the optional diffusion
|
||||
decoder separately: confirm NATTEN `na3d` is active, then inspect its
|
||||
per-block 3D RoPE construction and split QKV/SwiGLU projections.
|
||||
@@ -397,14 +414,17 @@ framework-specific optimization workflow.
|
||||
`bitexact_gate.py`, used by FLUX / GLM / Sana / Ernie fused norm sites.
|
||||
- Qwen-Image gating: `fuse_layernorm_scale_shift_gate_select01_kernel` and `fuse_residual_layernorm_scale_shift_gate_select01_kernel` through `fused_scale_shift_gate.py` and `qwen_image.py`.
|
||||
- Z-Image native norm modulation: `rmsnorm_scale` and `rmsnorm_tanh_residual`
|
||||
in `triton/native_bf16_rmsnorm.py`, with wrappers in `zimage.py` /
|
||||
`fused_gate_rmsnorm.py`. `zimage_native_norm.py` is QK-only.
|
||||
in `kernels/ops/diffusion/norm/native_bf16_rmsnorm_triton.py`, with callers
|
||||
in `zimage.py` and `sites/fused_gate_rmsnorm_site.py`.
|
||||
- HunyuanVideo VAE and LTX upsampler GroupNorm+SiLU: `apply_group_norm_silu` in `hunyuanvae.py` and `latent_upsampler.py`; default-eligible when wrapper guards pass.
|
||||
- MiniMax-H3 indexed modulation: `_modulate_scale_shift` and `_modulate_gate` in `minimax_h3.py`, backed by `triton/indexed_modulation.py`.
|
||||
- MiniMax-H3 indexed modulation: `_modulate_scale_shift` and `_modulate_gate` in `minimax_h3.py`, backed by `kernels/ops/diffusion/modulate/indexed_modulation_triton.py`.
|
||||
- MiniMax-H3 Ulysses relayout: `_usp_input_all_to_all_packed_qkv` and `usp_merge_heads` through `runtime/layers/usp.py`.
|
||||
- QK norm: `apply_qk_norm` used in `flux.py`, `flux_2.py`, `qwen_image.py`, `zimage.py`, `wanvideo.py`, `ltx_2.py`, `hunyuanvideo.py`.
|
||||
- QK norm + RoPE: `apply_qk_norm_rope` in `layernorm.py`; use this path when the model wants fused attention prep instead of separate QK norm and RoPE calls.
|
||||
- LTX2 split RoPE: `apply_ltx2_split_rotary_emb` in `ltx_2.py`.
|
||||
- LingBot Video MoE routing: `LingBotVideoRouter._group_limited_topk` uses
|
||||
`group_limited_topk` through the diffusion facade when the CUDA layout guard
|
||||
passes, otherwise it retains the eager reference chain.
|
||||
- LTX2 RMSNorm+modulate and FFN GELU epilogue under `quality="extra-high"` and `quality="high"`:
|
||||
`mark_ltx2_rms_norm_modulate_site` / `fused_ltx2_rms_norm_modulate` in
|
||||
`kernels/ops/diffusion/sites/ltx2_rmsnorm_modulate_site.py` (mount-based
|
||||
@@ -412,11 +432,11 @@ framework-specific optimization workflow.
|
||||
kernel is <=1 ULP off aten, so it is request-gated instead of verified),
|
||||
wired at the six `LTX2TransformerBlock` adaLN sites in `ltx_2.py`.
|
||||
- Shared residual-gate add: `ltx_2.py`, `sana.py`, and `sana_video.py` call `residual_gate_add` from
|
||||
`kernels/ops/diffusion/modulate/residual_gate_add_jit.py` directly for attention,
|
||||
`kernels/kda_kernels/residual_gate_add_jit.py` through the diffusion facade for attention,
|
||||
cross-attention, and MLP residual updates; SANA-Video's transposed residual
|
||||
uses the mixed-layout tiled kernel without an intermediate contiguous copy.
|
||||
- Wan causal VAE: `cat_pad_channels_last_3d` and `dup_up3d_add` in
|
||||
`wanvae.py`, backed by `triton/wan_causal_cache.py`.
|
||||
`wanvae.py`, backed by `kernels/ops/diffusion/layout/wan_causal_cache_triton.py`.
|
||||
- Varlen USP attention: `fused_pack_qkv` and `fused_scatter_to_padded` in `attention/layer.py`.
|
||||
- SANA packed projections: `to_qkv` and `to_kv` in `sana.py`.
|
||||
- Nunchaku fused GELU MLP: `_fused_gelu_mlp` in `flux.py` for quantized FLUX-family checkpoints.
|
||||
@@ -518,6 +538,9 @@ relying on any file path, flag, or claim about whether the work has merged.
|
||||
4. Use `apply_qk_norm` and ensure head_dim is in the supported list for fused QK norm.
|
||||
5. If using FlashInfer RoPE, avoid `pack qkv` and ensure Q/K are contiguous.
|
||||
6. For attention, follow `selector.py` priority; override with CLI only if needed.
|
||||
7. For LingBot-style grouped routing, reuse `group_limited_topk` only when its
|
||||
predicate accepts the exact group layout; keep the eager selection chain as
|
||||
the fallback.
|
||||
|
||||
**When Extending or Modifying Kernels**
|
||||
- Add `torch.library.custom_op` and `register_fake` for compile and meta support.
|
||||
|
||||
@@ -20,6 +20,18 @@ class LingBotVideoMLP(nn.Module):
|
||||
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
try:
|
||||
from sglang.kernels.ops.diffusion import (
|
||||
can_use_group_limited_topk as _can_use_group_limited_topk,
|
||||
)
|
||||
from sglang.kernels.ops.diffusion import (
|
||||
group_limited_topk as _fused_group_limited_topk,
|
||||
)
|
||||
except Exception: # pragma: no cover - triton/kernel unavailable
|
||||
_can_use_group_limited_topk = None
|
||||
_fused_group_limited_topk = None
|
||||
|
||||
|
||||
class LingBotVideoRouter(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -46,6 +58,18 @@ class LingBotVideoRouter(nn.Module):
|
||||
)
|
||||
|
||||
def _group_limited_topk(self, scores_for_choice: torch.Tensor) -> torch.Tensor:
|
||||
if (
|
||||
_can_use_group_limited_topk is not None
|
||||
and _fused_group_limited_topk is not None
|
||||
and self.n_group is not None
|
||||
and self.topk_group is not None
|
||||
and _can_use_group_limited_topk(
|
||||
scores_for_choice, self.n_group, self.topk_group, self.top_k
|
||||
)
|
||||
):
|
||||
return _fused_group_limited_topk(
|
||||
scores_for_choice, self.n_group, self.topk_group, self.top_k
|
||||
)
|
||||
seq_len = scores_for_choice.shape[0]
|
||||
experts_per_group = self.num_experts // self.n_group
|
||||
grouped = scores_for_choice.view(seq_len, self.n_group, experts_per_group)
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
"""Unit tests for the fused group-limited top-k used by the LingBot MoE router."""
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.diffusion import (
|
||||
can_use_group_limited_topk,
|
||||
group_limited_topk,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.cuda.is_available() or torch.version.hip is not None,
|
||||
reason="requires NVIDIA CUDA",
|
||||
)
|
||||
|
||||
|
||||
def _ref_group_limited_topk(scores_for_choice, num_experts, n_group, topk_group, top_k):
|
||||
seq_len = scores_for_choice.shape[0]
|
||||
epg = num_experts // n_group
|
||||
grouped = scores_for_choice.view(seq_len, n_group, epg)
|
||||
group_scores = grouped.topk(2, dim=-1)[0].sum(dim=-1)
|
||||
group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1]
|
||||
group_mask = torch.zeros_like(group_scores)
|
||||
group_mask.scatter_(1, group_idx, 1)
|
||||
score_mask = (
|
||||
group_mask.unsqueeze(-1).expand(seq_len, n_group, epg).reshape(seq_len, -1)
|
||||
)
|
||||
masked = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf"))
|
||||
return torch.topk(masked, k=top_k, dim=-1, sorted=False)[1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_experts,n_group,topk_group,top_k",
|
||||
[
|
||||
(128, 4, 2, 8), # LingBot Video production configuration
|
||||
(64, 8, 3, 8),
|
||||
(128, 8, 3, 8),
|
||||
(256, 8, 3, 8),
|
||||
(64, 8, 3, 4),
|
||||
(32, 4, 2, 4),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("seq_len", [1, 17, 128, 1024])
|
||||
def test_group_limited_topk_matches_reference(
|
||||
num_experts, n_group, topk_group, top_k, seq_len
|
||||
):
|
||||
torch.manual_seed(0)
|
||||
scores = torch.randn(seq_len, num_experts, device="cuda", dtype=torch.float32)
|
||||
ref = _ref_group_limited_topk(scores, num_experts, n_group, topk_group, top_k)
|
||||
fused = group_limited_topk(scores, n_group, topk_group, top_k)
|
||||
# Selection order is not guaranteed (sorted=False); compare per-row id sets.
|
||||
assert torch.equal(ref.sort(dim=-1)[0], fused.sort(dim=-1)[0])
|
||||
|
||||
|
||||
def test_group_limited_topk_handles_ties():
|
||||
# Uniform scores => every group/expert ties. Selection must stay in-range and
|
||||
# match the reference (both use first-max tie-breaking).
|
||||
seq_len, num_experts, n_group, topk_group, top_k = 8, 64, 8, 3, 8
|
||||
scores = torch.ones(seq_len, num_experts, device="cuda", dtype=torch.float32)
|
||||
ref = _ref_group_limited_topk(scores, num_experts, n_group, topk_group, top_k)
|
||||
fused = group_limited_topk(scores, n_group, topk_group, top_k)
|
||||
assert fused.shape == (seq_len, top_k)
|
||||
assert (fused >= 0).all() and (fused < num_experts).all()
|
||||
assert torch.equal(ref.sort(dim=-1)[0], fused.sort(dim=-1)[0])
|
||||
|
||||
|
||||
def test_group_limited_topk_handles_repeated_nonuniform_scores():
|
||||
# Router logits can contain repeated, nonuniform values after rounding or
|
||||
# saturation. This exercises duplicate maxima both within and across groups.
|
||||
torch.manual_seed(123)
|
||||
seq_len, num_experts, n_group, topk_group, top_k = 64, 64, 8, 3, 8
|
||||
scores = torch.randint(
|
||||
0,
|
||||
4,
|
||||
(seq_len, num_experts),
|
||||
device="cuda",
|
||||
dtype=torch.int32,
|
||||
).float()
|
||||
ref = _ref_group_limited_topk(scores, num_experts, n_group, topk_group, top_k)
|
||||
fused = group_limited_topk(scores, n_group, topk_group, top_k)
|
||||
assert torch.equal(ref.sort(dim=-1)[0], fused.sort(dim=-1)[0])
|
||||
|
||||
|
||||
def test_group_limited_topk_rejects_unsupported_group_layout():
|
||||
scores = torch.randn(1, 30, device="cuda", dtype=torch.float32)
|
||||
assert not can_use_group_limited_topk(scores, 3, 2, 4)
|
||||
with pytest.raises(ValueError, match="power-of-two experts per group"):
|
||||
group_limited_topk(scores, 3, 2, 4)
|
||||
|
||||
|
||||
def test_group_limited_topk_torch_compile_fullgraph():
|
||||
scores = torch.randn(17, 128, device="cuda", dtype=torch.float32)
|
||||
ref = _ref_group_limited_topk(scores, 128, 4, 2, 8)
|
||||
compiled = torch.compile(lambda x: group_limited_topk(x, 4, 2, 8), fullgraph=True)
|
||||
fused = compiled(scores)
|
||||
assert torch.equal(ref.sort(dim=-1)[0], fused.sort(dim=-1)[0])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
Reference in New Issue
Block a user