Jet-Long : Extension efficace de long contexte avec RoPE bifocale dynamique

Résumé

Les LLM modernes sont de plus en plus déployés dans des applications à long contexte telles que la génération augmentée par récupération, le codage au niveau du dépôt et les flux de travail agentiques dont les traces de raisonnement et d’outils accumulées poussent couramment l’entrée à un ordre de grandeur au-delà de la fenêtre de pré-entraînement, faisant de l’extension de contexte zero-shot la voie de déploiement dominante pour les points de contrôle à poids ouverts. La plupart des méthodes zero-shot existantes fixent un seul facteur de redimensionnement à l’avance, de sorte qu’un facteur agressif sacrifie la fidélité en contexte court tandis qu’un facteur conservateur échoue en contexte long. Nous proposons Jet-Long, une méthode zero-shot sans réglage qui associe une fenêtre locale fidèle à RoPE à une fenêtre longue dont le facteur de redimensionnement s’adapte dynamiquement à la longueur de séquence courante, retrouvant exactement le modèle de base pour les entrées courtes tout en extrapolant proprement pour les entrées longues. Une fusion d’attention par inclusion-exclusion et une rotation de correction RoPE en temps réel rendent la construction bifocale essentiellement gratuite à l’inférence ; fusionné en un seul noyau CuTe, le préremplissage en long contexte atteint jusqu’à 1,39 fois le débit FA2 sur H100 (approchant le FA4 uniquement Hopper), et la génération mono-lot induit un surcoût ≤ 4 % à chaque longueur. Sur Qwen3-1,7B/4B/8B jusqu’à 128K de contexte, Jet-Long dépasse RULER de +4,79/+2,18/+2,03 points de pourcentage par rapport à la baseline la plus forte à 1,7B/4B/8B, atteint la meilleure précision globale sur HELMET-RAG (un benchmark identifié par HELMET comme le prédicteur le plus efficace de la performance en long contexte aval) et obtient la plus faible perplexité PG-19. Jet-Long se généralise également aux architectures d’attention hybrides telles que Jet-Nemotron pour une amélioration supplémentaire du long contexte sans réentraînement, et reste résistant aux hyperparamètres pour faciliter le déploiement.

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.