背单词为什么不背词典：CANN上FlashAttention的分块逻辑

上个月有个实习生问我，为什么昇腾CANN的ops-transformer仓库里，FlashAttention算子比标准实现快那么多。我说你先想一个问题：背四级单词，你是把整本词典摊开从头背，还是一次看一页？

他说当然是看一页。我说对了，这就够了。

标准Attention：一次摊开整本词典

Attention做的事情，说白了就是让模型里每个token跟其他所有token"打个招呼"，算算谁跟谁最相关。

一句话有2048个token，每个token要跟另外2047个算相似度。2048 × 2048 = 约400万次计算，产生的中间矩阵必须存在显存里。float16存，每个元素2字节，这个矩阵大概8MB。

一个attention层8MB，Transformer有32层、32个head，中间结果加起来好几GB。反向传播还要存梯度，显存直接爆。

这就是标准Attention在昇腾NPU上跑长序列会OOM的原因——它非要把完整的注意力矩阵一次性算完存完，才肯往下走。就像你非要把整本词典摊在桌上，桌子不够大怎么办？办不了，直接报错。

FlashAttention：一页一页翻

FlashAttention做的事特别朴素：把巨大的注意力矩阵切成小块，每次只算一块，处理完就扔。

这就是翻词典的逻辑——桌子（昇腾NPU的片上存储）只要放得下一页（一个tile），就够了，不需要整本摊开。

ops-transformer里的实现流程：

📚 加载Q的一个tile到L1 Buffer 📚 逐块加载K的tile，跟Q做矩阵乘 → 分块attention分数 📚 对分数做在线softmax → 边算边归一化，不等全部算完 📚 用softmax结果加权V的tile → 分块输出 📚 输出写回HBM，中间的QK^T矩阵直接丢掉 📚 重复，直到所有Q的tile处理完

三个技术细节值得说。

在线softmax：不等人到齐就动筷子

标准softmax必须看到所有分数才能归一化，就像聚餐得等所有人到齐才动筷。FlashAttention的在线softmax允许你先吃，后面来人了再按比例调。

数学上用了一个"running max"技巧：每来一块新分数，跟之前最大值比较，更新全局最大值，然后把已经归一化的结果按新比例调整一遍。

结果跟标准softmax完全一致，但不需要把所有分数同时放在显存。

IO感知：搬运才是真瓶颈

昇腾达芬奇架构上，计算不是瓶颈——矩阵计算单元（Cube Unit）每秒几百TFLOPS。真正的瓶颈是数据搬运：从HBM搬到L1 Buffer，再搬回来。

标准Attention的搬运次数是O(N²)——大矩阵要读一遍、写一遍、再读一遍给softmax用。FlashAttention分块后降成O(N²/B)，B是tile大小。tile越大，搬运越少。

ops-transformer在Ascend 910上默认tile大小是128×128或64×128，平衡了L1容量和计算单元利用率。

融合kernel：砍掉中间商

标准Attention分三步：算QK^T → softmax → 加权V。每步之间数据要从L1写回HBM再读出来。

ops-transformer把这三步融合成一个kernel，全程留在L1 Buffer里：

// ops-transformer里的融合kernel示意 // 三步合一，中间结果不落HBM template <typename T> void FusedFlashAttention(KernelContext* ctx, const Tensor* Q, const Tensor* K, const Tensor* V, Tensor* out) { constexpr int BM = 64; // Q的行tile大小 constexpr int BN = 64; // K的列tile大小 constexpr int BD = 128; // head_dim，对齐128最高效 for (int m = 0; m < seq_len; m += BM) { // Q的一个tile加载到L1，全程不回HBM auto q_tile = ctx->LoadTileToL1(Q, m, BM, BD); // 累加器存softmax后的结果 auto acc = ctx->AllocL1<T>(BM * BN); auto row_max = ctx->AllocL1<float>(BM); auto row_sum = ctx->AllocL1<float>(BM); ZeroInit(acc, row_max, row_sum); for (int n = 0; n < seq_len; n += BN) { auto k_tile = ctx->LoadTileToL1(K, n, BN, BD); auto v_tile = ctx->LoadTileToL1(V, n, BN, BD); // 算QK^T的分块 → 直接softmax → 直接加权V // 三步在L1里一口气做完，不经过HBM auto scores = MatMul(q_tile, k_tile); Scale(scores, 1.0f / sqrt((float)BD)); OnlineSoftmax(scores, acc, row_max, row_sum, BM, BN); auto partial_out = MatMul(acc, v_tile); Accumulate(out + m, partial_out, BM, BD); } // 循环结束，acc和scores直接丢弃 // 那个巨大的N×N矩阵从头到尾没存在HBM过 } }

关键不在代码怎么写，而在为什么要这么写。注释里写的是"全程不回HBM"“三步一口气做完”——解释的是WHY。如果你去掉融合，每步之间多两次HBM读写，光IO开销就能让性能掉一半。

实测数据

我在Ascend 910上跑了ops-transformer的FlashAttention，对比标准实现：

短序列（512）FlashAttention反而慢一点，因为分块控制有开销。但超过2048之后，不只是快——是能跑和不能跑的区别。

调用方式

不想写kernel代码的话，PyTorch里一行就能调：

import torch from ascend_rs import flash_attention Q = torch.randn(1, 32, 4096, 128, device='npu', dtype=torch.float16) K = torch.randn(1, 32, 4096, 128, device='npu', dtype=torch.float16) V = torch.randn(1, 32, 4096, 128, device='npu', dtype=torch.float16) # 底层自动走ops-transformer的FlashAttention算子 out = flash_attention(Q, K, V, attn_scale=1.0 / (128 ** 0.5))

tile策略、在线softmax、融合kernel，全部自动处理。你只需要知道一件事：序列长的时候用这个，短的时候用标准的就行。

意外收获：FlashAttention论文的作者Tri Dao，本科在MIT学数学，博士在Stanford做系统优化。他后来去了一家叫Flashbots的公司做MEV研究。注意力机制和区块链抽水，本质上都是IO优化——怎么在有限的带宽里搬运最关键的数据。