Patch NPU序列化以及减少update idle 检查

python/sglang/srt/managers/scheduler.py

@@ -110,6 +110,7 @@
     UpdateWeightFromDiskReqInput,
     UpdateWeightsFromDistributedReqInput,
     UpdateWeightsFromTensorReqInput,
+    UpdateWeightsFromTensorReqOutput
 )
 from sglang.srt.managers.mm_utils import init_embedding_cache
 from sglang.srt.managers.schedule_batch import (
@@ -191,7 +192,7 @@
 GRAMMAR_TIMEOUT = float(os.environ.get("SGLANG_GRAMMAR_TIMEOUT", 300))
 
 _is_cpu = is_cpu()
-
+LAST_UPDATE = False
 
 @dataclass
 class GenerationBatchResult:
@@ -855,10 +856,13 @@ def event_loop_normal(self):
     def event_loop_overlap(self):
         """A scheduler loop that overlaps the CPU processing and GPU computation."""
         self.result_queue = deque()
+        global LAST_UPDATE
 
         while True:
             recv_reqs = self.recv_requests()
             self.process_input_requests(recv_reqs)
+            if LAST_UPDATE:
+                continue
 
             batch = self.get_next_batch_to_run()
             self.cur_batch = batch
@@ -1139,6 +1143,7 @@ def recv_requests(self) -> List[Req]:
         return recv_reqs
 
     def process_input_requests(self, recv_reqs: List):
+        global LAST_UPDATE
         for recv_req in recv_reqs:
             # If it is a health check generation request and there are running requests, ignore it.
             if is_health_check_generate_req(recv_req) and (
@@ -1175,6 +1180,10 @@ def process_input_requests(self, recv_reqs: List):
 
             output = self._request_dispatcher(recv_req)
             if output is not None:
+                if isinstance(output, UpdateWeightsFromTensorReqOutput):
+                    LAST_UPDATE = True
+                else:
+                    LAST_UPDATE = False
                 if isinstance(output, RpcReqOutput):
                     if self.recv_from_rpc is not None:
                         self.recv_from_rpc.send_pyobj(output)

python/sglang/srt/model_executor/model_runner.py

@@ -673,6 +673,9 @@ def _(data, dim):
                 server_args=self.server_args,
                 model_config=self.model_config,
             )
+            if is_npu():
+                from sglang.srt.patch_torch import register_sgl_tp_rank
+                register_sgl_tp_rank(self.gpu_id)
 
         min_per_gpu_memory = get_available_gpu_memory(
             self.device,

python/sglang/srt/patch_torch.py

@@ -19,22 +19,33 @@
 
 from sglang.srt.utils import is_npu
 
+if is_npu():
+    from torch_npu.multiprocessing import reductions as npu_reductions
+
+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.
+    # Currently, NPU does not support UUID. This is a temporary fix, with support expected in the fourth quarter.
     if is_npu():
-        return
-
-    if hasattr(reductions, "_reduce_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
+        '''
+        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
+
+        npu_reductions._rebuild_npu_tensor_original = npu_reductions.rebuild_npu_tensor
+        npu_reductions.rebuild_npu_tensor = _rebuild_npu_tensor_modified
+
+    else:
+        if hasattr(reductions, "_reduce_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()
 
@@ -44,6 +55,21 @@ def monkey_patch_torch_reductions():
 _REDUCE_TENSOR_ARG_DEVICE_INDEX = 6
 
 
+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 register_sgl_tp_rank(rank: int):
+    global SGLANG_TP_RANK
+    SGLANG_TP_RANK = rank
+
+
+def npu_verl_to_sglang(device: int):
+    assert SGLANG_TP_RANK is not None
+    return SGLANG_TP_RANK
+
+
 def _reduce_tensor_modified(*args, **kwargs):
     output_fn, output_args = reductions._reduce_tensor_original(*args, **kwargs)
     output_args = _modify_tuple(
@@ -85,4 +111,4 @@ def monkey_patch_torch_compile():
         import torch._higher_order_ops.auto_functionalize as af
 
         af.auto_functionalized_v2._cacheable = True
-        af.auto_functionalized._cacheable = True
+        af.auto_functionalized._cacheable = True