[Long Sequence Feat]support chunk prefill

vllm_ascend/attention/attention_v1.py

@@ -871,26 +871,64 @@ def _forward_decode_pcp_dcp(self, query: torch.Tensor,
             num_heads = self.num_heads
 
         # 1. Compute out&lse by "npu_fused_infer_attention_score"
-        attn_out, attn_lse = torch.ops.npu.npu_fused_infer_attention_score(
-            query.view(query.shape[0], 1, query.shape[1], query.shape[2]),
-            # [b,num_heads,head_size] -> [b,1,num_heads,head_size]
-            self.key_cache.view(self.key_cache.shape[0],
-                                self.key_cache.shape[1], -1),
-            self.value_cache.view(self.key_cache.shape[0],
-                                  self.key_cache.shape[1], -1),
-            num_heads=num_heads,
-            num_key_value_heads=self.num_kv_heads,
-            input_layout="BSND",
-            atten_mask=None,
-            scale=self.scale,
-            antiquant_mode=0,
-            antiquant_scale=None,
-            softmax_lse_flag=True,
-            block_table=attn_metadata.block_tables,
-            block_size=self.key_cache.shape[1],
-            actual_seq_lengths_kv=attn_metadata.decode_meta.
-            num_computed_tokens_of_pcp_dcp[:, self.pcp_rank, self.dcp_rank],
-        )
+        q_nope = query.view(query.shape[0], 1, query.shape[1], query.shape[2])
+        # [b,num_heads,head_size] -> [b,1,num_heads,head_size]
+        k_nope = self.key_cache.view(self.key_cache.shape[0],
+                                self.key_cache.shape[1], -1)
+        value = self.value_cache.view(self.key_cache.shape[0],
+                                  self.key_cache.shape[1], -1)
+        common_kwargs = {
+            'num_heads': num_heads,
+            'num_key_value_heads': self.num_kv_heads,
+            'input_layout': "BSND",
+            'atten_mask': None,
+            'scale': self.scale,
+            'antiquant_mode': 0,
+            'antiquant_scale': None,
+            'softmax_lse_flag': True,
+            'block_table': attn_metadata.block_tables,
+            'block_size': self.key_cache.shape[1],
+            "actual_seq_lengths_kv": attn_metadata.decode.num_computed_tokens_of_cp_dcp[:, self.cp_rank, self.dcp_rank],
+        }
+        graph_params = get_graph_params()
+        forward_context: ForwardContext = get_forward_context()
+        num_tokens = query.shape[0]
+        if forward_context.capturing:
+            stream = torch_npu.npu.current_stream()
+
+            event = torch.npu.ExternalEvent()
+            event.wait(stream)
+            event.reset(stream)
+            graph_params.events[num_tokens].append(event)
+
+            workspace = graph_params.workspaces.get(num_tokens)
+            if workspace is None:
+                workspace = torch_npu._npu_fused_infer_attention_score_get_max_workspace(
+                    q_nope, k_nope, value, **common_kwargs)
+                graph_params.workspaces[num_tokens] = workspace
+            attn_out = torch.empty_like(q_nope)
+            attn_lse = torch.empty((num_tokens, num_heads, 1, 1),
+                                      dtype=torch.float,
+                                      device=q_nope.device)
+
+            graph_params.attn_params[num_tokens].append(
+                (q_nope, k_nope, value, self.num_heads, self.num_kv_heads,
+                 self.scale, attn_metadata.block_tables, self.key_cache.shape[1],
+                 attn_metadata.decode.num_computed_tokens_of_cp_dcp[:, self.cp_rank, self.dcp_rank], workspace,
+                 attn_out, attn_lse, self.cp_rank, self.dcp_rank, self.dcp_size))
+            torch.npu.graph_task_group_begin(stream)
+            torch_npu.npu_fused_infer_attention_score.out(
+                q_nope,
+                k_nope,
+                value,
+                **common_kwargs,
+                workspace=workspace,
+                out=[attn_out, attn_lse])
+            handle = torch.npu.graph_task_group_end(stream)
+            graph_params.handles[num_tokens].append(handle)
+        else:
+            attn_out, attn_lse = torch_npu.npu_fused_infer_attention_score(
+                q_nope, k_nope, value, **common_kwargs)
 
         attn_out = attn_out.view(attn_out.shape[0], attn_out.shape[2],
                                  attn_out.shape[3])