Lecture 6 Kernels, Triton 课程笔记
本讲座:基准测试/分析 + 编写内核
总结
编程模型(PyTorch、Triton、PTX)与硬件之间的差距 => 性能奥秘
理解扩展的基准测试
用于理解 PyTorch 函数内部结构的分析(用内核触底)
看 PTX 汇编,了解 CUDA 内核的内部结构
编写函数的 5 种方法:manual、PyTorch、编译、CUDA、Triton
GeLU(按元素)、softmax(按行)、matmul(复聚合)
关键原则:组织计算以尽量减少读/写
关键思想:内核融合(仓库/工厂类比)、平铺(共享内存)
自动编译器(Triton、torch.compile)会随着时间的推移而变得更好
硬件
计算:流式多处理器 (SM) [A100:108]
内存:
DRAM [A100: 80GB] - 大而慢
二级缓存 [A100: 40MB]
L1 缓存 [A100:每 SM 192KB] - 小、快
执行模型
线程:进程单个索引(即 f(i))
线程块:(又名并发线程数组):在单个 SM 上调度
网格 :线块的集合
为什么使用线程块?共享内存。
直观理解:将读取相似数据的 f(i)分组在一起
线程块内的线程拥有共享内存(速度与 L1 缓存相当)[A100: 164KB]
可以在块内同步线程(用于读取/写入),但不能跨块同步
硬件与执行相互作用
线程块以波次形式调度到 SM 上。
问题:最后一波次线程块较少,导致部分 SM 空闲(低占用率)。
波形量化:使线程块数量能被 SM 数量整除。
经验法则:线程块的数量应该 >= 4 倍的 SM 数量
挑战:执行模型中隐藏了一些硬件方面的特性(例如,调度、SM 数量)。
算术强度:每字节的浮点运算次数
如果数值较高,表示操作是计算密集型(好)
如果数值较低,表示操作是内存密集型(坏)
一般规则:矩阵乘法是计算密集型,其他操作都是内存密集型
重要提示:请测试和分析您的代码!
您可以阅读规格表(营销材料)和论文
但性能取决于您的库版本、硬件和工作负载
因此,测试和分析您的代码是必不可少的。
示例计算:在多层感知机(MLP)上运行前向/后向传播。
run_mlp(dim=128, num_layers=16, batch_size=128, num_steps=5)
benchmarking() # 用时多久?profiling() # 时间花在哪里?
每次做修改时,请务必进行基准测试/性能分析!
矩阵乘法基准测试
首先,我们来基准测试方形矩阵的矩阵乘法。
在1024和2048上,时间几乎没有增加
因为在执行时这些矩阵乘法时存在恒定因子的开销,如从CPU传送到GPU、启动内核
让我们测试一下我们的多层感知机!
每次MLP执行大约需要五秒钟
性能分析
虽然基准测试关注的是端到端时间,但性能分析则关注时间的具体花费位置。
很明显:性能分析可以帮助你了解时间具体花费在哪些地方。
更深入:性能分析帮助你了解(被调用的内容)。
PyTorch 有一个内置的性能分析器 https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html
让我们对一些代码进行性能分析,看看幕后发生了什么。
CUDA
CUDA 是 C/C++ 的扩展,包含管理 GPU 的 API。
我们将编写一些函数f,当调用这个CUDA内核时,它将自动对向量或矩阵的所有元素调用f,然后我们将并行计算我们想要的一切
网格:
线程块的集合:numBlocks = (2, 4),blockDim = (1, 8)
线程块:
线程的集合:blockIdx = (0, 1)
线程:
单个操作单元:threadIdx = (0, 3)。
你编写线程执行的代码,并使用 (blockIdx, blockDim, threadIdx) 来确定要执行的操作。
设置 CUDA_LAUNCH_BLOCKING,以便如果有错误,CUDA 会告诉你哪里出了问题。
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
load_inline
函数使得编写 CUDA 代码并将其绑定到 Python 模块变得方便,以便立即使用。
CUDA 代码:包含完整的逻辑
实际内核所在,这将被发送到GPU,并进行计算,然后返回结果
这是一个包装器,将协调内核的启动,内核实际上存在于GPU中
TORCH_CHECK(x.device().is_cuda());//确保x存在于GPU设备中
TORCH_CHECK(x.is_contiguous());//确保x是连续的,位于连续的内存块中
// Allocate empty tensor
torch::Tensor y = torch::empty_like(x);
总结
-
编程模型(如PyTorch、Triton、PTX)与硬件的“差距”为何会导致性能差异?
不同编程模型对硬件细节的抽象程度不同:PyTorch封装了底层逻辑,易用但可能未充分利用硬件特性;Triton提供中间抽象,平衡灵活性与性能;PTX接近硬件指令,可控性最高但复杂度大。这种抽象差距导致对内存访问、线程调度等硬件细节的优化程度不同,最终体现为性能差异。 -
GPU执行模型中,线程、线程块、网格的关系及对硬件利用率的影响是什么?
线程是最小执行单元,负责单个索引操作;线程块是线程的集合,在单个SM上调度,共享块内共享内存;网格是线程块的集合。线程块数量需与SM数量匹配(建议≥4倍SM数),否则易因“最后一波次线程块不足”导致SM空闲,降低硬件利用率。 -
算术强度的定义及对GPU性能的指导意义是什么?
算术强度指每字节数据传输对应的浮点运算次数。数值高表示计算密集型(如矩阵乘法),能充分利用GPU算力;数值低表示内存密集型(如多数逐元素操作),易受内存带宽限制。优化需针对类型调整策略:计算密集型提升并行效率,内存密集型减少数据传输。 -
基准测试与性能分析的核心区别及各自作用?
基准测试关注端到端时间,用于评估整体性能(如矩阵乘法不同尺寸的耗时);性能分析聚焦时间分布,定位具体瓶颈(如哪部分内核、内存操作耗时最长)。二者结合:基准测试衡量优化效果,性能分析指导优化方向。 -
CUDA内核编写中,线程块设计需考虑哪些硬件特性?共享内存的关键作用是什么?
线程块设计需匹配SM资源:受限于共享内存大小(如A100每块164KB)、线程数上限,并需保证内存访问合并(匹配DRAM突发传输)。共享内存位于SM内部,速度接近L1缓存,核心作用是缓存块内重复访问的数据,减少全局内存读写,提升数据重用效率。