长上下文技术完全指南

核心论文：

RoFormer: Enhanced Transformer with Rotary Position Embedding（Su et al., 2021）arXiv:2104.09864

Extending Context Window of Large Language Models via Positional Interpolation（Chen et al., Meta, 2023）arXiv:2306.15595

YaRN: Efficient Context Window Extension of Large Language Models（Peng et al., EleutherAI, 2023）arXiv:2309.00071

FlashAttention-2（Dao, 2023）arXiv:2307.08691

Efficient Streaming Language Models with Attention Sinks（Xiao et al., MIT/Meta, 2023）arXiv:2309.17453

Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context（Team Google, 2024）arXiv:2403.05530

LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens（Ding et al., Microsoft, 2024）arXiv:2402.13753

一、长上下文的重要性与挑战
二、RoPE 的外推问题
三、位置插值（Position Interpolation）
四、YaRN：更精细的 RoPE 外推方案
五、LLaMA 3 的长上下文扩展：实践案例
六、FlashAttention 与长上下文：IO 优化基础
七、KV Cache 优化：处理长上下文的推理瓶颈
八、Gemini 1.5 的 100 万 Token 上下文
九、LongRoPE：超越 200 万 Token
十、长上下文评估方法
十一、长上下文的实际工程注意事项

一、长上下文的重要性与挑战

1.1 为什么需要长上下文

现代 LLM 应用中，长上下文能力是核心竞争力之一：

应用场景	所需上下文长度
普通对话	1K-8K tokens
长文档摘要	32K-128K tokens
整本书分析	128K-1M tokens
代码库理解（Claude Code）	100K-200K tokens
完整学术论文	50K-200K tokens
小时级视频（多模态）	500K-1M tokens

以 Claude Code 为例：一个中等规模的 Java 代码库（100 个文件，每个文件 500 行）大约有 200-400K tokens。如果上下文窗口不够，Agent 就无法同时”看到”整个代码库，必须反复读取文件，效率大幅下降。

1.2 标准 Transformer 为什么难以处理长序列

问题 1：注意力计算复杂度 $O(n^2)$

对于序列长度 $n$，注意力矩阵 $A \in \mathbb{R}^{n \times n}$ 的计算和存储复杂度都是 $O(n^2)$。

具体数字：

$n = 4096$（4K）：矩阵约 32MB（float16），计算可接受
$n = 32768$（32K）：矩阵约 2GB，开始困难
$n = 131072$（128K）：矩阵约 32GB，超过单张 GPU 显存

这就是为什么在 FlashAttention 出现之前，长上下文训练是不可行的。

问题 2：位置编码的外推能力（Position Extrapolation）

所有 LLM 都有一个”训练上下文长度”——在预训练时见过的最长序列。如果推理时序列超过这个长度，模型的位置编码需要外推到未见过的位置，通常会导致性能急剧下降。

这是长上下文的核心技术挑战：如何让模型的位置编码能够有效外推到更长的序列？

问题 3：KV Cache 的显存压力

在自回归推理时，需要缓存所有历史 token 的 Key 和 Value 向量（KV Cache）。对于 LLaMA 3 70B（80 层，GQA $h_{KV}=8$，$d_k=128$），上下文长度为 128K 时：

$$\text{KV Cache} = 2 \times 128000 \times 80 \times 8 \times 128 \times 2 \text{ bytes} \approx 42\text{ GB}$$

这还不算激活值。总显存需求（参数 + KV Cache）超过 100GB，需要多张 GPU。

二、RoPE 的外推问题

2.1 RoPE 的工作原理回顾

RoPE（Rotary Position Embedding，详见第 01 篇）通过旋转矩阵将位置信息注入 Query 和 Key 向量。对于位置 $m$ 的 token，其 Query 向量被旋转角度 $m\theta_i$（第 $i$ 对维度）：

$$q_m^{(2i, 2i+1)} = \begin{pmatrix} \cos m\theta_i & -\sin m\theta_i \\ \sin m\theta_i & \cos m\theta_i \end{pmatrix} q_m^{(2i, 2i+1)}$$

其中 $\theta_i = \text{base}^{-2i/d}$，LLaMA 原版使用 $\text{base} = 10000$。

注意力点积的性质：

$$q_m^T k_n = \text{Re}\left[\sum_{i=0}^{d/2-1} q_m^{(i)} \overline{k_n^{(i)}} e^{i(m-n)\theta_i}\right]$$

这说明注意力分数只依赖相对位置 $(m-n)$，这是 RoPE 的核心优点。

2.2 为什么 RoPE 难以直接外推

问题： 当推理时遇到的相对位置 $(m-n)$ 超过训练时见过的最大值时，$e^{i(m-n)\theta_i}$ 的旋转量超出了模型”见过”的范围，导致注意力分数的分布与训练时不同。

更精确地说：对于高频维度（$i$ 较小），$\theta_i$ 较大，一个小的位置增量就会导致大的旋转角度变化；对于低频维度（$i$ 较大），$\theta_i$ 较小，即使位置很远，旋转角度变化也较小。

当上下文长度超过训练长度时：

高频维度： 已经多次”转满一圈”（旋转角度 > $2\pi$），在训练内可能也转了很多圈，外推时行为相对稳定
低频维度： 在训练内只旋转了很小的角度，外推时这些维度第一次遇到大旋转量，模型没有见过

实验结果： LLaMA（训练长度 2048）在长度 2049 处的 perplexity 就已经开始上升，到 4096 时性能大幅下降。

三、位置插值（Position Interpolation）

论文： Extending Context Window of Large Language Models via Positional Interpolation（Chen et al., Meta, 2023）arXiv:2306.15595

3.1 核心思路

如果直接外推（extrapolation）不可行，那就通过压缩（interpolation）：将原始的 $[0, n]$ 位置范围映射到 $[0, L]$（训练长度），这样推理时的位置索引永远不会超出训练范围。

数学上，将位置索引从 $m$ 替换为 $m/s$（其中 $s = n/L$ 是缩放因子）：

$$m \rightarrow m' = m / s$$

这样，对于训练长度 $L = 2048$ 的模型，推理时遇到位置 $m = 4096$ 会被映射到 $m’ = 2048$——模型以为它在位置 2048，而实际上处理的是位置 4096 的 token。

直觉： 就像把一把卷尺压缩——原来 1 米的刻度，现在只占 0.5 米，测量范围翻倍，但精度也相应降低。

3.2 效果与代价

优点： 相比直接外推，位置插值可以平滑地扩展上下文长度。Meta 发现，只需少量的长上下文微调（约 1000 步，使用 $s=16$ 的长文档），就能让 LLaMA 7B 从 2K 扩展到 32K 上下文。

代价： 插值会”压缩”位置信息——相邻两个 token 的位置差从 1 变成 $1/s$，模型区分近距离 token 的位置关系的能力下降（短距离分辨率降低）。

对于 $s=8$（8 倍扩展），相邻两个 token 的旋转角度差变为原来的 1/8，对于高频维度来说这基本没问题（角度仍然可区分），但对于低频维度，近距离 token 几乎无法区分位置。

四、YaRN：更精细的 RoPE 外推方案

论文： YaRN: Efficient Context Window Extension of Large Language Models（Peng et al., EleutherAI, 2023）arXiv:2309.00071

4.1 YaRN 的核心洞察

位置插值的问题在于：它对所有频率维度使用相同的缩放因子 $s$。但不同频率维度的实际需求是不同的：

高频维度（$\theta_i$ 大，$i$ 小）： 在训练长度内，这些维度已经多次旋转满圈，外推时的行为与插值时类似，不需要压缩——直接外推即可
低频维度（$\theta_i$ 小，$i$ 大）： 这些维度在训练内旋转角度很小，直接外推到长序列时遇到未见过的大旋转量，需要插值处理

YaRN 的思路：根据维度的频率，差异化地处理外推和插值。

4.2 YaRN 的实现：NTK-Aware Interpolation

NTK（Neural Tangent Kernel）-Aware 缩放： 将 RoPE 的 base 值从 $\theta_0 = 10000$ 缩放：

$$\theta_i' = \text{base}'^{-2i/d}, \quad \text{base}' = \text{base} \cdot s^{d/(d-2)}$$

其中 $s$ 是上下文扩展倍数。

这个修改的效果：

对于高频维度（$i$ 小）：几乎不改变旋转角度
对于低频维度（$i$ 大）：将 base 值放大，相当于缩小旋转角度，类似插值效果

进一步优化：YaRN 的 NTK-by-Parts

YaRN 更进一步，对维度划分三个区域：

外推区（high frequency）： 维度旋转波长远小于训练上下文，直接外推
插值区（low frequency）： 维度旋转波长远大于训练上下文，做插值（缩放 $1/s$）
过渡区： 线性插值两种方案

数学公式： 对于维度 $i$，YaRN 的有效旋转角度为：

$$\theta_i' = \theta_i \cdot g(i, s)$$

其中 $g(i, s)$ 是从”不缩放”（外推区）到”缩放 $1/s$”（插值区）的线性过渡函数。

注意力温度缩放（Attention Temperature Scaling）：

YaRN 还发现，位置插值会改变注意力分数的分布（因为 Query 和 Key 的旋转角度变小了，点积的范数减小），需要对注意力分数乘以一个温度系数 $t \in [0.1, 0.2]$ 来调整：

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k} \cdot t}\right) V$$

这个细节对长上下文性能有显著影响（相差约 5% 的 passkey retrieval 准确率）。

4.3 YaRN 的效果

在 LLaMA 2 13B（训练长度 4096）上：

方法	扩展到 128K 后的 passkey retrieval 准确率
无扩展（直接使用）	约 0%（完全失效）
位置插值	约 45%
YaRN（NTK-by-Parts）	约 98%

YaRN 仅需约 400 步的长上下文微调（使用 128K 上下文的样本），就能达到接近完美的长距离检索能力。

五、LLaMA 3 的长上下文扩展：实践案例

来源： The Llama 3 Herd of Models（arXiv:2407.21783）

5.1 训练策略

LLaMA 3 的预训练分为两个阶段：

阶段一（标准预训练）： 使用 8K 上下文长度进行大部分训练（约 14.5T tokens）。较短的上下文有利于提高训练吞吐量（注意力计算量随上下文长度二次方增长）。

阶段二（长上下文扩展）： 在主预训练结束后，进行一个连续预训练阶段，逐步扩展上下文长度：

8K → 16K → 32K → 64K → 128K（每个阶段约数百 billion tokens）

关键技术：

RoPE theta 从 500K：LLaMA 3 将 RoPE 的 base 值从标准的 10K 提高到 500K，这大幅提高了低频维度的旋转周期，使其不需要额外的位置插值就能处理更长序列
逐步扩展：不是一步跳到 128K，而是逐步扩展，每步都用相应长度的数据训练一段时间

数据准备： 长上下文训练需要专门的长文档数据（普通的 web 数据通常被截断到 8K 以内）。Meta 专门收集了长格式文档（长学术论文、完整代码文件、长新闻报道等）用于这个阶段的训练。

5.2 RoPE Base 值与上下文长度的关系

为什么提高 RoPE base 值能帮助长上下文？

对于第 $i$ 对维度，旋转周期（完整旋转一圈所需的位置数）为：

$$T_i = \frac{2\pi}{\theta_i} = 2\pi \cdot \text{base}^{2i/d}$$

对于最低频维度（$i = d/2 - 1$）：

base = 10000：$T_{max} = 2\pi \times 10000 = 62832$ 个位置
base = 500000：$T_{max} = 2\pi \times 500000 = 3141593$ 个位置

当上下文长度超过旋转周期 $T_i$ 时，相同的旋转角度会出现在不同的绝对位置，造成”位置混淆”。

通过提高 base 值，使得所有维度的旋转周期都远大于实际使用的最大上下文长度（128K），从而避免位置混淆。

选择 500K 而非更大值的原因： base 值越大，相邻 token 的相对旋转角度越小，位置区分精度下降。500K 在 128K 上下文长度下达到了良好的平衡。

六、FlashAttention 与长上下文：IO 优化基础

（详细原理见第 01 篇第十二节，此处重点讨论长上下文场景）

FlashAttention 对长上下文的意义不只是速度，更是可行性：

上下文长度	标准注意力显存（BF16）	FlashAttention 显存（BF16）
4K	128 MB	约 4 MB
32K	8 GB	约 32 MB
128K	128 GB（超出单张 A100）	约 512 MB

FlashAttention 将注意力矩阵的显存从 $O(n^2)$ 降低到 $O(n)$（通过分块计算避免存储完整矩阵），使得 128K 上下文的训练在单张 A100 上成为可能（虽然还需要额外的 KV Cache 显存）。

6.1 上下文并行（Context Parallelism）

对于超过单张 GPU 显存的超长上下文（如 Gemini 的 1M tokens），还需要进一步的并行策略——上下文并行（Context Parallelism）：

将长序列分成若干段，分配给不同的 GPU，每个 GPU 负责计算其中一段的注意力。

挑战： 注意力机制需要每个位置访问所有其他位置的 K/V，上下文并行需要通过 GPU 间通信来收集远程 K/V。

Ring Attention（Liu et al., UC Berkeley, 2023） 是一种高效的上下文并行实现：将 GPU 排成”环形”，每个 GPU 保存一段的 Q/K/V，然后 K/V 在环形 GPU 之间流转，每个 GPU 在等待通信的同时处理已有的数据，实现计算和通信的重叠。

LLaMA 3 405B 训练时使用了 8 路上下文并行，与其他并行策略结合，支持 128K 上下文的大规模训练。

七、KV Cache 优化：处理长上下文的推理瓶颈

即使解决了训练时的长上下文问题，推理时的 KV Cache 仍然是主要瓶颈。

7.1 KV Cache 的精确计算

对于 LLaMA 3 70B（GQA，$h_{KV} = 8$，$d_k = 128$，80 层，BF16）：

$$\text{KV Cache per token per layer} = 2 \times h_{KV} \times d_k \times 2 \text{ bytes} = 2 \times 8 \times 128 \times 2 = 4096 \text{ bytes} = 4 \text{ KB}$$

$$\text{Total KV Cache (128K context, 80 layers)} = 128000 \times 80 \times 4 \text{ KB} = 40.96 \text{ GB}$$

加上模型参数（约 140 GB）：总显存需求约 181 GB，需要至少 3 张 A100 80GB。

7.2 KV Cache 量化

KVQuant（Hooper et al., UC Berkeley, 2024）arXiv:2401.18079：

将 KV Cache 从 BF16/FP16（16 bits）量化到 4 bits，显存减少 75%：

$$\text{KV Cache (4-bit, 128K, LLaMA 3 70B)} \approx 10 \text{ GB}$$（从 41 GB 降到约 10 GB）

代价：轻微的精度损失（通常 < 1% 的任务性能下降）。

7.3 Attention Sinks：流式长上下文

论文： Efficient Streaming Language Models with Attention Sinks（Xiao et al., MIT/Meta, 2023）arXiv:2309.17453

Attention Sink 现象的发现： 观察发现，无论输入是什么，最开始的几个 token（通常是 <BOS> 和前几个词）总是接收到远超其语义重要性的注意力权重。

这些 token 被称为”注意力汇聚点（Attention Sink）”。

为什么会出现 Attention Sink？

softmax 函数要求注意力权重之和为 1。当模型需要在某个位置”不关注任何特定内容”时（类似”保持原始输入”的残差），它把注意力权重倾倒到这些”垃圾收集”token 上，因为这比均匀分散权重更容易学习（初始 token 在所有层都存在，是稳定的汇聚点）。

StreamingLLM 的应用：

利用 Attention Sink 现象，可以实现”无限长度”的流式推理：只保留最初的几个 token（Attention Sinks）和最近的 $W$ 个 token 的 KV Cache，丢弃中间的历史 token。

这不能保留完整的历史（信息会丢失），但对于不需要引用远距离历史的流式对话场景，可以让上下文”滑动窗口”式地处理任意长的输入，而不需要无限增长的 KV Cache。

八、Gemini 1.5 的 100 万 Token 上下文

论文： Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context（Team Google, 2024）arXiv:2403.05530

8.1 技术概述

Gemini 1.5 Pro 支持 1M tokens 的上下文长度（约等于 10 本中等长度的书，或 11 小时的视频），并测试了 10M tokens 的极端情况。

已披露的关键技术（来自论文）：

MoE 架构：通过 MoE 在保持计算效率的同时扩大模型容量
Ring Attention 变体：用于长上下文的跨设备注意力并行
专门的长上下文预训练数据：包含大量长格式文档和视频字幕

未披露的细节： Google 没有公开 Gemini 1.5 的具体位置编码方案、MoE 配置参数等架构细节。

8.2 长上下文测试结果

“Needle in a Haystack”（大海捞针）测试：

将一条特定信息（”针”）随机插入到 100 万 token 的文档（”干草堆”）中，测试模型能否找到它。

Gemini 1.5 Pro 在整个 1M token 范围内的检索准确率 > 99%，在 10M token 时保持 > 92%（且在 10M 测试中也保持了多文件跨文档推理能力）。

特殊任务：”从一本书学一门语言”：

给 Gemini 1.5 一本 Kalamang 语（全球约 200 名使用者的巴布亚新几内亚语言）的语法书（约 500 页），然后测试它能否进行 Kalamang-英语翻译。

结果：Gemini 1.5 在没有任何预训练 Kalamang 数据的情况下，仅从这本语法书就学会了基础翻译。

这个测试展示了超长上下文 + 强大推理的结合：模型需要同时”理解”语法规则、记住词汇表、并正确应用到翻译任务中。

九、LongRoPE：超越 200 万 Token

论文： LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens（Ding et al., Microsoft, 2024）arXiv:2402.13753

9.1 核心发现

LongRoPE 的作者发现，YaRN 和位置插值方案的关键问题是：它们对 RoPE 所有维度使用相同的缩放因子（即使 NTK-by-Parts 也是按预定义规则划分），但不同 Transformer 层和不同维度对外推的容忍度是不同的。

实验发现：

靠近输入的低层（Layer 1-5）：对长距离位置外推更敏感
靠近输出的高层（Layer 20+）：对长距离外推容忍度更高
高频维度（$i$ 小）：可以大比例压缩
某些中频维度：几乎不需要压缩

9.2 非均匀缩放方法

LongRoPE 通过演化算法（Evolutionary Algorithm）搜索每个维度的最优缩放因子：

给定训练长度 $L$ 和目标长度 $L_{ext}$，寻找缩放向量 $\lambda = [\lambda_1, \lambda_2, \ldots, \lambda_{d/2}]$（每个维度有独立的缩放因子），使得在目标长度上的 perplexity 最小：

$$\min_\lambda \text{PPL}(L_{ext}, \lambda)$$

由于这个搜索空间是高维的（$d/2$ 维，如 2048 维），使用演化算法（而非梯度优化，因为 perplexity 相对缩放因子不可微）来搜索。

搜索成本远低于重新训练：对于一个 7B 模型，整个搜索过程只需要约 1000 步前向推理（无梯度计算）。

9.3 两步推进策略

LongRoPE 提出”两步推进”（Two-Step Progression）策略：

步骤一：扩展到目标长度的 8 倍

先用进化搜索扩展到目标长度的 8 倍，然后做少量微调（约 1000 步，使用长文档）。

步骤二：恢复短上下文性能

扩展到超长上下文后，模型在原始短上下文（如 2K-4K tokens）上的性能可能略有下降（因为位置缩放改变了近距离位置的表示）。

LongRoPE 为短上下文推理和长上下文推理分别维护两套缩放因子，在推理时根据实际序列长度动态选择。

结果： 基于 Phi-2（2.7B）和 LLaMA 2（7B），LongRoPE 将上下文从 4K 扩展到 2048K（200 万 tokens），在 “Needle in a Haystack” 测试中保持约 95% 的准确率。

十、长上下文评估方法

10.1 Needle in a Haystack（大海捞针）

最常见的长上下文测试，由 Greg Kamradt 提出：

将一条短信息（”针”）插入到长文本（”干草堆”）的随机位置
在文本末尾添加问题，要求模型回答需要定位”针”的问题
测试不同文档位置和不同上下文深度的检索准确率

局限性： 这个测试只测试单条信息的检索，不能反映多跳推理（需要结合多处信息）或复杂文档理解能力。

10.2 RULER（Realistic Long-Context Understanding and Retrieval）

论文： RULER: What’s the Real Context Size of Your Long-Context Language Models?（Hsieh et al., NVIDIA, 2024）arXiv:2404.06654

RULER 提出了比”大海捞针”更全面的长上下文测试集，包含：

单跳 QA（标准检索）
多跳 QA（跨多处信息）
聚合（计数、排序长列表中的元素）
多 key-value 关联（记住多对对应关系）

发现： 几乎所有声称支持 128K 上下文的模型，在 RULER 的复杂任务上，实际有效上下文长度远低于其标称值（通常有效长度只有标称值的 20-50%）。

10.3 Long Context MT-Bench

多轮长文档理解测试，测试在长文档基础上进行多轮问答的能力，接近实际使用场景。

十一、长上下文的实际工程注意事项

11.1 “Lost in the Middle”现象

论文： Lost in the Middle: How Language Models Use Long Contexts（Liu et al., Stanford, 2023）arXiv:2307.03172

研究发现，LLM 处理长上下文时，对文档开头和结尾的信息关注度显著高于中间部分——即使相关信息被放在中间，模型也难以有效利用。

实践建议：

将最重要的信息放在上下文的开头或结尾
对于 RAG（Retrieval Augmented Generation）系统，尽量将检索到的相关段落放在靠近问题的位置（即上下文末尾）

11.2 推理成本的二次方增长

虽然 FlashAttention 解决了显存问题，但注意力的计算复杂度仍然是 $O(n^2)$（只是 IO 复杂度降低了）。对于推理时间：

上下文长度	相对于 4K 的推理时间（prefill 阶段）
4K（基准）	1x
32K	64x（$8^2$）
128K	1024x（$32^2$）

这意味着：对于 128K 上下文的首个 token 生成（prefill 阶段），计算时间是 4K 的约 1000 倍。实际中通常通过批处理和计算优化（如 Flash Decoding）来缓解，但长上下文的延迟确实显著增加。

11.3 稀疏注意力作为未来方向

为彻底解决 $O(n^2)$ 问题，研究者探索了各种稀疏注意力（Sparse Attention）方案：

方法	核心思路	复杂度
Longformer（Beltagy et al., 2020）	局部窗口 + 全局 token 注意力	$O(n)$
BigBird（Zaheer et al., 2020）	随机 + 局部 + 全局注意力	$O(n)$
Reformer（Kitaev et al., 2020）	LSH 哈希分组注意力	$O(n \log n)$
Hyena/H3（Fu et al., 2022-23）	长卷积替代注意力	$O(n \log n)$
Mamba（Gu & Dao, 2023）	结构化状态空间模型（SSM）	$O(n)$

这些方案在特定任务上有效果，但迄今为止没有一个在通用语言建模上完全媲美标准注意力的表现，它们仍在持续研究中。