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

2026-07-08

AI Summary

Existing zero-shot long-context extrapolation methods struggle to simultaneously preserve fidelity on short contexts and ensure stability over long contexts. This work proposes a tuning-free, zero-shot extension approach that introduces a novel dynamic dual-focus RoPE mechanism, integrating a locally faithful attention window with a dynamically scaled distant window. The method further incorporates inclusion-exclusion attention fusion and on-the-fly RoPE correction, enabling efficient and stable context extrapolation with negligible inference overhead. Compatible with hybrid attention architectures, the proposed technique significantly outperforms current baselines on the Qwen3 model family: it achieves state-of-the-art performance on the HELME-RAG benchmark, attains the lowest perplexity on PG-19, delivers 1.39× the inference throughput of FlashAttention-2, and incurs no more than a 4% generation overhead.

Abstract

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. Most existing zero-shot methods fix a single rescaling factor up front, so an aggressive factor sacrifices short-context fidelity while a conservative one breaks down at long contexts. 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, 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.

Problem

• long-context extension
• zero-shot context extension
• RoPE
• large language models
• context window

Innovation

• Dynamic Bifocal RoPE
• Zero-shot Context Extension
• Tuning-free Long-context LLM
• Efficient Attention Merge
• CuTe Kernel Optimization