(WIP) perf(kimi-k3): tune small-M MoE decode - #1140
Draft
panditsa wants to merge 4 commits into
Draft
Conversation
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
force-pushed
the
sanket/wip-k3-small-m-moe
branch
from
August 19, 2026 20:21
2ebd8f8 to
377b289
Compare
panditsa
force-pushed
the
sanket/wip-k3-a8w4-select
branch
from
August 19, 2026 20:25
e689319 to
bcebe08
Compare
Contributor
|
Exact #1132 screen of the only M=8-applicable change here (the sigmoid-router warp-count line), applied independently to current main (
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. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Performance
4K/1K TP8/EP8 CUDA graphs, median of three:
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
Stack