A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 03:45 UTC
[Model Runner V2] Avoid prepare prefill kernel launch overhead (#34780)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Feb 18, 2026 at 00:49 UTC
c50e105a8843907c8c89f95ee29b8cc5e3935bae
2 files changed
+18
-14
vllm/v1/worker/gpu/model_runner.py
+11
-10
@@ -614,16 +614,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
614
query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
615
max_query_len = num_scheduled_tokens.max().item()
616
617
- # Get prefill tokens.
618
- prepare_prefill_inputs(
619
- self.input_buffers.input_ids,
620
- self.req_states.next_prefill_tokens,
621
- idx_mapping,
622
- query_start_loc,
623
- self.req_states.all_token_ids.gpu,
624
- self.req_states.prefill_len.gpu,
625
- self.req_states.num_computed_tokens.gpu,
626
- )
617
+ # Get prefill tokens if any.
618
+ if self.req_states.any_prefills(idx_mapping_np):
619
+ prepare_prefill_inputs(
620
+ self.input_buffers.input_ids,
621
+ self.req_states.next_prefill_tokens,
622
+ idx_mapping,
623
+ query_start_loc,
624
+ self.req_states.all_token_ids.gpu,
625
+ self.req_states.prefill_len.gpu,
626
+ self.req_states.num_computed_tokens.gpu,
627
+ )
628
629
# Prepare positions and seq_lens.
630
prepare_pos_seq_lens(
vllm/v1/worker/gpu/states.py
+7
-4
@@ -60,10 +60,7 @@ class RequestState:
60
61
# Last sampled tokens.
62
self.last_sampled_tokens = torch.zeros(
63
- self.max_num_reqs,
64
- 1,
65
- dtype=torch.int64,
66
- device=device,
63
+ self.max_num_reqs, 1, dtype=torch.int64, device=device
64
)
65
66
# Draft tokens.
@@ -118,3 +115,9 @@ class RequestState:
115
return
116
self.index_to_req_id.pop(req_idx, None)
117
self.free_indices.append(req_idx)
118
+
119
+ def any_prefills(self, idx_mapping_np: np.ndarray) -> bool:
120
+ return np.any(
121
+ self.num_computed_prefill_tokens[idx_mapping_np]
122
+ < self.prefill_len.np[idx_mapping_np]
123
+ )