Repository navigation
[ROCm][Perf] Kimi-K3: token-sharded residual stream for long prefills - #59693
vanshbhatia-amd wants to merge 12 commits into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Signed-off-by: Vansh Bhatia <vansh.bhatia@amd.com>
54da716 to
8156003
Compare
Both up-projection tails folded the shared output in with addmm_. On ROCm, hipBLASLt's C-accumulating bf16 GEMM faults at some row counts for these shapes: 2914-2925 rows per rank at 7168x3584 in the token-sharded tail (a ~23.4K-token prefill step at TP8), and bands from ~16.5K tokens at 896x3584 in the existing tail. Use a plain GEMM plus an add, which costs about 10 us per MoE layer at 3K rows. Signed-off-by: Vansh Bhatia <vansh.bhatia@amd.com>
- Reuse the column-parallel up-projection tail and reduce-scatter its output over tokens, so the path works with the row-sharded up_proj. - Gather pre-quantized AttnRes outputs as data plus per-token scale and rebuild them, keeping the FP8 input fusion. - Pass the sharded token count to the decoder layers and the MoE runner as arguments instead of module globals. - Only shard with DP == 1. - Drop the addmm_ workaround; the hipBLASLt fault it avoided only reproduces on the ROCm 10 preview image. - Test the ColumnParallelLinear up_proj layout and the quantized gather. Signed-off-by: Vansh Bhatia <vansh.bhatia@amd.com>
kliuae
left a comment
kliuae
left a comment
There was a problem hiding this comment.
Thanks for the turnaround. Looks good in general, just have some nits on setting reduce_results and test coverage.
As the latent is now all-reduced, instead of the reduce scatter before syncing in the upstream changes. Could you remeasure the performance numbers to check if the performance lift is still intact?
| out = self._run_self_attn(positions, hidden_states) | ||
| return sp.scatter_tokens(out) | ||
|
|
||
| def _run_moe_sp_shard_routed(self, hidden_shard: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
pre_routed depends on select_experts being exercised so that its output can be replaced. Is it possible to make sure that it's used, maybe adding a test that builds a small Kimi MoE module and compare the shard results with what un-sharded gets on the same rows, to make sure that this path doesn't route on the uninitialized logits?
There was a problem hiding this comment.
Added both, in adb2392. sp.pre_routed now asserts that the MoE forward called select_experts exactly once, so a backend that routes from router_logits directly, or routes per chunk, fails loudly instead of reading the placeholder logits.
test_shard_routed_moe_tp{4,8}_matches_unsharded_rows builds a small latent KimiMoE (column-parallel up-projection, a shared expert, a sigmoid router with correction bias and grouped top-k). It runs the layer unsharded on all tokens and through _run_moe_sp_shard_routed on each rank's token shard, and checks that each rank's output matches its rows of the unsharded result. It is parametrized over the Triton and AITER MoE backends. The production MXFP4 + AITER path also passes the assert end to end (TP8 + DCP8 serving run, 160/160 requests).
| ) | ||
|
|
||
| @contextlib.contextmanager | ||
| def _output_unreduced(self, module: nn.Module): |
There was a problem hiding this comment.
This forces reduce_results=False for all RowParallelLinear modules. Even though self_attn and mlp contains only the intended o-proj and down_proj as RowParallelLinear, for readability, can we make this more targeted for the intended modules? Something like
@staticmethod
@contextlib.contextmanager
def _output_unreduced(proj: RowParallelLinear):
# Row-parallel outputs stay partial so the caller can reduce-scatter.
assert isinstance(proj, RowParallelLinear) and proj.reduce_results
proj.reduce_results = False
try:
yield
finally:
proj.reduce_results = True
then call it with self.self_attn_o_proj and self.mlp.down_proj
There was a problem hiding this comment.
Done as suggested: _output_unreduced is now a staticmethod that takes the one RowParallelLinear and asserts it, called with self.self_attn.o_proj and self.mlp.down_proj.
…rojs - pre_routed now asserts the MoE forward routes exactly once, so the placeholder router logits can never be read. - Test a small latent KimiMoE: routing on the token shard returns this rank's rows of the unsharded layer at TP4/TP8, on Triton and AITER. - _output_unreduced takes the one RowParallelLinear to leave partial (attention o_proj, MLP down_proj) instead of every one in the module. Signed-off-by: Vansh Bhatia <vansh.bhatia@amd.com>
|
Thanks for the review! adb2392 addresses the two nits. On performance: the Test Result table in the description was re-measured on 2026-10-14, after the switch to the all-reduced latent, with the same nightly image for both arms. It shows +4.4% / +8.5% / +7.4% throughput and -15% / -10% / -10% TTFT at 10K/32, 50K/32 and 200K/8. The previous revision measured +5.7% / +11.1% / +9.3% at those points, so dropping the latent reduce-scatter costs about 1-2 points and most of the lift is intact. |
kliuae
left a comment
kliuae
left a comment
There was a problem hiding this comment.
LGTM, and thanks for updating the performance numbers
…e sharded MLP step Routed-experts capture and EPLB load recording see routing from select_experts, which under SP only runs on this rank's token shard, so should_shard now also requires neither to be enabled. Tests: TP4 padded gather/scatter, a token count that does not divide TP in the shard-routed MoE check, and _run_mlp_sp against the unsharded rows for a dense KimiMLP and a non-token-shardable KimiMoE. The sp.py docstring now notes the extra latent and top-k all-gathers in latent-MoE layers. Signed-off-by: Vansh Bhatia <vansh.bhatia@amd.com>
xaguilar-amd
left a comment
xaguilar-amd
left a comment
There was a problem hiding this comment.
Nice work!
Some comments/suggestions:
-
The token-sharded path is not gated on gfx950 so could you add gfx942 checks on the current head, same image with the flag off and with VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS=1024? Just to be sure that everything stays healthy there too. Please include the commands.
-
Accuracy in the description is from the previous revision. Please rerun it at this head, again with the same image on both arms and the flag set to 1024: GSM8K 5-shot, passkey out to ~190K, and the first-token top-1 check. Put the serve and eval commands in the description. A 20-shot case would be nice to have as well and would make the PR stronger.
-
For the MI355X sweep, the 2026-10-14 table is three high-concurrency points and reports throughput and TTFT only. The 8-point sweep is the earlier revision. On this head, could you add a few more concurrencies at one of the existing input lengths (10K or 50K), keep the 10K/32, 50K/32, and 200K/8 points to see how the code does there? Please report TPOT as well since it's an important metric. Same nightly image for the base and the PR, with the serve and vllm bench serve commands in the description.
|
Thanks! Reran both at this head (f0c6a4b), same Accuracy, stock vs PR (GSM8K is 400 questions, mean of 3 runs):
The 5-shot gap is run-to-run noise: two runs of the same arm already disagree on 7-10 of the 400 answers, and PR vs stock runs on 6-11. The two passkey misses are the same 8K cases on both arms. The one top-1 miss is a 50K prompt whose top-1 also differs between two stock runs. Sweep: kept 10K/32, 50K/32 and 200K/8 and added 50K at concurrency 1, 2, 4, 8 and 16, with TPOT. 8x MI355X, TP8, 1K output tokens, stock → PR. Throughput is (input + output) tok/s/GPU; TTFT and TPOT are means.
The PR is ahead at every point, +6.2% geomean throughput. The gain grows with concurrency at 50K, since more prefill tokens per step are sharded. TPOT drops because decode steps that share a batch with a long prefill finish sooner. |
|
On gfx942: 8x MI300X, same Unit tests: the PR's Kimi-K3 AMD tests (prefill SP, latent-MoE runner, KDA/MLA direct return, sequence parallel), including the TP4/TP8 multi-GPU cases: 78 passed. End to end: the MXFP4 Kimi-K3 checkpoint is ~1.5 TB, which doesn't fit in 8x 192 GB, so this uses Kimi-K3's config reduced to 24 layers (two attention-residual blocks; full attention every 4th layer, KDA elsewhere) and 64 experts, in bf16 with seeded dummy weights, flag off vs 1024. The dummy init was rescaled for the test (norm weights 1 ± 0.1, matrices with variance 1/fan_in), because the default ±1e-3 init gives uniform logits. Both arms use Triton MoE (
Every prefill step was above the threshold (no step between 33 and 1024 tokens), so the sharded path ran. Top-1 isn't a useful metric with random weights, since the top two logprobs are exact bf16 ties on most prompts. Commands: # unit tests
pytest -q tests/models/kimi_k3/test_amd_prefill_sp.py tests/models/kimi_k3/test_amd_latent_moe_runner.py \
tests/models/kimi_k3/test_amd_kda_direct_return.py tests/models/kimi_k3/test_amd_mla_direct_return.py \
tests/models/kimi_k3/test_sequence_parallel.py
# serve, reduced config; VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS=0 or 1024
VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=0 VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS=$SP \
vllm serve <Kimi-K3, 24 layers, 64 experts, no quantization_config> --served-model-name k3 \
--trust-remote-code --tensor-parallel-size 8 --load-format dummy --language-model-only \
--max-model-len 32768 --max-num-seqs 32 --gpu-memory-utilization 0.9 --no-enable-prefix-cachingThe probe sends 20 random-token prompts of 1K–16K one at a time, then 16 concurrent 3K-token prompts, greedy with 16 output tokens and |
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Vansh Bhatia <vansh.bhatia@amd.com>
|
Follow-up on gfx942 with the full model: 8x MI325X (256 GB), the real Kimi-K3 MXFP4 checkpoint, TP8, same image ( On gfx942 this image runs Kimi-K3's MXFP4 MoE through the
Single runs; the GSM8K differences are smaller than the run-to-run spread on MI355X. The passkey misses are the same conversational answers as on MI355X.
Separately, I merged main (3acec0a). The only conflict was the layer loop in |
|
@vanshbhatia-amd thanks for the extra experiments! I cannot see the serve commands on the description as you mentioned though. |
|
@xaguilar-amd sorry, they're in the Test Plan section of the description ("Serve (the PR arm adds ..."), but here are the exact MI325X ones. Serve, both arms (the PR arm first applies the PR's diff to the image's vllm, docker run --rm --device=/dev/kfd --device=/dev/dri --group-add video --group-add render \
--security-opt seccomp=unconfined --ipc=host --shm-size=128g --network host \
-v <Kimi-K3>:<Kimi-K3>:ro \
-e HSA_NO_SCRATCH_RECLAIM=1 -e HIP_FORCE_DEV_KERNARG=1 -e SAFETENSORS_FAST_GPU=1 \
-e VLLM_ROCM_USE_AITER=1 -e VLLM_ROCM_USE_AITER_MOE=1 \
vllm/vllm-openai-rocm@sha256:e0be00f1b99da608cc3c8571c7c86bae0c05f19e1e61a59c1d70aeea6bca6140 \
vllm serve <Kimi-K3> --served-model-name k3 --trust-remote-code \
--tensor-parallel-size 8 --max-model-len 204800 --max-num-seqs 32 \
--gpu-memory-utilization 0.95 --no-enable-prefix-cachingOn gfx942 the log shows Benchmark (one warmup first; vllm bench serve --backend vllm --model <Kimi-K3> --served-model-name k3 \
--trust-remote-code --port 8000 --dataset-name random \
--random-input-len $ISL --random-output-len 64 --random-range-ratio 0 \
--max-concurrency $C --num-prompts $N --ignore-eos --seed 42 \
--percentile-metrics ttft,tpot,itl,e2elGSM8K, passkey and the first-token prompts are the same requests as described under "Accuracy" in the description, against |
|
Hi @dllehr-amd @hongxiayang @tjtanaa - would love for a review on this! Thanks! |
|
LGTM |
Purpose
Scope: ROCm only, Kimi-K3 only, opt-in. All model changes are in
vllm/models/kimi_k3/amd/, which is only loaded on ROCm, and the new path onlyruns for Kimi-K3's attention-residual layers. It is off by default
(
VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS=0). NVIDIA, other models andCUDA-graph decode are unaffected.
With TP, every rank repeats the same per-token work (attention-residual mixing,
MoE router, latent down-projection, top-k) on all tokens. For eager steps with
at least
VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENStokens (1024 recommended),this PR shards the residual stream over tokens instead, padded to a multiple of
TP:
before the next block. For attention and dense MLP blocks the communication
volume is unchanged. Latent-MoE layers also all-gather the latent and the
top-k, about 17% more bytes per MoE layer at K3's sizes; the measured numbers
below include this.
all-reducing it.
It is off with PP, with DP > 1, during graph capture and specialized-decode
steps, when aux hidden states
are requested (EAGLE drafters), and with
--enable-return-routed-expertsorEPLB, since routing would only see this rank's token shard.
The NVIDIA Kimi-K3 path has something similar (
use_sequence_parallelinnvidia/model.py), but it only turns on with expert parallelism plus DP > 1or DeepGEMM MegaMoE, for every step. This PR covers plain TP on ROCm, and only
for long prefill steps.
Test Plan
tests/models/kimi_k3/test_amd_prefill_sp.py:TP4/TP8 for both the
ColumnParallelLinearand replicated up-projectionlayouts.
a dense
KimiMLP, a shard-routed latentKimiMoE, and a latentKimiMoEthat takes the all-gather fallback, with a token count that does not divide
TP.
(data plus per-token scale) are bit-exact, padded gather/scatter
round-trips at TP4/TP8, and the enable conditions are covered.
parallel tests: 78 passed on
vllm/vllm-openai-rocm:nightly(vllm 5555a2d).After merging main, with the attention-residual tests added, 94 pass on the
merged tree.
MI325X, the full Kimi-K3 MXFP4 checkpoint at TP8 matches stock on GSM8K 5/20-shot
and passkey, with the same first-token top-1 on 20/20 prompts up to 190K tokens,
and cuts TTFT by 5-9% (MoE runs through the
EMULATIONMXFP4 backend there, soabsolute numbers are not representative).
vllm/vllm-openai-rocm:nightlyimage (vllm 5555a2d); the PR arm appliesthis PR's diff to the image's vllm and sets the flag. The agentic TP8 + DCP8
config (prefix caching, CPU KV offload, 24576 batched tokens) also serves
160/160 requests with and without the flag.
Serve (the PR arm adds
-e VLLM_KIMI_K3_AMD_PREFILL_SP_MIN_TOKENS=1024):Benchmark (one warmup run first;
N = max(2 * C, 4)):Accuracy, against
/v1/completionswith temperature 0 on the same servers:first k train examples as
Question: ...\nAnswer: ...\n\n. Settings aremax_tokens320 and stop onQuestion:or three newlines, and the last numberis compared with the
####answer. 5-shot prompts are only about 550 tokens,so they are sent in synchronized waves of 32 to put every prefill step above
the threshold. 20-shot prompts (over 2K tokens) are sent at concurrency 32.
Each is run 3 times per arm.
8K/50K/100K/190K-token prompts (12 cases, 12 output tokens), plus 20 cases at
depths 0.1-0.9 (48 output tokens).
offsets at each of 2K/10K/50K/100K/190K tokens), sent one at a time with
logprobs=5, run twice per arm.Test Result
At this head, 2026-10-20, stock → PR. Throughput is (input + output)
tok/s/GPU; TTFT and TPOT are means.
TPOT drops because decode steps that share a batch with a long prefill finish
sooner.
The GSM8K differences are within run-to-run noise: two runs of the same arm
disagree on 7-10 of the 400 answers, and PR vs stock runs on 6-11 (mean 8.7).
On the agentic config at this head, 5-shot is 0.987 for both. The two passkey
misses are the same 8K cases on both arms, where the model answers
conversationally instead of with the key. The one top-1 miss is a 50K prompt
whose top-1 also differs between two stock runs.