[Bugfix][Model] Support variable-dimensional M-RoPE

tests/models/transformers/test_multimodal_mrope.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import torch
+
+from vllm.model_executor.models.transformers.multimodal import MultiModalMixin
+
+
+class _ImageOnlyMRoPEModel:
+    def get_rope_index(self, input_ids, image_grid_thw):
+        seq_len = input_ids.shape[-1]
+        positions = torch.arange(seq_len).view(1, 1, -1).expand(4, 1, -1)
+        return positions, torch.tensor([0])
+
+
+def test_get_mrope_input_positions_filters_unsupported_grid_kwargs():
+    mixin = SimpleNamespace(model=_ImageOnlyMRoPEModel())
+
+    positions, delta = MultiModalMixin.get_mrope_input_positions(mixin, [1, 2, 3], [])
+
+    assert positions.shape == (4, 3)
+    assert delta == 0

tests/transformers_utils/test_config.py

@@ -27,6 +27,46 @@
 )
 
 
+@pytest.mark.parametrize(
+    ("mrope_section", "expected_num_dims"),
+    [
+        ([16, 24, 24], 3),
+        ([16, 16, 16, 16], 4),
+    ],
+)
+def test_model_config_mrope_num_dims(mrope_section, expected_num_dims):
+    model_config = cast(
+        ModelConfig,
+        SimpleNamespace(
+            hf_config=SimpleNamespace(),
+            hf_text_config=SimpleNamespace(
+                rope_parameters={"mrope_section": mrope_section}
+            ),
+        ),
+    )
+
+    assert ModelConfig.mrope_num_dims.fget(model_config) == expected_num_dims
+
+
+def test_model_config_mrope_num_dims_from_thinker_config():
+    thinker_text_config = SimpleNamespace(
+        rope_parameters={"mrope_section": [16, 16, 16, 16]}
+    )
+    model_config = cast(
+        ModelConfig,
+        SimpleNamespace(
+            hf_config=SimpleNamespace(
+                get_text_config=lambda: SimpleNamespace(),
+                thinker_config=SimpleNamespace(text_config=thinker_text_config),
+            ),
+            hf_text_config=SimpleNamespace(),
+            uses_mrope=True,
+        ),
+    )
+
+    assert ModelConfig.mrope_num_dims.fget(model_config) == 4
+
+
 def test_glm5_next_accepts_deepseek_sparse_attention_layers():
     layer_types = ["linear_attention", "deepseek_sparse_attention"]
 

tests/v1/worker/test_rope_state.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from vllm.model_executor.models.interfaces import SupportsMRoPE
+from vllm.v1.worker.gpu.mm import rope
+
+
+class _MRoPEModel(SupportsMRoPE):
+    def get_mrope_input_positions(self, input_tokens, mm_features):
+        raise NotImplementedError
+
+
+def test_get_rope_state_uses_configured_mrope_dimensions():
+    model_config = SimpleNamespace(
+        uses_mrope=True,
+        mrope_num_dims=4,
+        uses_xdrope_dim=0,
+    )
+
+    with patch.object(rope, "RopeState") as rope_state_cls:
+        rope.get_rope_state(
+            model_config,
+            _MRoPEModel(),
+            max_num_reqs=2,
+            max_num_tokens=8,
+            max_model_len=16,
+            device=torch.device("cpu"),
+        )
+
+    rope_state_cls.assert_called_once_with(
+        num_dims=4,
+        has_delta=True,
+        max_num_reqs=2,
+        max_num_tokens=8,
+        max_model_len=16,
+        device=torch.device("cpu"),
+    )

vllm/config/model.py

@@ -1849,6 +1849,28 @@ def uses_alibi(self) -> bool:
     def uses_mrope(self) -> bool:
         return uses_mrope(self.hf_config)
 
+    @property
+    def mrope_num_dims(self) -> int:
+        thinker_config = getattr(self.hf_config, "thinker_config", None)
+        thinker_text_config = getattr(thinker_config, "text_config", None)
+
+        for hf_config in (
+            self.hf_text_config,
+            thinker_text_config,
+            self.hf_config,
+        ):
+            rope_parameters = getattr(hf_config, "rope_parameters", None)
+            if not isinstance(rope_parameters, dict):
+                continue
+
+            mrope_section = rope_parameters.get("mrope_section")
+            if isinstance(mrope_section, (list, tuple)):
+                return len(mrope_section)
+
+        # Preserve the historical three-dimensional allocation for custom
+        # configs that identify themselves as M-RoPE without exposing sections.
+        return 3 if self.uses_mrope else 0
+
     @property
     def uses_xdrope_dim(self) -> int:
         return uses_xdrope_dim(self.hf_config)

vllm/model_executor/models/transformers/multimodal.py

@@ -1235,18 +1235,32 @@ def get_mrope_input_positions(
         image_grid_thw = torch.stack(image_grid_thw) if image_grid_thw else None
         video_grid_thw = torch.stack(video_grid_thw) if video_grid_thw else None
 
-        # `get_rope_index` doesn't always accept arbitrary `kwargs`
-        kwargs = {}
-        if not hasattr(self, "_get_rope_index_accepts_mm_token_type_ids"):
+        # `get_rope_index` doesn't always accept arbitrary keyword arguments.
+        if not hasattr(self, "_get_rope_index_kwarg_names"):
             import inspect
 
             sig = inspect.signature(self.model.get_rope_index)
             params = sig.parameters
-            self._get_rope_index_accepts_mm_token_type_ids = (
-                "mm_token_type_ids" in params
-                or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
+            self._get_rope_index_kwarg_names = (
+                None
+                if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
+                else frozenset(params)
             )
-        if self._get_rope_index_accepts_mm_token_type_ids:
+
+        kwarg_names = self._get_rope_index_kwarg_names
+
+        def accepts_kwarg(name: str) -> bool:
+            return kwarg_names is None or name in kwarg_names
+
+        kwargs = {
+            name: value
+            for name, value in (
+                ("image_grid_thw", image_grid_thw),
+                ("video_grid_thw", video_grid_thw),
+            )
+            if accepts_kwarg(name)
+        }
+        if accepts_kwarg("mm_token_type_ids"):
             mm_token_type_ids = torch.zeros(len(input_tokens), dtype=torch.int)
             for feature in mm_features:
                 position = feature.mm_position
@@ -1257,8 +1271,6 @@ def get_mrope_input_positions(
 
         mrope_positions, mrope_position_delta = self.model.get_rope_index(
             input_ids=torch.tensor(input_tokens).unsqueeze(0),
-            image_grid_thw=image_grid_thw,
-            video_grid_thw=video_grid_thw,
             **kwargs,
         )
 

vllm/v1/worker/gpu/mm/rope.py

@@ -14,17 +14,17 @@
 class RopeState:
     """Unified state for multi-dimensional RoPE variants (M-RoPE, XD-RoPE).
 
-    M-RoPE: 3 dims, uses position delta for decode.
+    M-RoPE: dimensions are defined by the model config, uses position delta
+    for decode.
     XD-RoPE: 3 or 4 dims, delta is 0 (decode uses orig_pos for all dims).
 
     NOTE: `positions` is implemented with one additional dummy position on
     purpose to make it non-contiguous so that it can work with torch compile.
     See detailed explanation in
     https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923
 
-    NOTE: When M-RoPE is enabled, position ids are 3D regardless of the
-    modality of inputs. For text-only inputs, each dimension has identical
-    position IDs, making M-RoPE functionally equivalent to 1D-RoPE.
+    NOTE: For text-only inputs, each dimension has identical position IDs,
+    making M-RoPE functionally equivalent to 1D-RoPE.
     See page 5 of https://arxiv.org/abs/2409.12191
     """
 
@@ -143,7 +143,7 @@ def get_rope_state(
     if model_config.uses_mrope:
         assert isinstance(model, SupportsMRoPE)
         return RopeState(
-            num_dims=3,
+            num_dims=model_config.mrope_num_dims,
             has_delta=True,
             max_num_reqs=max_num_reqs,
             max_num_tokens=max_num_tokens,

vllm/v1/worker/gpu_model_runner.py

@@ -576,6 +576,7 @@ def __init__(
         # Multi-modal data support
         self.mm_registry = MULTIMODAL_REGISTRY
         self.uses_mrope = model_config.uses_mrope
+        self.mrope_num_dims = model_config.mrope_num_dims
         self.uses_xdrope_dim = model_config.uses_xdrope_dim
         self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
             model_config
@@ -872,13 +873,11 @@ def __init__(
             # with torch compile.
             # See detailed explanation in https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923
 
-            # NOTE: When M-RoPE is enabled, position ids are 3D regardless of
-            # the modality of inputs. For text-only inputs, each dimension has
-            # identical position IDs, making M-RoPE functionally equivalent to
-            # 1D-RoPE.
+            # NOTE: For text-only inputs, each dimension has identical position
+            # IDs, making M-RoPE functionally equivalent to 1D-RoPE.
             # See page 5 of https://arxiv.org/abs/2409.12191
             self.mrope_positions = self._make_buffer(
-                (3, self.max_num_tokens + 1), dtype=torch.int64
+                (self.mrope_num_dims, self.max_num_tokens + 1), dtype=torch.int64
             )
 
         # Only relevant for models using XD-RoPE (e.g, HunYuan-VL)