xx

Author: wangxiaoxin-sherie

vllm_ascend/attention/attention_v1.py

@@ -192,6 +192,7 @@ class AscendMetadata:
     seq_lens: torch.Tensor = None
     seq_lens_list: List[int] = None  # type: ignore
     actual_seq_lengths_q: List[int] = None  # type: ignore
+    query_start_loc_list: List[int] = None  # type: ignore
 
     query_start_loc: torch.Tensor = None
     query_lens: torch.Tensor = None
@@ -360,6 +361,7 @@ def build(
             num_actual_tokens_pcp_padded=num_actual_tokens_pcp_padded,
             block_tables=block_table,
             query_start_loc=query_start_loc,
+            query_start_loc_list=query_start_loc_cpu.tolist(),
             query_lens=query_lens,
             seq_lens=seq_lens,
             seq_lens_list=seq_lens.tolist(),
@@ -454,7 +456,7 @@ def full_graph_attention(self,
         forward_context: ForwardContext = get_forward_context()
         if forward_context.capturing:
             graph_params = get_graph_params()
-            query_start_loc = attn_metadata.actual_seq_lengths_q
+            query_start_loc = attn_metadata.query_start_loc_list
             seq_lens = attn_metadata.seq_lens_lis
             # Prepare tensors for attention output
             # TODO: Refactor this to step-level instead of layer-level