diff --git a/python/sglang/srt/hardware_backend/npu/utils.py b/python/sglang/srt/hardware_backend/npu/utils.py index 715af908e..478c73d24 100644 --- a/python/sglang/srt/hardware_backend/npu/utils.py +++ b/python/sglang/srt/hardware_backend/npu/utils.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Callable import torch from sglang.srt.environ import envs -from sglang.srt.utils import is_npu +from sglang.srt.utils import get_npu_memory_capacity, is_npu if TYPE_CHECKING: from sglang.srt.server_args import ServerArgs @@ -48,6 +48,29 @@ def set_default_server_args(args: "ServerArgs"): if args.page_size is None: args.page_size = 128 + # NPU memory settings + npu_mem = get_npu_memory_capacity() + if npu_mem <= 32 * 1024: + # Ascend 910B4,910B4_1 + # (chunked_prefill_size 4k, cuda_graph_max_bs 16 if tp < 4 else 64) + if args.chunked_prefill_size is None: + args.chunked_prefill_size = 4 * 1024 + if args.cuda_graph_max_bs is None: + if args.tp_size < 4: + args.cuda_graph_max_bs = 16 + else: + args.cuda_graph_max_bs = 64 + elif npu_mem <= 64 * 1024: + # Ascend 910B1,910B2,910B2C,910B3,910_9391,910_9392,910_9381,910_9382,910_9372,910_9362 + # (chunked_prefill_size 8k, cuda_graph_max_bs 64 if tp < 4 else 256) + if args.chunked_prefill_size is None: + args.chunked_prefill_size = 8 * 1024 + if args.cuda_graph_max_bs is None: + if args.tp_size < 4: + args.cuda_graph_max_bs = 64 + else: + args.cuda_graph_max_bs = 256 + # NPU does not support CustomAllReduce args.disable_custom_all_reduce = True diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index e79154665..bff7fbf60 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -653,17 +653,17 @@ class ServerArgs: # Set missing default values. self._handle_missing_default_values() + # Handle device-specific backends. + self._handle_hpu_backends() + self._handle_cpu_backends() + self._handle_npu_backends() + # Get GPU memory capacity, which is a common dependency for several configuration steps. gpu_mem = get_device_memory_capacity(self.device) # Handle memory-related, chunked prefill, and CUDA graph batch size configurations. self._handle_gpu_memory_settings(gpu_mem) - # Handle device-specific backends. - self._handle_hpu_backends() - self._handle_cpu_backends() - self._handle_npu_backends() - # Apply model-specific adjustments. self._handle_model_specific_adjustments() @@ -806,16 +806,6 @@ class ServerArgs: self.chunked_prefill_size = 2048 if self.cuda_graph_max_bs is None: self.cuda_graph_max_bs = 8 - elif is_npu() and gpu_mem < 32 * 1024: - # Atlas A2B4 - # (chunked_prefill_size 32k, cuda_graph_max_bs 16 if tp < 4 else 64) - if self.chunked_prefill_size is None: - self.chunked_prefill_size = 32768 - if self.cuda_graph_max_bs is None: - if self.tp_size < 4: - self.cuda_graph_max_bs = 16 - else: - self.cuda_graph_max_bs = 64 elif gpu_mem < 35 * 1024: # A10, 4090, 5090 # (chunked_prefill_size 2k, cuda_graph_max_bs 24 if tp < 4 else 80) @@ -839,16 +829,6 @@ class ServerArgs: self.cuda_graph_max_bs = 32 else: self.cuda_graph_max_bs = 160 - elif is_npu() and gpu_mem < 64 * 1024: - # Atlas A2 and Atlas A3 - # (chunked_prefill_size 32k, cuda_graph_max_bs 64 if tp < 4 else 128) - if self.chunked_prefill_size is None: - self.chunked_prefill_size = 32768 - if self.cuda_graph_max_bs is None: - if self.tp_size < 4: - self.cuda_graph_max_bs = 64 - else: - self.cuda_graph_max_bs = 128 elif gpu_mem < 90 * 1024: # H100, A100 # (chunked_prefill_size 8k, cuda_graph_max_bs 256 if tp < 4 else 512)