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 14, 05:30 UTC
[Misc] Minor code cleanup for _get_prompt_logprobs_dict (#23064)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Aug 17, 2025 at 18:16 UTC
8ea0c2753a273e24957ab4587c200a3254ebe970
1 file changed
+3
-4
vllm/v1/worker/gpu_model_runner.py
+3
-4
@@ -1722,7 +1722,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1722
# Compute prompt logprobs if needed.
1723
prompt_logprobs_dict = self._get_prompt_logprobs_dict(
1724
hidden_states[:num_scheduled_tokens],
1725
- scheduler_output,
1725
+ scheduler_output.num_scheduled_tokens,
1726
)
1727
1728
# Get the valid generated tokens.
@@ -2064,7 +2064,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2064
def _get_prompt_logprobs_dict(
2065
self,
2066
hidden_states: torch.Tensor,
2067
- scheduler_output: "SchedulerOutput",
2067
+ num_scheduled_tokens: dict[str, int],
2068
) -> dict[str, Optional[LogprobsTensors]]:
2069
num_prompt_logprobs_dict = self.input_batch.num_prompt_logprobs
2070
if not num_prompt_logprobs_dict:
@@ -2077,8 +2077,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2077
# maintainable loop over optimal performance.
2078
completed_prefill_reqs = []
2079
for req_id, num_prompt_logprobs in num_prompt_logprobs_dict.items():
2080
-
2081
- num_tokens = scheduler_output.num_scheduled_tokens[req_id]
2080
+ num_tokens = num_scheduled_tokens[req_id]
2081
2082
# Get metadata for this request.
2083
request = self.requests[req_id]