chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[NPU] fix for NPU memory settings logic (#15258)

Browse Source
This commit is contained in:
Even Zhou
2025-12-17 09:04:22 +08:00
committed by GitHub
parent c8c64876a7
commit 71cb90378b
2 changed files with 29 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
python/sglang/srt/hardware_backend/npu/utils.py
View File

@@ -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

30
python/sglang/srt/server_args.py
View File

@@ -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)