[feat]dcp pcp support aclgraph (#3731)

### What this PR does / why we need it?
dcp pcp support  full aclgraph, including mla attention_v1

- vLLM version: v0.11.0rc3
- vLLM main:
vllm-project/vllm@c9461e0

Signed-off-by: weiguihua2 <[email protected]>

Author: weiguihua2

tests/ut/attention/test_mla_v1.py

@@ -176,16 +176,30 @@ def test_ascend_mla_metadata_default(self):
 
 class TestAscendMLAMetadataBuilder(TestBase):
 
-    def test_ascend_mla_metadata_builder_default(self):
+    @patch('vllm.distributed.parallel_state.get_dcp_group')
+    @patch('vllm.distributed.parallel_state._DCP',
+           new_callable=lambda: MagicMock(spec=GroupCoordinator))
+    @patch("vllm.distributed.get_decode_context_model_parallel_world_size",
+           return_value=1)
+    def test_ascend_mla_metadata_builder_default(self, mock_get_dcp_size,
+                                                 mock_dcp, mock_get_dcp_group):
         mock_vllm_config = MagicMock()
         mock_vllm_config.model_config.max_model_len = 1024
         mock_vllm_config.model_config.get_head_size.return_value = 64
         mock_vllm_config.model_config.dtype = torch.float16
         mock_vllm_config.cache_config.block_size = 16
         mock_vllm_config.scheduler_config.max_num_seqs = 4
+        mock_vllm_config.scheduler_config.decode_max_num_seqs = 4
         mock_vllm_config.scheduler_config.chunked_prefill_enabled = False
         mock_device = 'cpu'
 
+        mock_dcp.world_size = 1
+        dcp_group = MagicMock(spec=GroupCoordinator)
+        dcp_group.rank_in_group = 0
+        dcp_group.world_size = 1
+        dcp_group.device_group = MagicMock()
+        mock_get_dcp_group.return_value = dcp_group
+
         mock_vllm_config.speculative_config = None
 
         ascend_config = MagicMock()
@@ -200,16 +214,31 @@ def test_ascend_mla_metadata_builder_default(self):
                 builder.chunked_prefill_enabled,
                 mock_vllm_config.scheduler_config.chunked_prefill_enabled)
 
-    def test_ascend_mla_metadata_builder_spec_decode(self):
+    @patch('vllm.distributed.parallel_state.get_dcp_group')
+    @patch('vllm.distributed.parallel_state._DCP',
+           new_callable=lambda: MagicMock(spec=GroupCoordinator))
+    @patch("vllm.distributed.get_decode_context_model_parallel_world_size",
+           return_value=1)
+    def test_ascend_mla_metadata_builder_spec_decode(self, mock_get_dcp_size,
+                                                     mock_dcp,
+                                                     mock_get_dcp_group):
         mock_vllm_config = MagicMock()
         mock_vllm_config.model_config.max_model_len = 1024
         mock_vllm_config.model_config.get_head_size.return_value = 64
         mock_vllm_config.model_config.dtype = torch.float16
         mock_vllm_config.cache_config.block_size = 16
         mock_vllm_config.scheduler_config.max_num_seqs = 4
+        mock_vllm_config.scheduler_config.decode_max_num_seqs = 4
         mock_vllm_config.scheduler_config.chunked_prefill_enabled = False
         mock_device = 'cpu'
 
+        mock_dcp.world_size = 1
+        dcp_group = MagicMock(spec=GroupCoordinator)
+        dcp_group.rank_in_group = 0
+        dcp_group.world_size = 1
+        dcp_group.device_group = MagicMock()
+        mock_get_dcp_group.return_value = dcp_group
+
         mock_spec_config = MagicMock()
         mock_spec_config.num_speculative_tokens = 3
         mock_vllm_config.speculative_config = mock_spec_config
@@ -226,16 +255,30 @@ def test_ascend_mla_metadata_builder_spec_decode(self):
                 builder.chunked_prefill_enabled,
                 mock_vllm_config.scheduler_config.chunked_prefill_enabled)
 
-    def test_reorder_batch(self):
+    @patch('vllm.distributed.parallel_state.get_dcp_group')
+    @patch('vllm.distributed.parallel_state._DCP',
+           new_callable=lambda: MagicMock(spec=GroupCoordinator))
+    @patch("vllm.distributed.get_decode_context_model_parallel_world_size",
+           return_value=1)
+    def test_reorder_batch(self, mock_get_dcp_size, mock_dcp,
+                           mock_get_dcp_group):
         ascend_config = MagicMock()
 
         mock_vllm_config = MagicMock()
         mock_vllm_config.model_config.max_model_len = 1024
         mock_vllm_config.cache_config.block_size = 16
         mock_vllm_config.scheduler_config.max_num_seqs = 4
+        mock_vllm_config.scheduler_config.decode_max_num_seqs = 4
         mock_vllm_config.scheduler_config.chunked_prefill_enabled = False
         mock_device = 'cpu'
 
+        mock_dcp.world_size = 1
+        dcp_group = MagicMock(spec=GroupCoordinator)
+        dcp_group.rank_in_group = 0
+        dcp_group.world_size = 1
+        dcp_group.device_group = MagicMock()
+        mock_get_dcp_group.return_value = dcp_group
+
         mock_vllm_config.speculative_config = None
 
         with patch("vllm_ascend.attention.mla_v1.get_ascend_config",

vllm_ascend/attention/attention_v1.py

@@ -865,26 +865,81 @@ 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.
+        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_meta.
             num_computed_tokens_of_pcp_dcp[:, self.pcp_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)
+                update_graph_params_workspaces(num_tokens,
+                                               weak_ref_tensors(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(
+                (weak_ref_tensors(q_nope), weak_ref_tensors(k_nope),
+                 weak_ref_tensors(value), self.num_heads, self.num_kv_heads,
+                 self.scale, attn_metadata.block_tables,
+                 self.key_cache.shape[1], attn_metadata.decode_meta.
+                 num_computed_tokens_of_pcp_dcp[:, self.pcp_rank,
+                                                self.dcp_rank],
+                 weak_ref_tensors(attn_out), weak_ref_tensors(attn_lse),
+                 self.pcp_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])