大語言模型推理優化技術綜述(The Art of LLM Inference) 原創
編者按: 如何將 LLM 的推理過程從“燒錢的無底洞”轉變為“高性能的生產力引擎”?本文深入剖析了提升 LLM 推理效率的五大核心技術:巧妙的 KV 緩存管理、Query-sparsity attention(QUEST)、推測解碼(使用 draft model 加速生成過程)、權重調度(通過 Flexgen 實現跨設備資源分配)以及系統級優化(FastServe 解決隊頭阻塞問題)。此外,作者還簡要介紹了多種常見優化方向,包括量化技術(AWQ、INT8、FP8 等)、Early Exit Inference(LITE)、注意力機制優化(FlashAttention、ROFormer)以及使用非自回歸 LLM 等。
作者 | Trung Thanh Tran
編譯 | 岳揚
如果您正在開發 AI 解決方案,并托管基于大語言模型(LLMs)的基礎模型,那么您就應該關注模型服務的成本。然而,資金并非唯一的考量因素。請相信,如果無法解決模型性能的難題,即便預算充足,LLMs serving 的實際效果仍會大打折扣。本文將探討如何將 LLM 的推理過程從「燒錢的無底洞」轉變為「高性能的生產力引擎」。
目錄
01 LLMs serving 面臨的一些挑戰
02 主題 1:巧妙的 KV 緩存管理
03 主題 2:Query-sparsity attention
04 主題 3:推測解碼
05 主題 4:權重調度
06 主題 5:系統級優化
07 其他主題
08 如何應用這些技術
01 LLMs serving 面臨的一些挑戰
LLMs 非常強大,但它們的特性使其難以高效服務。LLM 的推理過程包含兩個階段:
1) 預填充階段:當你輸入提示詞(上下文、對話歷史、問題等)時,模型會一次性處理所有 token。
2) 解碼階段:在初始的提示詞后,模型逐 token 生成內容,每個新 token 依賴于之前生成的 token。
舉一個易懂的類比:預填充階段如同下棋時擺棋盤(耗時較長),而解碼階段則像擺好棋后逐步下棋(每一步都很快)。
然而,LLMs serving(譯者注:將訓練好的大語言模型部署到實際應用中,以低延遲、高吞吐、資源高效的方式處理用戶請求的技術過程。) 并非輕而易舉,必須考慮以下問題:
Sparsity
在神經網絡(尤其是 FFN 模塊)中,大量神經元的激活值為零。跳過這些零激活值的神經元、僅計算非零元素可以大大節省運算時間。
LLM 中大量神經元的激活值為零,導致矩陣運算中存在大量零值。圖片來源[1]
內存帶寬限制與內存瓶頸
在 GPU 上傳輸數據往往超過數據計算的耗時。此外,大型模型(例如傳聞參數量達萬億的 ChatGPT)無法單卡裝載。
將當前最先進 LLM 的內存需求,與 GPU 的顯存容量進行對比。圖片來源:ChatGPT
低效調度——先到先得
LLM 通常需要同時處理多個請求。這會導致短請求(例如詢問天氣、時間或簡短的回答)被迫等待長請求完成。那么,平均響應時間幾乎完全由等待時間主導,而非實際計算時間。
你更快,但必須等待之前的請求先處理完。圖片來源:ChatGPT
Sequential Decoding(按順序進行解碼)
生成 token 時無法輕松實現并行處理。每次前向傳播只能產生一個 token(或一個小 batch)。當我們向 ChatGPT 請求長回復時,輸出內容往往是逐詞生成的。這就是為什么“流式輸出”(streaming output)的用戶體驗并不比等待完整答案一次性輸出更差。
逐步進行解碼。圖片來源:ChatGPT
KV Cache 增長
注意力機制需對整個序列的所有文本進行計算,這是 LLM 的推理過程中最核心且最耗時的操作。有趣的是,每當序列中生成新 token 時,系統會對過去的 token 重復大量相同的計算。鍵值緩存(KV Cache)技術通過存儲前幾步的關鍵信息來加速此過程(使用 KV Cache 可使 T4 GPU 上的 GPT2 推理速度提升 5 倍)。下圖展示了使用緩存與不使用的區別,但使用緩存也會額外占用內存。
解碼序列 [token 1, token 2, token 3, token 4] 時的 KV Cache 操作步驟。圖片來源[2]
實驗表明,KV(Key-Value)緩存的使用率在 20.4% 到 38.2% 之間。我用 Qwen-VL 2.0 模型對約 1 萬張圖片生成簡短描述(要求回答少于20字),發現速度比未使用 KV 緩存的版本快 20%。
這些特性看似棘手,但通過巧妙的工程化手段,反而能轉化為優勢。
02 主題 1:巧妙的 KV 緩存管理
Page attention
KV 緩存會占用大量內存。上下文越長,KV 緩存占用的內存越大。 例如,若某 LLM 的輸入長度為 2048 個 token,則需預留 2048 個詞槽(slots)。下圖說明了我提到的情況。
在圖中,2048 個詞槽被一個包含 7 個單詞的提示詞(“four, score, and, seven, years, ago, our”)占用了,后續生成的 4 個單詞(“fathers, brought, forth, ”)占用了第 8-11 個詞槽。這意味著仍有 2038 個詞槽被保留,但從未被使用過,這就產生了內存的內部碎片(internal fragmentation)。
每個推理步驟都會生成鍵值對(KV pairs),在使用注意力機制時必須緩存這些數據。KV 緩存通常以連續的塊(chunks)或頁(pages)的形式分配在內存中。當序列生成完成并釋放內存頁后,已釋放的頁可能不再連續。后續序列所需的內存大小可能無法恰好匹配現有空閑塊,導致內存中散布小型的空閑塊——即外部碎片(external fragmentation)。
受操作系統的內存管理啟發,Page Attention 機制也將數據組織為邏輯內存塊(logical memory block),并通過頁表(page table)進行監控,再將合適的頁(page)映射到物理內存中。具體實現如下:
1) Fixed-size Blocks(固定大小的內存塊) :PagedAttention 分配固定大小且相對較小的內存塊(稱為“頁(pages)”)來存儲 KV 緩存。
2) Shared Blocks(共享內存塊) :這些固定大小的內存塊可在不同請求間共享。
3) On-demand Allocation(按需進行分配) :隨著生成過程逐步分配內存塊,無需根據最大序列長度的估算預先分配。
LLM 中的分頁機制示意圖。Image by the author
支持多請求間共享內存塊的 LLM 分頁機制示意圖。Image by the author
Raddix tree KV cache
在計算機科學中,基數樹(radix tree,亦稱 radix trie、compact prefix tree 或 compressed trie)是一種優化了空間效率的字典樹(前綴樹),其將每個唯一的子節點與其父節點合并。
Raddix tree KV cache 是一種支持跨不同推理請求高效復用鍵值(KV)緩存的技術,尤其適用于多個請求共享共同前綴的場景。 通過將 KV 緩存組織為 Raddix 樹結構,可高效檢索緩存數據并在請求間共享。在下面的例子中,三個請求共享相同的前綴 "ABC"(存儲于父節點中),每個請求中的最后一個單詞則分別存儲在三個葉子節點。需注意:樹結構的運行時間復雜度為 O(nlogn),遠低于注意力計算的 O(n2)。
Raddix tree KV cache 示例
Compressed attention
多頭注意力機制(Multi-head Attention)[3]是 Transformer 模型(LLMs 的基石)的核心機制。每個注意力頭從不同視角分析文本:其中一個注意力頭關注主謂關系,另一個注意力頭解析詞匯特征,第三個注意力頭分析句子結構。這種多頭機制雖增強了模型的理解能力,但也導致每個注意力頭需要獨立的 KV 對。在實時文本處理或長序列場景中,這些獨立的 Key 和 Value 會占用大量內存。
分組查詢注意力機制(Group Query Attention, GQA)允許多個查詢(queries)共享同一組 Key 和 Value,從而減少所需 KV 對數量。多查詢注意力機制(Multi Query Attention, MQA) 則更為激進,僅用一組 KV 對服務所有查詢(queries)。
多頭注意力、多查詢注意力、分組查詢注意力對比圖
中國 AI 初創公司深度求索(DeepSeek)今年初發布了其 chatbot。該產品以高效、開源著稱,人們傳言他們的成功源于對 ChatGPT 生成數據的分析工作。然而,閱讀了他們的技術報告后,我發現其技術突破不僅僅局限于數據提取操作。DeepSeek 提出的 Flash Multi Latent Attention(Flash MLA) 通過低秩壓縮將 Key 和 Value 向下投影到更小維度的 latent vector 中,大幅減小了緩存體積。計算注意力時再將 latent vector 向上投影,且上投影矩陣權重與查詢矩陣權重"折疊"融合,進一步加速了注意力的計算。
多頭潛在注意力機制(MLA)示意圖。Image by the author
03 主題 2:Query-sparsity attention
QUEST: Query-Aware Sparsity for Efficient Long-Context LLM Inference
從 MIT 研究人員撰寫的論文《QUEST: Query-Aware Sparsity for Efficient Long-Context LLM Inference》[4]中,我們得知 Transformer 模型在推理過程中(尤其是注意力計算環節)常存在高稀疏性「譯者注:high sparsity,大部分神經元或注意力權重在計算過程中未被激活(值為零或接近零)」。這意味著大模型中并非全部的神經節點被激活。通過將高稀疏性(high sparsity)的特性應用于剪枝機制(pruning mechanism),我們能夠發現一種高效運行大模型的方法。下圖展示了 Transformer 模型各層的稀疏性統計數據。遺憾的是,第 3 層之后的模型層通常非常稀疏。極端情況下,某些模型層(如第 10 層)甚至達到了 100% 稀疏,這種現象等同于在運行大語言模型時多次乘以 0,從而產生零值輸出。
出現這種現象的原因很簡單:并非每個單詞都對當前上下文有貢獻。
例如,給定提示詞:"A is B. C is D. A is",模型應生成 "B"。這意味著只需要最關鍵的 token,而這很大程度上取決于查詢(queries)。因此該技術被命名為查詢感知的稀疏性算法(query-aware sparsity)。
Transformer 模型推理中的稀疏性估算。圖片來源[5]
了解這一特性后,QUEST 的核心策略就是定位對注意力計算最關鍵的數據塊。QUEST 將找出前 K 個數據塊。其算法流程直觀清晰(見下圖):
QUEST 獲取 top K 個關鍵數據塊進行注意力計算的流程
首先,對于每個數據塊,QUEST 會找出最小和最大鍵值(minimum and maximum keys)及其通道尺度上的數值(channel-wise values)。接著,query 會逐個元素地生成最大和最小鍵值。這種方法能夠大大減少所需的計算量 —— 即使 query 的符號(sign)變化,后續的乘積運算通常仍能得到最大值:當 query 符號為負時,乘以最小值必然得到最大輸出值,反之亦然。在獲取每個數據塊的最大值后,QUEST 僅篩選出與 query 最相關的 K 個關鍵 KV 塊。通過這一流程,計算量得以大幅降低。
最后一個關鍵問題是選擇恰當的 K 值,以避免模型性能下降。K 是一個需要通過實驗才能確定的超參數(hyperparameter)。在論文中,作者建議選擇 K=4096,可使模型性能保持在接近 100% 的水平。
以下是 K=4096 時的數據:
- PG19(一種教科書數據集)上的準確率 ≈ 完全達到了全局注意力(Full Attention)的基準準確率
- passkey retrieval 數據集上準確率 ≈ 100%
- LongBench 任務上的準確率 ≈ 在多數數據集上等效于全緩存(full cache)
04 主題 3:推測解碼
推測解碼(speculative decoding)對于大語言模型(LLM)推理的加速非常重要,該技術由 Andrej Karpathy[6] 提出,并由 Google 在 2022 年首次引入[7]。
該技術的核心思想非常簡單:
與其僅用龐大、緩慢但精確的模型(稱為 target model)逐詞生成,不如先用小型、快速但不太準確的模型(通常稱為 draft model)快速"推測"后續的多個 token。然后用大型模型驗證這些猜測的 token。 若大型模型認同小型模型的預測,則一次性接受所有結果(減少計算量)。如若不一致,就從分歧點開始回退(重新執行)。示意圖如下文所示。
draft model 可以是 Ngrams、1B 參數級別的模型,最高可達 3B 參數級別的模型。target model 則可以是數十億甚至數萬億參數規模的模型。
雖然使用兩個模型會消耗較多內存,且重復生成過程也比較耗時,但該技術的核心價值在于其卓越的實用性 —— 連 Gemini 這樣的頂級模型都已采用該技術(如下圖所示)。實際情況是,draft model 生成的 token 通常正確率很高,以至于 target model 無需修正結果。這是因為在現實語境中常見詞匯如"yes, this, is, and so on"出現頻率極高,即使小型語言模型也能輕松預測。通過并行驗證 draft model 生成的所有 token,而非逐詞進行自回歸解碼,這樣可以節省大量時間。
05 主題 4:權重調度
調度(Scheduling)的核心在于將模型權重平衡分配到物理機器的資源(包括 GPU、CPU 和硬盤)中。這種策略不僅能通過并行計算加快推理速度,還能讓 100B 參數級別的超大型模型在僅配備有 T4 GPU 的低配置 PC 上運行。
實現這一目標的關鍵在于兩個核心要素:
- 在 GPU、CPU 和硬盤間智能地加載/卸載模型權重
- 高效處理計算單元間的 I/O 數據傳輸
Flexgen
由 Stanford、UC Berkeley 和 CMU 聯合提出的 Flexgen[8],正是解決這兩個關鍵問題最具創新性的方案之一。
推理過程通常如下圖所示。需要處理的每個數據塊被定義為加載到模型層的一批數據,其中列方向按批次處理的,行方向則按照模型層維度處理。
我們定義有效路徑為滿足以下約束條件的遍歷(即計算)所有方格的路徑:
- 從左向右執行
- 所有數據必須位于同一設備
- 當前方塊的激活值(Activation) 需要等待其右側相鄰方塊完成計算后,才能被釋放或復用
- KV Cache需存儲至最右側數據計算完成時釋放
- 任意時刻設備上存儲的張量總大小不得超過該設備內存容量
如果我們有 n 個 token,每個 token 的數據將按順序加載和計算。每一層的權重僅在需要計算時加載,并在計算結束后立即卸載。由于 GPU 的計算速度極快(如閃電),而內存傳輸速度極慢(如蝸牛),頻繁的加載/卸載會帶來巨大時間開銷。
圖中每個方塊表示 GPU 對一個模型層(layer)的批處理計算(batch computation),同色方塊共享同一層的權重參數
Flexgen 通過將行掃描改為列掃描或之字形塊調度(zig-zag block schedule)進行優化:在無 I/O 開銷的情況下保留模型層權重,并為下一列保存激活值。在計算塊執行期間,Flexgen 實現三大操作的并行執行:加載下一層權重、存儲前一數據批次的激活值/KV 緩存、執行當前數據批次的計算,從而有效解決內存傳輸問題。
Flexgen 的另一個核心創新在于模型權重的硬件分布策略。
Flexgen 采用線性規劃策略搜索法(Linear Programming Policy Search)來尋找最優加載配置,使模型整體推理時間最小化。
其中:
- n:每個序列輸出 token 的數量
- ??:transformer 層數
- block size:每個計算塊處理的樣本量(數據批次大小 × 數據批次數量)
下圖展示了 Flexgen 在搭載了 T4 GPU 的機器上運行 OPT-30B 模型的配置示例。
論文對比了 HuggingFace 的 DeepSpeed 庫和 Accelerate 庫的性能:Flexgen 聲稱能達到 7.32 tokens/秒,而 DeepSpeed 為 1.57 tokens/秒,Accelerate 僅 0.62 tokens/秒。
06 主題 5:系統級優化
現有 LLM serving 系統(如 vLLM、Orca)通常采用先到先服務(FCFS,First-Come-First-Serve)機制和執行到完成為止(run-to-completion)的運行方式,這會導致隊頭出現阻塞(HOL) —— 簡單來說:長任務會延遲短任務的處理。這就造成了較高的排隊延遲,在實際工作負載中可達總延遲的 90%。 請看論文 FastServe[9] 中的統計數據:
注意:當我們提及長請求/短請求時,并非指提示詞(prompt)的長度,而是指生成首個 token 所需的時間。
time by execuation vs queuing. Image source[9]
該問題的解決方案是:允許中斷長請求,將已完成的部分存入緩存,保留未完成的部分稍后處理,然后切換至短請求。待短請求完成后,繼續運行長請求的剩余部分。此方案需實現具有多個不同優先級的隊列(multi-queue)。
但這一想法仍然存在一個問題:若高優先級隊列中存在多個長請求位于短請求之前,可能導致長請求被多次中斷才切換到短請求。這不僅增加長請求的處理時間,還會對緩存造成額外壓力。
FastServe 通過引入跳轉式多級反饋隊列(Skip-Join MLFQ)來解決這個問題:當系統收到請求時,會預估生成首個 token 所需的時間,據此將請求路由至合適優先級的隊列,避免干擾更短請求。此外,通過 KV 緩存管理,可在 GPU 處理先前的隊列時主動在隊列間遷移數據,進一步降低延遲。
07 其他主題
還有很多其他優化 LLM 推理的技術方向,本文不再詳述(因為這些技術方向非常常見,許多工程師每天都在使用),僅列舉技術方向與參考文獻:
量化(Quantization)
通過降低權重和激活值的精度(如從 FP16 降至 INT4 或 FP8)壓縮模型體積并提升推理速度,同時將精度損失降到最小。
- AWQ
通過激活值計算每個權重/通道的重要性分數進行基于激活值感知的權重量化。支持低比特推理(如 INT3)且無需重訓練。
- LLM.int8()
提出通過校準過程(Calibration)實現 Post-training INT8 矩陣乘法,支持 Transformer 推理且不會降低精度。
- SmoothQuant
通過跨層對齊激活值與權重的數值范圍,提升 post-training 量化效果。
- ZeroQuant / V2 / FP
使用校準過程(Calibration)和低秩補償(Low-rank Compensation)的低比特量化技術(INT4, FP4)。
- LLM-FP4
證明 FP4 表示法可在顯著提升推理速度的同時保持模型質量。
- WINT8
針對生產環境中的 MoE 模型的 INT8 量化方案。
- SpQR
將模型量化與 sparsity 相結合,實現近乎無損的 LLM 壓縮,適用于邊緣部署。
- FP8-LM
使用 FP8 格式訓練 Transformer,降低訓練與推理過程的內存和算力消耗。
- FP8 Formats
定義 NVIDIA 的 FP8 格式及其在深度學習推理/訓練中的應用。
Early Exit Inference
- LITE
在神經網絡的中間層添加預測能力,當置信度較高時,token 會提前退出,最高可節省 38% FLOPS。
注意力機制優化(Attention Optimization)
- FlashAttention 1, 2, 3
通過內存分塊實現快速、精確的注意力計算,速度與內存效率優于標準實現方式。
- ROFormer
引入旋轉位置編碼(Rotary Position Embedding),提升模型對長程依賴關系的泛化能力。
- StreamLLM
允許注意力在流式處理過程中動態適配新的輸入塊。
非自回歸的 LLM(Non-autoregressive LLMs)
- Diffusion-LM: Improving Controllable Text Generation
將擴散模型應用于文本生成的首個重要工作。
08 如何應用這些技術
vLLM[10] 是一個開源庫,能夠大大提升大語言模型(LLM)推理(運行模型)的速度和效率。
它由 UC Berkeley 的研究者開發,專注于實現 LLM 的高吞吐、低延遲服務。該庫最初基于 PageAttention 的核心思想,但如今已整合了前文提及的絕大多數優化技術。在我看來,vLLM 已成為大模型推理優化領域生態最活躍的開源社區之一。
以下是我使用 vLLM 調用 QwenVL 2.5 7B instruct 模型對圖片進行描述的示例代碼:
感謝你閱讀本文!這是我以 Fatima Fellowship 身份開展的研究工作,我與 Colorado School of Mines[11] 的 PhD candidate Ismet Dagl 博士合作,專注于提升邊緣設備上大語言模型(LLM)、視覺語言模型(LVM)及基礎模型的性能與內存優化。
About the author
Trung Thanh Tran
CTO of ClientScan | Co-founder of ??Takenote.ai?? | AI Researcher | | Data Scientist at Pixta Vietnam | IGI AI Book Writer
END
本期互動內容 ??
?在您的工作中,最希望出現哪種“開箱即用”的推理優化工具?歡迎在評論區分享~
文中鏈接
[1]??https://developer.nvidia.com/blog/accelerating-inference-with-sparsity-using-ampere-and-tensorrt/??
[3]??https://d2l.ai/chapter_attention-mechanisms-and-transformers/multihead-attention.html??
[4]??https://arxiv.org/abs/2406.10774??
[5]??https://medium.com/r?url=https%3A%2F%2Farxiv.org%2Fabs%2F2406.10774??
[6]??https://x.com/karpathy/status/1697318534555336961?lang=en??
[7]??https://arxiv.org/pdf/2211.17192??
[8]??https://arxiv.org/pdf/2303.06865??
[9]??https://arxiv.org/pdf/2305.05920??
[10]??https://docs.vllm.ai/en/latest/index.html??
[11]??https://cs.mines.edu/??
本文經原作者授權,由 Baihai IDP 編譯。如需轉載譯文,請聯系獲取授權。
原文鏈接:
??https://blog.gopenai.com/the-art-of-llm-inference-fast-fit-and-free-c9faf1190d78??
