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

Co-authored-by: cy <chenyang08056032@163.com>

python/sglang/srt/configs/model_config.py

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

python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py

@@ -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)
+                    if not self.use_alibi:
-                    attn_output = torch.empty(
+                        query = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
-                        (query.shape[0], layer.tp_q_head_num * layer.v_head_dim),
+                        attn_output = torch.empty(
-                        dtype=query.dtype,
+                            (query.shape[0], layer.tp_q_head_num * layer.v_head_dim),
-                        device=query.device,
+                            dtype=query.dtype,
-                    )
+                            device=query.device,
+                        )
 
-                    torch_npu._npu_flash_attention_qlens(
+                        torch_npu._npu_flash_attention_qlens(
-                        query=query,
+                            query=query,
-                        key_cache=k_cache,
+                            key_cache=k_cache,
-                        value_cache=v_cache,
+                            value_cache=v_cache,
-                        mask=self.mask,
+                            mask=self.mask,
-                        block_table=self.forward_metadata.block_tables,
+                            block_table=self.forward_metadata.block_tables,
-                        seq_len=self.forward_metadata.extend_seq_lens_cpu_int,
+                            seq_len=self.forward_metadata.extend_seq_lens_cpu_int,
-                        context_lens=self.forward_metadata.seq_lens_cpu_int,
+                            context_lens=self.forward_metadata.seq_lens_cpu_int,
-                        scale_value=layer.scaling,
+                            scale_value=layer.scaling,
-                        num_heads=layer.tp_q_head_num,
+                            num_heads=layer.tp_q_head_num,
-                        num_kv_heads=layer.tp_k_head_num,
+                            num_kv_heads=layer.tp_k_head_num,
-                        out=attn_output,
+                            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(
+                if not self.use_alibi:
-                    (num_tokens, layer.tp_q_head_num, layer.v_head_dim),
+                    attn_output = torch.empty(
-                    dtype=query.dtype,
+                        (num_tokens, layer.tp_q_head_num, layer.v_head_dim),
-                    device=query.device,
+                        dtype=query.dtype,
-                )
+                        device=query.device,
+                    )
 
-                torch_npu._npu_paged_attention(
+                    torch_npu._npu_paged_attention(
-                    query=query,
+                        query=query,
-                    key_cache=k_cache,
+                        key_cache=k_cache,
-                    value_cache=v_cache,
+                        value_cache=v_cache,
-                    num_heads=layer.tp_q_head_num,
+                        num_heads=layer.tp_q_head_num,
-                    num_kv_heads=layer.tp_k_head_num,
+                        num_kv_heads=layer.tp_k_head_num,
-                    scale_value=layer.scaling,
+                        scale_value=layer.scaling,
-                    block_table=self.forward_metadata.block_tables,
+                        block_table=self.forward_metadata.block_tables,
-                    context_lens=self.forward_metadata.seq_lens_cpu_int,
+                        context_lens=self.forward_metadata.seq_lens_cpu_int,
-                    out=attn_output,
+                        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:

python/sglang/srt/models/baichuan.py

@@ -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.total_num_kv_heads = self.total_num_heads
-        self.num_heads = self.total_num_heads // tensor_model_parallel_world_size
+        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()
+            alibi_slopes = alibi_slopes[head_start:head_end]
-
+            self.alibi_slopes = torch.tensor(
-            scaling = self.head_dim**-0.5
+                alibi_slopes, dtype=dtype, device="npu" if _is_npu else "cuda"
-            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),
             )
         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.attn_kwargs = {}
-                self.num_heads,
+        if self.position_embedding == "ALIBI" and _is_npu:
-                self.head_dim,
+            self.attn_kwargs["slopes"] = self.alibi_slopes
-                self.scaling,
-                num_kv_heads=self.num_kv_heads,
-                layer_id=layer_id,
-                quant_config=quant_config,
-                prefix=add_prefix("attn", prefix),
-            )
 
     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__(

python/sglang/test/ascend/gsm8k_ascend_mixin.py

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

test/registered/ascend/llm_models/test_ascend_baichuan2_13b_chat.py

@@ -4,17 +4,26 @@ from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin
 from sglang.test.ci.ci_register import register_npu_ci
 from sglang.test.test_utils import CustomTestCase
 
-register_npu_ci(
+register_npu_ci(est_time=400, suite="nightly-1-npu-a3", nightly=True)
-    est_time=400,
-    suite="nightly-1-npu-a3",
-    nightly=True,
-    disabled="The accuracy test result is 0.",
-)
 
 
-class TestMistral7B(GSM8KAscendMixin, CustomTestCase):
+class TestBaichuan(GSM8KAscendMixin, CustomTestCase):
     model = "/root/.cache/modelscope/hub/models/baichuan-inc/Baichuan2-13B-Chat"
-    accuracy = 0.00
+    accuracy = 0.48
+    other_args = [
+        "--trust-remote-code",
+        "--mem-fraction-static",
+        "0.8",
+        "--attention-backend",
+        "ascend",
+        "--disable-cuda-graph",
+        "--max-running-requests",
+        "128",
+        "--disable-radix-cache",
+        "--chunked-prefill-size",
+        "-1",
+    ]
+    gsm8k_num_shots = 1
 
 
 if __name__ == "__main__":