[Model Runner V2] Avoid prepare prefill kernel launch overhead (#34780)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

vllm/v1/worker/gpu/model_runner.py

@@ -614,16 +614,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
         max_query_len = num_scheduled_tokens.max().item()
 
-        # Get prefill tokens.
-        prepare_prefill_inputs(
-            self.input_buffers.input_ids,
-            self.req_states.next_prefill_tokens,
-            idx_mapping,
-            query_start_loc,
-            self.req_states.all_token_ids.gpu,
-            self.req_states.prefill_len.gpu,
-            self.req_states.num_computed_tokens.gpu,
-        )
+        # Get prefill tokens if any.
+        if self.req_states.any_prefills(idx_mapping_np):
+            prepare_prefill_inputs(
+                self.input_buffers.input_ids,
+                self.req_states.next_prefill_tokens,
+                idx_mapping,
+                query_start_loc,
+                self.req_states.all_token_ids.gpu,
+                self.req_states.prefill_len.gpu,
+                self.req_states.num_computed_tokens.gpu,
+            )
 
         # Prepare positions and seq_lens.
         prepare_pos_seq_lens(

vllm/v1/worker/gpu/states.py

@@ -60,10 +60,7 @@ class RequestState:
 
         # Last sampled tokens.
         self.last_sampled_tokens = torch.zeros(
-            self.max_num_reqs,
-            1,
-            dtype=torch.int64,
-            device=device,
+            self.max_num_reqs, 1, dtype=torch.int64, device=device
         )
 
         # Draft tokens.
@@ -118,3 +115,9 @@ class RequestState:
             return
         self.index_to_req_id.pop(req_idx, None)
         self.free_indices.append(req_idx)
+
+    def any_prefills(self, idx_mapping_np: np.ndarray) -> bool:
+        return np.any(
+            self.num_computed_prefill_tokens[idx_mapping_np]
+            < self.prefill_len.np[idx_mapping_np]
+        )