请添加图片描述
个人主页：ujianu

文章目录

前言

为什么同样一个矩阵乘法，在昇腾NPU上有的跑得飞快，有的却慢得像在老牛拉车？

答案往往不在算法本身，而在你用没用对模板。

仓库定位：catlass到底是什么？

看到"catlass"这个名字，估计不少人心里犯嘀咕：这不就是CUTLASS的昇腾版吗？

大错特错。

catlass 和 CUTLASS 是两个完全不同的东西。CUTLASS 是 NVIDIA GPU 上的矩阵乘模板库，而 catlass 是专门为昇腾 CANN 设计的高性能矩阵乘算子模板库——从根子上就是两条路。特别注意：catlass 的拼写是 c-a-t-l-a-s-s，绝不能写成 CUTLASS，这是两个完全独立的项目。

那 catlass 究竟是干什么的？

说白了，它是一个算子模板工厂。你可能听过这样一个说法：深度学习底层就是矩阵乘法——MatMul、Convolution、Attention，说到底都是大矩阵和小矩阵的各种排列组合。但怎么把这个"排列组合"写得快、写得省显存、写得适配昇腾 NPU 的硬件特性，就是另一回事了。

catlass 干的就是这件事：它把高性能矩阵乘算子的实现拆成一套可拆卸、可组装、可微调的模板。你不需要从零开始写一个算子，只需要找到合适的模板，往里填参数，就能编译出一个能在昇腾 NPU 上跑得飞起的矩阵乘算子。

设计哲学：分层抽象→白盒化组装→硬件特化

catlass 的设计哲学可以概括为三个递进的层次，每一层都在解决不同的问题：

分层抽象：将复杂的矩阵乘算子拆解为可独立设计和优化的层次结构。开发者可以只关注某一层的优化，而不必理解整个算子的实现细节。

白盒化组装：每个模板的每个部件都向开发者敞开。你可以复用已有部件、替换某个部件、甚至只修改部件的某一行代码，实现真正的"所见即所得"。

硬件特化：针对昇腾 NPU 的不同硬件型号（如 Ascend 910、Ascend 950）提供差异化的模板特化版本。同一套模板代码，换个配置就能适配不同硬件。

这就是 catlass 的核心定位：昇腾上的矩阵乘模板库，提供 TLA/MLA/FlashAttention 等高性能模板，支持分层抽象设计与硬件特化。

仓库地址：https://atomgit.com/cann/catlass

核心能力：三层递进的设计哲学

要理解 catlass，光知道它是"模板工厂"不够。你得理解它的设计哲学——三层递进：分层抽象设计→白盒化组装→硬件特化。

第一层：分层抽象设计

想象一下乐高积木。最基础的颗粒度是一个个凸起和凹槽，但你不会用单个凸起来拼乐高，你会用预制好的小块——门、窗、楼梯、屋顶。catlass 也是这么干的。

它把矩阵乘算子拆分成了多个抽象层次：

  • Epilogue 层（收尾层）：后处理操作，如激活函数、偏置加法、残差连接等
  • Kernel 层（计算核层）：核心矩阵乘计算逻辑，包括累加器、同步机制、流水线调度
  • Tensor Operator 层（算子层）：完整的算子接口，支持多核并行、负载均衡
  • Device 层（设备层）：设备级调度、内存管理、流式执行

每一层都可以单独设计和优化。这意味着什么？你可以只替换其中一层，而不用重写整个算子。比如，你的模型需要在矩阵乘后接一个 ReLU 激活，传统做法是写两个 kernel 再串起来；但在 catlass 里，你只需要在 Epilogue 层加一个 ReLU 模块，编译器会自动把它融合到矩阵乘 kernel 里，省去一次显存读写。

// Epilogue 层：融合激活函数示例
template <typename ElementOutput, typename ElementAccumulator>
struct EpilogueReLU {
    // 将累加结果写入输出前，先过一遍激活函数
    __device__ void operator()(ElementOutput* output, 
                               ElementAccumulator accumulator) {
        ElementAccumulator val = accumulator > 0 ? accumulator : 0;
        *output = static_cast<ElementOutput>(val);
    }
};

// 在模板中组合使用
using GemmWithReLU = GemmTemplate<
    float, float,                          // 输入输出类型
    EpilogueReLU<float, float>,           // 收尾层：ReLU激活
    Ascend910                             // 目标硬件
>;

第二层：白盒化组装

这是 catlass 最独特的地方。很多模板库是"黑盒"——你调用它，它给你返回一个结果，中间怎么实现的你管不着。catlass 不是，它是"白盒"。

白盒的意思是：每个模板的每个部件都对你敞开。你可以复用已有的部件，可以替换某个部件，甚至可以只修改部件的某一行代码。这种设计让算子开发者能够基于社区已有的优化成果，快速定制自己的高性能算子。

TLA 模板：张量层级注意力

TLA（Tensor-Layer Attention）模板是 catlass 为大模型推理优化的核心模板之一。它通过精细的张量分块策略，将 Attention 计算拆解为适合昇腾 NPU 硬件结构的小块，利用 Cube 单元的高吞吐量和 SRAM 的低延迟，实现显存访问量的大幅降低。

TLA 的核心思想是：将 Q、K、V 矩阵按头维度和序列维度分块，每次只加载一小块到 SRAM 中计算，计算完成后再加载下一块。这种"流水线化"的策略，使得大序列长度下的 Attention 计算不再受限于 HBM 带宽。

// TLA (Tensor-Layer Attention) 模板配置
template <typename Element>
struct TLAConfig {
    // 分块策略：Q/K/V 各自的块大小
    static constexpr int kBlockM = 64;    // Q 序列维度块大小
    static constexpr int kBlockN = 128;   // K/V 序列维度块大小
    static constexpr int kBlockK = 64;    // 头维度块大小
    
    // 数据布局：决定内存访问模式
    static constexpr Layout kLayoutQ = Layout::kRowMajor;
    static constexpr Layout kLayoutK = Layout::kColumnMajor;
    static constexpr Layout kLayoutV = Layout::kRowMajor;
    
    // 流水线策略
    static constexpr int kStageCount = 2;  // 双缓冲
    static constexpr bool kEnableSwizzle = true;  // 地址交错优化
};

// 实例化一个具体的 TLA 算子
using TLA_FP16 = TLAGemmTemplate<
    half,                                  // 计算精度
    TLAConfig<half>,                      // 配置参数
    Ascend910                             // 目标硬件
>;

TLA 模板适用于：大模型推理、长序列 Attention、需要融合 Softmax 的场景。

MLA 模板：多层注意力

MLA（Multi-Layer Attention）模板针对多头注意力机制进行了深度优化。与 TLA 的区别在于，MLA 更关注多个 Attention 头之间的并行计算效率，通过共享 K、V 矩阵的加载策略，减少重复的显存访问。

在多头注意力中，Q 矩阵通常按头维度拆分，但 K、V 矩阵对所有头是共享的。MLA 模板利用这一特性，将 K、V 的加载与多个头的计算流水线化：加载一次 K、V 块，计算多个头的 Q·K^T 和 Softmax，然后再加载下一块。这种策略在头数较多时（如 32 头、64 头）效果尤为明显。

// MLA (Multi-Layer Attention) 模板示意
template <typename Config>
class MLAAttention {
public:
    void compute(const Tensor& q, const Tensor& k, 
                 const Tensor& v, Tensor* output) {
        // 核心计算逻辑：多块并行 + 头间复用 K/V
        for (int tile_k = 0; tile_k < k_seq_tiles; ++tile_k) {
            // 加载一块 K/V（所有头共享）
            load_kv_tile(k, v, tile_k);
            
            // 多个头并行计算
            #pragma unroll
            for (int head = 0; head < num_heads; ++head) {
                compute_qk_attention(q, head, tile_k);
                apply_softmax(head);
                accumulate_output(output, head);
            }
            synchronize();
        }
    }
private:
    Config config_;
    TileCache cache_;  // SRAM 中的缓存
};

MLA 模板适用于：多头注意力计算、头数较多（≥16）的场景、需要最大化 Cube 利用率的任务。

FlashAttention 模板：极致显存优化

FlashAttention 是近年来深度学习领域最著名的优化之一。catlass 提供的 FlashAttention 模板，将这一优化思想完整移植到昇腾 NPU 上。

FlashAttention 的核心洞察是：Attention 计算的显存瓶颈在于 softmax 的归一化因子。传统实现需要把整个 Q·K^T 矩阵存到显存，再计算 softmax，这导致 O(N²) 的显存占用。FlashAttention 通过在线算法（online algorithm），将 softmax 的计算拆分为多个块，每块计算完立即累加到输出，最终只需要 O(N) 的显存。

catlass 的 FlashAttention 模板进一步针对昇腾 NPU 的硬件特性进行了优化：利用 Cube 单元的矩阵乘加能力，将 softmax 的指数运算和归一化操作向量化，利用 Vector 单元并行处理；同时，通过双缓冲策略，将计算与数据加载完全重叠。

// FlashAttention 模板配置
template <typename Element, typename Arch>
struct FlashAttentionConfig {
    // 分块大小：需要根据硬件 SRAM 容量调整
    static constexpr int kBlockM = 128;
    static constexpr int kBlockN = 64;
    
    // Softmax 计算策略
    static constexpr SoftmaxMode kSoftmaxMode = SoftmaxMode::kOnline;
    static constexpr float kSoftmaxScale = 1.0f / sqrtf(head_dim);
    
    // 硬件特化参数
    static constexpr size_t kSRAMSize = HardwareSpecialize<Arch>::SRAMCapacity;
    static constexpr bool kUseVectorUnit = true;
};

// 在线 Softmax 计算
template <typename Element>
__device__ void OnlineSoftmax(
    Element* output,           // 当前块的输出
    Element* accumulator,      // 全局累加器
    float* norm_factor,        // 归一化因子
    const Element* qk_block,   // Q·K^T 块
    int block_size
) {
    // 计算当前块的 max 和 sum
    float local_max = -INFINITY;
    float local_sum = 0.0f;
    
    for (int i = 0; i < block_size; ++i) {
        local_max = max(local_max, qk_block[i]);
    }
    for (int i = 0; i < block_size; ++i) {
        local_sum += expf(qk_block[i] - local_max);
    }
    
    // 更新全局归一化因子
    float old_norm = *norm_factor;
    float new_max = max(old_norm, local_max);
    *norm_factor = new_max;
    
    // 在线更新累加器
    float rescale_old = expf(old_norm - new_max);
    float rescale_new = expf(local_max - new_max);
    *accumulator = *accumulator * rescale_old + output * rescale_new * local_sum;
}

FlashAttention 模板适用于：超大序列长度（≥4096）的 Attention 计算、显存受限的场景、训练和推理均可。

第三层：硬件特化

到这里，设计还没完。昇腾 NPU 和其他芯片的硬件特性差异巨大——Cube 单元、Vector 单元、SRAM vs HBM 的访存延迟、数据搬运的方式…这些都是模板必须考虑的因素。

catlass 的第三层就是解决这个问题的。它为不同的昇腾硬件型号提供了差异化的模板特化版本。也就是说，同一套模板代码，换个配置就能适配 Ascend 910、Ascend 950 等不同型号。

// 硬件特化示例：针对 Ascend 910 的配置
template <>
struct HardwareSpecialize<Ascend910> {
    static constexpr int CubeUnits = 32;        // Cube 单元数量
    static constexpr int VectorUnits = 16;      // Vector 单元数量
    static constexpr size_t SRAMCapacity = 4 * 1024 * 1024;  // 4MB
    static constexpr bool EnableDoubleBuffer = true;
    static constexpr int MaxTileSize = 256;     // 最大分块尺寸
};

// 针对 Ascend 950 的配置
template <>
struct HardwareSpecialize<Ascend950> {
    static constexpr int CubeUnits = 64;        // 更多 Cube 单元
    static constexpr int VectorUnits = 32;      // 更多 Vector 单元
    static constexpr size_t SRAMCapacity = 8 * 1024 * 1024;  // 8MB
    static constexpr bool EnableDoubleBuffer = true;
    static constexpr int MaxTileSize = 512;     // 更大的分块尺寸
};

// 硬件自适应的分块策略
template <typename Arch>
constexpr int ComputeOptimalTileSize() {
    constexpr size_t sram = HardwareSpecialize<Arch>::SRAMCapacity;
    constexpr int cube = HardwareSpecialize<Arch>::CubeUnits;
    // 根据 SRAM 容量和计算单元数量，自动推导最优分块大小
    return min(sram / (sizeof(float) * 3), cube * 8);
}

硬件特化的好处是：开发者不需要关心硬件细节，模板会自动选择最优配置。但如果你是"硬核玩家"，也可以手动覆盖这些配置，榨干硬件的每一滴性能。

在架构中的位置：算子层的承上启下

理解了核心能力，再看 catlass 在 CANN 五层架构中的位置。

第1层：昇腾计算语言层 AscendCL
  ├─ 应用开发接口（推理/预处理/单算子）
  ├─ 图开发接口（统一构图/多框架支持）
  └─ 算子开发接口 Ascend C

第2层：昇腾计算服务层  ←← catlass 在这里
  ├─ AOL 算子库（NN/BLAS/DVPP/AIPP/HCCL/融合算子）
  ├─ AOE 调优引擎（OPAT/SGAT/GDAT/AMCT）
  └─ Framework Adaptor 框架适配器

第3层：昇腾计算编译层
  ├─ Graph Compiler 图编译器
  └─ BiSheng / ATC 编译器
...

catlass 位于第二层——昇腾计算服务层的 AOL 算子库。这位置有意思了：它是算子层的"基础设施"，上面被 ascend-boost-comm（算子公共平台）和 opbase（算子基础组件）支撑，下面服务着 ops-transformer、ops-nn、ops-blas 这些具体算子仓库。

换句话说，catlass 是算子的算子——它不直接提供能被业务代码调用的算子，而是提供给算子开发者一套编写高性能算子的模板。开发者使用 Ascend C 语言编写算子时，可以直接调用 catlass 提供的模板，快速实现高性能实现。

与其他仓库的关系

既然是"算子的算子"，catlass 和上下游仓库的关系就清楚了：

上游依赖：

  • opbase：算子基础组件，提供公共的数据结构、工具函数、错误处理等
  • ascend-boost-comm：算子公共平台，提供 M×N 算子复用的能力

下游调用：

  • ops-transformer：大模型进阶算子库，内部使用 catlass 的 TLA/MLA/FlashAttention 模板实现 Attention、FFN 等算子
  • ops-nn：神经网络基础算子库，使用 catlass 的矩阵乘模板优化 Convolution、Linear 等算子性能
  • ops-blas：线性代数基础算子库，提供轻量化 GEMM 调用

这里有一个容易混淆的点：catlass 和 ops-blas 的 GEMM 是什么关系？

简单说：互补。ops-blas 提供的是"开箱即用"的 GEMM 算子，适合通用场景；catlass 提供的是"定制高性能"的模板，适合需要极致优化的场景。当你的模型对矩阵乘性能极其敏感（比如大模型的 Attention 层），用 catlass 模板手写一个特化版本，往往能比直接调用 ops-blas 获得更高的计算效率。

# 实际使用场景对比

# 场景1：普通矩阵乘，直接用 ops-blas
import ops_blas
output = ops_blas.matmul(a, b)  # 简单快捷

# 场景2：Attention 计算，需要极致性能，用 catlass 模板
from catlass import FlashAttentionTemplate

# 配置模板参数
config = FlashAttentionConfig(
    block_m=128,
    block_n=64,
    enable_fusion=True  # 融合 Softmax
)

# 编译生成特化算子
flash_attn_op = FlashAttentionTemplate(config, arch="Ascend910")
output = flash_attn_op(q, k, v)  # 高性能执行

// ops-transformer 中使用 catlass 模板的示例
#include "catlass/flash_attention.h"

namespace ops_transformer {

class MultiHeadAttention {
public:
    MultiHeadAttention(int num_heads, int head_dim) {
        // 使用 catlass 的 FlashAttention 模板
        using FlashAttn = catlass::FlashAttention<
            half,                              // FP16 计算精度
            Ascend910,                        // 目标硬件
            catlass::ConfigOptimal            // 自动最优配置
        >;
        flash_attn_ = std::make_unique<FlashAttn>();
    }
    
    Tensor forward(const Tensor& q, const Tensor& k, const Tensor& v) {
        return flash_attn_->compute(q, k, v);
    }
    
private:
    std::unique_ptr<catlass::FlashAttention<half, Ascend910, catlass::ConfigOptimal>> flash_attn_;
};

}  // namespace ops_transformer

如何使用 catlass？

catlass 的使用流程可以概括为四个步骤：

步骤1：选择模板。根据你的算子类型，选择合适的模板（TLA、MLA、FlashAttention 或基础 GEMM 模板）。

步骤2：配置参数。设置分块大小、数据布局、融合策略等参数。catlass 提供了默认配置，也支持手动覆盖。

步骤3：编译算子。使用 Ascend C 编译器，将模板编译为可在昇腾 NPU 上执行的算子二进制。

步骤4：集成调用。将编译好的算子集成到你的模型或框架中，通过 AscendCL 接口调用。

# 步骤1：克隆 catlass 仓库
git clone https://atomgit.com/cann/catlass.git
cd catlass

# 步骤2：编译模板（以 FlashAttention 为例）
python build_template.py \
    --template flash_attention \
    --arch Ascend910 \
    --precision fp16 \
    --output ./build/flash_attn_ascend910.so

# 步骤3：在 Python 中调用编译好的算子
import torch
import custom_ops  # 编译生成的自定义算子库

q = torch.randn(1, 32, 4096, 128, device='npu')  # [batch, heads, seq, dim]
k = torch.randn(1, 32, 4096, 128, device='npu')
v = torch.randn(1, 32, 4096, 128, device='npu')

output = custom_ops.flash_attention(q, k, v)

总结：什么时候该用 catlass？

花了这么多篇幅拆解设计哲学，是时候回到最初的问题了：什么时候该用 catlass？

三条标准：

  1. 你要写自定义矩阵乘算子 —— 现有算子库满足不了你的需求，需要自己动手
  2. 现有算子性能不够 —— 通用 GEMM 能跑，但喂不饱你的模型对性能的饥渴
  3. 你想复用已有的优化经验 —— 社区已经验证过的高性能模板，为什么还要自己从头踩坑？

如果不满足这三条，老老实实用 ops-nn 或 ops-blas 就够了。模板是用来超越的，不是用来交学费的。

去 https://atomgit.com/cann/catlass 看看模板源码，自己试着组装一个能跑通的矩阵乘算子，才是真正的开场。

# 矩阵 # 线性代数
Logo
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

更多推荐

cover

昇腾CANN大模型推理加速：ATB加速库的Transformer层融合优化

avatar 鲲鹏昇腾开发者社区
cover

昇腾CANN elec-ops-inspection 仓：电力巡检AI算子实战

avatar 鲲鹏昇腾开发者社区
cover

昇腾CANN mat-chem-sim-pred 仓：材料化学AI模拟与预测实战

avatar 鲲鹏昇腾开发者社区