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

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md

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

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md

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

python/sglang/multimodal_gen/runtime/layers/moe.py

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