The Green Light: Validating CUDA Kernel Portability Across GPU Architectures

In the high-stakes world of large-scale ML inference optimization, few moments are as satisfying as the one captured in message 774 of this opencode session. After hours of wrestling with NaN crashes, PCIe bottlenecks, and architecture-gated code paths, the assistant issues a single, decisive bash command — a grep across three CUDA source files — and receives the answer it has been hoping for: zero matches for architecture-specific code. This message is a turning point, the moment where a risky patching strategy is validated and the path forward becomes clear.

The Message

The command is straightforward but surgically precise:

ssh root@10.1.230.174 "grep -c 'sm_90\|sm_100\|__SM_\|arch.*90\|arch.*100\|__CUDA_ARCH__' \
  /root/ml-env/lib/python3.12/site-packages/flashinfer/data/csrc/trtllm_allreduce.cu \
  /root/ml-env/lib/python3.12/site-packages/flashinfer/data/csrc/trtllm_allreduce_fusion.cu \
  /root/ml-env/lib/python3.12/site-packages/flashinfer/data/csrc/trtllm_moe_allreduce_fusion.cu"

The result: all three files return a count of zero. Not a single instance of sm_90, sm_100, __SM_, arch.*90, arch.*100, or __CUDA_ARCH__ exists in any of the three kernel source files. This is the green light the assistant needed.

The Journey to This Point

To understand why this simple grep matters so much, we must trace the path that led here. The assistant had been deploying the GLM-5-NVFP4 model — a massive Mixture-of-Experts (MoE) language model quantized to NVFP4 — across eight RTX PRO 6000 Blackwell GPUs. After resolving an initial NaN crash during decode (by selecting the correct NSA backends), the team achieved a baseline throughput of around 880 tokens per second. But GPU power draw hovered around 250W out of a 600W TDP, indicating severe underutilization.

The primary suspect was the allreduce operation — the communication step where gradients and activations are summed across GPUs. In a tensor-parallel configuration spanning eight GPUs, every MoE layer requires an allreduce to combine expert outputs. On a system where GPUs are connected via PCIe (especially in a virtualized Proxmox environment with multiple PCIe root complexes), this allreduce becomes a bottleneck. FlashInfer, the high-performance inference library powering SGLang, offers a solution: allreduce fusion, which overlaps the allreduce communication with ongoing computation, hiding latency.

The problem? FlashInfer's allreduce fusion was explicitly gated to only work on NVIDIA GPU architectures SM90 (Hopper, e.g., H100) and SM100 (datacenter Blackwell, e.g., B200). The RTX PRO 6000 Blackwell uses SM120 — the consumer/professional variant of the Blackwell architecture. Despite being built on the same fundamental Blackwell microarchitecture, SM120 was excluded from the fusion path.

The Patching Strategy

The assistant's first attempt was straightforward: add is_sm120_supported() checks alongside the existing is_sm90_supported() and is_sm100_supported() calls in SGLang's communicator.py and server_args.py. This seemed logical — the function is_sm120_supported() already existed and returned True on the target hardware. But when the server was restarted, it crashed immediately with a RuntimeError: No supported CUDA architectures found for major versions [9, 10].

The error originated deep inside FlashInfer's JIT compilation pipeline. The gen_trtllm_comm_module function in flashinfer/jit/comm.py explicitly called get_nvcc_flags_list(supported_major_versions=[9, 10]), which filtered the available CUDA architectures to only SM90 and SM100. SM120 was present in the compilation context's TARGET_CUDA_ARCHS list (detected from the actual GPU hardware), but it was filtered out by the version gate. The JIT compilation literally refused to compile the kernel for SM120.

The assistant reverted the changes and marked the approach as "BLOCKED (kernel not compiled for SM120)." But then the user intervened with a crucial directive in message 768: "Please think big and don't be afraid to fork/modify code." This changed the calculus entirely. The assistant was no longer limited to configuration changes — it could modify the library source code itself.

The Critical Assumption

With the user's blessing, the assistant patched flashinfer/jit/comm.py directly, changing supported_major_versions=[9, 10] to supported_major_versions=[9, 10, 12]. This was a bold move. The version gate existed for a reason — presumably because the TRT-LLM communication kernels contained architecture-specific code (PTX instructions, warp-level primitives, shared memory configurations) that only worked on SM90 and SM100. Adding SM120 without verification could cause silent data corruption, incorrect results, or kernel crashes at runtime.

But was this assumption correct? The assistant needed evidence. The version gate could be there for two reasons: either the kernels genuinely required SM90/SM100-specific features (like specific tensor core instructions or shared memory sizes), or the gate was simply conservative — the original authors hadn't tested on SM120 and chose to restrict to known-working architectures.

The Diagnostic

Message 774 is the diagnostic that answers this question. The assistant searches for any architecture-specific code in the three kernel source files:

trtllm_allreduce.cu — the core allreduce kernel
trtllm_allreduce_fusion.cu — the fused allreduce kernel that overlaps communication with computation
trtllm_moe_allreduce_fusion.cu — the MoE-specific allreduce fusion variant The grep pattern is carefully constructed to catch any form of architecture dependence: explicit architecture names (sm_90, sm_100), macro-based architecture checks (__SM_), pattern-matched architecture references (arch.*90, arch.*100), and the universal CUDA architecture macro (__CUDA_ARCH__). The -c flag returns a count rather than the matching lines, providing a quick yes/no answer. The result — zero matches across all three files — is profound. It means these kernels contain no SM-specific intrinsics, no architecture-gated code paths, no conditional compilation based on CUDA architecture version. They are written in portable CUDA C++ that should compile and execute correctly on any NVIDIA GPU architecture that supports the required features (concurrent kernel execution, device-side synchronization, and peer-to-peer memory access).

Assumptions and Risks

The assistant made several assumptions in this approach. First, that the absence of architecture-specific code in the source files guarantees correct execution on SM120. While the CUDA code itself may be portable, the compiled PTX (Parallel Thread Execution) ISA could still contain instructions that behave differently on SM120. NVIDIA's CUDA compiler (nvcc) compiles to an intermediate PTX representation, which is then optimized by the GPU driver for the specific architecture. Differences in warp scheduler behavior, shared memory bank conflicts, or cache hierarchy could still cause issues.

Second, the assistant assumed that the allreduce fusion kernels' synchronization primitives (likely using CUDA events or cooperative groups) would work correctly on SM120. The Blackwell consumer architecture (SM120) has different shared memory partitioning and warp scheduling than the datacenter variant (SM100). A kernel that relies on specific timing assumptions could deadlock or produce incorrect results.

Third, there was an implicit assumption that the version gate was purely conservative rather than based on known incompatibilities. In many open-source projects, architecture gates are added reactively — a developer tests on SM90, it works, so they gate it to SM90. Later, someone tests on SM100, it also works, so SM100 is added. SM120 may simply never have been tested.

The Thinking Process

The assistant's reasoning here reveals a methodical approach to risk assessment. Rather than blindly patching the version gate and hoping for the best, it first investigated the root cause of the crash (message 755), traced the error to the JIT compilation context (message 763-765), identified the specific version filter (message 766), and only then — after the user's encouragement — modified the library code. But even after patching, the assistant didn't immediately restart the server. It paused to verify the underlying assumption: that the kernels themselves are architecture-agnostic.

This is the hallmark of a careful engineer. The grep command is cheap (it runs in milliseconds), non-destructive (it only reads files), and provides high-value information. If the grep had found architecture-specific code, the assistant would have known that simply adding SM120 to the version gate was insufficient — the kernels would need actual code changes to support the new architecture. If it found nothing, the patched version gate had a reasonable chance of working.

Output Knowledge Created

This message creates critical knowledge: the TRT-LLM allreduce kernels in FlashInfer are architecture-portable at the source level. This knowledge has several implications:

The patched version gate is likely safe. Adding SM120 to supported_major_versions should produce a correctly compiled kernel.
The allreduce fusion bottleneck on SM120 is a policy limitation, not a technical one. The hardware is capable; it was simply excluded by convention.
Future architecture additions (SM130, SM200, etc.) may also work without kernel changes. The same pattern of patching the version gate could apply.
The performance characteristics remain unknown. Even if the kernel compiles and runs correctly, its performance on SM120 (with different shared memory sizes and warp schedulers) may differ from SM90/SM100. The allreduce fusion might not provide the same speedup.

The Broader Context

This message sits at the intersection of several themes in the session: the challenge of adapting datacenter-optimized code to consumer hardware, the willingness to fork and modify upstream dependencies, and the iterative cycle of hypothesis-testing-validation. The session had already seen the assistant modify SGLang's server_args.py, communicator.py, and now flashinfer's comm.py and compilation_context.py. Each modification was preceded by careful investigation and followed by verification.

The zero-result grep in message 774 doesn't just validate the patch — it embodies a philosophy. When faced with a "BLOCKED" status, the assistant didn't give up. It traced the blockage to its source, questioned the assumptions behind it, gathered evidence, and found a path forward. The user's encouragement to "think big" was the catalyst, but the technical execution — the precise diagnostic, the careful verification, the risk-aware patching — is what made the approach viable.

Conclusion

Message 774 is a small command with outsized significance. In the span of a single grep invocation, the assistant transforms a blocked optimization path into a viable one. The zero matches across three kernel files reveal that the allreduce fusion gate on SM120 is not a technical necessity but a policy choice — one that can be overridden with a simple source patch. This discovery unlocks the next phase of optimization: testing whether allreduce fusion actually improves throughput on SM120 hardware, and if so, by how much.

The message also serves as a reminder that in systems engineering, the most important questions are often the simplest: "Is this actually architecture-specific, or does someone just assume it is?" A well-crafted grep can sometimes tell you more than hours of speculation.