The Blackwell Compatibility Audit: Tracing SM120 Support Through SGLang's Codebase

Introduction

In the rapidly evolving landscape of AI inference infrastructure, the gap between cutting-edge hardware and production-ready software can be vast. When NVIDIA's Blackwell architecture (compute capability SM120) arrived with the RTX PRO 6000 GPUs, it promised generational leaps in performance for large language model serving. But promises on a spec sheet mean nothing if the software stack doesn't know how to talk to the hardware. This is the story of a systematic codebase investigation — a deep audit of the SGLang inference engine to determine whether its most critical performance optimizations would actually activate on Blackwell GPUs.

The investigation, conducted across a sub-session within a larger coding conversation, was triggered by a user deploying models on a machine equipped with eight RTX PRO 6000 Blackwell GPUs. The user posed two precise questions to the assistant:

  1. Does the apply_flashinfer_allreduce_fusion() function in SGLang's communicator module support SM120 (Blackwell)? The user wanted exact version checks and code snippets, not summaries.
  2. Is the enable_fused_moe_sum_all_reduce flag actually checked at runtime in the fused MoE code? Specifically, does it only apply during the "down" MoE GEMM, and is it activated during CUDA graph capture? These questions cut to the heart of production inference performance. Allreduce fusion and fused MoE sum-reduce are optimizations that can dramatically reduce latency by overlapping communication with computation. If they silently disabled themselves on Blackwell hardware, the user's deployment would leave significant performance on the table.

The Search for SM Capability Checks: A Multi-Phase Investigation

Phase 1: Reading the Primary Source

The assistant began by reading the communicator file directly ([msg 1]). The apply_flashinfer_allreduce_fusion() function at lines 93–103 of communicator.py contained the critical gate:

def apply_flashinfer_allreduce_fusion(batch_size: int):
    return (
        (_is_sm90_supported or _is_sm100_supported)
        and _is_flashinfer_available
        and batch_size > 0
        and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE
        and not is_dp_attention_enabled()
        and get_global_server_args().enable_flashinfer_allreduce_fusion
    )

The module-level flags _is_sm90_supported and _is_sm100_supported were set at lines 73–74, imported from sglang.srt.utils. The function checked for SM90 (Hopper) or SM100 (Blackwell B100/B200) support — but notably absent was any check for SM120 (Blackwell Ultra / RTX 5090 / RTX PRO 6000). This was the first red flag: the flashinfer allreduce fusion optimization would silently disable itself on SM120 hardware.

But the assistant needed to verify this by tracing the actual SM support functions. This is where the investigation became a multi-round detective story.

Phase 2: The Dead End of utils.py

The assistant assumed, reasonably, that the SM support functions would live in sglang/srt/utils.py — a common pattern in Python projects ([msg 2]). The grep returned an error: the file did not exist. This was the first incorrect assumption.

In many Python projects, utility functions are collected in a single utils.py module. SGLang, however, organizes its utilities as a package: sglang/srt/utils/ is a directory containing multiple modules. The assistant's mental model of the codebase structure was wrong, and the search had to adapt.

Phase 3: Broadening the Net

The assistant pivoted to a broader search strategy ([msg 3]). Using find piped through xargs grep, it searched every Python file in the sglang/srt/ directory for def is_sm90_supported, def is_sm100_supported, or def is_sm120_supported. The result was silence — no matches at all.

A second command searched the entire sglang/ package for any def is_sm pattern. This returned only one hit: is_sm120 as a class method in multimodal_gen/runtime/platforms/interface.py — a completely unrelated module. This was a false positive that nonetheless confirmed SM120 awareness existed somewhere in the codebase, just not in the expected form.

Phase 4: Following the Import Trail

The assistant changed tactics ([msg 4]). Instead of searching for function definitions, it searched for function references — any file that used is_sm90_supported, is_sm100_supported, or is_sm120_supported. This was a clever shift: even if the definitions were hard to find, the import statements would reveal the module path.

The search found that test files and attention backends imported these functions from sglang.srt.utils. This confirmed the module existed but wasn't at the expected file path. The assistant then searched for files named utils.py in the SGLang tree, finding several — cli/utils.py, benchmark/utils.py, jit_kernel/utils.py — but none in the sglang/srt/utils/ directory.

Phase 5: The Breakthrough — Finding common.py

The assistant combined two insights ([msg 5]): the import path sglang.srt.utils implied a package directory, and the missing def definitions suggested the functions might be defined using a different pattern. The command was simple:

find /root/sglang-main/ -name 'common.py' -path '*/sglang/srt/utils/*'

The result: /root/sglang-main/python/sglang/srt/utils/common.py. This was the missing link. The SM capability functions were not in a flat utils.py; they were in a common.py submodule within the utils package.

Phase 6: The def Search Fails Again

With the correct file identified, the assistant ran a targeted grep for def is_sm90_supported etc. on common.py ([msg 6]). The result was empty — no matches. At this point, an experienced developer would suspect that the functions weren't defined with standard def statements. The assistant tried the same def search across the entire SGLang source tree ([msg 7]), finding nothing. This confirmed the hypothesis: the functions were created dynamically.

Phase 7: The Final Discovery

The assistant changed strategy completely ([msg 8]). Instead of searching for def, it searched for the function names themselves as strings:

grep -rn 'is_sm90_supported\|is_sm100_supported\|is_sm120_supported' /root/sglang-main/python/sglang/srt/utils/ --include='*.py'

The output revealed the truth:

/root/sglang-main/python/sglang/srt/utils/common.py:251:is_sm120_supported = lru_cache(maxsize=1)(
/root/sglang-main/python/sglang/srt/utils/common.py:256:is_sm100_supported = lru_cache(maxsize=1)(
/root/sglang-main/python/sglang/srt/utils/common.py:261:is_sm90_supported = lru_cache(maxsize=1)(

The functions were defined using a functional programming pattern: lru_cache(maxsize=1)(partial(_check_cuda_device_version, ...)). This created memoized, partially-applied versions of a helper function — an elegant pattern that was invisible to naive grep for def statements.

The Core Logic: _check_cuda_device_version

With the functions located, the assistant read their implementations ([msg 9], [msg 10], [msg 11]). The core helper function at lines 223–231 of common.py was:

def _check_cuda_device_version(
    device_capability_majors: List[int], cuda_version: Tuple[int, int]
):
    if not is_cuda():
        return False
    return (
        torch.cuda.get_device_capability()[0] in device_capability_majors
        and tuple(map(int, torch.version.cuda.split(".")[:2])) >= cuda_version
    )

This function checks two conditions: whether the GPU's compute capability major version matches a list of allowed values (e.g., [12] for SM120), and whether the CUDA runtime version meets a minimum threshold (e.g., (12, 8) for CUDA 12.8). The individual support functions were created by partially applying this helper with specific parameters:

Answering Question 1: Does apply_flashinfer_allreduce_fusion() Support SM120?

The answer was definitive: no. The function checked for _is_sm90_supported or _is_sm100_supported, which mapped to compute capability majors [9] and [10] respectively. SM120 (major [12]) was not included. Despite is_sm120_supported and is_blackwell_supported existing in the codebase, neither was wired into the flashinfer allreduce fusion gate. This meant that on Blackwell RTX PRO 6000 GPUs, this optimization would silently disable itself — a significant finding for the user's deployment.

Tracing the Fused MoE Sum-All-Reduce Flag

The second question required tracing the enable_fused_moe_sum_all_reduce flag through the MoE code path. The assistant systematically traced this flag across multiple files ([msg 12] through [msg 22]).

The Server Argument Definition

The flag was defined in server_args.py at line 659:

enable_fused_moe_sum_all_reduce: bool = False

A simple boolean server argument, defaulting to False. The user would need to explicitly enable it.

The Runtime Check in fused_experts_impl()

Inside fused_moe.py (lines 471–476), the flag was evaluated with additional conditions:

use_fused_moe_sum_all_reduce = (
    get_global_server_args().enable_fused_moe_sum_all_reduce
    and (not no_combine)
    and (curr_topk_ids.shape[1] > 2)
    and (not use_int8_w8a16)
    and (not use_int4_w4a16)
)

The fusion was only active when: the server flag was enabled, no_combine was false, the top-k routing dimension was greater than 2, and the model wasn't using INT8 or INT4 quantization. These were sensible constraints — the fusion optimization only makes sense when there are multiple expert outputs to combine, and it's only implemented for FP8/FP16 paths.

The "Down" GEMM Only

The assistant confirmed that the fusion only applied during the "down" (w2) projection — the second linear transformation in an MoE layer. The first GEMM (up/gate projection via w1) did not pass fuse_sum_all_reduce. The second GEMM invocation (lines 583–621) used the fused path when enabled:

if use_fused_moe_sum_all_reduce:
    out_slice = out_hidden_states[begin_chunk_idx:end_chunk_idx]
    out_slice.zero_()

The output buffer was zeroed and passed directly to the kernel, which would accumulate results via atomic operations.

The Triton Kernel Implementation

The actual fusion logic lived in the Triton kernel at fused_moe_triton_kernels.py. The kernel accepted a FUSE_SUM_ALL_REDUCE compile-time constant (line 380). When enabled (lines 607–614), instead of writing to a per-expert intermediate buffer, the kernel used tl.atomic_add to accumulate results directly into the output buffer:

if FUSE_SUM_ALL_REDUCE:
    offs_token_out = offs_token // ROUTER_TOPK
    c_ptrs = (
        c_ptr + stride_cm * offs_token_out[:, None] + stride_cn * offs_cn[None, :]
    )
    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
    tl.atomic_add(c_ptrs, accumulator, mask=c_mask)

This fused the per-expert sum-reduction into the down-projection GEMM itself, eliminating the separate moe_sum_reduce step that would otherwise be required. The division offs_token // ROUTER_TOPK mapped each expert's output token position back to the original token index, effectively summing the contributions from all experts that routed to that token.

CUDA Graph Capture Interaction

The assistant checked whether the flag had any special interaction with CUDA graph capture — a technique for recording and replaying GPU operations to reduce launch overhead. The search across the entire codebase found no CUDA-graph-specific gating. There was no torch.cuda.is_current_stream_capturing() check, no is_cuda_graph parameter, and no special disable logic during graph capture. The flag was evaluated identically whether the code was running in normal execution or during CUDA graph capture.

This was good news: the Triton kernel's use of tl.atomic_add is CUDA-graph-safe, so there was no fundamental incompatibility. But it also meant there was no special handling — the fusion would work the same way in both modes.

The Architecture of the Fused MoE All-Reduce

The investigation revealed a multi-layered architecture for the fused MoE sum-all-reduce optimization:

  1. Server configuration layer (server_args.py): The enable_fused_moe_sum_all_reduce flag, defaulting to False, controlled by the user.
  2. Runtime logic layer (fused_moe.py): The fused_experts_impl() function evaluated the flag alongside quantization and routing constraints, passing fuse_sum_all_reduce only to the down-projection GEMM.
  3. Kernel wrapper layer (fused_moe_triton_kernels.py): The invoke_fused_moe_kernel function accepted the fuse_sum_all_reduce parameter and passed it as a Triton compile-time constant.
  4. Triton kernel layer: The actual GPU kernel used FUSE_SUM_ALL_REDUCE to conditionally enable atomic-add-based accumulation, fusing the sum-reduction into the GEMM. This layered design is characteristic of well-architected inference engines: configuration at the top, logic in the middle, and hardware-aware kernels at the bottom.

The SM Capability Check Architecture

Similarly, the SM capability checking system revealed a clean architecture:

  1. Core comparison function (_check_cuda_device_version): A generic helper that checked device capability majors and CUDA version.
  2. Partially-applied wrappers: is_sm90_supported, is_sm100_supported, is_sm120_supported, and is_blackwell_supported created via lru_cache + partial.
  3. Module-level flags: _is_sm90_supported, _is_sm100_supported in communicator.py cached the results at import time.
  4. Feature gates: Functions like apply_flashinfer_allreduce_fusion() used these flags to decide whether to enable optimizations. The gap was at step 4: apply_flashinfer_allreduce_fusion() only checked SM90 and SM100, missing SM120 entirely despite the infrastructure existing to support it.

Broader Implications

This investigation revealed several important patterns for anyone deploying AI inference on new hardware:

The gap between infrastructure and integration: SGLang's codebase had all the building blocks for SM120 support — the is_sm120_supported function existed, the is_blackwell_supported function covered it, and the CUDA version requirements were defined. But the feature gates in apply_flashinfer_allreduce_fusion() hadn't been updated to use them. This is a common pattern in rapidly evolving codebases: the foundation is laid, but not all consumers are updated.

The importance of tracing the full chain: A naive check might have stopped at finding is_sm120_supported in common.py and concluded that SM120 was supported. Only by tracing the full chain — from the utility function to the module-level flag to the feature gate — could the assistant determine that the support existed at the utility level but wasn't wired into the optimization.

The elegance of functional composition: The use of lru_cache and partial to create memoized capability checks is a clean pattern. It avoids repetitive boilerplate, ensures consistent caching behavior, and makes it easy to add new compute capability checks. But it also makes the code harder to find with naive text searches — a tradeoff that developers should be aware of.

Conclusion

The investigation produced a clear, actionable answer for the user deploying on Blackwell GPUs: