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
Migrate Phi4 inputs to TensorSchema (#23471)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Aug 31, 2025 at 23:05 UTC
437c3ce02615443ab166f4155028c1d81ee27c06
2 files changed
+129
-73
vllm/model_executor/models/phi4_multimodal.py
+72
-39
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import math
4
from collections.abc import Iterable, Mapping, Sequence
5
-from typing import Any, Literal, Optional, TypedDict, Union
5
+from typing import Annotated, Any, Literal, Optional, Union
6
7
import numpy as np
8
import torch
@@ -40,6 +40,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
40
from vllm.multimodal.profiling import BaseDummyInputsBuilder
41
from vllm.sequence import IntermediateTensors
42
from vllm.utils import is_list_of
43
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
44
45
from .idefics2_vision_model import Idefics2VisionTransformer
46
from .interfaces import MultiModalEmbeddings, SupportsLoRA, SupportsMultiModal
@@ -615,50 +616,90 @@ class Phi4MMAudioEmbedding(nn.Module):
616
return loaded_params
617
618
618
-class Phi4MMImagePixelInputs(TypedDict):
619
- type: Literal["pixel_values"]
620
- data: Union[torch.Tensor, list[torch.Tensor]]
619
+class Phi4MMImagePixelInputs(TensorSchema):
620
"""
622
- Shape:
623
- `(batch_size * num_images, 1 + num_patches, num_channels, height, width)`
624
-
625
- Note that `num_patches` may be different per batch and image,
626
- in which case the data is passed as a list instead of a batched tensor.
621
+ Dimensions:
622
+ - bn: Batch size * number of images
623
+ - p: Number of patches (1 + num_patches)
624
+ - c: Number of channels (3)
625
+ - h: Height of each image patch
626
+ - w: Width of each image patch
627
+ - nc: Number of crops
628
+ - H_mask: Height of attention mask
629
+ - W_mask: Width of attention mask
630
"""
631
629
- image_sizes: torch.Tensor
630
- """
631
- Shape: `(batch_size * num_images, 2)`
632
+ type: Literal["pixel_values"]
633
633
- This should be in `(height, width)` format.
634
- """
634
+ data: Annotated[
635
+ Union[torch.Tensor, list[torch.Tensor]],
636
+ TensorShape("bn", "p", 3, "h", "w", dynamic_dims={"p"}
637
+ ), # may be different per batch and image
638
+ ]
639
636
- num_img_tokens: list[int]
637
- """Shape: `(batch_size * num_images)`"""
640
+ image_sizes: Annotated[
641
+ torch.Tensor,
642
+ TensorShape("bn", 2), # (height, width)
643
+ ]
644
639
- image_attention_mask: torch.Tensor
640
- """Shape: `(batch_size * num_images, H_mask, W_mask)`"""
645
+ num_img_tokens: Annotated[
646
+ list[int],
647
+ TensorShape("bn"),
648
+ ]
649
650
+ image_attention_mask: Annotated[
651
+ torch.Tensor,
652
+ TensorShape("bn", "nc", 32, 32), # H_mask, W_mask
653
+ ]
654
643
-class Phi4MMImageEmbeddingInputs(TypedDict):
644
- type: Literal["image_embeds"]
645
- data: Union[torch.Tensor, list[torch.Tensor]]
646
- """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`
655
648
- `hidden_size` must match the hidden size of language model backbone.
656
+class Phi4MMImageEmbeddingInputs(TensorSchema):
657
"""
658
+ Dimensions:
659
+ - bn: Batch size * number of images
660
+ - f: Image feature size
661
+ - h: Hidden size (must match language model backbone)
662
+ """
663
+
664
+ type: Literal["image_embeds"]
665
+
666
+ data: Annotated[
667
+ Union[torch.Tensor, list[torch.Tensor]],
668
+ TensorShape("bn", "f", "h"),
669
+ ]
670
+
671
672
+class Phi4MMAudioFeatureInputs(TensorSchema):
673
+ """
674
+ Dimensions:
675
+ - bn: Batch size * number of audios
676
+ - f: Number of Mel filterbank bins (80)
677
+ - t: Time frames (M)
678
+ """
679
652
-class Phi4MMAudioFeatureInputs(TypedDict):
680
type: Literal["audio_features"]
654
- data: Union[torch.Tensor, list[torch.Tensor]]
655
- """Shape: `(batch_size * num_audios, 80, M)"""
681
682
+ data: Annotated[
683
+ Union[torch.Tensor, list[torch.Tensor]],
684
+ TensorShape("bn", "t", 80, dynamic_dims={"t"}),
685
+ ]
686
+
687
+
688
+class Phi4MMAudioEmbeddingInputs(TensorSchema):
689
+ """
690
+ Dimensions:
691
+ - b: Batch size
692
+ - n: Number of audios
693
+ - f: Audio feature size
694
+ - h: Hidden size (must match language model backbone)
695
+ """
696
658
-class Phi4MMAudioEmbeddingInputs(TypedDict):
697
type: Literal["audio_embeds"]
660
- data: NestedTensors
661
- """Shape: `(batch_size, num_audios, audio_feature_size, hidden_size)"""
698
+
699
+ data: Annotated[
700
+ NestedTensors,
701
+ TensorShape("b", "n", "f", "h"),
702
+ ]
703
704
705
Phi4MMImageInput = Union[Phi4MMImagePixelInputs, Phi4MMImageEmbeddingInputs]
@@ -1170,18 +1211,10 @@ class Phi4MultimodalForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
1211
return None
1212
1213
if audio_features is not None:
1173
- if not isinstance(audio_features, (torch.Tensor, list)):
1174
- raise ValueError("Incorrect type of audio features. "
1175
- f"Got type: {type(audio_features)}")
1176
-
1214
return Phi4MMAudioFeatureInputs(type="audio_features",
1215
data=flatten_bn(audio_features))
1216
1217
if audio_embeds is not None:
1181
- if not isinstance(audio_embeds, (torch.Tensor, list)):
1182
- raise ValueError("Incorrect type of audio embeds. "
1183
- f"Got type: {type(audio_embeds)}")
1184
-
1218
return Phi4MMAudioEmbeddingInputs(type="audio_embeds",
1219
data=audio_embeds)
1220
@@ -1259,7 +1292,7 @@ class Phi4MultimodalForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
1292
elif isinstance(image_sizes, torch.Tensor):
1293
image_sizes = image_sizes.flatten(0, 1)
1294
else:
1262
- raise ValueError("Incorrect image_attention_mask inputs")
1295
+ raise ValueError("Incorrect image_sizes inputs")
1296
1297
if isinstance(num_img_tokens, list):
1298
num_img_tokens = [
@@ -1269,7 +1302,7 @@ class Phi4MultimodalForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
1302
elif isinstance(num_img_tokens, torch.Tensor):
1303
num_img_tokens = num_img_tokens.flatten(0, 1).tolist()
1304
else:
1272
- raise ValueError("Incorrect image_attention_mask inputs")
1305
+ raise ValueError("Incorrect num_img_tokens inputs")
1306
1307
return Phi4MMImagePixelInputs(
1308
type="pixel_values",
vllm/model_executor/models/phi4mm.py
+57
-34
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import math
4
from collections.abc import Iterable, Mapping, Sequence
5
-from typing import Any, Literal, Optional, TypedDict, Union
5
+from typing import Annotated, Any, Literal, Optional, Union
6
7
import numpy as np
8
import torch
@@ -31,6 +31,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
31
from vllm.multimodal.profiling import BaseDummyInputsBuilder
32
from vllm.sequence import IntermediateTensors
33
from vllm.utils import is_list_of
34
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
35
36
from .idefics2_vision_model import Idefics2VisionTransformer
37
from .interfaces import MultiModalEmbeddings, SupportsLoRA, SupportsMultiModal
@@ -391,41 +392,71 @@ class Phi4MMImageEncoder(nn.Module):
392
return img_set_tensor
393
394
394
-class Phi4MMImagePixelInputs(TypedDict):
395
- type: Literal["pixel_values"]
396
- data: Union[torch.Tensor, list[torch.Tensor]]
395
+class Phi4MMImagePixelInputs(TensorSchema):
396
"""
398
- Shape:
399
- `(batch_size * num_images, 1 + num_patches, num_channels, height, width)`
400
-
401
- Note that `num_patches` may be different per batch and image,
402
- in which case the data is passed as a list instead of a batched tensor.
397
+ Dimensions:
398
+ - bn: Batch size * number of images
399
+ - p: Number of patches (1 + num_patches)
400
+ - c: Number of channels (3)
401
+ - h: Height of each image patch
402
+ - w: Width of each image patch
403
+ - nc: Number of crops
404
+ - H_mask: Height of attention mask
405
+ - W_mask: Width of attention mask
406
"""
407
405
- image_sizes: torch.Tensor
406
- """
407
- Shape: `(batch_size * num_images, 2)`
408
+ type: Literal["pixel_values"]
409
409
- This should be in `(height, width)` format.
410
- """
410
+ data: Annotated[
411
+ Union[torch.Tensor, list[torch.Tensor]],
412
+ TensorShape("bn", "p", 3, "h", "w", dynamic_dims={"p"}
413
+ ), # may be different per batch and image
414
+ ]
415
+
416
+ image_sizes: Annotated[
417
+ torch.Tensor,
418
+ TensorShape("bn", 2), # (height, width)
419
+ ]
420
412
- num_img_tokens: list[int]
413
- """Shape: `(batch_size * num_images)`"""
421
+ num_img_tokens: Annotated[
422
+ list[int],
423
+ TensorShape("bn"),
424
+ ]
425
415
- image_attention_mask: torch.Tensor
416
- """Shape: `(batch_size * num_images, H_mask, W_mask)`"""
426
+ image_attention_mask: Annotated[
427
+ torch.Tensor,
428
+ TensorShape("bn", "nc", 32, 32), # H_mask, W_mask
429
+ ]
430
431
419
-class Phi4MMAudioFeatureInputs(TypedDict):
432
+class Phi4MMAudioFeatureInputs(TensorSchema):
433
+ """
434
+ Dimensions:
435
+ - bn: Batch size * number of audios
436
+ - t: Time frames (M)
437
+ """
438
+
439
type: Literal["audio_features"]
421
- data: Union[torch.Tensor, list[torch.Tensor]]
422
- """Shape: `(batch_size * num_audios, 80, M)"""
440
+
441
+ data: Annotated[
442
+ Union[torch.Tensor, list[torch.Tensor]],
443
+ TensorShape("bn", "t", 80, dynamic_dims={"t"}),
444
+ ]
445
446
425
-class Phi4MMAudioEmbeddingInputs(TypedDict):
447
+class Phi4MMAudioEmbeddingInputs(TensorSchema):
448
+ """
449
+ Dimensions:
450
+ - b: Batch size
451
+ - n: Number of audios
452
+ - f: Audio feature size
453
+ - h: Hidden size (must match language model backbone)
454
+ """
455
type: Literal["audio_embeds"]
427
- data: NestedTensors
428
- """Shape: `(batch_size, num_audios, audio_feature_size, hidden_size)"""
456
+ data: Annotated[
457
+ NestedTensors,
458
+ TensorShape("b", "n", "f", "h"),
459
+ ]
460
461
462
Phi4MMAudioInputs = Union[Phi4MMAudioFeatureInputs, Phi4MMAudioEmbeddingInputs]
@@ -985,18 +1016,10 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
1016
return None
1017
1018
if audio_features is not None:
988
- if not isinstance(audio_features, (torch.Tensor, list)):
989
- raise ValueError("Incorrect type of audio features. "
990
- f"Got type: {type(audio_features)}")
991
-
1019
return Phi4MMAudioFeatureInputs(type="audio_features",
1020
data=flatten_bn(audio_features))
1021
1022
if audio_embeds is not None:
996
- if not isinstance(audio_embeds, (torch.Tensor, list)):
997
- raise ValueError("Incorrect type of audio embeds. "
998
- f"Got type: {type(audio_embeds)}")
999
-
1023
return Phi4MMAudioEmbeddingInputs(type="audio_embeds",
1024
data=audio_embeds)
1025
@@ -1074,7 +1097,7 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
1097
elif isinstance(image_sizes, torch.Tensor):
1098
image_sizes = image_sizes.flatten(0, 1)
1099
else:
1077
- raise ValueError("Incorrect image_attention_mask inputs")
1100
+ raise ValueError("Incorrect image_sizes inputs")
1101
1102
if isinstance(num_img_tokens, list):
1103
num_img_tokens = [
@@ -1084,7 +1107,7 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
1107
elif isinstance(num_img_tokens, torch.Tensor):
1108
num_img_tokens = num_img_tokens.flatten(0, 1).tolist()
1109
else:
1087
- raise ValueError("Incorrect image_attention_mask inputs")
1110
+ raise ValueError("Incorrect num_img_tokens inputs")
1111
1112
return Phi4MMImagePixelInputs(
1113
type="pixel_values",