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

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 _is_npu:
-        return
+    if not _is_npu:
+        if hasattr(reductions, "_reduce_tensor_original"):
+            return
+        reductions._reduce_tensor_original = reductions.reduce_tensor
+        reductions._rebuild_cuda_tensor_original = reductions.rebuild_cuda_tensor
 
-    if hasattr(reductions, "_reduce_tensor_original"):
-        return
+        reductions.reduce_tensor = _reduce_tensor_modified
+        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
-    reductions._rebuild_cuda_tensor_original = reductions.rebuild_cuda_tensor
-
-    reductions.reduce_tensor = _reduce_tensor_modified
-    reductions.rebuild_cuda_tensor = _rebuild_cuda_tensor_modified
-
-    reductions.init_reductions()
+        npu_reductions._rebuild_npu_tensor_original = npu_reductions.rebuild_npu_tensor
+        npu_reductions.rebuild_npu_tensor = _rebuild_npu_tensor_modified
 
 
 # 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(