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

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
 

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)