第一次接触昇腾CANN生态的人，面对几十个仓库很容易懵：哪个仓库放算子？哪个放驱动？哪个放编译器？FlashAttention这种热门算子又该去哪找？

答案很直接：ops-transformer。这个仓库是昇腾CANN大模型算子的主阵地，FlashAttention、RMSNorm、RoPE、SwiGLU——大模型推理训练用到的核心算子，全在这里。

仓库定位：算子层，不是框架层

昇腾CANN的软件栈从下到上分四层：硬件→驱动/CANN→算子→框架。ops-transformer处在算子层，它不负责模型调度和图编译，只负责把单个算子做到又快又省。

跟其他几个容易混淆的仓库做个区分：

仓库	层级	干什么
ops-transformer	算子层	大模型算子（FlashAttention、RoPE等）
catlass	算子模板层	算子开发模板，ops-*的底层依赖
ge	图引擎层	算子编排、图优化、算子融合
torch_npu	框架适配层	PyTorch→昇腾NPU的桥接

一句话：要改FlashAttention的算子实现，看ops-transformer；要写新算子，看catlass；要调图融合，看ge；要在PyTorch里调FlashAttention，装torch_npu。

仓库结构：FlashAttention藏在哪

克隆下来看目录结构：

git clone https://atomgit.com/cann/ops-transformer.git
cd ops-transformer

顶层目录：

ops-transformer/
├── opkernel/ # 算子内核实现（核心代码在这）
│ ├── flash_attention/ # FlashAttention算子
│ │ ├── flash_attention_score.cc # 前向核心逻辑
│ │ ├── flash_attention_score_grad.cc # 反向
│ │ └── flash_attention_score_tiling.cc # 分块策略
│ ├── rms_norm/ # RMSNorm算子
│ ├── rope/ # RoPE位置编码
│ └── ...
├── opplugin/ # 算子注册与接口层
│ └── flash_attention/
│ └── flash_attention_op.cc # GE算子注册
├── inc/ # 公共头文件
├── scripts/ # 编译脚本
└── cmake/ # 构建配置

FlashAttention的核心实现就在opkernel/flash_attention/下面，三个文件各有分工：

• flash_attention_score.cc：前向计算，分块策略+在线softmax的实现
• flash_attention_score_grad.cc：反向传播，重计算逻辑在这里
• flash_attention_score_tiling.cc：分块参数计算，决定每次算多大的块

核心逻辑：分块计算是怎么实现的

FlashAttention的精髓在分块（tiling）。昇腾NPU的Unified Buffer大约256KB，分块大小必须适配这个限制。tiling.cc做的事就是根据序列长度、头数、维度，算出最优的分块参数：

// ops-transformer/opkernel/flash_attention/flash_attention_score_tiling.cc
// 简化的分块策略示意（非原始代码）

uint32_t CalcTilingSize(uint32_t seq_len, uint32_t head_dim, uint32_t dtype_size) {
 // Unified Buffer约256KB，预留空间给输入输出和中间结果
 const uint32_t ub_size = 256 * 1024;
 
 // 每个分块需要的UB空间：
 // Q块 + K块 + V块 + 输出块 + softmax中间结果
 uint32_t block_elements = ub_size / (5 * dtype_size);
 
 // 分块的序列维度大小，向下对齐到128字节的倍数（昇腾NPU对齐要求）
 uint32_t block_seq = (block_elements / head_dim / 16) * 16;
 block_seq = std::min(block_seq, seq_len);
 
 return block_seq;
}

分块大小不是随便定的。太小了循环次数多，算子调度开销大；太大了UB装不下，直接溢出。tiling.cc的职责就是在UB容量和循环开销之间找平衡点。

前向核心：在线softmax的昇腾实现

flash_attention_score.cc里的在线softmax实现，核心思路跟论文一致，但昇腾NPU上有几个特殊处理：

// flash_attention_score.cc 核心循环（伪代码，示意流程）
for (uint32_t i = 0; i < num_q_blocks; i++) {
 // 取一块Q
 LoadQBlock(q_ptr, i * block_size, q_block);
 
 // 初始化在线softmax累加器
 row_max = NEG_INF; // 当前行最大值
 row_sum = 0.0; // 当前行指数和
 acc_out = 0.0; // 累加输出
 
 for (uint32_t j = 0; j < num_kv_blocks; j++) {
 LoadKVBlock(k_ptr, v_ptr, j * block_size, k_block, v_block);
 
 // 局部注意力分数
 local_scores = MatMul(q_block, k_block.T) * scale;
 
 // causal mask处理：昇腾NPU上用对角线偏移跳过
 // 不是先算出完整下三角矩阵再mask，而是直接跳过无效块
 if (IsCausalBlockSkipped(i, j, block_size)) {
 continue; // 跳过因果mask之外的块
 }
 
 // 在线softmax更新
 local_max = ReduceMax(local_scores);
 new_max = Max(row_max, local_max);
 
 // 重新缩放之前的累加结果
 correction = Exp(row_max - new_max);
 row_sum = row_sum * correction;
 acc_out = acc_out * correction;
 
 // 加上当前块
 local_weights = Exp(local_scores - new_max);
 local_sum = ReduceSum(local_weights);
 row_sum = row_sum + local_sum;
 acc_out = acc_out + MatMul(local_weights, v_block);
 
 row_max = new_max;
 }
 
 // 最终归一化
 out_block = acc_out / row_sum;
 StoreOutput(out_ptr, i * block_size, out_block);
}

昇腾NPU上的两个特殊处理：

1. causal mask用块级跳过，不是算出完整下三角矩阵再乘mask。FlashAttention分块计算天然适合块级跳过——如果整块都在mask之外，直接continue跳过，省掉整块计算量
2. 数据搬运和计算流水化。Q块加载完后，K/V块的加载和当前块的计算可以overlap，掩盖搬运延迟。这是昇腾达芬奇架构的特性——向量计算单元和DMA搬运引擎可以并行工作

算子注册：GE怎么找到FlashAttention

opplugin/flash_attention_op.cc负责把FlashAttention注册到GE图引擎里，让框架层能通过torch_npu.npu_flash_attention()调到它：

// 算子注册（伪代码，示意流程）
IMPLEMT_INFERFUNC(FlashAttentionScore, FlashAttentionScoreInfer) {
 // 输出shape推导：跟输入Q的shape一致
 auto q_shape = op.GetInputDesc("q").GetShape();
 auto out_shape = q_shape;
 op.UpdateOutputDesc("output", TensorDesc(out_shape, ...));
 return GRAPH_SUCCESS;
}

IMPLEMT_COMMON_INFERFUNC(FlashAttentionScoreInfer) {
 return FlashAttentionScoreInfer;
}

// 注册到GE算子库
REGISTER_CUSTOM_OP("FlashAttentionScore")
 .FrameworkType("onnxff")
 .OpImplPath("opplugin/flash_attention/")
 .InferShapeFn(FlashAttentionScoreInfer);

注册完后，torch_npu调用npu_flash_attention时，GE图引擎会根据算子名找到这个注册，把计算委托给opkernel/flash_attention/下的实现。

跟catlass的关系

catlass是昇腾的算子模板库，提供矩阵乘法、reduce、softmax等基础操作的模板化实现。ops-transformer里的FlashAttention算子，部分基础操作（比如分块矩阵乘、softmax reduce）会依赖catlass提供的模板。

catlass（算子模板）
 ↑ 依赖
ops-transformer（大模型算子）
 ↑ 依赖
torch_npu（框架适配）

开发新算子时：catlass提供"积木"，ops-transformer负责"搭房子"。FlashAttention就是用catlass的分块矩阵乘和reduce模板，搭出了完整的在线softmax+分块attention逻辑。

怎么编译和验证

改了FlashAttention的代码，需要重新编译才能生效：

cd ops-transformer
mkdir build && cd build

# CANN环境必须先source
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# CMake配置
cmake .. -DCANN_INSTALL_PATH=/usr/local/Ascend/ascend-toolkit \
 -DCMAKE_BUILD_TYPE=Release

# 编译FlashAttention算子（只编单个算子比全量编译快很多）
make flash_attention -j8

# 产物在output/目录下
ls output/opkernel/libflash_attention.so

编译完后替换torch_npu对应的so文件，或者在GE环境变量里指定自定义算子路径：

# 指定自定义算子库路径
export ASCEND_CUSTOM_OPP_PATH=/path/to/ops-transformer/output

验证改动是否生效，用前文的数值验证脚本对比FlashAttention输出和标准attention的差异，确认没有回归。

仓库全景图

把ops-transformer在昇腾CANN生态里的位置画清楚：

昇腾NPU硬件
 ↑
CANN驱动+运行时（Ascend 910固件、runtime、hccl）
 ↑
算子层 ──→ ops-transformer（大模型算子）
 │ catlass（算子模板）
 │ ops-ascendc（通用算子）
 ↑
图引擎 ──→ ge（图编排、算子融合）
 ↑
框架适配 ──→ torch_npu / msmodelzoo
 ↑
应用层 ──→ LLM推理/训练服务

ops-transformer处在算子层的"大模型"赛道，专门服务Transformer架构的算子需求。通用算子（卷积、池化等）在ops-ascendc，不在这里。

---

 克隆ops-transformer仓库，按opkernel/flash_attention/→opplugin/flash_attention_op.cc的顺序读代码，重点关注tiling.cc的分块策略和score.cc的在线softmax实现。如果想跑通编译流程，先确认CANN环境和cmake版本，再按文中的步骤只编译FlashAttention单个算子。

https://atomgit.com/cann/ops-transformer