來自NVIDIA、CMU和華盛頓大學的研究人員發布了「FlashInfer」：一個提供最先進LLM推理和服務的內核實現的內核庫

大型語言模型 (LLMs) 現在已經成為現代人工智慧應用中不可或缺的一部分，驅動著像聊天機器人和程式碼生成器等工具。然而，對這些模型的依賴增加也暴露了推理過程中的一些重要效率問題。注意力機制，例如 FlashAttention 和 SparseAttention，常常在多樣化的工作負載、動態輸入模式和 GPU 資源限制下表現不佳。這些挑戰，加上高延遲和記憶體瓶頸，顯示出需要一個更有效和靈活的解決方案來支持可擴展和響應迅速的 LLM 推理。

來自華盛頓大學 (University of Washington)、NVIDIA、Perplexity AI 和卡內基梅隆大學 (Carnegie Mellon University) 的研究人員開發了 FlashInfer，這是一個專為 LLM 推理量身定制的人工智慧庫和內核生成器。FlashInfer 提供了多種注意力機制的高效 GPU 內核實現，包括 FlashAttention、SparseAttention、PageAttention 和取樣。它的設計優先考慮靈活性和效率，解決了 LLM 推理服務中的關鍵挑戰。

FlashInfer 採用了區塊稀疏格式來有效處理異質的 KV-cache 存儲，並使用動態負載平衡調度來優化 GPU 的使用。它與流行的 LLM 服務框架如 SGLang、vLLM 和 MLC-Engine 集成，提供了一種實用且可適應的方法來提升推理性能。

技術特點與優勢

FlashInfer 引入了幾項技術創新：

• 全面的注意力內核：FlashInfer 支持多種注意力機制，包括預填充、解碼和附加注意力，確保與各種 KV-cache 格式的兼容性。這種適應性提高了單請求和批量服務場景的性能。
• 優化的共享前綴解碼：通過分組查詢注意力 (GQA) 和融合旋轉位置嵌入 (RoPE) 注意力，FlashInfer 在長提示解碼方面實現了顯著的速度提升，例如相較於 vLLM 的 Page Attention 實現，速度提高了 31 倍。
• 動態負載平衡調度：FlashInfer 的調度器能夠根據輸入變化動態調整，減少 GPU 空閒時間，確保高效利用。它與 CUDA Graphs 的兼容性進一步增強了其在生產環境中的適用性。
• 可自定義的 JIT 編譯：FlashInfer 允許用戶定義和編譯自定義的注意力變體為高性能內核。這一特性適應了專門的使用案例，例如滑動窗口注意力或 RoPE 轉換。

性能洞察

FlashInfer 在多個基準測試中顯示出顯著的性能提升：

• 延遲減少：該庫將每個標記之間的延遲減少了 29-69%，相較於現有解決方案如 Triton。這些增益在涉及長上下文推理和並行生成的場景中特別明顯。
• 吞吐量提升：在 NVIDIA H100 GPU 上，FlashInfer 在並行生成任務中實現了 13-17% 的速度提升，突顯了其在高需求應用中的有效性。
• 增強的 GPU 利用率：FlashInfer 的動態調度器和優化內核改善了帶寬和 FLOP 利用率，特別是在序列長度不均或均勻的場景中。

FlashInfer 在並行解碼任務中也表現優異，組合格式使得首次標記時間 (TTFT) 顯著減少。例如，在 Llama 3.1 模型 (70B 參數) 上的測試顯示，在特定配置下，TTFT 減少了高達 22.86%。

結論

FlashInfer 提供了一個實用且高效的解決方案來應對 LLM 推理的挑戰，顯著改善了性能和資源利用率。其靈活的設計和集成能力使其成為推進 LLM 服務框架的寶貴工具。通過解決關鍵的效率問題並提供強大的技術解決方案，FlashInfer 為更可及和可擴展的人工智慧應用鋪平了道路。作為一個開源項目，它邀請研究社群進一步合作和創新，確保持續改進和適應人工智慧基礎設施中出現的新挑戰。

查看論文和 GitHub 頁面。所有對這項研究的貢獻都歸功於這個項目的研究人員。此外，別忘了在 Twitter 上關注我們，並加入我們的 Telegram 頻道和 LinkedIn 群組。還有，加入我們的 60k+ ML SubReddit。

🚨 免費即將舉行的人工智慧網路研討會 (2025 年 1 月 15 日)：使用合成數據和評估智慧提升 LLM 準確性——加入這個網路研討會，獲取提升 LLM 模型性能和準確性的可行見解，同時保護數據隱私。