Patching the Gates: Enabling FlashInfer Allreduce Fusion for Blackwell SM120 GPUs

Introduction

In the high-stakes world of large language model inference, every microsecond counts. When deploying the massive GLM-5-NVFP4 model—a Mixture-of-Experts architecture with approximately 78 layers and FP4 quantization—across eight NVIDIA RTX PRO 6000 Blackwell GPUs, the assistant had already achieved an impressive throughput of ~3,740 tokens per second at 1024 concurrent requests. Yet something was deeply wrong: each GPU was drawing only ~250W of power, barely 42% of its 600W thermal design power (TDP). The hardware was loafing while the software struggled to keep it fed.

Message [msg 749] captures a pivotal moment in the debugging saga: the moment the assistant confirmed that it had successfully patched the sglang inference server to recognize the SM120 architecture (the compute capability of the RTX PRO 6000) in three critical code paths that had previously excluded it. This seemingly minor change—adding or is_sm120_supported() to a handful of conditional checks—represented the culmination of a deep investigation into why the GPUs were underutilized, and it set the stage for the next round of performance tuning.

The Investigation That Led to This Moment

The story begins in [msg 726], when a collaborator pointed out that the GPUs were drawing suspiciously low power. The assistant immediately verified this: nvidia-smi confirmed a 600W power limit per GPU, yet the instantaneous draw was only ~84W at idle and ~250W under load ([msg 727]). This was a clear signal that the GPU compute units were spending most of their time stalled, waiting for something.

The assistant hypothesized three possible bottlenecks: MoE expert GEMMs (FP4 quantized compute), attention operations (BF16 with FP8 KV cache), and allreduce communication over PCIe. To get definitive answers, it spawned a subagent task ([msg 729]) that performed a deep analysis of the GLM-5 forward pass, examining data types, kernel launches, and memory traffic at each stage.

The subagent's findings were revelatory ([msg 730]). The MoE expert GEMMs were correctly using FP4 (E2M1) precision with BF16 accumulation—the tensor cores were doing real work. But the critical bottleneck was the allreduce fusion: FlashInfer's ability to overlap allreduce communication with compute was completely disabled on SM120. The gatekeeping code in communicator.py at line 78 checked for _is_sm90_supported or _is_sm100_supported—but not _is_sm120_supported. Since the RTX PRO 6000 has SM120 compute capability (a consumer/professional variant of the Blackwell architecture, distinct from the datacenter SM100), every single allreduce operation—156 per forward pass (2 per layer × ~78 layers)—was serialized with compute. The GPU SMs sat idle while PCIe transfers completed.

The Three Code Gaps

The assistant identified three distinct places where SM120 was missing from the architecture checks:

  1. communicator.py line 78 — The runtime gate for apply_flashinfer_allreduce_fusion(). This function checks whether to enable fused allreduce, and SM120 was absent from the condition. This was the most critical fix because it directly controls whether allreduce operations overlap with computation at runtime.
  2. server_args.py line 1686 — The auto-enable logic for --enable-flashinfer-allreduce-fusion. When the server starts, it automatically enables allreduce fusion for supported architectures. SM120 was missing from this auto-detection, meaning even if a user manually passed the flag, the server might not honor it correctly.
  3. server_args.py line 1603 — The MoE runner backend auto-selection for GLM-5 models. This code path detects the GPU architecture and selects the appropriate quantization and MoE runner backend. Without SM120 support, the server would fall back to suboptimal defaults.

Applying the Fixes

Messages [msg 740] through [msg 747] show the assistant methodically applying four sed commands to patch these files. The approach was surgical and precise:

Fix 1 ([msg 740]): Added is_sm120_supported to the import list in communicator.py, inserting it right after is_sm100_supported in the import statement.

Fix 2 ([msg 741]): Added a cached module-level variable _is_sm120_supported = _is_cuda and is_sm120_supported() alongside the existing _is_sm90_supported and _is_sm100_supported variables. This ensures the architecture check is evaluated once at import time rather than on every function call.

Fix 3 ([msg 742]): Modified the allreduce fusion condition from (_is_sm90_supported or _is_sm100_supported) to (_is_sm90_supported or _is_sm100_supported or _is_sm120_supported). This is the core runtime fix.

Fix 4 ([msg 745]): Updated server_args.py line 1686 to include is_sm120_supported() in the auto-enable condition for allreduce fusion.

Fix 5 ([msg 747]): Updated server_args.py line 1603 to include is_sm120_supported() in the MoE runner backend auto-selection for GLM-5 models.

Message [msg 748] then verified all fixes by grepping for sm120 across both files, confirming that the patches were applied correctly and that the modified lines read as expected.

The Confirmation Message

Message [msg 749] is the status update that follows these verifications. It contains a todowrite call that marks three high-priority tasks as "completed":

Assumptions and Potential Pitfalls

The assistant made several assumptions in applying these patches:

  1. That SM120 is functionally equivalent to SM100 for allreduce fusion. The TRT-LLM communication kernels that FlashInfer relies on were written for SM90 (Hopper) and SM100 (datacenter Blackwell). SM120 (consumer/professional Blackwell) has a different shared memory size and potentially different synchronization primitives. Simply adding SM120 to the gate might enable the feature, but the underlying kernels might not work correctly—or might perform worse than the non-fused path.
  2. That the is_sm120_supported() function exists and returns True. The assistant verified this in [msg 739] by running a quick Python import test, confirming that the utility function was already defined in the sglang codebase. This was a critical prerequisite—if the function didn't exist, the patches would have caused import errors.
  3. That the server_args auto-detection changes would be sufficient. The assistant noted that since the server was being launched with explicit flags (--moe-runner-backend flashinfer_cutlass and --enable-flashinfer-allreduce-fusion), the server_args changes primarily helped for auto-detection scenarios. The communicator.py fix was the one that actually mattered at runtime. These assumptions would prove partially incorrect. In the subsequent testing (described in the chunk summary), enabling allreduce fusion on SM120 actually degraded performance to 236 tok/s with power dropping to 125W—suggesting that the TRT-LLM communication kernels had synchronization issues on SM120. The assistant had to revert the flashinfer kernel patches and explore alternative approaches.

Input Knowledge Required

To understand this message, the reader needs:

Output Knowledge Created

This message, combined with the patches it reports, creates several pieces of knowledge:

The Thinking Process

The assistant's reasoning in this sequence reveals a systematic debugging methodology:

  1. Observe the symptom: Low GPU power draw despite high utilization metrics.
  2. Formulate hypotheses: Three potential bottlenecks (MoE, attention, allreduce).
  3. Gather data: Subagent task performs deep analysis of the forward pass.
  4. Identify root cause: Allreduce fusion disabled on SM120.
  5. Trace the code paths: Find every location where SM120 is missing from architecture checks.
  6. Apply targeted fixes: Modify only the condition checks, not the underlying kernels.
  7. Verify correctness: Grep the patched files to confirm the changes.
  8. Update task tracking: Mark completed items and identify remaining work. This is classic root-cause debugging applied to a complex distributed inference system. The assistant resists the temptation to throw hardware at the problem (e.g., "buy faster GPUs") or apply superficial fixes (e.g., "increase power limit"). Instead, it traces the symptom to a specific conditional check in a specific file, applies a minimal fix, and prepares to measure the result.

Conclusion

Message [msg 749] appears, at first glance, to be a mundane status update—a todo list with three items checked off. But it represents the culmination of a deep investigation into GPU underutilization in large-scale model inference. The assistant identified that the sglang inference server had hardcoded architecture gates that excluded SM120—the compute capability of the RTX PRO 6000 Blackwell GPUs—from three critical features: allreduce fusion, MoE backend auto-selection, and fusion auto-enable.

The patches themselves are minimal: a few lines added to conditional checks. But their implications are significant. They represent the ongoing work of adapting a rapidly evolving inference stack (sglang) to a new GPU architecture (Blackwell SM120) that sits in an awkward middle ground—too new for the datacenter-focused SM100 kernels, too different from the Hopper SM90 kernels. The assistant's willingness to fork and patch upstream code, combined with its systematic debugging methodology, demonstrates the kind of deep systems engineering required to squeeze maximum performance from cutting-edge hardware.

The story does not end with a happy benchmark result. In fact, the subsequent testing would reveal that the allreduce fusion kernels themselves have issues on SM120, forcing the assistant to revert the changes and explore alternative approaches. But that is precisely what makes this message interesting: it captures the moment of applying a fix based on a correct diagnosis, before the fix's effectiveness is validated. It is a testament to the iterative nature of performance engineering, where identifying the right bottleneck is only half the battle—the other half is finding a working solution for your specific hardware.