Jet-Long: 高效的長上下文擴展與動態雙焦RoPE

摘要

現代大型語言模型（LLM）日益被部署於長上下文應用中，例如檢索增強生成、倉庫級程式碼編寫，以及代理型工作流程——這些場景中累積的推理與工具痕跡常使輸入長度超過預訓練窗口一個數量級，因此零樣本上下文擴展成為開放權重模型的主要部署路徑。現有的大多數零樣本方法會在事前固定單一重新縮放因子，導致激進的因子犧牲短上下文忠實度，而保守的因子則在長上下文時失效。我們提出Jet-Long，一種免調參的零樣本方法，其結合一個局部RoPE忠實窗口與一個長程窗口，後者的重新縮放因子會根據當前序列長度動態調整，從而在短輸入時精確還原基礎模型，並在長輸入時乾淨地進行外推。透過包含-排除注意力合併與即時RoPE校正旋轉，此雙焦點結構在推論時幾乎沒有額外成本；將其融合為單一CuTe內核後，長上下文預填充在H100上可達1.39倍FA2的吞吐量（接近僅限Hopper的FA4），且單批次生成在所有長度下僅增加不超過4%的開銷。在Qwen3-1.7B/4B/8B模型上，於128K上下文範圍內，Jet-Long在1.7B/4B/8B規模下分別比最強基線高出+4.79/+2.18/+2.03個百分點的RULER分數，在HELMET-RAG（被HELMET認定為預測下游長上下文表現最有效率的基準）上達到整體最佳準確率，並取得最低的PG-19困惑度。Jet-Long亦能泛化至混合注意力架構（如Jet-Nemotron），無需重新訓練即可進一步改善長上下文表現，且對超參數具有韌性，便於部署。

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.