[NPU]support model MiniCPM3-4B for npu (#16866)
This commit is contained in:
McZyWu
@@ -211,10 +211,16 @@ class AscendAttnBackend(AttentionBackend):
|
||||
if self.use_mla:
|
||||
self.kv_lora_rank = model_runner.model_config.kv_lora_rank
|
||||
self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim
|
||||
self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
|
||||
self.q_head_dim = (
|
||||
self.qk_rope_head_dim + model_runner.model_config.qk_nope_head_dim
|
||||
)
|
||||
if (
|
||||
"MiniCPM3ForCausalLM"
|
||||
in model_runner.model_config.hf_config.architectures
|
||||
):
|
||||
self.qk_nope_head_dim = (
|
||||
model_runner.model_config.hf_config.qk_nope_head_dim
|
||||
)
|
||||
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
|
||||
self.native_attn = TorchNativeAttnBackend(model_runner)
|
||||
self.graph_metadata = {}
|
||||
self.max_context_len = model_runner.model_config.context_len
|
||||
@@ -834,43 +840,82 @@ class AscendAttnBackend(AttentionBackend):
|
||||
assert (
|
||||
layer.qk_head_dim != layer.v_head_dim
|
||||
), "FIA only supports qk_head_dim != v_head_dim"
|
||||
|
||||
num_token_padding = q.shape[0]
|
||||
q, k, v = [
|
||||
data[: forward_batch.num_token_non_padded_cpu] for data in [q, k, v]
|
||||
]
|
||||
|
||||
q_nope, q_rope = q.split([layer.v_head_dim, self.qk_rope_head_dim], dim=-1)
|
||||
k_nope, k_rope = k.split([layer.v_head_dim, self.qk_rope_head_dim], dim=-1)
|
||||
|
||||
attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score(
|
||||
q_nope,
|
||||
k_nope,
|
||||
v,
|
||||
query_rope=q_rope,
|
||||
key_rope=k_rope,
|
||||
num_heads=layer.tp_q_head_num,
|
||||
input_layout="TND",
|
||||
atten_mask=self.fia_mask,
|
||||
sparse_mode=3,
|
||||
actual_seq_lengths=self.forward_metadata.seq_lens_list_cumsum,
|
||||
actual_seq_lengths_kv=self.forward_metadata.seq_lens_list_cumsum,
|
||||
scale=layer.scaling,
|
||||
next_tokens=0,
|
||||
)
|
||||
|
||||
attn_output = attn_output.reshape(-1, layer.tp_q_head_num, layer.v_head_dim)
|
||||
if num_token_padding != forward_batch.num_token_non_padded_cpu:
|
||||
attn_output = torch.cat(
|
||||
[
|
||||
attn_output,
|
||||
attn_output.new_zeros(
|
||||
num_token_padding - attn_output.shape[0],
|
||||
*attn_output.shape[1:],
|
||||
),
|
||||
],
|
||||
dim=0,
|
||||
if layer.v_head_dim in [256]:
|
||||
"""Currently, in NO_QUANT situation, qk_nope_head_dim == v_head_dim, and rope exists, v_head_dim only support 512 and 128"""
|
||||
kv_lora_rank = k.shape[-1] - self.qk_rope_head_dim
|
||||
kv_c, k_rope = k.split([kv_lora_rank, self.qk_rope_head_dim], dim=-1)
|
||||
if save_kv_cache:
|
||||
forward_batch.token_to_kv_pool.set_kv_buffer(
|
||||
layer, forward_batch.out_cache_loc, kv_c, k_rope
|
||||
)
|
||||
attn_output = q.new_empty(
|
||||
(q.shape[0], layer.tp_q_head_num, kv_lora_rank)
|
||||
)
|
||||
use_gqa = layer.tp_q_head_num != layer.tp_k_head_num
|
||||
|
||||
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
|
||||
)
|
||||
kv_cache = torch.cat([k_cache, v_cache], dim=-1)
|
||||
attn_output = self.native_attn._run_sdpa_forward_extend(
|
||||
q,
|
||||
attn_output,
|
||||
kv_cache.view(-1, layer.tp_k_head_num, layer.qk_head_dim),
|
||||
k_cache.view(-1, layer.tp_v_head_num, layer.v_head_dim),
|
||||
forward_batch.req_to_token_pool.req_to_token,
|
||||
forward_batch.req_pool_indices,
|
||||
forward_batch.seq_lens,
|
||||
forward_batch.extend_prefix_lens,
|
||||
forward_batch.extend_seq_lens,
|
||||
scaling=layer.scaling,
|
||||
enable_gqa=use_gqa,
|
||||
causal=True,
|
||||
)
|
||||
else:
|
||||
num_token_padding = q.shape[0]
|
||||
q, k, v = [
|
||||
data[: forward_batch.num_token_non_padded_cpu] for data in [q, k, v]
|
||||
]
|
||||
|
||||
q_nope, q_rope = q.split(
|
||||
[layer.v_head_dim, self.qk_rope_head_dim], dim=-1
|
||||
)
|
||||
k_nope, k_rope = k.split(
|
||||
[layer.v_head_dim, self.qk_rope_head_dim], dim=-1
|
||||
)
|
||||
|
||||
attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score(
|
||||
q_nope,
|
||||
k_nope,
|
||||
v,
|
||||
query_rope=q_rope,
|
||||
key_rope=k_rope,
|
||||
num_heads=layer.tp_q_head_num,
|
||||
input_layout="TND",
|
||||
atten_mask=self.fia_mask,
|
||||
sparse_mode=3,
|
||||
actual_seq_lengths=self.forward_metadata.seq_lens_list_cumsum,
|
||||
actual_seq_lengths_kv=self.forward_metadata.seq_lens_list_cumsum,
|
||||
scale=layer.scaling,
|
||||
next_tokens=0,
|
||||
)
|
||||
|
||||
attn_output = attn_output.reshape(
|
||||
-1, layer.tp_q_head_num, layer.v_head_dim
|
||||
)
|
||||
if num_token_padding != forward_batch.num_token_non_padded_cpu:
|
||||
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_mtp(
|
||||
@@ -1385,7 +1430,8 @@ class AscendAttnBackend(AttentionBackend):
|
||||
assert (
|
||||
self.graph_mode == False
|
||||
) # _npu_paged_attention_mla not support graph mode
|
||||
q = torch.cat([q, q_rope], dim=-1)
|
||||
if q_rope is not None:
|
||||
q = torch.cat([q, q_rope], dim=-1)
|
||||
query = q.view(-1, layer.tp_q_head_num, layer.head_dim)
|
||||
kv_c_and_k_pe_cache = torch.cat([kv_c, k_pe], dim=-1)
|
||||
kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
|
||||
|
||||
Reference in New Issue
Block a user