（2）V2 BWD

一個(gè)建議：如果你在閱讀本節(jié)中覺得很困惑，一定記得先去看V1的BWD部分，有非常詳細(xì)的推導(dǎo)介紹。看完再來看本節(jié)就很順暢了。

我們觀察到，在V2 BWD中，內(nèi)外循環(huán)的位置又換回來了，即還是KV外循環(huán)，Q內(nèi)循環(huán)，這是為什么呢？

我們知道在BWD的過程中，我們主要是求（為了求它們還需要求中間結(jié)果，我們來總結(jié)一下這些梯度都需要沿著哪些方向AllReduce：

：沿著i方向做AllReduce，也就是需要每行的結(jié)果加總

：沿著i方向做AllReduce，也就是需要每行的結(jié)果加總

：沿著j方向做AllReduce，也就是需要每列的結(jié)果加總

：只與當(dāng)前i,j相關(guān)

基于此，如果你還是保持Q外循環(huán)，KV外循環(huán)不變的話，這種操作其實(shí)是固定行，遍歷列的，那么在這些梯度中，只有從中受益了，K和V的梯度則進(jìn)入了別扭的循環(huán)（也意味著要往shared memory上寫更多的中間結(jié)果）；但如果你采用KV外循環(huán)，Q內(nèi)循環(huán)，這樣K和V都受益，只有Q獨(dú)自別扭，因此是一種更好的選擇。（S和P的計(jì)算不受循環(huán)變動(dòng)影響）。

前面說過，在BWD過程中讀寫我們要用全局的重新計(jì)算，計(jì)算公式如下：

但如此一來，我們就要從shared memory上同時(shí)讀取，似乎有點(diǎn)消耗讀寫。所以在V2中，我們只存儲(chǔ)，然后計(jì)算：

很容易發(fā)現(xiàn)這兩個(gè)計(jì)算是等價(jià)的，但V2的做法節(jié)省了讀寫量

好，現(xiàn)在我們就把V2相對(duì)于V1在計(jì)算原理上的改進(jìn)介紹完了。接下來我們總結(jié)一下V2相對(duì)于V1所有的改進(jìn)點(diǎn)。

二、V2相對(duì)V1的改進(jìn)點(diǎn)

之所以把這塊內(nèi)容放到“V2整體流程介紹”之后，是想讓大家在先理解V2是怎么做的基礎(chǔ)上，更好體會(huì)V2的優(yōu)點(diǎn)。

總體來說，V2從以下三個(gè)方面做了改進(jìn)：

置換內(nèi)外循環(huán)位置，同時(shí)減少非矩陣的計(jì)算量。（這兩點(diǎn)我們?cè)诘谝徊糠种幸呀o出詳細(xì)說明）

優(yōu)化Attention部分thread blocks的并行化計(jì)算，新增seq_len維度的并行，使SM的利用率盡量打滿。這其實(shí)也是內(nèi)外循環(huán)置換這個(gè)總體思想配套的改進(jìn)措施

優(yōu)化thread blocks內(nèi)部warp級(jí)別的工作模式，盡量減少warp間的通訊和讀取shared memory的次數(shù)。

第二和第三點(diǎn)都可以歸結(jié)為是cuda gemm層面的優(yōu)化，我們馬上來細(xì)看這兩點(diǎn)。

三、V2中的thread blocks排布

//gridDiminV1
//params.b=batch_size,params.h=num_heads
dim3grid(params.b,params.h);

//gridDiminV2
constintnum_m_block=(params.seqlen_q+Kernel_traits::kBlockM-1)/Kernel_traits::kBlockM;
dim3grid(num_m_block,params.b,params.h);

這段代碼整合自flash attention github下的cutlass實(shí)現(xiàn)，為了方便講解做了一點(diǎn)改寫。

這段代碼告訴我們：

在V1中，我們是按batch_size和num_heads來劃分block的，也就是說一共有batch_size * num_heads個(gè)block，每個(gè)block負(fù)責(zé)計(jì)算O矩陣的一部分

在V2中，我們是按batch_size，num_heads和num_m_block來劃分block的，其中num_m_block可理解成是沿著Q矩陣行方向做的切分。例如Q矩陣行方向長度為seqlen_q（其實(shí)就是我們熟悉的輸入序列長度seq_len，也就是圖例中的N），我們將其劃分成num_m_block份，每份長度為kBlockM（也就是每份維護(hù)kBlockM個(gè)token）。這樣就一共有batch_size * num_heads * num_m_block個(gè)block，每個(gè)block負(fù)責(zé)計(jì)算矩陣O的一部分。

為什么相比于V1，V2在劃分thread block時(shí)，要新增Q的seq_len維度上的劃分呢？

先說結(jié)論，這樣做的目的是盡量讓SM打滿。我們知道block是會(huì)被發(fā)去SM上執(zhí)行的。以1塊A100 GPU為例，它有108個(gè)SM，如果此時(shí)我們的block數(shù)量比較大（例如論文中所說>=80時(shí)），我們就認(rèn)為GPU的計(jì)算資源得到了很好的利用。現(xiàn)在回到我們的輸入數(shù)據(jù)上來，當(dāng)batch_size和num_heads都比較大時(shí)，block也比較多，此時(shí)SM利用率比較高。但是如果我們的數(shù)據(jù)seq_len比較長，此時(shí)往往對(duì)應(yīng)著較小的batch_size和num_heads，這是就會(huì)有SM在空轉(zhuǎn)了。而為了解決這個(gè)問題，我們就可以引入在Q的seq_len上的劃分。

看到這里你可能還是有點(diǎn)懵，沒關(guān)系，我們通過圖解的方式，來一起看看V1和V2上的thread block到底長什么樣。

3.1 V1 thread block

假設(shè)batch_size = 1，num_heads = 2，我們用不同的顏色來表示不同的head。

我們知道在Multihead Attention中，各個(gè)head是可以獨(dú)立進(jìn)行計(jì)算的，在計(jì)算完畢后將結(jié)果拼接起來即可。所以我們將1個(gè)head劃分給1個(gè)block，這樣就能實(shí)現(xiàn)block間的并行計(jì)算，如此每個(gè)block只要在計(jì)算完畢后把結(jié)果寫入自己所維護(hù)的O的對(duì)應(yīng)位置即可。

而每個(gè)block內(nèi)，就能執(zhí)行V1中的"KV外循環(huán)，Q內(nèi)循環(huán)”的過程了，這個(gè)過程是由block的再下級(jí)warp level層面進(jìn)行組織，thread實(shí)行計(jì)算的。這塊我們放在第四部分中講解。

3.2 V2 thread block

現(xiàn)在我們繼續(xù)假設(shè)batch_size = 1，num_heads = 2。

與V1不同的是，我們?cè)赒的seq_len維度上也做了切分，將其分成四份，即num_m_block = 4。所以現(xiàn)在我們共有124 = 8個(gè)block在跑。這些block之間的運(yùn)算也是獨(dú)立的，因?yàn)椋?/p>

head的計(jì)算是獨(dú)立的，所以紅色block和藍(lán)色block互不干擾

采用Q做外循環(huán)，KV做內(nèi)循環(huán)時(shí)，行與行之間的block是獨(dú)立的，因此不同行的block互相不干擾。

每個(gè)block從Q上加載對(duì)應(yīng)位置的切塊，同時(shí)從KV上加載head0的切塊，計(jì)算出自己所維護(hù)的那部分O，然后寫入O的對(duì)應(yīng)位置。

在這里你可能想問，為什么只對(duì)Q的seq_len做了切分，而不對(duì)KV的seq_len做切分呢？

在V2的cutlass實(shí)現(xiàn)中，確實(shí)也提供了對(duì)KV的seq_len做切分的方法。但除非你認(rèn)為SM真得打不滿，否則盡量不要在KV維度上做切分，因?yàn)槿绱艘粊恚煌腷lock之間是沒法獨(dú)立計(jì)算的（比如對(duì)于O的某一行，它的各個(gè)部分來自不同的block，為了得到全局的softmax結(jié)果，這些block的結(jié)果還需要匯總做一次計(jì)算）。

3.3 seq parallel不是V2特有

如果你看過V1的代碼，你會(huì)發(fā)現(xiàn)，其實(shí)在V1后期的版本中，也出現(xiàn)了seq維度的并行：

//V1seqparallel:csrc/flash_attn/src/fmha_fwd_launch_template.h
dim3grid(launch_params.params.b,launch_params.params.h,launch_params.params.num_splits);

//nums_splits計(jì)算方法
//Findthenumberofsplitsthatmaximizestheoccupancy.Forexample,ifwehave
//batch*n_heads=48andwehave108SMs,having2splits(efficiency=0.89)is
//betterthanhaving3splits(efficiency=0.67).However,wealsodon'twanttoomany
//splitsasthatwouldincurmoreHBMreads/writes.
//Sowefindthebestefficiency,thenfindthesmallestnumberofsplitsthatgets95%
//ofthebestefficiency.
//[2022-11-25]TD:Markthisas"inline"otherwiseweget"multipledefinition"error.
inlineintnum_splits_heuristic_fwd(intbatch_nheads,intnum_SMs,intctas_per_sm,intmax_splits){
floatmax_efficiency=0.f;
std::vectorefficiency;
efficiency.reserve(max_splits);
for(intnum_splits=1;num_splits<=?max_splits;?num_splits++)?{
????????float?n_waves?=?float(batch_nheads?*?num_splits)?/?(num_SMs?*?ctas_per_sm);
????????float?eff?=?n_waves?/?ceil(n_waves);
????????//?printf("num_splits?=?%d,?eff?=?%f
",?num_splits,?eff);
????????if?(eff?>max_efficiency){max_efficiency=eff;}
efficiency.push_back(eff);
}
for(intnum_splits=1;num_splits<=?max_splits;?num_splits++)?{
????????if?(efficiency[num_splits?-?1]?>0.95*max_efficiency){
//printf("num_splitschosen=%d
",num_splits);
returnnum_splits;
}
}
return1;
}

....
//可以發(fā)現(xiàn)num_splits也是由Q的seq_len維度切分來的
launch_params.params.num_splits=num_splits_heuristic_fwd(
launch_params.params.b*launch_params.params.h,dprops->multiProcessorCount,
ctas_per_sm,
/*max_splits=*/std::min(30,(launch_params.params.seqlen_q+M-1/M))
);

上圖代碼中的num_splits也是在由Q的seq_len維度切分來的。通過這段代碼，我猜想作者在V1后期引入seq_len維度切分的原因是：V1也需要解決seq_len過長時(shí)，batch_size和num_heads較小而造成SM打不滿的問題。

num_splits_heuristic_fwd這個(gè)函數(shù)的作用概括起來就是，我先提供一連串num_splits值的備選，然后由這個(gè)函數(shù)計(jì)算出每個(gè)備選值下SM的利用率。計(jì)算完之后，我先找到最高的利用率，然后再找出滿足利用率>=0.95 * max(利用率)的那個(gè)最小的num_split值，作為最終的選擇。

細(xì)心的你此時(shí)可能已經(jīng)觀察到了，雖然V1也引進(jìn)過seq parallel，但是它的grid組織形式時(shí)(batch_size, num_heads, num_m_blocks)，但V2的組織形式是(num_m_blocks, batch_size, num_heads)，這種順序調(diào)換的意義是什么呢？

直接說結(jié)論，這樣的調(diào)換是為了提升L2 cache hit rate。大家可以看下3.2中的圖（雖然block實(shí)際執(zhí)行時(shí)不一定按照?qǐng)D中的序號(hào)），對(duì)于同一列的block，它們讀的是KV的相同部分，因此同一列block在讀取數(shù)據(jù)時(shí)，有很大概率可以直接從L2 cache上讀到自己要的數(shù)據(jù)（別的block之前取過的）。

3.4 FWD和BWD過程中的thread block劃分

在3.1～3.3中，我們其實(shí)給出的是FWD過程中thread block的劃分方式，我們知道V2中FWD和BWD的內(nèi)外循環(huán)不一致，所以對(duì)應(yīng)來說，thread block的劃分也會(huì)有所不同，我們?cè)敿?xì)來看：

在圖中：

worker表示thread block，不同的thread block用不同顏色表示

整個(gè)大方框表示輸出矩陣O

我們先看左圖，它表示FWD下thread block的結(jié)構(gòu)。每一行都有一個(gè)worker，它表示O矩陣的每一行都是由一個(gè)thread block計(jì)算出來的（假設(shè)num_heads = 1），這就對(duì)應(yīng)到我們3.1～3.3中說的劃分方式。那么白色的部分表示什么呢？我們知道如果采用的是casual attention，那么有一部分是會(huì)被mask掉的，所以這里用白色來表示。但這不意味著thread block不需要加載白色部分?jǐn)?shù)據(jù)對(duì)應(yīng)的KV塊，只是說在計(jì)算的過程中它們會(huì)因被mask掉而免于計(jì)算（論文中的casual mask一節(jié)有提過）。

我們?cè)倏从覉D，它表示BWD下thread block的結(jié)構(gòu)，每一列對(duì)應(yīng)一個(gè)worker，這是因?yàn)锽WD中我們是KV做外循環(huán)，Q做內(nèi)循環(huán)，這種情況下dK, dV都是按行累加的，而dQ是按列累加的，少數(shù)服從多數(shù)，因此這里thread_block是按的列劃分的。

四、Warp級(jí)別并行

講完了thread block，我們就可以再下一級(jí)，看到warp level級(jí)別的并行了。左圖表示V1，右圖表示V2。不管是V1還是V2，在Ampere架構(gòu)下，每個(gè)block內(nèi)進(jìn)一步被劃分為4個(gè)warp，在Hopper架構(gòu)下則是8個(gè)warp。

在左圖（V1）中，每個(gè)warp都從shared memory上讀取相同的Q塊以及自己所負(fù)責(zé)計(jì)算的KV塊。在V1中，每個(gè)warp只是計(jì)算出了列方向上的結(jié)果，這些列方向上的結(jié)果必須匯總起來，才能得到最終O矩陣行方向上的對(duì)應(yīng)結(jié)果。所以每個(gè)warp需要把自己算出來的中間結(jié)果寫到shared memory上，再由一個(gè)warp（例如warp1）進(jìn)行統(tǒng)一的整合。所以各個(gè)warp間需要通訊、需要寫中間結(jié)果，這就影響了計(jì)算效率。

在左圖（V2）中，每個(gè)warp都從shared memory上讀取相同的KV塊以及自己所負(fù)責(zé)計(jì)算的Q塊。在V2中，行方向上的計(jì)算是完全獨(dú)立的，即每個(gè)warp把自己計(jì)算出的結(jié)果寫到O的對(duì)應(yīng)位置即可，warp間不需要再做通訊，通過這種方式提升了計(jì)算效率。不過這種warp并行方式在V2的BWD過程中就有缺陷了：由于bwd中dK和dV是在行方向上的AllReduce，所以這種切分方式會(huì)導(dǎo)致warp間需要通訊。

針對(duì)V2 warp切分影響B(tài)WD這點(diǎn)，作者在論文中依然給出了“BWD過程相比V1也有提升”的結(jié)論，針對(duì)這點(diǎn)，我在github issue上找到了一條作者的回復(fù)（在“安裝報(bào)錯(cuò)”組成的issue海洋里撈出的寶貴一條）：

最關(guān)鍵的可能是第1和第2點(diǎn)，關(guān)于第1點(diǎn)，我想作者應(yīng)該是說，之前需要反復(fù)讀取KV的數(shù)據(jù)，現(xiàn)在只用反復(fù)讀取Q的數(shù)據(jù)，因此從一定程度上節(jié)省了shared memory的讀寫次數(shù)。第2點(diǎn)理解起來有點(diǎn)復(fù)雜，個(gè)人覺得是將warp處理的tile劃分得更像方形。這樣做的好處是在做casual mask的時(shí)候可以方便寫代碼大塊丟掉被mask掉的tile（見論文casual masking部分），進(jìn)一步加速計(jì)算。第3點(diǎn)是關(guān)于一些底層的優(yōu)化，就不提了。

好！關(guān)于V2我們就介紹到這了，寫這篇文章的時(shí)候，我剛粗過了一遍triton的flash attention實(shí)現(xiàn)，以及掃了一下cutlass實(shí)現(xiàn)的入口。如果后續(xù)有時(shí)間，我會(huì)出一些源碼解讀的文章（從cuda gemm -> triton gemm -> triton flash attention，看，又給自己挖了一個(gè)坑）。如果出不了，那一定不是我鴿人，那肯定是我不會(huì)（沒錯(cuò)，就是這樣）。

審核編輯：黃飛