[Ascend]Support of piecewise graph compilation for prefill on NPU (#12287)

Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
Vladimir221
2025-12-11 16:10:07 +03:00
committed by GitHub
parent 5d804a3767
commit 27032cecd9
10 changed files with 289 additions and 40 deletions

View File

@@ -639,6 +639,9 @@ class ServerArgs:
self._handle_cpu_backends()
self._handle_npu_backends()
# Handle compilation config
self._handle_compilation_cfg()
# Apply model-specific adjustments.
self._handle_model_specific_adjustments()
@@ -951,6 +954,15 @@ class ServerArgs:
self.attention_backend = "intel_amx"
self.sampling_backend = "pytorch"
def _handle_compilation_cfg(self):
# NPU platform
if is_npu() and self.piecewise_cuda_graph_compiler != "eager":
logger.warning(
"At this moment Ascend platform only support prefill graph compilation with "
"piecewise_cuda_graph_compiler='eager', change piecewise_cuda_graph_compiler to 'eager'."
)
self.piecewise_cuda_graph_compiler = "eager"
def _handle_npu_backends(self):
if self.device == "npu":
from sglang.srt.hardware_backend.npu.utils import set_default_server_args