diff --git a/tests/model_executor/test_mrope_cache_memory.py b/tests/model_executor/test_mrope_cache_memory.py new file mode 100644 index 000000000..776e2e217 --- /dev/null +++ b/tests/model_executor/test_mrope_cache_memory.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.config.multimodal import MultiModalConfig +from vllm.model_executor.layers import rotary_embedding as rope + +pytestmark = pytest.mark.skip_global_cleanup + + +def _config(video: int | None, *, multimodal: bool = True) -> VllmConfig: + config = VllmConfig() + limits = {"image": 1} + if video is not None: + limits["video"] = video + config.model_config = SimpleNamespace( + multimodal_config=MultiModalConfig(limit_per_prompt=limits) + if multimodal + else None, + max_model_len=512, + ) + return config + + +def _rope(**kwargs): + return rope.get_rope( + head_size=256, + max_position=512, + dtype=torch.float16, + rope_parameters={ + "rope_type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25, + "mrope_section": [11, 11, 10], + "mrope_interleaved": True, + **kwargs, + }, + ) + + +@pytest.mark.parametrize(("video", "rows"), [(0, 512), (1, 2048), (None, 2048)]) +def test_video_capacity_and_exact_cache_prefix(monkeypatch, video, rows): + monkeypatch.setattr(rope, "_ROPE_DICT", {}) + with set_current_vllm_config(_config(1)): + reference = _rope() + with set_current_vllm_config(_config(video)): + actual = _rope() + assert actual.cos_sin_cache.shape == (rows, 64) + assert torch.equal(actual.cos_sin_cache, reference.cos_sin_cache[:rows]) + positions = torch.tensor([[0, 511], [511, 0], [510, 511]]) + query = torch.randn(2, 6 * 256, dtype=torch.float16) + key = torch.randn(2, 256, dtype=torch.float16) + expected = reference.forward_native(positions, query.clone(), key.clone()) + observed = actual.forward_native(positions, query.clone(), key.clone()) + assert all(torch.equal(x, y) for x, y in zip(expected, observed)) + + +def test_video_engine_does_not_reuse_bounded_cache(monkeypatch): + monkeypatch.setattr(rope, "_ROPE_DICT", {}) + with set_current_vllm_config(_config(0)): + bounded = _rope() + assert _rope() is bounded + with set_current_vllm_config(_config(1)): + full = _rope() + assert _rope() is full + assert full is not bounded + assert full.cos_sin_cache.shape[0] == 2048 + + +def test_text_only_engine_keeps_served_context(monkeypatch): + monkeypatch.setattr(rope, "_ROPE_DICT", {}) + config = _config(0, multimodal=False) + config.model_config.max_model_len = 1024 + with set_current_vllm_config(config): + assert _rope().cos_sin_cache.shape[0] == 1024 + + +def test_unknown_engine_keeps_standalone_capacity(monkeypatch): + monkeypatch.setattr(rope, "_ROPE_DICT", {}) + with set_current_vllm_config(VllmConfig()): + assert _rope().cos_sin_cache.shape[0] == 2048 + + +def test_yarn_retains_existing_frequency_and_capacity(monkeypatch): + monkeypatch.setattr(rope, "_ROPE_DICT", {}) + args = {"rope_type": "yarn", "factor": 2.0, "original_max_position_embeddings": 512} + with set_current_vllm_config(_config(1)): + reference = _rope(**args) + with set_current_vllm_config(_config(0)): + actual = _rope(**args) + assert actual is reference diff --git a/vllm/model_executor/layers/rotary_embedding/__init__.py b/vllm/model_executor/layers/rotary_embedding/__init__.py index 0a69b3b3f..fc345bb81 100644 --- a/vllm/model_executor/layers/rotary_embedding/__init__.py +++ b/vllm/model_executor/layers/rotary_embedding/__init__.py @@ -71,7 +71,26 @@ def get_rope( raise ValueError(f"{partial_rotary_factor=} must be between 0.0 and 1.0") rotary_dim = int(head_size * partial_rotary_factor) - key = ( + mrope_cache_positions = None + if ( + dual_chunk_attention_config is None + and scaling_type == "default" + and "mrope_section" in rope_parameters + ): + from vllm.config import get_current_vllm_config_or_none + + config = get_current_vllm_config_or_none() + if config is not None and config.model_config is not None: + model_config = config.model_config + mm_config = model_config.multimodal_config + if mm_config is None or mm_config.get_limit_per_prompt("video") == 0: + # Text/image positions fit the token window; video's temporal + # coordinates can require the existing four-times envelope. + mrope_cache_positions = min( + 4 * max_position, max(max_position, model_config.max_model_len) + ) + + key: tuple[Any, ...] = ( head_size, rotary_dim, max_position, @@ -80,6 +99,9 @@ def get_rope( dual_chunk_attention_args, dtype, ) + if mrope_cache_positions is not None: + # A video-enabled engine must never inherit a smaller cached table. + key += (mrope_cache_positions,) if key in _ROPE_DICT: return _ROPE_DICT[key] @@ -109,6 +131,7 @@ def get_rope( dtype, mrope_section=rope_parameters["mrope_section"], mrope_interleaved=rope_parameters.get("mrope_interleaved", False), + cache_max_position_num=mrope_cache_positions, ) elif "use_fope" in rope_parameters and rope_parameters["use_fope"]: extra_kwargs = { diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py index 3c946dd13..c41b3399d 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope.py @@ -213,6 +213,7 @@ class MRotaryEmbedding(RotaryEmbeddingBase): mrope_interleaved: bool = False, # YaRN parameters. *, + cache_max_position_num: int | None = None, scaling_factor: float | None = None, extrapolation_factor: float = 1, attn_factor: float = 1, @@ -235,7 +236,11 @@ class MRotaryEmbedding(RotaryEmbeddingBase): # In Qwen2.5-VL, the maximum index value is related to the duration of # the input video. We enlarge max_position_embeddings to 4 times to get # a larger the cos and sin cache. - self.cache_max_position_num = max_position_embeddings * 4 + self.cache_max_position_num = ( + max_position_embeddings * 4 + if cache_max_position_num is None + else cache_max_position_num + ) super().__init__( head_size, rotary_dim,