刚接触昇腾CANN生态的时候光是找FlashAttention算子在哪就花了不少时间。官方文档按功能模块划分仓库按层级划分两套逻辑对不上号——文档里写的是大模型算子优化仓库里是一个叫ops-transformer的目录。这种对应关系不搞清楚改代码都不知道去哪改。ops-transformer是昇腾CANN大模型算子的主阵地。FlashAttention、RMSNorm、RoPE、SwiGLU——大模型推理训练用到的核心算子全在这个仓库里。仓库在CANN生态里的位置昇腾CANN的软件栈从下到上分四层硬件层Ascend 910 NPU ↑ 驱动与Runtime层CANN基础组件固件、CCE编译器、调度器 ↑ 算子层ops-transformer / catlass / ops-ascendc ↑ 框架适配层torch_npu / msModelZoo / geops-transformer在算子层偏大模型方向。通用算子卷积、池化、BatchNorm等在ops-ascendc不在ops-transformer里。这两个仓库的边界要记清楚不然在ops-transformer里搜conv2d永远搜不到。ops-transformer跟catlass是上下游关系。catlass提供分块矩阵乘、在线softmax等基础操作的模板ops-transformer的FlashAttention依赖catlass的模板实现具体算法。开发新算子的时候catlass提供积木ops-transformer负责搭房子。目录结构FlashAttention藏在哪git clone https://atomgit.com/cann/ops-transformer.git cd ops-transformer find . -type f -name *.cc -o -name *.h | head -30顶层结构ops-transformer/ ├── opkernel/ # 算子内核实现改FlashAttention核心逻辑的地方 │ ├── flash_attention/ │ │ ├── flash_attention_score.cc # 前向分块在线softmax │ │ ├── flash_attention_score_grad.cc # 反向重计算 │ │ └── flash_attention_score_tiling.cc # 分块策略UB容量、块大小 │ ├── rms_norm/ │ │ └── rms_norm.cc │ └── rope/ │ └── rope.cc ├── opplugin/ # 算子注册GE怎么找到FlashAttention │ └── flash_attention/ │ └── flash_attention_op.cc ├── inc/ # 公共头文件 │ └── ops_kernel.h ├── scripts/ # 编译脚本 └── cmake/ └── FindCANN.cmake三个核心文件各有分工tiling.cc算分块参数score.cc写计算逻辑score_grad.cc写反向传播。tiling.cc分块大小不是随便定的FlashAttention的核心思想是分块计算——把Q、K、V切成小块每次只拿一小块算算完就扔中间结果。分块大小怎么定这是tiling.cc要做的事。昇腾NPU有个硬约束Unified Buffer约256KB。UB是达芬奇架构上的高速存储计算单元直接访问数据必须先从全局内存HBM搬到UB才能用。一次计算需要同时放得下Q块、K块、V块、输出块、softmax中间结果分块太大直接溢出。// ops-transformer/opkernel/flash_attention/flash_attention_score_tiling.cc // 简化分块策略伪代码 struct TilingParam { uint32_t block_m; // Q块在seq维的大小 uint32_t block_n; // K/V块在seq维的大小 uint32_t block_k; // head_dim }; TilingParam CalcTiling(uint32_t seq_len, uint32_t head_dim) { // 昇腾910的UB约256KB const uint32_t kUBSize 256 * 1024; // 字节 const uint32_t dtype_bytes 2; // FP16 2字节 // 单块需要的空间估算 // Q块(block_m × head_dim) K块(block_n × head_dim) // V块(block_n × head_dim) 输出块(block_m × head_dim) // softmax中间结果(block_m × block_n) uint32_t elements_per_block kUBSize / (5 * dtype_bytes); // 反推seq维的块大小向下取到16的倍数对齐约束 uint32_t block_seq elements_per_block / head_dim; block_seq (block_seq / 16) * 16; block_seq std::min(block_seq, seq_len); return {block_seq, block_seq, head_dim}; }分块策略不是越大越好。块大了UB装不下块小了循环次数多、算子调度开销大。tiling.cc的职责就是在硬件约束下找最优的块大小。昇腾910的UB容量、128字节对齐要求、Cube计算单元的tile形状——这些硬件参数全部要反映在分块逻辑里。score.cc前向计算的核心tiling.cc算出了块大小score.cc负责按这个分块做计算。核心逻辑就是论文里的在线softmax// flash_attention_score.cc 核心循环伪代码 for (uint32_t qi 0; qi num_q_blocks; qi) { // 加载一块Q到UB LoadTensor(cur_q, q_base, qi * block_m); // 初始化softmax累加器 float row_max -1e9f; float row_sum 0.0f; float* acc_out ub_buffer; for (uint32_t ki 0; ki num_kv_blocks; ki) { // 加载K/V块到UB与计算流水化 LoadTensor(cur_k, k_base, ki * block_n); LoadTensor(cur_v, v_base, ki * block_n); // 局部注意力分数Q × K^T // 达芬奇Cube单元做矩阵乘高吞吐 Gemm(cur_q, cur_k.T, local_scores, block_m, block_n, head_dim); Scale(local_scores, 1.0f / sqrtf(head_dim)); // causal mask达芬奇上用块级跳过不算下三角矩阵 // 如果整块都在mask外直接跳过省掉整块计算量 if (IsBlockOutsideCausal(qi, ki, block_m, block_n)) { continue; } // 在线softmax更新 float local_max ReduceMax(local_scores); float new_max fmax(row_max, local_max); // 关键缩放之前的累加结果 // exp(row_max - new_max)保证了数学等价性 float correction expf(row_max - new_max); row_sum * correction; Scale(acc_out, correction); // 加上当前块的贡献 float local_sum ReduceSum(Exp(local_scores - new_max)); row_sum local_sum; Accumulate(acc_out, Exp(local_scores - new_max), cur_v); row_max new_max; } // 归一化输出 Scale(acc_out, 1.0f / row_sum); StoreTensor(out_base, qi * block_m, acc_out); }昇腾NPU上有几个值得注意的细节causal mask用块级跳过。标准的因果mask是算出一个完整的下三角矩阵再逐元素乘。FlashAttention天然适合块级跳过——如果整块K/V都在当前Q的因果范围之外直接continue跳过这一块不浪费计算资源。数据搬运和计算流水化。加载下一块K/V的同时当前块的矩阵乘和softmax计算同步进行。达芬奇架构的DMA引擎和Cube计算单元独立运作overlap得好搬运时间可以被计算时间掩盖。数值稳定性。FP16下exp(x)在x88.7时溢出。每次更新最大值后用exp(row_max - new_max)重新缩放之前的累加结果把所有分数拉到同一个尺度——这行代码是整个在线softmax数值稳定的关键。score_grad.cc反向传播怎么重计算标准attention前向时存了scores和attn两个N×N矩阵反向时直接用。FlashAttention不存中间结果反向时重新算一遍——这就是重计算recomputation。// flash_attention_score_grad.cc 反向逻辑伪代码 void FlashAttentionGrad( Tensor grad_out, // 输出梯度 Tensor q, k, v, // 前向输入 Tensor out, # 前向输出必须存 Tensor grad_q, grad_k, grad_v // 输出梯度 ) { // 前向时存的两个东西out和row_max/row_sum // 没有存scores和attn反向时重新算 // 第1步算输出对out的梯度 // dP grad_out × V^T for each block: dP_block Gemm(grad_out_block, v_block.T); // softmax反向dS P * (dP - sum(dP * P)) dS_block SoftmaxBackward(dP_block, out_block); // 累积到grad_q grad_q_block Gemm(dS_block, K); // 第2步重计算Q×K^T再反向传播到grad_k和grad_v for each block: # 重新算前向的scores scores_block Gemm(Q_block, K_block.T) * scale; # 链式法则反向 grad_k_block Gemm(Q_block.T, dS_block); grad_v_block Gemm(dS_block.T, V_block); }重计算的代价是多算一遍前向但换来了显存从O(N²)降到O(N)。达芬奇架构的算力充沛显存带宽才是瓶颈——这个trade-off划算。算子怎么注册到GE写完计算逻辑还得让框架能调到它。opplugin/flash_attention_op.cc负责把FlashAttention注册到GE图引擎// flash_attention_op.cc 算子注册伪代码 IMPLEMT_INFERFUNC(FlashAttentionScore, FlashAttentionScoreInfer) { // GE需要知道输出的shape推导逻辑 // output.shape Q.shape auto q_shape op.GetInputDesc(0).GetShape().GetDims(); op.UpdateOutputDesc(output, TensorDesc(q_shape, FORMAT_ND, DT_FLOAT16)); return GRAPH_SUCCESS; } REGISTER_OP(FlashAttentionScore) .Input(q: float16) .Input(k: float16) .Input(v: float16) .Output(output: float16) .Attr(head_num: int) .Attr(input_layout: string) .Attr(scale: float) .InferShapeFunction(FlashAttentionScoreInfer);注册完成后torch_npu调用npu_flash_attention时GE根据算子名找到这个注册把计算委托给opkernel/flash_attention/下的实现。跟其他仓库的关系Ascend C达芬奇架构的C编程接口 ↑ catlass算子模板库GEMM、softmax、reduce ↑ ops-transformer算子实现FlashAttention、RMSNorm等 ↑ ge图引擎算子编排、图优化、算子融合 ↑ torch_npuPyTorch适配npu_flash_attention等接口开发FlashAttention变体时逻辑改ops-transformer分块策略改catlass的模板参数图优化改ge。编译和验证改完代码后的编译流程cd ops-transformer mkdir build cd build source /usr/local/Ascend/ascend-toolkit/set_env.sh cmake .. -DCANN_INSTALL_PATH/usr/local/Ascend/ascend-toolkit \ -DCMAKE_BUILD_TYPERelease # 只编译FlashAttention单个算子比全量编译快 make flash_attention -j8 # 验证编译产物 ls output/opkernel/libflash_attention*.so替换到torch_npu路径或指定自定义算子库路径后用数值验证脚本确认改动没有引入回归。克隆ops-transformer仓库按opkernel/flash_attention/→opplugin/flash_attention_op.cc的顺序读代码。tiling.cc的分块策略和score.cc的在线softmax是核心理解了这两块就抓住了FlashAttention昇腾实现的主线。https://atomgit.com/cann/ops-transformer
昇腾NPU上的FlashAttention藏在哪？ops-transformer仓库全景图
发布时间：2026/5/23 1:07:57
刚接触昇腾CANN生态的时候光是找FlashAttention算子在哪就花了不少时间。官方文档按功能模块划分仓库按层级划分两套逻辑对不上号——文档里写的是大模型算子优化仓库里是一个叫ops-transformer的目录。这种对应关系不搞清楚改代码都不知道去哪改。ops-transformer是昇腾CANN大模型算子的主阵地。FlashAttention、RMSNorm、RoPE、SwiGLU——大模型推理训练用到的核心算子全在这个仓库里。仓库在CANN生态里的位置昇腾CANN的软件栈从下到上分四层硬件层Ascend 910 NPU ↑ 驱动与Runtime层CANN基础组件固件、CCE编译器、调度器 ↑ 算子层ops-transformer / catlass / ops-ascendc ↑ 框架适配层torch_npu / msModelZoo / geops-transformer在算子层偏大模型方向。通用算子卷积、池化、BatchNorm等在ops-ascendc不在ops-transformer里。这两个仓库的边界要记清楚不然在ops-transformer里搜conv2d永远搜不到。ops-transformer跟catlass是上下游关系。catlass提供分块矩阵乘、在线softmax等基础操作的模板ops-transformer的FlashAttention依赖catlass的模板实现具体算法。开发新算子的时候catlass提供积木ops-transformer负责搭房子。目录结构FlashAttention藏在哪git clone https://atomgit.com/cann/ops-transformer.git cd ops-transformer find . -type f -name *.cc -o -name *.h | head -30顶层结构ops-transformer/ ├── opkernel/ # 算子内核实现改FlashAttention核心逻辑的地方 │ ├── flash_attention/ │ │ ├── flash_attention_score.cc # 前向分块在线softmax │ │ ├── flash_attention_score_grad.cc # 反向重计算 │ │ └── flash_attention_score_tiling.cc # 分块策略UB容量、块大小 │ ├── rms_norm/ │ │ └── rms_norm.cc │ └── rope/ │ └── rope.cc ├── opplugin/ # 算子注册GE怎么找到FlashAttention │ └── flash_attention/ │ └── flash_attention_op.cc ├── inc/ # 公共头文件 │ └── ops_kernel.h ├── scripts/ # 编译脚本 └── cmake/ └── FindCANN.cmake三个核心文件各有分工tiling.cc算分块参数score.cc写计算逻辑score_grad.cc写反向传播。tiling.cc分块大小不是随便定的FlashAttention的核心思想是分块计算——把Q、K、V切成小块每次只拿一小块算算完就扔中间结果。分块大小怎么定这是tiling.cc要做的事。昇腾NPU有个硬约束Unified Buffer约256KB。UB是达芬奇架构上的高速存储计算单元直接访问数据必须先从全局内存HBM搬到UB才能用。一次计算需要同时放得下Q块、K块、V块、输出块、softmax中间结果分块太大直接溢出。// ops-transformer/opkernel/flash_attention/flash_attention_score_tiling.cc // 简化分块策略伪代码 struct TilingParam { uint32_t block_m; // Q块在seq维的大小 uint32_t block_n; // K/V块在seq维的大小 uint32_t block_k; // head_dim }; TilingParam CalcTiling(uint32_t seq_len, uint32_t head_dim) { // 昇腾910的UB约256KB const uint32_t kUBSize 256 * 1024; // 字节 const uint32_t dtype_bytes 2; // FP16 2字节 // 单块需要的空间估算 // Q块(block_m × head_dim) K块(block_n × head_dim) // V块(block_n × head_dim) 输出块(block_m × head_dim) // softmax中间结果(block_m × block_n) uint32_t elements_per_block kUBSize / (5 * dtype_bytes); // 反推seq维的块大小向下取到16的倍数对齐约束 uint32_t block_seq elements_per_block / head_dim; block_seq (block_seq / 16) * 16; block_seq std::min(block_seq, seq_len); return {block_seq, block_seq, head_dim}; }分块策略不是越大越好。块大了UB装不下块小了循环次数多、算子调度开销大。tiling.cc的职责就是在硬件约束下找最优的块大小。昇腾910的UB容量、128字节对齐要求、Cube计算单元的tile形状——这些硬件参数全部要反映在分块逻辑里。score.cc前向计算的核心tiling.cc算出了块大小score.cc负责按这个分块做计算。核心逻辑就是论文里的在线softmax// flash_attention_score.cc 核心循环伪代码 for (uint32_t qi 0; qi num_q_blocks; qi) { // 加载一块Q到UB LoadTensor(cur_q, q_base, qi * block_m); // 初始化softmax累加器 float row_max -1e9f; float row_sum 0.0f; float* acc_out ub_buffer; for (uint32_t ki 0; ki num_kv_blocks; ki) { // 加载K/V块到UB与计算流水化 LoadTensor(cur_k, k_base, ki * block_n); LoadTensor(cur_v, v_base, ki * block_n); // 局部注意力分数Q × K^T // 达芬奇Cube单元做矩阵乘高吞吐 Gemm(cur_q, cur_k.T, local_scores, block_m, block_n, head_dim); Scale(local_scores, 1.0f / sqrtf(head_dim)); // causal mask达芬奇上用块级跳过不算下三角矩阵 // 如果整块都在mask外直接跳过省掉整块计算量 if (IsBlockOutsideCausal(qi, ki, block_m, block_n)) { continue; } // 在线softmax更新 float local_max ReduceMax(local_scores); float new_max fmax(row_max, local_max); // 关键缩放之前的累加结果 // exp(row_max - new_max)保证了数学等价性 float correction expf(row_max - new_max); row_sum * correction; Scale(acc_out, correction); // 加上当前块的贡献 float local_sum ReduceSum(Exp(local_scores - new_max)); row_sum local_sum; Accumulate(acc_out, Exp(local_scores - new_max), cur_v); row_max new_max; } // 归一化输出 Scale(acc_out, 1.0f / row_sum); StoreTensor(out_base, qi * block_m, acc_out); }昇腾NPU上有几个值得注意的细节causal mask用块级跳过。标准的因果mask是算出一个完整的下三角矩阵再逐元素乘。FlashAttention天然适合块级跳过——如果整块K/V都在当前Q的因果范围之外直接continue跳过这一块不浪费计算资源。数据搬运和计算流水化。加载下一块K/V的同时当前块的矩阵乘和softmax计算同步进行。达芬奇架构的DMA引擎和Cube计算单元独立运作overlap得好搬运时间可以被计算时间掩盖。数值稳定性。FP16下exp(x)在x88.7时溢出。每次更新最大值后用exp(row_max - new_max)重新缩放之前的累加结果把所有分数拉到同一个尺度——这行代码是整个在线softmax数值稳定的关键。score_grad.cc反向传播怎么重计算标准attention前向时存了scores和attn两个N×N矩阵反向时直接用。FlashAttention不存中间结果反向时重新算一遍——这就是重计算recomputation。// flash_attention_score_grad.cc 反向逻辑伪代码 void FlashAttentionGrad( Tensor grad_out, // 输出梯度 Tensor q, k, v, // 前向输入 Tensor out, # 前向输出必须存 Tensor grad_q, grad_k, grad_v // 输出梯度 ) { // 前向时存的两个东西out和row_max/row_sum // 没有存scores和attn反向时重新算 // 第1步算输出对out的梯度 // dP grad_out × V^T for each block: dP_block Gemm(grad_out_block, v_block.T); // softmax反向dS P * (dP - sum(dP * P)) dS_block SoftmaxBackward(dP_block, out_block); // 累积到grad_q grad_q_block Gemm(dS_block, K); // 第2步重计算Q×K^T再反向传播到grad_k和grad_v for each block: # 重新算前向的scores scores_block Gemm(Q_block, K_block.T) * scale; # 链式法则反向 grad_k_block Gemm(Q_block.T, dS_block); grad_v_block Gemm(dS_block.T, V_block); }重计算的代价是多算一遍前向但换来了显存从O(N²)降到O(N)。达芬奇架构的算力充沛显存带宽才是瓶颈——这个trade-off划算。算子怎么注册到GE写完计算逻辑还得让框架能调到它。opplugin/flash_attention_op.cc负责把FlashAttention注册到GE图引擎// flash_attention_op.cc 算子注册伪代码 IMPLEMT_INFERFUNC(FlashAttentionScore, FlashAttentionScoreInfer) { // GE需要知道输出的shape推导逻辑 // output.shape Q.shape auto q_shape op.GetInputDesc(0).GetShape().GetDims(); op.UpdateOutputDesc(output, TensorDesc(q_shape, FORMAT_ND, DT_FLOAT16)); return GRAPH_SUCCESS; } REGISTER_OP(FlashAttentionScore) .Input(q: float16) .Input(k: float16) .Input(v: float16) .Output(output: float16) .Attr(head_num: int) .Attr(input_layout: string) .Attr(scale: float) .InferShapeFunction(FlashAttentionScoreInfer);注册完成后torch_npu调用npu_flash_attention时GE根据算子名找到这个注册把计算委托给opkernel/flash_attention/下的实现。跟其他仓库的关系Ascend C达芬奇架构的C编程接口 ↑ catlass算子模板库GEMM、softmax、reduce ↑ ops-transformer算子实现FlashAttention、RMSNorm等 ↑ ge图引擎算子编排、图优化、算子融合 ↑ torch_npuPyTorch适配npu_flash_attention等接口开发FlashAttention变体时逻辑改ops-transformer分块策略改catlass的模板参数图优化改ge。编译和验证改完代码后的编译流程cd ops-transformer mkdir build cd build source /usr/local/Ascend/ascend-toolkit/set_env.sh cmake .. -DCANN_INSTALL_PATH/usr/local/Ascend/ascend-toolkit \ -DCMAKE_BUILD_TYPERelease # 只编译FlashAttention单个算子比全量编译快 make flash_attention -j8 # 验证编译产物 ls output/opkernel/libflash_attention*.so替换到torch_npu路径或指定自定义算子库路径后用数值验证脚本确认改动没有引入回归。克隆ops-transformer仓库按opkernel/flash_attention/→opplugin/flash_attention_op.cc的顺序读代码。tiling.cc的分块策略和score.cc的在线softmax是核心理解了这两块就抓住了FlashAttention昇腾实现的主线。https://atomgit.com/cann/ops-transformer
