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

[NPU][eagle3] support qwen eagle3 on NPU (#14820)

Browse Source
This commit is contained in:
Liwansi
2025-12-16 02:25:13 +08:00
committed by GitHub
parent 4901693110
commit 30da2f0598
11 changed files with 275 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/sglang/srt/configs/model_config.py
View File

@@ -236,8 +236,14 @@ class ModelConfig:
server_args: ServerArgs,
model_path: str = None,
model_revision: str = None,
is_draft_model: bool = False,
**kwargs,
):
quantization = (
server_args.speculative_draft_model_quantization
if is_draft_model
else server_args.quantization
)
return ModelConfig(
model_path=model_path or server_args.model_path,
trust_remote_code=server_args.trust_remote_code,
@@ -247,7 +253,7 @@ class ModelConfig:
is_embedding=server_args.is_embedding,
enable_multimodal=server_args.enable_multimodal,
dtype=server_args.dtype,
quantization=server_args.quantization,
quantization=quantization,
hybrid_kvcache_ratio=server_args.hybrid_kvcache_ratio,
model_impl=server_args.model_impl,
sampling_defaults=server_args.sampling_defaults,
@@ -255,6 +261,7 @@ class ModelConfig:
override_config_file=server_args.decrypted_config_file,
language_only=server_args.language_only,
encoder_only=server_args.encoder_only,
is_draft_model=is_draft_model,
**kwargs,
)

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

@@ -889,102 +889,163 @@ class AscendAttnBackend(AttentionBackend):
layer, forward_batch.out_cache_loc, k, v
)
c_kv, k_rope = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
if is_fia_nz():
k_rope_cache = _reshape_kv_for_fia_nz(
k_rope, layer.tp_k_head_num, self.qk_rope_head_dim, self.page_size
)
c_kv_cache = _reshape_kv_for_fia_nz(
c_kv, layer.tp_v_head_num, self.kv_lora_rank, self.page_size
)
else:
k_rope_cache = k_rope.view(
-1, layer.tp_k_head_num, self.page_size, self.qk_rope_head_dim
)
c_kv_cache = c_kv.view(
-1, layer.tp_v_head_num, self.page_size, self.kv_lora_rank
)
if not self.use_mla:
k_cache = forward_batch.token_to_kv_pool.get_key_buffer(
layer.layer_id
).view(-1, self.page_size, layer.tp_k_head_num * layer.qk_head_dim)
v_cache = forward_batch.token_to_kv_pool.get_value_buffer(
layer.layer_id
).view(-1, self.page_size, layer.tp_v_head_num * layer.v_head_dim)
query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim).contiguous()
if not self.graph_mode:
num_token_padding = query.shape[0]
query = query[: forward_batch.num_token_non_padded_cpu]
if self.forward_metadata.seq_lens_cpu_int is None:
actual_seq_lengths_kv = self.forward_metadata.seq_lens_cpu_list
else:
actual_seq_lengths_kv = (
self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist()
)
if forward_batch.forward_mode.is_draft_extend():
actual_seq_lengths = (
np.array(forward_batch.extend_seq_lens_cpu).cumsum().tolist()
)
else:
actual_seq_lengths = np.arange(
self.speculative_num_draft_tokens,
self.speculative_num_draft_tokens + query.shape[0],
self.speculative_num_draft_tokens,
)
q_nope = q.view(-1, layer.tp_q_head_num, self.kv_lora_rank).contiguous()
q_rope = q_rope.view(-1, layer.tp_q_head_num, self.qk_rope_head_dim)
if not self.graph_mode:
num_token_padding = q.shape[0]
q_nope = q_nope[: forward_batch.num_token_non_padded_cpu]
q_rope = q_rope[: forward_batch.num_token_non_padded_cpu]
if self.forward_metadata.seq_lens_cpu_int is None:
actual_seq_lengths_kv = self.forward_metadata.seq_lens_cpu_list
attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score(
query,
k_cache,
v_cache,
block_table=self.forward_metadata.block_tables,
block_size=self.page_size,
num_heads=layer.tp_q_head_num,
num_key_value_heads=layer.tp_k_head_num,
input_layout="TND",
atten_mask=self.mtp_mask,
scale=layer.scaling,
actual_seq_lengths=actual_seq_lengths,
actual_seq_lengths_kv=actual_seq_lengths_kv,
sparse_mode=3,
)
attn_output = attn_output.view(-1, layer.tp_q_head_num * layer.v_head_dim)
if (
not self.graph_mode
and forward_batch.num_token_non_padded_cpu != num_token_padding
):
attn_output = torch.cat(
[
attn_output,
attn_output.new_zeros(
num_token_padding - forward_batch.num_token_non_padded_cpu,
*attn_output.shape[1:],
),
],
dim=0,
)
return attn_output
else:
actual_seq_lengths_kv = (
self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist()
)
if forward_batch.forward_mode.is_draft_extend():
actual_seq_lengths = (
np.array(forward_batch.extend_seq_lens_cpu).cumsum().tolist()
)
else:
actual_seq_lengths = np.arange(
self.speculative_num_draft_tokens,
self.speculative_num_draft_tokens + q_nope.shape[0],
self.speculative_num_draft_tokens,
)
c_kv, k_rope = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
if is_fia_nz():
k_rope_cache = _reshape_kv_for_fia_nz(
k_rope, layer.tp_k_head_num, self.qk_rope_head_dim, self.page_size
)
c_kv_cache = _reshape_kv_for_fia_nz(
c_kv, layer.tp_v_head_num, self.kv_lora_rank, self.page_size
)
else:
k_rope_cache = k_rope.view(
-1, layer.tp_k_head_num, self.page_size, self.qk_rope_head_dim
)
c_kv_cache = c_kv.view(
-1, layer.tp_v_head_num, self.page_size, self.kv_lora_rank
)
workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(
q_nope,
c_kv_cache,
c_kv_cache,
query_rope=q_rope,
key_rope=k_rope_cache,
num_heads=layer.tp_q_head_num,
num_key_value_heads=layer.tp_k_head_num,
input_layout="TND",
scale=layer.scaling,
antiquant_mode=0,
antiquant_scale=None,
block_table=self.forward_metadata.block_tables,
block_size=self.page_size,
sparse_mode=3,
atten_mask=self.mtp_mask,
actual_seq_lengths=actual_seq_lengths,
actual_seq_lengths_kv=actual_seq_lengths_kv,
)
attn_output = torch.empty_like(q_nope, dtype=q.dtype, device=q.device)
softmax_lse = torch.empty(1, dtype=q.dtype, device=q.device)
torch_npu.npu_fused_infer_attention_score.out(
q_nope,
c_kv_cache,
c_kv_cache,
query_rope=q_rope,
key_rope=k_rope_cache,
num_heads=layer.tp_q_head_num,
num_key_value_heads=layer.tp_k_head_num,
input_layout="TND",
scale=layer.scaling,
antiquant_mode=0,
antiquant_scale=None,
block_table=self.forward_metadata.block_tables,
block_size=self.page_size,
sparse_mode=3,
atten_mask=self.mtp_mask,
actual_seq_lengths=actual_seq_lengths,
actual_seq_lengths_kv=actual_seq_lengths_kv,
workspace=workspace,
out=[attn_output, softmax_lse],
)
attn_output = attn_output.view(-1, layer.tp_q_head_num * layer.v_head_dim)
if (
not self.graph_mode
and forward_batch.num_token_non_padded_cpu != num_token_padding
):
attn_output = torch.cat(
[
attn_output,
attn_output.new_zeros(
num_token_padding - attn_output.shape[0], *attn_output.shape[1:]
),
],
dim=0,
q_nope = q.view(-1, layer.tp_q_head_num, self.kv_lora_rank).contiguous()
q_rope = q_rope.view(-1, layer.tp_q_head_num, self.qk_rope_head_dim)
if not self.graph_mode:
num_token_padding = q.shape[0]
q_nope = q_nope[: forward_batch.num_token_non_padded_cpu]
q_rope = q_rope[: forward_batch.num_token_non_padded_cpu]
if self.forward_metadata.seq_lens_cpu_int is None:
actual_seq_lengths_kv = self.forward_metadata.seq_lens_cpu_list
else:
actual_seq_lengths_kv = (
self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist()
)
if forward_batch.forward_mode.is_draft_extend():
actual_seq_lengths = (
np.array(forward_batch.extend_seq_lens_cpu).cumsum().tolist()
)
else:
actual_seq_lengths = np.arange(
self.speculative_num_draft_tokens,
self.speculative_num_draft_tokens + q_nope.shape[0],
self.speculative_num_draft_tokens,
)
workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(
q_nope,
c_kv_cache,
c_kv_cache,
query_rope=q_rope,
key_rope=k_rope_cache,
num_heads=layer.tp_q_head_num,
num_key_value_heads=layer.tp_k_head_num,
input_layout="TND",
scale=layer.scaling,
antiquant_mode=0,
antiquant_scale=None,
block_table=self.forward_metadata.block_tables,
block_size=self.page_size,
sparse_mode=3,
atten_mask=self.mtp_mask,
actual_seq_lengths=actual_seq_lengths,
actual_seq_lengths_kv=actual_seq_lengths_kv,
)
return attn_output
attn_output = torch.empty_like(q_nope, dtype=q.dtype, device=q.device)
softmax_lse = torch.empty(1, dtype=q.dtype, device=q.device)
torch_npu.npu_fused_infer_attention_score.out(
q_nope,
c_kv_cache,
c_kv_cache,
query_rope=q_rope,
key_rope=k_rope_cache,
num_heads=layer.tp_q_head_num,
num_key_value_heads=layer.tp_k_head_num,
input_layout="TND",
scale=layer.scaling,
antiquant_mode=0,
antiquant_scale=None,
block_table=self.forward_metadata.block_tables,
block_size=self.page_size,
sparse_mode=3,
atten_mask=self.mtp_mask,
actual_seq_lengths=actual_seq_lengths,
actual_seq_lengths_kv=actual_seq_lengths_kv,
workspace=workspace,
out=[attn_output, softmax_lse],
)
attn_output = attn_output.view(-1, layer.tp_q_head_num * layer.v_head_dim)
if (
not self.graph_mode
and forward_batch.num_token_non_padded_cpu != num_token_padding
):
attn_output = torch.cat(
[
attn_output,
attn_output.new_zeros(
num_token_padding - attn_output.shape[0],
*attn_output.shape[1:],
),
],
dim=0,
)
return attn_output
def forward_decode_graph(
self,

3
python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py
View File

@@ -37,6 +37,9 @@ class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner):
def _create_graph(self):
return torch.npu.NPUGraph()
def _cache_loc_dtype(self):
return torch.int32
def _capture_init(self, run_once_fn):
for _ in range(2):
torch.npu.synchronize()

39
python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py
View File

@@ -17,11 +17,13 @@ from __future__ import annotations
import logging
import threading
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Dict, Union
import numpy as np
import torch
from sglang.srt.configs.model_config import is_deepseek_nsa
from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
EAGLEDraftCudaGraphRunner,
@@ -46,6 +48,19 @@ if is_npu():
class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
def __init__(self, eagle_worker: EAGLEWorker):
super().__init__(eagle_worker)
self.update_attr_name = None
self.update_attr_type = None
self._init_arch_map()
def _init_arch_map(self):
self.attr_name: Dict[str, str] = {
AttentionArch.MLA: "actual_seq_lengths_kv",
AttentionArch.MHA: "context_lens",
}
self.attr_type: Dict[str, Union[list, torch.Tensor]] = {
AttentionArch.MLA: [],
AttentionArch.MHA: torch.Tensor(),
}
def _create_graph(self):
return torch.npu.NPUGraph()
@@ -63,12 +78,27 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
out = run_once_fn()
return out
def _get_update_attr_name(self, model_runner):
if self.bs < get_attention_tp_size():
return self.attr_name[AttentionArch.MLA]
return self.attr_name[model_runner.model_config.attention_arch]
def _get_update_attr_type(self, model_runner):
if self.bs < get_attention_tp_size():
return self.attr_type[AttentionArch.MLA]
return self.attr_type[model_runner.model_config.attention_arch]
def _replay_update(self, seq_lens):
if isinstance(self.update_attr_type, torch.Tensor):
seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32))
self.graphs[self.bs].update(
cpu_update_input=[{"actual_seq_lengths_kv": seq_lens}]
cpu_update_input=[{self.update_attr_name: seq_lens}]
)
def _replay(self, forward_batch: ForwardBatch):
self.update_attr_name = self._get_update_attr_name(self.model_runner)
self.update_attr_type = self._get_update_attr_type(self.model_runner)
if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * (
self.bs - self.raw_bs
@@ -79,3 +109,6 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
thread.join()
else:
self.graphs[self.bs].replay()
def _cache_loc_dtype(self):
return torch.int32

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

@@ -77,13 +77,19 @@ class NPUGraphRunner(CudaGraphRunner):
out = run_once_fn()
return out
def _get_update_attr_name(self, model_runner):
if self.bs < get_attention_tp_size():
def _get_update_attr_name(self, model_runner, forward_batch):
if (
self.bs < get_attention_tp_size()
or forward_batch.forward_mode.is_target_verify()
):
return self.attr_name[AttentionArch.MLA]
return self.attr_name[model_runner.model_config.attention_arch]
def _get_update_attr_type(self, model_runner):
if self.bs < get_attention_tp_size():
def _get_update_attr_type(self, model_runner, forward_batch):
if (
self.bs < get_attention_tp_size()
or forward_batch.forward_mode.is_target_verify()
):
return self.attr_type[AttentionArch.MLA]
return self.attr_type[model_runner.model_config.attention_arch]
@@ -139,8 +145,12 @@ class NPUGraphRunner(CudaGraphRunner):
self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions)
self.update_attr_name = self._get_update_attr_name(self.model_runner)
self.update_attr_type = self._get_update_attr_type(self.model_runner)
self.update_attr_name = self._get_update_attr_name(
self.model_runner, forward_batch
)
self.update_attr_type = self._get_update_attr_type(
self.model_runner, forward_batch
)
# Replay
if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
if forward_batch.forward_mode.is_target_verify():

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

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

19
python/sglang/srt/server_args.py
View File

@@ -111,6 +111,8 @@ QUANTIZATION_CHOICES = [
"modelslim", # for NPU
]
SPECULATIVE_DRAFT_MODEL_QUANTIZATION_CHOICES = [*QUANTIZATION_CHOICES, "unquant"]
ATTENTION_BACKEND_CHOICES = [
# Common
"triton",
@@ -431,6 +433,7 @@ class ServerArgs:
speculative_attention_mode: str = "prefill"
speculative_moe_runner_backend: Optional[str] = None
speculative_moe_a2a_backend: Optional[str] = None
speculative_draft_model_quantization: Optional[str] = None
# Speculative decoding (ngram)
speculative_ngram_min_match_window_size: int = 1
@@ -747,8 +750,15 @@ class ServerArgs:
# TODO: when extra_buffer is more verified, we can set the default path based on
# [overlap, non-overlap]
self.mamba_scheduler_strategy = "no_buffer"
# In speculative scenario:
# - If `speculative_draft_model_quantization` is specified, the draft model uses this quantization method.
# - Otherwise, the draft model defaults to the same quantization as the target model.
if self.speculative_draft_model_quantization is None:
self.speculative_draft_model_quantization = self.quantization
elif self.speculative_draft_model_quantization == "unquant":
self.speculative_draft_model_quantization = None
# Handle ModelScope model downloads
# Handle ModelScope model downloads
if get_bool_env_var("SGLANG_USE_MODELSCOPE"):
if not os.path.exists(self.model_path):
from modelscope import snapshot_download
@@ -3399,6 +3409,13 @@ class ServerArgs:
default=ServerArgs.speculative_moe_a2a_backend,
help="Choose the backend for MoE A2A in speculative decoding",
)
parser.add_argument(
"--speculative-draft-model-quantization",
type=str,
choices=SPECULATIVE_DRAFT_MODEL_QUANTIZATION_CHOICES,
default=ServerArgs.speculative_draft_model_quantization,
help="The quantization method for speculative model.",
)
# Speculative decoding (ngram)
parser.add_argument(

6
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
View File

@@ -88,7 +88,8 @@ class EAGLEDraftCudaGraphRunner:
self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32)
self.out_cache_loc = torch.zeros(
(self.max_num_token * self.speculative_num_steps,), dtype=torch.int64
(self.max_num_token * self.speculative_num_steps,),
dtype=self._cache_loc_dtype(),
)
self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
self.mrope_positions = torch.zeros(
@@ -132,6 +133,9 @@ class EAGLEDraftCudaGraphRunner:
f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
)
def _cache_loc_dtype(self):
return torch.int64
def can_run(self, forward_batch: ForwardBatch):
if self.require_mlp_tp_gather:
cuda_graph_bs = (

7
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
View File

@@ -92,7 +92,9 @@ class EAGLEDraftExtendCudaGraphRunner:
with torch.device(model_runner.device):
self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32)
self.out_cache_loc = torch.ones((self.max_num_token,), dtype=torch.int64)
self.out_cache_loc = torch.ones(
(self.max_num_token,), dtype=self._cache_loc_dtype()
)
self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
self.mrope_positions = torch.zeros(
(3, self.max_num_token), dtype=torch.int64
@@ -204,6 +206,9 @@ class EAGLEDraftExtendCudaGraphRunner:
def _create_graph(self):
return torch.cuda.CUDAGraph()
def _cache_loc_dtype(self):
return torch.int64
def _capture_init(self, run_once_fn):
for _ in range(2):
torch.cuda.synchronize()

3
python/sglang/srt/speculative/eagle_info_v2.py
View File

@@ -468,8 +468,6 @@ def assign_extend_cache_locs_func(
return out_cache_loc
elif _is_npu:
import sgl_kernel_npu # noqa: F401
out_cache_loc = torch.empty(
(batch_size * draft_token_num,),
dtype=torch.int32,
@@ -482,6 +480,5 @@ def assign_extend_cache_locs_func(
end_offset,
out_cache_loc,
)
out_cache_loc = out_cache_loc.to(dtype=torch.int64)
return out_cache_loc