Skip to content

(WIP) perf(kimi3): shard the TP MoE final projection - #1145

Draft
panditsa wants to merge 1 commit into
lightseekorg:sanket/wip-k3-tp-joint-moefrom
panditsa:sanket/wip-k3-sharded-finalize
Draft

(WIP) perf(kimi3): shard the TP MoE final projection#1145
panditsa wants to merge 1 commit into
lightseekorg:sanket/wip-k3-tp-joint-moefrom
panditsa:sanket/wip-k3-sharded-finalize

Conversation

@panditsa

@panditsa panditsa commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Summary

  • enable column-sharded routed up-projection weights for gfx950 TP8/EP1
  • reduce routed latents before projecting each local column shard for M>=8
  • fold each shard into the shared partial so the final all-reduce assembles the full output

Validation result

This WIP is not ready to land: the intended M=8 crossover regresses full-model performance.

Exact parent 30cac366 → head 8ad030df, 4K/1K TP8/EP1 CUDA graphs, median of three:

M Parent tok/s Head tok/s Change Parent TPOT Head TPOT
8 312.48 272.01 -13.0% 23.196 ms 25.688 ms

The three head repetitions agree closely. The result indicates that the extra latent reduction plus sharded projection/final assembly costs more than the replicated tail it replaces at M=8. The M>=8 selector should not be enabled without a different implementation or a newly measured crossover.

Tests

  • ordered projection-before-reduction behavior: 1 passed
  • AMD TP8-only platform selection: 1 passed
  • gfx950 strided shard update and CUDA-graph capture at M=8/16: 2 passed
  • all 48 measured A/B requests completed with valid lengths

Stack

WIP blocker

  • remove/disable the regressing runtime selection or rework the sharded tail and remeasure before landing

Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
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.

1 participant