原標(biāo)題：高效Attention引擎是怎樣煉成的？陳天奇團(tuán)隊(duì)FlashInfer打響新年第一槍！
文章來(lái)源：新智元
內(nèi)容字?jǐn)?shù)：6055字

FlashInfer：高效LLM推理引擎

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

1. LLM推理的挑戰(zhàn)

   傳統(tǒng)的Transformer模型在實(shí)際應(yīng)用中面臨諸多工程挑戰(zhàn)，例如多樣化的工作負(fù)載、個(gè)性化的Attention實(shí)現(xiàn)以及對(duì)延遲和吞吐量的嚴(yán)格要求。不同的推理階段（Prefill、Decode、多輪對(duì)話）和技術(shù)（Speculative Decoding）都帶來(lái)了不同的計(jì)算訪存比，影響GPU利用率。此外，各種LLM框架（如vLLM、SGLang）在KV cache存儲(chǔ)方面也存在差異，增加了框架的復(fù)雜性。

2. 塊稀疏格式：高效存儲(chǔ)KV Cache

   FlashInfer提出使用塊稀疏（BSR）矩陣來(lái)統(tǒng)一存儲(chǔ)各種形式的KV cache。BSR類似于CSR，但操作單元由單個(gè)元素變?yōu)閿?shù)據(jù)塊，從而提高GPU的利用率。對(duì)于Paged Attention、Radix Tree和Token Tree等不同形式的KV cache，BSR都能有效地減少冗余，并最大化Tensor Core的利用率。通過將稀疏塊數(shù)據(jù)加載到共享內(nèi)存并排布成密集格式，F(xiàn)lashInfer消除了硬件浪費(fèi)。

3. 塊并行與可組合性

   FlashInfer采用與BPT相同的塊并行方法分解KV cache，實(shí)現(xiàn)并行計(jì)算，并通過可組合特性支持不同大小塊的組合，在共享和全局內(nèi)存之間取得平衡，進(jìn)一步優(yōu)化性能。

4. 可定制的注意力機(jī)制

   FlashInfer為FlashAttention開發(fā)了CUDA/CUTLASS模板，支持不同架構(gòu)的Nvidia GPU，并提供可自定義的Attention機(jī)制。開發(fā)者可以通過插入自定義函數(shù)來(lái)實(shí)現(xiàn)個(gè)性化的Attention計(jì)算，例如ALiBi、Logits SoftCap、RoPE等，提高了框架的靈活性。

5. 高效的數(shù)據(jù)移動(dòng)和負(fù)載均衡

   FlashInfer采用高效的數(shù)據(jù)加載方法，將tiles從全局內(nèi)存轉(zhuǎn)移到共享內(nèi)存，并利用異步復(fù)制指令最大化內(nèi)存帶寬。同時(shí)，它根據(jù)硬件資源和工作負(fù)載選擇合適的塊大小，并通過負(fù)載均衡調(diào)度算法，根據(jù)用戶請(qǐng)求動(dòng)態(tài)調(diào)整資源分配，并兼容CUDAGraph。

6. 性能提升

   實(shí)測(cè)結(jié)果表明，F(xiàn)lashInfer顯著提升了LLM推理性能，token間延遲降低了29%-69%，長(zhǎng)上下文推理延遲降低了28%-30%，并行生成速度提高了13%-17%。

總而言之，F(xiàn)lashInfer通過巧妙地結(jié)合塊稀疏格式、塊并行計(jì)算、可組合性和負(fù)載均衡等技術(shù)，為構(gòu)建高效且可定制的LLM推理引擎提供了一種新的解決方案，顯著提升了LLM的推理效率。

聯(lián)系作者

文章來(lái)源：新智元
作者微信：
作者簡(jiǎn)介：智能+中國(guó)主平臺(tái)，致力于推動(dòng)中國(guó)從互聯(lián)網(wǎng)+邁向智能+新紀(jì)元。重點(diǎn)關(guān)注人工智能、機(jī)器人等前沿領(lǐng)域發(fā)展，關(guān)注人機(jī)融合、人工智能和機(jī)器人對(duì)人類社會(huì)與文明進(jìn)化的影響，領(lǐng)航中國(guó)新智能時(shí)代。