From efdf02a38a83b0c819efe5a867f38cbafce624eb Mon Sep 17 00:00:00 2001 From: loading66 <102725781+loading66@users.noreply.github.com> Date: Tue, 7 Jul 2026 17:17:29 +0800 Subject: [PATCH] [NPU]Add support --pre-warm-nccl (#30312) --- python/sglang/srt/model_executor/model_runner.py | 6 +++--- python/sglang/srt/server_args.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ea17a2558..758c8395c 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1328,7 +1328,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): if is_npu(): register_sgl_tp_rank(self.gpu_id) - # Pre-warm NCCL/RCCL to eliminate cold-start latency in first request + # Pre-warm NCCL/RCCL/HCCL to eliminate cold-start latency in first request # Controlled by --pre-warm-nccl flag (default: enabled on AMD GPUs) if self.server_args.pre_warm_nccl and ( self.tp_size > 1 or self.pp_size > 1 or self.moe_ep_size > 1 @@ -1336,14 +1336,14 @@ class ModelRunner(ModelRunnerKVCacheMixin): warmup_start = time.perf_counter() tp_group_handle = get_tp_group().device_group - # Single warmup all_reduce to initialize NCCL/RCCL communicator + # Single warmup all_reduce to initialize NCCL/RCCL/HCCL communicator warmup_tensor = torch.zeros(1, device=torch.cuda.current_device()) dist.all_reduce(warmup_tensor, group=tp_group_handle) current_platform.synchronize() warmup_elapsed = time.perf_counter() - warmup_start logger.info( - f"NCCL/RCCL warmup completed in {warmup_elapsed:.3f}s " + f"NCCL/RCCL/HCCL warmup completed in {warmup_elapsed:.3f}s " f"(tp_size={self.tp_size}, pp_size={self.pp_size}, ep_size={self.moe_ep_size})" ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 04d96d34d..2bc8b30c8 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -4594,10 +4594,10 @@ class ServerArgs: self.triton_attention_num_kv_splits = 16 def _handle_nccl_pre_warm(self): - # pre_warm_nccl is only used with CUDA or HIP hardware - if self.pre_warm_nccl and not (is_cuda() or is_hip()): + # pre_warm_nccl is only used with CUDA or HIP hardware or NPU hardware + if self.pre_warm_nccl and not (is_cuda() or is_hip() or is_npu()): logger.warning( - "pre_warm_nccl is only applicable for CUDA or HIP hardware. " + "pre_warm_nccl is only applicable for CUDA or HIP hardware or NPU hardware. " "Ignoring pre_warm_nccl setting on current hardware." ) self.pre_warm_nccl = False