01、概述

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

02、大模型推理的困境

注意力機制的局限

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

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

高延遲與內存瓶頸

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

03、FlashInfer的破局之道

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

技術特點與優(yōu)勢

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

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

2）優(yōu)化的共享前綴解碼

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

3）動態(tài)負載均衡調度

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

4）可定制的JIT編譯

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

性能提升實例

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

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

04、結語

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

參考：

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

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

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