昇腾CANN的算子“零件厂“：catlass仓库到底在生产什么

如果把昇腾NPU上的大模型算子比作一辆汽车FlashAttention是发动机RMSNorm是刹车片RoPE是方向盘——那catlass是什么是生产这些零件的模具和机床。第一次接触昇腾CANN生态的时候很容易忽略catlass。它不像ops-transformer那样直接提供FlashAttention这种成品算子也不像torch_npu那样可以直接调用。catlass藏在更底层提供的是算子开发的基础模板——分块矩阵乘、softmax归一化、数据搬运策略。FlashAttention里的在线softmax和分块计算底层都靠catlass的模板撑着。catlass不是CUTLASS这个名字容易让人误解。NVIDIA生态里有个著名的CUTLASSCUDA Templates for Linear Algebra Subroutines是英伟达的矩阵运算模板库。catlass的命名确实有致敬CUTLASS的意味但它跟CUTLASS没有代码层面的关系也不是CUTLASS的昇腾移植版。catlass是昇腾NPU原生的算子模板库专门针对达芬奇架构的硬件特性设计Unified Buffer的容量约束、Cube和Vector计算单元的流水编排、DMA搬运和计算的重叠。这些在CUDA上是另一套逻辑没法直接搬过来。仓库定位算子开发的基础设施昇腾CANN的算子生态分三层catlass在最底层应用算子层ops-transformerFlashAttention、RMSNorm等成品算子 ↑ 依赖 算子模板层catlass分块矩阵乘、reduce、softmax模板 ↑ 依赖 硬件抽象层Ascend C达芬奇架构的编程接口打个比方catlass提供砖头和水泥ops-transformer用这些材料盖房子。想在昇腾NPU上开发新的FlashAttention变体先得理解catlass怎么生产砖头。仓库结构模板在哪里git clone https://atomgit.com/cann/catlass.git cd catlass核心目录catlass/ ├── include/ │ └── catlass/ │ ├── gemm/ # 分块矩阵乘模板核心 │ │ ├── kernel/ │ │ │ ├── gemm_split_k.h # Split-K并行策略 │ │ │ └── gemm_batched.h # 批量矩阵乘 │ │ ├── thread/ │ │ │ └── mma.h # 矩阵乘累加 │ │ └── collective/ │ │ └── sm80_gemm.h # 分块搬运计算流水 │ ├── reduction/ # 归约操作模板 │ │ ├── thread/ │ │ │ └── reduce.h # 单块内reduce │ │ └── block/ │ │ └── reduce.h # 跨块reduce │ ├── softmax/ # softmax模板 │ │ └── online_softmax.h # 在线softmaxFlashAttention的基石 │ └── layout/ # 数据布局适配 │ └── layout.h # BNSD/BSND等格式转换 ├── examples/ # 使用示例 │ ├── flash_attention/ # 用catlass搭FlashAttention │ └── simple_gemm/ # 基础矩阵乘示例 └── tests/ # 单元测试online_softmax.h和gemm/是FlashAttention开发者最需要关注的两个模块。分块矩阵乘catlass的发动机FlashAttention的核心操作是Q×K^T和注意力权重×V本质上都是分块矩阵乘。catlass的GEMM模板把分块策略、数据搬运、计算流水全部封装好// catlass/include/catlass/gemm/kernel/gemm_split_k.h // 分块矩阵乘的简化示意伪代码 templatetypename ElementA, typename ElementB, typename ElementC struct GemmSplitK { // 分块参数决定一次算多大的块 struct TileShape { static const int kM 128; // M方向分块大小 static const int kN 128; // N方向分块大小 static const int kK 64; // K方向分块大小 }; // 整个GEMM分三阶段流水 void operator()(Params const params) { // 阶段1从GMEM加载A/B分块到L1/UB LoadTileFromGlobal(params.A, params.B, tile_a, tile_b); // 阶段2Cube单元做矩阵乘累加 // 阶段1和阶段2流水化加载第(i1)块的同时计算第i块 for (int k 0; k params.K / TileShape::kK; k) { LoadNextTile(...); // DMA搬运 ComputeCurrentTile(tile_a, tile_b, acc); // Cube计算 } // 阶段3结果从UB写回GMEM StoreTileToGlobal(params.C, acc); } };这里最关键的设计是搬运和计算的流水化。昇腾NPU有独立的DMA引擎和计算单元加载下一块数据的同时可以计算当前块两者并行。如果等加载完再算、算完再加载吞吐直接砍半。分块大小128×128×64不是随便定的。昇腾910的Cube单元一次能处理16×16的FP16矩阵乘128×128正好是8×8个Cube微操作UB装得下且对齐到128字节边界。在线softmaxFlashAttention的数学基石catlass里的online_softmax.h是在线softmax的模板实现FlashAttention的前向计算直接依赖它// catlass/include/catlass/softmax/online_softmax.h // 在线softmax模板伪代码示意流程 templatetypename T, int BlockSize struct OnlineSoftmax { struct State { T row_max; // 当前行最大值防溢出用 T row_sum; // 当前行指数和 T* acc_output; // 累加输出指针 }; // 每处理一个新块就调用一次Update static void Update(State state, T* new_scores, T* new_values, int block_len) { // 找新块的局部最大值 T local_max ReduceMax(new_scores, block_len); // 更新全局最大值 T new_max Max(state.row_max, local_max); // 关键重新缩放之前的累加结果 // 数学上等价于把所有分数放到同一个exp尺度下 T correction Exp(state.row_max - new_max); state.row_sum state.row_sum * correction; ScaleAccOutput(state.acc_output, correction); // 加上新块的贡献 // 减new_max防溢出FP16下exp(88.7)inf T* exp_scores ExpSub(new_scores, new_max, block_len); T local_sum ReduceSum(exp_scores, block_len); state.row_sum state.row_sum local_sum; AccumulateOutput(state.acc_output, exp_scores, new_values); state.row_max new_max; } // 所有块处理完后归一化 static void Finalize(State state, int seq_len) { ScaleAccOutput(state.acc_output, 1.0 / state.row_sum); } };correction Exp(state.row_max - new_max)这一行是在线softmax的精髓。标准softmax需要先扫一遍全局最大值再扫一遍算指数和。在线softmax把它变成增量的每来一个新块如果发现了更大的值就把之前所有结果按比例缩放回来。数学上完全等价但不需要存全局信息。FlashAttention的前向计算就是外层循环遍历Q分块内层循环遍历K/V分块每一步调catlass的OnlineSoftmax::Update更新累加结果。catlass跟ops-transformer的协作实际代码里ops-transformer的FlashAttention算子直接引用catlass的头文件ops-transformer/opkernel/flash_attention/ ├── flash_attention_score.cc │ #include catlass/gemm/kernel/gemm_split_k.h // 分块矩阵乘 │ #include catlass/softmax/online_softmax.h // 在线softmax │ #include catlass/reduction/block/reduce.h // 跨块reduce开发流程是这样的catlass定义OnlineSoftmax、GemmSplitK等模板ops-transformer在FlashAttention算子里实例化这些模板传入昇腾NPU的硬件参数UB大小、Cube规格等编译时模板展开生成针对910硬件优化的机器码这种分层设计的好处改FlashAttention的算法逻辑比如加个causal mask只改ops-transformer改分块策略或softmax的数值精度改catlass。互不干扰。从catlass搭建一个简化版FlashAttentioncatlass的examples目录里有一个简化版FlashAttention示例展示了怎么用模板搭出完整算子// catlass/examples/flash_attention/flash_attention_example.cc // 简化示意 templatetypename T void FlashAttentionForward( T* query, T* key, T* value, T* output, int batch, int heads, int seq_len, int dim, int block_size 128 ) { int num_blocks seq_len / block_size; for (int b 0; b batch; b) { for (int h 0; h heads; h) { for (int qi 0; qi num_blocks; qi) { // 取Q分块 T* q_block query offset(b, h, qi * block_size, 0); // 初始化在线softmax状态 typename OnlineSoftmaxT, 128::State state; state.row_max -1e9; state.row_sum 0; for (int ki 0; ki num_blocks; ki) { T* k_block key offset(b, h, ki * block_size, 0); T* v_block value offset(b, h, ki * block_size, 0); // 用catlass的GEMM模板算Q×K^T GemmSplitKT, T, T gemm; T local_scores[128 * 128]; gemm(q_block, k_block, local_scores, block_size, block_size, dim); // 缩放 Scale(local_scores, 1.0 / sqrt(dim)); // 用catlass的在线softmax更新 OnlineSoftmaxT, 128::Update(state, local_scores, v_block, block_size); } // 归一化输出 OnlineSoftmaxT, 128::Finalize(state, seq_len); // 写回output StoreOutput(output offset(b, h, qi * block_size, 0), state); } } } }这就是FlashAttention的骨架分块矩阵乘在线softmax。catlass把底层的搬运、对齐、流水化全部封装在模板里上层代码只需要关心算什么和怎么累加。跟其他仓库的边界catlass的职责边界需要搞清楚需求去哪个仓库调FlashAttention跑推理torch_npu改FlashAttention的算法逻辑ops-transformer改分块策略或softmax模板catlass写一个全新的算子catlass模板 ops-*注册调图融合优化gecatlass只负责怎么高效地在昇腾NPU上做基础运算不涉及具体的算子逻辑和图调度。##昇腾NPU上catlass的硬件适配catlass的模板设计有三个昇腾特有的约束UB容量约束Unified Buffer约256KB分块大小必须适配。catlass的TileShape默认128×128×64单个FP16分块约占32KB5个分块Q/K/V/Score/Output加中间结果约160KB留有余量128字节对齐昇腾NPU的DMA搬运要求数据起始地址128字节对齐。catlass的layout模块自动处理padding和对齐开发者不需要手动算偏移Cube和Vector的分工矩阵乘走Cube单元高吞吐标量运算和softmax走Vector单元。catlass的GEMM模板默认走CubeOnlineSoftmax走Vector两者通过UB传递数据克隆catlass仓库先看examples/flash_attention/的示例代码理解模板怎么搭出完整算子。然后读online_softmax.h的在线softmax实现这是FlashAttention的数学基石。https://atomgit.com/cann/catlass