xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-06-09 15:55:41 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Misc] Minor code cleanup for _get_prompt_logprobs_dict (#23064)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
This commit is contained in:
Woosuk Kwon 2025-08-17 18:16:03 -07:00 committed by GitHub
parent 0fc8fa751a
commit 8ea0c2753a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
vllm/v1/worker/gpu_model_runner.py
View File

@ -1722,7 +1722,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
# Compute prompt logprobs if needed. # Compute prompt logprobs if needed.
prompt_logprobs_dict = self._get_prompt_logprobs_dict( prompt_logprobs_dict = self._get_prompt_logprobs_dict(
hidden_states[:num_scheduled_tokens], hidden_states[:num_scheduled_tokens],
scheduler_output, scheduler_output.num_scheduled_tokens,
) )
# Get the valid generated tokens. # Get the valid generated tokens.
@ -2064,7 +2064,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
def _get_prompt_logprobs_dict( def _get_prompt_logprobs_dict(
self, self,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
scheduler_output: "SchedulerOutput", num_scheduled_tokens: dict[str, int],
) -> dict[str, Optional[LogprobsTensors]]: ) -> dict[str, Optional[LogprobsTensors]]:
num_prompt_logprobs_dict = self.input_batch.num_prompt_logprobs num_prompt_logprobs_dict = self.input_batch.num_prompt_logprobs
if not num_prompt_logprobs_dict: if not num_prompt_logprobs_dict:
@ -2077,8 +2077,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
# maintainable loop over optimal performance. # maintainable loop over optimal performance.
completed_prefill_reqs = [] completed_prefill_reqs = []
for req_id, num_prompt_logprobs in num_prompt_logprobs_dict.items(): for req_id, num_prompt_logprobs in num_prompt_logprobs_dict.items():
num_tokens = num_scheduled_tokens[req_id]
num_tokens = scheduler_output.num_scheduled_tokens[req_id]
# Get metadata for this request. # Get metadata for this request.
request = self.requests[req_id] request = self.requests[req_id]