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] Avoid accessing req_ids inside a loop (#23159)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Aug 19, 2025 at 02:39 UTC
21bcc8263f06cf84deaf3a5d7b11b479468eaf27
1 file changed
+6
-3
vllm/v1/worker/gpu_model_runner.py
+6
-3
@@ -1748,6 +1748,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1748
# NOTE(woosuk): As an exception, when using PP, the scheduler sends
1749
# the sampled tokens back, because there's no direct communication
1750
# between the first-stage worker and the last-stage worker.
1751
+ req_ids = self.input_batch.req_ids
1752
for req_idx, sampled_ids in enumerate(valid_sampled_token_ids):
1753
if not sampled_ids:
1754
continue
@@ -1763,7 +1764,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1764
start_idx:end_idx] = sampled_ids
1765
self.input_batch.num_tokens_no_spec[req_idx] = end_idx
1766
self.input_batch.num_tokens[req_idx] = end_idx
1766
- req_id = self.input_batch.req_ids[req_idx]
1767
+ req_id = req_ids[req_idx]
1768
req_state = self.requests[req_id]
1769
req_state.output_token_ids.extend(sampled_ids)
1770
@@ -1843,6 +1844,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1844
elif self.speculative_config.use_eagle():
1845
assert isinstance(self.drafter, EagleProposer)
1846
# TODO(woosuk): Refactor the loop.
1847
+ req_ids = self.input_batch.req_ids
1848
next_token_ids: list[int] = []
1849
for i, token_ids in enumerate(sampled_token_ids):
1850
if token_ids:
@@ -1851,7 +1853,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1853
else:
1854
# Partial prefill (rare case).
1855
# Get the next token id from the request state.
1854
- req_id = self.input_batch.req_ids[i]
1856
+ req_id = req_ids[i]
1857
req_state = self.requests[req_id]
1858
seq_len = (req_state.num_computed_tokens +
1859
scheduler_output.num_scheduled_tokens[req_id])
@@ -1914,6 +1916,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1916
sampled_token_ids: list[list[int]],
1917
) -> list[list[int]]:
1918
# TODO(woosuk): Optimize.
1919
+ req_ids = self.input_batch.req_ids
1920
draft_token_ids: list[list[int]] = []
1921
for i, sampled_ids in enumerate(sampled_token_ids):
1922
num_sampled_ids = len(sampled_ids)
@@ -1924,7 +1927,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1927
1928
# Skip requests that require sampling parameters that are not
1929
# supported with speculative decoding.
1927
- req_id = self.input_batch.req_ids[i]
1930
+ req_id = req_ids[i]
1931
if req_id in self.input_batch.spec_decode_unsupported_reqs:
1932
draft_token_ids.append([])
1933
continue