[Kimi-K3] AG-GEMM for Sequence Parallelism

[gau-nernst]

Purpose

This PR implements overlapping AG-GEMM with CUDA streams pipelining: compute local shard while sending the local activation shard to neighboring remote rank, and continue the rotation until all shards are computed. No custom kernel required, and technically we can adapt to any quantization kernels in the future.

Implementation details

• Inter-GPU transfer is done in push mode using cuMemcpyDtoDAsync -> CopyEngine-based, no SM usage. Hence, this only requires P2P, no multicast is needed.
• Synchronization is done via PyTorch's built-in handle.put_signal(next_rank, channel=channel) and handle.wait_signal(previous_rank, channel=step).
• We purposely don't have a cross-rank exit barrier, which is meant to prevent next invocation overrides recv buffers that previous GEMM might not have finished. Instead, we rely on the fact that there will be another collective between 2 consecutive AG-GEMM invocations acting as an implicit barrier.
• For data push, we don't issue all push at once in the beginning. Instead, we only push to the receiver of the current round, and wait for the expected incoming data to arrive before pushing to the next receiver. In microbenchmarks, this design seems to be a bit faster, likely because we avoid transport contention.

vLLM integration

The AG-GEMM logic alone is very straight forward: ag_gemm.py only has 150 LOC. The rest of the changes target some refactoring of SP collective ownership and initialization logic.

• KDA/MLA/MLP now owns SP collectives, regardless of whether the fused AG-GEMM/GEMM-RS are activated. This is more natural as the fusion boundary is within the same module, not leaking outside to DecoderLayer. When fused AG-GEMM and GEMM-RS are not used (either because it's not supported or disabled), KDA/MLA/MLP modules are now responsible to invoke the unfused AG/RS collectives.
  • For Kimi-Linear, DecoderLayer still executes SP collectives
• GEMM-RS is also enabled by default now. Previously I left it off by default in [Kimi-K3] Add GEMM-RS for sequence parallelism #52079 because I don't know a reliable way to determine whether multicast memory is available. However, the check is actually quite simple: see if multicast_ptr is not None [Kimi-K3] Extend GEMM-RS to GEMM-AR #53053 (thanks @wzhao18)
• This PR also separate GEMM-RS and GEMM-AR. Previously [Kimi-K3] Extend GEMM-RS to GEMM-AR #53053 combines them together since they share the same kernel code. However, separating them (i.e. initialization logic, dispatch logic) improves readability and logical flow: use_sequence_parallel now implies AG-GEMM and GEMM-RS, while run_gemm_ar controls GEMM-AR only
  • In the future, we may want to integrate GEMM-AR directly into RowParallelLinear so all models can benefit from it, and we don't need manual dispatch for each model.

Microbenchmark results

Using benchmarks/kernels/benchmark_kimi_k3_ag_gemm.py, TP8, GB300

Projection	Global M	Local M	N	K	AG + GEMM	AG-GEMM	Speedup
KDA	2,048	256	6,288	7,168	178.30 µs	202.21 µs	0.882×
KDA	3,072	384	6,288	7,168	290.51 µs	231.57 µs	1.255×
KDA	3,584	448	6,288	7,168	318.24 µs	246.77 µs	1.290×
KDA	4,096	512	6,288	7,168	351.97 µs	259.68 µs	1.355×
KDA	8,192	1,024	6,288	7,168	592.56 µs	469.36 µs	1.262×
KDA	16,384	2,048	6,288	7,168	1,154.70 µs	866.64 µs	1.332×
KDA	32,768	4,096	6,288	7,168	2,261.60 µs	1,747.81 µs	1.294×
MLA	2,048	256	3,648	7,168	142.30 µs	194.10 µs	0.733×
MLA	3,072	384	3,648	7,168	236.96 µs	220.64 µs	1.074×
MLA	3,584	448	3,648	7,168	245.15 µs	258.02 µs	0.950×
MLA	4,096	512	3,648	7,168	267.52 µs	244.05 µs	1.096×
MLA	8,192	1,024	3,648	7,168	476.66 µs	401.82 µs	1.186×
MLA	16,384	2,048	3,648	7,168	832.72 µs	684.16 µs	1.217×
MLA	32,768	4,096	3,648	7,168	1,628.88 µs	1,055.84 µs	1.543×

Component breakdown

Projection	Global M	Local M	N	K	GEMM only	AG only	AG-GEMM
KDA	2,048	256	6,288	7,168	111.31 µs	76.80 µs	202.21 µs
KDA	3,072	384	6,288	7,168	161.15 µs	142.32 µs	231.57 µs
KDA	3,584	448	6,288	7,168	185.33 µs	145.26 µs	246.77 µs
KDA	4,096	512	6,288	7,168	215.02 µs	147.26 µs	259.68 µs
KDA	8,192	1,024	6,288	7,168	399.31 µs	203.81 µs	469.36 µs
KDA	16,384	2,048	6,288	7,168	801.20 µs	362.53 µs	866.64 µs
KDA	32,768	4,096	6,288	7,168	1,576.83 µs	695.84 µs	1,747.81 µs
MLA	2,048	256	3,648	7,168	76.56 µs	76.18 µs	194.10 µs
MLA	3,072	384	3,648	7,168	106.78 µs	141.90 µs	220.64 µs
MLA	3,584	448	3,648	7,168	112.91 µs	144.18 µs	258.02 µs
MLA	4,096	512	3,648	7,168	130.48 µs	149.70 µs	244.05 µs
MLA	8,192	1,024	3,648	7,168	276.61 µs	205.82 µs	401.82 µs
MLA	16,384	2,048	3,648	7,168	478.93 µs	363.95 µs	684.16 µs
MLA	32,768	4,096	3,648	7,168	950.88 µs	692.45 µs	1,055.84 µs

At very large bs, the overlap is ideal, which makes sense because our GEMM efficiency is determined by local GEMM shape (M/8). Future work can improve performance for medium M by having a single GEMM kernel that perform in-kernel waiting.

E2E perf

VLLM_ALLREDUCE_USE_FLASHINFER: "1"
VLLM_USE_RUST_FRONTEND: "1"
VLLM_ENGINE_READY_TIMEOUT_S: "3600"

vllm serve {model}
--port {port}
-tp 8 -ep --moe-backend deep_gemm_mega_moe
--nnodes {nnodes}
--node-rank {node_rank}
--master-addr {master_addr}
--trust-remote-code
--load-format fastsafetensors
--no-enable-prefix-caching
--kv-cache-dtype fp8
--attention-config '{"mla_prefill_backend":"flashinfer","use_prefill_query_quantization":true}'
--reasoning-parser kimi_k3
--tool-call-parser kimi_k3
--enable-auto-tool-choice

VLLM_USE_RUST_BENCH: "1"

vllm bench serve
--backend openai-chat
--base-url http://{head_addr}:8000
--model {model}
--dataset-name speed-bench
--speed-bench-config throughput_8k
--speed-bench-max-input-len 8192
--speed-bench-output-len 1024
--ignore-eos
--num-warmups 5
--sweep-max-concurrency 1,4,16
--sweep-num-prompts-factor 10
--percentile-metrics "ttft,tpot,itl,e2el"
--metric-percentiles 50,90,99
--result-dir "{log_dir}/results"
--save-result

Baseline is 3bb19cd

Mixed prefill-decode (aggregated serving)

8k-1k

Prefill-only (PD serving)

8k-1

Test Plan

TP8

• GSM8K: 0.9651
• OCRBench: 89.6

TEP8+SP (MegaMoE)

• GSM8K: 0.9651
• OCRBench: 89.5

Test Result

---

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.

[gcanlin]

TEP8+SP (MegaMoE)

Hi, may I ask which moe backend do you plan to use? I recently found that DeepGEMM mega moe is slower than default FI non-mega moe.(But notice that I run mega moe on b200 by TP8 + EP8 + PP2). So I'm a bit curious whether it's also true on GB300 by TP8 + EP8. From my experience before, mega moe should be faster than non-mega.

[gau-nernst]

When --moe-backend deep_gemm_mega_moe is set, together with -tp 8 -ep, vLLM will enable sequence parallelism for Kimi-K3 automatically. This is not the case for TRTLLM MoE backend (default). So direct comparison might not be appropriate. Which exact configuration are you comparing?

Anyway, TEP8+SP with MegaMoE is meant for prefill. Its decode performance might not be competitive.

[mergify]

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

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