【AI专栏】图解Transformer - 第05章:LLM 推理工程

【AI专栏】图解Transformer - 第05章:LLM 推理工程 博主介绍程序喵大人35 - 资深C/C/Rust/Android/iOS客户端开发10年大厂工作经验嵌入式/人工智能/自动驾驶/音视频/游戏开发入门级选手《C20高级编程》《C23高级编程》等多本书籍著译者更多原创精品文章首发gzh见文末记得订阅专栏以防走丢C基础系列专栏C语言基础系列专栏C大佬养成攻略专栏C训练营个人网站【AI专栏】图解Transformer - 第04章LLM生成1. 图解 LLM 推理工程前四章把模型本身拆完了——token 怎么变向量、attention 怎么算、Block 怎么堆、生成循环怎么转。这一章换个视角同一个模型跑一次请求到底要付出什么工程成本你在聊天框打一段 prompt 发出去到第一个字冒出来再到整段回答写完——中间经历了 prefill、decode、KV Cache 增长、显存占用、attention 开销。每一步都有对应的资源消耗也都有对应的优化手段。GQA/MQA 减少缓存压力FlashAttention 减少数据搬运量化压缩权重体积Speculative Decoding 减少大模型串行调用。这些技术看着像独立的优化技巧但拆开都是在回答同一个问题怎么让一次请求跑得更快、更省显存。本章只看一个请求内部的成本。多个用户同时提问、排队、调度、continuous batching——那就是第六章的事啦。2. Prefill一口气读完整个 prompt用户发出 prompt 以后模型做的第一件事是 prefill把整段 prompt 一口气读完。Prompt 里可能有几十个、几百个甚至上千个 token。在 prefill 阶段这些 token 一起进入模型并行走完所有 Transformer Block。每一层 attention 都能同时看到 prompt 里所有可见位置——这跟 decode 阶段一次只处理一个 token 完全不同。Prefill 做完两件事一是算出 prompt 最后位置的 hidden state为生成第一个新 token 做准备二是建立初始的 KV Cache——把 prompt 中每一层每个位置的 Key 和 Value 都存下来后续 decode 直接复用。所以 prefill 的计算量跟 prompt 长度直接相关。prompt 越长这个阶段越重但它只跑一次。用户感受到的从发送到第一个字出现的等待时间很大一部分就花在这里。在推理服务的指标里这段等待叫 TTFTTime To First Token。prompt 很长的场景——比如把整篇文档丢给模型——TTFT 会明显变长因为 prefill 要并行处理的 token 数量更大。3. Decode每轮只生成一个新 tokenPrefill 结束以后模型进入 decode 循环。第四章讲过自回归生成——每一步只预测一个新 token选出来追加到序列末尾再继续预测下一个。从推理工程的角度看decode 每一步只有一个新 token 进入模型。它拿着自己的 Query去读取 KV Cache 里所有历史 token 的 Key 和 Value经过 attention、FFN、LM Head产出下一个 token。然后这个新 token 的 Key 和 Value 也被追加进 cache供后续步骤使用。Decode 的特点是串行。每一步依赖上一步的输出没法跳过也没法把未来的 step 提前算。这意味着生成 100 个 token 就要跑 100 次 decode step——每步的延迟直接累加到用户等待时间上。单步 decode 的计算量比 prefill 小很多只处理一个 token但它受带宽瓶颈影响更大每一步都要从显存里读取整个模型的权重和当前长度的 KV Cache。模型越大、上下文越长这个读取量越可观。4. Prefill vs Decode 时间轴把 prefill 和 decode 放在一条时间轴上看一次请求的推理过程分成两个阶段Prefill 阶段所有 prompt token 一起进入模型并行计算。这一步计算密集——大量矩阵乘法同时发生GPU 的算力被充分利用。计算完成后prompt 的 KV Cache 就绑定好了第一个新 token 也随之产出。Decode 阶段逐个 token 串行生成。每一步只处理一个新 token计算量小但每步都要从显存读取权重和不断增长的 KV Cache。GPU 算力利用率往往不高瓶颈更多在显存带宽上。两个阶段的计算形态差异很大。Prefill 是 compute-bound计算密集decode 是 memory-bandwidth-bound带宽密集。同一块 GPU跑 prefill 和跑 decode 时的资源瓶颈完全不同。KV Cache 在 prefill 之后出现在 decode 过程中逐步增长——每生成一个 tokencache 就多一行。5. KV Cache缓存历史 K/VKV Cache 是推理工程里最重要的优化之一。它存的是什么每一层 attention 里历史 token 的 Key 和 Value。回顾一下 attention 的计算当前 token 的 Query 要和所有可见 token 的 Key 做匹配再用权重去读取对应的 Value。如果没有 cache每个 decode step 都要重新算一遍所有历史 token 的 K 和 V——序列越长重复计算越多。KV Cache 把已经算过的 K 和 V 存在显存里。下一步生成时只需要计算新 token 的 K/V然后和 cache 里的历史拼在一起就行。历史部分直接复用不重新计算。注意cache 存的是 K 和 V不是 Q。每个 decode step 的 Query 是当前新 token 自己的每步都不同没法缓存。而历史 token 的 K/V 在当前 step 不会变——它们在之前的 step 已经算好了直接读就行。6. KV Cache 每层都有KV Cache 是 per-layer 的。模型有多少层 Transformer Block就有多少份独立的 KV Cache。第 1 层的 attention 有自己的历史 K/V第 2 层也有自己的第 N 层也有。decode 时新 token 沿着层逐步向上流动每到一层都要从该层自己的 cache 里读取历史 K/V 来做 attention。容易搞混的一个点KV Cache 是按层存放的各层之间的 cache 互相独立内容也不同。第 1 层存的是第 1 层 attention 投影出来的 K/V跟第 2 层存的完全不是一回事。这意味着 cache 的总量 每层的 cache 大小 × 层数。模型层数越深cache 占的显存越多。一个 32 层的模型就有 32 份独立的 KV Cache 同时占着显存。7. KV Cache 显存账本KV Cache 到底占多少显存可以用一个近似公式来估算KV Cache ≈ layers × seq_len × kv_heads × head_dim × 2 × bytes_per_element拆开看每个因子layers模型层数。32 层就是 32 份 cache。seq_len当前可见的序列长度。prompt 已生成的 token 数。kv_heads每层有多少组 K/V head。标准 MHA 下等于 attention head 数。head_dim每个 head 的维度。2Key 和 Value 各一份。bytes_per_element数据类型的字节数。FP16 是 2 字节FP32 是 4 字节。这些因子相乘量级上去得很快。以一个 32 层、32 个 KV heads、head_dim128、FP16 的模型为例序列长度 4096 时KV Cache 大约占 2 GB。序列拉到 32K就变成约 16 GB——这还只是一个请求的 cache没算模型权重本身。这个账本说明了一件事上下文长度是 KV Cache 显存的核心变量。序列越长cache 越大GPU 剩余的显存空间就越少。8. 上下文越长缓存越大把上下文长度从 2K 拉到 8K 再到 32KKV Cache 的增长是线性的——序列翻 4 倍cache 也翻 4 倍。这在实际使用中有很直接的感受把一整本手册丢给模型当上下文prompt 阶段就要建立很大的初始 cache。之后每生成一个 tokencache 又多一行。decode 每一步的 attention 也需要扫更长的历史 K/V——计算量和带宽需求同步增长。长上下文贵贵在两个地方一是 KV Cache 的显存占用持续增长二是 attention 每步需要匹配的历史 key 更多计算量也更大。这也是为什么很多模型虽然声称支持 128K 甚至更长的上下文窗口但实际使用时推理速度会显著下降、费用明显上升。上下文窗口的大小只是能不能放得下的问题真正的工程挑战在于放得下以后还跑得快不快。后面会看到GQA/MQA、FlashAttention、量化这些优化很多都是在对抗长上下文带来的资源压力。9. GQA / MQA减少 KV headsKV Cache 的显存公式里有一个因子是 kv_heads。标准的 Multi-Head AttentionMHA里每个 Query head 都配一个独立的 Key head 和 Value head——head 数一样多cache 也最大。GQAGrouped-Query Attention和 MQAMulti-Query Attention的思路是Query head 可以保持多个但 K/V head 不需要那么多。MQA 最激进所有 Query head 共用同一个 K/V head。KV Cache 直接缩到原来的 1/hh 是 query head 数。代价是只有一组 K/V 视角模型的表达灵活性可能下降。GQA 取折中把 Query head 分成几组每组共享一个 K/V head。比如 32 个 Query head 分成 8 组就只需要 8 个 K/V head。Llama 2 70B、Llama 3 等模型采用的就是 GQA。效果很直接K/V head 从 32 减到 8KV Cache 显存降到原来的四分之一attention 中 K/V 的带宽需求也同比下降。Query head 数量不变模型在 query 端的表达能力保持不变。10. Attention 复杂度长上下文为什么贵Attention 的核心操作是让 Query 和所有可见的 Key 做匹配。可见历史越长匹配次数越多。Prefill 阶段prompt 内所有 token 互相做 attention。n 个 token每个都要和其他 n 个位置算匹配——attention 矩阵是 n×n 的计算量大约和 n² 成正比。prompt 从 1K 拉到 4K这部分计算量大约变成原来的 16 倍。Decode 阶段每一步只有一个新 Query但它要看所有历史 Key。如果当前已经生成了 n 个 token这一步的 attention 就要做 n 次匹配。单步成本和已有上下文长度成正比——越往后生成单步越慢。要注意的是attention 的计算成本只是推理成本的一部分。权重加载、FFN 计算、KV Cache 读写也都有开销。把长上下文贵全部归因于 attention 是不准确的但 attention 确实是随上下文增长最快的那个环节。11. FlashAttention少搬数据FlashAttention 是 exact attention——它算出来的结果和标准 attention 一模一样不做任何近似。那它优化的是什么数据搬运。GPU 内部有两级存储HBM高带宽显存容量大但读写相对慢和 SRAM片上缓存容量很小但读写极快。标准 attention 的实现方式是先在 HBM 里把 QK^T 算出来存成一个巨大的 n×n attention 矩阵再读回来做 softmax再读 V 做加权求和——中间来回搬运了很多数据。FlashAttention 的做法是 IO-aware tiling把 Q、K、V 切成小块tile一次只把一小块搬进 SRAM在 SRAM 里算完局部 attention用 online softmax 技巧把局部结果逐步合并成最终结果避免在 HBM 里存那个巨大的 attention 矩阵。核心收益HBM 读写量大幅下降。实际跑起来FlashAttention 在长序列上能快 2-4 倍同时减少显存占用。12. FlashAttention 数据流拆开看 FlashAttention 的数据流HBM 里放着完整的 Q、K、V 矩阵。计算时不是一次性把所有数据拉进来而是分批把小块 Q tile、K tile、V tile 搬进 SRAM。SRAM 里一次只处理一小块。在这个小块上算 QK^T、做 scale、做 mask、算 softmax、加权 V——完整的 attention 流程在 tile 级别跑一遍。算完以后结果写回 HBMSRAM 腾出来给下一块用。关键技巧是 online softmax标准 softmax 需要先看完所有分数才能计算。FlashAttention 用了一种数值稳定的在线算法可以一边看一边更新 softmax 的分母不需要在 HBM 里存完整的 attention 矩阵。最终效果attention 的计算结果不变exact但中间不再需要 O(n²) 的 HBM 临时存储HBM 读写次数也大幅减少。13. 量化用更少 bit 存权重模型权重通常以 FP1616 bit或 BF16 存储。一个 7B 参数的模型FP16 下权重大约占 14 GB 显存。量化的思路是用更少的 bit 来表示权重。INT8 量化把每个参数从 16 bit 压到 8 bit显存占用减半。INT4 更激进压到 4 bit显存占用降到原来的四分之一。量化带来两个直接好处一是模型占的显存变少同一块 GPU 能跑更大的模型或者留更多空间给 KV Cache二是权重从显存搬到计算单元的带宽压力下降——同样的带宽单位时间能搬运更多参数。对于 decode 这种 memory-bandwidth-bound 的场景量化的加速效果尤其明显。每一步 decode 都要把整个模型权重从显存读一遍权重体积小了读取时间就短了。量化也可以应用在 KV Cache 上。有些框架支持把 cache 以 INT8 或更低精度存储进一步压缩长上下文的显存压力。14. 量化权衡速度、显存、精度量化不是免费午餐。精度方面用更少 bit 表示参数必然损失一些信息。量化后模型的输出可能和原始 FP16 模型有细微差异。差异大不大取决于量化方案、校准数据和模型本身。有些任务对精度敏感数学推理、代码生成有些相对宽容开放对话、摘要。速度方面量化降低了数据搬运量理论上会加速。但实际加速幅度取决于硬件是否支持对应精度的高效计算。比如 INT4 矩阵乘法需要专门的 kernel不是所有 GPU 都能跑出理想速度。有些量化方案甚至需要在推理时做反量化dequantize反而增加了计算步骤。选量化方案是一个工程决策。需要在目标硬件上实际测速、测精度、测任务表现然后做取舍。笼统地说量化一定好或量化一定有损都不准确——得看具体场景。一条工程原则measure first。先跑 benchmark再做决定。15. Speculative Decoding小模型先猜Decode 阶段最大的痛点是串行每步只产一个 token步数等于生成长度。大模型跑一步很贵100 个 token 就要跑 100 步。Speculative Decoding 的思路是用一个小模型draft model先快速猜一串候选 token。小模型参数少、跑得快一口气猜出 4-8 个 token 的成本可能还不到大模型跑一步的时间。猜完之后把这串候选 token 一起交给目标大模型target model验证。大模型一次前向传播就能并行检查这些候选——跟 prefill 类似多个 token 一起处理。通过验证的 token 直接采纳第一个被拒绝的位置由大模型的输出替换后面的候选作废。当小模型猜得准的时候比如猜 4 个对了 3 个一次验证就相当于跳过了 3 个 decode step。当猜得不准的时候最坏也只多浪费了小模型的计算时间。关键的理论保证通过特定的接受/拒绝策略Speculative Decoding 可以保持目标模型原始的输出分布。最终生成的文本质量跟直接用大模型生成是一样的。加速的来源是减少了昂贵的大模型串行调用次数。16. 单请求推理成本总览把整章串起来。一次请求从 prompt 到输出走的是这条路Prompt → Prefill并行处理建立 KV Cache→ Decode Loop逐 token 生成cache 逐步增长→ Output成本来源有四个维度计算prefill 阶段的大规模矩阵乘法decode 阶段每步的前向传播。显存模型权重占一块KV Cache 随上下文增长占另一块。带宽decode 每步要把权重和 cache 从 HBM 搬到计算单元。延迟prefill 决定 TTFTdecode 的逐步串行决定生成速度。对应的优化手段每个都在削减某一项或几项成本GQA / MQA减少 KV heads → 降低 cache 显存和带宽。FlashAttentiontiling online softmax → 减少 HBM 读写。量化低 bit 表示 → 降低权重显存和带宽。Speculative Decoding小模型先猜 → 减少大模型串行 decode 步数。这些优化互相不冲突可以同时使用。实际的推理系统往往是多种优化叠加的结果。到这里单个请求内部的成本和优化就讲完了。下一章进入多用户服务场景当很多用户同时发请求模型服务怎么排队、怎么调度、怎么把 GPU 利用率拉满。码字不易欢迎大家点赞关注评论谢谢