01、概述

在人工智能的浩瀚星河中，大型語言模型（LLMs）猶如一顆耀眼的明星，為現代AI應用注入了強勁動力。它們是聊天機器人、代碼生成器等工具背后的“大腦”，憑借海量數據和復雜算法，能夠理解、生成自然語言，甚至編寫代碼，展現出驚人的智能。然而，隨著對這些大模型依賴程度的加深，其推理過程中的諸多問題也逐漸浮出水面，成為制約其發展的瓶頸。

02、大模型推理的困境

注意力機制的局限

注意力機制是大模型的核心組件之一，它幫助模型聚焦于輸入信息中最重要的部分，從而提高處理效率和準確性。但現有的注意力機制，如FlashAttention和SparseAttention，在面對多樣化的任務負載、動態變化的輸入模式以及有限的GPU資源時，常常顯得力不從心。

• 多樣化任務負載：不同的應用場景對模型的要求各異。例如，在聊天機器人中，用戶的問題可能涉及多個話題，需要模型快速切換注意力焦點；而在代碼生成中，模型則需要關注代碼的邏輯結構和語法規范。現有的注意力機制往往難以靈活應對這種多變的任務需求。
• 動態輸入模式：現實世界中的輸入信息是不斷變化的。以在線客服為例，用戶在咨詢過程中可能會隨時插入新的問題或改變話題，導致輸入序列的長度和內容不斷變化。這種動態性使得模型難以提前規劃注意力分配，影響推理效率。
• GPU資源限制：大模型通常需要大量的計算資源，尤其是GPU。但在實際應用中，GPU資源往往是有限的，需要在多個任務或模型之間共享。現有的注意力機制在資源受限的情況下，容易出現計算瓶頸，導致推理速度大幅下降。

高延遲與內存瓶頸

• 高延遲：在一些實時性要求較高的場景，如自動駕駛輔助系統，模型需要在極短的時間內做出決策。然而，現有的大模型推理過程往往存在較高的延遲，無法滿足這種實時性需求。延遲過高會導致系統反應遲緩，影響用戶體驗和系統性能。
• 內存瓶頸：大模型通常包含數十億甚至數千億個參數，存儲這些參數需要大量的內存空間。在推理過程中，模型還需要存儲大量的中間計算結果，進一步增加了內存需求。現有的內存架構和管理方式難以高效地支持這種大規模的內存需求，導致推理速度受限，甚至出現內存溢出的情況。

03、FlashInfer的破局之道

為了解決上述問題，來自華盛頓大學、NVIDIA、Perplexity AI和卡內基梅隆大學的研究人員聯合開發了FlashInfer，這是一個專為大模型推理量身定制的AI庫和內核生成器。FlashInfer通過一系列創新的技術手段，為大模型推理帶來了顯著的性能提升和更高的靈活性。

技術特點與優勢

1）全面的注意力內核支持

FlashInfer支持多種注意力機制，包括預填充、解碼和追加注意力等，能夠與各種KV緩存格式兼容。這種廣泛的適應性使得FlashInfer在單請求和批量服務場景下都能發揮出色的性能。例如，在處理長文本生成任務時，FlashInfer能夠根據不同的注意力機制靈活調整計算策略，提高生成效率和質量。

2）優化的共享前綴解碼

FlashInfer采用了分組查詢注意力（GQA）和融合旋轉位置嵌入（RoPE）注意力等技術，實現了顯著的速度提升。以長提示解碼為例，FlashInfer在與vLLM的Page Attention實現相比，能夠達到31倍的速度提升。這種優化不僅加快了推理速度，還使得模型能夠更高效地處理復雜的長文本任務。

3）動態負載均衡調度

FlashInfer的調度器能夠根據輸入的變化動態調整，減少GPU的空閑時間，確保資源的高效利用。其與CUDA圖的兼容性進一步增強了其在生產環境中的適用性。例如，在處理多個并發請求時，FlashInfer能夠根據每個請求的計算需求和優先級，動態分配GPU資源，避免資源浪費和任務阻塞。

4）可定制的JIT編譯

FlashInfer允許用戶自定義并編譯特定的注意力變體，生成高性能的內核。這一功能為滿足特殊用例需求提供了極大的便利，如滑動窗口注意力或RoPE變換等。用戶可以根據自己的應用場景和需求，靈活地定制注意力機制，從而進一步提升推理性能。

性能提升實例

FlashInfer在多個基準測試中都展現出了顯著的性能提升：

• 延遲降低：與現有的解決方案如Triton相比，FlashInfer將令牌間延遲降低了29%至69%。在涉及長上下文推理和平行生成的場景中，這種延遲降低尤為明顯。例如，在長文本生成任務中，FlashInfer能夠更快地生成每個令牌，從而縮短整體推理時間。
• 吞吐量提升：在NVIDIA H100 GPU上，FlashInfer在平行生成任務中實現了13%至17%的速度提升。這表明FlashInfer在處理高需求應用時具有更高的效率，能夠更快地完成大規模的推理任務。
• GPU利用率增強：FlashInfer的動態調度器和優化內核提高了帶寬和FLOP利用率，尤其在處理傾斜或均勻序列長度的場景中表現突出。這意味著在有限的GPU資源下，FlashInfer能夠更充分地發揮其計算能力，提升整體性能。
• 并行解碼任務優勢：FlashInfer在并行解碼任務中也表現出色，其可組合的格式使得首次令牌時間（TTFT）大幅降低。以Llama 3.1模型（70B參數）為例，在特定配置下，FlashInfer將TTFT降低了22.86%。這使得在需要快速響應的應用場景中，如實時翻譯或語音識別，FlashInfer能夠更快地提供結果。

04、結語

FlashInfer為大模型推理提供了一個實用且高效的解決方案，顯著提升了性能和資源利用率。其靈活的設計和強大的集成能力，使其成為推動大模型服務框架發展的重要工具。通過解決關鍵的效率問題并提供強大的技術解決方案，FlashInfer為更易用、更可擴展的AI應用鋪平了道路。作為一個開源項目，FlashInfer還期待與研究社區的進一步合作與創新，以應對AI基礎設施中不斷出現的新挑戰，持續推動人工智能技術的進步。

參考：

1. ??https://arxiv.org/abs/2501.01005??
2. ??https://github.com/flashinfer-ai/flashinfer??

本文轉載自公眾號Halo咯咯 作者：基咯咯

原文鏈接：??https://mp.weixin.qq.com/s/LPKvpSs83691Ih_7TNNNMQ??