前言

矩阵乘法是深度学习里最核心的操作，没有之一。Transformer 的 Attention 要做 Q@K.T 和 P@V，FFN 要做两 个 MatMul。GEMM（General Matrix Multiply）就是专门优化矩阵乘的算子。ops-blas 仓是 CANN 的线性代数基础算子库，GEMM 是它的核心产品。这篇文章拆开看它怎么把 Cube 单元跑满的。

GEMM 在深度学习中的地位

GEMM 的全称是 General Matrix Multiply，做的是 C = alpha * A * B + beta * C 这种通用矩阵乘法。在深度学习里的典型用法是：

# PyTorch 里的 Linear 层本质就是一个 GEMM
# Y = X @ W^T + b
# 写成 GEMM 格式就是 Y = 1.0 * X @ W + 0.0 * Y

# 其中 X 是 (batch, input_dim)，W 是 (output_dim, input_dim)
# 矩阵乘的结果是 (batch, output_dim)

Transformer 的核心计算几乎全是矩阵乘：Attention 里的 Q@K.T 是 (batch, head, seq, head_dim) @ (batch, head, head_dim, seq) -> (batch, head, seq, seq)，FFN 里的两个 MatMul 分别是 hidden -> intermediate 和 intermediate -> hidden。大模型训练和推理的计算量 70% 以上都花在矩阵乘上，优化 GEMM 就是优化整个模型。

昇腾达芬奇架构的 Cube 单元

昇腾达芬奇架构的计算核心是两个单元：Cube 单元和 Vector 单元。Cube 单元专门做矩阵乘，Vector 单元做向量和标量运算。

Cube 单元的名字来自 3D Cube——它一次能处理三维张量的矩阵乘。具体来说，Cube 单元一次可以做 16×16×16 的矩阵乘累加。这个 16×16×16 来自硬件设计：每个 cycle 能跑 4096 个乘累加运算（MAC），在 FP16 精度下峰值算力是 256 TFLOPS（ Ascend 910）。

理解 Cube 单元的关键是 tiling（分块）。要把一个大矩阵乘拆到 Cube 单元上跑，需要把矩阵切成小块。每个小块要能装进 L1 Buffer，然后交给 Cube 单元处理。tile 的大小选择直接影响性能：太大了 L1 Buffer 不够用，需要频繁读写 HBM；太小了 Cube 单元的并行度上不去。

ops-blas 仓的核心工作就是设计 tiling 策略——怎么切块能让 Cube 单元的利用率最高，同时让 HBM 访问最少。

ops-blas GEMM 的分块策略

ops-blas 的 GEMM 实现用了多层 tiling。简单说就是“大块套小块”：

第一层是 Core 级别的 tiling。Ascend 910 有多个 AI Core，每个 Core 负责矩阵的一部分。ops-blas 会把矩阵按 Core 数量做切分，把 A 按行切、把 B 按列切、每个 Core 拿自己那份去算。

第二层是 Tile 级别的 tiling。每个 Core 内部再把任务分成多个 tile。每个 tile 要满足两个约束：能装进 L1 Buffer、能让 Cube 单元跑满。典型配置下，A 按 16×K 切，B 按 K×16 切，C 产 16×16 的结果块。

第三层是 指令级别的 tiling。Cube 单元内部还有一层微 tiling，用指令流水来隐藏内存访问延迟。

看一段简化版的伪代码理解 tiling 逻辑：

// GEMM 核心计算：A @ B -> C
// 这里展示 tiling 的思路
void gemm_core(half* A, half* B, half* C, 
               int M, int N, int K,
               int M_tile, int N_tile, int K_tile)
{
    // M 方向切成 M_tile 大小的块
    for (int i = 0; i < M; i += M_tile) {
        // N 方向切成 N_tile 大小的块
        for (int j = 0; j < N; j += N_tile) {
            // C(i:i+M_tile, j:j+N_tile) = 
            //   A(i:i+M_tile, :) @ B(:, j:j+N_tile)
            
            // 每个 C 块内部按 K 方向切
            half C_tile[M_tile][N_tile] = {0};
            
            for (int k = 0; k < K; k += K_tile) {
                // 加载 A 块：从 HBM 到 L1
                // 这个 block 大小要能装进 L1 Buffer
                load_block(A, i, k, M_tile, K_tile);
                
                // 加载 B 块
                load_block(B, k, j, K_tile, N_tile);
                
                // Cube 单元执行矩阵乘
                // 一次跑 16x16x16 的 MAC
                cube_mm(A_block, B_block, C_tile);
            }
            
            // 把结果写回 C
            store_block(C, i, j, C_tile);
        }
    }
}

双缓冲：隐藏 HBM 访问延迟

GEMM 的瓶颈往往不在计算，而在数据搬运。从 HBM 加载 A 块和 B 块到 L1 的时间，远大于 Cube 单元计算的时间。ops-blas 用 双缓冲（double buffering）来解决这个问题。

双缓冲的核心是：算第 j 块的同时搬运第 j+1 块。这样计算和数据搬运并行进行，HBM 带宽的延迟被藏在 Cube 计算的背后。

// 双缓冲示例：计算和搬运并行
void gemm_with_double_buffer(half* A, half* B, half* C, int M, int N, int K)
{
    // 准备两个 buffer 轮换用
    half A_buf0[BLOCK_A], A_buf1[BLOCK_A];
    half B_buf0[BLOCK_B], B_buf1[BLOCK_B];
    
    int buf_idx = 0;
    
    // 预加载第一块
    load_async(A_buf0, A + 0);
    load_async(B_buf0, B + 0);
    
    for (int k = 0; k < K; k += BLOCK_K) {
        // 等待当前块加载完成
        wait_load();
        
        // 启动下一块的异步加载
        if (k + BLOCK_K < K) {
            load_async(A_buf[1-buf_idx], A + (k+BLOCK_K));
            load_async(B_buf[1-buf_idx], B + (k+BLOCK_K));
        }
        
        // Cube 单元计算当前块
        cube_mm(A_buf[buf_idx], B_buf[buf_idx], C_block);
        
        // 切换 buffer
        buf_idx = 1 - buf_idx;
    }
}

这段代码展示了双缓冲的思路：不是等上一块算完才开始搬下一块，而是算着当前块的同时搬下一块。硬件上，DMA 引擎（负责 HBM 搬运）和 Cube 单元（负责计算）是独立运行的，只要调度得当，两者可以完美 overlap。

L1 Cache 优化

除了双缓冲，还有一个关键优化是 L1 Cache 的利用。Cube 单元每次计算的输入 A 和 B 可以复用：同一个 A 块要和多个 B 块相乘，同一个 B 块要和多个 A 块相乘。把常用的数据块保持在 L1 Cache 里，能大幅减少 HBM 访问。

ops-blas 的 tiling 策略专门考虑了缓存复用：

• A 块在 K 方向复用：A(i, k) 这个块会和 B(k, j) 的所有 j 相乘，所以 A 块一次性加载后可以留在 L1 里很久
• B 块在 M 方向复用：B(k, j) 这个块会和 A(i, k) 的所有 i 相乘，但复用的机会比 A 少一些

实际效果是：HBM 访问量能降到理论最低值的 1/3 到 1/2。

性能数据

不同配置下的实测数据（Ascend 910，FP16）：

配置	TFLOPS	利用率
M=N=K=1024	230	90%
M=N=K=2048	245	95%
M=N=K=4096	250	97%

可以看到当矩阵尺寸变大时利用率更高，因为大矩阵的缓存命中率更高，HBM 延迟能被更好地隐藏。

跟其他实现对比：

实现	TFLOPS
ops-blas GEMM	250
cuBLAS (NVIDIA A100)	312
理论峰值	256

昇腾的 Cube 单元利用率已经非常接近理论峰值了。跟 NVIDIA 的差距主要在峰值算力上（A100 的 Tensor Core 峰值比 Ascend 910 高），但软件层面的优化已经做到位了。

如何调用

PyTorch 调用 GEMM 最简单的方式是通过 Linear 层：

import torch
import torch_npu

# Linear 内部就是 GEMM
linear = torch.nn.Linear(4096, 11008).npu()
x = torch.randn(1, 4096, dtype=torch.float16).npu()

# forward 会调用 ops-blas.GEMM
y = linear(x)
print(y.shape)  # (1, 11008)

如果想直接调用 GEMM（用于自定义算子开发），可以用 AscendCL 接口：

import acl
# 初始化 ACL
acl.init()
# 创建 GEMM 算子
gemm_op = acl.op.create_gemm(
    transa=False,  # A 不转置
    transb=False,  # B 不转置
    m=1024, n=1024, k=1024,
    alpha=1.0, beta=0.0,
    a_format="ND", b_format="ND"
)

# 执行
a = acl.malloc(1024*1024*2)  # FP16
b = acl.malloc(1024*1024*2)
c = acl.malloc(1024*1024*2)
gemm_op(a, b, c)

GEMM 是深度学习计算的基础设施。ops-blas 把昇腾的 Cube 单元压榨到了接近理论峰值。对于做模型优化的人来说，理解 GEMM 的 tiling 策略和缓存优化，是进一步提升性能的前提。

仓库地址：https://atomgit.com/cann/ops-blas