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

原標題：高效Attention引擎是怎樣煉成的？陳天奇團隊FlashInfer打響新年第一槍！
文章來源：新智元
內容字數：6055字

FlashInfer：高效LLM推理引擎

新年伊始，陳天奇團隊發布了FlashInfer論文，旨在提升大語言模型（LLM）的推理效率。該論文詳細介紹了FlashInfer如何通過塊稀疏、可組合、可定制以及負載均衡等技術，在各種推理場景下顯著提升性能。

1. LLM推理的挑戰

   傳統的Transformer模型在實際應用中面臨諸多工程挑戰，例如多樣化的工作負載、個性化的Attention實現以及對延遲和吞吐量的嚴格要求。不同的推理階段（Prefill、Decode、多輪對話）和技術（Speculative Decoding）都帶來了不同的計算訪存比，影響GPU利用率。此外，各種LLM框架（如vLLM、SGLang）在KV cache存儲方面也存在差異，增加了框架的復雜性。

2. 塊稀疏格式：高效存儲KV Cache

   FlashInfer提出使用塊稀疏（BSR）矩陣來統一存儲各種形式的KV cache。BSR類似于CSR，但操作單元由單個元素變為數據塊，從而提高GPU的利用率。對于Paged Attention、Radix Tree和Token Tree等不同形式的KV cache，BSR都能有效地減少冗余，并最大化Tensor Core的利用率。通過將稀疏塊數據加載到共享內存并排布成密集格式，FlashInfer消除了硬件浪費。

3. 塊并行與可組合性

   FlashInfer采用與BPT相同的塊并行方法分解KV cache，實現并行計算，并通過可組合特性支持不同大小塊的組合，在共享和全局內存之間取得平衡，進一步優化性能。

4. 可定制的注意力機制

   FlashInfer為FlashAttention開發了CUDA/CUTLASS模板，支持不同架構的Nvidia GPU，并提供可自定義的Attention機制。開發者可以通過插入自定義函數來實現個性化的Attention計算，例如ALiBi、Logits SoftCap、RoPE等，提高了框架的靈活性。

5. 高效的數據移動和負載均衡

   FlashInfer采用高效的數據加載方法，將tiles從全局內存轉移到共享內存，并利用異步復制指令最大化內存帶寬。同時，它根據硬件資源和工作負載選擇合適的塊大小，并通過負載均衡調度算法，根據用戶請求動態調整資源分配，并兼容CUDAGraph。

6. 性能提升

   實測結果表明，FlashInfer顯著提升了LLM推理性能，token間延遲降低了29%-69%，長上下文推理延遲降低了28%-30%，并行生成速度提高了13%-17%。

總而言之，FlashInfer通過巧妙地結合塊稀疏格式、塊并行計算、可組合性和負載均衡等技術，為構建高效且可定制的LLM推理引擎提供了一種新的解決方案，顯著提升了LLM的推理效率。

聯系作者

文章來源：新智元
作者微信：
作者簡介：智能+中國主平臺，致力于推動中國從互聯網+邁向智能+新紀元。重點關注人工智能、機器人等前沿領域發展，關注人機融合、人工智能和機器人對人類社會與文明進化的影響，領航中國新智能時代。