矩阵乘法模板如何做到 92-98% 手写性能？深度拆解 catlass 的实现

1. 背景：为什么手写算子很难做到最优？

要理解catlass的价值，得先搞清楚"手写算子做到最优"有多难。

1.1 算子性能的天花板

NPU上算子性能的天花板是理论峰值：Cube单元的理论峰值FLOPS（比如Ascend 910是256 TFLOPS FP16）。

实际算子的性能，通常用理论峰值利用率来衡量：

利用率 = 实测FLOPS 理论峰值FLOPS \text{利用率} = \frac{\text{实测FLOPS}}{\text{理论峰值FLOPS}}利用率=理论峰值FLOPS实测FLOPS​

一个"好"的算子，利用率应该在85-95%。低于80%，说明有优化空间；高于95%，基本到硬件极限了。

1.2 手写算子的常见陷阱

手写Ascend C算子，新手（甚至有一定经验的开发者）容易踩这几个坑：

坑1：Tile大小选得不好

Tile是NPU上的一次计算块。Tile大小选得太小，Cube/Vector单元的并行度利用不充分；选得太大，片上存储放不下，频繁洒显存。

// 错误示例：tile大小选得不好__global__voidmatmul_bad(Tensor<float>A,Tensor<float>B,Tensor<float>C){// tile大小 = 32×32（太小，Cube单元没吃满）floatlocal_A[32][32];floatlocal_B[32][32];// ... 数据拷贝 + 计算 ...}// WHY: 这个tile大小(32×32)对于Ascend 910的Cube单元来说太小了，// Cube单元一次能处理128×128的矩阵块，// 用32×32的tile，Cube单元的并行度只利用了1/16。// 正确示例：tile大小选得合适__global__voidmatmul_good(Tensor<float>A,Tensor<float>B,Tensor<float>C){// tile大小 = 128×128（合适，Cube单元吃满）floatlocal_A[128][128];floatlocal_B[128][128];// ... 数据拷贝 + 计算 ...}// WHY: 128×128的tile能让Cube单元的并行度充分利用，// 同时也要考虑片上存储的大小（128×128×4 bytes × 2个矩阵 ≈ 128KB，// 对于Ascend 910的Local Memory来说是合适的）。

坑2：数据搬运和计算的并行度没做好

NPU支持计算和数据搬运并行（类似GPU的Compute + Memory Copy并行）。如果没做好，算子要等数据搬完才能算，或者算完要等数据写回去。

// 错误示例：计算和搬运串行__global__voidmatmul_serial(Tensor<float>A,Tensor<float>B,Tensor<float>C){// 阶段1：把A搬进片上copy_matrix(A,local_A);// 搬运// 阶段2：把B搬进片上copy_matrix(B,local_B);// 搬运// 阶段3：计算cube_matmul(local_A,local_B,local_C);// 计算// 阶段4：把结果写回去copy_matrix(local_C,C);// 写回}// WHY: 这个实现是纯串行的：搬A → 搬B → 算 → 写回。// NPU支持搬运和计算并行，应该让"搬A"和"算上一次的结果"并行。// 正确示例：计算和搬运流水线__global__voidmatmul_pipelined(Tensor<float>A,Tensor<float>B,Tensor<float>C){// 初始化：搬第0块copy_matrix(A[0],local_A);copy_matrix(B[0],local_B);for(inti=1;i<num_tiles;i++){// 计算第i-1块（用上一次搬进来的数据）cube_matmul(local_A,local_B,local_C);// 同时：搬第i块（和计算并行）async_copy_matrix(A[i],local_A_next);async_copy_matrix(B[i],local_B_next);// 等待搬运完成wait_all();// 交换指针swap(local_A,local_A_next);swap(local_B,local_B_next);}// 计算最后一块cube_matmul(local_A,local_B,local_C);// 写回结果copy_matrix(local_C,C);}// WHY: 这个实现用了double buffering + 流水线：// 计算第i-1块的同时，搬第i块的数据。// 这样Cube单元就不会闲着等数据。

1.3 手写算子做到最优需要多久？

根据我的经验：

catlass的价值就是：把"1-2周写到最优"这件事，压缩到20分钟（改模板参数 + 编译 + 测试）。

2. 原理：catlass的代码生成策略

catlass的核心是一个模板库：它把算子优化的专家知识，封装成一个个可调的模板参数。你改模板参数，它生成对应的Ascend C代码。

2.1 模板参数体系

catlass的MatMul模板，核心参数有这几个：

// catlass的MatMul模板（示意）template<// 1. Tile大小（影响Cube单元利用率）intTileM,intTileN,intTileK,// 2. 数据类型（FP16/BF16/FP32）typenameDataTypeA,typenameDataTypeB,typenameDataTypeC,// 3. 分块策略（影响显存访问模式）intBlockM,intBlockN,// 4. 流水线深度（影响计算和搬运并行度）intPipelineDepth,// 5. 预取策略（影响数据搬运效率）boolEnablePrefetchA,boolEnablePrefetchB>classMatMulTemplate{public:voidoperator()(Tensor<DataTypeA>A,Tensor<DataTypeB>B,Tensor<DataTypeC>C){// 生成的代码：根据模板参数，// 自动选择一个最优的Tile大小、分块策略、流水线深度、预取策略// ...}};

关键：这些模板参数不是让你瞎选的。catlass内置了一个代价模型（Cost Model），你给它一个MatMul的配置（M, N, K, dtype），它自动算出最优的模板参数组合。

2.2 代码生成流程

catlass的代码生成分三步：

步骤1：代价模型搜索

给你一个MatMul配置（M=1024, N=1024, K=1024, dtype=FP16），代价模型会搜索所有可能的模板参数组合，预测每个组合的性能。

fromcatlassimportMatMulTemplate,search_best_config# 搜索最优配置config=search_best_config(M=1024,N=1024,K=1024,dtype='fp16',device='ascend910')print(config)# 输出（示意）：# {# 'TileM': 128, 'TileN': 128, 'TileK': 64,# 'BlockM': 64, 'BlockN': 64,# 'PipelineDepth': 3,# 'EnablePrefetchA': True, 'EnablePrefetchB': True# }# WHY: 代价模型通过"模拟NPU执行"来预测性能。# 它会考虑：Tile大小 → Cube单元利用率；# 分块策略 → 显存访问模式（连续 vs 非连续）；# 流水线深度 → 计算和搬运的并行度；# 预取策略 → 数据搬运是否和前面计算重叠。

步骤2：代码生成

有了最优配置，catlass调用代码生成器（Code Generator），生成对应的Ascend C代码。

fromcatlassimportMatMulTemplate,generate_code# 生成Ascend C代码code=generate_code(template=MatMulTemplate,config=config,output_format='ascend_c')print(code[:500])# 打印前500个字符# 输出（示意）：# __global__ void matmul_optimized(Tensor<half> A, Tensor<half> B, Tensor<half> C) {# __shared__ half local_A[128][64];# __shared__ half local_B[64][128];# // ... 根据config生成的优化代码 ...# }# WHY: 代码生成器把模板参数"实例化"成具体的Ascend C代码。# 比如TileM=128 → local_A的大小是[128][64]；# PipelineDepth=3 → 生成3级流水线（double buffering + 预取）。

步骤3：编译 + 性能验证

生成的代码，调用Ascend C编译器（BiSheng/ATC）编译成NPU kernel，然后跑一个小的benchmark，验证性能是否达到预期。

fromcatlassimportcompile_and_verify# 编译 + 验证kernel=compile_and_verify(code=code,M=1024,N=1024,K=1024,dtype='fp16',verify_correctness=True,# 验证正确性（和PyTorch结果对比）verify_performance=True,# 验证性能（是否达到理论峰值的90%+）)print(f"利用率:{kernel.utilization():.1%}")# WHY: 编译后的验证很重要，因为代价模型的预测可能不准# （比如NPU的某些特殊指令延迟没建模好）。# 如果性能不达标，catlass会回退到"搜索下一个最优配置"。

3. 昇腾NPU上的代码生成策略

上一节讲的是通用原理，这一节深入昇腾NPU的硬件特性，看catlass如何利用这些特性做进一步的优化。

3.1 Cube单元专用优化

昇腾NPU的Cube单元，有几个特殊性质：

1. Cube单元只支持特定的矩阵大小：比如FP16的MatMul，Cube单元期望的输入是[16, 16] × [16, 16]的块（这个大小叫Cube Tile）
2. Cube单元有专用的数据通路：数据从Global Memory → Cube单元的寄存器，有专门的DMA通道（不经过Vector单元）

catlass在做代码生成时，会针对Cube单元的这些特性做优化：

// 针对Cube单元优化后的MatMul（示意，catlass生成）template<intTileM,intTileN,intTileK>__global__voidmatmul_cube_optimized(Tensor<half>A,Tensor<half>B,Tensor<half>C){// 1. 把Tile大小对齐到Cube Tile的倍数static_assert(TileM%16==0,"TileM must be multiple of 16");static_assert(TileN%16==0,"TileN must be multiple of 16");static_assert(TileK%16==0,"TileK must be multiple of 16");// 2. 用Cube单元专用的DMA指令搬数据cube_dma_load_a(local_A,A[block_idx*TileM,...]);cube_dma_load_b(local_B,B[...,block_idx*TileN]);// 3. 调用Cube单元的MatMul指令cube_matmul<half,16,16,16>(local_A,local_B,local_C);// 4. 用Cube单元专用的DMA指令写结果cube_dma_store_c(local_C,C[block_idx*TileM,...]);}// WHY: 这个优化后的代码，// 1. Tile大小对齐到Cube Tile (16×16)，Cube单元利用率100%；// 2. 用Cube专用的DMA（而不是通用的DMA），数据搬运更快；// 3. 调用Cube的MatMul指令（而不是用Vector单元模拟MatMul），// 计算速度快10-20倍。

3.2 Vector单元辅助优化

MatMul的计算，除了Cube单元做矩阵乘法，还有一部分工作是Vector单元做的：

1. Bias加法：C = A × B + bias（bias是Vector操作）
2. 激活函数：C = ReLU(A × B)（ReLU是Vector操作）
3. 类型转换：C_fp32 = A_fp16 × B_fp16（类型转换是Vector操作）

catlass在做代码生成时，会把Cube单元的计算和Vector单元的计算流水线化：

// Cube + Vector流水线化（示意，catlass生成）__global__voidmatmul_cube_vector_pipelined(Tensor<half>A,Tensor<half>B,Tensor<float>C,Tensor<float>bias){// 阶段1：Cube算MatMul（第i-1块）cube_matmul(local_A_prev,local_B_prev,local_C_prev);// 阶段2：Vector算Bias加法 + 激活（第i-1块，和Cube算第i块并行）vector_add_bias(local_C_prev,bias,local_C_bias);vector_relu(local_C_bias,local_C_relu);// 阶段3：Cube算MatMul（第i块，和Vector算第i-1块并行）cube_matmul(local_A,local_B,local_C);// ... 循环 ...}// WHY: Cube和Vector是两个独立的执行单元，可以并行。// 让Cube算第i块MatMul的同时，Vector算第i-1块的Bias+ReLU，// Cube和Vector的利用率都能接近100%。

3.3 显存层级优化

昇腾NPU的显存层级是：

Global Memory (显存, GB级, 慢) ↓ DMA搬运 Local Memory (片上存储, MB级, 快) ↓ 寄存器搬运 Cube/Vector寄存器 (KB级, 极快)

catlass在做代码生成时，会针对这个显存层级做Tile大小和分块策略的联合优化：

// 显存层级优化（示意，catlass生成）template<intTileM,intTileN,intTileK,intBlockM,intBlockN>__global__voidmatmul_memory_hierarchical(Tensor<half>A,Tensor<half>B,Tensor<half>C){// 1. Global → Local：按Tile大小搬（一次搬一个Tile）dma_load(A_global,local_A,TileM,TileK);dma_load(B_global,local_B,TileK,TileN);// 2. Local → 寄存器：按Block大小分块（一次搬一个Block到寄存器）for(intbm=0;bm<TileM;bm+=BlockM){for(intbn=0;bn<TileN;bn+=BlockN){// 搬Block到Cube寄存器cube_load_register(local_A[bm,...],reg_A,BlockM,TileK);cube_load_register(local_B[...,bn],reg_B,TileK,BlockN);// Cube计算（寄存器级别，极快）cube_matmul<BlockM,BlockN,TileK>(reg_A,reg_B,reg_C);// 写回Local Memorycube_store_local(reg_C,local_C[bm,bn],BlockM,BlockN);}}// 3. Local → Global：写回结果dma_store(local_C,C_global,TileM,TileN);}// WHY: 这个分层的代码，// 1. Global → Local 用DMA（快）// 2. Local → 寄存器 用Cube的专用指令（更快）// 3. 寄存器 → Cube计算（极快）// 显存层级的每个层级都用最优的数据通路。

4. 跟手写 Ascend C 的对比

这一节用实测数据对比"手写Ascend C"和"catlass模板生成"的性能差异。

4.1 测试环境

• 硬件：昇腾910 NPU（32GB显存）
• 软件：CANN 8.0, Ascend C 2.1, catlass 1.0
• 测试算子：MatMul（FP16, 各种M, N, K）

4.2 性能对比（理论峰值利用率）

我们测的是MatMul算子的理论峰值利用率（实测FLOPS / 理论峰值FLOPS）。

解读：catlass生成的MatMul，利用率在92-95%，比手写Ascend C高5-19%。而且矩阵越大，catlass的优势越小（因为大矩阵的优化空间更小，手写也能做到不错的效果）。

4.3 开发时间对比

解读：catlass不仅性能更好，开发效率也高得多。手写一个最优的MatMul要5-7天，catlass只要20-25分钟。

4.4 正确性验证

性能高但结果错，等于没用。我们验证了catlass生成的MatMul和PyTorch的MatMul的输出差异。

解读：catlass生成的MatMul，和PyTorch的结果非常接近（相对误差<0.02%），完全可以替代手写的算子。

5. 性能数据深度分析

上一节的对比是"手写 vs catlass"的整体效果。这一节深入一点，看catlass在不同场景下的性能表现。

5.1 不同数据类型的性能

catlass支持多种数据类型（FP16, BF16, FP32）。我们测了不同数据类型的MatMul性能。

解读：catlass在各种数据类型下都比手写快。FP16和BF16的利用率接近（因为Cube单元对它们的处理差不多），FP32的利用率稍低（因为FP32的计算更复杂，Cube单元的利用率难做到100%）。

5.2 不同矩阵形状的性能

实际的MatMul，不一定是方阵（M=N=K）。我们测了几种典型的非方阵形状。

解读：非方阵的利用率比方阵低（因为Tile的某些维度很小，Cube单元吃不满）。但catlass的优化仍然比手写好14-17%。

5.3 跟其他模板库的对比

学术界和工业界已经有不少算子模板库。我们拿catlass和几个有代表性的方案做对比：

解读：catlass在NPU上的性能是最优的（比TVM好，比CUTLASS更适配NPU）。开发时间也最短（20分钟 vs CUTLASS的30分钟或TVM的2-4小时）。

6. 使用技巧

最后一节，总结一些实际使用catlass时的技巧和坑点。

6.1 技巧1：先搜索最优配置，再生成代码

不要瞎猜模板参数。用catlass内置的search_best_config搜索最优配置。

fromcatlassimportsearch_best_config,MatMulTemplate,generate_code,compile_and_verify# 1. 搜索最优配置（耗时10-15分钟，要跑很多组参数的性能测试）best_config=search_best_config(M=2048,N=2048,K=2048,dtype='fp16',device='ascend910',search_space='exhaustive'# 穷举搜索（慢但准）)print(f"最优配置:{best_config}")# 2. 生成代码code=generate_code(MatMulTemplate,best_config)# 3. 编译 + 验证kernel=compile_and_verify(code,M=2048,N=2048,K=2048,dtype='fp16')print(f"利用率:{kernel.utilization():.1%}")# WHY: 穷举搜索能保证找到全局最优（只要搜索空间覆盖得够全）。# 如果搜索时间太长，可以用 search_space='heuristic'（启发式搜索，快但可能局部最优）。

6.2 技巧2：注意内存对齐

catlass生成的代码，要求输入张量是内存对齐的。如果没对齐，性能会下降10-20%。

importtorchimporttorch_npu# 不好的做法：张量没对齐A=torch.randn(1025,1024,dtype=torch.float16).npu()# 1025不是16的倍数B=torch.randn(1024,1024,dtype=torch.float16).npu()C=torch.zeros(1025,1024,dtype=torch.float16).npu()kernel(A,B,C)# 性能下降15%# 好的做法：张量对齐到16的倍数A=torch.randn(1024,1024,dtype=torch.float16).npu()# 1024是16的倍数B=torch.randn(1024,1024,dtype=torch.float16).npu()C=torch.zeros(1024,1024,dtype=torch.float16).npu()kernel(A,B,C)# 性能最优# WHY: Cube单元的DMA要求输入对齐到16字节（FP16）或32字节（FP32）。# 如果没对齐，DMA要做一个额外的"对齐拷贝"，性能下降。

6.3 技巧3：用profiling工具验证性能是否达标

catlass生成的kernel，怎么知道性能是否达到最优？

用NPU的profiling工具看Cube单元利用率：

# 用msprof抓profilingmsprof--output=./profiling--application="python test_catlass.py"# 查看Cube单元利用率msprof--export=on--output=./profiling|grep"Cube"# 输出（示意）：# Cube Utilization: 94.2%# - MatMul: 96.8%# - BiasAdd: 12.3% (Vector单元在做，Cube单元空闲)# WHY: 如果Cube利用率 < 90%，说明模板参数没选好（Tile大小不合适、流水线深度不够等）。# 这时候应该回到技巧1，重新搜索最优配置。

6.4 技巧4：注意动态形状的编译开销

如果模型的输入形状是动态的（比如NLP模型的变长序列），catlass要为每个不同的形状都生成和编译一次kernel，编译开销很大（每次20-30秒）。

catlass提供了一个形状范围声明的API，让你提前告诉它"可能的形状范围"，它会在初始化时就把这个范围内的所有kernel都编译好。

fromcatlassimportShapeRange,precompile_kernels# 声明形状范围shape_range=ShapeRange(M=[1,4,8,16,32,64,128,256,512,1024],N=[1024,2048,4096],K=[1024,2048,4096])# 预编译（耗时5-10分钟，但之后所有形状都能直接用）precompile_kernels(template=MatMulTemplate,shape_range=shape_range,dtype='fp16',device='ascend910')# WHY: 动态形状的模型（比如NLP的变长序列），# 如果每次都现场编译kernel，推理延迟会很高（每次20-30秒）。# 用precompile_kernels提前编译所有可能用到的kernel，# 运行时直接取用，没有编译开销。

总结

把这件事从头到尾捋一遍：

手写Ascend C算子要做到最优（利用率>90%），需要深度理解NPU的Cube/Vector并行度、内存层级、指令流水线、数据预取……这些知识的积累需要几个月甚至几年。

catlass的价值就是：把专家级的算子优化知识，封装成可复用的模板。

实测数据显示，catlass生成的MatMul算子，理论峰值利用率达到92-95%，比手写Ascend C高5-19%。而且开发时间从5-7天压缩到20-25分钟，效率提升99.7%。

catlass的核心技术分三层：

1. 模板参数体系：Tile大小、分块策略、流水线深度、预取策略
2. 代价模型：给定配置，自动搜索最优的模板参数组合
3. 代码生成器：根据最优配置，生成针对NPU硬件优化的Ascend C代码

仓库链接：https://atomgit.com/cann/catlass