很久之前笔者碰到过一个相对棘手的问题: 怎么优雅的在运行时阶段自动把一个tvm生成的CUDA C文件封装成一个Torch Function? 作者大概尝试过三种方法,哥各有利弊:
本文接下来再详细分享一下这三种方法的利弊与笔者做出的一些尝试与改进,包括怎么尽可能的减少torch cpp extension的编译overhead,怎么分析dlpack的overhead等,希望能够帮助到之后踩坑的同学以及供大家讨论 :) 代码放在CPPTorchExecutable
在Tile Language目前的设计理念中,我们提供了三种不同的编程接口,分别面向初学者、开发者和写Kernel的专家这三个等级,我们对Level 1的使用者的定义是不清楚目标硬件架构的开发者(例如不知道GPU的cache一般需要用到tiling,把数据层层cache),这个时候用户编写的程序只是一个简单的计算表达式,不包含任何schedule相关的信息(和tvm的设计一样),这种方法理论上来讲对于不太kernel的人是最舒服的,但是因为很多现存未解决的问题:
1. 从调优的角度来看,从计算到具体的schedule的自动调优时间一般要花比较久的时间,虽然一些工作例如MSRA System Research的Roller,利用硬件感知的白盒搜索空间生成缓解了这一问题, 请看BitBLAS对Roller的复现[Tutorial](https://github.com/microsoft/BitBLAS/blob/main/tutorials/1.fast_and_efficient_codegen.ipynb), 但知道并使用这种方法的人的人也不是很多
1. 从Kernel的表达能力角度来看,目前一些主流的复杂的计算并不能够通过schedule描述出来,例如Flash Attention, 虽然理想上应该可以通过推导 matmul+softmax+matmul 简单计算表达式的一种算子融合版本得到,但是这个问题从写文章的时间点来看仍然很困难。
1. 从社区来看,Schedule的代码虽然看起来还是很优雅的(从写schedule和理解schedule的角度来看, 例如bitblas对于matmul dequantize的[schedule模版](https://github.com/microsoft/BitBLAS/blob/main/bitblas/gpu/matmul_mma_dequantize.py)我个人觉得还是写的很通用的),schedule的魅力在于其从一个最原始不包含任何调度信息的计算表达式,在确保正确性不受影响的情况下一步步变换到最终的高性能调度形式。但是schedule实在太难学习和理解了,即使是一个会写schedule表达式的开发者(这部分玩家已经很少了),想要看明白我写的各种schedule模版,继续扩展也是非常非常困难的。其次,很复杂的计算,例如Flash Attention, 其因为设计本身是要在shared memory上进行算子融合所以计算是无法表达的,其次,即使是要强行写出特别针对Flash attention的多算子fuse版本的schedule模版,schedule本身的代码量可能会比cuda还要长(最后,同样受限于社区,tvm的生态逐渐变得不如triton,一些新feature例如tma这些的支持会有点滞后)
于是在搞bitblas的时候我就觉得这一套有点难受(, 于是觉得需要一个类似triton的东西,但是triton的限制也很大,例如不能显式声明内存,不能显式控制线程的行为等,这一点之后分享的时候再讨论讨论。总之目前bitblas的所有kernel实现都已经换成了Tile Lang,自己用起来非常舒服。
话说回来,我们对Level 2的使用者的定义是知道目标硬件架构的开发者,及在cuda上知道shared memory这个概念,知道在cuda上做tile要先把数据load到一个高速的缓存上再进行计算(有点类似triton的开发模式),本文我们以这种模式为例子介绍一下矩阵乘法Kernel的组成。
最后提一嘴Level 3, thread primitives允许用户完全控制每个线程的行为,写这一部分的代码其实就和写PyCUDA差不多了(但是支持多后端,例如HIP等),于是Level 3就是给完全的expert写的了,但本质上,经过LowerTileOP
和LayoutInference
这两个Pass之后,Level 2的代码也会被Lower到Level 3上。
所以非常值得一提的是,我们的设计中这三种语法可以出现在同一个program中,例如BitBLAS里的量化矩阵乘法的program中对于复杂的反量化部分,我们使用了thread primitives来精心控制每个线程的行为,以及利用ptx来加速精度转换的过程,显式在progam中调用mma 来在一些情况下在寄存器中做反量化等,其他的一些操作,例如数据拷贝和Pipeline仍然使用Level 2的编程方式(T.Pipelined, T.Copy)等,代码参考matmul_dequantize_tensorcore_finegrained.py。
翻译自: https://tilelang.tile-ai.cn/tutorials/debug_tools_for_tilelang.html
一个Tile Language程序(我们称为 program)到具体的硬件可执行文件的流程如下图所示,大致分为以下几步:1. 用户首先编写 Tile Language program。2. 程序会经过多个 Pass 的转换和优化处理(即 lower 阶段,相关代码位于 tilelang/engine/lower.py),最终生成中间代码,比如针对 CPU 的 LLVM 或 C 代码,或者针对 NVIDIA GPU 的 CUDA 代码等。3. 生成的中间代码会通过对应的编译器进一步编译,最终输出硬件可执行文件。
在这个过程中,用户可能会碰到大概三类问题:
1. Tile Language Program无法生成硬件可执行文件,也就是lower的过程中出现问题,我们可以归纳成生成问题。
1. 正确性问题,生成的可执行文件运行后,行为不符合预期。
1. 性能问题,执行文件的性能表现与硬件的理论值存在显著差距。
本文将重点讨论前两类问题的调试方法。至于性能问题的调优,则需要结合硬件厂商提供的性能分析工具(如 Nsight Compute、rocProf 等),通过分析具体的硬件指标进一步优化,我们将在后续文章中详细探讨。
接下来,我们以矩阵乘法(Matrix Multiplication)为例,使用 Tile Language 展示如何编写和调试相关程序。
不久之前的一篇分享里,我介绍了AMD CDNA架构(MI210, MI250, MI300)上的异步拷贝相关指令,在BitBLAS可以找到相关的实现,然而实际过程中发现AMD的异步拷贝指令的要求实际上要比那篇分享所写的更加苛刻,每个warp里的线程必须要求访问连续的数据,或者通过M0寄存器来控制每个线程的偏置。
一般来说,我们习惯这个指令就是明确的要load给定指针的一小块数据就行了,但是这个指令因为上述提到的两个限制就很难做到。经过笔者非常繁琐的Micro bencmark之后,笔者终于调教出了可以让每个线程Load给定数据块的写法,如下:
1 |
|
在这篇文章里,笔者填一下AMD Matrix Core的坑,介绍一下过去一个月里BitBLAS针对AMD的的高性能Matrix Core支持,在这篇文章里笔者将介绍一下MFMA(AMD版的MMA)。如何进行AMD Kernel的性能分析,及Profile一个AMD Kernel,最后我们介绍若干种绞尽了笔者脑汁的优化方法,完全利用好硬件的带宽(全都是128bits的内存访问指令,并且没有Memory bank conflict)。
这篇文章涉及到的算子有矩阵乘法和Flash Attention。本篇文章的实现在BitBLAS里, Codegen以及Swizzle等Layout变换依托于TVM, TVM可以帮助我们显式地操作一个数据的Layout,相比Triton更加灵活和可观。虽然AMD提供的文档十分有限,但是在这一个月里笔者参考了很多AMD开发人员提供的实现,例如Composable Kernel和Triton for ROCm,笔者从这些项目中收获良多。
本文假设读者对Nvidia GPU的编程有一定的了解,熟悉最基本的Tile优化程序的方法,以及Tensor Core的基本概念。
最近给BitBLAS添加了AMD的后端,发现AMD的异步拷贝等和Nvidia有很大的不同(但是FA3在MI300上需要用到这一个Feature),然而官方根本没有文档,只有Instruction Set,我在这里做一下自己的理解和解读,大部分内容是参考自这个Instruction Set。
Update your browser to view this website correctly. Update my browser now