3.4 矩阵编程：GEMM 与 AI Cube

Elementwise 算子（3.2）走的是 AI Vector 的路；而 Transformer 里占比最大的计算是矩阵乘——它走的是 AI Cube 的路。这一章用 GEMM（通用矩阵乘法）把矩阵编程讲透：数据怎么切、怎么进 L0A/L0B、怎么用 Matmul 接口，以及怎么调优。

📑 目录

1. 矩阵乘法的硬件执行路径
2. 矩阵编程的核心：Cube 的 Tiling
3. 使用 Matmul 高级接口
4. 深入 L0A / L0B / L0C 通路
5. GEMM 性能优化的核心手段
6. 与 CUDA GEMM 优化的对照
7. 从 GEMM 到 Attention：算子融合的预告
总结
延伸阅读

1. 矩阵乘法的硬件执行路径

先回顾第 1 章的知识：昇腾把矩阵乘交给 AI Cube，数据路径是：

GM（HBM）──DataCopy──▶ L0A / L0B ──▶ AI Cube 执行 ──▶ L0C（累积）──▶ UB ──▶ GM

和矢量算子的关键差异：

	矢量算子（3.2）	矩阵算子（本章）
计算单元	AI Vector	AI Cube
工作台	UB	L0A / L0B / L0C
单次计算粒度	一维向量	二维矩阵块
累积精度	通常按输出精度	常为 FP32 累积

所以 GEMM 算子代码里，你的心智模型不再是「一个元素一个元素」，而是「一块矩阵一块矩阵」。

2. 矩阵编程的核心：Cube 的 Tiling

矢量 Tiling 切的是「长度」，矩阵 Tiling 切的是「方块」。以 C[M×N] = A[M×K] × B[K×N] 为例：

L0A 能装下的一小块 A：尺寸记为 (tileM × tileK)
L0B 能装下的一小块 B：尺寸记为 (tileK × tileN)
L0C 累积一小块 C：尺寸 (tileM × tileN)

Tiling 就是决定 tileM / tileK / tileN，然后三重循环遍历所有块：

// 伪代码：遍历所有 C 的小块
for (int m = 0; m < M; m += tileM)
  for (int n = 0; n < N; n += tileN)
    for (int k = 0; k < K; k += tileK)
      // 搬 A[m:m+tileM, k:k+tileK] → L0A
      // 搬 B[k:k+tileK, n:n+tileN] → L0B
      // Cube: L0C += L0A × L0B

选择 tile 尺寸的约束：

L0A + L0B 的总大小不能超过片上容量；
理论上 tileK 越大，K 方向循环越少、数据复用越好；
但 tileK 受 L0B 容量限制，需要权衡；
同时考虑 L2 缓存：让同一块 A 在 N 方向被多个 tile 复用，减少 HBM 访问。

3. 使用 Matmul 高级接口

手搓 L0A/L0B 的细节繁琐且易错，Ascend C 提供了更高层的 Matmul 类接口：

#include "matmul.h"

Matmul<MatmulType<TPosition::GM, CubeFormat::ND>,
       MatmulType<TPosition::GM, CubeFormat::ND>,
       MatmulType<TPosition::GM, CubeFormat::ND>> mm;

// 输入 A、B，输出 C，均为 GM 上的张量
mm.SetInputA(aGlobal, M, K, ...);
mm.SetInputB(bGlobal, K, N, ...);
mm.SetOutput(cGlobal, M, N, ...);

mm.Execute();   // 内部完成 Tiling、L0A/L0B 搬运与 Cube 计算

Matmul 接口把「Tiling 策略、L0 缓冲管理、Cube 指令调度」封装起来。
想进一步优化，可以传入自定义 Tiling 参数，或用更底层的 Cube 指令接口。
建议路径：先用 Matmul 接口跑通，再逐步下钻到 L0 通路手写优化。

4. 深入 L0A / L0B / L0C 通路

如果目标是「手撕 GEMM」或理解 Matmul 内部做了什么，就需要认识 L0 通路：

L0A：A 矩阵分块缓冲（Cube 输入左矩阵）
L0B：B 矩阵分块缓冲（Cube 输入右矩阵）
L0C：结果累积缓冲（FP32 累积）
数据经 DataCopy/专门接口从 GM 或 L2 灌入 L0A/L0B，Cube 一条指令完成一个小矩阵乘，累积到 L0C，最后从 L0C 搬出（可能需要经过向量路径做激活/精度转换）到 UB 再写回 GM。

这组缓冲对应 GPU 世界里的 Tensor Core 输入/累积寄存器路径（mma.sync 的 A/B/C 寄存器）。手动管理 L0 通路 = 昇腾版的「手写 SASS 级 GEMM」，普通业务很少需要，但面试与极致调优时会遇到。

5. GEMM 性能优化的核心手段

昇腾 GEMM 优化思路与 CUDA 同构：

访存/搬运优化：A 在 M-K 循环内复用，B 在 N 循环内复用；连续、对齐搬运。
Tiling 调优：找到「片上利用率最高」的 tileM/tileK/tileN 组合。
多核并行：把 M 维度（或 N 维度）切成多块分给多个 AI Core。
双缓冲：搬运下一块的同时计算当前块，隐藏搬运延迟。
FP32 累积：Cube 内部用高精度累积，避免精度损失（对标 GPU 的 FP32 accumulate）。
算子融合：GEMM 后紧跟的 BiasAdd/激活/量化搬回 GM 前就地完成（见第 7 节）。

6. 与 CUDA GEMM 优化的对照

优化维度	CUDA GEMM	昇腾 GEMM
分块单元	Thread Block Tile（如 128×128）	Cube Tile（tileM/tileK/tileN）
片上缓冲	Shared Memory + 寄存器	L0A/L0B/L0C + UB
数据复用	循环分块 + 寄存器双缓冲	tile 分块 + L0/L2 复用
矩阵乘指令	`mma.sync` / Tensor Core	Cube 指令 / Matmul 接口
多核并行	Grid/Block 切分	按核切分 M/N 维度
调优工具	Nsight Compute	msprof（3.5）

结论：昇腾 GEMM 的优化空间分布与 CUDA 几乎一致——「搬运带宽、片上容量、并行切分」这三个杠杆在哪里都成立。

7. 从 GEMM 到 Attention：算子融合的预告

Transformer 里的计算链是「GEMM → 融合激活/归一化 → 下一个 GEMM」。昇腾上主流的优化方向是把整条链融合进单个算子：

GEMM 输出在 L0C 累积后，不急着写回 GM，而是直接在片上完成 Bias/激活/量化。
FlashAttention 系列在昇腾上的适配，本质上就是「QK^T GEMM → 缩放/Softmax → PV GEMM」的片上融合。
vLLM-Ascend 中的 Attention 实现，就是这套「融合算子 + PagedAttention」的组合。

给读者的提示：掌握了 GEMM + 矢量编程，你就掌握了昇腾上实现 FlashAttention 类算子的全部零件——这也是昇腾算子岗位面试的最高频考点。

总结

矩阵编程走 L0A/L0B → AI Cube → L0C 通路，Tiling 从「切长度」变成「切方块」。
先用 Matmul 高级接口跑通，再按需下钻到 L0 通路手动优化。
GEMM 优化三杠杆：搬运/访存、Tiling 容量利用、多核并行，与 CUDA 完全同构。
GEMM + 矢量编程 = FlashAttention 类融合算子的全部零件，是昇腾算子能力的分水岭。

搜索