第 9 章 · GEMM 深入

学习目标

• 把 Ch6 tiled matmul 推进到 2D 寄存器 tile，并用计时与 profiler 判断数据复用是否带来收益
• 第一次摸 Tensor Core（WMMA API），看到 fp16 算力如何碾压 fp32
• 对比 cuBLAS baseline，知道还差多少、差在哪
• 了解 CUTLASS / CuTe 与"手写到极致"的关系

前置知识

已完成 Ch06 与 Ch08，能解释 tiled GEMM、arithmetic intensity、warmup 和可重复计时。

核心概念

9.1 GEMM 在 LLM 里的地位

在 Transformer 中，GEMM 通常占据大部分 FLOPs。Transformer 每层主要算子：

• QKV projection：X (T, D) @ W_qkv (D, 3D) → (T, 3D)
• Attention output：(T, D) @ W_o (D, D)
• MLP up/down：(T, D) @ W_up (D, 4D) → (T, 4D)，再 (T, 4D) @ W_down (4D, D)

所以 prefill、大 batch decode 和 verifier 阶段通常优先优化 GEMM；但 batch=1 decode 还会被权重读取、KV cache 和 launch gap 限制，必须分阶段 profile。

关键代码

9.2 寄存器 Tile：从 Ch6 再进一步

Ch6 里每 thread 算 1 个 cell。瓶颈是每读 As/Bs 一次只用 2 个 FLOP。 如果每 thread 算 8×8 = 64 个 cell：每读 8+8 个 shared 值就做 64 个 FLOP，shared 访问减少 8×。

graph TD
    A["block 处理 C 的 128×128 tile
thread block = 16×16 = 256 threads"]
    A --> B["每 thread 持有 8×8 reg accumulator"]
    B --> C["K 维分 BK=16 切片
每片协作加载 128×16 + 16×128 到 shared"]
    C --> D["内层 K 循环 (BK=16):
读 8 个 As 行 + 8 个 Bs 列 → 64 个 fmadd"]
    D --> E["写回 C"]
    style B fill:#f3f1e8,stroke:#2f5d3a
    style D fill:#f3f1e8,stroke:#a86420

constexpr int BM = 128, BN = 128, BK = 16;
constexpr int TM = 8,   TN = 8;

__global__ void gemm_reg_tile(const float* A, const float* B, float* C, int M, int N, int K) {
    __shared__ float As[BM][BK], Bs[BK][BN];
    int ty = threadIdx.y, tx = threadIdx.x;
    int row0 = blockIdx.y * BM + ty * TM;
    int col0 = blockIdx.x * BN + tx * TN;
    float acc[TM][TN] = {0};

    for (int kt = 0; kt < K; kt += BK) {
        /* 协作 load A 的 128x16 + B 的 16x128 (每 thread load 8+8 个数) */
        __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();
    }
    /* write back acc → C */
}

关键洞察：内层 8×8 = 64 个 fmadd 全部在寄存器之间发生。BK=16 步内只触发 16×(8+8) = 256 次 shared load，每 thread 做 16×64 = 1024 个 FLOP。算术强度 4 FLOP/shared-byte。

9.3 Tensor Core — WMMA API

fp32 CUDA core 路线很难接近低精度 Tensor Core 的吞吐。要让 GEMM 跑到现代 GPU 的高吞吐区间，必须使用 Tensor Core，并确认 Nsight Compute 中 tensor pipe 真正活跃。

WMMA fragment 三件套

#include <mma.h>
using namespace nvcuda;

wmma::fragment<wmma::matrix_a, 16, 16, 16, __half, wmma::row_major> a_frag;
wmma::fragment<wmma::matrix_b, 16, 16, 16, __half, wmma::row_major> b_frag;
wmma::fragment<wmma::accumulator, 16, 16, 16, float>                c_frag;

wmma::fill_fragment(c_frag, 0.f);
for (int kt = 0; kt < K; kt += 16) {
    wmma::load_matrix_sync(a_frag, A + row*K + kt, K);
    wmma::load_matrix_sync(b_frag, B + kt*N + col, N);
    wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);   // <-- TC magic
}
wmma::store_matrix_sync(C + row*N + col, c_frag, N, wmma::mem_row_major);

每个 warp 共享一个 fragment，硬件保证整个 warp 32 个 lane 协作完成 16×16×16 = 4096 FMA，一次 mma_sync 指令吞 8 个时钟。

运行结果

9.4 cuBLAS baseline

调 cuBLAS 是判断"自己写得好不好"的标尺：

#include <cublas_v2.h>
cublasHandle_t h; cublasCreate(&h);
float alpha = 1.f, beta = 0.f;

// 注意：cuBLAS 是 column-major！
// row-major C(M,N) = A(M,K) @ B(K,N)
//   ≡  C^T(N,M) = B^T(N,K) @ A^T(K,M)
//   ≡  cublasSgemm(N, N, N, M, K, B, N, A, K, C, N)
cublasSgemm(h, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K,
            &alpha, dB, N, dA, K, &beta, dC, N);

性能数据

9.5 性能采集表（1024³ MatMul）

下面这张表留给目标 GPU 实测。记录 GPU、driver、CUDA、ARCH、矩阵维度、dtype、是否启用 Tensor Core，并用 Nsight Compute 对照 Tensor pipe 与 DRAM pipe。

本章手写 WMMA 还有大量优化空间（async copy、swizzle、warpgroup mma）。它的目标是理解 Tensor Core 调用路径，不是替代 cuBLAS/CUTLASS。

9.6 下一步技术（不在本章实现，但要知道存在）

• cp.async（sm_80+）：异步 global → shared，让 mainloop 拷贝与计算重叠
• Double / multi-stage buffer：2-4 块 shared 轮转，吞下 HBM 延迟
• swizzle layout：让 shared 读写避免 bank conflict 同时支持 TC 对齐
• mma.sync PTX（sm_80+）：比 WMMA 更细粒度，CUTLASS / FlashAttention 都用
• warpgroup MMA（sm_90+）：H100 的 wgmma，128 thread 协作大 tile
• TMA（sm_90+）：硬件 tensor 加载，不再写 load 循环

自检清单

练习题

1. 给 gemm_reg_tile 加 double buffer（开两块 shared，K-loop 中 prefetch 下一片）。
2. 用 __pipeline_memcpy_async（sm_80+）替换 shared 加载，看 GFLOPS 提升。
3. 把 WMMA 改成支持 bf16（把 __half 换成 __nv_bfloat16）。
4. 用 ncu 看 gemm_wmma 的 Tensor Active %，如果明显偏低，检查 tile、occupancy、shared bank conflict 和 global load 是否喂饱 Tensor Core。

9.9 工业实战：CUTLASS、低精度、autotune、生产 GEMM 决策树

9.9.1 CUTLASS — 不是库，是 GEMM 工厂

cuBLAS 是闭源黑盒，你不能改 kernel 内部行为；CUTLASS 是 NVIDIA 开源的 C++ 模板库，把 GEMM 拆成可拼装的组件：

// CUTLASS CUDA C++ template API（不是 Python CuTe DSL）
#include <cutlass/gemm/device/gemm_universal.h>

using Gemm = cutlass::gemm::device::GemmUniversal<
    cutlass::half_t, cutlass::layout::RowMajor,            // A: fp16 row-major
    cutlass::half_t, cutlass::layout::ColumnMajor,         // B: fp16 col-major
    cutlass::half_t, cutlass::layout::RowMajor,            // C: fp16 row-major
    float,                                                  // accumulator: fp32
    cutlass::arch::OpClassTensorOp,                        // Tensor Core
    cutlass::arch::Sm80,                                   // A100
    cutlass::gemm::GemmShape<128, 256, 32>,                // ThreadBlock tile
    cutlass::gemm::GemmShape<64, 64, 32>,                  // Warp tile
    cutlass::gemm::GemmShape<16, 8, 16>,                   // Instruction shape (mma.sync)
    cutlass::epilogue::thread::LinearCombinationRelu<...>, // epilogue: out = relu(alpha*acc + beta*C)
    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
    3                                                       // pipeline stages (软流水深度)
>;

Gemm gemm;
gemm({M, N, K}, alpha, A, lda, B, ldb, beta, C, ldc, ...);

看似冗长，但 CUTLASS 模板能把 tile、pipeline、swizzle 和 epilogue 交给成熟库生成。大公司生产代码大量基于 CUTLASS（TensorRT-LLM、FlashAttention、xformers），最终性能仍要用 profiler 选择具体配置。

CUTLASS profiler — 自动 autotune

不知道你的 size 该选哪个 tile？跑 profiler：

cd cutlass/build
./tools/profiler/cutlass_profiler \
    --kernels=cutlass_tensorop_h16816gemm \
    --m=4096 --n=4096 --k=4096 \
    --A=f16:row --B=f16:col --C=f16 \
    --accumulator-type=f32 \
    --providers=cutlass,cublas

# 输出: top 5 best-performing configurations

9.9.2 mma.sync vs WMMA：CUTLASS 用前者

本章用了 WMMA API（nvcuda::wmma），它是 fragment 级抽象。CUTLASS 用更底层的 mma.sync PTX 指令：

// mma.sync 长这样 (Ampere fp16):
asm volatile(
    "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
    "{%0,%1,%2,%3}, "                  // D (acc, 4 fp32)
    "{%4,%5,%6,%7}, "                  // A (4 fp16x2)
    "{%8,%9}, "                        // B (2 fp16x2)
    "{%0,%1,%2,%3};\n"                 // C (= D in/out)
    : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
    : "r"(a0), "r"(a1), "r"(a2), "r"(a3),
      "r"(b0), "r"(b1));

普通工程师不需要直接写 mma.sync——CUTLASS 模板已封装。但读 FlashAttention v2 源码必看懂。

9.9.3 低精度类型一览：fp16 / bf16 / tf32 / fp8 / int8 / int4

fp16 vs bf16：选哪个？

• 训练：用 bf16。同 fp32 的动态范围避免梯度下溢，需要的 loss scaling 少。
• 推理：fp16 略快（编码精度高，TC 路径更成熟），但 bf16 更稳，新代码倾向 bf16。
• FlashAttention：fp16 / bf16 都支持，看模型权重类型。

fp8 — H100 的 LLM 杀手锏

Llama 70B 这类模型使用 fp8 推理（权重或 KV cache 采用 fp8）时，相比 fp16 的主要变化是：

• 显存 / HBM 流量下降，decode 的 roofline 上限提高
• 低精度 Tensor Core 理论吞吐更高，prefill 和 verifier 更容易受益
• 精度要用模型自己的校准集、PPL 和业务指标回归确认

调用：

cublasGemmEx(h, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K,
             &alpha,
             dB_e4m3, CUDA_R_8F_E4M3, N,
             dA_e4m3, CUDA_R_8F_E4M3, K,
             &beta,
             dC_half, CUDA_R_16F, N,
             CUBLAS_COMPUTE_32F,
             CUBLAS_GEMM_DEFAULT);

9.9.4 production GEMM 决策树

graph TD
    Start["GEMM with M, N, K, dtype"] --> M1{"M 等于 1 ?"}
    M1 -->|"是"| GEMV["走 GEMV kernel
weight-only int4 量化"]
    M1 -->|"否"| Size{"min of M N 小于 64 ?"}
    Size -->|"是"| Thin["thin GEMM
每 SM 算一行或一列"]
    Size -->|"否"| K1{"K 大于 4 倍 min M N ?"}
    K1 -->|"是"| StreamK["CUTLASS Stream-K"]
    K1 -->|"否"| Dtype{"dtype 是 fp16 bf16 fp8 ?"}
    Dtype -->|"是"| TC["cublasGemmEx + TC
或 CUTLASS"]
    Dtype -->|"否 fp32"| Sgemm["cublasSgemm"]
    style GEMV fill:#f3f1e8,stroke:#a86420
    style TC fill:#f3f1e8,stroke:#2f5d3a

9.9.5 LLM 各层 GEMM 的实际 shape

Llama 7B (D=4096, n_head=32, d_head=128, d_ff=11008)，prefill 阶段 (T=2048, batch=1)：

同样模型 decode 阶段 (M=1)：所有 GEMM 都变 GEMV，FLOPs 减为 1/2048 但 HBM 流量不变 → 完全 memory-bound。

这就是 LLM 推理的"两个世界"：prefill 在 compute-bound，可以用大 batch + TC 充分利用算力；decode 在 memory-bound，只能靠量化 + KV cache 优化 + speculative decoding 救。

9.10 研究前沿（2025-2026）：FP4 GEMM、DeepSeek 原生 FP8、CUTLASS 4.6.1 / CuTe DSL（2026-07 快照）

9.10.1 FP4 GEMM — Blackwell 的杀招

2025 NVIDIA Blackwell 把"用 4-bit 浮点训练 / 推理"从研究变成产品。两种 fp4 微缩放格式：

FP4 GEMM 调用（CUTLASS 4.6.1 CUDA C++ path）

using Gemm = cutlass::gemm::collective::CollectiveBuilder<
    cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp,
    cutlass::nv_float4_t, cutlass::layout::RowMajor, 16,        // A: NVFP4, block 16
    cutlass::nv_float4_t, cutlass::layout::ColumnMajor, 16,     // B: NVFP4
    float, /*accumulator*/                                       // accumulator fp32
    cutlass::gemm::collective::StageCount<4>,                    // 4-stage pipeline
    ...
>::CollectiveOp;

性能判断方式：先按官方 dense/sparse 峰值算 roofline，再用 CUTLASS profiler 或引擎内 benchmark 采集实际 TFLOPS。 FP4 的核心价值是同时降低权重/激活字节数并使用 Blackwell Tensor Core；真实吞吐还取决于 scale 读取、block size、epilogue 融合和 batch shape。

9.10.2 DeepSeek-V3 原生 FP8 训练 — 验证可行性

2024.12 DeepSeek-V3 发布，用 fp8 训练 671B MoE 模型成功，是产业第一个公开复现的案例。技术要点：

• tile-wise 1×128 + block-wise 128×128 量化：activation per-token-per-tile, weight per-128-block；细粒度 scale 让精度损失 < bf16 训练 1%
• fp32 accumulator + 周期性 promote to bf16：避免 fp8 累加器溢出
• 关键算子（norm、softmax）保持 bf16：fp8 只用在 GEMM
• 训练成本：仅 ~558 万美元（H800 集群 2 个月），证明 fp8 路线经济可行

开源 fp8 训练框架推荐：

• TransformerEngine（NVIDIA）：Hopper+ 训练首选
• FBGEMM_GPU（Meta）：fp8 + 量化 GEMM
• vLLM W8A8 fp8：推理
• SGLang FP8：推理 + 训练混合

9.10.3 CUTLASS 4.6.1 与 CuTe DSL（2026-07 快照）

按 2026-07 的官方 overview 与 changelog，必须把两条入口分开理解：

• CUTLASS C++：继续提供 CUDA C++ template abstractions、collective builder 和 device kernels。
• Python CuTe DSL：用 layout algebra、JIT/compile API 与 pipeline abstractions 编写 kernels。
• 二者都覆盖 Blackwell 能力，但版本支持和示例范围不同；“CUTLASS 已全面改写成 DSL”是不准确的。
• 选择入口后仍需用 CUTLASS profiler 或端到端 workload 验证 shape、dtype 与 epilogue。

# CUTLASS Python 调用 (2025+)
import cutlass
plan = cutlass.op.Gemm(element_a=cutlass.DataType.bf16,
                       element_b=cutlass.DataType.bf16,
                       element_c=cutlass.DataType.bf16,
                       element_accumulator=cutlass.DataType.f32,
                       layout=cutlass.LayoutType.RowMajor)
plan.run(A, B, C, alpha=1.0, beta=0.0)        # 自动选最优 tile

9.10.4 W4A4 / W4A8KV4 — 极致量化推理（2024-2025）

关键论文：

• QServe (MIT, 2024)：W4A8KV4，把权重、激活、KV cache 都纳入量化设计；复现时要按论文 shape 和目标引擎重测
• QuaRot (ICML 2024)：用 Hadamard 旋转把激活的 outlier 平摊，支持 W4A4 几乎无损
• SmoothQuant（2022 老经典）：活了 4 年仍是 W8A8 默认
• FP8 vs INT8（2024 多篇）：结论是 fp8 精度略高、HW 支持更好，新硬件优先 fp8

9.10.5 训练 GEMM 的新范式：Microscaling + Transformer Engine 2

Hopper 上的 fp8 训练已经稳定（NVIDIA TE 1）。Blackwell 把它升级到动态精度训练：

• 每层 forward / backward自动选 fp4 / fp6 / fp8
• 梯度小的用 fp4，大的用 fp8
• 关键 norm / softmax 自动 fall back bf16
• 训练吞吐要按模型、序列长度、并行策略和 precision recipe 在目标集群上测

9.10.6 工业 GEMM 库格局（2026）

常见坑

• cuBLAS row/col-major 算错 → 结果是 C^T，调试一上午才发现
• WMMA 的 M/N/K 不是 16 倍数 → 编译能过运行结果 0
• fp16 输入没归一化 → 大量上溢，c_frag 直接 inf/nan
• shared mem 太大 (BM=BN=256) → kernel launch failed (resource exceeded)

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

WMMA Fragment、MMA PTX、wgmma、Warp Specialization、TMA + Swizzle

H100 TMA：硬件张量加载器与张量映射 (CUtensorMap)

9.6 节末尾提到 sm_90+ 的 TMA (Tensor Memory Accelerator)：一条指令把 global 张量 tile 异步搬到 shared memory，省掉手写 cp.async 循环。本节给出从主机端到 device 端的最小可运行骨架（基于 CUDA Programming Guide 13.2（核验：2026-07-20） §4.11.2.2）。

第一步：主机端 driver API 构造 CUtensorMap

TMA 不直接读全局指针，而是读一个 128 字节的 张量描述子。它由 driver API cuTensorMapEncodeTiled 在 host 上一次性构造，然后通过 __grid_constant__ 参数传入 kernel：

// host code — D=128, T_tile=64
    constexpr uint32_t rank = 2;
    uint64_t global_dim[rank]   = { (uint64_t)D, (uint64_t)T };        // (cols, rows)
    uint64_t global_stride[rank-1] = { D * sizeof(__half) };           // 单位: 字节
    uint32_t box_dim[rank]      = { 64, 64 };                           // shared 端 tile 形状
    uint32_t elem_stride[rank]  = { 1, 1 };

    CUtensorMap tmap{};
    CUresult res = cuTensorMapEncodeTiled(
        &tmap,
        CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_FLOAT16,
        rank,
        (void*)dQ,            // 全局起始地址
        global_dim,
        global_stride,
        box_dim,
        elem_stride,
        CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
        CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,                 // ← 关键, 见下
        CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
        CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
    

第二步：device 端用 cp_async_bulk_tensor 触发 TMA

#include <cuda/barrier>
    #include <cuda/ptx>
    namespace ptx = cuda::ptx;

    __global__ void gemm_tma(const __grid_constant__ CUtensorMap tmap_A,
                             const __grid_constant__ CUtensorMap tmap_B,
                             __half* C, int M, int N, int K)
    {
        // 1024 字节对齐 — swizzle 128B 模式要求
        __shared__ alignas(1024) __half As[64][64];
        __shared__ alignas(1024) __half Bs[64][64];

        #pragma nv_diag_suppress static_var_with_dynamic_init
        __shared__ cuda::barrier<cuda::thread_scope_block> bar;
        if (threadIdx.x == 0) init(&bar, blockDim.x);
        __syncthreads();

        cuda::barrier::arrival_token tok;
        if (threadIdx.x == 0) {
            int crd[2] = { blockIdx.x * 64, blockIdx.y * 64 };
            // 一条指令：把 64x64 fp16 tile 异步搬到 shared
            ptx::cp_async_bulk_tensor(
                ptx::space_shared, ptx::space_global,
                &As, &tmap_A, crd,
                cuda::device::barrier_native_handle(bar));
            tok = cuda::device::barrier_arrive_tx(bar, 1, sizeof(As));
        } else {
            tok = bar.arrive();
        }
        bar.wait(std::move(tok));     // TMA 完成后继续

        /* ... 在 As/Bs 上跑 wgmma / mma.sync ... */
    }
    

第三步：选 swizzle 模式避免 bank conflict

TMA 可让 shared 端按 swizzle 模式重新排列数据，使后续 wgmma/mma.sync 读取无 bank conflict。v13.1 PG 表 25 总结：

访问换算公式（128B 模式，PDF §4.11.2.2.5 表 24）：smem[y][x] ↔ smem[y][((y + offset) % 8) ^ x]，其中 offset = ((uintptr_t)smem_ptr / 128) % 8。这就是 CUTLASS / FA v3 内层访问看似 XOR 一通的来源——不是杂技，是按 swizzle 表展开。

Hopper Cluster + STAS：跨 thread block 的生产者-消费者流水

Hopper (sm_90) 引入了介于 grid 和 block 之间的 thread block cluster（最多 16 个 block，典型 2-8）。同一 cluster 内的 block 可：

• 互相访问对方的 shared memory（distributed shared memory, DSMEM）
• 用 cluster.sync() 同步
• 用 STAS（store async）异步把寄存器值写到远端 block 的 shared

这让 CUTLASS 4.6.1 / CuTe DSL（2026-07 快照） 的 collective builder 把 GEMM 的 "load K/V" 和 "compute" 分到不同 block，靠 cluster 内 DSMEM 接力，实现真正的多级流水。下面是 PG §4.11.3 中的环形生产者-消费者骨架（精简）：

#include <cooperative_groups.h>
    #include <cuda/barrier>
    #include <cuda/ptx>

    __global__ __cluster_dims__(8, 1, 1)
    void producer_consumer_kernel()
    {
        namespace cg  = cooperative_groups;
        namespace ptx = cuda::ptx;
        using barrier_t = cuda::barrier<cuda::thread_scope_block>;

        auto cluster = cg::this_cluster();
        __shared__ int       buffer[BLOCK_SIZE];
        __shared__ barrier_t filled, ready;

        if (threadIdx.x == 0) {
            init(&filled, 1);
            init(&ready,  BLOCK_SIZE);
        }
        cluster.sync();   // ← 保证所有 block 的远端 barrier 都已初始化

        int rk      = cluster.block_rank();
        int rk_next = (rk + 1) % 8;
        int rk_prev = (rk + 7) % 8;

        auto buffer_next = cluster.map_shared_rank(buffer, rk_next);
        auto bar_next    = cluster.map_shared_rank(
                              cuda::device::barrier_native_handle(filled), rk_next);

        for (int it = 0; it < 1000; ++it) {
            // 1) 生产者: 把寄存器值 STAS 到下一 block 的 shared buffer
            ptx::st_async(&buffer_next[threadIdx.x], rk, bar_next);

            if (threadIdx.x == 0) {
                ptx::mbarrier_arrive_expect_tx(
                    ptx::sem_release, ptx::scope_cluster, ptx::space_shared,
                    cuda::device::barrier_native_handle(filled), sizeof(buffer));
            }
            // 2) 消费者: 等本地 filled barrier (左邻居写来)
            while (!ptx::mbarrier_try_wait_parity(
                      cuda::device::barrier_native_handle(filled), it & 1)) { }

            /* ... 在 buffer 上做计算 ... */
        }
    }
    

在 GEMM 上的映射：

wgmma vs mma.sync：Hopper warpgroup MMA 的 PTX 真身

Ampere 的 mma.sync 由 一个 warp (32 lane) 协作完成 16×8×16 fp16；Hopper 的 wgmma.mma_async 由 一个 warpgroup (4 warps = 128 lane) 协作完成更大的 m×n×k 操作，且是 异步（提交后立即返回，靠 fence/commit_group 等待）。

典型 wgmma 触发序列（Hopper, fp16 输入 fp32 累加）：

// 1) TMA load A, B 到 shared (上节代码), 用 128B swizzle
    // 2) 构造 wgmma "descriptor" — 编码 shared 地址 + swizzle 模式
    //    libcu++ / CUTLASS 把这步封装在 SM90_GMMA::Desc

    // 3) issue 一组 wgmma
    asm volatile (
        "wgmma.fence.sync.aligned;\n"
        "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 "
        "{%0,%1,...,%63},"        // 64 个 fp32 累加器（warpgroup 持有）
        " %64, %65, "             // A, B 的 shared descriptor (uint64)
        " 1, 1, 1, 0, 0;\n"       // scale_d, scale_ab, im_layout 等
        "wgmma.commit_group.sync.aligned;\n"
        : /* outputs ... */
        : "l"(desc_A), "l"(desc_B));

    // 4) 等待这一组完成
    asm volatile ("wgmma.wait_group.sync.aligned 0;\n");
    

    graph LR
        H["Host: cuTensorMapEncodeTiled
构造 A/B 的 CUtensorMap"]
        H --> K["Kernel launch
__cluster_dims__(2,1,1)"]
        K --> P["Producer block:
cp.async.bulk.tensor (TMA)
→ DSMEM"]
        P --> C["Consumer block warpgroup:
wgmma.mma_async × N
fp32 累加在 regs"]
        C --> EP["Epilogue:
cp.async.bulk.tensor (TMA store)
→ HBM"]
        style P fill:#f3f1e8,stroke:#a86420
        style C fill:#f3f1e8,stroke:#2f5d3a
    

普通工程师写 GEMM 不需要手写这段 wgmma 内联汇编——用 CUTLASS 4.6.1 / CuTe DSL（2026-07 快照） 的 CollectiveBuilder<arch::Sm90, ...> 自动生成。但 9.10.2 提到 DeepSeek-V3 自己手写了 fp8 wgmma kernel（绕开 CUTLASS 模板膨胀编译时间），这段背景是必备。

下一章导览

第 10 章把 reduction 与数值稳定性组合成 Softmax、LayerNorm 和 RMSNorm。