Skip to content

(WIP) perf(kimi-k3): tune small-M MoE decode - #1140

Draft
panditsa wants to merge 4 commits into
lightseekorg:sanket/wip-k3-a8w4-selectfrom
panditsa:sanket/wip-k3-small-m-moe
Draft

(WIP) perf(kimi-k3): tune small-M MoE decode#1140
panditsa wants to merge 4 commits into
lightseekorg:sanket/wip-k3-a8w4-selectfrom
panditsa:sanket/wip-k3-small-m-moe

Conversation

@panditsa

@panditsa panditsa commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Summary

  • extend the joint latent/shared decode pipeline from M=1 through M=4
  • rebalance routing and W2 work for M=2/4
  • tune TP-local A16W4 tiles and split top-k route work across CTAs

Performance

4K/1K TP8/EP8 CUDA graphs, median of three:

M Prior equivalent stack tok/s Head tok/s Change Prior TPOT Head TPOT
2 84.77 93.39 +10.2% 22.864 ms 20.663 ms
4 152.23 156.34 +2.70% 24.954 ms 24.255 ms

The comparison uses the last equivalent EP8 boundary: intervening #1137/#1138 only add unselected A8W4 foundations and #1139 registers the exact TP8/EP1 shape. The measured EP8 execution path is unchanged until this PR.

Tests

  • BF16 SiTU M=1/2/3/4 coverage: 4 passed
  • router dispatch/per-token FP32 route coverage: 2 passed
  • fused RMSNorm + projection M=1/2/3/4 graph coverage: 4 passed
  • all 18 measured 4K/1K requests completed with valid lengths

Stack

Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Extend the latent decode registration, shared-down fusion, and RMSNorm projection support from one token through M=4.

Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Use wider router wave counts after M=1, finer W2 output tiles for M=2/4, and masked tail loads for 128-aligned intermediate widths.

Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Use narrower TP-local W13/W2 tiles and split top-k routes across CTAs before a final reduction. Keep the experimental A8W4 wrapper and load reordering out of this change.

Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
panditsa force-pushed the sanket/wip-k3-small-m-moe branch from 2ebd8f8 to 377b289 Compare August 19, 2026 20:21
@panditsa
panditsa force-pushed the sanket/wip-k3-a8w4-select branch from e689319 to bcebe08 Compare August 19, 2026 20:25
@jhinpan

jhinpan commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

Exact #1132 screen of the only M=8-applicable change here (the sigmoid-router warp-count line), applied independently to current main (9196fa27) with forced acceptance=1 at 131072/1024 TP8/EP8:

  • baseline: 36.80 ms TPOT, 19.7647 tok/s
  • router candidate: 36.77 ms TPOT, 19.7753 tok/s

That is only ~+0.05% throughput / 0.03 ms and is not an issue-level material win. The larger W13/W2/pipeline changes remain gated to M<=4 or the TP-local EP1 shape, so I will not extract this line into a separate #1132 PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants