Jet-Long: 具有动态双焦RoPE的高效长上下文扩展

为了最大化位置分辨率，\(G\) 是保持远程窗口内的位置角度完全落在预训练旋转网格上的最小整数。具体地，令 \(L\) 为当前序列长度，\(w_{\text{pretrained}}\) 为预训练上下文窗口，则分组大小定义为：

\[
G = \max\left(1, \left\lceil \frac{L}{w_{\text{pretrained}}} \right\rceil \right)
\]

对于查询位置 \(q\) 和键位置 \(k\)，远程窗口的位置映射为：
\[
f(x) = \left\lfloor \frac{x}{G} \right\rfloor
\]
这样，所有远程相对位置 \(\Delta = f(q) - f(k)\) 都落在预训练窗口内，即 \(|\Delta| \leq w_{\text{pretrained}}\)。当 \(L \leq w_{\text{pretrained}}\) 时，\(G=1\)，\(f(x)=x\)，退化为标准RoPE。

### 3.2 包含-排除注意力融合

Jet-Long 的注意力计算需要同时处理局部窗口（标准RoPE）和远程窗口（重映射RoPE）。直接分别计算会导致两倍的计算量。我们提出一种包含-排除融合策略，通过三次 FlashAttention 调用来高效合并两个窗口：

1. **全远程注意力**：将所有键-值对视为远程窗口（使用重映射RoPE），计算完整注意力分数。
2. **局部仅-标准RoPE**：仅对局部窗口内的键-值对使用标准RoPE，计算局部注意力分数。
3. **局部仅-重映射RoPE**：对相同的局部窗口键-值对使用重映射RoPE（与步骤1中的远程映射相同），计算局部注意力分数。

然后，最终的注意力输出为：
\[
\text{Attn}_{\text{final}} = \text{Attn}_{\text{remote}} + \text{Attn}_{\text{local, standard}} - \text{Attn}_{\text{local, remap}}
\]
其中后两项仅对局部窗口有效。这种组合等价于对每个查询，局部窗口内的键使用标准RoPE，远程窗口内的键使用重映射RoPE。由于所有操作都是矩阵运算，并且我们可以利用 LogSumExp 技巧来稳定 softmax，因此三次调用可以合并到单个 CuTe 内核中，在前缀填充阶段几乎不引入额外开销。

### 3.3 生成阶段的 KV 缓存修正

在自回归生成阶段，KV 缓存存储了所有已生成位置的键和值，位置坐标使用原始索引。当生成新 token 时，其查询位置 \(q\) 相对于缓存中的键位置 \(k\) 可能属于远程窗口（即 \(q - k > w_0\)）。为了在不修改缓存的情况下应用远程窗口的重映射，我们在每个注意力头中动态地对查询和键应用额外的旋转：

- 对于查询位置 \(q\)，计算所需的旋转偏移 \(\Delta q = f(q) - q\)（即重映射后的位置与原始位置的差）。
- 对于缓存中的每个键位置 \(k\)，计算类似的偏移 \(\Delta k = f(k) - k\)。

然后，在计算注意力分数之前，对查询和键向量分别旋转这些偏移量。由于 RoPE 是线性旋转，这相当于在点积之前对隐向量进行逐位置旋转，并且可以通过融合到 FlashAttention 内核中来实现，而无需修改缓存。这种“即时”修正使得生成阶段的开销仅来自额外的旋转计算，而 KV 缓存保持不变。

由于局部窗口内的键-值对不需要这种修正（它们使用标准RoPE），我们只需对超出局部窗口的键应用旋转即可。这可以通过注意力掩码来实现，进一步减少计算量。

### 3.4 实现细节

我们将 Jet-Long 实现为单个 CuTe 内核，包含以下步骤：
- **前缀填充**：使用包含-排除融合（三次 FlashAttention 调用）计算整个序列的注意力。
- **生成**：对每个新 token，计算其相对位置，对于超出局部窗口的键，在注意力计算中动态应用旋转修正。

内核通过 CUDA 编程优化，利用 Hopper 架构的特性（如 Tensor Core 和 shared memory）来最大化吞吐量。实验表明，相比于标准的 FlashAttention-2，前缀填充加速比达到 1.28×–1.39×（接近使用 Hopper 专用特性的 FlashAttention-4），而生成阶段的开销不超过 4%。

## 4 实验

我们评估 Jet-Long 在 Qwen3-1.7B/4B/8B 模型上的表现，上下文长度扩展至 128K。我们与最先进的零样本基线方法进行比较，包括 PI、NTK、Dynamic NTK、YaRN、Self-Extend 和 DCA。评估基准包括：
- **RULER**：一套包含 13 个合成任务的长上下文检索基准。
- **HELMET-RAG**：来自 HELMET 研究的最佳整体下游性能预测器。
- **PG-19**：长篇文本的困惑度评估。

此外，我们还测试了 Jet-Long 在混合注意力模型 Jet-Nemotron 上的迁移能力，验证其无需重新训练即可推广。

### 4.1 主要结果

在 Qwen3-1.7B 上，Jet-Long 在 RULER 所有任务的平均准确率上领先最强基线（Self-Extend）4.79 个百分点（pp），在 128K 上下文长度上达到 76.3% vs 71.5%。在 Qwen3-4B 和 8B 上，优势分别为 +2.18 pp 和 +2.03 pp。在 HELMET-RAG 上，Jet-Long 取得最佳总体准确率（与最佳基线持平或略优）。在 PG-19 上，Jet-Long 在长上下文范围内的困惑度最低，尤其是在 64K 和 128K 处显著优于其他方法。

### 4.2 鲁棒性与超参数敏感性

Jet-Long 只有一个超参数 \(w_0\)（局部窗口大小），默认设置为训练上下文窗口的大小（例如 32K）。我们发现，在合理范围内（16K–64K）调整 \(w_0\) 对性能影响很小，表明该方法对超参数选择具有鲁棒性。这简化了部署中的调参工作。

### 4.3 迁移到混合注意力架构

将 Jet-Long 应用于 Jet-Nemotron（包含标准注意力层和 NoPE 层的混合），无需修改任何模型权重。实验表明，在 Jet-Nemotron 上，Jet-Long 进一步提升了长上下文性能，与直接使用 NoPE 层相比，检索准确率提高了约 3%，验证了其通用性。

### 4.4 与连续频率插值的对比

我们进行消融实验，比较离散分组（如 Jet-Long）与连续频率插值（如 Dynamic NTK）。结果表明，离散分组在长上下文检索任务上平均高 1.5%，且对短上下文的保持更精确，因为其完全复用了预训练的旋转角度。

### 4.5 效率分析

前缀填充阶段：Jet-Long 的融合内核在 H100 上达到 1.28×–1.39× 的 FlashAttention-2 吞吐量，接近使用 FlashAttention-4 的性能。生成阶段：每条 token 的延迟增加不超过 4%，主要来自动态旋转计算，这部分被高效地融合到内核中。内存占用与标准注意力相同，因为 KV 缓存未修改。

## 5 结论

我们提出了 Jet-Long，一种无调优的零样本上下文扩展方法，通过动态双焦 RoPE 机制，在短输入时完全恢复基础模型行为，在长输入时实现干净的外推。包含-排除注意力融合和即时修正旋转使得推理开销几乎可忽略。在多个模型和基准上，Jet-Long 取得了最先进的性能，并且对单个超参数鲁棒，易于部署。未来工作将探索将 Jet-Long 扩展到更大上下文窗口（如 1M token），以及与其他架构（如状态空间模型）的结合。