Llama2架构与推理全解析：从原理到KV Cache与量化部署

Llama2 这个模型现在回头看可能已经不算最前沿了但在我心里它一直是位置很特别的一个开源大模型从“能跑通”走向“能商用”基本就是从它开始的。这两年我拿它做过问答、做过文本总结、也做过小规模的私有化部署每次回去翻它的架构和推理实现都能读出点新东西。这篇文章想做的事很直接——把 Llama2 的架构设计和完整推理过程掰开揉碎讲清楚。不堆术语每个设计为什么存在、推理时每一步到底在算什么、显存都花到哪去了全部串起来。适合刚接触大模型、想搞清楚模型内部机制的人也适合已经跑过推理、但对 KV Cache、采样参数、量化选型这些细节还比较模糊的开发者。看完你会发现大模型的推理没有想象中那么玄无非是“词变成向量、向量过几十层网络、再变回词”这样一个循环。1. 先把 Llama2 的骨架看清楚架构设计1.1 Decoder-only一条直路走到底Llama2 属于 decoder-only 架构的 Transformer。早期做翻译的 Transformer 是 encoder-decoder 结构encoder 把整个输入编码成一批上下文向量decoder 再基于这些向量逐步生成输出。而 Llama2 这类语言模型直接把 encoder 砍掉只剩 decoder模型拿到一段文本每一步只能看到当前位置之前的内容这就是常说的因果掩码causal mask。为什么这么设计核心原因是训练目标变了。翻译任务需要“先理解完整句子再逐词翻译”所以编码器是必要的而语言模型要学的是“根据前文预测下一个词”这是一个天然的自回归任务decoder-only 结构刚好匹配。更重要的是decoder-only 在训练和推理时的形态完全一致训练时一次性输入整段文本推理时也是一段一段往后接没有 encoder 和 decoder 之间的接口转换成本。Llama2 的主体就是一堆完全相同的 Transformer 层叠在一起7B 模型是 32 层13B 是 40 层70B 直接堆到 80 层。每一层内部其实只有两个核心模块——多头自注意力Multi-Head Self-Attention和前馈网络FFN外面再用残差连接包起来。残差连接的作用是给梯度留一条“高速公路”让 80 层这么深的网络也能训练得动。这里有个容易忽略的细节Llama2 用的是 Pre-Norm也就是先归一化再做子层计算而不是原始 Transformer 那种后置归一化。Pre-Norm 的好处是每一层的输入都被拉回到稳定范围深层网络的训练更稳对学习率的敏感度也低很多。1.2 RMSNorm去掉均值反而更稳归一化层在 Transformer 里几乎是标配作用是防止深层网络的数值分布失控。原始 Transformer 用的是 LayerNorm对向量求均值和方差然后做标准化。Llama2 用的 RMSNorm 则省掉了“减均值”这一步只按均方根RMS来缩放。RMSNorm 的计算很简单对输入的每个 token 向量计算所有维度的平方均值开根号然后让每个维度除以这个值最后乘上一个可学习的缩放系数。公式大致是y (x / sqrt(mean(x^2) eps)) * gamma为什么敢去掉均值研究表明对于 Transformer 这种结构真正起作用的主要是“把向量长度归一化”而均值偏移的影响相对小。省掉均值计算省的是实打实的算力在几十亿参数规模的模型上这个节省会被反复累加。而且 RMSNorm 不需要维护均值的统计量在训练和推理时都更省事。实际操作中还会在分母里加一个很小的 eps 防止除零这个值一般在 1e-5 到 1e-6 之间。有一个容易踩的坑RMSNorm 是对每个 token 的隐藏向量独立归一化不是对整个 batch 混着归一化。理解这一点对调试很有用——如果你在自定义实现里把归一化维度算错了早期可能看不出问题但模型一加深就会发散。1.3 SwiGLU给前馈网络装个门控原始 Transformer 的前馈网络很简单先升维过一个 ReLU 激活再降维回来。公式大致是 FFN(x) ReLU(xW1)W2。这种结构把隐藏维度扩大到 4 倍让模型有更丰富的非线性表达能力但 ReLU 的一个问题是它对负输入直接置零信息一刀切。Llama2 换成了 SwiGLU公式为SwiGLU(x) (silu(xW_gate) ⊙ xW_up) W_down其中 silu 是 x 乘以 sigmoid(x)也叫 SiLU 或 Swish 激活。这个结构里有两路输入一路经过激活函数变成“门控信号”另一路保持原始变换两者逐元素相乘最后再通过输出投影。门控机制让网络可以学习“哪些信息值得放行”而不是简单地把负值扔掉。这里顺带讲一个配置细节因为 SwiGLU 引入了额外的参数矩阵Llama2 的中间维度并不是标准的 4 倍隐藏维度而是大约隐藏维度的 2.7 倍左右。以 7B 模型为例隐藏维度是 4096中间层维度是 11008而不是常见的 16384。这也是看模型配置文件时容易困惑的地方——看到 intermediate_size 比 hidden_size × 4 小别以为是配置错了这是 SwiGLU 设计的一部分。1.4 RoPE旋转一下位置就出来了Transformer 的注意力本身是不感知词序的把“我喜欢你”和“你喜欢我”打乱成同样的一组词注意力计算出来的结果不会变。所以必须显式地把位置信息加进去。早期方案有两种加绝对位置的 sin/cos 编码或者让模型学习一个位置嵌入表。Llama2 用的是旋转位置编码RoPE。RoPE 的思路很巧妙不去“加”一个位置向量而是对注意力里的 Q 和 K 向量做“旋转”。具体做法是把每个向量按维度两两分组每一组看成一个二维平面上的坐标然后根据 token 的位置用旋转矩阵把这些坐标旋转一个角度。不同维度分组的旋转频率不同高频维度转得快、低频维度转得慢形成一套类似时钟指针的编码方式。RoPE 最妙的地方是数学性质两个旋转过的向量做点积结果只和它们的相对位置有关和绝对位置无关。也就是说模型天然就能学到“词 A 和词 B 隔了多远”的语义而不是死记每个位置长什么样。这个性质让 RoPE 在理论上具备一定外推能力——模型训练时见过的序列长度是 4096但在实践中把上下文稍微拉长一点RoPE 通常不会立刻崩掉。当然外推范围有限拉太长照样会退化这一点后面讲长上下文时还会提到。1.5 GQA让大模型跑得更快的注意力变体标准多头注意力MHA里每个 query 头都配有自己独立的 key 和 value 投影。这意味着多少个头就要算多少份 K、V显存和带宽开销都不小。Llama2 的 7B 和 13B 用的还是标准的 MHA但 70B 模型换成了分组查询注意力GQA。GQA 的逻辑很简单把注意力头分成若干组每组内的多个 query 头共享一组 key 和 value 头。70B 模型一共有 64 个 query 头但只有 8 组 K、V 头相当于 8 个 query 头共用一组 K、V。这样 K、V 的总量直接缩小到原来的八分之一。为什么要这么做推理阶段每个生成步骤都要把当前 KV Cache 加载到显存里参与计算K、V 头越少KV Cache 越小访存带宽的压力就越小推理速度直接受益。而且实验表明GQA 在质量上的损失远小于 MQA所有 query 头共用一个 KV 头却实实在在省下了一大块显存。可以说 GQA 是专门为大模型的高并发推理场景设计的这也是为什么后来很多新模型都跟进了这个设计。1.6 不同尺寸的具体配置对照把几个关键超参数放在一起看会更直观配置项7B13B70B层数324080隐藏维度409651208192注意力头数324064KV 头数32408中间层维度110081382428672词表大小320003200032000最大上下文409640964096词表大小统一是 32000用的分词器是基于 BPE 的子词方案。隐藏维度决定了一个 token 被表示成多长的向量7B 是 4096 维70B 是 8192 维这也是不同尺寸模型表达能力差异的来源之一。2. 一次完整推理从文本进来到 token 出去2.1 推理的两段式Prefill 与 Decode大模型推理不是拿到 prompt 就直接吐词而是分成两个明显不同的阶段。第一个阶段叫预填充Prefill把整个输入 prompt 一次性喂给模型所有 token 并行计算产生每个 token 对应的中间状态同时把 K、V 张量写入缓存。第二阶段叫解码Decode模型一次只生成一个 token每生成一个就把这个新 token 的结果追加到缓存里然后基于缓存预测下一个 token。这两个阶段的特性差别很大。Prefill 是并行计算吞吐高耗时主要取决于 prompt 的长度几百 token 的输入在一张中高端显卡上一般不到一秒就能算完Decode 是串行的每个 token 都要走一遍完整的前向传播速度通常在每 token 几十毫秒的量级而且后续 token 的预测依赖前面的结果没法并行。所以同样是生成 100 个 token 和生成 500 个 token后者的耗时几乎是前者的五倍这个线性增长是无法绕开的。理解这两段式对做工程很重要。用户感知到的“首字延迟”其实主要来自 Prefill“流畅度”则取决于 Decode 的每 token 耗时。优化手段也得分开看Prefill 慢可以靠更好的并行计算优化Decode 慢更可能是显存带宽和 KV Cache 的瓶颈而不是算力不够。2.2 token 化文本如何变成数字模型不认识字符只认识整数。输入文本进入模型前要先经过分词器变成 token id 序列。Llama2 用的是 BPE 算法词表 32000 个 token覆盖常见英文单词、子词、标点以及一部分中文和其他语言的字符。分词后还要拼接特殊的控制 token。原始 base 模型通常以 BOS序列开始token 开头以 EOS序列结束token 结尾对话模型则会在用户输入外面加上指令格式标记结构大致是“指令前缀 用户内容 指令结束符 回答”。这个模板在推理时一定要和模型训练时保持一致否则效果会明显变差这是很多人换模型后觉得“输出很怪”的原因之一。还有个实操细节中文和英文的 token 效率差别很大。英文一个常见单词通常就是一个 token而中文一个字可能对应一到两个 token遇到生僻字甚至更多。这意味着同样 4096 的上下文窗口中文能装下的有效内容大约只有英文的一半。设计 prompt 时要提前把指令、示例和预期回答的长度都算进去不然内容很容易被截断。2.3 主干计算嵌入、残差流与逐层变换token id 序列准备好后先查嵌入表。嵌入表是一张“词表大小 × 隐藏维度”的矩阵每一行就是一个 token 的向量表示。查表得到的嵌入向量序列就是模型的主干输入接下来要依次穿过所有 Transformer 层。每一层的计算可以简化成两步。第一步是注意力对输入做 RMSNorm然后投影出 Q、K、VQ 和 K 按位置旋转后计算注意力分数也就是 Q 和所有 K 的点积除以维度的平方根再用因果掩码把未来位置的分数设为负无穷softmax 归一化后和 V 加权求和最后经过输出投影。第二步是前馈网络注意力输出再做一次 RMSNorm经过门控投影、激活、线性变换后降回隐藏维度。每一步的输出都会和输入做残差相加这也是“残差流”叫法的由来。可以把这个过程理解成每一层的输入是一块“公共黑板”注意力负责在黑板上读取和整合全局信息前馈网络负责在局部做非线性变换两者写完再把自己的成果加回黑板。最后的输出头也是一层线性变换把隐藏向量映射回词表大小的 logits。需要说明的是Llama2 的嵌入表和输出头的参数是独立的并没有做权重绑定训练参数量会略大一点但表达上更灵活。下面用伪代码把主干流程写出来方便对照实现h embed(token_ids) # [batch, seq, hidden] for layer in layers: h h attn(rmsnorm(h), kv_cache) # 注意力子层 h h ffn(rmsnorm(h)) # 前馈子层 logits lm_head(rmsnorm(h)) # 输出映射2.4 KV Cache推理加速的核心KV Cache 可以说是 Llama2 推理加速最关键的一个设计很多新手第一次看推理代码时最容易在这里绕晕。没有 KV Cache 的话生成第 3 个 token 时模型需要从头再算一遍第 1、2 个 token 的 K 和 V生成第 100 个 token 时就要重复算前 99 个 token 的 K 和 V。计算量随序列长度呈平方级增长。KV Cache 的思路是把已经算过的 K 和 V 存下来生成新 token 时只算新 token 自己的 K 和 V然后拼接到缓存后面注意力只用新 Q 去和全部缓存的 K、V 做计算。这样单步计算量只和当前上下文长度相关不再重复算历史 token。KV Cache 的显存开销需要认真估算。公式大致是KV 缓存大小字节 2K 和 V 两份× 层数 × KV 头数 × 每头维度 × 序列长度 × 单元素字节数。以 7B 模型、4096 上下文、fp16 精度来算大约是 2GB13B 约 2.6GB70B 因为有 GQAKV 头从 64 降到 8缓存反而只有 0.67GB 左右。这正好体现了 GQA 的工程价值模型越大KV Cache 的节省越夸张。正因为 KV Cache 只在推理阶段存在它的内存是动态增长的。所以推理服务的内存占用不能只看模型权重大小还要看并发数乘上每个序列的缓存开销。这也是为什么长上下文场景下显存往往最先被 KV Cache 吃满而不是被模型权重吃满。2.5 从 logits 到采样最后一个 token 怎么定模型最后一层输出的是每个词表位置上的 logits也就是模型认为下一个 token 是某个词的概率对数。把这个 logits 经过 softmax 变成概率分布后还需要一个采样策略来决定到底选哪个 token。最简单的策略是贪心解码每次都选概率最高的 token。贪心的优点是稳定但缺点是容易陷入重复和呆板的输出。所以实际推理几乎都会引入随机性。最基础的是温度参数temperature把 logits 整体除以一个温度值再 softmax。温度大于 1 时分布变平输出更随机温度小于 1 时分布变尖输出更确定温度接近 0 时退化回贪心。温度之外还有两个常见的截断策略。Top-k 是只从概率最高的 k 个 token 里采样Top-p核采样是选择累计概率达到 p 的最小子集再采样。Top-p 比 top-k 更灵活因为它在分布集中时选得少、分布分散时选得多不会像 top-k 那样在极端情况下硬性裁掉可能合理的选项。为了避免重复还可以用重复惩罚把已经出现过的 token 的概率压低。把这些参数组合起来就完成了“从概率分布里挑一个”的决策。2.6 把整条链路串起来现在把整个过程连起来看一遍。输入“请写一首关于春天的短诗”先分词成 token id加上指令格式标记再补一个 BOS。进入 Prefill 阶段所有 token 一起过嵌入层和几十层 Transformer得到初步的注意力缓存。输出头算出第一个 logits 分布采样出一个 token比如“春”。这个 token 拼接回序列进入 Decode 阶段。模型只对“春”这个新 token 做前向计算拿到它的 K、V 并追加缓存再预测下一个 token“风”。如此循环直到模型输出 EOS token或者生成达到设定的最大长度。整个过程中模型每一步其实只在做一件事基于当前已有的全部上下文预测下一个最合理的 token。这也是所有自回归语言模型的基本范式。3. 实战部署显存、量化和生成参数3.1 选多大模型怎么算显存部署 Llama2 的第一步是决定用哪个尺寸。7B 适合 8GB 到 16GB 显存的设备13B 至少需要 24GB70B 则基本要几十 GB 显存或者多卡并行。但这只是模型权重本身实际推理的显存要按“权重 KV Cache 激活值额外开销”来估算。以 fp16 精度为例每个参数占 2 字节7B 模型权重约 14GB13B 约 26GB70B 约 140GB。加上 KV Cache 和计算过程中的中间张量单条请求的实际峰值会比纯权重高出 20% 到 50%。我自己部署时习惯的公式是总显存 ≈ 权重大小 × 1.2 KV Cache 大小 × 并发请求数。先按最坏情况估算再实际跑一遍监控比拍脑袋靠谱得多。多卡部署时还要注意负载均衡。模型权重可以在各卡之间平分但 KV Cache 通常会跟着注意力计算分布如果并行策略没配好容易出现一张卡爆显存、另一张卡闲着的情况。3.2 量化能省一半以上的内存如果显存不够优先考虑量化。量化就是把模型权重从 fp16 降到更低的精度比如 INT8、INT4用更少的位数表示参数代价是轻微的精度损失。INT8 量化一般把权重从 2 字节降到 1 字节模型显存直接减半精度损失通常很小很多场景看不大出来。INT4 量化更进一步权重只有约 0.5 字节7B 模型能压到 4GB 左右普通消费卡就能跑但精度损失会更明显尤其在一些对细节敏感的任务上会看到回答质量下降。实操中还有一种混合方案大部分层用 INT8少部分关键层比如输出头、注意力投影保持 fp16这样能在精度和显存之间找到平衡。量化的选择逻辑很简单先跑 fp16 的基线效果记录回答质量如果显存够就用 fp16不够就降 INT8再不够降 INT4。每降一档都要拿自己的测试集重新评估不要只看显存能不能塞下。量化工具的选择上社区主流的方案基本都成熟了但要特别注意量化后的模型是否还支持 KV Cache 的正常复用有些量化实现为了省事会把 KV Cache 也量化省了显存但可能影响精度要仔细看文档。3.3 采样参数怎么调才不“瞎说”采样参数直接影响输出质量而且不同任务的最优参数差别很大。对话任务我通常用温度 0.7、top-p 0.9、top-k 50、重复惩罚 1.1这个组合在大多数场景下效果比较稳。创作类任务可以把温度调到 0.9 到 1.1输出更有发散性代码生成、翻译、信息抽取这类任务温度反而要调低到 0.1 到 0.3需要的是确定性和准确率。重复惩罚是个双刃剑。设得太低比如 1.0 等于没设长文本容易陷入复读设得太高比如 1.5 以上模型会开始刻意回避常用词句子变得生硬不自然。我一般从 1.05 起步发现问题再慢慢加很少超过 1.3。另外建议把“最大生成长度”设成合理值。Llama2 有时候生成的回答已经结束了但不会输出 EOS token而是一直续写“总结”“总之”之类的废话直到撞上长度上限。设置一个合适的 max_new_tokens比如 512 或 1024既能控制成本也能避免这种尾大不掉的情况。3.4 一个可在本地跑的完整流程示例下面给一个完整的推理流程示例用伪代码风格写出关键步骤方便对照自己的框架落地# 1. 加载分词器和模型 tokenizer load_tokenizer(llama2-7b-chat) model load_model(llama2-7b-chat, precisionfp16) # 2. 构造 prompt注意带上指令格式 prompt [INST] user_input [/INST] # 3. token 化并转成模型输入 input_ids tokenizer.encode(prompt, add_bosTrue) input_ids truncate(input_ids, max_len4096) # 4. 预填充 循环生成 output_ids [] next_id model.prefill(input_ids) # 处理 prompt返回第一个预测 while len(output_ids) max_new_tokens and next_id ! eos_id: output_ids.append(next_id) next_id model.decode_one(next_id) # 逐 token 生成自动维护 KV Cache # 5. 转回文本 result tokenizer.decode(output_ids)注意第四步里 Prefill 和 Decode 的接口是分开的实际使用的主流推理框架一般会把两者封装好但理解这两个动作的边界对排查显存和速度问题至关重要。4. 踩坑记录与排查技巧4.1 显存不足OOM最典型的报错就是 CUDA out of memory。遇到这个按优先级做三件事先把 batch size 降到 1再把 max_new_tokens 或者输入长度降下来如果还不行就量化模型。很多人第一反应是换更大的显卡其实很多时候是 KV Cache 或者并发请求太多把显存挤爆了而不是权重本身的问题。还有一个很容易忽略的点同上一个序列的 KV Cache 在会话结束后有没有被及时释放。在写多轮对话服务时如果每轮都把全部历史对话塞进缓存内存会只增不减必须手动管理会话的缓存生命周期。4.2 推理速度慢速度慢要先定位瓶颈在 Prefill 还是 Decode。可以看两个指标首 token 延迟和后续平均每 token 延迟。首 token 慢通常是 Prefill 阶段并行计算没有发挥好或者输入太长了每 token 慢则更可能是 Decode 阶段的带宽瓶颈。Decode 阶段每生成一个 token 都要把 KV Cache 重新读一遍这是典型的访存密集型操作算力反而不是主瓶颈所以纯堆显卡算力对 Decode 提速有限。如果确认是 KV Cache 的问题可以考虑几个方向改用支持 Flash Attention 类优化的推理引擎这类优化能减少注意力矩阵的显存占用并提升缓存复用效率对大模型做张量并行把 KV Cache 分散到多张卡以及用 GQA 本身带来的红利——70B 模型的 KV Cache 小Decode 速度反而可能比同等规模的 MHA 模型快。还有一个我自己踩过的坑有些框架默认开了 CPU offload权重在显存和内存之间来回搬运速度会断崖式下降检查配置时先把这个关掉。4.3 输出重复和胡说重复问题的根源通常是两个贪心解码或采样参数设置太激进导致模型陷入概率高峰的循环以及上下文过长时模型对后面内容的注意力被分散。处理手段按顺序试先加重复惩罚到 1.1 左右再把 top-p 调低到 0.85 左右如果还不行考虑给 prompt 加“不要重复已给过的信息”这类约束或者换用更小的采样集合。“胡说”也就是幻觉这个问题在 Llama2 上很常见尤其是事实性问答。模型本质是在做“最合理的续写”而不是查数据库。缓解手段包括把问题具象化要求模型先给出推理过程再给结论用 few-shot 提供参考回答格式对事实敏感任务限制输出格式以降低自由发挥空间。任何方案都无法完全消除幻觉这一点要在产品设计时就留好人工复核的余地。4.4 中文效果与模板问题Llama2 的原版模型中文能力明显弱于英文原因主要是训练语料里中文占比不高。解决思路有几个方向对话场景尽量用 chat 版本并遵循官方模板效果远好于直接用 base 模型中文任务可以用 System Prompt 指明语言和风格更彻底的办法是在中文数据上做微调也就是现在常说的中文社区增强版模型。模板问题值得单独说。很多人从网上抄了一段 prompt 模板但模板和模型训练时用的不一致导致模型经常忽略指令、自说自话。遇到这种情况第一步就是核对模板对话格式的模型必须用对应的指令包裹格式base 模型则别加对话模板两者混用是最常见的出错原因。4.5 常见问题速查表现象常见原因处理思路显存溢出KV Cache/并发/精度过高降 batch、降长度、量化首 token 很慢Prefill 过长/并行不佳缩短输入、检查框架并行设置每 token 很慢访存带宽受限/CPU offload换优化引擎、关 offload、GQA 模型输出无限重复采样参数/极端分布加重复惩罚、调低 top-p回答事实错误模型幻觉约束格式、few-shot、人工复核未遵循指令模板不匹配核对对话模板中文质量差语料比例低换中文增强版/微调5. 最后说点实操体会我自己第一次把一个 7B 模型完整跑起来的时候最大的感受是架构图上薄薄的一层“KV Cache”实际推理时竟然是显存管理里最头疼的一块论文里一句“GQA 降低了内存占用”在 70B 模型上能直接决定你是否需要多买一块卡。纸上谈兵和真正动手之间差的就是这些细颗粒度的理解。如果你正准备开始做 Llama2 的部署或二次开发我的建议是不要一上来就追 70B 大模型先用 7B 把架构、推理、参数调优整条链路跑通记录每一步的显存和耗时数据。这些经验可以直接迁移到任何同架构的新模型上。等你对 KV Cache 的增减、采样参数的扰动、量化的精度变化都有了直觉再上大模型、多卡并行会顺手得多。这个模型虽然已经不是最新的但它背后这套“预训练 指令微调 推理优化”的经验体系到今天依然值得花时间吃透。