探秘Transformer系列之（24）--- KV Cache優(yōu)化
探秘Transformer系列之（24）--- KV Cache優(yōu)化
0x00 前言
目前的大型語言模型（LLM）服務(wù)系統(tǒng)采用KV Cache來避免在解碼階段重復(fù)計(jì)算鍵和值的投影。雖然這對于單個(gè)客戶端請求生成短序列而言是一個(gè)有效的解決方案，但是，面對多個(gè)客戶時(shí)，每個(gè)請求都保留自己的KV緩存，從而增加了推理過程中的總體KV緩存大小。另外，即使是針對單個(gè)客戶端請求，當(dāng)我們生成長序列或處理多輪對話時(shí)，KV Cache 依然會對推理性能造成極大的影響。比如，束搜索和平行采樣也被廣泛用于生成更好的輸出或?yàn)榭蛻籼峁┖蜻x選擇。這些技術(shù)也會像批處理推理一樣增加KV緩存的大小，因?yàn)樗鼈儠瑫r(shí)處理多個(gè)序列。
KV Cache問題的核心在于，KV Cache 占用了大量內(nèi)存和訪存帶寬，同時(shí)在生成階段引入了大量重復(fù)計(jì)算。從而阻止了我們處理或生成非常長的序列和處理大批量數(shù)據(jù)，無法最大化硬件效率。這就好比一個(gè)圖書館管理員每天忙于檢索每本書的信息，卻沒有時(shí)間去為用戶去拿書。而總體趨勢上，LLM 的窗口長度也在不斷增大，因此就出現(xiàn)一組主要矛盾，即：對不斷增長的 LLM 的窗口長度的需要與有限的 GPU 顯存之間的矛盾。
因此優(yōu)化 KV cache 就顯得非常必要，KV Cache 壓縮技術(shù)也成為了 LLM 推理領(lǐng)域的熱門研究方向。本文會帶領(lǐng)大家深入剖析KV Cache的各項(xiàng)特性，并詳細(xì)闡述了當(dāng)前用于優(yōu)化LLMs中KV Cache空間使用的各種方法，闡明了它們之間的相互關(guān)系，并比較它們的核心思想。
注意：因?yàn)镵V Cache優(yōu)化的內(nèi)容太多，因?yàn)槲覀儗⒎譃槿恼聛碜屑?xì)學(xué)習(xí)。
注：全部文章列表在這里，估計(jì)最終在35篇左右，后續(xù)每發(fā)一篇文章，會修改此文章列表。
cnblogs 探秘Transformer系列之文章列表
0x01 背景知識
1.1 度量指標(biāo)
推理加速一般可以從兩個(gè)層面體現(xiàn)：吞吐量與延遲。這兩個(gè)層面的指標(biāo)參見下圖。
生成式LLM可以用于各種具有不同SLO（ Service Level Objective/服務(wù)級別目標(biāo)）的任務(wù)。對于批處理任務(wù)（例如摘要），TTFT或TBT延遲指標(biāo)不如吞吐量重要。另一方面，對于對延遲敏感的任務(wù)（例如對話API），TTFT和TBT是更重要的指標(biāo)，其SLO更嚴(yán)格。
1.1.1 吞吐量
吞吐量 (Throughput) 是從系統(tǒng)的角度來看的指標(biāo)，或者說是 LLM 服務(wù)的成本指標(biāo)，表示每生成一個(gè) token 服務(wù)商需要支付的算力成本。吞吐量有多種評估方式，比如 tokens per second（tps），即推理服務(wù)器單位時(shí)間內(nèi)能處理針對所有用戶和請求生成的輸出token數(shù)。如果不僅考慮預(yù)填和解碼，還考慮內(nèi)存限制和上下文切換，則可以考慮基于會話的吞吐量，即在給定時(shí)間內(nèi)的并發(fā)用戶交互數(shù)量，是一個(gè)端到端的目標(biāo)。
吞吐量衡量的是所有用戶和請求中完成的請求或詞元的數(shù)量，因此忽略了時(shí)延要求。有研究人員引入了有效吞吐量（goodput），即每秒完成請求的遵守 SLO（TTFT 和 TPOT 要求）的數(shù)量。因?yàn)橛行掏铝坎东@了在 SLO 達(dá)成下的請求吞吐量 —— 因此既考慮了成本又考慮了服務(wù)質(zhì)量。
1.1.2 延遲
模型為用戶生成完整響應(yīng)所需的總時(shí)間可以使用這兩個(gè)指標(biāo)來計(jì)算：TTFT和TPOT。
首次token生成時(shí)間（Time To First Token，簡稱TTFT）
TTFT是在用戶輸入查詢的內(nèi)容后，模型生成第一個(gè)輸出token所需要的時(shí)間。在人機(jī)實(shí)時(shí)交互的過程中，讓用戶得到快速的響應(yīng)至關(guān)重要，返回結(jié)果越快，用戶體驗(yàn)越好。但對于離線工作負(fù)載則不太重要。在實(shí)踐中，我們一般會人為設(shè)置一個(gè)TTFT SLO（Service Level Objective）。舉例來說，假設(shè)我們定P90 TTFT SLO = 0.4s，那么這意味著我們對該系統(tǒng)的要求是：“90%的request的TTFT值都必須<=0.4”。
單個(gè)輸出詞元的生成時(shí)間（Time Per Output Token，簡稱TPOT）
TPOT 是推理系統(tǒng)根據(jù)用戶請求生成后續(xù)詞元所需要的平均時(shí)間，即總生成時(shí)間除以總生成 token 數(shù)。
該指標(biāo)與每個(gè)用戶如何感知模型的“速度”相對應(yīng)。延時(shí)較高會讓客戶陷入較長的等待時(shí)間，大大影響交互體驗(yàn)，但只要生成速度大于人類的閱讀速度就能獲得很好的用戶體驗(yàn)。例如，100 毫秒/token 的 TPOT 意味著每個(gè)用戶每秒 10 個(gè) token，或每分鐘約 450 個(gè)單詞，這比一般人的閱讀速度要快。
token之間的時(shí)間（TBT）
也有研究人員建議使用token之間的時(shí)間（TBT），因?yàn)?TBT 的定義則更加明確的指定為兩個(gè) token 生成間的延遲，因此基于 TBT 上限來設(shè)計(jì) Service Level Objective (SLO) 能更好的反映流式交互時(shí)的用戶體驗(yàn)。
TBT 和 TPOT 在以下情況下可能會表現(xiàn)不一致：
- 生成速度不均勻：如果模型在生成過程中速度不穩(wěn)定，TBT 會顯示出這種波動，而TPOT則會平均化這些差異。假如某些token生成特別慢，而其他的很快，TBT 會反映這種差異，TPOT則會給出一個(gè)平均值。
- 批處理效應(yīng)：某些模型可能會進(jìn)行批處理，一次生成多個(gè)token。這種情況下，TBT 可能會顯示為周期性的模式（批次內(nèi)快，批次間慢），而TPOT會平滑這種效果。
- 長序列生成：在生成長序列時(shí)，模型可能會因?yàn)樯舷挛脑鲩L而逐漸變慢，TBT 會清楚地顯示這種逐漸減速的趨勢，而TPOT只會反映整體的平均速度。
- 硬件資源波動：如果服務(wù)器負(fù)載不穩(wěn)定或存在資源競爭，可能會導(dǎo)致某些時(shí)刻的生成速度突然變慢，TBT 能夠捕捉到這些瞬時(shí)的波動，而TPOT則會將其平均化。
- 模型架構(gòu)特性：某些模型架構(gòu)可能在生成特定類型的token時(shí)更快或更慢。TBT 可以反映出這種因token類型而異的速度差異，TPOT則會給出一個(gè)整體的平均值。
- 緩存效應(yīng)：如果模型使用緩存來加速生成，可能會導(dǎo)致某些token生成得特別快。TBT 會顯示出這種不均勻的速度，而TPOT會平均化這種效果。
總之，TBT 提供了更細(xì)粒度的性能視圖，能夠反映出生成過程中的變化和波動。而TPOT則提供了一個(gè)整體的、平均的性能指標(biāo)。在分析模型性能時(shí)，最好同時(shí)考慮這兩個(gè)指標(biāo)，而不是只看常用的TPOT。
關(guān)注指標(biāo)
LLM的推理過程一般分為預(yù)填充（Prefill）和解碼（Decode）兩個(gè)階段。預(yù)填充階段負(fù)責(zé)處理輸入提示（Prompt）的完整內(nèi)容，計(jì)算量大但并行性高，同時(shí)生成第一個(gè)Token；解碼階段則通過自回歸方式逐個(gè)生成后續(xù)的Token，盡管單步計(jì)算量較小，但每個(gè)新Token的生成都必須反復(fù)訪問之前生成的所有Token對應(yīng)的KV緩存（Key-Value Cache）。
由于預(yù)填充和解碼階段具有不同的特征，人們用不同的指標(biāo)來衡量它們對應(yīng)的SLO。具體來說，預(yù)填充階段主要關(guān)注請求到達(dá)和生成第一個(gè)標(biāo)記之間的延遲，即TTFT。另一方面，解碼階段關(guān)注同一請求連續(xù)標(biāo)記生成之間的延遲，即TBT，或者TPOT。最大化總體有效吞吐量是LLM推理優(yōu)化主要目標(biāo)，這主要是 TTFT 和 TBT，簡單說就是prefilling階段耗時(shí)優(yōu)化+decoding階段耗時(shí)優(yōu)化。
1.2 內(nèi)存危機(jī)
LLM 推理服務(wù)的吞吐量指標(biāo)主要受制于顯存限制。有研究團(tuán)隊(duì)發(fā)現(xiàn)，現(xiàn)有系統(tǒng)由于缺乏精細(xì)的顯存管理方法而浪費(fèi)了 60% 至 80% 的顯存，浪費(fèi)的顯存主要來自 KV Cache。下圖展示了一個(gè)13B的模型在A100 40GB的gpu上做推理時(shí)的顯存占用分配（others表示forward過程中產(chǎn)生的激活大小，這些激活可以用完則廢，因此占據(jù)的顯存不大），從這張圖中我們可以直觀感受到推理中KV cache對顯存的占用。
內(nèi)存危機(jī)的原因主要是以下幾點(diǎn)：
- LLM中最昂貴的操作是矩陣乘法和softmax，它們本身是 memory-bound。
- LLM中的權(quán)重和激活 本身已經(jīng)就接近達(dá)到 memory capacity
- KV caching進(jìn)一步加劇了內(nèi)存容量的需求（隨著句子長度和batch size線性增長），而存儲容量是由序列長度和模型維度決定的。對于給定的LLM，隨著批大小和序列長度的增加，分配給KV緩存的內(nèi)存繼續(xù)線性增長，并在某些時(shí)候超過可用內(nèi)存容量。
- 長序列訪問瓶頸：隨著Token數(shù)量增加，每次生成新Token時(shí)需讀取的KV緩存數(shù)據(jù)量成線性增長，導(dǎo)致顯存帶寬壓力增大，延遲上升。
下圖則可以從細(xì)節(jié)來查看推理中KV cache對顯存的占用，以及推理的不同階段關(guān)注的指標(biāo)。其中：
- Model memory：模型內(nèi)存。模型大小保持不變，不受序列長度或批次大小的影響。
- Peak Momery：峰值內(nèi)存 = 模型內(nèi)存+kvcache。
- Latency：生成總耗時(shí) = TTFT+TPOT*生成的token數(shù)目。
從圖上可以看到，在生成過程中 Latency 和 Memory 的關(guān)系，具體分析如下。
-
綠色標(biāo)號1為待機(jī)階段，此時(shí)模型等待輸入，內(nèi)存是模型權(quán)重。
-
綠色標(biāo)號2為Prefill 階段，特點(diǎn)如下：
- 可以并行計(jì)算。淺藍(lán)色標(biāo)號1代表生成第一個(gè)token所需時(shí)間，該Latency比較低。
- 會一次生成大量 KV Cache，所以內(nèi)存在短期內(nèi)激增。
- 假設(shè) Batch size 為 1，序列長度越大，計(jì)算強(qiáng)度越大，通常都會位于 Compute Bound 區(qū)域。
- 預(yù)填充長輸入比短輸入需要更長的計(jì)算時(shí)間和GPU內(nèi)存；
-
綠色標(biāo)號3為Decoding 階段，其特點(diǎn)如下：
- 不能并行計(jì)算，天藍(lán)色標(biāo)號2代表Decoding 階段生成一個(gè) Token 的 Latency，該Latency 比較高，而且，隨著序列的變長，這個(gè) Latency 會逐漸增加。
- 顯存增加比較慢，但是顯存的增量與序列的長度成正比。
- 在預(yù)填充后，駐留在GPU HBM上的大型KV Cache大大限制了可服務(wù)的并發(fā)用戶數(shù)量；
- 在解碼過程中，從HBM到SM重復(fù)讀取KV Cache會導(dǎo)致顯著的上下文切換延遲。
- Batch size 越大，計(jì)算強(qiáng)度越大，理論性能峰值越大，但此時(shí)通常都會位于 Memory Bound 區(qū)域。
-
綠色標(biāo)號4是當(dāng)一個(gè)序列處理完之后，會釋放相應(yīng)的 KV Cache。此外，模型參數(shù)在不同序列間是共享的，而 KV Cache 則不是。
下圖給出了有和沒有KV緩存的OPT-6.7B推理的執(zhí)行時(shí)間和GPU內(nèi)存使用情況。x軸步長索引表示輸出序列長度。在沒有KV緩存的情況下，執(zhí)行時(shí)間迅速增加；在有KV緩存的情況下，只有新生成的令牌的注意力權(quán)重和分?jǐn)?shù)被計(jì)算為向量-矩陣乘法，在不同的步驟中，執(zhí)行時(shí)間幾乎保持不變。這種運(yùn)行時(shí)間的減少是以GPU內(nèi)存使用量為代價(jià)的，隨著時(shí)間的推移，GPU內(nèi)存使用量逐漸增加，這是由于KV張量的規(guī)模越來越大。
因此，如何優(yōu)化KV cache，節(jié)省顯存，提高推理吞吐量，就成了LLM推理框架需要解決的重點(diǎn)問題。
1.3 KV Cache問題
我們總結(jié)下傳統(tǒng)KV Cache中存在的問題。
- 顯存占用問題。隨著總序列長度（上下文窗口大?。┗蛞淮翁幚淼男蛄袛?shù)量（batch size，即吞吐量）的增長，KV Cache也會隨之線性增長。因此，KV緩存的大小實(shí)際上并沒有上限，而GPU內(nèi)存顯然是有限的。 這樣KV Cache緩存成為主要緩存占用。在更長的上下文中，KV Cache大小甚至超過了模型參數(shù)的大小。
- 內(nèi)存碎片化。每個(gè)新生成的Token都會產(chǎn)生對應(yīng)的
K和V向量，并存入緩存。在持續(xù)的\({K, k_n} \longrightarrow K\) 和 \({V, v_n} \longrightarrow V\)操作中，內(nèi)存被反復(fù)分配和釋放，導(dǎo)致大量碎片化問題。 - 顯存管理問題。隨著序列長度增加，緩存規(guī)模線性增長，導(dǎo)致管理復(fù)雜度隨之上升。比如，傳統(tǒng)策略（如LRU、LFU）難以適配LLM的訪問模式，需設(shè)計(jì)專用策略；緩存內(nèi)存需求隨序列長度和層數(shù)線性增長，易超出硬件限制，需優(yōu)化多存儲設(shè)備協(xié)作。面對變化的訪問模式，需自適應(yīng)緩存策略來進(jìn)行動態(tài)負(fù)載；分布式環(huán)境下，跨節(jié)點(diǎn)一致性和資源如何協(xié)調(diào)。
- 內(nèi)存帶寬瓶頸。隨著Token數(shù)量的增加，每次生成新Token時(shí)需讀取的KV緩存數(shù)據(jù)量呈線性增長，且無法批量處理KV Cache也加劇了嚴(yán)重的內(nèi)存帶寬瓶頸。
- GPU利用率低下。解碼階段從計(jì)算密集變成訪存密集了。這就導(dǎo)致GPU的算力，在解碼階段不是瓶頸了，要想加速解碼，得從訪存方面下手。但是，我們獲取大型權(quán)重矩陣和不斷增長的KV緩存，只是為了執(zhí)行微不足道的矩陣到向量的操作。我們最終花費(fèi)的時(shí)間更多地用于數(shù)據(jù)加載，而不是實(shí)際的數(shù)值計(jì)算，這顯然導(dǎo)致了GPU計(jì)算能力的低效利用。
0x02 總體思路
2.1 分類
考慮到GPU內(nèi)存的有限性，而減小內(nèi)存占用的策略是一箭三雕的，因?yàn)槟軌蚴刮覀兲岣哂布寐屎统杀拘б妫瑫r(shí)減少時(shí)延并增加吞吐量。因此，KV Cache的內(nèi)存壓力激發(fā)了許多不同方向的創(chuàng)新，比如新型注意力架構(gòu)（如MQA、GQA、SWA）、緩存壓縮策略（如H2O、Scissorhands、FastGen）、高效的內(nèi)存管理（如PagedAttention、RadixAttention），以及量化和存儲容量擴(kuò)展（如負(fù)載系統(tǒng)、單主機(jī)和多主機(jī)模型并行）。
但是如何給這些 KV Cache 優(yōu)化思路進(jìn)行分類？有些研究人員已經(jīng)給出了他們的答案，這些答案各有不同。
比如下圖來自論文“A Survey on Large Language Model Acceleration based on KV Cache Management”。該論文基于KV Cache對LLM推理時(shí)間與內(nèi)存需求的顯著影響，系統(tǒng)梳理了現(xiàn)有優(yōu)化策略，并將其劃分為三個(gè)層次：token級優(yōu)化、模型級優(yōu)化和系統(tǒng)級優(yōu)化。
- token級優(yōu)化是指通過專注于細(xì)粒度來提高KV緩存管理效率。在token級進(jìn)行仔細(xì)的選擇、組織和壓縮，不需要對原始模型進(jìn)行架構(gòu)更改或系統(tǒng)并行技術(shù)。優(yōu)化方法分為五類：KV Cache選擇、預(yù)算分配、合并、量化及低秩分解。這些方法通過直接操作token級數(shù)據(jù)，實(shí)現(xiàn)了計(jì)算效率與內(nèi)存需求的初步優(yōu)化，為后續(xù)更高層次的改進(jìn)奠定基礎(chǔ)。
- KV Cache選擇：側(cè)重于優(yōu)先排序和僅存儲最相關(guān)的token。優(yōu)先存儲對推理最重要的token，通過評估token相關(guān)性減少冗余緩存，提升內(nèi)存利用率。
- KV Cache預(yù)算分配：在token之間動態(tài)分配內(nèi)存資源，以確保在有限內(nèi)存約束下優(yōu)化緩存分布，確保高效利用。
- KV Cache合并：識別并合并相似或重疊的KV對來減少冗余存儲，進(jìn)一步壓縮緩存規(guī)模。
- KV Cache量化：通過降低KV對的存儲精度（如從Float32到Float16或更低比特）來最小化內(nèi)存占用，同時(shí)需平衡精度損失與性能收益。
- KV Cache低秩分解：利用低秩分解技術(shù)將KV矩陣分解為較小維度表示，壓縮緩存體積，同時(shí)保留關(guān)鍵信息。
- 模型級優(yōu)化是指從架構(gòu)設(shè)計(jì)角度優(yōu)化KV Cache管理，通過調(diào)整模型結(jié)構(gòu)或引入新機(jī)制減少緩存依賴，從根本上減少KV Cache的生成與使用需求，提升整體效率。具體策略包括：
- 注意力分組和共享：檢查鍵值對的冗余功能，并在Transformer層內(nèi)或跨Transformer層進(jìn)行分組和共享KV緩存，降低存儲需求。
- 架構(gòu)更改：設(shè)計(jì)新的注意力機(jī)制或構(gòu)建用于KV優(yōu)化的外部模塊，針對KV優(yōu)化重新構(gòu)造模型計(jì)算流程，減少對緩存的直接依賴。
- 非Transformer架構(gòu)：這些架構(gòu)采用了其他內(nèi)存高效的設(shè)計(jì)，如使用循環(huán)神經(jīng)網(wǎng)絡(luò)來優(yōu)化傳統(tǒng)Transformer中的KV緩存。
- 系統(tǒng)級優(yōu)化是指從底層硬件與資源調(diào)度角度優(yōu)化KV Cache管理，涉及內(nèi)存分配與任務(wù)調(diào)度等經(jīng)典問題，旨在最大化硬件資源利用率，以在多樣化計(jì)算環(huán)境中提升效率。具體技術(shù)包括：
- 內(nèi)存管理：通過架構(gòu)創(chuàng)新優(yōu)化內(nèi)存使用，如虛擬內(nèi)存適配、智能前綴共享和層感知資源分配，確保緩存高效存儲與訪問。
- 調(diào)度策略：通過前綴感知方法以最大限度地提高緩存重用率，采用搶占式技術(shù)實(shí)現(xiàn)公平上下文切換，或設(shè)計(jì)層特定機(jī)制進(jìn)行精細(xì)化緩存控制，以解決多樣化優(yōu)化目標(biāo)。
- 硬件感知設(shè)計(jì)：針對單/多GPU、I/O優(yōu)化、異構(gòu)計(jì)算及SSD存儲提出解決方案，提升KV Cache的訪問速度與存儲容量。
這一分類框架從微觀到宏觀，針對KV Cache管理的不同挑戰(zhàn)提出了具體技術(shù)路徑，為理解和設(shè)計(jì)高效加速方案提供了清晰的理論指引。
再比如下圖來自論文“Thus Spake Long-Context Large Language Model”。其主要分為Token Dropping，Token Merging，Layer-wise Sharing等幾類。這種分類方式更加直擊本質(zhì)。由于KV緩存的大小是由緩存序列長度、層數(shù)、KV頭數(shù)、特征維度數(shù)和存儲數(shù)據(jù)類型的乘積決定的，我們可以通過改進(jìn)這些因素中的每一個(gè)來優(yōu)化開銷。特別是，由于對序列長度的優(yōu)化討論最多，又可以將其分為token丟棄）和token合并。
本篇是借鑒了前人的思想，也加入了筆者自己的思考。筆者認(rèn)為，從整體上來講，KV cache主要分成3個(gè)方向進(jìn)行優(yōu)化：從公式角度進(jìn)行優(yōu)化；按照特性進(jìn)行優(yōu)化；按照階段進(jìn)行優(yōu)化。這些方向不一定是正交的。因?yàn)闀幸恍┓桨讣磸哪Ｐ徒Y(jié)構(gòu)特點(diǎn)出發(fā)，也是從公式角度出發(fā)進(jìn)行優(yōu)化。為了更好的分析，如果能從公式角度梳理的方案，筆者都將其劃分到公式角度。
2.2 從公式角度優(yōu)化
如上篇所示，KV cache的峰值顯存占用大小計(jì)算公式（字節(jié)）如下：
\(batch\_size \times sequence\_length \times head\_number \times head\_dimension \times layer\_number \times kv\_bits \times 2\)
以QWen為例，4K上下文和100K上下文的對比如下：
通過上述公式我們可以看到， LLM 推理服務(wù)為每個(gè) Token 分配了顯存空間，其顯存空間大小與批次大小，序列長度，頭維度、頭數(shù)量，模型層數(shù)、以及KVCache存儲的數(shù)據(jù)類型大小幾個(gè)維度相關(guān)。我們的目標(biāo)是在給定顯存空間，推理服務(wù)要處理的 Token 數(shù)盡可能的多。因此存在兩個(gè)優(yōu)化策略：
- 對 批次大小、序列長度這兩個(gè)維度進(jìn)行壓縮，用較少的 token 信息表示全量 token 信息；
- 對頭維度、頭數(shù)量，模型層數(shù)、以及KVCache存儲的數(shù)據(jù)類型幾個(gè)維度進(jìn)行壓縮，降低每個(gè) token 的顯存開銷。
我們逐一檢查各個(gè)項(xiàng)目來看看影響KV cache的具體因素，以及減少KV緩存內(nèi)存占用的方法。
注：“從公式角度優(yōu)化”基本對應(yīng)論文“A Survey on Large Language Model Acceleration based on KV Cache Management”中的"token級別優(yōu)化“。
2.2.1 難以利用的選項(xiàng)
我們發(fā)現(xiàn)，如下選項(xiàng)難以利用。
減少模型層數(shù)
有一部分研究通過比較輸入輸出的embedding的相似性，來判斷某一層是否重要，從而把一些不重要的層進(jìn)行剪枝。這可以同時(shí)達(dá)到減少模型weight和KVCache的作用。但是減少層數(shù)帶來的收獲并不多。
減少batch size
batch size和推理的子階段性能密切相關(guān)。比如，Prefill 階段在比較小 Batch Size 下可以獲得比較大的計(jì)算強(qiáng)度，相應(yīng)的吞吐也很高；而 Decoding 階段需要比較大的 Batch Size 才能獲得相對高的計(jì)算強(qiáng)度及吞吐。
雖然減少批大小可以減少KV緩存的內(nèi)存占用，降低時(shí)延，但這會使得單次運(yùn)行能夠處理的對象減少，也會降低硬件利用率，導(dǎo)致成本效率降低。所以，在多數(shù)情況下，我們并不想減少批大小，反而會盡可能地增加批次大小。
2.2.2 LLM推理加速的優(yōu)化方向
由此我們可以找到一些優(yōu)化KV cache的具體方向：
- 減少序列長度：針對超長 prompt，通過減少 kv cache 中 slot 的數(shù)量進(jìn)行壓縮，以此達(dá)到減少輸入輸出序列的長度（或者說是提示 + 完成部分）的目的。例如，對于長度為 128K tokens 的 prompt，僅挑選 1024 個(gè) tokens 的 kv cache 進(jìn)行存儲。
- 減少注意力頭個(gè)數(shù)：MQA（multi-query attention)、GQA（Grouped-query attention）通過減少kv cache的head個(gè)數(shù)減少顯存占用。
- 減少key_bits。因?yàn)镕P16 占2個(gè)bytes，所以優(yōu)化方式主要是量化，把模型參數(shù)由fp16轉(zhuǎn)換到int8/int4，每個(gè)參數(shù)占用的字節(jié)也由2byte轉(zhuǎn)換到1byte/0.5byte。在保持 kv cache slot 數(shù)量不變的情況下，將數(shù)據(jù)格式從 fp16 壓縮到 int8 或 int4 等低精度格式。
- 減少頭維度。DeepSeek V2的MLA引入了類似LoRA的想法，有效地減少了KV頭的大小。
- 優(yōu)化 KV cache 的顯存管理：目前GPU上KV Cache的有效存儲率較低，可以通過類似PagedAttention的方法進(jìn)行優(yōu)化。 因?yàn)椴皇钦坏?，我們后續(xù)把內(nèi)存管理部分也劃分到”按照特性進(jìn)行優(yōu)化“部分進(jìn)行分析。
另外，論文“A Survey on Large Language Model Acceleration based on KV Cache Management”中總結(jié)了基于模型級別的優(yōu)化，這種優(yōu)化通過設(shè)計(jì)新的變換器架構(gòu)或機(jī)制優(yōu)化KV Cache復(fù)用，通常需對模型進(jìn)行重訓(xùn)練或微調(diào)。模型級別的優(yōu)化和本文分類不是正交的，因此筆者把本文涉及的方案添加在圖上。
2.2 按照特性進(jìn)行優(yōu)化
前文有提到，Prefilling和decoding的特性不同，比如：
- prefilling階段對計(jì)算資源的需求量極大，哪怕是小批次的預(yù)填充任務(wù)，甚至單個(gè)較長的預(yù)填充任務(wù)，都足以使GPU的計(jì)算能力達(dá)到飽和。Prefilling階段是計(jì)算瓶頸，隨著輸入長度增加，耗時(shí)超線性增加（吞吐量下降，如果耗時(shí)線性增加，那么平均吞吐量不變）。
- 與此相對，解碼任務(wù)則需要更大的批大小才能充分利用計(jì)算資源。decoding階段是傳輸瓶頸，隨著bs的增大，耗時(shí)略有增加（主要是因?yàn)閿?shù)據(jù)增多導(dǎo)致傳輸耗時(shí)增多，計(jì)算耗時(shí)基本認(rèn)為不變），吞吐量大概是隨著輸入長度而線性增加。
下圖給出了prefill和decoding階段中，KV Cache的使用。在初始化階段（對應(yīng)于第一次迭代），LLM為輸入提示中的每個(gè)令牌生成鍵值緩存。在隨后的解碼階段，LLM只需要計(jì)算一個(gè)新生成的令牌的查詢、密鑰和值，利用預(yù)先計(jì)算的鍵值緩存來逐步簡化該過程。因此，與初始化階段（即第一次迭代）相比，解碼階段迭代的執(zhí)行時(shí)間通常較小。
因此，我們要針對這些特性對KV Cache進(jìn)行不同的處理。這部分與論文“Thus Spake Long-Context Large Language Model”中的系統(tǒng)級別優(yōu)化有部分重疊。
2.2.1 prefill
針對prefill階段計(jì)算強(qiáng)度高的特點(diǎn)有幾種優(yōu)化思路，比如并行生成，把prompt切分等。我們將在本文后續(xù)章節(jié)中進(jìn)行深入分析。
2.2.2 Decoding
雖然一次 Decoding Step 計(jì)算量線性增加，但由于是訪存瓶頸，因此計(jì)算時(shí)間不會明顯增加，如果 Decoding Step 數(shù)量減少的比例更多，那么就很有可能降低整個(gè)請求的時(shí)延。因此需要想辦法增加Decoding的請求數(shù)目或者起到類似效果。針對 Decoding 階段計(jì)算強(qiáng)度比較低的情況，有兩種優(yōu)化的思路：
- 請求間優(yōu)化。在不同請求間 Batching 的方式（Continuous Batching），可以增大 Decoding 階段的 Batch Size，進(jìn)而更充分發(fā)揮 GPU 算力，但是其無法降低單個(gè)請求的 Latency，而且如果用戶請求比較少，比如只有單一用戶，那么永遠(yuǎn)無法實(shí)現(xiàn) Batching。
- 請求內(nèi)優(yōu)化。在單一請求內(nèi)一次驗(yàn)證多個(gè) Decoding Step（Speculative Decoding），雖然一次 Decoding Step 計(jì)算量線性增加，但由于是訪存瓶頸，因此計(jì)算時(shí)間不會明顯增加，如果 Decoding Step 數(shù)量減少的比例更多，那么就很有可能降低整個(gè)請求的時(shí)延。
我們會在后續(xù)文章進(jìn)行詳細(xì)分析。
2.2.3 分離與否
由于Prefilling和decoding的特性不同，計(jì)算量差異非常大，導(dǎo)致兩者容易會對彼此造成干擾。比如：連續(xù)批處理任務(wù)中，新加入的prefill任務(wù)會顯著增加了decoding任務(wù)的時(shí)延，導(dǎo)致處于解碼階段的請求在每次預(yù)填充請求進(jìn)入系統(tǒng)時(shí)都會“卡住”。因此，有些實(shí)現(xiàn)會將預(yù)填充和解碼解耦到不同的GPU中，并為每個(gè)階段定制并行策略。這自然解決了上述兩個(gè)問題：
- 預(yù)填充和解碼之間沒有干擾，使得兩個(gè)階段都能更快地達(dá)到各自的SLO。
- 資源分配和并行策略解耦，從而為預(yù)填充和解碼量身定制優(yōu)化策略。
當(dāng)然，也有研究人員對分離策略提出了質(zhì)疑，認(rèn)為盡管PD分離提供了靈活性，但是導(dǎo)致了嚴(yán)重的GPU資源浪費(fèi)的問題。具體來說，運(yùn)行計(jì)算密集型Prefill階段的GPU的HBM容量和帶寬的利用率較低。內(nèi)存密集型的Decode階段則面臨算力資源利用率低的問題。因此，另外一個(gè)流派強(qiáng)調(diào)需要融合或者算力遷移（比如，將 Decode 階段內(nèi)存密集的 Attention 計(jì)算操作卸載一部分到 Prefill 節(jié)點(diǎn)上執(zhí)行），我們會在后續(xù)文章進(jìn)行詳細(xì)分析。
2.2.4 內(nèi)存管理
此處可以從兩個(gè)維度進(jìn)行思考。
首先，我們從系統(tǒng)維度來看。目前很多推理框架實(shí)現(xiàn)中存在顯存碎片化和顯存過度保留的問題，顯存利用率不高。由于序列長度具有高度可變和不可預(yù)測的特點(diǎn)，因此，KV緩存的內(nèi)存需求是未知的，很難提前為KV cache量身定制存儲空間。因此，很多推理框架會預(yù)先為一條請求開辟一塊連續(xù)的矩形存儲空間。然而，這樣的分配方法很容易引起“gpu顯存利用不足”的問題，進(jìn)而影響模型推理時(shí)的吞吐量。因此，系統(tǒng)級別的內(nèi)存管理會通過架構(gòu)創(chuàng)新優(yōu)化內(nèi)存使用，如虛擬內(nèi)存適配、智能前綴共享和層感知資源分配，確保緩存高效存儲與訪問。我們會在本文介紹。
下圖給出了綜述論文對于現(xiàn)有內(nèi)存管理方案的一些總結(jié)。以vLLM與PagedAttention和vTensor為例的架構(gòu)設(shè)計(jì)，采用了經(jīng)典的操作系統(tǒng)原理，創(chuàng)建了靈活的動態(tài)內(nèi)存分配系統(tǒng)，通過復(fù)雜的映射和虛擬內(nèi)存抽象優(yōu)化了物理內(nèi)存的使用。ChunkAttention和MemServe等前綴感知設(shè)計(jì)通過組織數(shù)據(jù)結(jié)構(gòu)來實(shí)現(xiàn)高效的緩存重復(fù)數(shù)據(jù)刪除和共享公共前綴，從而進(jìn)一步優(yōu)化了這種方法，從而提高了內(nèi)存利用率和計(jì)算效率。這些創(chuàng)新共同展示了通過內(nèi)存管理顯著增強(qiáng)LLM服務(wù)的潛力。
其次，論文“Thus Spake Long-Context Large Language Model”提出了把KV Cache看作是記憶管理的一部分，這是從業(yè)務(wù)角度來看。盡管KV緩存優(yōu)化在實(shí)踐中努力實(shí)現(xiàn)更長的上下文，但本質(zhì)上，它是在效率和性能之間尋求平衡。緩存優(yōu)化并沒有試圖突破LLM的能力上限，因?yàn)樗]有改變上下文信息的組織形式?；跇?biāo)準(zhǔn)KV緩存機(jī)制的長上下文LLMs仍然面臨一些限制，包括只讀訪問和一次性讀取所有信息的需求，這使其不適合更復(fù)雜的場景。這促使將記憶管理納入LLMs，將KV緩存視為一種特定的記憶實(shí)例。這種記憶管理相較于KV緩存最直觀的改進(jìn)在于其更加靈活的訪問方式，避免了一次性讀取全部KV緩存。MemTrans將預(yù)訓(xùn)練階段的KV緩存存儲在外部存儲中，以便在推理時(shí)提供更相關(guān)的信息。MemLong將這一概念擴(kuò)展到長上下文，通過存儲上下文片段的KV緩存，并根據(jù)相關(guān)性檢索KV對來指導(dǎo)推理。本文不對此方面進(jìn)行解讀。
2.2.5 調(diào)度
這部分主要是針對KV Cache特點(diǎn)實(shí)施的一些調(diào)度策略，主要是三種途徑：
-
前綴感知調(diào)度策略。通過前綴感知方法以最大限度地提高緩存重用率，以BatchLLM和RadixAttention為代表；
-
搶占式和面向公平的調(diào)度。采用搶占式技術(shù)實(shí)現(xiàn)公平上下文切換以FastServe和FastSwitch為代表。
-
特定于層的分層調(diào)度方法。針對層的特點(diǎn)進(jìn)行精細(xì)化緩存控制，以解決多樣化優(yōu)化目標(biāo)。LayerKV和CachedAttention是典型代表。
下圖展示了這些分類。
另外，PD分離后，如何最大效率的調(diào)度Prefilling服務(wù)器和Decoding服務(wù)器是進(jìn)一步優(yōu)化的核心所在。這里有兩個(gè)要點(diǎn)：
- 如何優(yōu)化KV Cache的傳輸。因?yàn)橄鄬τ?jì)算資源來說，存儲更廉價(jià)和充足，所以KV Cache會存儲在系統(tǒng)的不同部分中，然后通過網(wǎng)絡(luò)進(jìn)行傳輸。
- Decoding階段如何最大化batch（但是太大的batch size也會導(dǎo)致平均單個(gè)token的生成時(shí)間變長）。
2.3 按照階段來優(yōu)化
下圖是論文"Keep the Cost Down: A Review on Methods to Optimize LLM’s KV Cache Consumption"提出來的劃分方式，
該論文按LLM的時(shí)間順序展開，從訓(xùn)練階段到部署階段，再到post-training（后訓(xùn)練）階段對KV Cache的優(yōu)化方案進(jìn)行了劃分。
- 訓(xùn)練階段主要包括可用于模型預(yù)訓(xùn)練的KV Cache壓縮方法。在這個(gè)階段，模型具有最大的可塑性。此階段的主要調(diào)整是模型架構(gòu)，它在保留Attention的優(yōu)秀特性的同時(shí)，將生成的鍵和值向量的大小減小到四分之一甚至更小，也可以改變架構(gòu)，完全放棄具有二次復(fù)雜性的注意力。這些方法通常是最有效的，但它們不適合施加在現(xiàn)有模型上，以及計(jì)算量較低的場景。
- 部署階段主要包括用不同的框架來優(yōu)化KV Cache。這些方法不會對KV Cache本身進(jìn)行大量修改，但可以在相同的環(huán)境中顯著優(yōu)化其效率。
- post-training階段主要包括針對KV Cache的實(shí)時(shí)優(yōu)化方法，主要包括三種方法：刪除、合并和量化。在推理過程中，從模型輸入提示級別到嵌入級別，再到KV Cache級別，都可以進(jìn)行壓縮優(yōu)化。
0x03 按照公式優(yōu)化
我們接下來通過具體案例來剖析如何依據(jù)以下公式的特點(diǎn)來進(jìn)行優(yōu)化。
\(batch\_size \times sequence\_length \times head\_number \times head\_dimension \times layer\_number \times kv\_bits \times 2\)
3.1 減少注意力頭的數(shù)量
前文提到，單個(gè) Token 的 KVCache 由四大因素決定：Layers、NumHead、HeadDim、以及 KVCache 數(shù)據(jù)類型決定的。所以壓縮每個(gè) Token 的 KVCache 大小就有四大方向，我們首先看看如何通過減少NumHead來進(jìn)行優(yōu)化。
對于給定的模型架構(gòu)，模型大小主要受控于層數(shù)和頭數(shù)，減少注意力頭的數(shù)量可能意味著選擇一個(gè)更小的模型。然而，如果仔細(xì)觀察，我們會發(fā)現(xiàn)只需要減少鍵頭和值頭的數(shù)量，因?yàn)椴樵冾^的數(shù)量不會影響KV緩存的大小。與層數(shù)維度優(yōu)化類似，減少頭數(shù)會顯著影響LLMs的表示能力。為了保留性能，頭數(shù)維度優(yōu)化通常依賴于共享策略。這正是MQA（多查詢注意力）和GQA（分組查詢注意力）架構(gòu)的核心思想。這些多頭注意力（MHA）變體的唯一目的是用更少的注意力頭訓(xùn)練基礎(chǔ)模型，從而減小KV緩存的大小。然而，這些模型通常在針對特定任務(wù)進(jìn)行微調(diào)后使用。也有一些KV Cache壓縮的算法會去profile哪些head對生成更加重要，進(jìn)而為它們分配更多的cache空間，那些不重要的就減少空間或者干脆不分配。
此外，微調(diào)現(xiàn)有模型還可以進(jìn)一步優(yōu)化頭數(shù)的大小。例如，SHA計(jì)算頭權(quán)重矩陣之間的余弦相似度，并將相似的頭分組以共享單個(gè)KV緩存。DHA采用質(zhì)心對齊方法計(jì)算頭之間的相似度，通過線性融合相似頭的KV緩存，有效地將MHA壓縮為GQA。
我們將在后續(xù)專門用一篇文章來介紹MQA和GQA。
3.2 減少每個(gè)參數(shù)占用的字節(jié)
此方向就是量化技術(shù)在KV Cache中的應(yīng)用。KV Cache量化不僅能顯著降低推理過程中的顯存占用，更重要的是在壓縮顯存占用的同時(shí)，能夠很好地保持模型的推理性能。
我們會在后面介紹量化時(shí)候統(tǒng)一進(jìn)行學(xué)習(xí)。
3.3 減少注意力頭維度
在很多模型家族中，無論模型的大小如何，注意力頭的隱藏維度都可能保持不變。因此，難以通過降低注意力頭維度來優(yōu)化KV Cache，除非更換模型。而且也可能是因?yàn)殡[藏維度大小已經(jīng)是128，如果再縮小會影響效果。
想要減少頭維度，一種可能的方法就是低秩分解（Low-Rank Decomposition），畢竟LORA就是利用的這個(gè)思想。DeepSeek V2的MLA就引入了類似LoRA的想法，有效地減少了KV頭的大小。ECH對分組的頭權(quán)重矩陣應(yīng)用基于SVD的低秩分解，實(shí)現(xiàn)了與GQA類似的KV壓縮效果，但不同的是，它不采用平均融合。Neurocache對頭矩陣應(yīng)用低秩壓縮，并在注意力計(jì)算中使用最相似的緩存。
而且也有針對/利用channel的稀疏化。比如Double Sparsity 將 Token 稀疏性和 Channel 稀疏性相結(jié)合：
- Token 稀疏性側(cè)重于利用重要的 Token 來計(jì)算 Attention，也就是 Sequence 維度的稀疏。
- Channel 稀疏性是使用重要的 Channel 來識別重要的 Token，也就是 Hidden 維度的稀疏。
我們將在后續(xù)文章單獨(dú)介紹MLA。此處介紹針對channel的稀疏化。
3.2.1 Double Sparsity
論文”Post-Training Sparse Attention with Double Sparsity”在針對長序列場景使用 Token 稀疏化的同時(shí)，也額外引入了 Channel 稀疏化，將 Token 稀疏性和 Channel 稀疏性相結(jié)合來獲得不錯(cuò)的加速。
洞見
作者的關(guān)鍵見解是：Channel 的稀疏性通常是相對靜態(tài)的，可以通過離線校準(zhǔn)的方式來執(zhí)行，從而能夠準(zhǔn)確和高效的識別重要的 Token，使推理更高效。此外，這種方法可以與 Offload 結(jié)合，以顯著減少內(nèi)存占用。
方案
Double Sparsity是一種訓(xùn)練后稀疏注意力機(jī)制。其整體的方案比較簡單。
Offline Calibration
首先通過離線的方式（Offline Calibration）識別出重要的 Channel。
具體而言如下圖所示，借鑒 AWQ 中識別重要 Channel 的方式，使用少量樣本數(shù)據(jù)，通過 Offline 的方式識別哪些 Channel 更加重要。這里有幾種不同的策略：
- Activation Outlier：根據(jù) Activation 中的 Outlier 來識別重要的 Channel，紅色部分。
- Q-Outlier：根據(jù) Query 中的 Outlier 來識別重要的 Channel，綠色部分。
- K-Outlier：根據(jù) Key 中的 Outlier 來識別重要 Channel，藍(lán)色部分。
- QK-Outlier：根據(jù) Query * Key 來識別重要 Channel，紫色部分。
實(shí)時(shí)計(jì)算
離線操作之后，會進(jìn)行實(shí)際計(jì)算。具體分為 3 步：
- 根據(jù)重要 Channel 的 Index 提取 Query，同樣要連續(xù)存儲。
- 使用稀疏但連續(xù)存儲的 Query 與 Key Cache 計(jì)算 Attention Score。
- 根據(jù) Attention Score 選出 Top-k 的重要 Token（Key/Value Cache），然后使用這些重要的 Token 進(jìn)行 Attention 計(jì)算。
如下圖所示。
Double Sparsity-Offload
從上面也可以看出，Double Sparsity 方案雖然可以在一定程度上減少顯存 IO，但需要額外存儲 Label Cache，實(shí)際上是增加了顯存占用。為了解決這個(gè)問題，作者提出了將全量的 Key/Value Cache Offload 到 CPU 內(nèi)存中，在需要的時(shí)候動態(tài)的加載到顯存中，也就是 Double Sparsity-Offload。
然而，如果每次使用 Label Cache 獲得 Top-k 的 Token 后再去加載對應(yīng)的 Key/Value Cache 到顯存中，其延時(shí)可能比較大，并且無法與計(jì)算 overlap。為了解決這個(gè)問題，作者引入了 Double Buffer，其核心思想是：在計(jì)算當(dāng)前 Transformer Layer 的時(shí)候可以異步的預(yù)取下一個(gè) Layer 的相關(guān) Token 的 Key/Value Cache。為什么可以提前加載下一層的 Key/Value Cache 呢？因?yàn)樽髡邔Ρ认噜弻拥?Attention Head 發(fā)現(xiàn)，除了前兩層以及最后兩層外，后續(xù)相鄰層的相似性都很高，甚至達(dá)到 95% 以上，因此也就可以使用當(dāng)前 Layer 來近似下一個(gè) Layer。
3.4 減少Layers
3.4.1 總體思路
需要注意的是，此處不是真正的減少層數(shù)，而是只緩存模型部分層的KV，比如，逐層僅保留關(guān)鍵的上下文鍵和值，合并相鄰層的KV Cache等。
對于層維度，基本假設(shè)是，一些任務(wù)可能不需要完整深度的計(jì)算，在預(yù)填充過程中跳過一些層可能會同時(shí)減少預(yù)填充的浮點(diǎn)運(yùn)算次數(shù)和KV緩存的大小。我們接下來通過幾個(gè)實(shí)際案例進(jìn)行學(xué)習(xí)。
3.4.2 LayerSkip
LayerSkip的總體思路是跳過一些層然后驗(yàn)證，其實(shí)也是投機(jī)推理的思路。在訓(xùn)練階段，LayerSkip對不同層應(yīng)用不同的丟失率（dropout rate），所有Transformer層共享同一出口的提前退出損失（early exit loss）。在推理階段，因?yàn)橐呀?jīng)通過訓(xùn)練提高了早期層的提前退出準(zhǔn)確性，因此無需增加輔助層；此外，LayerSkip引入一種自我推測解碼方法，在早期層退出后，使用剩余層進(jìn)行驗(yàn)證和糾正，從而減少內(nèi)存占用并提高效率。
動機(jī)
LayerSkip 作者在設(shè)計(jì) LayerSkip 前做了一個(gè)實(shí)驗(yàn)，觀察整個(gè) Transformer 每一層到底完成什么事情。因?yàn)樵趬嚎s前，需要問自己一個(gè)問題：需要壓縮的對象是否存在冗余信息？是否需要全量執(zhí)行？如果可以，應(yīng)該怎么在保證模型效果不會有大幅降低的前提下提升推理的性能？所以，LayerSkip實(shí)驗(yàn)的目的就是找出最關(guān)鍵的層。
作者使用 LLaMA 7B 測試 HumanEval 代碼測試集來測試代碼生成。作者在 LLaMA -7B 每層 TransformerLayer 后都接入 LM Head 層，統(tǒng)計(jì)整個(gè)生成過程中，每一層 TransformerLayer 生成每一個(gè) Token 的概率分布。實(shí)驗(yàn)舉例如下。
實(shí)驗(yàn)統(tǒng)計(jì)結(jié)果如下圖所示。通過圖上表格可以觀察到兩個(gè)現(xiàn)象：
- 前面層生成的結(jié)果大部分和最終結(jié)果幾乎毫無關(guān)系，主要原因是前面的層的權(quán)重和 LM Head權(quán)重沒有適配過，所以生成效果很差；越到后面的層，生成的結(jié)果越接近最終結(jié)果，最終不斷收斂到最終結(jié)果。而且，中間層有時(shí)會“改變主意”，例如，在第07層時(shí)模型早已預(yù)測出標(biāo)記“range”，但在第22至26層間改變了主意，之后又重新確定為“range”。
- 大部分情況下不需執(zhí)行所有層，提前退出計(jì)算也可獲得最終正確的結(jié)果。作者發(fā)現(xiàn)，大多數(shù)時(shí)候，最終的 token 在最終前幾層就已經(jīng)被預(yù)測出來。在32層的LLM中，一個(gè)token平均需要23.45層被預(yù)測出來。如果有一個(gè)零計(jì)算消耗的預(yù)測器能夠完美預(yù)測在哪一層退出，我們可以降低26%的計(jì)算量。因此，有必要使LLM用更少的層來預(yù)測每個(gè)token以及如何“改變主意”。作者希望能夠降低對模型后面層的依賴，只用后面這些層來預(yù)測困難的 token。避免像上面的例子中，用了32層來預(yù)測 for 循環(huán)的 “for”這種情況。
方案
所以，針對以上兩個(gè)現(xiàn)象，作者提出 LayerSkip，利用“早退”+ 投機(jī)采樣優(yōu)化推理性能。其中，早退層（Early-exiting）是在模型的非最終的ransformer層進(jìn)行 Unembedding 操作，相當(dāng)于跳過了剩余的Transformer層。具體貢獻(xiàn)為：
- 訓(xùn)練側(cè)：早退層沒有和 LMHead 適配，導(dǎo)致生成效果很差，所以在訓(xùn)練時(shí)候需要前面的層適應(yīng)早退，讓模型具有一定的直接生成 Token 的能力。具體而言，訓(xùn)練時(shí)針對早退層添加了一個(gè) Early Exit Loss，增強(qiáng)早退層的生成能力。
- 推理側(cè)：采用自投機(jī)解碼策略，利用早退層推理速度快的特點(diǎn)來生成多個(gè) draft tokens，然后驗(yàn)證階段跑全量層，將通過驗(yàn)證的 draft tokens 返回。下圖右側(cè)所示的是自投機(jī)解碼策略整體架構(gòu)。淺綠色節(jié)點(diǎn)表示 緩存到顯存的狀態(tài) KVCache，深綠色節(jié)點(diǎn)表示需要進(jìn)行計(jì)算，透明節(jié)點(diǎn)表示跳過計(jì)算。通過下圖可知，僅早退層及之前的層需要 KVCache，后續(xù)的層都不需要，所以 LayerSkip 有兩大優(yōu)化：加速預(yù)測性能；降低顯存開銷，潛在提升服務(wù)最大吞吐。
3.4.3 YOCO
LayerSkip主要是用“裁剪”層方式，僅用部分層的 KVCache 完成全量結(jié)果的計(jì)算。除了“裁剪”層方式，還可以通過“共享”層方式壓縮 KVCache，或者說是跨層推理（ Cross-Attention 機(jī)制）。
論文”You Only Cache Once: Decoder-Decoder Architectures for Language Models“提出了YOCO方案，進(jìn)行KV Cache跨層共享。選擇固定的前幾層來產(chǎn)生KV Cache。
YOCO 構(gòu)建了一個(gè)由兩個(gè)解碼器模塊組成的雙解碼器架構(gòu)：自解碼器和交叉解碼器。自解碼器有效地對全局鍵值緩存進(jìn)行編碼，而交叉解碼器通過交叉注意重用這些緩存。整個(gè)模型的執(zhí)行過程和 Decoder Only 的 Transformer 模型類似，但選擇固定的前幾層來產(chǎn)生KV Cache，只保留一層全局的 KV Cache。YOCO的計(jì)算流程還使預(yù)填充能夠提前退出，從而在不改變最終輸出的情況下實(shí)現(xiàn)更快的預(yù)填充階段，降低GPU內(nèi)存使用。 這種設(shè)計(jì)可以大大降低 GPU 顯存的需求，同時(shí)保持了全局注意力能力。
下圖給出了論文“A Survey on Large Language Model Acceleration based on KV Cache Management”對層內(nèi)模型改造方案的總結(jié)，可以看到Y(jié)OCO的特點(diǎn)。
模型架構(gòu)
與 LLaMa 等 Decoder Only 模型不同，YOCO 采用 Decoder-Decoder 模型架構(gòu)，之所以命名為Decoder-Decoder架構(gòu)，是因?yàn)檫@兩個(gè)Decoder的含義不是完全一致的。模型架構(gòu)如圖所示。
YOCO 的兩大Decoder是：負(fù)責(zé)產(chǎn)生global KV Cache的Self-Decoder 和 負(fù)責(zé)使用global KV Cache的Cross-Decoder。
Self-Decoder的特點(diǎn)如下：
- YOCO 的前 L/2 層由 Self-Decoder 層堆疊，由最后一層（第 L/2 層）Self-Decoder 層輸出 global KV Cache，也就是只有一層有全局 KV Cache。
- Self-Decoder和常見的Decoder是一樣的。但是，如果僅是將后一半的 Transformer Block 替換為 Cross-Decoder，那么顯存節(jié)約和計(jì)算節(jié)約的上限也是確定的，最多節(jié)約一半。因此 Self-Decoder 的設(shè)計(jì)也是非常關(guān)鍵的，要起到大幅節(jié)約顯存和降低計(jì)算量的目的。YOCO選擇了Efficient Self-Attention，比如Slide Window Attention和YOCO作者之前其它論文提出的 Multi-Scale Retention。其只用保存窗口內(nèi)的 KV Cache 即可。其計(jì)算量和 KV Cache 大小與窗口大小或 Chunk Size 有關(guān)，而不是與序列長度有關(guān)。
- Self-Decoder產(chǎn)生的KV Cache會直接被后續(xù)的Cross-Decoder使用。
Cross-Decoder的特點(diǎn)如下：
- YOCO 的后 L/2 層由 Cross-Decoder 層堆疊。
- Cross-Decoder本身不產(chǎn)生KV Cache，而是由第一個(gè) Cross-Decoder 層加載使用Self-Decoder產(chǎn)生的KVCache。也就是后續(xù)所有 L/2 層的 Cross Attention 的 KV Cache 都是相同的。在 Decoding Stage，每個(gè) Step 中都會把第 L/2 層新生成的 KV Cache 擴(kuò)充到 Global KV Cache 中。
- Cross-Decoder使用KV Cache做交叉注意力。
訓(xùn)練和推理
由于只有一層 KVCache，所以推理時(shí) Prefill 和 Decode 流程均和傳統(tǒng)的 Transformer 不同。YOCO在推理階段，可以節(jié)省大量Prefill的耗時(shí)，
- Prefill 階段：基于YOCO的特性，Prefill階段可以跳過首Token的生成，也就是跳過Cross-Decoder。在傳統(tǒng)的 Transformer 結(jié)構(gòu)的模型中，每一層都會計(jì)算 KVCache，所以在 Prefill 階段每一層都會運(yùn)行，并且生成 KVCache；而 YOCO 只有 Self-Decoder 才會生成 KVCache，并且生成下一個(gè) Token 取決于當(dāng)前 Token 的值，所以在 Prefill 時(shí)所有輸入的 Token 經(jīng)過 Self-Decoder 計(jì)算后，僅最后一個(gè) Token 需要運(yùn)行后續(xù) Cross-Decoder，完成首 Token 的生成。
- Decode 階段：在該階段，模型輸入只有當(dāng)前一個(gè) Token，輸出下一個(gè) Token。所以在 Self-Decoder 僅對當(dāng)前Token 進(jìn)行 Attention 計(jì)算，類似對一個(gè) Token 進(jìn)行 Prefill，生成該 Token 的 KVCache，然后在 Cross-Decoder 進(jìn)行Decode Attention計(jì)算，生成下一個(gè) Token。
效果
下圖為 YOCO 與原始 Transformer 相比 KV Cache 的占用對比，以及計(jì)算量對比，其中假設(shè)序列長度為 N，transformer層數(shù)為 L，D是隱狀態(tài)維度。
- Prefill階段的KV Cache顯存需求從O(LND)下降到O((N+L)D)。其中，O(ND) 對應(yīng) Global KV Cache，只有一層，前 N/2 層使用了滑動窗口 Attention 或 Multi-Scale Retention，對應(yīng)的 KV Cache 為 O(DL)。
- Prefill耗時(shí)則從O(LN^2D)下降為O(LND)，即耗時(shí)從N平方復(fù)雜度變?yōu)榫€性復(fù)雜度。因?yàn)镻refill 階段只用計(jì)算前一半的 Self-Decoder Layer 來產(chǎn)生 KV Cache，后一半的 Cross-Decoder Layer 直接生成即可，所以Prefill 階段計(jì)算量降低一半。
3.4.4 CLA
CLA（跨層注意力）：該方法將產(chǎn)生KV Cache的層交替分布在模型不同深度的層，然后通過在相鄰層之間共享鍵和值頭來擴(kuò)展GQA和MQA的思想，進(jìn)一步減少了KV緩存中的冗余。與MQA相比，CLA可以再減少2倍KV Cache大小，在不改變計(jì)算復(fù)雜性的情況下顯著提高了內(nèi)存效率。
論文”Reducing Transformer Key-Value Cache Size with Cross-Layer Attention“提出了CLA（ Cross-Layer Attention），即KV Cache跨層注意力，和YOCO思路不謀而合。和YOCO不一樣的是，CLA并不是選擇固定的前幾層來產(chǎn)生KV Cache（比如YOCO，使用的是前L/2層），而是將產(chǎn)生KV Cache的層交替分布在模型不同深度的層，然后鄰近層復(fù)用附近層產(chǎn)生的KV Cache進(jìn)行Attention計(jì)算。
CLA可以選擇多種跨層共享KV Cache的方式，比較靈活，也比較通用。如下圖左所示。最簡單的方式就是隔層共享，稱作 CLA2。實(shí)際也可以每 3 層共享，稱作 CLA3，如下圖右所示。其中C表示多少層共享一份KV Cache，比如C=2表示CLA2模式，相鄰的2個(gè)層共享一份KV Cache，則KV Cache總量可以節(jié)省一半。即，CLA2 顯存減小 2x，CLA3 顯存減小 3x。
下圖使用傳統(tǒng)注意力設(shè)計(jì)的Transformer（左）和使用CLA的Transformer（右）中兩個(gè)連續(xù)層的示意圖。當(dāng)使用傳統(tǒng)注意力時(shí)，每一層都會計(jì)算自己單獨(dú)的K和V激活，這些激活必須在自回歸解碼期間按層緩存。當(dāng)使用CLA時(shí)，CLA僅計(jì)算模型中一部分層的鍵/值預(yù)測；沒有鍵/值投影的層中的注意力塊重用了先前層的KV激活。只有具有鍵/值投影的層的子集對KV緩存有貢獻(xiàn)，與在每一層中應(yīng)用單獨(dú)鍵/值預(yù)測的傳統(tǒng)架構(gòu)相比，CLA可以減少內(nèi)存占用。
另外，CLA提供的是一種通用的跨層KV Cache共享思路，其本質(zhì)上和MQA/GQA等層內(nèi)KV Cache共享的方式是不沖突的，即MQA和GQA是在單個(gè)層內(nèi)跨查詢頭共享鍵/值頭，而CLA為層間KV Cache共享。因此，CLA可以和MQA/GQA結(jié)合一起使用。
3.4.5 MLKV
MLKV（多層鍵值/Multi-Layer Key-Value）引入了一種跨多個(gè)transformer 層的簡單KV頭共享機(jī)制。MLKV在一個(gè)層內(nèi)使用與MQA相同的單個(gè)KV頭，但它也與多個(gè)層共享這個(gè)KV頭。這種極端策略將緩存大小減少到正常GQA策略的近1%，實(shí)驗(yàn)表明MLKV仍然具有相當(dāng)?shù)男阅堋?/p>
從創(chuàng)新點(diǎn)上看，MLKV和YOCO都是CLA的特例。MLKV和CLA一樣，都是做跨層的KV Cache共享，但MLKV主要是對MQA做更加極端的擴(kuò)展，也就是MQA+跨層KV Cache共享，具體如下圖所示。
下圖則給出了MLKV和其它幾個(gè)方案的區(qū)別。
- MHA：原始的 Multi Head Attention，每一層的每一個(gè) Head 都有獨(dú)立的 K 和 V。
- MQA：Multi Query Attention，每一層的所有 Head 共享 K 和 V.
- GQA：Grouped Query Attention，MHA 和 MQA 的折衷，每一層的 Head 分為多組，每一組共享 K 和 V.
- MLKV：多個(gè)層共享 K 和 V，并且可以與上述 MQA 和 GQA 兼容。
下圖則做了進(jìn)一步細(xì)化。
3.5 減少序列長度
因此這個(gè)方向內(nèi)容較多，我們將在本文的下一節(jié)來介紹如何減少序列長度。
0x04 按照特性進(jìn)行優(yōu)化
我們在本節(jié)來對“按照特性進(jìn)行優(yōu)化”進(jìn)行深入學(xué)習(xí)。
4.1 針對Prefill特性進(jìn)行優(yōu)化
相比于更耗時(shí)的Decode階段，Prefill有一個(gè)更加突出的問題：長上下文。過長的prompt會給顯存帶來壓力。因此，研究人員針對此特性進(jìn)行相關(guān)優(yōu)化。
4.1.1 Chunking
Chunking把prompt切成若干chunk，每次只喂給模型1個(gè)chunk，更新1次KV Cache。因?yàn)闆]有把所有token一起計(jì)算，Chunking犧牲了Prefill計(jì)算的并行性。但是卻可以緩解顯存壓力。
下圖是chunking的展示，輸入的prompt為The cat sat on the mat and saw the dog go to，同時(shí)chunk_size = cache_window = sliding_window = 4，即chunk、cache的尺寸都和滑動窗口的尺寸保持一致。這樣無論是從計(jì)算邏輯還是空間利用率上，都是更好的選擇。句子就被分成三塊："the cat sat on"，"the mat and saw”，“the dog go to”。
圖上的行方向表示query，列方向表示key和value，整個(gè)0/1數(shù)據(jù)塊表示mask矩陣，即行方向的query和列方向的key、value是否做注意力計(jì)算。圖下方的Past表示前一個(gè)時(shí)刻的cache，已經(jīng)被當(dāng)前cache刷新。Cache表示當(dāng)前真正的cache，Current表示即將被更新進(jìn)cache的Xk, Xv值。
當(dāng)我們遍歷到某個(gè)chunk時(shí)，我們?nèi)〕霎?dāng)前的cache和這個(gè)chunk做attention計(jì)算，然后再把這個(gè)chunk相關(guān)的KV值按Rolling Buffer Cache的方式更新進(jìn)這個(gè)cache中。
4.1.2 KV-Runahead
KV-Runahead的主要思路就是讓前面的層多做點(diǎn)，后面層少做點(diǎn)，借此加速提示階段。
動機(jī)
KV-Runahead的關(guān)鍵觀察是，由于鍵值緩存（KV-cache），擴(kuò)展階段生成標(biāo)記的速度比提示階段更快。因此，KV-Runahead通過協(xié)調(diào)多個(gè)進(jìn)程來填充 KV-cache 并最小化第一個(gè)標(biāo)記時(shí)間 (TTFT)，從而并行化提示階段。
雙重利用KV Cache方案有幾個(gè)主要好處。
- 由于KV Cache旨在利用因果注意力圖，我們自動最小化了計(jì)算和通信。
- 由于KV Cache已經(jīng)存在于擴(kuò)展階段，因此KV-Runahead易于實(shí)現(xiàn)。
- 我們進(jìn)一步提出了上下文級負(fù)載平衡，以處理因果注意力導(dǎo)致的不均勻KV Cache生成，并優(yōu)化TTFT。
方案
下圖展示了KV-Runahead與張量/序列并行推理（TSP/Tensor+Sequence Parallel Inference）的比較。
- TSP：TSP的重點(diǎn)是并行化。通過均勻的上下文劃分，(Q, K, V) 在每個(gè)進(jìn)程上獨(dú)立計(jì)算。然后，執(zhí)行collective操作all-gather以交換K和V到所有進(jìn)程，使得QKT可以均勻分布。層中的所有計(jì)算都是對稱的，全局同步的，但是TSP沒有利用LLM推理中的因果性。
- KV-Runahead：以不均勻的上下文分區(qū)開始，(Q, K, V)在每個(gè)進(jìn)程上獨(dú)立計(jì)算。因此，每個(gè)進(jìn)程計(jì)算(Q, K, V)中的不同數(shù)量的條目。KV-Runahead簡單地從每個(gè)進(jìn)程填充KV Cache并交給鏈中的下一個(gè)進(jìn)程，以模擬擴(kuò)展階段。結(jié)果是，只有最后一個(gè)進(jìn)程將擁有完整的(K, V)，但每個(gè)進(jìn)程仍然可以輸出與Q形狀相同的A，驅(qū)動下一層。由于后續(xù)進(jìn)程等待KV緩存就緒，各層將異步通信，并且用戶上下文被不均勻地劃分（用于負(fù)載平衡）以最小化TTFT。由于KV Cache本身是建立在因果性之上的，KV-Runahead可以自動最小化上三角的計(jì)算并減少Q(mào)KT的點(diǎn)積數(shù)量。
KV-Runahead通過利用多個(gè)進(jìn)程來為最后一個(gè)進(jìn)程填充KV Cache，從而實(shí)現(xiàn)并行推理。因此，KV-Runahead需要良好的上下文分區(qū)（建議進(jìn)行離線搜索以找到最佳分區(qū)方案，然后將結(jié)果存儲在分區(qū)查找表中）來平衡每個(gè)計(jì)算進(jìn)程的KV Cache量并最小化TTFT。一旦分區(qū)完成，每個(gè)進(jìn)程將基于前一個(gè)進(jìn)程的KV Cache進(jìn)行計(jì)算。具體來說，當(dāng)前進(jìn)程必須等待前一個(gè)進(jìn)程傳來的所需KV Cache（即注意圖(b)中的層錯(cuò)位），通過本地點(diǎn)對點(diǎn)通信而不是通過all-gather進(jìn)行全局同步，形成依賴鏈。
下圖是使用BLAS矩陣矩陣乘法的QKT計(jì)算覆蓋率。通過將每個(gè)上下文分區(qū)鏈接到KV緩存，我們可以逼近（approximate ）下三角部分并最小化不必要的點(diǎn)積。注意力的上三角部分將被掩蓋，以強(qiáng)制因果關(guān)系。
4.2 內(nèi)存管理優(yōu)化
4.2.1 問題
現(xiàn)有的大語言模型（LLM）服務(wù)系統(tǒng)在管理KV緩存內(nèi)存方面效率不高。這主要是因?yàn)樗鼈儼袺V緩存存儲在連續(xù)的內(nèi)存空間中（大多數(shù)深度學(xué)習(xí)框架要求張量必須存儲在連續(xù)的內(nèi)存中）。然而，與傳統(tǒng)深度學(xué)習(xí)工作負(fù)載中的張量不同，KV緩存有其獨(dú)特的特點(diǎn)（每個(gè)請求都需要大量的內(nèi)存，且這些內(nèi)存的需求是動態(tài)變化的）：
- 它隨著模型生成新詞而動態(tài)增長和收縮。而且，每個(gè)請求的KV緩存增長緩慢，即每次迭代增加一個(gè)token。
- 它的生命周期和長度是未知的。無法預(yù)先預(yù)判KV緩存的總大小。
這些特點(diǎn)使得現(xiàn)有系統(tǒng)的管理方式顯得非常低效。在現(xiàn)有系統(tǒng)的塊預(yù)分配方案中，具體存在以下幾種主要的內(nèi)存浪費(fèi)來源：為潛在最大序列長度過度分配導(dǎo)致的內(nèi)部碎片化，為未來token保留的插槽、以及來自內(nèi)存分配器（如伙伴分配器）的外部碎片化。
- 為潛在最大序列長度過度分配導(dǎo)致的內(nèi)部碎片化。對應(yīng)下圖上的圓形標(biāo)號1。
- 由于請求的總序列長度事先不確定，為了在連續(xù)空間中存儲請求的KV緩存，這些系統(tǒng)通常會預(yù)測請求的最大可能序列長度，從而提前為一個(gè)請求靜態(tài)分配一塊連續(xù)的存儲空間，以容納最大的序列長度。而不管實(shí)際的輸入或最終的輸出長度。因?yàn)檎埱蟮膶?shí)際長度往往比最大長度短很多，這種分配中的一部分幾乎肯定永遠(yuǎn)不會被使用，并且也無法被其他請求利用，最終被浪費(fèi)（即內(nèi)部內(nèi)存碎片化）。由于內(nèi)部碎片而浪費(fèi)了大量的GPU內(nèi)存。因此，這些系統(tǒng)的吞吐量較差，無法支持大批量大小。
- KV緩存是矩形的：你分配一個(gè)大的矩形內(nèi)存，其中一個(gè)維度是批量大?。Ｐ鸵淮慰梢蕴幚淼淖畲笮蛄袛?shù)），另一個(gè)維度是允許用戶使用的最大序列長度。當(dāng)一個(gè)新序列進(jìn)入時(shí)，你會為這個(gè)用戶分配整行內(nèi)存，假設(shè)預(yù)測每個(gè)用戶最大上下文長度是2048，則需要為每個(gè)用戶預(yù)先分配一個(gè)可支持 2048 個(gè) tokens 緩存的顯存空間。如果用戶實(shí)際使用的上下文長度低于2048，這種方法顯然會造成存儲資源的大量浪費(fèi)。
- 提前為未來token保留的插槽。對應(yīng)下圖上的圓形標(biāo)號2。即使實(shí)際序列長度事先已知，預(yù)分配的方式仍然低效。因?yàn)槊總€(gè)請求的KV-cache每次增加一個(gè)標(biāo)記，而在請求的生命周期內(nèi)，整個(gè)內(nèi)存塊都會被保留，由于內(nèi)存逐漸被消耗，但內(nèi)存塊被保留到請求的整個(gè)生命周期，較短的請求仍然無法使用尚未使用的內(nèi)存塊。
- 來自內(nèi)存分配器（如伙伴分配器）的外部碎片化和其它損耗。對應(yīng)下圖上的圓形標(biāo)號3。
- 由于KV Cache是一個(gè)巨大的矩陣，且必須占用連續(xù)內(nèi)存，操作系統(tǒng)只分配大的連續(xù)內(nèi)存，如果不同請求的預(yù)分配大小不同，勢必有很多小的內(nèi)存空間被浪費(fèi)掉，導(dǎo)致外部內(nèi)存碎片化。
- 在 Paged Attention 之前，業(yè)界主流 LLM 推理框架在 KV Cache 管理方面均存在一定的低效。HuggingFace Transformers 庫中，KV Cache 是隨著執(zhí)行動態(tài)申請顯存空間，由于 GPU顯存分配耗時(shí)一般都高于 CUDA kernel 執(zhí)行耗時(shí)，因此動態(tài)申請顯存空間會造成極大的時(shí)延開銷，且會引入顯存碎片化。
- 現(xiàn)有系統(tǒng)無法利用內(nèi)存共享的機(jī)會。LLM服務(wù)常常使用高級解碼算法，比如并行采樣和束搜索，這些算法會為每個(gè)請求生成多個(gè)輸出。在這些場景中，請求包含的多個(gè)序列可以部分共享其KV緩存。然而，現(xiàn)有系統(tǒng)中，序列的KV緩存存儲在不同的連續(xù)空間中，無法實(shí)現(xiàn)內(nèi)存共享。
另外，盡管已經(jīng)提出了壓縮作為解決碎片化的潛在方案，但在一個(gè)對性能敏感的LLM服務(wù)系統(tǒng)中執(zhí)行壓縮是不現(xiàn)實(shí)的，因?yàn)镵V緩存非常龐大。即使有壓縮，每個(gè)請求的預(yù)分配塊空間也阻礙了現(xiàn)有內(nèi)存管理系統(tǒng)中特定解碼算法的內(nèi)存共享。
4.2.2 Paged Attention
PagedAttention (vLLM) 借鑒了操作系統(tǒng)中的分頁（Page）方式進(jìn)行存儲，將KV Cache分為固定大小塊，存儲在非連續(xù)的物理內(nèi)存中，這樣可以為vLLM提供更靈活的分頁內(nèi)存管理。通過使用相對較小的小塊并按需分配，PagedAttention緩解了內(nèi)部內(nèi)存碎片化問題。此外，由于所有小塊大小相同，它還消除了外部內(nèi)存碎片化問題。PagedAttention也可以使同一請求的不同序列甚至不同請求之間可以在小塊級別共享內(nèi)存，實(shí)現(xiàn)了內(nèi)存共享。
下圖給出了因?yàn)閮?nèi)部碎片造成的內(nèi)存浪費(fèi)，陰影框表示KVtoken占用的內(nèi)存，而白色框表示已分配但未使用的內(nèi)存。
- Orca（頂部）提前為KV緩存保留一塊連續(xù)的虛擬內(nèi)存。分配的大小對應(yīng)于模型支持的最大序列長度，例如32K。由于模型生成的 token 通常遠(yuǎn)少于最大限制，因此由于內(nèi)部碎片而浪費(fèi)了大量GPU內(nèi)存。
- vLLM（底部）通過動態(tài)內(nèi)存分配減輕碎片化。
實(shí)現(xiàn)
PagedAttention的靈感來源于操作系統(tǒng)中的虛擬內(nèi)存和分頁技術(shù)。操作系統(tǒng)將內(nèi)存劃分為固定大小的頁面，并將用戶程序的邏輯頁面映射到物理頁面。連續(xù)的邏輯頁面可以對應(yīng)非連續(xù)的物理內(nèi)存頁面，使得用戶程序可以像訪問連續(xù)內(nèi)存一樣訪問內(nèi)存。PagedAttention將KV緩存組織成固定大小的KV塊，類似于虛擬內(nèi)存中的頁面。而vLLM實(shí)現(xiàn)了一個(gè)類似虛擬內(nèi)存的系統(tǒng)，通過復(fù)雜的映射機(jī)制管理這些塊。這種架構(gòu)將邏輯和物理KV塊分開，通過跟蹤映射關(guān)系和填充狀態(tài)的塊表實(shí)現(xiàn)動態(tài)內(nèi)存分配和靈活的塊管理。具體特點(diǎn)如下。
- 為了能夠動態(tài)分配物理內(nèi)存，PagedAttention 將 KV Cache 的布局從連續(xù)的虛擬內(nèi)存更改為非連續(xù)的虛擬內(nèi)存 。
- PagedAttention將KV緩存分割為固定大小且相對較小的內(nèi)存塊，稱為塊，每個(gè)塊可以包含固定數(shù)量的詞元的鍵和值。這些塊類似于進(jìn)程虛擬內(nèi)存中的頁面，因此我們可以像操作系統(tǒng)管理虛擬內(nèi)存那樣靈活地管理KV緩存。PagedAttention每次為一個(gè)塊分配內(nèi)存。這些小塊不需要存儲在連續(xù)的空間中，可以在GPU和CPU內(nèi)存之間進(jìn)行交換，并且必要時(shí)可以在不同的請求之間共享。
- vLLM不是預(yù)先保留KV緩存內(nèi)存的最大序列長度，而是只在需要時(shí)分配請求所需的內(nèi)存，即允許生成過程中不斷為用戶追加顯存，而不是提前分配。通過 token 塊粒度的顯存管理，系統(tǒng)可以精確計(jì)算出剩余顯存可容納的 token 塊的個(gè)數(shù)，配合 Dynamic Batching 技術(shù)，即可避免系統(tǒng)發(fā)生顯存溢出的問題。
- PagedAttention實(shí)現(xiàn)一個(gè)塊表來將這些非連續(xù)的分配拼接在一起。也實(shí)現(xiàn)了一個(gè)高效的顯存管理器，通過精細(xì)化管理顯存，實(shí)現(xiàn)了在物理非連續(xù)的顯存空間中以極低的成本存儲、讀取、新增和刪除鍵值向量。
PagedAttention涉及兩層地址轉(zhuǎn)換。首先，該框架通過塊表跟蹤KV Cache塊的虛擬內(nèi)存地址。其次，GPU通過操作系統(tǒng)管理的頁表執(zhí)行虛擬到物理地址的轉(zhuǎn)換。具體如下圖所示。
架構(gòu)
vLLM的架構(gòu)如下圖所示，具體分為量部分：
- Scheduler。vLLM采用集中式調(diào)度器來協(xié)調(diào)分布式GPU worker的執(zhí)行和對應(yīng)的物理KV緩存內(nèi)存。
- KV Cache Manager。KV緩存管理器通過PagedAttention以分頁方式有效地管理KV緩存。
KV Cache管理器
PagedAttention首先加載模型，這樣就知道剩余多少空間給KV Cache，后續(xù)會用內(nèi)存塊填充這些空間。這些塊可以包含最多16或32個(gè)token。
系統(tǒng)會動態(tài)維護(hù)一個(gè)顯存塊的鏈表——每個(gè)請求的邏輯和物理KV塊之間的映射關(guān)系。每個(gè)塊表?xiàng)l目記錄了邏輯塊的相應(yīng)物理塊及其填充位置的數(shù)量，通過該方式即可實(shí)現(xiàn)將地址不連續(xù)的物理塊串聯(lián)在一起統(tǒng)一管理。當(dāng)一個(gè)請求來到之后，系統(tǒng)會為這個(gè)請求分配一小塊顯存，這一小塊顯存通常只夠容納 8 個(gè)字符，當(dāng)請求生成了 8 個(gè)字符之后，系統(tǒng)會追加一塊顯存，請求可以把結(jié)果再寫到這塊顯存里面。當(dāng)生成的長度不斷變長時(shí)，系統(tǒng)會不斷地給用戶追加顯存塊的分配，并動態(tài)更新顯存塊列表。將邏輯和物理KV塊分開使得vLLM可以動態(tài)增長KV緩存內(nèi)存，而無需提前為所有位置預(yù)留空間，這消除了現(xiàn)有系統(tǒng)中大部分的內(nèi)存浪費(fèi)。
下圖是算法。PagedAttention將每個(gè)序列的鍵值（KV）緩存劃分為KV塊。每個(gè)塊包含一定數(shù)量的token的鍵和值向量。例如，鍵塊\(K_j\)包含的是第j個(gè)塊中的鍵向量，而值塊\(V_j\)包含的是第j個(gè)塊中的值向量。在注意力計(jì)算過程中，PagedAttention內(nèi)核分別識別和獲取不同的KV塊。例如，當(dāng)處理查詢標(biāo)記為“forth”時(shí)，內(nèi)核將查詢向量\(q_i\)與每個(gè)塊中的鍵向量\(K_j\)相乘，以計(jì)算注意力得分\(A_{ij}\)。然后，它將這些得分與每個(gè)塊中的值向量\(V_j\)相乘，得出最終的注意力輸出。
基本場景的解碼算力
下圖是vLLM是如何執(zhí)行PagedAttention，并在單個(gè)輸入序列的解碼過程中管理內(nèi)存的示例?？偟膩碚f，對于每個(gè)解碼迭代，vLLM首先選擇一組候選序列進(jìn)行批處理，并為新需求的邏輯塊分配物理塊。然后，vLLM將當(dāng)前迭代的所有輸入標(biāo)記（即提示階段請求的所有標(biāo)記和生成階段請求的最新標(biāo)記）串聯(lián)為一個(gè)序列，并將其輸入到LLM中。
- vLLM在初始階段不需要為最大可能生成的序列長度預(yù)留內(nèi)存。相反，它只保留了足夠容納提示計(jì)算期間生成的KV緩存所需的KV塊。在本例中，提示有7個(gè)標(biāo)記，因此vLLM將前兩個(gè)邏輯KV塊（0和1）映射到兩個(gè)物理KV塊（7和1）。在預(yù)填充步驟中，vLLM使用傳統(tǒng)的自注意力算法生成提示和第一個(gè)輸出標(biāo)記的KV緩存。然后，vLLM將前4個(gè)標(biāo)記的KV緩存存儲在邏輯塊0中，后續(xù)3個(gè)標(biāo)記存儲在邏輯塊1中。剩余的位置留待后續(xù)自回歸生成階段使用。此處對應(yīng)下圖上標(biāo)號1。
- 在第一個(gè)自回歸解碼步驟中，vLLM使用PagedAttention算法在物理塊7和1上生成新的標(biāo)記。由于最后一個(gè)邏輯塊中還有一個(gè)空位，新生成的KV緩存被存儲在那里，并且更新了塊表中的已填充記錄。
- 在第二個(gè)自回歸解碼步驟中，由于最后一個(gè)邏輯塊已滿，vLLM將新生成的KV緩存存儲在一個(gè)新的邏輯塊中；vLLM為其分配一個(gè)新的物理塊（物理塊3）并將此映射存儲在塊表中。
劣勢
然而，為了能夠動態(tài)分配物理內(nèi)存，PagedAttention也增加了軟件復(fù)雜性、導(dǎo)致可移植性問題、冗余和低效。
- 需要重新編寫注意力核心（GPU代碼），以便注意力核心通過索引查找到KV緩存（虛擬連續(xù)對象）的所有元素。
- 增加了軟件復(fù)雜性和冗余（CPU代碼）。PagedAttention還迫使開發(fā)人員在服務(wù)框架中實(shí)現(xiàn)內(nèi)存管理器，使其負(fù)責(zé)（解）分配KV緩存并跟蹤動態(tài)分配的KV緩存塊的位置。這種方法實(shí)質(zhì)上是在用戶代碼中重新實(shí)現(xiàn)了按需分頁-這是一個(gè)操作系統(tǒng)功能。
- 引入了性能開銷。PagedAttention在關(guān)鍵路徑上增加了運(yùn)行時(shí)開銷。首先，它需要GPU核心執(zhí)行與從非連續(xù)內(nèi)存塊獲取KV緩存相關(guān)的額外代碼。我們展示了這在許多情況下可能會使注意力計(jì)算速度減慢超過10%。其次，用戶空間內(nèi)存管理器可能會增加CPU開銷，導(dǎo)致額外的10%成本。
4.2.3 vAttention
vAttention作者認(rèn)為保留KV Cache的虛擬內(nèi)存連續(xù)性對于減少LLM部署中的軟件復(fù)雜性和冗余至關(guān)重要，并主張利用操作系統(tǒng)中現(xiàn)有的虛擬內(nèi)存抽象來管理動態(tài)KV Cache內(nèi)存，從而實(shí)現(xiàn)簡化部署和更高性能?；谶@些思考，vAttention作者提出了vAttention，其特點(diǎn)如下。
- 將KV Cache存儲在連續(xù)的虛擬內(nèi)存中，而不是事先分配物理內(nèi)存。作為對比，PagedAttention在虛擬地址上是不連續(xù)的。
- 為KV Cache保留連續(xù)的虛擬空間，并按需分配物理內(nèi)存。
- 在不需要在服務(wù)系統(tǒng)中實(shí)現(xiàn)內(nèi)存管理器的情況下，實(shí)現(xiàn)了更高的可移植性，同時(shí)可以無縫重用現(xiàn)有的GPU核心。使注意力核心開發(fā)人員不必顯式支持分頁，并避免在服務(wù)框架中重新實(shí)現(xiàn)內(nèi)存管理，實(shí)現(xiàn)了對各種注意力核的無縫動態(tài)內(nèi)存管理。
洞察
vAttention對LLM服務(wù)系統(tǒng)有一些洞察。
-
觀察1：在每次迭代中，KV Cache的內(nèi)存需求是可以預(yù)先知道的。這是由于自回歸解碼每次只生成一個(gè)標(biāo)記。因此，隨著每次迭代，請求的KV Cache內(nèi)存占用量會均勻地增加一個(gè)標(biāo)記，直到請求完成。
-
觀察2：KV Cache不需要高內(nèi)存分配帶寬。在所有層中，單個(gè)標(biāo)記的內(nèi)存占用通常為幾十到幾百KB。此外，每次迭代持續(xù)10到100毫秒，這意味著每秒最多需要幾兆字節(jié)的內(nèi)存。雖然批處理可以提高系統(tǒng)的吞吐量，但每秒生成的標(biāo)記數(shù)量在達(dá)到一定批大小后就會趨于平穩(wěn)。這意味著內(nèi)存分配帶寬需求在大批大小下也會飽和。
設(shè)計(jì)概述
vAttention目標(biāo)是通過為現(xiàn)有的kernel添加動態(tài)內(nèi)存分配支持來提高效率和可移植性。為了實(shí)現(xiàn)這一目標(biāo)，vAttention利用系統(tǒng)支持的動態(tài)內(nèi)存分配，而不是在用戶空間中實(shí)現(xiàn)分頁。
vAttention基于將虛擬內(nèi)存和物理內(nèi)存分別分配的能力進(jìn)行構(gòu)建。具體來說，vAttention預(yù)先在虛擬內(nèi)存中為KV-cache分配一個(gè)大的連續(xù)緩沖區(qū)（類似于基于保留的分配器），而將物理內(nèi)存的分配推遲到運(yùn)行時(shí)，即只在需要時(shí)分配物理內(nèi)存（類似于PagedAttention）。這樣，vvAttention在不浪費(fèi)物理內(nèi)存的情況下保留了KV Cache的虛擬連續(xù)性。這種方法之所以可行，是因?yàn)閮?nèi)存容量和碎片化僅對物理內(nèi)存構(gòu)成限制，而虛擬內(nèi)存是充裕的，例如，現(xiàn)代64位系統(tǒng)為每個(gè)進(jìn)程提供128TB的用戶管理虛擬地址空間。
-
預(yù)留虛擬內(nèi)存。由于虛擬內(nèi)存是豐富的，我們預(yù)先分配足夠大的虛擬內(nèi)存空間，以容納所需支持的最大批次大小的KV-cache（可配置）。虛擬內(nèi)存緩沖區(qū)的數(shù)量：每個(gè)LLM中的每個(gè)層都維護(hù)自己的K和V張量：我們將它們分別稱為K-cache和V-cache。我們?yōu)镵-cache和V-cache分別分配單獨(dú)的虛擬內(nèi)存緩沖區(qū)。
-
按需分配物理內(nèi)存。vAttention優(yōu)先按頁一次性分配物理內(nèi)存，僅當(dāng)一個(gè)請求使用完其先前分配的所有物理內(nèi)存頁時(shí)才分配。
下圖展示了vAttention如何通過按需分配物理內(nèi)存來動態(tài)管理 KV Cache 的內(nèi)存，同時(shí)通過延遲回收策略優(yōu)化內(nèi)存使用，以便在新的請求到達(dá)時(shí)可以高效地重新利用已經(jīng)分配的物理內(nèi)存。此處是處理模型的一層中的K-cache（或V-cache），假設(shè)最大批次大小為兩個(gè)。其余的K-cache和V-cache緩沖區(qū)在所有層中以相似的方式進(jìn)行管理。具體分為五個(gè)步驟：
- (a)：虛擬內(nèi)存中包含了兩個(gè)請求（R1和R2）的虛擬張量，但還沒有進(jìn)行任何物理內(nèi)存分配。淺色陰影區(qū)域代表了沒有任何物理頁支持的虛擬張量部分。
- (b)：R1被分配了一頁物理內(nèi)存。這個(gè)分配使得虛擬張量的一部分（淺色陰影區(qū)域）映射到了實(shí)際的物理內(nèi)存上（深色陰影區(qū)域）。
- (c)：R1被分配了兩頁物理內(nèi)存，而R2被分配了一頁物理內(nèi)存。此時(shí)，R1的虛擬張量有兩部分映射到了物理內(nèi)存上，而R2的虛擬張量有一部分映射到了物理內(nèi)存上。
- (d)：R1已經(jīng)完成了它的任務(wù)，但vAttention沒有立即回收其內(nèi)存（延遲回收）。因此，R1的虛擬張量依舊占據(jù)物理內(nèi)存。
- (e)：當(dāng)新的請求R3到達(dá)時(shí)，vAttention將R1的張量分配給R3，因?yàn)檫@些張量已經(jīng)有物理內(nèi)存支持。
4.3 調(diào)度
這部分主要是針對KV Cache特點(diǎn)實(shí)施的一些調(diào)度策略，主要是三種途徑：
-
前綴感知調(diào)度策略。
-
搶占式和面向公平的調(diào)度。
-
特定于層的分層調(diào)度方法。
下圖給出了一些常見方案的特點(diǎn)。
4.3.1 Prefix-aware Scheduling
在傳統(tǒng)的基于LRU的緩存管理系統(tǒng)中，共享的KV上下文可能會被過早地被驅(qū)逐，或在內(nèi)存中不必要地?cái)U(kuò)展。Prefix-aware Scheduling則會基于前綴進(jìn)行調(diào)度，即具有相同前綴的請求被放在一起調(diào)度。
- BatchLLM通過全局前綴識別和基于共享公共KV緩存內(nèi)容的動態(tài)規(guī)劃算法來將共享前綴的請求放在一起調(diào)度，這樣可以最大化KV Cache復(fù)用，減少緩存生命周期。
- RadixAttention圍繞基數(shù)樹結(jié)構(gòu)實(shí)現(xiàn)了緩存的快速匹配、插入和替換。在RadixAttention中，緩存的令牌和正在運(yùn)行的請求共享同一個(gè)內(nèi)存池，然后結(jié)合LRU驅(qū)逐和引用計(jì)數(shù)來管理內(nèi)存池。
下圖給出了BatchLLM的總體流程。
-
首先，對大批量輸入 Prompt 進(jìn)行分析，并識別出共同的前綴部分，這一步驟在 Request 調(diào)度之前完成。
-
BatchLLM 利用前綴樹來識別和表示 Prompt 之間的共同前綴。然而，直接使用前綴樹的第一層前綴作為共同前綴可能導(dǎo)致前綴復(fù)用效果不佳。因此，BatchLLM 設(shè)計(jì)了一種基于前綴樹的動態(tài)規(guī)劃算法，以最大化一級 Token 的復(fù)用。
-
隨后，為了簡化 KV 復(fù)用并縮短公共前綴在內(nèi)存中的生命周期，BatchLLM 將共享相同前綴的請求集中調(diào)度。具體而言，Batch LLM將 Request 分組，一個(gè)前綴組對應(yīng)一組共享相同前綴的 Request 的集合。這些組將成為 Request 調(diào)度的基本單元。在調(diào)度各組之前，BatchLLM 還會根據(jù)前綴 Prefill 長度與估計(jì)的 Decoding 長度的比例對組進(jìn)行重新排序，并推遲那些前綴 Prefill 較長而 Decoding 較短的組。
-
接著，充分考慮 KV 內(nèi)存使用情況來實(shí)現(xiàn) Batching，旨在更好地混合 Decoding 步驟與前綴 Prefill Chunk，以增大整體 Token Batch 規(guī)模。
-
此外，BatchLLM 還實(shí)現(xiàn)了融合 Attention Kernel，以減少長尾效應(yīng)及 Kernel 啟動開銷。
4.3.2 Preemptive and Fairness-oriented scheduling
FastServe實(shí)施了一種主動的（proactive ）KV緩存管理策略。FastServe 的特色如下：
- 創(chuàng)新調(diào)度策略：FastServe 引入了 Skip-Join Multi-Level Feedback Queue（MLFQ）調(diào)度器，利用 LLM 推理的半信息無關(guān)特性，有效避免頭阻塞問題。
- 優(yōu)先級跳過機(jī)制：根據(jù)輸入長度和初始化階段的時(shí)間，將任務(wù)分配到合適的隊(duì)列，減少了不必要的降級操作。
- 高效內(nèi)存管理：為了應(yīng)對 GPU 內(nèi)存限制，F(xiàn)astServe 設(shè)計(jì)了主動 KV 緩存管理機(jī)制，在低優(yōu)先級任務(wù)的緩存即將滿時(shí)，主動將狀態(tài)卸載到主機(jī)內(nèi)存，并在需要時(shí)重新加載。
- 數(shù)據(jù)傳輸優(yōu)化：通過異步內(nèi)存操作和流水線技術(shù)，顯著提高了數(shù)據(jù)傳輸效率。
下圖給出了FastServe的架構(gòu)。
下圖上方給出了MLFQ，下方給出了各個(gè)調(diào)度算法的執(zhí)行時(shí)間流。
FastSwitch 引入了一個(gè)公平導(dǎo)向的 KV 緩存調(diào)度系統(tǒng)，解決了在 LLM 服務(wù)中搶占式調(diào)度的開銷挑戰(zhàn)。FastSwitch 有三個(gè)關(guān)鍵機(jī)制：通過智能緩存移動調(diào)度提高 I/O 利用率，在上下文切換期間最小化 GPU 空閑時(shí)間，以及在多輪對話中消除冗余 I/O 操作。FastSwitch 與傳統(tǒng)的基于塊的 KV 緩存內(nèi)存策略不同，傳統(tǒng)策略優(yōu)先考慮內(nèi)存效率，但是該策略以碎片化和粒度限制為代價(jià)。FastSwitch 實(shí)施了一種平衡的方法，在促進(jìn)更平滑的上下文切換的同時(shí)保持高效的內(nèi)存使用。這種集成調(diào)度方法實(shí)現(xiàn)了動態(tài)優(yōu)先級調(diào)整以確保公平性，同時(shí)最小化上下文切換對性能的影響。
下圖給出了FastSwitch 的總體結(jié)構(gòu)。
4.3.3 Layer-specific and Hierarchical Scheduling
LayerKV
為了解決在大上下文 LLM 服務(wù)中日益增長的首次令牌生成時(shí)間（TTFT）延遲挑戰(zhàn)，LayerKV 引入了一種逐層（layer-wise）KV 緩存調(diào)度方法，其貢獻(xiàn)在于其細(xì)粒度的、層特定（layer-specific）的 KV 緩存塊分配和管理策略，與傳統(tǒng)的整體緩存管理方法不同。通過實(shí)施層特定的 KV 塊調(diào)度和卸載機(jī)制，LayerKV 實(shí)現(xiàn)了更高效的內(nèi)存利用，并減少了在長上下文窗口中競爭有限的 GPU KV 緩存塊時(shí)通常發(fā)生的排隊(duì)延遲。LayerKV 還使用了一個(gè) SLO 感知調(diào)度器，該調(diào)度器根據(jù)服務(wù)水平目標(biāo)優(yōu)化緩存分配決策，允許跨模型層動態(tài)管理內(nèi)存資源。下圖給出了LayerKV的總體結(jié)構(gòu)。
CachedAttention
因?yàn)槎噍唽υ挼臅r(shí)候每一輪都會重復(fù)使用前面所有的緩存數(shù)據(jù)，為了解決這個(gè)問題，讓KVCache 緩存最大限度地留在顯存中，減少offload再reload的過程，CachedAttention 引入了一種分層調(diào)度方法。該方法具體包括三層策略：
- 層特定的預(yù)加載。通過調(diào)度器感知的獲取和淘汰策略協(xié)調(diào) KV 緩存在存儲層次結(jié)構(gòu)中的移動。
- 異步保存。使 I/O 操作與 GPU 計(jì)算重疊。
- 智能緩存放置決策。該策略基于調(diào)度器提示，以確保頻繁訪問的 KV 緩存位于更快的內(nèi)存層。
CachedAttention 還提出了一種新穎的位置編碼解耦機(jī)制，通過有效的截?cái)嗖呗苑乐乖谏舷挛拇翱谝绯銎陂g KV 緩存失效。
下圖給出了總體架構(gòu)。
下圖展示了分層預(yù)加載。
下圖展示了異步保存。
下圖給出了位置編碼解耦。
(a) 是絕對位置編碼，絕對位置編碼 Key 和 Query 直接永遠(yuǎn)都是一樣的位置編碼。
(b) 是相對位置編碼。
(c) 是可以解耦合 Key 和 Query 的位置編碼。通過在RPE中嵌入位置編碼之前簡單地移動緩存KV的時(shí)間，CachedAttention可以將沒有嵌入位置編碼的KV存儲在AttentionStore中。在CachedAttention中重用KV時(shí)，KV嵌入了新的位置編碼，并進(jìn)一步用于以下推理。
ALISA
ALISA 引入了一個(gè)雙層 KV 緩存調(diào)度框架，將算法稀疏性與系統(tǒng)級優(yōu)化相結(jié)合。在算法級別，稀疏窗口注意力機(jī)制識別并優(yōu)先考慮用于注意力計(jì)算的最重要令牌，創(chuàng)建全局動態(tài)和局部靜態(tài)稀疏模式的混合，顯著減少 KV 緩存內(nèi)存需求。在系統(tǒng)級別，其三相（three-phase）令牌級動態(tài)調(diào)度器管理 KV 張量分配并優(yōu)化緩存與重新計(jì)算之間的權(quán)衡。調(diào)度器根據(jù)令牌的重要性和系統(tǒng)資源約束動態(tài)決定在 GPU 內(nèi)存中緩存哪些令牌以及重新計(jì)算哪些令牌。
下圖給出了稀疏窗口注意力（Sparse Window Attention）的計(jì)算。
下圖給出了動態(tài)調(diào)度機(jī)制。
0xFF 參考
聊聊大模型推理內(nèi)存管理之 CachedAttention/MLA 刀刀寧
Y. Xiong, H. Wu, C. Shao, Z. Wang, R. Zhang, Y. Guo, J. Zhao, K. Zhang, and Z. Pan, “Layerkv: Optimizing large language model serving with layer-wise kv cache management,” 2024. [Online]. Available: https://arxiv.org/abs/2410.00428
B. Gao, Z. He, P. Sharma, Q. Kang, D. Jevdjic, J. Deng, X. Yang, Z. Yu, and P. Zuo, “Cost-efficient large language model serving for multi-turn conversations with cachedattention,” 2024. [Online]. Available: https://arxiv.org/abs/2403.19708
Y. Zhao, D. Wu, and J. Wang, “Alisa: Accelerating large language model inference via sparsity-aware kv caching,” 2024. [Online]. Available: https://arxiv.org/abs/2403.17312
R. Shahout, C. Liang, S. Xin, Q. Lao, Y. Cui, M. Yu, and M. Mitzenmacher, “Fast inference for augmented large language models,” 2024. [Online]. Available: https://arxiv.org/abs/2410.18248
BatchLLM：離線大 Batch 場景 LLM 推理優(yōu)化 程序猿李巡天
聊聊大模型推理系統(tǒng)之 FastServe 機(jī)智流
GPU加速LLM論文夸贊｜InfiniGen通過動態(tài)KV Cache管理高效實(shí)現(xiàn)LLM的生成式推理：動機(jī) GPUer
ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition 趙尚春
Modular Attention Reuse for Low-Latency Inference(@yale.edu)
論文: 《Efficiently Programming Large Language Models using SGLang》
ZigZagKV，KVCache80%的壓縮率可以達(dá)到近乎無損的效果
Efficient Prompt Caching via Embedding Similarity(@UC Berkeley)
Modular Attention Reuse for Low-Latency Inference(@yale.edu)
33 倍 LLM 推理性能提升：Character.AI 的最佳實(shí)踐
[CLA] Reducing Transformer Key-Value Cache Size with Cross-Layer Attention
[GQA] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
[KV Cache優(yōu)化]??MQA/GQA/YOCO/CLA/MLKV筆記: 層內(nèi)和層間KV Cache共享 DefTruth
[LLM 推理服務(wù)優(yōu)化] DistServe速讀——Prefill & Decode解耦、模型并行策略&GPU資源分配解耦 阿杰
[LLM]KV cache詳解 圖示，顯存，計(jì)算量分析，代碼 莫笑傅立葉
[LLM性能優(yōu)化]聊聊長文本推理性能優(yōu)化方向 阿杰
[MLKV] MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding
[MQA] Fast Transformer Decoding: One Write-Head is All You Need
[Prefill優(yōu)化][萬字]??原理&圖解vLLM Automatic Prefix Cache(RadixAttention): 首Token時(shí)延優(yōu)化 DefTruth
[YOCO] You Only Cache Once: Decoder-Decoder Architectures for Language Models
A Survey on Efficient Inference for Large Language Models
Cascade Inference: Memory Bandwidth Efficient Shared Prefix Batch Decoding
Continuous Batching：一種提升 LLM 部署吞吐量的利器 幻方AI
DejaVu：KV Cache Streaming For LLM
Efficient Streaming Language Models with Attention Sinks
GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
https://lmsys.org/blog/2024-01-17-sglang/
Hydragen: High-Throughput LLM Inference with Shared Prefixes(@Stanford University)
Inference without Interference：混合下游工作負(fù)載的分離式LLM推理服務(wù) 竹言見智
kimi chat大模型的200萬長度無損上下文可能是如何做到的？ [GiantPandaCV](javascript:void(0)??
KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache
LaViT：這也行，微軟提出直接用上一層的注意力權(quán)重生成當(dāng)前層的注意力權(quán)重 | CVPR 2024 VincentLee
LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding
LLaMa 原理+源碼——拆解 (KV-Cache, Rotary Positional Embedding, RMS Norm, Grouped Query Attention, SwiGLU) Yuezero_
LLM Inference Series: 3. KV caching unveiled
LLM 推理優(yōu)化探微 (2) ：Transformer 模型 KV 緩存技術(shù)詳解 Baihai IDP
LLM 推理優(yōu)化探微 (3) ：如何有效控制 KV 緩存的內(nèi)存占用，優(yōu)化推理速度？ Baihai IDP
LLM推理入門指南②：深入解析KV緩存 OneFlow
LLM推理加速3：推理優(yōu)化總結(jié)Mooncake/AttentionStore/vllm0.5/cache優(yōu)化 etc
MiniCache 和 PyramidInfer 等 6 種優(yōu)化 LLM KV Cache 的最新工作 AI閑談
MLKV：跨層 KV Cache 共享，降低內(nèi)存占用
MLSys 2024 Keyformer：選擇關(guān)鍵token，降低50%KV 緩存 MLSys2024 [MLSys2024](javascript:void(0)??
MODEL TELLS YOU WHAT TO DISCARD: ADAPTIVE KV CACHE COMPRESSION FOR LLMS
Mooncake (1): 在月之暗面做月餅，Kimi 以 KVCache 為中心的分離式推理架構(gòu) ZHANG Mingxing
Mooncake (2)：Kimi “潑天的流量”怎么接，分離架構(gòu)下基于預(yù)測的調(diào)度策略 ZHANG Mingxing
Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving 手抓餅熊
Orca: How to Serve Large-scale Transformer Models WXB506
OSDI22 ORCA：LLM Serving系統(tǒng)優(yōu)化經(jīng)典之作，動態(tài)batch，吞吐提升36倍 MLSys2024
OSDI24最新力作：InfiniGen大模型推理高效KV緩存管理框架，吞吐提升3倍，準(zhǔn)確性提升32% MLSys2024
SnapKV: KV Cache 稀疏化，零微調(diào)加速長序列 LLM 推理 AI閑談
SnapKV: LLM Knows What You are Looking for Before Generation
Token 稀疏化 + Channel 稀疏化：16 倍 LLM 推理加速 AI閑談
Transformers KV Caching Explained
TriForce：KV Cache 稀疏化+投機(jī)采樣，2.3x LLM 無損加速
vAttention：用于在沒有Paged Attention的情況下Serving LLM BBuf
YOCO：長序列優(yōu)化，1M 長文本10x LLM 推理加速 AI閑談
You Only Cache Once: Decoder-Decoder Architectures for Language Models
ZHANG Mingxing：Mooncake (1): 在月之暗面做月餅，Kimi 以 KVCache 為中心的分離式推理架構(gòu)
【手撕LLM- KV Cache】為什么沒有Q-Cache ?? 小冬瓜AIGC
一文讀懂KVCache 游凱超
為什么加速LLM推斷有KV Cache而沒有Q Cache？ DefTruth
全面解析 LLM 推理優(yōu)化：技術(shù)、應(yīng)用與挑戰(zhàn) AI閑談
刀刀寧：筆記：Llama.cpp 代碼淺析（一）：并行機(jī)制與KVCache
刀刀寧：筆記：Llama.cpp 代碼淺析（二）：數(shù)據(jù)結(jié)構(gòu)與采樣方法
分析 GPT 模型自回歸生成過程中的 KV Cache 優(yōu)化技術(shù) sudit
圖解Mixtral 8 * 7b推理優(yōu)化原理與源碼實(shí)現(xiàn) 猛猿
圖解大模型計(jì)算加速系列：分離式推理架構(gòu)1，從DistServe談起 猛猿
圖解大模型計(jì)算加速系列：分離式推理架構(gòu)2，模糊分離與合并邊界的chunked-prefills 猛猿
圖解大模型計(jì)算加速系列：分離式推理架構(gòu)2，模糊分離與合并邊界的chunked-prefills 猛猿
大模型并行推理的太祖長拳：解讀Jeff Dean署名MLSys 23杰出論文 方佳瑞
大模型推理優(yōu)化--Prefill階段seq_q切分 kimbaol
大模型推理優(yōu)化實(shí)踐：KV cache復(fù)用與投機(jī)采樣 米基 [阿里技術(shù)](javascript:void(0)??
大模型推理優(yōu)化技術(shù)-KV Cache [吃果凍不吐果凍皮](https://www.zhihu.com
/people/liguodong-iot)
大模型推理加速：KV Cache Sparsity(稀疏化)方法 歪門正道
大模型推理加速：看圖學(xué)KV Cache 看圖學(xué)
大模型推理性能優(yōu)化之KV Cache解讀 Young
大白話解說Continous Batching Gnuey Iup
微軟 MInference：百萬 Token 序列，10x 加速
思辨｜H2O: 從緩存視角解決LLM長上下文問題 DSTS Center [數(shù)據(jù)空間技術(shù)與系統(tǒng)](javascript:void(0)??
打造高性能大模型推理平臺之Prefill、Decode分離系列（一）：微軟新作SplitWise，通過將PD分離提高GPU的利用率 哆啦不是夢
打造高性能大模型推理平臺之Prefill、Decode分離系列（一）：微軟新作SplitWise，通過將PD分離提高GPU的利用率 哆啦不是夢
月之暗明的月餅：以KVCache為中心，分離式推理架構(gòu)，LLM吞吐能力飆升75%
極限探索：LLM推理速度的深度解析 MLSys2024
符堯博士-全棧Transformer推理優(yōu)化第二季：部署長上下文模型-翻譯 強(qiáng)化學(xué)徒
阿杰：vLLM & PagedAttention 論文深度解讀（一）—— LLM 服務(wù)現(xiàn)狀與優(yōu)化思路
Efficient Streaming Language Models with Attention Sinks
PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference
MiniCache: KV Cache Compression in Depth Dimension for Large Language Models
TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding
ZigZagKV，KVCache80%的壓縮率可以達(dá)到近乎無損的效果 杜凌霄
SnapKV: LLM Knows What You are Looking for Before Generation
FastGen：Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs
[Prompt Cache]: Modular Attention Reuse for Low-Latency Inference(@yale.edu)
[Cache Similarity]: Efficient Prompt Caching via Embedding Similarity(@UC Berkeley)
[Shared Prefixes]: Hydragen: High-Throughput LLM Inference with Shared Prefixes(@Stanford University)
[ChunkAttention]: ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition(@microsoft.com)
ClusterKV，接近無損的KV Cache壓縮方法 杜凌霄
圖解大模型計(jì)算加速系列：分離式推理架構(gòu)2，模糊分離與合并邊界的chunked-prefills 猛猿
《基于KV Cache管理的LLM加速研究綜述》精煉版 常華Andy
[KV Cache優(yōu)化]??MQA/GQA/YOCO/CLA/MLKV筆記: 層內(nèi)和層間KV Cache共享 DefTruth
Venkat Raman. Essential math & concepts for llm inference, 2024. URL https://venkat.e u/essential-math-concepts-for-llm-inference#heading-insights-from-model-lat ency-amp-understanding-hardware-utilization-on-modern-gpus.
Yao Fu. Challenges in deploying long-context transformers: A theoretical peak performance analysis. arXiv preprint arXiv:2405.08944, 2024
A Survey on Large Language Model Acceleration based on KV Cache Management
浙公網(wǎng)安備 33010602011771號