基於Ascend C的FlashAttention運算元效能最佳化最佳實踐

华为云开发者联盟發表於2024-06-12

本文分享自華為雲社群《基於Ascend C的FlashAttention運算元效能最佳化最佳實踐》，作者：昇騰CANN。

LLM的Attention部分處理給計算系統帶來巨大的計算和訪存壓力。業界先後出現FlashAttention、FlashAttention2等演算法，透過計算等價和切分有效降低HBM資料訪問量。

昇騰異構計算架構CANN針對昇騰AI處理器的片上記憶體和快取大小，以及資料搬運通路，基於Ascend C運算元程式語言最佳化實現FlashAttention融合運算元，充分利用片上快取，提升Attention處理效能。根據實測，在一些典型場景中CANN的FlashAttention運算元相比小運算元取得了5倍以上的效能提升，開發者可直接呼叫相關運算元API介面使能大模型極致效能最佳化。

本文針對FlashAttention反向融合運算元的效能最佳化方案展開介紹，並透過最佳化實現了典型場景4倍左右的效能提升，希望對開發者最佳化此類基於Ascend C開發的融合運算元帶來啟發。

FlashAttention演算法簡介

在主流大模型網路模型中，大量使用典型的Multi-Head Attention結構，帶來了巨大的計算和記憶體開銷。其執行過程中，矩陣乘和softmax結果存放在片上記憶體會帶來巨大的記憶體消耗，訪存效能嚴重下降，甚至會導致模型無法正常執行，同時網路中的矩陣和向量計算序列執行，也會導致硬體算力發揮受限。

史丹佛的Tri DAO提出了FlashAttention融合運算元，其原理是對attention處理過程進行切分和計算等價，使得attention的多個步驟在一個運算元中完成，並且透過多重迴圈、每次處理一小部分資料，以近似流式的方式訪問片上記憶體，減少了片上記憶體訪問的總資料量，並能夠將計算和資料搬運更好的重疊隱藏。

圖片1.png

對於self-attention來講，Q(Query), K(Key), V(Value)三個矩陣均來自同一輸入，首先我們要計算Q與K之間的點乘，然後為了防止其結果過大，會除以一個尺度標度圖片1.png ，其中捕獲4.PNG 為一個query和key向量的維度。再利用Softmax操作將其結果歸一化為機率分佈，然後再乘以矩陣V就得到權重求和的表示。該操作可以表示為：

圖片4.png

注意力的正向計算公式為：

圖片5.png

為方便表達，以變數S和P表示計算公式：

圖片6.png

圖片7.png

圖片8.png

注意力的反向計算公式為：

圖片9.png

圖片10.png

圖片11.png

圖片12.png

昇騰CANN基於Ascend C程式語言實現了FlashAttention正反向融合運算元，其中反向運算元計算流程可參考下圖所示：

圖片13.png

本案例對FlashAttention反向運算元進行了效能最佳化，主要涉及的最佳化手段包括tiling基本塊大小調整，核間負載均衡，CV流水並行，MTE2流水最佳化以及FixPipe流水最佳化等，並在Atlas A2訓練系列產品/Atlas 800I A2推理產品驗證平臺下收益4倍左右的效能提升。下面以如下兩個輸入場景為例，介紹整個最佳化過程。

第一個場景的輸入維度資訊為：B=1，N1=12，N2=12，S1=6144，S2=6144，D=128，並且為casual場景，casual場景即atten_mask的形狀為下三角。

圖片14.png

第二個場景的輸入維度資訊為：B=24，N1=5，N2=5，S1=9216，S2=9216，D=64，不帶atten_mask和drop_mask輸入。

tiling基本塊調整

根據以往最佳化的經驗，迴圈間可能存在一些不必要的頭開銷，迴圈越多效能可能越差；滿足UB最大空間限制的情況下，UB切分的基本塊越大，迴圈越少，運算元中透過InitBuffer介面分配UB buffer大小。

pipe->InitBuffer(ubBuffer, 120 * 1024);   
pipe->InitBuffer(tmpBuffer, 30 * 1024);   
pipe->InitBuffer(vecClc3, 8 * 1024);

如上程式碼所示，InitBuffer介面的第二個參數列示buffer佔用的大小，所有buffer大小的和即為佔用的總空間。這裡120 * 1024 + 30 * 1024 + 8 * 1024 = 158KB < UB Size，沒有充分利用UB空間。

接下來試圖透過調整tiling基本塊進行效能最佳化，在滿足UB空間大小夠用的情況下，tiling基本塊切分的越大越好。下圖為最佳化前按照(64, 128)切分計算，總共需要迴圈計算32次：

圖片15.png

考慮到UB空間沒有用滿，基本塊調整到(128, 128)，如下圖最佳化後只需迴圈計算16次，切分後運算元效能提升一倍：

圖片16.png

CV流水並行

從流水圖可以看到，可以看出兩側的流水都存在大段的空隙（圖中綠色為vector部分流水，橙色為cube側流水），CV之間流水很大程度上未並行，需要考慮CV流水最佳化。

圖片17.png

由於FAG運算元中cube計算比vector計算快且存在依賴性，同時為了減少CV之間的通訊次數，透過快取機制實現讓matmul提前計算多塊，這裡的快取機制指的是將mm一次性計算多個基本塊快取到GM上。如下程式碼中，SetTail設定的SingleM和SingleN大小為BaseM，BaseN的倍數，即matmul一次發起多個基本塊的計算，實現matmul結果的快取，vector側分多次取matmul的結果。

mm3.SetTail(s2CvExtend, -1, preS1Extend);   
mm3.SetTensorA(mulWorkSpaceGm[pingpongIdx * coreNum * cubeBaseMN + cBlockIdx * cubeBaseMN], true);  
mm3.SetTensorB(queryGm[mm2aTensorOffsetCv]);   
mm3.template IterateAll<false>(dkWorkSpaceGm[bTensorOffsetCv], true);

下圖是實現mm1、mm2和mm3快取的流水圖，綠色的vector流水與橙色的cube流水均變得更密集，並行度提高，cv的間隔減小，提升了運算元效能：

圖片18.png

基於快取mm1/mm2/mm3的最佳化後，在本輪Vector等Cube流水的間隔，插入下一輪迴圈的Vector計算，這樣使Vector流水與Cube流水之間的並行度更高，反映到流水圖中為Vector計算更密集：

圖片19.png

相關最佳化點實現虛擬碼如下所示：

mm1計算; 
dropout(); 
Sub(); 
dropout(); // 下一輪迴圈的Vector計算 
Sub();  // 下一輪迴圈的Vector計算 
mm2計算; 
Softmax(); 
AttenMask(); 
...

核間負載均衡

對於上述場景一，casual場景下可能存在核間分佈不均勻的情況，如下圖經過atten mask掩碼後，紅色部分是運算元需要計算的部分，綠色無需計算；如果不按照基本塊的個數來分核，按照第一根軸的大小8（行）來分核，假設平均分到9個核上，每個核做ceil(8 / 9) = 1行，則第一個核只需做1個基本塊，但是第8個核需要做8個基本塊的計算，出現嚴重的負載不均衡：

圖片20.png

因此需要考慮將紅色塊均勻分到多個核上計算，儘量實現每個核的計算量均勻，負載均衡。最佳化後，紅色塊總共36個基本塊，均分到每個核上，每個核的計算量為4塊，效能提升一倍。

圖片21.png

FixPipe流水最佳化

透過對場景一的Profilling資料進行分析可以看到，aic_fixpipe_ratio佔比極高，佔比高達81%，出現了很嚴重的bound：

圖片22.png

同時，CAModel工具列印發現存在很多異常的128B搬運，經過程式碼排查，發現workspace地址未512B對齊。程式碼實現中使用SetGlobalBuffer介面設定workspace的起始地址，如果起始地址不是按照512B對齊，搬運效率會很低，可以強制地址512B對齊來避免這個情況，下面程式碼中ADDR_ALIGN_SIZE即為512：

// init workspace address   
syncGlobal.SetGlobalBuffer((__gm__ int32_t*)workspace);   
uint64_t workspaceOffsets = SYNC_GLOBAL_WORKSPACE_SIZE;   
dqWorkSpaceGm.SetGlobalBuffer((__gm__ float*)workspace + workspaceOffsets / sizeof(T2));   
workspaceOffsets = (workspaceOffsets + qPostBlockTotal * sizeof(float) + ADDR_ALIGN_SIZE) / ADDR_ALIGN_SIZE * ADDR_ALIGN_SIZE;  dkWorkSpaceGm.SetGlobalBuffer((__gm__ float*)workspace + workspaceOffsets / sizeof(T2));   
workspaceOffsets = (workspaceOffsets + kvPostBlockTotal * sizeof(float) + ADDR_ALIGN_SIZE) / ADDR_ALIGN_SIZE * ADDR_ALIGN_SIZE;  dvWorkSpaceGm.SetGlobalBuffer((__gm__ float*)workspace + workspaceOffsets / sizeof(T2));   
workspaceOffsets = (workspaceOffsets + kvPostBlockTotal * sizeof(float) + ADDR_ALIGN_SIZE) / ADDR_ALIGN_SIZE * ADDR_ALIGN_SIZE;  
// matmul1 and matmul2 workspace size   
matmulWorkspaceSize = cubeBaseMN * sizeof(float);  
mm1WorkspaceGm.SetGlobalBuffer((__gm__ T2*)(workspace + workspaceOffsets + cBlockIdx * matmulWorkspaceSize));  mm2WorkspaceGm.SetGlobalBuffer((__gm__ T2*)(workspace + workspaceOffsets + coreNum * matmulWorkspaceSize + cBlockIdx * matmulWorkspaceSize));   // drop workspace offset   
workspaceOffsets = (workspaceOffsets + coreNum * cubeBaseMN * sizeof(float) * INPUT_NUMS + ADDR_ALIGN_SIZE) / ADDR_ALIGN_SIZE * ADDR_ALIGN_SIZE;   
dropWorkSpaceGm.SetGlobalBuffer((__gm__ T1*)workspace + workspaceOffsets / sizeof(T1));    
// mul workspace offset   
workspaceOffsets = (workspaceOffsets + coreNum * cubeBaseMN * sizeof(half) * 2 + ADDR_ALIGN_SIZE) / ADDR_ALIGN_SIZE * ADDR_ALIGN_SIZE;   
mulWorkSpaceGm.SetGlobalBuffer((__gm__ T1*)workspace + workspaceOffsets / sizeof(T1));

修改程式碼，workspace地址經過512B對齊後，fixpipe時間減半：

圖片23.png

MTE2流水最佳化

從場景二採集的profiling和打點圖來看，mte2_ratio佔比高，cube MTE2出現了明顯bound，且部分MTE2搬運時間異常。

圖片24.png

圖片25.png

將輸入資料排布格式從BSH更改為BNSD後，資料搬運連續，不需要跳地址讀取資料，搬運效率提升一倍，部分異常搬運時長降低了一半。

最佳化方案效能收益

調整tiling基本塊：理論評估vector切塊越大，計算和搬運迴圈次數越少，同時能夠充分利用搬運頻寬和vector算力。基本塊大小從(64, 128)增大到(128, 128)後，效能提升一倍，實測與理論分析一致。
CV流水並行：CV流水掩蓋的時間即為提升的效能，符合預期的收益。
核間負載均衡：最佳化前負載最多的核的計算量減少的倍數，即為預期提升的效能；案例中最佳化前負載最多的核的計算量大小為8塊，最佳化後為4塊，實際效能提升一倍，符合預期的收益。
FixPipe最佳化：從Profiling資料看出FixPipe佔比8，最佳化後佔比0.55，實測運算元效能提升45%，與理論分析一致。
MTE2最佳化：從Profiling資料看出MTE2佔比52，最佳化後佔比減少一半，實測運算元效能提升30%，與理論分析一致。

開發者在對基於Ascend C開發的融合運算元進行效能最佳化時，可參考此案例中的最佳化思路。