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, 14:45 UTC
[Perf] Optimize Preparing Inputs for GPU Model Runner (#16484)
Signed-off-by: snowcharm <snowcharmqq@gmail.com> Co-authored-by: Nick Hill <nhill@redhat.com>
SnowCharm committed
Apr 12, 2025 at 22:54 UTC
93e5f3c5fb4a4bbd49610efb96aad30df95fca66
1 file changed
+4
-8
vllm/v1/worker/gpu_model_runner.py
+4
-8
@@ -484,14 +484,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
484
self.input_batch.block_table.commit(num_reqs)
485
486
# Get the number of scheduled tokens for each request.
487
- # TODO: The Python loop can be slow. Optimize.
488
- num_scheduled_tokens = np.empty(num_reqs, dtype=np.int32)
489
- max_num_scheduled_tokens = 0
490
- for i, req_id in enumerate(self.input_batch.req_ids):
491
- num_tokens = scheduler_output.num_scheduled_tokens[req_id]
492
- num_scheduled_tokens[i] = num_tokens
493
- max_num_scheduled_tokens = max(max_num_scheduled_tokens,
494
- num_tokens)
487
+ req_ids = self.input_batch.req_ids
488
+ tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
489
+ num_scheduled_tokens = np.array(tokens, dtype=np.int32)
490
+ max_num_scheduled_tokens = max(tokens)
491
492
# Get request indices.
493
# E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]