上礼拜有个同学问我：“catlass 到底干嘛的？跟 ops-transformer 啥关系？”

我给他打了个比方：ops-transformer 是现成的成品菜，开袋即食；catlass 是菜谱加厨具，你自己照着做。成品菜方便但口味固定，菜谱灵活但得自己动手。

今天我们就动手——用 catlass 的算子模板，写一个跑在昇腾NPU上的 FlashAttention 算子。

环境准备

先确认你手上有这些东西：

• 一台搭载 Ascend 910 的服务器
• 昇腾CANN 8.0 及以上版本（CANN 8.0 对 FlashAttention 算子做了专项优化）
• Ascend C 开发环境已配好
• catlass 仓库代码拉到本地

# 拉代码
git clone https://atomgit.com/cann/catlass.git
cd catlass

⚠️ 踩坑预警：CANN 版本低于 8.0 的，catlass 里部分模板编译会报 undefined symbol。别折腾了，先升级 CANN。

第一步：理解 catlass 的模板结构

catlass 不是一堆算子实现，而是一套算子模板框架。它的目录结构大概长这样：

catlass/
├── include/
│ └── catlass/ # 核心模板头文件
│ ├── gemm/ # GEMM模板（最基础）
│ ├── epilogue/ # 尾巴操作模板（bias/relu等）
│ └── layout/ # 内存布局模板
├── examples/ # 示例
└── test/ # 测试用例

核心思想：GEMM是万物的底座。 FlashAttention 里最吃计算的部分——QK^T 和 softmax(QK^T)V——本质上都是矩阵乘法。catlass 把 GEMM 模板做好，你在上面拼装就行。

这跟 CUTLASS（NVIDIA 的算子模板库）思路一模一样。但 catlass 是给昇腾达芬奇架构定制的，内存层次、计算单元调度全不一样。

第二步：找到 FlashAttention 相关的模板

catlass 本身提供 GEMM 模板，FlashAttention 需要用 GEMM + 自定义 epilogue（把 softmax 融进去）。

#include "catlass/gemm/device/gemm_universal.h"
// 这就是catlass的核心入口——通用GEMM模板
// 不直接叫flash_attention，你得自己拼

拼装思路：

1. 第一段 GEMM：Q × K^T → 得到注意力分数矩阵
2. 自定义 epilogue：在线 softmax（分块归一化）
3. 第二段 GEMM：softmax 输出 × V → 得到最终注意力输出

关键点在于 epilogue 的自定义。catlass 的 GEMM 模板支持你塞自己的 epilogue functor，在矩阵乘法算完的瞬间立刻做后处理，不用把中间结果写回 HBM。

// 自定义epilogue：在线softmax + 缩放
struct FlashAttentionEpilogue {
 template<typename Element, int Rows, int Cols>
 __aicore__ inline void operator()(
 LocalTensor<Element>& output,
 LocalTensor<Element>& accumulator,
 float scale) {
 // accumulator里是QK^T的原始结果
 // 先缩放
 Muls(accumulator, accumulator, scale, Rows * Cols);
 // 在线softmax——这块是核心，分块归一化
 // 不展开，catlass/examples/里有参考实现
 // ...
 // 结果写进output
 }
};

⚠️ 踩坑预警： epilogue functor 里不要调用 DataCopy 往 HBM 写数据。epilogue 的设计意图就是在片上完成所有操作。你写了就打破了 FlashAttention 省显存的核心逻辑。

第三步：组装完整的 FlashAttention kernel

两个 GEMM 拼起来，中间用自定义 epilogue 衔接：

// 这不是完整的可编译代码，是结构示意
// 完整实现参考 catlass/examples/

// 第一段：QK^T
auto gemm_qk = GemmUniversal<
 ElementA, LayoutA, // Q的元素类型和布局
 ElementB, LayoutB, // K的元素类型和布局 
 ElementC, // 输出类型
 FlashAttentionEpilogue // 自定义尾巴
>();

// 第二段：softmax_output × V
auto gemm_sv = GemmUniversal<
 ElementA, LayoutA, // softmax结果
 ElementB, LayoutB, // V
 ElementC, // 最终输出
 DefaultEpilogue // 这段不需要特殊尾巴
>();

// 跑起来
gemm_qk(Q, K_t, attention_scores, scale);
gemm_sv(attention_scores, V, output, 1.0f);

两个 GEMM 之间的 attention_scores 存在哪？ 这就是 catlass 模板灵活的地方——你可以让它留在 L1 Buffer 里，不写回 HBM。这样显存占用就从 O(N²) 变成 O(N)。

但这需要你自己控制 L1 Buffer 的分配和生命周期。catlass 给你工具，不替你做决定。

第四步：编译和验证

# 用asc-devkit编译
cd catlass/examples/flash_attention
mkdir build && cd build
cmake .. -DCANN_INSTALL_PATH=/usr/local/Ascend/ascend-toolkit
make -j8

⚠️ 踩坑预警：cmake 找不到 CANN 路径的，手动 export 一下：

export ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest

编译完成后跑验证：

./flash_attention_example --seq_len=2048 --head_dim=128

看到输出 PASSED 就对了。

如果结果不对，先查 tile size。catlass 默认的 tile 配置是 128×128，适合大部分场景。但如果你的 head_dim 不是 128 的倍数，GEMM 模板会自动 padding，结果可能有精度差异。把 head_dim 对齐到 128 再跑一次试试。

catlass 和 ops-transformer 的关系

跑通之后你可能会问：那我直接用 ops-transformer 不就完了，干嘛要 catlass？

因为灵活性。ops-transformer 提供的 FlashAttention 是固定实现，参数和流程都封装好了。但如果你要做以下事情：

• 自定义注意力机制（比如 sliding window attention）
• 在 softmax 后面插一个自定义操作
• 针对 Ascend 910 的特定 workload 调整 tile size 和流水线策略

那就得用 catlass 从模板级别开始定制。写完之后，你的自定义算子可以直接被 ops-transformer 或 ATB 调用。

一句话说就是：ops-transformer 管能用，catlass 管好用。

下一步可以试的

1. 把上面的 FlashAttention kernel 改成 FlashAttention-2（加上并行化策略）
2. 用 catlass 的 epilogue 模板实现 sliding window attention
3. 对比你写的 kernel 和 ops-transformer 内置实现的性能差异

代码和模板都在这：

https://atomgit.com/cann/catlass

意外收获：catlass 的 GEMM 模板不止能写 attention。MoE 里的 expert routing、BLAS 里的大规模矩阵乘，底层全是 GEMM。搞懂 catlass 的模板机制，基本上昇腾NPU上的高性能算子你都能写。