【FlashAttention-V4，非官方】FlashDecoding++

1. Introdcution

為了提高softmax并行性，之前方法（FlashAttention、FlashDecoding）將計(jì)算過程拆分，各自計(jì)算partial softmax結(jié)果，最后需要通過同步操作來更新partial softmax結(jié)果。例如FlashAttention每次計(jì)算partial softmax結(jié)果都會更新之前的結(jié)果，而FlashDecoding是在最后統(tǒng)一更新所有partial softmax結(jié)果。

本文在A100 GPU上分析了輸入長度為1024的情況，這種同步partial softmax更新操作占Llama2-7B推理的注意力計(jì)算的18.8%。（本文沒說是FlashAttention還是FlashDecoding的結(jié)果，個(gè)人認(rèn)為FlashDecoding的同步更新代價(jià)并不大，應(yīng)該遠(yuǎn)小于18.8%）

這是LLM推理加速的第一個(gè)挑戰(zhàn)。此外，本文還提出了兩個(gè)挑戰(zhàn)：

在解碼階段，F(xiàn)lat GEMM操作的計(jì)算資源未得到充分利用。這是由于解碼階段是按順序生成token（一次只生成一個(gè)token），GEMM操作趨于flat-shape，甚至batch size等1時(shí)變成了GEMV（General Matrix-Vector Multiplication），具體看論文Figure 2。當(dāng)batch size較小時(shí)（e.g., 8），cublas和cutlass會將矩陣填充zeros以執(zhí)行更大batchsize（e.g., 64）的GEMM，導(dǎo)致計(jì)算利用率不足50%。

動態(tài)輸入和固定硬件配置影響了LLM推理的性能。例如，當(dāng)batch size較小時(shí)，LLM推理的解碼過程是memory-bounded，而當(dāng)batch size較大時(shí)是compute-bounded。

針對這3個(gè)問題，本文分別提出了對應(yīng)優(yōu)化方法：

Asynchronized softmax with unified max value.FlashDecoding++為分塊softmax計(jì)算設(shè)置了一個(gè)共享的最大值。這樣可以獨(dú)立計(jì)算partial softmax，無需同步更新。

Flat GEMM optimization with double buffering.FlashDecoding++只將矩陣大小填充到8，對比之前針對flat-shaped GEMM設(shè)計(jì)的為64，提高了計(jì)算利用率。論文指出，具有不同shape的flat GEMMs面臨的瓶頸也不同，于是進(jìn)一步利用雙緩沖等技術(shù)提高kernel性能。

Heuristic dataflow with hardware resource adaption.FlashDecoding++同時(shí)考慮了動態(tài)輸入和硬件配置，針對LLM推理時(shí)數(shù)據(jù)流進(jìn)行動態(tài)kernel優(yōu)化。

下圖展示了以上3種方法的示意圖：

2. Backgrounds

LLM推理中的主要操作如下圖所示：linear projection(①和⑤)、attention(②、③和④)和feedforward network(⑥)。為簡單起見，這里忽略了position embedding、non-linear activation、mask等操作。本文將LLM推理時(shí)對Prompt的處理過程稱為prefillphase，第二階段預(yù)測過程稱為decodephase。這兩個(gè)階段的算子基本一致，主要是輸入數(shù)據(jù)的shape是不同的。由于decodephase一次只處理一個(gè)令牌（batch size=1，或batch size很?。虼溯斎刖仃囀莊lat-shape matrices（甚至是vectors），參見下圖Decode phase部分中和KV Cache拼接的紅色向量。

LLM推理中的另一個(gè)問題就是Softmax算子，其需要計(jì)算并存儲所有全局?jǐn)?shù)據(jù)，并且數(shù)據(jù)量隨著數(shù)據(jù)長度成平方增長，存在內(nèi)存消耗高和低并行性等問題。一般計(jì)算流程如下：

3. Asynchronized Softmax with Unified Maximum Value如下

圖b所示，F(xiàn)lashAttention和FlashDecoding對softmax操作進(jìn)行了分塊處理，但是塊與塊之間需要進(jìn)行同步（主要是局部最大值）。本文發(fā)現(xiàn)這種同步操作的開銷約為20%。因此，作者希望去除同步操作，也就是獨(dú)立計(jì)算出partial softmax結(jié)果。

4. Flat GEMM Optimization with Double Buffering

Decoding階段的過程主要由GEMV（batch size=1）或flat GEMM（batch size>1）。GEMV/GEMM運(yùn)算可以用M、N、K來表示，其中兩個(gè)相乘矩陣的大小分別為M × K和K × N。一般LLM推理引擎利用Tensor Core使用cuBLAS和CUTLASS等庫來加速。盡管Tensor Core適合處理M = 8的GEMM，但這些庫為了隱藏memory latency，通常將M維度平鋪到64。然而，decodephase的GEMV或flat GEMM的M通遠(yuǎn)小于64，于是填充0到64，導(dǎo)致計(jì)算利用率低下。

為了隱藏memory access latency，本文引入了double buffering技術(shù)。具體來說就是在共享內(nèi)存中分配兩個(gè)buffer，一個(gè)buffer用于執(zhí)行當(dāng)前tile的GEMM計(jì)算，同時(shí)另一個(gè)buffer則加載下一個(gè)tile GEMM所需的數(shù)據(jù)。這樣計(jì)算和內(nèi)存訪問是重疊的，本文在N較大時(shí)采取這種策略，下圖為示意圖。

5. Heuristic Dataflow with Hardware Resource Adaption

影響LLM推理性能的因素有很多：（a）動態(tài)輸入。batch size和輸入序列長度的變化造成了工作負(fù)載變化。（b）模型多樣性。主要指模型結(jié)構(gòu)和模型大小。（c）GPU能力不同。例如內(nèi)存帶寬、緩存大小和計(jì)算能力。（d）工程優(yōu)化。

雖然這些因素構(gòu)建了一個(gè)很大的搜索空間，但LLM中不同layer的同質(zhì)性大大減少了算子優(yōu)化的搜索空間。例如，prefillphase和decodephase中有4個(gè)GEMV/GEMM操作（K、Q、V投影、O投影、2個(gè)FFN），都可以表示為[M, K]和N x K，對應(yīng)了四種[N, K]組合，如下圖所示。此外，prefillphase的M與輸入序列長度和batch size有關(guān)，decodephase的M只與batch size有關(guān)。

本文根據(jù)不同的M, K, N選取FastGEMV、flat GEMM（本文方法）、CUTLASS。

個(gè)人總結(jié)

這篇文章沒有FlashAttention和FlashDecoding驚艷，個(gè)人覺得FlashDecoding的同步處理代價(jià)不大，而且本文中動態(tài)調(diào)整softmax方法也引入了判斷、終止和分支跳轉(zhuǎn)等操作。另一個(gè)Double Buffering就是內(nèi)存優(yōu)化常用的乒乓buffer，也沒什么新東西。

不過話說回來，如今在tranformer架構(gòu)不變的情況，LLM加速只能靠這些工程手段去優(yōu)化，的確也有不錯(cuò)效果。還是很有價(jià)值的。