Fix max_num_reqs error

Signed-off-by: anon189Ty <Stari_Falcon@outlook.com>

Author: anon189Ty

vllm_ascend/compilation/acl_graph.py

@@ -259,6 +259,9 @@ def update_mla_attn_params(update_stream, forward_context, runtime_shape,
                 spec_multiple * (i + 1)
                 for i in range(runtime_shape // spec_multiple)
             ]
+        elif forward_context.is_mtp_model:
+            seq_lens_list = seq_lens_list + [0] * (len(actual_seq_lengths) -
+                                                   len(seq_lens_list))
         else:
             seq_lens_list = seq_lens_list + [0] * (runtime_shape -
                                                    len(seq_lens_list))

vllm_ascend/spec_decode/mtp_proposer.py

@@ -75,19 +75,22 @@ def __init__(
         self.use_sparse = hasattr(vllm_config.model_config.hf_config,
                                   "index_topk")
 
-        self.actual_seq_lengths_q = list(
-            range(1, self.runner.max_num_tokens + 1, 1))
-        self.query_start_loc = torch.zeros(self.runner.max_num_reqs + 1,
-                                           dtype=torch.int32,
-                                           device=self.device)
-        self.query_start_loc_cpu = torch.zeros(self.runner.max_num_reqs + 1,
-                                               dtype=torch.int32,
-                                               device="cpu",
-                                               pin_memory=True)
+        # self.actual_seq_lengths_q = list(
+        #     range(1, self.runner.max_num_tokens + 1, 1))
+        self.query_start_loc = torch.zeros(
+            self.runner.max_num_reqs * (self.num_speculative_tokens + 1) + 1,
+            dtype=torch.int32,
+            device=self.device)
+        self.query_start_loc_cpu = torch.zeros(
+            self.runner.max_num_reqs * (self.num_speculative_tokens + 1) + 1,
+            dtype=torch.int32,
+            device="cpu",
+            pin_memory=True)
         self.slot_mapping = torch.zeros(self.runner.max_num_tokens,
                                         dtype=torch.int32,
                                         device=self.device)
-        self.seq_lens_cpu = torch.zeros(self.runner.max_num_reqs,
+        self.seq_lens_cpu = torch.zeros(self.runner.max_num_reqs *
+                                        (self.num_speculative_tokens + 1),
                                         dtype=torch.int32,
                                         device="cpu",
                                         pin_memory=True)
@@ -175,7 +178,6 @@ def dummy_run(self,
         elif aclgraph_runtime_mode == CUDAGraphMode.FULL:
             assert with_prefill is False, \
                 "Full decode graph only supports uniform batch now."
-            num_reqs = num_tokens
             max_seq_lens = self.runner.model_config.max_model_len
             self.seq_lens_cpu[:num_reqs] = max_seq_lens
             self.seq_lens_cpu[num_reqs:] = 0
@@ -184,7 +186,7 @@ def dummy_run(self,
                     self.runner.input_batch.
                     num_computed_tokens_cpu_tensor[:num_reqs])
                 query_start_loc = torch.tensor(
-                    [0] + self.actual_seq_lengths_q[:num_reqs],
+                    [0] + self.runner.actual_seq_lengths_q[:num_reqs],
                     device=self.runner.device,
                     dtype=torch.int32)
                 self.query_start_loc[:num_reqs + 1].copy_(query_start_loc)
@@ -207,7 +209,7 @@ def dummy_run(self,
                     spec_attn_mask=self.runner.spec_attn_mask,
                     attn_state=self.runner.attn_state,
                     decode_token_per_req=self.runner.decode_token_per_req,
-                    cos=self.runner.cos,  # 考虑mrope，是否可以共用？
+                    cos=self.runner.cos,
                     sin=self.runner.sin,
                 )
 
@@ -350,7 +352,8 @@ def generate_token_ids(self,
             block_table=attn_metadata.block_tables,
             sampling_metadata=sampling_metadata,
             token_indices=accepted_token_indices,
-            scheduler_output=scheduler_output)
+            scheduler_output=scheduler_output,
+            num_scheduled_tokens=num_scheduled_tokens)
         spec_token_ids = draft_token_ids.tolist()
         return spec_token_ids
 
@@ -416,12 +419,16 @@ def _prepare_inputs(
         batch_size = num_rejected_tokens.shape[0]
         self.query_start_loc[:batch_size + 1].copy_(cu_num_tokens[:batch_size +
                                                                   1])
+        self.query_start_loc[batch_size + 1:].fill_(0)
         self.query_start_loc_cpu[:batch_size + 1].copy_(
             self.query_start_loc[:batch_size + 1], non_blocking=True)
+        self.query_start_loc_cpu[batch_size + 1:].fill_(0)
         target_positions_len = target_positions.shape[0]
         self.positions[:target_positions_len].copy_(target_positions)
+        self.positions[target_positions_len:].fill_(0)
         target_slot_mapping_len = target_slot_mapping.shape[0]
         self.slot_mapping[:target_slot_mapping_len].copy_(target_slot_mapping)
+        self.slot_mapping[target_slot_mapping_len:].fill_(0)
 
         return cu_num_tokens, token_indices, target_token_ids, target_positions, target_hidden_states, target_slot_mapping
 
@@ -443,7 +450,8 @@ def _propose(
             block_table: torch.Tensor,
             sampling_metadata: SamplingMetadata,
             token_indices=None,
-            scheduler_output: SchedulerOutput = None) -> torch.Tensor:
+            scheduler_output: SchedulerOutput = None,
+            num_scheduled_tokens: int = 0) -> torch.Tensor:
         num_tokens = target_token_ids.shape[0]
         batch_size = next_token_ids.shape[0]
         last_token_indices = cu_num_tokens[1:] - 1
@@ -489,6 +497,30 @@ def _propose(
         seq_lens = seq_lens.int()
         seq_lens_len = seq_lens.shape[0]
         self.seq_lens_cpu[:seq_lens_len].copy_(seq_lens, non_blocking=True)
+        self.seq_lens_cpu[seq_lens_len:].fill_(0)
+
+        if self.torchair_graph_enabled:
+            # torchair mode can reuse self.runner.num_tokens_across_dp
+            num_tokens_across_dp = self.runner.num_tokens_across_dp
+            with_prefill = self.runner.with_prefill
+        elif self.vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs(
+        ):
+            (num_input_tokens, num_tokens_across_dp, with_prefill,
+             _) = self.runner._sync_metadata_across_dp(
+                 num_scheduled_tokens, self.runner.with_prefill, False)
+        else:
+            # torch mode need to update num_tokens_across_dp
+            # TODO: adapt enable_dbo later
+            (num_input_tokens, num_tokens_across_dp, with_prefill,
+             _) = self.runner._sync_metadata_across_dp(
+                 num_input_tokens, self.runner.with_prefill, False)
+
+        self.vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs(
+        ):
+            graph_pad_size = num_input_tokens
+        else:
+            graph_pad_size = self.runner.graph_pad_size
+
         common_attn_metadata = AscendCommonAttentionMetadata(
             query_start_loc=self.query_start_loc[:batch_size + 1],
             query_start_loc_cpu=self.query_start_loc_cpu[:batch_size + 1],
@@ -504,7 +536,7 @@ def _propose(
             attn_mask=self.runner.attn_mask,
             spec_attn_mask=self.runner.spec_attn_mask,
             attn_state=self.runner.attn_state,
-            graph_pad_size=self.runner.graph_pad_size,
+            graph_pad_size=graph_pad_size,
             decode_token_per_req=self.runner.decode_token_per_req,
             num_computed_tokens_cpu=None,
             seq_lens=None)
@@ -522,20 +554,8 @@ def _propose(
             attn_metadata = self.runner.attn_metadata_builder.build(
                 0, common_attn_metadata, self.runner.get_model())
 
-        self.positions[:num_tokens] = target_positions
         self.hidden_states[:num_tokens] = target_hidden_states
 
-        if not self.torchair_graph_enabled:
-            # torch mode need to update num_tokens_across_dp
-            # TODO: adapt enable_dbo later
-            (num_input_tokens, num_tokens_across_dp, with_prefill,
-             _) = self.runner._sync_metadata_across_dp(
-                 num_input_tokens, self.runner.with_prefill, False)
-        else:
-            # torchair mode can reuse self.runner.num_tokens_across_dp
-            num_tokens_across_dp = self.runner.num_tokens_across_dp
-            with_prefill = self.runner.with_prefill
-
         moe_comm_type = self.runner._select_moe_comm_method(
             num_input_tokens, with_prefill)