📅 创建时间: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/cuDNN | CUTLASS |
|---|---|---|
| FP8 量化(E4M3/E5M2) | ❌ 不支持 | ✅ 完全支持 |
| 融合自定义激活函数(SiLU、GELU) | ❌ 只支持 ReLU/Sigmoid | ✅ SkEpilogue 可编程 |
| 非标准数据布局(行主序、Channel-Last) | ❌ 固定列主序 | ✅ Iterator 可定制 |
| 混合精度(INT8 输入 + FP16 累加) | ✅ 支持 | ✅ 支持,且更灵活 |
| Warp-level 细粒度控制 | ❌ 不开放 | ✅ 完全可控 |
| 自定义规约操作 | ❌ 不支持 | ✅ 通过 Epilogue fusion 实现 |
CUTLASS 的核心理念:
cuBLAS = 高性能黑盒 (你不能改)
CUTLASS = 高性能白盒 (每个原子操作都可替换)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.3 CUTLASS 2.x vs 3.x 对比
| 特性 | CUTLASS 2.x | CUTLASS 3.x |
|---|---|---|
| SM 架构 | Volta/Ampere (sm_70-sm_89) | Hopper (sm_90) + Ada |
| Tensor Core | Tensor Core (16×16) | WGMMA (Hopper 专用) |
| Warp-level GEMM | cutlass::gemm::warp::Gemm | TmaWarpSpecialized |
| 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 倍2.2 Iterator——遍历矩阵的迭代器
Iterator 是 CUTLASS 的核心抽象:把矩阵遍历逻辑从计算逻辑中解耦出来。
// 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; // 移动到下一个元素
}
}
};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 = 一个线程处理的一个小矩阵块,存在寄存器中。
// 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_frag2.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,只增加少量寄存器压力第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)3.2 Level 3: Warp-level GEMM
// 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;
}
};3.3 Level 4: MMA Instruction Emulation vs Actual MMA
// 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第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% 内存带宽4.2 CUTLASS 的 Epilogue fusion 实现
// 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;
}
}
};4.3 常用 Epilogue 模式
| Epilogue 类型 | 公式 | 用途 |
|---|---|---|
LinearCombination | C = α*(A@B) + β*C_prev | 通用 scale + 残差连接 |
LinearCombinationRelu | C = max(0, α*(A@B) + β*C_prev) | ReLU 激活 |
LinearCombinationSigmoid | C = 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(宽动态范围)5.2 FP8 GEMM 的 Scaling Factor
FP8 的精度有限,需要 per-tensor 或 per-channel scaling 来保持数值稳定:
// 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);5.3 CUTLASS FP8 实现
// 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
>
;第6节 CUTLASS 3.x Hopper 优化——WGMMA + Ping-Pong
6.1 TMA (Transfer Memory Accelerator)
Hopper 引入了 TMA,这是硬件级别的异步内存传输单元:
// 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();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第7节 实战——写一个自定义 GEMM Kernel
7.1 完整的 CUTLASS Kernel 实例化流程
# 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},
)7.2 CUTLASS vs cuBLAS vs WMMA API 对比
| 维度 | cuBLAS | CUTLASS | WMMA 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) │ │
│ └─────────────────────────────────────────────────────────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────────────┘升华
┌────────────────────────────────────────────────────────────────────────────┐
│ 🚀 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 性能 │
│ │
└────────────────────────────────────────────────────────────────────────────┘"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 数据)
学习状态:🟡 开始学习