[NPU] fix: reach torch>=2.8 CUDA memory-pool APIs lazily via torch._C (#29100)

Co-authored-by: Alex Nails <alex.nails@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/srt/arg_groups/pipeline.py

@@ -199,6 +199,7 @@ def run_resolution_pipeline(server_args: Any) -> None:
         handle_mps_backends,
         handle_nccl_pre_warm,
         handle_npu_backends,
+        handle_symm_mem_device_support,
         handle_xpu_backends,
     )
 
@@ -207,6 +208,9 @@ def run_resolution_pipeline(server_args: Any) -> None:
     handle_npu_backends(server_args)
     handle_mps_backends(server_args)
     handle_xpu_backends(server_args)
+    # Must precede handle_gpu_memory_settings: its symm-mem prealloc default
+    # keys off enable_symm_mem.
+    handle_symm_mem_device_support(server_args)
 
     # OOT platform plugins set fields directly (an interface this tree
     # does not own); the diff records what they applied.

python/sglang/srt/arg_groups/platform_hook.py

@@ -76,6 +76,20 @@ def handle_nccl_pre_warm(server_args: Any):
         declare_resolution(server_args, "_handle_nccl_pre_warm", pre_warm_nccl=False)
 
 
+def handle_symm_mem_device_support(server_args: Any):
+    cfg = resolving_view(server_args)
+    # The symm-mem allocator compiles a CUDA plugin and links -lnccl, so off
+    # CUDA/HIP (e.g. Ascend NPU) it fails deep in a build step rather than here.
+    if cfg.enable_symm_mem and not (is_cuda() or is_hip()):
+        logger.warning(
+            "--enable-symm-mem is not supported on non CUDA/HIP devices "
+            "(NCCL symmetric memory is unavailable). Disabling symmetric memory."
+        )
+        declare_resolution(
+            server_args, "_handle_symm_mem_device_support", enable_symm_mem=False
+        )
+
+
 def handle_xpu_backends(server_args: Any):
     cfg = resolving_view(server_args)
     if cfg.device == "xpu":

python/sglang/srt/distributed/device_communicators/pynccl_allocator.py

@@ -6,12 +6,10 @@ import traceback
 from contextlib import nullcontext
 
 import torch
-from torch.cuda.memory import (
-    CUDAPluggableAllocator,
-    _cuda_beginAllocateCurrentThreadToPool,
-    _cuda_endAllocateToPool,
-    _cuda_releasePool,
-)
+
+# The private _cuda_* pool APIs are absent before torch 2.8; the call sites below
+# reach them via torch._C.<name> so torch 2.7 (Ascend NPU) can still import this.
+from torch.cuda.memory import CUDAPluggableAllocator
 
 from sglang.srt.distributed.parallel_state import GroupCoordinator
 from sglang.srt.environ import envs
@@ -189,6 +187,11 @@ def get_nccl_mem_pool() -> torch.cuda.MemPool:
     All groups share the same pool to avoid memory fragmentation.
     Comm registration is handled at context exit time.
     """
+    assert after_2_8_0, (
+        "--enable-symm-mem requires torch>=2.8 "
+        "(torch._C._cuda_beginAllocateCurrentThreadToPool was added there)."
+    )
+
     global _allocator, _mem_pool, _cur_device, _register_func
     if _allocator is None:
         import torch.utils.cpp_extension
@@ -279,7 +282,9 @@ class SymmetricMemoryContext:
                     _cur_device, _graph_pool_id
                 )
 
-        _cuda_beginAllocateCurrentThreadToPool(self._device_index, self._pool_id)
+        torch._C._cuda_beginAllocateCurrentThreadToPool(
+            self._device_index, self._pool_id
+        )
 
         global _active_symmetric_memory_context
         _active_symmetric_memory_context = self
@@ -287,8 +292,8 @@ class SymmetricMemoryContext:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        _cuda_endAllocateToPool(self._device_index, self._pool_id)
-        _cuda_releasePool(self._device_index, self._pool_id)
+        torch._C._cuda_endAllocateToPool(self._device_index, self._pool_id)
+        torch._C._cuda_releasePool(self._device_index, self._pool_id)
         # Register all unregistered segments
         # with the current comm
         self._register_segments_for_comm()