From 14bb15ffd7faf509f409941e6f8fb2acca28cdbe Mon Sep 17 00:00:00 2001 From: xiaguan <751080330@qq.com> Date: Wed, 19 Aug 2026 16:49:58 +0000 Subject: [PATCH] perf(k3): block-parallel router top-k replaces the serial TileLang scan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The TileLang router_topk_batched kernel ran the whole biased top-k as a serial scan on thread 0 — TOPK x E iterations per row, 65.7us per launch at 896 experts, ~6ms of every EP16 verify round across the 92 expert layers. Replace it with a hand-written CUDA kernel whose selection is a block-parallel argmax per round while every arithmetic step keeps the retired kernel's spelling, so the outputs are bit-identical: f32 sigmoid via plain expf, bias added in f32, strict less-than comparisons with the lowest-index tie-break (each thread scans its stride ascending; the reduction prefers the strictly greater value and breaks equal values to the lower index, which is exactly the serial first-match), the denominator accumulated in selection order, and the division-then-scale weight normalization. The TileLang factory, its generator plan and the AOT launcher entry are retired; batch and expert count become plain runtime launch values (no per-bucket instantiation). Gates on the pruned 224-expert checkpoint: golden decode 13/13 (mega + masked chain — any selection divergence would blow past the 2.0-ULP noise floor), spec_verify 6/6, lib tests green. Perf: 65.7us -> 10.3us per launch (bucket-one golden replay capture, 240 launches, max 11.2us); an EP16 verify step's 92 router launches drop from ~6.0ms to ~0.95ms. Co-Authored-By: Claude Fable 5 Signed-off-by: xiaguan <751080330@qq.com> --- pegainfer-k3/kernels/generate.py | 56 +------- pegainfer-k3/kernels/tilelang_defs.py | 54 +------ pegainfer-kernels/build.rs | 4 - pegainfer-kernels/csrc/k3/k3_router_topk.cu | 151 ++++++++++++++++++++ pegainfer-kernels/src/ffi/k3.rs | 20 +++ pegainfer-kernels/src/ffi/k3_tilelang.rs | 15 -- pegainfer-kernels/src/ops/k3/mod.rs | 2 + pegainfer-kernels/src/ops/k3/router_topk.rs | 82 +++++++++++ pegainfer-kernels/src/ops/k3_tilelang.rs | 70 +-------- 9 files changed, 265 insertions(+), 189 deletions(-) create mode 100644 pegainfer-kernels/csrc/k3/k3_router_topk.cu create mode 100644 pegainfer-kernels/src/ops/k3/router_topk.rs diff --git a/pegainfer-k3/kernels/generate.py b/pegainfer-k3/kernels/generate.py index 350d101e..8fab1c14 100644 --- a/pegainfer-k3/kernels/generate.py +++ b/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, ] diff --git a/pegainfer-k3/kernels/tilelang_defs.py b/pegainfer-k3/kernels/tilelang_defs.py index d360c951..4bbaf771 100644 --- a/pegainfer-k3/kernels/tilelang_defs.py +++ b/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): diff --git a/pegainfer-kernels/build.rs b/pegainfer-kernels/build.rs index a36b3787..23561414 100644 --- a/pegainfer-kernels/build.rs +++ b/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", diff --git a/pegainfer-kernels/csrc/k3/k3_router_topk.cu b/pegainfer-kernels/csrc/k3/k3_router_topk.cu new file mode 100644 index 00000000..1bddfbe7 --- /dev/null +++ b/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 +#include + +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<<>>( + s, bias, static_cast(rs), idx, wts, num_experts, + topk); + return map_cuda_error(cudaGetLastError()); + PEGAINFER_FFI_GUARD_END(CUDA_ERROR_UNKNOWN) +} + +} // extern "C" diff --git a/pegainfer-kernels/src/ffi/k3.rs b/pegainfer-kernels/src/ffi/k3.rs index 9249b528..000d55ef 100644 --- a/pegainfer-kernels/src/ffi/k3.rs +++ b/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 diff --git a/pegainfer-kernels/src/ffi/k3_tilelang.rs b/pegainfer-kernels/src/ffi/k3_tilelang.rs index 1ea23a19..bd9f4f8f 100644 --- a/pegainfer-kernels/src/ffi/k3_tilelang.rs +++ b/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; diff --git a/pegainfer-kernels/src/ops/k3/mod.rs b/pegainfer-kernels/src/ops/k3/mod.rs index 2121cda9..33804b6e 100644 --- a/pegainfer-kernels/src/ops/k3/mod.rs +++ b/pegainfer-kernels/src/ops/k3/mod.rs @@ -6,6 +6,7 @@ mod flash_mla_prefill; mod mega_moe; mod mla_paged; mod moe_chain; +mod router_topk; pub use deepgemm::*; pub use flash_kda::*; @@ -13,3 +14,4 @@ pub use flash_mla_prefill::*; pub use mega_moe::*; pub use mla_paged::*; pub use moe_chain::*; +pub use router_topk::*; diff --git a/pegainfer-kernels/src/ops/k3/router_topk.rs b/pegainfer-kernels/src/ops/k3/router_topk.rs new file mode 100644 index 00000000..17c30001 --- /dev/null +++ b/pegainfer-kernels/src/ops/k3/router_topk.rs @@ -0,0 +1,82 @@ +//! Kimi-K3 MoE router: sigmoid scores plus biased top-k selection. +//! +//! Hand-written replacement for the retired TileLang `router_topk_batched` +//! kernel, whose serial thread-0 scan cost ~65us per launch at 896 experts. +//! The selection is a block-parallel argmax per round with the serial +//! kernel's exact arithmetic and lowest-index tie-break, so the outputs are +//! bit-identical — see `csrc/k3/k3_router_topk.cu` for the argument. Batch is +//! a plain launch dimension (no per-bucket instantiation), but callers still +//! run the compiled buckets: every other kernel in the step is bucket-shaped. + +use core::ffi::c_void; + +use anyhow::Result; +use anyhow::anyhow; +use anyhow::ensure; +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtr; +use cudarc::driver::DevicePtrMut; +use half::bf16; + +use crate::ffi; +use crate::tensor::DeviceContext; + +/// Sigmoid router plus biased top-k over already-merged f32 score rows. +/// +/// The weights come from the *un-biased* scores, are normalized with a +/// `+1e-20` guard and scaled by the bf16 routed scale `rs`. Ties break to the +/// lowest expert index. +#[allow(clippy::too_many_arguments)] +pub fn k3_router_topk_batched_launch( + ctx: &DeviceContext, + b: usize, + num_experts: usize, + topk: usize, + s: &CudaSlice, + bias: &CudaSlice, + rs: &CudaSlice, + idx: &mut CudaSlice, + wts: &mut CudaSlice, +) -> Result<()> { + ensure!(b > 0, "K3 router needs rows"); + ensure!( + topk <= num_experts, + "K3 router topk={topk} exceeds the expert count {num_experts}" + ); + ensure!( + s.len() >= b * num_experts + && bias.len() >= num_experts + && !rs.is_empty() + && idx.len() >= b * topk + && wts.len() >= b * topk, + "K3 router buffers too small for b={b}, experts={num_experts}, topk={topk}: \ + s {}, bias {}, rs {}, idx {}, wts {}", + s.len(), + bias.len(), + rs.len(), + idx.len(), + wts.len() + ); + let (s_ptr, _s_guard) = s.device_ptr(&ctx.stream); + let (bias_ptr, _bias_guard) = bias.device_ptr(&ctx.stream); + let (rs_ptr, _rs_guard) = rs.device_ptr(&ctx.stream); + let (idx_ptr, _idx_guard) = idx.device_ptr_mut(&ctx.stream); + let (wts_ptr, _wts_guard) = wts.device_ptr_mut(&ctx.stream); + unsafe { + ffi::k3_router_topk_cuda( + s_ptr as *const f32, + bias_ptr as *const f32, + rs_ptr as *const c_void, + idx_ptr as *mut i32, + wts_ptr as *mut f32, + i32::try_from(b)?, + i32::try_from(num_experts)?, + i32::try_from(topk)?, + crate::tensor::active_cu_stream(ctx), + ) + } + .result() + .map_err(|err| { + anyhow!("K3 router_topk (B={b}, E={num_experts}, TOPK={topk}) launch failed: {err}") + }) +} diff --git a/pegainfer-kernels/src/ops/k3_tilelang.rs b/pegainfer-kernels/src/ops/k3_tilelang.rs index e16abe27..2680d785 100644 --- a/pegainfer-kernels/src/ops/k3_tilelang.rs +++ b/pegainfer-kernels/src/ops/k3_tilelang.rs @@ -3,11 +3,13 @@ //! //! The set covers one whole K3 decode step that is not a GEMM or attention — //! norms and the bf16 landings of the framework GEMMs, the KDA convolution and -//! delta rule, the MoE router and expert combine, the situ activation and the +//! delta rule, the expert combine, the situ activation and the //! attention-residual mix. Dense projections are served by cuBLASLt, the -//! routed experts by the DeepGEMM masked grouped-GEMM chain, and MLA decode by -//! the hand-written absorbed paged kernel (`ops::k3::mla_paged`), so neither a -//! GEMV nor an attention family lives here. The wrappers keep the certified kernels' operand names, so an +//! routed experts by the DeepGEMM masked grouped-GEMM chain, MLA decode by +//! the hand-written absorbed paged kernel (`ops::k3::mla_paged`), and the MoE +//! router top-k by the hand-written parallel-argmax kernel +//! (`ops::k3::router_topk`), so neither a GEMV nor an attention family lives +//! here. The wrappers keep the certified kernels' operand names, so an //! executor written against them reads like the Python engine's launch //! sequence. //! @@ -584,66 +586,6 @@ pub fn k3_o_norm_gate_batched_launch( ) } -/// Sigmoid router plus biased top-k over already-merged f32 score rows. -/// -/// The weights come from the *un-biased* scores, are normalized with a -/// `+1e-20` guard and scaled by the bf16 routed scale `rs`. Ties break to the -/// lowest expert index. -#[allow(clippy::too_many_arguments)] -pub fn k3_router_topk_batched_launch( - ctx: &DeviceContext, - b: usize, - num_experts: usize, - topk: usize, - s: &CudaSlice, - bias: &CudaSlice, - rs: &CudaSlice, - idx: &mut CudaSlice, - wts: &mut CudaSlice, -) -> Result<()> { - check_bucket(b)?; - ensure!( - topk <= num_experts, - "K3 router topk={topk} exceeds the expert count {num_experts}" - ); - ensure!( - s.len() >= b * num_experts - && bias.len() >= num_experts - && !rs.is_empty() - && idx.len() >= b * topk - && wts.len() >= b * topk, - "K3 router buffers too small for b={b}, experts={num_experts}, topk={topk}: \ - s {}, bias {}, rs {}, idx {}, wts {}", - s.len(), - bias.len(), - rs.len(), - idx.len(), - wts.len() - ); - let (s_ptr, _s_guard) = s.device_ptr(&ctx.stream); - let (bias_ptr, _bias_guard) = bias.device_ptr(&ctx.stream); - let (rs_ptr, _rs_guard) = rs.device_ptr(&ctx.stream); - let (idx_ptr, _idx_guard) = idx.device_ptr_mut(&ctx.stream); - let (wts_ptr, _wts_guard) = wts.device_ptr_mut(&ctx.stream); - let rc = unsafe { - ffi::k3_router_topk_batched( - s_ptr as *const f32, - bias_ptr as *const f32, - rs_ptr as *const c_void, - idx_ptr as *mut i32, - wts_ptr as *mut f32, - b as i32, - num_experts as i32, - topk as i32, - ctx.stream.cu_stream(), - ) - }; - check( - rc, - &format!("K3 router_topk_batched (B={b}, E={num_experts}, TOPK={topk})"), - ) -} - /// Score the `blocks + 1` attention-residual candidates of every row: a /// weightless RMS normalization then a dot with the fused f32 scoring vector. #[allow(clippy::too_many_arguments)]