Skip to content

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

🕒 Published (v1): 2026-07-08 06:23 UTC · Source: Arxiv · Venue: ACL 2026 · link

Why this paper was selected

Dynamic RoPE for long agentic context; directly addresses harness token-budget and context management

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 to extend LLM context windows far beyond the pretraining limit by dynamically setting a position-aliasing group size \(G = \max(1, \lceil L / w_\text{pretrained} \rceil)\) that adapts to the current sequence length, pairing a RoPE-faithful local window with a dynamically rescaled remote window. It avoids KV-cache rewrites via an on-the-fly correction rotation and achieves exact two-window attention via inclusion–exclusion over three FlashAttention calls fused into a single CuTe kernel. On Qwen3 1.7B/4B/8B at 128K context it outperforms all zero-shot baselines on RULER and PG-19 perplexity while adding ≤4% generation latency over FA2.

Problem

Deployed LLMs are pretrained at moderate windows (4K–32K tokens) but must handle inputs an order of magnitude longer in RAG, repository coding, and agentic workflows. Zero-shot extension methods (YaRN, Self-Extend, DCA) fix a single rescaling factor at deployment time, forcing a tradeoff: an aggressive factor sacrifices short-context fidelity; a conservative one breaks at long context. Length-adaptive variants (AdaGroPE, LaMPE) relax this but require per-model fitted schedules. No prior method simultaneously recovers the base model exactly within the pretraining window, keeps all remote rotation angles on the model's discrete training grid, and achieves this without any fitted parameters or KV-cache recomputation across length boundaries.

Method

Dynamic bifocal RoPE (Jet-Long). For a query at position \(q\) and key at \(k\), attention scores use two regimes: $\(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}\)$ The remote mapping is floor-division aliasing \(f(x) = \lfloor x / G \rfloor\) with group size \(G = \max(1, \lceil L / w_\text{pretrained} \rceil)\), the smallest integer that keeps every aliased position inside the pretrained rotation grid. When \(L \le w_\text{pretrained}\), \(G=1\) and the method is identical to the base model.

KV-cache stability. The cache stores only base-position keys (\(\text{RoPE}(x_k, k)\)). When \(G\) changes during generation, correction offsets \(\Delta_q = \lfloor q/G \rfloor - q\) and \(\Delta_k = \lfloor k/G \rfloor - k\) are applied on-the-fly to the active query and cached keys via RoPE's additive angle composition \(R_a R_b = R_{a+b}\); the physical cache is never rewritten.

Inclusion–exclusion prefill. Three FlashAttention calls—(A) sliding-window with base RoPE, (B) full causal with remote mapping, (C) sliding-window with remote mapping—are combined via LogSumExp weights \(W_X = \exp(\ell_X - M)\) where \(M = \max(\ell_A, \ell_B, \ell_C)\): $\(O_\text{final} = \frac{W_A O_A + W_B O_B - W_C O_C}{W_A + W_B - W_C}\)$ The subtraction cancels remote-rotated local-window contributions exactly. All three calls are fused into a single CuTe kernel; accumulation is in FP32 to prevent catastrophic cancellation.

Key Contributions

  • Parameter-free dynamic group size \(G\): analytically derived from \(L\) and \(w_\text{pretrained}\), no per-model fitting, reduces exactly to the base model at \(L \le w_\text{pretrained}\).
  • On-the-fly KV correction rotation: enables streaming generation across length boundaries without cache invalidation or recomputation.
  • Inclusion–exclusion prefill kernel: exact distance-based routing without a boolean mask matrix; fused CuTe kernel reaches 1.28–1.39Ă— FA2 prefill throughput on H100 (approaching Hopper-only FA4), ≤4% overhead at generation.
  • Hyperparameter robustness: single hyperparameter \(w_0 \in \{512, 1024, 2048, 4096\}\) stays within 2 pp of the per-row best across all tested lengths.
  • Architecture generalization: transfers unchanged to hybrid softmax+linear-attention architectures (Jet-Nemotron) without retraining.

Results

  • RULER (average over 13 tasks, 4K–128K): Jet-Long leads the strongest zero-shot baseline by +4.79 pp (1.7B), +2.18 pp (4B), +2.03 pp (8B); leads by +7.31 pp (1.7B) and +6.07 pp (8B) at \(L=65536\) specifically.
  • HELMET-RAG: best overall at 4B and 8B; 0.73 pp behind Self-Extend at 1.7B; average 51.08 vs 50.95 (Self-Extend) vs 49.93 (DCA).
  • PG-19 perplexity: lowest geometric-mean perplexity across all three model sizes; at 128K, 11.41/9.62/8.51 for 1.7B/4B/8B vs. bare base collapse to 71.00/104.66/79.37.
  • Hybrid Jet-Nemotron: +10.01 pp average RULER at 2B, +11.31 pp at 4B; at 128K, +25.24 pp (2B) and +27.49 pp (4B) over the bare hybrid base.
  • Prefill throughput: 1.28–1.39Ă— FA2 on H100 (fused CuTe kernel); ≤4% generation overhead at every tested length.
  • Ablation—aliasing vs. interpolation: discrete aliasing wins by 6.99 pp at 64K and 4.30 pp at 128K on Qwen3-1.7B; advantage narrows at extreme lengths on specific tasks.

Limitations

  • Evaluated only on Qwen3 (1.7B/4B/8B) and Jet-Nemotron; generalization to other model families (Llama, Mistral, Phi, etc.) is not demonstrated in the text.
  • HELMET evaluation is restricted to the RAG sub-benchmark; other HELMET sub-tasks are not reported.
  • The aliasing-vs.-interpolation ablation suggests that at lengths beyond 128K, frequency interpolation may overtake discrete aliasing on some tasks; the method's behavior at extreme (>128K) contexts is not characterized.
  • All experiments run on a single H100; multi-GPU and ring-attention integration are not discussed.
  • The inclusion–exclusion kernel requires FP32 accumulation to prevent cancellation; half-precision pipelines may need special handling.
  • No chat/instruction-tuned model evaluation; results are on base checkpoints only.

Relevance to Harnesses / Meta-Harnesses

Agentic harnesses accumulate long contexts across reasoning traces, tool call results, retrieved documents, and multi-turn histories — exactly the 100K+ token inputs Jet-Long targets. A harness deploying open-weight LLMs can drop Jet-Long in as a zero-shot, tuning-free inference wrapper with no checkpoint changes, no fine-tuning budget, and near-FA2 throughput, directly extending the effective context budget available to each agent invocation. The KV-cache correction mechanism — which preserves cache validity across decode steps even as \(G\) changes — is particularly relevant for generation-heavy harness loops (streaming tool outputs, partial reasoning) where cache recomputation would stall throughput. Harness designers evaluating HELMET-RAG as a proxy for RAG-heavy workloads have a validated methodology here for selecting long-context strategies without running costly full downstream evals.