前言

昇腾NPU作为华为面向AI计算的核心硬件平台，其矩阵乘法算子开发长期面临定制难度高、开发周期长的困境。CANN软件栈在此背景下推出CATLASS算子模板库，将通用矩阵乘法（GEMM）计算逻辑进行分层抽象与模板化，使得算子代码可复用、可替换、可局部修改。CATLASS的核心价值在于：开发者不再需要从零手写硬件亲和的底层流水编排，而是通过组装预置模板组件，快速构建针对特定shape和精度需求的高性能算子实现。在定制shape场景下，CATLASS实现能达到标杆性能的零点九八到一点二倍，这一数据已经过大量实测验证。从社区版本迭代来看，v1.0.0到v1.5.0之间，CATLASS逐步覆盖了基础Matmul、Batched Matmul、Grouped Matmul、量化Matmul、Flash Attention推理、StreamK Matmul、稀疏Matmul、W4A4量化等众多算子形态，同时新增了对下一代昇腾硬件Ascend 950PR和Ascend 950DT的支持。这些演进并非简单的功能堆叠，而是模板库分层设计理念在更多实际场景中的持续验证与拓展。CANN生态通过CATLASS这种模板化工程方案，把昇腾NPU硬件能力转化为开发者可便捷调用的算子组装能力，缩短了从算法设计到算子落地的距离。

CATLASS的分层架构设计哲学

CATLASS的全称是CANN Templates for Linear Algebra Subroutines，中文名为昇腾算子模板库。其设计出发点是对GEMM类算子的计算过程进行结构化拆解。在Transformer架构中，矩阵乘法计算占据重要比重，不同场景和不同优化点的实现变种众多，算法演进过程中会诞生大量新的定制化开发诉求，难以事先枚举覆盖。直接基于硬件能力定制开发GEMM类算子时，开发者面临的最大痛点是：算子代码与硬件细节高度耦合，修改一处牵动全局，不同优化策略之间无法共享代码逻辑。

CATLASS采用分层模块化设计解决这一痛点。整个GEMM计算被解耦为多个层级：Device层负责Host侧适配与Kernel调度，Kernel层负责Block间循环与后处理组合，Block层负责单个Process内的主循环计算，Tile层负责硬件指令级的分片搬运与MMAD运算。每个层级的模板组件都可以独立替换、局部修改，而不影响其他层级的逻辑。

这种分层抽象并非简单地把代码拆成多个文件。它的设计依据来自昇腾NPU硬件架构本身的层级特性。昇腾NPU的AI Core内部存在多级缓冲区（全局内存GM、L1 Buffer、L0A/L0B/L0C Buffer），数据从GM搬运到L1再搬运到L0进行MMAD计算，这一搬运路径天然对应了Block层和Tile层的职责划分。Block层管理L1 Buffer上的数据分块策略，Tile层管理L0 Buffer上的硬件指令调用。硬件流水线（MTE1、MTE2、FixPipe等）之间的数据依赖和同步操作，也在Block层和Tile层得到对应处理。

模板化方式提取各层共性逻辑的同时，保留了必要的差异化扩展能力。算法框架中的特定步骤延迟到子类实现，子类能够在不改变算法整体结构的情况下灵活重定义关键步骤。这种模板模式（Template Method Pattern）在C++模板元编程中的体现，正是CATLASS命名的由来——它不是一套固定算子集合，而是一套可组装、可特化的模板框架。

分层架构带来的工程收益体现在多个方面。代码复用率大幅提高——同一个TileMmad实现可被不同DispatchPolicy的BlockMmad组合使用，同一个BlockMmad可被不同Epilogue的Kernel组合使用。修改隔离性得到保证——调整L1分块策略只需修改TileShape参数，不会影响Kernel层和Device层的逻辑。跨平台适配成本降低——新硬件平台接入时只需扩展Tile层的特化实现，上层逻辑保持共享。

GEMM三层嵌套循环与CATLASS组件映射

矩阵乘法的经典算法可以用三层嵌套循环描述。外层两重循环对应Block维度的M和N方向分块，这两重循环在昇腾NPU上通过多个AI Core的并行执行来实现——代码上并不显式表达为for循环，而是通过BlockIdx区分不同核心处理的数据块。内层循环对应K方向上的tile迭代，每一次迭代包含从GM到L1的数据搬运、从L1到L0的数据搬运、以及L0上的MMAD计算。

CATLASS将这三层嵌套循环映射到四层API组件。Device层对应整个算子的Host侧入口，Kernel层对应Block间循环编排，Block层对应K方向主循环，Tile层对应单次MMAD指令调用。用户组装内核时遵循严格的实例化顺序：创建特化的Block层MMAD组件、指定Block层后处理类型、指定数据走位方式（BlockScheduler）、在Kernel层将MMAD和后处理组合、将Kernel放入Device适配器。

// 第一步：创建所需的特化Block层mmad参数
using DispatchPolicy = Gemm::MmadAtlasA2Pingpong<true>;
using L1TileShape = GemmShape<128, 256, 256>;
using L0TileShape = GemmShape<128, 256, 64>;
using AType = Gemm::GemmType<ElementA, LayoutA>;
using BType = Gemm::GemmType<ElementB, LayoutB>;
using CType = Gemm::GemmType<ElementC, LayoutC>;

using BlockMmad = Gemm::Block::BlockMmad<DispatchPolicy,
    L1TileShape, L0TileShape, AType, BType, CType>;

// 第二步：指定Block层的后处理类型
using BlockEpilogue = void;

// 第三步：指定计算时的数据走位方式
using BlockScheduler = typename Gemm::Block::GemmIdentityBlockSwizzle<>;

// 第四步：在kernel层将mmad和后处理组合
using MatmulKernel = Gemm::Kernel::BasicMatmul<
    BlockMmad, BlockEpilogue, BlockScheduler>;

// 第五步：将kernel放入device适配器
using Matmul = Catlass::Gemm::Device::DeviceGemm<MatmulKernel>;

这五步实例化顺序严格对应了从底层硬件计算单元到上层Host调用的组装链路。BlockMmad是计算核心，必须先确定其调度策略和TileShape；BlockEpilogue可选配置后处理逻辑；BlockScheduler决定多核间数据分块映射；Kernel层将Block计算与后处理组合为完整内核逻辑；Device层封装Host侧参数解析和Kernel启动。这种自底向上的组装顺序确保每一层的参数选择都有下层约束作为依据，避免参数不匹配导致的编译或运行错误。

DeviceGemm作为最外层适配器，承担了Host侧Tiling参数解析、GM地址管理、Kernel启动参数构造等职责。开发者在使用DeviceGemm时，只需传入矩阵A、B、C的GM地址和Tiling描述，DeviceGemm自动完成Kernel参数打包和调度。这种封装方式将Host侧的样板代码降到最低，开发者可以将精力集中在Block层和Tile层的参数选择上。

Kernel层的BasicMatmul实现了Block维度的MN循环逻辑。它通过BlockScheduler获取当前Block负责的MN分块坐标，之后再调用BlockMmad执行该分块的K方向主循环，再调用BlockEpilogue执行后处理。BasicMatmul的代码逻辑在不同算子形态间保持一致——无论是基础Matmul还是Grouped Matmul，MN循环的框架结构不变，差异通过BlockMmad和BlockScheduler的模板参数体现。

DispatchPolicy调度策略深度剖析

DispatchPolicy是BlockMmad最重要的模板参数之一，它决定了Block层主循环的具体实现方式。CATLASS采用基于标签的调度策略类型来特例化Block层MMAD实现，这避免了代码重复，使通用代码更容易编写，并提供了清晰的扩展点供用户插入定制实现。

CATLASS目前提供四种DispatchPolicy，分别针对不同性能需求场景：

MmadAtlasA2Pingpong在L1和L0A/B设置Pingpong Buffer机制，通过双缓冲区交替使用实现数据搬运与计算并行。STAGES参数控制缓冲区片数，ENABLE_UNIT_FLAG参数控制是否启用Mmad运算与L0C结果拷出全局内存的细粒度并行。这是最基础的调度策略，适用于00_basic_matmul、01_batched_matmul、03_matmul_add等简单场景。

MmadAtlasA2Preload在Pingpong基础上增加了ShuffleK策略和Block间预加载能力。ShuffleK策略通过重新排列K方向上的计算顺序来优化Cache利用率，Block间预加载允许当前Block的计算与下一个Block的数据搬运提前并行执行。

MmadAtlasA2PreloadAsync引入了nBuffer机制，支持L1和L0A/L0B/L0C上配置不同数量的缓冲区片，同时支持ShuffleK策略、Block间预加载和group间预加载。PRELOAD_STAGES参数控制GM到L1数据读取多少次后启动L1到L0的数据搬运和Mmad计算，L1_STAGES参数控制L1缓冲区片数量。这种异步调度策略适用于Grouped Matmul等需要多group计算的场景。

struct MmadAtlasA2PreloadAsync {
    static constexpr uint32_t PRELOAD_STAGES = 1;
    static constexpr uint32_t L1_STAGES = 2;
    static constexpr uint32_t L0A_STAGES = 2;
    static constexpr uint32_t L0B_STAGES = 2;
    static constexpr uint32_t L0C_STAGES = 1;
    static constexpr bool ENABLE_UINT_FLAG = false;
    static constexpr bool ENABLE_SHUFFLE_K = true;
};

多级缓冲区片数配置让开发者能精确控制每个缓冲区层级的资源占用。L0C设为1片是因为MMAD结果在L0C上只需暂存一轮，计算完成后立即拷出到GM或经FixPipe处理；L0A和L0B设为2片是为了支持Pingpong交替搬运；L1设为2片配合PRELOAD_STAGES等于1实现预取一轮后启动计算。这种参数化设计让同一套模板代码能适配不同的内存预算和并行需求，无需为每种配置重写实现逻辑。

MmadAtlasA2PreloadAsyncWithCallback在PreloadAsync基础上增加了aic和aiv之间同步命令的callback机制，用户可以将自定义同步逻辑以callback形式传入Block层，由Block层决定调用时机。这种设计在量化Matmul等需要aic与aiv协作的场景中尤为关键，因为量化计算中aic侧的反量化操作和aiv侧的MMAD计算之间存在精确的数据依赖关系。

四种DispatchPolicy形成了从简单到复杂的递进关系。Pingpong满足基础场景，Preload增加了ShuffleK和Block间预取，PreloadAsync进一步支持nBuffer和group间预取，PreloadAsyncWithCallback在此基础上提供自定义同步机制。开发者可以根据算子复杂度选择合适的策略层级，不必为简单算子承担复杂策略的配置开销。

TileShape分块参数与硬件缓冲区资源适配

TileShape参数由三元组表示，分别是M、N、K维度方向上的分块大小。L1TileShape对应L1 Buffer上的基本块尺寸，L0TileShape对应L0 Buffer上的基本块尺寸。这两个参数的选择直接影响计算性能和内存利用率。

L1TileShape的选择需要考虑L1 Buffer的容量限制。昇腾NPU的AI Core上L1 Buffer大小有限（Atlas A2系列约为1.5MB），分块尺寸过大会导致缓冲区溢出，分块尺寸过小则无法充分利用MMAD指令的计算能力。以GemmShape<128, 256, 256>为例，当矩阵A为FP16数据类型时，单个L1分块占用128乘以256乘以2字节等于64KB，矩阵B同样占用128KB（256乘以256乘以2字节），合计192KB，远小于L1 Buffer容量，为Pingpong双缓冲留出了充足空间。

L0TileShape的选择需要与MMAD指令的硬件支持对齐。昇腾NPU的Cube单元支持特定尺寸的MMAD指令，L0上的分块尺寸必须是这些指令尺寸的整数倍。GemmShape<128, 256, 64>的K维度设为64，对应了Cube单元一次MMAD计算的K方向计算量，M和N方向则根据L1分块和L0缓冲区容量综合确定。

分块参数的选择还与DispatchPolicy耦合。Pingpong策略下L1需要2片缓冲区，L1TileShape占用的空间需要满足2片不超过L1总容量；PreloadAsync策略下需要更多片数，对L1TileShape的约束更加严格。这种参数间的耦合关系正是CATLASS模板化组装的挑战所在——每个参数都不是孤立选择，而是在硬件资源、调度策略、计算需求三者之间寻找平衡。

// TileShape与缓冲区容量约束的工程计算示例
// L1 Buffer容量约1.5MB (Atlas A2)
// Pingpong STAGES=2时，需要为A和B各预留2片空间
// FP16数据类型下:
// A矩阵单片 = L1TileShape.M * L1TileShape.K * sizeof(FP16)
// B矩阵单片 = L1TileShape.K * L1TileShape.N * sizeof(FP16)
// 总占用 = STAGES * (A单片 + B单片)
// GemmShape<128, 256, 256>:
//   A单片 = 128 * 256 * 2 = 65536 bytes = 64KB
//   B单片 = 256 * 256 * 2 = 131072 bytes = 128KB
//   总占用 = 2 * (64KB + 128KB) = 384KB < 1536KB, 约束满足
constexpr uint32_t l1_total_bytes = 1536 * 1024;
constexpr uint32_t a_tile_bytes = 128 * 256 * 2;
constexpr uint32_t b_tile_bytes = 256 * 256 * 2;
constexpr uint32_t total_stages_occupancy = 2 * (a_tile_bytes + b_tile_bytes);
static_assert(total_stages_occupancy < l1_total_bytes, "L1 overflow");

TileShape参数与缓冲区容量之间的约束关系必须在编译期确定，C++模板参数的静态特性恰好满足这一需求。将TileShape作为模板参数而非运行时参数，使得编译器能在编译期检查缓冲区容量约束是否满足，避免了运行时缓冲区溢出导致的计算错误或设备异常。static_assert等编译期断言机制进一步强化了这一安全保障，开发者在编译阶段就能发现参数配置错误，而非在运行时遇到难以排查的设备异常。

L0TileShape的K维度选择也具有工程讲究。K维度过大时，单次MMAD计算的数据量超过L0A和L0B的容量，无法一次性完成；K维度过小时，MMAD指令的启动开销占比过高，计算效率下降。64这个数值对应了Cube单元在FP16数据类型下的最优K方向计算粒度，既充分利用了L0A/L0B的存储空间，又避免了过细粒度带来的指令调度开销。

Swizzle策略与多核数据走位编排

BlockScheduler负责决定多核间的数据分块映射方式，即不同AI Core处理哪些M和N方向的矩阵分块。CATLASS提供了GemmIdentityBlockSwizzle等基础策略，也支持更复杂的Swizzle编排。

Swizzle策略的核心目标是在多核并行场景下优化GM访问的带宽利用率。当多个AI Core按照简单的行列顺序分配分块时，相邻Core可能在同一时刻访问GM上相近的地址区域，导致Bank冲突和带宽拥塞。Swizzle策略通过重新排列分块到Core的映射顺序，使得同一时刻活跃的Core访问的GM地址区域尽可能分散，从而减少Bank冲突。

以一个8乘以8的分块网格为例，Identity Swizzle按行优先顺序将分块分配给Core 0、1、2、3……，同一行上的分块连续映射到相邻Core，这些Core在同一时刻需要读取矩阵B的相同K方向数据段，造成GM读取热点。经过Swizzle重排后，同一行上的分块被分散到不相邻的Core，这些Core读取矩阵B的不同数据段，GM读取压力更加均匀。

CATLASS设计文档中的Swizzle策略说明详细介绍了不同Swizzle模式对分块顺序的影响。Identity Swizzle保持原始行列顺序，适用于分块数量较少或GM带宽充足的场景；更复杂的Swizzle模式则在大规模并行计算中带来明显的带宽优化效果。

Swizzle策略的选择与DispatchPolicy中的ShuffleK策略相互配合。ShuffleK影响K方向上的计算顺序，Swizzle影响MN平面上的分块映射，两者共同决定了数据访问模式和计算顺序的时空分布。在Flash Attention等需要精细流水编排的算子中，这种双重控制尤为重要，因为Flash Attention的Q和K分块计算结果需要按特定顺序传递给后续步骤，Swizzle编排既要考虑带宽优化，又要满足数据依赖约束。

量化Matmul与FixPipe随路量化机制

量化推理是当前AI部署的关键技术路径，CATLASS从v1.3.0版本开始支持FixPipe随路量化，并在v1.5.0版本新增了W8A8 Per-Token量化、Per-Group和Per-Block量化Matmul TLA等示例。

FixPipe是昇腾NPU上的一个硬件流水线单元，它能在MMAD计算结果从L0C拷出到GM的过程中，同步执行量化、反量化、类型转换等操作。这种随路处理方式避免了额外的内存往返开销——计算结果不需要先写回GM再单独调用量化算子处理，而是在拷出路径上一步完成。

在W8A8量化Matmul场景中，矩阵A和矩阵B以INT8格式存储，计算时需要先将INT8数据反量化到FP16或BF16进行MMAD，再将FP32累加结果量化回INT8格式输出。CATLASS通过TileCopy中的量化与反量化配置和Epilogue中的FixPipe处理，将这一流程编排为计算与量化的融合流水线。TileCopy在数据从GM搬运到L1或从L1搬运到L0的过程中执行反量化，FixPipe在结果从L0C拷出到GM的过程中执行量化。

Per-Token量化对每个Token使用独立的量化参数（scale和zero_point），Per-Group量化对每组元素使用共享的量化参数，Per-Block量化则对整个矩阵分块使用统一参数。不同量化粒度的选择影响量化参数的存储开销和精度损失。Per-Token量化精度最高但参数存储开销最大，Per-Block量化参数存储开销最小但精度损失相对较大。CATLASS的模板化设计让开发者能通过配置GemmType参数灵活切换量化粒度，无需重写计算流水逻辑。

v1.5.0版本新增的动态W8A8 Per-Token量化示例（103_dynamic_optimized_quant_matmul_per_token_basic）展示了在Matmul泛化工程框架下实现Per-Token量化的完整流程。这个示例中，量化参数作为额外的输入张量传入Kernel，TileCopy在搬运矩阵A和B的数据时，根据每个Token对应的scale和zero_point执行反量化，MMAD计算在反量化后的高精度数据上进行。

W4A4量化是量化场景的极致压缩。矩阵A以4-bit权重存储，矩阵B同样4-bit，计算时反量化到更高精度进行MMAD。4-bit量化将内存占用压缩到原来的四分之一，但反量化操作引入了额外计算开销。CATLASS通过FixPipe随路反量化将这一开销最小化，使得W4A4量化在推理场景中具备实用价值。v1.4.0版本引入的38_w4a4_matmul_per_token_per_channel_dequant示例展示了这一机制的完整实现。

---

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