[Kimi K3] Support sequence parallelism with pipeline parallelism

[gcanlin]

Purpose

Enable Kimi K3 model-level sequence parallelism (SP) when pipeline parallelism (PP) is also enabled. This lets the DeepGEMM MegaMoE path run with the production TP8/PP2/EP8 topology instead of disabling SP whenever PP > 1.

The implementation keeps tokens sequence-sharded across a PP boundary:

• Select SP for Kimi K3 when EP is enabled, TP > 1, and the MoE backend is deep_gemm_mega_moe, including PP > 1. Non-MegaMoE behavior is unchanged.
• The first PP stage shards the embedded tokens. Later PP stages consume the already-local hidden_states and residual rows instead of sharding them a second time.
• Add a model-specific PP transport override for hidden_states and residual. The default PP optimization assumes TP-replicated tensors and reconstructs them with a TP all-gather on the receiver. SP tensors are distinct across TP ranks, so matching TP ranks now send their shards directly without that all-gather.
• Preserve the received local leading dimension in the V2 and legacy GPU model runners. Full-sized persistent buffers are still used for profiling and CUDA graph capture, then narrowed exactly once inside the model.
• Record the non-last-PP output row count per CUDA graph descriptor so replay returns an SP-local IntermediateTensors view.
• Delegate the transport override through the multimodal Kimi K3 wrapper and add focused tests for backend selection, PP-stage shard handling, and PP transport behavior.

Performance-sensitive details

• Attention still gathers the full token sequence at the start of each layer; the O-projection/GEMM-RS returns the local shard. MoE, shared-expert work, and PP activation transport remain sequence-sharded.
• Avoiding full-token reconstruction at the PP boundary removes an otherwise redundant TP all-gather and prevents a second SP split on the receiving stage.
• CUDA graph and runtime PP buffers use the actual local row count. Sending a full capture buffer would erase the communication benefit and break replay shapes.
• The first unseen short-prompt concurrency shape can incur a one-time 6--10-second MegaMoE preparation cost. The benchmark below warms at least one complete target-concurrency wave before collecting three measured waves.

Test Plan

Unit tests

pytest -q tests/models/kimi_k3/test_sequence_parallel.py \
  -k "not sharded_sequence_parallel_mlp_matches_replicated"

Result: 30 passed, 1 deselected. The nine newly added parameterized cases all pass.

Two-node serving

Run the following command on both 8xB200 nodes. Set NODE_RANK=0 on the first node and NODE_RANK=1 plus HEADLESS=--headless on the second node.

export MODEL=/path/to/Kimi-K3
export MASTER_ADDR=<rank-0-host>
export HOST_IP=<this-host>
export NODE_RANK=<0-or-1>
export HEADLESS=<empty-or---headless>

PYTHONPATH=/path/to/DeepGEMM \
VLLM_HOST_IP="$HOST_IP" \
VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT=1 \
VLLM_KIMI_K3_GEMM_RS=1 \
VLLM_USE_RUST_FRONTEND=1 \
VLLM_USE_V2_MODEL_RUNNER=1 \
vllm serve "$MODEL" \
  --served-model-name Kimi-K3 \
  --trust-remote-code \
  --max-model-len 1048576 \
  --gpu-memory-utilization 0.84 \
  --load-format fastsafetensors \
  --kv-cache-dtype fp8 \
  --attention-config '{"use_prefill_query_quantization":true,"mla_prefill_backend":"flashinfer"}' \
  --no-disable-hybrid-kv-cache-manager \
  --no-enable-flashinfer-autotune \
  --host 0.0.0.0 \
  --port 8691 \
  --tensor-parallel-size 8 \
  --pipeline-parallel-size 2 \
  --nnodes 2 \
  --node-rank "$NODE_RANK" \
  --master-addr "$MASTER_ADDR" \
  --master-port 29991 \
  --distributed-executor-backend mp \
  --enable-expert-parallel \
  --all2all-backend deepep_v2 \
  --decode-context-parallel-size 1 \
  --dcp-comm-backend a2a \
  --no-enable-prefix-caching \
  --max-num-seqs 128 \
  --max-num-batched-tokens 32768 \
  --max-cudagraph-capture-size 128 \
  --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}' \
  --kernel-config '{"moe_backend":"deep_gemm_mega_moe"}' \
  $HEADLESS

For the FI+EP baseline, remove --kernel-config. For the FI/no-EP baseline, also remove --enable-expert-parallel and --all2all-backend deepep_v2. All other arguments remain identical.

Prefill benchmark

Output length is fixed at one token, so total-token throughput is effectively prefill throughput.

run_point() {
  input_len="$1"
  concurrency="$2"
  warmups=4
  if ((concurrency > warmups)); then warmups="$concurrency"; fi

  vllm bench serve \
    --backend openai-chat \
    --base-url http://127.0.0.1:8691 \
    --endpoint /v1/chat/completions \
    --model Kimi-K3 \
    --tokenizer "$MODEL" \
    --trust-remote-code \
    --dataset-name random \
    --random-input-len "$input_len" \
    --random-output-len 1 \
    --random-range-ratio 0 \
    --num-prompts $((concurrency * 3)) \
    --num-warmups "$warmups" \
    --request-rate inf \
    --max-concurrency "$concurrency" \
    --ignore-eos \
    --temperature 0 \
    --percentile-metrics ttft,e2el \
    --metric-percentiles 50,90,95,99 \
    --seed 2026
}

for c in 1 8 32 64 128; do run_point 128 "$c"; done
for c in 1 2 4 8; do run_point 8192 "$c"; done
run_point 32768 1

Test Results

Hardware: 2 nodes x 8 B200, TP8/PP2, EP8 where enabled. Serving limits are max_num_seqs=128 and max_num_batched_tokens=32768. Every arm completed 747 measured requests with zero failures.

Total token throughput

Input	Concurrency	FI/no-EP tok/s	FI+EP tok/s	MegaMoE+SP tok/s	Mega vs FI+EP	Mega vs no-EP
128	1	962	1,025	1,381	+34.7%	+43.5%
128	8	3,529	3,791	4,840	+27.7%	+37.2%
128	32	12,401	12,963	20,346	+57.0%	+64.1%
128	64	19,472	18,044	26,332	+45.9%	+35.2%
128	128	22,291	21,757	31,118	+43.0%	+39.6%
8192	1	20,450	19,769	27,859	+40.9%	+36.2%
8192	2	22,518	21,658	30,634	+41.4%	+36.0%
8192	4	24,734	32,564	46,948	+44.2%	+89.8%
8192	8	35,767	34,268	49,473	+44.4%	+38.3%
32768	1	21,688	21,160	29,259	+38.3%	+34.9%

Mean TTFT

Input	Concurrency	FI/no-EP	FI+EP	MegaMoE+SP	Mega vs FI+EP	Mega vs no-EP
128	1	225.3 ms	211.5 ms	157.0 ms	-25.8%	-30.3%
128	8	459.8 ms	429.2 ms	315.6 ms	-26.5%	-31.4%
128	32	547.2 ms	519.0 ms	336.3 ms	-35.2%	-38.5%
128	64	700.9 ms	743.8 ms	522.3 ms	-29.8%	-25.5%
128	128	1,232.8 ms	1,265.4 ms	884.7 ms	-30.1%	-28.2%
8192	1	404.7 ms	418.7 ms	297.0 ms	-29.1%	-26.6%
8192	2	668.8 ms	695.3 ms	491.6 ms	-29.3%	-26.5%
8192	4	1,254.9 ms	920.0 ms	638.4 ms	-30.6%	-49.1%
8192	8	1,679.9 ms	1,755.4 ms	1,214.6 ms	-30.8%	-27.7%
32768	1	1,541.2 ms	1,579.6 ms	1,142.3 ms	-27.7%	-25.9%

The FI/no-EP 8192/c4 point has a repeatable scheduling/PP performance cliff when four chat-templated prompts cross the 32768-token step budget. A separate six-wave rerun measured 24,143 tok/s versus 24,734 tok/s in the paired three-wave run.

Accuracy results will be added after the full GSM8K evaluation completes.

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

[ZJY0516]

Suggested change

"(TP=%d, PP=%d, MoE backend=%s).",

[ZJY0516]

self.language_model should always has get_pp_intermediate_tensor_all_gather_overrides?

[ZJY0516]

not sure if it's appropriate

[ZJY0516]

I prefer this

output_token_counts = {
    tensor.shape[0]
    for tensor in intermediate_tensors.tensors.values()
}
assert len(output_token_counts) == 1
output_num_tokens = next(iter(output_token_counts))

And we should have a util function for this beause it's used more than once in this PR

[mergify]

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

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

[ZJY0516]

/ci run

[github-actions]

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

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• /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.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[github-actions]

✅ Triggered Buildkite CI #88833 for commit c43e0215f1ad.

[gcanlin]

/ci retry

[github-actions]

✅ The previous CI build is still running: https://buildkite.com/vllm/ci/builds/88833

[gcanlin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88883 for commit c7f52aca17ad.