第 6 章 · Shared Memory & Tile

学习目标

• 理解为什么 naive matmul 是 memory-bound（算术强度 0.25 FLOP/B）
• 用 shared memory tile 把全局读次数降一个量级
• 看清 bank conflict 是什么，用 padding 消除
• 知道下一步（寄存器 tile、double buffer）的方向

前置知识

已完成 Ch05，能分析合并访问、shared-memory bank 和 global-memory 流量。

核心概念

6.1 Naive Matmul 为什么慢？

源码：matmul_naive.cu。每个 output cell 由一个 thread 跑一个 K 长度的内积：

__global__ void matmul_naive(const float* A, const float* B, float* C,
                             int M, int N, int K) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    float acc = 0.f;
    for (int k = 0; k < K; ++k)
        acc += A[row * K + k] * B[k * N + col];
    C[row * N + col] = acc;
}

算术强度分析：每个 cell 做 2K FLOPs（K 次乘 + K 次加），读 2K floats = 8K bytes。 比率 0.25 FLOP/Byte。A100 的 roofline：

• HBM 带宽 1.5 TB/s → 0.25 × 1500 = 375 GFLOPS 上限（受内存限制）
• FP32 算力 19.5 TFLOPS → 想用满还要 52× 的算术强度

所以 naive kernel 往往先撞到内存 roofline，而不是峰值算力。优化的核心目标是：把数据读进 shared memory 一次、多次复用，再在目标 GPU 上测量 achieved throughput。

关键代码

6.2 Tiled MatMul — shared memory 复用

把 C 切成 32×32 的 tile，每个 block 算一个 tile。K 维度也切：

graph LR
    subgraph "K-loop: K/BK 个 tile"
      direction TB
      L1["拉 A 的 32×32 slice → As"]
      L2["拉 B 的 32×32 slice → Bs"]
      L3["__syncthreads()"]
      L4["32 次 fmadd 到 reg acc"]
      L5["__syncthreads()"]
      L1 --> L2 --> L3 --> L4 --> L5
    end
    style L4 fill:#f3f1e8,stroke:#2f5d3a

源码：matmul_tiled.cu。核心：

constexpr int BM = 32, BN = 32, BK = 32;

__global__ void matmul_tiled(const float* A, const float* B, float* C,
                             int M, int N, int K) {
    __shared__ float As[BM][BK];
    __shared__ float Bs[BK][BN];

    int row = blockIdx.y * BM + threadIdx.y;
    int col = blockIdx.x * BN + threadIdx.x;
    float acc = 0.f;

    for (int kt = 0; kt < K; kt += BK) {
        As[threadIdx.y][threadIdx.x] = A[row * K + kt + threadIdx.x];
        Bs[threadIdx.y][threadIdx.x] = B[(kt + threadIdx.y) * N + col];
        __syncthreads();
        #pragma unroll
        for (int k = 0; k < BK; ++k)
            acc += As[threadIdx.y][k] * Bs[k][threadIdx.x];
        __syncthreads();
    }
    C[row * N + col] = acc;
}

关键收益：每个 A、B 元素被 32 个线程读，但只从 global memory 读 1 次。算术强度从 0.25 提升到 8（提高 32×），roofline 上限拉到 12 TFLOPS。

运行结果

先对所有矩阵形状执行 CPU/GPU 对拍，再把计时结果写入同一硬件、同一精度的实验记录。

性能数据

6.3 继续往上爬：四个进阶方案

6.2 的 32×32 tile 通过 shared-memory 复用提高算术强度，但距离具体 GPU 上的最佳实现还有多级优化空间。 继续优化有四条路，本节做概念预览，第 9 章会给完整代码与 Nsight 分析。 理解这四条路的思路，6.3 末尾的结果对照表才看得懂。

6.3.1 寄存器 tile：每 thread 算 4×4 或 8×8 个 cell

当前 tiled 版本：每 thread 在内层循环里读 2 个 shared 值（一个 As 行、一个 Bs 列），做 1 个 fmadd。 shared 读取与算力的比例 = 1 FLOP / 2 shared-load = 0.5，shared mem 仍是瓶颈。

如果每 thread 算 4×4 = 16 个 cell：

• 读 4 个 As 列 + 4 个 Bs 行 = 8 shared values（保存在寄存器里）
• 做 16 次 fmadd
• 比例飙到 16 / 8 = 2 FLOP/shared-load，再翻 4 倍

constexpr int BM = 64, BN = 64, BK = 8;
constexpr int TM = 4,  TN = 4;       // 每 thread 算 4×4
float acc[TM][TN] = {0};             // 16 个 reg 累加器

for (int kt = 0; kt < K; kt += BK) {
    /* 协作把 A 的 64×8 + B 的 8×64 加载到 shared */
    __syncthreads();
    #pragma unroll
    for (int k = 0; k < BK; ++k) {
        float a_reg[TM], b_reg[TN];
        #pragma unroll for (int i=0;i<TM;++i) a_reg[i] = As[ty*TM + i][k];
        #pragma unroll for (int j=0;j<TN;++j) b_reg[j] = Bs[k][tx*TN + j];
        #pragma unroll for (int i=0;i<TM;++i)
        #pragma unroll for (int j=0;j<TN;++j) acc[i][j] += a_reg[i] * b_reg[j];
    }
    __syncthreads();
}

这就是本章练习题 #1（starter）要写的版本。 BM=BN=64, BK=8 时 block 大小 16×16=256 thread，每 thread 持有 16 个 reg； 扩到 BM=BN=128, TM=TN=8 还能再快一倍——代价是寄存器用得更狠（接近 256/thread 上限）。 第 9 章会推到这个版本。

6.3.2 直接调 cuBLAS（baseline）

cuBLAS 是 NVIDIA 闭源的 BLAS 实现，跟着 CUDA Toolkit 一起装。 它包含 cublasSgemm（fp32 矩阵乘）、cublasGemmEx（任意精度）等。 是衡量"我自己写得怎么样"的标尺——能跑到 cuBLAS 的 80%+ 就算优秀。

cuBLAS 之所以快，是因为底层用了几种自己手写很难做到的工程：

• autotuner：对同一个 SGEMM 准备几十个 kernel（不同 tile size、warp 切分），运行时按 M/N/K 形状自动选最优
• 软流水（software pipelining）：让 K-loop 的"下一片 load"与"当前片计算"重叠
• vectorized load（float4）：一次取 16 字节而非 4 字节，减少访存指令数
• swizzled shared layout：避免 bank conflict 同时支持 Tensor Core 对齐

#include <cublas_v2.h>
cublasHandle_t h; cublasCreate(&h);
float alpha = 1, beta = 0;
// 注意 cuBLAS 是 column-major (沿用 Fortran/LAPACK)
// row-major C(M,N) = A(M,K) @ B(K,N) 等价于 C^T = B^T @ A^T,
// 所以传参顺序"看似乱反": (N, M, K, B, A, C).
cublasSgemm(h, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K,
            &alpha, dB, N, dA, K, &beta, dC, N);

本仓库 code/ch09_gemm/gemm_cublas.cu 就是这个 baseline。

6.3.3 启用 Tensor Core（fp16 / bf16 / tf32）

到目前为止我们都在用 CUDA core——一周期吐一个 fp32 FMA。 Volta（2017）开始 NVIDIA 在 SM 里塞了Tensor Core——专门做 16×16×16 矩阵乘加的硬件， 一条指令完成 4096 个 FMA。算力跳跃式提升：

触发 Tensor Core 需要满足：

1. 数据类型为 fp16 / bf16 / tf32 / int8 / fp8（不能是纯 fp32）
2. M / N / K 对齐到 8 或 16 的倍数
3. 调 cublasGemmEx 并把 compute type 设为 Tensor Core 类型

cublasGemmEx(h, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K,
             &alpha,
             dB_half, CUDA_R_16F, N,
             dA_half, CUDA_R_16F, K,
             &beta,
             dC_half, CUDA_R_16F, N,
             CUBLAS_COMPUTE_16F,           // <-- 关键: 让 cuBLAS 走 TC 路径
             CUBLAS_GEMM_DEFAULT_TENSOR_OP);

第 9 章会手写 WMMA API 让你看到 Tensor Core 的 "魔法" 是怎么发生的——其实就是 16×16 的 fragment + mma.sync 指令。 后面 LLM 推理章节（11/12/14）所有 GEMM 都默认走 TC 路径。

6.3.4 终极：CUTLASS / 手写 WMMA（建立可测 baseline）

cuBLAS 是闭源黑盒。CUTLASS（NVIDIA 开源的 C++ 模板库）把上面这些技巧都暴露成 type traits，让你能：

• 自定义 tile 切分（mma shape、warp shape、thread shape 三级）
• 组合 cp.async（sm_80+ 异步 global→shared 拷贝）
• 插入自定义 epilogue（GEMM 后接 bias add / activation 融合）

CUTLASS 是许多 NVIDIA GPU kernels 的重要构件。模板代码量、可维护性与性能取决于具体集成；用 CUTLASS profiler 和同 shape 的 cuBLAS baseline 比较。

6.3.5 把上面五条路放一起看（1024³ MatMul · T4）

6.4 Bank Conflict — Shared Memory 的隐藏陷阱

shared memory 物理上分 32 个 bank（32-bit-word 交错）。一个 warp 同时访问 shared 时：

• 32 lane 访问 32 个不同 bank → 一周期完成 ✅
• 多个 lane 访问同一 bank 的不同 word → 串行化（n-way conflict）❌
• 所有 lane 访问同一 word → 广播，0 代价 ✅

经典案例：__shared__ float tile[32][32]，访问 tile[k][threadIdx.x] 时一切正常；但访问 tile[threadIdx.x][k] 时（按列），32 个 lane 落到同一 bank → 32-way conflict。

对照实验：transpose.cu

消除方法 1：+1 padding（教学版，够用但浪费 ~3% shared）

// 没 padding: tile[ty][k] 和 tile[ty+1][k] 落到同一 bank
__shared__ float tile[32][32];

// 有 padding: 列方向 + 1 → 每行长度变 33, 错开 bank
__shared__ float tile[32][33];

消除方法 2：XOR swizzle（工业版，CUTLASS 在用）

+1 padding 的问题：① 浪费 1/32 的 shared 容量；② 跟 Tensor Core 要求的 fragment 对齐冲突（fragment 行必须是 16 / 32 字节倍数）；③ 一旦 BK 不是固定 32，padding 量要重新算。生产代码用XOR swizzle：

// 写入时把 col 与 row 高位 XOR 一下, 让相邻 row 的同 col 落到不同 bank
// (CUTLASS 的 cutlass::layout::TensorOpMultiplicand 自动选择)
auto swizzle_idx = [](int row, int col) {
    return row * 32 + (col ^ ((row >> 0) & 7));
    //                       ^^^^^^^^^^^^^^^^
    //                       每 8 行循环错位一次
};
tile[swizzle_idx(ty, tx)] = global_A[...];

swizzle 可避免 padding 空间，同时保留 Tensor Core 对齐，并为行列访问设计无冲突映射。代价是布局更复杂，所以通常用 CUTLASS 模板生成并通过 bank-conflict metric 验证。

读写两侧都要看

Bank conflict 不只在 load 时发生，store（写 shared）也会。常见坑：转置类 kernel 写入是 coalesced，读出按列就 32-way conflict——这是 6.4 的 transpose 例子。写一次读一次的对称性是 shared 布局设计的灵魂。

6.5 工业实战：教程到生产之间的 8 个套路

6.5.1 向量化加载（float4 / half2）— 最低投入回报最高

标量 As[i][j] = A[...] 与 vectorized load 会产生不同的 load 指令和 transaction 组合。 对齐、数据类型、lane 地址与 cache line 决定实际合并结果；用 SASS 和 requested/actual bytes 验证。

// ❌ 慢：每 thread 一个 float
for (int i = tid; i < BM*BK; i += blockDim.x*blockDim.y)
    As[i / BK][i % BK] = A_global[...];

// ✅ 快：每 thread 一个 float4，指令数 1/4，HBM 事务合并
const float4* A4 = reinterpret_cast<const float4*>(A_global);
float4*       S4 = reinterpret_cast<float4*>(&As[0][0]);
for (int i = tid; i < (BM*BK)/4; i += blockDim.x*blockDim.y)
    S4[i] = A4[base + i];

要求：地址必须自然对齐 16 字节；tile 尺寸的最低维（通常是 BK）是 4 的倍数。fp16 时用 __half2 或 uint4（一次搬 8 个 fp16）。

练习：固定 GPU、CUDA、shape、dtype 与 warmup，比较标量和 vectorized load 的 指令数、global-load efficiency 与 kernel 时间；先 allclose，再判断差异是否超过计时噪声。

6.5.2 软流水 + cp.async（sm_80+）— 把 HBM 延迟藏起来

教程 K-loop 是串行的：load → sync → compute → sync → 下一轮。GPU 在 load HBM 那 ~200 个周期里计算单元闲着。Ampere 引入 cp.async PTX 指令让 global → shared 拷贝异步触发：

// 同步版 (教程):
As[i][j] = A[...];                          // LDG → STS, 阻塞
__syncthreads();
compute(As);                                 // 等 load 完才能开始

// 异步版 (生产):
__pipeline_memcpy_async(&As[i][j], &A[...], sizeof(float4));
__pipeline_commit();                         // 提交但不等
compute(As_prev);                            // 用上一轮已到位的 As, 计算与 load 并行
__pipeline_wait_prior(0);                    // 真正要用新 As 时才 wait
__syncthreads();

配合 double buffer（开两块 shared 奇偶轮交替）实现完整软流水：

串行 (教程)               时间轴 →
  load_tile_0      ━━━━
  compute_tile_0       ━━━━
  load_tile_1              ━━━━
  compute_tile_1                ━━━━
  load_tile_2                       ━━━━
  compute_tile_2                          ━━━━

软流水 (生产, 重叠)       时间轴 →
  load_tile_0      ━━━━
  load_tile_1          ━━━━                ← prefetch
  compute_tile_0       ━━━━                ← 同时 compute
  load_tile_2              ━━━━
  compute_tile_1           ━━━━
  compute_tile_2               ━━━━

总时间从 12 缩到 8 单位 (节省 33%) — 长流水越长收益越大

CUTLASS 的多阶段流水会预取多个 tile，以重叠内存访问与 MMA；stage count 受 shared memory、 register、tile shape 和 occupancy 限制。Ch9 用 profiler 比较不同配置。

6.5.3 解锁 48 KB+ shared memory

CUDA 默认每 block 最多 48 KB shared，但 sm_80+ SM 物理上有 164–228 KB。要解锁更多必须显式申请：

size_t shm_bytes = 100 * 1024;
cudaFuncSetAttribute(my_kernel,
                     cudaFuncAttributeMaxDynamicSharedMemorySize,
                     shm_bytes);
my_kernel<<<grid, block, shm_bytes, stream>>>(...);

// kernel 内用 extern shared，运行时定大小:
extern __shared__ float dyn_shm[];
float* As = dyn_shm;
float* Bs = dyn_shm + BM * BK;

典型生产配置：

• A100 fp16 GEMM tile 128×128 + double buffer：~64 KB shared
• H100 fp16 GEMM tile 256×128 + 4-stage：~196 KB shared
• FlashAttention 长 context：~96 KB shared

典型坑：忘了 cudaFuncSetAttribute 直接申请 100 KB → too many resources requested for launch。Nsight 也不会告诉你具体原因，只能凭经验。

6.5.4 Epilogue 融合 — LLM 性能的真正胜负手

LLM 推理里几乎每个 GEMM 都接着做点别的：

Y = activation(GEMM(X, W) + bias)        // 标准 FFN 第 1 层
Y = RMSNorm(GEMM(X, W) + residual)        // attention output
Y = silu(GEMM(X, W_gate)) ⊙ GEMM(X, W_up) // SwiGLU
Y = dequant(int8_GEMM(X_int8, W_int8) * scale)  // 量化推理

教程做法：3-5 个独立 kernel，中间结果每次进出 HBM。memory-bound 算子被 HBM 带宽锁死。生产做法：把 bias / activation / norm / quantize 都融合到 GEMM 的写回阶段（epilogue）：

// GEMM mainloop 结束后，acc[TM][TN] 还在寄存器里
// 直接在这里做 bias + activation 再写 HBM, 省掉一次 N×D 的 HBM 来回
#pragma unroll
for (int i = 0; i < TM; ++i)
#pragma unroll
for (int j = 0; j < TN; ++j) {
    float v = acc[i][j];
    v += bias[col0 + j];                    // bias broadcast
    v = v * (1.f / (1.f + __expf(-v)));     // SiLU
    C[(row0 + i) * N + col0 + j] = v;       // 一次性写出
}

对 memory-bound 的小算子，融合 epilogue 可能减少中间张量读写；收益取决于 shape、dtype 与后续操作。CUTLASS 把 epilogue 设计成可拼接的模板（cutlass::epilogue::collective），应与未融合路径做正确性和 profile 对照。

6.5.5 GEMV (M=1)：LLM decode 完全不能用 tile

到目前为止讨论假设 M、N、K 都比较大。但 LLM 推理 decode 阶段，每步只生成 1 个 token：

hidden state shape  : (1, D)         ← M=1 !
weight              : (D, 4D)
output              : (1, 4D)

这是 GEMV 不是 GEMM。tile 策略完全不适用：

• 开一个 32×32 tile，31 行是空的 → 利用率 1/32
• 瓶颈是权重的 HBM 读带宽，不是算力（Tensor Core 完全用不上）
• 7B 模型 decode：先按权重/KV dtype 计算单步读取字节，再除以目标 GPU 的可达带宽，最后用实测校正 roofline 估算

工业方案是专用 GEMV kernel：

• 每个 SM 负责权重的一段列，并行加载
• warp 内 32 lane 做 D 维内积，__shfl_xor_sync reduce
• 关键武器：weight-only 量化（W4A16）——权重 INT4，激活 fp16；HBM 读量降 4×，单步 token 时间从 10 ms 降到 3 ms

vLLM、TensorRT-LLM、llama.cpp 都有专门的 GEMV 实现路径。本仓库 Ch13 会让你自己写一个 fp16 GEMV，Ch14 capstone 默认走这条路。

6.5.6 Split-K / Stream-K — 处理"细长"形状

常规 tile 策略假设 M、N、K 大致相当。但下面这种 GEMM 让 SM 大半闲着：

M = 128, N = 128, K = 16384
→ tile 128×128 → 只 1 个 output tile → 1 个 block → A100 上 108 个 SM 用 1 个，其他 107 个浪费

Split-K：把 K 维拆 8 份，每份独立算 partial sum 写不同 buffer，最后再 reduce。block 数 ×8，SM 跑满。 Stream-K（NVIDIA 2023）：进一步把 K 切到比 block 还细的单位，按 SM 调度均衡负载——是当前 CUTLASS 默认策略，比 Split-K 更适应不规则 size。

什么时候用：min(M, N) < 256 且 K > 2048 时强烈推荐。生产代码会有形状感知逻辑：

if (M * N < 16 * 16 * num_sms && K > 2048) {
    launch_stream_k_kernel(...);   // 长 K
} else if (M == 1) {
    launch_gemv_kernel(...);       // decode
} else {
    launch_gemm_kernel(...);       // 普通
}

6.5.7 持久化 kernel — 杀掉 launch overhead

kernel 启动包含 CPU→GPU 命令传递与 scheduler 分发；当单步 decode 触发许多短 kernel 时，launch gap 可能成为端到端瓶颈，应由 Nsight Systems 实测。

持久化 kernel：让 block 长期驻留 SM，从 host 准备的 work queue 拉任务循环处理，整个 decoding loop 只发一次 launch：

__global__ void persistent_engine(WorkItem* queue, int* head, int n_items) {
    while (true) {
        int my_id = atomicAdd(head, 1);
        if (my_id >= n_items) break;
        process(queue[my_id]);
    }
}

推理引擎会组合 persistent-style kernels 与 CUDA Graph（见 Ch8）来减少 host launch gaps。端到端收益取决于 batch、graph bucket、CPU scheduler 和 kernel 时长，应从 timeline 测量。

6.5.8 Hopper TMA — 硬件级 tile 加载

H100 引入 TMA (Tensor Memory Accelerator)：一条指令告诉硬件 "把 global memory 的 64×64 tile 拷到 shared mem 的某处"，硬件自动处理跨页、跨 bank 的 swizzle、边界，CPU 完全不写 load 循环。

等价代码（PTX 简化）：

// pre-Hopper: 几十行手写 load 循环
for (int i = tid; i < BM*BK/4; i += 256)
    cp_async(&As[...], &A[...]);

// Hopper TMA: 1 条指令
cuda::ptx::cp_async_bulk_tensor_2d(&As[0][0], tma_descriptor, {row0, col0});

FlashAttention v3 几乎完全用 TMA 替代手写 load。这是 Hopper 上 fp16 GEMM 能跑到 700+ TFLOPS（接近 989 TF peak）的核心机制。

6.6 工业 tile size 选择速查表

记忆比推导实用。下面是 production GEMM kernel 在不同场景下的常用 tile 配置（基于 CUTLASS profiler + 真实部署）：

tile 选择的 4 条经验法则

1. BM × BN ≤ 16K（fp32）或 ≤ 32K（fp16）：保证两块 shared + accumulator 装得下 SM 的 100-200 KB
2. BK ≥ 8：BK 太小 → 软流水间隙小，HBM 延迟藏不住；BK 太大 → shared 占用多挤压 occupancy
3. grid 数 ≥ SM 数 × 2：让 GPU 有多波 (wave) work 可调度，避免尾部 SM 闲置（H100 有 132 SM → 至少 264 个 block）
4. per-thread tile TM × TN ≤ 64：超过会触发寄存器 spill 到 local memory（实际是 global 私有分区，奇慢）。Nsight 的 "Register Spills" 行非零就是这个

6.7 实战案例：一次 fused QKV GEMM 的调优历程

下面是待采集的 GEMM 优化记录模板，目标 workload 为 fused QKV projection： M=2048、N=2304、K=768。仓库没有对应 GPU artifact，因此不填入数字；在目标 GPU 上用相同输入和 correctness gate 逐步采集到 benchmarks/reports/ch06/<gpu>-<date>.md。

分析：

• 每一步先 allclose，再单独记录变化，避免把多个改动合成一个无法归因的数字。
• 同时记录编译参数、GPU、CUDA、clock/power 状态、warmup/repeat 与 shape。
• Tensor Core、vectorized load、async copy 与 epilogue fusion 的收益会随 workload 改变。
• 最终同时报告与 cuBLAS baseline 的正确性、时间和差距，不预设可达到的比例。

自检清单

练习题

1. 01_2d_block_tile_starter.cu：实现"每 thread 算 4×4 个 cell"的版本，目标 GFLOPS 提升到 tiled 版的 2×。
2. 给 matmul_tiled 加 __shared__ float As[BM][BK+1] 看 bank conflict 减少情况（用 Nsight Compute 看 l1tex__data_bank_conflicts_pipe_lsu_mem_shared）。
3. 调 BK 从 8 到 64，记录 GFLOPS 曲线——为啥 BK=32 是甜点？（提示：BK 过小流水间隙小，BK 过大 shared 挤压 occupancy）。
4. 把 transpose 的 TILE 从 32 改 16/64，看带宽变化。
5. 工业级：把练习 #1 的 kernel 改成 vec-load 版（float4）。预期再 +30%。
6. 工业级：把 epilogue 融合到练习 #1 中——在写回前加一个 ReLU 或 +bias。对比 "GEMM+独立 bias_add kernel" 的总耗时。
7. 挑战：让你的 kernel 支持 M、N、K 任意值（不是 32 的倍数）。这是 CUTLASS 几千行模板代码要解决的核心问题。

6.10 研究前沿（2025-2026）：Warp Specialization、WGMMA、TMEM GEMM

6.10.1 Warp Specialization — FA v3 + DeepSeek FlashMLA 的核心模式

传统 GEMM/Attention kernel：block 内所有 warp 做同样的事（load → compute → load → compute），通过 __syncthreads 协调。瓶颈：load 时 compute 单元闲、compute 时 load 单元闲。

warp specialization 把 block 内 warp 分两组：

• Producer warps（通常 1-2 个）：专门 cp.async / TMA 加载下一阶段 K/V tile
• Consumer warps（其余 4-8 个）：专门做 wgmma 计算
• 两组通过 mbarrier（硬件 barrier）单向同步：producer signal "tile ready"，consumer signal "tile consumed"

__global__ void warp_spec_gemm(...) {
    extern __shared__ __align__(128) uint8_t smem[];
    __shared__ uint64_t mbar_load[N_STAGES], mbar_compute[N_STAGES];
    if (threadIdx.x == 0)
        for (int i = 0; i < N_STAGES; ++i) {
            mbarrier_init(&mbar_load[i],    1);
            mbarrier_init(&mbar_compute[i], 1);
        }
    __syncthreads();

    int warp_id = threadIdx.x / 32;
    if (warp_id == 0) {                          // PRODUCER
        for (int k = 0; k < K_iters; ++k) {
            int slot = k % N_STAGES;
            mbarrier_wait(mbar_compute[slot]);   // 等 consumer 用完上一轮
            tma_load_async(smem + slot, K_global + k);
            mbarrier_arrive(mbar_load[slot]);
        }
    } else {                                     // CONSUMERS (warps 1..N-1)
        for (int k = 0; k < K_iters; ++k) {
            int slot = k % N_STAGES;
            mbarrier_wait(mbar_load[slot]);      // 等 producer load 完
            wgmma_async(acc, smem + slot, ...);
            wgmma_commit();
            wgmma_wait<0>();
            mbarrier_arrive(mbar_compute[slot]);
        }
    }
}

FA v3 的论文结果属于论文指定的 H100、dtype、shape 与软件栈；本章只借其 producer/consumer 结构解释 WGMMA/TMA pipeline，不把论文数字当作本仓库测量。

6.10.2 WGMMA — Hopper 的 warp-group MMA

Hopper 引入 wgmma（warp-group MMA）：4 个 warp（128 lane）协作执行一个大 MMA，shape 可达 m64n256k16（fp16）。

对比：

关键变化：wgmma 是异步的——发射后立即返回，warp 可以继续干别的，accumulator 在后台写完。这就是软流水的硬件版。

6.10.3 ThunderKittens — 100 行写 attention 的秘诀

Stanford Hazy Research 2024 推出的 C++ 嵌入式 DSL，把上面这堆 wgmma + mbarrier + TMA 抽象成"tile 操作"：

// ThunderKittens 完整 attention forward (单 head, 极简, ~80 行)
#include "kittens.cuh"
using namespace kittens;

template <int D>
__global__ void tk_attn(...) {
    auto q = make_rt<bf16, 64, D>();         // register tile (block 内)
    auto k = make_st<bf16, 64, D>();         // shared tile (block 内)
    auto v = make_st<bf16, 64, D>();
    auto o = make_rt<float, 64, D>();
    auto m = make_col_vec<float, 64>();      // running max
    auto l = make_col_vec<float, 64>();      // running sum

    load_async(q, Q_global);
    zero(o); neg_inf(m); zero(l);

    for (int kt = 0; kt < n_kv_tiles; ++kt) {
        load_async(k, K_global + kt);
        load_async(v, V_global + kt);
        kittens::wait();                       // 等 cp.async / TMA

        auto s = mm<TransB>(q, k);             // 自动 wgmma
        auto m_new = max(m, row_max(s));
        sub_row(s, m_new);
        exp(s);
        l = exp(m - m_new) * l + row_sum(s);
        o = exp(m - m_new) * o + mm(s, v);     // 自动 wgmma
        m = m_new;
    }
    div_row(o, l);
    store(O_global, o);
}

ThunderKittens 可用较紧凑的抽象表达 attention 原型；与 CUTLASS 的性能和代码量比较必须固定 feature set、shape、dtype、编译参数与 correctness gate。

6.10.4 Blackwell TMEM GEMM — TCGEN05

Blackwell（sm_100）的 GEMM 范式再次变化：

// Hopper:
wgmma_async(acc_reg, A_smem, B_smem);    // acc 在 register

// Blackwell:
tcgen05_mma_async(acc_tmem, A_smem, B_smem);   // acc 在 TMEM (新)
// ... 其他事 ...
tcgen05_ld(acc_reg, acc_tmem);           // 用之前才读到 register

TMEM 可让部分 accumulator 不长期占用 register file，但 register、occupancy 与最终吞吐的变化依赖 tile 和 compiler allocation。配合 microscaling 与 2-CTA MMA 后仍需在 Blackwell 上 profile： TODO(on GPU)。

6.10.5 实战推荐路径（2026）

关键洞察：从 Hopper 开始，"手写 kernel 接近 cuBLAS"已经几乎不可能（cuBLAS 也是 CUTLASS 编译出来的）。生产决策：

• 标准 GEMM → 直接调 cuBLAS / cuBLASLt
• 带 fused epilogue → CUTLASS 模板拼装
• 新算子原型 → Triton 或 ThunderKittens
• 极致 attention → 看 FA v3+ 源码学
• 本章 + Ch9 的手写 → 是理解原理的途径，不再是部署路线

常见坑

• tile 边界处理漏掉 → 在非整数倍 size 上 GG，用 if 兜底或者把 As/Bs 初始化为 0
• 少 __syncthreads → 偶发错误，不可复现
• shared array 过大没用 cudaFuncSetAttribute → too many resources requested 报错且不告诉你原因
• 转置时块坐标也要换：x2 = blockIdx.y * TILE + threadIdx.x
• vec load 时地址未对齐 → nvcc 静默回退到 scalar load，性能没涨
• per-thread tile 过大 (TM×TN > 64) → reg spill 到 local mem，性能暴跌，Nsight "Stack Frame Spill" 行非 0
• epilogue 里做 fp16 *= scale，scale 是 fp32 → 隐式转换，整 epilogue 跑 fp32 单元，性能减半
• “我跑得比 cuBLAS 快！” → 先排查 warmup、同步、正确性与 shape cherry-pick，并扩展多个代表性 size
• 给 grid 开 num_blocks = num_sms 但 kernel 还要 sync 整个 grid → 死锁。grid sync 要用 cooperative launch + cooperative_groups

6.12 CUDA 官方手册精讲（CUDA Programming Guide 13.2（核验：2026-07-20））

Bank 冲突诊断、memcpy_async、TMA、DSMEM、Pipeline

用 Nsight 量化 bank conflict：硬件 32-bank 模型与对照实验

讲完 padding 与 swizzle，怎么知道自己写的版本到底有没有冲突？必须用 Nsight Compute 的硬件计数器。先复习硬件模型：

关键计数器（Nsight Compute UI 里搜索 bank_conflicts 即可）：

ncu --metrics \
      l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld.sum,\
      l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_st.sum,\
      l1tex__data_pipe_lsu_wavefronts_mem_shared_op_ld.sum,\
      l1tex__data_pipe_lsu_wavefronts_mem_shared_op_st.sum \
      ./my_app

每条 shared load/store 指令理想是 1 wavefront（一周期完成）。如果 conflict 数非零，wavefront 数会等比例放大。判定公式：

有效冲突 = bank_conflicts.sum / wavefronts.sum
               ↑ 0  → 完美
               ↑ 31 → 32-way 冲突（最坏情况）

在 fp64 / int2 这种 8 字节访问场景，Volta+ 还允许显式切到 8-byte bank 模式，把"两个 32-bit lane 共享同一 bank"的情况合并：

// 默认 4-byte bank：fp64 数组按列访问 → 16-way conflict
    __shared__ double tile_d[16][16];

    // 切到 8-byte bank：相邻 lane 用相邻 8-byte bank，冲突减半
    cudaDeviceSetSharedMemConfig(cudaSharedMemBankSizeEightByte);

用 cuda::pipeline 写多阶段软流水（标准 C++ API）

6.5.2 给的 __pipeline_memcpy_async 是 C 风格底层原语，CUDA 12+ 推荐改用 <cuda/pipeline> 头里的 cuda::pipeline —— 它是 libcu++ 提供的 C++ 抽象， 可以在 block / warp / thread 任意 scope 上工作，配合 cuda::memcpy_async 自动管理"几阶段流水"。

核心数据结构：

• cuda::pipeline_shared_state<scope, N>：一个 N 阶段流水的共享状态，要么 thread-local 要么放 shared memory
• cuda::pipeline：每个线程持有的 pipeline handle
• 用 producer_acquire() / producer_commit() 推送一片；用 consumer_wait() / consumer_release() 取走一片

#include <cuda/pipeline>
    #include <cooperative_groups.h>
    namespace cg = cooperative_groups;

    constexpr int N_STAGES = 4;            // 4 阶段：当前算第 0 片时，1/2/3 在路上
    constexpr int BM = 128, BN = 128, BK = 32;

    __global__ void matmul_pipeline(const float* A, const float* B, float* C,
                                    int M, int N, int K) {
        auto block = cg::this_thread_block();

        __shared__ float As[N_STAGES][BM][BK];
        __shared__ float Bs[N_STAGES][BK][BN];

        // 多阶段共享状态：每阶段占一组 barrier
        __shared__ cuda::pipeline_shared_state<cuda::thread_scope_block, N_STAGES> pstate;
        auto pipe = cuda::make_pipeline(block, &pstate);

        int n_iters = K / BK;

        // 1) 预热：先 prefetch 前 N_STAGES-1 片
        for (int s = 0; s < N_STAGES - 1 && s < n_iters; ++s) {
            pipe.producer_acquire();
            cuda::memcpy_async(block, &As[s][0][0], &A[/* tile (blockIdx.y, s) */],
                               cuda::aligned_size_t<16>(BM * BK * 4), pipe);
            cuda::memcpy_async(block, &Bs[s][0][0], &B[/* tile (s, blockIdx.x) */],
                               cuda::aligned_size_t<16>(BK * BN * 4), pipe);
            pipe.producer_commit();
        }

        float acc[8][8] = {0};

        // 2) 稳态：每轮"提交下一片 + 算当前片"完全重叠
        for (int it = 0; it < n_iters; ++it) {
            int slot_cur  = it % N_STAGES;
            int slot_next = (it + N_STAGES - 1) % N_STAGES;

            // 提交未来的一片（it + N-1）
            int future = it + N_STAGES - 1;
            if (future < n_iters) {
                pipe.producer_acquire();
                cuda::memcpy_async(block, &As[slot_next][0][0], &A[/* (by, future) */],
                                   cuda::aligned_size_t<16>(BM * BK * 4), pipe);
                cuda::memcpy_async(block, &Bs[slot_next][0][0], &B[/* (future, bx) */],
                                   cuda::aligned_size_t<16>(BK * BN * 4), pipe);
                pipe.producer_commit();
            }

            // 等当前片到位
            pipe.consumer_wait();
            block.sync();
            // ... 用 As[slot_cur] / Bs[slot_cur] 做内层 mma ...
            pipe.consumer_release();
        }
    }

对比 __pipeline_memcpy_async：

动态 shared 与 cudaFuncSetAttribute 进阶：占用率 vs 容量的权衡

6.5.3 教你怎么把 shared memory 申请超 48KB；进一步要知道SM 内 L1/Shared 的统一物理 SRAM是分配出来的，多吃 shared 等于少吃 L1。cudaFuncSetAttribute 提供两个相关属性：

// 1) 解锁 100 KB 动态 shared（必须在第一次 launch 之前调）
    size_t shm_bytes = 100 * 1024;
    CUDA_CHECK(cudaFuncSetAttribute(
        matmul_pipeline,
        cudaFuncAttributeMaxDynamicSharedMemorySize,
        shm_bytes));

    // 2) 告诉 driver: 我希望尽量多 shared 而不是 L1
    CUDA_CHECK(cudaFuncSetAttribute(
        matmul_pipeline,
        cudaFuncAttributePreferredSharedMemoryCarveout,
        cudaSharedmemCarveoutMaxShared));

    // 3) launch 时第三个参数 = shm_bytes
    matmul_pipeline<<<grid, block, shm_bytes, stream>>>(...);

kernel 内部用 extern __shared__ 拿一段，自己切：

extern __shared__ __align__(16) unsigned char dyn_shm[];

    float* As = reinterpret_cast<float*>(dyn_shm);
    float* Bs = reinterpret_cast<float*>(dyn_shm + sizeof(float) * BM * BK * N_STAGES);
    // __align__(16) 保证 float4 vec-load 可用

真实数字（来自 CUTLASS profiler）：

Hopper TMA 完整用法：从 Host 端 cuTensorMapEncodeTiled 到 Kernel 内 cp.async.bulk.tensor

6.5.8 把 TMA 一笔带过。这里给出"从 host 到 device 的完整最小可运行例子"，让你明白这条指令背后有多少 boilerplate——这也是为什么 CUTLASS / ThunderKittens 要做封装。

第 1 步：在 Host 端构造一个 CUtensorMap（描述符）

#include <cuda.h>        // 必须用 driver API
    #include <cudaTypedefs.h>

    // global 矩阵：M × K, row-major, fp16
    constexpr int M = 4096, K = 4096;
    constexpr int BM = 128, BK = 32;          // tile / "box"

    CUtensorMap tmap_A{};
    uint64_t globalDim[2]    = { (uint64_t)K, (uint64_t)M };       // 列在前
    uint64_t globalStride[1] = { K * sizeof(__half) };             // 行步
    uint32_t boxDim[2]       = { (uint32_t)BK, (uint32_t)BM };
    uint32_t elemStride[2]   = { 1, 1 };

    CUresult res = cuTensorMapEncodeTiled(
        &tmap_A,                                     // out
        CU_TENSOR_MAP_DATA_TYPE_FLOAT16,
        /*tensorRank=*/2,
        A_global_ptr,                                // 实际 global 指针
        globalDim, globalStride, boxDim, elemStride,
        CU_TENSOR_MAP_INTERLEAVE_NONE,
        CU_TENSOR_MAP_SWIZZLE_128B,                  // ← 128B swizzle (默认 fp16 GEMM 用)
        CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
        CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
    assert(res == CUDA_SUCCESS);

第 2 步：把 CUtensorMap 作为 __grid_constant__ 传给 kernel

__global__ void gemm_tma(const __grid_constant__ CUtensorMap tmap_A,
                             const __grid_constant__ CUtensorMap tmap_B,
                             __half* C, int M, int N, int K) {
        // 1) 准备 mbarrier
        __shared__ alignas(8) uint64_t bar;
        __shared__ alignas(128) __half As[BM][BK];   // 128B 对齐, 给 swizzle 用
        __shared__ alignas(128) __half Bs[BK][BN];

        if (threadIdx.x == 0) {
            cuda::ptx::mbarrier_init(&bar, 1);   // 期待 1 次 arrive (TMA 完成)
        }
        __syncthreads();

        // 2) 一条指令发起 BM×BK 的 fp16 tile 拷贝
        int32_t coords_A[2] = { 0 /*k0*/, blockIdx.y * BM /*m0*/ };
        if (threadIdx.x == 0) {
            cuda::ptx::cp_async_bulk_tensor(
                cuda::ptx::space_shared,
                cuda::ptx::space_global,
                &As[0][0],
                &tmap_A,
                coords_A,
                cuda::device::barrier_native_handle(bar));
            // 告诉 barrier "等 BM*BK*sizeof(half) 字节"
            cuda::ptx::mbarrier_arrive_expect_tx(
                cuda::ptx::sem_release,
                cuda::ptx::scope_cta,
                cuda::ptx::space_shared,
                &bar, sizeof(As));
        }

        // 3) 等 TMA 完成 (try_wait_parity 是非阻塞 spin)
        int parity = 0;
        while (!cuda::ptx::mbarrier_try_wait_parity(&bar, parity)) {}

        // 4) 此时 As 里已经有数据 + 自带 128B swizzle 排布, 可以直接 wgmma
        // wgmma_async(acc, As, Bs);
    }

关键收益（vs cp.async）

Cluster + DSMEM：跨 block 共享 shared memory（Hopper）

Hopper 引入 Thread Block Cluster（CGA）：最多 8 个 block 组成一个簇，簇内 block 之间的 shared memory 可以互相读写，称为 Distributed Shared Memory（DSMEM）。可用容量受 cluster size、每 block 动态 shared 与 occupancy 限制，应查询属性并验证 launch 配置。

用法核心三步：

1. 在 kernel 上加 __cluster_dims__(N, 1, 1) 声明簇形状，或 launch 时通过 cudaLaunchAttributeClusterDimension 设
2. 用 cg::this_cluster() 拿到 cluster handle
3. 用 cluster.map_shared_rank(local_ptr, peer_rank) 把本 block 的 shared 指针映射成"peer block 上同样位置的 shared"

#include <cuda/barrier>
    #include <cooperative_groups.h>
    namespace cg = cooperative_groups;

    // 簇 = 4 个 block。每个 block 算 C 的 BM × BN，但 K 维 load 摊到 4 个 block
    __global__ __cluster_dims__(4, 1, 1)
    void cluster_gemm(const __half* A, const __half* B, __half* C,
                      int M, int N, int K) {
        using barrier_t = cuda::barrier<cuda::thread_scope_block>;

        auto cluster = cg::this_cluster();
        auto block   = cg::this_thread_block();
        int  rank    = cluster.block_rank();          // 0..3
        int  csize   = cluster.num_blocks();          // 4

        __shared__ alignas(128) __half As[BM][BK];
        __shared__ barrier_t bar;
        if (block.thread_rank() == 0) init(&bar, csize);  // 期 csize 次 arrive
        cluster.sync();   // 等所有 block 的 barrier init 完

        // 每 block 各自负责 K 维的 1/4
        int k0 = rank * (K / csize);
        int k1 = (rank + 1) * (K / csize);

        for (int kt = k0; kt < k1; kt += BK) {
            // 1) 本地 load 一片
            // (vector load 略)

            // 2) 把自己 load 完的 As "发给"右邻居 block
            int peer = (rank + 1) % csize;
            __half* As_peer = cluster.map_shared_rank(&As[0][0], peer);
            // 这里是真·跨 block 的 shared 写！硬件走 DSMEM 链路（NVLink 内部）
            for (int i = block.thread_rank(); i < BM*BK; i += block.size()) {
                As_peer[i] = As[i / BK][i % BK];
            }

            // 3) cluster-wide barrier 等大家都写完
            cluster.sync();

            // 现在 As 里有当前 block 的本地一片 + 来自左邻居的一片 → 可以做更大 K 的 mma
        }
    }

    graph TB
        subgraph "Hopper SM Cluster (4 blocks)"
          direction LR
          B0["Block 0
shared: As[0..]"] -.->|map_shared_rank| B1["Block 1
shared: As[0..]"]
          B1 -.->|map_shared_rank| B2["Block 2
shared: As[0..]"]
          B2 -.->|map_shared_rank| B3["Block 3
shared: As[0..]"]
          B3 -.->|map_shared_rank| B0
        end
        note["cluster.sync() 同步整簇"]
        style B0 fill:#f3f1e8,stroke:#2f5d3a
        style B1 fill:#f3f1e8,stroke:#2f5d3a
        style B2 fill:#f3f1e8,stroke:#2f5d3a
        style B3 fill:#f3f1e8,stroke:#2f5d3a
    

下一章导览

第 7 章把相同的 tile、warp 通信和分层合并思路用于 reduction、scan 与 histogram。