[LoRA][Refactor] Unify multimodal LoRA token count hooks (#55071)
Signed-off-by: linitra24 <renshuang.zhou@daocloud.io>
linitra24 committed
Sep 14, 2026 at 21:19 UTC
9d4d9aa5bced3fc84c9e7c89b1257f206fbff6e2
23 files changed
+279
-298
tests/models/multimodal/processing/test_llava_next_video.py
+22
-15
@@ -14,7 +14,7 @@ from ...utils import build_model_context
14
15
16
class _StubModel:
17
- """Carries only the two attributes the token helpers read from
17
+ """Carries only the attributes the token-count method reads from
18
`self`, so the real methods can be exercised without constructing
19
the full `nn.Module` (vision tower, language model, etc.)."""
20
@@ -22,18 +22,15 @@ class _StubModel:
22
pooled_grid_length: int
23
24
25
-get_num_mm_encoder_tokens = (
26
- LlavaNextVideoForConditionalGeneration.get_num_mm_encoder_tokens
27
-)
28
-get_num_mm_connector_tokens = (
29
- LlavaNextVideoForConditionalGeneration.get_num_mm_connector_tokens
25
+get_mm_lora_token_counts = (
26
+ LlavaNextVideoForConditionalGeneration.get_mm_lora_token_counts
27
)
28
29
30
@pytest.mark.parametrize("model_id", ["llava-hf/LLaVA-NeXT-Video-7B-hf"])
31
def test_num_mm_tokens_match_real_config(model_id):
32
"""The stored grid lengths must match what `__init__` derives from
36
- the real HF config, and the two helpers must invert each other's
33
+ the real HF config, and the LoRA token counts must preserve the
34
frame-level scaling exactly."""
35
ctx = build_model_context(model_id, limit_mm_per_prompt={"video": 1})
36
config = ctx.model_config.hf_config
@@ -49,10 +46,13 @@ def test_num_mm_tokens_match_real_config(model_id):
46
for num_frames in (1, 2, 8, 16, 32):
47
num_video_tokens = num_frames * pooled_grid_length**2
48
52
- encoder_tokens = get_num_mm_encoder_tokens(stub, num_video_tokens)
49
+ encoder_tokens, connector_tokens = get_mm_lora_token_counts(
50
+ stub,
51
+ modality="video",
52
+ mm_kwargs=None,
53
+ num_mm_embeds=num_video_tokens,
54
+ )
55
assert encoder_tokens == num_frames * patch_grid_length**2
54
-
55
- connector_tokens = get_num_mm_connector_tokens(stub, encoder_tokens)
56
assert connector_tokens == num_video_tokens
57
58
@@ -72,10 +72,13 @@ def test_num_mm_tokens_roundtrip(patch_grid_length, pooled_grid_length, num_fram
72
73
num_video_tokens = num_frames * pooled_grid_length**2
74
75
- encoder_tokens = get_num_mm_encoder_tokens(stub, num_video_tokens)
75
+ encoder_tokens, connector_tokens = get_mm_lora_token_counts(
76
+ stub,
77
+ modality="video",
78
+ mm_kwargs=None,
79
+ num_mm_embeds=num_video_tokens,
80
+ )
81
assert encoder_tokens == num_frames * patch_grid_length**2
77
-
78
- connector_tokens = get_num_mm_connector_tokens(stub, encoder_tokens)
82
assert connector_tokens == num_video_tokens
83
84
@@ -84,5 +87,9 @@ def test_num_mm_tokens_zero():
87
stub.patch_grid_length = 24
88
stub.pooled_grid_length = 12
89
87
- assert get_num_mm_encoder_tokens(stub, 0) == 0
88
- assert get_num_mm_connector_tokens(stub, 0) == 0
90
+ assert get_mm_lora_token_counts(
91
+ stub,
92
+ modality="video",
93
+ mm_kwargs=None,
94
+ num_mm_embeds=0,
95
+ ) == (0, 0)
vllm/model_executor/models/blip2.py
+15
-20
@@ -22,6 +22,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
22
from vllm.multimodal import MULTIMODAL_REGISTRY
23
from vllm.multimodal.inputs import (
24
MultiModalFieldConfig,
25
+ MultiModalKwargsItem,
26
MultiModalKwargsItems,
27
)
28
from vllm.multimodal.parse import MultiModalDataItems
@@ -687,28 +688,22 @@ class Blip2ForConditionalGeneration(
688
tower_model="vision_model",
689
)
690
690
- def get_num_mm_encoder_tokens(
691
+ def get_mm_lora_token_counts(
692
self,
692
- num_image_tokens: int,
693
- ) -> int:
694
- if num_image_tokens <= 0:
695
- return 0
696
- assert num_image_tokens % self.config.num_query_tokens == 0, (
693
+ *,
694
+ modality: str,
695
+ mm_kwargs: MultiModalKwargsItem | None,
696
+ num_mm_embeds: int,
697
+ ) -> tuple[int, int | None]:
698
+ del modality, mm_kwargs
699
+ if num_mm_embeds <= 0:
700
+ return 0, 0
701
+ assert num_mm_embeds % self.config.num_query_tokens == 0, (
702
"The number of image tokens must be a multiple of "
703
"the number of query tokens."
704
)
700
- num_images = num_image_tokens / self.config.num_query_tokens
701
- return num_images * self._vision_tokens_per_image
702
-
703
- def get_num_mm_connector_tokens(
704
- self,
705
- num_vision_tokens: int,
706
- ) -> int:
707
- if num_vision_tokens <= 0:
708
- return 0
709
- assert num_vision_tokens % self._vision_tokens_per_image == 0, (
710
- "The number of vision tokens must be a multiple of "
711
- "the number of tokens per image."
705
+ num_images = num_mm_embeds // self.config.num_query_tokens
706
+ return (
707
+ num_images * self._vision_tokens_per_image,
708
+ num_images * self.config.num_query_tokens,
709
)
713
- num_images = num_vision_tokens / self._vision_tokens_per_image
714
- return num_images * self.config.num_query_tokens
vllm/model_executor/models/dots_ocr.py
+10
-6
@@ -54,6 +54,7 @@ from vllm.model_executor.models.utils import (
54
)
55
from vllm.model_executor.models.vision import get_vit_attn_backend
56
from vllm.multimodal import MULTIMODAL_REGISTRY
57
+from vllm.multimodal.inputs import MultiModalKwargsItem
58
from vllm.sequence import IntermediateTensors
59
from vllm.transformers_utils.configs.dotsocr import DotsOCRConfig, DotsVisionConfig
60
from vllm.utils.tensor_schema import TensorSchema, TensorShape
@@ -718,13 +719,16 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA
719
720
return image_embeds.split(sizes)
721
721
- def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
722
- merge_size = self.vision_tower.spatial_merge_size
723
- return num_image_tokens * (merge_size**2)
724
-
725
- def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
722
+ def get_mm_lora_token_counts(
723
+ self,
724
+ *,
725
+ modality: str,
726
+ mm_kwargs: MultiModalKwargsItem | None,
727
+ num_mm_embeds: int,
728
+ ) -> tuple[int, int | None]:
729
+ del modality, mm_kwargs
730
merge_size = self.vision_tower.spatial_merge_size
727
- return num_vision_tokens // (merge_size**2)
731
+ return num_mm_embeds * (merge_size**2), num_mm_embeds
732
733
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
734
image_input = self._parse_and_validate_image_input(**kwargs)
vllm/model_executor/models/gemma3_mm.py
+11
-37
@@ -19,6 +19,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
19
from vllm.multimodal import MULTIMODAL_REGISTRY
20
from vllm.multimodal.inputs import (
21
MultiModalFieldConfig,
22
+ MultiModalKwargsItem,
23
MultiModalKwargsItems,
24
)
25
from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
@@ -686,43 +687,16 @@ class Gemma3ForConditionalGeneration(
687
tower_model="vision_tower",
688
)
689
689
- def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
690
- """
691
- Calculate the number of tokens output by the vision encoder.
692
-
693
- The vision encoder processes images into patch embeddings. For Gemma3,
694
- the relationship between prompt placeholder tokens and actual vision
695
- encoder output tokens depends on the patch grid size.
696
-
697
- Args:
698
- num_image_tokens: Number of image placeholder tokens in the prompt
699
- (typically mm_tokens_per_image per image)
700
-
701
- Returns:
702
- Number of tokens output by the vision encoder
703
- """
704
- # For Gemma3, the vision encoder outputs tokens_per_side x tokens_per_side
705
- # tokens per image. Since num_image_tokens represents the number of
706
- # connector output tokens (mm_tokens_per_image = 256), and tokens_per_side
707
- # is sqrt(256) = 16, we need to account for the token expansion.
708
- # Based on empirical testing, the multiplier of 16 works correctly.
709
- return num_image_tokens * 16
710
-
711
- def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
712
- """
713
- Calculate the number of tokens output by the multimodal connector.
714
-
715
- The connector applies projection and normalization but maintains the
716
- token count for Gemma3.
717
-
718
- Args:
719
- num_vision_tokens: Number of tokens from vision encoder
720
-
721
- Returns:
722
- Number of tokens after connector processing
723
- """
724
- # The Gemma3 connector maintains a 1:1 token mapping
725
- return num_vision_tokens
690
+ def get_mm_lora_token_counts(
691
+ self,
692
+ *,
693
+ modality: str,
694
+ mm_kwargs: MultiModalKwargsItem | None,
695
+ num_mm_embeds: int,
696
+ ) -> tuple[int, int | None]:
697
+ del modality, mm_kwargs
698
+ tower_tokens = num_mm_embeds * 16
699
+ return tower_tokens, tower_tokens
700
701
def get_encoder_cudagraph_config(self):
702
from vllm.v1.worker.encoder_cudagraph_defs import (
vllm/model_executor/models/glm4_1v.py
+9
-11
@@ -83,6 +83,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
83
from vllm.multimodal.inputs import (
84
MultiModalFeatureSpec,
85
MultiModalFieldConfig,
86
+ MultiModalKwargsItem,
87
MultiModalKwargsItems,
88
VideoItem,
89
)
@@ -2378,19 +2379,16 @@ class Glm4vForConditionalGeneration(
2379
tower_model="visual.",
2380
)
2381
2381
- def get_num_mm_encoder_tokens(
2382
+ def get_mm_lora_token_counts(
2383
self,
2383
- num_image_tokens: int,
2384
- ) -> int:
2385
- merge_size = self.config.vision_config.spatial_merge_size
2386
- return num_image_tokens * (merge_size**2)
2387
-
2388
- def get_num_mm_connector_tokens(
2389
- self,
2390
- num_vision_tokens: int,
2391
- ) -> int:
2384
+ *,
2385
+ modality: str,
2386
+ mm_kwargs: MultiModalKwargsItem | None,
2387
+ num_mm_embeds: int,
2388
+ ) -> tuple[int, int | None]:
2389
+ del modality, mm_kwargs
2390
merge_size = self.config.vision_config.spatial_merge_size
2393
- return num_vision_tokens // (merge_size**2)
2391
+ return num_mm_embeds * (merge_size**2), num_mm_embeds
2392
2393
2394
@MULTIMODAL_REGISTRY.register_processor(
vllm/model_executor/models/h2ovl.py
+15
-14
@@ -14,7 +14,7 @@ from transformers import PretrainedConfig
14
15
from vllm.model_executor.layers.quantization import QuantizationConfig
16
from vllm.multimodal import MULTIMODAL_REGISTRY
17
-from vllm.multimodal.inputs import BatchedTensorInputs
17
+from vllm.multimodal.inputs import BatchedTensorInputs, MultiModalKwargsItem
18
from vllm.multimodal.parse import (
19
ImageEmbeddingItems,
20
ImageProcessorItems,
@@ -178,16 +178,17 @@ class H2OVLChatModel(InternVLChatModel):
178
prefix=prefix,
179
)
180
181
- def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
182
- if num_image_tokens <= 0 or self.num_image_token <= 0:
183
- return 0
184
-
185
- num_patches = num_image_tokens // self.num_image_token
186
- return num_patches * (self.patch_tokens + 1)
187
-
188
- def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
189
- if num_vision_tokens <= 0 or self.num_image_token <= 0:
190
- return 0
191
-
192
- num_patches = num_vision_tokens // (self.patch_tokens + 1)
193
- return num_patches * self.num_image_token
181
+ def get_mm_lora_token_counts(
182
+ self,
183
+ *,
184
+ modality: str,
185
+ mm_kwargs: MultiModalKwargsItem | None,
186
+ num_mm_embeds: int,
187
+ ) -> tuple[int, int | None]:
188
+ del modality, mm_kwargs
189
+ if num_mm_embeds <= 0 or self.num_image_token <= 0:
190
+ return 0, 0
191
+
192
+ num_patches = num_mm_embeds // self.num_image_token
193
+ tower_tokens = num_patches * (self.patch_tokens + 1)
194
+ return tower_tokens, num_patches * self.num_image_token
vllm/model_executor/models/idefics3.py
+9
-14
@@ -39,6 +39,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
39
from vllm.multimodal import MULTIMODAL_REGISTRY
40
from vllm.multimodal.inputs import (
41
MultiModalFieldConfig,
42
+ MultiModalKwargsItem,
43
MultiModalKwargsItems,
44
)
45
from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems
@@ -837,20 +838,14 @@ class Idefics3ForConditionalGeneration(
838
tower_model="model.vision_model",
839
)
840
840
- def get_num_mm_encoder_tokens(
841
+ def get_mm_lora_token_counts(
842
self,
842
- num_image_tokens: int,
843
- ) -> int:
844
- hf_config = self.config
845
- scale_factor = hf_config.scale_factor
846
-
847
- return num_image_tokens * scale_factor**2
848
-
849
- def get_num_mm_connector_tokens(
850
- self,
851
- num_vision_tokens: int,
852
- ) -> int:
843
+ *,
844
+ modality: str,
845
+ mm_kwargs: MultiModalKwargsItem | None,
846
+ num_mm_embeds: int,
847
+ ) -> tuple[int, int | None]:
848
+ del modality, mm_kwargs
849
hf_config = self.config
850
scale_factor = hf_config.scale_factor
855
-
856
- return num_vision_tokens // scale_factor**2
851
+ return num_mm_embeds * scale_factor**2, num_mm_embeds
vllm/model_executor/models/interfaces.py
+10
-17
@@ -31,7 +31,7 @@ import torch
31
import torch.nn as nn
32
from torch import Tensor
33
from transformers.models.whisper.tokenization_whisper import LANGUAGES
34
-from typing_extensions import Self, TypeIs
34
+from typing_extensions import Self, TypeIs, deprecated
35
36
from vllm.logger import init_logger
37
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -420,23 +420,16 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol):
420
421
yield
422
423
- def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
424
- """
425
- Implement this function to enable LoRA support
426
- for the tower module of the multi-modal model.
427
- Given the number of image tokens, output the number of
428
- multi-modal encoder tokens.
429
- """
430
- ...
423
+ @deprecated(
424
+ "get_num_mm_encoder_tokens is deprecated; use get_mm_lora_token_counts instead."
425
+ )
426
+ def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: ...
427
432
- def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
433
- """
434
- Implement this function to enable LoRA support
435
- for the connector module of the multi-modal model.
436
- Given the number of vision tokens, output the number of
437
- multi-modal connector tokens.
438
- """
439
- ...
428
+ @deprecated(
429
+ "get_num_mm_connector_tokens is deprecated; use "
430
+ "get_mm_lora_token_counts instead."
431
+ )
432
+ def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: ...
433
434
def get_mm_lora_token_counts(
435
self,
vllm/model_executor/models/internvl.py
+15
-13
@@ -29,6 +29,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
29
from vllm.multimodal.inputs import (
30
BatchedTensorInputs,
31
MultiModalFieldConfig,
32
+ MultiModalKwargsItem,
33
MultiModalKwargsItems,
34
)
35
from vllm.multimodal.parse import (
@@ -914,19 +915,20 @@ class InternVLChatModel(
915
tower_model="vision_model",
916
)
917
917
- def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
918
- if num_image_tokens <= 0 or self.num_image_token <= 0:
919
- return 0
920
-
921
- num_patches = num_image_tokens // self.num_image_token
922
- return num_patches * (self.patch_tokens + 1)
923
-
924
- def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
925
- if num_vision_tokens <= 0 or self.num_image_token <= 0:
926
- return 0
927
-
928
- num_patches = num_vision_tokens // (self.patch_tokens + 1)
929
- return num_patches * self.num_image_token
918
+ def get_mm_lora_token_counts(
919
+ self,
920
+ *,
921
+ modality: str,
922
+ mm_kwargs: MultiModalKwargsItem | None,
923
+ num_mm_embeds: int,
924
+ ) -> tuple[int, int | None]:
925
+ del modality, mm_kwargs
926
+ if num_mm_embeds <= 0 or self.num_image_token <= 0:
927
+ return 0, 0
928
+
929
+ num_patches = num_mm_embeds // self.num_image_token
930
+ tower_tokens = num_patches * (self.patch_tokens + 1)
931
+ return tower_tokens, num_patches * self.num_image_token
932
933
# -- SupportsEncoderCudaGraph protocol methods --
934
vllm/model_executor/models/keye.py
+10
-6
@@ -40,6 +40,7 @@ from vllm.multimodal.inputs import (
40
ImageItem,
41
MultiModalFeatureSpec,
42
MultiModalFieldConfig,
43
+ MultiModalKwargsItem,
44
MultiModalKwargsItems,
45
VideoItem,
46
)
@@ -1495,13 +1496,16 @@ class BaseKeyeModule(nn.Module, SupportsMultiModal):
1496
tower_model="visual.",
1497
)
1498
1498
- def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
1499
- merge_size = self.config.vision_config.spatial_merge_size
1500
- return num_image_tokens * merge_size**2
1501
-
1502
- def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
1499
+ def get_mm_lora_token_counts(
1500
+ self,
1501
+ *,
1502
+ modality: str,
1503
+ mm_kwargs: MultiModalKwargsItem | None,
1504
+ num_mm_embeds: int,
1505
+ ) -> tuple[int, int | None]:
1506
+ del modality, mm_kwargs
1507
merge_size = self.config.vision_config.spatial_merge_size
1504
- return num_vision_tokens // merge_size**2
1508
+ return num_mm_embeds * merge_size**2, num_mm_embeds
1509
1510
1511
@MULTIMODAL_REGISTRY.register_processor(
vllm/model_executor/models/lfm2_vl.py
+10
-8
@@ -35,6 +35,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
35
from vllm.multimodal import MULTIMODAL_REGISTRY
36
from vllm.multimodal.inputs import (
37
MultiModalFieldConfig,
38
+ MultiModalKwargsItem,
39
MultiModalKwargsItems,
40
)
41
from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
@@ -1284,12 +1285,13 @@ class Lfm2VLForConditionalGeneration(
1285
tower_model="vision_tower",
1286
)
1287
1287
- def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
1288
- downsample_factor = self.config.downsample_factor
1289
-
1290
- return num_image_tokens * downsample_factor**2
1291
-
1292
- def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
1288
+ def get_mm_lora_token_counts(
1289
+ self,
1290
+ *,
1291
+ modality: str,
1292
+ mm_kwargs: MultiModalKwargsItem | None,
1293
+ num_mm_embeds: int,
1294
+ ) -> tuple[int, int | None]:
1295
+ del modality, mm_kwargs
1296
downsample_factor = self.config.downsample_factor
1294
-
1295
- return num_vision_tokens // downsample_factor**2
1297
+ return num_mm_embeds * downsample_factor**2, num_mm_embeds
vllm/model_executor/models/llava.py
+9
-14
@@ -28,6 +28,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
28
from vllm.multimodal.cache import BaseMultiModalProcessorCache
29
from vllm.multimodal.inputs import (
30
MultiModalFieldConfig,
31
+ MultiModalKwargsItem,
32
MultiModalKwargsItems,
33
)
34
from vllm.multimodal.parse import (
@@ -739,18 +740,12 @@ class LlavaForConditionalGeneration(
740
tower_model="vision_tower",
741
)
742
742
- def get_num_mm_encoder_tokens(
743
+ def get_mm_lora_token_counts(
744
self,
744
- num_image_tokens: int,
745
- ) -> int:
746
- # LLaVA's vision encoder outputs one token per patch without
747
- # spatial merging or pixel shuffle
748
- return num_image_tokens
749
-
750
- def get_num_mm_connector_tokens(
751
- self,
752
- num_vision_tokens: int,
753
- ) -> int:
754
- # LLaVA's MLP projector outputs the same number of tokens
755
- # as it receives from the vision encoder (1:1 mapping)
756
- return num_vision_tokens
745
+ *,
746
+ modality: str,
747
+ mm_kwargs: MultiModalKwargsItem | None,
748
+ num_mm_embeds: int,
749
+ ) -> tuple[int, int | None]:
750
+ del modality, mm_kwargs
751
+ return num_mm_embeds, num_mm_embeds
vllm/model_executor/models/llava_next_video.py
+21
-21
@@ -17,6 +17,7 @@ from vllm.model_executor.models.clip import CLIPVisionModel
17
from vllm.multimodal import MULTIMODAL_REGISTRY
18
from vllm.multimodal.inputs import (
19
MultiModalFieldConfig,
20
+ MultiModalKwargsItem,
21
MultiModalKwargsItems,
22
)
23
from vllm.multimodal.parse import (
@@ -494,31 +495,30 @@ class LlavaNextVideoForConditionalGeneration(
495
tower_model="vision_tower",
496
)
497
497
- def get_num_mm_encoder_tokens(
498
+ def get_mm_lora_token_counts(
499
self,
499
- num_video_tokens: int,
500
- ) -> int:
500
+ *,
501
+ modality: str,
502
+ mm_kwargs: MultiModalKwargsItem | None,
503
+ num_mm_embeds: int,
504
+ ) -> tuple[int, int | None]:
505
+ del modality, mm_kwargs
506
+
507
# Invert the spatial pooling done by `vision_resampler`: each frame
508
# contributes `pooled_grid_length ** 2` tokens after the language
509
# model's placeholder count, but `patch_grid_length ** 2` tokens
510
# when it leaves the vision encoder.
511
pooled_tokens_per_frame = self.pooled_grid_length**2
506
- if num_video_tokens <= 0 or pooled_tokens_per_frame <= 0:
507
- return 0
508
-
509
- num_frames = num_video_tokens // pooled_tokens_per_frame
510
- return num_frames * (self.patch_grid_length**2)
511
-
512
- def get_num_mm_connector_tokens(
513
- self,
514
- num_vision_tokens: int,
515
- ) -> int:
516
- # The projector is length-preserving; it runs after
517
- # `vision_resampler` has already pooled each frame down to
518
- # `pooled_grid_length ** 2` tokens.
512
patch_tokens_per_frame = self.patch_grid_length**2
520
- if num_vision_tokens <= 0 or patch_tokens_per_frame <= 0:
521
- return 0
522
-
523
- num_frames = num_vision_tokens // patch_tokens_per_frame
524
- return num_frames * (self.pooled_grid_length**2)
513
+ if (
514
+ num_mm_embeds <= 0
515
+ or pooled_tokens_per_frame <= 0
516
+ or patch_tokens_per_frame <= 0
517
+ ):
518
+ return 0, 0
519
+
520
+ num_frames = num_mm_embeds // pooled_tokens_per_frame
521
+ return (
522
+ num_frames * patch_tokens_per_frame,
523
+ num_frames * pooled_tokens_per_frame,
524
+ )
vllm/model_executor/models/mllama4.py
+13
-15
@@ -60,6 +60,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
60
from vllm.multimodal import MULTIMODAL_REGISTRY
61
from vllm.multimodal.inputs import (
62
MultiModalFieldConfig,
63
+ MultiModalKwargsItem,
64
MultiModalKwargsItems,
65
)
66
from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
@@ -1281,21 +1282,18 @@ class Llama4ForConditionalGeneration(
1282
tower_model="vision_model.",
1283
)
1284
1284
- def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
1285
+ def get_mm_lora_token_counts(
1286
+ self,
1287
+ *,
1288
+ modality: str,
1289
+ mm_kwargs: MultiModalKwargsItem | None,
1290
+ num_mm_embeds: int,
1291
+ ) -> tuple[int, int | None]:
1292
+ del modality, mm_kwargs
1293
vision_config = self.config.vision_config
1294
patches_per_chunk = Mllama4ProcessingInfo.get_patch_per_chunk(vision_config)
1287
- if num_image_tokens <= 0 or patches_per_chunk <= 0:
1288
- return 0
1289
- raw_patches = (vision_config.image_size // vision_config.patch_size) ** 2
1290
- num_chunks = num_image_tokens // patches_per_chunk
1291
- # Encoder processes raw_patches + 1 (CLS) per chunk
1292
- return num_chunks * (raw_patches + 1)
1293
-
1294
- def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
1295
- vision_config = self.config.vision_config
1295
+ if num_mm_embeds <= 0 or patches_per_chunk <= 0:
1296
+ return 0, 0
1297
raw_patches = (vision_config.image_size // vision_config.patch_size) ** 2
1297
- if num_vision_tokens <= 0:
1298
- return 0
1299
- num_chunks = num_vision_tokens // (raw_patches + 1)
1300
- patches_per_chunk = Mllama4ProcessingInfo.get_patch_per_chunk(vision_config)
1301
- return num_chunks * patches_per_chunk
1298
+ num_chunks = num_mm_embeds // patches_per_chunk
1299
+ return num_chunks * (raw_patches + 1), num_chunks * patches_per_chunk
vllm/model_executor/models/moondream3.py
+10
-5
@@ -47,6 +47,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
47
from vllm.multimodal import MULTIMODAL_REGISTRY
48
from vllm.multimodal.inputs import (
49
MultiModalFieldConfig,
50
+ MultiModalKwargsItem,
51
MultiModalKwargsItems,
52
)
53
from vllm.multimodal.parse import ImageSize, MultiModalDataItems
@@ -1104,11 +1105,15 @@ class Moondream3ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
1105
def get_language_model(self) -> nn.Module:
1106
return self.text
1107
1107
- def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
1108
- return num_image_tokens
1109
-
1110
- def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
1111
- return num_vision_tokens
1108
+ def get_mm_lora_token_counts(
1109
+ self,
1110
+ *,
1111
+ modality: str,
1112
+ mm_kwargs: MultiModalKwargsItem | None,
1113
+ num_mm_embeds: int,
1114
+ ) -> tuple[int, int | None]:
1115
+ del modality, mm_kwargs
1116
+ return num_mm_embeds, num_mm_embeds
1117
1118
def _split_pixel_values(
1119
self,
vllm/model_executor/models/openvla.py
+14
-6
@@ -24,7 +24,11 @@ from vllm.model_executor.models.interfaces import (
24
SupportsPP,
25
)
26
from vllm.multimodal import MULTIMODAL_REGISTRY
27
-from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
27
+from vllm.multimodal.inputs import (
28
+ MultiModalFieldConfig,
29
+ MultiModalKwargsItem,
30
+ MultiModalKwargsItems,
31
+)
32
from vllm.multimodal.parse import (
33
ImageEmbeddingItems,
34
ImageProcessorItems,
@@ -487,11 +491,15 @@ class OpenVLAForActionPrediction(nn.Module, SupportsMultiModal, SupportsPP):
491
tower_model="vision_backbone",
492
)
493
490
- def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
491
- return num_image_tokens
492
-
493
- def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
494
- return num_vision_tokens
494
+ def get_mm_lora_token_counts(
495
+ self,
496
+ *,
497
+ modality: str,
498
+ mm_kwargs: MultiModalKwargsItem | None,
499
+ num_mm_embeds: int,
500
+ ) -> tuple[int, int | None]:
501
+ del modality, mm_kwargs
502
+ return num_mm_embeds, num_mm_embeds
503
504
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
505
def maybe_rename_vision_weights(
vllm/model_executor/models/paligemma.py
+10
-5
@@ -14,6 +14,7 @@ from vllm.logger import init_logger
14
from vllm.multimodal import MULTIMODAL_REGISTRY
15
from vllm.multimodal.inputs import (
16
MultiModalFieldConfig,
17
+ MultiModalKwargsItem,
18
MultiModalKwargsItems,
19
)
20
from vllm.multimodal.parse import (
@@ -401,8 +402,12 @@ class PaliGemmaForConditionalGeneration(
402
tower_model="vision_tower",
403
)
404
404
- def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
405
- return num_image_tokens
406
-
407
- def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
408
- return num_vision_tokens
405
+ def get_mm_lora_token_counts(
406
+ self,
407
+ *,
408
+ modality: str,
409
+ mm_kwargs: MultiModalKwargsItem | None,
410
+ num_mm_embeds: int,
411
+ ) -> tuple[int, int | None]:
412
+ del modality, mm_kwargs
413
+ return num_mm_embeds, num_mm_embeds
vllm/model_executor/models/pixtral.py
+11
-9
@@ -42,6 +42,7 @@ from vllm.model_executor.models.utils import WeightsMapper
42
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems
43
from vllm.multimodal.inputs import (
44
MultiModalFieldConfig,
45
+ MultiModalKwargsItem,
46
MultiModalKwargsOptionalItems,
47
NestedTensors,
48
)
@@ -623,17 +624,18 @@ class PixtralForConditionalGeneration(
624
tower_model="vision_encoder",
625
)
626
626
- def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
627
- if getattr(self, "patch_merger", None) is None:
628
- return num_image_tokens
629
- merge_size = self.vision_args.spatial_merge_size
630
- return num_image_tokens * (merge_size**2)
631
-
632
- def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
627
+ def get_mm_lora_token_counts(
628
+ self,
629
+ *,
630
+ modality: str,
631
+ mm_kwargs: MultiModalKwargsItem | None,
632
+ num_mm_embeds: int,
633
+ ) -> tuple[int, int | None]:
634
+ del modality, mm_kwargs
635
if getattr(self, "patch_merger", None) is None:
634
- return num_vision_tokens
636
+ return num_mm_embeds, num_mm_embeds
637
merge_size = self.vision_args.spatial_merge_size
636
- return num_vision_tokens // (merge_size**2)
638
+ return num_mm_embeds * (merge_size**2), num_mm_embeds
639
640
641
# Vision encoder
vllm/model_executor/models/qwen2_5_vl.py
+9
-14
@@ -70,6 +70,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
70
from vllm.multimodal.inputs import (
71
MultiModalFeatureSpec,
72
MultiModalFieldConfig,
73
+ MultiModalKwargsItem,
74
MultiModalKwargsItems,
75
)
76
from vllm.multimodal.parse import MultiModalDataItems
@@ -2050,21 +2051,15 @@ class Qwen2_5_VLForConditionalGeneration(
2051
tower_model="visual.",
2052
)
2053
2053
- def get_num_mm_encoder_tokens(
2054
+ def get_mm_lora_token_counts(
2055
self,
2055
- num_image_tokens: int,
2056
- ) -> int:
2057
- hf_config = self.config
2058
- vision_config = hf_config.vision_config
2059
- merge_size = vision_config.spatial_merge_size
2060
-
2061
- return num_image_tokens * merge_size**2
2062
-
2063
- def get_num_mm_connector_tokens(
2064
- self,
2065
- num_vision_tokens: int,
2066
- ) -> int:
2056
+ *,
2057
+ modality: str,
2058
+ mm_kwargs: MultiModalKwargsItem | None,
2059
+ num_mm_embeds: int,
2060
+ ) -> tuple[int, int | None]:
2061
+ del modality, mm_kwargs
2062
hf_config = self.config
2063
vision_config = hf_config.vision_config
2064
merge_size = vision_config.spatial_merge_size
2070
- return num_vision_tokens // merge_size**2
2065
+ return num_mm_embeds * merge_size**2, num_mm_embeds
vllm/model_executor/models/qwen2_vl.py
+9
-14
@@ -66,6 +66,7 @@ from vllm.multimodal.inputs import (
66
ImageItem,
67
MultiModalFeatureSpec,
68
MultiModalFieldConfig,
69
+ MultiModalKwargsItem,
70
MultiModalKwargsItems,
71
VideoItem,
72
)
@@ -1747,21 +1748,15 @@ class Qwen2VLForConditionalGeneration(
1748
tower_model="visual.",
1749
)
1750
1750
- def get_num_mm_encoder_tokens(
1751
+ def get_mm_lora_token_counts(
1752
self,
1752
- num_image_tokens: int,
1753
- ) -> int:
1754
- hf_config = self.config
1755
- vision_config = hf_config.vision_config
1756
- merge_size = vision_config.spatial_merge_size
1757
-
1758
- return num_image_tokens * merge_size**2
1759
-
1760
- def get_num_mm_connector_tokens(
1761
- self,
1762
- num_vision_tokens: int,
1763
- ) -> int:
1753
+ *,
1754
+ modality: str,
1755
+ mm_kwargs: MultiModalKwargsItem | None,
1756
+ num_mm_embeds: int,
1757
+ ) -> tuple[int, int | None]:
1758
+ del modality, mm_kwargs
1759
hf_config = self.config
1760
vision_config = hf_config.vision_config
1761
merge_size = vision_config.spatial_merge_size
1767
- return num_vision_tokens // merge_size**2
1762
+ return num_mm_embeds * merge_size**2, num_mm_embeds
vllm/model_executor/models/qwen3_asr.py
+10
-6
@@ -67,6 +67,7 @@ from vllm.multimodal.inputs import (
67
AudioItem,
68
MultiModalFeatureSpec,
69
MultiModalFieldConfig,
70
+ MultiModalKwargsItem,
71
MultiModalKwargsItems,
72
)
73
from vllm.multimodal.parse import (
@@ -625,16 +626,19 @@ class Qwen3ASRForConditionalGeneration(
626
tower_model=["audio_tower."],
627
)
628
628
- def get_num_mm_encoder_tokens(self, num_audio_tokens: int) -> int:
629
- """Return the number of tokens processed by the audio tower encoder.
630
-
631
- Required for LoRA support on the tower module.
632
- """
629
+ def get_mm_lora_token_counts(
630
+ self,
631
+ *,
632
+ modality: str,
633
+ mm_kwargs: MultiModalKwargsItem | None,
634
+ num_mm_embeds: int,
635
+ ) -> tuple[int, int | None]:
636
+ del modality, mm_kwargs
637
# For Qwen3-ASR, the audio tower produces one embedding per audio
638
# placeholder token inserted into the prompt (no additional
639
# merge/downsample step like vision towers). Therefore, the encoder
640
# token budget is identity.
637
- return num_audio_tokens
641
+ return num_mm_embeds, None
642
643
@classmethod
644
def get_speech_to_text_config(
vllm/model_executor/models/qwen3_vl.py
+8
-14
@@ -3077,24 +3077,18 @@ class Qwen3VLForConditionalGeneration(
3077
tower_model="visual.",
3078
)
3079
3080
- def get_num_mm_encoder_tokens(
3080
+ def get_mm_lora_token_counts(
3081
self,
3082
- num_image_tokens: int,
3083
- ) -> int:
3084
- hf_config = self.config
3085
- vision_config = hf_config.vision_config
3086
- merge_size = vision_config.spatial_merge_size
3087
-
3088
- return num_image_tokens * merge_size**2
3089
-
3090
- def get_num_mm_connector_tokens(
3091
- self,
3092
- num_vision_tokens: int,
3093
- ) -> int:
3082
+ *,
3083
+ modality: str,
3084
+ mm_kwargs: MultiModalKwargsItem | None,
3085
+ num_mm_embeds: int,
3086
+ ) -> tuple[int, int | None]:
3087
+ del modality, mm_kwargs
3088
hf_config = self.config
3089
vision_config = hf_config.vision_config
3090
merge_size = vision_config.spatial_merge_size
3097
- return num_vision_tokens // merge_size**2
3091
+ return num_mm_embeds * merge_size**2, num_mm_embeds
3092
3093
3094
@lru_cache
vllm/model_executor/models/ultravox.py
+19
-14
@@ -29,6 +29,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
29
from vllm.multimodal import MULTIMODAL_REGISTRY
30
from vllm.multimodal.inputs import (
31
MultiModalFieldConfig,
32
+ MultiModalKwargsItem,
33
MultiModalKwargsItems,
34
NestedTensors,
35
)
@@ -373,9 +374,9 @@ def _build_chunk_attn_metadata(
374
valid/padding boundary (equivalent to the key-padding mask the HF
375
implementation uses), while every row still flows through the
376
(potentially LoRA-wrapped) linears, keeping the per-chunk token counts
376
- constant as required by `get_num_mm_encoder_tokens` /
377
- `get_num_mm_connector_tokens`. Queries at padding positions produce
378
- (garbage) outputs, which are trimmed by `audio_token_len` downstream.
377
+ constant as required by `get_mm_lora_token_counts`. Queries at padding
378
+ positions produce (garbage) outputs, which are trimmed by `audio_token_len`
379
+ downstream.
380
"""
381
batch_size = feature_lens.shape[0]
382
starts = np.arange(batch_size, dtype=np.int64) * seq_len
@@ -684,7 +685,7 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
685
686
# LoRA on the tower/connector requires per-item token counts that are
687
# predictable from the placeholder count alone (see
687
- # `get_num_mm_encoder_tokens`), so pad every audio chunk's mel input
688
+ # `get_mm_lora_token_counts`), so pad every audio chunk's mel input
689
# to the tower's full context instead of the batch's max length.
690
self.pad_audio_to_max_context = bool(
691
lora_config is not None and lora_config.enable_tower_connector_lora
@@ -769,22 +770,26 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
770
self.config.audio_config.max_source_positions / self.config.stack_factor
771
)
772
772
- def get_num_mm_encoder_tokens(self, num_audio_tokens: int) -> int:
773
+ def get_mm_lora_token_counts(
774
+ self,
775
+ *,
776
+ modality: str,
777
+ mm_kwargs: MultiModalKwargsItem | None,
778
+ num_mm_embeds: int,
779
+ ) -> tuple[int, int | None]:
780
+ del modality, mm_kwargs
781
# With `pad_audio_to_max_context` (enforced when tower/connector LoRA
782
# is enabled), the tower's LoRA-wrapped linears always run on
783
# `max_source_positions` conv-downsampled tokens per chunk, regardless
784
# of the valid frame count.
777
- num_chunks = math.ceil(num_audio_tokens / self._get_max_tokens_per_chunk())
778
- return num_chunks * self.config.audio_config.max_source_positions
779
-
780
- def get_num_mm_connector_tokens(self, num_encoder_tokens: int) -> int:
785
+ num_chunks = math.ceil(num_mm_embeds / self._get_max_tokens_per_chunk())
786
+ tower_tokens = num_chunks * self.config.audio_config.max_source_positions
787
# The connector runs on the frame-stacked tower output. Stacking pads
788
# each chunk to a multiple of `stack_factor`, so this is
789
# ceil(max_source_positions / stack_factor) tokens per chunk (188 for
784
- # whisper's 1500), not `num_encoder_tokens // stack_factor` (187).
785
- max_source_positions = self.config.audio_config.max_source_positions
786
- num_chunks = num_encoder_tokens // max_source_positions
787
- return num_chunks * self._get_max_tokens_per_chunk()
790
+ # whisper's 1500), not `tower_tokens // stack_factor` (187).
791
+ connector_tokens = num_chunks * self._get_max_tokens_per_chunk()
792
+ return tower_tokens, connector_tokens
793
794
def _audio_features_to_embeddings(
795
self,
@@ -866,7 +871,7 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
871
# Pad every chunk to the tower's full context so the number of
872
# tokens processed by the tower/connector per chunk is constant,
873
# keeping the LoRA token mappings computed by
869
- # `get_num_mm_encoder_tokens`/`get_num_mm_connector_tokens` exact.
874
+ # `get_mm_lora_token_counts` exact.
875
# Attention to the extra padding is masked via `audio_lens`.
876
# Over-long inputs are not trimmed here; the tower raises on them.
877
pad_len = self.audio_tower.max_context_length - audio_features.shape[-1]