[Kimi-K3] Add option to shard the shared expert instead of replicating

[tlrmchlsmth]

Shard the shared-expert weights to avoid redundant work. In order to do so we need to re-distribute the activations, which adds an extra AllGather and ReduceScatter.

This has been tested e2e to work and speed up decodes. So far, preliminary perf measurements show this is a decode win, and for prefill it's a throughput vs kv cache size tradeoff.

FYI the performance and memory footprint improvements are overstated since all of the numbers are from running @mgoin's 75% expert-sparse Kimi K3

Generated stuff below

---

Problem

Under sequence-parallel MoE, KimiMLP threads use_sequence_parallel straight through as disable_tp on both projections (model.py), so the dense and shared-expert MLPs are replicated on every rank. Each rank streams the entire weight to serve only its own token shard.

For Kimi-K3's shared experts (hidden 7168, shared intermediate 6144):

gate_up  7168 x 12288 x 2B = 176.2 MB
down     6144 x  7168 x 2B =  88.1 MB
                           = 264.2 MB per layer per rank, replicated
x ~92 MoE layers           = ~22.6 GiB per rank

The routed experts are EP-sharded and the latent projections are deliberately replicated, but the shared experts are sharded by nothing — they sit outside the EP machinery, and disable_tp removes the only sharding they had.

This is visible from two independent directions:

1. Profiling. On a decode step the shared-expert GEMM pair is the second-largest line item after the MegaMoE routed GEMM. The kernels are healthy (~5.9 TB/s); the problem is byte count, not efficiency.
2. Memory budgets. A deployment spec whose weight budget assumed dense/TP + experts/(TP*DP) under-predicted actual per-GPU usage by ~28 GiB, of which the replicated shared experts account for ~23 GiB.

Corroborating detail: KIMI_K3_PROJECTIONS in low_latency_gemm.py contains shared_gate_up_proj / shared_down_proj entries only at TP-sharded shapes (TP4/TP8/TP16). The replicated shapes have no entry, so under SP they fall through to cuBLAS — the table was tuned against the non-SP path, where these layers are TP-sharded.

Note this replicate-under-SP pattern is a vLLM-wide convention (DeepSeek-V4 does the same, with an explicit comment). Kimi-K3 is an outlier only because its shared expert is unusually large.

Change

Add VLLM_KIMI_K3_SHARD_SP_MLP (default off). Under the flag, KimiMLP TP-shards both projections and:

if self.shard_sequence_parallel:
    x = sp_all_gather(x)          # each rank owns disjoint tokens; gather the full set
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(x)          # RowParallel, partial over intermediate
if self.shard_sequence_parallel:
    x = sp_reduce_scatter(x)      # sums across TP and restores sequence sharding

The reduce-scatter does double duty, so the block still ends with one collective per direction rather than an all-reduce plus a re-shard.

Deliberately not enabled on the FusedMoE path, which hands the shared experts to the runner and assumes the replicated layout.

Trade-off and why it is opt-in

This trades weight bandwidth and resident memory for two collectives per layer. That only wins at low token counts, so it is off by default and intended for decode instances under P/D disaggregation, which can commit to the sharded layout and never see prefill-sized batches. Prefill instances should leave it off.

One known limitation: the flag is global, but prefill and decode want opposite answers. In a P/D deployment it can be scoped per role via the env var. In an aggregated deployment a single flag cannot be right for both — the dense MLP would shard on prefill, on the losing side of the crossover.

Testing

pytest tests/models/kimi_k3/test_sequence_parallel.py     # 20 passed
pre-commit run --files vllm/models/kimi_k3/nvidia/model.py vllm/envs.py \
    tests/models/kimi_k3/test_sequence_parallel.py         # all hooks pass

The added test models disjoint token shards per rank, which is what makes the failure mode visible: an earlier version of this change sharded both projections over their output dim and all-gathered the feature slices, on the reasoning that no reduction was then needed. That is wrong under sequence parallelism — each rank computes its intermediate chunk for its own tokens, so no rank ever holds every chunk for any single token. It ran cleanly and produced degenerate output. A test that models every rank holding the same tokens passes against that broken design; this one does not.

Evaluation

Measured on 4x GB200 nodes (16 GPUs), Kimi-K3 pruned checkpoint, P/D disaggregated, decode role TP4/DP2/EP8 with the deep_gemm_mega_moe backend, max_num_seqs=128. Load: 64 concurrent, 9k ISL / 256 OSL, 7 min. Torch profiler on one DP rank, 5 iterations.

	flag off	flag on	delta
Weights per rank	94.62 GiB	76.71 GiB	−17.91 GiB
GPU KV cache	2,988,288 tok	3,679,187 tok	+23.1%
Decode kernel time	27,881 us/iter	25,052 us/iter	−10.1%
ITL p50	24.49 ms	21.86 ms	−10.7%
ITL p99	39.88 ms	38.15 ms	−4.4%
Output tok/s	925.9	902.2	−2.6%
TTFT p50	8777 ms	9713 ms	+10.7%
GSM8K (1300 problems)	~40%	40.7%	—
Errors	0	0	—

In the profile the replicated shared gate_up kernel (29.64 us, n=460) disappears entirely, and the lamport all-gather / reduce-scatter counts double (n=470 -> 935, n=465 -> 930) exactly as the design implies.

Honest reading of these numbers:

• The memory result is the main benefit and is solid: −17.9 GiB per rank converts directly into +23% KV cache at gpu_memory_utilization=0.95.
• The decode-latency result is real: kernel time −10.1% and ITL −10.7% are independent measurements agreeing to 0.6%.
• End-to-end throughput did not improve (−2.6%). TTFT moved +10.7%, which this change cannot cause (it is decode-scoped), so that is either run-to-run variance or a scheduling effect of the larger KV cache. The deployment appears prefill-bound at this operating point.
• One run per arm, several hours apart. The −2.6% / +10.7% deltas are within plausible run-to-run variance; the ITL number is trusted because the trace corroborates it independently.
• Measured at TP4 on a pruned checkpoint from a development build, so these numbers do not describe a shipped Kimi-K3 configuration. Only the flag-on/flag-off delta is meaningful.

Not done: no congestion-terminated concurrency sweep, so the saturation knee is unlocated; no measurement at the batch sizes where replication should win, so the crossover is estimated (~700 tokens at TP8) rather than measured.

Not a duplicate

Checked before opening:

gh pr list --repo vllm-project/vllm --state open --search "kimi k3 shared expert"
gh pr list --repo vllm-project/vllm --state open --search "shared experts sequence parallel"
gh pr list --repo vllm-project/vllm --state open --search "sequence parallel MoE tensor parallel"
gh issue list --repo vllm-project/vllm --search "shared expert replicated"

The only open Kimi-K3 PRs are #50404 (MLA with disabled context parallelism) and #50319 (ROCm gfx942); neither touches the MoE block or sequence-parallel weight layout. Nothing open addresses shared-expert replication under SP.

AI assistance

This change was developed with AI assistance (Claude). The analysis, implementation, and the cluster measurements above were produced with it. A human submitter has reviewed the diff and is responsible for defending it.

[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.

[jeejeelee]

NIT

Suggested change

	"VLLM_KIMI_K3_SHARD_SP_MLP": lambda: bool(
	"VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT": lambda: bool(

[tlrmchlsmth]

much better, agreed will change it :)

[jeejeelee]

QQ：How about combining this RS and the sequential AG into one AR?

[WoosukKwon]

How's the perf compared to pure TP? 👀

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @tlrmchlsmth.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[tlrmchlsmth]

ugh - sorry for the slop

[tlrmchlsmth]

Need to follow this up with support for the non-mega-moe case.
For DP+EP we may also want to add the ability to do this only within a node for the cross-node TP case (but that will be annoying since we'll need to create a new parallel group).

[tlrmchlsmth]

Some results running Kimi K3 on GB200 DP8xTP4 using the DecodeBenchConnector + uniform random expert routing:

Replicated:

(EngineCore_DP1 pid=369) INFO 08-02 15:48:44 [kv_cache_utils.py:2235] GPU KV cache size: 2,913,390 tokens

TP-sharded:

(EngineCore_DP1 pid=374) INFO 08-02 18:46:07 [kv_cache_utils.py:2235] GPU KV cache size: 3,605,817 tokens

Pure decode performance on GB200 DP8xTP4 Kimi K3 + DecodeBenchConnector:

concurrency	Replicated Out Tok/s	Sharded Out Tok/s	Speedup
64	1824.5075	2143.2081	0.1746776048
128	3451.477	3768.6133	0.09188422811
256	5661.4198	5940.4618	0.04928834283
512	8432.3926	8731.1662	0.03543165198
768	10249.9561	10624.3745	0.03652878084
1024	11199.3843	11880.0562	0.06077761793
1280	10931.4287	11079.1001	0.01350888379
1536	10658.8695	10634.9697	-0.002242245296