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 1be763233..565d8b606 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -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( diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index b50f2374a..064d55743 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -756,8 +756,8 @@ class Phi3LongRoPEScaledRotaryEmbedding(nn.Module): key: torch.Tensor, offsets: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - query = query.view(*query.shape[:-1], -1, self.head_size) - key = key.view(*key.shape[:-1], -1, self.head_size) + query = query.unflatten(1, (-1, self.head_size)) + key = key.unflatten(1, (-1, self.head_size)) k = self.original_max_position_embeddings long_prompt_offset = ( diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 7a89012f2..48e83831e 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -406,7 +406,9 @@ class ModelRunnerKVCacheMixin: dtype=self.kv_cache_dtype, kv_lora_rank=self.model_config.kv_lora_rank, qk_rope_head_dim=self.model_config.qk_rope_head_dim, - index_head_dim=self.model_config.index_head_dim, + index_head_dim=( + self.model_config.index_head_dim if is_nsa_model else None + ), layer_num=self.num_effective_layers, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, diff --git a/python/sglang/srt/models/minicpm3.py b/python/sglang/srt/models/minicpm3.py index 821dfa98a..9755a6f6b 100644 --- a/python/sglang/srt/models/minicpm3.py +++ b/python/sglang/srt/models/minicpm3.py @@ -239,7 +239,9 @@ class MiniCPM3AttentionMLA(nn.Module): original_shapes = [q_pe.shape, k_pe.shape] q_pe, k_pe = self.rotary_emb( - positions, q_pe.reshape(q_pe.shape[0], -1), k_pe.reshape(k_pe.shape[0], -1) + positions, + q_pe.reshape(-1, q_pe.shape[1] * q_pe.shape[2]), + k_pe.reshape(-1, k_pe.shape[1] * k_pe.shape[2]), ) q_pe, k_pe = q_pe.view(original_shapes[0]), k_pe.view(original_shapes[1]) q_input[..., self.kv_lora_rank :] = q_pe diff --git a/test/registered/ascend/llm_models/test_ascend_minicpm3_4b.py b/test/registered/ascend/llm_models/test_ascend_minicpm3_4b.py new file mode 100644 index 000000000..f53692999 --- /dev/null +++ b/test/registered/ascend/llm_models/test_ascend_minicpm3_4b.py @@ -0,0 +1,30 @@ +import unittest + +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) + + +class TestMiniCPM3(GSM8KAscendMixin, CustomTestCase): + model = "/root/.cache/modelscope/hub/models/OpenBMB/MiniCPM3-4B" + accuracy = 0.69 + other_args = [ + "--trust-remote-code", + "--mem-fraction-static", + "0.8", + "--attention-backend", + "ascend", + "--disable-cuda-graph", + "--disable-radix-cache", + "--disable-overlap-schedule", + "--max-running-requests", + "128", + "--chunked-prefill-size", + "-1", + ] + + +if __name__ == "__main__": + unittest.main()