[pref] calculate local_total_toks in build meata

Signed-off-by: F.Liu <liufeng248@huawei.com>

Author: F.Liu

vllm_ascend/attention/attention_cp.py

@@ -152,6 +152,7 @@ def build(
                                                                          dcp_rank]
                 actual_seq_lengths_kv = torch.cumsum(
                     local_chunked_kv_lens_rank, dim=0).tolist()
+                local_total_toks = local_chunked_kv_lens_rank.sum()
                 chunked_req_mask = self._get_chunked_req_mask(
                     local_context_lens_allranks)
                 local_chunk_starts = torch.zeros(
@@ -181,7 +182,8 @@ def build(
                         cp_kv_recover_idx_for_chunk=cp_kv_recover_idx_for_chunk,
                         kv_inverse_idx_for_chunk=kv_inverse_idx_for_chunk,
                         batch_chunk_seq_mask=batch_chunk_seq_mask,
-                        chunk_seq_mask_filtered_indices=chunk_seq_mask_filtered_indices
+                        chunk_seq_mask_filtered_indices=chunk_seq_mask_filtered_indices,
+                        local_total_toks=local_total_toks.item()
                     )
             attn_mask_seqlens = common_long_seq_metadata.attn_mask_seqlens
             head_attn_nomask_seqlens = common_long_seq_metadata.head_attn_nomask_seqlens
@@ -672,7 +674,7 @@ def _compute_prefill_context(self, query: torch.Tensor,
 
         local_chunked_kv_lens_rank = local_chunked_kv_lens[:, self.pcp_rank,
                                                            self.dcp_rank]
-        total_toks = local_chunked_kv_lens_rank.sum()
+        total_toks = prefill_metadata.chunked_context.local_total_toks
 
         key, value = self._load_kv_for_chunk(attn_metadata, kv_cache,
                                              local_chunked_kv_lens_rank, query,

vllm_ascend/attention/utils.py

@@ -63,6 +63,7 @@ class ChunkedContextMetadata:
         cp_kv_recover_idx_for_chunk: Optional[list[int]] = None
         kv_inverse_idx_for_chunk: Optional[list[int]] = None
         batch_chunk_seq_mask: Optional[list[bool]] = None
+        local_total_toks: int = None
 
     """ Prefill Specific Metadata for Ascend"""
     pcp_metadata: Optional[AscendPCPMetadata] = None