Make longrope compileable by jiqing-feng · Pull Request #49555 · huggingface/transformers

jiqing-feng · 2026-10-13T05:14:45Z

longrope_frequency_update branches on a tensor (if seq_len > original_max_position_embeddings). So fullgraph=True fails on every longrope model, and the default static-cache auto-compile graph-breaks on every step. This PR picks the long or short frequency set with torch.where instead.

It also fixes two longrope bugs:

The long-frequency cache used the wrong key, so it was recomputed every forward, or raised UnboundLocalError with per-layer-type rope_parameters.
yarn_apply_mscale read rope_parameters["factor"] unconditionally, so openbmb/MiniCPM3-4B (no factor) failed to load with KeyError.

A phi3 test covers both branches under fullgraph=True and fails on main. Eager logits are bitwise identical to main. No new test failures.

A100, torch==2.14.0+cu130, bf16, 128 new tokens, tok/s before → after:

	Phi-3.5-mini	MiniCPM3-4B
eager	43.4 → 44.0	16.3 → 16.8
auto-compiled	122.1 → 129.7	79.8 → 83.0
`fullgraph=True`	fails → 128.6	fails → 82.8

Reproduction

import sys

import torch

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation.configuration_utils import CompileConfig


model_id = sys.argv[1]
device = torch.accelerator.current_accelerator().type

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16).to(device).eval()

inputs = tokenizer("The capital of France is", return_tensors="pt").to(device)
with torch.no_grad():
    eager = model.generate(**inputs, max_new_tokens=16, do_sample=False, disable_compile=True)
    out = model.generate(
        **inputs,
        max_new_tokens=16,
        do_sample=False,
        cache_implementation="static",
        compile_config=CompileConfig(fullgraph=True),
    )
print(tokenizer.decode(out[0], skip_special_tokens=True))
print("tokens identical to eager:", torch.equal(eager, out))

Before:

# Phi-3.5-mini-instruct
torch._dynamo.exc.Unsupported: Data-dependent branching
  File "transformers/modeling_rope_utils.py", line 64, in longrope_frequency_update
    if seq_len > original_max_position_embeddings:

# MiniCPM3-4B
  File "transformers/models/minicpm3/modeling_minicpm3.py", line 188, in yarn_apply_mscale
    scaling_factor = rope_parameters["factor"]
KeyError: 'factor'

After: both models generate and print tokens identical to eager: True.

Benchmark

python bench.py MODEL_ID {eager,compiled,fullgraph}

import sys
import time

import torch

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from transformers.generation.configuration_utils import CompileConfig


# One mode per process: torch.compile state leaks between configs otherwise.
MODES = {
    # `cache_implementation="static"` auto-compiles on CUDA, so the eager baseline needs `disable_compile`
    "eager": {"disable_compile": True},
    "compiled": {"cache_implementation": "static"},
    "fullgraph": {"cache_implementation": "static", "compile_config": CompileConfig(fullgraph=True)},
}
model_id, mode = sys.argv[1], sys.argv[2]
device = torch.accelerator.current_accelerator().type
NEW_TOKENS, WARMUP, RUNS = 128, 2, 5

config = AutoConfig.from_pretrained(model_id)
rp = config.rope_parameters
# Only needed to load MiniCPM3-4B before this PR. This is the value longrope derives when `factor` is absent.
rp.setdefault("factor", config.max_position_embeddings / rp["original_max_position_embeddings"])

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, config=config, dtype=torch.bfloat16).to(device).eval()
inputs = tokenizer("The capital of France is", return_tensors="pt").to(device)


def run():
    return model.generate(
        **inputs, max_new_tokens=NEW_TOKENS, min_new_tokens=NEW_TOKENS, do_sample=False, **MODES[mode]
    )


for _ in range(WARMUP):
    run()
torch.cuda.synchronize()

times = []
for _ in range(RUNS):
    t0 = time.perf_counter()
    run()
    torch.cuda.synchronize()
    times.append(time.perf_counter() - t0)

median = sorted(times)[RUNS // 2]
print(f"{mode:10s} {median:.3f} s   {NEW_TOKENS / median:6.1f} tok/s")

Logits check

Run before and after, then compare the saved tensors with torch.equal.

import sys

import torch

from transformers import AutoConfig, AutoModelForCausalLM


model_id, out_path = sys.argv[1], sys.argv[2]
device = torch.accelerator.current_accelerator().type

config = AutoConfig.from_pretrained(model_id)
rp = config.rope_parameters
# Only needed to load MiniCPM3-4B before this PR. This is the value longrope derives when `factor` is absent.
rp.setdefault("factor", config.max_position_embeddings / rp["original_max_position_embeddings"])
model = AutoModelForCausalLM.from_pretrained(model_id, config=config, dtype=torch.bfloat16).to(device).eval()

orig = rp["original_max_position_embeddings"]
g = torch.Generator().manual_seed(0)
logits = {}
with torch.no_grad():
    # short stays on `short_factor`, long crosses `original_max_position_embeddings` and uses `long_factor`
    for tag, n in [("short", 512), ("long", orig + 512)]:
        ids = torch.randint(100, config.vocab_size - 100, (1, n), generator=g).to(device)
        logits[tag] = model(ids, use_cache=False).logits[0, -64:].float().cpu()
torch.save(logits, out_path)

`longrope_frequency_update` picks between the short and long scaling factors based on the sequence length. `seq_len` is derived from `position_ids` and is therefore always a tensor, so the check is data-dependent branching and fullgraph compilation fails on the very first forward pass. Unlike the other dynamic RoPE types, the choice genuinely depends on the sequence length and cannot be resolved at initialization. Compute both sets of frequencies once and select between them with `torch.where`, which is traceable and keeps the selected values identical to before. This also fixes two latent bugs in the same block. The cache lookup used `f"{layer_type}_long_inv_freq"` instead of the prefix, so for a non-nested RoPE it looked up `"None_long_inv_freq"`, which never exists and the long factors were recomputed on every forward pass. And `long_inv_freq` was never moved to `device`, unlike `original_inv_freq`. `openbmb/MiniCPM3-4B` is blocked earlier: `yarn_apply_mscale` reads `rope_parameters["factor"]` unconditionally, but `factor` is optional and a longrope config does not necessarily carry it, so building the model raises `KeyError`. Move the read inside the branch that needs it. That helper is shared through `modular`, so the models deriving from it are regenerated.

Phi3Config copies `original_max_position_embeddings` (default 4096) into `rope_parameters`, so seq_len=16 never exceeded it and the long factors were never exercised.

On CPU, `.to(device)` returns the same nn.Buffer, so setattr goes through Module.__setattr__ and reads `.persistent`, which dynamo cannot trace. The buffer already follows model.to(), so the reassignment is not needed.

Rocketknight1 · 2026-10-14T13:49:36Z

LGTM but it's a change to the core RoPE code so cc core maintainers before merge! @ArthurZucker @Cyrilvallez @vasqu

vasqu · 2026-10-14T15:58:05Z

cc @zucchini-nlp if you could take a first look here

zucchini-nlp · 2026-10-15T06:23:28Z

+        # selecting branchlessly keeps this traceable, `seq_len` is a tensor
+        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)


from what i see longrope changes the freq of all positions when seq_len goes beyond config. Doesnt look equivalent with current code if we just keep original freqs as is

It still does. seq_len is max(position_ids) + 1, a scalar, so torch.where picks the whole long or short set for every position. It is the same if/else as before, just without the Python branch. I ran main and this PR on a tiny Phi3 (short, then long, then short again): inv_freq and cos/sin are bitwise identical, and in the long call even position 1 uses the long frequencies. I also reworded the comment so this is clearer.

Script

import hashlib import torch from transformers import Phi3Config from transformers.models.phi3.modeling_phi3 import Phi3RotaryEmbedding torch.manual_seed(0) head_dim, ompe = 16, 8 cfg = Phi3Config( hidden_size=64, num_attention_heads=4, max_position_embeddings=64, original_max_position_embeddings=ompe, rope_parameters={ "rope_type": "longrope", "rope_theta": 10000.0, "short_factor": [1.0 + 0.1 * i for i in range(head_dim // 2)], "long_factor": [2.0 + 0.3 * i for i in range(head_dim // 2)], "original_max_position_embeddings": ompe, }, ) rope = Phi3RotaryEmbedding(cfg) x = torch.randn(1, 1, 32, head_dim) def run(seq_len): cos, sin = rope(x[:, :, :seq_len], torch.arange(seq_len)[None]) return rope.inv_freq.clone(), cos, sin def h(t): return hashlib.md5(t.contiguous().numpy().tobytes()).hexdigest()[:10] short_freq, cos_s, _ = run(ompe) long_freq, cos_l, sin_l = run(2 * ompe) back_freq, _, _ = run(ompe) print("inv_freq short", h(short_freq), "long", h(long_freq), "back", h(back_freq)) print("cos/sin long ", h(cos_l), h(sin_l)) print("pos1 cos short vs long equal:", torch.equal(cos_s[0, 1], cos_l[0, 1])) ref = torch.outer(torch.arange(2 * ompe).float(), long_freq) ref = torch.cat([ref, ref], -1).cos() * rope.attention_scaling print("long call: all positions use long freqs:", torch.allclose(cos_l[0], ref, atol=1e-6))

Same output on main and this PR:

inv_freq short ac8ef83cd5 long 6a1a3cdc3e back ac8ef83cd5 cos/sin long 569afb9d1e 159ce6875c pos1 cos short vs long equal: False long call: all positions use long freqs: True

ah right right, i imaged that was a mask over positions greater/smaller 😄 Thanks for adding hte comment

zucchini-nlp · 2026-10-15T06:24:00Z

    def dynamic_frequency_update(self, position_ids, device, layer_type=None):
        """


do you want to try and make dynamic rope compilable? 😄

Happy to look at it, but I'd rather do it in a separate PR. Dynamic rope keeps state: it grows the cached length and resets it when the input gets short again. So it needs more care than longrope, and I'd like to keep this PR small since it touches core RoPE code. We can discuss the approach there.

Wait, it might be possible with little changes. Let me try, I will give you feedback once I figured it out.

Done in 0ed7ffb

…to HEAD

ArthurZucker

Very nice contribution ty!
I'd just make sure that small inputs are not affected w.o compile (perf wise)? but I'm sure it should be fine!

ArthurZucker · 2026-10-16T00:46:46Z

-        scaling_factor = rope_parameters["factor"]
        if mscale_all_dim:
-            mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
+            # `factor` is optional for some rope types, only read it when it is needed


we don't need the comment everywhere!

isnt that modular dragging it around 😆 oh yeah, maybe not always

Removed. It came from modular_deepseek_v2.py, so I dropped it there and regenerated the 13 modeling files.

zucchini-nlp · 2026-10-16T01:03:58Z

+        inv_freq, self.attention_scaling = rope_init_fn(
+            self.config,
+            device,
+            seq_len=max_seq_len_cached,
+            layer_type=layer_type,
+        )


only thing is that we recompute freqs each time at reset or smaller seq length, can you confirm the perfs pls

It does slow down eager: ~155 → ~430 µs per rotary_emb call, +0.1–1.2 ms/token in generate. I reverted the dynamic part (743408a) to keep this PR to longrope, and will look into dynamic separately.

A100, torch 2.14, bf16, eager, two alternating rounds (round 1 / round 2), dynamic RoPE forced with factor=4.0, max_position_embeddings=1024:

dynamic RoPE rotary_emb µs/call, main → PR generate ms/token, main → PR

Qwen2.5-0.5B, 32-token prompt 154 / 158 → 443 / 423 16.45 / 16.70 → 17.66 / 16.92

Qwen2.5-0.5B, 1536-token prompt 153 / 158 → 470 / 420 16.58 / 16.95 → 17.66 / 17.06

Llama-3.1-8B, 32-token prompt 160 / 170 → 431 / 429 22.22 / 22.71 → 22.82 / 22.83

Llama-3.1-8B, 1536-token prompt 159 / 174 → 430 / 429 22.60 / 23.48 → 23.77 / 23.59

zucchini-nlp · 2026-10-16T01:05:31Z

+    def test_longrope_compiled_matches_eager(self, seq_len):
+        config = self.model_tester.get_config()
+        config.original_max_position_embeddings = 8
+        config.max_position_embeddings = 32
+        rotary_dim = config.hidden_size // config.num_attention_heads // 2
+        config.rope_parameters = {
+            "rope_type": "longrope",
+            "rope_theta": 10000.0,
+            "short_factor": [1.0] * rotary_dim,
+            "long_factor": [2.0] * rotary_dim,
+            "original_max_position_embeddings": 8,
+            "factor": 4.0,
+        }
+        model = Phi3ForCausalLM(config).to(torch_device).eval()
+        input_ids = torch.arange(1, seq_len + 1, device=torch_device).unsqueeze(0)
+
+        with torch.no_grad():
+            expected = model(input_ids, use_cache=False).logits


dynamic matcthes eager in compile test missing 😄

If you wish, feel free to split dynamic into separate PR - just dont want to block you from getting longrope merged. With dynamic we'll need to check perfs, looks sus from the code to re-compute freq all the time

It was test_dynamic_rope_compiled_matches_eager in tests/utils/test_modeling_rope_utils.py, now reverted together with the dynamic change.

jiqing-feng · 2026-10-16T02:37:47Z

Very nice contribution ty! I'd just make sure that small inputs are not affected w.o compile (perf wise)? but I'm sure it should be fine!

Not slower in eager: short inputs are on par, long inputs are faster since main recomputed the long freqs on every forward (the cache key bug fixed here). Same setup, Phi-3.5-mini:

longrope	`rotary_emb` µs/call, main → PR	`generate` ms/token, main → PR
32-token prompt	157 / 148 → 146 / 147	24.61 / 22.52 → 24.28 / 23.41
4608-token prompt	351 / 316 → 151 / 148	26.86 / 24.68 → 26.21 / 25.39

generate varies ~±1 ms between rounds here, so those are within noise.

Script

python bench.py MODEL_ID [--dynamic]

import statistics
import sys
import time

import torch

from transformers import AutoConfig, AutoModelForCausalLM


name = sys.argv[1]
cfg = AutoConfig.from_pretrained(name)
if "--dynamic" in sys.argv:
    cfg.rope_parameters = {"rope_type": "dynamic", "rope_theta": cfg.rope_parameters["rope_theta"], "factor": 4.0}
    cfg.max_position_embeddings = 1024
model = AutoModelForCausalLM.from_pretrained(name, config=cfg, dtype=torch.bfloat16, device_map="cuda").eval()
rope = model.config.rope_parameters
orig = rope.get("original_max_position_embeddings", model.config.max_position_embeddings)

rot = model.model.rotary_emb
x = torch.zeros(1, 1, 8, dtype=torch.bfloat16, device="cuda")
g = torch.Generator().manual_seed(0)
NEW = 128
short_len, long_len = 32, orig + 512


def bench_rotary(pos, n=2000, rounds=7):
    pid = torch.tensor([[pos]], device="cuda")
    res = []
    for _ in range(rounds):
        for _ in range(50):
            rot(x, pid)
        torch.cuda.synchronize()
        t = time.perf_counter()
        for _ in range(n):
            rot(x, pid)
        torch.cuda.synchronize()
        res.append((time.perf_counter() - t) / n * 1e6)
    return statistics.median(res)


def bench_generate(ids, reps):
    times = []
    for i in range(reps + 2):
        torch.cuda.synchronize()
        t = time.perf_counter()
        model.generate(ids, max_new_tokens=NEW, min_new_tokens=NEW, do_sample=False)
        torch.cuda.synchronize()
        if i >= 2:
            times.append((time.perf_counter() - t) / NEW * 1e3)
    return statistics.median(times)


with torch.no_grad():
    print(f"rotary_emb, short: {bench_rotary(short_len - 1):.1f} us/call")
    print(f"rotary_emb, long : {bench_rotary(long_len - 1):.1f} us/call")
    short = torch.randint(100, model.config.vocab_size - 100, (1, short_len), generator=g).cuda()
    long = torch.randint(100, model.config.vocab_size - 100, (1, long_len), generator=g).cuda()
    print(f"generate, short  : {bench_generate(short, 10):.2f} ms/token")
    print(f"generate, long   : {bench_generate(long, 6):.2f} ms/token")

zucchini-nlp

Thanks lets merge and do dynamic in follow-up then

HuggingFaceDocBuilderDev · 2026-10-19T01:56:28Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

zucchini-nlp · 2026-10-19T02:06:19Z

The test is failing on main and hopefully fixed soon, I will try to merge tmrw if it's not urgent. Otherwise ping Anton/Arthur to force it

qgallouedec · 2026-10-19T19:01:34Z

This closes #31339

github-actions · 2026-10-20T03:37:06Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: axk1, axk2, bailing_hybrid, deepseek_v2, deepseek_v3, deepseek_v32, glm4_moe_lite, glm_moe_dsa, hy_v4, longcat_flash, minicpm3, mistral4, phi3, youtu

github-actions · 2026-10-20T04:00:45Z

CI recap

Dashboard: View test results in Grafana
Latest run: 39365126936:2
Result: success | Jobs: 5 | Tests: 374 | Failures: 0 | Duration: 1m 5s

jiqing-feng added 4 commits October 13, 2026 03:33

Make the longrope compile test reach the long branch

ed7aa84

Phi3Config copies `original_max_position_embeddings` (default 4096) into `rope_parameters`, so seq_len=16 never exceeded it and the long factors were never exercised.

Don't reassign the original_inv_freq buffer inside forward

cfa265b

On CPU, `.to(device)` returns the same nn.Buffer, so setattr goes through Module.__setattr__ and reads `.persistent`, which dynamo cannot trace. The buffer already follows model.to(), so the reassignment is not needed.

Merge branch 'main' into fix-longrope-compile

aa04f1c

jiqing-feng marked this pull request as ready for review October 14, 2026 03:00

github-actions Bot requested review from ArthurZucker and Rocketknight1 October 14, 2026 03:01

zucchini-nlp reviewed Oct 15, 2026

View reviewed changes

jiqing-feng added 4 commits October 15, 2026 14:33

Clarify that longrope swaps the whole frequency set

ce640f6

Merge branch 'main' of https://github.com/huggingface/transformers in…

f431d9b

…to HEAD

Regenerate bailing_hybrid after merging main

4221422

Make dynamic rope compileable

0ed7ffb

jiqing-feng changed the title ~~Make longrope compileable~~ Make longrope and dynamic RoPE compileable Oct 15, 2026

ArthurZucker approved these changes Oct 16, 2026

View reviewed changes

zucchini-nlp reviewed Oct 16, 2026

View reviewed changes

jiqing-feng added 2 commits October 16, 2026 09:55

Drop redundant comment in yarn_apply_mscale

5cddc50

Move dynamic RoPE compile support to a separate PR

743408a

jiqing-feng changed the title ~~Make longrope and dynamic RoPE compileable~~ Make longrope compileable Oct 16, 2026

zucchini-nlp approved these changes Oct 19, 2026

View reviewed changes

Merge branch 'main' into fix-longrope-compile

7a24437

zucchini-nlp enabled auto-merge October 19, 2026 01:46

zucchini-nlp and others added 2 commits October 19, 2026 22:46

Merge branch 'main' into fix-longrope-compile

7164fed

Merge branch 'main' into fix-longrope-compile

b4bcd58

Merge branch 'main' into fix-longrope-compile

97ba00b

jiqing-feng and others added 2 commits October 20, 2026 09:22

Merge branch 'main' into fix-longrope-compile

f43a0b7

Merge branch 'main' into fix-longrope-compile

f9eb338

zucchini-nlp added this pull request to the merge queue Oct 20, 2026

Merged via the queue into huggingface:main with commit 823a6ea Oct 20, 2026
116 checks passed

This was referenced Oct 21, 2026

feat: RoPE-related typing improvements #41967

Closed

[serge] Fix 4 integration tests for model minicpm3 failing with other (other (4)) #49041

Closed

stevhliu mentioned this pull request Oct 23, 2026

[docs] Translation workflow #49889

Closed

hmellor mentioned this pull request Nov 8, 2026

Update the longrope inv_freq buffer in place #50480

Merged

3 of 5 tasks

		def dynamic_frequency_update(self, position_ids, device, layer_type=None):
		"""

dynamic RoPE	`rotary_emb` µs/call, main → PR	`generate` ms/token, main → PR
Qwen2.5-0.5B, 32-token prompt	154 / 158 → 443 / 423	16.45 / 16.70 → 17.66 / 16.92
Qwen2.5-0.5B, 1536-token prompt	153 / 158 → 470 / 420	16.58 / 16.95 → 17.66 / 17.06
Llama-3.1-8B, 32-token prompt	160 / 170 → 431 / 429	22.22 / 22.71 → 22.82 / 22.83
Llama-3.1-8B, 1536-token prompt	159 / 174 → 430 / 429	22.60 / 23.48 → 23.77 / 23.59

Conversation

jiqing-feng commented Oct 13, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Rocketknight1 commented Oct 14, 2026

Uh oh!

vasqu commented Oct 14, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zucchini-nlp Oct 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jiqing-feng commented Oct 16, 2026

Uh oh!

zucchini-nlp left a comment

Choose a reason for hiding this comment

Uh oh!

HuggingFaceDocBuilderDev commented Oct 19, 2026

Uh oh!

zucchini-nlp commented Oct 19, 2026

Uh oh!

qgallouedec commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 20, 2026

Uh oh!

github-actions Bot commented Oct 20, 2026

CI recap

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

jiqing-feng commented Oct 13, 2026 •

edited by github-actions Bot

Loading

zucchini-nlp Oct 16, 2026 •

edited

Loading