在CUDA的天下，OpenAI开源GPU编程语言Triton，同时支持N卡和A卡( 二 )

文章图片
CUDAvsTriton编译器优化对比。
编程模型
在所有可用的领域专用语言和JIT编译器中， Triton或许与Numba最相似：内核被定义为修饰过的Python函数，并与实例网格上不同的program_id的同时启动。但不同之处值得注意：如下图代码片段所示， Triton通过对block的操作来展示intra-instance并行，此处block是维数为2的幂的数组，而不是单指令多线程（SIMT）执行模型。如此一来， Triton高效地抽象出了与CUDA线程block内的并发相关的所有问题（比如内存合并、共享内存同步/冲突、张量核心调度）。

文章图片
Triton中的向量加法。
虽然这对embarrassingly并行（即element-wise）计算可能没什么帮助，但是可以简化更复杂的GPU程序的开发。例如，在融合softmax核的情况下，对于每个输入张量X∈R^M×N来说，每个实例对给定输入张量的不同行进行归一化。这种并行化策略的标准CUDA实现可能难以编写，需要线程之间的显式同步，因为这种策略并发地减少X的同一行。而Triton很大程度上消除了这种复杂性，每个内核实例加载感兴趣的行，并使用类似NumPy的原语顺序对其进行规范化。
importtritonimporttriton.languageastl@triton.jitdefsoftmax(Y,stride_ym,stride_yn,X,stride_xm,stride_xn,M,N):#rowindexm=tl.program_id(0)#colindices#thisspecifickernelonlyworksformatricesthat#havelessthanBLOCK_SIZEcolumnsBLOCK_SIZE=1024n=tl.arange(0,BLOCK_SIZE)#thememoryaddressofalltheelements#thatwewanttoloadcanbecomputedasfollowsX=X+m*stride_xm+n*stride_xn#loadinputdata;padout-of-boundselementswith0x=tl.load(X,mask=n<N,other=-float('inf'))#computenumerically-stablesoftmaxz=x-tl.max(x,axis=0)num=tl.exp(z)denom=tl.sum(num,axis=0)y=num/denom#writebacktoYY=Y+m*stride_ym+n*stride_yntl.store(Y,y,mask=n<N)importtorch#Allocateinput/outputtensorsX=torch.normal(0,1,size=(583,931),device='cuda')Y=torch.empty_like(X)#SPMDlaunchgridgrid=(X.shape[0],)#enqueueGPUkernelsoftmax[grid](Y,Y.stride(0),Y.stride(1),X,X.stride(0),X.stride(1),X.shape[0],X.shape[1])在Triton中融合softmax
TritonJIT把X、Y当作指针而不是张量。最重要的是， softmax这种特殊实现方式在整个规范化过程中保持SRAM中X的行不变，从而在适用时最大限度地实现数据重用（约32K列）。这与PyTorch的内部CUDA代码不同，后者使用临时内存使其更通用，但速度明显变慢（见下图）。

文章图片
融合softmax、M=4096的A100性能。
Torch(v1.9)JIT较低的性能突出了从高级张量操作序列自动生成CUDA代码的难度。
@torch.jit.scriptdefsoftmax(x):x_max=x.max(dim=1)[0]z=x-x_max[:,None]numerator=torch.exp(x)denominator=numerator.sum(dim=1)returnnumerator/denominator[:,None]融合softmax与TorchJIT
矩阵乘法
能够为元素操作（element-wiseoperation）和规约操作（reductionoperation）编写融合内核是很重要的，但考虑到神经网络中矩阵乘法的重要性，这还不够。事实证明， Triton在这些方面表现很好，仅用大约25行Python代码就能达到最佳性能。相比之下， CUDA效率就没有那么高了。

文章图片

文章图片
Triton中的矩阵乘法。
手写矩阵乘法内核的一个重要优点是它们可以根据需要进行定制，以适应其输入（例如切片）和输出（例如LeakyReLU）的融合变换。假如不存在Triton这样的系统，那么对于没有出色的GPU编程专业知识的开发人员来说，矩阵乘法内核将很难大改。