fix

Author: wangyao-i

vllm_ascend/attention/attention_v1.py

@@ -530,7 +530,6 @@ def __init__(
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
         self.key_cache = None
         self.value_cache = None
-        self.block_size = VllmConfig().cache_config.block_size
         self.pcp_size = get_prefill_context_model_parallel_world_size(
         ) if prefill_context_parallel_enable() else 1
         self.pcp_rank = get_prefill_context_model_parallel_rank(
@@ -575,12 +574,12 @@ def _forward_prefill_no_cache(
             output,_ = torch_npu.npu_fused_infer_attention_score_v2(
                 query[:num_tokens],
                 key[:num_tokens],
-                atten_mask=maks.to(torch.bool),
+                atten_mask=mask.to(torch.bool),
                 actual_seq_qlen=attn_metadata.query_lens.cumsum(0),
-                actual_seq_kvlen=attn_metadata.seq_lens.cumcum(0),
+                actual_seq_kvlen=attn_metadata.seq_lens.cumsum(0),
                 num_query_heads=self.num_heads,
                 num_key_value_heads=self.num_kv_heads,
-                imput_layout="TND",
+                input_layout="TND",
                 softmax_scale=self.scale
             )
             assert output is not None
@@ -615,24 +614,24 @@ def _forward_prefill_cache_hit(
 
         if is_A5():
             compress_mask = compress_mask.to(torch.bool)
-            key = self.key_cache.transpos(1,2)
-            value = slef.value_cache.transpose(1,2)
+            key = self.key_cache.transpose(1,2)
+            value = self.value_cache.transpose(1,2)
             block_size = self.block_size
 
-            output, _ = troch_npu.npu_fused_infer_attention_score_v2(
+            output, _ = torch_npu.npu_fused_infer_attention_score_v2(
                 query=query,
                 key=key,
                 value=value,
                 block_table=block_table,
-                atten_mask=mask,
+                atten_mask=compress_mask,
                 actual_seq_qlen=attn_metadata.query_lens.cumsum(0),
                 actual_seq_kvlen=attn_metadata.seq_lens,
                 num_query_heads=self.num_heads,
                 num_key_value_heads=self.num_kv_heads,
                 softmax_scale=self.scale,
                 spare_mode=2, #spare_mode=2时，代表leftupCausal模式的mask
                 block_size=block_size,
-                imput_layout="TND"
+                input_layout="TND"
             )
             return output
 
@@ -768,23 +767,24 @@ def _forward_decode_only(
             else:
                 if is_A5(): 
                     batch_size = attn_metadata.query_lens.shape[0]
-                    hidden_szie = self.num_heads * self.head_size
-                    query = query[:batch_szie]
+                    hidden_size = self.num_heads * self.head_size
+                    query = query[:batch_size]
                     query = query.view(batch_size, 1, hidden_size)
                     block_size = self.key_cache.shape[1]
                     key = self.key_cache.flatten(2, 3).contiguous()
+                    value = self.value_cache.flatten(2, 3).contiguous()
                     ori_output = output
-                    output, _ = torch_nup.npu_fused_infer_attention_score_v2(
+                    output, _ = torch_npu.npu_fused_infer_attention_score_v2(
                         query=query,
                         key=key,
                         value=value,
-                        actual_seq_kvlen=attn_metadata.seq_len,
+                        actual_seq_kvlen=attn_metadata.seq_lens,
                         num_query_heads=self.num_heads,
                         num_key_value_heads=self.num_kv_heads,
-                        block_table=attn_metadata.block_tables[:batch_szie],
+                        block_table=attn_metadata.block_tables[:batch_size],
                         block_size=block_size,
                         softmax_scale=self.scale,
-                        inpt_layout="BSH"
+                        input_layout="BSH"
                     )
                     output = output.view(-1, self.num_heads, self.head_size)
                     ori_output[:batch_size] = output[:batch_size]
@@ -859,9 +859,9 @@ def _forward_v1_style(
                 num_query_heads=self.num_heads,
                 num_key_value_heads=self.num_kv_heads,
                 block_table=attn_metadata.block_tables[:attn_metadata.query_lens.shape[0]],
-                 block_size=self.key_cache.shape[1],
+                block_size=self.key_cache.shape[1],
                 softmax_scale=self.scale,
-                imput_layout="TND"
+                input_layout="TND"
             )
             return output
         output, _ = torch_npu.npu_fused_infer_attention_score(
@@ -1611,23 +1611,24 @@ def forward(
                 if is_A5(): # 这里代码变动较大需要重新适配
                     num_token = slots.shape[0]
                     torch_npu.npu_scatter_a_kv_cache(
-                        key=key[:num_tokens],
-                        value=value[:num_tokens],
-                        slot_mapping=slots,
+                        key=key[self.pcp_size * num_decode_tokens:attn_metadata.num_actual_tokens_pcp_padded],
+                        value=value[self.pcp_size * num_decode_tokens:attn_metadata.num_actual_tokens_pcp_padded],
+                        slot_mapping=slot_mapping[self.pcp_size * num_decode_tokens:attn_metadata.num_actual_tokens_pcp_padded]
                         out=(self.key_cache, slef.value_cache)
                     )
-                torch_npu._npu_reshape_and_cache(
-                    key=key[self.pcp_size * num_decode_tokens:attn_metadata.
-                            num_actual_tokens_pcp_padded],
-                    value=value[self.pcp_size *
-                                num_decode_tokens:attn_metadata.
+                else:
+                    torch_npu._npu_reshape_and_cache(
+                        key=key[self.pcp_size * num_decode_tokens:attn_metadata.
                                 num_actual_tokens_pcp_padded],
-                    key_cache=self.key_cache,
-                    value_cache=self.value_cache,
-                    slot_indices=attn_metadata.
-                    slot_mapping[self.pcp_size *
-                                 num_decode_tokens:attn_metadata.
-                                 num_actual_tokens_pcp_padded])
+                        value=value[self.pcp_size *
+                                    num_decode_tokens:attn_metadata.
+                                    num_actual_tokens_pcp_padded],
+                        key_cache=self.key_cache,
+                        value_cache=self.value_cache,
+                        slot_indices=attn_metadata.
+                        slot_mapping[self.pcp_size *
+                                    num_decode_tokens:attn_metadata.
+                                    num_actual_tokens_pcp_padded])
 
         if self.pcp_size * self.dcp_size > 1:
             intermediate_output = self._forward_pcp_dcp(

vllm_ascend/attention/mla_v1.py

@@ -932,6 +932,8 @@ def _compute_prefill_context(
         cache_k_pe = kv_c_and_k_pe_cache[1]
         num_heads = cache_k_pe.size(2)
         latent_kv_dim = kv_c_and_k_pe_cache[0].size(-1)
+        all_prefix_lse = [prefix_lse.view(-1)]
+        all_prefix_output = [prefix_output.view(-1, q_nope.shape[-1])]
         # token -> request mapping for building per-token masks when CP>1
         seq_len1 = torch.tensor(prefill_metadata.query_lens,
                                 dtype=torch.int32,
@@ -990,17 +992,27 @@ def _compute_prefill_context(
                                        rope_dim,
                                        dtype=q_nope.dtype,
                                        device=q_nope.device)
-
-                    torch_npu.atb.npu_paged_cache_load(
-                        cache_kv_c,
-                        cache_k_pe,
-                        prefill_metadata.block_table,
-                        seq_len2_rank.to(q_nope.device),
-                        seq_starts=
-                        context_starts_rank,  # slot offsets of current chunk in current iteration
-                        key=kv_c_normed,
-                        value=k_pe,
-                    )
+                    if is_A5():
+                        torch_npu.npu_gather_pa_kv_cache(
+                            cache_kv_c,
+                            cache_k_pe,
+                            prefill_metadata.block_table,
+                            context_seq_len_npu,
+                            key=kv_c_normed,
+                            value=k_pe,
+                            seq_offset=prefill_metadata.chunked_context.starts[i],
+                        )
+                    else:
+                        torch_npu.atb.npu_paged_cache_load(
+                            cache_kv_c,
+                            cache_k_pe,
+                            prefill_metadata.block_table,
+                            seq_len2_rank.to(q_nope.device),
+                            seq_starts=
+                            context_starts_rank,  # slot offsets of current chunk in current iteration
+                            key=kv_c_normed,
+                            value=k_pe,
+                        )
                     seq_len2 = seq_len2_rank.to(q_nope.device)
                 else:
                     # If current rank has no tokens to process, create empty tensors

vllm_ascend/distributed/llmdatadist_c_mgr_connector.py

@@ -503,7 +503,7 @@ def get_device_info(self, global_rank_table, device_filter, device_type):
             and device_filter(d.get("device_id", ""))
         ]
         if len(device_list) <= self.pcp_rank * self.tp_size + self.tp_rank:
-            retunr None
+            return None
         device_info = device_list[self.pcp_rank * self.tp_size + self.tp_rank]
         return device_info
 
@@ -531,7 +531,7 @@ def read_agent_metadata(self, global_rank_table):
                 agent_metadata = LLMDataDistCMgrAgentMetadataA5(
                     server_id=server_id_,
                     device_id=device_id_,
-                    device_ip=device_ip_,
+                    device_ip=device_id_,
                     cluster_id=cluster_id_,
                     level_list = level_list_,
                 )

vllm_ascend/ops/fused_moe/experts_selector.py

@@ -198,7 +198,7 @@ def _select_experts_with_fusion_ops(
     if not use_grouped_topk and custom_routing_function is None and scoring_func == "softmax":
         if is_A5():
             # A5 MOCK
-            new_shape = router_logits.shape[-1] + (topk,)
+            new_shape = router_logits.shape[:-1] + (topk,)
             topk_weights = torch.ones(new_shape, dtype=router_logits.dtype, device=router_logits.device)
             topk_ids = torch.zeros(topk_weights.shape, dtype=torch.int32, device=router_logits.device)
         else :

vllm_ascend/ops/rotary_embedding.py

@@ -421,8 +421,8 @@ def forward_oot(
                 query.dtype)  # type: ignore
 
         if is_A5(): # A5不支持npu_mrope算子，这里需要使用小算子替换
-            return
-        
+        return query, key
+                
         query, key = torch_npu.npu_mrope(positions,
                                          query.contiguous(),
                                          key.contiguous(),

vllm_ascend/worker/model_runner_v1.py

@@ -1076,8 +1076,8 @@ def _make_attention_mask(self, seq_lens, position,
         # Prefill without cache situation.
         elif attn_state == AscendAttentionState.PrefillNoCache:
             if is_A5():
-                mas_seq_len = max(seq_lens, default=0)
-                max_seq_len = (max_seq_len + self.block_szie - 1) // self.block_size * self.block_size
+                max_seq_len = max(seq_lens, default=0)
+                max_seq_len = (max_seq_len + self.block_size - 1) // self.block_size * self.block_size
                 new_element = torch.tensor([max_seq_len])
                 seq_lens = torch.cat([seq_lens, new_element], dim =0)
                 return self.attn_mask_builder.get_attn_mask(max_seq_len, self.dtype, self.device).to(torch.bool)

vllm_ascend/worker/worker_v1.py

@@ -338,7 +338,7 @@ def compile_or_warm_up_model(self) -> None:
             self.model_runner.capture_model()
         # Call ATB matmul to warm up; otherwise, the first operation (ReshapeAndCache)
         # may cause performance degradation at runtime.
-        if ~is_A5():
+        if not is_A5():
             self._warm_up_atb()
         # Reset the seed to ensure that the random state is not affected by
         # the model initialization and profiling.