Opt-in two-stream + launch-elision latency optimizations for Kimi-K2.6-NVFP4 decode on B300 (TP=4)

[jinhuang12]

Summary

This PR adds three opt-in decode-latency optimizations for nvidia/Kimi-K2.6-NVFP4 running TP=4 on NVIDIA B300 SXM6 (Blackwell, sm_103) with the FlashInfer-MLA attention + FlashInfer-TRTLLM MoE backends, NVFP4 weights / FP8-E4M3 paged KV cache, CUDA graphs (FULL_AND_PIECEWISE) + torch.compile. Each optimization sits behind its own environment variable that defaults to off, so with all three flags unset the code path is unchanged from base — they are guarded fall-throughs, not changes to the default execution path. All changes are pure Python under vllm/; no C++/CUDA is touched, so no rebuild is required.

This branch is rebased onto current main. The optimization commit was developed and measured on vLLM v0.22.1 (commit 0decac0d96c42b49572498019f0a0e3600f50398); the rebase onto main applied with no logic changes (the shared-experts and MoE-runner sites were re-anchored onto upstream's renamed routed_experts.forward_monolithic dispatch API and the SharedExperts.__init__ signature). The numbers below were measured on v0.22.1 and have not been re-validated on main.

This work was validated only on B300 (sm_103) with the FlashInfer-MLA + FlashInfer-TRTLLM MoE backends; no claim is made about other architectures or backends. The optimizations are scoped to those code paths.

Optimizations

• MoE FP4 preamble launch-elision — VLLM_MOE_FP4_PREAMBLE_FUSION=1, lossless (bit-exact). On the FlashInfer-TRTLLM MoE backend, elides two redundant per-MoE-layer launches: (1) the torch.zeros/FillFunctor that pre-zeros the FP4 scale-factor buffer before cvt_fp16_to_fp4_sf_major (the non-swizzled sf_major conversion kernel writes every scale-factor address unconditionally, so the pre-zero is dead work — replaced with torch.empty), and (2) the per-layer e_score_correction_bias.to(bfloat16) routing-bias cast, hoisted to module init so it runs once instead of per decode step. With the flag off, torch.zeros and the per-layer cast are retained, so behavior is bit-identical.

• Two-stream MLA decode pipelining — VLLM_MLA_TWO_STREAM_DECODE=1, lossy in bits, accuracy-preserving. The MLA wrapper allocates an auxiliary CUDA stream + event pair and hands it to MLAAttention via set_aux_stream(). The kv-path tail (concat_and_cache_mla) is dispatched onto the aux stream while the q-path head (the W_UK_T BMM) runs on the default stream; an event record/wait protocol guarantees the cache write completes before the attention kernel reads it. Applied at both the direct-call path and the compiled opaque-op path (unified_mla_kv_cache_update / unified_mla_attention_with_output). With the flag off the aux stream is None and dispatch falls through sequentially with no behavioral change.

• Event-based MoE shared/routed two-stream overlap — VLLM_MOE_SHARED_EXPERTS_TWO_STREAM=1, lossy in bits, accuracy-preserving. Replaces the existing Stream.wait_stream-based shared-expert overlap — which collapses to sequential execution under CUDA-graph capture (all MoE kernels land on a single stream) — with a torch.cuda.Event record/wait pattern (per-microbatch for DBO). The shared-expert layer runs on a side stream while the routed FP4 BMMs run on the default stream; the shared experts are launched early (before the routed BMMs) when MULTI_STREAM_OVERLAPPED is selected so the routed work is captured as parallel in the CUDA-graph DAG, then the default stream joins. With the flag off the legacy wait_stream ordering is bit-identical.

Fixed-batch latency

vllm bench latency on nvidia/Kimi-K2.6-NVFP4, 4× B300, TP=4 / DP=1, --max-num-seqs=8 --trust-remote-code, OSL=1000, 5 iters per shape, CUDA graphs (FULL_AND_PIECEWISE) + torch.compile enabled. A = base (vLLM v0.22.1, all flags unset); B = all three flags enabled. Both arms run the same base commit 0decac0d9 with an identical workload; the only difference is the patched files (A imports a clean tree, B imports the patched tree with the three env flags set). The grid sweeps input length × batch size in a single model load per arm. OTPS = output tokens/sec; TPOT = time per output token (both decode-path, from per-request metrics).

Input len	Batch	A E2E (s)	B E2E (s)	E2E Δ	A OTPS	B OTPS	A TPOT (ms)	B TPOT (ms)
2000	1	6.117	5.948	−2.8%	166.2	171.1	6.02	5.85
2000	8	9.250	8.691	−6.0%	892.7	957.0	8.96	8.36
10000	1	6.540	6.379	−2.5%	159.1	163.7	6.29	6.11
10000	4	8.750	8.464	−3.3%	492.1	510.1	8.13	7.84
10000	8	11.074	10.634	−4.0%	801.4	838.7	9.98	9.54

Every shape improves on E2E, OTPS, and TPOT; the E2E gain ranges −2.5% to −6.0%. The improvement is decode-dominated (decode is 93–97% of E2E at OSL=1000), consistent with the two-stream MLA/MoE overlap and the per-MoE-layer launch-elision acting on the decode path. On significance: using a two-sample difference-of-means test (n=5 per arm), all five shapes' E2E deltas exceed twice the standard error of the difference. Under the stricter single-arm test (delta vs twice the baseline's per-iteration standard deviation), four of five shapes clear it; the 10000/bs1 shape (Δ=0.161 s vs 2σ_A=0.181 s) is the one exception — its E2E gain is within the baseline's single-run spread, though its OTPS/TPOT move in the same direction as the others. Prefill latency deltas are small and mixed across shapes (within their own per-iteration spread), so the end-to-end gains come from decode. No shape regresses on any metric.

Correctness

GSM8K (greedy decoding, full 1319-question test split), B300, vLLM v0.22.1, A/B against the base config under the production-parity environment:

• Base 92.95% (1226/1319); all three flags enabled 93.10% (1228/1319); delta +0.15pp (within the 1.0pp gate).

The preamble launch-elision is bit-exact (no accuracy risk); the two stream-overlap changes are not bit-exact but introduce only floating-point order-of-add churn from running independent op chains on separate streams, with accuracy held within tolerance.

Changes

7 files changed, +426 / −29, all Python under vllm/ — no C++/CUDA, no rebuild required.

• Flag registration: vllm/envs.py (three env vars, each defaulting to "0"/off).
• MoE FP4 preamble launch-elision: vllm/_custom_ops.py, vllm/model_executor/models/deepseek_v2.py.
• Two-stream MLA pipelining: vllm/model_executor/layers/mla.py, vllm/model_executor/layers/attention/mla_attention.py.
• Event-based MoE shared/routed overlap: vllm/model_executor/layers/fused_moe/runner/shared_experts.py, vllm/model_executor/layers/fused_moe/runner/moe_runner.py.

AI-assisted contribution

These optimizations were produced by an automated process driving Claude (Anthropic) under human review: each candidate was profiled, implemented behind an opt-in flag, and kept only after passing a correctness check (a bit-exact / torch.equal argument for the lossless change, a before/after GSM8K A/B for the two stream-overlap changes) and a before/after end-to-end latency benchmark. All numbers in this PR were measured on the real B300 deployment hardware with CUDA graphs and torch.compile enabled, and a human reviewed every change before submission. With the flags unset the default execution path is unchanged.

[mergify]

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

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