Systematic Optimization of Blackwell GPUs: From Allreduce Dead Ends to CUDA 13

Introduction

The path to high-performance LLM inference on cutting-edge hardware is rarely a straight line. It is a labyrinth of dead ends, incremental gains, and strategic pivots. This article chronicles a critical phase in an ongoing optimization effort for an 8× NVIDIA RTX PRO 6000 Blackwell GPU system connected via PCIe, deployed on Ubuntu 24.04. The work described here represents a systematic, evidence-driven journey through multiple allreduce optimization approaches—most of which failed—leading to a pivotal discovery about CUDA graph batch size tuning and ultimately a strategic pivot to upgrading CUDA to version 13 to unlock Blackwell-native optimizations.

The stakes were high. The team was deploying the GLM-5-NVFP4 model, a Mixture-of-Experts (MoE) architecture following the DeepSeek-V2 design, using SGLang as the inference engine. The hardware—8× RTX PRO 6000 Blackwell GPUs—represented the latest NVIDIA architecture (compute capability SM120), but the PCIe interconnect (rather than NVLink or NVSwitch) created a severe communication bottleneck. Every optimization decision had to account for this constraint.

The Allreduce Optimization Gauntlet

The core challenge was straightforward in principle but brutal in practice: Tensor Parallelism (TP) across 8 GPUs requires allreduce operations to synchronize activations after each transformer layer. On a PCIe-connected system, these allreduces become the dominant bottleneck. The team set out to test every available optimization approach to reduce this cost.

FlashInfer Allreduce Fusion: Blocked by Architecture

The first approach was FlashInfer's allreduce fusion, which promises to fuse multiple allreduce operations into a single kernel launch, reducing launch overhead and improving bandwidth utilization. The assistant configured the environment and attempted to enable this feature, only to discover that FlashInfer's JIT compiler does not support the SM120 (Blackwell) architecture. The compilation failed with architecture-specific errors, immediately closing this avenue.

This failure was particularly frustrating because FlashInfer is a well-established library in the SGLang ecosystem, and its allreduce fusion has been shown to provide significant speedups on earlier GPU generations (H100, A100). The Blackwell architecture's novelty meant that many software libraries had not yet been updated to support it.

Custom Allreduce Kernel: 2× Slower Than NCCL

Undeterred, the team turned to a custom allreduce kernel implementation. The hypothesis was that a kernel specifically optimized for the PCIe topology—perhaps using a different communication schedule or reducing the number of synchronization points—could outperform NCCL's generic ring allreduce.

The results were stark. The custom kernel produced only 38 tokens per second—more than 2× slower than NCCL's ring allreduce, which achieved approximately 82 tok/s. The root cause was massive PCIe bus contention. The all-to-all communication pattern used by the custom kernel required every GPU to communicate with every other GPU simultaneously, saturating the PCIe switches and creating a bottleneck far worse than NCCL's ring-based approach, which sequences communications to avoid contention.

This result was a humbling reminder that NCCL's ring allreduce, despite being a general-purpose solution, is highly optimized for exactly this kind of hardware topology. The ring pattern minimizes the number of simultaneous transfers, keeping PCIe links from being overwhelmed.

Torch Symmetric Memory: Another SM120 Casualty

The third approach was PyTorch's symmetric memory feature, which enables efficient peer-to-peer GPU memory access. This feature can reduce the need for explicit allreduce operations by allowing GPUs to directly read from each other's memory.

Unfortunately, Torch symmetric memory also failed because the SM120 architecture is not present in its architecture lookup table. PyTorch's implementation of symmetric memory requires architecture-specific code paths, and Blackwell support had not yet been added. This was another casualty of being on the leading edge of GPU hardware.

Expert Parallelism with FlashInfer A2A: Assertion Error and OOM

The fourth and most ambitious approach was Expert Parallelism (EP) using FlashInfer's all-to-all (A2A) backend. EP fundamentally changes the communication pattern for MoE layers: instead of allreducing activations across all GPUs (as in TP), each GPU holds a subset of experts, and tokens are routed to the GPU that owns the relevant expert via all-to-all communication. This eliminates the allreduce in the MoE layers entirely, replacing it with two all-to-all operations (dispatch and combine).

This approach was the subject of a dedicated subagent research session [1][2][3][4][5], which thoroughly investigated how EP is configured in SGLang. The research revealed that EP requires TP (it operates within the TP group), that the ep_size is automatically set to tp_size when any A2A backend is enabled, and that six backends are available: none, deepep, mooncake, mori, ascend_fuseep, and flashinfer [4]. The recommended configuration for NVIDIA GPUs with DeepSeek-V2-style models was --moe-a2a-backend deepep with --deepep-mode auto [4].

However, when the team attempted to deploy EP with the FlashInfer A2A backend on the Blackwell GPUs, they hit an assertion error followed by an out-of-memory (OOM) condition. The assertion error suggested that the FlashInfer A2A implementation had internal assumptions about GPU architecture or memory layout that were violated by the Blackwell system. The OOM indicated that the all-to-all dispatch buffers required more memory than was available, possibly due to the large expert count in the GLM-5-NVFP4 model.

The Verdict: NCCL Ring Remains King

After systematically testing and eliminating four optimization approaches, the team arrived at a sobering conclusion: NCCL's ring allreduce remains the best allreduce strategy for this hardware configuration. Despite being a general-purpose solution, NCCL's ring algorithm is better suited to PCIe-connected GPUs than any of the specialized alternatives tested. The ring pattern's sequential communication avoids the bus contention that plagues all-to-all patterns on PCIe.

This conclusion was not reached lightly. Each approach was tested with careful benchmarking, and the results were documented in the optimization plan. The team now had a clear baseline: NCCL ring allreduce at approximately 82 tok/s, with no immediately obvious path to improvement through alternative allreduce strategies.

The Breakthrough: Tuning CUDA Graph Batch Size

While the allreduce optimization approaches were failing, a different line of investigation yielded a significant and unexpected gain. The team experimented with the --cuda-graph-max-bs parameter, which controls the maximum batch size for which CUDA graphs are captured and reused.

CUDA graphs are a powerful optimization that captures a sequence of GPU kernel launches into a single graph object, which can then be replayed with minimal CPU overhead. This is particularly beneficial for the decode phase of LLM inference, where the same sequence of operations is repeated for each token. However, CUDA graphs consume GPU memory to store the captured operations, and larger batch sizes require more memory.

The default value of --cuda-graph-max-bs was 512. The team discovered that reducing this to 128 freed significant GPU memory that was previously reserved for CUDA graph capture. This freed memory could then be used for the KV cache, which is the primary memory consumer during inference and directly impacts the maximum batch size and throughput.

The result was a throughput improvement from 82 to 89.5 tok/s—a 9% gain achieved through a single parameter change. This was the first unambiguous success in the optimization effort, and it came not from a complex new algorithm or library, but from tuning an existing parameter to better match the hardware's memory constraints.

This discovery was particularly valuable because it required no code changes, no new library installations, and no architectural workarounds. It was a pure configuration optimization that could be applied immediately to the production deployment.

The EAGLE-3 Speculative Decoding Challenge

With the baseline throughput improved to 89.5 tok/s, the team turned to speculative decoding as the next potential optimization avenue. Speculative decoding uses a smaller, faster draft model to generate candidate tokens, which are then verified by the target model in parallel. The EAGLE-3 framework is a state-of-the-art approach that uses a lightweight draft model trained to predict the target model's hidden states.

The results were disappointing. EAGLE-3 speculative decoding reached only 54.1 tok/s, well below the 89.5 tok/s baseline. The bottleneck was the verify pass, where the target model must process all candidate tokens in parallel. This verify pass required approximately 30 milliseconds for 122 NCCL allreduces—one allreduce per transformer layer per verification step.

The fundamental problem was that speculative decoding's verify pass is inherently communication-bound on this hardware. Even though the draft model generates tokens quickly, the verification step requires the full target model forward pass, including all the allreduce operations that the team had been trying to optimize. The allreduce bottleneck that limited baseline throughput was not eliminated by speculative decoding—it was simply deferred to the verification step.

This result was particularly frustrating because speculative decoding is supposed to improve throughput by trading compute for communication: the draft model does less work per token, but the verification step still requires the full model. On a system where communication (not compute) is the bottleneck, this trade-off is unfavorable.

The Strategic Pivot: CUDA 13 Upgrade

With allreduce optimization approaches exhausted and speculative decoding underperforming, the team faced a strategic decision. The root cause of most failures was the same: the Blackwell SM120 architecture was too new for the software ecosystem to support fully. FlashInfer's JIT compiler didn't support it. Torch symmetric memory didn't support it. The FlashInfer A2A backend hit assertion errors on it.

The user proposed a bold solution: upgrade CUDA to version 13, which has native SM120 support. If the CUDA toolkit itself supported Blackwell natively, then all the libraries built on top of it—FlashInfer, PyTorch, SGLang—could potentially leverage that support.

The assistant investigated the feasibility of this upgrade and found encouraging results:

  1. Driver compatibility: The NVIDIA driver already installed (version 590.48.01) supports CUDA 13.1. The driver is forward-compatible with newer CUDA toolkits, so no driver change was needed.
  2. CUDA toolkit: The currently installed toolkit was CUDA 12.8. Upgrading to CUDA 13.1 would provide native SM120 support, potentially unblocking all the Blackwell-native optimizations that had been unavailable.
  3. PyTorch support: PyTorch nightly builds provide cu130 wheels, meaning PyTorch can be installed with CUDA 13 support directly from the nightly channel.
  4. SGLang kernel support: The sgl-kernel package has a dedicated cu130 index, indicating that SGLang's CUDA kernel implementations are being built for CUDA 13.
  5. FlashInfer support: FlashInfer also supports CUDA 13, which could unblock both the allreduce fusion and the A2A backend that had failed on SM120. The upgrade path would enable several optimizations that were previously blocked: - FlashInfer allreduce fusion: With CUDA 13's native SM120 support, FlashInfer's JIT compiler could generate Blackwell-compatible code, potentially reducing allreduce latency. - Torch symmetric memory: PyTorch's symmetric memory feature could work on Blackwell, enabling more efficient peer-to-peer GPU communication. - Expert Parallelism: The FlashInfer A2A backend and DeepEP could potentially work correctly on Blackwell, offering the all-to-all communication pattern that might better suit the PCIe topology. - Other Blackwell-native optimizations: CUDA 13 includes architecture-specific optimizations for SM120 that could improve kernel performance across the board. The key insight was that the CUDA 12.8 toolkit, while functional, was limiting the software stack to a compatibility layer that didn't leverage Blackwell's capabilities. Upgrading to CUDA 13 would allow the entire stack to run in a Blackwell-native mode, potentially eliminating the architecture compatibility issues that had blocked every optimization attempt.

The Subagent Research: Expert Parallelism Configuration

As part of this optimization journey, a subagent was spawned to conduct deep research into Expert Parallelism configuration in SGLang [1][2][3][4][5]. This research was critical because EP represented a fundamentally different approach to parallelism that might better suit the PCIe-connected Blackwell system.

The subagent's investigation revealed several key findings:

EP requires TP, not replaces it: When any A2A backend is enabled, SGLang automatically sets ep_size = tp_size. This means EP operates within the TP group rather than replacing it. The attention layers still use tensor parallelism, while the MoE layers use expert parallelism [4].

Six A2A backends are available: The MoeA2ABackend enum defines six options: none, deepep, mooncake, mori, ascend_fuseep, and flashinfer. For NVIDIA GPUs with DeepSeek-V2-style models, the relevant choices are deepep and flashinfer [2][4].

CUDA graph compatibility is mode-dependent: For the deepep backend, CUDA graphs are disabled in normal mode but can be used in low_latency mode. The auto mode intelligently selects normal for prefill and low_latency for decode [4].

The communication pattern changes fundamentally: With EP, the allreduce in the MoE layer is eliminated and replaced by two all-to-all operations (dispatch and combine). This changes the communication from "everyone talks to everyone" to "each GPU sends tokens to the specific GPU that has the expert it needs" [4].

The recommended launch command was:

python -m sglang.launch_server \
    --model-path <model> \
    --tp 8 \
    --moe-a2a-backend deepep \
    --deepep-mode auto \
    --trust-remote-code

This research was thorough and well-structured, but the runtime testing of EP with the FlashInfer A2A backend ultimately failed due to assertion errors and OOM conditions on the Blackwell hardware. The CUDA 13 upgrade path might unblock this as well.

Synthesis: What Was Learned

The optimization effort in this chunk produced several important lessons:

1. Hardware constraints dominate optimization choices. The PCIe interconnect was the single most important factor determining which optimization strategies were viable. Approaches that work well on NVLink-connected systems (like all-to-all communication patterns) fail on PCIe due to bus contention. The NCCL ring allreduce, despite being a general-purpose solution, outperformed specialized alternatives because its sequential communication pattern respects PCIe's bandwidth limitations.

2. Being on the leading edge of GPU hardware has costs. The Blackwell SM120 architecture was too new for many software libraries to support. FlashInfer, Torch symmetric memory, and the FlashInfer A2A backend all failed due to architecture incompatibility. This is a common pattern in ML infrastructure: the latest hardware offers the best performance potential, but the software ecosystem takes time to catch up.

3. Configuration tuning can yield significant gains without code changes. The discovery that reducing --cuda-graph-max-bs from 512 to 128 improved throughput by 9% was the most impactful finding in this chunk. It required no new libraries, no code modifications, and no architectural changes—just understanding how CUDA graph memory allocation interacts with KV cache memory pressure.

4. Speculative decoding is not a panacea for communication-bound systems. The EAGLE-3 results demonstrated that when communication is the bottleneck, speculative decoding's verify pass simply inherits the same bottleneck. The draft model's speed advantage is negated by the verification step's communication requirements.

5. Strategic pivots are sometimes necessary. After systematically testing and eliminating four optimization approaches, the team recognized that the fundamental limitation was the CUDA toolkit version, not any specific library or algorithm. The pivot to CUDA 13 represents a bet that upgrading the foundation will unblock multiple optimization paths simultaneously.

Conclusion

This chunk of the optimization journey was characterized by systematic experimentation, careful measurement, and strategic decision-making. The team tested and eliminated FlashInfer allreduce fusion, custom allreduce kernels, Torch symmetric memory, and Expert Parallelism with FlashInfer A2A—all due to the combination of Blackwell's new SM120 architecture and the PCIe interconnect. The one clear win was the discovery that reducing --cuda-graph-max-bs improved throughput by 9%, a pure configuration optimization that could be applied immediately.

The EAGLE-3 speculative decoding results were disappointing but informative, confirming that the allreduce bottleneck in the verify pass remains the dominant constraint. The strategic pivot to upgrading CUDA to version 13 offers a path forward that could unblock multiple optimization avenues simultaneously, potentially enabling FlashInfer fusion, Torch symmetric memory, and Expert Parallelism to finally work on the Blackwell hardware.

The subagent research into Expert Parallelism configuration [1][2][3][4][5] provided a thorough understanding of how EP works in SGLang, even though runtime testing was blocked by architecture compatibility issues. This research will become immediately actionable once the CUDA 13 upgrade is complete.

The work in this chunk exemplifies the iterative nature of ML infrastructure optimization: test hypotheses, measure results, learn from failures, and pivot when the evidence demands it. The path to high-performance inference on Blackwell GPUs is not a straight line, but each dead end provides information that guides the next step.