[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>

python/sglang/kernels/ops/diffusion/README.md

@@ -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

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

@@ -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",

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

@@ -0,0 +1 @@
+"""Diffusion-specific routing kernels."""

python/sglang/kernels/ops/diffusion/routing/group_limited_topk_triton.py

@@ -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)