为什么 GEMM 是 CUDA 的必修课

通用矩阵乘法：

$$C_{M\times N}=A_{M\times K}B_{K\times N}$$

每个$C_{ij}$ 是长度$K$ 的点积，需要约$2K$ FLOPs（乘和加）。GEMM 同时包含 CUDA 优化的全部核心：

• 二维线程映射；
• 合并访存；
• shared/register 多级分块；
• 同步与双缓冲；
• 资源占用和 occupancy；
• 混合精度与 Tensor Core；
• 与成熟库的基准比较。

朴素 GEMM

假设 row-major：

__global__ void gemm_naive(const float* A, const float* B, float* C,
                           int M, int N, int K) {
  int col = blockIdx.x * blockDim.x + threadIdx.x;
  int row = blockIdx.y * blockDim.y + threadIdx.y;
  if (row >= M || col >= N) return;

  float acc = 0.0f;
  for (int k = 0; k < K; ++k)
    acc = fmaf(A[row * K + k], B[k * N + col], acc);
  C[row * N + col] = acc;
}

同一 warp 计算相邻 col 时：

• B[k*N + col] 是连续读取；
• A[row*K + k] 对同一行的线程是相同地址，可借助缓存/广播；
• 但每个输出元素仍重复读取大量 A/B，跨输出的显式复用不足。

Tiled GEMM

让一个 block 负责 C 的一个 tile，并分批把 A、B tile 放入 shared：

A                              B
┌───────────┐                  ┌───────────────┐
│           │                  │   B tile      │
│  A tile ──┼── K tile ───────┼───────────┐   │
│           │                  │           │   │
└───────────┘                  └───────────┼───┘
                                          ▼
                                  C output tile

K 方向循环：load A_tile + B_tile → barrier → accumulate → barrier

教学版 16×16：

template<int T>
__global__ void gemm_tiled(const float* A, const float* B, float* C,
                           int M, int N, int K) {
  __shared__ float As[T][T];
  __shared__ float Bs[T][T];

  int row = blockIdx.y * T + threadIdx.y;
  int col = blockIdx.x * T + threadIdx.x;
  float acc = 0.0f;

  for (int k0 = 0; k0 < K; k0 += T) {
    int ak = k0 + threadIdx.x;
    int bk = k0 + threadIdx.y;
    As[threadIdx.y][threadIdx.x] =
        (row < M && ak < K) ? A[row * K + ak] : 0.0f;
    Bs[threadIdx.y][threadIdx.x] =
        (bk < K && col < N) ? B[bk * N + col] : 0.0f;
    __syncthreads();

    #pragma unroll
    for (int k = 0; k < T; ++k)
      acc = fmaf(As[threadIdx.y][k], Bs[k][threadIdx.x], acc);
    __syncthreads();
  }

  if (row < M && col < N) C[row * N + col] = acc;
}

这里每个 A/B tile 从 global 读取一次，却被 block 内多个线程复用。边界元素填 0，保证所有线程仍能到达 barrier。

算术强度为什么提高

忽略 C 写回，朴素地为$T\times T$ 个输出完成一个 K tile：

• 计算：约$2T^3$ FLOPs；
• 从 global 读 A/B：$2T^2$ 个元素；
• FP32 字节数：$8T^2$ B。

算术强度近似：

$$AI \approx \frac{2T^3}{8T^2}=\frac{T}{4}\;\text{FLOP/B}$$

Tile 越大，片外复用越高；但 shared memory、寄存器、线程数和边界浪费也随之变化，所以不能无限增大。

生产级 GEMM 的多级分块

教学 kernel 每线程只算一个输出。高性能 GEMM 通常是：

Device/grid tile：把输出矩阵分给 thread blocks
  └─ CTA tile：一个 block 计算较大 C tile
      ├─ shared memory tile：缓存 A/B，可能多 stage 双缓冲
      └─ warp tile：每 warp 负责 C 子块
          └─ thread/register tile：每线程持有多个累加器
              └─ Tensor Core MMA 或标量 FMA

关键技术：

1. Register tiling：每线程算多个 C 元素，A/B 寄存器值复用多次。
2. Vectorized access：合适对齐下减少 load/store 指令。
3. Double/multi buffering：计算 tile k 时预取 k+1。
4. Swizzle/layout：兼顾 global coalescing、shared bank 与 MMA fragment 布局。
5. Split-K：K 很大或 M/N 太小时，把 K 也并行切分，再归约 partial C。
6. Epilogue fusion：在 accumulator 写回前融合 bias、activation、scale、quantize。

每项优化都在“复用/并行度”和“寄存器/shared/同步成本”之间交换资源。

Tensor Core

Tensor Core 执行小矩阵的 MMA（matrix multiply-accumulate）操作：

$$D=A\times B+C$$

它不是一个线程独占的“大号乘法器”，而通常由 warp 或更大协作组共同提供 fragment。支持的 shape、输入/累加类型、布局、对齐与专属指令随 Compute Capability 变化；当前官方 Compute Capability 表列出 FP64、TF32、BF16、FP16、FP8、FP6、FP4、INT8、INT4 等支持矩阵，但不是每代、每个型号都支持全部类型。

使用优先级：

1. cuBLAS/cuBLASLt、cuDNN 等 CUDA-X 库；
2. CUTLASS/CuTe 等模板与 DSL；
3. 高层矩阵 API；
4. inline PTX/架构专属指令，仅在确有必要时。

官方 Compute Capability 章节也建议优先通过 CUDA-X 或 CUTLASS 使用 Tensor Core，而不是直接绑定底层指令。

精度不是数据类型名称这么简单

矩阵计算至少要区分：

误差来自舍入、溢出/下溢、归约顺序、FMA、非确定性原子、fast math 和量化尺度。验证应使用贴近业务的数据分布，不只用 [0,1] 均匀随机数。

常用容差形式：

$$|y-\hat y| \le \text{atol}+\text{rtol}|\hat y|$$

还应观察最大误差、平均误差、ULP、NaN/Inf、下游任务指标，以及误差随 K 或序列长度的增长。

cuBLAS 与 cuBLASLt

cuBLAS 是 BLAS API 的 GPU 实现。传统 BLAS 约定常以 column-major 表述，而 C/C++ 业务数据常为 row-major，必须仔细处理转置标志和 leading dimension。不要靠“交换 A/B 好像能跑”来蒙对，先写清内存解释。

cuBLASLt 提供更灵活的矩阵布局、算法选择、workspace 和 epilogue 融合，适合现代 GEMM 工作流。基本工程模式：

创建 handle
描述 A/B/C/D 的 shape、dtype、layout、leading dimension
描述 matmul：transpose、compute type、scale type、epilogue
查询 heuristic algorithms + workspace
warmup / benchmark / correctness
缓存所选算法（仍需处理设备、shape、版本变化）

库调用也是异步进入指定 stream；handle 的 stream、math mode、pointer mode 与生命周期必须明确。

CUTLASS：从模板组合高性能 GEMM

CUTLASS 把生产级 GEMM 分解为 device/CTA/warp/instruction 层级的模板组件，覆盖数据移动、layout、MMA、pipeline 与 epilogue。它适合：

• 需要库中没有的 fused epilogue；
• 特殊布局或量化格式；
• 研究新架构 tile/pipeline；
• 需要接近库性能，又不愿从 PTX 开始。

代价是模板复杂度、编译时间与架构版本耦合。先用 cuBLASLt 建立上限，再决定是否需要 CUTLASS。

Attention 是数据移动问题

标准 scaled dot-product attention：

$$S=QK^T/\sqrt d,\qquad P=\operatorname{softmax}(S),\qquad O=PV$$

朴素实现会把$N\times N$ 的 score/probability 矩阵写入 HBM，再读回来。长序列时，中间矩阵的数据移动和容量成为核心问题。

Online softmax

对一段值维护当前最大值$m$ 和归一化和$l$。加入新值$x$：

$$m'=\max(m,x)$$

$$l'=l\,e^{m-m'}+e^{x-m'}$$

两个分块统计量$(m_1,l_1)$、$(m_2,l_2)$ 可合并：

$$m=\max(m_1,m_2),\quad l=l_1e^{m_1-m}+l_2e^{m_2-m}$$

因此可以流式处理 score tile，不必先保存整行再做 max/sum。

FlashAttention 的 I/O 主线

for each Q tile:
  Q 留在片上
  for each K/V tile:
    加载 K/V tile
    计算局部 QK^T
    在线更新 row max / normalizer
    按新基准重缩放旧 output accumulator
    累加 softmax(score) × V
  归一化并写回 O tile

核心收益不是减少了 Attention 的主要算术复杂度，而是通过 tiling、online normalization 和融合，避免把完整$N^2$ 中间矩阵往返 HBM。实际实现还要处理 causal mask、padding、dropout、反向传播、数值缩放、不同 head dimension 和硬件专属 pipeline。

卷积与 GEMM

卷积可通过 im2col 转为 GEMM，但显式展开会增加中间内存。高性能库可能使用 implicit GEMM、direct convolution、Winograd、FFT 等算法。选择取决于 kernel size、stride、dilation、channel、batch、dtype、workspace 和确定性要求。

这再次说明：不要仅凭算法名字判断性能。库的 heuristic/autotune 需要在真实 shape 上比较，并缓存结果。

GEMM 基准的正确方式

报告至少包含：

• $M,N,K$、batch、layout、transpose；
• 输入、累加、输出与 compute type；
• alpha/beta、epilogue；
• GPU、时钟/功耗环境、Toolkit/库版本；
• workspace 与算法；
• 预热、重复次数、统计量；
• correctness tolerance；
• 实际 FLOP/s 与 cuBLAS/cuBLASLt 基线。

常用计算：

$$\text{FLOPs}\approx 2MNK$$

如果 beta != 0、epilogue 或 batch 有额外操作，应在说明中列出。Tensor Core 的宣传峰值可能带稀疏、特定精度或特殊统计口径，不能直接拿来当任意 kernel 的分母。

学习项目

1. naive FP32 GEMM：验证索引和边界。
2. shared tiled GEMM：测算术强度和 global traffic 变化。
3. 每线程 2×2/4×4 register tile：观察 register 与 occupancy。
4. 双缓冲：观察 barrier 和 long scoreboard。
5. 对照 cuBLAS：记录性能差距，而不是把“超过库”当作初学目标。
6. 用 cuBLASLt 融合 bias + activation。
7. 实现 online softmax，再画出 attention tile 的 HBM 读写账本。

官方参考

• Compute Capabilities and Tensor Core Types
• cuBLAS Documentation
• CUTLASS Documentation
• CUDA C++ Best Practices: Shared Memory in Matrix Multiplication