高效Attention引擎是怎樣煉成的?陳天奇團隊FlashInfer打響新年第一槍!

新年第一天,FlashInfer在arxiv打響第一槍。

作者團隊來自華盛頓大學、英偉達、Perplexity AI和CMU,曾開發了TVM、XGBoost,同時也是MXNET作者之一的陳天奇,也位列其中。

FlashInfer實現了高效的注意力引擎,利用塊稀疏和可組合格式來解決KV cache存儲異構問題,優化了內存訪問並減少冗餘。
它還提供了可定製的注意力模板,通過即時編譯(JIT)來適應各種Attention的設置。
另外,FlashInfer的負載平衡調度算法,可以根據用戶請求的動態變化進行調整,同時又能夠與靜態配置的CUDAGraph保持兼容。

雖然論文是剛放出來,但FlashInfer早已實際應用到SGLang、vLLM和MLC-Engine等流行的LLM Serving框架中。
實測表明,FlashInfer能夠在各種推理場景中顯著提升內核性能。
與最先進的解決方案相比,FlashInfer將token間延遲降低了29%-69%,將長上下文推理的延遲降低了28%-30%,使並行生成的速度提高了13%-17%。

「對如何在LLM Serving框架中構建高效且可定製的注意力引擎感到好奇嗎?快來看看Flashlnfer的最新論文吧,瞭解所有酷炫的想法。」
高效注意力引擎是怎樣煉成的
LLM推理現狀
原始Transformer本身的計算很簡單,但要更好的為人民服務,就需要解決許多工程上的問題。
實際應用中,服務端會面臨多樣的工作負載,還有個性化的Attention實現,要滿足延遲吞吐量等指標,又要能夠充分發揮硬件的能力。
比如,LLM的推理可以分類為下面三種情況:

Prefill階段拿到最開始的Prompt,填充kv cache;Decode階段則是一個query計算出一個輸出;存在多輪對話或者使用投機推理(Speculative Decoding)時,又可以有多個query向量並行計算。

不同的輸入帶來了不同的計算訪存比,給GPU的利用造成麻煩。
另一方面,當今流行的不同框架和方法在存儲kv cache時,存在很大差異。
比如vLLM使用的Paged Attention,參照操作系統中分頁管理內存的方式,將kv cache切成一個個block,邏輯上連續而物理上不連續。

而SGLang採用的Radix Attention是一棵前綴樹,不同query共享的kv cache(比如系統提示)存儲在同一個節點。

再看Speculative Decoding,多個草稿小模型生成的token序列通常組織成token tree的形式,表示成矩陣時是稀疏的。

還有下面這種,通過重要性計算只選擇topk個kv cache參與Attention計算,同樣是稀疏矩陣的形式。

個性的Attention這麼多,框架想要滿足所有人的要求,kv cache應該怎麼存儲?
塊稀疏格式出馬
作者表示,所有這些都可以抽像成塊稀疏(Block Compressed Sparse Row,BSR)矩陣。
BSR跟通常用來存儲稀疏矩陣的CSR(Compressed Sparse Row)很像,如下圖所示,CSR只需存儲矩陣元素的非零值、在行中的下標、以及每行的偏移。

——但是這麼稀碎的存儲格式對顯卡來說很不友好,於是做一些改進:
相同的存儲方式,對應到BSR,只需將操作的最小單元由元素(單個數據)換成Block(一塊數據),塊大小為(Br,Bc)。

這樣一來,顯卡就可以直接加載整塊的數據(比如16 * 16),來填滿自己的tensor core,進行矩陣乘法。
有了BSR,我們再來看下之前提到的幾種Attention。
Paged Attention的block可以直接對應到BSR的block,如下圖所示,讓Bc = 1。而query tile size(Br)設為4保證矩陣乘法的硬件利用率。

再看本身是樹結構的Radix Tree和Token Tree,表示成矩陣時相當稀疏的,如果BSR使用比較大的方塊(比如16 * 16)會造成很大的浪費。

此時可以設置Bc = 1,如下圖所示,再使用合適的向量長度,就可以消除大部分的冗餘,同時顯卡也支持下圖這種形式的稀疏計算。

有研究表明,稀疏的矩陣只需讓最後一維保持連續,即可享受到tensor core的加速。
在FlashInfer中,系統從global memory中加載稀疏塊數據,並在shared memory中將其排布成密集格式,tensor core可以直接計算這些數據,不會產生硬件的浪費。

塊並行與可組合
另外,FlashInfer採用與Block-Parallel Transformer (BPT)相同的形式分解kv cache(同時也是RingAttention和Flash-Decoding的實現依據)。


對於相同的query,需要計算的完整kv cache被分塊,每塊的Attention計算可以獨立(並行)進行,只需記錄一個統計量(這裏是LSE),用於最後的彙聚。

學會這個之後,再來看一下FlashInfer提供的可組合特性,當query之間存在可共享的kv cache時,可以把存儲表示成不同大小block的組合。

能夠共享的部分使用大的block,讓請求可以通過shared memory訪問;不能共享的部分使用小的block,在global memory中訪問。
拿捏定製化Attention
作者為FlashAttention開發了CUDA/CUTLASS模板,專為密集矩陣和塊稀疏矩陣設計,兼容從Turing到Hopper架構的英偉達GPU。
其中,FlashAttention2用於Ada(sm89)之前的架構,FlashAttention3用於Hopper架構。

實現的主要改進包括:增強將稀疏塊加載到共享內存中的能力、擴展塊大小配置、針對分組查詢注意力(GQA)優化內存訪問模式、以及可自定義的注意力。
下面就說一下這個可自定義的注意力。
我們實際中接觸到的各種LLM,會在算子中實現自己個性化的部分。受FlexAttention啟發,FlashInfer提供了在Attention計算的不同階段插入自定義函數的機制。
比如下面這個用於大模型位置編碼的ALiBi,可以利用FlashInfer提供的接口輕鬆實現。

還有Gemma和Grok需要的Logits SoftCap、常見的RoPE位置編碼、滑動窗口注意力機制等,都可以通過這種方式無痛解決。

數據移動
FlashInfer的注意力模板支持任意的塊大小,由於塊可能與張量核心形狀不對齊,所以需要專門的數據加載方法,也就是前面講過的,將tiles從分散的全局內存轉移到連續的共享內存。
單個MMA指令可以指定可以塊稀疏矩陣中的不同塊作為Tensor core的輸入,下圖展示了FlashInfer如何將tiles加載到共享內存中:

對於稀疏的KV-Cache,地址使用BSR矩陣的indices數組計算;而密集的KV-Cache地址使用row index仿射變換。
KV-Cache的最後一個維度保持連續(也就是head_size),以適應GPU緩存行大小的合併內存訪問。這裏使用寬度為128B的異步複製指令LDGSTS來最大化內存帶寬。
儘管Hopper架構中的TMA(Tensor Memory Accelerator)可以進一步加速數據移動,但它不支持非仿射內存訪問模式,所以只在連續的情況下使用TMA。

另外,為了適應不同的計算訪存比,FlashInfer提供具有不同塊大小的FlashAttention2內核,並根據硬件資源和工作負載強度來選擇合適的尺寸。
負載均衡
下圖展示了FlashInfer運行時調度程序的工作流程。

Attention kernel不會直接產生最終輸出,因為一些長KV被分成了多個chunk,每個chunk的部分輸出存儲在用戶提供的工作區緩衝區中,最終輸出需要所有chunk進行聚合。
FlashInfer實現了高效的注意力合成算子,可以處理變長聚合。每個CTA的工作隊列,以及部分輸出與最終輸出之間的映射,都由調度程序規劃。
先由CPU計算規劃信息,之後FlashInfer會將計劃信息異步複製到GPU工作區緩衝區的特定區域,用作注意力內核的輸入。
FlashInfer保證Attention和contraction內核與CUDAGraph兼容。二者都使用持久內核,並且編譯後網格大小是固定的,這意味著內核在每個生成步驟中都以相同的網格大小啟動。
作者為工作區緩衝區的每個部分設置了固定偏移量,以存儲部分輸出和計劃信息,確保傳遞給內核的指針對於每個生成步驟都是相同的,滿足CUDAGraph的要求。
參考資料:
https://arxiv.org/abs/2501.01005
https://x.com/tqchenml/status/1875123214919328039