企业官网建设流程全解析

做推理引擎优化的同学，99%的时间都盯在Prefill和Decode这两个阶段上。很多人一上手就急着折腾decode阶段的KV Cache量化、投机采样，但真正遇到线上事故时才发现，往往是一段超长Prompt进来，整个Prefill阶段把GPU卡了十几秒，后面的请求全部排队。这个阶段里decoder层的多头注意力单元，数据流的形态和decode阶段完全不一样——它是整段Prompt一次性灌进来的，Q、K、V一次性算完，注意力矩阵会达到S×S的量级，访存模式和计算密度都跟decode有着本质区别。把这条数据流吃透了，再去调Batch策略、去选FlashAttention的切块方式、去做PD分离，才算是踩在了点子上。

本文就顺着Prefill阶段的decoder层多头注意力单元，把从输入张量到最终输出之间的每一步数据流动、每一个中间tensor的形状变化和内存去向，完整拆一遍。同时会讲讲为什么prefill是计算密集而decode是访存密集，以及长Prompt场景下真正值得优化的关键路径在哪里。

1. Prefill阶段为什么值得单独研究

1.1 Prefill阶段在整个推理链路中的位置

一个标准的大模型推理请求，从用户输入到最终吐出完整回答，会经历两个截然不同的阶段：Prefill（预填充）和Decode（逐token生成）。Prefill负责把用户输入的Prompt完整地过一遍模型，得到每个token对应的KV Cache，以及第一个输出token的logits；Decode阶段则拿着这些KV Cache，逐token地往后生成，每生成一个token就追加一组K和V到Cache里。

以GPT这类decoder-only架构为例，输入一个长度为S的Prompt，Prefill阶段在每一个decoder层里，都要对S个token做完整的自注意力计算。注意这里的关键点：S个token是同时进入模型的，不存在decode阶段那种"必须等前一个token算完才能算下一个"的依赖关系。因果掩码处理的是每个query只能看到当前位置及之前的位置，但这只是数值上的掩码约束，并不构成计算上的串行依赖——因为整个S×S的注意力矩阵可以一次性并行算出来，只需要把上三角位置填上负无穷即可。

正是这个"一次性并行"的特性，决定了Prefill阶段的计算特征：它是计算密集的，GPU的算力能打得很满；而decode阶段每个step只处理1个token，变成了访存密集——每算一个token都要把整个权重和KV Cache读一遍，算力利用率反而上不去。所以如果只盯着decode阶段优化，而prefill阶段卡在S×S的注意力矩阵上，整个系统的P95延迟和吞吐上限都会被它拖死。

1.2 Prefill和Decode在数据流上的本质差异

我习惯把这俩阶段的数据流差异归纳成一张表，做调度分析时可以直接拿出来对照：

这张表对着看，就能理解为什么Prefill阶段要把注意力矩阵的中间结果管理得极其谨慎：如果老老实实把[B, H, S, S]的score矩阵落到显存里，S=4096、H=32、B=1时，仅这一个张量就是32×4096×4096×4字节≈2GB，纯属于"算完就扔"的临时产物。这是FlashAttention这类kernel在prefill阶段显得尤为重要的直接原因。

1.3 Prefill的数据流对系统指标的影响面

Prefill阶段的数据流设计不只是影响单次请求的时延，它对整个推理系统的稳定性影响远比表面上看起来大。由于Prefill阶段S个token并行计算，单个请求对显存和算力的需求是突发式的——一段4096长度的prompt，其计算量相当于decode阶段4096个step的总和，但全部压缩在几秒内完成。

这会造成一个经典的调度难题：在线推理服务同时遇到一个长Prompt请求和一堆短的decode请求时，如果prefill不切割，它就会长时间霸占整张GPU，decode阶段的请求全部排队，延迟直接起飞。这也是Chunked Prefill（PD分离）技术的核心动机——把长Prompt的prefill切成小块，和decode请求交错执行，让算力利用率和延迟曲线平滑下来。

理解了Prefill在系统层面的地位，接下来就该钻到decoder层的多头注意力单元内部，看看数据流具体是怎么一步步走的。

2. 多头注意力单元的数据流追踪：从输入到输出

2.1 输入张量与初始化状态

把一个token序列送入decoder层时，数据流的第一站是输入特征张量。假设batch size为B，Prompt长度为S，隐藏维度为D，那么这一层的输入张量X的形状是[B, S, D]。对于GPT这类模型，通常在进入Attention子层之前会先过一个RMSNorm或LayerNorm，处理后的特征张量形状不变，仍然是[B, S, D]。

在未启用GQA/MQA的标准MHA（Multi-Head Attention）配置下，这一步的数据是完整D维的，可以为后续的QKV投影提供足够的特征信息。这里有一个值得留意的工程细节：很多推理框架在prefill阶段会把输入张量先重新排列成"连续的内存块"，把[S, D]每一行尽量对齐，避免后续多头切分时出现大量非连续内存访问。别看这个细节不起眼，S特别长时，非连续访问会让Attention Kernel的访存效率掉一截。

2.2 通过QKV投影生成Query、Key、Value

接下来是数据流的核心环节：QKV投影。输入X经过一个拼接权重矩阵W_qkv，一次性生成三组向量：

• 数学上就是X @ W_qkv，其中W_qkv的形状是[D, 3D]
• 输出的形状是[B, S, 3D]，需要拆成三份[B, S, D]
• 实际工程中大部分框架会一次投影、一次切分，而不是分开做三次matmul

我自己在理解这个环节时最喜欢用的类比是"复印室"：输入X是借阅登记本，每个token是登记本上的一行记录，QKV投影像三台复印机，把同一行记录按不同规则复印出三份，分别用于"我要查什么"（Q）、"我有什么可被查"（K）、"我实际给出了什么内容"（V）。

在数据流的视角下，这一步的关键在于：prefill阶段一次性把S个token的Q、K、V全部算出来了，而decode阶段每步只需要为一个新token算一组Q、K、V。所以prefill阶段的计算量天然是S倍的decode单步，但好处是matmul的batch维度和序列维度都很大，GPU的Tensor Core可以打得非常满。

2.3 多头切分的形状变化与内存布局

QKV投影完成后，数据流进入多头切分阶段。这一步不是计算，而是纯视图变换，学术上叫reshape。以标准MHA为例，把[B, S, D]的Q拆成[B, S, H, d_k]，其中H是注意力头数，d_k = D / H是每个头的向量维度，然后通常再转置成[B, H, S, d_k]，让每个头的数据在内存上保持连续。

转置这个动作背后有个很大的坑：如果直接做transpose再参与后续matmul，内存布局可能非连续，矩阵乘法的效率会大打折扣。所以工程上常见做法是让Kernel直接基于stride处理，或者在切分时就用torch.view/reshape配合contiguous把内存布局理顺。很多新手在复现MHA时，盯着shape看半天没毛病，一跑性能就拉胯，大概率就是栽在这里。

多头切分后的内存结构是这样：

上面这行代码是标准MHA的中间状态。《attention is all you need》原文里，每个头的计算是：head_i = Attention(Q_i, K_i, V_i) = softmax(Q_i·K_i^T / sqrt(d_k)) · V_i而所有头最终是拼在一起，过一次输出投影W_o，shape重新回到[B, S, D]。

2.4 注意力分数的计算、掩码与Softmax

多头切分完毕后，进入注意力单元最核心的数学计算环节。对于每一个头，Q矩阵与K矩阵做点积并缩放：

• 公式：score_i = Q_i @ K_i^T / sqrt(d_k)
• 结果：形状[B, H, S, S]的注意力分数矩阵
• 缩放因子sqrt(d_k)是为了防止维度增大后点积数值过大，把softmax推入饱和区

这步在标准实现中会产生一个S×S的大矩阵，这在decode阶段问题不大（每步只算1×S），但在prefill阶段会是显存杀手。所以绝大多数推理框架在prefill阶段不会真的实例化这个完整矩阵，而是用FlashAttention分块计算，把score矩阵切成小块，算完一块的softmax就去更新输出，不落全局显存。

分数矩阵算完后要做两件事：因果掩码和softmax归一化。因果掩码就是把上三角位置全部置为负无穷（工程上常用-inf或一个极大负数），确保第i个token只attend到0到i的位置，这是decoder自回归性质的要求。然后对每一行做softmax，得到归一化的注意力权重矩阵。注意softmax是沿S维度（也就是key侧）做归一化，语义是"当前query对前面所有key的注意力权重之和为1"。

在FlashAttention的实现中，这里还有个精妙的细节：为了不保存完整score矩阵，softmax会被改造成online softmax——分块处理时维护每行已见过的最大值m和归一化常数l，每次来一个块就增量更新。这个技巧是FlashAttention能够在prefill阶段省显存的关键数学基础，后面优化部分会再展开。

2.5 加权求和、多头合并与输出投影

softmax归一化后的注意力权重矩阵与V矩阵做加权求和，得到每个头对应的上下文向量：

• 公式：context_i = attention_weight_i @ V_i
• 结果：形状[B, H, S, d_k]，这就像一个"聚合后的特征视图"

接着进入多头合并阶段：将[B, H, S, d_k]转置回[B, S, H, d_k]，再reshape成[B, S, D]，把H个头的输出拼接起来。这一步不需要额外计算，只是内存布局调整，但同样要注意contiguous问题，否则后续输出投影的matmul性能会受影响。

最后再过一次输出投影权重W_o，形状[D, D]，输出仍为[B, S, D]。到这里，一个decoder层内多头注意力单元的数据流就走完了。后续还有一个残差连接——把注意力输出加上原始输入X，再过FFN和第二个Norm，才真正完成整个decoder层的处理。这条数据流放在prefill阶段，最鲜明的特征就是每一步的张量都具有S这个完整维度，任何中间结果如果被完整实例化，随S增长都是平方级爆显存。

3. Prefill阶段优化关键路径：让数据流不再拥堵

3.1 FlashAttention分块：把S×S的显存炸弹拆掉

前面提到，标准MHA实现里score矩阵的形状是[B, H, S, S]，S以千为单位时这基本就是显存灾难。FlashAttention解决这个问题的思路非常直白：不计算和保存完整的attention矩阵，而是将Q、K、V按块遍历。

以其中一个注意力头为例，把Q按行切成若干块，每块大小是B_r个query；K、V按列切成大小B_c的块。对每个Q块：

• 遍历所有KV块，计算局部Q_block @ K_block^T / sqrt(d_k)
• 对局部分数块做掩码、online softmax（记录当前行最大值和归一化常数）
• 用局部softmax权重对V_block加权，增量更新输出块

整个过程只用维护一个小尺寸的分数块（B_r × B_c），输出的output块也是增量写入。这个设计和prefill阶段是天然契合的：S越大，分块省下的显存越可观；同时因为块内部全是标准矩阵乘法，GPU的Tensor Core利用率很高。这也是为什么现在的推理框架在prefill阶段无一例外都用FlashAttention或类似kernel。

对于特别长的Prompt，还有一个进阶方案：把FlashAttention的Q、K、V切块跟序列并行结合起来。既然Q的每一块都可以独立算注意力，那完全可以把S维度切到多个设备上，各自算局部注意力，最后把KV部分做一次all-gather或all-to-all。这个分布式的数据流设计在后文的"张量并行下的通信"里会展开讲。

3.2 KV Cache写入：prefill阶段的数据沉淀

在注意力计算的同时，prefill阶段还有一个关键的"副产物"要落盘：KV Cache。正是这一段Cache，让decode阶段可以不必重新计算历史token的K和V。

具体到数据流，prefill阶段需要把每一层、每一个注意力头的K和V保存下来。以标准MHA为例，KV Cache的总shape为[2, num_layers, B, H, S, d_k]。S=4096、H=32、d_k=128、层数L=32时，单个请求KV Cache仅一层的K就是32×4096×128×2字节≈32MB，整个模型约1GB。虽然prefill只写一次，写法的连续性很好，但KV Cache的内容直接决定了后续decode的访存压力，所以现在大家都在做GQA/MQA来削减KV Cache体积。

有个工程上的重要细节：KV Cache通常使用预分配的内存池，prefill阶段写入时是一次性连续写入，速度快且碎片少；但如果用torch.cat动态追加，长序列场景下会有大量内存搬移和碎片。所以在实际部署时，KV Cache往往采用类似PagedAttention的块式管理，把连续的逻辑Cache映射到不连续的物理页上，prefill阶段逐个块填充，decode阶段逐页追加。

3.3 Chunked Prefill与PD分离：平滑数据流的冲击

前面讲系统调度时提过，长Prompt的prefill是突发式计算。为了不让这个突发拖垮decode，现在的主流方案是Chunked Prefill——把prefill阶段的长序列以块为单位切割，每个块是一个独立的注意力计算单元，多个块之间共享KV Cache但保持因果掩码范围。

具体的数据流变化是这样的：假设Prompt长度S=8192，设置chunk size为512，那么模型不把8192个token一次性全部送入，而是按512个token一批逐批前向计算。每一批算完，KV Cache就多写512行；同时这一批的最后一个token位置可以直接输出logits，用来做decode的第一个token。这样prefill和decode就能在同一批次里交错执行，GPU算力被匀速消耗，而不是被一个长请求脉冲式打满。

这个方案在工程实现上的要点是：每个chunk在计算注意力时，score矩阵只有chunk_size × (之前的KV总长)那么大，不会出现8192×8192的爆显存问题；同时因果mask只需要约束当前chunk内的可见范围，历史KV直接拼接在K/V前面即可。整套数据流像是"滚动窗口"，兼顾了显存控制和算力平滑。

3.4 GQA/MQA下的多头数据流变化

既然提到KV Cache和注意力计算，就不能不说GQA（Grouped Query Attention）和MQA（Multi-Query Attention）。在标准MHA下，每个Q head有一个对应的K和V head，KV Cache大小跟Q头数H成正比。GQA把KV head数减少到G（G远小于H），每个KV head被一组H/G个Q head共享；MQA更极端，所有Q head共享同一份K和V。

这个改变对prefill阶段数据流的影响有三点：

• Q的投影矩阵变宽：W_q[出D]维不变，还是每个Q head各算各的
• K和V的投影矩阵变窄：只需要投影G份，而不是H份，prefill时的matmul计算量明显下降
• Attention分数阶段存在隐式广播：每个KV head被多个Q head复用，数据从"每家一份"变成"分时共享"

这组数据流改造，换来的是decode阶段KV Cache访存量的下降，代价是Q head对KV的读取逻辑会变成group内多次复用。硬件上self-attention score的计算仍然按Q head为单位，所以GQA在做FlashAttention时，通常会先把KV head在内部广播到对应Q group的块里，让kernel层面的数据流保持连续。

4. 工程实测：瓶颈定位与排查实录

4.1 长Prompt触发Prefill耗时陡增

我遇到过最典型的线上问题是：一个几千token的检索式Prompt进来，prefill阶段耗时从几百毫秒直接飙到十几秒，后面所有decode请求延迟全部雪崩。排查时第一件事是打点，把prefill阶段在每个decoder层的耗时切出来看，发现耗时集中在Attention Kernel上，FFN占比反而正常。

顺着数据流定位，问题非常典型：序列长度S已经超出了FlashAttention单块处理的最优尺寸，score矩阵切块后的kernel启动开销变大，同时长序列下L2 Cache命中率下降，数据反复从HBM搬运。这种情况下最有效的两个手段：一是把FlashAttention的块大小（B_r和B_c）调大，让每个block多算一些，减少kernel启动次数；二是直接上Chunked Prefill，把一次长prefill拆成多个短prefill，让每一段数据流的长度都落在硬件最舒服的区间，同时也避免了一次请求把整机的计算资源全部吸干。

4.2 Prefill阶段显存峰值的两个隐蔽来源

除了score矩阵这个显存杀手之外，prefill阶段的显存峰值还有两个非常隐蔽的来源。第一个是第一层FFN的激活值。有些优化只在attention部分做FlashAttention，但FFN部分的中间激活还是完整保存下来，S很长时[B, S, 4D]的中间激活体积非常可观。解决方法是把FFN也用类似的分块策略处理，或者干脆用激活重计算，以额外计算换显存。

第二个隐蔽来源是多头切分时的转置副本。如果没有做好内存布局，transpose可能触发一次完整的数据拷贝，[B, H, S, d_k]这个tensor是S量级的大块头，多拷贝几次显存直接翻倍。这块排查起来最坑，看显存分析工具只能看到一个大tensor，很难意识到是布局拷贝导致的。一个实用的检查手段是看Kernel的耗时占比与理论FLOPs的比值，如果matmul耗时远高于理论值，优先怀疑非连续内存访问。

4.3 张量并行场景下的Prefill通信瓶颈

模型大了以后必然要上张量并行（TP），把多头注意力拆到多张卡上。标准做法是Attention的Q、K、V投影各自在自己卡上做，但计算attention之前需要把K、V进行all-gather，让每张卡上拿到完整的KV；算完attention后，输出还要过一次all-reduce，保证下一层的输入在每张卡上一致。

在prefill阶段，这个通信组件的开销会被S放大：因为KV的序列维度就是S，all-gather的数据量直接跟S成正比。S越长，通信耗时占比越高，有时能达到整体prefill耗时的30%以上。实测下来有效的优化就这么几招：

• 让all-gather跟当前分片的attention计算重叠起来，先算自己手上已有的KV部分，边算边收别的卡
• 在QKV投影之前先做一次序列维度切分，让每张卡先算一部分token的Q，再交换KV，减少单次通信量
• 如果用了GQA，通信的压力会明显下降，因为要交换的KV head变少了

把TP的通信数据流理清楚后，那种"明明加卡但是prefill更慢"的诡异现象就好解释了——加卡后通信开销的增长压过了计算分片带来的收益，尤其S不足够长时，甚至会出现负优化。

4.4 Prefill阶段常见问题速查表

以下是一份基于实际运维沉淀的排查速查表，可以贴在工位边上看：

4.5 实测中的两个额外心得

最后分享两个别人不太会写进文档里的经验。

第一个是关于chunk size的选择。Chunked Prefill的chunk size不是越小越好——太小了每个chunk的attention计算矩阵都偏瘦，GPU算力打不满；太大了又起不到平滑调度和降低显存峰值的作用。实测下来，在A100/A800上，chunk大小在256到1024之间表现比较均衡。判断标准就是看GPU的MFU：如果整个prefill阶段的算力利用率掉到40%以下，大概率是chunk太小了。

第二个是关于prefill阶段要不要降低精度。现在很多框架支持FP8/BF16混合精度跑prefill。我的建议是score矩阵和softmax阶段保持BF16精度，QKV投影和输出投影可以放心用FP8，因为这两个环节本身就是对噪声容忍度较高的线性变换。但softmax之前如果已经用了FP8，长序列时数值误差容易累积，表现为生成质量波动。这条是踩了几次坑才总结出来的，建议做精度实验时重点盯attention score的数值范围。

从我个人的实际经验看，Prefill阶段的数据流是最适合"用手推一遍shape"的优化训练场。把每个中间tensor从哪里来、到哪里去、有没有被完整实例化、有没有触发非连续访问都过一遍，基本就能把推理性能的底摸透八成。后续无论你要上PD分离、切分布式并行还是压缩KV Cache，都是在为这条数据流上的某个环节做减负。把这套基本功练扎实了，再激进的技术方案都不怕跑偏。