★ FEATURED ARTICLE LLM 推理圈子里大家聊得最多的是 decode 阶段的采样策略、KV cache 淘汰和投机解码Prefill 经常被一句话带过。之前我也是这习惯直到做一个长文档问答项目一批 8k 长度 prompt 的首 token 延迟卡在 3 秒上下怎么调都下不去才被迫回过头把 Prefill 阶段 decoder 层内部的数据流整个过了一遍。之后我把 QKV 投影、多头注意力、KV cache 写入这些环节逐个测才发现 prefill 和 decode 在内存访问、算力消耗上完全是两种性格。这篇文章就从一个实战者视角把 Prefill 阶段 decoder 层多头注意力单元的数据流转完整写清楚每一层输入长什么样、注意力头内部怎么拆和并、哪些中间结果吃显存、哪些算子决定延迟以及工程上最容易出问题的几个点。想看懂推理框架的显存规划和性能优化Prefill 这条数据流是最好的切入点。1. Prefill 阶段在全链路中的定位与性能模型1.1 Prefill 和 decode同一套权重两条完全不同的执行路径一次完整的大模型生成可以分为两段。Prefill 输入整段 prompt在一次前向计算里把所有位置过一遍decode 每次只输入当前已生成的 token逐个往前推。我们通常说的“首 token 延迟”主要由 prefill 决定“每秒生成 token 数”由 decode 决定。但数据流上两个阶段不只是“一次算一批”和“一个个算”的区别。Prefill 阶段输入张量是 (s, h)s 是 prompt 长度h 是隐藏维度decode 阶段输入张量是 (1, h)如果带 batch 也就是 (n_batch, h)但每个 batch 通常只放新生成的 token。矩阵乘法的形状里有一个维度从 s 变成 1算术强度完全变样。同样是 QKV 投影prefill 里权重矩阵被序列的每个位置重复使用FLOPS 随 s 线性增长decode 里权重矩阵只被 1 个位置读到数据搬运量没变计算量却掉到原来的 1/s。这导致 prefill 是计算密集compute-bounddecode 是访存密集memory-bound。这个差异直接决定优化方向。decode 阶段我们习惯把权重做低比特、压缩 KV cache、减少每次 kernel 启动的数据搬运prefill 阶段如果只做这些收益有限因为它并不缺带宽缺的是把 SM 上的矩阵乘单元喂满。反过来prefill 阶段那些优化比如增大 batch、提高算术强度、算子融合、flash attention拿到 decode 阶段也未必有效因为 decode 的计算量太少融合得再漂亮也救不了带宽瓶颈。1.2 算力与带宽的平衡点怎么估算在 GPU 上判断一个大算子该往哪个方向优化会先看它的算术强度。A100 的 float16 tensor core 算力大约 312 TFLOPSHBM 带宽大约 2TB/s这意味着每个字节的数据需要约 156 次浮点运算才能让算力和带宽打平。对 prefill 的 QKV 投影输入激活 (s, h) 和权重 (h, 3h) 从 HBM 读入输出 (s, 3h) 写回计算量约 2×s×h×3h FLOP。当 s 很大时除了权重那部分固定开销算术强度近似跟 s 成正比很容易超过 156 比 1所以 QKV 投影在长 prefill 里理论上可以逼近峰值算力。注意力部分就不一样。QK^T 和 PV 如果不做融合需要把 (s, n_head, s) 的分数矩阵写回 HBM 再读出来显存中间量会先爆炸。flash attention 的核心贡献不只是省显存它把 attention 的中间结果留在 SRAM 里反复复用HBM 流量从 O(s^2) 量级降到 O(s)这才是长序列下又快又稳的根本原因。这些内容后面会展开。这一节给整篇文章定调Prefill 是一条有明确算力上限的路径任何让中间数据在 HBM 和 SM 之间多走几趟的实现都是潜在性能杀手。2. Decoder 层整体数据流从 token 序列到隐藏状态2.1 Embedding 与层输入输出格式输入侧先做 embedding 查表prompt 的 token IDs 长度为 s查表后变成形状 (s, h) 的浮点张量。h 通常来自模型配置常见开源模型用 4096 或 5120。这个 (s, h) 张量在层与层之间一直保持不因为多头而要临时改成三维直到进入注意力单元内部才会按 head 拆开。以 LLaMA 类开源模型结构为例每个 decoder 层有两个子块多头注意力和 FFN。数据流伪代码如下# 伪代码仅用于说明数据流 x_res x # (s, h) x rms_norm(x, w_norm) # (s, h) q, k, v split(x w_qkv) # 各 (s, h)之后按头拆分 o attention(q, k, v) # (s, h) x_res x_res o # 残差 x rms_norm(x_res, w_norm2) # (s, h) x silu(x w_gate) * (x w_up) # (s, 4h) x x w_down # (s, h) x_res x_res x为了少启动几个 kernelQKV 投影的权重通常按 [h, 3h] 拼在一起一次矩阵乘输出 (s, 3h)再在逻辑上切三段。FFN 里的 up 和 gate 投影也常常合并成一次输出 2×4h 的大矩阵乘。这些合并都服务同一个目的让一个大 GEMM 替代多个小 GEMM。GPU 上矩阵乘的效率跟张量形状的规整度强相关小 kernel 除了调度开销还会让 tensor core 的利用率下降。我早先 debug 时见过一个现象同一份模型把三个 GEMM 拼成一个后矩阵乘吞吐从 45% 升到 80% 以上。数据流逻辑没变变的只是内存布局和 kernel 粒度但结果差异巨大。2.2 注意力单元里的 QKV 投影与多头拆分继续往注意力单元内部走。从 (s, 3h) 按列切成 Q、K、V 之后每个都是 (s, h)。多头拆分在数学上就是 reshape 成 (s, n_head, head_dim)注意顺序它把一个 head 的所有维度放在最后一维。如果我们更关心每个 head 内部的数据流可以把 Q 看成 n_head 个 (s, head_dim) 矩阵K、V 同理。这里有个常见的误区很多人以为 QKV 投影之后张量就变成了五维六维其实物理上它还是连续的显存区域只是 kernel 通过 stride 把它解释成多个头。这个“解释”方式非常关键不同框架的头拆分顺序可能不同有的是 (s, n_head, head_dim)有的是 (s, head_dim, n_head)后续所有矩阵乘的访存模式都会跟着变。2.3 多头注意力内部的张量变换与残差回写这一步是 Prefill 阶段数据流最核心的一段。对每个 query headQ 形状是 (s, head_dim)同头对应的 K 转置后是 (head_dim, s)两者乘出注意力分数矩阵形状 (s, s)。这个 (s, s) 矩阵每一行代表一个 query token 对全部 key token 的权重。语言模型里的因果关系要求每个位置只能看当前位置及之前位置所以矩阵上三角要掩码掉。之后 softmax 归一化再用这个权重和 V 矩阵形状 s×head_dim做加权求和得到当前 head 的输出 (s, head_dim)。所有 head 的输出再拼回 (s, h)过 W_o 输出投影到与输入相同的隐藏维度最后加回残差流。这个单元的输出形状和输入完全一致所以 decoder 层可以无限堆叠这是 transformer 能做成深层结构的一个前提。这里有一个容易忽略的点Q 与 K 的矩阵乘在数学上要求 K 按 head 转置但这个转置不一定真的要在显存里搬一次数据。有的实现通过 kernel 里调整 stride 来完成逻辑转置有的实现则做物理转置。到底是物理转置还是逻辑转置会影响矩阵乘的访存模式和 tiling 效率后面“工程实现”里会再展开。3. 逐算子拆解多头注意力单元的数据流3.1 从 QKV 投影到 QK^T 缩放点积以预处理后的张量形态来看Q 投影输出形状是 (s, n_head, head_dim)算上 batch 时会在前面多一维我们先把 batch 隐藏掉。每个 query head 的矩阵 Q_head 形状是 (s, head_dim)同 head 的 K_head 形状是 (s, head_dim)。分数矩阵 S Q_head K_head.T再除以 sqrt(head_dim)防止内积随维度增大而过大。head_dim 常见值是 128所以缩放系数约 11.31。因果掩码有两种实现方式。一种是“数据流式”构造一个 (s, s) 的 mask允许位置是 1不允许位置是 -inf让 mask 参与 attention 的加法再用 softmax 压掉。另一种是“控制流式”在融合 kernel 里根据 query 块和 key 块的索引直接跳过或只计算块内部分不把 mask 作为显存里的矩阵。后一种在长序列的 prefill 里优势明显因为掩码不会变成额外的数据搬运。很多老代码优化前用的是前者跑到 8k 就慢得没法看。从内存量级上看如果完全不融合一个 8k 长度、32 个 head 的 prompt每层注意力分数矩阵 S 在 fp16 下就是 8192×8192×32×2 字节接近 4GB。这个数字随着序列变长还是平方级增长所以任何不做分块的注意力实现在长 prompt 上都撑不住。3.2 在线 Softmax 与加权求和softmax 的朴素实现需要先知道每行最大值和 exp 和所以通常要把 S 扫描两遍。问题是在 flash attention 分块逻辑里S 矩阵是分块算出来的不能等全部算完再做 softmax因此要用 online softmax对每一块维护局部最大值和局部 exp 和新的块和旧块合并时对旧结果做重缩放。这样做的好处是只需要一遍分块循环不需要把完整的 S 写回 HBM。加权求和这一步 D softmax(S) VV_head 形状是 (s, head_dim)。在分块注意力视角下V 块被读取的次数和 key 块被读取的次数一致。如果分块尺寸选得合适Q/K/V 块都能放进 SRAM块内部的点乘和加权在片内完成只有最终累积结果和 KV cache 内容需要写回 HBM。这也是为什么同样数学含义的注意力融合实现的 HBM 流量能比朴素实现低一个量级。我实际用 profiler 看 A100 上的表现时未融合的注意力在长序列下 HBM 读写量经常是融合实现的 5 到 10 倍而且差距随序列长度继续拉大。HBM 流量是最直观的指标先看它基本能判断一个注意力实现的好坏。3.3 输出投影与多头的合并多头输出合并指的是把 (s, n_head, head_dim) 恢复成 (s, h)再过一个输出投影 W_o。合并时内存布局有讲究标准做法是把 n_head 放在隐藏维度的前半部分也就是 (s, n_head, head_dim) 直接平铺成 (s, h)这样 W_o 的列和头顺序一一对应输出投影可以当作普通 GEMM。有些框架内部把头维度和 head_dim 的顺序反过来或者为对齐 SIMD 宽度加 padding这时 W_o 的权重布局也要跟着调。你直接看到的 tensor shape 可能不是数学推导里最直觉的那种这是为了性能做的妥协。输出投影之后的数据流和 FFN 一样过残差、做 RMSNorm、进 MLP最后再回残差。如果多层叠加prefill 的数据流就是把这些步骤重复 L 次每一层之间要清掉中间激活。很多推理框架的显存峰值就卡在这种“每层临时 buffer 是否及时释放”的地方。3.4 MHA 和 GQA 对 KV cache 数据量的影响注意力结构还有一个和 KV cache 直接相关的设计选择。早期模型普遍用 MHAmulti-head attention每个 query head 都有独立的 K/V head。后来长上下文场景越来越多GQAgrouped-query attention流行起来多个 query head 共享一组 K/V head。假设一个模型 h4096n_head32head_dim128。若是 MHAKV cache 每层每个 token 需要保存 K、V 两组每组 32×128 个 fp16 数值合计 16KB。若用 GQA 把 n_kv_head 压到 8同一项降到 4KB。在 32 层、32k 上下文的场景里MHA 的 KV cache 约 16GBGQA 约 4GB。省下来的不只是显存Prefill 阶段把整条 prompt 的 K/V 写入 DRAM 的字节数同样省 4 倍decode 阶段每次读取 K/V 的带宽也成比例下降。这也是为什么新的推理框架在配置模型时会把 n_kv_head 作为一等公民来处理。数据流分析如果不区分 query head 和 kv head后面的显存预估会出现 4 倍级别的误差。4. 工程实现中的关键选型与实操要点4.1 用大 GEMM 承载 QKV 投影减少 kernel 启动和碎片前面讲过 QKV 投影在 prefill 里是 tensor core 的主战场。实际工程落地时权重拼接、激活拼接、矩阵乘 kernel 选择都要围绕“一次大 GEMM”来设计。权重矩阵如果逻辑上分成 W_q、W_k、W_v 三个物理上也分开存那数据流就是三次 (s,h)×(h,h) 矩阵乘每次都要重新读一遍输入激活 x且 kernel 启动三次。拼成 (h, 3h) 后一次 GEMM 解决输入 x 只需从 HBM 读一次。对 4096 维模型输出激活是 3×s×h中间张量变大但峰值显存在分块执行下仍然可控。此外GEMM 的 M 维也就是 seq 长度要尽量规整。flash attention 和矩阵乘 kernel 在分块 tiling 时如果 M 正好是 block 的整数倍尾部边界检查可以简化SM 占用也更整齐。所以工程上常见做法是对 prompt 做 padding或者至少把 kernel 的 M 维对齐到 64 或 128 的倍数哪怕多算一点 padding 都值。这个小细节常常被忽略但它对算力利用率的影响可以到两位数百分点。4.2 KV cache 布局按层连续、按 block 连续还是按 head 交错KV cache 的物理布局直接影响 prefill 写入和 decode 读取的带宽。三种常见布局按层连续每一层一个连续大 buffer内部按 token 顺序存 K 和 V适合离线处理但在服务场景里多序列内存管理不便。按 block 连续以固定 block 为单位分配每个 block 内部存若干 token 的 K/V配合分页调度方便前缀共享和续跑是目前主流推理框架的选择。按头交错把 head 维度放在外层或用跨步访问方便某些 kernel 的向量化加载但需要相应 kernel 支持。从数据流角度真正要注意的是“写入时不需要额外搬运”。如果 KV cache 布局和 flash attention kernel 的输出布局不一致要么 kernel 的 epilogue 做一次 scatter要么单独再走一个 copy kernel。前者可以融合后者就是纯浪费。我在调试时见过一个实现注意力计算已经融合了但 KV cache 写入单独用了 memcpy导致 prefill 的显存写带宽多了一倍开销去掉这次拷贝后性能立竿见影。4.3 Flash Attention 的分块与掩码控制流Flash attention 在 prefill 阶段对长序列几乎是必需品。它把 Q 分成 query blockK/V 分成 kv block做双层循环。对每个 query block累加一个 O 块和两个在线 softmax 状态量。因果掩码在控制流里的处理是关键如果 kv block 索引整体大于 query block 索引这整块 key 在当前 query 之后可以直接跳过如果 kv block 索引小于 query block整块都能参与不需要掩码相等时块内按列掩码。这个逻辑看似简单但在实现中要处理 block 边界和 padding稍不留神多算了 padding 位置的 keymask 就会漏。一个实用的经验是把 seq_len 手动对齐到 block_size 的整数倍比在 kernel 里处理任意余数稳定得多。另一个实操点是 head_dim 和 SRAM 的关系。SRAM 有限flash attention 的 block 大小不能随便调。head_dim128、dtypefp16 时一个头的一行是 256 字节一个 128×128 的 Q/K/V 块大约是 96KB加上 O 和 softmax 状态SM 上的 SRAM 只能放下少数几个块。block 太大SM 并行度下降block 太小循环开销增加。不同 GPU 上的最优值不太一样建议用 profiling 数据来定不要照抄别人配置。4.4 激活峰值与 KV cache 的显存博弈很多人一提到长 prompt 的显存第一反应是 KV cache但 prefill 阶段的峰值显存其实常常由中间激活决定。在朴素逐层执行里QKV 投影的输出 (s, 3h)在 fp16 下s8192、h4096 时单层接近 192MB。如果框架没有及时释放32 层叠加的激活会迅速膨胀。工程上普遍做法是“按层复用一块 activation buffer”因为有残差连接前后层依赖并不需要所有激活同时驻留除非你要做分层 offload。此外prefill 期间 KV cache 并不是空的它会随着每一层推进写满整段 prompt 的数据。KV cache 本身、激活、权重、临时 buffer 加在一起就是 OOM 现场最常见的组合。排查时不要只看缓存总量要按“权重 激活峰值 KV cache 增长 临时中间量”四类分别统计。实际操作中我会在框架里加几个计数器分别统计这四类数据OOM 时看看到底是哪一类的增长曲线出了问题。5. 排查实录几个由数据流引起的真实问题5.1 QKV 权重布局错误导致的算力浪费我有一段时间用一份从训练框架直接导出的权重跑推理长 prompt 的 prefill 性能一直上不去profiler 显示 QKV 投影的 tensor core 利用率只有 30% 出头。逐步排查发现权重矩阵的存储顺序是 [n_head, head_dim] 的行主序而推理框架的 GEMM 为了向量化加载期望的是 [head_dim, n_head] 的列主序布局。不匹配时矩阵乘内核只能走低速的通用路径或者先对每个 token 的激活做逻辑转置。把权重在加载时做一次预转置后QKV 投影时间几乎砍半整个 prefill 延迟降了 28%。这个案例让我的习惯变成换框架、换导出脚本时第一件事不是看算力而是看权重 layout 是否和 kernel 匹配。数据流示意图画出来很干净但物理内存里一个转置就能让性能天差地别。5.2 长 Prompt 的 Prefill 利用率反而低另一类常见问题是prompt 长度不够长时算力利用率反而上不去。原因是 GEMM 的 M 维太小SM 调度出现“波浪”不满一轮。举例GPU 上一个计算任务可能有 100 多个 SM如果 GEMM 按 128×128 的 tile 切分总 tile 数只有几十个SM 就会有一批空载。处理办法是尽量把请求一起 batch或者对短 prompt 走 decode 式逐个 token 处理只把超过一定长度的请求走 prefill 专用路径。这里的数据流逻辑是算力利用率取决于矩阵乘的形状而不只是模型大小同样的权重在不同 seq 长度下性能差异会非常明显。我在项目里给 prefill 设的阈值是 512 token小于这个值直接走 decode 路径大于这个值才走完整 prefill 路径效果比统一处理好了不少。5.3 常见问题速查表现象可能原因优先排查方向显存峰值异常高激活 buffer 未复用 / 注意力中间量未融合看 QKV 输出和激活峰值检查 flash attention 是否启用算力利用率低于预期GEMM 形状不规整 / 权重 layout 不匹配profiling 看 QKV kernel检查序列 padding 和权重转置8k 以上 prompt 明显变慢注意力矩阵打满显存 / 未用分块确认 flash attention 是否生效观察 HBM 读写量KV cache 占用不合预期MHA 与 GQA 配置错误 / 多序列碎片检查 n_kv_head核算每 token 字节数prefill 首 token 延迟波动大存在逐 token 处理的退化路径看是否有小请求走入 decode 式循环5.4 一点工程习惯做 Prefill 优化我最后的建议是“先画数据流再上 profiler”。把每一层的张量形状、中间量落在显存里的大小、kernel 融合的位置画清楚很多问题不用 profiling 就已经暴露了。我做过几次优化后发现真正的热点往往不是最深奥的算法而是某个没有必要的转置、一个没有对齐的 seq_len、一次不该出现的 HBM 回写。数据流图画出来这些问题都会自己跳到眼前。像这次复盘里最值钱的几个改动全是沿着这条数据流找出来的而不是对着大模型一通瞎调参。
阅读完成 · 觉得有帮助？
