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
[Bugfix] should use stack instead of concat (#22972)
Signed-off-by: 947132885 <947132885@qq.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
947132885 committed
Aug 17, 2025 at 16:46 UTC
fe0411fc6fa32cebeacd3a3aef87a591e7309c45
1 file changed
+12
-2
vllm/model_executor/models/transformers.py
+12
-2
@@ -694,6 +694,17 @@ class TransformersForCausalLM(TransformersBase):
694
return logits
695
696
697
+def flatten_and_concat(x: list[torch.Tensor]) -> torch.Tensor:
698
+ """Flatten until a list of tensors can be concatenated then do concat"""
699
+
700
+ def _can_concat(x: list[torch.Tensor]):
701
+ return len(set(map(lambda _x: _x.shape[1:], x))) == 1
702
+
703
+ if _can_concat(x):
704
+ return torch.concat(x)
705
+ return flatten_and_concat(flatten_bn(x))
706
+
707
+
708
@MULTIMODAL_REGISTRY.register_processor(
709
MultiModalProcessor,
710
info=MultiModalProcessingInfo,
@@ -766,8 +777,7 @@ class TransformersForMultimodalLM(TransformersForCausalLM, SupportsMultiModal):
777
if isinstance(pixel_values, torch.Tensor):
778
pixel_values = flatten_bn(pixel_values).to(self.dtype)
779
elif is_list_of(pixel_values, torch.Tensor):
769
- pixel_values = flatten_bn(flatten_bn(pixel_values),
770
- concat=True).to(self.dtype)
780
+ pixel_values = flatten_and_concat(pixel_values).to(self.dtype)
781
else:
782
raise ValueError(
783
f"Unsupported pixel_values type {type(pixel_values)}. "