Source-linked AI summary
Jet-Long: Efficient Long-Context Extension with Dynamic Bifocal RoPE
Haozhan Tang, Zerui Wang, Yuxian Gu, Song Han, Han Cai
TL;DR
Long-context deployment requires zero-shot extension beyond pretrained windows, but fixed rescaling factors trade short-context fidelity against long-context coverage. Jet-Long dynamically combines local RoPE-faithful and compressed long-range windows, achieving superior RULER, HELMET-RAG, and PG-19 results up to 128K on Qwen3 models.
Problem
Zero-shot context extension must support arbitrary longer inputs because long-context training is expensive and can degrade short-context behavior.
Method
Jet-Long dynamically pairs a local RoPE-faithful window with a discretely grouped long-range window using an analytic, parameter-free rescaling factor.
Results
Jet-Long is best across RULER and PG-19 columns and leads the strongest RULER baseline by 4.79, 2.18, and 2.03 pp across Qwen3 sizes.
Takeaways & Limitations
Jet-Long provides a tuning-free context-extension approach that remains effective across Qwen3 sizes, benchmarks, and hybrid attention architectures.
Takeaways & Limitations
Jet-Long requires a softmax-with-RoPE base, limiting its natural extensions to compatible attention architectures.
Abstract
from arXiv · showhide
Modern LLMs are increasingly deployed in long-context applications such as retrieval-augmented generation, repository-level coding, and agentic workflows whose accumulated reasoning and tool traces routinely push the input an order of magnitude past the pretraining window, making zero-shot context extension the dominant deployment path for open-weight checkpoints. The dominant zero-shot methods (YaRN, Self-Extend, DCA) fix a single rescaling factor up front, so an aggressive factor sacrifices short-context fidelity while a conservative one breaks down at long contexts; recent length-aware variants adapt the mapping, but with a fitted or distance-dependent schedule. We propose Jet-Long, a tuning-free zero-shot method that pairs a local RoPE-faithful window with a long-range window whose rescaling factor adapts dynamically to the current sequence length via a parameter-free analytic schedule, recovering the base model exactly at short inputs while extrapolating cleanly at long ones. An inclusion-exclusion attention merge and an on-the-fly RoPE correction rotation make the bifocal construction essentially free at inference; fused into a single CuTe kernel, long-context prefill reaches up to $1.39\times$ FA2 throughput on H100 (approaching the Hopper-only FA4), and single-batch generation incurs $\le 4\%$ overhead at every length. On Qwen3-1.7B/4B/8B up to 128K context, Jet-Long leads RULER by $+4.79$/$+2.18$/$+2.03$ pp over the strongest baseline at 1.7B/4B/8B, achieves the best overall accuracy on HELMET-RAG (a benchmark identified by HELMET as the most efficient predictor of downstream long-context performance) and attains the lowest PG-19 perplexity. Jet-Long also generalizes to hybrid attention architectures such as Jet-Nemotron for further long-context improvement without retraining, and remains hyperparameter-resilient for ease of deployment.
1. Introduction
Long-context deployments accumulate inputs beyond pretrained windows, while direct long-context training is costly and can degrade short-context behavior. Jet-Long addresses this with a tuning-free bifocal extension whose analytically dynamic compression preserves native-context behavior and supports longer inputs without per-model fitting.
- Motivation: Agentic workflows and coding tasks routinely accumulate 100K+ tokens from reasoning, tool use, source code, execution traces, and tool outputs.
- Motivation: Long-context training remains expensive because attention compute stays quadratic, long-context data is scarce, and fine-tuning can degrade short-context behavior.
- Method: Jet-Long pairs a local RoPE-faithful window with a long-range window whose rescaling factor adapts analytically to sequence length and the pretrained window.
- Method: Unlike YaRN, Self-Extend, and DCA, Jet-Long avoids fixing a single grouping size or factor up front and requires no per-model fitting.
- Method: G=max(1,⌈L/w_pretrained⌉) gives Jet-Long minimal integer remote compression, keeps remote rotations in-distribution, and reproduces the base model exactly within its native context.
2. Related Work
Long-context extrapolation is challenged by position out-of-distribution and attention diffusion, motivating zero-shot remapping and attention-correction strategies. Jet-Long belongs to grouped-position methods but derives its group size dynamically, exactly matching the base model within the pretrained window.
- Motivation: RoPE extrapolation fails through out-of-distribution rotation angles and attention diffusion that disperses probability mass while under-attending to middle-context evidence.These correspond to position-OOD and the U-shaped positional attention bias described for dense softmax attention.
- Zero-shot strategies: Zero-shot extensions address these problems by interpolating RoPE positions or frequencies and by applying attention penalties or temperature scaling.Jet-Long specifically targets position-OOD through dynamic aliasing onto the pretrained rotation grid.
- Prior methods: Frequency-rescaling methods include PI, NTK-aware scaled RoPE, Dynamic NTK, and YaRN, while Self-Extend and DCA reuse in-distribution grouped or chunked positions.Dynamic NTK adapts scaling to current sequence length; Self-Extend combines neighbor and grouped windows, while DCA uses asymmetric cross-chunk indices.
- Length-adaptive remapping: Length-adaptive methods condition remapping on input length, using distance-dependent reuse, fitted regional compression, or logistic growth of Self-Extend’s group size.The cited methods are AdaGroPE, LaMPE, and SELF, respectively.
- Jet-Long’s position mapping: Jet-Long derives group size from current sequence length, uses identity mapping within the pretrained window, and preserves the base model exactly for L≤w_pretrained.Its correction rotation also preserves a base-position KV cache when the group size changes during generation.
3. Methodology
Jet-Long combines a RoPE-faithful local window with an analytically compressed remote window whose grouping adapts to sequence length, recovering the base model within its pretraining window. Its inference machinery preserves the KV cache and implements exact distance routing with FlashAttention-compatible operations.
- Dynamic position mapping: The dynamic group size G = max(1, ⌈L/w_pretrained⌉) applies minimum compression analytically as sequence length L grows.This discrete grouping keeps remote RoPE angles in-distribution and maximizes positional resolution without fitted parameters.
- Bifocal attention: Jet-Long uses a bifocal decomposition with classic RoPE in local window w0 and remapped positions in the remote window.The local branch preserves pretraining behavior exactly, while the remote branch maps positions back into the training range.
- Dynamic position mapping: When L ≤ w_pretrained, f(x) = x, so Jet-Long reduces exactly to the unmodified base model.At longer contexts, the remote mapping compresses positions while the local window remains RoPE-faithful.
- Cache-preserving inference: On-the-fly correction rotations reconstruct remote queries and keys in registers, leaving the physical KV cache unchanged during generation.This avoids rewriting or recomputing the cache when G changes across length boundaries.
- Efficient attention merge: Three FlashAttention calls merged by inclusion–exclusion realize exact distance-based routing without materializing a quadratic attention matrix.The merge uses LogSumExp statistics and retains FlashAttention’s memory efficiency and near-FA2 throughput.
4. Experiments
Experiments on Qwen3 models extended from a 32K native window to 128K show that Jet-Long improves long-context accuracy and perplexity while preserving in-window behavior. It also generalizes to hybrid attention, remains hyperparameter-resilient, and achieves near-baseline generation with improved fused-kernel prefill throughput.
- Main results: 4.79, 2.18, and 2.03 pp are Jet-Long’s RULER leads over the strongest baseline at 1.7B, 4B, and 8B, respectively.Jet-Long is best on every RULER and PG-19 column and best on HELMET-RAG at 4B and 8B.
- Main results: 11.41 / 9.62 / 8.51 PG-19 perplexity at 128K beats all extrapolation methods for Qwen3-1.7B / 4B / 8B, while Base reaches 71.00 / 104.66 / 79.37.Within 32K, Jet-Long is mathematically equivalent to Base; beyond 32K, Base collapses while Jet-Long remains lowest among extrapolation methods.
- Main results: 7.31 pp and 6.07 pp are Jet-Long’s aggregated RULER gains over the strongest baseline at 65,536 tokens for Qwen3-1.7B and 8B.Jet-Long is best or tied on 8 of 13 tasks at both sizes, with large leads on Multi-Key NIAH and Variable Tracking.
- Hybrid architectures: +10.01 pp and +11.31 pp are Jet-Long’s average RULER gains on hybrid Jet-Nemotron-2B and 4B, respectively.At 128K, Jet-Long scores 33.78 and 33.14 versus Base’s 8.54 and 5.65, while matching Base within the 32K window.
- Robustness and mapping: Within 1 pp on per-model averages, every w0 ∈ {512, 1024, 2048, 4096} remains near the per-row best, supporting w0=2048 without deployment tuning.The w0=0 control collapses RULER to near-zero, while w0=8192 incurs a 1.4–2.1 pp gap because the local window consumes more context.
- Inference efficiency: 1.28–1.39× FA2 prefill throughput is achieved past 32K, while generation stays ≥0.96× FA2 at every length with ≤4% residual overhead.The fused CuTe kernel removes multi-launch and per-token correction costs; within 32K, all configurations match FA2 within ±1%.
5. Conclusion
Jet-Long is a tuning-free, zero-shot context-extension method that combines dynamic bifocal RoPE with an inference-efficient fused implementation. It demonstrates strong long-context performance and can extend to additional softmax-with-RoPE and hybrid attention architectures.
- Core method: Jet-Long combines a local RoPE-faithful window with a dynamically rescaled long-range window for tuning-free zero-shot context extension.An inclusion–exclusion attention merge and on-the-fly correction rotation in a fused CuTe kernel make the construction essentially free at inference.
- Results: +4.79/+2.18/+2.03 pp at 1.7B/4B/8B: Jet-Long shows superior RULER performance on Qwen3 models at context lengths up to 128K.The reported model sizes are Qwen3-1.7B/4B/8B.
- Extensions and limitations: Jet-Long addresses position-OOD at the RoPE level, while architectural alternatives naturally target the complementary attention-diffusion failure mode.The method can extend to Multi-head Latent Attention, sparse attention, and architectures interleaving softmax with sparse or linear-attention layers, beyond Jet-Nemotron.
A. Metric definitions
This section defines percentage-point gaps for RULER and HELMET-RAG, distinguishes them from relative percent change, and specifies perplexity evaluation and aggregation for PG-19.
- Percentage points: Percentage-point gaps quantify absolute differences between percentage scores and are used for RULER and HELMET-RAG accuracy comparisons.They differ from relative percent change, (a1 − a2)/a2.
- Perplexity: Perplexity is the exponentiated mean negative log-likelihood per token, with lower values indicating better performance.It is computed from conditional token probabilities under a language model.
- Perplexity: PG-19 uses anchored growing-window evaluation at 4K–128K contexts, scoring the next 1024 tokens and averaging perplexity across 100 evaluation books.The evaluated context lengths are 4K, 8K, 16K, 32K, 64K, 96K, and 128K.
- Geometric-mean aggregation across lengths: The PG-19 Avg column reports the geometric mean across seven context lengths.This aggregates the exponentiated metric by uniformly averaging underlying log-likelihoods across lengths.
B. Full inference efficiency results
Across Qwen3-1.7B, 4B, and 8B on H100, Jet-Long CuTe preserves near-FA2 generation speed while recovering and exceeding FA2 for long-context prefill. The fused kernel removes the generation penalty of unfused Jet-Long across the full sweep.
- Full inference efficiency results: 1.28–1.45× long-context prefill speedups over FA2 are achieved by Jet-Long CuTe across all three Qwen3 sizes.The sweep compares FA2, FA4, Jet-Long (unfused), and Jet-Long CuTe, with speedups measured against FA2 at matching lengths.
- Full inference efficiency results: ≥0.96× generation throughput versus FA2 is maintained by Jet-Long CuTe at every length and model size.Inside the 32K native window, Jet-Long CuTe matches FA2; at longer contexts, it recovers and surpasses FA2 for prefill.
- Full inference efficiency results: The fused CuTe kernel eliminates the real generation cost incurred by the unfused multi-launch Jet-Long variant.FA4 is omitted from generation rows because no H100 generation kernel has been released for it.
C. Baseline configurations
The study compares four zero-shot baselines using fixed configurations across Qwen3-1.7B/4B/8B and evaluation lengths from 4K to 128K. All Qwen3 models share a 32,768-token pretrained context window and RoPE base θ=106.
- Baseline configurations: Four zero-shot baselines use a single configuration each throughout Tables 1–4.These configurations are summarized in Table 9.
- Baseline configurations: The configurations remain constant across Qwen3-1.7B/4B/8B and evaluation lengths of 4K–128K.Table 9 specifies the baseline hyperparameters used for these comparisons.
- Baseline configurations: All Qwen3 models have a pretrained context window of 32,768 tokens and RoPE base θ=106.These model properties apply to Qwen3-1.7B, Qwen3-4B, and Qwen3-8B.