一个 8K prompt 进入模型时，Attention 要处理大量 query 与 key 的配对。生成下一个 token 时，query 缩到一行，历史 K/V 却仍然接近 8K 行。前者要让数据尽量在片上复用，后者要想办法让更多计算单元一起读历史。这两个场景可以共用在线 Softmax 的数学，却往往需要不同的任务划分。
1. 从 Transformer 到 IO 问题
以 decoder-only Transformer 的一个 head 为例。设 query 长度为 \(L_q\)，KV 长度为 \(L_k\)，Q、K、V 的 head 维度均为 \(d\)。Causal mask 对未来位置取负无穷，对可见位置取零。Vaswani 等，NeurIPS 2017
按一次乘加计 2 FLOPs，两个矩阵乘在全矩形上合计约 \(4L_qL_kd\) FLOPs，另有 mask、指数和归约。一次处理长为 \(L\) 的完整 prompt 时，因果三角形只有 \(L(L+1)/2\) 个有效位置。能否跳过被遮挡的计算，取决于内核是否跳过全遮挡 tile，以及如何处理对角 tile。
朴素实现把分数和概率写入全局内存，随后由下一个算子读回。以 8,192 token、32 个 Q head 为例，即使每个元素只占两字节，一份完整矩形中间张量也占 4 GiB。两个矩阵乘之间的访存开销，足以抵消把 GEMM 单独做快的收益。
图 1，原创示意。Prefill 沿 Q 分配行块，每个行块遍历自己的可见 KV。Decode 沿 KV 切分，同一行 query 由多个任务协作。
2. MHA、MQA、GQA 与 MLA 改了哪一维
MHA 为每个 Q head 单独投影 K、V。MQA 让所有 Q head 共用一组 K、V，GQA 则让每组 Q head 共用一组。设组大小为 \(g=H_q/H_{kv}\)，第 \(h\) 个 Q head 访问第 \(\lfloor h/g\rfloor\) 个 KV head。Q 仍然因 head 而异，Softmax 和输出也仍然逐 head 计算。32 个 Q head、8 个 KV head 时，GQA 把 K/V 投影和缓存宽度降到同维度 MHA 的四分之一，QK 和 PV 的配对数依旧由 32 个 Q head 决定。MQA GQA
在 Decode 里，把同组的四个 Q 排成四行，就可以用一个 \([4,128]\times[128,T]\) 消费同一份 K。Prefill 则能把多个 query 位置也放进矩阵的行维，利用率更容易做高。把共享 KV 展开成 32 份会失去带宽复用。反过来，MQA 只有一个 KV group，也会改变可用于分核和张量并行的轴。
MLA 的矩阵吸收
以 DeepSeek-V2 的维度为例，令 \(H=128,d_h=128,d_c=512,d_r=64\)，用列向量表示。将历史 token 压到 \(c_j=W_{DKV}x_j\)，内容 key、value 为 \(W_{UK,h}c_j,W_{UV,h}c_j\)。独立的 RoPE 路径生成跨 head 共享的 \(k_j^R\)。对 query 的内容和位置分量 \(q_{i,h}^C,q_{i,h}^R\)，吸收前后可写成
两项分数先相加，再统一缩放和 Softmax。缩放仍取模型原始的 head 维度。历史位置的旋转随 \(j\) 变化，所以位置分支单独保留，内容上投影才可移到 query 一侧。BF16 的持久缓存每层每 token 是 \((512+64)\times2=1,152\) B。DeepSeek-V2，§2.1
图 2，原创。两条路径读取同一种持久缓存，决定性区别在于历史上投影是否进入本步计算。
下表用 \((M,N,K)\) 标记 GEMM，按请求计，除特别标注外均为单 head。设首次 Prefill 长度 4,096，Decode 可见长度 8,192。
| 步骤 | Prefill，展开内容 K/V | Decode，吸收到 query 与输出 |
|---|---|---|
| KV 降投影，所有 head 共享 | \((4096,512,d_{\rm model})\) | \((1,512,d_{\rm model})\)，只算新增位置 |
| Q 降投影，再上投影，合计所有 head | \((4096,d_q,d_{\rm model})\)，再 \((4096,24576,d_q)\) | 对应的 \(M=1\) |
| K/V 内容上投影，合计所有 head | 各 \((4096,16384,512)\) | K 移至 Q 侧 \((1,512,128)\)，V 移至汇总后 \((1,128,512)\) |
| 共享 RoPE key 投影 | \((4096,64,d_{\rm model})\)，再旋转 | \((1,64,d_{\rm model})\)，只追加当前 key |
| 内容与位置 QK | 可拼接为 \((4096,4096,192)\)，按 causal tile 执行 | 内容 \((1,8192,512)\)，位置 \((1,8192,64)\)，可融合为一条 576 宽的逻辑规约 |
| PV | \((4096,128,4096)\) | \((1,512,8192)\)，随后还原输出维度 |
Prefill 展开一次后能让大量 query 重用内容 K/V。Decode 只增加一行 Q，避免展开整个历史更有价值，但每个 Q head 的扫描规约从 128 增到 512，PV 的输出维也增大。假设 MAC 阵列的 K tile 宽 256，512 可以整分为两段，128 和 64 在简单 padding 映射下分别占半个和四分之一个 tile。实际利用率还受小 \(M\)、独立任务能否打包，以及矩阵布局和精度模式影响。MLA 要把多个 head 对同一 latent tile 的访问安排在一起，才能将容量收益变成带宽收益。选择 GQA 的模型可以保留较简单的投影、并行切分和成熟的部署路径，架构取舍需要对比质量和端到端吞吐。DeepSeek-V2
一个容易忽略的对照是，BF16 下 8 个 128 维 KV head 的 GQA 每 token 每层用 4,096 B，上述 MLA 用 1,152 B。如果 GQA 收到 2 个 KV head，就只有 1,024 B。此时应对照模型质量、TP 复制、实际读量和 GEMM 形状。模型实例也已经分化。Kimi Linear 按层混合 KDA 与全局 MLA，GLM-5 则用 DSA 缩小长上下文主 Attention 的访问集合。Kimi Linear GLM-5
YOCO 在层间共享
YOCO 的 self-decoder 先用滑窗或递推类机制编码序列，随后 cross-decoder 的多层 cross-attention 读同一份全局 K/V。普通逐层缓存是每层产生并保留自己的全局 K/V。YOCO 把这份全局历史的生成移到两个 decoder 之间，各 cross 层仍有各自的 query。Prefill 可以先完成整段 self-decoder 的编码，再由 cross-decoder 处理得到输出所需的位置。Decode 则每步更新 self-decoder 状态并追加共享 K/V，多个 cross 层仍会消费历史。层间共享首先节省驻留容量，跨层的 HBM 重读取决于执行和片上复用。YOCO
3. 稀疏 Attention 的成本在选、取、算三个阶段
滑窗先用位置规则限定每个 query 的候选，Prefill 从二次配对降到 \(O(LW)\)，Decode 读取范围稳定。块稀疏将选择单位放大到 tile，牺牲一点选择精度以换取连续搬运与规则 GEMM。逐 token Top-K 的粒度最细。如果先生成完整 QK 再取 Top-K，已经支付了主 QK 的计算量，只能省后续部分工作。Longformer
NSA 同时维护压缩、选择和滑窗三个分支。压缩分支从块摘要获得全局信号，选择分支读取一批完整的历史块，滑窗提供局部细节，各分支通过门控合成。按块选择可以用连续 DMA 送给矩阵阵列，多个共享同一 KV group 的 query head 可以复用选中的块。代价是压缩分支、滑窗分支和分支合并仍有计算成本。训练和 Prefill 要处理随 query 变化的块选择，Decode 则要让索引与历史追加同时更新。Native Sparse Attention
DSA 用独立的 Lightning Indexer 给可见历史打分。按 DeepSeek-V3.2 的形式，indexer 对不同头的点积做 ReLU，再用 query 相关权重合并，然后从候选里选 Top-K。
主 MLA 只对 \(\mathcal J_i\) 计算 QK、Softmax 和 PV。若 8,192 个历史位置中选 2,048 个，主 Attention 配对数是全读的四分之一，indexer 仍扫 8,192 个位置，也需要保留独立的 key cache。Prefill 若逐 query 扫全部因果候选，indexer 配对数仍随长度二次增长，省下的是昂贵主路径的配对。Decode 应分别计量 indexer 的矩阵计算、Top-K 归约、按索引访问页池，以及主路径的 QK/PV。DeepSeek-V3.2 vLLM 的 DSA 实现
这里两个 \(b\) 都是每位置字节数，公式假设无重复读和额外对齐。若 token 分布在大量不相邻的页，真实 DMA 流量会更高。按物理页整理已选索引可以改善合并读取，前提是索引、K/V 与位置语义始终对应。矩阵内核面对的是收集后的稠密小 tile，跳过哪些 tile 主要靠运行时调度。阵列的固定结构化稀疏能力无法直接吃任意 token 索引。
图 3，原创。索引路径扫描整个可见前缀，主路径按选择结果取 latent 页。
4. 线性、混合与长上下文
Linear Transformer 把相似度替换成可分解特征映射。设 \(\phi(k)\) 为特征向量，因果递推累积 \(A_i=A_{i-1}+\phi(k_i)v_i^{\mathsf T}\)、\(z_i=z_{i-1}+\phi(k_i)\)。输出为
Decode 每步更新固定形状的状态，因而不必扫描完整历史；Prefill 可以做分块前缀归约。状态是特征维乘 value 维的矩阵，批量少时反复读写状态和执行外积仍会压低阵列利用率。特征映射改变了模型的归一化核，模型质量需随任务验证。Linear Transformer
RetNet 引入带衰减的 retention，提供并行、递推及 chunkwise 形式。Lightning Attention-2 用块内矩阵计算加块间状态传递处理带衰减的线性 Attention。对 NPU，块内可组织为 Cube GEMM，跨块状态更新、衰减和精度累积要与搬运重叠。块太小容易被同步开销主导，太大则提高片上临时存储和块内工作量。RetNet Lightning Attention-2
混合 Attention 可以让一部分层走全局或稀疏 Attention，其他层走局部窗或线性递推。同层也可有多个分支，NSA 就保留了多路信息。全局层提供直接检索，状态层压低 Decode 对历史的扫描需求。对运行时而言，一次请求同时拥有 token 页和固定形状状态，峰值容量、调度和迁移都要逐类核算。MiniMax-01 NSA
Ring Attention 将序列分片，每个设备负责自己的 Q 块，沿环接收 K/V 块并更新在线 Softmax。对有效分片仍计算完整可见集合，收益来自分摊容量和让通信与计算交叠。链路延迟或末尾 tile 不均衡时，环上的最慢设备会限制进度。PagedAttention 改变的是 KV 的物理布局，StreamingLLM 则保留 sink 和最近窗口，改变可见历史。三者可以解决不同的长上下文瓶颈。Ring Attention PagedAttention StreamingLLM
5. 在线 Softmax 如何跨 tile 合并
对一个 query，将分数按 KV tile 切开。每段保存最大值 \(m\)、指数和 \(\ell\)、未归一化输出向量 \(u\)。分数记为 \(s_j\)，value 记为 \(v_j\)。Milakov 与 Gimelshein，2018
合并两段 A、B 时，先选公共基准 \(m'=\max(m_A,m_B)\)，再把两段的状态缩放到同一个基准。
例如两个分数分别为 2、4，对应标量 value 为 10、20。每段的局部指数和都是 1，但合并时第一段的权重应乘 \(e^{-2}\)，所以结果约为 18.81。平均两个局部输出会得到 15。这个区别也解释了 split-KV 为什么要带着归一化信息做归约。
在实数算术中，归约顺序可以自由组合，结果等价于完整 Softmax。浮点加法次序会产生舍入差异。实现还要把全遮挡 tile 当成空贡献，防止计算负无穷减负无穷。在因果 Prefill 中，整个上三角 tile 可以直接跳过。
6. FlashAttention 逐代解决了什么
| 工作 | 主要改动 | 观察对象 |
|---|---|---|
| FlashAttention，2022 | KV 分块、在线归一化，将分数和概率留在片上。反向用保存的统计量重算局部概率 | HBM 中间读写、片上容量 |
| FlashAttention-2，2023 | 减少非矩阵运算，沿序列增加 thread block 并行，块内按 Q 分配工作以减少 warp 间通信 | 占用率、共享内存通信 |
| FlashAttention-3，2024 | 针对 Hopper 采用异步搬运与矩阵运算，让不同阶段交叠，并研究 FP8 路径的数值误差 | Tensor Core 空转、Softmax 依赖、低精度误差 |
| FlashAttention-4，2026 | 针对 Blackwell 重做异步 MMA 流水，引入指数近似和条件重缩放，缓解矩阵吞吐与非矩阵操作增长不均的问题 | 指数吞吐、共享内存流量、反向归约 |
密集 FlashAttention 仍然计算可见的 Q–K 配对。它可以与 GQA、MLA 和分页缓存结合，也能作为规则稀疏 tile 的执行基础。每种缓存表示占多少内存，见《KV Cache 技术演进》。
7. Prefill 与 Decode 用同一套公式，走不同的数据路径
Prefill 一次处理 prompt 的多行 Q，通常按 Q tile 分工，复用每个 tile 内的 K/V。首次 Prefill 的 KV 写出建立了后续 Decode 的历史。若已有长度为 \(P\) 的前缀，再处理长度为 \(C\) 的 chunk，有效配对数为 \(CP+C(C+1)/2\)。切小 chunk 可以限制单次调度占用，但更短的 Q 轴也可能牺牲复用与矩阵单元利用率。
Decode 每条活跃序列通常只有一行 Q。长上下文、小 batch 下，单纯沿 batch 和 head 分工可能填不满设备。Flash-Decoding 沿 KV 轴分段，每段独立产生局部输出和归一化统计量，最后合并。KV 很短或并发足够时，额外的部分结果写回和归约可能抵消切分收益。Flash-Decoding，2023
用 GQA 给 Decode 做一笔理想账。设 Q head 数为 \(H_q\)，KV head 数为 \(H_{kv}\)，每个缓存元素占 \(b\) 字节。假设同组 Q head 复用一次加载的 KV，忽略页表、权重、Q/O 和归约流量，Attention 主体的算术强度约为
32 个 Q head、8 个 KV head、BF16 时，这个估计是 4 FLOP/byte。如果每个 Q head 都重新从 HBM 读取同组 KV，估计会接近 1 FLOP/byte。对照设备同一精度下的有效算力与有效带宽之比，可以判断 Attention 内核的瓶颈。完整模型还会读权重和执行其他算子，短上下文或小 batch 时应单独核算权重流量与启动开销。
8. 昇腾 NPU 上的 Cube、Vector 和搬运流水
下面以 Ascend C 的 Cube/Vector 分工描述设计思路。以一个 Q tile 为单位，先把 Q 和首段 K 搬进片上缓冲，由 Cube 做 \(QK^{\mathsf T}\)。随后在 Vector 路径执行缩放、mask、指数、行归约和重缩放，再由 Cube 计算概率 tile 与 V 的乘积。片上空间和同步条件允许时，可以让下一段 K/V 的搬运与当前段计算交叠，并用双缓冲管理交替读写。Cube 与 Vector 的交接路径、可用缓冲区和同步粒度需要按具体芯片与实现核对。Ascend C 硬件架构抽象
- Prefill 的调优点是 Q/KV tile 形状、对角 tile 的负载差异，以及概率 tile 交接时的布局转换。增大 tile 能提升复用，也可能因中间缓冲占用过大挤掉双缓冲。
- Decode 的调优点是为短 Q 找到足够多的独立任务。先检查 batch、KV group 和页粒度能提供多少并行度，再调 KV 切分。分片太细会增加页表访问、尾块浪费和跨核归约。把同组 Q 放进能共享 K/V 的任务，可用实测 HBM 读量验证效果。
- 接口边界取决于设备和 CANN 版本。CANN 8.5 的
aclnnFusedInferAttentionScoreV4文档区分 Prompt 与 Incremental 路径，并列出blockTable等分页参数。移植时先核对产品、输入布局、数据类型和实际长度约束，再比较融合路径与参考实现。昇腾 CANN 接口文档
做内核对比时，固定模型、dtype、head 配置和 causal mask，扫 prompt 长度、历史长度及活跃请求数。Prefill 记录 kernel latency、Cube/Vector 活跃度和实际 HBM 流量；Decode 再增加页表命中场景、KV 分段数、归约用时和每步延迟。服务层同时看首 token 延迟（TTFT）和 token 间延迟（ITL），避免长 prompt 的调度收益掩盖 Decode 的尾延迟变化。
9. 用同一组形状推导六条路线
下面把模型层、内核层与服务层放回同一个推导里。固定一个请求，Q head 数 \(H_q=32\)，KV head 数 \(H_{kv}=8\)，head 维度 \(d=128\)，历史长度 \(T=8192\)，缓存为 BF16。各方法若改变了这些条件，会在对应推导中单独标出。
9.1 从 GQA 的一组 Q 推到 MLA
取一组四个 Q head。Decode 时，组内 query 可以排成 \(Q_g\in\mathbb R^{4\times128}\)，共享的 \(K_g,V_g\in\mathbb R^{8192\times128}\)。于是
共享 K/V 后，四行 Q 仍产生四行分数。若两行分数是 \([2,0]\) 和 \([0,2]\)，两个 value 在某通道上为 10、20，那么输出分别约为 11.19、18.81。把多个 head 的 Softmax 混成一次归约就改变了计算语义。
每组一次读取的 K/V 是 \(2\times8192\times128\times2=4\) MiB，八组共 32 MiB。两个 GEMM 总量约 \(4H_qTd=134217728\) FLOPs。若共享加载生效，算术强度为 4 FLOP/B；若四个 Q head 各自重读，就退到约 1 FLOP/B。容量节省由表示决定，带宽节省还要由任务组织兑现。
再换成 MLA。用行向量记法，历史 latent 为 \(C\in\mathbb R^{T\times d_c}\)，单 head 上投影为 \(U_K,U_V\in\mathbb R^{d_c\times d}\)。先忽略位置分支：
左侧先展开历史，右侧先变换 query、再聚合 latent。令 \(d_c=512\)，Decode 左侧的历史 K 上投影需要约 \(2T d_cd\) FLOPs，右侧 query 变换只需 \(2d_cd\)。但后续 QK 的规约维从 128 变成 512，PV 的结果宽度也变成 512。短上下文时，新增投影和调度开销可能更显眼；长上下文时，减少历史展开和读流量更有机会获益。实际比较还要加回共享 RoPE key，并使用模型规定的缩放因子。
YOCO 又改变了层间的状态来源。假设 cross-decoder 有 \(L_c\) 层，传统每层独立全局 KV 需要 \(L_cTb_{\rm token}\) 字节，YOCO 共享全局部分的载荷变成 \(Tb_{\rm token}\)，另外加 self-decoder 的状态。每个 cross 层的 Q 仍随层变化，所以该层的分数和输出都要重新计算。若共享历史大于片上容量，逐层执行仍可能多次从 HBM 读取同一份 KV。YOCO
9.2 从窗口到动态稀疏选择
滑窗长度为 \(W\) 时，第 \(i\) 个 query 只允许读取最近的 \(\min(i,W)\) 个位置。长度为 \(T\)、且 \(T\ge W\) 的完整 Prefill 配对数为
取 \(W=1024,T=8192\)，配对数为 7,864,832；完整因果 Attention 为 33,558,528。两者的比值约为 23.4%。这个比例只描述配对数。对角 tile、边缘 padding、投影和其他层仍然需要计算。
块稀疏把窗口外的访问扩展到若干离散块。若每块 \(P\) 个 token、选 \(r\) 块，读取预算约为 \(rP\)。逐 token Top-K 能把预算精确到 \(K\) 个位置，但可能访问 \(K\) 个不同物理页。选中的 token 数相同，DMA 事务数量可以相差很大。
NSA 用压缩块提供全局线索，再访问选择块，并保留局部窗口。DSA 的轻量索引路径扫描候选，主路径读取 Top-K。设索引时间为 \(t_I\)，选择和地址整理为 \(t_S\)，主稀疏 Attention 为 \(t_A\)，则收益条件是
例如密集路径 1 ms，主稀疏路径 0.3 ms，索引和选择合计 0.8 ms，最后仍是 1.1 ms。这里是说明临界条件的假设数字。实际 profile 要分开测索引 GEMM、Top-K、页读取和主 QK/PV，才能定位下一步该优化什么。
9.3 把线性 Attention 的全序列求和变成递推
设特征映射维度为 \(r\)，value 维度为 \(d_v\)。定义相似度 \(\phi(q_i)^{\mathsf T}\phi(k_j)\)，则因果输出可写成
把只依赖 query 的部分提出求和，就得到
\(A_i\) 的形状是 \(r\times d_v\)。令两维都是 128，以 FP32 保存，一份状态为 64 KiB，另加归一化向量。它不会随历史长度增长，但 head 数、层数和 batch 都会增加总状态量。一次 Decode 更新是外积加法，随后是 query 与状态矩阵的乘法，小 batch 时需要特别关注状态读写和小矩阵利用率。Linear Transformer
带常数衰减的递推 \(A_i=\gamma A_{i-1}+k_iv_i^{\mathsf T}\) 展开后是 \(\sum_{j\le i}\gamma^{i-j}k_jv_j^{\mathsf T}\)。因此一个长度 \(C\) 的 chunk 可以写成
这一展开解释了并行 Prefill 与串行 Decode 的联系。块内用矩阵计算汇总增量，块间传递状态。RetNet 和 Lightning Attention 的具体归一化、位置处理与衰减形式还需遵循各自模型，但分块推导给出了同一类硬件组织方式。RetNet Lightning Attention-2
9.4 混合模型的状态怎样算
设共有 \(L\) 层，其中 \(L_f\) 层保留完整 KV，其余层用固定状态。每个 full 层每 token 载荷为 \(b_f\)，每个递推层状态为 \(b_r\)，那么一个请求的状态量近似为
只要还存在 full 层，总量就仍随 \(T\) 增长。取 \(L=32,L_f=8,b_f=4096\) B，完整历史部分在 8K 时是 256 MiB，再加 24 层的递推状态。与 32 层全为该 GQA 配置的 1 GiB 相比，可以直接看出改了哪一项。投机解码或请求分叉时，递推状态还需要支持快照或回滚，token 页和状态矩阵的生命周期都要由调度器管理。混合比例与状态尺寸应从模型配置读取。Kimi Linear
9.5 精确 Attention 的 tile 为什么可以任意合并
在线 Softmax 保存三项状态。对已处理集合 \(A\)，保存最大值 \(m_A\)、指数和 \(\ell_A\)、加权 value 和 \(u_A\)。新块的最大值 \(m_B\) 可能更大，旧状态的指数基准就必须同步改变。
设 \(m_A=2,m_B=4\)，旧块的指数项原先相对 2 计算，现在相对 4 就整体乘 \(e^{-2}\)。因此先做重缩放，再相加，最后除以总指数和，就保留了原始 Softmax 的比例。这使 Prefill 的串行 tile 遍历、Decode 的分段归约和 Ring Attention 的跨设备合并可以共用同一数学状态。
FlashAttention 避免把完整 \(T\times T\) 分数写回 HBM。以 query tile 高度 \(B_q\)、KV tile 长度 \(B_k\) 计，片上要同时容纳 Q/K/V tile、局部分数或概率、行统计量和输出累积，双缓冲又会增加部分占用。tile 大小应从可用片上存储反推，再用有效吞吐验证。FA2 优化工作划分，FA3/4 进一步针对异步矩阵指令和非矩阵运算组织流水，各代策略都需要映射到具体硬件。FlashAttention FA2 FA3 FA4
9.6 长上下文如何跨设备
Ring Attention 把 Q 固定在本设备，把 K/V 分块沿环传递。收到一个块后更新本地 query 的在线状态，直到覆盖需要的历史。若每轮计算时间为 \(t_c\)，通信时间为 \(t_n\)，理想双缓冲下稳态每轮接近 \(\max(t_c,t_n)\)，另加启动与收尾。小块可能被通信时延主导，大块则要求更多片上和设备缓冲。
Decode 的 query 很短，单轮计算不足以遮住通信时，分片 KV 的读取与归约就更关键。分页决定本地数据怎样定位，StreamingLLM 的保留策略决定哪些历史仍可见，Ring 决定这些数据怎样跨设备移动。设计组合时先明确可见集合，再选择物理布局和分布方式。Ring Attention
