perf: add push-based allreduce for small tensor reductions - #44891
alexm-redhat wants to merge 9 commits into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review to trigger a review and subscribe this PR to future pushes, or @claude review once for a one-time review.
Tip: disable this comment in your organization's Code Review settings.
|
This pull request has merge conflicts that must be resolved before it can be |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![depthfirst-app[bot]](https://avatars.githubusercontent.com/in/1021002?s=60&v=4){kind=small}
|
|
||
| // Allocate storage | ||
| storage_bytes_ = push_signal_bytes() + push_buffer_total_bytes(); | ||
| cudaMalloc(&storage_, storage_bytes_); |
There was a problem hiding this comment.
⚪ Severity: LOW
All CUDA API calls (cudaMalloc, cudaMemset, cudaIpcGetMemHandle, cudaIpcOpenMemHandle) lack error checking. The existing custom_all_reduce.cu wraps these with AT_CUDA_CHECK(). If cudaMalloc fails, storage_ remains null, and the subsequent cudaMemset and all kernel operations write through an invalid pointer, causing undefined behavior.
Helpful? Add 👍 / 👎
💡 Fix Suggestion
Suggestion: Wrap all CUDA API calls with AT_CUDA_CHECK() to match the existing pattern used in custom_all_reduce.cu. This applies to cudaMalloc, cudaMemset, cudaGetDevice, cudaDeviceGetAttribute, cudaIpcGetMemHandle, cudaIpcOpenMemHandle, and other CUDA calls throughout this file. You will also need to add #include <ATen/cuda/Exceptions.h> at the top of the file to make AT_CUDA_CHECK available. At the immediate location (lines 44-48), change cudaMalloc(...) to AT_CUDA_CHECK(cudaMalloc(...)) and cudaMemset(...) to AT_CUDA_CHECK(cudaMemset(...)). Apply the same wrapping to all other unchecked CUDA calls in the file (lines 37, 39-40, 56, 60, 67, 78, and the destructor calls).
There was a problem hiding this comment.
Fixed — added CUDACHECK wrapping for all CUDA API calls.
| // Verify input fits in push buffer | ||
| const int64_t input_bytes = | ||
| static_cast<int64_t>(sizeof(T)) * num_elements; | ||
| assert(input_bytes <= push_buffer_bytes_); |
There was a problem hiding this comment.
⚪ Severity: LOW
This buffer-overflow guard uses assert(), which is compiled out under the default RelWithDebInfo build (-DNDEBUG). If all_reduce() is called with an input exceeding push_buffer_bytes_ (the Python all_reduce method has no independent size check), the kernel performs volatile NVLink stores past the allocated push buffer into adjacent GPU memory on all peers.
Helpful? Add 👍 / 👎
💡 Fix Suggestion
Suggestion: Replace the assert() with a runtime check that is not removed by -DNDEBUG. Since this .cuh header does not include Torch headers, the simplest fix is to replace the assert with an explicit if check that throws a std::runtime_error. You will also need to add #include <stdexcept> to the existing includes at the top of the file (or alternatively, include <torch/all.h> and use TORCH_CHECK). Replace line 99:
assert(input_bytes <= push_buffer_bytes_);with:
if (input_bytes > push_buffer_bytes_) {
throw std::runtime_error(
"push_all_reduce: input (" + std::to_string(input_bytes) +
" bytes) exceeds push buffer capacity (" +
std::to_string(push_buffer_bytes_) + " bytes)");
}Also add #include <stdexcept> near line 14 alongside the other standard library includes.
There was a problem hiding this comment.
Fixed — replaced assert with a runtime buffer overflow guard that is not compiled out under NDEBUG.
| if (i == rank_) { | ||
| peer_storage_[i] = storage_; | ||
| } else { | ||
| cudaIpcOpenMemHandle(&peer_storage_[i], peer_handles[i], |
There was a problem hiding this comment.
⚪ Severity: LOW
The cudaIpcOpenMemHandle return value is unchecked. If it fails (e.g. invalid/corrupted IPC handle), peer_storage_[i] is left as null/undefined. The kernel then performs NVLink volatile stores to these invalid addresses via push_impl, corrupting arbitrary remote GPU memory.
Helpful? Add 👍 / 👎
💡 Fix Suggestion
Suggestion: Wrap the cudaIpcOpenMemHandle call with a CUDA error-checking macro to detect and handle failures. The repo already defines a CUDACHECK macro in custom_all_reduce.cuh. Either include that header or define an equivalent macro in push_all_reduce.cuh, then wrap the call:
CUDACHECK(cudaIpcOpenMemHandle(&peer_storage_[i], peer_handles[i],
cudaIpcMemLazyEnablePeerAccess));Note: the same error-checking gap applies to other unchecked CUDA calls in this file (cudaIpcGetMemHandle at line 67, cudaMalloc/cudaMemset in the constructor, cudaFree in the destructor). Consider wrapping all of them consistently.
⚠️ Experimental Feature: This code suggestion is automatically generated. Please review carefully.
| cudaIpcOpenMemHandle(&peer_storage_[i], peer_handles[i], | |
| CUDACHECK(cudaIpcOpenMemHandle(&peer_storage_[i], peer_handles[i], | |
| cudaIpcMemLazyEnablePeerAccess)); |
There was a problem hiding this comment.
Fixed — added CUDACHECK wrapping for cudaIpcOpenMemHandle and all other unchecked CUDA calls.
|
@ilmarkov would be good to have a quick sanity check |
ilmarkov
left a comment
ilmarkov
left a comment
There was a problem hiding this comment.
Thank you for the PR!
In this review I focused on the integration of the kernel postponing the review of the copy-pasted kernel (assume it works fine in Sglang).
Some comments:
- I would suggest is to add the kernel to
benchmark_device_communicators.pyto compare to the other implementations and possibly tune the thresholds on other architectures. VLLM_DISABLE_PUSH_ALLREDUCEnot registered inenvs.py- Update with the latest main and add the method to
_log_all_reduce_backend_selection
| @@ -81,6 +80,7 @@ def __init__( | |||
| register_nccl_symmetric_ops(self.pynccl_comm) | |||
|
|
|||
| self.ca_comm: CustomAllreduce | None = None | |||
| self.push_ar_comm = None | |||
There was a problem hiding this comment.
Maybe, type annotation to be consistent with other communicators.
There was a problem hiding this comment.
Done — added type annotation for consistency with other communicators.
|
|
||
| # Push threshold map: world_size -> buffer_bytes | ||
| # From SGLang's tuned thresholds for sm100 (B200) | ||
| PUSH_THRESHOLD_SM100 = { |
There was a problem hiding this comment.
Looks like it is applied unconditionally to all archs
There was a problem hiding this comment.
Fixed — added architecture-specific threshold selection.
There was a problem hiding this comment.
Fixed — added cleanup.
This is a pretty strange workload - what is the performance in a more realistic scenario? |
6d32045 to
7e7e30c
Compare
|
The test was done for a KV cache long workload which Roberto said is more typical for DeepSeek V4. Will test more use-cases. |
|
@ilmarkov Thank you for the review!
|
Port SGLang's push-based 2-buffer allreduce protocol into vLLM as a new communicator backend for small-message reductions. The push protocol eliminates the two explicit cross-GPU NVLink barrier round-trips used by the existing barrier-based CustomAllreduce, replacing them with a sentinel-based data arrival detection mechanism and double-buffered epoch alternation. Key advantages over the barrier-based approach: - Zero barriers: data arrival IS the synchronization (positive-zero sentinel) - Single NVLink round-trip instead of two barrier exchanges + remote reads - All SMs active (SM_count CTAs vs 2 CTAs) for higher NVLink bandwidth - No cudaMemcpy to IPC staging buffer in eager mode - PDL (griddepcontrol) support for kernel overlap on sm_90+ The new PushAllReduce is inserted in the CudaCommunicator dispatch chain above the existing CustomAllreduce for messages below a size threshold (~720 KB at TP=8). Larger messages continue to use the barrier-based path. The existing CustomAllreduce code is not modified. Measured results on DeepSeek-V4-Pro (61 layers, TP=8, 8x NVIDIA B200, BS=1, decode with ISL=4, OSL=33024): - Throughput: +2.14% (84.06 vs 82.30 tokens/s) - TPOT: -2.09% (11.90 vs 12.15 ms/token) Correctness verified via lm_eval gsm8k 5-shot with no regression (exact_match delta within statistical noise). The feature can be disabled at runtime via VLLM_DISABLE_PUSH_ALLREDUCE=1 to fall back to the barrier-based path. Signed-off-by: Alexander Matveev <amatveev@redhat.com> Signed-off-by: Alexander Matveev <alexm-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com>
…uard - Add PUSH_AR_CUDACHECK macro wrapping all CUDA API calls (cudaGetDevice, cudaDeviceGetAttribute, cudaMalloc, cudaMemset, cudaIpcGetMemHandle, cudaIpcOpenMemHandle) to match the CUDACHECK pattern in custom_all_reduce.cuh - Replace assert(input_bytes <= push_buffer_bytes_) with a runtime std::runtime_error check that is not compiled out under -DNDEBUG - Add #include <stdexcept> and #include <string> for the runtime check Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Alexander Matveev <alexm-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com>
Fix CodeQL security warning by binding test helper sockets to "localhost" instead of "" (all interfaces). These sockets are only used for finding a free port for torch distributed init in tests. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Alexander Matveev <alexm-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com>
- Add PushAllReduce | None type annotation on push_ar_comm to be consistent with other communicator fields (ca_comm, qr_comm, etc.) - Add push_ar_comm.close() + None assignment in destroy() method to match the cleanup pattern for other communicators - Add lazy import of PushAllReduce alongside other communicator imports Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Alexander Matveev <alexm-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com>
The push threshold map was labeled as sm100-specific but applied unconditionally to all architectures. Now: - PUSH_THRESHOLD_SM100 is only used on Blackwell (compute capability 10.x) - PUSH_THRESHOLD_DEFAULT provides conservative 512 KB thresholds for architectures without tuned values - _THRESHOLD_BY_ARCH maps GPU major compute capability to threshold tables - A log message is emitted when falling back to conservative defaults Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Alexander Matveev <alexm-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com>
Register the push allreduce feature toggle env var in the central envs.py registry so it is validated on startup and follows the standard vllm env var pattern. Default is False (push allreduce enabled); set to 1 to disable. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Alexander Matveev <alexm-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com>
Replace direct os.environ.get(_DISABLE_ENV_VAR) == "1" check with envs.VLLM_DISABLE_PUSH_ALLREDUCE to use the centrally registered env var from envs.py, which provides validation and caching. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Alexander Matveev <alexm-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com>
- Add PushAllReduce to benchmark_device_communicators.py for comparing against other allreduce implementations - Add PUSH_AR to _log_all_reduce_backend_selection in cuda_communicator.py for visibility in dispatch logging Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Alexander Matveev <alexm-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com>
b833264 to
fd44100
Compare
- Apply ruff, clang-format formatting fixes - Replace torch.cuda.set_device/device_count/synchronize with torch.accelerator equivalents per project convention Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Alexander Matveev <alexm-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com> Signed-off-by: <>
|
@alexm-redhat Looks good, thanks! Could you please add the allreduce benchmark results to the description to verify that the hardcoded constants actually make sense? |
|
@ilmarkov ok will do, yeah in process of running more tests |
|
This pull request has merge conflicts that must be resolved before it can be |
|
@ilmarkov @tlrmchlsmth here is a summary of extensive nightly benchmarks. Everything seems to be working, also verified thresholds. Push AllReduce Benchmark Results — Complete SummaryHardware: 8x NVIDIA B200 (183 GiB each), NVLink
1. Batch=1 across ISL/OSL configurationsDeepSeek V4 Flash (TP=4, 4x B200)
DeepSeek V4 Pro (TP=8, 8x B200)
2. Batch size scaling (ISL/OSL=1024/1024)DeepSeek V4 Flash (TP=4, 4x B200) — threshold = 2 MB
DeepSeek V4 Pro (TP=8, 8x B200) — threshold = 720 KB
Key Findings
|
4 participants
Summary
This PR adds a push-based 2-buffer allreduce protocol to vLLM, targeting small-message reductions during decode (e.g., the per-layer attention and MoE allreduces in DeepSeek-V4 models). The existing barrier-based
CustomAllreduceuses two explicit cross-GPU NVLink barrier round-trips per call, which dominate latency for small tensors. The push protocol eliminates these barriers entirely by having each rank push its data to all remote GPUs via NVLink volatile stores and then polling locally for arrival using a positive-zero sentinel mechanism with double-buffered epoch alternation.The new
PushAllReducecommunicator is inserted into theCudaCommunicatordispatch chain above the existingCustomAllreducefor messages below a dynamic size threshold (~720 KB at TP=8 on B200). Messages above this threshold continue to use the barrier-based path. No existingCustomAllreducecode is modified. The implementation includes a consolidated CUDA kernel header vendored from SGLang's proven push kernel with minimal macro substitutions (namespace, C++17 compatibility, AOT compilation guards), a C++ host-side manager class following vLLM patterns, torch custom op bindings, and a Python communicator wrapper with CUDA graph support.The kernel uses all available SMs (vs 2 CTAs in the barrier-based approach), supports PDL (
griddepcontrol) overlap on sm_90+ architectures, and avoids thecudaMemcpyto IPC staging buffer required by the barrier-based path in eager mode. CUDA graph safety is ensured by placing the epoch counter in device memory (read live during replay) with a fixed grid size equal to the SM count.Performance Results
Workload: decode, BS=1, ISL=4, OSL=33024, KV cache FP8. Baseline = same patched code with push allreduce disabled via env var (falls back to barrier-based
CustomAllreduce).Per-call improvement on V4-Pro (122 allreduce calls per forward pass across 61 layers): ~2.09 us/call.
Correctness (lm_eval)
Benchmark: gsm8k 5-shot (1319 prompts). All deltas are within 1 standard error (~0.006) and statistically indistinguishable (z-scores < 1.96). The sign reversal between configurations confirms random variance, not a systematic shift.
Feature Toggle
VLLM_DISABLE_PUSH_ALLREDUCEVLLM_DISABLE_PUSH_ALLREDUCE=1AllReduce push-based 2-buffer protocol is ENABLEDAllReduce push-based 2-buffer protocol is DISABLED (env override)When disabled,
PushAllReduce.__init__returns immediately without allocating any GPU resources.should_use()returnsFalsefor all inputs, andCudaCommunicatorfalls through to the barrier-basedCustomAllreduce.Test Coverage
24 tests included in the patch, all passing on 2x NVIDIA B200:
should_use()predicate logic, bit-exact integer correctness vs NCCL, float correctness, positive-zero sentinel handling, 1000-call epoch alternation, thread count selection, threshold boundary, out-of-place semantics, dtype coverage (fp16/bf16/fp32), multi-layer simulation (122 ARs x 3 steps), lifecycle/cleanup, warmup path, asymmetric rank data, identical-values edge case