前阵子帮一个团队排查prefill阶段的显存溢出问题，模型是标准的decoder-only结构，prompt长度一超过两千个token，单卡就撑不住。我第一反应是去看attention实现，果然，代码里老老实实把注意力分数矩阵构建成了完整形状。那一刻几乎可以断定，问题就出在那张中间张量上。

这个经历其实很典型。很多人聊大模型推理时，能把attention公式倒背如流，但要真问到"prefill阶段、decoder层、多头注意力单元里，每个张量从进来到出去，形状怎么变、步长怎么走、中间结果落在显存哪里"——就含糊了。这篇内容就是把这条数据流完整拆开，按张量级别过一遍，并且说清楚每一步背后的工程取舍。适合正在做推理优化、想自己写attention kernel、或者单纯想把transformer源码读到"形状级"程度的人。

1. Prefill在推理请求中的位置：为什么这个阶段值得单独拆开看

1.1 一条请求里的两段式生命周期

所有decoder-only架构的模型，无论是GPT、LLaMA还是各类开源变体，生成一条回复都要经历两个阶段。第一阶段叫prefill（预填充），把用户输入的整段prompt一次性塞进模型，从第一层跑到最后一层，算出所有层的中间激活，同时把每一层的Key和Value写成缓存。第二阶段叫decode（解码），每次只生成一个新token，带着新token重新走一遍decoder栈，借助缓存查询历史信息，得到下一个token的概率分布。

这两个阶段不是同一个问题的两种写法，而是数据流形态完全不同的两种计算。prefill阶段，prompt里的几千个token可以同时并行计算，所有token之间相互独立地过一个完整的前向过程，本质是一场大规模矩阵乘法，属于计算密集型。decode阶段，一次只处理一个token，但这次token要跟历史上所有token做注意力交互，核心操作变成从缓存里把K/V读出来，再跟当前的Q做内积，属于访存密集加低算术强度。

所以你会发现一个现象：很多推理框架里，prefill和decode用的是完全不同的kernel。flash attention生态里也专门分出了prefill版本和decode版本，两者的分块策略、线程映射、缓存读取逻辑都不一样。不理解两者的数据流差异，就很难理解为什么不能共用一个kernel。

1.2 先把"坐标"定下来：全文通用的形状约定

讲数据流必须有一个共同的坐标系统，否则后面说"形状变化"就是空对空。我用一组接近主流开源模型的参数作为固定配置，下文所有定量讨论都基于这组值。

这套配置的含义是：隐藏维度被拆成32个独立的头，每个头在128维的子空间里做注意力。模型的所有decoder层共享这套维度设定。后文凡是出现[B,S,H]、[B,N,S,D_h]这类符号，都可以直接套入这组数字估算显存和计算量。比如一张完整的注意力分数矩阵形状是[B,N,S,S]，fp16下就是4×32×2048×2048×2字节，约1GiB。这个数字你最好记住，因为它是prefill显存压力的主要来源之一。

1.3 Prefill与Decode在数据流上的根本差异

把两阶段的核心张量形状摆在一起看，差异是一目了然的。

同样的模型参数，prefill的计算量随S的平方增长，decode的计算量随S线性增长。这也是为什么长上下文场景下prefill延迟暴涨——不是模型变大，而是attention矩阵本身被S平方放大。理解了这一层，再看任何优化方案，无非是两条路：要么掐掉中间大矩阵的落盘，要么改变它的分块方式。后面聊flash attention时会再次回到这个话题。

2. 输入侧数据流：从token序列到QKV张量

2.1 decoder层输入的长相与前置归一化

一条prompt在进入第一个decoder层之前，先要做token embedding。原始的token ID数组形状是[B,S]，每个元素是一个整数索引。查完词表之后，变成一个[B,S,H]的浮点张量。注意，这里有一个经常被新手忽略的点：从这一刻起，一直到模型输出，大部分张量都保持"batch在最前面"的排列。这是现代框架的默认选择，因为对GEMM来说，[B,S,H]可以直接拉平成[B*S,H]，作为一个大矩阵送入cuBLAS，而batch first的布局让这一步不需要任何显式转换。

进入每一个decoder层之后，第一个操作是LayerNorm或RMSNorm。对数据流来说，这一步的形状不变，仍然是[B,S,H]。它在做的事情是沿着每个token的H维做归一化，把4096个元素的分布拉回均值为零、方差受限的状态。

为什么要归一化？因为多头注意力里的内积对输入尺度极其敏感。如果输入向量的模长不稳定，Q和K的点积方差就会大幅波动，softmax之后的结果要么趋近于one-hot要么趋近于均匀分布，注意力就废了。所以norm本质是给后续矩阵运算提供了一个数值上稳定的输入环境。

2.2 QKV投影：一次大矩阵乘法的两种做法

归一化之后，输入被送入QKV投影。这是decoder层里第一处真正的矩阵乘法：输入x与一个权重矩阵相乘，得到Q、K、V三个张量。形状上，x是[B,S,H]，Q/K/V各为[B,S,D]，其中D就是隐藏注意力维度，通常等于H。以我们这套参数，H=D=4096。

实现上存在两条路线。老一点的写法是三个独立的Linear层，每个都是[B,S,H]乘[H,D]，得到三个[B,S,D]。新一点的写法是把三个投影的权重拼成一个[3D,H]的大矩阵，一次GEMM得到[B,S,3D]，再在内存里切成三块。从计算量上看，两者一模一样；但从数据流上看，融合写法有几个实打实的好处：kernel启动次数从三次变成一次，显存分配从三块变成一块，GEMM的规模变大之后GPU利用率往往也更高。

不过融合写法有一个隐藏的坑：切分出来的Q、K、V是对同一块内存的不同切片视图，它们的"连续性"取决于切分方向。如果切分发生在最后一维，也就是[B,S,3D]沿3D方向切成三个[B,S,D]，那么每个切片在内存上依然连续，后续操作可以放心。如果切分发生在batch或序列方向，事情就麻烦了。所以绝大多数框架都选择最后一维拼接权重，这个选择本身就是数据流约束下的自然结果。

2.3 分头操作：view、reshape与permute到底动了什么

现在到了多头注意力最关键的一步：把Q/K/V从"一个4096维的大向量"拆成"32个128维的头"。这一步在数学上很简单，但在数据流上有讲究。

先说形状。Q现在是[B,S,D]=[B,S,4096]。要拆成32个头，每个头128维，先view成[B,S,N,D_h]=[B,S,32,128]，然后permute成[B,N,S,D_h]=[B,32,S,128]。

这里有一个工程上的关键区分。view和reshape基本只改变形状解释，不移动数据，但它们要求原张量内存连续。permute是真正调整维度顺序的操作，它会改变每个维度的步长。permute执行完，数据在内存里还是原来的字节，但逻辑顺序变了，通常是"非连续"的，很多底层算子没法直接处理这种布局，于是需要调用contiguous()把数据真正搬一遍。

所以，在分头这个环节，大多数朴素PyTorch实现里会有一处隐蔽的整张拷贝。这个拷贝有多大？还是用前面的参数，Q/K/V三个张量各是[B,S,D]=4×2048×4096，fp16下每个64MB，三个就是192MB。如果每一层都要为了分头把数据搬到连续布局，几十层下来，这笔开销相当可观。很多自研kernel就是从这里开始做优化的：不在主存里构建[B,N,S,D_h]的连续版本，而是让attention内核直接按索引从[B,S,D]里读取对应的头片段，用"布局感知"替代"搬数据"。

2.4 位置编码的注入位置

现代模型大多数使用RoPE旋转位置编码，它作用于Q和K上。从数据流角度，RoPE注入发生在分头之后、注意力打分之前，针对[B,N,S,D_h]中的最后一维做操作。具体做法是把D_h维两两配对，每个位置根据绝对位置计算旋转角度，对向量做旋转。

一个值得注意的工程点是：旋转角度只跟位置有关，跟batch、head无关，所以实现上可以预计算一个[1,1,S,D_h]的cos/sin表，然后沿batch和head广播。这里的带宽开销通常较小，但如果D_h很大且S很长，RoPE的逐元素操作也会成为一个小热点。量化下来，如果你是手写融合kernel，把RoPE和QKV投影或者attention的读取阶段融合，能省掉一次[B,N,S,D_h]的完整读写。生产级框架里这基本是标配。

3. 注意力打分的流动路径：分数矩阵、掩码与Softmax

3.1 QK^T这一步在数据流上到底做了什么

多头注意力进入核心计算的第一步，是Q和K做内积。形状上是这样：

Q: [B, N, S, D_h] K: [B, N, S, D_h] K^T （最后两维转置）: [B, N, D_h, S] score = Q @ K^T: [B, N, S, S]

从GEMM的角度看，它等价于B×N个独立的[S,D_h]乘[D_h,S]的小矩阵乘法，拼在一起变成[B×N,S,S]。微观上，score[b,h,i,j]是第b个样本、第h个头、第i个查询位置与第j个键位置的向量内积。内积越大，这两个位置的向量方向越一致，模型就认为它们越相关。

这一层的计算量是2×B×N×S×S×D_h。代进前面的参数，大约137 GFLOPs。这个数字在整个attention里占比非常大，而且随着S增长是平方级膨胀。这也是为什么长prompt的prefill那么慢。

但更麻烦的是中间结果本身。完整score矩阵[B,N,S,S]在fp16下约1GiB。请记住，这只是单个decoder层、单个中间张量。一个三四十层的模型，如果每一层都在HBM里落一张这样的矩阵，显存是绝对扛不住的。朴素实现里，这个矩阵还会被后续的mask、softmax逐个读取，实际访存成本还要再乘上好几倍。

3.2 因果掩码：数据流上的"只改一个角"

自动回归模型要求每个位置只能看自己和自己左边的位置，不能偷看未来。这就是causal mask的由来。

数据流上，如果构造完整的掩码张量，它和score同shape，是[B,N,S,S]或者广播成[1,1,S,S]的上三角矩阵。操作方式是在softmax之前，把上三角的合法区域之外的分数加上一个绝对值很大的负数，这样softmax之后这些位置的权重会趋近于零。

一个工程细节值得展开：完整的掩码矩阵也需要1GiB显存，所以成熟的实现几乎不会真的构建它。常见的做法有两种。一种是在写kernel时，对每个query行，只让有效的key索引参与计算，用行索引和列索引的比较替代掩码矩阵。另一种是flash attention里更常见的做法，把掩码信息拆到分块内部，每个block只关心自己能看到的范围，超出范围不加载。理解score的数据流之后，你会发现掩码不一定要以"完整矩阵"的形式存在，它本质上是作用在每个score元素上的一个判断条件。

3.3 Softmax的两次遍历与统计量的流动

score矩阵经过掩码之后，进入softmax。这一步沿着最后一个维度，也就是key方向的S维，对每一行做归一化。

softmax的经典计算流程是：先找出这一行的最大值，把所有元素减去这个最大值，再求指数，再求和，最后逐元素除以这个和。朴素实现在数据流上意味着两到三次完整遍历。同样是1GiB的score矩阵，读一遍、再读一遍，访存成本立刻翻倍。如果直接落盘到fp16并且跟后面的乘法分开计算，这里会成为明显的带宽瓶颈。

这里还涉及一个非常经典的数值技巧。Q和K的维度是D_h=128，两个独立128维向量做内积，结果的方差大约是D_h，也就是说score的数值范围被放大到一个不利于fp16计算的区间。直接在这个数值上做exp，很容易溢出或者精度崩坏。所以标准做法是在QK^T之后除以根号D_h，把方差拉回1附近。这个scale操作在数据流上只是一个逐元素缩放，通常不会单独生成一个大矩阵，而是融合进矩阵乘法的epilogue或者flash attention的scale阶段完成。

3.4 P和V相乘：从"注意"到"提取"

softmax之后的概率矩阵P，形状仍然是[B,N,S,S]。接下来它和V做乘法：P @ V，其中V是[B,N,S,D_h]，得到每个头的输出[B,N,S,D_h]。

这一层在数学上是在做什么呢？对第i个query位置来说，它输出的是一个加权和：把所有位置j的V向量乘上权重P[i,j]，再做累加。也就是说，模型不是从历史位置上"选"一个出来，而是把历史位置的信息按匹配度加权融合到一起。匹配度高的位置贡献大，匹配度低的位置几乎不贡献。

从GEMM视角看，这一层和QK^T的计算量是对等的，同样是2×B×N×S×S×D_h，约137 GFLOPs。两个大矩阵乘法相加，整个注意力核心的计算量就有270多GFLOPs，这还没算后面的输出投影。理解了这层，你就能明白为什么flash attention要把这些操作融合在分块里——每一层独立落盘的代价不只是显存，还有两次完整矩阵写、两次完整矩阵读的带宽开销。

4. 输出侧数据流：多头回归单头、输出投影与残差

4.1 从多头世界回归单头世界

注意力核心计算结束后，每个头都产出了一个局部视图，形状是[B,N,S,D_h]。下一步要把它们合并回一个大的表示，供输出投影使用。

合并的流程是：先执行permute，把[B,N,S,D_h]变换成[B,S,N,D_h]，也就是说把序列维度恢复到第二维；然后执行view或者reshape，把最后两个维度拼起来，变成[B,S,D]，其中D=N×D_h=4096。

为什么要按这个顺序合并？因为输出投影的权重形状是[D,H]。它期望的输入排列是：对于一个固定的序列位置，所有头的输出应该拼成一个连续的4096维向量。如果保持"头的维度在中间"，投影权重需要做额外的索引变换才能正确对应。所以必须先拼头，再做投影。

但这里几乎必然伴随一次内存搬运。permute之后张量是非连续布局，要做连续性修正，就得把整个attn_out拷一份。这个拷贝的大小是4×2048×4096×2字节，64MB。单看一次拷贝还能接受，但每一层都来一次，叠加起来就是可观的延迟。不少优化实现的思路，是把"合并头+输出投影"当成一个大GEMM来规划：先将各个头的输出在逻辑上按正确的布局排列，让分头输出直接在GEMM的B矩阵里充当对应片段，省掉显式转置。

4.2 输出投影：最后一公里GEMM

将合并后的[B,S,D]与输出权重W_o（形状[D,H]）相乘，得到[B,S,H]。这一步的数学意义，是把32个头的独立结论做一次全局混合。前面说过，每个头在注意力计算阶段是彼此隔离的，模型需要让它们有机会互相交换信息，W_o就是那个交换层。

计算量上，这一步是2×B×S×D×H，代进数值大约是275 GFLOPs。注意，这个数字比单个QK^T或P@V都大，因为D和H都是4096，一个token在这个投影里要做的乘加次数等于4096×4096。很多人只盯着attention score的大矩阵，却忽略了输出投影其实也是prefill阶段的重要计算热点。

工程实现上，输出投影通常就是一个普通GEMM，没有太多花活。它值得注意的地方在于内存复用：prefill阶段B和S都很大，输出投影的中间输入[B,S,D]是64MB量级，如果能复用前面释放的缓冲区，而不是每层都重新分配，对控制显存碎片很有帮助。

4.3 残差连接：数据流上最便宜的一个"加"

输出投影之后是残差连接：x = x + attn_out。x是这一层输入时的形状，[B,S,H]。attn_out经过输出投影后也回到[B,S,H]。两者逐元素相加，形状不变。

残差连接在数据流上的成本几乎为零，但意义重大。它让信息在深层网络中有一条稳定的"高速公路"，即便attention的输出出现了某种极端分布，原始输入仍然可以原样保留到下一层。对推理阶段来说，残差完全可以融合进输出投影GEMM的epilogue里，在结果写回显存之前，就把输入加进去，这样既不需要额外的读写，也不需要为一个加法单独启动一个kernel。

紧接着的下一次归一化，也就是第二个norm，通常也会和残差加法融合成一个kernel。从带宽角度看，这样的融合能减少一次中间张量的整读整写。prefill阶段序列很长，这类逐元素kernel虽然单个不贵，但数量多、规模大，累计起来相当可观。能融合的地方尽量融合，这是工程优化里性价比很高的一类改动。

4.4 一个完整的Decoder层形状流转小结

把一个decoder层里多头注意力单元的数据流从头到尾排一遍，就是下面这张表。建议在你调试推理框架时把它贴在旁边，遇到形状对不上或者显存异常，可以直接对着查。

输入 x [B, S, H] norm(x) [B, S, H] QKV投影 x @ W_qkv^T -> [B, S, 3*D] 拆分 Q/K/V [B, S, D] * 3 分头+调序 q/k/v -> [B, N, S, D_h] QK^T q @ k^T -> [B, N, S, S] scale + causal mask [B, N, S, S] softmax [B, N, S, S] P @ V [B, N, S, D_h] 合并头 permute+view -> [B, S, D] 输出投影 attn_out @ W_o^T -> [B, S, H] 残差相加 x + attn_out -> [B, S, H]

这张表每一条都值得在显存profile里找到对应的算子。如果某个阶段多了一次意外的形状转换，或者某个中间张量比预期大了一倍，顺着这张表去查，定位速度会快很多。

5. KV Cache的写入与跨阶段交接：Prefill最后的数据归宿

5.1 为什么prefill阶段就要把KV写进缓存

刚才的注意力数据流里，K和V在attention内部被算出来之后，不止用于当前的注意力计算，还有一个去向：写入KV Cache。

原因在于decode阶段。当模型生成第100个token时，它需要跟前面99个历史token做注意力，第101个token时需要跟100个历史token做注意力。如果每次都从第一层重新计算所有历史token的K/V，计算量会随着生成长度线性爆炸。所以推理框架会选择在prefill阶段就把整段prompt的K/V都算好、存好，后续decode阶段直接读取，只计算新token自己的K/V并追加写入。

从数据流上看，KV Cache的写入点和K/V张量被生成的点是重合的。也就是说，在投影得到K/V之后、参与注意力计算之前，或者跟注意力计算同步，就可以把K/V的副本写入cache区域。现代kernel更倾向于让cache写入和注意力计算融合，在读取K/V时直接向cache对应槽位写一份，避免再单独做一次显存拷贝。

5.2 KV Cache的形状、步长与内存布局

Cache布局需要考虑和attention kernel的访存模式匹配。常见的做法是：

K Cache: [B, N, max_seq_len, D_h] V Cache: [B, N, max_seq_len, D_h]

prefill阶段写的是这个cache的前S个槽位。以我们这组参数估算，单层KV Cache的大小是4（batch）×32（头）×8192（预分配的最大长度）×128（头维度）×2字节×2（K和V两份），大约是256MiB。如果你的模型有32层，单是这个KV Cache就占8GiB以上。这还没算prefill阶段临时中间激活的开销。

布局选择的权衡很有意思。如果采用[B,N,L,D_h]，那么同一个位置的D_h维是连续的，适合按tile加载；如果改用[B,N,D_h,L]，那么长度维度连续，对某些按query行扫描的decode kernel更友好。此外，生产级的continuous batching场景会把cache切成分页块（block），每个请求的KV存到若干不连续的物理块上。prefill阶段一次性申请足够多的块，decode阶段随着生成逐步追加。从纯数据流的角度，逻辑上的连续长度S被映射成了物理上的多个块，attention kernel需要维护一个block table才能正确寻址。

5.3 Prefill与Decode的交接细节

prefill结束后，模型持有的核心状态是每一层的KV Cache，以及当前已生成的序列内容。decode阶段每次生成一个token，都需要执行这样的数据流：

新token经过embedding变成[B,1,H]，进入第一个decoder层。norm之后做QKV投影，Q形态是[B,N,1,D_h]，K和V也是[B,N,1,D_h]。注意力计算时，读取cache里前S个位置以及自己这个新位置，构成完整的K/V序列[B,N,S+1,D_h]。score矩阵此时是[B,N,1,S+1]，远没有prefill阶段那么巨大。

这里有一个值得留意的工程点：如果cache保存的是RoPE之后的K，那么新写入cache的K也必须做RoPE，保证新旧位置处于同一坐标空间。如果cache保存的是RoPE之前的K，那么decode阶段读取K后还要在attention kernel内部重复做一次RoPE，这会增加解码延迟但换来更灵活的长度扩展能力。两种方案各有取舍，从数据流的角度看，它实际是"是否在cache写入时就固化位置信息"的一个选择。

此外，prefill和decode的显存分配策略天然不一样：prefill是一次性给整段prompt的cache填满，decode是一个token一个token地追加。如果cache的block大小不够对齐，可能产生大量碎片。比较好的实践是，在做prefill之前就根据当前prompt长度预分配足够的block，并且预留decode阶段可能的增长空间，避免中途频繁搬迁。

6. 从数据流视角看工程优化：几个值得下手的点

6.1 形状变换的成本比你想象得高

petite hidden gem是"形状变换看着不产生计算量，但会产生数据移动"。permute本身不拷贝，但后续一旦需要连续布局，contiguous()就是一次完整的读和写。前面算过，attention输出[B,S,D]每个64MB，split、permute、merge这一套组合拳打下来，单层可能有几百MB的数据搬移。

我在实际优化中见过非常典型的情况：profile里看到很大的时间占比落在permute或者.contiguous()上，而不是矩阵乘法。这种情况通常不是算子本身慢，而是上游没有保持一个让GEMM舒服的布局。改进思路有两个。一是尽量把形状变换放在小张量上，比如先从[B,S,3D]里切出Q/K/V，再做分头；二是尽量在自定义kernel里用索引访问替代显式变换，让GEMM自己处理非连续视图，这需要你对cuBLAS的布局支持和内存步长比较熟悉。

6.2 Flash Attention从数据流角度做了什么

flash attention在prefill阶段的核心，就是不再构建完整的[B,N,S,S]分数矩阵。它的做法是把Q、K、V按块读入SRAM，在片上完成分块的score计算、mask、softmax和P@V，通过online softmax增量更新统计量，最终只把每个位置的输出写回HBM。

这个重排带来的收益是决定性的。显存占用从1GiB级别的分数矩阵降到了几十到几百KB级的片上缓冲；HBM带宽从"写score再反复读score"变成了"只读Q/K/V各一遍"。在很多prefill实测中，切换到flash attention之后，单层attention的HBM流量能下降一个数量级，延迟自然也跟着降。

理解这一点之后，你就不会把flash attention当成一个黑盒。它的tile大小设计、每个block内如何规避掩码、如何在SRAM里保留partially updated softmax统计量，本质上都是围绕"让中间矩阵不落HBM"这一条数据流原则展开的。prefill因为S已知，可以做更激进的分块规划；S特别长时，还需要二次分块，避免单个score分块都超出SRAM。

6.3 显存规划与并行量的估算

最后分享一个我在实践中反复用的思路：拿到数据流形状之后，先在纸上估算每个关键张量的显存和带宽，不要先盲目调参。prefill阶段最核心的起伏变量是batch B和prompt长度S。

以[B,N,S,S]的score矩阵为例，B翻倍，显存翻倍；S翻倍，该张量显存翻四倍。当S很大时，减少batch、降低并发量往往比增加batch得到更高的有效吞吐，因为过大的中间矩阵会导致显存溢出和频繁换入换出。相反，S较短时，加大batch能让GEMM规模更大，GPU利用率更高。这种取舍不是玄学，只要把数据流形状列出来，代入显存容量反推一个合理的batch上限，很快就能得出结论。

实际项目里，我通常会把所有关键中间张量的shape、stride、dtype和生命周期全部整理进一份文档，等profile数据出来之后，逐个对照。凡是和形状预期不一致的算子，基本就是优化入口或者隐藏bug所在。这个方法帮我处理过不少诡异的性能问题，也让我在给同事讲解attention内核时，不用一遍遍重复"这里有一张隐形的中间矩阵"。

6.4 一些值得保留的工程习惯

写到最后，分享几个从数据流视角沉淀下来的工程习惯。第一个，在改动attention代码之前，把这条flow从头到尾写一遍，并且标出每个节点的张量形状和内存量。很多设计决策，比如是否融合QKV、是否显式做contiguous、cache用哪种布局，都会在这个环节暴露出来。第二个，遇到prefill阶段显存爆炸，优先怀疑[B,N,S,S]相关的中间张量，它几乎永远是最大的"意外消费者"。第三个，如果是GQA或MQA架构，不必另起炉灶，把上面思路里的N换成KV头的数量即可：Q仍然是所有头，K/V只有少量共享头，数据流整体框架完全一致，只是KV Cache的写放大变小了。

我个人在多次排查和调优之后最大的感受是：attention的公式是透明的，但对工程真正起决定作用的是形状、步长、连续性和生命周期这些"数据流的边角料"。把这些边角料理顺，很多性能问题甚至不用调kernel就能消除大半。希望这篇拆解能帮你把这条链路彻底看穿。