2026-04-22 · architecture / ai-infra
【大模型基础设施工程】11:推理引擎基础
从 Prefill/Decode 两阶段、KV Cache、Continuous Batching 到 PD 分离,系统讲清楚大模型推理的工程基础。



























2022 年 GPT-3.5 上下文 4K token,2024 年 Claude 3 做到
200K,Gemini 1.5 Pro 冲到 1M、实验版本 2M;2025 年
GPT-4.1、Kimi K2、Qwen2.5-1M 等陆续把 1M
级上下文推向生产。对 infra 工程师而言,这不是简单”把
max_position 调大”就能解决的问题——注意力的
O(n²) 计算、KV cache 的 O(n)
显存、位置编码外推的精度崩塌、长样本稀缺导致的训练困难、长序列的通信与负载不均,每一环都需要重构。
本篇系统梳理长上下文的工程栈:从 RoPE、YaRN、LongRoPE 等位置编码扩展,到 Mamba / 线性 attention / NSA 等计算复杂度优化,再到 Ring Attention / Ulysses 序列并行、MLA / StreamingLLM 等 KV 压缩策略,最后落到评测(NIAH、RULER、LongBench)与 Agent 场景的 KV 复用。
| 场景 | 长度区间 | 代表任务 |
|---|---|---|
| 单轮 QA | 2–8K | 客服、闲聊 |
| 文档 QA | 16–128K | 论文、合同、财报摘要 |
| 代码库理解 | 64–500K | “把我这个 monorepo 重构一下” |
| 多轮对话 / 陪伴 | 32K–1M | 长期记忆、角色扮演 |
| Agent 轨迹 | 128K–2M | ReAct、浏览器 Agent、代码 Agent 的 trace + tool output |
| 视频 / 音频理解 | 200K–10M | Gemini 1.5 视频 token 化 |
| 长 CoT 推理 | 32K–256K | o1 / DeepSeek-R1 的 thinking trace |
注意到一个变化:早期长上下文的 killer app 是”一次性喂一本书”,现在 Agent 轨迹和长思维链成为更关键的需求——模型输出自己也要消耗上下文。
工程上永远有 “RAG vs long-context” 的争论。真实答案是两者互补:
生产上普遍的做法是:RAG 做召回 + long-context 模型做 reasoning,第 17–18 篇会详细讨论。
| 模型 | 官方上下文 | 备注 |
|---|---|---|
| Gemini 2.5 Pro | 2M | Google,多模态;1M 为默认 |
| GPT-4.1 / GPT-5 | 1M | OpenAI,2024 Q4 之后 |
| Claude 3.5 / 3.7 Sonnet | 200K | Anthropic;企业版 500K |
| Kimi K2 / Moonshot | 2M(128K 推理稳态) | 月之暗面 |
| Qwen2.5-Turbo-1M | 1M | 阿里,YaRN + 稀疏 |
| Qwen3-256K | 256K | 阿里,默认 32K 可扩展 |
| DeepSeek-V3 / R1 | 128K | 深度求索 |
| GLM-4-Long | 1M | 智谱 |
| Llama 3.1 / 3.3 | 128K | Meta |
| Mistral Large 2 | 128K | Mistral |
| Yi-1.5-9B-Long | 200K | 零一万物 |
| MiniMax-01 | 4M | MiniMax,Lightning attention |
注意:官方宣称长度 ≠ 实际有效长度。Needle-in-a-Haystack 和 RULER 都显示,几乎所有 1M+ 模型在 256K 以上精度明显下滑;MiniMax-01 由于采用线性 attention,4M 在长距离检索任务上还有不小差距。
Transformer 本身对序列顺序不敏感,位置编码(Positional Encoding, PE)是注入顺序信息的关键。长上下文工程的第一个战场就是 PE。
绝对位置编码(如 BERT 的可学习 PE、Transformer 原论文的 sinusoidal)直接给每个位置一个 embedding,长度外推非常差——训练时没见过的位置,embedding 无从查起。
RoPE(Rotary Position Embedding, Su et al. 2021) 把位置以”旋转”形式注入到 Q、K 向量:对每对通道 \((x_{2i}, x_{2i+1})\) 按角度 \(\theta_i = \text{base}^{-2i/d}\) 旋转 \(m\theta_i\)(\(m\) 为位置索引)。
伪代码:
def rope(x, position_ids, base=10000.0):
# x: [B, H, T, D], D must be even
half = x.shape[-1] // 2
freqs = base ** (-torch.arange(0, half, device=x.device) * 2 / x.shape[-1])
angles = position_ids[:, None] * freqs[None, :] # [T, half]
cos, sin = angles.cos(), angles.sin()
x1, x2 = x[..., :half], x[..., half:]
return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)RoPE 的优点: - 相对位置:\(q_m^\top k_n\) 只依赖 \(m-n\)。 - 无额外参数。 - 在 LLaMA、Qwen、DeepSeek、Mistral 等几乎所有现代开源模型里使用。
缺点:直接外推到训练长度以外效果差——高频分量在超长位置下摆动剧烈,注意力分布”混乱”。
最简单的扩展:把 base(θ,常设为
10000)改大,例如 500000 或 1000000。
# LLaMA 3 把 base 从 10000 改到 500000,上下文扩到 8K→128K直观理解:base 越大,低频分量周期越长,高位置也不至于绕太多圈。但单纯 scaling 会让”相对位置感”被稀释,短文本精度下降,因此通常搭配 continued pretraining 微调。
Meta 2023 年提出的 Position Interpolation:推理时把位置索引 \(m\) 线性压缩到训练范围内:\(m' = m \cdot L_{\text{train}} / L_{\text{new}}\)。等价于把 RoPE 频率按比例缩小。
简单粗暴、但精度损失明显,尤其在 4× 以上外推。
社区(blog、bloc97、emozilla)很快发现:高频分量外推还不错,低频分量需要插值。NTK-Aware 的做法是按频率分段缩放 base:
\[\text{base}' = \text{base} \cdot s^{d/(d-2)}, \quad s = L_{\text{new}} / L_{\text{train}}\]
实际效果比 PI 好很多。Dynamic NTK 则是推理时根据当前实际长度动态调整 scaling 因子,避免短文本退化。
YaRN(Yet another RoPE extensioN, Peng et al. 2023) 在 NTK-Aware 基础上做了三件事:
YaRN 已成为开源事实标准:Qwen、Mistral、Yi、DeepSeek 长上下文分支都有 YaRN 版本。
配置示例(HuggingFace Transformers):
{
"rope_scaling": {
"type": "yarn",
"factor": 4.0,
"original_max_position_embeddings": 8192,
"attention_factor": 1.0,
"beta_fast": 32,
"beta_slow": 1
},
"max_position_embeddings": 32768
}微软 LongRoPE(2024) 用进化搜索找到每个频率通道的最佳 scaling 因子(不是解析函数,而是 per-channel 的 search),配合短-长两阶段微调,把模型扩到 2048K(2M)。Phi-3-mini-128K 等模型使用该方案。
ALiBi(Attention with Linear Biases):在 attention logits 上直接加 \(-m \cdot |i-j|\) 的斜率偏置,\(m\) 是头特定常数。完全没有”位置外推”问题,外推性好,但实测长距离语义建模比 RoPE 弱,主要用于 BLOOM、MPT。
CoPE(Contextual Position Encoding, Meta 2024):位置不再单调递增,而是由模型自己决定”哪些 token 该算位置单位”——例如句号计数、tool call 计数。对 Agent 长对话可能有用,目前尚未大规模落地。
位置编码只是”让模型能索引到远处”。真正决定能不能跑 1M 的是注意力计算与显存。
标准 attention: - 计算 \(QK^\top\) 与 \(\text{softmax}(\cdot)V\):浮点数 \(\Theta(n^2 d)\)。 - KV cache:推理时每层保存 \(n \cdot d_{\text{kv}}\) 个元素,对 70B 模型每 token 数百 KB,1M token 就是几百 GB。
第一项决定prefill 延迟,第二项决定能不能装下。
Linear Attention(Katharopoulos 2020):把 softmax 替换为 \(\phi(Q)\phi(K)^\top\),利用结合律先算 \(\phi(K)^\top V\) 得到 \(d \times d\) 的状态矩阵,再和 \(\phi(Q)\) 相乘,复杂度变 \(O(nd^2)\)。
Performer:用随机特征近似 softmax。 Linformer:对 K、V 做低秩投影到固定长度 \(k\)。 Reformer:LSH 分桶。
这些方法在小模型上有效,但大规模 LLM 预训练后质量普遍落后 softmax attention,工业落地有限。
State Space Model(SSM):用连续系统 \(h'(t) = Ah(t) + Bx(t), y = Ch\) 的离散化,每 token 用固定大小状态 \(h \in \mathbb{R}^d\) 累积,推理是 O(n) 且 KV cache 常数大小。
SSM 的痛点是纯 SSM 模型检索能力弱——在 NIAH 上显著差于 transformer,因为固定大小状态会遗忘。
工业界的共识:纯线性/SSM 不够,纯 attention 太贵,混合是最佳点。
经验规律:每 6–8 层放 1 层 full attention,其余用 SSM 或 linear,就能保住 NIAH 精度并大幅降本。
Longformer / BigBird(2020):每个 token 只看前后 \(w\) 个 token(local),加少量 global token 看全局。简单、有效、不需重训。
Mistral SWA(Sliding Window
Attention):window=4096,配合 rolling
KV buffer,7B 模型原生 32K。
注意:SWA 上下文容量是 \(w
\cdot
L\)(层数),信息逐层向下游传递,并非硬截断。Mistral
7B 的”有效上下文”其实接近
4096 × 32 = 131K,但在长距离检索上仍弱于 full
attention。
DeepSeek 2025 年 NSA 是目前最受关注的”训练原生稀疏”方案,核心思想:
三分支并行、输出加权和。关键是 三个分支都从预训练开始就存在,不是事后稀疏化,因此模型学会了稀疏注意力的路由——这是”Native”的含义。
工程上 NSA 的收益: - 64K 以上,端到端 prefill / decode 加速 2–10×。 - 精度在 LongBench / RULER 上与 full attention 持平。
微软 MInference(2024) 走的是推理时稀疏:把 attention 分类为 A-shape、Vertical-Slash、Block-Sparse 三种模式,推理时动态识别并用 sparse kernel 计算。不需要重训,落地门槛低,已集成进 vLLM。
PyTorch 2.5+ 的 FlexAttention 提供”写 Python
描述 mask / score_mod,自动编译成 Triton
kernel”,让自定义稀疏 attention 不用手写
CUDA。对研究与定制化非常友好:
from torch.nn.attention.flex_attention import flex_attention, create_block_mask
def sliding_window(b, h, q_idx, kv_idx):
return (q_idx - kv_idx).abs() <= 1024
block_mask = create_block_mask(sliding_window, B=None, H=None, Q_LEN=32768, KV_LEN=32768)
out = flex_attention(q, k, v, block_mask=block_mask)| 方法 | 是否需重训 | 精度 | 加速 |
|---|---|---|---|
| NSA | 是(原生) | 持平 | 2–10× |
| MInference | 否 | 轻微降 | 2–5× |
| H2O / SnapKV(decode 端) | 否 | 轻微降 | 显存 2–4× |
| FlexAttention 自定义 | 取决于 mask | 取决于 | 可大 |
预训练语料里,大多数文档 < 4K token,>32K 的样本极少。直接在目标长度 pretrain 太贵又数据不足,业界普遍采用两阶段:
LLaMA 3.1 的做法:base 在 8K 训练,之后分 6 个阶段把 window 逐步扩到 128K,每阶段加 800M token。Qwen2.5-1M 的做法类似,但最后一段 1M 上下文只用很少数据、靠 YaRN + 稀疏 attention 外推。
真实长文本不足,工程上大量依赖合成:
关键:没有 needle 类训练样本,长上下文检索精度会很差,仅靠 next-token loss 无法让模型学会”在 500K 位置精确回忆某个实体”。
varlen 接口原生支持。训练和推理 prefill 阶段,单张 GPU 放不下 1M token 的激活。序列并行(Sequence Parallel, SP) 应运而生。
UC Berkeley 2023 提出,核心思想:把序列切分到 N
张卡,每张卡持有 n/N 段 Q, K, V;计算 attention
时卡间环形传递 K, V,每张卡累积自己的 Q
和所有 KV 块的结果,通信与计算重叠。
step 0: GPU_i has (Q_i, K_i, V_i), compute local Attn(Q_i, K_i, V_i)
step 1: recv K_{i-1}, V_{i-1} from left, compute Attn(Q_i, K_{i-1}, V_{i-1})
...
step N-1: finished all N-1 peer blocks
显存降为 \(O(n/N)\),通信量 \(O(n)\) per GPU,可与 Flash Attention 结合(Blockwise Flash)。
All-to-all 式序列并行:切 seq
维,但在进入 attention 前做一次 all-to-all,把
“seq 切分”转成 “head 切分”——每张卡拿到全序列但只负责一部分
head;attention 算完再 all-to-all 转回 seq 切分。
2 × seq × hidden / N),对带宽要求低。N ≤ num_heads。| 方案 | 通信模式 | 显存 | 适用长度 | 实现 |
|---|---|---|---|---|
| Ring Attention | p2p ring | \(O(n/N)\) | 极长(256K+) | Megatron-LM CP, DS-Ulysses Ring |
| Ulysses | all-to-all | \(O(n/N)\) | 中长(32K–256K) | DeepSpeed |
| Context Parallel(Megatron) | Ring 变体 | \(O(n/N)\) | 极长 | Megatron-LM |
| Sequence Parallel(Megatron 原版) | 只切 layer norm / dropout | 小幅 | 配合 TP | Megatron-LM |
2024 之后生产方案普遍是 Context Parallel = Ring Attention + Flash Attention,Megatron-LM / NeMo / ColossalAI / veScale 都有实现。
训练搞定了只是开始。推理端的挑战更多,因为要在延迟、吞吐、显存之间实时取舍。
给定 1M token 的 prompt,单卡 prefill 要几分钟,根本无法接受。生产上:
1M token × 70B 模型的 KV 大约 300 GB,单机 8×H100 勉强能装,但要和权重抢显存。方案:
长上下文推理最大的工程红利。很多场景共享前缀: - 系统 prompt 几 KB 但被所有请求共享; - RAG 把检索文档拼前面,同一文档多次命中; - Agent 在多步之间,前 N-1 步的 context 是 N 步的前缀。
vLLM --enable-prefix-caching、SGLang
RadixAttention、TensorRT-LLM KV Reuse
都做自动前缀匹配,命中后 prefill
时间归零,对 50K+ 的系统 prompt 收益巨大(可以把首
token 延迟从 10s 降到 100ms)。
# vLLM 开启前缀缓存并扩大 KV pool
vllm serve Qwen/Qwen2.5-7B-Instruct-1M \
--enable-prefix-caching \
--max-model-len 1048576 \
--kv-cache-dtype fp8 \
--block-size 16 \
--tensor-parallel-size 41M token 的 KV 是长上下文推理的最大显存开销,降 KV 是硬刚需。
DeepSeek-V2/V3 的 MLA:不保存每个头独立的 K, V,而是保存一个 低秩 latent(约 512 维),推理时再线性投回 K、V。
效果: - KV cache 约为 MHA 的 1/6–1/10(具体看秩选择)。 - 精度几乎不损失。 - 1M 上下文在同样显存下完全可行。
MLA 是 DeepSeek-V3 能跑 128K 且推理经济的关键之一。见前作第 7 / 11 篇。
GQA(Grouped Query Attention):多个 Q 头共享一组 KV 头,LLaMA 2/3、Qwen、Mistral 都用,KV 减小 4–8×。 MQA(Multi-Query):所有头共享单组 KV,PaLM、Falcon 用,更激进但质量略损。
GQA 是当前 transformer 的默认选项,配合长上下文几乎必选。
MIT StreamingLLM(2023):观察到 transformer 对最前面几个 token(“attention sink”)有稳定的高注意力权重。滑窗推理时只要保留前 N 个 sink token + 最近 W 个 token,就能在 4M+ 流式生成下不崩。
典型
sink=4, window=2048,可以在不改模型的情况下跑”永不结束”的对话。缺点:中间内容被丢弃,不是真正的长上下文,只适合对话、流式场景。
H2O(Heavy-Hitters Oracle):统计每个 token 的累积 attention score,保留 heavy hitter + recent。 SnapKV:在 prefill 结束时,对每个 head 选 top-k KV 保留。 PyramidKV:不同层用不同预算,浅层保留多、深层少。
这些方法都是 decode 端 KV eviction,不需要重训,配 vLLM 可 2–4× 减显存、精度损失小。
见上一篇,INT8 / FP8 / INT4 KV 是最简单粗暴的节省方式。长上下文场景下 FP8 KV 基本无损、INT4 KV 有 2–5% 精度损失,多与上述稀疏方案叠加使用。
把一句话(needle)插到长文档(haystack)任意位置,让模型回答。画
(上下文长度 × needle 位置) → 准确率
热力图。
NIAH 过于简单:只考一句话精确回忆,很多方法(SnapKV、StreamingLLM)都能刷高分,但不代表模型真的能在该长度下 reasoning。
NVIDIA 2024 推出,13 个子任务,包括: - Multi-needle retrieval:同一 haystack 多 needle; - Variable tracking:追踪变量赋值链; - Aggregation:统计某类元素数量; - QA:多跳 QA。
RULER 更能暴露模型真实长上下文能力。业界普遍以 “RULER > 85% 的最长长度” 作为”有效上下文”。典型结果: - LLaMA 3.1-70B 官方 128K,RULER 有效 ~64K。 - Qwen2.5-72B-128K,RULER 有效 ~64K–96K。 - Gemini 1.5 Pro 1M,RULER 有效 ~512K。
真实长文任务:代码库 QA、多文档摘要、长篇阅读理解。工程选型时建议跑 LongBench + RULER 两者组合。
Liu et al. 2023 指出:把重要信息放在长文档中间,模型准确率显著低于放在开头或结尾。解法: - 训练层面:balanced needle 分布; - 推理层面:re-ranking 检索结果放到首尾; - 模型层面:YaRN、NTK scaling 间接缓解。
Agent 是长上下文的最大用户,也是最大挑战。
一个浏览器 Agent 平均每步消耗: - HTML snapshot:10–30K token; - tool output:1–10K; - reasoning / thinking:1–5K;
10 步就 100K+,50 步轻松上 500K。而且90% 的 context 是上一步已经处理过的重复内容。
生产级 Agent 推理:
Mooncake、LMCache、KVRouter 等都在做这件事;vLLM 1.0 计划引入分布式 KV store。
当 Agent 轨迹接近窗口上限时: - Summarization:让小模型压缩前 N-K 步; - Selective drop:drop 掉非关键 tool output(如已经查过的 HTML); - Hierarchical memory:把旧轨迹压成更抽象的 memory。
这一块没有银弹,框架(LangGraph、Agno、AutoGen)各有实现,下一篇也会涉及。
下面展示一个最小但完整的 workflow:拿一个 RoPE base 10000 的 7B 模型,通过修改 config + 少量 continued pretrain,把有效上下文从 8K 扩到 32K。
# patch_config.py
import json, sys
from pathlib import Path
path = Path(sys.argv[1])
cfg = json.loads((path / "config.json").read_text())
cfg["rope_scaling"] = {
"type": "yarn",
"factor": 4.0,
"original_max_position_embeddings": 8192,
"attention_factor": 1.0,
"beta_fast": 32,
"beta_slow": 1,
}
cfg["max_position_embeddings"] = 32768
(path / "config.json").write_text(json.dumps(cfg, indent=2))
print("patched.")从 Pile / SlimPajama 筛 ≥ 24K token 的文档,或拼接同作者文档到 24–32K。目标 2B token。
from datasets import load_dataset
ds = load_dataset("cerebras/SlimPajama-627B", split="train", streaming=True)
ds = ds.filter(lambda x: len(x["text"]) > 120_000) # 粗估 24K+ token
ds = ds.take(100_000)用 DeepSpeed Ulysses + Flash Attention:
deepspeed --num_gpus 8 train.py \
--model_name_or_path ./model-yarn \
--dataset_path ./long_samples.jsonl \
--max_seq_len 32768 \
--sequence_parallel_size 4 \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 16 \
--learning_rate 2e-5 \
--max_steps 1000 \
--save_steps 500 \
--bf16 \
--deepspeed ds_zero3_ulysses.json典型参数:2e-5 lr、1000 step、global batch 128、warmup 5%。千万不要用 pretrain 的 lr,YaRN 微调是精细调整,过大 lr 会毁短上下文能力。
# NIAH quick check
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tok = AutoTokenizer.from_pretrained("./model-yarn-32k")
mdl = AutoModelForCausalLM.from_pretrained("./model-yarn-32k", torch_dtype=torch.bfloat16, device_map="auto")
def niah(ctx_len, depth):
filler = "The quick brown fox jumps over the lazy dog. " * 4000
needle = "The secret passphrase is ALPACA-42-NEPTUNE."
ins_at = int(len(filler) * depth)
text = filler[:ins_at] + needle + filler[ins_at:]
text = tok.decode(tok(text, max_length=ctx_len-64, truncation=True)["input_ids"])
q = text + "\n\nQuestion: What is the secret passphrase?\nAnswer:"
ids = tok(q, return_tensors="pt").to(mdl.device)
out = mdl.generate(**ids, max_new_tokens=32, do_sample=False)
return tok.decode(out[0, ids.input_ids.shape[1]:])
for L in [8192, 16384, 32768]:
for d in [0.1, 0.5, 0.9]:
print(L, d, niah(L, d))vllm serve ./model-yarn-32k \
--max-model-len 32768 \
--rope-scaling '{"type":"yarn","factor":4.0,"original_max_position_embeddings":8192}' \
--enable-prefix-caching \
--tensor-parallel-size 2注意:短上下文(< 8K)请求在 YaRN 模型上性能会轻微下降(~2–4%)。如果短请求占比高,考虑部署 base + YaRN 双模型由网关分流(第 22 篇)。
总结一张决策表:
| 需求 | 推荐路线 |
|---|---|
| 32K 以下,只要能跑 | GQA + RoPE base 调整,不用 YaRN |
| 32K–200K,预算有限 | 现有 YaRN 版开源模型(Qwen2.5、LLaMA 3.1) + vLLM + 前缀缓存 |
| 128K+ 训练预算充足 | 短 pretrain → YaRN continued pretrain → Ring/Ulysses SP |
| 1M+ 商用 | API:Gemini / GPT-4.1 / Claude;自建:DeepSeek-V3 + MLA + NSA,或 Qwen2.5-Turbo-1M |
| Agent 长轨迹 | 前缀缓存 + KV offload + 轨迹压缩,上下文 128K 足够 |
| 无限流式对话 | StreamingLLM + sink |
| 视频 / 多模态 1M | Gemini API,或自建 Mamba-hybrid(Jamba、MiniMax-01) |
关键原则:别盲目追求百万上下文。RULER 有效长度才是真实能力,对大多数业务 64–128K 已经绰绰有余;把钱花在 RAG 质量、prefix caching、Agent 工程上回报更大。
长上下文工程是 LLM infra 最”立体”的一块:
下一篇进入 RAG 工程全景,聊长上下文最重要的互补方案:检索、切块、ranker、融合策略与生产系统。
把当前热点继续串成多页阅读,而不是停在单篇消费。
2026-04-22 · architecture / ai-infra
从 Prefill/Decode 两阶段、KV Cache、Continuous Batching 到 PD 分离,系统讲清楚大模型推理的工程基础。
2026-04-22 · architecture / ai-infra
系列收官:从 2022 到 2026 的四年拐点出发,梳理推理时 Scaling、世界模型、Agentic OS、专用芯片、架构创新、端侧、成本腰斩、合规八大趋势,并给出工程师成长路径与 25 篇索引。
2026-04-22 · architecture / ai-infra
面向工程师的大模型基础设施开篇地图,覆盖 2022 到 2026 的工程分水岭、五层工程栈、训练与推理的工程差异、中国与全球行业版图以及成本曲线。
2026-04-22 · architecture / ai-infra
从 CPU 与 GPU 的架构差异出发,讲清楚 SM、Warp、Tensor Core、HBM、NVLink 的工程含义,并结合 Roofline、FlashAttention 与国产算力栈,给出大模型工程师能直接上手的 GPU 心智模型。
此内容由惯性聚合(RSS阅读器)自动聚合整理,仅供阅读参考。 原文来自 — 版权归原作者所有。