[Memory] Bound video-disabled MRoPE caches to the position envelope (#1019)

* [Memory] Bound TP4 prefill communication buffers to engine capacity

Assisted-by: Codex
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>

* Revert "[Memory] Bound TP4 prefill communication buffers to engine capacity"

This reverts commit dccb97db1b8e2a428639a4c130092cb8b03f22d0.

Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>

* [Memory] Bound video-disabled MRoPE caches to the position envelope

Assisted-by: Codex
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>

---------

Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>

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

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 = {

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,