[Fix] modelrunner support

Signed-off-by: Jingchun Gao <[email protected]>

Author: Jingchun Gao

vllm/v1/worker/gpu_model_runner.py

@@ -1001,6 +1001,24 @@ def _get_pcp_metadata(
         >>> kv_for_tail_indices r0 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, ..., 23]
         >>> r1 [0, 1, 2, 3, 4, 5, 8, 9, ..., 19]
         """
+        if len(q_lens) == 0:
+            return PrefillContextParallelMetadata(
+                allgather_restore_idx=allgather_restore_idx,
+            )
+
+        def _get_partial_kv_idx(kv_len_per_pcp_chunk):
+            kv_partial_len = pcp_chunk_sizes * kv_len_per_pcp_chunk
+            kv_partial_indptr = np.zeros(len(kv_partial_len) + 1)
+            kv_partial_indptr[1:], kv_partial_arange = self._get_cumsum_and_arange(kv_partial_len)
+            kv_parial_indices = kv_partial_arange + np.repeat(
+                kv_start_loc,
+                kv_partial_len,
+            )
+            return kv_partial_indptr, kv_parial_indices
+
+        def _to_tensor(data, **kwargs):
+            return {k: torch.from_numpy(v).to(**kwargs) for k, v in data.items()}
+
         pcp_chunk_sizes = q_lens // 2  
         q_indptr = np.zeros(len(pcp_chunk_sizes) + 1)
         q_indptr[1:], q_chunk_arange = self._get_cumsum_and_arange(pcp_chunk_sizes)
@@ -1021,41 +1039,23 @@ def _get_pcp_metadata(
         kv_start_loc = np.roll(np.cumsum(kv_lens), 1)
         kv_start_loc[0] = 0
         # kv_for_q_head
-        kv_head_len = pcp_chunk_sizes * (self.pcp_rank + 1)
-        kv_for_head_indptr = np.zeros(len(kv_head_len) + 1)
-        kv_for_head_indptr[1:], kv_nomask_head_arange = self._get_cumsum_and_arange(kv_head_len)
-        kv_for_head_indices = kv_nomask_head_arange + np.repeat(
-            kv_start_loc,
-            kv_head_len,
-        )
+        kv_for_head_indptr, kv_for_head_indices = _get_partial_kv_idx(self.pcp_rank + 1)
         # kv_for_q_tail
-        kv_tail_len = pcp_chunk_sizes * (2 * self.pcp_world_size - self.pcp_rank)
-        kv_for_tail_indptr = np.zeros(len(kv_tail_len) + 1)
-        kv_for_tail_indptr[1:], kv_nomask_tail_arange = self._get_cumsum_and_arange(kv_tail_len)
-        kv_for_tail_indices = kv_nomask_tail_arange + np.repeat(
-            kv_start_loc,
-            kv_tail_len,
+        kv_for_tail_indptr, kv_for_tail_indices = _get_partial_kv_idx(
+            2 * self.pcp_world_size - self.pcp_rank
         )
-        
-        head_tail_indices = {
+
+        head_tail_indices = _to_tensor({
             "q_head": q_head_indices,
             "q_tail": q_tail_indices,
-            "kv_head": kv_for_head_indices,
+            "kv_head": kv_for_head_indices, 
             "kv_tail": kv_for_tail_indices,
-        }
-        head_tail_indptr = {
+        }, device=self.device, dtype=torch.int64, non_blocking=True)
+        head_tail_indptr = _to_tensor({
             "q": q_indptr,
             "kv_head": kv_for_head_indptr,
             "kv_tail": kv_for_tail_indptr
-        }
-        for key, value in head_tail_indices.items():
-            head_tail_indices[key] = torch.from_numpy(value).to(
-                device=self.device, dtype=torch.int64, non_blocking=True
-            )
-        for key, value in head_tail_indptr.items():
-            head_tail_indptr[key] = torch.from_numpy(value).to(
-                dtype=torch.int64
-            )
+        }, dtype=torch.int64)
 
         q_full_indices = torch.cat([head_tail_indices["q_head"], head_tail_indices["q_tail"]])
         q_full_indices = q_full_indices.to(torch.float32).argsort().to(torch.int32)
@@ -1074,7 +1074,10 @@ def _get_pcp_metadata(
 
     def _update_tokens_for_pcp(
         self,
-        tokens: np.ndarray
+        tokens: np.ndarray,
+        dummy_input: bool = False,
+        num_reqs: int | None = None,
+        num_decode_reqs: int | None = None,
     ) -> tuple[np.ndarray, np.ndarray, PrefillContextParallelMetadata]:
         """
         If prefill context parallelism is enabled, we will update
@@ -1104,13 +1107,14 @@ def _update_tokens_for_pcp(
         >>> self.pcp_allgather_resotre_idx
         [0, 9, 1, 2, 10, 11, 12, 13, 3, 4, 5, 6, 14, 15, 16, 17, 7, 8]
         """
-        num_reqs = self.input_batch.num_reqs
+        if not dummy_input:
+            num_reqs = self.input_batch.num_reqs
+            num_decode_reqs = sum(
+                self.input_batch.num_computed_tokens_cpu[:num_reqs]
+                >= self.input_batch.num_prompt_tokens[:num_reqs]
+            )
         self.num_pcp_pads_cpu[:num_reqs] = 0
-
-        num_decode_reqs = sum(
-            self.input_batch.num_computed_tokens_cpu[:num_reqs]
-            >= self.input_batch.num_prompt_tokens[:num_reqs]
-        )
+        
         num_decode_tokens = sum(tokens[:num_decode_reqs])
 
         num_padded_scheduled_tokens = np.ceil(
@@ -1175,11 +1179,11 @@ def get_current_rank_positions(
         self.pcp_allgather_restore_idx.copy_to_gpu(all_positions.shape[0])
         return (
             pcp_tokens[:num_reqs],
-            positions, 
+            positions,
             self._get_pcp_metadata(
                 pcp_tokens[num_decode_reqs:],
                 num_padded_scheduled_tokens[num_decode_reqs:],
-                self.pcp_allgather_restore_idx.gpu[:all_positions.shape[0]]
+                self.pcp_allgather_restore_idx.gpu[: all_positions.shape[0]]
             )
         )
 
@@ -1474,11 +1478,14 @@ def _prepare_inputs(
 
         # Record the index of requests that should not be sampled,
         # so that we could clear the sampled tokens before returning
-        discard_requests_mask = (
-            self.input_batch.num_computed_tokens_cpu[:num_reqs]
-            + num_scheduled_tokens * self.pcp_world_size 
-            - self.num_pcp_pads_cpu[:num_reqs]
-        ) < num_tokens_np
+        if self.pcp_world_size > 1:
+            discard_requests_mask = (
+                self.input_batch.num_computed_tokens_cpu[:num_reqs]
+                + num_scheduled_tokens * self.pcp_world_size 
+                - self.num_pcp_pads_cpu[:num_reqs]
+            ) < num_tokens_np
+        else:
+            discard_requests_mask = self.seq_lens.np[:num_reqs] < num_tokens_np
         discard_request_indices = np.nonzero(discard_requests_mask)[0]
         self.num_discarded_requests = len(discard_request_indices)
         self.discard_request_indices.np[: self.num_discarded_requests] = (
@@ -1702,7 +1709,7 @@ def _build_attention_metadata(
                 num_logits_indices=num_logits_indices,
                 causal=True,
                 encoder_seq_lens=encoder_seq_lens,
-                cp_local_seq_lens=self.cp_local_seq_lens.gpu[:num_reqs]
+                dcp_local_seq_lens=self.cp_local_seq_lens.gpu[:num_reqs]
                 if self.total_cp_world_size > 1
                 else None,
                 pcp_metadata=pcp_metadata,
@@ -2872,7 +2879,7 @@ def execute_model(
                 use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
                 attn_metadata, spec_decode_common_attn_metadata = (
                     self._build_attention_metadata(
-                        total_num_scheduled_tokens=total_num_scheduled_tokens,
+                        total_num_scheduled_tokens=total_num_scheduled_tokens if self.pcp_world_size == 1 else num_scheduled_tokens_np.sum(),
                         max_num_scheduled_tokens=max_num_scheduled_tokens,
                         num_reqs=num_reqs,
                         ubatch_slices=ubatch_slices,
@@ -3811,6 +3818,16 @@ def _dummy_run(
         assert sum(num_scheduled_tokens_list) == num_tokens
         assert len(num_scheduled_tokens_list) == num_reqs
         num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
+        pcp_metadata = None
+        if self.pcp_world_size > 1 and force_attention:
+            num_decode_reqs = sum(num_scheduled_tokens == 1)
+            num_scheduled_tokens[:num_reqs], _, pcp_metadata = \
+                self._update_tokens_for_pcp(
+                    num_scheduled_tokens[:num_reqs],
+                    dummy_input=True,
+                    num_reqs=num_reqs,
+                    num_decode_reqs=num_decode_reqs,
+                )
         total_num_scheduled_tokens = int(num_scheduled_tokens.sum())
         num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
 
@@ -3828,7 +3845,7 @@ def _dummy_run(
             uniform_decode=uniform_decode,
             num_scheduled_tokens_per_request=num_scheduled_tokens,
         )
-        num_tokens_after_padding = num_tokens
+        num_tokens_after_padding = num_tokens if self.pcp_world_size == 1 else total_num_scheduled_tokens
         if num_tokens_across_dp is not None:
             dp_rank = self.parallel_config.data_parallel_rank
             num_tokens_after_padding = int(num_tokens_across_dp[dp_rank])
@@ -3854,11 +3871,12 @@ def _dummy_run(
             self.query_start_loc.copy_to_gpu()
 
             attn_metadata, _ = self._build_attention_metadata(
-                total_num_scheduled_tokens=num_tokens,
+                total_num_scheduled_tokens=total_num_scheduled_tokens,
                 max_num_scheduled_tokens=max_query_len,
                 num_reqs=num_reqs,
                 ubatch_slices=ubatch_slices,
                 for_cudagraph_capture=True,
+                pcp_metadata=pcp_metadata if self.pcp_world_size > 1 else None,
             )
 
         with self.maybe_dummy_run_with_lora(