GPU推理提速4倍，256K上下文全球最長：無問芯穹刷新大模型優(yōu)化記錄

　　本文來自于微信公眾號(hào) 機(jī)器之心(ID:almosthuman2014)，作者:機(jī)器之心。

　　想用大模型賺錢?這個(gè)實(shí)力強(qiáng)勁的新面孔決定先把推理成本打下來。

　　大模型業(yè)務(wù)到底多燒錢?前段時(shí)間，華爾街日報(bào)的一則報(bào)道給出了參考答案。

　　報(bào)道顯示，微軟的 GitHub Copilot 業(yè)務(wù)(背后由 OpenAI 的 GPT 大模型支撐)雖然每月收費(fèi)10美元，但平均還是要為每個(gè)用戶倒貼20美元。可見當(dāng)前 AI 服務(wù)提供商們正面臨著嚴(yán)峻的經(jīng)濟(jì)賬挑戰(zhàn) —— 這些服務(wù)不僅構(gòu)建成本高昂，運(yùn)營成本也非常高。

　　有人比喻說:「使用 AI 總結(jié)電子郵件，就像是讓蘭博基尼送披薩外賣。」

　　對此，OpenAI 算過一筆更詳細(xì)的賬:當(dāng)上下文長度為8K 時(shí)，每1K 輸入 token 的成本為3美分，輸出的成本為6美分。目前，OpenAI 擁有1.8億用戶，每天收到的查詢數(shù)量超過1000萬次。這樣算來，為了運(yùn)營 ChatGPT 這樣的模型，OpenAI 每天都需要在必要的計(jì)算硬件上投入大約700萬美元，可以說是貴得嚇人。

　　降低 LLM 的推理成本勢在必行，而提升推理速度成為一條行之有效的關(guān)鍵路徑。

　　實(shí)際上，研究社區(qū)已經(jīng)提出了不少用于加速 LLM 推理任務(wù)的技術(shù)，包括 DeepSpeed、FlexGen、vLLM、OpenPPL、FlashDecoding 和 TensorRT-LLM 等。這些技術(shù)自然也各有優(yōu)勢和短板。其中，F(xiàn)lashDecoding是 FlashAttention 作者、斯坦福大學(xué)團(tuán)隊(duì)的 Tri Dao 等人在上個(gè)月提出的一種 state-of-the-art 方法，它通過并行加載數(shù)據(jù)，大幅提升了 LLM 的推理速度，被認(rèn)為極具潛力。但與此同時(shí)，它也引入了一些不必要的計(jì)算開銷，因此依然存在很大的優(yōu)化空間。

　　為了進(jìn)一步解決問題，近日，來自無問芯穹(Infinigence-AI)、清華大學(xué)和上海交通大學(xué)的聯(lián)合團(tuán)隊(duì)提出了一種新方法 FlashDecoding++，不僅能帶來比之前方法更強(qiáng)的加速能力(可以將 GPU 推理提速2-4倍)，更重要的是還同時(shí)支持 NVIDIA 和 AMD 的 GPU!它的核心思想是通過異步方法實(shí)現(xiàn)注意力計(jì)算的真正并行，并針對「矮胖」矩陣乘優(yōu)化加速 Decode 階段的計(jì)算。

　　將 GPU 推理提速2-4倍，

　　FlashDecoding++ 是怎么做到的?

　　LLM 推理任務(wù)一般為輸入一段文字(token)，通過 LLM 模型計(jì)算繼續(xù)生成文字或其他形式的內(nèi)容。

　　LLM 的推理計(jì)算可被分為 Prefill 和 Decode 兩個(gè)階段，其中 Prefill 階段通過理解輸入文字，生成第一個(gè) token;Decode 階段則順序輸出后續(xù) token。在兩個(gè)階段，LLM 推理的計(jì)算可被分為注意力計(jì)算和矩陣乘計(jì)算兩個(gè)主要部分。

　　對于注意力計(jì)算，現(xiàn)有工作如 FlashDecoding 切分注意力計(jì)算中的 softmax 算子實(shí)現(xiàn)并行加載數(shù)據(jù)。這一方法由于需要在不同部分 softmax 同步最大值，在注意力計(jì)算中引入了20% 的計(jì)算開銷。而對于矩陣乘計(jì)算，在 Decode 階段，左乘矩陣多表現(xiàn)為「矮胖」矩陣，即其行數(shù)一般不大(如 <=8)，現(xiàn)有 LLM 推理引擎通過補(bǔ)0將行數(shù)擴(kuò)充到64從而利用 Tensor Core 等架構(gòu)加速，從而導(dǎo)致大量的無效計(jì)算(乘0)。

　　為解決上述問題，「FlashDecoding++」的核心思想在于，通過異步方法實(shí)現(xiàn)注意力計(jì)算的真正并行，并針對「矮胖」矩陣乘優(yōu)化加速 Decode 階段的計(jì)算。

　　異步并行部分 softmax 計(jì)算

　　先前工作對每個(gè)部分 softmax 計(jì)算求輸入最大值作為縮放系數(shù)，避免 softmax 計(jì)算中 e 指數(shù)的溢出，這就導(dǎo)致了不同部分 softmax 計(jì)算的同步開銷(圖1(a)(b))。

　　「FlashDecoding++」指出，對于大部分 LLM，其 softmax 的輸入分布較為集中。如圖2所示，Llama2-7B 的 softmax 輸入99.99% 以上集中在 [-16.8，6.5] 這個(gè)區(qū)間。因此，「FlashDecoding++」提出在部分 softmax 計(jì)算時(shí)使用一個(gè)固定的最大值(圖1(c))，從而避免了不同部分 softmax 計(jì)算間的頻繁同步。而當(dāng)小概率發(fā)生的輸入超出給定范圍時(shí)，「FlashDecoding++」對這一部分的 softmax 計(jì)算退化為原先的計(jì)算方法。

　　「矮胖」矩陣乘的優(yōu)化

　　由于 Decode 階段的輸入為一個(gè)或幾個(gè) token 向量，因此該階段的矩陣乘表現(xiàn)為「矮胖」形狀。以矩陣 A×B=C 為例，A 與 B 矩陣的形狀為 M×K 與 K×N，「矮胖」矩陣乘即 M 較小的情況�！窮lashDecoding++」指出「矮胖」矩陣乘一般緩存受限，并提出雙緩存機(jī)制等優(yōu)化手段進(jìn)行加速(圖3)。

　　此外，「FlashDecoding++」進(jìn)一步指出，在 LLM 推理階段，針對特定模型，N 和 K 的取值固定。因此，「FlashDecoding++」會(huì)根據(jù) M 的大小，自適應(yīng)選取矩陣乘的最優(yōu)實(shí)現(xiàn)。

　　將 GPU 推理提速2-4倍

　　目前，「FlashDecoding++」可以實(shí)現(xiàn) NVIDIA 與 AMD 等多款 GPU 后端的 LLM 推理加速(圖5)。通過加速 Prefill 階段的首 token 生成速度，以及 Decode 階段每個(gè) token 的生成速度，「FlashDecoding++」可以在長、短文本的生成上均取得加速效果。相較于 FlashDecoding，「FlashDecoding++」在 NVIDIA A100上的推理平均加速37%，并在 NVIDIA 和 AMD 的多 GPU 后端上相較于 Hugging Face 實(shí)現(xiàn)加速多達(dá)2-4倍。

　　AI 大模型創(chuàng)業(yè)新秀:無問芯穹

　　該研究的三位共同一作分別是無問芯穹首席科學(xué)家、上海交通大學(xué)副教授戴國浩博士，無問芯穹研究實(shí)習(xí)生、清華大學(xué)碩士生洪可，無問芯穹研究實(shí)習(xí)生、上海交通大學(xué)博士生許珈銘。通訊作者為上海交通大學(xué)戴國浩教授和清華大學(xué)電子工程系主任汪玉教授。

　　創(chuàng)立于2023年5月的無問芯穹，目標(biāo)是打造大模型軟硬件一體化最佳解決方案，目前 FlashDecoding++ 已被集成于無問芯穹的大模型計(jì)算引擎「Infini-ACC」中。在「Infini-ACC」的支持下，無問芯穹正在開發(fā)一系列大模型軟硬件一體化的解決方案，其中包含大模型「無穹天權(quán)(Infini-Megrez)」、軟硬件一體機(jī)等。

　　據(jù)了解，「Infini-Megrez」在處理長文本方面表現(xiàn)非常出色，將可處理的文本長度破紀(jì)錄地提升到了256k token，實(shí)測處理大約40萬字的一整本《三體3:死神永生》也不成問題。這是當(dāng)前的大模型所能處理的最長文本長度。

　　此外，「Infini-Megrez」大模型在 CEval (中)、MMLU (英)、CMMLU (中)、AGIEval 等數(shù)據(jù)集上均取得了第一梯隊(duì)算法性能，并依托「Infini-ACC」計(jì)算引擎持續(xù)進(jìn)化中。