GEMM 工程教程：从三重循环到 BF16 / FP8 Tensor Core

0. 学习地图

GEMM 的优化可以看成一个不断减少“无效数据移动”、提高“计算单元饱和度”的过程。

如果只记一条主线：先让每个输出 tile 复用 A/B 数据，再让每个线程保留多个 accumulator，最后让 Tensor Core 在数据搬运期间持续工作。

所以一个成熟 GEMM kernel 的数据流一般是：HBM 读入 A/B tile，到 shared memory 做 CTA 级复用，再进入 warp fragment / register，最后用 Tensor Core 累加到 register accumulator。

0.1 阅读路径

这份文档同时服务新手和已经有 CUDA 基础的读者。建议先按自己的目标选路径，不需要第一次就从头读到尾。

0.2 读前准备

这份文档的目标不是让你背 CUDA 名词，而是让你能看懂一个生产级 GEMM kernel 的性能结构，并能判断一个优化点为什么有效。

0.2.1 术语速查

1. 数学定义和形状

标准 GEMM 写作：

C = alpha * A * B + beta * C

教学里先忽略 alpha 和 beta，关注核心乘法：

A: [M, K]
B: [K, N]
C: [M, N]

C[m, n] = sum(A[m, k] * B[k, n] for k in 0..K-1)

计算量近似为 2 * M * N * K FLOPs，因为一次 multiply-add 通常算 2 FLOPs。

1.1 为什么 GEMM 值得专门优化

GEMM 的计算量随 M*N*K 增长，但输入输出数据量只随 M*K + K*N + M*N 增长。如果 tile 设计合理，同一份 A/B 数据可以被重复使用很多次，kernel 就更容易接近计算峰值。

这也是 roofline 模型的核心：当 arithmetic intensity 低时，瓶颈是内存带宽；当它足够高时，瓶颈才可能变成 Tensor Core / CUDA core 吞吐。

1.2 Layout 先于 kernel

矩阵在内存里通常是 row-major 或 column-major。对 GPU 来说，线程束访问连续地址时才能合并成更少的 memory transaction。一个数学上等价的转置或 layout 选择，可能直接决定 global load 是否 coalesced。

2. 一般写法：三重循环

最直接的写法如下：

for m in range(M):
  for n in range(N):
    acc = 0.0
    for k in range(K):
      acc += A[m, k] * B[k, n]
    C[m, n] = acc

2.1 循环顺序会改变缓存行为

即使在 CPU 上，m-n-k、m-k-n、k-m-n 的性能也可能差很多，因为 B 的访问是连续还是跨 stride，会影响 cache line 利用率。GPU 上这个问题更尖锐：warp 是否 coalesced 比单个线程的局部性更重要。

// row-major B[K, N] 下，n 连续时 B[k, n] 连续
for m:
  for k:
    a = A[m, k]
    for n:
      C[m, n] += a * B[k, n]

这段 CPU 伪代码已经开始表达“复用 A、连续读 B”的思想。GPU kernel 的 tiling 只是把这个思想推进到 CTA、warp、register 三个层级。

3. Naive GPU Kernel：一个线程算一个输出元素

第一版 GPU 化通常是：每个 thread 负责一个 C[m,n]。

__global__ void matmul_naive_bf16(A, B, C, M, N, K) {
  int m = blockIdx.y * blockDim.y + threadIdx.y;
  int n = blockIdx.x * blockDim.x + threadIdx.x;
  if (m >= M || n >= N) return;

  float acc = 0.0f;
  for (int k = 0; k < K; ++k) {
    acc += float(A[m * K + k]) * float(B[k * N + n]);
  }
  C[m * N + n] = bf16(acc);
}

这比 CPU 三重循环更并行，但离高性能仍很远。主要问题通常不是乘加数量，而是 A/B 数据缺少 CTA 内复用，导致相对计算量产生了过多 HBM 访问。

3.1 Naive kernel 的访存问题

假设一个 block 有很多线程同时计算相邻的 C[m,n]。它们在同一个 k 上往往会重复读取同一个 A[m,k]，并读取一段连续的 B[k,n]。B 的访问可能还能合并，A 的读取却会在多个线程之间重复发生。

3.2 第一个性能判断

如果 profile 里看到 global load throughput 很高，但 Tensor Core 或 FMA 利用率很低，通常说明 kernel 被内存喂不饱。这时不要急着调 block size，先问：A/B 是否被重复从 HBM 读了太多次？

3.3 从零能跑的 CUDA GEMM demo

只看伪代码很难建立性能直觉。这里提供一个单文件 CUDA demo：同一个 FP32 GEMM 同时实现 naive kernel 和 shared-memory tiled kernel，并用 CUDA event 计时。它不是生产级 Tensor Core kernel，而是用来亲眼看到“重复读 HBM”和“CTA 内复用 A/B tile”的差别。

nvcc -O3 -std=c++17 -arch=sm_80 gemm-demo.cu -o gemm-demo
./gemm-demo
./gemm-demo 1024 1024 1024 50

-arch=sm_80 只是示例，应该替换成你的 GPU 架构。这个 demo 只依赖 CUDA runtime，不依赖 CUTLASS、Triton 或 cuBLAS。

__global__ void matmul_tiled_kernel(const float* A, const float* B, float* C,
                                    int M, int N, int K) {
  __shared__ float As[16][16];
  __shared__ float Bs[16][16];
  int n = blockIdx.x * 16 + threadIdx.x;
  int m = blockIdx.y * 16 + threadIdx.y;
  float acc = 0.0f;

  for (int k0 = 0; k0 < K; k0 += 16) {
    As[threadIdx.y][threadIdx.x] = in_bounds_A ? A[m * K + k0 + threadIdx.x] : 0.0f;
    Bs[threadIdx.y][threadIdx.x] = in_bounds_B ? B[(k0 + threadIdx.y) * N + n] : 0.0f;
    __syncthreads();
    for (int kk = 0; kk < 16; ++kk) acc += As[threadIdx.y][kk] * Bs[kk][threadIdx.x];
    __syncthreads();
  }
  if (m < M && n < N) C[m * N + n] = acc;
}

// gemm-demo.cu
// Minimal CUDA GEMM lab: naive one-thread-per-output vs shared-memory tiling.
//
// Build on a CUDA machine:
//   nvcc -O3 -std=c++17 -arch=sm_80 gemm-demo.cu -o gemm-demo
//
// Run:
//   ./gemm-demo
//   ./gemm-demo 1024 1024 1024 50
//   ./gemm-demo 1024 1024 1024 50 --skip-check

#include <cuda_runtime.h>

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <vector>

#define CUDA_CHECK(call)                                                     \
  do {                                                                       \
    cudaError_t err__ = (call);                                               \
    if (err__ != cudaSuccess) {                                               \
      std::cerr << "CUDA error at " << __FILE__ << ":" << __LINE__ << ": " \
                << cudaGetErrorString(err__) << std::endl;                   \
      std::exit(EXIT_FAILURE);                                                \
    }                                                                        \
  } while (0)

constexpr int TILE = 16;

__global__ void matmul_naive_kernel(const float* __restrict__ A,
                                    const float* __restrict__ B,
                                    float* __restrict__ C,
                                    int M,
                                    int N,
                                    int K) {
  int n = blockIdx.x * blockDim.x + threadIdx.x;
  int m = blockIdx.y * blockDim.y + threadIdx.y;

  if (m >= M || n >= N) {
    return;
  }

  float acc = 0.0f;
  for (int k = 0; k < K; ++k) {
    acc += A[m * K + k] * B[k * N + n];
  }
  C[m * N + n] = acc;
}

__global__ void matmul_tiled_kernel(const float* __restrict__ A,
                                    const float* __restrict__ B,
                                    float* __restrict__ C,
                                    int M,
                                    int N,
                                    int K) {
  __shared__ float As[TILE][TILE];
  __shared__ float Bs[TILE][TILE];

  int tx = threadIdx.x;
  int ty = threadIdx.y;
  int n = blockIdx.x * TILE + tx;
  int m = blockIdx.y * TILE + ty;

  float acc = 0.0f;
  int k_tiles = (K + TILE - 1) / TILE;

  for (int tile = 0; tile < k_tiles; ++tile) {
    int a_col = tile * TILE + tx;
    int b_row = tile * TILE + ty;

    As[ty][tx] = (m < M && a_col < K) ? A[m * K + a_col] : 0.0f;
    Bs[ty][tx] = (b_row < K && n < N) ? B[b_row * N + n] : 0.0f;
    __syncthreads();

    #pragma unroll
    for (int kk = 0; kk < TILE; ++kk) {
      acc += As[ty][kk] * Bs[kk][tx];
    }
    __syncthreads();
  }

  if (m < M && n < N) {
    C[m * N + n] = acc;
  }
}

void init_matrix(std::vector<float>& x, int rows, int cols, int salt) {
  for (int r = 0; r < rows; ++r) {
    for (int c = 0; c < cols; ++c) {
      int v = (r * 17 + c * 13 + salt) % 97;
      x[r * cols + c] = static_cast<float>(v - 48) / 48.0f;
    }
  }
}

void cpu_reference(const std::vector<float>& A,
                   const std::vector<float>& B,
                   std::vector<float>& C,
                   int M,
                   int N,
                   int K) {
  for (int m = 0; m < M; ++m) {
    for (int n = 0; n < N; ++n) {
      float acc = 0.0f;
      for (int k = 0; k < K; ++k) {
        acc += A[m * K + k] * B[k * N + n];
      }
      C[m * N + n] = acc;
    }
  }
}

float max_abs_error(const std::vector<float>& ref, const std::vector<float>& got) {
  float err = 0.0f;
  for (size_t i = 0; i < ref.size(); ++i) {
    err = std::max(err, std::abs(ref[i] - got[i]));
  }
  return err;
}

template <typename Kernel>
float time_kernel(Kernel kernel,
                  dim3 grid,
                  dim3 block,
                  const float* d_A,
                  const float* d_B,
                  float* d_C,
                  int M,
                  int N,
                  int K,
                  int iters) {
  cudaEvent_t start;
  cudaEvent_t stop;
  CUDA_CHECK(cudaEventCreate(&start));
  CUDA_CHECK(cudaEventCreate(&stop));

  kernel<<<grid, block>>>(d_A, d_B, d_C, M, N, K);
  CUDA_CHECK(cudaGetLastError());
  CUDA_CHECK(cudaDeviceSynchronize());

  CUDA_CHECK(cudaEventRecord(start));
  for (int i = 0; i < iters; ++i) {
    kernel<<<grid, block>>>(d_A, d_B, d_C, M, N, K);
  }
  CUDA_CHECK(cudaEventRecord(stop));
  CUDA_CHECK(cudaEventSynchronize(stop));
  CUDA_CHECK(cudaGetLastError());

  float ms = 0.0f;
  CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop));
  CUDA_CHECK(cudaEventDestroy(start));
  CUDA_CHECK(cudaEventDestroy(stop));
  return ms / static_cast<float>(iters);
}

double gflops(int M, int N, int K, float ms) {
  double flops = 2.0 * static_cast<double>(M) * static_cast<double>(N) * static_cast<double>(K);
  return flops / (static_cast<double>(ms) * 1.0e6);
}

int parse_arg(char** argv, int index, int fallback) {
  if (argv[index] == nullptr) {
    return fallback;
  }
  int value = std::atoi(argv[index]);
  return value > 0 ? value : fallback;
}

bool has_flag(int argc, char** argv, const char* flag) {
  for (int i = 1; i < argc; ++i) {
    if (std::strcmp(argv[i], flag) == 0) {
      return true;
    }
  }
  return false;
}

int main(int argc, char** argv) {
  int M = argc > 1 ? parse_arg(argv, 1, 512) : 512;
  int N = argc > 2 ? parse_arg(argv, 2, 512) : 512;
  int K = argc > 3 ? parse_arg(argv, 3, 512) : 512;
  int iters = argc > 4 ? parse_arg(argv, 4, 20) : 20;
  bool skip_check = has_flag(argc, argv, "--skip-check");

  std::cout << "GEMM demo shape: M=" << M << " N=" << N << " K=" << K
            << " iters=" << iters << std::endl;
  if (skip_check) {
    std::cout << "correctness: skipped (--skip-check)" << std::endl;
  }

  std::vector<float> h_A(static_cast<size_t>(M) * K);
  std::vector<float> h_B(static_cast<size_t>(K) * N);
  std::vector<float> h_ref(skip_check ? 0 : static_cast<size_t>(M) * N);
  std::vector<float> h_naive(static_cast<size_t>(M) * N);
  std::vector<float> h_tiled(static_cast<size_t>(M) * N);

  init_matrix(h_A, M, K, 7);
  init_matrix(h_B, K, N, 19);
  if (!skip_check) {
    cpu_reference(h_A, h_B, h_ref, M, N, K);
  }

  float* d_A = nullptr;
  float* d_B = nullptr;
  float* d_C = nullptr;
  size_t bytes_A = h_A.size() * sizeof(float);
  size_t bytes_B = h_B.size() * sizeof(float);
  size_t bytes_C = static_cast<size_t>(M) * N * sizeof(float);

  CUDA_CHECK(cudaMalloc(&d_A, bytes_A));
  CUDA_CHECK(cudaMalloc(&d_B, bytes_B));
  CUDA_CHECK(cudaMalloc(&d_C, bytes_C));
  CUDA_CHECK(cudaMemcpy(d_A, h_A.data(), bytes_A, cudaMemcpyHostToDevice));
  CUDA_CHECK(cudaMemcpy(d_B, h_B.data(), bytes_B, cudaMemcpyHostToDevice));

  dim3 block(TILE, TILE);
  dim3 grid((N + TILE - 1) / TILE, (M + TILE - 1) / TILE);

  float naive_ms = time_kernel(matmul_naive_kernel, grid, block, d_A, d_B, d_C, M, N, K, iters);
  CUDA_CHECK(cudaMemcpy(h_naive.data(), d_C, bytes_C, cudaMemcpyDeviceToHost));
  float naive_err = skip_check ? NAN : max_abs_error(h_ref, h_naive);

  float tiled_ms = time_kernel(matmul_tiled_kernel, grid, block, d_A, d_B, d_C, M, N, K, iters);
  CUDA_CHECK(cudaMemcpy(h_tiled.data(), d_C, bytes_C, cudaMemcpyDeviceToHost));
  float tiled_err = skip_check ? NAN : max_abs_error(h_ref, h_tiled);

  std::cout << "kernel, avg_ms, GFLOP/s, max_abs_error" << std::endl;
  std::cout << "naive, " << naive_ms << ", " << gflops(M, N, K, naive_ms)
            << ", " << naive_err << std::endl;
  std::cout << "tiled, " << tiled_ms << ", " << gflops(M, N, K, tiled_ms)
            << ", " << tiled_err << std::endl;

  CUDA_CHECK(cudaFree(d_A));
  CUDA_CHECK(cudaFree(d_B));
  CUDA_CHECK(cudaFree(d_C));
  return 0;
}

运行输出会包含 correctness、max absolute error、平均耗时和 GFLOP/s。做 profiling 时可以加 --skip-check 跳过 CPU reference，避免时间线里混入 CPU 端三重循环。不要把 demo 数字当成硬件上限：naive/tiled 只是教学 baseline，真正的 BF16/FP8 高性能路径还要进入 Tensor Core、pipeline、epilogue 和库级调参。

nsys profile -o gemm_demo_timeline --trace=cuda,nvtx,osrt ./gemm-demo 1024 1024 1024 20 --skip-check
ncu --set full --kernel-name regex:matmul_ ./gemm-demo 1024 1024 1024 20 --skip-check

4. Block Tiling：让一个 CTA 负责一个 C tile

优化的第一步是把 C 切成 tile。比如一个 CTA 负责 128 x 128 的 C tile。

for block_m in range(0, M, BM):
  for block_n in range(0, N, BN):
    C_tile[BM, BN] = 0
    for block_k in range(0, K, BK):
      A_tile = A[block_m:block_m+BM, block_k:block_k+BK]
      B_tile = B[block_k:block_k+BK, block_n:block_n+BN]
      C_tile += A_tile @ B_tile

这里的关键变化是：A_tile 会被同一个 C tile 的多个列复用，B_tile 会被多个行复用。

4.1 CTA tile、warp tile、MMA tile 是三层结构

不要把 BM/BN/BK 理解成一个孤立参数。实际 kernel 通常会把一个 CTA tile 再切给多个 warp，每个 warp 再执行多个 MMA tile。

CTA tile:   128 x 128 x 64
Warp tile:   64 x  64 x 64   // 一个 CTA 里多个 warp 分工
MMA tile:    16 x   8 x 16   // 具体 Tensor Core 指令形状示意

CTA tile 太小，数据复用少，调度开销相对变大；CTA tile 太大，shared memory 和 register 使用量会上升，occupancy 下降。高性能 GEMM 本质上是在这些约束之间找平衡点。

4.2 为什么 BK 很关键

BK 决定每个 mainloop stage 搬多少 K 维数据。BK 大会增加每次搬运的数据和 shared memory 占用，也可能让单个 stage 的计算更饱满；BK 小则 pipeline 粒度细，但 loop overhead 和同步占比可能更高。

例如 BF16 每个元素 2 bytes，BM=128, BN=128, BK=64, stages=2 时，只存 A/B tile 就大约需要 (128*64 + 64*128)*2*2 = 65536 bytes shared memory，还没算 padding 或特殊 layout。

4.3 BM / BN / BK 怎么量化选择

这确实是 GEMM 最难的参数之一，因为它不是单目标优化。一个 tile shape 同时决定数据复用、CTA 数量、shared memory 占用、寄存器数量、tail 浪费、warp 分工和 Tensor Core 指令排布。

比较实用的方法不是“背一个固定组合”，而是先用资源约束过滤候选，再用 profile 选择。

4.4 一个可执行的选型流程

1. 先根据 dtype 和架构选 MMA 指令族，例如 BF16 Tensor Core 的基本 MMA 形状。
2. 让 BM/BN/BK 都对齐到 MMA 形状和 vectorized load 粒度，避免主路径里到处都是 tail。
3. 估算 shared memory：把 stages 算进去，再留出 padding / swizzle 的空间。
4. 估算寄存器：accumulator fragment 是大头，再加 A/B fragment、地址、predicate、epilogue 临时变量。
5. 检查 CTA 并行度：如果 ceil(M/BM) * ceil(N/BN) 太小，即使单 CTA 很快，整卡也可能吃不满。
6. 用 profile 比较候选，观察 Tensor Core 利用率、HBM 带宽、shared conflict、occupancy 和 local memory。

// 候选生成的思路，不是固定答案
for BM, BN in [(64, 64), (64, 128), (128, 64), (128, 128), (128, 256)]:
  for BK in [32, 64, 128]:
    if not mma_aligned(BM, BN, BK): continue
    if smem_bytes(BM, BN, BK, stages) > smem_budget: continue
    if estimated_registers(BM, BN, BK) > register_budget: continue
    benchmark(BM, BN, BK)

4.5 不同 shape 的倾向

4.6 Split-K：当 M/N 并行度不够时，拆 K 维

前面的 block tiling 默认一个 C_tile[BM, BN] 由一个 CTA 沿完整 K 维算完。这个设计在大 M、大 N 的 GEMM 上通常很好，因为 ceil(M/BM) * ceil(N/BN) 已经能提供足够多 CTA。但在小 M、大 K，或者 M/N tile 数量很少的场景，整张卡可能只有很少 CTA 可调度，Tensor Core 明明很快，SM 却吃不满。

Split-K 的想法是：不要让一个 CTA 独占完整 K 维，而是把 K 维切成多个 slice，让多个 CTA 同时计算同一个 C tile 的不同 partial sum，最后再把 partial sum reduce 成最终 C。

// 普通 tiling：一个 CTA 算完整 K
for cta_m, cta_n:
  acc = 0
  for k0 in range(0, K, BK):
    acc += A[cta_m, k0] @ B[k0, cta_n]
  C[cta_m, cta_n] = acc

// Split-K：多个 CTA 分别算 K slice，然后再 reduce
for cta_m, cta_n, split_id:
  k_begin, k_end = split_range(split_id, K)
  partial = 0
  for k0 in range(k_begin, k_end, BK):
    partial += A[cta_m, k0] @ B[k0, cta_n]
  Workspace[split_id, cta_m, cta_n] = partial

reduce Workspace[:, cta_m, cta_n] -> C[cta_m, cta_n]

4.6.1 Split-K 解决的不是复用，而是并行度

Split-K 不会神奇地减少总计算量。它通常还会增加额外的 workspace 写读和 reduction 成本。它真正解决的是 CTA 数量不足：当 M/N tile 太少时，把 K 拆开可以制造更多独立 work，让更多 SM 同时工作。

4.6.2 两种 reduction 路线

Split-K 的关键不是前半段 partial GEMM，而是后半段怎么合并 partial sums。常见有两条路线。

4.6.3 Split-K 的选择条件

Split-K 应该是一个 profile-driven 的开关，而不是默认打开。一个简单判断是先算普通 tiling 的 CTA 数：

如果 num_cta 已经远大于 SM 数，Split-K 往往只会引入 reduction overhead；如果 num_cta 明显小于 SM 数，且 K 很大，Split-K 才可能值得尝试。

4.6.4 Split-K 和 Stream-K 不一样

有些库文档还会提到 Stream-K。它和 Split-K 都是在 K 维和 tile 调度上做文章，但目标更偏向负载均衡和持续填满 SM。Split-K 更容易理解为“把一个 C tile 的 K 维拆给多个 worker”；Stream-K 通常会把全局 GEMM 的 tile/K 工作以更流式的方式分配给 worker，减少尾部不均衡。学习时可以先掌握 Split-K，再把 Stream-K 看成更工程化的调度策略。

5. Shared Memory：把 A/B tile 搬到片上

Block tiling 只是逻辑复用。要真正减少 HBM 访问，通常需要把当前 A_tile 和 B_tile 搬到 shared memory。

// 伪代码
for k0 in range(0, K, BK):
  shared_A = load_global_to_shared(A_tile)
  shared_B = load_global_to_shared(B_tile)
  __syncthreads()

  for kk in range(BK):
    acc += shared_A[row, kk] * shared_B[kk, col]
  __syncthreads()

这样 block 内线程会从 shared memory 复用 A/B 数据，而不是每次都打到 HBM。

5.1 Coalesced global load

从 HBM 搬到 shared 时，首先要让 global memory load 合并。常见做法是让每个线程搬连续的 16 bytes，例如 float4、int4 或等价 vectorized load。这里的 16B 是单线程指令宽度和自然对齐要求；32 个 lane 合起来会请求 512B，但这 512B 不会成为一个单独的 512B transaction。当前 CUDA 文档对 Compute Capability 6.0 及之后的通用解释，是硬件按覆盖到的 32B segment 数量合并请求。

// 概念示意：每个线程搬 16B
int lane = threadIdx.x % 32;       // 当前线程在 warp 内的编号
int vec_id = block_thread_id;      // CTA 内连续的线性线程编号

const char* src = reinterpret_cast<const char*>(global_ptr) + vec_id * 16;
char* dst = reinterpret_cast<char*>(shared_ptr) + vec_id * 16;

// debug 时检查实际地址，而不只是检查 cudaMalloc 返回的基地址
assert((reinterpret_cast<uintptr_t>(src) & 0xF) == 0);
assert((reinterpret_cast<uintptr_t>(dst) & 0xF) == 0);

uint4 x = *reinterpret_cast<const uint4*>(src);  // 16B global load
*reinterpret_cast<uint4*>(dst) = x;              // 16B shared store

uint4 由 4 个 32-bit 无符号整数组成，因此 sizeof(uint4) = 4 * 4B = 16B。vec_id 每增加 1，地址增加 16B；只要起始地址 16B 对齐，连续线程的每次 vector load 就都满足 16B 对齐。上面的 lane 在简化代码里没有参与地址计算，它只是提示真实 kernel 通常还会用 lane id 设计 warp 内的数据分工。

5.1.1 从矩阵下标判断 16B 对齐

对 row-major 矩阵，线程从 A[row, col] 开始搬运时，有效字节地址是：

如果 base(A) 已经 16B 对齐，那么只需继续判断 (row * lda + col) * sizeof(T) 是否为 16 的倍数。注意：设备分配返回的基地址通常已经满足较强对齐，但 subview、切片、tile 起点和非整齐的 leading dimension 仍然可能破坏有效地址的对齐。

// A 是 row-major BF16，base(A) 已满足 16B 对齐
const __nv_bfloat16* src = A + row * lda + k0;
bool aligned_16b =
    (reinterpret_cast<uintptr_t>(src) & 0xF) == 0;

// lda=128、k0 是 8 的倍数：任意 row 都满足 16B 对齐
// (row * 128 + k0) * 2B 是 16B 的倍数

// lda=130、row=1、k0=0：第二行行首没有 16B 对齐
// (1 * 130 + 0) * 2B = 260B，260 % 16 = 4

因此，高性能 kernel 往往会把 lda padding 到 vector width 的整数倍，让 mainloop 走对齐的 16B load；对矩阵开头、结尾或 K-tail 中不满足条件的元素，再使用 masked load、标量 load，或者单独的 fallback kernel。

5.1.2 16B、32B、128B 和 512B 分别是什么

先给结论：warp 在这段代码里总共请求 512B，但 512B 只是 payload，不是 transaction size，所以没有“warp 首地址必须 512B 对齐”的规则。对齐要求应该跟它所属的层级一起读。

5.1.2.1 为什么 512B payload 不要求 512B 对齐

一条 warp load 发出后，coalescer 会统计 32 个 lane 的地址覆盖了多少个 transaction segment，而不是把所有 lane 拼成一个 512B 原子 transaction。对于连续的 16B-per-lane 访问：

只要 base % 32 == 0，512B 连续范围就正好覆盖 512 / 32 = 16 个 transaction。把 base 从 128B 对齐进一步提高到 512B 对齐，并不会把 16 个 transaction 继续减少。

5.1.2.2 那为什么 GEMM 里仍经常看到 128B 对齐

128B 属于更高一级的布局选择，而不是由 32 lanes × 16B 推导出来的 warp 对齐要求。它常见有三个原因：

1. Cache-line / segment 边界：部分旧架构和文档用 128B L1 segment 解释 coalescing；即使当前数据访问单元按 32B 计算，128B 对齐也能让四个连续 32B sector 自然落入同一条 128B line。
2. Tile 和 pipeline stage 易于管理：GEMM 会反复搬运 A/B tile。让每行或每个 stage 从 128B 边界开始，地址计算、cache-line 利用、shared layout 和异步流水线通常更规整。
3. 历史 API 性能指导：部分旧版 cuda::memcpy_async 文档曾建议 global/shared 都按 128B 对齐。当前 Cooperative Groups 文档则推荐 16B，因此这不是跨版本、跨 API 的统一硬规则。

base=0x020 是一个有用的反例：它已经 32B 对齐，因此仍是最少的 16 个 32B transaction；但 512B 范围会跨 5 条 128B line。也就是说，128B 对齐可能改善 cache-line/tile 组织，却不等于它是 coalescing 的最低门槛。

5.1.2.3 Async copy 应该检查哪种对齐

低层 cp.async 家族按线程支持 4B、8B、16B copy，指针至少要满足所选 copy size 的自然对齐。对于本节每线程 16B 的写法，首先必须证明 global source 和 shared destination 都满足 16B 对齐。

至于 tile 是否进一步 32B、128B 或更高对齐，要看具体 CUDA 版本、API、架构以及使用的是 cp.async、Cooperative Groups、TMA 还是其他 copy atom。不要把某个 API 某一版文档的“最佳性能建议”改写成所有 warp load 的硬件要求。

5.1.2.4 实际写 kernel 时按这个顺序判断

1. 检查每个 lane 的实际地址：16B vector load 必须满足 addr % 16 == 0。
2. 检查 lane 间地址是否连续，确保请求能被 coalescer 合并。
3. 在当前 32B transaction 模型下，检查 warp 首地址是否满足 base % 32 == 0，避免首尾多覆盖一个 segment。
4. 把 128B 当作常见的 tile/cache-line 强化对齐，再通过 profile 判断它是否有额外收益。
5. 只有具体 copy API、TMA descriptor 或自定义 allocator 明确要求时，才继续要求 256B、512B 等更强对齐；不要从 warp payload 大小反推对齐要求。

5.2 Bank conflict 是什么

Bank conflict 不是“多个线程同时访问 shared memory”这么简单。真正的判断条件是：同一个 warp 发出一次 shared memory 指令时，是否有多个 lane 访问了同一个 bank 中的不同地址。如果是，硬件就要把这次请求拆成多轮；如果 32 个 lane 的地址落到 32 个不同 bank，则这些 bank 可以并行服务。

5.2.1 Bank 到底是什么：先看 shared memory 的硬件模型

从 CUDA 程序看，shared memory 是一个普通的线性地址空间；你声明的是数组或 byte buffer，并不会手工选择 bank。硬件为了让一个 warp 的多个地址能够并行访问，把这片片上 SRAM 在逻辑上组织成 32 个可同时工作的带宽分区，这些分区就是 bank。

可以把一次 shared load 想成下面的流水：warp 提供最多 32 个 lane address；地址解码器从每个地址里算出 bank id 和 bank 内位置；路由网络把请求送到 32 个 bank。不同 bank 可以并行返回，同一 bank 的不同位置则需要拆成多轮。

Bank 不是只保存一个 word。每个 bank 内部还有很多位置。连续地址会横向铺到 B0...B31，铺满后再回到 B0 的下一个 depth。下面只画前 8 个 bank；真实硬件同一行还有 B8...B31。

BF16/FP16 元素只有 2B，所以两个相邻元素位于同一个 32-bit bank word；FP8/INT8 则一个 word 可容纳四个元素。真实 kernel 常使用 packed/vector load、ldmatrix 或更宽 transaction，分析时要把实际指令拆成它触及的 32-bit bank word，不能机械地把“一个元素”当成“一个 bank”。

把路由结果分成三种情况，Bank conflict 的定义就很直观：

5.2.2 地址怎样映射到 bank

因此前 32 个 32-bit word 分别进入 bank 0 到 bank 31；第 32 个 word 又回到 bank 0。Bank 是交错编址，不是“bank 0 存一整段地址、bank 1 再存下一整段”。

5.2.3 把问题放回 GEMM：一个 warp 读取 A 的同一列

考虑一个教学用的 warp 映射。把两个 BF16 打包成一个 32-bit word，shared A tile 记作 A_pack[32][32]；它在 K 方向覆盖 64 个 BF16。warp 的 lane r 负责输出列 n0 上的不同行，因此在某个固定 k_pair 上执行：

当 k_pair=0 时，lane 0...31 要读取 A tile 的同一逻辑列 A_pack[0...31][0]。B 侧所有 lane 读取的是完全相同的 B_pack[0][n0]，可以 broadcast；真正出问题的是 A 侧“同列、不同地址”的读取。

5.2.4 普通 row-major layout 为什么产生 32-way conflict

A_pack[32][32] 的一行正好是 32 个 32-bit word。普通 row-major 下：

固定 k_pair=0 后，无论 lane 是多少，bank 都是 0。每个 lane 读的是不同 row，因此地址不同，不能 broadcast；硬件必须把一次 warp request 拆成最多 32 轮 conflict-free request，这就是 32-way bank conflict。

5.2.5 Swizzle：不改逻辑矩阵，只改物理地址

Swizzle 的核心是对 shared memory 的地址做一个可逆置换。写入和读取都使用同一映射，所以数学上仍然是原来的 A[row][col]；只是逻辑列不再总是落在相同的物理列和 bank。为了把原理画清楚，先使用最简单的 32-bit word 级 XOR：

// 教学映射：tile width 是 32 个 32-bit word
physical_col = logical_col ^ row;

// global -> shared：把逻辑元素写到置换后的物理位置
smem[row][physical_col] = A_pack[row][logical_col];

// shared -> register：用同一映射找回原来的逻辑元素
x = smem[row][logical_col ^ row];

Swizzle 后，在固定 k_pair 时：

对固定的 k_pair，XOR 是一个一一映射：lane 0...31 会得到 0...31 的某种排列，不会有两个 lane 落到同一个 bank。以 k_pair=0 为例，physical column 就等于 lane id：

5.2.6 为什么生产 kernel 的 Swizzle 看起来更复杂

physical_col = logical_col ^ row 是为了讲清 bank 置换的 32-bit word 级模型。真实 Tensor Core GEMM 还要同时满足三类约束：

1. global → shared 通常按 16B vector copy，不能随意打散一个 vector 内的元素。
2. shared → register 可能使用 ldmatrix、TMA 或架构特定 fragment layout，lane 的读取模式不是简单的一 lane 一 word。
3. A/B 的 dtype、major mode、MMA tile 和 stage layout 不同，参与 XOR 的地址 bit 也不同。

因此 CUTLASS/CuTe 的 Swizzle 通常表现为“保留最低若干对齐 bit，再把一组高位 bit XOR 到地址低位”。它本质上仍是同一件事：保持 vector alignment，同时把消费者会同时访问的地址分散到不同 bank。

进一步阅读：CUDA C++ Best Practices Guide：Shared Memory and Memory Banks，以及 CUTLASS CuTe Swizzle API。

5.3 Padding 什么时候有用

padding 是最简单的破冲突方法之一：如果每行长度刚好导致 stride 命中同一组 bank，就给 shared tile 的 leading dimension 多加几个元素。

// 不是固定要 +1，而是用 padding 改变行间 stride
__shared__ bf16 smem_a[BM][BK + PAD];

padding 的代价是 shared memory 占用增加，并且可能影响 vectorized store/load 对齐。因此它适合教学和简单 kernel；在生产级 Tensor Core kernel 里，更常见的是专门设计 shared layout。

6. Register Tiling：每个线程算多个输出元素

如果一个 thread 只算一个 C[m,n]，每次加载 A/B 后只产生一个 accumulator 的贡献，计算密度不够。Register tiling 会让每个 thread 持有多个 accumulator，例如 4 x 4 的小块。

float acc[TM][TN] = {0};

for k0 in range(0, K, BK):
  load A/B tile
  for kk in range(BK):
    a_frag[TM] = load A rows for this thread
    b_frag[TN] = load B cols for this thread
    for i in range(TM):
      for j in range(TN):
        acc[i][j] += a_frag[i] * b_frag[j]

这一步提升 arithmetic intensity，但也会增加寄存器压力。寄存器太多会降低 occupancy，反而可能变慢。

6.1 accumulator 为什么放寄存器

每个 C[m,n] 都要沿 K 累加。如果每次 partial sum 都写回 shared 或 global，再读出来继续累加，带宽和延迟都会爆炸。所以高性能 GEMM 会让 accumulator 长时间留在寄存器里，直到 K 维全部处理完后再进入 epilogue。

例如一个线程维护 8x4 个 FP32 accumulator，仅 accumulator 就需要 32 个寄存器。再加上 fragment 和临时变量，很容易达到 80、120 甚至更多寄存器。

6.2 寄存器越多不一定越快

更多 register tiling 会提高单线程计算密度，但会降低一个 SM 上能同时驻留的 warp 或 CTA 数。如果 occupancy 太低，kernel 对内存延迟、指令调度空洞、同步等待的隐藏能力会下降。

6.3 Register tiling 和 warp-level 分工

在 Tensor Core kernel 里，寄存器不只是普通标量数组。A/B fragment 会按照 MMA 指令要求分布在 warp 的多个 lane 上，accumulator fragment 也由整个 warp 共同表示。单个线程看到的是自己那部分 fragment，整个 warp 合起来才是一个完整的 MMA tile。

7. BF16 Tensor Core：真正的高吞吐路径

在 NVIDIA GPU 上，BF16 GEMM 的高性能路径通常不是普通 CUDA core FMA，而是 Tensor Core MMA 指令。

// 概念模型，不是完整可编译代码
for k0 in range(0, K, BK):
  load A/B tiles into shared memory
  for warp_k in range(0, BK, mma_k):
    a_frag = ldmatrix_or_load_matrix_A(shared_A)
    b_frag = ldmatrix_or_load_matrix_B(shared_B)
    acc_frag = mma_sync(acc_frag, a_frag, b_frag)

BF16 常见策略是：输入 A/B 为 BF16，Tensor Core 做 BF16 multiply，accumulator 使用 FP32。最后写回时再转成 BF16 或 FP32，取决于输出需求。

7.1 MMA 指令在做什么

Tensor Core 指令可以理解为 warp 级别的小矩阵乘法：从 A fragment 和 B fragment 读入一小块矩阵，乘加到 accumulator fragment。它不是每个线程独立做完整矩阵乘法，而是 warp 内 32 个 lane 共同完成。

// 概念公式
D_fragment = A_fragment @ B_fragment + C_fragment

具体指令形状会随架构变化，例如 m16n8k16 这类形状表示一次 MMA 覆盖的 M/N/K 小块大小。kernel 的 warp tile 通常由很多个 MMA 指令拼起来。

7.2 ldmatrix 和 shared layout 的关系

Tensor Core 前的数据搬运通常不是普通标量 load，而是按照矩阵片段形式从 shared memory 取数。ldmatrix 一类指令要求 warp 中各 lane 以特定模式读取 shared 地址，所以第 5 章的 swizzle 和 padding 会直接影响 Tensor Core 喂数效率。

7.3 BF16 的数值路径

BF16 只有 7 个显式 fraction bit，但 exponent 范围与 FP32 相同。实践里的常见安全路径是 BF16 输入、FP32 accumulator，最后按需求 cast 到 BF16/FP16/FP32。除非目标硬件和误差预算明确允许更低精度累加，否则 K 维 reduction 通常应保留 FP32 accumulator。

bf16 A, bf16 B
fp32 acc = mma_bf16(A, B, fp32 acc)
output = cast_or_fuse(acc)

8. Pipeline 和 Async Copy：让加载与计算重叠

如果每次都同步加载 A/B tile，再计算，再加载下一块，Tensor Core 会等待内存。高性能 kernel 会做 pipeline。

实际实现里会使用 double buffering 或 multi-stage buffering：当 Tensor Core 计算当前 stage 时，异步加载下一 stage 的 A/B。

// 概念模型：允许下一 stage 保持 in flight
prefetch stage 0
commit copy group

for k_stage in range(num_k_stages):
  if has_next_stage:
    prefetch next stage
    commit copy group

  wait until current stage is ready
  CTA barrier for current buffer
  mma current stage
  CTA barrier before reusing current buffer
  swap current / next buffer

wait for remaining copies
store C tile

把 mainloop 组织成“搬运下一块、计算当前块”的稳态流水，是 CUTLASS、Triton 生成代码和 BLASLt 类库中普遍存在的高性能方向；具体使用 cp.async、TMA 还是其他 copy primitive，取决于架构与实现。

8.1 Double buffering 的直觉

只有一个 shared buffer 时，kernel 必须“加载完再算，算完再加载”。double buffering 准备两份 shared buffer：计算 stage 0 时加载 stage 1，计算 stage 1 时加载 stage 2。这样内存延迟被 mainloop 的 MMA 指令部分隐藏。

load smem[0]
for k_stage:
  load smem[next]       // async copy
  compute smem[current] // mma
  swap current/next

8.2 stage 数不是越多越好

更多 stage 可以覆盖更长内存延迟，但每多一个 stage 都会增加 shared memory 占用，也可能减少可驻留 CTA 数。生产 kernel 常在 2、3、4 stage 之间调参，取决于矩阵形状、GPU 架构和每个 CTA 的 tile 大小。

8.3 同步点要尽量少但必须正确

shared memory 被多个线程协作写入和读取，因此同步不能随意删。优化的方向通常是用异步 copy 的 wait group、barrier 或 pipeline abstraction 精确表达“哪一 stage 已经可读”，而不是在每个小步骤后粗暴同步。

// 稳态示意：此时 current group 和 next group 都在队列中
cp_async_copy(smem[next], gmem[next]);
cp_async_commit_group();

cp_async_wait_group(1);  // 最多保留 1 组未完成：current 已就绪，next 可继续在途
__syncthreads();         // CTA 内线程现在可以安全读取 current buffer
mma_on(smem[current]);
__syncthreads();         // 覆盖这个 buffer 前，确保所有消费者都已读完

// pipeline drain 时再用 wait_group(0) 等完最后一组

wait_group(0) 会等待所有已提交 group；如果在发出 next stage 后立刻这样做，就会把预取也等完，失去加载与计算的重叠。稳态常用 wait_group(1) 只保证当前最老的 group 已完成，同时允许下一组继续在途；最后 drain 才等到 0。生产代码通常用 CUTLASS pipeline、CUDA barrier 或架构专用 primitive 管理这些状态。wait_group 管 copy group 的完成边界，barrier 管 CTA 线程之间的可见性和 buffer 复用边界，两者不能互相替代。

9. Epilogue：累加之后还没结束

mainloop 负责把 A @ B 累加到 FP32 accumulator，但真实模型里的 GEMM 往往还要做 bias、activation、scale、residual、cast、量化和写回。这部分通常叫 epilogue。

// 概念 epilogue
float x = acc[m, n];
x = alpha * x + beta * old_c[m, n];
x = x + bias[n];
x = activation(x);
C[m, n] = cast_to_output_dtype(x);

9.1 为什么 epilogue 要融合

如果 mainloop 写出 FP32 C，再启动另一个 kernel 做 bias/activation/cast，就会多一次 global write 和 read。对大 GEMM 来说这可能还可以接受；对 LLM 推理里的小 M、decode、MoE grouped GEMM，epilogue 的内存流量和 kernel launch overhead 会非常明显。

9.2 写回也要 coalesced

不要只关心 A/B 的 load。C 的 store 如果是分散写、非对齐写，也会浪费带宽。很多 kernel 会让 warp 或 CTA 的输出 layout 匹配全局内存中的连续 N 维，这样写回更容易合并。

10. Profiling 工具和方法

评估 GEMM 不要只看 wall time。wall time 告诉你“慢了”，profiling 才告诉你“为什么慢”。CUDA 性能分析通常分两层：先用 Nsight Systems 看时间线和 kernel 调度，再用 Nsight Compute 深挖单个 kernel 的硬件指标。

10.1 工具怎么选

10.2 推荐工作流

1. 先用业务 benchmark 固定 shape、batch、dtype、warmup、iteration，避免 profile 的对象不稳定。
2. 用 nsys 看时间线：确认慢的是 GEMM kernel 本身，还是 kernel launch、同步、memcpy、NCCL、CPU 调度。
3. 挑出最耗时或最可疑的 GEMM kernel，用 ncu 采集详细指标。
4. 用指标判断瓶颈：Tensor Core 没打满、HBM 压力高、shared bank conflict、寄存器溢出、occupancy 太低、tail 浪费。
5. 每次只改一个变量：BM/BN/BK、stage 数、warp 数、layout、epilogue 融合、split-K 或 grouped 调度。

10.3 Nsight Systems 怎么用

nsys 适合先看全局时间线。它不会告诉你 shared memory bank conflict 细节，但能快速判断“是不是 GEMM kernel 自己慢”。

# 采集端到端 timeline
nsys profile \
  -o gemm_timeline \
  --trace=cuda,nvtx,osrt,cublas,cudnn \
  --force-overwrite=true \
  ./your_benchmark

# 只有 benchmark 调用了 cudaProfilerStart/Stop 时，才按 API 范围截取
nsys profile -o gemm_range \
  --trace=cuda,nvtx,osrt \
  --capture-range=cudaProfilerApi \
  --capture-range-end=stop \
  ./your_benchmark

# 也可以直接 profile Python
nsys profile -o torch_gemm --trace=cuda,nvtx,osrt python bench.py

10.4 Nsight Compute 怎么用

ncu 用来深挖单个 kernel。它采集 counter 会显著拖慢程序，所以一般只 profile 少量 iteration，并用 kernel 名称或 launch 序号过滤。

# 基础采集：先用常用 section
ncu --set full \
  --target-processes all \
  --kernel-name regex:gemm \
  -o gemm_kernel \
  ./your_benchmark

# 更轻量：只采集常用分析 section
ncu \
  --section SpeedOfLight \
  --section MemoryWorkloadAnalysis \
  --section Occupancy \
  --section SchedulerStats \
  --section WarpStateStats \
  --target-processes all \
  ./your_benchmark

10.5 GEMM 里最该看的指标

至少看这些指标：

10.6 指标怎么解读

10.7 一个实用排查顺序

1. 先确认结果正确，尤其是 tail block、非整除 K、转置 layout、dtype cast。
2. 看是否走到 Tensor Core 指令路径，而不是退回普通 FMA。
3. 看 HBM load/store 是否异常高，判断 tile 复用和 epilogue 融合是否有效。
4. 看 shared bank conflict，判断 shared layout 是否喂得动 Tensor Core。
5. 看寄存器数量、local memory、occupancy，判断 tile 是否过大。
6. 最后再微调 tile size、stage 数、warp 数和 split-K 等策略。

10.8 PyTorch 场景怎么定位到 GEMM

如果 GEMM 来自 PyTorch 或推理框架，先用 PyTorch Profiler 或 NVTX 标记把高层 op 和底层 kernel 对上，再用 nsys/ncu 深挖。

import torch
from torch.profiler import profile, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    with_stack=True,
) as prof:
    run_model_once()

print(prof.key_averages().table(
    sort_by="cuda_time_total",
    row_limit=20,
))

看到慢 op 后，再在关键区域加 NVTX range，nsys 里就能按阶段定位。

torch.cuda.nvtx.range_push("prefill_gemm")
run_prefill()
torch.cuda.nvtx.range_pop()

10.9 小 M 场景的特殊性

LLM decode 常见 M 很小，甚至 batch token 数只有几十或更少。这时大 CTA tile 可能浪费，输出 tile 数也可能少到无法铺满所有 SM。工程上可能需要 split-K、persistent kernel、grouped GEMM，或者把多个请求/专家合并调度。

10.9.1 Persistent kernel：CTA 算完一个 tile 后不退出

普通 GEMM 通常按输出 tile 启动大量 CTA，一个 CTA 算完自己负责的 tile 就结束。Persistent kernel 则只启动接近硬件可驻留数量的 CTA，让这些 CTA 在 kernel 生命周期内循环领取后续 tile；因此“persistent”表示 CTA 在一次 kernel launch 内持续工作，不是 kernel 永远不退出。

CTA 0: tile 0 -> exit
CTA 1: tile 1 -> exit
CTA 2: tile 2 -> exit
...
下一批工作再次 launch

while scheduler.has_work():
  tile = scheduler.next_tile()
  compute_gemm_tile(tile)

// 所有 tile 完成后才退出

当多个 tile、problem 或请求被合进同一次 launch 时，它可以摊薄 kernel launch 和 CTA 建立/回收开销；scheduler 也能把剩余 tile 交给先空闲的 CTA，改善尾部负载均衡。某些实现还可以跨多个 tile 复用调度状态或权重相关 metadata。对单次、单个 GEMM 而言，persistent 本身并不会减少 launch 次数。

10.9.2 Grouped GEMM：把多个不同 GEMM 作为一组调度

Grouped GEMM 的输入不是一对 A/B，而是一组 GEMM problem。每个 problem 可以有自己的指针、M/N/K、stride 和输出地址；kernel 把所有 problem 的输出 tile 放进同一个逻辑工作集合，再分给 CTA。

// 四个 expert 实际收到的 token 数不同，因此 M 不同
GEMM 0: [ 1, K] @ [K, N] -> [ 1, N]
GEMM 1: [ 3, K] @ [K, N] -> [ 3, N]
GEMM 2: [17, K] @ [K, N] -> [17, N]
GEMM 3: [64, K] @ [K, N] -> [64, N]

grouped_gemm([problem0, problem1, problem2, problem3])

如果分别 launch 四个 kernel，前三个 GEMM 可能都只有很少 CTA，launch 和尾部空转占比很高。Grouped GEMM 用一次 launch 暴露整组 tile，scheduler 可以让空闲 CTA 从下一个 expert/problem 继续取工作。它尤其适合 MoE，因为每个 expert 的 token 数动态变化，problem shape 天然不规则。

10.9.3 两个概念是什么关系

代价也很直接：device scheduler、problem metadata 读取和 tile 查找会产生额外开销；不同 problem 的 K 相差很大时，CTA 工作量仍可能失衡。因此 grouped/persistent 适合小而多、shape 不规则的 workload，不代表单个标准大 GEMM 也应该使用。

11. FP8 GEMM 怎么做

FP8 不是简单地把 BF16 输入换成 FP8。FP8 的动态范围和精度都更紧，需要 scale 管理。

11.1 FP8 的基本形态

11.2 scale 是 FP8 GEMM 的核心

通常会把真实值映射到 FP8 可表达范围：

x_fp8 = quantize(x / scale)
x_real ≈ dequantize(x_fp8) * scale

scale 粒度可以不同：

• per-tensor scale：实现简单，但容易被 outlier 影响。
• per-row / per-column scale：必须注明广播轴；weight 常按输出通道，activation 常按 token/row。
• per-block scale：例如每 32/64/128 个元素一个 scale，兼顾精度和吞吐。

11.3 FP8 GEMM 的计算路径

// 简化路径：scale_a / scale_b 在该输出 tile 的整个 K reduction 上恒定
A_fp8, scale_A = quantize(A_bf16_or_fp32)
B_fp8, scale_B = quantize(B_bf16_or_fp32)

acc_fp32 = fp8_tensor_core_mma(A_fp8, B_fp8)
C = acc_fp32 * scale_A * scale_B
C_out = cast_or_quantize(C)

只有当 scale_A * scale_B 对整个 K reduction 都不变时，scale 才能从求和中提出，在 mainloop 结束后统一应用。高性能实现会把 scale load、MMA、scale apply、输出 cast 或 requantize 融合起来，减少额外内存读写。

11.4 从 BF16 kernel 迁移到 FP8 时要改什么

1. global memory layout 要容纳 FP8 数据和 scale 数据。
2. load path 要加载 FP8 tile，同时加载对应 scale。
3. MMA 指令路径从 BF16 MMA 换成 FP8 MMA。
4. accumulator 通常仍用 FP32。
5. 按 scale 粒度决定缩放位置：K 维恒定的 scale 可放到 epilogue；沿 K 变化的 block scale 必须进入 mainloop / block-scaled MMA。
6. epilogue 还要处理 bias、activation、输出 dtype，以及可能的 output scale。
7. 如果输出继续 FP8，还要做 requantize，并按约定读取或写出新的 output scale / amax。

11.5 scale 粒度会改变 kernel 设计

per-tensor scale 在整个 reduction 上恒定时，可以在 epilogue 里乘一个标量，最简单但最容易受 outlier 牵制。所谓 per-channel 必须说明轴：权重按输出通道量化时，scale 常沿 N 维；activation 按 token/row 量化时，scale 常沿 M 维。这些 scale 若不随 K 变化，仍可在输出 tile 上应用。per-block scale 一旦包含 K 维分块，scale load 和缩放就进入 mainloop，不能再被当作普通 epilogue 标量。

11.6 FP8 的 accumulator 和输出

FP8 GEMM 的常见安全路径仍使用 FP32 accumulator，但具体 accumulator 模式由硬件指令和库配置决定。原因和 BF16 一样：K 维累加误差需要受控。输出可以是 BF16/FP16/FP32，也可以重新量化成 FP8；如果输出 FP8，epilogue 还要决定 output scale，并处理 clipping 和 rounding。

// 仍假设 input scale 在整个 K reduction 上恒定
acc_fp32 = fp8_mma(A_fp8, B_fp8)
y_fp32 = acc_fp32 * scale_a * scale_b + bias
out_fp8 = quantize(y_fp32 / scale_out)

即使 input scale 可以留到 K reduction 之后处理，scale apply、bias/activation、output scale 和再量化仍会增加 epilogue 的指令与访存；因此只比较 FP8/BF16 的 MMA 峰值，不能推导完整 kernel 的加速比。

11.7 per-block scale 的 layout 和索引

per-block scale 的难点不是多乘一个标量，而是 scale tensor 必须和 tile 形状、K-loop 顺序一起设计。下面用一个简化例子：A 的 scale 覆盖 64 x 64 子块，B 的 scale 也覆盖 64 x 64 子块。对某个固定的输出 scale-block (mb, nb)，每个 K-block 都有不同的 scale pair。

// mb / kb / nb 是 64x64 scale block 的索引
acc_fp32[mb, nb] = 0

for kb in range(num_k_scale_blocks):
  partial_fp32 = fp8_mma(
      A_fp8[mb, kb],
      B_fp8[kb, nb])

  scale_a = ScaleA[mb, kb]
  scale_b = ScaleB[kb, nb]
  acc_fp32[mb, nb] += partial_fp32 * scale_a * scale_b

y[mb, nb] = epilogue(acc_fp32[mb, nb])

公式的重点是 scale multiplication 位于 K-block 求和内部。实现上不一定真的生成一块独立的 partial_fp32：支持 block-scaled MMA 的架构可以把 data fragment 和 scale-factor fragment 一起交给 MMA；其他实现可能在 mainloop 中提升、缩放并累加 partial fragment。数学语义必须相同。

如果 scale block 和 CTA/MMA tile 不整齐对齐，kernel 会多出 predicate、额外 scale load，甚至跨 block 的 scale 拼接与 fragment 重排。实际工程常让 scale layout 和 mainloop 的消费粒度对齐，牺牲一点布局自由度，换取更简单的地址计算和更连续的 scale load。

11.8 requantize：clipping、rounding、amax

如果 GEMM 输出还要写成 FP8，epilogue 需要把 FP32 accumulator 重新映射回 FP8 范围。这个过程通常包含三件事：先按约定除以 output scale，再对目标 FP8 格式的可表达范围做 clipping，最后按硬件或库的策略 rounding。output scale 可以预先给定，也可以根据当前或历史 amax 动态更新；两种路径的同步和可复现性成本不同。

11.9 为什么 FP8 epilogue 可能成为瓶颈

BF16 GEMM 的 epilogue 可能只是 bias、activation 和 cast；FP8 epilogue 往往还要处理可从 K reduction 提出的 input scale、output scale、clipping/rounding，并可能写 amax 或 auxiliary tensor。沿 K 变化的 A/B block scale 属于 mainloop 成本，不应误记到 epilogue。mainloop 因为 FP8 Tensor Core 变快以后，剩余 epilogue 工作在总时间里的占比反而可能上升。

下面是 profile 的解释框架，不是某张 GPU 上的实测比例。对融合 kernel，mainloop/epilogue 占比通常需要结合 source-level counter、指令区间，或用“保留 mainloop、逐项关闭 epilogue 功能”的变体对照来估算，不能指望一个通用 NCU 指标直接给出百分比。

12. CuTe / CUTLASS DSL：NVIDIA 生产实现的一条主力路径

前面讲的是 GEMM kernel 的原理：tile、shared memory、register、Tensor Core、pipeline、epilogue。在 NVIDIA 生产实现中，常见路径不是从裸 CUDA 一行行维护所有地址计算，而是使用 CuTe / CUTLASS 3.x API 这类 DSL 和模板库。

12.1 CuTe 在解决什么问题

CuTe 的核心抽象是 Tensor = pointer + layout，以及用 Shape、Stride、Layout、TiledMMA、TiledCopy 描述数据如何被 CTA、warp 和 MMA 指令消费。它让你把“逻辑上的矩阵 tile”和“物理上的 shared/global/register 布局”分开表达。

12.1.1 CuTe 极简心智模型

读 CuTe 代码时，先不要从模板类型硬啃。把它看成三层：Tensor 保存指针和 layout，Layout 把逻辑坐标映射到物理 offset，TiledMMA 描述 warp 如何把 fragment 喂给 Tensor Core。

// CuTe 风格示意，不是完整 kernel
Tensor gA = make_tensor(make_gmem_ptr(A), make_layout(make_shape(M, K), make_stride(K, 1)));
Tensor tileA = local_tile(gA, make_shape(BM, BK), make_coord(block_m, block_k));

// Layout 决定 (m, k) 如何变成地址 offset
// TiledMMA 决定 warp/lane 如何持有 A/B/C fragment
auto tiled_mma = make_tiled_mma(mma_atom, warp_layout);
auto thr_mma = tiled_mma.get_thread_slice(threadIdx.x);

12.2 CUTLASS 3.x 的分层

所以工程上常见的做法是：先用 CUTLASS 的 CollectiveBuilder 组合出一个 kernel，再对 tile shape、cluster shape、schedule 和 epilogue fusion 做调参。只有当 builder 覆盖不了需求时，才下沉到更底层的 CuTe kernel。

12.3 Hopper FP8 GEMM：完整 CUTLASS 3.x collective 代码

下面是一份教学版 Hopper FP8 GEMM。它保留了完整 host 侧流程：定义 CUTLASS/CuTe kernel 类型、分配 A/B/C/D、初始化输入、构造 Gemm::Arguments、申请 workspace、检查支持性、初始化并运行。真实项目应以 CUTLASS 仓库中的 examples/54_hopper_fp8_warp_specialized_gemm、examples/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling 和 Python CuTe DSL 示例为准。

// hopper_fp8_gemm.cu
// Teaching version: CUTLASS 3.x + CuTe shapes for Hopper FP8 GEMM.
// It is intentionally compact, but keeps the complete host-side run path.
//
// Typical compile shape:
//   nvcc -std=c++17 -arch=sm_90a -I${CUTLASS_ROOT}/include \
//        -I${CUTLASS_ROOT}/tools/util/include hopper_fp8_gemm.cu -o hopper_fp8_gemm

#include <iostream>

#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/device_memory.h"
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/host/tensor_compare.h"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/epilogue/thread/activation.h"
#include "cutlass/epilogue/fusion/operations.hpp"
#include "cute/tensor.hpp"

using namespace cute;

using ElementA = cutlass::float_e4m3_t;
using ElementB = cutlass::float_e4m3_t;
using ElementC = cutlass::half_t;
using ElementD = cutlass::half_t;
using ElementAccumulator = float;
using ElementCompute = float;
using ElementAux = ElementD;
using ElementAmax = float;
using ElementBias = float;

using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutD = cutlass::layout::ColumnMajor;
using LayoutAux = LayoutD;

// Alignment is expressed in elements. For FP8, 128 bits means 16 elements.
constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementA>::value;
constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;
constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;

using ArchTag = cutlass::arch::Sm90;
using OperatorClass = cutlass::arch::OpClassTensorOp;

// This is the CTA-level tile discussed earlier as BM/BN/BK.
using TileShape = Shape<_128, _128, _128>;
using ClusterShape = Shape<_1, _2, _1>;

using KernelSchedule =
    cutlass::gemm::KernelTmaWarpSpecializedCooperative;
using EpilogueSchedule =
    cutlass::epilogue::TmaWarpSpecializedCooperative;

using EpilogueTile = cutlass::epilogue::collective::EpilogueTileAuto;
using FusionOperation =
    cutlass::epilogue::fusion::ScaledLinCombPerRowBiasEltActAmaxAux<
        LayoutAux,
        cutlass::epilogue::thread::ReLU,
        ElementD,
        ElementCompute,
        ElementAux,
        ElementAmax,
        ElementBias,
        ElementC>;

using CollectiveEpilogue =
    typename cutlass::epilogue::collective::CollectiveBuilder<
        ArchTag, OperatorClass,
        TileShape, ClusterShape, EpilogueTile,
        ElementAccumulator, ElementCompute,
        ElementC, LayoutC, AlignmentC,
        ElementD, LayoutD, AlignmentD,
        EpilogueSchedule,
        FusionOperation
    >::CollectiveOp;

using CollectiveMainloop =
    typename cutlass::gemm::collective::CollectiveBuilder<
        ArchTag, OperatorClass,
        ElementA, LayoutA, AlignmentA,
        ElementB, LayoutB, AlignmentB,
        ElementAccumulator,
        TileShape, ClusterShape,
        cutlass::gemm::collective::StageCountAutoCarveout<
            static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))
        >,
        KernelSchedule
    >::CollectiveOp;

using GemmKernel =
    cutlass::gemm::kernel::GemmUniversal<
        Shape<int, int, int, int>,
        CollectiveMainloop,
        CollectiveEpilogue
    >;

using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
using StrideA = typename GemmKernel::StrideA;
using StrideB = typename GemmKernel::StrideB;
using StrideC = typename GemmKernel::StrideC;
using StrideD = typename GemmKernel::StrideD;

template <class Element, class Layout>
void fill_random(cutlass::HostTensor<Element, Layout>& tensor, int seed) {
  int bits = cutlass::sizeof_bits<Element>::value;
  double scope = bits <= 8 ? 2.0 : 8.0;
  cutlass::reference::host::TensorFillRandomUniform(
      tensor.host_view(), seed, scope, -scope, 0);
  tensor.sync_device();
}

int main() {
  int M = 4096;
  int N = 4096;
  int K = 4096;
  int L = 1;

  auto stride_A = cutlass::make_cute_packed_stride(
      StrideA{}, cute::make_shape(M, K, L));
  auto stride_B = cutlass::make_cute_packed_stride(
      StrideB{}, cute::make_shape(N, K, L));
  auto stride_C = cutlass::make_cute_packed_stride(
      StrideC{}, cute::make_shape(M, N, L));
  auto stride_D = cutlass::make_cute_packed_stride(
      StrideD{}, cute::make_shape(M, N, L));

  cutlass::HostTensor<ElementA, LayoutA> tensor_A({M, K});
  cutlass::HostTensor<ElementB, LayoutB> tensor_B({K, N});
  cutlass::HostTensor<ElementC, LayoutC> tensor_C({M, N});
  cutlass::HostTensor<ElementD, LayoutD> tensor_D({M, N});

  fill_random(tensor_A, 2026);
  fill_random(tensor_B, 2027);
  fill_random(tensor_C, 2028);
  tensor_D.sync_device();

  typename Gemm::Arguments args{
    cutlass::gemm::GemmUniversalMode::kGemm,
    {M, N, K, L},
    {tensor_A.device_data(), stride_A,
     tensor_B.device_data(), stride_B},
    {
      {}, // fusion args; filled below
      tensor_C.device_data(), stride_C,
      tensor_D.device_data(), stride_D
    }
  };

  // Epilogue computes roughly:
  // D = activation(alpha * scale_a * scale_b * accumulator
  //              + beta * scale_c * C + bias)
  auto& fusion = args.epilogue.thread;
  fusion.alpha = 1.0f;
  fusion.beta = 0.0f;
  fusion.scale_a = 1.0f;
  fusion.scale_b = 1.0f;
  fusion.scale_c = 1.0f;
  fusion.scale_d = 1.0f;
  fusion.scale_aux = 1.0f;
  fusion.bias_ptr = nullptr;
  fusion.aux_ptr = nullptr;
  fusion.amax_D_ptr = nullptr;
  fusion.amax_aux_ptr = nullptr;

  Gemm gemm;
  size_t workspace_size = Gemm::get_workspace_size(args);
  cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);

  cutlass::Status status = gemm.can_implement(args);
  if (status != cutlass::Status::kSuccess) {
    std::cerr << "GEMM configuration is not supported\n";
    return -1;
  }

  status = gemm.initialize(args, workspace.get());
  if (status != cutlass::Status::kSuccess) {
    std::cerr << "GEMM initialization failed\n";
    return -1;
  }

  status = gemm.run();
  if (status != cutlass::Status::kSuccess) {
    std::cerr << "GEMM run failed\n";
    return -1;
  }

  // cudaDeviceSynchronize();
  tensor_D.sync_host();
  std::cout << "Hopper FP8 GEMM finished\n";
  return 0;
}

12.4 blockwise scaling FP8 要多什么

上面的代码是“普通 FP8 A/B 输入”的主线骨架。LLM 推理中更常见的是 blockwise scaling：A/B 除了 FP8 数据外，还要带 scale tensor。CUTLASS Hopper 示例会把 A 和 scale-A 组合成一个 tuple layout，把 B 和 scale-B 组合成一个 tuple layout，mainloop 在 MMA 前把 scale 带入。

// 概念差异：普通 FP8
ElementA, LayoutA, AlignmentA
ElementB, LayoutB, AlignmentB

// blockwise scaling FP8
ElementA, cute::tuple<LayoutA, LayoutSFA>, AlignmentA
ElementB, cute::tuple<LayoutB, LayoutSFB>, AlignmentB

在 Hopper blockwise FP8 示例里，scale layout 通常由 tile shape 推导，例如 sm90_trivial_blockwise_scale_config(TileShape{}) 再得到 LayoutSFA/LayoutSFB。这样 scale 的分块方式和 BM/BN/BK 绑定，避免 epilogue 或 mainloop 里临时做复杂索引。

12.5 Python CuTe DSL：运行入口与源码版本

Python CuTe DSL 的完整 GEMM kernel 不是几十行教学代码，而是包含 TMA descriptor、shared layout、pipeline barrier、TMEM、MMA、epilogue、reference check、benchmark 和 cold-L2 workspace 的工程实现。本节保留可运行命令，并从固定版本中摘取关键路径；完整文件仍以官方仓库为准。

普通 FP8 dense GEMM 的完整运行命令：

cd ${CUTLASS_ROOT}

python examples/python/CuTeDSL/cute/blackwell/kernel/dense_gemm/dense_gemm.py \
  --ab_dtype Float8E4M3FN \
  --c_dtype Float16 \
  --acc_dtype Float32 \
  --a_major k \
  --b_major n \
  --c_major n \
  --mma_tiler_mn 128,128 \
  --cluster_shape_mn 1,1 \
  --mnkl 4096,4096,4096,1 \
  --use_tma_store \
  --warmup_iterations 1 \
  --iterations 10

blockscaled FP8 / MXF8 GEMM 的完整运行命令：

cd ${CUTLASS_ROOT}

python examples/python/CuTeDSL/cute/blackwell/kernel/blockscaled_gemm/dense_blockscaled_gemm_persistent.py \
  --a_dtype Float8E4M3FN \
  --b_dtype Float8E4M3FN \
  --sf_dtype Float8E8M0FNU \
  --sf_vec_size 32 \
  --c_dtype Float16 \
  --a_major k \
  --b_major n \
  --c_major n \
  --mma_tiler_mn 128,128 \
  --cluster_shape_mn 1,1 \
  --mnkl 4096,4096,4096,1 \
  --warmup_iterations 1 \
  --iterations 10

这里的 --a_major k 表示 A 的 K 维连续，等价于常见 row-major A；--b_major n 表示 B 的 N 维连续，也等价于常见 row-major B。参数名写的是逻辑连续维，不是直接写 RowMajor/ColumnMajor，所以读命令时要先结合矩阵形状 A[M,K]、B[K,N] 判断。CuTe DSL 仍在快速演进；实际运行前应先对当前 checkout 执行 python .../dense_gemm.py --help，确认参数名与支持的 dtype。

12.6 普通 FP8 dense kernel：从 host 到 device

dense_gemm.py 并不是“Python 调一个现成 GEMM”。Python 代码先把 dtype、layout、tile 和 cluster 组织成 CuTe 对象，cute.compile 再把这些静态信息编译成目标 GPU kernel。下面的命令把 A/B dtype 设成 FP8；同一个 DenseGemmKernel 类也能被其他受支持 dtype 专门化。

12.6.1 run()：先实例化，再编译

a_tensor, b_tensor, c_tensor, a_torch_cpu, b_torch_cpu, c_torch_cpu, c_torch_gpu = (
    create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype)
)

# Build GEMM object
gemm = DenseGemmKernel(
    acc_dtype, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, use_tma_store
)

# Check if configuration can be implemented
can_implement = gemm.can_implement(a_tensor, b_tensor, c_tensor)
if not can_implement:
    raise ValueError(
        f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, "
        f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, "
        f"use_tma_store = {use_tma_store}"
    )
max_active_clusters = utils.HardwareInfo().get_max_active_clusters(
    cluster_shape_mn[0] * cluster_shape_mn[1]
)
compiled_gemm = cute.compile(gemm, a_tensor, b_tensor, c_tensor, current_stream)

if not skip_ref_check:
    compiled_gemm(a_tensor, b_tensor, c_tensor, current_stream)
    compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance)

查看固定版本中的原始行

create_tensors 同时准备 PyTorch tensor 和带动态 layout 的 cute.Tensor view。随后构造的 DenseGemmKernel 还不是已编译 kernel；它保存 accumulator dtype、MMA tiler、cluster shape 和 store 策略。can_implement 先验证 dtype、连续维对齐、tile/cluster 组合与输出边界，cute.compile 才依据对象配置及 A/B/C tensor 的类型和 layout 生成可调用的 compiled_gemm。

cute.compile(gemm, ...) 以带 @cute.jit 的 gemm.__call__ 为编译入口：编译器据此专门化 layout、TMA atom、grid 和 kernel launch stub，并产出 compiled_gemm。第一次 compiled_gemm(...) 用来做 reference check，benchmark 则复用这一编译结果。也就是说，Python 在这里主要负责配置、编译和调度，矩阵主循环仍在 @cute.kernel 生成的 GPU 代码里。

12.6.2 @cute.jit __call__：把 tensor 和 tile 变成 TMA 描述

# Setup TMA load for A
a_op = sm100_utils.cluster_shape_to_tma_atom_A(
    self.cluster_shape_mn, tiled_mma.thr_id
)
a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
    a_op,
    a,
    a_smem_layout,
    self.mma_tiler,
    tiled_mma,
    self.cluster_layout_vmnk.shape,
    internal_type=(
        cutlass.TFloat32 if a.element_type is cutlass.Float32 else None
    ),
)

# Setup TMA load for B
b_op = sm100_utils.cluster_shape_to_tma_atom_B(
    self.cluster_shape_mn, tiled_mma.thr_id
)
b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
    b_op,
    b,
    b_smem_layout,
    self.mma_tiler,
    tiled_mma,
    self.cluster_layout_vmnk.shape,
    internal_type=(
        cutlass.TFloat32 if b.element_type is cutlass.Float32 else None
    ),
)

a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout)
b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout)
self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size

查看固定版本中的原始行

make_tiled_tma_atom_A/B 把五类信息绑定在一起：global tensor、单 stage 的 shared layout、MMA tile、TiledMMA 的线程组织，以及 cluster layout。返回的 tma_atom_a 是 copy 操作描述，tma_tensor_a 是按照该描述重解释后的 global tensor view；二者不是已经搬入 shared memory 的数据。

self.num_tma_load_bytes 也不是“预估带宽”。后面的 TMA pipeline 会把它作为 barrier 的 transaction byte count：只有该 stage 对应的 A/B 字节都到达，消费者才能把 stage 判为 ready。A/B multicast 是否启用，则由 cluster shape 和 TMA atom 的选择共同决定。

12.6.3 Mainloop：等待 stage，发出 tcgen05 MMA，再释放 stage

    if is_leader_cta:
        # Conditionally wait for AB buffer full
        consumer_handle = ab_consumer.wait_and_advance(peek_ab_full_status)

        # tCtAcc += tCrA * tCrB
        num_kblks = cute.size(tCrA, mode=[2])
        for kblk_idx in cutlass.range(num_kblks, unroll_full=True):
            kblk_crd = (None, None, kblk_idx, consumer_handle.index)

            cute.gemm(
                tiled_mma, tCtAcc, tCrA[kblk_crd], tCrB[kblk_crd], tCtAcc
            )
            # Enable accumulate on tCtAcc after first kblock
            tiled_mma.set(tcgen05.Field.ACCUMULATE, True)

        # Async arrive AB buffer empty
        consumer_handle.release()

    if k_tile_idx + 1 < k_tile_cnt - prefetch_k_tile_cnt:
        # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1
        peek_ab_empty_status = ab_producer.try_acquire()

    if k_tile_idx + 1 < k_tile_cnt and is_leader_cta:
        # Peek (try_wait) AB buffer full for k_tile = k_tile + 1
        peek_ab_full_status = ab_consumer.try_wait()

# Async arrive accumulator buffer full
if is_leader_cta:
    acc_pipeline.producer_commit(acc_producer_state)

查看固定版本中的原始行

这段代码位于 producer 已经发出 TMA load 之后。wait_and_advance 返回当前可消费的 stage；consumer_handle.index 选择 shared-memory ring buffer 中对应的 A/B fragment。内层 kblk_idx 再把一个 CTA K tile 拆成 tcgen05 MMA 能直接处理的 K fragment。

tCtAcc 的名字里是 t，因为 Blackwell 版本把 accumulator 放在 TMEM，而不是普通线程寄存器数组。第一次 cute.gemm 后设置 ACCUMULATE=True，后续 MMA 才读旧 accumulator 并继续累加。consumer_handle.release() 把当前 A/B stage 归还 producer；K-loop 完成后，acc_pipeline.producer_commit 再通知 epilogue accumulator 已就绪。

12.6.4 Epilogue：TMEM → register → shared → global

subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
for subtile_idx in cutlass.range(subtile_cnt):
    #
    # Load accumulator from tensor memory buffer to register
    #
    tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)]
    cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)

    #
    # Perform epilogue op on accumulator and convert to C type
    #
    acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load()
    acc_vec = epilogue_op(acc_vec.to(self.c_dtype))
    tRS_rC.store(acc_vec)

    #
    # Store C to shared memory
    #
    c_buffer = subtile_idx % self.num_c_stage
    cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)])
    # Fence and barrier to make sure shared memory store is visible to TMA store
    cute.arch.fence_proxy(
        "async.shared",
        space="cta",
    )
    pipeline.sync(barrier_id=1)

    # TMA store C to global memory
    if warp_idx == 0:
        cute.copy(
            tma_atom_c, bSG_sC[(None, c_buffer)], bSG_gC[(None, subtile_idx)]
        )
        # Fence and barrier to make sure TMA store is completed to recollect C buffer
        c_pipeline.producer_commit()
        c_pipeline.producer_acquire()
    pipeline.sync(barrier_id=1)

# Wait for C store complete
c_pipeline.producer_tail()

查看固定版本中的原始行

TMA store 不能直接读取 TMEM accumulator，因此输出要经过两次显式重排：tiled_copy_t2r 先把一个 epilogue subtile 从 TMEM 搬到 register，tiled_copy_r2s 再把转换后的结果写入 shared C buffer。epilogue_op 和 output dtype 转换发生在 register vector 上。

fence_proxy("async.shared") 与 CTA barrier 保证 shared store 对 TMA 可见，warp 0 随后发出 shared-to-global TMA store。PipelineTmaStore 控制 C buffer 何时可以复用。这里也能看到：即使 mainloop 已经结束，输出重排、类型转换和异步 store 仍是一条独立流水线。

12.7 Blockscaled FP8：SFA/SFB 如何进入 mainloop

普通 dense 示例的 kernel 参数只有 A、B、C；如果模型使用沿 K 分块的 scale，SFA/SFB 就不能留到 epilogue。Blackwell blockscaled kernel 把 scale factor 作为独立 tensor，与 A/B 一起经过 TMA、shared pipeline 和 TMEM，再由 block-scaled tcgen05 MMA 在每个 K tile 内消费。

12.7.1 Warp specialization：192 个线程不是做同一件事

self.acc_dtype = cutlass.Float32
self.sf_vec_size = sf_vec_size
self.use_2cta_instrs = mma_tiler_mn[0] == 256
self.cluster_shape_mn = cluster_shape_mn
# K dimension is deferred in _setup_attributes
self.mma_tiler = (*mma_tiler_mn, 1)

self.cta_group = (
    tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
)

self.occupancy = 1
# Set specialized warp ids
self.epilog_warp_id = (
    0,
    1,
    2,
    3,
)
self.mma_warp_id = 4
self.tma_warp_id = 5
self.threads_per_warp = 32
self.threads_per_cta = self.threads_per_warp * len(
    (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
)

查看固定版本中的原始行

一个 CTA 有 6 个 warp：warp 0-3 负责 epilogue，warp 4 专门发出 MMA，warp 5 专门发出 TMA。TMA warp 可以继续为后续 K stage 或后续 output tile 搬数据，而 MMA warp 消费当前 stage；四个 epilogue warp 则并行处理 TMEM accumulator 的输出 subtile。这是 warp specialization，不是把 192 个线程平均分给一个标量循环。

kernel 还是 persistent 的：TMA、MMA 和 epilogue 角色各自维护 StaticPersistentTileScheduler 状态，在完成一个 output tile 后继续领取下一 tile。三个角色通过 A/B pipeline、accumulator pipeline、TMEM allocation barrier 和 epilogue barrier 对齐生命周期。

12.7.2 SFA/SFB 有自己的逻辑 layout 和 TMA atom

# Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout
# ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL)
sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
    a_tensor.shape, self.sf_vec_size
)
sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)

# ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
    b_tensor.shape, self.sf_vec_size
)
sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)

tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
    self.a_dtype,
    self.b_dtype,
    self.a_major_mode,
    self.b_major_mode,
    self.sf_dtype,
    self.sf_vec_size,
    self.cta_group,
    self.mma_inst_shape_mn,
)

查看固定版本中的原始行

tile_atom_to_shape_SF 根据 A/B 的逻辑 shape 和 sf_vec_size 构造 scale tensor：SFA 与 A 的 M/K 分块对应，SFB 与 B 的 N/K 分块对应。它们不是把一个一维 scale 数组随便挂到 GEMM 参数上；layout 必须表达“哪个 scale 覆盖哪段 K 数据”，并匹配 block-scaled MMA 的 scale-factor atom。

# Setup TMA load for SFA
sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(
    self.cluster_shape_mn, tiled_mma.thr_id
)
sfa_smem_layout = cute.slice_(
    self.sfa_smem_layout_staged, (None, None, None, 0)
)
tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
    sfa_op,
    sfa_tensor,
    sfa_smem_layout,
    self.mma_tiler,
    tiled_mma,
    self.cluster_layout_vmnk.shape,
    internal_type=cutlass.Int16,
)

# Setup TMA load for SFB
sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(
    self.cluster_shape_mn, tiled_mma.thr_id
)
sfb_smem_layout = cute.slice_(
    self.sfb_smem_layout_staged, (None, None, None, 0)
)
tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
    sfb_op,
    sfb_tensor,
    sfb_smem_layout,
    self.mma_tiler_sfb,
    tiled_mma_sfb,
    self.cluster_layout_sfb_vmnk.shape,
    internal_type=cutlass.Int16,
)

查看固定版本中的原始行

SFA 复用 A 方向的 cluster multicast 规则，SFB 使用面向 B/列方向的专用 cluster_shape_to_tma_atom_SFB。两者最终都生成独立 TMA atom 和 tensor view。这里的 internal_type=cutlass.Int16 描述 TMA/SMEM 搬运时的内部表示，不等于把 scale 的逻辑数值类型改成 INT16；逻辑类型仍由 self.sf_dtype 决定。

12.7.3 Producer：一个 stage 必须同时等到 A/B/SFA/SFB

for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
    # Conditionally wait for AB buffer empty
    ab_pipeline.producer_acquire(
        ab_producer_state, peek_ab_empty_status
    )

    # TMA load A/B/SFA/SFB
    cute.copy(
        tma_atom_a,
        tAgA_slice[(None, ab_producer_state.count)],
        tAsA[(None, ab_producer_state.index)],
        tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
        mcast_mask=a_full_mcast_mask,
    )
    cute.copy(
        tma_atom_b,
        tBgB_slice[(None, ab_producer_state.count)],
        tBsB[(None, ab_producer_state.index)],
        tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
        mcast_mask=b_full_mcast_mask,
    )
    cute.copy(
        tma_atom_sfa,
        tAgSFA_slice[(None, ab_producer_state.count)],
        tAsSFA[(None, ab_producer_state.index)],
        tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
        mcast_mask=sfa_full_mcast_mask,
    )
    cute.copy(
        tma_atom_sfb,
        tBgSFB_slice[(None, ab_producer_state.count)],
        tBsSFB[(None, ab_producer_state.index)],
        tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
        mcast_mask=sfb_full_mcast_mask,
    )

查看固定版本中的原始行

四次 cute.copy 使用同一个 stage barrier。该 kernel 在创建 pipeline 时把 A + B + SFA + SFB 的字节数都计入 tx_count，因此消费者不会在“数据到了、scale 还没到”的中间状态开始 MMA。ab_producer_state.index 选择 ring buffer 的物理 stage，count 则选择逻辑上的下一个 K tile。

12.7.4 Consumer：scale 从 SMEM 进入 TMEM，并成为 MMA operand

for k_tile in range(k_tile_cnt):
    if is_leader_cta:
        # Conditionally wait for AB buffer full
        ab_pipeline.consumer_wait(
            ab_consumer_state, peek_ab_full_status
        )

        #  Copy SFA/SFB from smem to tmem
        s2t_stage_coord = (
            None,
            None,
            None,
            None,
            ab_consumer_state.index,
        )
        tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
        tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
        cute.copy(
            tiled_copy_s2t_sfa,
            tCsSFA_compact_s2t_staged,
            tCtSFA_compact_s2t,
        )
        cute.copy(
            tiled_copy_s2t_sfb,
            tCsSFB_compact_s2t_staged,
            tCtSFB_compact_s2t,
        )

        # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB
        tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0)
        tile_crd = (None, None, None, ab_consumer_state.index)
        cute.gemm(
            tiled_mma,
            tCtAcc,
            [tCrA[tile_crd], tCtSFA],
            [tCrB[tile_crd], tCtSFB_mma],
            tCtAcc,
        )

查看固定版本中的原始行

MMA warp 先等待当前 A/B stage ready，再用两次 S2T copy 把这个 stage 的 SFA/SFB 从 shared memory 放进 TMEM scale layout。随后 cute.gemm 接收的不是单独 A/B fragment，而是 [A, SFA] 和 [B, SFB] 两个 paired operand。

ACCUMULATE 按 k_tile 更新，说明每一轮 K tile 都使用与该轮匹配的 scale factor，再累加到 FP32 TMEM accumulator。这正是 §11.7 公式 Σ partial[kb] × scale_A[kb] × scale_B[kb] 的硬件实现；把 SFA/SFB 留到完整 K reduction 之后再乘一次，在这里既不符合代码，也不符合数学语义。

12.8 读 Python CuTe DSL 时先分清这些对象

13. Triton GEMM：用 Python 写 tile kernel

Triton 位于“直接调用库”和“手写 CUDA/CuTe kernel”之间：它保留了 tile、program id、mask、accumulator、autotune 这些关键控制点，但把线程级地址计算、LLVM/MLIR lowering 和具体后端代码生成隐藏掉。对教学来说，Triton 很适合把前面讲的 BM/BN/BK、tail mask、tl.dot 和 accumulator 连起来。

13.1 一个 Triton matmul kernel 对应前文哪几层

13.2 教学版 Triton GEMM

下面这段代码保留了 Triton GEMM 的核心结构：一个 program 负责一个 BLOCK_M x BLOCK_N 的 C tile，沿 K 维分块循环，每轮加载 A/B 子块，用 tl.dot 累加，最后写回 C。真实项目还会加 grouped ordering、autotune、dtype 选择、epilogue fusion 和 benchmark。

import torch
import triton
import triton.language as tl


@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
    b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for k_block in range(0, tl.cdiv(K, BLOCK_K)):
        a = tl.load(
            a_ptrs,
            mask=(offs_m[:, None] < M) &
                 (offs_k[None, :] < K - k_block * BLOCK_K),
            other=0.0,
        )
        b = tl.load(
            b_ptrs,
            mask=(offs_k[:, None] < K - k_block * BLOCK_K) &
                 (offs_n[None, :] < N),
            other=0.0,
        )
        acc += tl.dot(a, b)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    tl.store(c_ptrs, acc, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))


def matmul(a, b):
    assert a.is_cuda and b.is_cuda
    assert a.ndim == 2 and b.ndim == 2
    assert a.dtype == b.dtype
    assert a.dtype in (torch.float16, torch.bfloat16)
    M, K = a.shape
    K2, N = b.shape
    assert K == K2
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)

    grid = (triton.cdiv(M, 128) * triton.cdiv(N, 128),)
    matmul_kernel[grid](
        a, b, c, M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BLOCK_M=128, BLOCK_N=128, BLOCK_K=64,
        num_warps=4, num_stages=4,
    )
    return c

M/N/K 和 stride 作为运行时参数传入，避免每个 shape 都生成一份专门化 kernel；BLOCK_M/N/K、num_warps、num_stages 才是编译期 meta-parameter。tl.store 会按 c_ptr 的元素类型把 FP32 accumulator 转成 FP16 或 BF16。

13.3 从 program_id 到 C tile

上面的 grid 是一维的：pid 先除以 N 方向 tile 数得到 pid_m，再取模得到 pid_n。这和 CUDA 里用二维 blockIdx.y/blockIdx.x 没有本质区别，只是 Triton 教程常用一维 program id 来方便做 grouped ordering。

13.4 Autotune 调什么

Triton 官方 matmul 教程会给出多组 triton.Config，CUDA 和 HIP 后端的推荐候选也不同。经验上，不要把 NVIDIA 上的 config 原样搬到 AMD；同一个 BLOCK_M/N/K，寄存器、LDS、wavefront、MFMA 路径和编译器 lowering 都可能让最优点变化。

13.5 Triton 和 CUTLASS / 手写 CUDA 的取舍

14. hipBLAS / hipBLASLt：AMD ROCm 里的 GEMM 库路线

在 ROCm 生态里，GEMM 不只有一种入口。hipBLAS 更像 BLAS API 的可移植门面：应用调用 hipBLAS，下面可以接 AMD 的 rocBLAS，也可以在 CUDA 平台接 cuBLAS。hipBLASLt 则是更灵活的 matmul API，重点放在 layout、algorithm、heuristic、workspace、tuning 和 epilogue 上。

hipBLASLt 不是 rocBLAS 调用链的下一层，Triton/CK/AITER 也不是 hipBLASLt 的后端。它们是应用可以分别选择、并在同一框架中并存的实现路线。

14.1 先把几个名字分清

14.2 hipBLAS：标准 GEMM baseline

hipBLAS 的优势是接口熟悉、迁移成本低。注意 BLAS API 的传统约定偏 column-major；如果你的张量是 row-major，常见处理方式是交换 A/B、调整转置标志，或在上层 layout 上显式处理 stride。

hipblasHandle_t handle;
hipblasCreate(&handle);

const float alpha = 1.0f;
const float beta = 0.0f;

// BLAS uses column-major convention. For row-major C[M, N] = A[M, K] * B[K, N],
// many examples call GEMM as column-major C^T[N, M] = B^T[N, K] * A^T[K, M].
hipblasSgemm(
    handle,
    HIPBLAS_OP_N, HIPBLAS_OP_N,
    N, M, K,
    &alpha,
    B, N,
    A, K,
    &beta,
    C, N);

hipblasDestroy(handle);

这段代码适合作 baseline，不适合解释高性能 GEMM 的全部细节。真实性能来自库内部选到了什么 kernel、对应什么 dtype、layout、stride、workspace 和架构，而不是来自这几行 API 本身。

14.3 hipBLASLt：更像生产 matmul 的入口

LLM 推理里的 GEMM 往往不是裸乘法：可能有 bias、activation、scale、输出 dtype 转换，也可能要为固定 shape 做离线调优。hipBLASLt 的心智模型就是把一次 matmul 拆成 descriptor、matrix layout、preference、heuristic result 和实际 launch。

create hipblasLt handle

create matmul descriptor
  - compute type
  - transpose flags
  - epilogue mode
  - bias / auxiliary pointer if needed

create matrix layouts for A, B, C, D
  - dtype
  - rows / cols
  - leading dimension
  - batch stride if batched

create preference
  - max workspace bytes
  - algorithm search constraints

query heuristic algorithms
allocate workspace
run hipblasLtMatmul(...)
destroy descriptors

14.4 和 CUDA / Triton / CUTLASS 路线怎么比较

14.5 在 LLM 推理里怎么选

14.6 ROCm 上怎么 profile

与 §10 的 CUDA 工作流类似，ROCm 也应先分清“端到端时间线问题”和“单 kernel 硬件瓶颈”。rocprofv3 适合采集 HIP API、kernel dispatch、memory copy 和指定硬件 counter；rocprof-compute 在这些数据之上提供 Speed-of-Light、memory hierarchy、roofline 和结果对比。

# 先看 HIP runtime、kernel dispatch 和 memory copy 时间线
rocprofv3 --runtime-trace -- ./your_benchmark

# 再采集单 kernel 分析；工具会多次运行 workload 收集 counter
rocprof-compute profile \
  --name gemm_profile \
  --no-roof \
  -- ./your_benchmark

# 分析采集结果；实际 SoC 子目录由工具生成
rocprof-compute analyze -p workloads/gemm_profile/<soc>/

15. 实战 Checklist

• 先确认 shape：M/N/K 是否足够大，是否存在小 M 或 tail block。
• 先用 cuBLAS / cuBLASLt / CUTLASS / Triton / hipBLAS / hipBLASLt 建 baseline。
• BF16 先检查 Tensor Core / MFMA matrix-core path 是否生效，accumulator 是否符合精度预算。
• 调 tile size 时同时看 shared memory、register pressure 和 occupancy。
• FP8 先定 scale 粒度，再定 kernel layout；不要把 scale 当成事后补丁。
• CUDA 先用 nsys 判断是不是 kernel 本身慢，再用 ncu 看单 kernel；ROCm 对应先用 rocprofv3，再用 rocprof-compute。
• profile 时把 mainloop 和 epilogue 分开看，很多 FP8 kernel 瓶颈会转移到 epilogue。

16. 常见误区

17. 参考资料

继续深入时，优先看官方文档和官方示例，因为 CUDA/CUTLASS 的接口和硬件细节会随版本更新。

• NVIDIA CUDA C++ Programming Guide：memory coalescing、shared memory、async copy、alignment。
• NVIDIA Nsight Systems User Guide：端到端 CUDA timeline、CPU/GPU overlap、kernel launch、memcpy、NCCL。
• NVIDIA Nsight Compute Profiling Guide：单 kernel 的 Tensor Core、memory workload、occupancy、warp stall、shared bank conflict。
• NVIDIA CUTLASS 文档：CUTLASS GEMM、CuTe layout、collective mainloop/epilogue 与 CuTe DSL。
• CUTLASS Hopper FP8 示例：examples/54_hopper_fp8_warp_specialized_gemm、examples/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling。
• Python CuTe DSL 示例：dense_gemm.py、dense_blockscaled_gemm_persistent.py；另见 CUTLASS block-scaled GEMM 教程。
• Inside NVIDIA GPUs: Anatomy of high performance matmul kernels：一篇很适合作进一步阅读的 matmul 长文，覆盖 GPU 架构、PTX/SASS、warp tiling、Hopper TMA 和异步 Tensor Core pipeline。
• Triton official matrix multiplication tutorial：Triton matmul、autotune、CUDA/HIP config 示例。
• AMD hipBLAS documentation：hipBLAS 作为 BLAS marshaling library 的定位和 API。
• AMD hipBLASLt documentation：GEMM、heuristic、tuning、workspace、Stream-K、datatype 和 samples。
• AMD rocBLAS documentation：ROCm 底层 BLAS 库。
• Composable Kernel user guide：ROCm 的 kernel 组件和 template 路线。
• ROCprofiler-SDK：Using rocprofv3 与 ROCm Compute Profiler：ROCm 时间线、counter、Speed-of-Light 和 roofline 分析。
• Composable Kernel：AMD GPU LDS and Bank Conflicts：LDS bank、wave 分相访问、padding 与 XOR preshuffle。
• ROCm AITER：面向 AI workload 的高性能算子与 kernel 集合。