Repository navigation
[kimi-k3] Rank-local expert ids for the EP routed MoE (upstream #45544) - #37
Merged
Venkat2811 merged 1 commit intoOct 27, 2026
Merged
Conversation
Venkat2811
force-pushed
the
venkat/n5884e4fe-03-dcp-pack-device
branch
from
October 27, 2026 12:20
1e37e20 to
e9a9607
Compare
Venkat2811
force-pushed
the
venkat/n5884e4fe-04-rank-local-ids
branch
from
October 27, 2026 12:20
7ac60c5 to
5abf7e0
Compare
Venkat2811
force-pushed
the
venkat/n5884e4fe-03-dcp-pack-device
branch
from
October 27, 2026 15:28
e9a9607 to
60cf309
Compare
…gl-project#45544) Layer 4 of the Kimi-K3 nightly line on upstream main 5884e4f. Under expert parallelism each rank's trtllm-gen routed MoE computes only its own experts (112 of 896 at EP8), but FlashInfer picks its routing kernel from num_experts. K3's 896 experts select the Tier<896, 16> specialisation, which is above the dyn-block kernel's 512-expert limit, so decode steps of 5 to 16 tokens run routingIndicesClusterKernel, about 11 to 12.6 us per MoE layer on 8x B300 (once per MoE layer, so about 1 ms per decode step). The fused route+quant launch (route_quant_fused, decode batches up to 64 tokens) already writes the packed (id << 16 | bf16 weight) ids the routed MoE consumes. When the model stages the rank's expert range, its pack epilogue now writes the rank's experts as 0..count-1 and any other expert in top-k slot k as count + k (unique per row), and the runner passes num_experts = round_up(count + top_k, 4) with local_expert_offset 0 (128 at EP8: Tier<256, 16>, dyn-block kernel up to 16 tokens, a smaller cluster kernel above). The global ids (out_i, which TopKOutput carries on) are unchanged, and the same (token, local expert, weight) entries reach the expert GEMMs in the same order, so the MoE output is bitwise unchanged for the same GEMM tactic. Plumbing: - route_radix.cuh / route_quant_fused.cuh / moe_route_quant_fused.py: optional (local_expert_offset, num_local_experts) for the packed output only; 0 keeps the global ids. MAX_TOKENS (the fused launch's token cap) is public. - route_quant_handoff.py: stage(x, local_experts) takes the rank's range; rank_local_num_experts() decides whether it applies (a proper sub-range that shrinks the routing problem); the range is recorded with the packed tensor and read back once with take_local_experts(routing), keyed by that tensor. - kimi_k3.py: stages the range Mxfp4MoEMethod hands the runner, read per forward because bind_full_expert_weights can change it after construction. - flashinfer_trtllm.py: passes the rank-local num_experts and offset 0 for such packed ids, after asserting the range equals quant_info's. Autotune coverage: FlashInfer's MoE tuner key includes num_experts. When the decode-shaped autotune forward runs above the fused route's 64-token cap (decode CUDA graphs past 64 rows), it takes the unfused routing and would only tune the global key, leaving every rank-local decode call on FlashInfer's untuned default tactic. During tuning, when the runner sees unfused routing and rank-local ids apply, it also runs FlashInfer on up to 64 of that forward's routes, relabelled as the fused route would, so the decode key is tuned from real routes. Serving paths are unchanged. Measured on 8x B300, Kimi-K3 TP8 DCP8 EP8, one node: routing kernel per MoE layer 11.1 / 12.2 / 12.6 us -> 4.3 / 4.8 / 7.1 us at 8 / 16 / 24 decode rows; decode output throughput (bench_one_batch_server, 8192 in / 512 out) +6.1 / +4.2 / +1.0 % at batch 8 / 16 / 32 with the expert-GEMM tactics held at main's fresh picks, and +6.3 / +10.7 / -5.0 % with a fresh autotune per arm (the tuner picks different GEMM tactics for the new key). Tests: the kernel's rank-local ids against a torch oracle (EP 2/8/16, first, middle and last rank, 1 to 64 tokens) with everything else it writes unchanged; the handoff bookkeeping (keyed, consumed once, global ids when the range does not apply); and FlashInfer's routed MXFP4 MoE output bitwise equal between global and rank-local ids (EP 2/4/8/16, 1 to 64 tokens, SM100). Source: e251e3d on the kimi-k3-nightly-3c560089 line (fork PR #20), a mirror of upstream sgl-project#45544 (head fc02c94). One conflict, in moe_route_quant_fused.covered(): upstream sgl-project#44690 dropped the x.shape[0] == scores.shape[0] check (a mismatch is a caller bug the launcher reports); kept that, and took the PR's public MAX_TOKENS rename. Checked against sgl-project#44480 (merged 20 Oct, StandardDispatcher maps expert ids for runners that need local ids under EP): flashinfer_trtllm and flashinfer_mxfp4 stay in its keep-global-ids list, so nothing else remaps for this runner and the packed ids are made rank-local once, in the fused route epilogue. The range the model stages (moe_ep_rank * num_local_experts, num_local_experts) is the one the MXFP4 quant methods hand the runner at 5884e4f (mxfp4.py, mxfp4_flashinfer_trtllm_moe.py), and the runner asserts the two are equal. The runner function this changes is untouched by upstream since 3c56008. Tests: GPU only (test_rank_local_routed_moe.py, test_route_quant_prerequisites.py launch the JIT kernels; on a CPU host the new file does not import on the parent and skips or fails for want of CUDA here, as the existing cases in test_route_quant_prerequisites.py do on the base). On one B300 GPU on the 3c56008 line: test_rank_local_routed_moe.py 1/1 at this layer (2/2 on 23 Oct counted layer 6's padded-row case) and test_route_quant_prerequisites.py 5/5. They rerun on GPU with this branch's image. Upstream: sgl-project#45544 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Venkat2811
force-pushed
the
venkat/n5884e4fe-04-rank-local-ids
branch
from
October 27, 2026 15:28
5abf7e0 to
d29175e
Compare
Venkat2811
changed the base branch from
venkat/n5884e4fe-03-dcp-pack-device
to
kimi-k3-nightly-5884e4fe
October 27, 2026 15:28
Venkat2811
marked this pull request as ready for review
October 27, 2026 15:28
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.
Layer 4 of the Kimi-K3 serving line on upstream main 5884e4fec6 (base branch
kimi-k3-nightly-5884e4fe, thenightly-dev-cu13-20261027-5884e4feimage). Stacked: this PR's base is the previous layer's branch.What: under EP the fused route+quant epilogue writes rank-local packed expert ids, so FlashInfer routes a 128-expert problem (EP8) instead of 896. Commit 5abf7e05a4: mirror of upstream #45544 (as e251e3d32b, fork PR #20). One conflict: upstream #44690 dropped a shape check in
covered(); kept that, took the PR's publicMAX_TOKENS.Why: about 1 ms per decode step on B300 (routing kernel 11-12.6 us -> 4.3-7.1 us per MoE layer at 8-24 rows; 23 Oct line: -0.75 ms decode step, mean of five nodes).
Double remap check (#44480):
flashinfer_trtllmandflashinfer_mxfp4stay in its keep-global-ids list, so the dispatcher never remaps for this runner; the staged range equals what the MXFP4 quant methods hand the runner, and the runner asserts it.Tests: GPU only (
test_rank_local_routed_moe.py,test_route_quant_prerequisites.py); on one B300 GPU on this line: on the parent the new file does not import (no rank_local_num_experts) and the prerequisite file fails 2 (unexpected local_experts / local_expert_offset); with the commit 6 passed. On CPU the new file skips.Risk: medium (decode hot path); output bitwise unchanged for the same GEMM tactic by construction and test.
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): 🚫 Run #40531310894
Latest PR Test (Extra): 🚫 Run #40531310854
Latest PR Test (AMD ROCm 10): 🚫 Run #40531311311