The Last Patch: Completing SM120 Support in SGLang's Torch Symmetric Memory

In the long arc of optimizing speculative decoding on an 8× RTX PRO 6000 Blackwell GPU system, few moments capture the essence of systems engineering better than message [msg 5371]. It is a brief, almost anticlimactic message — a single sed command and a grep verification — yet it represents the culmination of a multi-hour debugging odyssey. The message reads:

Need to add SM120 here too. PCIe topology probably won't support multimem (that's NVLink), but let's add it anyway for completeness: ``bash ssh root@10.1.230.174 'sed -i "s/ 10: \[6, 8\],/ 10: [6, 8],\n 12: [6, 8],/" /root/sglang/python/sglang/srt/distributed/device_communicators/torch_symm_mem.py && grep -A5 "_WORLD_SIZES_MULTIMEM" /root/sglang/python/sglang/srt/distributed/device_communicators/torch_symm_mem.py | head -10' _WORLD_SIZES_MULTIMEM = { 9: [4, 6, 8], 10: [6, 8], 12: [6, 8], } ``

To understand why this message exists, we must trace the chain of events that led to it. The project had just completed a major CUDA stack upgrade from version 12.8 to 13.0.1 — a delicate operation that involved navigating ABI compatibility issues, rebuilding PyTorch, sgl-kernel, and FlashInfer against the new toolkit. The payoff was immediate: the baseline inference throughput jumped from 89.5 tok/s to 92.6 tok/s, a 3.5% improvement. But the real prize was unblocking two Blackwell-native optimizations that had been dead ends under CUDA 12: FlashInfer allreduce fusion and Torch symmetric memory.

The Two Dead Ends

Under CUDA 12, both optimizations had failed catastrophically. FlashInfer allreduce fusion crashed with No supported CUDA architectures found for major versions [9, 10] — the prebuilt wheels simply had no SM120 (Blackwell) support. Torch symmetric memory crashed with a cryptic KeyError: 12 deep in SGLang's distributed communication layer. The CUDA 13 upgrade was supposed to fix both.

And indeed, after the upgrade, FlashInfer allreduce fusion worked. The server started, benchmarks ran, and the feature was validated — even if it didn't improve baseline throughput (allreduces are already overlapped with compute in the single-stream case), it was expected to help the EAGLE-3 verify pass where 122 tiny allreduces dominate latency.

Torch symmetric memory, however, still crashed. The KeyError: 12 persisted. This was the critical clue: the error was not a CUDA toolkit issue but an SGLang code issue. The assistant traced the error to torch_symm_mem.py, which queries torch.cuda.get_device_capability(device)[0] — returning 12 for Blackwell GPUs (compute capability 12.0) — and then uses that value as a dictionary key. The dictionary, TORCH_SYMM_MEM_ALL_REDUCE_MAX_SIZES, defined in all_reduce_utils.py, only had entries for compute capabilities 9 (Hopper/Ada) and 10 (SM100). Capability 12 was missing.

The First Fix

In message [msg 5368], the assistant patched all_reduce_utils.py by adding a 12 entry to TORCH_SYMM_MEM_ALL_REDUCE_MAX_SIZES, using the same values as capability 10: 64 MB for 2 and 4 GPUs, 128 MB for 6 and 8 GPUs. This was a straightforward fix — a missing key in a lookup table. But the assistant, thorough as always, then asked: "Now let me also check if there are other places that need the capability 12 mapping" ([msg 5369]).

The Second Dictionary

The grep revealed a second dictionary: _WORLD_SIZES_MULTIMEM, defined in torch_symm_mem.py itself. This dictionary maps compute capability to a list of world sizes that support the "multimem" allreduce path — a faster, single-shot allreduce that uses NVIDIA's symmetric memory (symm_mem) feature. If the current world size is in the list for the detected compute capability, the code takes the multimem path; otherwise, it falls back to a two-shot approach.

The dictionary had entries for capability 9 (world sizes [4, 6, 8]) and capability 10 (world sizes [6, 8]). Capability 12 was absent. If the assistant had stopped after the first fix, the server would still crash — just at a different line, with the same KeyError: 12.

The Reasoning in Message 5371

This brings us to the subject message. The assistant identifies the second dictionary and prepares to patch it. But then something interesting happens: the assistant pauses to reason about whether the patch is even useful.

PCIe topology probably won't support multimem (that's NVLink), but let's add it anyway for completeness.

This is a moment of genuine engineering judgment. The multimem allreduce path requires NVLink — NVIDIA's high-speed GPU-to-GPU interconnect. The system in question uses 8× RTX PRO 6000 GPUs connected via PCIe, not NVLink. (The GPUs are spread across two PCIe switches in a Proxmox container, as established earlier in the session.) Without NVLink, the multimem path would likely fail at runtime even if the dictionary lookup succeeds. The assistant knows this.

So why add the entry? Three reasons. First, correctness: the immediate problem is a KeyError crash, and adding the entry prevents that crash regardless of whether the multimem path ultimately works. Second, future-proofing: if the hardware configuration changes (e.g., NVLink-connected Blackwell GPUs become available), the code will work without modification. Third, completeness: the assistant had already added SM120 to TORCH_SYMM_MEM_ALL_REDUCE_MAX_SIZES in all_reduce_utils.py; leaving _WORLD_SIZES_MULTIMEM unpatched would be inconsistent and would cause a confusing crash for anyone who enabled both features.

The Execution

The fix itself is a single sed command, executed over SSH on the remote server:

sed -i "s/        10: \[6, 8\],/        10: [6, 8],\n        12: [6, 8],/" /root/sglang/python/sglang/srt/distributed/device_communicators/torch_symm_mem.py

This inserts a new line 12: [6, 8], after the existing 10: [6, 8], line. The values [6, 8] are copied from capability 10, meaning that on Blackwell systems, the multimem path is enabled for 6-GPU and 8-GPU configurations — the same as SM100. (Capability 9 also includes [4], but the assistant chose not to add that for capability 12, perhaps because 4-GPU Blackwell configurations are uncommon or because the PCIe topology concern made it moot.)

The grep verification confirms the patch was applied correctly, showing the updated dictionary with all three entries.

Input Knowledge Required

To understand this message, one needs several layers of knowledge. First, CUDA compute capability numbering: SM120 corresponds to Blackwell GPUs (compute capability 12.0), while Hopper/Ada are 9.x and SM100 is 10.x. Second, the architecture of SGLang's distributed communication: the torch_symm_mem.py module implements allreduce using PyTorch's symmetric memory operations, with two code paths (multimem and two-shot) selected based on the GPU topology. Third, the difference between NVLink and PCIe interconnects: NVLink enables direct GPU-to-GPU memory access (needed for multimem), while PCIe requires data to pass through the CPU and system memory. Fourth, the specific error signature (KeyError: 12) and how it propagates through Python's dictionary lookup.

Output Knowledge Created

The patch creates a corrected SGLang installation that can start with --enable-torch-symm-mem on Blackwell GPUs without crashing. It also documents (through the commit history and the assistant's reasoning) the assumption that SM120 should use the same multimem world sizes as SM100. Perhaps most importantly, it captures the engineering judgment that even though the feature may not work on PCIe-connected Blackwell GPUs, the code should not crash — it should either work or gracefully degrade.

The Broader Significance

Message [msg 5371] is the final piece of a larger puzzle. The CUDA 13 upgrade, the all_reduce_utils.py patch, and this torch_symm_mem.py patch together form a complete set of changes needed to enable Blackwell-native optimizations in SGLang. After this patch, the assistant was able to restart the server with both --enable-flashinfer-allreduce-fusion and --enable-torch-symm-mem and achieve the breakthrough result: EAGLE-3 speculative decoding throughput jumping from a net-negative 54.1 tok/s to a net-positive 96.1 tok/s — a 77.6% improvement.

The message also illustrates a recurring pattern in systems engineering: the second bug. The first fix (adding SM120 to TORCH_SYMM_MEM_ALL_REDUCE_MAX_SIZES) was necessary but not sufficient. The real skill is not just fixing the first error you see, but asking "are there other places that need the same fix?" — and then methodically checking. The assistant's grep for _WORLD_SIZES_MULTIMEM\|device_capability in message [msg 5369] is what turned a partial fix into a complete one.

In the end, this small message — a single sed command with a thoughtful comment about PCIe and NVLink — encapsulates what makes the session successful: not just the ability to execute technical steps, but the judgment to know when a fix is worth doing even if its benefits are uncertain, and the thoroughness to ensure no second dictionary is left behind.