[feature] support pcp + mtp in full graph

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

Author: zhangsicheng5

vllm_ascend/attention/mla_v1.py

@@ -433,17 +433,9 @@ def build(
                 common_attn_metadata.block_table_tensor[:graph_pad_size])
         else:
             block_table = (common_attn_metadata.block_table_tensor[:num_reqs])
-        # NOTE: Currently, MTP-fullgraph is incompatibility pcp
-        if self.pcp_size > 1:
-            num_decodes_flatten = num_decodes * self.decode_threshold
-            block_table = common_attn_metadata.block_table_tensor[:
-                                                                  num_decodes_flatten
-                                                                  +
-                                                                  num_prefills]
         if num_actual_tokens_pcp_padded is None:
             num_actual_tokens_pcp_padded = num_actual_tokens
 
-        # NOTE: Currently, MTP-fullgraph is incompatibility pcp
         slot_mapping = common_attn_metadata.slot_mapping[:
                                                          num_actual_tokens_pcp_padded]
         input_positions = common_attn_metadata.positions[:
@@ -466,6 +458,13 @@ def build(
         seq_lens = common_attn_metadata.seq_lens_cpu[:num_reqs]
         num_computed_tokens_cpu = (seq_lens - query_lens)
 
+        if self.pcp_size * self.dcp_size > 1:
+            num_decodes_flatten = query_lens[:num_decodes].sum().item()
+            block_table = common_attn_metadata.block_table_tensor[:
+                                                                  num_decodes_flatten
+                                                                  +
+                                                                  num_prefills]
+
         prefill_metadata = None
         chunked_context_metadata = None
         if num_prefills > 0:
@@ -530,8 +529,9 @@ def build(
                 if self.dcp_size * self.pcp_size > 1:
                     if num_computed_tokens_of_pcp_dcp is not None:
                         local_context_lens_allranks = torch.tensor(
-                            num_computed_tokens_of_pcp_dcp[reqs_start:num_reqs]
-                        ).reshape(-1, self.dcp_size * self.pcp_size)
+                            num_computed_tokens_of_pcp_dcp[
+                                num_decodes_flatten:]).reshape(
+                                    -1, self.dcp_size * self.pcp_size)
                     # Note(qcs): The max local context lengths
                     # padded to `cp_local_block_size`.
                     padded_local_context_lens_cpu = (cdiv(
@@ -617,7 +617,7 @@ def build(
                 cos=cos,
                 pcp_metadata=pcp_metadata,
             )
-            if self.pcp_size > 1:
+            if self.pcp_size * self.dcp_size > 1:
                 prefill_metadata.block_table = block_table[
                     num_decodes_flatten:, ...]
 
@@ -630,13 +630,12 @@ def build(
             max_seq_lens = seq_lens[:num_decodes].max().item()
             seq_lens = seq_lens[:num_decodes]
             input_positions = input_positions[:num_decode_tokens]
-            if self.pcp_size > 1:
+            if self.pcp_size * self.dcp_size > 1:
                 # For pcp + spec decode, we flatten seq_lens and block_table
                 # to avoid irregular spec_attn_mask shape
                 block_table = block_table[:num_decodes_flatten, ...]
             else:
                 block_table = block_table[:num_decodes, ...]
-            # NOTE: Currently, MTP-fullgraph is incompatibility pcp
             # NOTE: Maybe this block_table change can be removed when graph_pad_size > 1.
             if graph_pad_size > num_decodes and \
                     self.speculative_config.disable_padded_drafter_batch:
@@ -646,8 +645,7 @@ def build(
             if num_computed_tokens_of_pcp_dcp is not None:
                 # [bs, pcp_size, dcp_size]
                 num_computed_tokens_of_cp_dcp_array = np.array(
-                    num_computed_tokens_of_pcp_dcp)[:num_decodes *
-                                                    self.decode_threshold]
+                    num_computed_tokens_of_pcp_dcp)[:num_decodes_flatten]
 
                 cp_seq_len = num_computed_tokens_of_cp_dcp_array[:,
                                                                  self.pcp_rank,

vllm_ascend/spec_decode/mtp_proposer.py

@@ -255,6 +255,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(
@@ -344,7 +351,7 @@ def generate_token_ids(self,
                 )
 
         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
@@ -381,7 +388,6 @@ def generate_token_ids(self,
                     query_start_loc_pcp_full[:num_reqs + 1]
             if self.speculative_config.disable_padded_drafter_batch:
                 assert isinstance(sampled_token_ids, list)
-                # NOTE: Currently, MTP-fullgraph is incompatibility with pcp
                 token_indices_to_sample = None
                 common_attn_metadata, token_indices =\
                     self._prepare_inputs(
@@ -592,28 +598,35 @@ 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
+        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
+            query_lens_d = self.runner.query_lens[:num_decode_reqs]
+            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:]
@@ -749,6 +762,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.

vllm_ascend/worker/block_table.py

@@ -80,7 +80,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),

vllm_ascend/worker/model_runner_v1.py

@@ -496,7 +496,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
             dtype=torch.int32,
             device=self.device)
         self.num_actual_tokens_pcp_padded = 0
-        if self.speculative_config and self.pcp_size > 1:
+        if self.speculative_config and self.pcp_size * self.dcp_size > 1:
             self.input_ids_pcp_full = torch.zeros(self.max_num_tokens,
                                                   dtype=torch.int32,
                                                   device=self.device)
@@ -1738,7 +1738,7 @@ def _prepare_inputs(
             self.num_accepted_tokens.np[num_reqs:].fill(1)
             self.num_accepted_tokens.copy_to_gpu()
 
-        if self.speculative_config and self.pcp_size > 1:
+        if self.speculative_config and self.pcp_size * self.dcp_size > 1:
             self._generate_pcp_mtp_input(
                 num_reqs, scheduler_output.total_num_scheduled_tokens,
                 scheduler_output.num_scheduled_tokens)
@@ -1820,28 +1820,29 @@ def _prepare_inputs(
                 prefill_context_parallel_metadata=long_seq_metadata,
             )
 
-            if self.speculative_config and self.pcp_size > 1:
+            if self.speculative_config and self.pcp_size * self.dcp_size > 1:
                 # For pcp + spec decode, we flatten block_table
                 # to avoid irregular spec_attn_mask shape, e.g.,
                 # num_decode_req=2, num_prefill_req=3, num_speculative_tokens=1,
                 # ori block_table: # [d0, d1, p0, p1, p2]
                 # (num_reqs_d + num_reqs_p, max_num_blocks),
                 # flattened block_table: [d0, d0, d1, d1, p0, p1, p2]
                 # (num_reqs_d * decode_threshold + num_reqs_p, max_num_blocks),
-                ori_query_lens = self.query_start_loc_pcp_full_cpu[1:num_reqs+1] - \
-                    self.query_start_loc_pcp_full_cpu[:num_reqs]
+                ori_query_lens = self.query_start_loc_pcp_full[1:num_reqs+1] - \
+                    self.query_start_loc_pcp_full[:num_reqs]
                 num_prefill_reqs = (ori_query_lens
                                     > self.decode_threshold).sum().item()
                 num_decode_reqs = num_reqs - num_prefill_reqs
-                num_decode_reqs_flatten = num_decode_reqs * self.decode_threshold
+                num_decode_reqs_flatten = \
+                    ori_query_lens[:num_decode_reqs].sum().item()
                 blk_table_tensor[
                     num_decode_reqs_flatten:num_decode_reqs_flatten +
                     num_prefill_reqs].copy_(
                         blk_table_tensor[num_decode_reqs:num_decode_reqs +
                                          num_prefill_reqs].clone())
                 blk_table_tensor[:num_decode_reqs_flatten].copy_(
                     blk_table_tensor[:num_decode_reqs].repeat_interleave(
-                        self.decode_threshold, dim=0))
+                        ori_query_lens[:num_decode_reqs], dim=0))
                 common_attn_metadata.block_table_tensor = \
                     blk_table_tensor[:num_decode_reqs_flatten + num_prefill_reqs]
 
@@ -2788,7 +2789,7 @@ def _build_dummy_attn_metadata(
                     sin=self.sin,
                     prefill_context_parallel_metadata=long_seq_metadata,
                 )
-                if self.pcp_size > 1:
+                if self.pcp_size * self.dcp_size > 1:
                     common_attn_metadata.block_table_tensor = \
                         block_table_tensor[:num_reqs * self.decode_threshold]
                 attn_state = AscendAttentionState.DecodeOnly
@@ -4250,8 +4251,8 @@ def _get_cp_local_seq_lens(
     def _generate_pcp_metadata(self, total_num_scheduled_tokens):
         # In dummy run num_reqs == 0, update it from seq_lens
         num_reqs = self.input_batch.num_reqs or self.query_lens.size(0)
-        num_decodes = sum(self.input_batch.num_computed_tokens_cpu[:num_reqs]
-                          >= self.input_batch.num_prompt_tokens[:num_reqs])
+        num_decodes = (self.query_lens <= self.decode_threshold).sum().item()
+        num_prefills = num_reqs - num_decodes
         num_actual_tokens_pcp_padded = total_num_scheduled_tokens * self.pcp_size
         self.num_actual_tokens_pcp_padded = num_actual_tokens_pcp_padded
         long_seq_metadata = None
@@ -4269,16 +4270,41 @@ def _generate_pcp_metadata(self, total_num_scheduled_tokens):
                 dtype=torch.int32,
             )
             # For pcp + spec decode, we flatten seq_lens
-            # to avoid irregular spec_attn_mask shape
+            # to avoid irregular spec_attn_mask shape.
+            # Same as block_table, we flatten decode seq_lens to query_lens,
+            # and keep prefill seq_lens unchanged.
             for decode_idx in range(self.decode_threshold):
                 num_computed_tokens_of_pcp_dcp[
                     self.decode_threshold - 1 - decode_idx::self.decode_threshold] = \
                     self._get_cp_local_seq_lens(
-                        torch.tensor(context_lens),
+                        torch.tensor(context_lens) - decode_idx,
                         self.pcp_size,
                         self.dcp_size,
                         self.parallel_config.cp_kv_cache_interleave_size,
                     )
+            if self.decode_threshold > 1:
+                num_computed_tokens_of_pcp_dcp_list = []
+                if num_decodes:
+                    num_decodes_flatten = \
+                        self.query_lens[:num_decodes].sum().item()
+                    if self.query_lens[:num_decodes].min().item(
+                    ) == self.decode_threshold:
+                        decode_flatten_idx = list(range(num_decodes_flatten))
+                    else:
+                        decode_flatten_idx = []
+                        for req_id in range(num_decodes):
+                            offset = (req_id + 1) * self.decode_threshold
+                            decode_flatten_idx += \
+                                list(range(offset - self.query_lens[req_id], offset))
+                    num_computed_tokens_of_pcp_dcp_list.append(
+                        num_computed_tokens_of_pcp_dcp[decode_flatten_idx])
+                if num_prefills:
+                    num_computed_tokens_of_pcp_dcp_list.append(
+                        num_computed_tokens_of_pcp_dcp[
+                            (num_decodes + 1) * self.decode_threshold -
+                            1::self.decode_threshold])
+                num_computed_tokens_of_pcp_dcp = torch.cat(
+                    num_computed_tokens_of_pcp_dcp_list, dim=0)
             long_seq_metadata = AscendPrefillContextParallelMetadata(
                 num_actual_tokens_pcp_padded=num_actual_tokens_pcp_padded,
                 num_computed_tokens_of_pcp_dcp=num_computed_tokens_of_pcp_dcp.