From 592f5ecd53c5ec9ff44b7e10355f897c713ea824 Mon Sep 17 00:00:00 2001 From: Charles xu Date: Sat, 5 Sep 2026 19:48:14 +0800 Subject: [PATCH] [Bugfix][Model] Support variable-dimensional M-RoPE Derive the number of M-RoPE position channels from the configured sections, including thinker text configs, instead of assuming three dimensions. Also filter optional grid arguments against each Transformers model's get_rope_index signature. Assisted-by: OpenAI GPT-5 Signed-off-by: Charles xu --- .../transformers/test_multimodal_mrope.py | 24 +++++++++++ tests/transformers_utils/test_config.py | 40 ++++++++++++++++++ tests/v1/worker/test_rope_state.py | 42 +++++++++++++++++++ vllm/config/model.py | 22 ++++++++++ .../models/transformers/multimodal.py | 30 +++++++++---- vllm/v1/worker/gpu/mm/rope.py | 10 ++--- vllm/v1/worker/gpu_model_runner.py | 9 ++-- 7 files changed, 158 insertions(+), 19 deletions(-) create mode 100644 tests/models/transformers/test_multimodal_mrope.py create mode 100644 tests/v1/worker/test_rope_state.py diff --git a/tests/models/transformers/test_multimodal_mrope.py b/tests/models/transformers/test_multimodal_mrope.py new file mode 100644 index 000000000000..c2e018f4784c --- /dev/null +++ b/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 diff --git a/tests/transformers_utils/test_config.py b/tests/transformers_utils/test_config.py index 17bb031e863d..8024aa5311fa 100644 --- a/tests/transformers_utils/test_config.py +++ b/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"] diff --git a/tests/v1/worker/test_rope_state.py b/tests/v1/worker/test_rope_state.py new file mode 100644 index 000000000000..5b54d42f4e9a --- /dev/null +++ b/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"), + ) diff --git a/vllm/config/model.py b/vllm/config/model.py index d2bb65ea42c0..ff2a001192d1 100644 --- a/vllm/config/model.py +++ b/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) diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index befc2ffffeec..005b326d502a 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/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, ) diff --git a/vllm/v1/worker/gpu/mm/rope.py b/vllm/v1/worker/gpu/mm/rope.py index e5de28223473..57fcd234dd1f 100644 --- a/vllm/v1/worker/gpu/mm/rope.py +++ b/vllm/v1/worker/gpu/mm/rope.py @@ -14,7 +14,8 @@ 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 @@ -22,9 +23,8 @@ class RopeState: 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, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 6d35dc73ca0e..13e30f527dc0 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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)