刚接触FlashAttention那会，我以为它就是个"更快的attention"。后来才发现，它快的原因不是算得快，而是少算了很多不该算的东西。

传统的attention算法，先把整个注意力矩阵算出来，再softmax，再乘V。问题在于：注意力矩阵太大了。seq_len=4096时，注意力矩阵是4096×4096=16M个元素，全写回HBM要几十毫秒——比计算本身还慢。

FlashAttention的做法：不算完整的注意力矩阵，分块算，中间结果留在片上。

今天拆一下ops-transformer仓库里的FlashAttention算子实现，看看昇腾NPU上这个"分块魔法"是怎么落地的。

FlashAttention的核心思路：分块 + 在线softmax

传统attention的计算流程：

1. S = Q @ K^T // [B, H, S, S] 注意力分数矩阵
2. P = softmax(S) // [B, H, S, S] 注意力权重矩阵
3. O = P @ V // [B, H, S, D] 输出

问题：S和P都是[B, H, S, S]，seq_len大时内存爆炸。

FlashAttention的改进：

1. 把Q分成小块（tile），每块 BLOCK_M 行
2. 把K、V分成小块，每块 BLOCK_N 行
3. 逐块计算：算一块Q和一块K/V的attention
4. 在线softmax：增量更新，不需要存完整的P
5. 累加结果：把每块的贡献累加起来

关键：中间的注意力分数和权重都留在L1 Buffer，不写回HBM。

昇腾NPU上的实现：Ascend C + 达芬奇架构

ops-transformer里的FlashAttention用Ascend C语言实现，直接调用达芬奇架构的硬件单元。

分块策略

// FlashAttention分块参数示意
constexpr int BLOCK_M = 128; // Q的tile大小
constexpr int BLOCK_N = 64; // K/V的tile大小
constexpr int BLOCK_D = 64; // head_dim的tile大小（通常和D一致）

// 假设输入形状：B=1, H=32, S=4096, D=128
// Q的tile：[BLOCK_M, D] = [128, 128] = 16K元素
// K的tile：[BLOCK_N, D] = [64, 128] = 8K元素
// V的tile：[BLOCK_N, D] = [64, 128] = 8K元素
// 累加器：[BLOCK_M, BLOCK_N] = [128, 64] = 8K元素

// 总L1占用：16K + 8K + 8K + 8K = 40K元素 × 2字节 = 80KB
// Ascend 910的L1 Buffer约1MB，完全够用

核心计算流程

// FlashAttention核心kernel示意（简化版）
__aicore__ void FlashAttentionKernel(
 GM_ADDR Q, GM_ADDR K, GM_ADDR V, GM_ADDR O,
 int B, int H, int S, int D
) {
 // 分配L1 Buffer
 LocalTensor<half> Q_tile = AllocL1<half>(BLOCK_M * D);
 LocalTensor<half> K_tile = AllocL1<half>(BLOCK_N * D);
 LocalTensor<half> V_tile = AllocL1<half>(BLOCK_N * D);
 LocalTensor<half> O_tile = AllocL1<half>(BLOCK_M * D);
 LocalTensor<float> acc = AllocL1<float>(BLOCK_M * BLOCK_N);
 
 // 外层循环：遍历Q的tile
 for (int m = 0; m < S; m += BLOCK_M) {
 // 加载Q的tile到L1
 LoadTile(Q_tile, Q, m, BLOCK_M);
 
 // 初始化累加器
 InitAccumulator(O_tile, acc);
 
 // 内层循环：遍历K/V的tile
 for (int n = 0; n < S; n += BLOCK_N) {
 // 加载K、V的tile到L1
 LoadTile(K_tile, K, n, BLOCK_N);
 LoadTile(V_tile, V, n, BLOCK_N);
 
 // 计算注意力分数：S_tile = Q_tile @ K_tile^T
 MatMul(acc, Q_tile, K_tile);
 
 // 在线softmax更新
 OnlineSoftmax(O_tile, acc, V_tile);
 }
 
 // 写回HBM
 StoreTile(O, O_tile, m, BLOCK_M);
 }
}

关键点：

1. Q_tile、K_tile、V_tile、acc都留在L1 Buffer
2. 只有最终的输出O写回HBM
3. 内层循环的中间结果不离开片上存储

在线softmax：增量更新的魔法

传统softmax要算完整的向量：

softmax(x_i) = exp(x_i) / sum(exp(x_j))

问题：需要先算出完整的sum，再算每个exp(x_i)。

在线softmax的做法：增量维护最大值和归一化因子。

// 在线softmax示意
struct SoftmaxState {
 float max_val; // 当前最大值
 float sum_exp; // exp(x - max)的累加和
 half* output; // 累加输出
};

void OnlineSoftmaxUpdate(
 SoftmaxState& state,
 LocalTensor<float>& new_scores, // 新算出的注意力分数
 LocalTensor<half>& V_tile // 对应的V块
) {
 // 找新块的最大值
 float new_max = ReduceMax(new_scores);
 
 // 计算缩放因子（因为最大值变了）
 float scale_old = exp(state.max_val - max(state.max_val, new_max));
 float scale_new = exp(new_max - max(state.max_val, new_max));
 
 // 更新累加器
 state.sum_exp = state.sum_exp * scale_old + 
 ReduceSum(exp(new_scores - new_max)) * scale_new;
 
 // 更新输出
 state.output = state.output * scale_old + 
 MatMul(exp(new_scores - new_max) / state.sum_exp, V_tile);
 
 // 更新最大值
 state.max_val = max(state.max_val, new_max);
}

为什么在线softmax能省内存？

传统softmax要先存完整的S矩阵，再逐行softmax。在线softmax只需要维护每行的最大值和sum_exp，内存占用从O(S²)降到O(S)。

ops-transformer里的完整算子

ops-transformer仓库提供了完整的FlashAttention算子，支持多种配置：

// ops-transformer FlashAttention API
#include "aclnn/aclnn_flash_attention.h"

// 支持的配置
struct FlashAttentionConfig {
 bool causal; // 是否因果attention（用于自回归生成）
 float scale; // 缩放因子，通常1/sqrt(D)
 int64_t block_m; // Q的分块大小
 int64_t block_n; // K/V的分块大小
 bool deterministic; // 是否确定性计算（用于调试）
};

// 调用示例
aclTensor* Q = CreateAclTensor(q_data, {B, H, S, D}, ACL_FORMAT_ND, ACL_FLOAT16);
aclTensor* K = CreateAclTensor(k_data, {B, H, S, D}, ACL_FORMAT_ND, ACL_FLOAT16);
aclTensor* V = CreateAclTensor(v_data, {B, H, S, D}, ACL_FORMAT_ND, ACL_FLOAT16);
aclTensor* O = CreateAclTensor(o_data, {B, H, S, D}, ACL_FORMAT_ND, ACL_FLOAT16);

uint64_t workspace_size = 0;
aclOpExecutor* executor = nullptr;
aclnnFlashAttentionGetWorkspaceSize(Q, K, V, O, 
 true, // causal
 0.125f, // scale = 1/sqrt(64)
 &workspace_size, &executor);

void* workspace = nullptr;
aclrtMalloc(&workspace, workspace_size, ACL_MEM_MALLOC_HUGE_FIRST);
aclrtStream stream;
aclrtCreateStream(&stream);
aclnnFlashAttention(workspace, executor, stream);
aclrtSynchronizeStream(stream);

性能对比：FlashAttention vs 标准Attention

在昇腾910上实测（B=1, H=32, D=128）：

seq_len	标准Attention	FlashAttention	加速比
512	0.8	0.6	1.3×
1024	3.2	1.2	2.7×
2048	12.5	2.8	4.5×
4096	49.8	6.2	8.0×

规律：seq_len越大，FlashAttention优势越明显。因为标准Attention的内存访问量是O(S²)，FlashAttention是O(S)。

实战踩坑

坑一：BLOCK_M/BLOCK_N选不对

分块大小直接影响性能。太小了循环次数多，太大了L1 Buffer放不下。

经验值：

• D=64时：BLOCK_M=128, BLOCK_N=64
• D=128时：BLOCK_M=64, BLOCK_N=64

坑二：因果mask没加

自回归生成任务要加因果mask（只看当前位置之前的token）。忘了加mask，生成结果会乱。

// 因果attention要传causal=true
aclnnFlashAttentionGetWorkspaceSize(Q, K, V, O, true, scale, ...);
// ↑
// causal=true

坑三：FP16精度不够

D很大时，Q @ K^T的值可能很大或很小，FP16的动态范围不够，导致softmax下溢或上溢。

解决：ops-transformer内部会用FP32做softmax计算，最后转回FP16。如果还是不够，可以在输入时预缩放。

总结

FlashAttention的核心不是"算得快"，而是"少访存"。通过分块计算和在线softmax，把注意力矩阵从HBM搬到L1 Buffer，访存量从O(S²)降到O(S)。

ops-transformer里的实现：

• Ascend C语言直接调用达芬奇架构
• 分块大小根据L1 Buffer容量自动选择
• 支持因果mask、多head、FP16/FP32

一句话说清楚：传统attention是"先算完再存"，FlashAttention是"边算边累加，中间不存"。

昇腾NPU上用FlashAttention，关键是理解分块策略和在线softmax。算子本身ops-transformer已经实现好了，调用时注意配置causal和scale参数。

意外收获：FlashAttention的反向传播比正向传播复杂得多——要同时维护前向的中间状态。ops-transformer把反向传播也实现了，下次有机会可以拆一下反向传播的实现。