新闻宏观经济使用TileLang设计高性能GPU内核:Tensor-Core GEMM、融合Softmax、FlashAttention与自动调优

使用TileLang设计高性能GPU内核:Tensor-Core GEMM、融合Softmax、FlashAttention与自动调优

作者: MarkTechPost·

要点速览

  • TileLang使开发者能够在Python中以tile层级编写优化的GPU内核,而编译器自动处理线程映射、内存布局、同步和CUDA指令生成。
  • 教程逐步实现并验证了向量加法、Tensor-Core矩阵乘法、带有偏置和GELU的融合GEMM尾操作、逐行softmax和FlashAttention等内核,并与PyTorch基线进行对比。
  • TileLang的自动调优装饰器在tile大小、流水线深度和线程数范围内搜索,自动识别最优的架构相关配置,解决了理想参数在不同GPU代际(如Ampere和Hopper)之间变化的挑战。
  • 教程中展示的内核融合技术通过在写入最终输出之前在寄存器驻留的累加器中完成偏置加法和GELU激活等操作,减少了中间全局内存流量。
  • FlashAttention实现处理查询、键和值tile,无需实例化完整的注意力分数矩阵,通过在线softmax更新将内存占用从二次空间降低到线性空间。
  • 生成代码检查、设备端打印和内置分析器共同为开发者提供了对编译器生成的Tensor-Core操作、异步拷贝和同步屏障的可视化能力,用于调试和优化。
使用TileLang设计高性能GPU内核:Tensor-Core GEMM、融合Softmax、FlashAttention与自动调优

TileLang是一种高级Python领域特定语言(DSL),用于通过TVM设计和编译面向性能的GPU内核。随着大语言模型和其他Transformer架构对计算需求不断推高,编写能够充分利用Tensor Core——NVIDIA专用的矩阵乘法加速器——的自定义GPU内核,对于训练和推理效率变得愈发关键。然而,编写此类内核传统上需要深厚的CUDA C++、warp级编程和硬件特定内存层次结构的专业知识。TileLang弥补了这一空白,让开发者在tile层级表达计算,而由编译器处理底层细节。本教程全面介绍了TileLang的各项功能,从环境验证和建立可复用的基准测试与数值验证工具开始,逐步实现向量加法、分块Tensor-Core矩阵乘法、调度探索、融合GEMM尾操作、逐行softmax和FlashAttention。

在整个教程中,开发者直接使用TileLang的共享内存tile、寄存器片段、流水线循环、并行迭代原语、归约操作和Tensor-Core GEMM运算符。编译器自动处理线程映射、内存布局、同步、向量化和底层CUDA指令生成。内核与PyTorch和cuBLAS基线进行基准对比,并检查生成的CUDA源代码、评估内存和计算吞吐量,以及使用自动调优来确定依赖于架构的内核配置。TileLang可在GitHub上获取。

教程首先配置Google Colab CUDA环境,安装TileLang(带有nightly回退机制),并导入所需的PyTorch和TileLang模块。定义了可复用的基准测试、验证和报告工具,用于测量内核延迟并通过相对误差比较数值输出。随后实现一个TileLang向量加法内核,在GPU上执行,与PyTorch的带宽进行对比,并检查编译器生成的CUDA源代码。

接下来,实现一个分块Tensor-Core矩阵乘法内核,将输入tile从全局内存经过共享内存传递到寄存器片段。手动控制tile维度、流水线阶段、线程数和L2 swizzling,同时由TileLang生成Tensor-Core指令、同步和内存传输逻辑。对多种调度配置进行基准测试,验证其数值精度,并确定性能最优的架构相关内核配置。完整的教程代码可在此处获取。

矩阵乘法内核随后被扩展,通过将偏置加法和GELU激活直接融合到寄存器驻留的累加器中。这种方法通过在写入最终输出张量之前完成尾操作来减少中间全局内存流量。像这样的内核融合是一种减少内存带宽瓶颈的成熟技术,而内存带宽瓶颈往往主导大规模神经网络工作负载的延迟。融合实现与PyTorch的即时执行模式进行了对比。此外,还使用片段级最大值和求和归约实现了一个逐行softmax内核,将归一化过程主要保留在寄存器中。

实现了一个融合FlashAttention前向内核,用于处理查询、键和值tile,而无需在全局内存中实例化完整的注意力分数矩阵。使用运行最大值、归一化求和、重缩放因子和分块Tensor-Core矩阵乘法来应用在线softmax更新。FlashAttention最初由斯坦福大学的研究人员提出,现已成为一种广泛采用的技术,可将自注意力的内存占用从二次空间降低到线性空间,目前已集成到包括PyTorch缩放点积注意力API在内的主流框架中。因果和非因果注意力均与PyTorch缩放点积注意力进行了验证,并比较了延迟和计算吞吐量。

定义了一个跨矩阵tile大小、K块维度、流水线深度和线程数的自动调优搜索空间,并过滤掉超出共享内存预算的配置。使用TileLang的自动调优装饰器为相同的矩阵乘法工作负载编译、基准测试、验证和缓存多个内核调度。执行选定的内核,与PyTorch进行输出验证,并报告实现的延迟和Tensor-Core吞吐量。这种自动搜索解决了GPU内核开发中的一个实际挑战:最优tile大小和流水线深度在不同的GPU架构(如Ampere和Hopper)之间存在差异,使得在面向多个硬件代际时手动调优变得脆弱。

通过设备端打印、生成CUDA代码检查和内置内核分析器引入了TileLang的调试和内省工作流程。检查了编译器发出的标记,包括Tensor-Core操作、异步拷贝、同步屏障和矩阵加载指令。所有教程章节被组织到一个容错运行器中,记录执行状态、报告计时信息,并打印一份简洁的TileLang编程参考。

该教程展示了TileLang如何将tile级Python程序转换为优化的GPU内核,而无需手动管理线程索引、warp级数据布局、Tensor-Core指令或异步内存屏障。实现和验证的内核涵盖带宽受限的逐元素操作、计算密集型GEMM工作负载、融合神经网络尾操作、寄存器驻留归约和在线softmax注意力。探索还展示了块维度、共享内存消耗、流水线深度、线程数、tile形状和L2 swizzling如何在不同GPU架构上影响性能。生成代码检查、设备端调试、性能分析和自动调度搜索共同建立了一个完整的开发、验证、基准测试和优化自定义TileLang内核的工作流程。