xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-10 01:55:01 +08:00
Code Issues Packages Projects Releases Wiki Activity

perf: Avoid copying inputs_embeds tensors to GPU unless prompt_embeds is enabled (#25739)

Browse Source 
Signed-off-by: Andrew Sansom <andrew@protopia.ai>
This commit is contained in:
Andrew Sansom 2025-09-26 03:18:09 -05:00 committed by GitHub
parent e84e0735c7
commit d48f4d6daf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 15 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/v1/worker/gpu_model_runner.py
View File

@ -836,6 +836,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
if self.input_batch.prev_sampled_token_ids is None:
# Normal scheduling case
self.input_ids.copy_to_gpu(total_num_scheduled_tokens)
if self.enable_prompt_embeds:
self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens)
self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens)
return
@ -863,6 +864,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
# If not all requests are decodes from the last iteration,
# We need to copy the input_ids_cpu to the GPU first.
self.input_ids.copy_to_gpu(total_num_scheduled_tokens)
if self.enable_prompt_embeds:
self.inputs_embeds.copy_to_gpu(total_num_scheduled_tokens)
self.is_token_ids.copy_to_gpu(total_num_scheduled_tokens)
if num_commmon_tokens == 0:
@ -878,6 +880,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
self.input_batch.prev_sampled_token_ids[:num_commmon_tokens,
0],
non_blocking=True)
if self.enable_prompt_embeds:
self.is_token_ids.gpu[:num_commmon_tokens] = True
return
# Upload the index tensors asynchronously
@ -978,6 +981,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
0,
token_indices_tensor,
out=self.input_ids.cpu[:total_num_scheduled_tokens])
if self.enable_prompt_embeds:
is_token_ids = self.input_batch.is_token_ids.flatten()
torch.index_select(
is_token_ids,