Skip to content
Gains Summary
Main Navigation 首页 / Home
C++ 编程 / C++ Programming
系统与高性能 / Systems & Performance
Web 开发 / Web Development
人工智能 / Artificial Intelligence
工业软件 / Industrial Software
其他内容 / Other Topics
C++ 编程 / C++系统与性能 / SystemsWeb 开发 / Web人工智能 / AI工业软件 / Industrial

外观

Sidebar Navigation

← 人工智能 / Artificial Intelligence

AI 编译器 / AI Compilers

1. AI 编译器全景——为什么模型需要编译器 / The AI Compiler Landscape and Why Models Need Compilers

2. 编译原理速通——面向 ML 工程师的核心概念 / Compiler Fundamentals for Machine Learning Engineers

3. 中间表示基础——理解 IR 层级与 lowering 链路 / Intermediate Representation Levels and Lowering Pipelines

4. 计算图的构建与表示 / Building and Representing Computational Graphs

5. MLIR 架构、方言与渐进式降级 / MLIR Architecture, Dialects, and Progressive Lowering

6. 算子语义、广播、归约与形状推导 / Operator Semantics, Broadcasting, Reduction, and Shape Inference

7. 模型前端格式:ONNX、TFLite、HLO 与 SavedModel / Model Frontend Formats: ONNX, TFLite, HLO, and SavedModel

8. 图优化 Pass——经典优化在 ML 中的应用 / Graph Optimization Passes for Machine Learning

9. 算子融合——编译器最重要的性能优化 / Operator Fusion as a Core Compiler Optimization

10. 内存规划——Buffer 分配与生命周期管理 / Memory Planning, Buffer Allocation, and Lifetime Management

11. Layout 优化——数据排布转换与内存效率 / Layout Optimization for Data Movement and Memory Efficiency

12. 动态 Shape——符号分析与形状处理 / Dynamic Shapes, Symbolic Analysis, and Shape Processing

13. 硬件约束下的操作调度 / Operation Scheduling Under Hardware Constraints

14. 从模板、DSL 到 IR 降级的代码生成架构 / Code Generation Architectures from Templates and DSLs to IR Lowering

15. CPU 后端:SIMD、分块与多线程 / CPU Backends with SIMD, Tiling, and Multithreading

16. CUDA 后端:合并访存与 Tensor Core / CUDA Backends, Memory Coalescing, and Tensor Cores

17. NPU 后端:脉动阵列与端侧 AI 生态 / NPU Backends, Systolic Arrays, and Edge AI Ecosystems

18. Kernel 性能基础:Roofline 与 Occupancy / Kernel Performance Fundamentals with Roofline and Occupancy

19. CUTLASS 与分层 GEMM 模板 / CUTLASS and Hierarchical GEMM Templates

20. TVM Tensor Expression 与计算调度分离 / TVM Tensor Expressions and Compute-Schedule Separation

21. 使用 Triton 编写高性能 GPU Kernel / Triton for High-Performance GPU Kernels in Python

22. 基于成本模型与实测搜索的自动调度 / Automatic Scheduling with Cost Models and Measurement-Based Search

23. XLA 内部机制:HLO、融合与 SPMD / XLA Internals, HLO, Fusion, and SPMD

24. Torch-MLIR:从 PyTorch 算子到 MLIR 方言 / Torch-MLIR from PyTorch Operators to MLIR Dialects

25. torch.compile:Dynamo、AOTAutograd、Inductor 与 Triton / Torch Compile with Dynamo, AOTAutograd, Inductor, and Triton

26. 从 MLIR 经 LLVM 降级到机器码 / Lowering from MLIR Through LLVM to Machine Code

27. 量化——低精度推理的工程实践 / Engineering Low-Precision Inference with Quantization

28. 分布式编译与训练——多设备编排的编译器支持 / Compiler Support for Distributed Training and Multi-Device Orchestration

29. 生产调试——真实问题的编译器视角排查 / Production Debugging from the Compiler Perspective

30. 未来方向——AI 编译器的新挑战与机遇 / Future Challenges and Opportunities for AI Compilers

本页目录

📅 创建时间:2026-06-03 🏷️ 标签:#CUTLASS #GEMM #WarpLevel #TensorCore #FP8 #INT8 #WMMA #NVIDIA #CuBlas 📚 前置知识:[[/04-ai/01-llm-engineering/07-llm-evolution]](LLM 发展脉络)[[17-kernel-primer]](Kernel 开发入门) 📚 相关知识:[[15-backend-cuda]](CUDA 后端)[[20-triton]](Triton)


CUTLASS 与分层 GEMM 模板 / CUTLASS and Hierarchical GEMM Templates ​

┌──────────────────────────────────────────────────────────────────────────────┐ │ 📖 场景:需要写一个 FP8 Transformer Attention Kernel │ ├──────────────────────────────────────────────────────────────────────────────┤ │ 你需要写一个 FP8 的 Transformer Attention kernel: │ │ - 矩阵维度:2048×2048,batch=16,heads=12 │ │ - 需要用 Hopper FP8 Tensor Core(WGMMA 指令) │ │ - cuBLAS 不支持 FP8,官方说用 CUTLASS │ │ - 看了 CUTLASS 文档,template 嵌套了 8 层,参数几十个,完全不知道怎么下手 │ └──────────────────────────────────────────────────────────────────────────────┘

第1节 CUTLASS 设计哲学——分层 GEMM 模板,原子操作可组合 ​

1.1 为什么需要 CUTLASS ​

NVIDIA 提供 cuBLAS 和 cuDNN 库,但它们是黑盒:你传入矩阵形状,库返回结果,无法定制。对于以下场景,cuBLAS/cuDNN 就无能为力了:

场景cuBLAS/cuDNNCUTLASS
FP8 量化(E4M3/E5M2)❌ 不支持✅ 完全支持
融合自定义激活函数(SiLU、GELU)❌ 只支持 ReLU/Sigmoid✅ SkEpilogue 可编程
非标准数据布局(行主序、Channel-Last)❌ 固定列主序✅ Iterator 可定制
混合精度(INT8 输入 + FP16 累加)✅ 支持✅ 支持,且更灵活
Warp-level 细粒度控制❌ 不开放✅ 完全可控
自定义规约操作❌ 不支持✅ 通过 Epilogue fusion 实现

CUTLASS 的核心理念:

cuBLAS = 高性能黑盒 (你不能改)
CUTLASS = 高性能白盒 (每个原子操作都可替换)
1
2

1.2 CUTLASS 的分层架构 ​

┌─────────────────────────────────────────────────────────────────────────────┐
│                        CUTLASS 分层架构                                    │
│                                                                             │
│  Level 4: MMA Instruction   ──► 最低层:硬件指令层面的矩阵乘指令               │
│            (MMA, WGMMA)         WMMA: Warp-level MMA (sm_70+)              │
│                               WGMMA: Hopper Warp-group MMA (sm_90)          │
│                                    ┌──────────────┐                          │
│  Level 3: Warp-level GEMM    ──► │  WarpMma     │  一个 warp 处理         │
│            (Warp-level)            │  (8×8 or 16×16 tiles)                │
│                                    └──────────────┘                          │
│                                        ┌──────────────────┐                   │
│  Level 2: Threadblock GEMM    ──►    │  ThreadblockMma │  一个 SM 处理      │
│            (SM-level)                 │  (64×64 tiles)  │  多个 warp 协作    │
│                                        └──────────────────┘                   │
│                                            ┌──────────────────────┐          │
│  Level 1: Collective GEMM         ──►     │  CollectiveGemm    │  多 SM     │
│            (Device-level)                 │  (多个 threadblock)  │  协作      │
│                                            └──────────────────────┘          │
│  Level 0: Mainloop/Epilogue    ──►    ┌──────────────────────┐             │
│            (Kernel 入口)               │  Gemm kernel        │  包含       │
│                                          │  (mainloop+epilogue) │ mainloop   │
│                                          └──────────────────────┘ 循环和    │
│                                                                       epilogue│
└─────────────────────────────────────────────────────────────────────────────┘
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

1.3 CUTLASS 2.x vs 3.x 对比 ​

特性CUTLASS 2.xCUTLASS 3.x
SM 架构Volta/Ampere (sm_70-sm_89)Hopper (sm_90) + Ada
Tensor CoreTensor Core (16×16)WGMMA (Hopper 专用)
Warp-level GEMMcutlass::gemm::warp::GemmTmaWarpSpecialized
Ping-Pong 双缓冲❌ 手动✅ 自动 (CUTE 库)
** Cooperative Loading**❌ 手动✅ TMA (Transfer Memory Accelerator)
FP8 支持基础支持完整支持 (E4M3/E5M2)

第2节 核心概念详解 ​

2.1 MMA (Matrix Multiply Accumulate) ​

MMA = Matrix Multiply Accumulate,是 GPU 上矩阵乘加运算的硬件指令名称。

传统 WMMA (Ampere and earlier):
  - 一个 warp (32 threads) 协同执行一次 MMA
  - 处理 16×16×16 的矩阵乘累加
  - M=N=K=16 (FP32累加) 或 M=N=K=8 (FP64)
  - 一次 WMMA: 2×16×16 = 512 FLOPs (乘加)

Hopper WGMMA (sm_90):
  - 一个 warp-group (4 warps = 128 threads) 执行一次 WGMMA
  - 处理 16×16×16 或 64×16×16 (取决于指令)
  - 一次 WGMMA: 2×64×16 = 2048 FLOPs
  - 比 WMMA 快 4 倍
1
2
3
4
5
6
7
8
9
10
11

2.2 Iterator——遍历矩阵的迭代器 ​

Iterator 是 CUTLASS 的核心抽象:把矩阵遍历逻辑从计算逻辑中解耦出来。

cpp
// CUTLASS Iterator 示例:行主序 vs 列主序

// 场景:你有一个矩阵 A,需要按行主序遍历
// 问题:A 在内存中是列主序(cuBLAS 约定),你不想手动算索引

// CUTLASS 的解决方案:Iterator 封装了索引计算
// 你只需要说"我要从 (i,j) 开始"和"步长是多少",Iterator 自动处理

// 伪代码展示 Iterator 模式
template <typename Layout_, typename Tile_, typename Threadblock_>
class GemmIteratorA {
    // 内部状态
    LongIndex pointer_;      // 当前内存地址
    LongIndex stride_i_;     // 行方向步长(元素个数)
    LongIndex stride_k_;     // K 方向步长(用于 GEMM 中 A[i,k])
    
    // operator++ 移动到下一个位置
    // 自动处理:行主序/列主序、tile 边界、bank conflict 避免
    
    // 核心接口
    void operator++() {
        // 行主序 Layout
        if (layout_is_row_major) {
            pointer_ += stride_k_;  // 移动到下一个 K 索引
        } else {
            // 列主序 Layout
            pointer_ += 1;  // 移动到下一个元素
        }
    }
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

CUTLASS 预定义的 Iterator 类型:

Iterator用途支持 Layout
GemmStreamA遍历 A 矩阵(行主序)RowMajor, ColumnMajor
GemmStreamB遍历 B 矩阵(列主序)RowMajor, ColumnMajor
GemmIdentitySwizzle避免 shared memory bank conflict自动 swizzle
GlobalTensor直接从 global memory 加载-

2.3 Fragment——寄存器级的矩阵分块 ​

Fragment = 一个线程处理的一个小矩阵块,存在寄存器中。

cpp
// WMMA Fragment 的概念
// 在 sm_70+ 上,WMMA 操作一次处理 16×16 的矩阵
// 但每个线程只处理其中一小块(fragment)

// 示例:处理 16×16 WMMA 时,每个线程负责
//  - A: 16×16 矩阵,每个线程负责一列或一个分块
//  - B: 16×16 矩阵,每个线程负责一行或一个分块
//  - C: 16×16 结果,每个线程负责一列或一个分块

// 使用 wmma::fragment 声明
#include <mma.h>

wmma::fragment<wmma::matrix_a, 16, 8, 16, half, wmma::row_major> a_frag;
// a_frag: 16×16 的 A 矩阵分块,half 类型,row_major 排列
//         存在寄存器中

wmma::fragment<wmma::matrix_b, 16, 8, 16, half, wmma::col_major> b_frag;
// b_frag: 16×16 的 B 矩阵分块,half 类型,col_major 排列

wmma::fragment<wmma::accumulator, 16, 8, 16, float> c_frag;
// c_frag: 累加器,必须是 float 或 int 类型(TF32/FP32累加)

// 执行一次矩阵乘累加
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
// c_frag = A × B + c_frag
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

2.4 Epilogue——融合 Element-wise 操作 ​

Epilogue = GEMM 之后的 element-wise 操作,如激活函数、bias 加法、scale、quantize。

┌──────────────────────────────────────────────────────────┐
│                    GEMM 主循环                          │
│                                                          │
│   for k in 0..K:                                        │
│       C_tile += A_tile × B_tile                          │
│                                                          │
├──────────────────────────────────────────────────────────┤
│                    Epilogue Fusion                       │
│                                                          │
│   C_tile = activation(C_tile)         # GELU/SiLU/Sigmoid│
│   C_tile = C_tile + bias               # Bias Add       │
│   C_tile = C_tile * scale              # 动态 Scale     │
│   C_tile = quantize(C_tile)            # FP32→INT8      │
│                                                          │
│  融合后的优势:减少 memory traffic,只写一次 global memory │
└──────────────────────────────────────────────────────────┘

CUTLASS Epilogue fusion:
  融合的 element-wise 操作在寄存器中完成
  不用额外 kernel launch,只增加少量寄存器压力
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

第3节 GEMM 分层实现 ​

3.1 Level 2: Threadblock-level GEMM ​

Threadblock 是 GPU 的基本调度单位。一个 Threadblock 分配到一个 SM 上执行。

Threadblock GEMM 工作原理:

1. 加载 A tile 到 shared memory (shape: ThreadblockM×K)
   - 多个线程协同加载,每个线程负责一部分
   - 使用 cooperative loading(原子操作确保无重复加载)

2. 加载 B tile 到 shared memory (shape: K×ThreadblockN)
   - 同上

3. 主循环:
   for k_tile in 0..K step K_tile_size:
       // Warp-level GEMM:每个 warp 计算一部分 C
       for warp in warps_in_threadblock:
           warp.mma_sync(A[warp_row], B[warp_col], C[warp_tile])
       // 移动到下一个 k tile
       advance A, B

4. Warp 间同步 (__syncthreads)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

3.2 Level 3: Warp-level GEMM ​

cpp
// CUTLASS 中 Warp-level GEMM 的简化概念
// 每个 warp 负责 C 矩阵的一个 tile (例如 64×64)

template <typename GemmShape_, typename IteratorA_, typename IteratorB_>
class WarpGemmMma {
private:
    // Warp 内 32 个线程协作
    // 每个线程持有 A 和 B 的一个 fragment
    
    // Mma 操作一次处理 16×16 块
    // 64×64 = 4×4 个 16×16 块
    // 需要 16 次 WMMA 操作完成整个 64×64 tile
    
public:
    // 执行 warp-level 矩阵乘
    CUTLASS_DEVICE
    void operator()(
        FragmentC& d,           // 输出分块
        FragmentA a_frag,       // A 分块
        FragmentB b_frag,       // B 分块
        FragmentC acc_frag      // 累加器(初始为 0)
    ) {
        // 循环执行 4×4 = 16 次 16×16 WMMA
        #pragma unroll
        for (int m = 0; m < 4; ++m) {
            for (int n = 0; n < 4; ++n) {
                wmma::mma_sync(
                    acc_frag[m][n],   // 累加到 acc_frag[m][n]
                    a_frag[m],        // A[m][:]
                    b_frag[n],        // B[:][n]
                    acc_frag[m][n]    // 累加
                );
            }
        }
        d = acc_frag;
    }
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

3.3 Level 4: MMA Instruction Emulation vs Actual MMA ​

cpp
// CUTLASS 提供了两种实现路径

// 路径 1: 使用硬件 MMA 指令 (sm_70+)
#if defined(CUTLASS_USE_WMMA)
    // 编译时选择 WMMA 路径
    // 在 Ampere (sm_80) 上使用 tensorcore__simt 模式
#endif

// 路径 2: 纯 SIMT 模拟 (所有架构)
#if defined(CUTLASS_USE_SIMT)
    // 在不支持 Tensor Core 的 GPU 上也能运行
    // 性能较低,但正确性有保证
#endif

// 对于 Hopper (sm_90),使用 WGMMA
#if defined(CUTLASS_USE_WGMMA)
    // WGMMA 有更严格的 alignment 要求:
    // - A: 16-byte alignment
    // - B: 16-byte alignment
    // - M, N: 16 的倍数
    // - K: 16 的倍数(或者使用 TMA 自动处理边界)
#endif
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

第4节 Epilogue Fusion——融合 Element-wise 操作 ​

4.1 为什么 Epilogue Fusion 重要 ​

不融合:
  kernel 1: C = A @ B           // GEMM kernel
  kernel 2: C = activation(C)    // activation kernel,需要再读一次 C
  kernel 3: C = C + bias        // bias kernel,又要读一次 C

  总共读 C 两次,写 C 两次 → 浪费内存带宽

融合后:
  kernel 1: C = activation(A @ B + bias)
  
  只读 C 一次,写 C 一次 → 节省 50% 内存带宽
1
2
3
4
5
6
7
8
9
10
11

4.2 CUTLASS 的 Epilogue fusion 实现 ​

cpp
// CUTLASS 的 Epilogue fusion 通过 SkEpilogue 实现
// Sk = "Swizzle" + "Kernel",swizzle 用于避免 bank conflict

// 使用 CUTLASS 预定义的 Epilogue 模板
#include <cutlass/epilogue/thread/linear_combination.h>

// LinearCombination: C = alpha * (A @ B) + beta * C
// 用于:scale + bias + activation 的组合

// 定义 Epilogue
using Epilogue = cutlass::epilogue::threadblock::
    Epilogue<...>;

// 或者自定义 Epilogue functor
template <typename ElementOutput_, int kThreads_>
class CustomEpilogue {
public:
    struct Params {
        float alpha;
        float beta;
        ElementOutput_* bias;  // 可选的 bias 指针
        ActivationType activation;  // GELU, SiLU, ReLU
    };
    
    CUTLASS_DEVICE
    void operator()(
        ElementOutput_* output,    // 写入 global memory 的地址
        int output_idx,            // 线性索引
        ElementOutput_* output_end,
        ElementOutput_* source,     // GEMM 结果
        Params params
    ) {
        // 每个 thread 处理多个元素
        for (; output < output_end; ++output, ++source) {
            // 读取 GEMM 结果
            float val = static_cast<float>(*source);
            
            // 加上 bias
            if (params.bias) {
                val += static_cast<float>(params.bias[output_idx]);
            }
            
            // 应用激活函数
            val = apply_activation(val, params.activation);
            
            // 写回 global memory
            *output = static_cast<ElementOutput_>(val);
            ++output_idx;
        }
    }
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

4.3 常用 Epilogue 模式 ​

Epilogue 类型公式用途
LinearCombinationC = α*(A@B) + β*C_prev通用 scale + 残差连接
LinearCombinationReluC = max(0, α*(A@B) + β*C_prev)ReLU 激活
LinearCombinationSigmoidC = 1/(1+exp(-(α*(A@B) + β*C_prev)))Sigmoid 激活
Custom你定义的任何 element-wise 操作GELU, SiLU, LayerNorm

第5节 FP8 支持——E4M3 vs E5M2 ​

5.1 FP8 格式详解 ​

FP8 有两种格式,NVIDIA Hopper 都支持:

E4M3 (主要用于权重和激活):
  - 4-bit 指数 + 3-bit 尾数 + 1-bit 符号
  - 范围: ±240 (指数偏移)
  - 精度较高,范围适中
  - 适合:Weight、Activation

E5M2 (主要用于梯度和大动态范围):
  - 5-bit 指数 + 2-bit 尾数 + 1-bit 符号
  - 范围: ±57344 (更大的动态范围)
  - 精度较低,范围更大
  - 适合:Gradient、某些特殊激活

格式对比:
           E4M3           E5M2
  最小正数: 2^-9 ≈ 0.002   2^-14 ≈ 0.00006
  最大正数: 448            57344
  精度(尾数): 3 bits      2 bits
  
  E4M3 ≈ float8(标准精度)
  E5M2 ≈ float8(宽动态范围)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

5.2 FP8 GEMM 的 Scaling Factor ​

FP8 的精度有限,需要 per-tensor 或 per-channel scaling 来保持数值稳定:

cpp
// FP8 GEMM 需要维护 scaling factor
// 目的:将输入值缩放到 FP8 可表示的范围内

// 两种 scaling 策略:

// 1. Per-Tensor Scaling (简单,全局一个 scale)
float scale_a = compute_scale(tensor_a, Element_e4m3);
// 对整个 A 矩阵用一个 scale
// 实现:scale = max(abs(A)) / max_representable_e4m3

// 2. Per-Channel Scaling (更精细,每列/行一个 scale)
// 用于 Transformer 中的 Q/K/V 投影
std::vector<float> scale_a(A.columns);  // 每列一个 scale
for (int c = 0; c < A.columns; ++c) {
    scale_a[c] = compute_scale(A[:, c], Element_e4m3);
}

// 3. Per-Block Scaling (折中方案)
// 每个 block 一个 scale,平衡精度和开销
float scale_a_block = compute_scale(A[block_rows, block_cols], Element_e4m3);
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

5.3 CUTLASS FP8 实现 ​

cpp
// CUTLASS 3.x 中 FP8 GEMM 的声明
#include <cutlass/gemm/gemm.h>
#include <cutlass/gemm/device/gemm_universal_adapter.h>
#include <cutlass/epilogue/collective/collective_builder.hpp>

// 定义 FP8 GEMM 的类型
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<
    cutlass::gemm::GemmUniversal<
        256,   // Threadblock M tile
        128,   // Threadblock N tile
        256,   // K tile
        cutlass::half_t,       // A 类型 (FP16 输入)
        cutlass::layout::RowMajor,
        cutlass::half_t,       // B 类型 (FP16 输入)
        cutlass::layout::ColumnMajor,
        float,                 // C 类型 (FP32 累加)
        cutlass::layout::RowMajor,
        float,                 // D 类型 (FP32 输出)
        
        // 使用 Hopper FP8 Tensor Core
        cutlass::arch::OpClassTensorOp,
        cutlass::arch::Sm90,   // Hopper architecture
        
        // Epilogue: 支持 scale + activation fusion
        cutlass::epilogue::collective::CollectiveBuilder<
            cutlass::arch::Sm90,
            cutlass::arch::OpClassTensorOp,
            256, 128, float,   // Epilogue tile shape
            float, float,
            cutlass::epilogue::thread::LinearCombination<half_t, float>
        >::CollectiveOp,
        
        // Mainloop: 使用 WGMMA
        cutlass::gemm::collective::CollectiveBuilder<
            cutlass::arch::Sm90,
            cutlass::arch::OpClassTensorOp,
            256, 128, 256,     // Tile shape
            cutlass::half_t,
            cutlass::layout::RowMajor,
            cutlass::half_t,
            cutlass::layout::ColumnMajor,
            cutlass::gemm::collective::StageCount<4>,  // Ping-pong 缓冲
            cutlass::gemm::collective::KernelScheduleAuto
        >::CollectiveOp
    >
;
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

第6节 CUTLASS 3.x Hopper 优化——WGMMA + Ping-Pong ​

6.1 TMA (Transfer Memory Accelerator) ​

Hopper 引入了 TMA,这是硬件级别的异步内存传输单元:

cpp
// TMA 的优势:
// 1. 异步传输:内存传输和计算可以重叠
// 2. 自动处理 alignment 和 boundary
// 3. 多维索引:直接指定 box 的 (x,y,z) 坐标

// CUTLASS 中使用 TMA
// 通过 cutlass::TmaDescriptor 配置 TMA

// TMA 加载 A tile 到 shared memory
// 这是一个异步操作,不阻塞 warp
tma_descriptor_a.load(
    shared_memory_base,     // 目标地址 (smem)
    {block_m, block_k},     // box size
    {tile_m, tile_k},       // tile position in source
    threadIdx.x, threadIdx.y
);
// warp 可以继续做其他事,不需要等待

// 同步点:需要使用 cp_async_wait_group(0) 等待 TMA 完成
cp_async_wait_group(0);
__syncthreads();
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

6.2 Ping-Pong 双缓冲 ​

Ping-Pong 双缓冲是 掩盖内存延迟 的经典技术:

Ping-Pong 双缓冲原理:

┌────────────────────────────────────────────────────────────┐
│ Time →                                                    
│                                                            
│  Buffer A          Buffer B                                 
│  ┌─────────┐                                                  
│  │ Compute │ ← 当前正在用的 buffer                         
│  │         │                                                 
│  └─────────┘                                                  
│        ↓ 完成后切换                                         
│  ┌─────────┐                                                  
│  │ Load B  │ ← 同时开始加载下一个 buffer                     
│  └─────────┘                                                  
│                                                            
│  理想情况:Load 和 Compute 完全重叠                         
│  内存延迟被 100% 掩盖                                      
└────────────────────────────────────────────────────────────┘

CUTLASS 3.x 的双缓冲实现:

// StageCount<4> = 使用 4 个 shared memory buffer
// 可以同时持有: 1个当前计算用 + 2个预加载 + 1个预热
using StageCount = cutlass::gemm::collective::StageCount<4>;

// CTA 会自动管理 double buffering
// warp 1 在计算 buffer A 时,warp 2 在加载 buffer B
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

第7节 实战——写一个自定义 GEMM Kernel ​

7.1 完整的 CUTLASS Kernel 实例化流程 ​

python
# Python-like pseudocode 展示 CUTLASS kernel 实例化
# 实际是 C++,这里用 Python 语法便于理解

# STEP 1: 定义矩阵形状
ThreadblockM = 128
ThreadblockN = 128
ThreadblockK = 64

WarpM = 64
WarpN = 64
WarpK = 64

# InstructionShape for WMMA: 16×16×16
InstructionM = 16
InstructionN = 16
InstructionK = 16

# STEP 2: 选择数据类型和 Layout
# A: RowMajor, FP16
# B: ColumnMajor, FP16
# C/D: RowMajor, FP32 (累加器)

ElementA = "half"    # __half (FP16)
ElementB = "half"    # __half
ElementC = "float"   # float (累加器)
ElementD = "float"   # float (输出)

LayoutA = "RowMajor"
LayoutB = "ColumnMajor"
LayoutC = "RowMajor"

# STEP 3: 定义 Epilogue (融合 bias + GELU)
class BiasGeluEpilogue:
    """
    D = GELU(A @ B + bias)
    """
    def __init__(self, bias_tensor, alpha=1.0, beta=0.0):
        self.bias = bias_tensor
        self.alpha = alpha
        self.beta = beta
    
    def apply(self, C_fragment, thread_idx):
        # 先加 bias
        C_fragment = C_fragment + self.bias[thread_idx]
        # 应用 GELU: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
        C_fragment = gelu(C_fragment)
        return C_fragment * self.alpha + C_fragment * self.beta

# STEP 4: 组装 kernel
GemmKernel = CutlassGemmKernel(
    tile_shape = (ThreadblockM, ThreadblockN, ThreadblockK),
    warp_shape = (WarpM, WarpN, WarpK),
    instruction_shape = (InstructionM, InstructionN, InstructionK),
    element_a = ElementA,
    element_b = ElementB,
    element_c = ElementC,
    element_d = ElementD,
    layout_a = LayoutA,
    layout_b = LayoutB,
    layout_c = LayoutC,
    epilogue = BiasGeluEpilogue,
)

# STEP 5: 实例化并运行
gemm = GemmKernel.instantiate()

# 运行 GEMM
result = gemm.run(
    A = input_a,      # shape: (M, K)
    B = input_b,      # shape: (K, N)
    C = output_c,     # shape: (M, N)
    bias = bias_tensor,
    alpha = 1.0,
    beta = 0.0,
    stream = cuda_stream,
)

# 或者使用 CUTLASS 提供的便捷 API
result = cutlass.gemm(
    A=input_a, B=input_b, C=output_c,
    threadblock_shape=(128, 128, 64),
    warp_shape=(64, 64, 64),
    epilogue_functor=BiasGeluEpilogue,
    epilogue_arguments={"bias": bias_tensor},
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

7.2 CUTLASS vs cuBLAS vs WMMA API 对比 ​

维度cuBLASCUTLASSWMMA API
易用性⭐⭐⭐⭐⭐ 简单调用⭐⭐ 模板复杂⭐⭐⭐ 需要手写调度
定制能力⭐ 不支持⭐⭐⭐⭐⭐ 完全可定制⭐⭐⭐ 部分可定制
性能⭐⭐⭐⭐⭐ 极优化⭐⭐⭐⭐ 接近 cuBLAS⭐⭐⭐ 取决于实现
FP8 支持❌ 无✅ 完整✅ 手动实现
Epilogue Fusion❌ 有限✅ 完全可编程❌ 手动实现
学习曲线低高中
适用场景标准 GEMM定制 GEMM + Fusion极致性能优化

7.3 CUTLASS 模板层次图 ​

┌─────────────────────────────────────────────────────────────────────────────┐
│                                                                             │
│   ┌─────────────────────────────────────────────────────────────────────┐   │
│   │  CollectiveEpilogue (多 SM 协作的 Epilogue)                          │   │
│   │   - 包含多个 Epilogue tile                                            │   │
│   │   - 处理整个 threadblock 的 output                                    │   │
│   └─────────────────────────────────────────────────────────────────────┘   │
│                                    │                                       │
│                                    ▼                                       │
│   ┌─────────────────────────────────────────────────────────────────────┐   │
│   │  ThreadblockEpilogue (单个 SM 的 Epilogue)                           │   │
│   │   - 处理 threadblock 级别的 post-processing                          │   │
│   │   - 例如: scale, bias, activation                                    │   │
│   └─────────────────────────────────────────────────────────────────────┘   │
│                                    │                                       │
│                                    ▼                                       │
│   ┌─────────────────────────────────────────────────────────────────────┐   │
│   │  CollectiveMainloop (多 SM 协作的 GEMM 主循环)                      │   │
│   │   - Cooperative Load: 多个 SM 协同加载 A, B                         │   │
│   │   - TMA Support (Hopper): 硬件异步传输                              │   │
│   └─────────────────────────────────────────────────────────────────────┘   │
│                                    │                                       │
│                                    ▼                                       │
│   ┌─────────────────────────────────────────────────────────────────────┐   │
│   │  ThreadblockMma (单个 SM 的 GEMM)                                   │   │
│   │   - 包含多个 WarpMma                                                │   │
│   │   - 管理 shared memory 缓冲                                         │   │
│   └─────────────────────────────────────────────────────────────────────┘   │
│                                    │                                       │
│                                    ▼                                       │
│   ┌─────────────────────────────────────────────────────────────────────┐   │
│   │  WarpMma (单个 Warp 的 GEMM)                                        │   │
│   │   - 使用 Tensor Core 或 WMMA 指令                                   │   │
│   │   - 处理 warp-level 的矩阵乘                                        │   │
│   └─────────────────────────────────────────────────────────────────────┘   │
│                                    │                                       │
│                                    ▼                                       │
│   ┌─────────────────────────────────────────────────────────────────────┐   │
│   │  SmemIterator / GmemIterator (矩阵遍历迭代器)                        │   │
│   │   - 封装索引计算逻辑                                                │   │
│   │   - 支持多种内存布局 (RowMajor, ColumnMajor, Swizzled)              │   │
│   └─────────────────────────────────────────────────────────────────────┘   │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

升华 ​

┌────────────────────────────────────────────────────────────────────────────┐
│                     🚀 CUTLASS 使用核心原则                                  │
├────────────────────────────────────────────────────────────────────────────┤
│                                                                            │
│  ① 先用 cuBLAS,再用 CUTLASS:只有 cuBLAS 满足不了时才考虑 CUTLASS           │
│     CUTLASS 的复杂度需要充分理由才值得承担                                   │
│                                                                            │
│  ② 从高层 API 入手:先用 CollectiveEpilogue/CollectiveMainloop 的组合        │
│     深入理解后再去改写底层 Iterator 和 Fragment                              │
│                                                                            │
│  ③ FP8 需要 Scaling 策略:没有正确的 scaling,FP8 精度会严重损失            │
│     Per-channel scaling 是 Transformer 的标准做法                            │
│                                                                            │
│  ④ Epilogue Fusion 是性能关键:融合激活函数可以减少 50% 内存带宽            │
│     但要注意寄存器压力,不要融合太多操作                                      │
│                                                                            │
│  ⑤ Hopper 的 TMA + Ping-Pong 是关键优化:                                   │
│     使用 CUTLASS 3.x + sm_90 获取 WGMMA 性能                               │
│                                                                            │
└────────────────────────────────────────────────────────────────────────────┘
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

"AI 可查 vs 必须理解"清单 ​

必须理解(不理解就等于不会):

  • 🔴 CUTLASS 分层架构:Level 1-4 各自负责什么,为什么要分层
  • 🔴 Iterator 模式:把矩阵遍历和计算逻辑解耦的意义
  • 🔴 Fragment 的作用:为什么需要在寄存器级别持有矩阵分块
  • 🔴 Epilogue Fusion 的原理:融合和不融合的内存访问差异
  • 🔴 WGMMA vs WMMA 的区别:Hopper 的 WGMMA 改变了什么

AI 可查(知道去哪查就行):

  • ✅ CUTLASS 具体 API 的参数签名(CUTLASS 官方文档)
  • ✅ cutlass::gemm::GemmUniversal 的完整模板参数列表(GitHub 源码)
  • ✅ E4M3/E5M2 的具体编码数值(IEEE 标准文档)
  • ✅ Hopper TMA 指令的具体语法(PTX 文档)
  • ✅ 特定 GEMM shape 的最优 threadblock/warp 配置(benchmark 数据)

学习状态:🟡 开始学习

最后更新于:

Pager
上一篇18. Kernel 性能基础:Roofline 与 Occupancy / Kernel Performance Fundamentals with Roofline and Occupancy
下一篇20. TVM Tensor Expression 与计算调度分离 / TVM Tensor Expressions and Compute-Schedule Separation

持续记录,持续成长

Copyright © Tidenflow