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

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

Browse Source 
Co-authored-by: ZhengdQin <zhengdqin@gmail.com>
This commit is contained in:
XDaoHong
2025-12-16 09:23:51 +08:00
committed by GitHub
parent 3ffa260474
commit 4733fcff1f
7 changed files with 243 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
View File

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

40
python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py
View File

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

85
python/sglang/srt/hardware_backend/npu/memory_pool_npu.py
View File

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

35
python/sglang/srt/layers/rotary_embedding.py
View File

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

19
python/sglang/srt/utils/common.py
View File

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