[tech-cow]

Motivation

Fixes #28999.

SGLang won't start on Ascend NPU (torch 2.7.0 + torch_npu):

ImportError: cannot import name '_cuda_beginAllocateCurrentThreadToPool' from 'torch.cuda.memory'

pynccl_allocator.py pulls four private torch.cuda.memory memory-pool names in at module scope. Two of them (_cuda_beginAllocateCurrentThreadToPool, _cuda_endAllocateToPool) only landed in torch 2.8 — 2.7 spelled the second one _cuda_endAllocateCurrentStreamToPool. A from ... import (a, b, c) is all-or-nothing, so one missing name takes the whole statement down with it. And this module gets imported eagerly all over the place — model_runner, linear, dp_attention, fp8, the MoE/quant layers, even the NPU graph runners — so the ImportError fires long before anything gets a chance to dispatch on device.

Modifications

Turns out we didn't need a guarded import at all. torch.cuda.memory re-exports these straight from torch._C, and this same file was already reaching two of them that way, inside the very methods that were using the imported bindings. So the remaining call sites just go through torch._C.<name> like their siblings, and the import-time dependency disappears. No try/except, and no None sentinels that could quietly turn a genuinely broken CUDA install into a confusing NoneType is not callable somewhere further down.

CUDAPluggableAllocator stays a plain import — it's a Python class that's been there since torch 2.0, and it was only collateral damage of sharing an import statement with the private names.

Two guards on top, so the unsupported paths fail on purpose instead of by accident:

• --enable-symm-mem now warns and turns itself off on non CUDA/HIP devices, the same way --debug-cuda-graph already does. NCCL symmetric memory compiles a CUDA plugin and links -lnccl, so on NPU it could only ever have blown up deep inside a build step. Nothing rejected the flag there before — that path was "safe" purely because nobody passed it.
• get_nccl_mem_pool asserts torch>=2.8 up front, instead of letting you find out via a bare AttributeError halfway through a context manager.

The test is now a static AST check that no torch>=2.8-only symbol appears in a module-scope import. It catches the regression on any torch version, needs no GPU or NPU, and doesn't mutate shared torch module state.

This version is @alexnails' — it replaces the try/except + None fallback the PR originally carried.

Accuracy Tests

N/A — import-time only. Every CUDA call site invokes the exact same C function it did before, just looked up on torch._C instead of through a module-scope binding.

Speed Tests and Profiling

N/A — no runtime or critical-path change.

Checklist

• Follow the SGLang code style guidance.
• Add unit tests (test/registered/unit/distributed/test_pynccl_allocator_import.py).
• Documentation update — N/A.
• Accuracy/speed benchmarks — N/A.

---

CI States

Latest PR Test (Base): ❌ Run #33275009508
Latest PR Test (Extra): ❌ Run #33275009302
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33275009607

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[JosephNew]

We hit the exact same ImportError on Ascend 910B4 (torch 2.7.0 + torch_npu), and the guarded-import direction here looks right — consistent with @ronhuafeng's point that all four symbols must be handled together, not just one.

We have an Ascend 910B environment and are glad to validate this on real NPU (which I understand neither @ronhuafeng nor @tech-cow could do from a CUDA-only setup). One note on full end-to-end validation: once this import fix lets SGLang start, inference on Ascend also trips a separate FlagGems bug in argmax (a duplicate @libentry()), which we just filed + PR'd at flagos-ai/FlagGems#4466 / flagos-ai/FlagGems#4467. Once both land, we can confirm SGLang starts + serves + runs inference end-to-end on 910B and report back here.

Re @ronhuafeng's second point (an explicit error when symmetric memory is actually requested on a non-CUDA backend): the current None fallback is safe only because the path is gated behind --enable-symm-mem; a clear "NCCL symmetric memory unavailable on this backend" error at the call site would be a nice follow-up hardening.

[alexnails]

https://github.com/alexnails/sglang/tree/fix/pynccl-allocator-lazy-cuda-symbols

[alexnails]

Setting None could cause odd behavior that would not make sense to bad CUDA install, wdyt of branch change?

[tech-cow]

Agreed, your branch is the better fix. I've pushed it onto this PR as your commit and updated the title and description to match.

What convinced me: torch.cuda.memory re-exports these straight from torch._C, and this file was already reaching two of them that way, inside the same methods that used the imported bindings. The module-scope import was never load-bearing, so dropping it is a runtime no-op. It also drops the None sentinel that would turn a broken CUDA install into a confusing NoneType is not callable further down. That was your point. CUDAPluggableAllocator stays a plain import; it's been there since torch 2.0 and only got swept into the fallback because it shared an import statement with the private names.

The two guards also close @ronhuafeng's and @JosephNew's second point here instead of deferring it. On a non CUDA/HIP device, --enable-symm-mem now warns and disables itself; before, it would have failed deep inside load_inline + -lnccl. And get_nccl_mem_pool asserts torch>=2.8, so no more AttributeError mid-context-manager.

One small note, not a blocker. With that assert in place, the after_2_8_0 else-branches in __enter__/__exit__ (_cuda_endAllocateCurrentStreamToPool / _cuda_beginAllocateToPool) are dead code. They already were, I think: the unconditional _cuda_beginAllocateCurrentThreadToPool call right below means symm-mem could never have worked on 2.7 anyway. Happy to strip them here or leave it for a follow-up, your call.

@JosephNew, the 910B validation offer is still welcome and it carries over unchanged: this fixes the same import path either way, so a run against the new commit is just as valid. Watching flagos-ai/FlagGems#4467 for the argmax side.