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 7, 14:00 UTC
[BugFix] Fix 3D rope in transformers backend (#35097)
Signed-off-by: raushan <raushan@huggingface.co> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Raushan Turganbay committed
Feb 27, 2026 at 19:34 UTC
fd6de37fcafe9540ed821256877127df75d74db8
1 file changed
+12
-2
vllm/model_executor/models/transformers/multimodal.py
+12
-2
@@ -218,7 +218,7 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
218
if "mm_token_type_ids" in processed_data
219
else "token_type_ids"
220
)
221
- mm_token_type_ids = processed_data.pop(token_type_key)
221
+ mm_token_type_ids = processed_data.get(token_type_key)
222
223
# We can infer vLLM style placeholder from token type ids, if we split
224
# it for each input `mm_data`.
@@ -353,6 +353,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
353
354
num_image_patches = kwargs.pop("num_image_patches")
355
kwargs.pop("token_type_ids", None) # used only in `forward`
356
+ kwargs.pop("mm_token_type_ids", None) # used only in `model.get_rope_index`
357
358
if pixel_values is not None:
359
# ROCm: Force math SDP backend for vision encoder to avoid accuracy issues
@@ -443,6 +444,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
444
{
445
"image_grid_thw",
446
"video_grid_thw",
447
+ "mm_token_type_ids",
448
"second_per_grid_ts",
449
"audio_feature_lengths",
450
"use_audio_in_video",
@@ -451,7 +453,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
453
if any(
454
v
455
for k, v in kwargs.items()
454
- if k not in {"image_grid_thw", "video_grid_thw"}
456
+ if k not in {"image_grid_thw", "mm_token_type_ids"}
457
):
458
raise NotImplementedError(
459
"Transformers modeling backend only supports images."
@@ -459,6 +461,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
461
462
image_grid_thw = kwargs.get("image_grid_thw", [])
463
video_grid_thw = kwargs.get("video_grid_thw", [])
464
+ mm_token_type_ids = kwargs.get("mm_token_type_ids")
465
466
image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)(
467
image_grid_thw
@@ -467,10 +470,17 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
470
video_grid_thw
471
)
472
473
+ # In v4 `get_rope_index` doesn't have wildcard `kwargs`, and
474
+ # can't accept arbitrary args, even if its value is `None`
475
+ kwargs = {}
476
+ if mm_token_type_ids:
477
+ kwargs["mm_token_type_ids"] = torch.cat(mm_token_type_ids)
478
+
479
mrope_positions, mrope_position_delta = self.model.get_rope_index(
480
input_ids=torch.tensor(input_tokens).unsqueeze(0),
481
image_grid_thw=image_grid_thw,
482
video_grid_thw=video_grid_thw,
483
+ **kwargs,
484
)
485
486
mrope_positions = mrope_positions[:, 0]