perf(k3): block-parallel router top-k replaces the serial TileLang scan

pegainfer-k3/kernels/generate.py

@@ -6,10 +6,9 @@
 
     k3_rms_norm_rbs_batched.cu      k3_conv_silu_batched.cu
     k3_land_batched.cu              k3_kda_core_batched.cu
-    k3_land_rms_norm_rbs_batched.cu k3_router_topk_batched.cu
-    k3_add2_batched.cu              k3_attnres_scores_batched.cu
-    k3_mul_sigmoid_batched.cu       k3_attnres_mix_batched.cu
-    k3_situ_batched.cu
+    k3_land_rms_norm_rbs_batched.cu k3_attnres_scores_batched.cu
+    k3_add2_batched.cu              k3_attnres_mix_batched.cu
+    k3_mul_sigmoid_batched.cu       k3_situ_batched.cu
 
 The batch size is a static compile-time dimension, so a single-stream step is
 served by the `B = 1` instantiation of the same family — its per-row spelling
@@ -230,11 +229,6 @@
     f"(const {BF16}* __restrict__ G2, const float* __restrict__ Go, "
     f"{BF16}* __restrict__ Out, const {BF16}* __restrict__ X)"
 )
-ROUTER_PARAMS = (
-    "(const float* __restrict__ Bias, int* __restrict__ Idx, "
-    f"const {BF16}* __restrict__ Rs, const float* __restrict__ S, "
-    "float* __restrict__ Wts)"
-)
 SCORES_PARAMS = (
     f"(const {BF16}* __restrict__ Bl, const {BF16}* __restrict__ Ps, "
     "float* __restrict__ Sc, const float* __restrict__ Sw)"
@@ -923,49 +917,6 @@ def plan_o_norm_gate() -> Plan:
     )
 
 
-def plan_router_topk() -> Plan:
-    insts = []
-    for experts in EXPERTS:
-        for batch in B_CHUNK_BUCKETS:
-            insts.append(Inst(
-                family="router_topk",
-                order=len(insts),
-                label=f"router_topk_batched E={experts} B={batch}",
-                factory="router_topk_batched",
-                args=(experts, TOPK, batch, THREADS),
-                num_params=5,
-                params=ROUTER_PARAMS,
-                symbol=f"k3_router_topk_b{batch}_e{experts}_topk{TOPK}_kernel",
-                grid=(batch,),
-                threads=THREADS,
-                guard=f"b == {batch} && num_experts == {experts} && topk == {TOPK}",
-                call_args=("Bias", "Idx", _bf16("Rs"), "S", "Wts"),
-            ))
-    return Plan(
-        stem=_STEM.format("router_topk"),
-        signature=(
-            "k3_router_topk_batched(\n"
-            "    const float* S,\n"
-            "    const float* Bias,\n"
-            "    const void* Rs,\n"
-            "    int* Idx,\n"
-            "    float* Wts,\n"
-            "    int b,\n"
-            "    int num_experts,\n"
-            "    int topk,\n"
-            "    cudaStream_t stream)"
-        ),
-        doc=(
-            "// Sigmoid router plus biased top-k over already-merged f32 score rows,\n"
-            "// one row per block. The selection is a serial O(topk * num_experts)\n"
-            "// scan by thread 0 with a lowest-index tie-break; the weights are\n"
-            "// gathered from the un-biased scores, the denominator carries the\n"
-            "// +1e-20 guard, and the result is scaled by the bf16 routed scale Rs."
-        ),
-        insts=tuple(insts),
-    )
-
-
 def plan_attnres_scores() -> Plan:
     insts = []
     for blocks in ATTNRES_NB:
@@ -1057,7 +1008,6 @@ def plan_attnres_mix() -> Plan:
     plan_conv_silu,
     plan_kda_core,
     plan_o_norm_gate,
-    plan_router_topk,
     plan_attnres_scores,
     plan_attnres_mix,
 ]

pegainfer-k3/kernels/tilelang_defs.py

@@ -1,7 +1,7 @@
 """Vendored TileLang kernel definitions for the K3 batched decode step.
 
 This file is a **verbatim** subset of the certified upstream kernel module: the
-shared prologue and the eleven batched kernel factories, copied character for
+shared prologue and the ten batched kernel factories, copied character for
 character. Nothing here is re-spelled, re-indented or "cleaned up", and no
 kernel body is edited to fit this repository. The upstream module is the
 authority on what these kernels compute; it carries the bitwise parity gates
@@ -65,58 +65,6 @@ def _compile(prim):
 # --------------------------------------------------------------------------- #
 
 
-@lru_cache(maxsize=None)
-def router_topk_batched(E: int, TOPK: int, B: int, threads: int = 256):
-    """Batched version of ``router_topk``. The input S (B, E) holds **f32
-    score rows** (output of the framework-side f32 GEMM -- the authored
-    spelling already casts to f32 before the matmul; the ascending f32 merge
-    over the SK segments in the bs=1 version is equivalent to the row being
-    merged upstream). One row per block; sigmoid, bias add, serial top-k
-    selection, un-biased gather, and normalization times routed_scale are
-    word-for-word identical to the bs=1 version (lowest-index tie-break, Bias
-    taken in f32)."""
-    EP = ((E + threads - 1) // threads) * threads
-
-    @T.prim_func
-    def main(
-        S: T.Tensor((B, E), ACC),
-        Bias: T.Tensor((E,), ACC),
-        Rs: T.Tensor((1,), DT),
-        Idx: T.Tensor((B, TOPK), "int32"),
-        Wts: T.Tensor((B, TOPK), ACC),
-    ):
-        with T.Kernel(B, threads=threads) as bb:
-            scores = T.alloc_shared((E,), ACC)
-            biased = T.alloc_shared((E,), ACC)
-            best = T.alloc_var(ACC)
-            bi = T.alloc_var("int32")
-            den = T.alloc_var(ACC)
-            for e in T.Parallel(EP):
-                with T.If(e < E):
-                    with T.Then():
-                        scores[e] = T.sigmoid(S[bb, e])
-                        biased[e] = T.sigmoid(S[bb, e]) + Bias[e].astype(ACC)
-            T.sync_threads()
-            if T.get_thread_binding() == 0:
-                den = 0.0
-                for t in T.serial(TOPK):
-                    best = NEG
-                    bi = 0
-                    for e in T.serial(E):
-                        with T.If(biased[e] > best):
-                            with T.Then():
-                                best = biased[e]
-                                bi = e
-                    Idx[bb, t] = bi
-                    Wts[bb, t] = scores[bi]
-                    biased[bi] = NEG
-                    den += scores[bi]
-                for t in T.serial(TOPK):
-                    Wts[bb, t] = Wts[bb, t] / (den + 1e-20) * Rs[0].astype(ACC)
-
-    return _compile(main)
-
-
 @lru_cache(maxsize=None)
 def attnres_scores_batched(NB: int, BC: int, H: int, B: int, eps: float,
                            threads: int = 256):

pegainfer-kernels/build.rs

@@ -1535,10 +1535,6 @@ const K3_TILELANG_LAUNCHERS: &[(&str, &str)] = &[
         "k3_o_norm_gate_batched",
         "const void*, const void*, const float*, void*, int, int, int",
     ),
-    (
-        "k3_router_topk_batched",
-        "const float*, const float*, const void*, int*, float*, int, int, int",
-    ),
     (
         "k3_attnres_scores_batched",
         "const void*, const void*, const float*, float*, int, int, int",

pegainfer-kernels/csrc/k3/k3_router_topk.cu

@@ -0,0 +1,151 @@
+// Kimi-K3 MoE router: sigmoid scores plus biased top-k selection.
+//
+// Replaces the retired TileLang `router_topk_batched` kernel, which ran the
+// whole top-k as a serial scan on thread 0 (TOPK x E iterations per row —
+// 65.7us per launch at E=896, ~6ms of every EP16 verify round). The selection
+// here is a block-parallel argmax per round, but every arithmetic step is the
+// retired kernel's spelling so the outputs are bit-identical:
+//
+//   scores[e] = 1 / (1 + expf(0 - s[e]))          (plain expf, f32 division)
+//   biased[e] = scores[e] + bias[e]               (bias read in f32)
+//   TOPK rounds of argmax over `biased` with strict `<` comparison — the
+//     first index attaining the maximum wins, i.e. ties break to the lowest
+//     expert index; the winner's biased score is set to -1e30 and its
+//     *un-biased* score joins the weight row and the denominator, in
+//     selection order;
+//   wts[t] = (wts[t] / (den + 1e-20)) * (float)rs  (division, then scale)
+//
+// The parallel argmax preserves the serial tie-break exactly: each thread
+// scans its stride-256 subsequence ascending with strict `<` (keeping its
+// local first maximum), and the shuffle/shared reduction prefers the
+// strictly-greater value, breaking equal values to the lower index — the
+// result is the minimum index among the global maxima, which is precisely
+// the index the serial first-match scan selects. All comparisons are on
+// identical f32 values, so no summation-order freedom exists anywhere.
+//
+// Deterministic and CUDA-graph safe: fixed reduction order, no allocation,
+// no host readback; grid is (b) with everything else read from device
+// tensors.
+
+#include "../common.cuh"
+#include "../shared/ffi_guard.cuh"
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+
+namespace {
+
+constexpr int kThreads = 256;
+constexpr int kWarps = kThreads / WARP_SIZE;
+constexpr float kNeg = -1.0e30f;
+
+__global__ void router_topk_kernel(const float* __restrict__ s,
+                                   const float* __restrict__ bias,
+                                   const __nv_bfloat16* __restrict__ rs,
+                                   int* __restrict__ idx,
+                                   float* __restrict__ wts, int num_experts,
+                                   int topk) {
+  extern __shared__ float smem[];
+  float* scores = smem;                // [num_experts]
+  float* biased = smem + num_experts;  // [num_experts]
+  __shared__ float red_v[kWarps];
+  __shared__ int red_i[kWarps];
+
+  const int bb = blockIdx.x;
+  const int tid = threadIdx.x;
+  const int lane = tid & (WARP_SIZE - 1);
+  const float* srow = s + (size_t)bb * num_experts;
+
+  for (int e = tid; e < num_experts; e += kThreads) {
+    const float sig = 1.0f / (1.0f + expf(0.0f - srow[e]));
+    scores[e] = sig;
+    biased[e] = sig + bias[e];
+  }
+  __syncthreads();
+
+  float den = 0.0f;  // only thread 0's copy accumulates
+  for (int t = 0; t < topk; ++t) {
+    float best = kNeg;
+    int bi = 0;
+    for (int e = tid; e < num_experts; e += kThreads) {
+      if (best < biased[e]) {
+        best = biased[e];
+        bi = e;
+      }
+    }
+    for (int off = WARP_SIZE / 2; off > 0; off >>= 1) {
+      const float ov = __shfl_down_sync(0xffffffffu, best, off);
+      const int oi = __shfl_down_sync(0xffffffffu, bi, off);
+      if (ov > best || (ov == best && oi < bi)) {
+        best = ov;
+        bi = oi;
+      }
+    }
+    if (lane == 0) {
+      red_v[tid / WARP_SIZE] = best;
+      red_i[tid / WARP_SIZE] = bi;
+    }
+    __syncthreads();
+    if (tid == 0) {
+      best = red_v[0];
+      bi = red_i[0];
+      for (int w = 1; w < kWarps; ++w) {
+        if (red_v[w] > best || (red_v[w] == best && red_i[w] < bi)) {
+          best = red_v[w];
+          bi = red_i[w];
+        }
+      }
+      idx[(size_t)bb * topk + t] = bi;
+      wts[(size_t)bb * topk + t] = scores[bi];
+      biased[bi] = kNeg;
+      den += scores[bi];
+    }
+    // The winner's knock-out (and red_* reuse) must land before the next scan.
+    __syncthreads();
+  }
+
+  if (tid == 0) {
+    const float rsf = __bfloat162float(rs[0]);
+    for (int t = 0; t < topk; ++t) {
+      wts[(size_t)bb * topk + t] =
+          (wts[(size_t)bb * topk + t] / (den + 1e-20f)) * rsf;
+    }
+  }
+}
+
+CUresult map_cuda_error(cudaError_t err) {
+  if (err == cudaSuccess) return CUDA_SUCCESS;
+  if (err == cudaErrorInvalidValue || err == cudaErrorInvalidDevicePointer) {
+    return CUDA_ERROR_INVALID_VALUE;
+  }
+  if (err == cudaErrorMemoryAllocation) return CUDA_ERROR_OUT_OF_MEMORY;
+  if (err == cudaErrorNotSupported) return CUDA_ERROR_NOT_SUPPORTED;
+  return CUDA_ERROR_LAUNCH_FAILED;
+}
+
+}  // namespace
+
+extern "C" {
+
+// Sigmoid router plus biased top-k over merged f32 score rows `s [b, E]`,
+// with `bias [E]` f32 and the bf16 routed scale `rs [1]`. Writes
+// `idx [b, topk]` i32 and `wts [b, topk]` f32. Shapes are runtime values —
+// no per-bucket instantiation; shared memory holds 2*E f32.
+CUresult k3_router_topk_cuda(const float* s, const float* bias, const void* rs,
+                             int* idx, float* wts, int b, int num_experts,
+                             int topk, cudaStream_t stream) {
+  PEGAINFER_FFI_GUARD_BEGIN
+  if (s == nullptr || bias == nullptr || rs == nullptr || idx == nullptr ||
+      wts == nullptr || b <= 0 || num_experts <= 0 || topk <= 0 ||
+      topk > num_experts) {
+    return CUDA_ERROR_INVALID_VALUE;
+  }
+  const size_t shmem = 2 * (size_t)num_experts * sizeof(float);
+  router_topk_kernel<<<b, kThreads, shmem, stream>>>(
+      s, bias, static_cast<const __nv_bfloat16*>(rs), idx, wts, num_experts,
+      topk);
+  return map_cuda_error(cudaGetLastError());
+  PEGAINFER_FFI_GUARD_END(CUDA_ERROR_UNKNOWN)
+}
+
+}  // extern "C"

pegainfer-kernels/src/ffi/k3.rs

@@ -9,6 +9,8 @@
 //!
 //! See `csrc/k3/k3_deepgemm_fp8_fp4_grouped_sm100.cu`.
 
+use core::ffi::c_void;
+
 use cudarc::driver::sys::CUresult;
 use cudarc::driver::sys::CUstream;
 
@@ -153,6 +155,24 @@ unsafe extern "C" {
         stream: CUstream,
     ) -> CUresult;
 
+    /// Sigmoid router plus biased top-k over merged f32 score rows
+    /// (`csrc/k3/k3_router_topk.cu`): `s [b, num_experts]` f32,
+    /// `bias [num_experts]` f32, the bf16 routed scale `rs [1]`; writes
+    /// `idx [b, topk]` i32 and `wts [b, topk]` f32. Block-parallel argmax
+    /// with the serial kernel's lowest-index tie-break; shapes are runtime
+    /// values (no per-bucket instantiation).
+    pub fn k3_router_topk_cuda(
+        s: *const f32,
+        bias: *const f32,
+        rs: *const c_void,
+        idx: *mut i32,
+        wts: *mut f32,
+        b: i32,
+        num_experts: i32,
+        topk: i32,
+        stream: CUstream,
+    ) -> CUresult;
+
     // --- fused MegaMoE (see `csrc/k3/k3_mega_moe_sm100.cu`) ---
 
     /// Token-count alignment the MegaMoE API enforces on

pegainfer-kernels/src/ffi/k3_tilelang.rs

@@ -151,21 +151,6 @@ unsafe extern "C" {
         stream: CUstream,
     ) -> i32;
 
-    /// Sigmoid router plus biased top-k over merged f32 score rows
-    /// `S [b, num_experts]`, with `Bias [num_experts]` f32 and the bf16 routed
-    /// scale `Rs [1]`. Writes `Idx [b, topk]` i32 and `Wts [b, topk]` f32.
-    pub fn k3_router_topk_batched(
-        s: *const f32,
-        bias: *const f32,
-        rs: *const c_void,
-        idx: *mut i32,
-        wts: *mut f32,
-        b: i32,
-        num_experts: i32,
-        topk: i32,
-        stream: CUstream,
-    ) -> i32;
-
     /// Attention-residual candidate scoring: weightless RMS normalization then
     /// a dot with the fused f32 scoring vector `Sw [h]`. `Ps [b, h]` is the
     /// running prefix sum, `Bl [b, blocks, h]` that row's snapshot history;

pegainfer-kernels/src/ops/k3/mod.rs

@@ -6,10 +6,12 @@ mod flash_mla_prefill;
 mod mega_moe;
 mod mla_paged;
 mod moe_chain;
+mod router_topk;
 
 pub use deepgemm::*;
 pub use flash_kda::*;
 pub use flash_mla_prefill::*;
 pub use mega_moe::*;
 pub use mla_paged::*;
 pub use moe_chain::*;
+pub use router_topk::*;