From 3723a1ea96df967f7cef21ef6ea5b4af6c529807 Mon Sep 17 00:00:00 2001 From: Frank <132732282+YongfengX@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:56:54 +0800 Subject: [PATCH] [Bugfix][Transformers backend] Fix CUDA graph capture for dynamic-rope models (#58064) Signed-off-by: YongfengX Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Kevin H. Luu --- .buildkite/test_areas/models_basic.yaml | 3 +- .../models/transformers/test_dynamic_rope.py | 86 +++++++++++++++++++ .../models/transformers/base.py | 3 + .../models/transformers/utils.py | 44 +++++++++- 4 files changed, 132 insertions(+), 4 deletions(-) create mode 100644 tests/models/transformers/test_dynamic_rope.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 3e0d840c6f..32559f7c69 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -253,6 +253,7 @@ steps: - tests/models/test_adapters.py - tests/models/test_qwen3_5_mtp_config.py - tests/models/transformers/fusers/ + - tests/models/transformers/test_dynamic_rope.py - tests/model_executor/layers/test_activation.py - tests/models/test_language_model_cache_is_weak.py - tests/models/deepseek_v32/test_sequence_parallel.py @@ -261,7 +262,7 @@ steps: - tests/models/transformers/test_layer_registry.py device: cpu-small commands: - - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py + - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ models/transformers/test_dynamic_rope.py model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py - label: ':amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics' key: dsv41-mono-decode-numerics-amd diff --git a/tests/models/transformers/test_dynamic_rope.py b/tests/models/transformers/test_dynamic_rope.py new file mode 100644 index 0000000000..e47d652846 --- /dev/null +++ b/tests/models/transformers/test_dynamic_rope.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for skipping dynamic rope's device sync where it can never rescale.""" + +from types import SimpleNamespace + +import pytest +import torch +from transformers import HunYuanDenseV1Config, LlamaConfig +from transformers.models.hunyuan_v1_dense.modeling_hunyuan_v1_dense import ( + HunYuanDenseV1RotaryEmbedding, +) +from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding + +from vllm.model_executor.models.transformers.utils import ( + can_enable_torch_compile, + disable_noop_dynamic_rope, +) + +MAX_POSITION_EMBEDDINGS = 128 +CONFIG_KWARGS = dict( + hidden_size=64, + num_attention_heads=4, + num_hidden_layers=1, + max_position_embeddings=MAX_POSITION_EMBEDDINGS, +) +DYNAMIC = {"rope_type": "dynamic", "factor": 4.0, "rope_theta": 1e4} + + +def make_vllm_config(hf_config, max_model_len: int) -> SimpleNamespace: + model_config = SimpleNamespace(hf_config=hf_config, max_model_len=max_model_len) + return SimpleNamespace(model_config=model_config) + + +@pytest.mark.parametrize( + "rope_parameters,max_model_len,compilable", + [ + (DYNAMIC, MAX_POSITION_EMBEDDINGS, True), + (DYNAMIC, MAX_POSITION_EMBEDDINGS + 1, False), + # The longrope layer type still syncs + ( + {"full_attention": DYNAMIC, "sliding_attention": {"rope_type": "longrope"}}, + MAX_POSITION_EMBEDDINGS, + False, + ), + ], +) +def test_can_enable_torch_compile(rope_parameters, max_model_len, compilable): + """Dynamic rope only blocks compilation when it can still rescale.""" + config = SimpleNamespace( + max_position_embeddings=MAX_POSITION_EMBEDDINGS, + rope_parameters=rope_parameters, + ) + config.get_text_config = lambda: config + vllm_config = make_vllm_config(config, max_model_len) + assert can_enable_torch_compile(vllm_config) is compilable + + +@pytest.mark.parametrize( + "config,rotary_cls", + [ + (LlamaConfig(**CONFIG_KWARGS, rope_parameters=DYNAMIC), LlamaRotaryEmbedding), + # Hunyuan's dynamic init also applies `alpha`, which must be kept + ( + HunYuanDenseV1Config( + **CONFIG_KWARGS, head_dim=16, rope_parameters=DYNAMIC | {"alpha": 1e3} + ), + HunYuanDenseV1RotaryEmbedding, + ), + ], +) +def test_disable_noop_dynamic_rope(monkeypatch, config, rotary_cls): + """The update no longer syncs and the frequencies are unchanged.""" + inputs = torch.zeros(1, 8, 64), torch.arange(8)[None, :] + rotary_emb = rotary_cls(config) + expected = rotary_emb(*inputs) + + vllm_config = make_vllm_config(config, MAX_POSITION_EMBEDDINGS) + disable_noop_dynamic_rope(rotary_emb, vllm_config) + + def no_sync(*args, **kwargs): + raise AssertionError("dynamic rope update still runs") + + monkeypatch.setattr(torch, "max", no_sync) + actual = rotary_emb(*inputs) + assert all(torch.equal(e, a) for e, a in zip(expected, actual)) diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 6f4cbdc923..7f890034f2 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -66,6 +66,7 @@ from vllm.model_executor.models.transformers.fusers.attention import ( from vllm.model_executor.models.transformers.utils import ( attrsetter, can_enable_torch_compile, + disable_noop_dynamic_rope, get_feature_request_tip, init_on_device_without_buffers, log_replacement, @@ -182,6 +183,8 @@ class Base( self._create_hf_to_vllm_mapper() # Remove layers not on this pipeline parallel rank self.pipeline_parallel() + # Skip the device sync of dynamic rope where it can never rescale + disable_noop_dynamic_rope(self.model.get_decoder(), vllm_config) # Substitute remaining layers with vLLM's layers as needed self.recursive_replace() # Create attention instances for KV cache allocation diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py index a39dd8f73b..d36ff2015b 100644 --- a/vllm/model_executor/models/transformers/utils.py +++ b/vllm/model_executor/models/transformers/utils.py @@ -357,13 +357,51 @@ def get_feature_request_tip( return f"{prefix}{tip}" +def dynamic_rope_update_is_noop(vllm_config: "VllmConfig") -> bool: + """Whether `max_model_len` keeps Transformers' dynamic rope from ever rescaling. + + `@dynamic_rope_update` only recomputes `inv_freq` once positions grow past + `max_position_embeddings`, which vLLM never feeds beyond `max_model_len`. + """ + text_config = vllm_config.model_config.hf_config.get_text_config() + max_position_embeddings = getattr(text_config, "max_position_embeddings", None) + if max_position_embeddings is None: + return False + return vllm_config.model_config.max_model_len <= max_position_embeddings + + def can_enable_torch_compile(vllm_config: "VllmConfig") -> bool: """Callable to be passed to `@support_torch_compile`'s `enable_if` argument. Defaults to `True` but is disabled in the following situations: - - The model uses dynamic rope scaling. + - The model uses dynamic rope scaling that can still rescale within + `max_model_len` (otherwise `disable_noop_dynamic_rope` removes it). """ text_config = vllm_config.model_config.hf_config.get_text_config() - # Dynamic rope scaling is not compatible with torch.compile - return all(rp["rope_type"] != "dynamic" for rp in iter_rope_parameters(text_config)) + rope_types = {rp["rope_type"] for rp in iter_rope_parameters(text_config)} + if "dynamic" not in rope_types: + return True + # Layer types using longrope still sync, even once dynamic rope is disabled + return "longrope" not in rope_types and dynamic_rope_update_is_noop(vllm_config) + + +def disable_noop_dynamic_rope(model: nn.Module, vllm_config: "VllmConfig") -> None: + """Skip `@dynamic_rope_update` in rotary modules where it can never rescale. + + The update syncs on `torch.max(position_ids)` every forward, which breaks + torch.compile and CUDA graph capture. Relabelling the module's `rope_type` as + "default" skips it while keeping the `inv_freq` computed at init. + """ + if not dynamic_rope_update_is_noop(vllm_config): + return + for module in model.modules(): + rope_type = getattr(module, "rope_type", None) + if not hasattr(module, "max_seq_len_cached") or rope_type is None: + continue + if isinstance(rope_type, dict): + module.rope_type = { + k: "default" if v == "dynamic" else v for k, v in rope_type.items() + } + elif rope_type == "dynamic": + module.rope_type = "default"