LLM 上下文窗口扩展：RoPE、ALiBi 与长文本建模技术

上下文窗口是大语言模型最关键的架构参数之一。它决定了模型在一次推理中能够「看到」和「记住」的信息量。一个 128K 上下文的模型可以读取约 10 万汉字的文档，而一个 1M 上下文的模型可以吞下一整本长篇小说。

上下文窗口的重要性体现在三个维度：第一，理解长文档——法律合同、科研论文、代码仓库的完整理解都需要超大上下文；第二，长对话历史——客服 Agent 和虚拟助手需要记住数轮对话的上下文；第三，RAG 增强检索——检索到的相关片段越多，上下文窗口需要越大。

Transformer 架构对上下文窗口的根本限制来自自注意力机制的二次方复杂度。注意力矩阵的大小是序列长度的平方——1024 长度的序列需要约 100 万次注意力计算，而 1M 长度的序列需要 1 万亿次计算。这就是为什么扩展上下文窗口不仅仅是「训练更长的数据」，而是需要在架构层面做深度优化。

位置编码是上下文窗口扩展的核心瓶颈。Transformer 本身不具备序列位置的概念，需要位置编码告诉模型「这个 token 在序列的哪个位置」。不同的位置编码方案对长度外推能力有巨大影响——有些方案在超过训练长度时性能急剧下降，有些则能平滑地外推到更长序列。

要理解上下文窗口扩展，必须先理解位置编码的本质问题。Transformer 的自注意力机制是对所有 token 对计算注意力权重——它天然是对称的，即「A 注意到 B」和「B 注意到 A」的注意力分数相同。但自然语言是有顺序的：「猫追狗」和「狗追猫」含义完全不同。

绝对位置编码是最直观的方案——给序列中每个位置分配一个唯一的向量。早期 Transformer 使用正弦余弦函数生成位置向量：位置 p 的第 i 个维度是 sin(p 除以 10000 的 2i 除以 d 次方)。这种方案的优点是计算简单，但缺点是外推能力差——模型在训练时只见过位置 0 到 2047，推理时给它位置 4096，它完全不知道这意味着什么。

相对位置编码的哲学不同：不关心「这个 token 在第几个位置」，而是关心「这个 token 距离那个 token 有多远」。注意力计算时，模型查询的是「我前面的第 3 个 token 是什么」而非「第 7 个 token 是什么」。这种方案天然具有更好的长度外推能力，因为相对距离的概念不随序列总长度变化。

现代 LLM 的主流位置编码方案有三种：旋转位置编码 RoPE（GPT 系列、Llama、Qwen 等）、注意力线性偏置 ALiBi（T5、MPT 等）、以及可学习位置编码（BERT、ViT）。它们在长度外推能力上有显著差异。

旋转位置编码（Rotary Position Embedding, RoPE） 是目前大语言模型中使用最广泛的位置编码方案。GPT 系列、Llama、Qwen、Baichuan、ChatGLM 等主流模型都采用了 RoPE。它的核心思想非常优雅：将位置信息编码为复数域中的旋转角度，通过旋转操作将位置注入注意力计算中。
RoPE 的数学本质 ：对于序列中位置 m 和位置 n 的两个 token，RoPE 让它们对应的查询向量 Q 和键向量 K 在复数空间中旋转 m 和 n 个角度。当计算注意力分数时（Q 和 K 的点积），旋转角度的差值恰好等于 m 减 n——也就是两个 token 的相对距离。这样，注意力分数天然就包含了相对位置信息，同时保留了绝对位置的可学习性。 RoPE 的维度分组策略：RoPE 将向量的维度分成若干组（每组 2 个维度），每组使用不同的旋转频率。低频组（前面维度）旋转慢，捕获长距离依赖；高频组（后面维度）旋转快，捕获短距离细节。 这种多尺度的设计使得 RoPE 既能关注邻近 token 的精细关系，也能关注远距离 token 的全局结构。RoPE 的优势在于它的数学优雅性和工程效率：旋转操作可以通过简单的矩阵乘法实现，不需要额外的显存开销；它与标准注意力机制完全兼容，不需要修改注意力计算的核心逻辑；它的相对位置编码性质使得它比绝对位置编码有更好的长度外推能力。但 RoPE 也有局限：当序列长度远超训练长度时（例如训练 4K 但推理 128K），旋转角度会超出模型见过的范围，导致注意力质量下降。这就是为什么需要 NTK 插值、YaRN 等外推增强技术来弥补 RoPE 的不足。

import torch
import math

def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
    """
    应用 RoPE 位置编码到 Q 和 K 向量
    
    Args:
        q: 查询向量 [batch, seq_len, num_heads, head_dim]
        k: 键向量   [batch, seq_len, num_heads, head_dim]
        cos: 旋转余弦矩阵 [seq_len, head_dim]
        sin: 旋转正弦矩阵 [seq_len, head_dim]
        position_ids: 位置索引 [batch, seq_len]
    """
    # 根据位置索引获取对应的 cos/sin
    cos = cos[position_ids].unsqueeze(2)  # [batch, seq_len, 1, head_dim]
    sin = sin[position_ids].unsqueeze(2)
    
    # RoPE 核心：复数旋转
    # q_rot = q * cos + rotate_half(q) * sin
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    
    return q_embed, k_embed

def rotate_half(x):
    """将向量后半部分与前半部分交换并取反"""
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

def precompute_freqs_cis(dim, max_seq_len, base=10000.0):
    """预计算 RoPE 频率常量（复数形式的 cos/sin）"""
    freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    t = torch.arange(max_seq_len, dtype=freqs.dtype)
    freqs = torch.outer(t, freqs)  # [max_seq_len, dim/2]
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # 转为复数
    return freqs_cis  # [max_seq_len, dim/2]

ALiBi（Attention with Linear Biases）是一种完全不同的位置编码方案。它不通过旋转或嵌入来注入位置信息，而是直接在注意力分数上减去一个与相对距离成正比的线性偏置。

ALiBi 的核心公式非常简洁：注意力分数等于 Q 乘以 K 的转置，再减去 m 乘以 j 减 i——其中 m 是一个预设的斜率，j 减 i 是查询 token 和键 token 之间的相对距离。距离越远的 token 对，注意力分数被减去得越多，这天然地鼓励模型关注局部上下文。

ALiBi 的最大优势是零外推成本：因为它使用的相对距离不依赖于序列的绝对长度，所以无论推理时的序列有多长，注意力计算方式完全一致。这意味着一个在 1024 长度上训练的 ALiBi 模型可以直接在 1M 长度上推理，而不会出现性能骤降。

ALiBi 的斜率设计：不同注意力头使用不同的斜率——靠近的头（捕获短距离依赖）使用大斜率，远离的头（捕获长距离依赖）使用小斜率。具体公式为 m 头等于 1 除以 2 的 p 次方，其中 p 是与头序号相关的参数。这种多尺度设计与 RoPE 的维度分组有异曲同工之妙。

ALiBi 的局限性在于它的归纳偏置过于强烈：它强制模型偏好短距离注意力，这可能导致在需要全局理解的任务（如长文档摘要、跨段落推理）上表现不佳。此外，ALiBi 的线性衰减不如 RoPE 的周期性旋转灵活——RoPE 可以表示「位置 100 和位置 200 的关系」与「位置 200 和位置 300 的关系」是相似的（因为旋转角度差相同），但 ALiBi 只能表达「距离越远越不重要」。

import torch

def build_alibi_slopes(n_heads, max_seq_len=2048):
    """
    构建 ALiBi 斜率矩阵
    
    Args:
        n_heads: 注意力头数量
        max_seq_len: 最大序列长度
    
    Returns:
        alibi_bias: 形状 [n_heads, max_seq_len, max_seq_len]
    """
    # 计算每个头的斜率：1/2^p
    # p 按照几何级数分配给不同头
    x = (2 ** 8) ** (1 / n_heads)
    slopes = torch.tensor([1 / (x ** i) for i in range(1, n_heads + 1)])
    
    # 构建相对距离矩阵
    distance = torch.arange(max_seq_len).unsqueeze(1) - torch.arange(max_seq_len).unsqueeze(0)
    # distance[i][j] = i - j (查询位置 - 键位置)
    
    # 只有 j <= i 的位置有效（因果注意力）
    # 对于无效的 j > i，设置为负无穷
    alibi_bias = slopes.unsqueeze(1).unsqueeze(2) * distance.clamp(max=0).unsqueeze(0)
    # alibi_bias[head, i, j] = slope_head * min(0, i - j)
    
    return alibi_bias

def alibi_attention(q, k, v, alibi_bias):
    """
    应用 ALiBi 的注意力计算
    
    Args:
        q, k, v: 标准的注意力输入
        alibi_bias: 预计算的 ALiBi 偏置矩阵
    """
    d_k = q.size(-1)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
    
    # 关键步骤：直接加上 ALiBi 偏置（注意是减法，所以加负值）
    scores = scores + alibi_bias
    
    attn_weights = torch.softmax(scores, dim=-1)
    return torch.matmul(attn_weights, v)

NTK 插值（Neural Tangent Kernel Interpolation） 是一种不需要重新训练就能让 RoPE 模型支持更长上下文的技术。它的核心洞察来自理论深度学习：神经网络的 NTK 描述了模型在训练过程中的动态行为——如果我们在推理时对输入做某种变换，而变换后的输入在 NTK 意义下与训练数据「足够接近」，那么模型应该仍然能给出合理输出。
NTK 插值的具体操作 ：对 RoPE 的频率基数 base 乘以一个缩放因子 s（通常 s 大于 1）。这等价于将旋转角度压缩——原来旋转 2pi 弧度需要 10000 步，现在只需要 10000 除以 s 步。角度压缩的效果是，模型在推理时「以为」自己处理的是更短的序列，从而激活了在训练时学到的注意力模式。 NTK 插值的三种变体 ： 静态 NTK 插值 ：整个模型使用统一的缩放因子。实现简单，但可能导致短距离注意力质量下降——因为所有频率都被压缩了。 动态 NTK 插值 ：根据输入序列的实际长度动态调整缩放因子。短序列不缩放，长序列适度缩放。这种方法在保持短距离注意力质量的同时，为长序列提供了外推支持。 混合 NTK 插值：只对部分维度（通常是高频维度）做缩放，低频维度保持原始频率。 这种方案兼顾了短距离和长距离的注意力质量。NTK 插值的效果：对于在 4K 长度上训练的 Llama-2-7B 模型，使用 NTK 插值可以在 32K 长度上推理，困惑度下降约 30%。对于更长序列（如 128K），可能需要结合 YaRN 技术。NTK 插值的局限性：它是一种启发式技术，缺乏严格的理论保证。缩放因子 s 的选择需要经验调优——s 太小，外推效果不足；s 太大，短距离注意力质量受损。

YaRN（Yet another RoPE extensioN method） 是对 NTK 插值的改进方案，由微软研究团队提出。它在 NTK 的基础上增加了两个关键改进：动态温度缩放和注意力修正。
YaRN 的第一项改进：动态温度缩放。NTK 插值使用固定的缩放因子 s，而 YaRN 引入了温度参数 T，使得频率缩放不再是简单的线性除法，而是通过一个平滑的温度函数：缩放后的频率等于原始频率乘以 1 加上 s 减 1 除以 1 加上 T 的函数。这个温度函数的作用是让缩放更加平滑——低频维度几乎不缩放，中频维度适度缩放，高频维度大幅缩放。
YaRN 的第二项改进：注意力修正。YaRN 发现，即使频率缩放正确，超长长序列的注意力质量仍然会下降——因为注意力权重的分布会随着序列长度变化。YaRN 在计算注意力分数后，对权重分布做一个修正：将注意力权重乘以一个与序列长度相关的修正因子，使得长序列上的注意力分布更接近训练时的分布。YaRN 的效果：在 Llama-2-7B 上，使用 YaRN 可以在 128K 长度上推理，困惑度比原始 NTK 插值再降低约 15%。在需要超长上下文的任务（如长文档问答、代码仓库理解）上，YaRN 是目前效果最好的外推方案之一。YaRN 与 NTK 插值的对比：

当上下文窗口从 4K 扩展到 128K 甚至 1M 时，KV Cache 的显存占用成为一个严峻的系统级挑战。KV Cache 存储了所有已生成 token 的键和键向量——对于 70B 参数、128 头、128 头维的模型，每个 token 的 KV Cache 约为 128 乘以 128 乘以 2 乘以 2 字节（FP16）等于 64KB。128K 长度的序列需要约 8GB KV Cache 显存，而 1M 长度则需要约 64GB——仅 KV Cache 就占满了整张 H100 的显存。

KV Cache 优化的核心思路是减少每个 token 的缓存大小或减少缓存的 token 数量。

量化 KV Cache：将 KV Cache 从 FP16 量化为 INT8 甚至 INT4，显存减少 50% 到 75%。量化引入的精度损失在注意力计算中通常可接受——因为注意力分数最终会被 softmax 归一化，微小的精度差异会被平滑掉。

KV Cache 压缩：只缓存关键 token 的 KV 向量——例如每个 64 个 token 只缓存一个，其余通过插值近似。这种方法可以将 KV Cache 压缩 32 倍，但会在需要精确定位的任务上（如代码补全、命名实体识别）产生质量问题。

分页 KV Cache：借鉴操作系统的虚拟内存机制，将 KV Cache 分为「热」（最近使用的 token）和「冷」（历史 token）两部分。热部分保留在 GPU 显存中，冷部分换出到 CPU 内存或 SSD 上。这种方案可以在不牺牲质量的前提下，支持远超单卡显存容量的上下文长度。

滑动窗口注意力：只关注最近的 N 个 token（如 4K），忽略更早的 token。这对于对话场景特别有效——用户通常只关心最近的几轮对话。滑动窗口可以将 KV Cache 的显存占用固定为 O（窗口大小），而不是 O（序列长度）。

import torch

class QuantizedKVCache:
    """量化 KV Cache 实现，将 FP16 缓存量化为 INT8 以节省显存"""
    
    def __init__(self, max_seq_len, num_heads, head_dim, device="cuda"):
        # 存储量化后的 KV
        self.k_cache_int8 = torch.zeros(max_seq_len, num_heads, head_dim, 
                                         dtype=torch.int8, device=device)
        self.v_cache_int8 = torch.zeros(max_seq_len, num_heads, head_dim,
                                         dtype=torch.int8, device=device)
        # 存储量化参数
        self.k_scales = torch.zeros(max_seq_len, num_heads, device=device)
        self.k_zeros = torch.zeros(max_seq_len, num_heads, device=device)
        self.v_scales = torch.zeros(max_seq_len, num_heads, device=device)
        self.v_zeros = torch.zeros(max_seq_len, num_heads, device=device)
        self.length = 0
    
    def append(self, k_new, v_new):
        """添加新 token 的 KV 并量化"""
        seq_len = k_new.size(0)
        
        # INT8 量化：scale = max / 127
        k_max = k_new.abs().max(dim=-1, keepdim=True).values.clamp(min=1e-5)
        k_scale = k_max / 127.0
        k_zero = torch.zeros_like(k_scale)
        k_int8 = (k_new / k_scale).round().clamp(-128, 127).to(torch.int8)
        
        v_max = v_new.abs().max(dim=-1, keepdim=True).values.clamp(min=1e-5)
        v_scale = v_max / 127.0
        v_zero = torch.zeros_like(v_scale)
        v_int8 = (v_new / v_scale).round().clamp(-128, 127).to(torch.int8)
        
        # 存储
        self.k_cache_int8[self.length:self.length+seq_len] = k_int8
        self.k_scales[self.length:self.length+seq_len] = k_scale
        self.k_zeros[self.length:self.length+seq_len] = k_zero
        self.v_cache_int8[self.length:self.length+seq_len] = v_int8
        self.v_scales[self.length:self.length+seq_len] = v_scale
        self.v_zeros[self.length:self.length+seq_len] = v_zero
        self.length += seq_len
    
    def get(self, end_idx):
        """获取并反量化 KV 向量"""
        k_int8 = self.k_cache_int8[:end_idx]
        k_scale = self.k_scales[:end_idx]
        k_dequant = k_int8.to(torch.float16) * k_scale
        
        v_int8 = self.v_cache_int8[:end_idx]
        v_scale = self.v_scales[:end_idx]
        v_dequant = v_int8.to(torch.float16) * v_scale
        
        return k_dequant, v_dequant

上下文窗口扩展的技术正在快速演进。2024 年主流模型的上下文窗口是 32K 到 128K，2025 年已达到 200K 到 1M，2026 年出现了「无限长上下文」的概念。

无限长上下文的核心思路是放弃传统的二次方注意力机制，改用线性复杂度的替代方案：

线性注意力（Linear Attention）：通过核函数近似 softmax 注意力，将注意力复杂度从 O（N 平方）降低到 O（N）。代表性方案有 Performer（使用随机特征映射）、Linear Transformer（使用核技巧）。这些方案的缺点是近似精度有限——在某些任务上可能不如标准注意力。

状态空间模型（State Space Models, SSM）：Mamba、RWKV 等模型用状态空间方程替代注意力机制，天然支持 O（N）复杂度的序列建模，且上下文长度在理论上无上限。2026 年，混合架构（Attention 加 SSM）成为研究热点——在需要精确定位的层使用注意力，在需要长距离建模的层使用 SSM。

检索增强上下文（Retrieval-Augmented Context）：将上下文存储在向量数据库中，只在需要时检索相关片段，而不是将所有内容一次性输入模型。这种方案本质上是 RAG 的升级版——从「检索文档片段」进化为「检索注意力权重」。

分块注意力（Sliding Window + Global Tokens）：结合局部滑动窗口和少量全局 token（如段落开头、文档摘要），在 O（N 乘窗口大小）的复杂度下，既保留了局部精细注意力，又捕获了全局结构信息。Longformer 和 BigBird 是这一方向的代表性工作。

上下文窗口扩展的未来方向：

第一，硬件协同设计。新的 GPU 架构（如 NVIDIA H200、B200）针对长序列注意力做了硬件级优化——更大的 HBM 容量、更高的内存带宽、专用的注意力加速单元。软件优化与硬件协同，才能充分发挥长上下文的能力。

第二，自适应上下文管理。未来的模型将根据输入内容的信息密度动态调整有效上下文——高信息密度区域分配更多注意力资源，低信息密度区域（如重复内容、格式标记）分配较少资源。

第三，跨模态上下文统一。当模型需要同时处理文本、图像、音频、视频时，上下文窗口的概念需要扩展到多模态空间——一个视频帧可能等价于数百个文本 token 的信息量。如何公平地分配上下文资源，是多模态 LLM 面临的新挑战。

本领域的必读论文：

RoPE 原始论文（Su 等，2021）是旋转位置编码的开创性工作，详细推导了旋转操作如何将相对位置信息注入注意力计算。

ALiBi 论文（Press 等，2022）展示了线性偏置方案在长度外推上的优势，是 RoPE 之外最重要的位置编码方案。

NTK-Aware Scaled RoPE（Chen 等，2023）首次将 NTK 理论应用于 RoPE 外推，为不需要重新训练的上下文扩展提供了理论基础。

YaRN 论文（Peng 等，2023）在 NTK 基础上增加了温度缩放和注意力修正，是目前效果最好的 RoPE 外推方案。

本领域的实用工具：

vLLM 是目前最流行的 LLM 推理框架，原生支持 PagedAttention 和 KV Cache 优化，是长文本推理的首选部署方案。

FlashAttention 通过 IO 感知的注意力计算，将长序列注意力的显存占用从 O（N 平方）降低到 O（N），是长上下文推理的必备技术。

Unsloth 提供了 LLM 训练和推理的加速方案，支持超长上下文的 NTK 插值和 YaRN 外推。

延伸阅读方向：

注意力机制的替代方案（状态空间模型、线性注意力、RWKV）。

KV Cache 的系统级优化（分页、量化、压缩、预填充优化）。

长文本评估基准（如 Needle in a Haystack、LongBench、InfiniteBench）。

多模态上下文建模（图文混合、视频-语言联合建模）。