diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index f9d5a831d..b3a49f8a4 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/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 diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 565d8b606..538fea08c 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/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) - 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: diff --git a/python/sglang/srt/models/baichuan.py b/python/sglang/srt/models/baichuan.py index 6e8d3b4a3..84596ba1f 100644 --- a/python/sglang/srt/models/baichuan.py +++ b/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.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__( diff --git a/python/sglang/test/ascend/gsm8k_ascend_mixin.py b/python/sglang/test/ascend/gsm8k_ascend_mixin.py index bcf3391f3..5257d8191 100644 --- a/python/sglang/test/ascend/gsm8k_ascend_mixin.py +++ b/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, diff --git a/test/registered/ascend/llm_models/test_ascend_baichuan2_13b_chat.py b/test/registered/ascend/llm_models/test_ascend_baichuan2_13b_chat.py index 31aceaac2..4472662ca 100644 --- a/test/registered/ascend/llm_models/test_ascend_baichuan2_13b_chat.py +++ b/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( - est_time=400, - suite="nightly-1-npu-a3", - nightly=True, - disabled="The accuracy test result is 0.", -) +register_npu_ci(est_time=400, suite="nightly-1-npu-a3", nightly=True) -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__":