[NPU]Fix for ipc handle with npu (#14138)

python/sglang/srt/model_executor/model_runner.py

@@ -182,7 +182,10 @@ from sglang.srt.utils.offloader import (
     get_offloader,
     set_offloader,
 )
-from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
+from sglang.srt.utils.patch_torch import (
+    monkey_patch_torch_reductions,
+    register_sgl_tp_rank,
+)
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 from sglang.srt.utils.weight_checker import WeightChecker
 from sglang.srt.weight_sync.tensor_bucket import (
@@ -759,6 +762,8 @@ class ModelRunner:
                 server_args=self.server_args,
                 model_config=self.model_config,
             )
+            if is_npu():
+                register_sgl_tp_rank(self.gpu_id)
 
         min_per_gpu_memory = get_available_gpu_memory(
             self.device,

python/sglang/srt/utils/patch_torch.py

@@ -21,24 +21,42 @@ from sglang.srt.utils.common import is_npu
 
 _is_npu = is_npu()
 
+if _is_npu:
+    from torch_npu.multiprocessing import reductions as npu_reductions
+
+    def _rebuild_npu_tensor_modified(*args):
+        args = _modify_tuple(args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, npu_verl_to_sglang)
+        return npu_reductions._rebuild_npu_tensor_original(*args)
+
+    def npu_verl_to_sglang(device: int):
+        assert (
+            SGLANG_TP_RANK is not None
+        ), "SGLANG_TP_RANK is not registered. Please call register_sgl_tp_rank() first."
+        return SGLANG_TP_RANK
+
+
+SGLANG_TP_RANK = None
+
 
 def monkey_patch_torch_reductions():
     """Monkey patching before Torch https://github.com/pytorch/pytorch/pull/149248 is fixed"""
 
-    # Currently, NPU does not support UUID. This has been temporarily commented out, with support expected in the fourth quarter.
+    if not _is_npu:
-    if _is_npu:
+        if hasattr(reductions, "_reduce_tensor_original"):
-        return
+            return
+        reductions._reduce_tensor_original = reductions.reduce_tensor
+        reductions._rebuild_cuda_tensor_original = reductions.rebuild_cuda_tensor
 
-    if hasattr(reductions, "_reduce_tensor_original"):
+        reductions.reduce_tensor = _reduce_tensor_modified
-        return
+        reductions.rebuild_cuda_tensor = _rebuild_cuda_tensor_modified
+        reductions.init_reductions()
+    else:
+        # FIXME: This is a temp patch for npu as HDK does not support device uuid for now
+        if hasattr(npu_reductions, "_rebuild_npu_tensor_original"):
+            return
 
-    reductions._reduce_tensor_original = reductions.reduce_tensor
+        npu_reductions._rebuild_npu_tensor_original = npu_reductions.rebuild_npu_tensor
-    reductions._rebuild_cuda_tensor_original = reductions.rebuild_cuda_tensor
+        npu_reductions.rebuild_npu_tensor = _rebuild_npu_tensor_modified
-
-    reductions.reduce_tensor = _reduce_tensor_modified
-    reductions.rebuild_cuda_tensor = _rebuild_cuda_tensor_modified
-
-    reductions.init_reductions()
 
 
 # The signature has not been changed for years, and we will not need this when the next version is released,
@@ -46,6 +64,11 @@ def monkey_patch_torch_reductions():
 _REDUCE_TENSOR_ARG_DEVICE_INDEX = 6
 
 
+def register_sgl_tp_rank(rank: int):
+    global SGLANG_TP_RANK
+    SGLANG_TP_RANK = rank
+
+
 def _reduce_tensor_modified(*args, **kwargs):
     output_fn, output_args = reductions._reduce_tensor_original(*args, **kwargs)
     output_args = _modify_tuple(