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, 08:30 UTC
[Core] Automatically cast multi-modal input dtype (#18756)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
May 27, 2025 at 23:45 UTC
696259ca0180c4357cf437a334aaf0966be5cb4b
16 files changed
+91
-44
vllm/model_executor/models/deepseek_vl2.py
+1
-3
@@ -210,9 +210,7 @@ class DeepseekVL2MultiModalProcessor(
210
dict(prompt=prompt, **mm_data),
211
mm_kwargs,
212
)
213
- target_dtype = self.info.ctx.model_config.dtype
214
- pixel_values = processed_outputs.pop("pixel_values").to(
215
- target_dtype)
213
+ pixel_values = processed_outputs["pixel_values"]
214
# split pixel values into patches corresponding to each image
215
images_spatial_crop = processed_outputs["images_spatial_crop"]
216
patches_per_image = [
vllm/model_executor/models/gemma3_mm.py
-5
@@ -263,11 +263,6 @@ class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]):
263
mm_data,
264
mm_kwargs,
265
)
266
- if "pixel_values" in processed_outputs:
267
- # Cast pixel values to model dtype already here,
268
- # so we need to transfer less data to the GPU
269
- processed_outputs["pixel_values"] = processed_outputs[
270
- "pixel_values"].to(self.info.ctx.model_config.dtype)
266
267
# HF processor pops the `num_crops` kwarg, which is needed by vLLM
268
if (images := mm_data.get("images")) is not None:
vllm/multimodal/inputs.py
+7
-1
@@ -746,11 +746,17 @@ class MultiModalKwargs(UserDict[str, NestedTensors]):
746
batched_inputs: BatchedTensorInputs,
747
*,
748
device: torch.types.Device,
749
+ dtype: Optional[torch.dtype] = None,
750
) -> BatchedTensorInputs:
751
json_inputs = cast(JSONTree[torch.Tensor], batched_inputs)
752
753
+ def maybe_cast_dtype(x: torch.Tensor):
754
+ # This mimics the behavior of transformers.BatchFeature
755
+ return x.to(dtype=dtype) if x.is_floating_point() else x
756
+
757
json_mapped = json_map_leaves(
753
- lambda x: x.to(device, non_blocking=True),
758
+ # NOTE: Cast the dtype before sending it to device
759
+ lambda x: maybe_cast_dtype(x).to(device=device, non_blocking=True),
760
json_inputs,
761
)
762
vllm/spec_decode/draft_model_runner.py
+5
-2
@@ -294,8 +294,11 @@ class TP1DraftModelRunner(ModelRunnerWrapperBase):
294
inputs_embeds=None,
295
positions=model_input.input_positions,
296
intermediate_tensors=intermediate_tensors,
297
- **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
298
- device=self.device),
297
+ **MultiModalKwargs.as_kwargs(
298
+ multi_modal_kwargs,
299
+ dtype=self.model_runner.model_config.dtype,
300
+ device=self.device,
301
+ ),
302
**model_execute_kwargs,
303
)
304
vllm/v1/worker/gpu_model_runner.py
+9
-3
@@ -929,8 +929,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
929
encoder_outputs = []
930
for grouped_mm_inputs in grouped_mm_inputs_list:
931
batched_mm_inputs = MultiModalKwargs.batch(grouped_mm_inputs)
932
- batched_mm_inputs = MultiModalKwargs.as_kwargs(batched_mm_inputs,
933
- device=self.device)
932
+ batched_mm_inputs = MultiModalKwargs.as_kwargs(
933
+ batched_mm_inputs,
934
+ dtype=self.model_config.dtype,
935
+ device=self.device,
936
+ )
937
938
# Run the encoder.
939
# `curr_group_outputs` is either of the following:
@@ -1874,7 +1877,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1877
batched_dummy_mm_inputs = MultiModalKwargs.batch(
1878
[dummy_mm_kwargs] * max_num_mm_items)
1879
batched_dummy_mm_inputs = MultiModalKwargs.as_kwargs(
1877
- batched_dummy_mm_inputs, device=self.device)
1880
+ batched_dummy_mm_inputs,
1881
+ dtype=self.model_config.dtype,
1882
+ device=self.device,
1883
+ )
1884
1885
# Run multimodal encoder.
1886
dummy_encoder_outputs = self.model.get_multimodal_embeddings(
vllm/v1/worker/tpu_model_runner.py
+10
-4
@@ -652,8 +652,11 @@ class TPUModelRunner(LoRAModelRunnerMixin):
652
encoder_outputs = []
653
for grouped_mm_inputs in grouped_mm_inputs_list:
654
batched_mm_inputs = MultiModalKwargs.batch(grouped_mm_inputs)
655
- batched_mm_inputs = MultiModalKwargs.as_kwargs(batched_mm_inputs,
656
- device=self.device)
655
+ batched_mm_inputs = MultiModalKwargs.as_kwargs(
656
+ batched_mm_inputs,
657
+ dtype=self.model_config.dtype,
658
+ device=self.device,
659
+ )
660
661
# Run the encoder.
662
# `curr_group_outputs` is either of the following:
@@ -1435,8 +1438,11 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1438
1439
batched_dummy_mm_inputs = MultiModalKwargs.batch([dummy_mm_kwargs] *
1440
batch_size)
1438
- return MultiModalKwargs.as_kwargs(batched_dummy_mm_inputs,
1439
- device=self.device)
1441
+ return MultiModalKwargs.as_kwargs(
1442
+ batched_dummy_mm_inputs,
1443
+ dtype=self.model_config.dtype,
1444
+ device=self.device,
1445
+ )
1446
1447
1448
def _get_req_paddings(min_req_size: int, max_req_size: int) -> list[int]:
vllm/worker/cpu_enc_dec_model_runner.py
+5
-2
@@ -297,8 +297,11 @@ class CPUEncoderDecoderModelRunner(
297
model_input.encoder_input_tokens,
298
"encoder_positions":
299
model_input.encoder_input_positions,
300
- **MultiModalKwargs.as_kwargs(model_input.multi_modal_kwargs or {},
301
- device=self.device),
300
+ **MultiModalKwargs.as_kwargs(
301
+ model_input.multi_modal_kwargs or {},
302
+ dtype=self.model_config.dtype,
303
+ device=self.device,
304
+ ),
305
"intermediate_tensors":
306
intermediate_tensors,
307
}
vllm/worker/cpu_model_runner.py
+4
-1
@@ -628,7 +628,10 @@ class CPUModelRunner(CPUModelRunnerBase[ModelInputForCPUWithSamplingMetadata]):
628
multimodal_kwargs = {}
629
if model_input.multi_modal_kwargs is not None:
630
multimodal_kwargs = MultiModalKwargs.as_kwargs(
631
- model_input.multi_modal_kwargs, device=self.device)
631
+ model_input.multi_modal_kwargs,
632
+ dtype=self.model_config.dtype,
633
+ device=self.device,
634
+ )
635
execute_model_kwargs = {}
636
if previous_hidden_states is not None:
637
execute_model_kwargs.update(
vllm/worker/cpu_pooling_model_runner.py
+5
-2
@@ -50,8 +50,11 @@ class CPUPoolingModelRunner(
50
model_input.input_tokens,
51
"positions":
52
model_input.input_positions,
53
- **MultiModalKwargs.as_kwargs(model_input.multi_modal_kwargs or {},
54
- device=self.device),
53
+ **MultiModalKwargs.as_kwargs(
54
+ model_input.multi_modal_kwargs or {},
55
+ dtype=self.model_config.dtype,
56
+ device=self.device,
57
+ ),
58
**cross_enc_kwargs,
59
"intermediate_tensors":
60
intermediate_tensors,
vllm/worker/enc_dec_model_runner.py
+7
-3
@@ -202,9 +202,13 @@ class EncoderDecoderModelRunner(GPUModelRunnerBase[EncoderDecoderModelInput]):
202
encoder_input_ids=model_input.encoder_input_tokens,
203
encoder_positions=model_input.encoder_input_positions,
204
intermediate_tensors=intermediate_tensors,
205
- **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
206
- device=self.device),
207
- **seqlen_agnostic_kwargs)
205
+ **MultiModalKwargs.as_kwargs(
206
+ multi_modal_kwargs,
207
+ dtype=self.model_config.dtype,
208
+ device=self.device,
209
+ ),
210
+ **seqlen_agnostic_kwargs,
211
+ )
212
213
logits = self.model.compute_logits(hidden_or_intermediate_states,
214
model_input.sampling_metadata)
vllm/worker/model_runner.py
+5
-2
@@ -1845,8 +1845,11 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
1845
inputs_embeds=model_input.inputs_embeds,
1846
positions=model_input.input_positions,
1847
intermediate_tensors=intermediate_tensors,
1848
- **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
1849
- device=self.device),
1848
+ **MultiModalKwargs.as_kwargs(
1849
+ multi_modal_kwargs,
1850
+ dtype=self.model_config.dtype,
1851
+ device=self.device,
1852
+ ),
1853
**seqlen_agnostic_kwargs,
1854
**model_kwargs,
1855
)
vllm/worker/multi_step_neuron_model_runner.py
+5
-2
@@ -70,8 +70,11 @@ class MultiStepNeuronModelRunner(NeuronModelRunner):
70
input_ids=model_input.input_tokens,
71
positions=model_input.input_positions,
72
input_block_ids=model_input.input_block_ids,
73
- **MultiModalKwargs.as_kwargs(model_input.multi_modal_kwargs or {},
74
- device=self.device),
73
+ **MultiModalKwargs.as_kwargs(
74
+ model_input.multi_modal_kwargs or {},
75
+ dtype=self.model_config.dtype,
76
+ device=self.device,
77
+ ),
78
)
79
80
output = self.model.sample(
vllm/worker/multi_step_neuronx_distributed_model_runner.py
+5
-2
@@ -49,8 +49,11 @@ class MultiStepNeuronxDistributedModelRunner(NeuronxDistributedModelRunner):
49
positions=model_input.input_positions,
50
input_block_ids=model_input.input_block_ids,
51
sampling_params=sampling_params,
52
- **MultiModalKwargs.as_kwargs(model_input.multi_modal_kwargs or {},
53
- device=self.device),
52
+ **MultiModalKwargs.as_kwargs(
53
+ model_input.multi_modal_kwargs or {},
54
+ dtype=self.model_config.dtype,
55
+ device=self.device,
56
+ ),
57
)
58
59
output = self.model.sample(
vllm/worker/neuron_model_runner.py
+10
-6
@@ -378,9 +378,11 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
378
positions=model_input.input_positions,
379
input_block_ids=model_input.input_block_ids,
380
sampling_params=sampling_params,
381
- **MultiModalKwargs.as_kwargs(model_input.multi_modal_kwargs
382
- or {},
383
- device=self.device),
381
+ **MultiModalKwargs.as_kwargs(
382
+ model_input.multi_modal_kwargs or {},
383
+ dtype=self.model_config.dtype,
384
+ device=self.device,
385
+ ),
386
)
387
elif current_platform.use_transformers_neuronx():
388
# [TODO] validate on-device sampling
@@ -389,9 +391,11 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
391
input_ids=model_input.input_tokens,
392
positions=model_input.input_positions,
393
input_block_ids=model_input.input_block_ids,
392
- **MultiModalKwargs.as_kwargs(model_input.multi_modal_kwargs
393
- or {},
394
- device=self.device),
394
+ **MultiModalKwargs.as_kwargs(
395
+ model_input.multi_modal_kwargs or {},
396
+ dtype=self.model_config.dtype,
397
+ device=self.device,
398
+ ),
399
)
400
401
# Compute the logits only if the on-device sampling is turned off as
vllm/worker/pooling_model_runner.py
+7
-3
@@ -119,10 +119,14 @@ class PoolingModelRunner(
119
input_ids=model_input.input_tokens,
120
positions=model_input.input_positions,
121
intermediate_tensors=intermediate_tensors,
122
- **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
123
- device=self.device),
122
+ **MultiModalKwargs.as_kwargs(
123
+ multi_modal_kwargs,
124
+ dtype=self.model_config.dtype,
125
+ device=self.device,
126
+ ),
127
**cross_enc_kwargs,
125
- **seqlen_agnostic_kwargs)
128
+ **seqlen_agnostic_kwargs,
129
+ )
130
131
if (self.observability_config is not None
132
and self.observability_config.collect_model_forward_time):
vllm/worker/xpu_model_runner.py
+6
-3
@@ -562,9 +562,12 @@ class XPUModelRunner(ModelRunnerBase[ModelInputForXPUWithSamplingMetadata]):
562
input_ids=model_input.input_tokens,
563
positions=model_input.input_positions,
564
intermediate_tensors=intermediate_tensors,
565
- **MultiModalKwargs.as_kwargs(model_input.multi_modal_kwargs
566
- or {},
567
- device=self.device))
565
+ **MultiModalKwargs.as_kwargs(
566
+ model_input.multi_modal_kwargs or {},
567
+ dtype=self.model_config.dtype,
568
+ device=self.device,
569
+ ),
570
+ )
571
# Compute the logits in the last pipeline stage.
572
if not get_pp_group().is_last_rank:
573
return hidden_or_intermediate_states