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-bound），decode 是访存密集（memory-bound）。
这个差异直接决定优化方向。decode 阶段我们习惯把权重做低比特、压缩 KV cache、减少每次 kernel 启动的数据搬运；prefill 阶段如果只做这些，收益有限，因为它并不缺带宽，缺的是把 SM 上的矩阵乘单元喂满。反过来，prefill 阶段那些优化，比如增大 batch、提高算术强度、算子融合、flash attention，拿到 decode 阶段也未必有效，因为 decode 的计算量太少，融合得再漂亮也救不了带宽瓶颈。
1.2 算力与带宽的平衡点怎么估算
在 GPU 上判断一个大算子该往哪个方向优化，会先看它的算术强度。A100 的 float16 tensor core 算力大约 312 TFLOPS，HBM 带宽大约 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为了少启动几个 kernel，QKV 投影的权重通常按 [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 head，Q 形状是 (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) @ V，V_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 直接相关的设计选择。早期模型普遍用 MHA（multi-head attention），每个 query head 都有独立的 K/V head。后来长上下文场景越来越多，GQA（grouped-query attention）流行起来：多个 query head 共享一组 K/V head。
假设一个模型 h=4096，n_head=32，head_dim=128。若是 MHA，KV cache 每层每个 token 需要保存 K、V 两组，每组 32×128 个 fp16 数值，合计 16KB。若用 GQA 把 n_kv_head 压到 8，同一项降到 4KB。在 32 层、32k 上下文的场景里，MHA 的 KV cache 约 16GB，GQA 约 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 block，K/V 分成 kv block，做双层循环。对每个 query block，累加一个 O 块和两个在线 softmax 状态量。因果掩码在控制流里的处理是关键：如果 kv block 索引整体大于 query block 索引，这整块 key 在当前 query 之后，可以直接跳过；如果 kv block 索引小于 query block，整块都能参与，不需要掩码；相等时，块内按列掩码。这个逻辑看似简单，但在实现中要处理 block 边界和 padding，稍不留神多算了 padding 位置的 key，mask 就会漏。
一个实用的经验是：把 seq_len 手动对齐到 block_size 的整数倍，比在 kernel 里处理任意余数稳定得多。另一个实操点是 head_dim 和 SRAM 的关系。SRAM 有限，flash attention 的 block 大小不能随便调。head_dim=128、dtype=fp16 时，一个头的一行是 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 下，s=8192、h=4096 时单层接近 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 回写。数据流图画出来，这些问题都会自己跳到眼前。像这次复盘里最值钱的几个改动，全是沿着这条数据流找出来的，而不是对着大模型一通瞎调参。
