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

accuracy enhancement for baichuan2-13B for npu (#16868)

Browse Source 
Co-authored-by: cy <chenyang08056032@163.com>
This commit is contained in:
McZyWu
2026-01-26 16:14:35 +08:00
committed by GitHub
parent 6756cf18c6
commit 2734b23481
5 changed files with 241 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -464,6 +464,9 @@ class ModelConfig:
):
setattr(self.hf_text_config, "head_dim", self.head_dim)
elif "BaichuanForCausalLM" in self.hf_config.architectures:
self.use_alibi = self.hf_config.hidden_size != 4096
self.attention_arch = AttentionArch.MHA
self.num_attention_heads = self.hf_text_config.num_attention_heads

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

@@ -221,6 +221,8 @@ class AscendAttnBackend(AttentionBackend):
else:
self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
self.q_head_dim = self.qk_rope_head_dim + self.qk_nope_head_dim
else:
self.use_alibi = getattr(model_runner.model_config, "use_alibi", False)
self.native_attn = TorchNativeAttnBackend(model_runner)
self.graph_metadata = {}
self.max_context_len = model_runner.model_config.context_len
@@ -405,6 +407,121 @@ class AscendAttnBackend(AttentionBackend):
def get_cuda_graph_seq_len_fill_value(self):
return 0
def _generate_alibi_bias(
self,
seq_len: int,
slopes: torch.Tensor,
num_heads: int,
device: torch.device,
dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
position_point = (
torch.arange(seq_len).view(1, 1, -1).expand(num_heads, -1, -1).to(device)
)
alibi = slopes.view(-1, 1, 1) * position_point
alibi_bias = alibi.view(num_heads, 1, seq_len).to(device).to(dtype)
return alibi_bias
def generate_alibi_bias(
self,
q_seq_len: int,
kv_seq_len: int,
slopes: torch.Tensor,
num_heads: int,
device: torch.device,
is_extend: bool = True,
dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
MAX_LEN_ALB = 5000
max_seq_len = max(kv_seq_len, q_seq_len, MAX_LEN_ALB)
if getattr(self, "alibi_bias", None) is None:
self.alibi_bias = self._generate_alibi_bias(
max_seq_len, slopes, num_heads, device, dtype
)
if getattr(self, "super_mask", None) is None:
super_mask = torch.ones(size=(1, max_seq_len, max_seq_len), dtype=dtype)
super_mask = super_mask.float().fill_(float("-inf")).type_as(super_mask)
super_mask = torch.triu(super_mask, 1).to(device)
self.super_mask = super_mask
if is_extend:
return (
self.alibi_bias[:, :q_seq_len, :kv_seq_len]
+ self.super_mask[:, :q_seq_len, :kv_seq_len]
)
else:
return self.alibi_bias[:, :q_seq_len, :kv_seq_len]
def attn_alibi(
self,
q,
k_cache,
v_cache,
block_tables,
seq_lens,
query_lens,
scale_value,
num_heads,
slopes,
is_extend,
):
curr = 0
num_prompts = query_lens.shape[0]
head_size = k_cache.shape[3]
head_size_v = v_cache.shape[3]
block_size = k_cache.shape[1]
attn_output = []
for i in range(num_prompts):
seq_len = seq_lens[i].item()
block_table = block_tables[i]
j = torch.arange(seq_len, device=block_table.device)
block_number = block_table[j // block_size]
block_offset = j % block_size
k = k_cache[block_number, block_offset]
v = v_cache[block_number, block_offset]
k = k.view(seq_len, num_heads, head_size)
v = v.view(seq_len, num_heads, head_size_v)
if is_extend:
q_len = query_lens[i].item()
query = q[curr : curr + q_len]
else:
q_len = 1
query = q[curr : curr + 1]
query = query.to(torch.float32)
query = query * scale_value
query = query.permute(1, 0, 2)
k = k.permute(1, 2, 0)
score = torch.bmm(query, k)
score = score.to(torch.float32)
if slopes is not None:
alibi_bias = self.generate_alibi_bias(
q_seq_len=q_len,
kv_seq_len=seq_len,
slopes=slopes,
num_heads=num_heads,
device=q.device,
is_extend=is_extend,
dtype=query.dtype,
)
score = score + alibi_bias
score = torch.max(score, torch.tensor(torch.finfo(score.dtype).min))
p = torch.nn.functional.softmax(score, dim=-1)
v = v.permute(1, 0, 2)
out = torch.bmm(p, v)
out = out.permute(1, 0, 2)
out = out.reshape(-1, num_heads * head_size_v)
attn_output.append(out)
curr += q_len
attn_output = torch.cat(attn_output, dim=0).to(q.dtype).to(q.device)
attn_output = attn_output.view(-1, num_heads * head_size)
return attn_output
def do_cp_balance_attn(
self,
q_nope,
@@ -594,6 +711,7 @@ class AscendAttnBackend(AttentionBackend):
k_rope: Optional[torch.Tensor] = None,
topk_indices: Optional[torch.Tensor] = None,
sinks: Optional[torch.Tensor] = None,
slopes: Optional[torch.Tensor] = None,
):
if topk_indices is not None:
return self.forward_sparse(
@@ -689,26 +807,40 @@ class AscendAttnBackend(AttentionBackend):
causal = False
if layer.qk_head_dim <= 128 and causal:
query = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
attn_output = torch.empty(
(query.shape[0], layer.tp_q_head_num * layer.v_head_dim),
dtype=query.dtype,
device=query.device,
)
if not self.use_alibi:
query = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
attn_output = torch.empty(
(query.shape[0], layer.tp_q_head_num * layer.v_head_dim),
dtype=query.dtype,
device=query.device,
)
torch_npu._npu_flash_attention_qlens(
query=query,
key_cache=k_cache,
value_cache=v_cache,
mask=self.mask,
block_table=self.forward_metadata.block_tables,
seq_len=self.forward_metadata.extend_seq_lens_cpu_int,
context_lens=self.forward_metadata.seq_lens_cpu_int,
scale_value=layer.scaling,
num_heads=layer.tp_q_head_num,
num_kv_heads=layer.tp_k_head_num,
out=attn_output,
)
torch_npu._npu_flash_attention_qlens(
query=query,
key_cache=k_cache,
value_cache=v_cache,
mask=self.mask,
block_table=self.forward_metadata.block_tables,
seq_len=self.forward_metadata.extend_seq_lens_cpu_int,
context_lens=self.forward_metadata.seq_lens_cpu_int,
scale_value=layer.scaling,
num_heads=layer.tp_q_head_num,
num_kv_heads=layer.tp_k_head_num,
out=attn_output,
)
else:
attn_output = self.attn_alibi(
q=q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim),
k_cache=k_cache,
v_cache=v_cache,
block_tables=self.forward_metadata.block_tables,
seq_lens=self.forward_metadata.seq_lens_cpu_int,
query_lens=self.forward_metadata.extend_seq_lens_cpu_int,
scale_value=layer.scaling,
num_heads=layer.tp_q_head_num,
slopes=slopes,
is_extend=True,
)
else:
if layer.qk_head_dim != layer.v_head_dim:
attn_output = q.new_empty(
@@ -1270,6 +1402,7 @@ class AscendAttnBackend(AttentionBackend):
k_rope: Optional[torch.Tensor] = None,
topk_indices: Optional[torch.Tensor] = None,
sinks: Optional[torch.Tensor] = None,
slopes: Optional[torch.Tensor] = None,
):
if is_mla_preprocess_enabled():
# MLAPO does saving kv_cache
@@ -1356,23 +1489,37 @@ class AscendAttnBackend(AttentionBackend):
else:
query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
num_tokens = query.shape[0]
attn_output = torch.empty(
(num_tokens, layer.tp_q_head_num, layer.v_head_dim),
dtype=query.dtype,
device=query.device,
)
if not self.use_alibi:
attn_output = torch.empty(
(num_tokens, layer.tp_q_head_num, layer.v_head_dim),
dtype=query.dtype,
device=query.device,
)
torch_npu._npu_paged_attention(
query=query,
key_cache=k_cache,
value_cache=v_cache,
num_heads=layer.tp_q_head_num,
num_kv_heads=layer.tp_k_head_num,
scale_value=layer.scaling,
block_table=self.forward_metadata.block_tables,
context_lens=self.forward_metadata.seq_lens_cpu_int,
out=attn_output,
)
torch_npu._npu_paged_attention(
query=query,
key_cache=k_cache,
value_cache=v_cache,
num_heads=layer.tp_q_head_num,
num_kv_heads=layer.tp_k_head_num,
scale_value=layer.scaling,
block_table=self.forward_metadata.block_tables,
context_lens=self.forward_metadata.seq_lens_cpu_int,
out=attn_output,
)
else:
attn_output = self.attn_alibi(
q=query,
k_cache=k_cache,
v_cache=v_cache,
block_tables=self.forward_metadata.block_tables,
seq_lens=self.forward_metadata.seq_lens_cpu_int,
query_lens=torch.ones(num_tokens, dtype=torch.int32),
scale_value=layer.scaling,
num_heads=layer.tp_q_head_num,
slopes=slopes,
is_extend=False,
)
return attn_output.view(num_tokens, layer.tp_q_head_num * layer.v_head_dim)
else:
if save_kv_cache:

68
python/sglang/srt/models/baichuan.py
View File

@@ -46,7 +46,9 @@ from sglang.srt.layers.vocab_parallel_embedding import (
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.utils import add_prefix
from sglang.srt.utils import add_prefix, is_npu
_is_npu = is_npu()
def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
@@ -124,20 +126,19 @@ class BaiChuanAttention(nn.Module):
max_position_embeddings: int = 8192,
quant_config: Optional[QuantizationConfig] = None,
layer_id: int = 0,
dtype: Optional[torch.dtype] = torch.bfloat16,
prefix: str = "",
):
super().__init__()
self.hidden_size = hidden_size
tensor_model_parallel_world_size = get_tensor_model_parallel_world_size()
tp_size = get_tensor_model_parallel_world_size()
self.total_num_heads = num_heads
assert self.total_num_heads % tensor_model_parallel_world_size == 0
self.num_heads = self.total_num_heads // tensor_model_parallel_world_size
self.total_num_kv_heads = self.total_num_heads
assert self.total_num_heads % tp_size == 0
self.head_dim = hidden_size // self.total_num_heads
self.postion_embedding = position_embedding
self.position_embedding = position_embedding
self.rope_theta = rope_theta
self.max_position_embeddings = max_position_embeddings
self.total_num_kv_heads = self.num_heads
if self.total_num_kv_heads >= tp_size:
# Number of KV heads is greater than TP size, so we partition
# the KV heads across multiple tensor parallel GPUs.
@@ -147,6 +148,7 @@ class BaiChuanAttention(nn.Module):
# the KV heads across multiple tensor parallel GPUs.
assert tp_size % self.total_num_kv_heads == 0
self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
self.num_heads = self.num_kv_heads
# pylint: disable=invalid-name
self.W_pack = QKVParallelLinear(
@@ -156,30 +158,36 @@ class BaiChuanAttention(nn.Module):
self.total_num_heads,
bias=False,
quant_config=quant_config,
prefix=add_prefix("W_pack", prefix),
)
self.o_proj = RowParallelLinear(
self.total_num_heads * self.head_dim,
hidden_size,
bias=False,
quant_config=quant_config,
prefix=add_prefix("o_proj", prefix),
)
self.scaling = self.head_dim**-0.5
self.attn = RadixAttention(
self.num_heads,
self.head_dim,
self.scaling,
num_kv_heads=self.num_kv_heads,
layer_id=layer_id,
quant_config=quant_config,
prefix=add_prefix("attn", prefix),
)
# Create the alibi slopes and slice them.
if self.postion_embedding == "ALIBI":
if self.position_embedding == "ALIBI":
tp_rank = get_tensor_model_parallel_rank()
head_start = tp_rank * self.num_heads
head_end = (tp_rank + 1) * self.num_heads
alibi_slopes = _get_alibi_slopes(self.total_num_heads)
alibi_slopes = alibi_slopes[head_start:head_end].tolist()
scaling = self.head_dim**-0.5
self.attn = RadixAttention(
self.num_heads,
self.head_dim,
scaling,
num_kv_heads=self.num_kv_heads,
layer_id=layer_id,
quant_config=quant_config,
prefix=add_prefix("attn", prefix),
alibi_slopes = alibi_slopes[head_start:head_end]
self.alibi_slopes = torch.tensor(
alibi_slopes, dtype=dtype, device="npu" if _is_npu else "cuda"
)
else:
self.rotary_emb = get_rope(
@@ -188,16 +196,10 @@ class BaiChuanAttention(nn.Module):
max_position=self.max_position_embeddings,
base=self.rope_theta,
)
self.scaling = self.head_dim**-0.5
self.attn = RadixAttention(
self.num_heads,
self.head_dim,
self.scaling,
num_kv_heads=self.num_kv_heads,
layer_id=layer_id,
quant_config=quant_config,
prefix=add_prefix("attn", prefix),
)
self.attn_kwargs = {}
if self.position_embedding == "ALIBI" and _is_npu:
self.attn_kwargs["slopes"] = self.alibi_slopes
def forward(
self,
@@ -207,9 +209,9 @@ class BaiChuanAttention(nn.Module):
) -> torch.Tensor:
qkv, _ = self.W_pack(hidden_states)
q, k, v = qkv.chunk(chunks=3, dim=-1)
if self.postion_embedding != "ALIBI":
if self.position_embedding != "ALIBI":
q, k = self.rotary_emb(positions, q, k)
attn_output = self.attn(q, k, v, forward_batch)
attn_output = self.attn(q, k, v, forward_batch, **self.attn_kwargs)
output, _ = self.o_proj(attn_output)
return output
@@ -292,6 +294,8 @@ class BaiChuanModel(nn.Module):
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
config.hidden_size,
org_num_embeddings=config.vocab_size,
prefix=add_prefix("embed_tokens", prefix),
)
self.layers = nn.ModuleList(
[
@@ -342,7 +346,9 @@ class BaiChuanBaseForCausalLM(nn.Module):
"gate_up_proj",
"down_proj",
]
embedding_modules = {}
embedding_modules = {
"embed_tokens": ["embed_tokens"],
}
embedding_padding_modules = []
def __init__(

3
python/sglang/test/ascend/gsm8k_ascend_mixin.py
View File

@@ -22,6 +22,7 @@ class GSM8KAscendMixin(ABC):
"ascend",
"--disable-cuda-graph",
]
gsm8k_num_shots = 5
@classmethod
def setUpClass(cls):
@@ -52,7 +53,7 @@ class GSM8KAscendMixin(ABC):
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
num_shots=self.gsm8k_num_shots,
data_path=None,
num_questions=200,
max_new_tokens=512,