[Bugfix] Transformers backend: recompute `mm_token_type_ids` per request for M-RoPE (#46552)
Signed-off-by: Gonzague de Carpentier <decarpentierg@gmail.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Gonzague de Carpentier committed
Jun 28, 2026 at 17:19 UTC
7544286b04a860fcfb98725345f2b43a5023b844
1 file changed
+13
-18
vllm/model_executor/models/transformers/multimodal.py
+13
-18
@@ -55,6 +55,8 @@ if TYPE_CHECKING:
55
56
logger = init_logger(__name__)
57
58
+_MODALITY_TO_TOKEN_TYPE_ID = {"image": 1, "video": 2, "audio": 3}
59
+
60
61
class MultiModalProcessingInfo(BaseProcessingInfo):
62
def get_supported_mm_limits(self):
@@ -206,9 +208,8 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
208
)
209
210
# For gemma3 we check `token_type_ids` as the key
209
- mm_token_type_ids = processed_data.get(
210
- "mm_token_type_ids", processed_data.pop("token_type_ids", None)
211
- )
211
+ mm_token_type_ids = processed_data.pop("token_type_ids", None)
212
+ mm_token_type_ids = processed_data.pop("mm_token_type_ids", mm_token_type_ids)
213
214
# We can infer vLLM style placeholder from token type ids, if we split
215
# it for each input `mm_data`.
@@ -377,7 +378,6 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
378
return None
379
380
num_image_patches = kwargs.pop("num_image_patches")
380
- kwargs.pop("mm_token_type_ids", None) # used only in `model.get_rope_index`
381
382
if pixel_values is not None:
383
# ROCm: Force math SDP backend for vision encoder to avoid accuracy issues
@@ -468,24 +468,18 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
468
{
469
"image_grid_thw",
470
"video_grid_thw",
471
- "mm_token_type_ids",
471
"second_per_grid_ts",
472
"audio_feature_lengths",
473
"use_audio_in_video",
474
},
475
)
477
- if any(
478
- v
479
- for k, v in kwargs.items()
480
- if k not in {"image_grid_thw", "mm_token_type_ids"}
481
- ):
476
+ if any(v for k, v in kwargs.items() if k not in {"image_grid_thw"}):
477
raise NotImplementedError(
478
"Transformers modeling backend only supports images."
479
)
480
481
image_grid_thw = kwargs.get("image_grid_thw", [])
482
video_grid_thw = kwargs.get("video_grid_thw", [])
488
- mm_token_type_ids = kwargs.get("mm_token_type_ids")
483
484
image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)(
485
image_grid_thw
@@ -494,8 +488,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
488
video_grid_thw
489
)
490
497
- # In v4 `get_rope_index` doesn't have wildcard `kwargs`, and
498
- # can't accept arbitrary args, even if its value is `None`
491
+ # `get_rope_index` doesn't always accept arbitrary `kwargs`
492
kwargs = {}
493
if not hasattr(self, "_get_rope_index_accepts_mm_token_type_ids"):
494
import inspect
@@ -507,11 +500,13 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
500
or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
501
)
502
if self._get_rope_index_accepts_mm_token_type_ids:
510
- if mm_token_type_ids:
511
- kwargs["mm_token_type_ids"] = torch.cat(mm_token_type_ids)
512
- else:
513
- shape = (1, len(input_tokens))
514
- kwargs["mm_token_type_ids"] = torch.zeros(*shape, dtype=torch.int)
503
+ mm_token_type_ids = torch.zeros(len(input_tokens), dtype=torch.int)
504
+ for feature in mm_features:
505
+ position = feature.mm_position
506
+ offset, length = position.offset, position.length
507
+ mm_token_type_id = _MODALITY_TO_TOKEN_TYPE_ID[feature.modality]
508
+ mm_token_type_ids[offset : offset + length] = mm_token_type_id
509
+ kwargs["mm_token_type_ids"] = mm_token_type_ids.unsqueeze(0)
510
511
mrope_positions, mrope_position_delta = self.model.get_rope_index(
512
input_ids=torch.tensor(input_tokens).unsqueeze(0),