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
perf: Avoid copying inputs_embeds tensors to GPU unless prompt_embeds is enabled (#25739)
Signed-off-by: Andrew Sansom <andrew@protopia.ai>
Andrew Sansom committed
Sep 26, 2025 at 03:18 UTC
d48f4d6daf7cf62568f84cc50b3ba1fd18ad8c82
1 file changed
+15
-11
vllm/v1/worker/gpu_model_runner.py
+15
-11
@@ -836,8 +836,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
836
if self.input_batch.prev_sampled_token_ids is None:
837
# Normal scheduling case
838
self.input_ids.copy_to_gpu(total_num_scheduled_tokens)
839
- self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens)
840
- self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens)
839
+ if self.enable_prompt_embeds:
840
+ self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens)
841
+ self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens)
842
return
843
844
# Async scheduling case, where some decode requests from the previous
@@ -863,8 +864,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
864
# If not all requests are decodes from the last iteration,
865
# We need to copy the input_ids_cpu to the GPU first.
866
self.input_ids.copy_to_gpu(total_num_scheduled_tokens)
866
- self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens)
867
- self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens)
867
+ if self.enable_prompt_embeds:
868
+ self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens)
869
+ self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens)
870
if num_commmon_tokens == 0:
871
# No requests in common with the previous iteration
872
# So input_ids_cpu will have all the input ids.
@@ -878,7 +880,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
880
self.input_batch.prev_sampled_token_ids[:num_commmon_tokens,
881
0],
882
non_blocking=True)
881
- self.is_token_ids.gpu[:num_commmon_tokens] = True
883
+ if self.enable_prompt_embeds:
884
+ self.is_token_ids.gpu[:num_commmon_tokens] = True
885
return
886
# Upload the index tensors asynchronously
887
# so the scatter can be non-blocking.
@@ -978,12 +981,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
981
0,
982
token_indices_tensor,
983
out=self.input_ids.cpu[:total_num_scheduled_tokens])
981
- is_token_ids = self.input_batch.is_token_ids.flatten()
982
- torch.index_select(
983
- is_token_ids,
984
- 0,
985
- token_indices_tensor,
986
- out=self.is_token_ids.cpu[:total_num_scheduled_tokens])
984
+ if self.enable_prompt_embeds:
985
+ is_token_ids = self.input_batch.is_token_ids.flatten()
986
+ torch.index_select(
987
+ is_token_ids,
988
+ 0,
989
+ token_indices_tensor,
990
+ out=self.is_token_ids.cpu[:total_num_scheduled_tokens])
991
992
# Because we did not pre-allocate a massive prompt_embeds CPU tensor on
993
# the InputBatch, we need to fill in the prompt embeds into the expected