[feature] support pcp + mtp in full graph

Signed-off-by: zhangsicheng5 <zhangsicheng5@huawei.com>

Author: zhangsicheng5

vllm_ascend/attention/utils.py

@@ -42,6 +42,10 @@ class AscendPrefillContextParallelMetadata:
 
     pcp_prefill_mask: torch.Tensor = None
 
+    query_lens_pcp_full_cpu: torch.Tensor = None
+
+    max_query_len_pcp_full: int = None
+
 
 @dataclass
 class AscendCommonAttentionMetadata:
@@ -135,10 +139,15 @@ def filter_chunked_req_indices(
 def split_decodes_and_prefills(
     common_attn_metadata: AscendCommonAttentionMetadata,
     decode_threshold: int = 1,
+    query_lens: torch.Tensor = None,
+    max_query_len: int = None,
 ) -> tuple[int, int, int, int]:
     """
     Assuming a reordered batch, finds the boundary between prefill and decode
     requests.
+    While pcp > 1, query_lens is split across pcp ranks,
+    In this case we pass in the original query_lens and max_query_len
+    to distinguish prefills and decodes.
 
     Args:
         common_attn_metadata: AscendCommonAttentionMetadata object containing the
@@ -151,15 +160,17 @@ def split_decodes_and_prefills(
         num_decode_tokens: The number of tokens in the decode requests.
         num_prefill_tokens: The number of tokens in the prefill requests.
     """
-    max_query_len = common_attn_metadata.max_query_len
+    max_query_len = common_attn_metadata.max_query_len \
+        if max_query_len is None else max_query_len
     num_reqs = common_attn_metadata.num_reqs
     num_tokens = common_attn_metadata.num_actual_tokens
     query_start_loc = common_attn_metadata.query_start_loc_cpu
 
     if max_query_len <= decode_threshold:
         return num_reqs, 0, num_tokens, 0
 
-    query_lens = query_start_loc[1:] - query_start_loc[:-1]
+    query_lens = (query_start_loc[1:] - query_start_loc[:-1]) \
+        if query_lens is None else query_lens
     is_prefill = query_lens > decode_threshold
     if not torch.any(is_prefill):
         return num_reqs, 0, num_tokens, 0

vllm_ascend/compilation/acl_graph.py

@@ -369,7 +369,10 @@ def update_attn_dcp_pcp_params(update_stream, forward_context, runtime_shape):
 
 def update_mla_attn_dcp_pcp_params(update_stream, forward_context,
                                    runtime_shape):
-    graph_params = get_graph_params()
+    if forward_context.is_mtp_model:
+        graph_params = get_mtp_graph_params()
+    else:
+        graph_params = get_graph_params()
     # FIXME: Behold! We are using a temporary hack here to update the args
     # for each layer's attention op in the graph.
     with torch.npu.stream(update_stream):

vllm_ascend/spec_decode/mtp_proposer.py

@@ -33,6 +33,7 @@
 from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
 from vllm_ascend.compilation.acl_graph import (ACLGraphWrapper,
                                                set_mtp_graph_params,
+                                               update_mla_attn_dcp_pcp_params,
                                                update_mla_attn_params)
 from vllm_ascend.spec_decode.interface import Proposer, SpecDcodeType
 from vllm_ascend.utils import (ProfileExecuteDuration, lmhead_tp_enable,
@@ -99,6 +100,7 @@ def __init__(
         self.pcp_size = self.runner.pcp_size
         self.dcp_size = self.runner.dcp_size
         self.pcp_rank = self.runner.pcp_rank
+        self.dcp_rank = self.runner.dcp_rank
 
         self.attn_metadata_builder: Optional[AttentionMetadataBuilder] = None
         self.draft_indexer_metadata_builder: Optional[
@@ -272,6 +274,13 @@ def dummy_run(self,
                     cos=self.runner.cos,
                     sin=self.runner.sin,
                 )
+                if self.pcp_size * self.dcp_size > 1:
+                    # update long_seq related params and flatten block_table
+                    common_attn_metadata.prefill_context_parallel_metadata=\
+                        self.runner.long_seq_metadata
+                    common_attn_metadata.block_table_tensor = \
+                        self.runner.input_batch.block_table[0].get_device_tensor()[
+                            :num_reqs * self.decode_threshold]
 
                 builder = self.runner.attn_groups[0][0].get_metadata_builder()
                 attn_metadata_mtp = builder.build_for_graph_capture(
@@ -317,9 +326,15 @@ def dummy_run(self,
                 if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL and \
                     not forward_context.capturing:
                     if self.vllm_config.model_config.use_mla:
-                        update_mla_attn_params(
-                            self.update_stream, forward_context, num_tokens,
-                            self.vllm_config.speculative_config)
+                        if self.pcp_size * self.dcp_size > 1:
+                            update_mla_attn_dcp_pcp_params(
+                                self.update_stream, forward_context,
+                                num_tokens)
+                        else:
+                            update_mla_attn_params(
+                                self.update_stream, forward_context,
+                                num_tokens,
+                                self.vllm_config.speculative_config)
                 if self.enable_shared_expert_dp:
                     positions = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
                         positions, True)
@@ -373,7 +388,7 @@ def generate_token_ids(self,
                                                  valid_sampled_tokens_count)
 
         req_scheduled_tokens = scheduler_output.num_scheduled_tokens
-        if self.pcp_size > 1:
+        if self.pcp_size * self.dcp_size > 1:
             long_seq_metadata = self.runner.long_seq_metadata
             input_ids_pcp_full = self.runner.input_ids_pcp_full
             query_start_loc_pcp_full = self.runner.query_start_loc_pcp_full
@@ -409,7 +424,6 @@ def generate_token_ids(self,
                 common_attn_metadata.query_start_loc = \
                     query_start_loc_pcp_full[:num_reqs + 1]
             if self.speculative_config.disable_padded_drafter_batch:
-                # NOTE: Currently, MTP-fullgraph is incompatibility with pcp
                 token_indices_to_sample = None
                 common_attn_metadata, token_indices =\
                     self._prepare_inputs(
@@ -642,28 +656,36 @@ def _propose(
         self.input_ids[last_token_indices] = next_token_ids
 
         # update pcp related params
-        if self.pcp_size > 1:
+        if self.pcp_size * self.dcp_size > 1:
             assert long_seq_metadata is not None
             common_attn_metadata.prefill_context_parallel_metadata = long_seq_metadata
+            ori_last_token_indices = last_token_indices.clone()
+            query_lens_d = self.runner.query_lens[:num_decode_reqs]
+        if self.pcp_size > 1:
             # 1. preprocess decode/prefill input_ids & target_hidden_states
             # decode input_ids: keep unchanged
             # decode target_hidden_states: remove padding
             # prefill input_ids: add padding and pcp split
             # prefill target_hidden_states: pcp split
-            num_tokens_d = num_decode_reqs * self.decode_threshold
+            num_tokens_d = query_lens_d.sum().item()
             num_tokens_d_padded = num_tokens_d * self.pcp_size
             input_ids_d = self.input_ids[:num_tokens_d]
             input_ids_p = self.input_ids[num_tokens_d:num_tokens]
             target_hidden_states_d_padded = \
                 target_hidden_states[:num_tokens_d_padded]
             if num_tokens_d:
                 # remove padding (from pcp all-gather) in decode part
-                target_hidden_states_d = target_hidden_states_d_padded.reshape(
-                    [
-                        num_decode_reqs, self.decode_threshold * self.pcp_size,
-                        -1
-                    ])[:, :self.decode_threshold, :].reshape(
-                        [num_tokens_d, -1])
+                mask_start_loc = torch.cat([
+                    torch.tensor([0], dtype=torch.int32),
+                    torch.cumsum(query_lens_d * self.pcp_size, dim=0)[:-1]
+                ])
+                mask_len = query_lens_d
+                mask = []
+                for req_id in range(num_decode_reqs):
+                    mask += list(
+                        range(mask_start_loc[req_id],
+                              mask_start_loc[req_id] + mask_len[req_id]))
+                target_hidden_states_d = target_hidden_states_d_padded[mask]
             else:
                 target_hidden_states_d = target_hidden_states_d_padded
             target_hidden_states_p = target_hidden_states[num_tokens_d_padded:]
@@ -798,10 +820,15 @@ def _propose(
                     forward_context = get_forward_context()
                     if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL:
                         if self.vllm_config.model_config.use_mla:
-                            update_mla_attn_params(
-                                self.update_stream, forward_context,
-                                num_input_tokens,
-                                self.vllm_config.speculative_config)
+                            if self.pcp_size * self.dcp_size > 1:
+                                update_mla_attn_dcp_pcp_params(
+                                    self.update_stream, forward_context,
+                                    num_input_tokens)
+                            else:
+                                update_mla_attn_params(
+                                    self.update_stream, forward_context,
+                                    num_input_tokens,
+                                    self.vllm_config.speculative_config)
 
                     if self.enable_shared_expert_dp:
                         hidden_states = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
@@ -820,6 +847,8 @@ def _propose(
                     (0, max_num_reqs_across_dp - num_indices))
 
             if self.pcp_size > 1:
+                # remove graph padding before all_gather
+                hidden_states = hidden_states[:num_tokens]
                 hidden_states = get_pcp_group().all_gather(hidden_states, 0)
                 hidden_states = torch.index_select(
                     hidden_states, 0, self.runner.
@@ -860,6 +889,51 @@ def _propose(
                 last_token_indices = self.arange[:batch_size]
                 if getattr(attn_metadata_i, "num_decode_tokens", 0):
                     attn_metadata_i.num_decode_tokens = batch_size
+                if self.pcp_size * self.dcp_size > 1:
+                    positions = target_positions[ori_last_token_indices]
+                    # For pcp/dcp, tokens are split across different cp ranks,
+                    # so we can not simply update slot_mapping by += 1.
+                    # Instead, we pre-allocate mtp slot_mapping in model_runner
+                    # (_generate_pcp_mtp_input), and use updated slot_indices
+                    # to get corresponding slot_mapping in each step.
+                    num_reject_tokens = torch.tensor(
+                        self.runner.cu_num_tokens_pcp_full,
+                        dtype=torch.int32).to(
+                            self.device) - ori_last_token_indices - 1
+                    num_accept_tokens = \
+                        query_lens_d.to(self.device) - num_reject_tokens
+                    ori_seq_len = attn_metadata_i.seq_lens
+                    mtp_slot_mapping = self.runner.mtp_slot_pad
+
+                    # slot_mapping index base offset:
+                    # scheduled tokens + pre-allocated mtp tokens + accepted tokens
+                    slot_idx_base = (
+                        torch.cat([
+                            torch.tensor(
+                                [0], dtype=torch.int32, device=self.device),
+                            (torch.cumsum(query_lens_d, dim=0)[:-1] *
+                             self.pcp_size).to(self.device)
+                        ]) +
+                        torch.arange(num_decode_reqs, device=self.device) *
+                        (self.num_speculative_tokens - 1) * self.pcp_size +
+                        (num_accept_tokens - 1) * self.pcp_size)
+                    slot_indices_list = []
+                    for req_id in range(num_decode_reqs):
+                        slot_indices_list.append(
+                            torch.arange(slot_idx_base[req_id],
+                                         slot_idx_base[req_id] + self.pcp_size,
+                                         device=self.device))
+                    slot_indices = torch.cat(slot_indices_list, dim=0)
+
+                    # fold block_table (restore it to original size before flattened)
+                    block_indices = torch.cat([
+                        torch.tensor([0], dtype=torch.int32),
+                        torch.cumsum(query_lens_d, dim=0)[:-1]
+                    ])
+                    attn_metadata_i.decode.block_table[:batch_size] = \
+                        attn_metadata_i.decode.block_table[block_indices]
+                    attn_metadata_i.decode.block_table = \
+                        attn_metadata_i.decode.block_table[:batch_size]
 
             input_ids = draft_token_ids_list[-1].int()
             positions += 1
@@ -906,13 +980,40 @@ def _propose(
             # Otherwise, the KV cache will be inadvertently updated with the
             # padding tokens.
             slot_mapping += 1
+            if self.pcp_size > 1:
+                exceeds_max_model_len = exceeds_max_model_len.repeat_interleave(
+                    slot_mapping.size(0) // exceeds_max_model_len.size(0))
             slot_mapping.masked_fill_(exceeds_max_model_len, PADDING_SLOT_ID)
 
             # copy inputs to buffer for cudagraph
             self.input_ids[:batch_size] = input_ids
             self.positions[:batch_size] = clamped_positions
             self.hidden_states[:hidden_states.shape[0]] = hidden_states
-            attn_metadata_i.slot_mapping[:batch_size] = slot_mapping
+            if self.pcp_size * self.dcp_size > 1:
+                # update local seq_len and batch_seq_mask
+                num_computed_tokens_of_pcp_dcp = self.runner._get_cp_local_seq_lens(
+                    ori_seq_len + step + 1,
+                    self.pcp_size,
+                    self.dcp_size,
+                    self.runner.parallel_config.cp_kv_cache_interleave_size,
+                )
+                cp_seq_len = \
+                    num_computed_tokens_of_pcp_dcp[:, self.pcp_rank, self.dcp_rank]
+                batch_seq_mask = (cp_seq_len == 0)
+                builder.batch_seq_mask_buf[:batch_seq_mask.shape[0]].copy_(
+                    batch_seq_mask, non_blocking=True)
+                batch_seq_mask = builder.batch_seq_mask_buf[:batch_seq_mask.
+                                                            shape[0]]
+                cp_seq_len = torch.where(cp_seq_len == 0, 1, cp_seq_len)
+                attn_metadata_i.decode.cp_seq_len = cp_seq_len
+                attn_metadata_i.decode.batch_seq_mask = batch_seq_mask
+                # update slot_mapping
+                slot_indices += self.pcp_size
+                slot_mapping = mtp_slot_mapping[slot_indices]
+                attn_metadata_i.slot_mapping[:batch_size *
+                                             self.pcp_size] = slot_mapping
+            else:
+                attn_metadata_i.slot_mapping[:batch_size] = slot_mapping
             if self.speculative_config.disable_padded_drafter_batch:
                 self.positions[batch_size:num_input_tokens] = 0
                 self.input_ids[batch_size:num_input_tokens] = 0

vllm_ascend/worker/block_table.py

@@ -74,7 +74,7 @@ def __init__(self,
             logical_table_size = max_num_blocks_per_req
 
         duplicate_size = 1
-        if self.pcp_world_size > 1:
+        if self.pcp_world_size * self.dcp_world_size > 1:
             duplicate_size += num_speculative_tokens
         self.block_table = torch.zeros(
             (max_num_reqs * duplicate_size, logical_table_size),