[Perf] [NVIDIA] qwen3.5: Add follow-up megakernel kernels

.github/workflows/unit-test.yml

@@ -14,10 +14,11 @@ concurrency:
   cancel-in-progress: true
 
 env:
-  # Includes the hybrid preparation/views API from sgl-project/sglang#45406.
+  # Includes the hybrid preparation/views API from sgl-project/sglang#45406 and the
+  # speculative runtime hooks from sgl-project/sglang#47201.
   # Pin source until a published nightly contains this interface; updating the
   # pin is a normal PR and must pass the static mirror/contract tests.
-  SGLANG_COMMIT: aa3e34aaebf880d72482cea74dbbfef54da728e2
+  SGLANG_COMMIT: 8df0cd51e1afbcf2a67358ab559f4a9046ff5782
 
 jobs:
   unit:

README.md

@@ -18,8 +18,8 @@ The package loads as an SGLang [general plugin](https://github.com/sgl-project/s
 | Qwen3.5-397B-A17B (NVFP4) | GB300 (SM103) | 8 (two trays, one NVLink domain) | MTP target verify steps of 7 rows (batch 1) |
 | Kimi K3 (packed MXFP4) | GB300 (SM103) | 8 | One-token decode steps of 1–8, 16, 32 and 64 rows |
 
-For DeepSeek-V4.1-Flash and Qwen3.5, larger verify steps, prefill and extend run
-SGLang's standard path. Their draft specs provide the DSpark proposal and the
+For DeepSeek-V4.1-Flash, larger verify steps, prefill and extend run
+SGLang's standard path. Qwen3.5 also offers the explicit GDN prefill slot below. Their draft specs provide the DSpark proposal and the
 Qwen3.5 EAGLE entries, respectively. For Kimi K3, the plugin covers the 93-layer
 language model (24 MLA and 69 KDA layers) at the listed widths; every other call
 runs the standard path.
@@ -50,25 +50,61 @@ The step is the plugin's own code over the weights SGLang loaded (`deepseek_v41/
 
 With the megakernel on, the default 39 slots run: the ones above, the attention chain, hop B's triggers, and the DSpark draft step's 16 slots. `SGLANG_NVIDIA_MEGAKERNEL_SLOTS` (comma-separated) overrides that list for tests and bisecting; empty runs every step function's stock code. The opt-in slots run only when listed there: `target_lm_head` (the target's LM head on the low-M kernel at 1-8 rows), `lm_head_pinned_tactic` (both low-M LM heads on one fixed tactic per row count, untuned), `hop_b_quad_trigger` (hop B's finalize-publish launches its consumer right after its own wait, batch 1 only), `hop_b_cap16` and `draft_embed`; listing `sm148` there admits the 148 SMs of a B300 instead of the 152 required otherwise.
 
-Qwen3.5 (`qwen3_5/runtime`) runs the same way: per layer the stage boundary (`boundary/`), the GatedDeltaNet or the attention (`gdn/`, `attention/`) and the MoE (`moe/`), each with a stock implementation that reproduces SGLang's verify step bit for bit. It reads the running forward's metadata and pools through SGLang's EAGLE binding (`sglang.srt.kernel_plugins.nvidia.eagle_binding`). Its slots run only when listed in `SGLANG_NVIDIA_MEGAKERNEL_SLOTS`, and each is bit for bit SGLang's:
+Qwen3.5 (`qwen3_5/runtime`) runs the same way: per layer the stage boundary (`boundary/`), the GatedDeltaNet or the attention (`gdn/`, `attention/`) and the MoE (`moe/`), each with a stock implementation that reproduces SGLang's verify step bit for bit. It reads the running forward's metadata and pools through SGLang's EAGLE binding (`sglang.srt.kernel_plugins.nvidia.eagle_binding`). With the megakernel enabled, its shape-bounded bitwise paths run by default. `SGLANG_NVIDIA_MEGAKERNEL_SLOTS` overrides that policy for testing and bisecting; an empty value runs the stock paths.
 
 | Slot | Rows | Kernel |
 | --- | --- | --- |
-| `route` | 7 | The router GEMM and the NVFP4 quantization of the MoE input in one launch |
-| `mkm` (with `route`) | 7 | On `route`'s outputs, the top-10, FC1, SwiGLU, the NVFP4 requantization and FC2 in one launch |
-| `gdn_verify` | 7 | The GatedDeltaNet verify chain (conv update, split, ReplaySSM verify, gated norm with the FP8 `out_proj` input) with its launch choices changed |
-| `gdn_pdl` (with `gdn_verify`) | 7 | Its ReplaySSM verify launched as a programmatic dependent of the conv |
-| `hop` | 7 | The layer boundary's fused all-reduce + residual add + RMSNorm points on an NVLink multicast hop |
+| `route_topk`, `mkm_m32` … `mkm_m512` | 1–896 | Materialized stock routing and the shape-selected fused wide routed-expert engines |
+| `hop`, `hop_fp8` | 1–15 | The layer boundary's fused all-reduce + residual add + RMSNorm, also writing the next static-FP8 input |
+| `multi_request` | 2–128 requests | The target graphs keep every offered request count instead of padding to a coarser width |
+| `draft_loop`, `draft_extend`, `draft_multi`, `draft_embed` | 1–128 requests | EAGLE draft entries, exact sharded top-1 and embedding exchange; the pre-FC norms write directly into the FC input |
 
-These slots change the arithmetic (numeric: each sums in another order than SGLang's op, within the tolerances `spec.py` declares):
+The remaining kernel families are opt-in through that one slot policy. These change the arithmetic within the tolerances `spec.py` declares:
 
 | Slot | Rows | Kernel |
 | --- | --- | --- |
+| `route`, `route_gemm`, `mkm`, `mkm_wide` | 1–32 | Fused router fronts and the small routed-expert engines |
+| `gdn_verify`, `gdn_pdl`, `gdn_conv` | 7–896 | Explicit numeric verify chain and dependent variants; FP8 norm uses the v2 relative-L2 bound of 2e-2 |
+| `gdn_core` | 7–224 | ReplaySSM verify, gated norm and FP8 output in one launch |
+| `hop_mid` (with `hop`) | 16–28 | Extends the one-shot boundary mailbox through four requests instead of returning those widths to the stock boundary |
 | `gdn_push` (with `gdn_verify`, `hop`) | 7 | The GatedDeltaNet `out_proj` as the hop's GEMM, whose epilogue pushes the partial into the hop's mailbox |
-| `gdn_front` (with `gdn_verify`, `gdn_pdl`) | 7 (one request) | The GatedDeltaNet `in_proj_qkvz` (static FP8) and `in_proj_ba` (BF16) projections, the conv update and an L2 prefetch of the verify's state in one launch, then the ReplaySSM verify and the gated norm writing `out_proj`'s FP8 input, each launched dependent on the one before |
+| `gdn_front`, `gdn_trt_front` | 7–28 | Fused GatedDeltaNet projections and conv before the dependent verify |
 | `route_shared` (with `route`) | 1–8 | The shared expert's static-FP8 gate/up in the router front's launch, so the shared expert on the alternate stream waits for that launch only |
-
-The sampling and acceptance after the verify, and the GatedDeltaNet extend of prefill, stay SGLang's own code. In an SGLang that has the switches, `SGLANG_ENABLE_EAGLE_POST_VERIFY_CUDA_GRAPH=1` runs the sampling and acceptance in CUDA graphs, and `SGLANG_ENABLE_GDN_PREFILL_GRAPH_EXTEND=1` captures the GatedDeltaNet extend inside the breakable prefill CUDA graphs.
+| `route_lin`, `skv2`, `mka`, `mka_push` | shape-bounded | Router/attention overlap, split-KV attention and fused attention boundaries |
+| `lowm_in`, `lowm_out`, `lowm_se`, `lowm_push` | shape-bounded | Validated low-row projection and shared-expert launches |
+| `draft_moe`, `draft_q1`, `draft_xchg`, `draft_hop`, `draft_fin`, `draft_front`, `draft_skv2` | shape-bounded | Fused MTP draft kernels and boundaries |
+
+The tuned Qwen3.5 TP8/TEP8 serving policy explicitly includes `hop_mid` whenever it
+includes `hop`.  It remains opt-in rather than part of the bitwise default policy;
+an explicit `SGLANG_NVIDIA_MEGAKERNEL_SLOTS` value replaces, rather than extends,
+that default list.
+
+The TP8 plan (`topology.plan`, aggregated and PD decode) also runs `route_lin`
+whenever the policy holds the slots it runs through (`route`, `mkm`, `hop`,
+`gdn_push`, `mka`, `mka_push`, and `lowm_se` with `route_shared`): its router logits
+stay within 1.05x stock's RMS error against FP64 on the same quantized operands.
+The same TP8 plan adds `mkm_wide_bal` for `mkm_m256` / `mkm_m512`: balanced
+CTA assignment at 129–512 verify rows, with bitwise-identical outputs.
+TEP8, DEP8, TP2 and TP4 run exactly the requested slots their topology supports.
+
+TEP8's routed body takes its NVFP4 MoE tactic per token bucket from a fixed table
+(`moe/tactics.py`) on FlashInfer 0.7.1, SM103 and the routed experts' MoERunner
+signature; any other configuration or bucket asks the autotuner, and the log names
+each width's source. `SGLANG_NVIDIA_MEGAKERNEL_MOE_TACTICS` switches the table; unset,
+it follows `SGLANG_ENABLE_NV_MEGA_RUNTIME_OPTIMIZATION`.
+
+The router-only `route`, `route_topk` and `route_gemm` slots support TP2, TP4, TP8 and TEP8 with the TRTLLM runner and no all-to-all. PD decode TP2 uses the same local row counts: batch 6/10 gives 42/70 verify rows. These fronts leave the expert GEMMs with the stock runner, so TP2 needs no fused expert engine. `route_shared` keeps its TP8/TEP8 projection bounds; DEP8 keeps stock MoE dispatch and routing.
+
+`mka` uses the fused projection kernels through 14 verify rows on TP8 and 8 on
+TP4. At wider rows, and at every TP2 width, it keeps SGLang's projection GEMMs
+and fuses Q/K normalization, RoPE and the FP8 KV write, then the output gate
+and static FP8 quantization. TP2 has 16 query heads and one KV head per rank;
+TP4 has 8/1 and TP8 has 4/1. The same front path serves PD decode TP2 and DEP8 (32/2 per rank).
+`multi_request` coverage is required for more than one request. `skv2` keeps
+its one-request, one-KV-head scope; `skv2_cp` additionally requires replicated
+KV heads across ranks.
+
+The sampling and acceptance after the verify stay SGLang's own code. `gdn_prefill_split_conv` is an explicit bitwise slot for ordinary text prefill on TP1/PP4 with the data-only prefill binding: it splits the native conv by Q/K/V channels above 8192 channels, retaining native recurrence, norm and state writes. Unsupported layouts use SGLang; an empty slot policy calls the original prefill forward directly. In an SGLang that has the switches, `SGLANG_ENABLE_EAGLE_POST_VERIFY_CUDA_GRAPH=1` runs the sampling and acceptance in CUDA graphs, and `SGLANG_ENABLE_GDN_PREFILL_GRAPH_EXTEND=1` captures the GatedDeltaNet extend inside the breakable prefill CUDA graphs.
 
 ### Design
 
@@ -124,7 +160,6 @@ The three environment variables are required by the DeepSeek-V4.1 spec; without
 Qwen3.5, on each of the two trays (`--node-rank 0` and `1`):
 
 ```bash
-SGLANG_NVIDIA_MEGAKERNEL_SLOTS=route,mkm,gdn_verify,gdn_pdl,hop \
 python -m sglang.launch_server \
   --model-path nvidia/Qwen3.5-397B-A17B-NVFP4-V2 \
   --tp 8 --nnodes 2 --node-rank 0 --dist-init-addr <tray-0 host>:<port> \

sgl_nvidia_mega/deepseek_v41/runtime/executor.py

@@ -113,6 +113,7 @@ class MegakernelDecode:
     draft_extend = None
     begin_draft_round = None
     graph_variants = None
+    speculative_runtime_hooks = None
 
     def __init__(
         self,

sgl_nvidia_mega/kimi_k3/spec.py

@@ -169,6 +169,7 @@ class _Built:
     draft_extend = None
     begin_draft_round = None
     graph_variants = None
+    speculative_runtime_hooks = None
 
     def __init__(self, model: Any, resources: Any, widths: tuple[int, ...]):
         from .executor import KimiK3MegakernelExecutor

sgl_nvidia_mega/qwen3_5/__init__.py

@@ -29,3 +29,26 @@
     hidden_act="silu",
     rms_norm_eps=1e-6,
 )
+
+CP_LITE_HEAD_DIM = 256
+
+
+def cp_lite_topology(
+    attention_tp_size: int,
+    *,
+    num_key_value_heads: int = TEXT_CONFIG["num_key_value_heads"],
+    num_attention_heads: int = TEXT_CONFIG["num_attention_heads"],
+    head_dim: int = TEXT_CONFIG["head_dim"],
+    kernel_head_dim: int = CP_LITE_HEAD_DIM,
+) -> tuple[int, int] | None:
+    """Return ``(CP size, local Q heads)`` when CP-lite supports the topology."""
+    tp = attention_tp_size
+    if tp <= num_key_value_heads or tp % num_key_value_heads:
+        return None
+    if head_dim != kernel_head_dim:
+        return None
+    cp = tp // num_key_value_heads
+    q_heads = num_attention_heads // tp
+    if cp not in (2, 4, 8) or q_heads * cp > 16:
+        return None
+    return cp, q_heads

sgl_nvidia_mega/qwen3_5/attention/__init__.py

@@ -47,15 +47,17 @@
     "srt/models/qwen3_5.py:Qwen3_5AttentionDecoderLayer.self_attention": "db0f4d98e0fb1abf",
 }
 
-ATTENTION_SLOTS = ("skv2", "mka", "mka_push")
+ATTENTION_SLOTS = ("skv2", "skv2_cp", "mka", "mka_push", "route_lin")
 
 
 class Bound(msgspec.Struct, frozen=True):
     """The attention package's build result (``st.bound["attention"]``)."""
 
     skv2: bool = False  # slot skv2 serves every full-attention layer
     mka: bool = False
+    mka_wide: bool = False  # mka's front path also serves steps of several requests
     mka_push: bool = False
+    route_lin: bool = False  # router-before-all-reduce packed for every layer
     reason: str = ""  # why a requested slot is not built
 
 
@@ -75,22 +77,42 @@ def build(model, ctx) -> Bound:
         for i in range(m.start_layer, m.end_layer)
         if isinstance(m.layers[i], Qwen3_5AttentionDecoderLayer)
     ]
+    reasons = []
+    skv2_on, mka_on = "skv2" in SLOTS, "mka" in SLOTS
     for i, layer in layers:
-        reason = ""
-        if "skv2" in SLOTS:
+        if skv2_on:
             reason = split_kv.unsupported(layer.attn, ctx.tokens_per_request)
-        if not reason and "mka" in SLOTS:
+            if reason:
+                reasons.append(f"layer {i}: {reason}")
+                skv2_on = False
+        if mka_on:
             reason = (
                 mka.unsupported(layer)
                 if model.is_mrope_enabled
                 else "no M-RoPE positions (MKA reads them)"
             )
+            if reason:
+                reasons.append(f"layer {i}: {reason}")
+                mka_on = False
+    if "skv2_cp" in SLOTS:
+        reason = split_kv.cp_reason()
         if reason:
-            return Bound(reason=f"layer {i}: {reason}")
-    if "mka" in SLOTS:
-        for _, layer in layers:
-            mka.pack(layer)
-    return Bound(skv2="skv2" in SLOTS, mka="mka" in SLOTS, mka_push="mka_push" in SLOTS)
+            reasons.append(f"slot skv2_cp: {reason}")
+    route_lin = mka_on and "route_lin" in SLOTS
+    if mka_on:
+        for i, layer in layers:
+            reason = mka.pack(layer, route_lin=route_lin)
+            if reason:
+                reasons.append(f"slot route_lin: layer {i}: {reason}")
+                route_lin = False
+    return Bound(
+        skv2=skv2_on,
+        mka=mka_on,
+        mka_wide=mka_on,
+        mka_push=mka_on and "mka_push" in SLOTS,
+        route_lin=route_lin,
+        reason="; ".join(reasons),
+    )
 
 
 def pool_reasons(view, binding, skv2: bool = True, mka: bool = False) -> list[str]:
@@ -111,16 +133,34 @@ def step(layer, x, st):
         from sgl_nvidia_mega.qwen3_5.attention import mka
 
         return mka.step(layer, x, st)
-    q, k, v, gate = layer._prepare_qkv_gate(
-        positions=st.positions, hidden_states=x, forward_batch=st.fb
-    )
-    if st.slots.skv2:
-        from sgl_nvidia_mega.qwen3_5.attention import split_kv
+    from sgl_nvidia_mega.qwen3_5 import lowm
+
+    x8 = lowm.take_fp8(st, x) if st.slots.hop_fp8 else None
+    if st.slots.lowm_in or x8 is not None:
+        q, k, v, gate = lowm.prepare_qkv_gate(
+            layer, st.positions, x, x8, low_m=st.slots.lowm_in
+        )
+    else:
+        q, k, v, gate = layer._prepare_qkv_gate(
+            positions=st.positions, hidden_states=x, forward_batch=st.fb
+        )
+    from sgl_nvidia_mega.qwen3_5.attention import split_kv
 
+    if st.slots.skv2 or split_kv.cp_multi(q.shape[0]):
         o = split_kv.attend(layer.attn, q, k, v, st)
     else:
         o = layer.attn(q, k, v, st.fb)
     if layer.attn_output_gate:
         o = fused_sigmoid_mul(o, gate, inplace=True)
+    if st.slots.lowm_oproj:
+        from sgl_nvidia_mega.qwen3_5 import lowm
+        from sgl_nvidia_mega.qwen3_5.boundary import hop as hop_step
+
+        hop = None
+        if st.slots.lowm_push and hop_step.takes_moe_input(
+            st.bound["boundary"].hop, layer, st, o.shape[0]
+        ):
+            hop = st.bound["boundary"].hop
+        return lowm.row_parallel(layer.o_proj, o, o.dtype, "o_proj", hop)
     out, _ = layer.o_proj(o)
     return out