阶段 5 · 工程化与 AI15 / 15约 16 分钟
下一步：从会写 kernel 到吃透 GPU
Tensor Core、CUTLASS、Triton、FlashAttention 的阅读顺序与判断标准
学完这一课你会
- 了解 Tensor Core 编程的三个层次
- 知道什么时候该用 Triton，什么时候必须回到 CUDA
- 拿到一份有依据的进阶阅读清单
走完前面五个阶段，你已经具备读懂绝大多数开源 CUDA kernel 的基础。接下来的路会分叉，取决于你想解决什么问题。这一课不教新语法，而是给你一张地图。
Tensor Core：现代 GPU 的算力主体
一个残酷的事实：A100 的 FP32 算力是 19.5 TFLOPS，但 Tensor Core 的 BF16 算力是 312 TFLOPS——差 16 倍。如果你在做矩阵乘而没用上 Tensor Core，就等于只用了这张卡不到 7% 的能力。编程方式有三个层次，抽象程度递减、控制力递增。
| 层次 | 接口 | 适合谁 |
|---|---|---|
| 库 | cuBLAS / cuDNN | 标准操作，直接调，性能已是最优 |
| 模板库 | CUTLASS | 需要定制的 GEMM（融合 epilogue、特殊数据类型） |
| intrinsic | wmma / mma.sync PTX | 研究、极致优化、CUTLASS 覆盖不到的形状 |
CUDA C++
1#include <mma.h>2using namespace nvcuda;3 4// 一个 warp 协作完成 16x16x16 的矩阵乘累加5__global__ void wmmaGemm(const half* A, const half* B, float* C,6 int M, int N, int K) {7 // fragment 是分布在 warp 32 个线程寄存器里的分块，布局由硬件定义8 wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a;9 wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::col_major> b;10 wmma::fragment<wmma::accumulator, 16, 16, 16, float> acc;11 12 wmma::fill_fragment(acc, 0.0f);13 14 int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;15 int warpN = blockIdx.y * blockDim.y + threadIdx.y;16 17 for (int k = 0; k < K; k += 16) {18 wmma::load_matrix_sync(a, A + warpM * 16 * K + k, K);19 wmma::load_matrix_sync(b, B + k * N + warpN * 16, N);20 wmma::mma_sync(acc, a, b, acc); // 这一条指令用的就是 Tensor Core21 }22 23 wmma::store_matrix_sync(C + warpM * 16 * N + warpN * 16, acc, N,24 wmma::mem_row_major);25}Triton：什么时候该用
Python
1import triton2import triton.language as tl3 4@triton.jit5def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):6 pid = tl.program_id(0) # 相当于 blockIdx.x7 offs = pid * BLOCK + tl.arange(0, BLOCK) # 一次处理一整块，不是一个元素8 mask = offs < n # 边界由 mask 表达9 10 x = tl.load(x_ptr + offs, mask=mask)11 y = tl.load(y_ptr + offs, mask=mask)12 tl.store(out_ptr + offs, x + y, mask=mask)13 14# 没有 threadIdx，没有 __syncthreads，没有共享内存管理——15# 这些由 Triton 编译器根据 BLOCK 自动决定。| CUDA C++ | Triton | |
|---|---|---|
| 抽象粒度 | 单个线程 | 一个 block 的数据块 |
| 共享内存 | 手动分配、手动同步 | 编译器自动管理 |
| 合并访问 | 自己保证 | 编译器负责 |
| Tensor Core | 手写 WMMA / MMA | tl.dot 自动映射 |
| 开发效率 | 低 | 高 |
| 性能天花板 | 无上限 | 多数情况 80%~95%，特殊场景受限 |
| 可移植性 | 仅 NVIDIA | NVIDIA + AMD |
进阶阅读清单
- 1.CUDA C++ Programming Guide — 官方文档，不用通读，当字典查。特别是 Performance Guidelines 和 Compute Capabilities 附录。
- 2.CUDA C++ Best Practices Guide — 比 Programming Guide 更实用，优化建议按收益排序。
- 3.CUTLASS 源码 — 学习工业级 GEMM 的组织方式。从
examples/里的基础例子入手，别一上来啃include/cutlass/gemm/。 - 4.FlashAttention 论文 + 源码 — 融合优化的教科书级案例。先读懂在线 softmax 的推导，再看 kernel。
- 5.Nsight Compute 的 Kernel Profiling Guide — 把每个指标的定义搞清楚，profiling 效率会翻倍。
- 6.PTX ISA 文档 — 遇到看不懂的指令时查，不用通读。
- 7.Volkov, Better Performance at Lower Occupancy (GTC 2010) — 理解 ILP 与 occupancy 权衡的经典材料，至今不过时。
找练手项目
- ▸手写 GEMM 打榜：从朴素版一路优化到接近 cuBLAS，这是最好的综合练习。目标定在 cuBLAS 的 80%。
- ▸实现一个 FlashAttention：把 QK^T、softmax、乘 V 融合进一个 kernel，用在线 softmax 避免物化注意力矩阵。
- ▸给 PyTorch 写一个融合算子：比如 fused LayerNorm + 残差，在真实模型里测端到端收益。
- ▸参加 kernel 优化比赛：像 KernelBench 这类基准会给出明确的题目和排行榜，反馈循环很快。
- ▸读一个推理引擎的 kernel 目录：vLLM、llama.cpp 的 CUDA 后端、TensorRT-LLM，都是很好的真实代码样本。
自测
先自己在心里回答一遍，再展开对照。答不上来的说明这一段值得重读。