From eca938a96414176f978143f3de5a41f830c41529 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sun, 8 Nov 2026 19:12:33 +0000 Subject: [PATCH] Update the longrope inv_freq buffer in place `longrope_frequency_update` re-registered `inv_freq` on every forward, replacing the buffer with a new tensor. CUDA graphs replay from the buffer's original address, so the replaced tensor is never read and vLLM rejects the module under CUDA graphs. Copy into the existing buffer instead. Co-Authored-By: Claude Opus 5.5 --- src/transformers/modeling_rope_utils.py | 3 ++- tests/models/phi3/test_modeling_phi3.py | 25 +++++++++++++++++++++++-- 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/src/transformers/modeling_rope_utils.py b/src/transformers/modeling_rope_utils.py index cdae3ca6fdd0..8281a0d9d648 100644 --- a/src/transformers/modeling_rope_utils.py +++ b/src/transformers/modeling_rope_utils.py @@ -77,7 +77,8 @@ def longrope_frequency_update(self, position_ids, device, layer_type=None): # `seq_len` is a scalar, so this picks the long or short set for all positions, like an if/else but traceable inv_freq = torch.where(seq_len > original_max_position_embeddings, long_inv_freq, original_inv_freq) - self.register_buffer(f"{prefix}inv_freq", inv_freq, persistent=False) + # Update in place so the buffer keeps its address, which CUDA graphs replay from + getattr(self, f"{prefix}inv_freq").copy_(inv_freq) def dynamic_frequency_update(self, position_ids, device, layer_type=None): """ diff --git a/tests/models/phi3/test_modeling_phi3.py b/tests/models/phi3/test_modeling_phi3.py index a5b36e9ed56a..85e457f57ad4 100644 --- a/tests/models/phi3/test_modeling_phi3.py +++ b/tests/models/phi3/test_modeling_phi3.py @@ -132,8 +132,7 @@ def test_longrope_generation_crossing_boundary(self, sliding_window): expected = model(prefix, use_cache=False).logits[:, -1] torch.testing.assert_close(logits, expected, rtol=1e-4, atol=1e-4) - @parameterized.expand([(8,), (16,)]) - def test_longrope_compiled_matches_eager(self, seq_len): + def _get_longrope_config(self): config = self.model_tester.get_config() config.original_max_position_embeddings = 8 config.max_position_embeddings = 32 @@ -146,6 +145,11 @@ def test_longrope_compiled_matches_eager(self, seq_len): "original_max_position_embeddings": 8, "factor": 4.0, } + return config + + @parameterized.expand([(8,), (16,)]) + def test_longrope_compiled_matches_eager(self, seq_len): + config = self._get_longrope_config() model = Phi3ForCausalLM(config).to(torch_device).eval() input_ids = torch.arange(1, seq_len + 1, device=torch_device).unsqueeze(0) @@ -154,6 +158,23 @@ def test_longrope_compiled_matches_eager(self, seq_len): compiled = torch.compile(model, fullgraph=True)(input_ids, use_cache=False).logits torch.testing.assert_close(compiled, expected, rtol=1e-4, atol=1e-4) + def test_longrope_updates_inv_freq_in_place(self): + # CUDA graphs replay from the buffer's address, so switching frequency sets must not replace it + config = self._get_longrope_config() + rotary_emb = Phi3ForCausalLM(config).to(torch_device).model.rotary_emb + inv_freq = rotary_emb.inv_freq + short_inv_freq = rotary_emb.original_inv_freq.clone() + x = torch.zeros(1, device=torch_device) + + with torch.no_grad(): + rotary_emb(x, torch.arange(16, device=torch_device)[None]) + self.assertIs(rotary_emb.inv_freq, inv_freq) + torch.testing.assert_close(inv_freq, rotary_emb.long_inv_freq) + + rotary_emb(x, torch.arange(8, device=torch_device)[None]) + self.assertIs(rotary_emb.inv_freq, inv_freq) + torch.testing.assert_close(inv_freq, short_inv_freq) + @slow @require_torch