01、概述

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

02、大模型推理的困境

注意力機(jī)制的局限

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

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

高延遲與內(nèi)存瓶頸

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

03、FlashInfer的破局之道

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

技術(shù)特點(diǎn)與優(yōu)勢

1）全面的注意力內(nèi)核支持

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

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

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

3）動(dòng)態(tài)負(fù)載均衡調(diào)度

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

4）可定制的JIT編譯

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

性能提升實(shí)例

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

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

04、結(jié)語

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

參考：

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

本文轉(zhuǎn)載自公眾號Halo咯咯 作者：基咯咯

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