[Kimi K3][Kernel] Fuse BF16 shared experts into latent MegaMoE tail

[gcanlin]

Summary

This PR adds a Kimi K3-specific SM100 path that fuses its BF16 shared expert with FP8/FP4 DeepGEMM MegaMoE and removes the standalone shared-output TP reduction from the latent-MoE tail.

Kimi K3 cannot use the same final accumulation as DeepSeek V4 in #53040: routed MegaMoE produces a 3584-wide latent tensor, while the shared expert produces a 7168-wide hidden tensor. The routed result still needs RMSNorm and a 3584 -> 7168 up-projection before the two branches can be added. The non-sequence-parallel path therefore uses two kernels:

1. DeepGEMM schedules routed FP8/FP4 MegaMoE, BF16 shared L1/L2, and routed RMSNorm together, then destination-scatters shared L2 fragments into a symmetric workspace.
2. The Kimi K3 CuTe DSL tail computes the routed up-projection, reduces the published shared fragments in FP32, adds them in the epilogue, and multicasts the final BF16 output.

This update also supports sequence-sharded shared experts. Routed tokens remain SP-local, while the BF16 shared branch consumes the gathered token set. DeepGEMM publishes each source rank's full-hidden partial directly to the owning token rank in [local_token, source_rank, hidden] order. A local Triton reduction followed by BF16 addmm sums the TP partials and combines them with the routed up-projection.

Sequence-parallel implementation

• DeepGEMM routed and shared work may use independent token counts: routed M is SP-local and shared M is the gathered global token count.
• fp8_fp4_mega_moe_bf16_shared_sp_rs publishes each TP intermediate shard to the destination token owner without materializing a full shared output or running the former shared-output sp_reduce_scatter.
• The shared-input all-gather overlaps the routed down-projection for large prefill batches.
• On TP8 B200, local M >= 256 uses a targeted NCCL NVLS symmetric-memory all-gather. It reuses one output scratch buffer sized for the maximum local token count (about 448 MiB/rank for M=4096, TP8, H=7168).
• The targeted API does not enable symmetric memory globally. TP16, small batches, decode-sized batches, batch-invariant mode, unsupported NCCL versions, and unavailable allocators retain the existing SP collective.
• Unsupported DeepGEMM APIs, layouts, devices, and topologies fail closed to the existing shared-expert path. VLLM_DISABLE_KIMI_K3_MEGAMOE_SHARED_EXPERT_FUSION=1 remains an emergency rollback.

PR dependency and benchmark stack

This PR does not merge #54347. The new commit is directly based on this PR's previous head. #54347 independently enables Kimi K3 model-level SP across a PP boundary. The TP8 x PP2 x EP8 measurements below were collected from a temporary benchmark branch with this PR stacked on #54347; after #54347 lands, this PR's SP-sharded shared path becomes available for that production topology. Without #54347, the existing non-SP behavior and supported SP topology remain unchanged.

DeepGEMM dependency: deepseek-ai/DeepGEMM#416.

Test Plan

Serving setup

The matched tests used two nodes with 8 NVIDIA B200 GPUs per node, TP8 x PP2 x EP8, max_num_seqs=128, max_num_batched_tokens=32768, exact random lengths, seed=2026, greedy decoding, four warmups, and FULL_DECODE_ONLY CUDA Graphs.

VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT=1 \
VLLM_KIMI_K3_GEMM_RS=1 \
PYTHONPATH=/path/to/DeepGEMM \
vllm serve /path/to/Kimi-K3 \
  --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"}' \
  --tensor-parallel-size 8 \
  --pipeline-parallel-size 2 \
  --nnodes 2 \
  --node-rank "$NODE_RANK" \
  --master-addr "$MASTER_ADDR" \
  --master-port 29992 \
  --distributed-executor-backend mp \
  --enable-expert-parallel \
  --all2all-backend deepep_v2 \
  --decode-context-parallel-size 1 \
  --dcp-comm-backend a2a \
  --kernel-config '{"moe_backend":"deep_gemm_mega_moe"}' \
  --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"}'

The SHARD=1 baseline adds only VLLM_DISABLE_KIMI_K3_MEGAMOE_SHARED_EXPERT_FUSION=1. All other server and benchmark arguments are identical.

vllm bench serve \
  --backend openai-chat \
  --base-url http://127.0.0.1:8692 \
  --endpoint /v1/chat/completions \
  --model Kimi-K3 \
  --tokenizer /path/to/Kimi-K3 \
  --trust-remote-code \
  --dataset-name random \
  --random-input-len "$INPUT_LEN" \
  --random-output-len "$OUTPUT_LEN" \
  --random-range-ratio 0 \
  --num-prompts "$NUM_PROMPTS" \
  --num-warmups 4 \
  --request-rate inf \
  --max-concurrency "$CONCURRENCY" \
  --ignore-eos \
  --temperature 0 \
  --seed 2026

Test Results

All requests completed successfully. Baseline and fused runs used the same node pair, image, model, scheduler limits, and request stream.

SP-sharded prefill: 8192 input / 1 output, concurrency 8, 128 requests

Run	SHARD=1 baseline total tok/s	Fused + targeted gather total tok/s	Change
1	53,543.40	54,523.92	+1.83%
2	53,084.73	54,776.67	+3.19%
Mean	53,314.06	54,650.29	+2.51%

Mean median TTFT decreases from 1230.43 ms to 1186.76 ms (-3.55%). The throughput range within each arm is below 1%.

SP-sharded decode: 128 input / 1024 output, concurrency 32, 32 requests

Run	SHARD=1 baseline output tok/s	Fused output tok/s
1	780.46	773.91
2	783.86	799.72
3	-	800.06

Including the first cold decode run, fused throughput averages 791.23 tok/s versus the baseline's 782.16 tok/s (+1.16%). Comparing the second steady-state run gives +2.02%, and the third fused run reproduces 800.06 tok/s. The targeted NVLS gather is not selected for this decode-sized M, so this result measures the shared-fusion/tail path rather than the new gather backend.

Targeted gather and overlap microbenchmark

TP8 on one B200 node; times are microseconds. Overlap measures shared-input gather running concurrently with routed down-projection.

Local M	Existing SP gather	Targeted NVLS gather	Existing overlap	Targeted overlap
256	66.144	52.807	100.859	66.382
512	131.222	94.274	161.784	97.987
1024	193.851	178.330	242.483	182.370
2048	364.418	320.485	437.690	315.810
4096	719.214	584.504	843.550	586.877

At local M=4096, the overlapped segment decreases by 30.4%. Below M=256 the existing lower-overhead SP gather remains faster, which is why the targeted path is thresholded.

SP-local tail microbenchmark

The current two-kernel tail consists of a Triton source-rank reduction and BF16 addmm with the routed 3584 -> 7168 up-projection.

Local M	Complete tail	Reduction only	addmm only
1	32.137 us	17.006 us	19.290 us
32	30.861 us	16.700 us	18.859 us
128	31.166 us	16.780 us	18.876 us

DeepGEMM producer microbenchmark

The TP8 producer benchmark uses Kimi K3 matrix widths (routed H/I=3584/3072, shared H/I per TP rank=7168/768), local routed M=8, global shared M=64, and 100 measured iterations. It takes the maximum rank time.

Run	Dense shared output	SP token-owner publication	Change
1	42.487 us	42.944 us	+1.08%
2	42.368 us	43.002 us	+1.50%
Mean	42.428 us	42.973 us	+1.29%

The direct remote publication overhead is about 0.55 us; it replaces downstream shared-output materialization and SP reduction.

Original non-SP results

The original TP8 x PP2 published-tail path remains covered by the previous measurements:

Concurrency, 128 input / 512 output	Main/native	Fused	Change
16	435.04 tok/s	470.62 tok/s	+8.18%
32	745.05 tok/s	767.16 tok/s	+2.97%
64	1,277.16 tok/s	1,335.88 tok/s	+4.60%

Correctness

• vLLM SP gather backend selection: 4 parameterized cases passed (TP8 threshold, below threshold, TP16 fallback, and batch-invariant fallback).
• DeepGEMM single-GPU dense/publication bitwise regression passed.
• The TP8 producer harness generated a different shared partial on every source rank and compared every destination workspace [local_token, source_rank, hidden] value bitwise against dense output; it passed twice with Kimi K3 matrix widths.
• The TP8 vLLM tail/gather harness passed routed up-projection, published-workspace reduction, and targeted-gather correctness.
• Ruff lint/format, mypy, vLLM pre-commit hooks, and git diff --check passed.

Accuracy

The original fused non-SP path was evaluated on the complete GSM8K v3 test split with lm-eval 0.4.12: 1,319 samples, 5-shot chat-template evaluation, multi-turn few-shot examples, greedy decoding, and a 1,500-token generation limit. The new SP-sharded extension is covered by the TP8 bitwise and serving regressions above; this full-dataset run was not repeated for the temporary #54347 benchmark stack.

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9682	±	0.0048
		strict-match	5	exact_match	↑	0.9689	±	0.0048

The exact scores are 1277/1319 and 1278/1319 respectively.

Contribution notes

AI assistance was used during implementation, debugging, review, benchmark analysis, and preparation of this PR description. The human submitter must review and understand every changed line and verify the final commands/results before submission.

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

[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

[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