[Perf][MoE] Eliminate staging from NCCL symmetric reduce-scatter

[samnordmann]

Summary

Remove both staging copies from the modular MoE symmetric reduce-scatter path, building on #46703:

Before: expert output → copy to registered input → ReduceScatter → copy to final output
After:  expert kernel / top-k reduction → registered input → ReduceScatter → final output

The optimization requires VLLM_USE_NCCL_SYMM_MEM=1, uniform token counts, and NCCL >= 2.29.2. Unsupported configurations retain their existing fallbacks.

Implementation

• Write expert output or contiguous top-k reduction directly into the registered combine buffer, and let NCCL write into the caller-owned final output.
• Reuse geometrically grown buffers, retaining old allocations for CUDA-graph replay and separating DBO microbatches.
• Preserve deterministic collectives, uneven-routing fallbacks, and CPU/ROCm behavior.

NCCL supports registered-send/nonregistered-receive ReduceScatter in ncclSymkPickKernel.

Performance

Measured on 4×GB200 with NCCL 2.30.7. Operator latency and end-to-end throughput measure different effects; their percentage improvements are not additive.

Reduce-scatter latency

BF16, hidden size 8192; latency in microseconds, lower is better. Latency reduction is (staged − PR) / staged.

Output/rank	Staged symmetric	Direct symmetric (PR)	Latency reduction
0.5 MiB	42.4	35.6	16.0%
1 MiB	43.1	37.1	13.9%
2 MiB	42.3	37.4	11.6%
4 MiB	51.6	43.5	15.7%
8 MiB	84.8	76.7	9.6%

Removing staging reduces symmetric reduce-scatter latency by about 10–16% over this range. Ordinary NCCL remains faster at some sizes.

Nemotron Ultra end-to-end

Ultra 550B A55B NVFP4, TP1/DP4/EP4, default AG/RS backend, symmetric memory enabled in both baseline and PR. All changes below use the same convention: PR / baseline − 1, computed from geometric means. Higher throughput and lower TPOT/TTFT are better; concurrency is global.

Workload (input / output tokens)	Concurrency	Output throughput	Mean TPOT	Mean TTFT
Long-prefill: 50,000 / 2,048	16	+8.91%	−9.38%	−0.36%
Short-context: ~1,024 / 512	32	+1.35%	−1.35%	−0.98%
Short-context: ~1,024 / 512	128	+2.04%	−1.05%	−6.46%
Short-context: ~1,024 / 512	512	+1.71%	−4.16%	+3.42%

The long-prefill workload benefits most; its TTFT change is inconclusive. At short-context concurrency 32, throughput and TPOT improve in all six paired runs. Higher-concurrency results are mixed: TTFT worsens in all six runs at concurrency 512. These are two distinct workloads, not one concurrency-scaling curve.

E2E results compare baseline 5414b4e with PR revision 622c1cc; the current rebased head has not been benchmarked. The short-context sweep uses common instrumentation and disables Mamba stochastic rounding in both arms, so it is not a native-default deployment claim. No Super or multi-node scaling claim is made.

Validation

• Local pre-commit checks, mypy, Python compilation, and git diff --check pass.
• Validated revision 3fbfae6: 13 checks pass on 4×GB200, with no skips (10 repository tests and 3 supplemental checks). Coverage includes caller-owned outputs, uneven routing, CPU/mocked-AITER regressions, buffer growth, and BF16/FP32 CUDA-graph replay after allocator churn.
• No separate model-accuracy suite was run; the E2E performance results above remain from the earlier revision.

AI tooling assisted inspection, testing, and analysis; the author remains responsible for the change.

[mergify]

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

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

[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, @samnordmann.

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

[github-actions]

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

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
• /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.

[mgoin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #91701 for commit 81b740460276.

[mergify]

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

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

[janbernloehr]

Validated 0802d9096949 against the GB200 DP4+EP serving crash reported in #59189. Both unpatched controls reproduced the failure; applying this PR’s runtime changes allowed both models to complete every tested concurrency.

Both arms used the same vllm/vllm-openai:nightly image, pulled September 29: vLLM 0.30.1rc1.dev312+g36768d1bf. The treatment applied only the vllm/ delta from merge-base dfc8e0f3e2ad to PR head 0802d9096949 using git apply against the installed package. Package, library, and system-version comparisons showed no differences between arms.

Model	Unpatched result	With this PR
nvidia/DeepSeek-R1-0528-NVFP4-v2 (25a138f)	Passed concurrency 1, then crashed at 128	Passed 1, 128, 512, 1024, 2048
nvidia/Kimi-K2-Thinking-NVFP4 (d427ad5)	Crashed at the first tested concurrency, 2048	Passed 2048, 128, 8

Both controls failed with RuntimeError: NCCL error: unhandled cuda error along this path:

_reduce_scatter_symm_mem → _get_symm_scratch → ncclCommWindowRegister

All requests completed in the patched runs, with no NCCL errors. Output throughput, in the concurrency order above:

• DeepSeek-R1: 76.7 / 4,643 / 12,335 / 16,427 / 20,124 tok/s
• Kimi-K2-Thinking: 7,728 / 4,253 / 581 tok/s

The environment was one node with 4×GB200, CUDA 13.0, driver 580.178.04, PyTorch 2.13.0+cu130, NCCL 2.30.7, and FlashInfer 0.7.0. Both models used TP1/PP1/DP4 with expert parallelism and:

VLLM_USE_NCCL_SYMM_MEM=1
NCCL_NVLS_ENABLE=1
NCCL_CUMEM_ENABLE=1

--kv-cache-dtype fp8
--gpu-memory-utilization 0.95
--max-model-len 10240
--max-num-batched-tokens 8192
--compilation_config.max_cudagraph_capture_size 2048
--compilation_config.cudagraph_mode FULL_DECODE_ONLY
--moe-backend flashinfer_trtllm
--async-scheduling
--no-enable-prefix-caching

Load was streaming completions through aiperf with 2,000 input tokens, exactly 1,000 output tokens, ignore_eos=true, and 5 × concurrency requests per point.

This validates resolution of the reproduced crash for these configurations. Coverage is one complete paired run per model; an earlier DeepSeek pair showed the same outcome. No model-accuracy evaluation was performed, and the selected collective path was not instrumented. The crashing controls prevent a complete throughput comparison against the unpatched baseline.

Drafted with Codex (GPT-6).

[samnordmann]

@mgoin is it ready to merge ?

[vllm-agent]

CI selector (shadow): 91 test steps (142 jobs) instead of 83 (106 jobs)

Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works.

Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.

steps (jobs)	Today's rules	Selector	Would skip	Would add
NVIDIA, CPU and others	83 (106)	91 (142)	33 (42)	41 (78)
AMD mirrors	89 (126)	73 (102)	40 (58)	24 (34)

Selector would run (91)

• amd-fp8-moe-kernels-mi355
• amd-lm-eval-small-models-harness
• amd-qwen3-next-mtp-async-eplb-accuracy
• arm-cpu-test ×3
• async-engine-inputs-utils-worker
• basic-correctness-sleep-mode
• basic-models-test-other-cpu
• basic-models-tests-extra-initialization ×14
• basic-models-tests-other
• batch-invariance-b200
• batch-invariance-h100
• cpu-kernel-tests ×2
• cpu-multi-modal-model-tests-n ×4
• cpu-params-env-tokenizers-parser
• cpu-spec-decode-tests
• crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus
• distributed-comm-ops
• distributed-compile-comm-4-gpus
• distributed-compile-unit-tests-2xh100
• distributed-dp-tests-2-gpus
• distributed-dp-tests-4-gpus
• distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
• distributed-model-tests-2-gpus ×3
• distributed-nixlconnector-pd-accuracy-4-gpus
• distributed-tests-8xh100
• distributed-torchrun-examples-4-gpus
• distributed-torchrun-shutdown-tests-2-gpus
• dp-ep-distributed-nixlconnector-pd-accuracy-tests-4-gpus
• e2e-core-1-gpu
• e2e-core-large-memory
• elastic-ep-scaling-test
• entrypoints-integration-api-server ×4
• entrypoints-integration-api-server-generate
• entrypoints-integration-api-server-openai-completion
• eplb-algorithm
• eplb-execution
• fault-tolerance-e2e-2xh100
• fusion-and-compile-unit-tests-2xb200
• fusion-e2e-quick-h100
• fusion-e2e-tp2-ar-rms-dynamo-partition-amd
• fusion-e2e-tp2-ar-rms-inductor-partition-amd
• fusion-e2e-tp2-b200
• fusion-e2e-tp2-quick-h100
• gemm-rs-ar-2xb200
• hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus
• hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus
• kernels-b200 ×3
• kernels-core-operation-test ×3
• kernels-deepgemm-test-h100
• kernels-fusedmoe-layer-test-2-b200s
• kernels-fusedmoe-layer-test-2-h100s
• kernels-mamba-test
• kernels-minimax-reduce-rms-test-2-gpus
• kernels-moe-test ×5
• kernels-root-misc-test-b200
• kimi-k3-prefix-cache-4xb200 ×2
• kv-offload-large
• kv-offload-medium
• language-models-tests-granite-l4-compatibility
• language-models-tests-hybrid ×2
• language-models-tests-standard
• lm-eval-small-models
• lora ×4
• lora-tp-distributed ×4
• model-executor
• multi-modal-models-standard-4-other-whisper
• pipeline-context-parallelism-4-gpus
• plugin-tests-2-gpus
• push-nixlconnector-pp-prefill-pd-accuracy-4-gpus
• pytorch-compilation-passes-unit-tests
• pytorch-compilation-unit-tests
• pytorch-compilation-unit-tests-h100
• pytorch-fullgraph-test
• quantization ×4
• quantized-models-test
• quantized-moe-test-b200
• rayexecutorv2-4-gpus
• replayssm-e2e
• rust-frontend-distributed
• sharded-rdt-weight-transfer
• spec-decode-draft-model ×4
• spec-decode-eagle-1-deepseek-qwen
• spec-decode-mtp-deepseek-mimo
• v1-attention-b200 ×2
• v1-attention-h100-mi300 ×2
• v1-core
• v1-executor-worker
• v1-kv-connectors ×4
• v1-kv-offload
• v1-others-cpu
• v1-spec-decode

Would skip (today's rules run them) (33)

• ascend-npu-test
• basic-correctness ×2
• basic-correctness-cpu-offload
• basic-correctness-cumem
• basic-correctness-prefetch-offload
• basic-models-tests-initialization
• benchmarks-cli-test
• cpu-distributed-tests-dp-tp
• cpu-distributed-tests-pp-tp
• cpu-language-generation-and-pooling-model-tests ×3
• distributed-compile-rpc-tests-2-gpus
• entrypoints-integration-api-server-openai-chat_completion
• entrypoints-integration-llm
• entrypoints-integration-multimodal
• entrypoints-integration-pooling
• entrypoints-integration-responses-api
• entrypoints-integration-speech_to_text
• kernels-fla-ops-test-b200
• kernels-mhc-test-b200
• metrics-tracing-2-gpus
• multi-modal-models-standard-1-qwen2
• multi-modal-models-standard-2-qwen3-gemma
• multi-modal-models-standard-3-llava-qwen2-vl
• multi-modal-processor ×4
• multi-modal-processor-cpu ×4
• pytorch-compilation-dynamic-shapes
• pytorch-fullgraph-cudagraph-l4-compatibility
• regression
• samplers-multimodal-beam-search
• samplers-test
• v1-logits-oracle
• v1-metrics-lmeval
• v1-sample

Would add (today's rules do not run them) (41)

• amd-lm-eval-small-models-harness (Python record)
• arm-cpu-test ×3 (code map)
• basic-models-tests-extra-initialization ×14 (Python record)
• cpu-kernel-tests ×2 (code map)
• cpu-multi-modal-model-tests-n ×4 (code map)
• cpu-params-env-tokenizers-parser (code map)
• cpu-spec-decode-tests (code map)
• crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus (Python record)
• distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus (Python record)
• distributed-model-tests-2-gpus ×3 (code map)
• distributed-nixlconnector-pd-accuracy-4-gpus (Python record)
• dp-ep-distributed-nixlconnector-pd-accuracy-tests-4-gpus (Python record)
• e2e-core-1-gpu (Python record)
• e2e-core-large-memory (Python record)
• eplb-algorithm (code map)
• eplb-execution (code map)
• fusion-and-compile-unit-tests-2xb200 (code map)
• gemm-rs-ar-2xb200 (code map)
• hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus (Python record)
• hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus (Python record)
• kernels-b200 ×3 (code map)
• kernels-core-operation-test ×3 (code map)
• kernels-mamba-test (code map)
• kernels-minimax-reduce-rms-test-2-gpus (code map)
• kimi-k3-prefix-cache-4xb200 ×2 (Python record)
• kv-offload-large (Python record)
• kv-offload-medium (Python record)
• lm-eval-small-models (Python record)
• lora ×4 (code map)
• plugin-tests-2-gpus (code map)
• push-nixlconnector-pp-prefill-pd-accuracy-4-gpus (Python record)
• quantization ×4 (code map)
• quantized-models-test (code map)
• rayexecutorv2-4-gpus (code map)
• replayssm-e2e (Python record)
• sharded-rdt-weight-transfer (code map)
• spec-decode-draft-model ×4 (Python record)
• spec-decode-eagle-1-deepseek-qwen (Python record)
• spec-decode-mtp-deepseek-mimo (Python record)
• v1-attention-b200 ×2 (code map)
• v1-attention-h100-mi300 ×2 (code map)

AMD mirrors: would skip (40)

• basic-correctness ×2
• basic-correctness-cpu-offload
• basic-correctness-cumem
• basic-correctness-prefetch-offload
• basic-models-tests-extra-initialization ×14
• basic-models-tests-initialization
• batch-invariance-h100
• benchmarks-cli-test
• distributed-compile-rpc-tests-2-gpus
• entrypoints-integration-llm
• entrypoints-integration-multimodal
• entrypoints-integration-pooling
• entrypoints-integration-responses-api
• entrypoints-integration-speech_to_text
• fusion-and-compile-unit-tests-2xb200
• fusion-e2e-tp2-ar-rms-config-sweep-h100
• fusion-e2e-tp2-asynctp-config-sweep-h100
• kernels-fla-ops-test-b200
• kernels-mhc-test-b200
• language-models-tests-extra-standard ×2
• metrics-tracing-2-gpus
• multi-modal-models-standard-1-qwen2
• multi-modal-models-standard-2-qwen3-gemma
• multi-modal-models-standard-3-llava-qwen2-vl
• multi-modal-processor ×4
• openai-api-correctness
• platform-tests
• pytorch-compilation-dynamic-shapes
• pytorch-fullgraph-cudagraph-l4-compatibility
• regression
• samplers-multimodal-beam-search
• samplers-test
• spec-decode-eagle-2-llama3-qwen-vl-other
• spec-decode-mtp-gemma4
• spec-decode-mtp-qwen3-5
• spec-decode-ngram-suffix
• spec-decode-speculators
• v1-logits-oracle
• v1-metrics-lmeval
• v1-sample

AMD mirrors: would add (24)

• crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus (Python record)
• distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus (Python record)
• distributed-nixlconnector-pd-accuracy-4-gpus (Python record)
• dp-ep-distributed-nixlconnector-pd-accuracy-tests-4-gpus (Python record)
• e2e-core-1-gpu (Python record)
• e2e-core-large-memory (Python record)
• eplb-algorithm (code map)
• eplb-execution (code map)
• hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus (Python record)
• hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus (Python record)
• kernels-core-operation-test ×3 (code map)
• kernels-mamba-test (code map)
• kv-offload-large (Python record)
• kv-offload-medium (Python record)
• lm-eval-small-models (Python record)
• lora ×4 (code map)
• plugin-tests-2-gpus (code map)
• push-nixlconnector-pp-prefill-pd-accuracy-4-gpus (Python record)
• quantization ×4 (code map)
• quantized-models-test (code map)
• rayexecutorv2-4-gpus (code map)
• sharded-rdt-weight-transfer (code map)
• v1-attention-b200 ×2 (code map)
• v1-attention-h100-mi300 ×2 (code map)

_{8 changed files · base 18c8a65edd · head 360b01c043 · Python record: build 92838 at 92044241a0 · kernel record: table 9204424 (build 92838), map 9204424 · not counted: 10 build steps, 5 A100 steps the generator no longer emits, 95 optional steps the selector would also run}

[mgoin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #92954 for commit 360b01c04328.