[WoosukKwon]

Overview

Fuse Beam's TP all-reduce with its sandwich norms and, for TRTLLM-Gen NvFP4 MoE layers, the top-k finalize. TP only; replicated-token EP deferral is in #62073, which builds on this PR. Follow-up to #62074 (merged). The kernel is built ahead of time in csrc/beam/, with no runtime JIT and no CUDA source in the model folder.

Claims

• csrc/beam/allreduce_norm.cu: MNNVL one-shot/two-shot all-reduce with Beam's FP32 norm-residual-norm epilogue and an optional MoE finalize before the reduction (keeping the local BF16 rounding). Exposed as torch.ops._C.beam_allreduce_norm, plus beam_allreduce (two-shot all-reduce only) for the large-batch path, which then runs the Triton boundary kernel.
• csrc/mnnvl_lamport.cuh: the subset of FlashInfer's MNNVL Lamport device helpers (trtllm_mnnvl_allreduce.cuh, v0.7.2, the version vLLM pins) that the kernel uses, vendored with its license header.
• Beam defers the TRTLLM-Gen NvFP4 top-k finalize into the fused all-reduce under TP, using the existing FusedMoEConfig.defer_moe_finalize(). No shared MoE code changes.
• Same speed as the previous JIT build of this kernel, and 6–16% faster at hidden 8192 from 64 tokens up (microbenchmarks below).

Validation

Unit tests on GB200 (4 GPUs):

pytest tests/models/test_beam_construction.py tests/models/test_beam_numerics.py \
  tests/kernels/moe/test_trtllm_nvfp4_moe.py

178 passed, 3 skipped (8-GPU case; two upstream FlashInfer tile-config skips). This covers the 2- and 4-rank all-reduce and MoE finalize numerics.

TP8 across two GB200 nodes (one MNNVL clique): the all-reduce numerics worker from test_beam_numerics.py (batch changes, one-shot/two-shot, large-batch path, CUDA-graph replay) passes at hidden 5120 and 8192 when launched with torchrun --nnodes=2. The 8-rank MoE-finalize case was not run (its test is single-node).

FlashInfer version: these runs used an environment with FlashInfer 0.7.0.post1 installed. trtllm_mnnvl_allreduce.cuh is byte-identical in v0.7.2, and the flashinfer.comm.mnnvl allocator APIs used here are unchanged (v0.7.2 only adds NVML init/shutdown around a capability probe).

Build: _C_stable_libtorch built for SM100 with CUDA 13.2; csrc/beam/allreduce_norm.cu also compiles for sm_90a. Not run on Hopper.

Kernel microbenchmark, GB200, BF16, each op captured 20× in a CUDA graph; µs per call, slowest rank. "JIT" is the previous FlashInfer-JIT build of the same kernel, "AOT" is this one. Fused all-reduce + residual + RMSNorm:

config	1 tok	8 tok	64 tok	256 tok	1024 tok	4096 tok
TP4, hidden 5120 (JIT → AOT)	5.39 → 5.75*	5.48 → 5.43	8.70 → 8.78	11.71 → 11.67	34.06 → 33.71	138.0 → 139.2
TP4, hidden 8192	6.34 → 6.31	6.34 → 6.21	9.97 → 9.34	19.00 → 16.05	64.52 → 54.28	251.2 → 226.8
TP8, hidden 5120	6.53 → 6.48	6.42 → 6.39	9.00 → 8.97	12.89 → 12.88	35.52 → 35.53	140.3 → 140.4
TP8, hidden 8192	7.86 → 7.78	7.98 → 7.91	10.59 → 9.88	19.82 → 16.67	68.27 → 57.19	265.1 → 225.8

TP4 is the median of 3 runs; TP8 is the mean of two runs with the measurement order swapped. *Whichever kernel is measured first in a process reads about 0.4 µs slow at 1 token; with the order reversed, TP4/5120 at 1 token is 5.79 (JIT) → 5.36 (AOT). The kernel code is unchanged: register usage matches, and at hidden 5120 the AOT kernel on the JIT path's workspace matches the JIT kernel within 0.07 µs. The hidden-8192 gap appears only with the JIT path's own workspace sized for 8192 tokens (it vanishes at 4096); its cause was not isolated.

Performance: TODO — end-to-end before/after against #62074 (standard TP all-reduce), TP8 batch-1 decode and 8 × 2K prefill on GB200.

Details

• Why vendor FlashInfer's helpers instead of reusing vLLM's MNNVL Lamport kernels: on GB200, vLLM's Lamport all-gather (csrc/custom_all_gather_reduce_scatter.cuh) is about 10% slower at TP4 and 25–40% slower at TP8 than FlashInfer's one-shot all-reduce, which does the same communication plus a reduction (TP8, 64 tokens × 5120: 16.7 vs 11.9 µs). Its reduce-scatter + all-gather pair is 1.3–1.8× slower than FlashInfer's single-launch two-shot. A likely cause is that vLLM waits on peer ranks one at a time, while FlashInfer polls them all at once; not verified with a patched kernel. Building on vLLM's primitives would have regressed Beam, so the device code is kept as-is.
• Workspace: FlashInfer's McastGPUBuffer still maps the multicast memory. The Lamport protocol, its initialization and its generation flags are vLLM's, so the build no longer depends on FlashInfer's headers or installed version. Torch symmetric memory works too, but measured 0.2–0.3 µs slower per call at 1–4 tokens (even from 8 tokens up). Rounding it to the 512 MB multicast granularity didn't help. Beam's GB200 path already requires FlashInfer for its MoE.
• Build: compiled only for SM90/SM100 (VLLM_ENABLE_BEAM_ALLREDUCE), with --fmad=false on that file only, so the FP32 boundary math is not contracted into FMAs. BeamAllReduceWorkspace.create returns None (standard TP all-reduce) when the op is not built, multicast is unavailable on any rank, or the shape/TP size is unsupported.
• CI coverage: tests/models/test_beam_numerics.py runs on a CPU agent in Buildkite, so the GPU and multi-GPU cases skip in CI. They were run on GB200 as above.
• Scope: an earlier revision also added defer_moe_finalize(allow_ep=True) to FusedMoEConfig for TP=1, EP>1. It moved to [Beam] Enable CuTeDSL NVFP4 MegaMoE #62073, which extends it to MegaMoE under DP/EP, so this PR leaves the shared MoE config untouched.
• Not a duplicate: no other open PR touches Beam's all-reduce.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[mergify]

Documentation preview: https://vllm--62366.org.readthedocs.build/en/62366/

[mergify]

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

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

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[WoosukKwon]

/ci run

[github-actions]

✅ Triggered Buildkite CI #96512 for commit 731958837ee9.