perf(k3): block-parallel router top-k replaces the serial TileLang scan - #939
perf(k3): block-parallel router top-k replaces the serial TileLang scan#939xiaguan wants to merge 1 commit into
Conversation
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 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 14bb15ffd7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| @@ -0,0 +1,151 @@ | |||
| // Kimi-K3 MoE router: sigmoid scores plus biased top-k selection. | |||
There was a problem hiding this comment.
This commit has no Signed-off-by: trailer, so the repository's DCO check will reject it regardless of the code's correctness; recreate or rebase the commit with --signoff before submitting it.
AGENTS.md reference: AGENTS.md:L184-L184
Useful? React with 👍 / 👎.
| 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 |
There was a problem hiding this comment.
After removing this TileLang family, the authoritative inventories still describe router_topk_batched as generated TileLang code: pegainfer-k3/kernels/README.md still reports eleven families, 420 instantiations, and the removed launcher, while pegainfer-kernels/KERNELS.md:141 still names the retired .cu artifact and serial implementation. This now misdirects anyone diagnosing or regenerating the K3 kernel set, so update those records alongside the implementation change.
AGENTS.md reference: AGENTS.md:L168-L170
Useful? React with 👍 / 👎.
1 participant
What
Replaces the TileLang
router_topk_batchedkernel — whose biased top-k ran as a serial scan on thread 0 (TOPK×E = 16×896 iterations per row, 65.7µs per launch, ~6ms of every EP16 verify round across the 92 expert layers) — with a hand-written CUDA kernel (csrc/k3/k3_router_topk.cu) doing a block-parallel argmax per selection round.Every arithmetic step keeps the retired kernel's exact spelling, so outputs are bit-identical:
expf(no fast-math), bias added in f32<comparisons with the lowest-index tie-break: each thread scans its stride-256 subsequence ascending, the shuffle/shared reduction prefers the strictly-greater value and breaks equal values to the lower index — exactly the serial first-match selection(w / (den + 1e-20)) * (float)rsunchangedThe TileLang factory, its generator plan, and the AOT launcher entry are retired (the vendored
tilelang_defs.pystays verbatim minus the removed factory). Batch and expert count become plain runtime launch values — no per-bucket instantiation.Correctness
Gates on the pruned 224-expert checkpoint (any selection divergence would blow the golden logits far past the calibrated 2.0-ULP noise floor):
golden_decode13/13 (mega + masked-chain transports)spec_verify6/6Perf
🤖 Generated with Claude Code