第1章 LLM基础原理
(续 3)
上下文窗口:物理限制、KV Cache 与工程应对策略
1. 引言:上下文窗口是 LLM 工程的第一性原理
当你向一个 LLM 发出请求,无论是简单的问答、复杂的代码生成,还是多轮 Agent 任务,有一条规则始终成立:模型只能"看到"当前上下文窗口内的内容。它没有长期记忆,没有外部状态,没有隐式的"理解历史"------一切都压缩进那个有限的 token 序列里。
这不是一个可以被绕过的工程限制,而是 Transformer 架构的根本性约束。理解它的物理成因、内部机制和工程应对,是从"能用 LLM"到"用好 LLM"的核心跨越。
本节将沿着"物理原理 → KV Cache 机制 → 工程瓶颈 → 扩展技术 → 实战策略"的完整链路展开,不仅解释 what ,更深入 why 和 how。
本节阅读地图
- 第 2 节:上下文窗口的精确定义与组成结构
- 第 3 节:为什么窗口有上限------三个维度的物理约束
- 第 4 节:KV Cache 原理、显存占用计算与优化
- 第 5 节:长上下文质量陷阱与 Lost in the Middle
- 第 6 节:扩展上下文窗口的主流技术路线
- 第 7 节:工程案例与最佳实践
2. 上下文窗口的精确定义
2.1 token 是什么单位
在深入上下文窗口之前,必须先建立对 token 的精确认知。Token 是 LLM 处理文本的最小单位,通常是子词(subword)片段,介于字符和单词之间。上下文窗口的大小以 token 计,而非字符、单词或字节。
一个关键的工程直觉:
- 英文:约 1 token ≈ 4 个字符 ≈ 0.75 个单词
- 中文:约 1 token ≈ 1.5--2 个汉字(主流英文优先词表)
- 代码:约 1 token ≈ 3--5 个字符(含缩进空白)
- Emoji 或罕见字符:可能需要 3--6 个 token
这意味着,一个标称 128K token 窗口的模型,实际能容纳的英文大约是 96,000 个单词(约 200 页 A4),但中文则大约只有 64,000--85,000 个汉字。
2.2 上下文窗口的精确边界
**上下文窗口(Context Window)**是模型在一次 forward pass 中能处理的最大 token 总量,包含输入和输出的全部内容。其组成如下:
scss
┌─────────────────────────────────────────────────────┐
│ Context Window │
│ ┌──────────────┐ ┌────────────┐ ┌─────────────┐ │
│ │ System Prompt│ │ History │ │ User Input │ │
│ │ (固定层) │ │ (历史对话) │ │ (当前轮) │ │
│ └──────────────┘ └────────────┘ └─────────────┘ │
│ ┌──────────────┐ ┌────────────────────────────┐ │
│ │ RAG Context │ │ Model Output │ │
│ │ (检索注入) │ │ (模型生成部分) │ │
│ └──────────────┘ └────────────────────────────┘ │
│ │
│ total_tokens = sum(all above) ≤ max_context_window │
└─────────────────────────────────────────────────────┘
两个容易混淆的概念必须区分:
| 概念 | 含义 | 典型值(Claude 3.5 Sonnet) | 是否可配置 |
|---|---|---|---|
| Context Window(上下文窗口) | 模型单次处理的最大 token 总量(输入+输出) | 200,000 tokens | 否(架构决定) |
| max_output_tokens(最大输出长度) | 单次请求模型最多生成的 token 数 | 8,192 tokens(默认) | 是(请求参数) |
二者的约束关系:input_tokens + output_tokens ≤ context_window。若输入已消耗 195,000 tokens,即使 max_output_tokens 设为 8,192,实际最多只能生成 5,000 tokens。
常见误区:混淆窗口大小与质量
窗口越大不等于质量越好。标称 128K 窗口的模型,在实际长文档任务中的表现可能远不如 32K 窗口的精调模型。窗口是容量上限,不是质量保证。后文的「Lost in the Middle」现象将深入解释这一点。
2.3 主流模型窗口对比(2025)
| 模型 | 上下文窗口 | 最大输出 | 定价(输入/输出,$/M tokens) | 备注 |
|---|---|---|---|---|
| GPT-4o | 128K | 16K | 2.5/10 | 支持 prompt caching |
| Claude 3.5 Sonnet | 200K | 8K | 3/15 | 支持 extended thinking |
| Claude 3.7 Sonnet | 200K | 64K(thinking 模式) | 3/15 | 思考 token 单独计费 |
| Gemini 1.5 Pro | 1M | 8K | 1.25--2.5 / 5--10 | 超过 128K 价格翻倍 |
| Gemini 2.0 Flash | 1M | 8K | 0.1/0.4 | 极低成本长上下文 |
| LLaMA 3.1 405B | 128K | 8K | 开源自部署 | 需 ~8×A100 80GB |
| Qwen2.5 72B | 128K | 8K | 开源自部署 | 中文友好,SentencePiece |
| DeepSeek V3 | 128K | 8K | 0.14/0.28 | MoE 架构,极低推理成本 |
3. 物理限制:上下文窗口上限的三个维度
上下文窗口的上限并非人为限制,而是由三个独立的物理约束共同决定的。理解每一个约束的成因,是选择正确工程解法的前提。
3.1 计算复杂度约束:O(n²) 的 Attention
Transformer 的核心是 Self-Attention 机制。对于长度为 n 的输入序列,标准 Self-Attention 需要计算每个 token 与其他所有 token 的注意力得分:
ini
Attention(Q, K, V) = softmax(QK^T / √d_k) · V
其中 QK^T 是 n×n 的矩阵,计算量 ∝ n²
序列长度对比:
n = 1K → 1× 基准
n = 2K → 4×
n = 8K → 64×
n = 32K → 1024×
n = 128K → 16384×(相对 1K 序列)
n = 1M → 1,000,000×(理论,实际不可行)
这个 O(n²) 的复杂度不仅体现在计算量(FLOP),还体现在显存占用:注意力矩阵本身需要 n×n 的显存来存储中间结果。
FlashAttention:解决显存瓶颈,不降低理论复杂度
FlashAttention(Dao et al. 2022)通过将 Q/K/V 分块(tiling)载入 SRAM 进行分块 softmax 计算,避免将完整的 n×n 注意力矩阵写回 HBM(高带宽显存)。理论 FLOP 复杂度仍是 O(n²),但显存 I/O 从 O(n²) 降至 O(n),使单卡训练 100K+ 序列成为可能。FlashAttention-2 在 A100 上比 PyTorch 原生实现快 2--4 倍,峰值显存降低 5--10 倍。
3.2 显存约束:KV Cache 的线性增长
在生成(decode)阶段,模型需要缓存每一层每个已生成 token 的 Key 和 Value 向量,这就是 KV Cache。KV Cache 的显存占用与序列长度呈线性关系:
ini
KV Cache 大小 = 2 × num_layers × seq_len × num_kv_heads × head_dim × batch_size × bytes
以 LLaMA-3 70B(GQA 配置)为例:
- num_layers = 80
- num_kv_heads = 8(GQA 分组)
- head_dim = 128
- FP16 = 2 bytes
单请求 (batch=1),序列 4K token:
2 × 80 × 4096 × 8 × 128 × 1 × 2 = 1.34 GB
单请求,序列 128K token:
2 × 80 × 131072 × 8 × 128 × 1 × 2 = 42.9 GB ← 超过任何单卡容量
这就是为什么 128K 上下文在商业 API 上可用,但在私有化部署时几乎需要多卡张量并行。KV Cache 往往是大模型推理中显存占用最大的部分,在长上下文场景下甚至超过模型权重本身。
3.3 位置编码约束:训练长度的外推困境
Transformer 通过位置编码(Positional Encoding)让模型感知 token 的顺序。现代 LLM 普遍使用 RoPE(Rotary Position Embedding,旋转位置编码),它将位置信息编码为旋转矩阵,嵌入到 Q/K 的内积中:
diff
RoPE 核心:
对位置 m 处的 token,将其 Q/K 向量乘以旋转矩阵 R(m)
R(m) 的不同维度 d 对应不同旋转频率:θ_d = 10000^(-2d/D)
效果:
- Q·K 的内积只依赖相对位置 (m-n),不依赖绝对位置
- 模型天然支持平移不变性
问题:
- 训练时最长见过 L_train = 4096 token
- 推理时序列长度 L_infer = 32768 token
- 位置 m > L_train 的旋转角度模型从未见过
- 注意力分数分布退化,性能骤降
这就是 **位置编码外推(Extrapolation)**问题。训练时见过的最长序列决定了模型能可靠处理的上下文长度,超过这个边界就需要特殊技术干预。
三重约束叠加
上下文窗口的上限由最紧张的约束决定。实践中,显存约束通常先于计算复杂度成为瓶颈------因为 KV Cache 线性增长,而计算可以通过 FlashAttention 优化。位置编码约束则决定了质量下界。三者共同画出了实际可用的上下文边界。
4. KV Cache:原理、计算与优化
4.1 为什么需要 KV Cache
LLM 生成文本是自回归(autoregressive)的过程:每次只生成一个 token,将其追加到输入序列后,再生成下一个。没有 KV Cache 的情况下,生成第 t 个 token 需要对前 t-1 个 token 全部重新计算 Q、K、V:
python
# 无 KV Cache 的朴素实现 - 注意:这只是说明原理,实际不这样写
def generate_without_cache(model, tokens):
"""
每步都对完整序列做 forward pass
第 100 步需要处理 100 个 token,重算前 99 个的 K/V
总计算量:O(1 + 2 + ... + n) = O(n²)
"""
generated = list(tokens)
for _ in range(max_new_tokens):
# 每次都把完整序列塞进去
logits = model(torch.tensor(generated))
# 只取最后一个位置的 logits
next_token = sample(logits[-1])
generated.append(next_token)
return generated
有了 KV Cache,每步只需计算新 token 的 Q/K/V,与缓存的历史 K/V 进行 attention:
python
import torch
import torch.nn.functional as F
from dataclasses import dataclass, field
from typing import Optional
@dataclass
class KVCacheLayer:
"""单层 KV Cache 的存储结构"""
# shape: (batch_size, num_kv_heads, seq_len, head_dim)
k: Optional[torch.Tensor] = None
v: Optional[torch.Tensor] = None
def update(self, new_k: torch.Tensor, new_v: torch.Tensor) -> tuple:
"""
追加新 token 的 K/V 并返回完整的历史 K/V
new_k/new_v shape: (batch, num_kv_heads, 1, head_dim) - 单个新 token
"""
if self.k is None:
# 第一个 token,直接赋值
self.k = new_k
self.v = new_v
else:
# 沿 seq_len 维度拼接
self.k = torch.cat([self.k, new_k], dim=2)
self.v = torch.cat([self.v, new_v], dim=2)
return self.k, self.v
@property
def seq_len(self) -> int:
"""当前缓存的序列长度"""
return 0 if self.k is None else self.k.shape[2]
def memory_bytes(self) -> int:
"""估算当前 KV Cache 占用的字节数"""
if self.k is None:
return 0
# K 和 V 各一份
return 2 * self.k.numel() * self.k.element_size()
def attention_with_kv_cache(
query: torch.Tensor, # (batch, num_heads, 1, head_dim) - 当前新 token 的 Q
key_new: torch.Tensor, # (batch, num_kv_heads, 1, head_dim) - 当前新 token 的 K
value_new: torch.Tensor, # (batch, num_kv_heads, 1, head_dim) - 当前新 token 的 V
cache: KVCacheLayer, # 当前层的 KV Cache
head_dim: int,
num_q_heads: int,
num_kv_heads: int,
) -> tuple[torch.Tensor, KVCacheLayer]:
"""
使用 KV Cache 加速的 attention 计算
核心思想:只计算新 token 的 Q/K/V,历史 K/V 从缓存中读取
"""
# 1. 将新 token 的 K/V 追加到缓存
full_k, full_v = cache.update(key_new, value_new)
# full_k shape: (batch, num_kv_heads, seq_len+1, head_dim)
# 2. 如果使用 GQA/MQA,需要将 KV 头数扩展匹配 Q 头数
if num_q_heads != num_kv_heads:
# GQA: 每组 Q 头共享同一组 KV
# repeat_factor = Q 头数 / KV 头数
repeat_factor = num_q_heads // num_kv_heads
# (batch, num_kv_heads, seq_len, head_dim) ->
# (batch, num_q_heads, seq_len, head_dim)
full_k = full_k.repeat_interleave(repeat_factor, dim=1)
full_v = full_v.repeat_interleave(repeat_factor, dim=1)
# 3. 计算注意力得分 - query 只有 1 个位置,full_k 有 seq_len+1 个位置
# (batch, num_heads, 1, head_dim) @ (batch, num_heads, head_dim, seq_len+1)
# = (batch, num_heads, 1, seq_len+1)
scale = head_dim ** -0.5
attn_scores = torch.matmul(query, full_k.transpose(-2, -1)) * scale
# 4. 因果 attention:新 token 可以 attend 所有历史 token(无需 mask)
# 因为 decode 阶段 query 只有 1 个 token,它天然只看历史
attn_weights = F.softmax(attn_scores, dim=-1)
# 5. 加权求和得到输出
# (batch, num_heads, 1, seq_len+1) @ (batch, num_heads, seq_len+1, head_dim)
# = (batch, num_heads, 1, head_dim)
output = torch.matmul(attn_weights, full_v)
return output, cache
def estimate_kv_cache_gb(
num_layers: int,
seq_len: int,
num_kv_heads: int,
head_dim: int,
batch_size: int = 1,
dtype_bytes: int = 2, # FP16 = 2 bytes, INT8 = 1 byte
) -> float:
"""
估算 KV Cache 总显存占用(GB)
公式:2 * num_layers * seq_len * num_kv_heads * head_dim * batch_size * dtype_bytes
- 乘以 2:K 和 V 各一份
- 随 seq_len 和 batch_size 线性增长
"""
total_bytes = (
2 # K and V
* num_layers
* seq_len
* num_kv_heads
* head_dim
* batch_size
* dtype_bytes
)
return total_bytes / (1024 ** 3) # 转换为 GB
# 实际模型 KV Cache 估算示例
configs = {
"LLaMA-2 7B (MHA)": {
"num_layers": 32,
"num_kv_heads": 32, # MHA:每个 Q 头都有独立 KV
"head_dim": 128,
},
"LLaMA-3 70B (GQA)": {
"num_layers": 80,
"num_kv_heads": 8, # GQA:8 组共享,大幅节省
"head_dim": 128,
},
"GPT-4 (MQA 假设)": {
"num_layers": 96,
"num_kv_heads": 1, # MQA:极致节省
"head_dim": 128,
},
}
for model_name, cfg in configs.items():
for seq_len in [4096, 32768, 131072]:
gb = estimate_kv_cache_gb(
num_layers=cfg["num_layers"],
seq_len=seq_len,
num_kv_heads=cfg["num_kv_heads"],
head_dim=cfg["head_dim"],
batch_size=1,
)
# 输出格式:模型名 | 序列长度 | KV Cache GB
print(f"{model_name:<25} | seq={seq_len:>7} | KV Cache: {gb:.2f} GB")
4.2 KV Cache 的三重瓶颈
KV Cache 引入了三个在生产环境中频繁触碰的瓶颈:
| 瓶颈类型 | 成因 | 影响 | 工程解法 |
|---|---|---|---|
| 显存容量 | KV Cache 随 seq_len × batch_size 线性增长 | 单卡最大并发数受限,长上下文时可能 OOM | GQA/MQA、KV 量化、PagedAttention |
| 显存带宽 | decode 每步都要从 HBM 读取完整 KV Cache | 吞吐量(tokens/s)瓶颈,I/O bound | 增大 batch size、FlashDecoding、投机解码 |
| 显存碎片 | 不同请求序列长度不一,预分配内存浪费 | 实际显存利用率 < 20%(naive 实现) | PagedAttention(vLLM) |
4.3 MHA/GQA/MQA:减少 KV 头数
减少 KV Cache 最直接的架构级优化是减少 Key/Value 的头数:
- MHA(Multi-Head Attention):每个 Q 头都有独立的 K、V。num_kv_heads = num_q_heads。KV Cache 最大。
- MQA(Multi-Query Attention,Shazeer 2019):所有 Q 头共享同一组 K、V。num_kv_heads = 1。KV Cache 节省 num_heads 倍,但注意力多样性下降。
- GQA(Grouped Query Attention,Ainslie 2023):折中方案,将 Q 头分成 g 组,每组共享一对 K、V。num_kv_heads = g。LLaMA-3、Mistral、Qwen 均采用。
ini
以 LLaMA-3 70B 为例(序列 4K,batch=1,FP16):
MHA (64 Q heads = 64 KV heads):
2 × 80 × 4096 × 64 × 128 × 2 = 10.7 GB
GQA (64 Q heads, 8 KV heads):
2 × 80 × 4096 × 8 × 128 × 2 = 1.34 GB ← 节省 8×
MQA (64 Q heads, 1 KV head):
2 × 80 × 4096 × 1 × 128 × 2 = 0.17 GB ← 节省 64×
4.4 PagedAttention:消除显存碎片
PagedAttention(Kwon et al. 2023,vLLM 核心创新)借鉴操作系统虚拟内存分页的思想,将 KV Cache 组织为固定大小的"页"(默认 16 token/页),用逻辑→物理页映射表按需分配:
less
传统 KV Cache 管理(预分配连续内存):
┌──────────────────────────────────┐
│ 请求A: [KV][KV][KV][ ][ ][ ][ ] │ ← 按 max_length 预留,浪费 4/7
│ 请求B: [KV][KV][ ][ ][ ][ ][ ] │ ← 浪费 5/7
│ 请求C: [KV][KV][KV][KV][KV][ ] │ ← 浪费 1/6
└──────────────────────────────────┘
实际利用率:通常 < 20%
PagedAttention(按页动态分配):
物理显存池:[page0][page1][page2][page3][page4][page5][page6][page7]...
请求A 逻辑页表: 0→page0, 1→page3, 2→page7
请求B 逻辑页表: 0→page1, 1→page5
请求C 逻辑页表: 0→page2, 1→page4, 2→page6
利用率:接近 100%(只有最后一页可能有碎片)
额外收益:
- Prefix Sharing:相同 system prompt 的请求共享物理页,无需重复计算
- Copy-on-Write:beam search 多候选共享父节点的 KV 页
vLLM 在 PagedAttention 基础上实现了 Continuous Batching,将吞吐量相比朴素实现提升 2--24 倍。这是目前生产推理服务的事实标准。
4.5 KV Cache 量化
KV Cache 量化将 FP16 K/V 压缩为 INT8 甚至 INT4,可节省 2--4 倍显存,但存在精度风险:K 矩阵中存在数值异常值(outliers),集中在少数通道,均匀量化会导致这些通道截断,影响长距离依赖的建模精度。
KV 量化的陷阱
实验数据显示,INT4 KV Cache 在短文本(< 4K tokens)任务上困惑度几乎无损,但在 NIAH 长上下文检索任务中,准确率可能下降 5--15%。工程决策:短文本 QA 场景 INT8 KV 风险可接受;长文档摘要、多轮 Agent 场景建议谨慎,必须跑 NIAH 回归再上线。
4.6 Prompt Caching(前缀缓存)
Prompt Caching 是 KV Cache 在服务侧的一种特殊利用方式:对于具有相同前缀(如固定 system prompt)的多个请求,服务端将前缀的 KV Cache 序列化到存储,后续请求命中时跳过前缀的 prefill 计算:
ini
"""
Prompt Caching 工程最佳实践
Anthropic API 示例(Claude 3.5 Sonnet)
命中条件:前缀 token 序列必须字节级精确匹配
TTL:约 5 分钟(Anthropic),会话级(部分 OpenAI 模型)
成本降低:命中时输入 token 成本降低约 90%
延迟降低:TTFT(首 token 延迟)降低 30--70%
"""
import anthropic
# 固定的系统提示(每次请求保持不变,可命中缓存)
SYSTEM_PROMPT = """
你是一个专业的代码审查助手,具有以下能力:
1. 识别安全漏洞(SQL 注入、XSS、路径穿越等)
2. 发现性能瓶颈(N+1 查询、无效循环、内存泄漏)
3. 检查代码规范(命名、注释、结构)
4. 提供具体的修复建议和代码示例
审查标准:遵循 OWASP Top 10 安全标准,Google 工程实践。
输出格式:严重程度(Critical/High/Medium/Low)+ 问题描述 + 修复建议。
""" * 10 # 模拟一个较长的 system prompt(约 2000 tokens)
# 每次调用时,system_prompt 内容不变 → 可命中缓存
# 只有 user message 变化
def review_code(code_snippet: str, filename: str) -> str:
client = anthropic.Anthropic()
response = client.messages.create(
model="claude-3-5-sonnet-20241022",
max_tokens=2048,
system=SYSTEM_PROMPT, # 固定前缀,命中缓存后成本降 90%
messages=[
{
"role": "user",
# 动态内容放在末尾,不影响前缀匹配
"content": f"请审查以下代码文件 `{filename}`:\n\n```\n{code_snippet}\n```"
}
]
)
# 检查缓存命中情况(Anthropic API 返回)
usage = response.usage
cache_read = getattr(usage, 'cache_read_input_tokens', 0)
cache_write = getattr(usage, 'cache_creation_input_tokens', 0)
if cache_read > 0:
# 命中缓存:只有 cache_read 而无 cache_write
print(f"✅ Prompt Cache 命中: 节省 {cache_read} token 的 prefill 计算")
elif cache_write > 0:
# 首次请求:创建缓存
print(f"📝 Prompt Cache 写入: {cache_write} tokens 已缓存")
return response.content[0].text
# 关键:破坏缓存命中的反模式
def bad_practice_breaks_cache(code_snippet: str) -> str:
"""
反模式:在 system prompt 中插入动态内容
每次请求的 system prompt 前缀不同 → 永远无法命中缓存
"""
import datetime
client = anthropic.Anthropic()
# ❌ 错误:timestamp 让每次 system prompt 都不同
bad_system = f"""
当前时间:{datetime.datetime.now().isoformat()} ← 破坏缓存!
你是代码审查助手...
"""
response = client.messages.create(
model="claude-3-5-sonnet-20241022",
max_tokens=2048,
system=bad_system,
messages=[{"role": "user", "content": code_snippet}]
)
return response.content[0].text
5. Lost in the Middle:长上下文不等于长上下文利用
5.1 现象描述
2023 年,Liu 等人发表了颇具影响力的论文《Lost in the Middle: How Language Models Use Long Contexts》。研究发现:当关键信息被放置在长上下文的中间位置时,模型检索该信息的准确率显著低于放置在开头或结尾的情况,呈现出明显的 U 型曲线。
matlab
20 篇文档检索实验(GPT-3.5,问答任务):
文档位置 | 准确率
---------|--------
最开头 | 78% ← 首因效应(Primacy Bias)
第 5 位 | 65%
第 10 位 | 52% ← 中间最低点
第 15 位 | 60%
最末尾 | 81% ← 近因效应(Recency Bias)
U 型曲线:首位和末位差不多,中间显著下降
5.2 机制解释
这一现象有三个叠加的机制:
- 注意力汇聚(Attention Sink)效应:研究发现,几乎所有注意力头都会给序列最开头的几个 token(通常是 BOS token)分配异常高的注意力权重,形成"注意力汇"(attention sink)。这些 token 本身没有语义价值,但 softmax 的归一化特性使模型需要一个"垃圾桶"来吸收多余权重,BOS token 因为始终在场而成为默认汇聚点。
- 近因偏好(Recency Bias):decoder-only 模型在生成答案时,距离 query 最近的 token 的梯度路径最短,注意力权重自然偏高。这也是 RAG 中应将最相关文档放在用户 query 正上方(即末尾)的数学依据。
- 位置编码精度衰减:RoPE 对相对距离较小的 token 对编码精度更高,远距离 token 的旋转角度累积误差较大,影响相似度计算的精确性。
5.3 NIAH:大海捞针测试
**NIAH(Needle in a Haystack,大海捞针)**是长上下文能力的标准压力测试,已成为模型选型的必跑基准:
python
"""
NIAH 测试框架实现
评测矩阵:横轴 = 文档长度,纵轴 = 针的相对位置
理想模型:全绿(100% 准确率)
"""
import random
import anthropic
from typing import Callable
# 测试"针"(目标信息)
NEEDLE = "Claude 的暗号是:紫色大象在月球上喝茶。"
# 生成"草堆"(无关填充文本)
def generate_haystack(target_length_tokens: int) -> str:
"""生成约 target_length_tokens token 的无关内容"""
# 用重复段落填充到目标长度
paragraph = """
在现代软件工程中,微服务架构已经成为构建大规模分布式系统的主流范式。
每个微服务负责一个独立的业务领域,通过 REST API 或消息队列进行通信。
服务发现、负载均衡、熔断器等模式是微服务架构的基础组件。
容器化技术(Docker/Kubernetes)使得微服务的部署和扩展变得更加便捷。
"""
# 简单估算:中文约 1.5 字符/token
target_chars = target_length_tokens * 1.5
repetitions = int(target_chars / len(paragraph)) + 1
return (paragraph * repetitions)[:int(target_chars)]
def run_niah_test(
model: str,
doc_lengths: list[int],
needle_positions: list[float], # 0.0 = 开头, 1.0 = 结尾
call_api: Callable,
) -> dict:
"""
运行 NIAH 测试并返回准确率矩阵
Returns:
results[doc_length][position] = is_correct (bool)
"""
results = {}
for doc_len in doc_lengths:
results[doc_len] = {}
haystack = generate_haystack(doc_len)
for position in needle_positions:
# 将针插入草堆的指定位置
insert_idx = int(len(haystack) * position)
context = (
haystack[:insert_idx]
+ f"\n\n【重要信息】{NEEDLE}\n\n"
+ haystack[insert_idx:]
)
# 构造查询
prompt = f"""
以下是一段文档内容:
{context}
---
问题:Claude 的暗号是什么?请直接回答暗号内容,不要解释。
"""
# 调用模型
response = call_api(model=model, prompt=prompt)
# 检查是否包含正确答案
is_correct = "紫色大象在月球上喝茶" in response
results[doc_len][position] = is_correct
print(
f"doc_len={doc_len:>7} | position={position:.1%} | "
f"{'✅' if is_correct else '❌'}"
)
return results
def visualize_niah_results(results: dict) -> str:
"""
将 NIAH 结果可视化为 ASCII 热图
绿色(✅)= 正确,红色(❌)= 错误
"""
doc_lengths = sorted(results.keys())
positions = sorted(results[doc_lengths[0]].keys())
lines = ["位置\\长度 " + " ".join(f"{l//1000}K" for l in doc_lengths)]
for pos in positions:
row = f"{pos:.0%} "
for dl in doc_lengths:
row += "✅ " if results[dl][pos] else "❌ "
lines.append(row)
return "\n".join(lines)
# 实际选型建议:
# - 跑完 NIAH 矩阵,检查是否存在红色区域(特别是 40%--70% 位置)
# - 对比多个模型,选中间位置仍保持 >85% 准确率的模型
# - 若预算有限,考虑 RAG 策略(把相关文档放末尾)代替盲目扩窗
5.4 工程应对策略
针对 Lost in the Middle 问题,实际工程中有以下缓解手段:
| 策略 | 原理 | 适用场景 | 代价 |
|---|---|---|---|
| 关键信息置于末尾 | 利用近因效应,最相关文档紧贴 user query | RAG、few-shot | 无额外成本 |
| Reranker 精排 | 减少文档数量,只保留最相关的 3--5 篇 | RAG 管道 | 增加一次 cross-encoder 推理延迟 |
| 上下文压缩 | LLM 先摘要,减少噪声 | 长文档问答 | 增加一次 LLM 调用成本 |
| 结构化标注 | 给文档编号、添加标签,引导模型引用 | 多文档综合 | 轻微 token 增加 |
| 冗余放置 | 关键信息在开头和结尾各放一次 | 安全要求高的场景 | token 成本翻倍 |
6. 扩展上下文窗口的技术路线
将一个训练在 4K 窗口的模型扩展到 128K,是工业界过去两年的核心工程挑战之一。主要有以下几条技术路线:
6.1 位置插值(Position Interpolation,PI)
Chen et al. 2023 提出,可以将推理时的位置索引线性压缩到训练范围内。若训练最长序列 L_train = 4096,推理目标长度 L_infer = 32768,则将位置索引 m 缩放为 m × (L_train / L_infer):
arduino
原始推理(超出训练范围):
token at position 32768 → 模型从未见过,性能崩溃
位置插值:
token at position 32768 → 映射到位置 32768 × (4096/32768) = 4096
模型"以为"序列只有 4096 长 → 在训练范围内,稳定
问题:
所有维度的频率都被等比压缩 1/8
高频维度(区分相邻 token)的分辨率降低
对短距离依赖任务有轻微损失
需要在长序列上进行少量微调(约 1000 步)来适应压缩后的分布
6.2 NTK-Aware Scaling
NTK-Aware Scaling(Reddit 社区提出,后经学界验证)的核心洞察:RoPE 的不同频率维度不应该被均匀压缩。低频维度(处理长距离关系)应该多外推,高频维度(处理短距离关系)应该少压缩。
具体做法:增大 RoPE 的 base 参数(原始通常为 10000),相当于对低频维度做更多的外推:
ini
原始 RoPE:base = 10000
θ_d = 10000^(-2d/D) 每个维度的旋转频率
NTK-Aware Scaling(目标扩窗倍数 k=8):
new_base = base × k^(D/(D-2))
= 10000 × 8^(D/(D-2))
≈ 10000 × 8 = 80000(约估)
效果:
- 高频维度(小 d)变化轻微 → 保持短距离分辨率
- 低频维度(大 d)被拉伸 → 支持更长的相对距离
优点:可以免微调或极少量微调直接使用
缺点:需要为目标扩窗倍数调整 base,超出预设倍数效果下降
6.3 YaRN(Yet another RoPE extensioN method)
YaRN(Peng et al. 2023)在 NTK-Aware 基础上进一步引入了两个改进:
- 频率选择性插值:对不同频率维度分别采用插值(高频)或外推(低频),而非统一处理。
- 注意力温度缩放 :长序列时注意力 softmax 的 logits 方差增大,引入温度系数 t 补偿:
scale = 0.1 × ln(s) + 1(s 为扩窗倍数)。
| 方案 | 是否需要微调 | 核心手段 | 代表应用 | PPL 损失(相对 PI) |
|---|---|---|---|---|
| Position Interpolation (PI) | 需要(~1000 步) | 线性缩放位置索引 | LLaMA 早期扩窗 | 基准 |
| NTK-Aware Scaling | 可免调 | 调整 RoPE base | Mistral 7B 早期 | -0.3(更好) |
| YaRN | 需要(~400 步) | NTK + 频率选择 + 温度 | LLaMA-3 扩窗 | -0.8(更好) |
| LongRoPE | 需要 | 非均匀位置插值 | Phi-3 Medium | -1.0(最好) |
| 直接长上下文预训练 | 完整重训 | 训练时直接用长序列 | GPT-4o、Claude 3 | 0(无损失,成本最高) |
6.4 滑动窗口注意力(SWA)与 StreamingLLM
**SWA(Sliding Window Attention)**让每个 token 只能 attend 其前 W 个 token,将 O(n²) 降至 O(n·W),同时 KV Cache 固定在 W×d 而非 n×d:
diff
Mistral 7B 的 SWA 配置:
- W = 4096(局部窗口大小)
- 32 层
- 有效感受野 = 32 × 4096 ≈ 131K token(通过多层叠加)
KV Cache 固定大小 = 4096 token(无论序列多长)
代价:
- 对跨越窗口边界的直接依赖有损
- 不适合需要精确引用远距离文本的任务
StreamingLLM(Xiao et al. 2023)发现了 attention sink 现象,并利用它实现了无限长流式生成:始终保留序列最开头的 4 个 token 的 KV(作为 attention sink),加上最近 W 个 token 的 KV,使滑动窗口不崩溃:
css
StreamingLLM KV Cache 结构:
[BOS][.][.][.] + [...最近 W 个 token...]
attention sink 滑动窗口
总 KV Cache 大小:4 + W(固定)
支持:无限长流式输入,不 OOM,不崩溃
用途:
- 在线客服长会话
- 实时字幕生成
- 无限流式文档处理
限制:
- 不能精确引用超过 W 个 token 以前的内容
- 适合"当前轮"理解,不适合"跨越历史"推理
6.5 扩窗后的质量验证流程
扩窗不是一步操作,而是一个需要严格验证的流程:
python
"""
长上下文扩窗质量验证 Checklist
扩窗后必须通过以下所有验证才能上线
"""
VALIDATION_SUITE = {
# 1. 长上下文能力验证(核心)
"niah_full_matrix": {
"描述": "大海捞针全矩阵测试",
"覆盖": "所有目标长度 × 所有插入位置(0%--100%)",
"通过标准": "中间位置(30%--70%)准确率 ≥ 85%",
"工具": "niah_eval.py(自定义)",
},
"niah_multi_needle": {
"描述": "多针测试(关键信息分散在不同位置)",
"覆盖": "同时检索 3--5 条独立信息",
"通过标准": "全部正确率 ≥ 70%",
},
"long_doc_qa": {
"描述": "长文档问答(SCROLLS、∞Bench)",
"通过标准": "F1 ≥ 目标模型的 -5%(与原窗口长度基准对比)",
},
# 2. 短文本回归(防止遗忘)
"mmlu_regression": {
"描述": "MMLU 标准 benchmark",
"通过标准": "与扩窗前相比,下降 ≤ 1%",
},
"humaneval_regression": {
"描述": "HumanEval 代码生成",
"通过标准": "pass@1 下降 ≤ 2%",
},
# 3. 性能基准
"latency_ttft": {
"描述": "首 token 延迟(TTFT)在目标上下文长度下",
"通过标准": "TTFT P99 ≤ SLA 要求",
},
"throughput": {
"描述": "在目标 batch size 下的 token/s 吞吐量",
"通过标准": "tokens/s ≥ 上线前的 80%",
},
"kv_cache_memory": {
"描述": "最大序列长度下 KV Cache 显存占用",
"通过标准": "不超过可用显存的 70%(留 30% 余量)",
},
}
# CI 流水线:扩窗 → 验证 → 灰度 → 全量
PIPELINE_STAGES = [
"位置编码调整(YaRN/NTK)",
"长序列微调(< 1000 steps)",
"运行 validation_suite(全部通过才继续)",
"A/B 测试(5% 流量灰度,关注长上下文任务指标)",
"全量上线",
]
7. 真实工程案例
案例一:某 SaaS 企业的代码审查 Agent------KV Cache 显存爆炸排查
背景(已脱敏)
某中型 SaaS 公司(员工约 300 人)构建了一个代码审查 Agent,使用私有化部署的 LLaMA-3 70B(4×A100 80GB,张量并行),为工程师提供代码审查服务。系统上线初期正常,随着使用量增加开始频繁出现 OOM(Out of Memory)崩溃。
故障现象:服务在高并发时段(早 10 点、下午 2 点)出现 CUDA OOM,每次重启需约 3 分钟,影响工程师体验。监控显示显存利用率在故障前 10 分钟从正常的 65% 快速爬升到 95%+。
排查过程:
python
"""
排查工具:显存追踪脚本(事后加入 CI 监控)
"""
import torch
import psutil
from vllm import LLM, SamplingParams
from vllm.engine.metrics import RayPrometheusStatLogger
def diagnose_kv_cache_usage(llm_engine):
"""
采集 KV Cache 使用情况,用于诊断 OOM
此函数在 OOM 排查期间每 30 秒被调用一次
"""
# 获取 KV Cache block 使用情况(vLLM 暴露的指标)
stats = llm_engine.get_engine_stats()
gpu_kv_cache_usage = stats.get("gpu_cache_usage_perc", 0)
num_running_requests = stats.get("num_running_requests", 0)
num_waiting_requests = stats.get("num_waiting_requests", 0)
print(f"""
KV Cache 使用率: {gpu_kv_cache_usage:.1%}
正在处理请求数: {num_running_requests}
等待中请求数: {num_waiting_requests}
""")
# 警戒阈值:KV Cache 使用率 > 80% 时触发告警
if gpu_kv_cache_usage > 0.8:
print("⚠️ KV Cache 使用率过高,即将触发 OOM!")
# 触发告警(接入公司内部 Slack/PagerDuty)
trigger_alert(
severity="WARNING",
message=f"KV Cache 使用率达到 {gpu_kv_cache_usage:.1%},"
f"当前 {num_running_requests} 个请求,"
f"{num_waiting_requests} 个等待"
)
return gpu_kv_cache_usage
def analyze_request_distribution():
"""
分析请求的 token 分布,找到 OOM 根因
"""
# 从日志中提取最近 1000 条请求的 token 数
# 发现:99% 的请求 input < 8K,但有约 1% 的请求 input > 60K
# 这 1% 的请求独占了 > 50% 的 KV Cache 显存!
token_stats = {
"P50_input_tokens": 2048,
"P95_input_tokens": 7500,
"P99_input_tokens": 62000, # ← 异常:有人把整个代码仓库塞进去了
"max_input_tokens": 98000,
}
return token_stats
根因:发现约 1% 的请求是工程师将整个代码仓库(有时超过 50K tokens)作为上下文传入,这些超长请求会独占大量 KV Cache block,导致正常请求被饿死,最终触发 OOM。
解决方案(三层防御):
python
"""
三层 KV Cache OOM 防御策略
Layer 1: 请求层限制
Layer 2: vLLM 调度优化
Layer 3: 动态 KV 预算控制
"""
from fastapi import HTTPException
import tiktoken
# === Layer 1: 请求层限制 ===
MAX_INPUT_TOKENS = 32768 # 硬性上限:32K tokens
def validate_request_tokens(system_prompt: str, messages: list) -> int:
"""
在请求到达 LLM 引擎前,先检查 token 数
超出上限直接拒绝,避免占用宝贵的 KV Cache
"""
# 使用近似计数(节省成本,精确值不关键)
enc = tiktoken.get_encoding("cl100k_base")
total_tokens = len(enc.encode(system_prompt))
for msg in messages:
total_tokens += len(enc.encode(msg.get("content", "")))
if total_tokens > MAX_INPUT_TOKENS:
raise HTTPException(
status_code=400,
detail={
"error": "input_too_long",
"message": f"输入 {total_tokens} tokens 超过上限 {MAX_INPUT_TOKENS},"
f"请使用 /summarize 接口先压缩代码,或拆分为多次请求",
"input_tokens": total_tokens,
"max_tokens": MAX_INPUT_TOKENS,
}
)
return total_tokens
# === Layer 2: vLLM 调度优化 ===
VLLM_CONFIG = {
# 为超长请求设置更低优先级
"max_num_seqs": 64, # 最大并发序列数
"max_model_len": 32768, # 截断超出此长度的请求
"gpu_memory_utilization": 0.85, # 只用 85% 显存给 KV Cache(留 15% 给模型权重和激活)
# PagedAttention 配置
"block_size": 16, # 每页 16 tokens
"swap_space": 4, # 4 GB CPU 内存作为 KV Cache overflow
# Continuous Batching:动态调度,不等最长序列完成
"scheduler_type": "continuous",
}
# === Layer 3: 动态 KV 预算控制 ===
def build_context_with_budget(
system_prompt: str,
conversation_history: list,
current_query: str,
code_context: str,
budget_tokens: int = 16384, # 为代码审查场景设置的预算
) -> str:
"""
动态构建上下文,确保总 token 数不超过预算
分配策略:
- system_prompt(固定):约 1000 tokens → 命中 prompt cache
- recent_history(最近 6 轮):约 3000 tokens → 保证对话连贯性
- code_context:剩余预算 → 重要代码放末尾(避免 Lost in Middle)
- current_query:约 200 tokens → 固定
- output_reserve:4096 tokens → 给模型足够的输出空间
"""
enc = tiktoken.get_encoding("cl100k_base")
# 固定分配
output_reserve = 4096
query_tokens = len(enc.encode(current_query))
system_tokens = len(enc.encode(system_prompt))
# 计算可用于代码上下文的预算
history_budget = min(3000, budget_tokens // 4)
code_budget = (
budget_tokens
- system_tokens
- query_tokens
- history_budget
- output_reserve
)
# 截断代码上下文(保留末尾,因为最末尾靠近 query,注意力最强)
code_tokens = enc.encode(code_context)
if len(code_tokens) > code_budget:
# 保留末尾(更重要的部分)
code_tokens = code_tokens[-code_budget:]
code_context = enc.decode(code_tokens)
code_context = "...(代码已截断,显示最近部分)...\n" + code_context
# 截断历史对话(保留最近几轮)
trimmed_history = trim_history_to_budget(
conversation_history,
history_budget,
enc
)
return f"{system_prompt}\n\n{trimmed_history}\n\n{code_context}\n\n{current_query}"
结果:上线三层防御策略后,OOM 频率从每日约 3--5 次降至 0。P95 延迟从 8.2 秒降至 4.1 秒(因为不再有超长请求占用 KV Cache)。显存利用率稳定在 70--80%。
案例二:某内容平台的长文档摘要服务------从"塞满窗口"到"智能分层"
背景(已脱敏)
某内容平台(日活约 500 万)为作者提供长文档智能摘要功能,处理对象包括法律合同(50--300 页)、学术论文、行业报告等。初版方案直接使用 Gemini 1.5 Pro 1M 窗口,将整个文档塞入,效果差且成本高昂。
问题量化:
- 成本:平均每份 200 页合同约 150K tokens,每次摘要成本约 0.30(2/M tokens × 150K)
- 质量:第三方人工评估,"关键条款遗漏率"约 28%(多集中在文档中间部分)
- 延迟:TTFT(首 token 延迟)P99 约 22 秒(150K prefill 时间)
重构方案:分层上下文策略:
javascript
/**
* 长文档摘要服务 - 分层处理策略
* 将"塞满窗口"改为"智能分层"
*
* 策略:
* 1. 文档结构分析 → 提取目录/章节边界
* 2. 并行摘要各章节(短窗口,成本低)
* 3. 二次摘要合并(将章节摘要合成全局摘要)
* 4. 关键条款检测(对高风险章节单独精读)
*/
const Anthropic = require("@anthropic-ai/sdk");
const client = new Anthropic();
/**
* 估算文本的 token 数(粗略估算,用于预算控制)
* 中文:约 1.5 字符/token;英文:约 4 字符/token
* @param {string} text - 输入文本
* @param {string} lang - 语言('zh' 或 'en')
* @returns {number} 估算 token 数
*/
function estimateTokens(text, lang = "zh") {
const charsPerToken = lang === "zh" ? 1.5 : 4;
return Math.ceil(text.length / charsPerToken);
}
/**
* 按章节分割文档
* 检测常见章节标记:第一章、1.、一、等
* @param {string} docText - 完整文档文本
* @returns {Array<{title: string, content: string, index: number}>}
*/
function splitIntoChapters(docText) {
// 匹配常见章节标记(简化版,实际需更鲁棒的规则)
const chapterPattern =
/^(第[一二三四五六七八九十百\d]+[章节条款]|[(\(]?\d+[)\)]\s|[一二三四五六七八九十]+、)/gm;
const chapters = [];
const matches = [...docText.matchAll(chapterPattern)];
for (let i = 0; i < matches.length; i++) {
const start = matches[i].index;
const end = i + 1 < matches.length ? matches[i + 1].index : docText.length;
chapters.push({
title: matches[i][0].trim(),
content: docText.slice(start, end).trim(),
index: i,
tokenEstimate: estimateTokens(docText.slice(start, end), "zh"),
});
}
return chapters;
}
/**
* 摘要单个章节(短窗口请求,成本低)
* @param {string} title - 章节标题
* @param {string} content - 章节内容
* @param {string} docType - 文档类型(用于调整摘要策略)
* @returns {Promise<{summary: string, keyPoints: string[], riskLevel: string}>}
*/
async function summarizeChapter(title, content, docType = "合同") {
// 使用更便宜的模型处理单章节(8K 窗口足够)
const response = await client.messages.create({
model: "claude-3-5-haiku-20241022", // 比 Sonnet 便宜约 5×
max_tokens: 512,
system: `你是${docType}分析专家。请对章节内容进行简洁摘要,并识别关键条款和风险点。
输出 JSON 格式:{"summary": "摘要", "keyPoints": ["要点1", "要点2"], "riskLevel": "low/medium/high"}`,
messages: [
{
role: "user",
// 将章节放在末尾(靠近输出,注意力最强)
content: `章节:${title}\n\n${content}`,
},
],
});
try {
return JSON.parse(response.content[0].text);
} catch {
// 解析失败时返回降级结果
return {
summary: response.content[0].text.slice(0, 200),
keyPoints: [],
riskLevel: "unknown",
};
}
}
/**
* 并行处理所有章节,然后合并生成全局摘要
* 这是核心的"分层"策略
* @param {string} docText - 完整文档
* @param {string} docType - 文档类型
* @returns {Promise<string>} 最终全局摘要
*/
async function hierarchicalSummarize(docText, docType = "合同") {
// 1. 分割文档
const chapters = splitIntoChapters(docText);
console.log(`文档共 ${chapters.length} 章节,总估算 ${estimateTokens(docText, "zh")} tokens`);
// 2. 并行摘要各章节(关键:并行而非串行,降低延迟)
const chapterSummaries = await Promise.all(
chapters.map((chapter) =>
summarizeChapter(chapter.title, chapter.content, docType)
)
);
// 3. 识别高风险章节(需要精读)
const highRiskChapters = chapters.filter(
(_, i) => chapterSummaries[i].riskLevel === "high"
);
// 4. 对高风险章节进行精读(使用更强的模型)
const riskAnalyses = await Promise.all(
highRiskChapters.map((chapter) =>
client.messages.create({
model: "claude-3-5-sonnet-20241022", // 高风险章节用强模型
max_tokens: 1024,
system: `你是法律风险分析专家。请对以下${docType}章节进行深度风险分析,
识别潜在的法律风险、不平等条款、模糊表述等。`,
messages: [
{
role: "user",
content: chapter.content,
},
],
})
)
);
// 5. 合并所有章节摘要生成全局摘要
// 注意:将高风险章节摘要放在末尾(靠近问题,注意力最强)
const regularSummaries = chapterSummaries
.filter((_, i) => chapterSummaries[i].riskLevel !== "high")
.map((s, i) => `${chapters[i].title}: ${s.summary}`)
.join("\n");
const riskSummaries = highRiskChapters
.map((chapter, i) => `⚠️ ${chapter.title}: ${riskAnalyses[i].content[0].text}`)
.join("\n");
const finalResponse = await client.messages.create({
model: "claude-3-5-sonnet-20241022",
max_tokens: 2048,
system: `你是${docType}摘要专家。请根据各章节摘要,生成结构化的全文摘要,重点突出关键条款和风险点。`,
messages: [
{
role: "user",
// 将风险章节放最后(末尾位置注意力更强)
content: `各章节摘要如下:\n\n${regularSummaries}\n\n高风险章节详细分析:\n${riskSummaries}\n\n请生成全文综合摘要。`,
},
],
});
return finalResponse.content[0].text;
}
// 使用示例
// const summary = await hierarchicalSummarize(contractText, "劳动合同");
结果对比:
| 指标 | 原方案(塞满 1M 窗口) | 新方案(分层处理) | 改善幅度 |
|---|---|---|---|
| 每份文档成本 | $0.30 | $0.04(Haiku 章节 + Sonnet 合并) | -87% |
| 关键条款遗漏率 | 28% | 11%(高风险章节精读) | -61% |
| P99 延迟 | 22 秒 | 6.5 秒(并行章节 + 短 prefill) | -70% |
| 月账单(10 万份/月) | $30,000 | $4,000 | -87% |
8. 上下文管理工程策略:从预算控制到分层架构
8.1 Token 预算控制
将上下文窗口视为有限预算,而非无限空间,是多轮对话和 Agent 系统设计的第一原则。
python
"""
生产级 Token 预算管理器
适用于多轮对话、RAG、Agent 等场景
"""
import tiktoken
from dataclasses import dataclass
from enum import Enum
from typing import Optional
class Priority(Enum):
"""上下文组件的优先级(越高越优先保留)"""
CRITICAL = 1 # system prompt、当前 query:绝不裁剪
HIGH = 2 # 最近几轮对话:尽量保留
MEDIUM = 3 # RAG 检索结果:按预算裁剪
LOW = 4 # 远期历史摘要:最先压缩
@dataclass
class ContextComponent:
"""上下文组件:有名字、内容、优先级"""
name: str
content: str
priority: Priority
is_cached: bool = False # 是否为 prompt cache 候选(固定前缀)
class TokenBudgetManager:
"""
Token 预算管理器
使用方式:
1. 设置总 budget
2. 按优先级添加各组件
3. 调用 build() 得到在预算内的最终 prompt
"""
def __init__(self, model: str, total_budget: int, output_reserve: int = 4096):
"""
Args:
model: 模型名称(用于选择 tokenizer)
total_budget: 上下文总 token 预算
output_reserve: 为模型输出预留的 token 数
"""
self.model = model
# 实际可用于输入的 token 预算
self.input_budget = total_budget - output_reserve
self.components: list[ContextComponent] = []
# 根据模型选择 tokenizer
# 注意:对于非 OpenAI 模型,cl100k 是合理的近似
try:
self.enc = tiktoken.encoding_for_model(model)
except KeyError:
# 未知模型使用 cl100k 近似(中文内容误差可达15%,英文通常 < 5%)
self.enc = tiktoken.get_encoding("cl100k_base")
def count(self, text: str) -> int:
"""计算文本的 token 数"""
return len(self.enc.encode(text))
def add_component(
self,
name: str,
content: str,
priority: Priority,
is_cached: bool = False,
) -> "TokenBudgetManager":
"""添加上下文组件(链式调用)"""
self.components.append(
ContextComponent(
name=name,
content=content,
priority=priority,
is_cached=is_cached,
)
)
return self # 支持链式调用
def build(self) -> tuple[str, dict]:
"""
按优先级构建最终 prompt
Returns:
(prompt_text, budget_report)
"""
# 按优先级排序(高优先级先放)
sorted_components = sorted(
self.components,
key=lambda c: c.priority.value
)
included = []
used_tokens = 0
report = {}
for component in sorted_components:
token_count = self.count(component.content)
if used_tokens + token_count <= self.input_budget:
# 预算充足,直接放入
included.append(component)
used_tokens += token_count
report[component.name] = {
"status": "included",
"tokens": token_count,
}
elif component.priority == Priority.CRITICAL:
# CRITICAL 组件:必须放入,即使超预算(上层应避免此情况)
included.append(component)
used_tokens += token_count
report[component.name] = {
"status": "forced_included",
"tokens": token_count,
"warning": "超出预算,但 CRITICAL 优先级强制放入",
}
else:
# 预算不足:尝试截断放入一半
remaining = self.input_budget - used_tokens
if remaining > 100: # 至少放 100 token 有意义
tokens = self.enc.encode(component.content)
truncated_tokens = tokens[:remaining]
truncated_content = self.enc.decode(truncated_tokens)
truncated_component = ContextComponent(
name=component.name,
content="...(内容已截断)...\n" + truncated_content,
priority=component.priority,
)
included.append(truncated_component)
used_tokens += remaining
report[component.name] = {
"status": "truncated",
"tokens_requested": token_count,
"tokens_used": remaining,
}
else:
# 完全放不下
report[component.name] = {
"status": "excluded",
"tokens": token_count,
}
# 拼接最终 prompt(CRITICAL 和 is_cached 的放最前面)
cached_parts = [c for c in included if c.is_cached]
dynamic_parts = [c for c in included if not c.is_cached]
prompt = "\n\n".join(
c.content for c in cached_parts + dynamic_parts
)
report["_summary"] = {
"total_input_budget": self.input_budget,
"total_used": used_tokens,
"utilization": f"{used_tokens / self.input_budget:.1%}",
}
return prompt, report
# ── 使用示例:多轮 Agent 上下文构建 ──────────────────────────────────
def build_agent_context(
system_prompt: str,
conversation_history: list[dict],
retrieved_docs: list[str],
current_query: str,
model: str = "claude-3-5-sonnet-20241022",
total_budget: int = 200_000,
) -> tuple[str, dict]:
"""
构建 Agent 的上下文
遵循分层原则:
- system_prompt 置顶(命中 prompt cache)
- 历史摘要(早期压缩)
- 最近 6 轮对话(全量保留)
- RAG 检索结果(按预算截断)
- 当前 query(末尾,注意力最强)
"""
# 将历史对话转为文本
recent_history_text = "\n".join(
f"{msg['role'].upper()}: {msg['content']}"
for msg in conversation_history[-12:] # 最近 6 轮(12 条消息)
)
# 将 RAG 结果合并
rag_text = "\n\n---\n\n".join(
f"[参考文档 {i+1}]\n{doc}"
for i, doc in enumerate(retrieved_docs[:5]) # 最多 5 篇
)
manager = TokenBudgetManager(
model=model,
total_budget=total_budget,
output_reserve=8192, # 为模型输出预留 8K
)
# 按优先级添加组件
# is_cached=True 的组件应放最前面(利用 prompt cache)
(
manager
.add_component("system", system_prompt, Priority.CRITICAL, is_cached=True)
.add_component("recent_history", recent_history_text, Priority.HIGH)
.add_component("rag_context", rag_text, Priority.MEDIUM)
.add_component("current_query", current_query, Priority.CRITICAL)
)
return manager.build()
8.2 滚动摘要策略(Rolling Summary)
对于长时间运行的 Agent,历史对话会持续增长并耗尽 token 预算。滚动摘要策略在历史超过阈值时自动触发压缩:
css
滚动摘要工作流:
轮次 1-10:完整保存历史(1,200 tokens)
↓ 超过阈值(2,000 tokens)
轮次 11:触发摘要
LLM 压缩轮次 1-8 → 摘要(300 tokens)
保留轮次 9-10 完整原文(300 tokens)
结构:[摘要:1-8] + [原文:9-10] + [新轮次:11]
↓
轮次 20:再次触发
LLM 压缩摘要 + 原文 9-18 → 新摘要(400 tokens)
保留轮次 19-20 完整
双轨设计(用于精确回溯):
主轨(上下文):[滚动摘要] + [最近 N 轮原文]
副轨(向量库):完整历史对话 → embedding → 向量检索
需要远期内容时:向量检索 + 注入上下文
8.3 多模型接入的 Tokenizer 适配
不同模型使用不同的 tokenizer,同一段文本 token 数可能相差 20--30%。多模型接入时必须为每个模型维护独立的 token 计数器:
typescript
/**
* 多模型 Tokenizer 适配器
* 不同模型使用不同 tokenizer,token 数可能相差 20-30%
*/
interface TokenCounter {
/** 计算文本的 token 数 */
count(text: string): number;
/** 将文本截断到指定 token 数 */
truncate(text: string, maxTokens: number): string;
}
/** OpenAI 系模型(使用 tiktoken) */
class OpenAITokenCounter implements TokenCounter {
private encoding: any; // tiktoken.Encoding
constructor(modelOrEncoding: string = "cl100k_base") {
// 动态导入 tiktoken(避免在不使用时加载)
// const { encoding_for_model, get_encoding } = await import("tiktoken");
// this.encoding = encoding_for_model(modelOrEncoding);
this.encoding = null; // 占位,实际使用时初始化
}
count(text: string): number {
// 实际实现:return this.encoding.encode(text).length;
// 近似实现(不依赖 tiktoken 时的 fallback):
return Math.ceil(text.length / 4); // 英文约 4 字符/token
}
truncate(text: string, maxTokens: number): string {
// 实际实现:
// const tokens = this.encoding.encode(text);
// return this.encoding.decode(tokens.slice(0, maxTokens));
return text.slice(0, maxTokens * 4); // 近似 fallback
}
}
/** Anthropic Claude(使用官方 API 计数或近似) */
class AnthropicTokenCounter implements TokenCounter {
count(text: string): number {
// Claude 使用 SentencePiece,中文比 cl100k 更高效
// 实际值:调用 /v1/messages count_tokens 端点
// 近似值:中文约 1.5 字符/token,英文约 4 字符/token
const chineseRatio = (text.match(/[\u4e00-\u9fff]/g) || []).length / text.length;
const avgCharsPerToken = chineseRatio > 0.5 ? 1.8 : 4;
return Math.ceil(text.length / avgCharsPerToken);
}
truncate(text: string, maxTokens: number): string {
const chineseRatio = (text.match(/[\u4e00-\u9fff]/g) || []).length / text.length;
const avgCharsPerToken = chineseRatio > 0.5 ? 1.8 : 4;
return text.slice(0, Math.floor(maxTokens * avgCharsPerToken));
}
}
/** Qwen/中文友好模型(SentencePiece,中文效率高) */
class QwenTokenCounter implements TokenCounter {
count(text: string): number {
// Qwen 词表包含大量中文词条(约 10 万条)
// 中文约 1.2 字符/token(显著优于 cl100k 的 1.5)
const chineseRatio = (text.match(/[\u4e00-\u9fff]/g) || []).length / text.length;
const avgCharsPerToken = chineseRatio > 0.5 ? 1.2 : 3.5;
return Math.ceil(text.length / avgCharsPerToken);
}
truncate(text: string, maxTokens: number): string {
const chineseRatio = (text.match(/[\u4e00-\u9fff]/g) || []).length / text.length;
const avgCharsPerToken = chineseRatio > 0.5 ? 1.2 : 3.5;
return text.slice(0, Math.floor(maxTokens * avgCharsPerToken));
}
}
/** Token Counter 工厂 */
const TOKEN_COUNTERS: Record<string, TokenCounter> = {
"gpt-4o": new OpenAITokenCounter("o200k_base"),
"gpt-4-turbo": new OpenAITokenCounter("cl100k_base"),
"claude-3-5-sonnet-20241022": new AnthropicTokenCounter(),
"claude-3-5-haiku-20241022": new AnthropicTokenCounter(),
"qwen2.5-72b-instruct": new QwenTokenCounter(),
};
function getTokenCounter(model: string): TokenCounter {
const counter = TOKEN_COUNTERS[model];
if (!counter) {
console.warn(`未知模型 ${model},使用 cl100k 近似计数(误差可能 ±20%)`);
return new OpenAITokenCounter("cl100k_base");
}
return counter;
}
/**
* 模型无关的上下文预算控制
* 切换模型时只需换 model 参数,预算计算自动适配
*/
function fitContextToBudget(
model: string,
components: Array<{ text: string; label: string }>,
budgetTokens: number
): string {
const counter = getTokenCounter(model);
let totalTokens = 0;
const included: string[] = [];
for (const { text, label } of components) {
const tokens = counter.count(text);
if (totalTokens + tokens <= budgetTokens) {
included.push(text);
totalTokens += tokens;
} else {
// 放剩余预算内的截断版本
const remaining = budgetTokens - totalTokens;
if (remaining > 50) {
const truncated = counter.truncate(text, remaining);
included.push(`[${label} 已截断]\n${truncated}`);
totalTokens += remaining;
}
break; // 预算耗尽
}
}
return included.join("\n\n");
}
8.4 上下文管理完整最佳实践清单
| 场景 | 策略 | 关键细节 |
|---|---|---|
| 固定 system prompt | 前置 + prompt caching | 不插入动态内容(时间戳等),保证前缀不变 |
| 多轮对话历史 | 近期全保留 + 远期滚动摘要 | 保留最近 6--10 轮;超过阈值触发 LLM 摘要 |
| RAG 检索结果 | Reranker 精排 + 末尾注入 | 最多 3--5 篇;最相关文档放末尾,靠近 query |
| 长文档处理 | 分层摘要(Chapter → Global) | 单章节用廉价模型;高风险章节精读 |
| Token 计数 | 模型独立 tokenizer | 每个模型一套计数器;中文按语言比例修正 |
| OOM 防护 | 请求层限制 + PagedAttention | 硬性 max_input_tokens 上限;vLLM 分页管理 |
| 质量验证 | NIAH 矩阵测试 | 选型和扩窗后必跑;中间位置(40%--70%)是关键 |
9. 前沿进展:突破上下文窗口限制的新方向
9.1 Mamba 与 SSM:线性复杂度的竞争者
Mamba(Gu & Dao 2023)代表的状态空间模型(State Space Models,SSM)以线性计算复杂度处理序列,通过维护一个固定大小的隐状态来"压缩"历史信息,而非存储完整 KV Cache。其优势在于:
- 计算复杂度:O(n) 而非 O(n²)
- 推理时显存:O(1)(固定状态大小)而非 O(n)
- 超长序列吞吐量显著优于 Transformer
当前局限:在需要精确回溯远距离信息的任务(如精确引用、多跳推理)上,表现仍不如 Transformer。Hybrid 架构(Jamba、Zamba)将 Mamba 层与 Attention 层交错,试图融合两者优势。
9.2 Memory-Augmented Transformers
外部记忆增强方案将超出窗口的内容卸载到外部存储(向量数据库、键值存储),按需检索:
markdown
Memory-Augmented 架构:
当前上下文窗口(Working Memory)
↕ 注意力机制
外部记忆(Long-term Memory)
↕ 近似最近邻检索(ANN)
[embedding1, embedding2, ...] ← 历史 KV 的压缩表示
优点:理论上支持无限历史
缺点:检索引入延迟;难以保证检索完整性;
ANN 检索有召回率损失
代表工作:
- MemGPT(Park et al. 2023):分层记忆管理
- LongMem(Wang et al. 2023):缓存历史层激活
- Infini-Attention(Munkhdalai et al. 2024):压缩 attention over past
9.3 推测解码(Speculative Decoding)
推测解码(Leviathan et al. 2022)通过小模型(draft model)并行预测多个 token,再由大模型(target model)一次性验证,将 decode 阶段从纯串行变为部分并行:
vbnet
标准解码:
T1 → T2 → T3 → T4 → T5(串行,5步)
推测解码:
Step 1: Draft model 并行预测 T1' T2' T3' T4'(快)
Step 2: Target model 验证并接受前 k 个合法 token
(若 T1'=T1, T2'=T2 验证通过,T3' 不对则停止)
Step 3: 接受了 T1, T2,生成 T3(1次大模型 forward 完成 3个 token)
效果:大模型调用次数不变,但每次调用产出 2-4 个 token
Speed-up:通常 2-3× TPOT 加速
10. 工程踩坑汇总
坑 1:用 tiktoken 计数用于 Claude,导致超窗口
tiktoken 的 cl100k 编码和 Claude 的 tokenizer 对同一文本可能相差 10--20%。特别是中文混合场景,差异更大。安全做法:用 Anthropic 官方 API 的 count_tokens 端点,或乘以 1.2 的安全系数。
坑 2:在 system prompt 里插入时间戳,破坏 prompt cache
f"当前时间:{datetime.now()}" 这一行会使每次请求的 system prompt 前缀不同,彻底禁用 prompt caching,成本上升 3--5×。所有动态内容(时间、用户 ID、会话 ID)必须放在 user message,不能在 system prompt。
坑 3:不设 max_output_tokens,允许模型无限生成
不设 max_output_tokens 时,模型可能生成数千 token 的冗余内容,成本按输出 token 计费(通常比输入贵 3--5×)。每个请求必须根据任务场景设置合理的 max_output_tokens 上限。
坑 4:标称窗口大就选哪个,忽略 NIAH 测试
几个真实案例:某开源模型标称 128K 窗口,但在 NIAH 测试中 60--80K 中间位置准确率仅 40%;另一个 32K 窗口模型反而全绿。选型必须实测,不能相信宣传。
坑 5:切换模型时未更新 token 计数逻辑
从 GPT-4(cl100k,50K vocab)切换到 Qwen 2.5(SentencePiece,150K vocab),相同中文文本 token 数可能减少 20%。如果复用了 GPT-4 的计数逻辑,历史裁剪阈值会比实际保守 20%,导致不必要的内容丢失。
坑 6:私有化部署时忽略 KV Cache 显存预算
部署 LLaMA-3 70B 时,模型权重本身约 140 GB(FP16)。在 4×A100 80GB 上,可用显存为 320 GB,模型占 140 GB,只剩 180 GB 给 KV Cache。若支持 128K 上下文,单请求 KV Cache 约 430 GB(远超剩余显存)。必须在部署前算清 KV Cache 预算,决定最大支持的上下文长度和并发数。
💡 关键要点
- 上下文窗口是硬约束:由计算复杂度(O(n²) Attention)、显存(KV Cache 线性增长)、位置编码(训练长度外推)三重约束共同决定,不可绕过。
- KV Cache 是推理的显存瓶颈 :计算公式为
2 × layers × seq_len × kv_heads × head_dim × batch × bytes,随序列长度和并发数线性增长。GQA/MQA 可节省 8--64 倍,PagedAttention 将利用率从 <20% 提升到 90%+。 - 大窗口不等于好质量:Lost in the Middle 现象表明,关键信息在中间位置时模型检索准确率显著下降(U 型曲线)。NIAH 测试是评估长上下文能力的标准基准,选型必跑。
- 扩窗技术路线按成本排序:NTK-Aware Scaling(免调)→ YaRN(少量微调)→ 直接长上下文预训练(成本最高)。工程上优先使用 YaRN,配合严格的 NIAH 验证。
- Token 预算是第一工程原则:将上下文视为有限预算而非无限空间;系统提示前置命中 prompt cache;重要内容放末尾(近因效应);多模型接入时模型独立 tokenizer。
- Prompt Caching 是成本优化的最大杠杆:固定前缀命中缓存后成本降 90%,但要求前缀字节序列精确匹配------绝不在 system prompt 中插入动态内容。
- 分层上下文策略优于"塞满窗口":按优先级分层(固定层/摘要层/近期层/检索层),配合动态预算控制和 Reranker,往往比直接使用最大窗口成本更低、质量更好。