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, 07:15 UTC
[Core] Cast multimodal input in hf processor (#18862)
Signed-off-by: Lukas Geiger <lukas.geiger94@gmail.com>
Lukas Geiger committed
Jun 4, 2025 at 04:24 UTC
1409ef913446aa282f6426efbb0ed02a59320467
15 files changed
+25
-25
vllm/inputs/registry.py
+24
-2
@@ -4,9 +4,12 @@ from collections.abc import Mapping
4
from dataclasses import dataclass
5
from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union
6
7
+import torch
8
from transformers import BatchFeature, PretrainedConfig, ProcessorMixin
9
from typing_extensions import TypeVar
10
11
+from vllm.jsontree import JSONTree, json_map_leaves
12
+from vllm.logger import init_logger
13
from vllm.transformers_utils.processor import cached_processor_from_config
14
from vllm.transformers_utils.tokenizer import AnyTokenizer
15
from vllm.utils import resolve_mm_processor_kwargs
@@ -21,6 +24,8 @@ _T = TypeVar("_T")
24
_C = TypeVar("_C", bound=PretrainedConfig, default=PretrainedConfig)
25
_P = TypeVar("_P", bound=ProcessorMixin, default=ProcessorMixin)
26
27
+logger = init_logger(__name__)
28
+
29
30
@dataclass(frozen=True)
31
class InputContext:
@@ -134,7 +139,7 @@ class InputProcessingContext(InputContext):
139
hf_processor: ProcessorMixin,
140
data: Mapping[str, object],
141
kwargs: Mapping[str, object] = {},
137
- ) -> BatchFeature:
142
+ ) -> Union[BatchFeature, JSONTree]:
143
"""
144
Call `hf_processor` on the prompt `data`
145
(text, image, audio...) with configurable options `kwargs`.
@@ -154,8 +159,25 @@ class InputProcessingContext(InputContext):
159
allow_var_kwargs=True,
160
)
161
162
+ def maybe_cast_dtype(x):
163
+ # This mimics the behavior of transformers.BatchFeature
164
+ if isinstance(x, torch.Tensor) and x.is_floating_point():
165
+ return x.to(dtype=self.model_config.dtype)
166
+ return x
167
+
168
try:
158
- return hf_processor(**data, **merged_kwargs, return_tensors="pt")
169
+ output = hf_processor(**data, **merged_kwargs, return_tensors="pt")
170
+ # this emulates output.to(dtype=self.model_config.dtype)
171
+ cast_output = json_map_leaves(maybe_cast_dtype, output)
172
+ if isinstance(output, BatchFeature):
173
+ return BatchFeature(cast_output)
174
+
175
+ logger.warning_once(
176
+ f"{type(hf_processor).__name__} did not return `BatchFeature`. "
177
+ "Make sure to match the behaviour of `ProcessorMixin` when "
178
+ "implementing custom processors.")
179
+ return cast_output
180
+
181
except Exception as exc:
182
msg = (f"Failed to apply {type(hf_processor).__name__} "
183
f"on data={data} with kwargs={merged_kwargs}")
vllm/multimodal/inputs.py
+1
-7
@@ -747,17 +747,11 @@ class MultiModalKwargs(UserDict[str, NestedTensors]):
747
batched_inputs: BatchedTensorInputs,
748
*,
749
device: torch.types.Device,
750
- dtype: Optional[torch.dtype] = None,
750
) -> BatchedTensorInputs:
751
json_inputs = cast(JSONTree[torch.Tensor], batched_inputs)
752
754
- def maybe_cast_dtype(x: torch.Tensor):
755
- # This mimics the behavior of transformers.BatchFeature
756
- return x.to(dtype=dtype) if x.is_floating_point() else x
757
-
753
json_mapped = json_map_leaves(
759
- # NOTE: Cast the dtype before sending it to device
760
- lambda x: maybe_cast_dtype(x).to(device=device, non_blocking=True),
754
+ lambda x: x.to(device=device, non_blocking=True),
755
json_inputs,
756
)
757
vllm/spec_decode/draft_model_runner.py
-1
@@ -297,7 +297,6 @@ class TP1DraftModelRunner(ModelRunnerWrapperBase):
297
intermediate_tensors=intermediate_tensors,
298
**MultiModalKwargs.as_kwargs(
299
multi_modal_kwargs,
300
- dtype=self.model_runner.model_config.dtype,
300
device=self.device,
301
),
302
**model_execute_kwargs,
vllm/v1/worker/gpu_model_runner.py
-2
@@ -957,7 +957,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
957
batched_mm_inputs = MultiModalKwargs.batch(grouped_mm_inputs)
958
batched_mm_inputs = MultiModalKwargs.as_kwargs(
959
batched_mm_inputs,
960
- dtype=self.model_config.dtype,
960
device=self.device,
961
)
962
@@ -1951,7 +1950,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1950
[dummy_mm_kwargs] * max_num_mm_items)
1951
batched_dummy_mm_inputs = MultiModalKwargs.as_kwargs(
1952
batched_dummy_mm_inputs,
1954
- dtype=self.model_config.dtype,
1953
device=self.device,
1954
)
1955
vllm/v1/worker/tpu_model_runner.py
-2
@@ -718,7 +718,6 @@ class TPUModelRunner(LoRAModelRunnerMixin):
718
batched_mm_inputs = MultiModalKwargs.batch(grouped_mm_inputs)
719
batched_mm_inputs = MultiModalKwargs.as_kwargs(
720
batched_mm_inputs,
721
- dtype=self.model_config.dtype,
721
device=self.device,
722
)
723
@@ -1560,7 +1559,6 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1559
batch_size)
1560
return MultiModalKwargs.as_kwargs(
1561
batched_dummy_mm_inputs,
1563
- dtype=self.model_config.dtype,
1562
device=self.device,
1563
)
1564
vllm/worker/cpu_enc_dec_model_runner.py
-1
@@ -300,7 +300,6 @@ class CPUEncoderDecoderModelRunner(
300
model_input.encoder_input_positions,
301
**MultiModalKwargs.as_kwargs(
302
model_input.multi_modal_kwargs or {},
303
- dtype=self.model_config.dtype,
303
device=self.device,
304
),
305
"intermediate_tensors":
vllm/worker/cpu_model_runner.py
-1
@@ -630,7 +630,6 @@ class CPUModelRunner(CPUModelRunnerBase[ModelInputForCPUWithSamplingMetadata]):
630
if model_input.multi_modal_kwargs is not None:
631
multimodal_kwargs = MultiModalKwargs.as_kwargs(
632
model_input.multi_modal_kwargs,
633
- dtype=self.model_config.dtype,
633
device=self.device,
634
)
635
execute_model_kwargs = {}
vllm/worker/cpu_pooling_model_runner.py
-1
@@ -53,7 +53,6 @@ class CPUPoolingModelRunner(
53
model_input.input_positions,
54
**MultiModalKwargs.as_kwargs(
55
model_input.multi_modal_kwargs or {},
56
- dtype=self.model_config.dtype,
56
device=self.device,
57
),
58
**cross_enc_kwargs,
vllm/worker/enc_dec_model_runner.py
-1
@@ -205,7 +205,6 @@ class EncoderDecoderModelRunner(GPUModelRunnerBase[EncoderDecoderModelInput]):
205
intermediate_tensors=intermediate_tensors,
206
**MultiModalKwargs.as_kwargs(
207
multi_modal_kwargs,
208
- dtype=self.model_config.dtype,
208
device=self.device,
209
),
210
**seqlen_agnostic_kwargs,
vllm/worker/model_runner.py
-1
@@ -1848,7 +1848,6 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
1848
intermediate_tensors=intermediate_tensors,
1849
**MultiModalKwargs.as_kwargs(
1850
multi_modal_kwargs,
1851
- dtype=self.model_config.dtype,
1851
device=self.device,
1852
),
1853
**seqlen_agnostic_kwargs,
vllm/worker/multi_step_neuron_model_runner.py
-1
@@ -73,7 +73,6 @@ class MultiStepNeuronModelRunner(NeuronModelRunner):
73
input_block_ids=model_input.input_block_ids,
74
**MultiModalKwargs.as_kwargs(
75
model_input.multi_modal_kwargs or {},
76
- dtype=self.model_config.dtype,
76
device=self.device,
77
),
78
)
vllm/worker/multi_step_neuronx_distributed_model_runner.py
-1
@@ -52,7 +52,6 @@ class MultiStepNeuronxDistributedModelRunner(NeuronxDistributedModelRunner):
52
sampling_params=sampling_params,
53
**MultiModalKwargs.as_kwargs(
54
model_input.multi_modal_kwargs or {},
55
- dtype=self.model_config.dtype,
55
device=self.device,
56
),
57
)
vllm/worker/neuron_model_runner.py
-2
@@ -395,7 +395,6 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
395
adapter_ids=model_input.adapter_ids,
396
**MultiModalKwargs.as_kwargs(
397
model_input.multi_modal_kwargs or {},
398
- dtype=self.model_config.dtype,
398
device=self.device,
399
),
400
)
@@ -408,7 +407,6 @@ class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
407
input_block_ids=model_input.input_block_ids,
408
**MultiModalKwargs.as_kwargs(
409
model_input.multi_modal_kwargs or {},
411
- dtype=self.model_config.dtype,
410
device=self.device,
411
),
412
)
vllm/worker/pooling_model_runner.py
-1
@@ -122,7 +122,6 @@ class PoolingModelRunner(
122
intermediate_tensors=intermediate_tensors,
123
**MultiModalKwargs.as_kwargs(
124
multi_modal_kwargs,
125
- dtype=self.model_config.dtype,
125
device=self.device,
126
),
127
**cross_enc_kwargs,
vllm/worker/xpu_model_runner.py
-1
@@ -565,7 +565,6 @@ class XPUModelRunner(ModelRunnerBase[ModelInputForXPUWithSamplingMetadata]):
565
intermediate_tensors=intermediate_tensors,
566
**MultiModalKwargs.as_kwargs(
567
model_input.multi_modal_kwargs or {},
568
- dtype=self.model_config.dtype,
568
device=self.device,
569
),
570
)