Skip to content

Jet-Long: Efficient Long-Context Extension with Dynamic Bifocal RoPE

🕒 Published (v1): 2026-07-08 00:00 UTC · Source: HuggingFace · link

Why this paper was selected

Song Han + Han Cai (MIT efficient-ML); dynamic long-context RoPE critical for agentic trace lengths

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Jet-Long is a tuning-free, zero-shot method for extending LLM context beyond the pretraining window by pairing a local RoPE-faithful attention window with a remote window whose rescaling factor adapts dynamically to the current sequence length. It avoids the short-context fidelity vs. long-context reach tradeoff inherent in fixed-factor methods, and a fused CuTe kernel makes the bifocal construction nearly free at inference. On Qwen3-1.7B/4B/8B extended to 128K, it leads the strongest zero-shot baseline on RULER by up to +4.79 pp with ≤4% generation overhead.

Problem

Modern LLMs have two failure modes when extrapolating beyond their pretraining window: (i) position out-of-distribution — low-frequency RoPE components produce rotation angles never seen during training; (ii) attention diffusion — softmax probability mass disperses over irrelevant tokens at long range. Existing zero-shot methods (NTK, YaRN, Self-Extend, DCA) fix a single rescaling factor or grouping size up front, forcing a tradeoff: an aggressive factor degrades short-context fidelity; a conservative one fails at the target long context. No prior method adapts the factor dynamically to sequence length while preserving the base model exactly within its native window.

Method

Jet-Long decomposes attention into two windows per query–key pair:

\[S(q,k) = \begin{cases} \text{RoPE}(x_q,q)^\top \text{RoPE}(x_k,k) & \text{if } q-k \le w_0, \\ \text{RoPE}(x_q,f(q))^\top \text{RoPE}(x_k,f(k)) & \text{if } q-k > w_0. \end{cases}\]

Dynamic group size. The remote mapping is floor-division \(f(x) = \lfloor x/G \rfloor\) with \(G = \max\!\left(1, \lceil L / w_{\text{pretrained}} \rceil\right)\). This is the smallest integer that keeps all remote positions within the pretrained rotation grid, maximizing positional resolution while guaranteeing every remote angle is in-distribution. When \(L \le w_{\text{pretrained}}\), \(G=1\) and the method reduces exactly to the base model.

KV cache invariant. Rather than rewriting cached keys when \(G\) changes during generation, the cache stores uncompressed absolute-position keys. An on-the-fly correction rotation is applied using offsets \(\Delta_q = f(q)-q\), \(\Delta_k = f(k)-k\), exploiting the additive-angle property of RoPE (\(R_a R_b = R_{a+b}\)). This leaves the cache unchanged across length boundaries.

Inclusion–exclusion prefill. Exact distance-based routing without a boolean mask is achieved by combining three FlashAttention calls via inclusion–exclusion with LogSumExp statistics:

\[O_{\text{final}} = \frac{W_A O_A + W_B O_B - W_C O_C}{W_A + W_B - W_C}\]

where \(W_X = \exp(\ell_X - M)\) and \(M = \max(\ell_A, \ell_B, \ell_C)\). All three passes are fused into a single CuTe kernel.

Key Contributions

  • Dynamic bifocal RoPE (Jet-Long): tuning-free zero-shot context extension with a sequence-length-adaptive group size \(G\), reducing to the unmodified base model within the native window.
  • On-the-fly KV correction rotation: maintains a cache-invariant that avoids recomputing or invalidating the KV cache when \(G\) changes, enabling streaming generation across length boundaries.
  • Inclusion–exclusion prefill kernel: fuses three FlashAttention passes into one CuTe kernel, achieving 1.28–1.39Ă— FA2 prefill throughput on H100 past 32K (approaching Hopper-only FA4 at 1.53Ă—) and ≤4% generation overhead at every length.
  • Generalization to hybrid architectures: Jet-Long applies unchanged to Jet-Nemotron (interleaved softmax + linear-attention layers) without retraining.
  • Hyperparameter resilience: the single hyperparameter \(w_0\) (local window size) is shown to be robust across \(\{512, 1024, 2048, 4096\}\), staying within 2 pp of per-row best.

Results

All results are on Qwen3-1.7B/4B/8B-Base (32K native window) extended to up to 128K, against baselines: Base, DNTK, YaRN, DCA, Self-Extend.

  • RULER (avg over 13 tasks, 7 lengths 4K–128K): Jet-Long 73.93 / 83.02 / 86.74 at 1.7B/4B/8B vs. strongest baseline 69.14 (DNTK) / 80.84 (Self-Extend) / 84.71 (Self-Extend); margins +4.79 / +2.18 / +2.03 pp.
  • HELMET-RAG (avg over 4 sub-tasks): Jet-Long 42.28 / 53.61 / 57.34 at 1.7B/4B/8B; best at 4B and 8B, 0.73 pp behind Self-Extend at 1.7B.
  • PG-19 perplexity (geometric mean, lower is better): Jet-Long 11.72 / 9.85 / 8.73 at 1.7B/4B/8B, lowest among all methods; Base collapses at 128K (ppl 71.00 / 104.66 / 79.37).
  • Per-task at 65K: best or tied on 8/13 tasks at 1.7B and 8B; largest leads on MK-NIAH-2 at 1.7B (61.80 vs. 20.00) and MK-NIAH-3 at 8B (70.80 vs. 39.80).
  • Hybrid (Jet-Nemotron): RULER avg lifts from 42.93 → 52.94 (+10.01 pp) at 2B and 42.16 → 53.47 (+11.31 pp) at 4B over 7 lengths.
  • Prefill throughput: fused kernel 1.28–1.39Ă— FA2 past 32K on H100 (FA4 reaches 1.53Ă— with Hopper-specific code); generation ≥0.96Ă— FA2 (≤4% overhead) at every length.

Limitations

  • Addresses only the position-OOD failure mode; attention diffusion is not directly targeted and remains a complementary open problem addressed by architectural alternatives.
  • Evaluated up to 128K tokens; behavior of discrete aliasing vs. frequency interpolation at contexts beyond 128K is noted as unresolved (the ablation shows aliasing advantage shrinks with larger \(G\)).
  • Throughput benchmarks are single-H100, single-batch generation; multi-GPU or batched-generation efficiency not reported.
  • Only one family of base models (Qwen3) and one hybrid (Jet-Nemotron) tested; generalization to other architectures (e.g., Llama, Mistral) is not empirically verified.
  • The \(w_0=0\) ablation confirms the local window is necessary, but optimal \(w_0\) at very long contexts (>128K) is not explored.

Relevance to Agentic AI / LLM Agents

Agentic LLMs that interleave reasoning, planning, and tool use accumulate context (tool traces, execution outputs, multi-turn history) that routinely exceeds 100K tokens — the paper explicitly motivates Jet-Long by this pressure point. A zero-shot, tuning-free extension method that preserves short-context fidelity is particularly valuable for deployed open-weight agents, where a single checkpoint must serve variable-length sessions without per-deployment fine-tuning. The ≤4% generation overhead and cache-invariant design mean agents can operate at extended contexts without throughput regression or cache invalidation on length boundary crossings. For RAG-heavy agentic pipelines, the strong HELMET-RAG results (an application-grounded benchmark) directly validate the method's utility in retrieval-augmented agent workflows.