一、背景

我們?cè)谥暗奈恼轮泻唵谓榻B了 Sample Packing 相關(guān)的技術(shù)方案及涉及的問題，也在看其中 Attention 計(jì)算帶來的各種挑戰(zhàn)。機(jī)緣巧合正好看到一篇文章試圖解決相應(yīng)的 Attention 計(jì)算問題，這里進(jìn)行簡單介紹。

對(duì)應(yīng)的論文為：[2409.15097] Efficiently Dispatching Flash Attention For Partially Filled Attention Masks

相關(guān)工作可以參考我們之前的文章：

• ??Sample Packing：長序列 LLM 訓(xùn)練的 Attention 問題及優(yōu)化??
• ??Sample Packing 綜述：LLM 效果與效率的 Tradeoff??
• ??萬字綜述 10+ 種 LLM 投機(jī)采樣推理加速方案??
• ??LLM 投機(jī)解碼 & 美杜莎(Medusa)實(shí)現(xiàn)??
• ??LLM 推理的 Attention 計(jì)算和 KV Cache 優(yōu)化：PagedAttention、vAttention 等??
• ??大規(guī)模分布式 AI 模型訓(xùn)練系列——序列并行??

二、摘要

Transformer 已廣泛應(yīng)用于各種應(yīng)用場景，會(huì)產(chǎn)生稀疏或部分填充的注意力矩陣，比如旨在降低 Attention 二次復(fù)雜度的 Attention Mask、Sample Packing 或其他技術(shù)引入的稀疏 Attention，例如用于 Medusa 投機(jī)采樣快速驗(yàn)證的 Tree Attention。盡管這些矩陣具有固定的稀疏性，但 SOTA 的 FlashAttention 仍然將其當(dāng)做稠密矩陣處理。

本文中，作者提出了 Binary Block Masking (BinBlkMsk) ，通過使 FlashAttention 具備 Mask 感知能力來增強(qiáng) FlashAttention。作者也進(jìn)一步提出了兩種優(yōu)化方案，在真實(shí)場景的 Attention Mask 實(shí)驗(yàn)表明，運(yùn)行速度可以提升 9x。

• 一種針對(duì)具有連續(xù)非零 Mask：Dense Binary Block Masking。
• 一種針對(duì)極度稀疏 Mask：Binary Block Masking with RCM。?

PS：其實(shí)上述的 Baseline 有問題，F(xiàn)lashAttention 可以通過 Varlen 方案支持 Sample Packing 引入的 Block Diagonal Mask，避免當(dāng)做稠密矩陣計(jì)算。此外，Pytorch 官方也發(fā)布了 FlexAttention，可以支持各種 Attention Mask 變體，本文作者也并沒有對(duì)比。

三、方案

3.1 Binary Block Masking

以 Sample Packing 引入的 Block Diagonal Mask 為例，如下圖所示，Max Sequence Length 為 16，拼接了 3 個(gè) Sample，長度分別為 4,5,7。因此會(huì)形成一個(gè) 16x16 的 Global Mask，其中包含 3 個(gè) Causal Mask。

在 FlashAttention 實(shí)際的計(jì)算中會(huì)將其劃分為不同的 Block 計(jì)算。本文提出的 Binary Block Masking 是通過一個(gè)額外的 Binary Block Matrix 來標(biāo)記哪些 Block 全是 0，而哪些中有非 0 值。如下圖，假設(shè) Block 的大小為 2x2，則會(huì)生成一個(gè) 8x8 Binary Matrix，如下圖紅框和藍(lán)框?qū)?yīng)的位置為 1，其他位置為 0。在計(jì)算時(shí)，如果 Binary Block Matrix 中為 0 的 Mask 可以忽略不計(jì)算。

3.2 Dense Binary Block Masking

在自然語言場景中，通常上述的 Attention Mask 的 Block 中會(huì)存在很多連續(xù)全 1 的 Block，只在邊界的 Block 里既有 0 又有 1。對(duì)于全為 1 的 Block，在計(jì)算時(shí)不用再讀取 Attention Mask，減少訪存并提升計(jì)算效率。因此作者使用兩個(gè)額外的數(shù)組 total_ones 和 offset 來標(biāo)識(shí)這些連續(xù)全 1 Block，如下圖所示，total_ones[7] 為 2，表示最后一行有 2 個(gè)全 1 Block；offset[7] 為 5，表示最后一行連續(xù)全 1 Block 的起始索引為 5。

• total_ones 存儲(chǔ)全 1 Block 的個(gè)數(shù)。
• offset 存儲(chǔ) Binary Block Matrix 中每行連續(xù)全 1 Block 的起始位置。
• 右圖 Binary Block Matrix 中紅色 1表示 Block 中全是 1，黑色 1 表示 Block 中部分為 1。?

PS：實(shí)際計(jì)算中，對(duì)于所有 Layer，所有 Attention Head 都只需計(jì)算一次。

3.3 Binary Block Masking with RCM

在處理非常稀疏的注意力掩碼時(shí)，如果使用標(biāo)準(zhǔn)的 Binary Block Masking 方法，可能會(huì)存在效率問題。這是因?yàn)榧词拐麄€(gè)塊中只有一個(gè)非零值，也必須處理整個(gè)塊，導(dǎo)致計(jì)算資源的浪費(fèi)。作者使用 Reverse Cuthill-McKee (RCM) 算法來有效地重新組織 Mask 矩陣的結(jié)構(gòu)，從而減少必須處理的 Block 的數(shù)量。如下圖 Figure 3 所示：

如下圖所示，作者展示了一個(gè)合成掩碼的結(jié)果，其中 RCM 預(yù)處理減少了 90% 的 Block 數(shù)量，從而獲得顯著的性能提升。Binary Block Masking with RCM 的速度明顯快于 Binary Block Masking。

四、實(shí)驗(yàn)和結(jié)果

4.1 配置

作者主要評(píng)估三種 Mask 類型，包括 Medusa（[2401.10774] Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads） 里面的 Tree Attention Mask，Sample Packing 中的 Block Diagonal Mask 以及 LongFormer（[2004.05150] Longformer: The Long-Document Transformer） 中的稀疏 Mask。測試時(shí)使用 Batch Size 為 4，32 個(gè) Attention Head。FlashAttention 中的 Block Size 固定為 (128, 32)。

此外，測試的 FlashAttention 方案并不是開源的 FlashAttention 實(shí)現(xiàn)，而是作者基于 Triton 實(shí)現(xiàn)的，以方便對(duì)比。

4.2 預(yù)處理

如下圖 Table 1 所示，作者使用 Triton 實(shí)現(xiàn)了基于 Attention Mask 生成 Binary Block Mask 的 Kernel，并測試了相應(yīng)的時(shí)間。雖然 Binary Block Mask 的生成時(shí)間和 Batch Size 為 1，1 個(gè) Head 的計(jì)算時(shí)間差不多，但是整體只用計(jì)算一次，影響比較小。

4.3 Medusa Tree Attention

如下圖 Figure 3 所示為對(duì) Medusa Tree Attention 的優(yōu)化效果。左圖為 Tree Attention 的 Mask，其中 K 表示 Medusa 中 Head 個(gè)數(shù)，sk 表示每個(gè) Head 候選 Token 個(gè)數(shù)。(b) 表示固定候選 Token 為 3，不同 Head 個(gè)數(shù)下的速度；(c) 表示固定 4 個(gè) Head，不同候選 Token 的速度。

• Naive Attention Masking：所有 Block 都計(jì)算。
• Bash FlashAttention：Causal Mask，只計(jì)算下三角的 Block，當(dāng) Head 或者候選 Token 很多時(shí) FlashAttention 的 Causal Mask 時(shí)間大概是 Naive 的一半（計(jì)算量大約是一半）。
• Binary Block Masking：只計(jì)算部分 Block，速度遠(yuǎn)快于上述兩種方案。?

4.4 Sample Packing Block Diagonal Mask

作者同樣評(píng)估了 Sample Packing 場景的 Block Diagonal Mask，如下圖（a）表示標(biāo)準(zhǔn) Causal Mask 的組合；（b）表示 Input 雙向 Attention，而 Output Causal Mask 的情況。從（b）和（c）可以得出類似 Tree Mask 的結(jié)論，當(dāng)序列比較長時(shí)，本文的 Binary Block Masking 可以帶來 9x 加速。不過需要說明的是：Base FlashAttention 的時(shí)間幾乎是 Naive Attention Masking 的一半，也可以推導(dǎo)出作者使用的 FlashAttention 方案確實(shí)計(jì)算了整個(gè)下三角 Block，而沒有使用 Varlen 方案。

4.5 LongFormer Attention Mask

作者同樣評(píng)估了 LongFormer 場景的各種稀疏 Mask，如下圖分別為 Sliding Window Attention，Dilated Attention 以及 Global Attention 等場景，可以得出類似的結(jié)論。需要說明的是，此時(shí) Base FlashAttention 相比 Naive Attention 的提升比較有限。

4.6 Pytorch FlexAttention

Pytorch 在 2.5.0 版本引入了 FlexAttention（FlexAttention: The Flexibility of PyTorch with the Performance of FlashAttention），可以很容易支持各種 Attention Mask 變種，比如標(biāo)準(zhǔn) Causal Mask、Sliding Window 、Prefix Mask 以及 Block Diagonal Mask 等，相比 FlashAttention 靈活了很多，本文提到的各種 Attention 基本也都能實(shí)現(xiàn)。

如下表所示，我們同樣做了 Sample Packing 的實(shí)驗(yàn)，隨著 Sequence 中 Sample 分布的不同，計(jì)算的耗時(shí)甚至可能差 10x，與本文中 9x 的提升也能對(duì)上。

五、參考鏈接

1. ??https://arxiv.org/abs/2409.15097??
2. ??https://arxiv.org/abs/2401.10774??
3. ??https://arxiv.org/abs/2004.05150??
4. ??https://pytorch.org/blog/flexattention/???

本文轉(zhuǎn)載自 ??AI閑談??，作者： AI閑談