前言

通過前兩章對于triton的簡單介紹，相信大家已經能夠通過從源碼來安裝triton，同時通過triton提供的language前端寫出自己想要的一些計算密集型算子。這章開始，我們通過構建一套比較標準的batch gemm的benchmark，來看看目前這些主流的代碼生成工具，高性能模板庫，與廠商提供的vendor library的差距。因為只有明確了目前的差距，后期關于針對性的優化才能做到點上。這一章，我將使用一個batch的gemm作為例子，來看看triton目前對其的優化能力。選batch gemm的原因是因為目前的LLM中不可避免會有對應的attention操作，而attention操作中，核心的計算密集型算子就是batch的gemm，如果你能夠對batch的gemm有一個很好的優化思路，那么在MLSys中大部分的算子優化類的工作對你來說將不會顯得那么無從下手。

通過Triton實現一個batch GEMM算子

在triton的官方tutorial中給出了如何使用triton的language api來實現gemm的算子，在上一章的最后，我也給出了對應的例子以及他通過和調用torch.matmul實現的gemm在3090上的性能比較。最終可以發現，針對某些size的gemm，triton在TFLOPS這個指標層面是能夠超過cublas的實現，但是后面我通過nsight system對每個kernel的具體執行時間進行了profiling，發現在torch.matmul或者torch.bmm底層所調用的cuBLAS的kernel并不是對應輸入輸出datatype以及computetype中最快的那個。所以，這樣的比較就顯得有些沒有意義。不過，沒事，這對我們建立起如何優化一個計算密集型算子來說是一個不錯的入門。

其實想要通過triton實現一個batch的gemm非常簡單，我們只需要將triton中原先例子里的tl.program_id(axis=0)，在這個program_id上再添加一個axis來表示batch維度的并行就可以了，然后針對每個數組的變化由單batch到多batch，只用增加一個大小為矩陣size的stride偏置即可，這種實現方式其實也是cuBLAS中cublasGemmStridedBatched命名的得來。具體的代碼如下所示:

@triton.jit
defmatmul_kernel(
#Pointerstomatrices
A_ptr,B_ptr,C_ptr,
#Matrixdimensions
B,M,N,K,
#Thestridevariablesrepresenthowmuchtoincreasetheptrbywhenmovingby1
#elementinaparticulardimension.E.g.stride_amishowmuchtoincreasea_ptr
#bytogettheelementonerowdown(AhasMrows)
stride_ab,stride_am,stride_ak,
stride_bb,stride_bk,stride_bn,
stride_cb,stride_cm,stride_cn,
#Meta-parameters
BLOCK_SIZE_M:tl.constexpr,BLOCK_SIZE_N:tl.constexpr,BLOCK_SIZE_K:tl.constexpr,
GROUP_SIZE_M:tl.constexpr,
ACTIVATION:tl.constexpr,
):
pid=tl.program_id(axis=0)
offs_b=tl.program_id(axis=1)
num_pid_m=tl.cdiv(M,BLOCK_SIZE_M)
num_pid_n=tl.cdiv(N,BLOCK_SIZE_N)
num_pid_k=tl.cdiv(K,BLOCK_SIZE_K)
num_pid_in_group=GROUP_SIZE_M*num_pid_n
group_id=pid//num_pid_in_group
first_pid_m=group_id*GROUP_SIZE_M
group_size_m=min(num_pid_m-first_pid_m,GROUP_SIZE_M)
pid_m=first_pid_m+(pid%group_size_m)
pid_n=(pid%num_pid_in_group)//group_size_m

offs_m=pid_m*BLOCK_SIZE_M+tl.arange(0,BLOCK_SIZE_M)
offs_n=pid_n*BLOCK_SIZE_N+tl.arange(0,BLOCK_SIZE_N)
offs_k=tl.arange(0,BLOCK_SIZE_K)

A_ptr=A_ptr+(offs_b*stride_ab+offs_m[:,None]*stride_am+offs_k[None,:]*stride_ak)
B_ptr=B_ptr+(offs_b*stride_bb+offs_k[:,None]*stride_bk+offs_n[None,:]*stride_bn)

#initializeanditerativelyupdateaccumulator
acc=tl.zeros((BLOCK_SIZE_M,BLOCK_SIZE_N),dtype=tl.float32)
forkinrange(0,K,BLOCK_SIZE_K):

a=tl.load(A_ptr)
b=tl.load(B_ptr)

acc+=tl.dot(a,b)

A_ptr+=BLOCK_SIZE_K*stride_ak
B_ptr+=BLOCK_SIZE_K*stride_bk

c=acc.to(tl.float16)
C_ptr=C_ptr+(offs_b*stride_cb+offs_m[:,None]*stride_cm+offs_n[None,:]*stride_cn)
c_mask=(offs_b

	

	然后寫一個簡單的單元測試，確保通過triton寫出來的kernel能夠和torch.matmul/torch.bmm對上即可。

	
torch.manual_seed(0)
a=torch.randn((4,512,512),device='cuda',dtype=torch.float16)
b=torch.randn((4,512,512),device='cuda',dtype=torch.float16)
torch_output=torch.bmm(a,b)
triton_output=matmul(a,b,activation=None)
print(f"triton_output={triton_output}")
print(f"torch_output={torch_output}")
iftorch.allclose(triton_output,torch_output,atol=1e-2,rtol=0):
print("TritonandTorchmatch")
else:
print("TritonandTorchdiffer")


	

	其實triton的language語法確實很簡單，相比較cuda來說，它能夠幫我們快速驗證一些idea，同時給出比cublas性能相當的算子。如果你想要用CUDA從0開始實現一個batch GEMM并且調用tensor core，借助shared memory，register files去幫你加速運算或者優化data movement，那么這個過程是非常需要一定的高性能計算和架構的經驗，你才可能拿到和cuBLAS的kernel接近的性能。OK，有了triton的具體kernel實現，接下來其實就是要去寫一個triton需要被調優的模版，需要triton從你定義的這個比較小的搜索空間中，去得到對應的最優解，從而作為本次batch gemm的最優實現，我在autotuner這塊并沒有花太大的精力去改進，依舊GEMM例子中的模版拿來作為一個參考，具體代碼如下:

	
@triton.autotune(
configs=[
triton.Config({'BLOCK_SIZE_M':128,'BLOCK_SIZE_N':256,'BLOCK_SIZE_K':64,'GROUP_SIZE_M':8},num_stages=3,num_warps=8),
triton.Config({'BLOCK_SIZE_M':64,'BLOCK_SIZE_N':256,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=4,num_warps=4),
triton.Config({'BLOCK_SIZE_M':128,'BLOCK_SIZE_N':128,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=4,num_warps=4),
triton.Config({'BLOCK_SIZE_M':128,'BLOCK_SIZE_N':64,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=4,num_warps=4),
triton.Config({'BLOCK_SIZE_M':64,'BLOCK_SIZE_N':128,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=4,num_warps=4),
triton.Config({'BLOCK_SIZE_M':128,'BLOCK_SIZE_N':32,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=4,num_warps=4),
triton.Config({'BLOCK_SIZE_M':64,'BLOCK_SIZE_N':32,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=5,num_warps=2),
triton.Config({'BLOCK_SIZE_M':32,'BLOCK_SIZE_N':64,'BLOCK_SIZE_K':32,'GROUP_SIZE_M':8},num_stages=5,num_warps=2),
],
key=['M','N','K'],
)


	

	然后通過調用Triton的do_bench就可以將你寫的算子跑起來了，do_bench處在python/triton/testing.py下，其中會對每個kernel進行25次的warm_up和100次iteration，最后會根據你設置的分位數得到一個相對穩定的性能。切記，在測試每個kernel的運行情況的時候，需要將GPU的頻率鎖在最高頻，通過下面的代碼就可以做到，由于我用到的A10，A10最大頻率在1695 MHz

	
sudonvidia-smi--lock-gpu-clocks=1695,1695


	

	這是通過對fp16的輸入，acc_type = fp32，最終輸出為fp16的batch gemm (16x4096x4096, 16x4096x4096)

	通過nsight system + nvtx就可以看到每個kernel的具體實現情況:

	img

	添加圖片注釋，不超過 140 字（可選）

	使用torch.bmm/torch.matmul來實現batch-gemm，其中調用的kernel名字為ampere_fp16_s1688gemm_fp16_256x64_Idg8_f2f_stages_32x1_nn，該kernel運行的時間是46.059ms

	那么，當我們運行triton的時候，通過同樣的方式來得到同樣迭代次序的kernel，nsight分析如下

	img

	該kernel的名字為matmul_kernel_0d1d2d3d4d5d6d7d8d9c10d11d12c13d14d15c，運行時間為35.067ms

	當然通過torch.matmul調用的cuBLAS這個算子，顯然不是我們想要的那個，我們就需要去深入到cuBLAS的具體文檔，翻一翻，找出其最快的API。在后面的benchmark中，我選用了cublasHgemmStridedBatched和cublasGemmStrideBatchedEx這兩個API來分別實現batch GEMM。通過cublasHgemmStridedBatched啟動kernel名字為ampere_h16816gemm_256x128_Idg8_stages_32x3_nn，其運行時間為30.330ms

	img

	通過cuBLAS的cublasGemmStridedBatchedEx API構建算子性能標準

	在cuBLAS中，針對batch gemm的實現有很多種方式，我也踩了不少坑。第一次調用成了cublasHgemmStridedBatched，該kernel的性能其實是不如cublasGemmStridedBatchedEx，因為cublasGemmStridedBatchedEx給了一個cublasGemmAlgo_t algo的參數，該參數可以幫我們選擇對應batch gemm的不同實現，關于algo又具有如下這么多種:

	
CUBLAS_GEMM_DEFAULT,
CUBLAS_GEMM_ALGO0,
CUBLAS_GEMM_ALGO1,
CUBLAS_GEMM_ALGO2,
CUBLAS_GEMM_ALGO3,
CUBLAS_GEMM_ALGO4,
CUBLAS_GEMM_ALGO5,
CUBLAS_GEMM_ALGO6,
CUBLAS_GEMM_ALGO7,
CUBLAS_GEMM_ALGO8,
CUBLAS_GEMM_ALGO9,
CUBLAS_GEMM_ALGO10,
CUBLAS_GEMM_ALGO11,
CUBLAS_GEMM_ALGO12,
CUBLAS_GEMM_ALGO13,
CUBLAS_GEMM_ALGO14,
CUBLAS_GEMM_ALGO15,
CUBLAS_GEMM_ALGO16,
CUBLAS_GEMM_ALGO17,
CUBLAS_GEMM_DFALT_TENSOR_OP,
CUBLAS_GEMM_ALGO0_TENSOR_OP,
CUBLAS_GEMM_ALGO1_TENSOR_OP,
CUBLAS_GEMM_ALGO2_TENSOR_OP,
CUBLAS_GEMM_ALGO3_TENSOR_OP,
CUBLAS_GEMM_ALGO4_TENSOR_OP,
CUBLAS_GEMM_ALGO18,
CUBLAS_GEMM_ALGO19,
CUBLAS_GEMM_ALGO20,
CUBLAS_GEMM_ALGO21,
CUBLAS_GEMM_ALGO22,
CUBLAS_GEMM_ALGO23,
CUBLAS_GEMM_ALGO5_TENSOR_OP,
CUBLAS_GEMM_ALGO6_TENSOR_OP,
CUBLAS_GEMM_ALGO7_TENSOR_OP,
CUBLAS_GEMM_ALGO8_TENSOR_OP,
CUBLAS_GEMM_ALGO9_TENSOR_OP,
CUBLAS_GEMM_ALGO10_TENSOR_OP,
CUBLAS_GEMM_ALGO11_TENSOR_OP,
CUBLAS_GEMM_ALGO12_TENSOR_OP,
CUBLAS_GEMM_ALGO13_TENSOR_OP,
CUBLAS_GEMM_ALGO14_TENSOR_OP,
CUBLAS_GEMM_ALGO15_TENSOR_OP,


	

	其中，帶有_TENSOR_OP后綴的則為調用tensor core來加速運算的。看到這么多種實現，不要慌，通過一個for-loop的遍歷，就可以方便的找到速度最快的那一個，然后對應就可以得到TFLOPS，對應實現如下:

	
floatmin_time=0xffff;
cublasGemmAlgo_talgo_index;
for(constauto&algo:algoList){
floattotal_time=0.0;
for(inti=0;i(algo));
cudaEventRecord(end,0);
cudaEventSynchronize(end);
floatelapsed_time;
cudaEventElapsedTime(&elapsed_time,start,end);
total_time+=elapsed_time;
}
floatcurrent_time=total_time/iteration;
std::cout<

	

	通過CUTLASS實現batch GEMM算子

	CUTLASS這里就不花過多的篇幅進行介紹了，知乎上有很多比較詳細的文章，建議做GPU性能優化的同學都能夠好好研究下CUTLASS，不得不說，CUTLASS的抽象層級做的確實很好，通過暴露出對應的C++模版，就可以通過這些模版組合成很多工程開發實際中可以跑的很快的算子，而且相比于直接寫CUDA嵌入PTX的匯編來說，開發的難易程度也被很大程度的降低，同時能帶來和cuBLAS肩比肩的效果。在本次benchmark的構建中，我使用的是2.9.1版本的CUTLASS，在編譯的時候一定要打開所有的kernel，然后通過下面的命令進行配置:

	
1.gitclonehttps://github.com/NVIDIA/cutlass.git
2.gitcheckoutv2.9.1
3.exportCUDACXX=/usr/local/cuda/bin/nvcc
4.mkdirbuild&&cdbuild
5.cmake..-DCUTLASS_NVCC_ARCHS=80-DCUTLASS_LIBRARY_KERNELS=all
6.makecutlass_profiler-j16


	

	然后我們可以通過使用cutlass_profiler來找到目前CUTLASS中針對應尺寸算子的TFLOPS最優的那個實現。這里直接使用如下代碼就可以得到CUTLASS對應的實現，同時只要在對應的workload添加不同尺寸的GEMM。

	Triton, CUTLASS, cuBLAS性能對比

	通過上述的講解，我們將所有的輸入和計算過程與cublasGemmStridedBatchedEx中的參數對齊，輸入為fp16，輸出為fp16，Accumulator_type也改為fp16。在triton中需要將如下代碼進行替換:

	
#acc=tl.zeros((BLOCK_SIZE_M,BLOCK_SIZE_N),dtype=tl.float32)
acc=tl.zeros((BLOCK_SIZE_M,BLOCK_SIZE_N),dtype=tl.float16)

#acc+=tl.dot(a,b)
acc+=tl.dot(a,b,out_dtype=tl.float16)


	

	然后把他們全部畫出來，縱坐標表示的TFLOPS，橫坐標對應矩陣的shape，batch=16。我們可以看出來，目前我這個版本的tirton代碼其實性能并不是很好，原因有很多，這個后面我給大家慢慢分析，最重要的其實就是triton.autotune中那些參數的選取和設定，以及后端的一些優化。cublasGemmStridedBatchedEx中最快的那個algo可以看出來目前基本上占據了領先位置，也就是為什么會被稱為目前GPU上去做計算密集型算子優化的上屆，CUTLASS在某些尺寸上的batch gemm還是表現的很優秀的，但是距離最快的cublasGemmStridedBatchedEx仍然有一些差距，不過只能說CUTLASS的優化真的牛逼，至少我知道目前國內很多HPC的組在開發對應的kernel的時候，都是選擇直接魔改拼接CUTLASS的組件來加快整個開發流程。

	img

	總結

	通過上述對batch gemm性能的分析，我們可以看出來triton距離cuBLAS的性能還有一定的距離要走，在后續的教程中，我們將結合Triton Dialect, TritonGPU Dialect, 以及Triton中autotuner作為核心組件來對Triton的所有優化過程中有一個清晰的認識。以及通過編譯手段，一步一步來逼近cuBLAS的性能，甚至超越他。

	

	   審核編輯：彭靜