[ROCm][Perf] Kimi-K3 Shard Latent MoE up-projection for ROCm path

[kliuae]

Purpose

For Kimi-K3, the default routed up-projection is a ReplicatedLinear where every TP rank holds the same weight and computes idential full projection where only 1/8 of the work is needed. On the CUDA path, there is LatentMoERunner that shards the workload in a column parallel manner. However, currently the ROCm path lacks this and does not pass a runner_cls to the FusedMoeFactory so it never reaches it.

This PR adds a ROCm LatentMoERunner counterpart and the up projection sharding strategy, mirroring the CUDA path's Tier 2, to allow each rank to project only its slice of the hidden states and accumulate on the partial output. The GEMM's beta-add epilogue folds the two together and the final all-reduce stitches the routed shards when it reduces the shared output. This improves throughput by 3~5%.

Test Plan

lm_eval with gsm8k on MI355X

Server command

VLLM_ROCM_USE_AITER=1 \
vllm serve moonshotai/Kimi-K3 \
  --tensor-parallel-size 8 \
  --trust-remote-code \
  --kv-cache-dtype fp8 \
  --max-model-len 16384 \
  --max-num-seqs 128 \
  --max-num-batched-tokens 16384 \
  --gpu-memory-utilization 0.93 \
  --block-size 128 \
  --no-enable-prefix-caching \
  --compilation-config '{"custom_ops":["+fused_rms_norm_gated"]}' \
  --enable-auto-tool-choice \
  --tool-call-parser kimi_k3 \
  --reasoning-parser kimi_k3

Benchmark command

vllm bench serve --model moonshotai/Kimi-K3  --trust-remote-code  --dataset-name random --random-input-len $ISL --random-output-len $OSL --ignore-eos --num-prompts $PROMPTS --max-concurrency $CONC

lm_eval command

lm_eval --model local-completions --model_args model=moonshotai/Kimi-K3,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False,trust_remote_code=True  --tasks gsm8k --num_fewshot 5 --seed 42

Test Result

lm_eval gsm8k (PR)

20-shot

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	20	exact_match	↑	0.9636	±	0.0052
		strict-match	20	exact_match	↑	0.9636	±	0.0052

E2E performance on MI355X

ISL/OSL	Conc	Config	Total tokens/s	Delta (%)	Total reqs/s	Delta (%)	Output tokens/s	Delta (%)	mean TTFT (ms)	Delta (%)	mean TPOT (ms)	Delta (%)	mean ITL (ms)	Delta (%)
8k/1k	8	main	2273.22		0.247		252.58		1888.15		29.848		29.848
8k/1k	8	PR	2384.89	+4.91	0.259	+4.86	264.99	+4.91	1783.93	-5.52	28.465	-4.63	28.465	-4.63
8k/1k	16	main	3721.83		0.404		413.54		3562.22		35.217		35.217
8k/1k	16	PR	3870.19	+3.99	0.420	+3.96	430.02	+3.99	3366.21	-5.50	33.926	-3.67	33.926	-3.67
8k/1k	32	main	5337.48		0.579		593.05		4975.48		49.088		49.088
8k/1k	32	PR	5522.93	+3.47	0.599	+3.45	613.66	+3.48	4717.98	-5.18	47.527	-3.18	47.527	-3.18
8k/1k	64	main	6960.93		0.755		773.44		7304.13		75.552		75.552
8k/1k	64	PR	7173.25	+3.05	0.778	+3.05	797.03	+3.05	6929.49	-5.13	73.469	-2.76	73.469	-2.76
8k/1k	128	main	8479.18		0.920		942.13		11767.46		124.15		124.15
8k/1k	128	PR	8760.63	+3.32	0.951	+3.37	973.40	+3.32	11161.10	-5.15	120.38	-3.04	120.38	-3.04

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[tjtanaa]

LGTM

[github-actions]

✅ @kliuae, CI is now available for this PR.

• /ci run starts a CI build.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.

[tjtanaa]

/ci run

[github-actions]

✅ Triggered Buildkite CI #82784 for commit 6afeacdcffe7.