Jet-Long: Extensión eficiente de contexto largo con RoPE bifocal dinámico

Resumen

Los LLMs modernos se implementan cada vez más en aplicaciones de contexto largo, como generación aumentada por recuperación, codificación a nivel de repositorio y flujos de trabajo agentivos cuyos registros acumulados de razonamiento y uso de herramientas empujan rutinariamente la entrada un orden de magnitud más allá de la ventana de preentrenamiento, lo que convierte la extensión de contexto cero-shot en la ruta de implementación dominante para los puntos de control de pesos abiertos. La mayoría de los métodos cero-shot existentes fijan un único factor de reescalado de antemano, por lo que un factor agresivo sacrifica la fidelidad en contextos cortos, mientras que uno conservador falla en contextos largos. Proponemos Jet-Long, un método cero-shot sin ajuste que empareja una ventana local fiel a RoPE con una ventana de largo alcance cuyo factor de reescalado se adapta dinámicamente a la longitud actual de la secuencia, recuperando exactamente el modelo base en entradas cortas y extrapolando limpiamente en entradas largas. Una fusión de atención por inclusión-exclusión y una rotación de corrección RoPE sobre la marcha hacen que la construcción bifocal sea esencialmente gratuita en inferencia; fusionada en un solo kernel CuTe, la precarga de contexto largo alcanza hasta 1.39 veces el rendimiento de FA2 en H100 (acercándose al FA4 exclusivo de Hopper), y la generación de un solo lote incurre en una sobrecarga ≤4% en cualquier longitud. En Qwen3-1.7B/4B/8B con contexto de hasta 128K, Jet-Long supera a RULER en +4.79/+2.18/+2.03 pp sobre la línea base más fuerte en 1.7B/4B/8B, logra la mejor precisión global en HELMET-RAG (un punto de referencia identificado por HELMET como el predictor más eficiente del rendimiento en contextos largos aguas abajo) y alcanza la menor perplejidad en PG-19. Jet-Long también se generaliza a arquitecturas de atención híbridas, como Jet-Nemotron, para mejorar aún más el contexto largo sin reentrenamiento, y se mantiene resistente a hiperparámetros para facilitar su implementación.

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.