【GPU 算子工程】全景:算子工程在 AI 计算栈的位置
从框架一行 matmul 到 PTX/SASS,拆开 AI 计算栈的分层:框架算子、算子库、手写 kernel、编译器生成。回答工程师什么时候才需要自己写或调 kernel,以及本系列的实验环境与方法。
发布来自土法炼钢兴趣小组的知识、笔记、进展和应用。主题包括数据结构和算法、编程语言、网络安全、密码学等。
共 22 篇文章 · 返回首页
从框架一行 matmul 到 PTX/SASS,拆开 AI 计算栈的分层:框架算子、算子库、手写 kernel、编译器生成。回答工程师什么时候才需要自己写或调 kernel,以及本系列的实验环境与方法。
讲清 grid/block/warp 如何映射到 SM,SIMT 执行与 32 线程 warp 的本质,分支发散为何昂贵(实测 1.7 倍),以及 occupancy 的含义。建立一切 GPU 性能优化的硬件直觉。
从向量加法到归一化,讲清 CUDA kernel 的结构:全局索引计算、grid-stride loop、__syncthreads 同步、launch 配置选择与错误检查。实测 block 大小对带宽的影响,给出安全默认值。
global memory 合并访问与 shared memory bank conflict 是 GPU 访存优化的两大主题。实测跨步访问让有效带宽从 412 跌到 90 GB/s,32 路 bank conflict 让 shared 访问慢 11 倍。讲清成因与规避方法。
occupancy 是 SM 驻留 warp 与上限之比,由寄存器、shared memory、block 限制决定。实测访存密集 kernel 在约 33% occupancy 就饱和带宽,更高 occupancy 无益,并解释寄存器溢出为何让高 occupancy 反而变慢。
Roofline 用算术强度把算子定位到性能上限曲线,回答优化该往算力还是访存使劲。在 RTX 3060 Ti 上实测扫描算术强度,得到经验屋顶线:脊点约 36 FLOP/byte,低强度区贴带宽、高强度区逼近 FP32 峰值 86%。
Nsight Systems 看时间线找哪个 kernel 值得优化,Nsight Compute 看单 kernel 的 SM/内存吞吐、stall reason、occupancy 定位瓶颈。讲清两者分工、关键指标含义,以及没有 GUI 时用 CUDA event 计时的轻量替代方法。
归约是协作类算子的入门。实测三种 block 内归约树:发散+bank conflict 75ms、顺序寻址 44ms、warp shuffle 22ms。同时揭示单遍归约受访存限制时这些优化为何不可见,以及 scan 的并行思路。
GEMM 是 GPU 算子优化的标杆。在 RTX 3060 Ti 上实测四个版本:朴素 990、shared tiling 1309、寄存器分块 64 达 4447、128 达 6375 GFLOP/s(峰值 39%)。讲清每一步优化提高的是什么,以及为什么数据复用是关键。
Tensor Core 把矩阵乘做进专用硬件。实测 RTX 3060 Ti 的 FP16 Tensor 吞吐达 72.8 TFLOP/s,约 FP32 峰值的 4.5 倍。讲清 MMA 指令、wmma fragment API、数据布局与精度要求,以及为什么喂数据才是真正的瓶颈。
CUTLASS 用分层模板把 GEMM 拆成 device/kernel/threadblock/warp/instruction 五层,CuTe 用统一的 Layout 代数描述张量在各级存储的布局。讲清这套抽象如何在不手写 PTX 的前提下把 Tensor Core 喂到接近峰值。
归约类算子是 memory-bound 的典型。讲 softmax 的数值稳定写法(减最大值、在线 softmax)、LayerNorm 的 Welford 单遍方差,以及逐元素融合:实测把 scale+bias+GELU 三个 kernel 融成一个,提速 2.94 倍。
FlashAttention 把注意力重写成分块的在线 softmax,不落地 N×N 分数矩阵,用重算换访存。本文推导算法、给出实测正确的简化实现(误差 4e-7、避免 16.8MB 分数矩阵),并引用原论文的加速与显存数据。
Triton 用 tile(block of pointers)抽象替代 CUDA 的单线程视角,把合并访问、shared 管理、bank conflict 交给编译器,配合 autotune 自动搜配置。讲清它的编程模型、与手写 CUDA 的能力边界,以及为什么它成了算子开发主力。
融合通过减少中间结果的 HBM 往返提速 memory-bound 算子。实测逐元素链融合的加速比随链长线性增长(k=16 时 16.8 倍)。讲清逐元素融合、归约融合、GEMM epilogue 融合,以及什么时候不该融合。
低精度既省显存带宽又提算力。实测 FP16 逐元素算子比 FP32 快 1.81 倍。讲清量化的对称/非对称、per-tensor/per-channel 粒度、反量化时机、INT8 dp4a 与 Tensor Core 路径,以及精度对齐的工程坑。
多卡训练/推理中,通信不与计算重叠就是纯开销。讲三个层次的重叠:kernel 内 cp.async 异步加载、kernel 间 stream 并发、分布式里 NCCL collective 与反向计算的重叠,以及 SM 资源争抢的代价。
单个 kernel 到可维护算子库的工程问题:按 shape/dtype/arch 选 kernel 的 dispatch、autotune 结果缓存、AOT 与 JIT(NVRTC 运行时编译)的取舍。以本系列实际用的 NVRTC JIT 流程为例。
GPU kernel 的 bug 分两类:内存/竞态错误和数值错误。讲 compute-sanitizer 查越界与 race、对参考实现做容差对齐(实测 FlashAttention 误差 4e-7)、浮点非结合性带来的不可复现,以及常见同步陷阱。
算子工程的前沿方向:Hopper 的 TMA 异步搬运与 wgmma、Blackwell 的更低精度、ThunderKittens 等 tile 级库降低门槛、Triton/MLIR 的编译器自动生成算子。本系列测试卡为 Ampere,相关特性为引用与前瞻,明确标注。
从 GPU 执行模型与内存层次出发,系统讲解如何写出并调优高性能 CUDA 算子:访存合并、occupancy、Roofline、Nsight 调优,reduction/GEMM/Tensor Core/FlashAttention 核心算子实现,以及 Triton、CUTLASS、kernel fusion 与算子库工程。
拆开 GPU 的存储金字塔:寄存器、shared memory、L1/L2、global memory 的容量、带宽与延迟量级。用实测展示 L2 命中(约 3.4 TB/s)与 DRAM(约 400 GB/s)相差近一个数量级,解释为什么数据放哪决定算子性能。