diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index e8f7e051b..e9f9d868f 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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( diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py index b4bc8d4d4..b338d32a3 100644 --- a/vllm/v1/worker/gpu/states.py +++ b/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] + )