Jet-Long: Effiziente Langkontexterweiterung mit dynamischer bifokaler RoPE

Zusammenfassung

Moderne große Sprachmodelle (LLMs) werden zunehmend in Langkontext-Anwendungen wie Retrieval-Augmented Generation, Repository-Level Coding und agentischen Workflows eingesetzt, deren akkumulierte Reasoning- und Tool-Traces die Eingabe routinemäßig um eine Größenordnung über das Vortrainingsfenster hinaus verlängern, sodass die Null-Kontext-Erweiterung (Zero-Shot Context Extension) zum dominierenden Deployment-Pfad für Open-Weight-Checkpoints wird. Die meisten existierenden Null-Kontext-Methoden legen im Voraus einen einzelnen Reskalierungsfaktor fest, sodass ein aggressiver Faktor die Kurzkontext-Treue beeinträchtigt, während ein konservativer bei langen Kontexten versagt. Wir schlagen Jet-Long vor, eine tuning-freie Null-Kontext-Methode, die ein lokal RoPE-treues Fenster mit einem langreichweitigen Fenster paart, dessen Reskalierungsfaktor sich dynamisch an die aktuelle Sequenzlänge anpasst – das Basis-Modell wird bei kurzen Eingaben exakt wiederhergestellt, während bei langen sauber extrapoliert wird. Eine Inklusions-Exklusions-Aufmerksamkeitsfusion und eine on-the-fly RoPE-Korrekturrotation machen die bifokale Konstruktion zur Inferenzzeit praktisch kostenlos; in einen einzigen CuTe-Kernel integriert, erreicht der Langkontext-Prefill einen bis zu 1,39-fachen FA2-Durchsatz auf H100 (nahe dem Hopper-only FA4), und die Single-Batch-Generierung verursacht bei jeder Länge einen Overhead von ≤ 4 %. Bei Qwen3-1,7B/4B/8B bis zu 128K Kontext führt Jet-Long RULER um +4,79/+2,18/+2,03 Prozentpunkte gegenüber der stärksten Baseline bei 1,7B/4B/8B an, erzielt die beste Gesamtgenauigkeit auf HELMET-RAG (einem Benchmark, den HELMET als den effizientesten Prädiktor für nachgelagerte Langkontext-Leistung identifizierte) und erreicht die niedrigste PG-19-Perplexität. Jet-Long generalisiert zudem auf hybride Aufmerksamkeitsarchitekturen wie Jet-Nemotron zur weiteren Langkontext-Verbesserung ohne Nachtraining und bleibt hyperparameter-resistent für eine einfache Bereitstellung.

English

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.39times 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.