Jet-Long: Extensão Eficiente de Contexto Longo com RoPE Bifocal Dinâmico

Resumo

Modelos de linguagem grandes (LLMs) modernos são cada vez mais implantados em aplicações de contexto longo, como geração aumentada por recuperação, codificação em nível de repositório e fluxos de trabalho agentivos, cujos traços acumulados de raciocínio e ferramentas rotineiramente empurram a entrada para uma ordem de grandeza além da janela de pré-treinamento, tornando a extensão de contexto zero-shot a via de implantação dominante para checkpoints de pesos abertos. A maioria dos métodos zero-shot existentes fixa um único fator de reescalonamento de antemão, de modo que um fator agressivo sacrifica a fidelidade de contexto curto, enquanto um conservador falha em contextos longos. Propomos Jet-Long, um método zero-shot sem ajuste fino que emparelha uma janela local fiel à RoPE com uma janela de longo alcance cujo fator de reescalonamento se adapta dinamicamente ao comprimento atual da sequência, recuperando o modelo base exatamente em entradas curtas e extrapolando de forma limpa em entradas longas. Uma fusão de atenção por inclusão-exclusão e uma rotação de correção de RoPE em tempo real tornam a construção bifocal essencialmente gratuita na inferência; fundida em um único kernel CuTe, a pré-carga de contexto longo atinge até 1,39 vezes o throughput FA2 em H100 (aproximando-se do FA4 apenas para Hopper), e a geração de lote único incorre em ≤ 4% de sobrecarga em todos os comprimentos. No Qwen3-1.7B/4B/8B com até 128K de contexto, Jet-Long lidera o RULER em +4,79/+2,18/+2,03 pp sobre a linha de base mais forte em 1.7B/4B/8B, alcança a melhor precisão geral no HELMET-RAG (um benchmark identificado pelo HELMET como o preditor mais eficiente de desempenho downstream em contexto longo) e obtém a menor perplexidade no PG-19. Jet-Long também se generaliza para arquiteturas de atenção híbridas, como Jet-Nemotron, para melhoria adicional de contexto longo sem re-treinamento, e permanece resiliente a hiperparâmetros para facilidade de implantação.

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.