矩阵乘是所有深度学习计算的根。Attention、全连接、卷积展开——归根到底都是矩阵乘。ops-blas 是 CANN 里专门做高性能 GEMM（General Matrix Multiply）的算子库，核心目标是把昇腾 NPU 的 Cube 单元利用率拉到 90% 以上。

ops-blas 和 ops-nn 的 MatMul 有重叠，但分工不同：ops-nn 侧重神经网络场景的融合算子（MatMul+Bias+GELU），ops-blas 侧重纯计算的 GEMM 极致性能。如果 ops-nn 是一辆带自动挡的轿车，ops-blas 就是赛道上拆了空调减重到极致的 F1。

矩阵乘的硬件分块

昇腾 NPU 的 Cube 单元一次算 16×16 的矩阵乘。对 M×K 的矩阵 A 和 K×N 的矩阵 B，理论上需要 (M/16) × (N/16) × (K/16) 次 Mmad 指令。分块的粒度直接影响计算效率：

// ops-blas/gemm/gemm_tiling_strategy.cpp（简化）

struct GemmTiling {
    int M_block;   // A 的行分块
    int N_block;   // B 的列分块
    int K_block;   // 公共维分块
};

GemmTiling select_tiling(int M, int N, int K, DataType dtype) {
    // L1 缓存可容纳的元素数（考虑双缓冲需要 3 个 buffer）
    constexpr int L1_CAPACITY = 192 * 1024 / sizeof(float);
    
    // A 块: M_block × K_block
    // B 块: K_block × N_block  
    // C 块: M_block × N_block
    // 总共必须 <= L1_CAPACITY

    // 经验公式: M_block ≈ N_block ≈ sqrt(L1_CAPACITY / 3)
    int optimal = (int)sqrt(L1_CAPACITY / 3.0);

    // 对齐到 16 的倍数（Cube 限制）
    int Mb = min((M + 15) / 16 * 16, optimal / 16 * 16);
    int Nb = min((N + 15) / 16 * 16, optimal / 16 * 16);
    int Kb = min((K + 15) / 16 * 16, optimal / 16 * 16);

    // 如果 M 很小（batch=1 的 FC 层），调整策略
    if (M < 64) {
        // M 小、N 大：增加 M 块大小，减少 load 的次数
        Mb = min(M, optimal / 4 * 16);
    }
    if (N < 64) {
        // N 小、M 大：增加 N 块大小
        Nb = min(N, optimal / 4 * 16);
    }

    return {Mb, Nb, Kb};
}

分块策略直接影响 Cube 利用率。分块太小，Mmad 指令开销占比高；分块太大，L1 缓存溢出，HBM 开始频繁换入换出。select_tiling 的核心约束就是 L1 容量——三个 buffer（A 块、B 块、C 块）加起来不能超过 L1 的容量，否则 Cube 会停下来等数据。

GEMM 的大中小三条优化路径

ops-blas 根据矩阵形状走三条不同的优化路径：

// ops-blas/gemm/gemm_dispatch.cpp

enum GemmPath {
    GEMM_LARGE,    // M,N >= 1024: 大矩阵
    GEMM_MEDIUM,   // M,N 在 64-1023: 中等矩阵
    GEMM_TINY      // M,N < 64: 小矩阵（比如 batch 中的单样本）
};

GemmPath classify(int M, int N, int K) {
    if (M >= 1024 && N >= 1024) return GEMM_LARGE;
    if (M < 64 || N < 64) return GEMM_TINY;
    return GEMM_MEDIUM;
}

Status gemm(const Tensor& A, const Tensor& B, Tensor& C) {
    auto [M, K] = A.shape();
    auto [K2, N] = B.shape();
    auto path = classify(M, N, K);

    switch (path) {
        case GEMM_LARGE:
            // 大矩阵：多级 tiling + packing
            // 对 A 和 B 做内存排布重排（packing），
            // 让 Cube 单元始终按连续地址读取
            return gemm_large(A, B, C);

        case GEMM_MEDIUM:
            // 中等矩阵：标准双缓冲流水
            return gemm_medium(A, B, C);

        case GEMM_TINY:
            // 小矩阵：batch 合并策略
            // 把多个小矩阵的 GEMM 合并成一个大矩阵乘
            // 比如把 32 个 [64, 768]×[768, 768] 合并成 [2048, 768]×[768, 768]
            return gemm_batched(A, B, C);
    }
}

大矩阵路径（packing）：对矩阵 A 和 B 做内存重排，把不连续的列主序或行主序重新排列成 Cube 单元需要的连续片。这个操作叫 packing，需要一次额外的 HBM 拷贝，但只要矩阵足够大，packing 的时间会被 Cube 的峰值吞吐摊平。

中等矩阵路径（双缓冲）：不额外 packing，用双缓冲流水让搬运和计算并行。搬运下一块数据的同时 Cube 在算当前块。

小矩阵路径（batch 合并）：单个小矩阵乘的开销全在启动延迟（kernel launch + HBM 初始化），实际计算时间微乎其微。ops-blas 把多个小矩阵的 GEMM 调用合并成一次 batch GEMM——比如 32 个 [64,768]×[768,768] 合并成 [2048,768]×[768,768]，延迟从 32 次降到 1 次。

Batched GEMM 的实战

LLaMA 的 attention 层里每个头独立做 Q×K^T，32 个头就是 32 个独立的矩阵乘。ops-blas 的 batched GEMM 在这种场景下有天然优势：

// ops-blas/batched_gemm/batched_gemm_kernel.cpp

__aicore__ void batched_gemm_kernel(
    GlobalTensor<float>& C,   // [batch, M, N]
    GlobalTensor<float>& A,   // [batch, M, K]
    GlobalTensor<float>& B,   // [batch, K, N]
    int batch, int M, int N, int K
) {
    // 把所有 batch 的矩阵拼接成大矩阵：
    // A_merged: [batch*M, K]  等价于把所有 batch 沿 M 维拼接
    // B_merged: [K, N]        所有 batch 共享同一个 B
    // C_merged: [batch*M, N]  输出同样沿 M 维拼接

    int MergedM = batch * M;
    auto tiling = select_tiling(MergedM, N, K);

    for (int i = 0; i < MergedM; i += tiling.M_block) {
        for (int j = 0; j < N; j += tiling.N_block) {
            // 标准 GEMM tiling 逻辑
            // batch 维被合并到了 M 维
            // Cube 单元一次 Mmad 处理 16 个元素，不再受 batch 拆分的影响
            LocalTensor<float> a_tile(tiling.M_block, tiling.K_block);
            LocalTensor<float> b_tile(tiling.K_block, tiling.N_block);
            LocalTensor<float> c_tile(tiling.M_block, tiling.N_block);

            DataCopy(a_tile, A[i * K], {tiling.M_block, tiling.K_block});
            DataCopy(b_tile, B[j], {tiling.K_block, tiling.N_block});

            Mmad(c_tile, a_tile, b_tile, {tiling.M_block, tiling.N_block, tiling.K_block});

            DataCopy(C[i * N + j], c_tile, {tiling.M_block, tiling.N_block});
        }
    }
}

合并后的效果：32 个独立的小 GEMM 变成 1 个大 GEMM，Cube 单元始终保持满负荷，不用等 kernel launch。

踩坑一：GEMM 维度的顺序

C = A × B 的维度约定有三种常见记法：M×K × K×N = M×N（标准）、N×K × K×M 和 K×M × M×N。Pytorch 侧 torch.nn.functional.linear 用的是 M×K × K×N，但 torch.matmul 可以接受任意维度组合。混淆维度顺序是高频错误。

错误写法：

import torch, torch_npu

# 错误：维度传反了
# 想要计算 [B, M, K] × [B, K, N] = [B, M, N]
A = torch.randn(8, 1024, 768).npu().half()  # [B, M, K]
B = torch.randn(8, 1024, 768).npu().half()  # [B, N, K] -- 错了! K 应该在最后

# torch.matmul 会做广播但维度对应不对
# 实际计算的是 [8, 1024, 768] × [8, 768, 1024]
# K=768, N=1024 被当成内部维相乘，输出形状 [8, 1024, 1024] 而非预期的 [8, 1024, 768]
C = torch.matmul(A, B)

正确写法：

# 正确：明确 K 维度位置
A = torch.randn(8, 1024, 768).npu().half()   # [B, M, K]
B = torch.randn(8, 768, 1024).npu().half()   # [B, K, N]

C = torch.matmul(A, B)  # [B, 1024, 1024]

# 或者用 transpose 显式控制
B_transposed = B.transpose(-1, -2)            # [B, 1024, 768] -> [B, 768, 1024]
C = torch.matmul(A, B_transposed)             # [B, 1024, 1024]

# 验证：shape 是否正确
assert C.shape == (8, 1024, 1024), f"期望 (8, 1024, 1024), 得到 {C.shape}"

踩坑二：Batch GEMM 的 stride 对齐

Batched GEMM 合并多个矩阵拼成大矩阵时，A 的各 batch 在 HBM 上是否连续存储决定了合并后是否需要额外的内存拷贝。

错误写法：

# 错误：A 的每个 batch 之间 stride 不对齐
A_list = [torch.randn(i+1, 768).npu().half() for i in range(32)]
# A_list[0] shape: [1, 768], A_list[1] shape: [2, 768]
# 不同 batch 的 M 维不等长——不能直接拼接
# 强行拼接会引入 padding 开销

A_padded = torch.nn.utils.rnn.pad_sequence(A_list, batch_first=True)
# padding 填了额外的 0 行，参与无效计算，浪费算力

正确写法：

# 正确：保证所有 batch 的 M 维等长，拼接后 stride 连续
A = torch.stack([
    torch.randn(64, 768) for _ in range(32)
]).npu().half()  # [32, 64, 768] -- 所有 batch M 恒定

B = torch.randn(32, 768, 768).npu().half()

# 方式一：直接 batched matmul
C = torch.bmm(A, B)  # [32, 64, 768], ops-blas 内部自动合并

# 方式二：手动合并，显式控制 stride
A_merged = A.reshape(32 * 64, 768)     # [2048, 768]
B_merged = B.reshape(32 * 768, 768)    # [24576, 768]

C_merged = torch.mm(A_merged, B_merged[:768])  # 如果用相同的 B
# 注意：B 的 batch 不同时需要不同的拼接方式

C++ 侧原理：ops-blas 的 batched GEMM 内部用 lda（leading dimension of A）参数来表示 A 在 HBM 上的实际列步长。如果各 batch 的 M 不等长，lda 每 batch 不同，合并策略失效。

踩坑三：Packing 和转置的隐式开销

Packing 需要对 A 或 B 做额外的内存拷贝和重排。如果矩阵太小，packing 的开销占了总时间的大部分，反而比不 packing 慢。

错误写法：

# 错误：小矩阵也强制 packing
A = torch.randn(64, 128).npu().half()
B = torch.randn(128, 64).npu().half()

# 64x64 的 GEMM 计算量很小，packing 的 HBM 拷贝开销占了大头
# 1. packing: 拷贝 A 和 B 到连续的 Cube-friendly 内存布局: ~3us
# 2. GEMM 计算: ~0.5us（Cube 几乎是瞬间算完的）
# packing 占比 85% -- 得不偿失

# 这种小矩阵应该走 gemm_batched 合并，而不是单独调用

正确写法：

# 正确：收集多个小 GEMM 一起跑
A_batch = []
B_batch = []
for _ in range(256):
    A_batch.append(torch.randn(64, 128).npu().half())
    B_batch.append(torch.randn(128, 64).npu().half())

# 合并后一次 batch GEMM
A_stacked = torch.stack(A_batch)   # [256, 64, 128]
B_stacked = torch.stack(B_batch)   # [256, 128, 64]
C = torch.bmm(A_stacked, B_stacked)  # [256, 64, 64]

# 合并后总计算量是大矩阵规模，packing 占比降到 5% 以下

性能实测

Ascend 910 上，FP16 GEMM 表现：

矩阵尺寸	路径	耗时	Cube 利用率
4096×4096×4096	GEMM_LARGE	1.82 ms	92%
1024×768×768	GEMM_MEDIUM	0.28 ms	78%
64×768×768	GEMM_TINY(单独)	0.12 ms	35%
512×64×768×768	GEMM_TINY(batch合并)	2.10 ms	91%

小矩阵单独调 Cube 利用率只有 35%——kernel launch 和 HBM 初始化占了大头。512 个合并成 batch 后利用率回到 91%，证明合并策略对小矩阵场景是刚需。

---

GEMM 是个数学上极其简单的操作——只有乘法和加法。复杂度全在工程上：多大的块、要不要 packing、怎么处理不规则的 batch。ops-blas 的三路径分法（大矩阵 packing + 中等矩阵双缓冲 + 小矩阵 batch 合并）覆盖了从 attention 到 FC 层的全部矩阵乘场景，Cube 利用率 90% 以上是这套策略体系的目标，不是单个 kernel 的能力。