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, 06:00 UTC
[Core] Batch multi modal input using pinned memory (#19169)
Signed-off-by: Lukas Geiger <lukas.geiger94@gmail.com>
Lukas Geiger committed
Jun 10, 2025 at 07:44 UTC
319cb1e351736890e1615ab827639d7fd05fc988
2 files changed
+18
-7
vllm/multimodal/inputs.py
+14
-5
@@ -680,7 +680,8 @@ class MultiModalKwargs(UserDict[str, NestedTensors]):
680
return self._items_by_modality.keys()
681
682
@staticmethod
683
- def _try_stack(nested_tensors: NestedTensors) -> NestedTensors:
683
+ def _try_stack(nested_tensors: NestedTensors,
684
+ pin_memory: bool = False) -> NestedTensors:
685
"""
686
Stack the inner dimensions that have the same shape in
687
a nested list of tensors.
@@ -697,7 +698,9 @@ class MultiModalKwargs(UserDict[str, NestedTensors]):
698
if isinstance(nested_tensors, (int, float)):
699
return torch.tensor(nested_tensors)
700
700
- stacked = [MultiModalKwargs._try_stack(t) for t in nested_tensors]
701
+ stacked = [
702
+ MultiModalKwargs._try_stack(t, pin_memory) for t in nested_tensors
703
+ ]
704
if not is_list_of(stacked, torch.Tensor, check="all"):
705
# Only tensors (not lists) can be stacked.
706
return stacked
@@ -713,10 +716,16 @@ class MultiModalKwargs(UserDict[str, NestedTensors]):
716
# The tensors have incompatible shapes and can't be stacked.
717
return tensors_
718
716
- return torch.stack(tensors_)
719
+ outputs = torch.empty(len(tensors_),
720
+ *tensors_[0].shape,
721
+ dtype=tensors_[0].dtype,
722
+ device=tensors_[0].device,
723
+ pin_memory=pin_memory)
724
+ return torch.stack(tensors_, out=outputs)
725
726
@staticmethod
719
- def batch(inputs_list: list["MultiModalKwargs"]) -> BatchedTensorInputs:
727
+ def batch(inputs_list: list["MultiModalKwargs"],
728
+ pin_memory: bool = False) -> BatchedTensorInputs:
729
"""
730
Batch multiple inputs together into a dictionary.
731
@@ -738,7 +747,7 @@ class MultiModalKwargs(UserDict[str, NestedTensors]):
747
item_lists[k].append(v)
748
749
return {
741
- k: MultiModalKwargs._try_stack(item_list)
750
+ k: MultiModalKwargs._try_stack(item_list, pin_memory)
751
for k, item_list in item_lists.items()
752
}
753
vllm/v1/worker/gpu_model_runner.py
+4
-2
@@ -962,7 +962,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
962
963
encoder_outputs = []
964
for grouped_mm_inputs in grouped_mm_inputs_list:
965
- batched_mm_inputs = MultiModalKwargs.batch(grouped_mm_inputs)
965
+ batched_mm_inputs = MultiModalKwargs.batch(
966
+ grouped_mm_inputs, pin_memory=self.pin_memory)
967
batched_mm_inputs = MultiModalKwargs.as_kwargs(
968
batched_mm_inputs,
969
device=self.device,
@@ -1989,7 +1990,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1990
).multi_modal_data
1991
1992
batched_dummy_mm_inputs = MultiModalKwargs.batch(
1992
- [dummy_mm_kwargs] * max_num_mm_items)
1993
+ [dummy_mm_kwargs] * max_num_mm_items,
1994
+ pin_memory=self.pin_memory)
1995
batched_dummy_mm_inputs = MultiModalKwargs.as_kwargs(
1996
batched_dummy_mm_inputs,
1997
device=self.device,