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 15, 18:30 UTC
[V1] Optimize the overhead of rewinding (#14905)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Mar 16, 2025 at 20:19 UTC
faa02757307583f2c5557ff23cb41f1db4f1f29c
1 file changed
+5
-6
vllm/v1/worker/gpu_model_runner.py
+5
-6
@@ -1032,17 +1032,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1032
1033
# TODO(woosuk): The following loop can be slow since it iterates over
1034
# the requests one by one. Optimize.
1035
- for i, req_id in enumerate(self.input_batch.req_ids):
1035
+ for i, generator in self.input_batch.generators.items():
1036
+ req_id = self.input_batch.req_ids[i]
1037
req_state = self.requests[req_id]
1038
seq_len = (req_state.num_computed_tokens +
1039
scheduler_output.num_scheduled_tokens[req_id])
1040
if seq_len < req_state.num_tokens:
1040
- # Ignore the sampled token.
1041
+ # Ignore the sampled token for partial prefills.
1042
# Rewind the generator state as if the token was not sampled.
1042
- generator = self.input_batch.generators.get(i)
1043
- if generator is not None:
1044
- # This relies on cuda-specific torch-internal impl details
1045
- generator.set_offset(generator.get_offset() - 4)
1043
+ # This relies on cuda-specific torch-internal impl details
1044
+ generator.set_offset(generator.get_offset() - 4)
1045
1046
# NOTE: GPU -> CPU Sync happens here.
1047
# Move as many CPU operations as possible before this sync point.