[Perf] Bound M-RoPE cache by certified position ranges

[lesj0610]

Purpose

M-RoPE currently uses a legacy 4x maximum both for position semantics and for the physical cosine/sine cache. The larger range is required by models whose video timestamps can exceed the token sequence length, but it over-allocates the physical cache for models whose M-RoPE positions are proven to stay within the configured sequence bound.

This change adds an explicit, fail-closed position-capacity contract:

• A concrete model class must opt in to sequence-bounded M-RoPE positions. The capability is read off __dict__, so a declaration inherited from a parent does not certify a derived position implementation. Qwen3_5MoeForConditionalGeneration is a subclass of the certified dense entrypoint and stays on the legacy allocation.
• The YaRN semantic maximum remains unchanged, preserving frequency and correction-range math; only the physical cache row count can be reduced.
• Multimodal pruning can push positions past the reduced sequence length, so ModelConfig falls back to the legacy allocation whenever pruning is enabled.
• The resolved physical capacity is included in both the RoPE instance key and the model compilation hash.
• Only the dense Qwen3_5ForConditionalGeneration path is certified here. Qwen3.5 MoE, Qwen2/2.5-VL, Qwen3-VL, Omni, Transformers-delegated, and out-of-tree models retain the legacy path.

Models without a bound keep delegating to the base and YaRN cache implementations, so their caches stay bit-for-bit identical. Bounded models build the cache from cache_max_position_num from torch.arange onwards rather than materializing the semantic-sized cache and slicing it: the cache is created on the accelerator inside the model loader's target_device context, so slicing afterwards would not lower the construction peak this bound exists to reduce, and would add a copy on top of it.

There are no model-name, quantization, GPU, FlashInfer, or environment-specific branches.

Correctness coverage

• Unbounded instances are asserted byte-identical to the delegated base and YaRN implementations, for both the unscaled and YaRN paths.
• Bounded caches are asserted to be a bitwise prefix of the legacy cache, for the unscaled, YaRN, and interleaved paths. The interleaved subclass previously had its own row count overwritten back to 4x by its parent.
• The YaRN semantic maximum is asserted unchanged between bounded and unbounded instances.
• A non-positive bound is rejected.
• _ROPE_DICT is exercised under an isolated dict: the same parameters with and without a bound must not alias, and the unbounded entry must still be reused.
• Capability resolution covers a plain class, a certified class, a subclass that only inherits the attribute, and a subclass that explicitly opts out. The Qwen3_5ForConditionalGeneration / Qwen3_5MoeForConditionalGeneration pair is asserted to be a real subclass relationship with opposite capability results.
• ModelConfig.mrope_cache_max_position covers certified/uncertified crossed with pruning enabled/disabled, and the bound is asserted to change compute_hash().
• The _ModelInfo registry cache is exercised through inspect_model_cls(): a cache file written before the field existed is re-inspected and rewritten once, and the following call is a warm hit that does not re-inspect the model class. The field intentionally carries no dataclass default.

Validation

Base commit e25c586b90, torch 2.13.0+cu130 / CUDA 13.0, GPU NVIDIA CMP 170HX compute capability (8, 0).

.venv/bin/python -m pytest \
  tests/kernels/core/test_mrope_cache_bounds.py \
  tests/models/test_mrope_position_bound_capability.py \
  tests/config/test_mrope_cache_bound_config.py -q

CUDA_VISIBLE_DEVICES=0 .venv/bin/python -m pytest tests/kernels/core/test_mrope.py -q
.venv/bin/python -m pytest tests/models/test_registry.py -q

• New tests: 25 passed.
• tests/kernels/core/test_mrope.py: 50 passed.
• tests/models/test_registry.py: 381 passed, 4 failed. All four (HCXVisionForCausalLM, Dots3NoteForCausalLM, Dots3NoteMTPModel, KananaVForConditionalGeneration) reproduce on unmodified e25c586b90, so they are pre-existing baseline failures rather than regressions.
• ruff check, ruff format --check, and the full pre-commit hook set (including mypy for Python 3.10, SPDX, and DCO sign-off) passed.

Isolated M-RoPE cache-construction measurement, using the served checkpoint's own rope parameters (head_size=256, partial_rotary_factor=0.25 so rotary_dim=64, mrope_interleaved=True, mrope_section=[11, 11, 10], rope_theta=10000000, max_position_embeddings=262144, bound = max_model_len = 262144, bfloat16). Each case clears _ROPE_DICT, calls torch.cuda.empty_cache() and reset_peak_memory_stats(), builds under with torch.device("cuda:0") to mirror the loader's target_device context, then synchronizes and takes the delta of memory_allocated() and max_memory_allocated():

	cache shape	resident	cache-construction peak
legacy	(1048576, 64)	128.0 MiB	644.0 MiB
bounded	(262144, 64)	32.0 MiB	161.0 MiB
delta		-96.0 MiB	-483.0 MiB

Both figures match the arithmetic: legacy 4 + 128 + 128 + 128 + 256 = 644 MiB and bounded 1 + 32 + 32 + 32 + 64 = 161 MiB of t, freqs, cos, sin, and the concatenated result, with a final bfloat16 buffer of 128 MiB versus 32 MiB.

These are allocated-memory figures for the M-RoPE cache construction in isolation, not a whole-model initialization peak. The serving-level KV capacity impact was not measured: it depends on engine configuration, block alignment, and memory state, which would add more noise than signal to this change's evidence.

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Summary by CodeRabbit

• New Features

  • Added bounded M-RoPE cache sizing for supported models, reducing physical cache growth while preserving position limits.
  • Added capability detection for models with sequence-bounded M-RoPE positions.
  • Enabled the Qwen3.5 model family to use bounded M-RoPE caching when pruning is disabled.
  • Added validation for invalid cache bounds and ensured distinct configurations use separate caches.

• Bug Fixes

  • Preserved legacy behavior for unsupported or pruned models.
  • Improved cache isolation and consistency across standard, YaRN, and interleaved RoPE configurations.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 15aeb67c-2fef-4c73-a1ac-94f08f3fcaf3

📥 Commits

Reviewing files that changed from the base of the PR and between b26039b and cd31a80.

📒 Files selected for processing (11)

• tests/config/test_mrope_cache_bound_config.py
• tests/kernels/core/test_mrope_cache_bounds.py
• tests/models/test_mrope_position_bound_capability.py
• vllm/config/model.py
• vllm/model_executor/layers/rotary_embedding/__init__.py
• vllm/model_executor/layers/rotary_embedding/mrope.py
• vllm/model_executor/layers/rotary_embedding/mrope_interleaved.py
• vllm/model_executor/models/interfaces.py
• vllm/model_executor/models/qwen3_5.py
• vllm/model_executor/models/qwen3_next.py
• vllm/model_executor/models/registry.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.

---

📝 Walkthrough

Walkthrough

Changes

M-RoPE cache bounds

Layer / File(s)	Summary
Capability detection and bound selection vllm/model_executor/models/interfaces.py, vllm/model_executor/models/registry.py, vllm/config/model.py, vllm/model_executor/models/qwen3_5.py, tests/models/*, tests/config/*	The model registry records explicit sequence-bounded M-RoPE support. ModelConfig returns max_model_len only for certified, unpruned models and includes the bound in its hash.
Bounded cache construction vllm/model_executor/layers/rotary_embedding/mrope.py, vllm/model_executor/layers/rotary_embedding/mrope_interleaved.py, tests/kernels/core/test_mrope_cache_bounds.py	M-RoPE validates positive bounds and separates semantic position limits from physical cache size. Standard, YaRN, and interleaved caches support bounded construction while preserving unbounded behavior.
Factory and model integration vllm/model_executor/layers/rotary_embedding/__init__.py, vllm/model_executor/models/qwen3_next.py, tests/config/test_mrope_cache_bound_config.py, tests/kernels/core/test_mrope_cache_bounds.py	get_rope accepts the cache bound, includes it in cache keys, and passes it to M-RoPE constructors. Qwen3Next supplies the value from ModelConfig.

Estimated code review effort: 4 (Complex) | ~45 minutes

Merge Risk: ⚪ Minimal · up to cd31a

The change bounds M-RoPE cache allocation only for explicitly certified model paths while preserving legacy behavior elsewhere, with targeted correctness and regression checks reported as passing. No actionable merge-blocking risk remains beyond normal checks and review.

Sequence Diagram(s)

sequenceDiagram
  participant ModelConfig
  participant Qwen3NextAttention
  participant get_rope
  participant MRotaryEmbedding
  participant RopeRegistry
  ModelConfig->>Qwen3NextAttention: mrope_cache_max_position
  Qwen3NextAttention->>get_rope: pass cache bound
  get_rope->>RopeRegistry: key lookup including cache bound
  get_rope->>MRotaryEmbedding: construct with cache bound
  MRotaryEmbedding-->>get_rope: bounded rotary cache

Loading

Suggested reviewers: hmellor, gcanlin

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the primary change: bounding the M-RoPE cache using certified position ranges.
Docstring Coverage	✅ Passed	Docstring check was indeterminate for this PR — some files could not be analyzed in time. Not blocking.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch agent/mrope-cache-position-bounds

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[lesj0610]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.