Jet-Long: 使用动态双焦点RoPE的高效长上下文扩展

TL;DR Jet-Long 提出零样本、免训练的上下文扩展方法，结合局部 RoPE 忠实窗口与动态缩放长程窗口，以极低推理开销实现短上下文精确复现和长上下文稳定外推，在 RULER、HELMET-RAG 等基准上大幅超越现有方案。

问题

领域关注点

LLM 正大规模部署于长上下文场景——RAG、仓库级编码、智能体工作流等——输入长度常超出预训练窗口一个数量级。零样本上下文扩展（zero-shot context extension）成为开源权重模型的主流部署路径，无需重新训练即可扩展窗口。

现有方法局限

现有零样本方案（如 YaRN、Self-Extend、DCA）预先固定单一缩放因子，导致两难：激进因子损害短上下文保真度，保守因子在长序列上失效。近期长度感知映射（如 LAMPE）引入距离依赖或拟合调度，但仍依赖额外校准或非解析设计，缺乏真正的动态自适应。此外，多数方法并未将推理效率纳入核心设计，长上下文预填充时 GPU 吞吐下降明显。

技术挑战与重要性

RoPE 位置外推的核心难点在于：

1. 序列长度超越训练分布后，位置编码出现 OOD；
2. 长上下文下注意力分数趋于扩散，位置偏差加重。 零样本扩展需同时满足 短上下文无损还原 与 长上下文稳定外推，且不能引入过高的推理开销。业界对免调优、高性能的方法需求迫切，因为重训练大模型成本高昂，而动态缩放策略能兼顾不同长度输入，是工程落地的关键。

行业类比

类似边缘计算中的自适应模型压缩——根据设备算力动态选择精度，而非一刀切裁剪，确保低延迟时不丢太多精度，高负载时仍能运行。

核心洞察

• 动态双焦旋转位置编码通过长度感知的解析调度动态调整外推因子，从根本上消除了固定因子在短上下文保真度与长上下文外推之间的固有矛盾。不同于 YaRN、Self-Extend 等固定因子方法，Jet-Long 在序列较短时保持原始 RoPE，完全恢复基础模型行为；随着序列增长，外推因子平滑增加，无需手动设定折中值，零样本性能在全长度范围保持最优。
• 包含排除注意力融合与即时 RoPE 校正旋转使双焦注意力的构建几乎零开销，单批次生成在所有长度下开销不超过 4%，长上下文预填充吞吐量可达 FlashAttention-2 的 1.39 倍，逼近仅限 Hopper 架构的 FlashAttention-4。与需要离线校准或显著增加计算量的方法不同，该设计通过单个 CuTe 核融合实现高效推理，为实际部署扫清障碍。
• 作为无调优零样本方法，Jet-Long 可直接应用于任何开放权重检查点，且对超参数不敏感，大幅降低了长上下文部署的门槛。它还泛化到混合注意力架构（如 Jet-Nemotron），无需重新训练即可进一步提升长序列建模能力，显示出较强的通用性和工程友好性。

方法

Jet-Long 是一种免训练的零样本长上下文扩展方法，核心是将局部的原始 RoPE 窗口与长程动态缩放窗口相结合，使模型在短输入时完全恢复基座性能，同时在长输入上平滑外推。

输入与基础设定

• 输入为序列长度 L 以及预训练的 RoPE 基础频率配置。
• 方法假设模型在短上下文（预训练窗口内）已充分学习，只需处理超出预训练窗口的输入。

关键模块

1. 双焦窗口划分

   • 局部窗口（大小为 w₀）：使用未经修改的 RoPE，保持对短上下文的忠实度。
   • 长程窗口：对其应用一个随序列长度动态变化的外推因子 s(L)，该因子由无参数解析调度给出，无需拟合。

2. 动态外推因子

   • 当 L 较小时，s(L)=1（无缩放），模型行为与原模型完全一致；当 L 增大时，s(L) 逐渐增大，实现平滑外推，避免固定因子在处理不同长度时权衡的难题。

3. 注意力合并与校正旋转

   • 包含–排除注意力合并：将两个窗口的注意力输出高效融合，避免冗余计算。
   • 即时 RoPE 校正旋转：对缓存的键向量施加旋转变换，使长程窗口的键与局部窗口的查询在点积时对齐，保证合并后的注意力分布正确。

4. 高效内核实现

   • 整个双焦注意力机制被融合进一个 CuTe 内核，在 H100 GPU 上长上下文预填充吞吐量最高可达 FlashAttention-2 的 1.39×，单批次生成开销 ≤ 4%。

输出

• 输出为适用于当前序列的注意力输出，无论在预训练窗口内还是远超窗口，都能保持高质量。

与同类方法的差异

与 YaRN、Self-Extend 等固定单一缩放因子的方法不同，Jet-Long 利用动态因子同时保持短上下文保真度和长上下文外推能力，且无需任何调参或训练，是一种即插即用的部署方案。

实验

实验设计

评估基于 Qwen3-1.7B/4B/8B 三个规模，上下文长度扩展至 128K，对比基线包括 YaRN、Self-Extend、DCA 等主流零样本方法。使用三组长上下文基准：RULER（综合能力探针）、HELMET-RAG（被 HELMET 验证为下游长上下文性能最高效预测器）以及 PG-19 困惑度。推理效率在 H100 GPU 上测量，与 FlashAttention-2 (FA2) 及 Hopper 专用 FA4 对比。

关键发现

• 长上下文精度：Jet-Long 在 RULER 上全面领先，1.7B/4B/8B 分别比最强基线高 +4.79/+2.18/+2.03 pp；HELMET-RAG 取得最佳总体准确率；PG-19 困惑度最低。
• 推理效率：通过 inclusion–exclusion 注意力合并与 on-the-fly RoPE 纠正旋转，长上下文 prefill 吞吐达 FA2 的 1.39×（接近 FA4）；单 batch 生成所有长度下开销 ≤4%。
• 消融验证：局部窗口大小 w0 影响短上下文保真度，频率插值与位置别名策略各有优势；动态外推因子解析调度对超参数不敏感，易于部署。

基线对比解读

固定重缩放因子方法（如 YaRN）面临短上下文退化与长上下文崩溃的取舍；Jet-Long 的双焦 RoPE 设计用局部忠实窗口保留短序列基模型精度，长程窗口随序列长度自适应缩放，从根本上解决了这一矛盾。与长度感知变体（LAMPE 等）不同，Jet-Long 的无参数解析调度无需拟合或距离依赖映射，部署更简单。推理融合 kernel 将双流程开销降至几乎为零，使其成为实用的长上下文零样本扩展方案。

行业影响

落地场景

Jet-Long 提供了一种零样本、无需微调的长上下文扩展方法，直接作用于现成的 open-weight 模型，尤其适合以下产品形态：

• 检索增强生成 (RAG)：处理超长文档或大量检索片段时，能维持高召回和回答质量，适合法律文书分析、金融研报解读。
• 代理工作流 (Agentic Workflow)：长链工具调用、多步推理产生的累积上下文常超过预训练窗口，Jet-Long 可避免中途遗忘或性能退化。
• 仓库级代码理解：大型代码库的全局分析、跨文件重构建议，需要模型在 128K token 以上仍保持精确的代码语义捕捉。
• 多轮对话系统：客服、教育辅导等场景，长对话历史能保持连贯性而不必频繁裁剪或摘要。

商业价值

• 降本：无需额外训练或数据准备，可直接复用已有模型权重，极大降低长上下文适配成本；同时其单 batch 生成开销仅增加 ≤4%，长上下文预填充在 H100 上可达 1.39× FlashAttention-2 吞吐（接近仅 Hopper 架构可用的 FA4），意味着在同等硬件下能服务更多用户或更长的输入。
• 增收/体验提升：在电商问答、内容审核等业务中，长上下文保持高准确率可减少人工复核成本，加快响应速度，提升用户留存和转化；对于代码助手产品，能覆盖更复杂的开发场景，拓宽付费墙。
• 易部署：超参数鲁棒，无需针对不同序列长度调 rescale 因子，降低维护复杂度，适合大规模线上服务。

与现有工作流的集成

Jet-Long 作为零样本 patch，可无缝嵌入现有推理框架：

• 推理引擎集成：作者提供单个 CuTe 内核实现，可将 bifocal attention（局部 RoPE-忠实窗口 + 动态 rescale 的长程窗口）直接替换标准 attention。对于使用 vLLM / TensorRT-LLM 的团队，可通过自定义 attention 算子或插件方式加载，无需修改模型权重。
• 与 RAG Pipeline 配合：在检索后，将检索结果拼接成长 prompt 送入 Jet-Long 增强的 LLM，避免因 prompt 超长导致位置编码失真。
• 混合架构扩展：论文展示了与 Jet-Nemotron 等 hybrid attention 架构的兼容性，说明其不只局限于纯 dense attention，可适用于更前沿的稀疏或混合注意力模型。

具体场景示例：

1. 电商平台多模态搜索：用户上传商品图片并附带多段文字描述，系统调用多个模块（图像理解、文本检索、属性抽取）生成上千 token 的上下文，由部署了 Jet-Long 的聊天模型统一整合，生成精准的推荐理由和对比说明，提高购买转化率。
2. 医疗临床决策支持：电子病历、检验报告、文献摘要等混合长文本输入，要求模型在 128K 窗口内准确定位信息。使用 Jet-Long 增强的医疗 LLM 可在不重新训练的情况下直接处理全量病历，给出诊断建议，降低漏诊风险，同时减少因上下文截断造成的重复训练成本。

局限

• 仅评估了 Qwen3 系列模型：实验集中在 Qwen3-1.7B/4B/8B，缺少在 Llama、Gemma 等主流架构上的泛化验证。不同模型的预训练上下文和 RoPE 配置（如基频 or 旋转方式）可能影响双焦窗口的动态缩放行为，尤其当基模型使用非标准位置编码时，Jet-Long 的零样本扩展能力可能下降。
• 超长上下文（>128K）未覆盖：论文仅在 128K 内测试，但当前应用已向 256K 甚至 1M tokens 演进。局部窗口大小固定，长距离窗口的缩放因子随序列长度线性增长，在极长序列下易导致位置失真累积或注意力偏向长程，削弱检索与推理精度。
• 效率增益依赖算子级融合与硬件：预填充吞吐提升通过 CuTe 编写的融合 kernel 实现，在 H100 (Hopper) 上接近 FA4，但该优化与硬件架构强耦合。迁移至其他 GPU（如 Ampere）需重写算子并可能损失部分加速，且生成阶段的 4% 开销在极低延迟场景仍有优化空间。