diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index de2ad078d19c..7c2fc5cab623 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 2cb27bcbefc1..784f6038821a 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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, diff --git a/python/sglang/srt/patch_torch.py b/python/sglang/srt/patch_torch.py index e1bbdd5b7af1..10ee0280d9e7 100644 --- a/python/sglang/srt/patch_torch.py +++ b/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 \ No newline at end of file