From 4733fcff1fe9f5b402922a4bb472a0a171559b28 Mon Sep 17 00:00:00 2001 From: XDaoHong Date: Tue, 16 Dec 2025 09:23:51 +0800 Subject: [PATCH] [Feature] npu support enable_torch_compile for torchair backend (#13410) Co-authored-by: ZhengdQin --- .../npu/attention/ascend_backend.py | 11 +- .../npu/graph_runner/npu_graph_runner.py | 40 ++++++- .../hardware_backend/npu/memory_pool_npu.py | 85 ++++++++++++--- python/sglang/srt/layers/rotary_embedding.py | 35 +++--- python/sglang/srt/utils/common.py | 19 +--- .../test_ascend_compile_graph_tp1_bf16.py | 102 ++++++++++++++++++ test/srt/run_suite.py | 1 + 7 files changed, 243 insertions(+), 50 deletions(-) create mode 100644 test/srt/ascend/test_ascend_compile_graph_tp1_bf16.py diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index ee3b41402..c05988994 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -248,6 +248,7 @@ class AscendAttnBackend(AttentionBackend): self.req_to_token = model_runner.req_to_token_pool.req_to_token self.graph_mode = False self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False") + self.enable_torch_compile = model_runner.server_args.enable_torch_compile self.speculative_num_draft_tokens = ( model_runner.server_args.speculative_num_draft_tokens ) @@ -1255,7 +1256,7 @@ class AscendAttnBackend(AttentionBackend): topk_indices, ) - if self.graph_mode: + if self.graph_mode and (not self.enable_torch_compile): return self.forward_decode_graph( q, k, @@ -1276,6 +1277,12 @@ class AscendAttnBackend(AttentionBackend): k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) v_cache = forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id) if self.use_fia: + if self.forward_metadata.seq_lens_cpu_int is None: + actual_seq_len_kv = self.forward_metadata.seq_lens_cpu_list + else: + actual_seq_len_kv = ( + self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist() + ) attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score( q.view( forward_batch.batch_size, @@ -1295,7 +1302,7 @@ class AscendAttnBackend(AttentionBackend): atten_mask=None, block_size=self.page_size, block_table=self.forward_metadata.block_tables, - actual_seq_lengths_kv=self.forward_metadata.seq_lens_cpu_int, + actual_seq_lengths_kv=actual_seq_len_kv, scale=layer.scaling, ) else: diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py index 838c5727d..d089fcc63 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py @@ -18,16 +18,24 @@ from __future__ import annotations import logging import os import threading +from contextlib import contextmanager from pathlib import Path from typing import TYPE_CHECKING, Dict, Optional, Union import numpy as np import torch +import sglang from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa +from sglang.srt.distributed.parallel_state import GroupCoordinator from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner -from sglang.srt.utils import is_npu +from sglang.srt.utils import ( + empty_context, + get_bool_env_var, + get_compiler_backend, + is_npu, +) is_npu = is_npu() @@ -44,15 +52,36 @@ from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +@contextmanager +def patch_model_npu( + model: torch.nn.Module, + enable_compile: bool, + num_tokens: int, + tp_group: GroupCoordinator, +): + if enable_compile: + backend = get_compiler_backend("npugraph_ex") + yield torch.compile( + torch.no_grad()(model.forward), + fullgraph=True, + dynamic=False, + backend=backend, + ) + else: + yield model.forward + + class NPUGraphRunner(CudaGraphRunner): """A NPUGraphRunner runs the forward pass of a model with npu graph and torch.compile.""" def __init__(self, model_runner: ModelRunner): + sglang.srt.model_executor.cuda_graph_runner.patch_model = patch_model_npu super().__init__(model_runner) self.update_attr_name = None self.update_attr_type = None self.model_runner = model_runner self._init_arch_map() + self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False") def _init_arch_map(self): self.attr_name: Dict[str, str] = { @@ -68,7 +97,12 @@ class NPUGraphRunner(CudaGraphRunner): return torch.npu.NPUGraph() def _capture_graph(self, graph, pool, stream, run_once_fn): - with torch.npu.graph( + if self.enable_torch_compile: + skip_guard_context = torch.compiler.set_stance(skip_guard_eval_unsafe=True) + else: + skip_guard_context = empty_context() + + with skip_guard_context, torch.npu.graph( graph, pool=pool, stream=stream, @@ -81,6 +115,7 @@ class NPUGraphRunner(CudaGraphRunner): if ( self.bs < get_attention_tp_size() or forward_batch.forward_mode.is_target_verify() + or self.use_fia ): return self.attr_name[AttentionArch.MLA] return self.attr_name[model_runner.model_config.attention_arch] @@ -89,6 +124,7 @@ class NPUGraphRunner(CudaGraphRunner): if ( self.bs < get_attention_tp_size() or forward_batch.forward_mode.is_target_verify() + or self.use_fia ): return self.attr_type[AttentionArch.MLA] return self.attr_type[model_runner.model_config.attention_arch] diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index 47ec111ba..41fb2aae2 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -9,6 +9,7 @@ from sglang.srt.mem_cache.memory_pool import ( MLATokenToKVPool, get_tensor_size_bytes, ) +from sglang.srt.utils import get_bool_env_var if TYPE_CHECKING: from sglang.srt.layers.radix_attention import RadixAttention @@ -16,6 +17,37 @@ if TYPE_CHECKING: class NPUMHATokenToKVPool(MHATokenToKVPool): + def __init__( + self, + size: int, + page_size: int, + dtype: torch.dtype, + head_num: int, + head_dim: int, + layer_num: int, + device: str, + enable_memory_saver: bool, + start_layer: Optional[int] = None, + end_layer: Optional[int] = None, + enable_alt_stream: bool = True, + enable_kv_cache_copy: bool = False, + ): + self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False") + super().__init__( + size, + page_size, + dtype, + head_num, + head_dim, + layer_num, + device, + enable_memory_saver, + start_layer, + end_layer, + enable_alt_stream, + enable_kv_cache_copy, + ) + def _create_buffers(self): with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): # [size, head_num, head_dim] for each layer @@ -37,6 +69,19 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): self.k_buffer = self.kv_buffer[0] self.v_buffer = self.kv_buffer[1] + if self.use_fia: + self.k_buffer = [] + self.v_buffer = [] + for i in range(self.layer_num): + k_buffer_layer = self.kv_buffer[0][i].view( + -1, 1, self.head_num, self.head_dim + ) + v_buffer_layer = self.kv_buffer[1][i].view( + -1, 1, self.head_num, self.head_dim + ) + self.k_buffer.append(k_buffer_layer) + self.v_buffer.append(v_buffer_layer) + # for disagg def get_contiguous_buf_infos(self): # layer_num x [seq_len, head_num, head_dim] @@ -89,18 +134,34 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): if self.store_dtype != self.dtype: cache_k = cache_k.view(self.store_dtype) cache_v = cache_v.view(self.store_dtype) - loc = loc.to(torch.int32) - torch_npu._npu_reshape_and_cache( - key=cache_k, - value=cache_v, - key_cache=self.k_buffer[layer_id - self.start_layer].view( - -1, self.page_size, self.head_num, self.head_dim - ), - value_cache=self.v_buffer[layer_id - self.start_layer].view( - -1, self.page_size, self.head_num, self.head_dim - ), - slot_indices=loc, - ) + + if self.use_fia: + k_buffer_layer = self.k_buffer[layer_id - self.start_layer] + v_buffer_layer = self.v_buffer[layer_id - self.start_layer] + + torch_npu.npu_scatter_nd_update_( + k_buffer_layer, + loc.view(-1, 1), + cache_k.view(-1, 1, self.head_num, self.head_dim), + ) + torch_npu.npu_scatter_nd_update_( + v_buffer_layer, + loc.view(-1, 1), + cache_v.view(-1, 1, self.head_num, self.head_dim), + ) + else: + loc = loc.to(torch.int32) + torch_npu._npu_reshape_and_cache( + key=cache_k, + value=cache_v, + key_cache=self.k_buffer[layer_id - self.start_layer].view( + -1, self.page_size, self.head_num, self.head_dim + ), + value_cache=self.v_buffer[layer_id - self.start_layer].view( + -1, self.page_size, self.head_num, self.head_dim + ), + slot_indices=loc, + ) class NPUMLATokenToKVPool(MLATokenToKVPool): diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index 73202b7b7..0385c4466 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -269,27 +269,22 @@ class RotaryEmbedding(CustomOp): fused_set_kv_buffer_arg is None ), "fused_set_kv_buffer_arg is not supported for npu implementation" - if get_bool_env_var("SGLANG_ENABLE_TORCH_COMPILE"): - return self.forward_native( - positions, query, key, offsets, fused_set_kv_buffer_arg - ) - else: + rotary_mode = "half" + if self.is_neox_style: rotary_mode = "half" - if self.is_neox_style: - rotary_mode = "half" - else: - rotary_mode = "interleave" - mrope_section = [0, 0, 0] - query_out, key_out = torch_npu.npu_mrope( - positions, - query, - key, - self.cos_sin_cache, - self.head_size, - mrope_section=mrope_section, - rotary_mode=rotary_mode, - ) - return query_out, key_out + else: + rotary_mode = "interleave" + mrope_section = [0, 0, 0] + query_out, key_out = torch_npu.npu_mrope( + positions, + query, + key, + self.cos_sin_cache, + self.head_size, + mrope_section=mrope_section, + rotary_mode=rotary_mode, + ) + return query_out, key_out def forward_cpu( self, diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index b4fef6ed6..0e9b4ce1c 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -1933,16 +1933,7 @@ def get_device_capability(device_id: int = 0) -> Tuple[int, int]: return major, minor -def get_npu_compiler_config(): - config = { - "frozen_parameter": True, - "tiling_schedule_optimize": True, - "topology_sorting_strategy": "StableRDFS", - } - return config - - -def get_compiler_backend() -> str: +def get_compiler_backend(mode=None) -> str: if hasattr(torch, "hpu") and torch.hpu.is_available(): return "hpu_backend" @@ -1957,10 +1948,10 @@ def get_compiler_backend() -> str: "Please install torchair for torch.compile support on NPU." ) compiler_config = CompilerConfig() - predefined_config = get_npu_compiler_config() - for k, v in predefined_config.items(): - setattr(compiler_config.experimental_config, k, v) - + compiler_config.mode = "max-autotune" + if mode == "npugraph_ex": + compiler_config.mode = "reduce-overhead" + compiler_config.debug.run_eagerly = True npu_backend = torchair.get_npu_backend(compiler_config=compiler_config) return npu_backend diff --git a/test/srt/ascend/test_ascend_compile_graph_tp1_bf16.py b/test/srt/ascend/test_ascend_compile_graph_tp1_bf16.py new file mode 100644 index 000000000..deaf4a5e0 --- /dev/null +++ b/test/srt/ascend/test_ascend_compile_graph_tp1_bf16.py @@ -0,0 +1,102 @@ +import os +import unittest +from types import SimpleNamespace +from urllib.parse import urlparse + +from sglang.srt.utils import kill_process_tree +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + is_in_ci, + popen_launch_server, + run_bench_offline_throughput, +) + +TEST_MODEL_MATRIX = { + "Qwen/Qwen2.5-7B-Instruct": { + "accuracy": 0.84, + "latency": 150, + "output_throughput": 30, + }, +} + +os.environ["ASCEND_USE_FIA"] = "true" + + +class TestAscendTp1Bf16(CustomTestCase): + + @classmethod + def setUpClass(cls): + cls.models = TEST_MODEL_MATRIX.keys() + cls.base_url = DEFAULT_URL_FOR_TEST + cls.url = urlparse(DEFAULT_URL_FOR_TEST) + cls.common_args = [ + "--trust-remote-code", + "--mem-fraction-static", + 0.6, + "--attention-backend", + "ascend", + "--disable-radix-cache", + "--enable-torch-compile", + "--watchdog-timeout", + 30000, + ] + + def test_a_gsm8k(self): + for model in self.models: + with self.subTest(model=model): + print(f"##=== Testing accuracy: {model} ===##") + + process = popen_launch_server( + model, + self.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + *self.common_args, + ], + ) + + try: + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=1319, + max_new_tokens=512, + parallel=32, + host=f"http://{self.url.hostname}", + port=int(self.url.port), + ) + + metrics = run_eval_few_shot_gsm8k(args) + self.assertGreaterEqual( + metrics["accuracy"], + TEST_MODEL_MATRIX[model]["accuracy"], + ) + finally: + kill_process_tree(process.pid) + + def test_b_throughput(self): + for model in self.models: + with self.subTest(model=model): + print(f"##=== Testing throughput: {model} ===##") + + output_throughput = run_bench_offline_throughput( + model, + [ + *self.common_args, + ], + ) + + print(f"##=== {model} throughput: {output_throughput} ===##") + + if is_in_ci(): + self.assertGreater( + output_throughput, + TEST_MODEL_MATRIX[model]["output_throughput"], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 1f50ef0ff..c90a4e5b8 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -368,6 +368,7 @@ suite_ascend = { TestFile("ascend/test_ascend_hicache_mha.py", 400), TestFile("ascend/test_ascend_sampling_backend.py", 400), TestFile("ascend/test_ascend_tp1_bf16.py", 400), + TestFile("ascend/test_ascend_compile_graph_tp1_bf16.py", 400), ], "per-commit-2-npu-a2": [ TestFile("ascend/test_ascend_graph_tp2_bf16.py", 400),