[Feature] npu support enable_torch_compile for torchair backend (#13410)

Co-authored-by: ZhengdQin <zhengdqin@gmail.com>

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:

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]

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(
+        if self.use_fia:
-            key=cache_k,
+            k_buffer_layer = self.k_buffer[layer_id - self.start_layer]
-            value=cache_v,
+            v_buffer_layer = self.v_buffer[layer_id - self.start_layer]
-            key_cache=self.k_buffer[layer_id - self.start_layer].view(
+
-                -1, self.page_size, self.head_num, self.head_dim
+            torch_npu.npu_scatter_nd_update_(
-            ),
+                k_buffer_layer,
-            value_cache=self.v_buffer[layer_id - self.start_layer].view(
+                loc.view(-1, 1),
-                -1, self.page_size, self.head_num, self.head_dim
+                cache_k.view(-1, 1, self.head_num, self.head_dim),
-            ),
+            )
-            slot_indices=loc,
+            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):

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"):
+        rotary_mode = "half"
-            return self.forward_native(
+        if self.is_neox_style:
-                positions, query, key, offsets, fused_set_kv_buffer_arg
-            )
-        else:
             rotary_mode = "half"
-            if self.is_neox_style:
+        else:
-                rotary_mode = "half"
+            rotary_mode = "interleave"
-            else:
+        mrope_section = [0, 0, 0]
-                rotary_mode = "interleave"
+        query_out, key_out = torch_npu.npu_mrope(
-            mrope_section = [0, 0, 0]
+            positions,
-            query_out, key_out = torch_npu.npu_mrope(
+            query,
-                positions,
+            key,
-                query,
+            self.cos_sin_cache,
-                key,
+            self.head_size,
-                self.cos_sin_cache,
+            mrope_section=mrope_section,
-                self.head_size,
+            rotary_mode=rotary_mode,
-                mrope_section=mrope_section,
+        )
-                rotary_mode=rotary_mode,
+        return query_out, key_out
-            )
-            return query_out, key_out
 
     def forward_cpu(
         self,

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():
+def get_compiler_backend(mode=None) -> str:
-    config = {
-        "frozen_parameter": True,
-        "tiling_schedule_optimize": True,
-        "topology_sorting_strategy": "StableRDFS",
-    }
-    return config
-
-
-def get_compiler_backend() -> 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()
+        compiler_config.mode = "max-autotune"
-        for k, v in predefined_config.items():
+        if mode == "npugraph_ex":
-            setattr(compiler_config.experimental_config, k, v)
+            compiler_config.mode = "reduce-overhead"
-
+            compiler_config.debug.run_eagerly = True
         npu_backend = torchair.get_npu_backend(compiler_config=compiler_config)
         return npu_backend
 

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

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