探秘Transformer系列之（26）— KV Cache優化 之 PD分離or合并

0x00 概述

在大模型的推理過程中，通常可以將任務分為兩個階段：Prefill 階段處理所有輸入的 Token，生成第一個輸出 Token，并生成 KVCache。Decode 利用 KVCache 進行多輪迭代，每輪生成一個 Token。由于 Prefill 階段并行處理許多 Token，因此是計算密集型的，其延遲通過首 Token 時延（TTFT）來衡量。相比之下，Decode 階段由于頻繁加載不斷增長的 KV Cache 而成為內存密集型，其時延通過 TPOT 來衡量。

因為Pefill 階段和Decode 階段的特性迥異，很難找到完美的方案同時滿足兩個階段的需求。研究人員也八仙過海，各顯神通。本文就帶領大家來學習下學術界和工業界的各種方案。

---

注：

• 全部文章列表在這里，估計最終在35篇左右，后續每發一篇文章，會修改此文章列表。探秘Transformer系列之文章列表
• 本系列是對論文、博客和代碼的學習和解讀，借鑒了很多網上朋友的文章，在此表示感謝，并且會在參考中列出。因為本系列參考文章太多，可能有漏給出處的現象。如果原作者發現，還請指出，我在參考文獻中進行增補。

探秘Transformer系列之（1）：注意力機制

探秘Transformer系列之（2）—總體架構

探秘Transformer系列之（3）—數據處理

探秘Transformer系列之（4）— 編碼器 & 解碼器

探秘Transformer系列之（5）— 訓練&推理

探秘Transformer系列之（6）— token

探秘Transformer系列之（7）— embedding

探秘Transformer系列之（8）— 位置編碼

探秘Transformer系列之（9）— 位置編碼分類

探秘Transformer系列之（10）— 自注意力

探秘Transformer系列之（11）— 掩碼

探秘Transformer系列之（12）— 多頭自注意力

探秘Transformer系列之（13）— FFN

探秘Transformer系列之（14）— 殘差網絡和歸一化

探秘Transformer系列之（15）— 采樣和輸出

探秘Transformer系列之（16）— 資源占用)

探秘Transformer系列之（17）— RoPE（上）

探秘Transformer系列之（17）— RoPE（下）

探秘Transformer系列之（18）— FlashAttention

探秘Transformer系列之（19）----FlashAttention V2 及升級版本

探秘Transformer系列之（20）— KV Cache

探秘Transformer系列之（21）— MoE

探秘Transformer系列之（22）— LoRA

探秘Transformer系列之（23）— 長度外推

探秘Transformer系列之（24）— KV Cache優化

探秘Transformer系列之（25）— KV Cache優化之處理長文本序列

---

0x01 背景知識

1.1 自回歸&迭代

當執行文本生成任務時候，基于Transformer架構的自回歸語言模型會接受一個文本序列作為輸入，模型通過生成連續的輸出標記來完成序列，這是一個迭代的過程。我們把模型從接受請求輸入到輸出一個token的操作定義為模型的一個迭代。對于每個請求，模型會逐步生成輸出序列的各個部分，直到生成停止標記或達到最大序列長度為止。這意味著每次模型前向傳遞時，都會獲得一個額外的輸出token。例如，如果我們以句子“加利福尼亞的首府是什么:”作為提示，它將需要進行十次前向傳遞才能得到完整的響應，即[“S”, “a”, “c”, “r”, “a”, “m”, “e”, “n”, “t”, “o”]。即LLM一個請求可能會包括多個迭代。

下圖示例展示了一個三層GPT模型的架構，其中節點和邊分別表示Transformer層和層間依賴關系。模型按照節點上的編號順序執行Transformer層，同一層的節點使用相同的模型參數，以同一顏色（不同深淺顏色）填充。生成的輸出標記反饋到模型中，生成下一個輸出標記，形成順序逐個推理的過程。

在示例中，推理過程包括三個迭代（對應圖上三個綠色標號）。這三個迭代又可以分為兩個階段：

• 初始化階段。該階段通常只包括一個迭代，負責通過并行處理所有輸入token來生成第一個輸出token。對應圖上就是第一次迭代（“iter 1”）一次性接受所有輸入標記（“I think this”），生成下一個標記（“is”）。其實可以看到，此處的迭代對應的就是prefill階段。
• 增量階段。該階段通常包括多個迭代，每個迭代會接受前一次迭代的輸出標記生成下一個標記（每次迭代只能處理一個標記）。對應下圖，本階段包括后兩次迭代（“iter 2”和“iter 3”）。在這個示例中，“iter 3”是最后一個迭代，因為它生成特殊的結束標記“”，終止輸出生成。此處的每個迭代對應的就是decode階段。

總結下，我們管請求做完prefill產出第一個token的過程叫1次iteration，請求每做一次decode也被稱為1次iteration。

1.2 KV Cache

下圖給出了prefill和decoding階段KV Cache的作用。(a)的頂部是自回歸LLM推理的示例，底部是transformer層中的模塊。(b)給出了KV Cache的操作。在預填充階段，所有輸入token同時處理，并存儲生成的中間KV張量，用深色標記。s和d表示輸入序列長度和KV張量的隱藏維數。在解碼階段，存儲的深色KV張量被檢索出來。輸入的Q、K、V張量由淺色標記。輸入Q張量與輸入K和存儲的K張量的合并結果（concatenation ）進行相乘，然后是整個注意力權重的softmax。注意力權重進一步與輸入V和存儲的V張量的合并結果（concatenation ）進行相乘，以生成新的結果。然后，將輸入的K和V張量再存儲在KV Cache中。對每個token都重復此過程。

• Prefill 階段會并行處理輸入 prompt 的所有 token，輸出第一個token，并生成用于未來解碼的KV Cache。因為輸入序列長度很長，所以計算開銷大，很小的 batch size 就會讓 GPU 打滿。增大 batch size，prefill 階段單個 token 開銷幾乎是不變的。
• （開啟 KV Cache）的 decode 階段使用先前的KV Cache以自回歸方式逐步生成新的token。decode 階段在每個自回歸步僅僅會生成一個 token。雖然decode 階段輸入的序列長度一直是 1，但是需要反復讀取 KV Cache，故而 IO 開銷很大。Decode 時，單個 token 的開銷顯著大于 prefill 階段。需要把 decode 階段的 batch size 配置得非常大才有可能占滿 GPU ，但會因為 KV Cache 讀寫開銷太大而變得不現實，所以 decode 階段的 GPU 利用率很低。

0x02 靜態批處理

2.1 調度策略

前面提到，LLM 正在改變整個行業在其服務中采用 AI 技術的方式，但 LLM serving 的成本卻仍然很高。為降低 serving 成本，目前許多公司專注于最大化整體 LLM serving 系統的整體吞吐量（throughput），即每秒服務的請求數量（或 rps）。幾乎所有流行的 LLM serving 引擎都把吞吐量作為比較性能的主要指標。

為了提高吞吐率，人們會采用批處理技術。批處理是一種將多個數據樣本一起傳遞給模型進行處理的技術。啟用批處理后，執行引擎會將來自多個請求的輸入張量合并成一個大的輸入張量，然后將其送入模型進行推理。由于加速器更喜歡大型輸入張量而不是小型張量，因此相比于逐個處理單個樣本，批處理可以在一次計算中同時處理多個樣本。這樣可以更有效地利用計算資源，提高計算速度。在LLM推理中，批處理的優勢主要體現在以下幾個方面：

• 減少模型參數加載次數： 在不使用批處理的情況下，每次處理一個輸入序列都需要加載一次模型參數，而加載模型參數通常會消耗大量的內存帶寬。使用批處理可以在一次加載模型后多次使用這些參數來處理多個輸入序列，從而減少了加載的次數。

• 提高內存帶寬和計算資源的利用率： GPU的內存帶寬是有限的資源，通過批處理，可以更有效地讓這些內存帶寬在計算過程中得到更充分的利用，從而可以更有效地利用計算資源，提高計算速度。

傳統的批處理方法被稱為靜態批處理，是因為在這種方法中，批處理的大小在批次中有所請求的推理完成之前保持不變。在靜態批處理中，當新請求在當前批次執行過程中到達時，引擎會保持這些新請求在請求隊列中，直到當前批次中所有請求都完成。即服務系統和執行引擎之間的交互僅在以下兩種情況下發生：當服務系統在空閑的引擎上安排下一批請求；或者當引擎完成當前批次的處理后。其具體流程如下：

• 如果當前執行引擎空閑，Scheduler（調度器）從隊列中獲取一批請求構建成一個batch（圖中的$x_1$ ："I Think"和 $x_2$ ：“I love”），對于下圖中標號1。
• Scheduler將batch交給Execution Engine（執行引擎）去做推理，對于下圖中標號2。
• Execution Engine通過運行模型的多次迭代來處理接收到的batch，對于下圖中標號3。
• 當Execution Engine把這個batch的請求都處理完之后（Execution Engine分別生成“this is great”和“you”作為請求$x_1$和$x_2$的響應。），會統一返回給Scheduler，對于下圖中標號4。
• Scheduler會從請求隊列中再次獲取一批請求構建一個新batch。

2.2 問題

我們很容易發現靜態批處理策略的劣勢：它難以有效地處理不同生成長度的序列。靜態批處理可能導致GPU在批處理中的不同序列的生成長度不同時被低效利用。因為不同的序列可能會在批處理中的不同迭代步驟中完成生成，而靜態批處理會等待所有序列完成生成后才開始處理新的序列。這導致了在等待最后一個序列完成生成之前，GPU可能會被低效利用的情況。因為通過自回歸方式處理單個請求時可能需要運行多次模型，而不同請求所生成的文本長度不一致，需要的迭代次數很可能不同，這就使得一些請求提前完成而其他請求仍在進行中。一旦有一個生成時間較長的序列存在，其他生成時間較短的序列將被迫等待、不能立即返回給客戶端，即增加了給客戶帶來得延遲，也導致GPU的計算資源無法充分利用，浪費GPU的計算能力。如上圖所示，$x_2$提前完成，但是$x_1$ 尚未完成迭代，所以$x_2$無法立即返回。而在上次分派之后到達的新請求（$x_3,x_4$）也必須等待當前batch中的$x_1$處理完畢才能繼續分發，這顯著增加了排隊時間。

因此，我們目前需要解決的問題如下：

• 問題1：如何處理提前完成的請求（early-?nished requests）。不同請求所生成的文本長度不一致，短文本請求只能等待長文本請求處理完畢。而因為不易預測生成文本長度，故此難以把相同長度的文本組成一個batch一并處理。因此需要一個將已生成結束的請求從Batch中移除并提前返回結果的機制。
• 問題2：如何釋放資源、處理延遲加入的請求（late-joining requests）。隊列中的新請求只能等到當前batch結束之后才能被scheduler進行分發，等待時間過長。因此需要一個將新請求插入到當前處理Batch中的機制。

需要注意的是，提前完成和延遲加入請求的問題在模型訓練中并不會出現；訓練過程通過使用Teacher Forcing技術在單次迭代中完成整個批次的處理。

2.3 原因分析

有兩個原因導致出現上述問題，分別可以從宏觀和微觀角度來進行分析。

2.3.1 宏觀

從宏觀角度來說，問題本質在于“基于Transformer的生成模型的多迭代特性”和“以請求為粒度的調度策略”之間的沖突。多迭代特性意味著不同請求可能需要不同次數的迭代，某些請求的迭代次數可能會很少，可以提前處理完畢并返回給用戶。而與靜態批處理的這種特點對應，傳統推理架構的調度策略是以請求為粒度的調度（Request-Level Schedule）。以請求為粒度的調度策略缺少靈活性，無法更改當前正在處理的請求batch，比批次中其他請求更早完成的請求無法返回客戶端，而新到達的請求則必須等到當前批次完全處理完畢。這就導致當前的請求級調度機制無法有效處理具有多次迭代特性的工作負載。

對指標的影響

從指標角度來看更為清晰，即靜態批處理通過犧牲TTFT和throughput的方式來確保TPOT。如前所述，在一個batch中，prefill階段通常作為一個迭代通過并行處理所有輸入標記來實現，這說明后續迭代都是decode階段。因為必須把當前batch的請求全部處理完成才能返回，所以靜態批處理可以確保Execution Engine始終在進行decode操作，這樣TPOT時間可以得到保證。但是因為在這個batch做推理的過程中不能接受新的請求，導致吞吐量下降，進而導致新請求的TTFT指標下降（新請求只能在請求隊列中等待，無法在第一時間執行prefill操作）。

對流水線的影響

另外，以請求為粒度的調度也會增加流水線中的氣泡。當模型尺寸過大時，人們常常采用多張GPU進行并行推理，而流水線并行是常見方式之一。流水線并行的優勢是只涉及層間激活的通信，對帶寬要求相對較小。下圖中，模型被切分為三層，分別部署在Partition1~3這三張GPU上。A和B是兩個序列，下標代表本序列正在進行第幾次模型迭代。假設A和B都在decode階段，因為decode階段是自回歸逐一輸出token，因此A和B需要在第一次迭代完成之后，才能進入下一次迭代。這樣就造成了GPU上的空閑時間（氣泡）。

2.3.2 微觀

由于預填充和解碼階段共享LLM權重和工作內存，現有的LLM服務系統通常將這兩個階段放在一個GPU上，并通過批量處理預填充和解碼步驟來最大化整個系統的吞吐量（每秒跨所有用戶和請求生成的token數）。而出現問題的本質原因在于，當前的LLM系統對LLM預填充和解碼階段所展示的不同特征以及計算模式一無所知，也不清楚這些模型對SLO的不同影響。預填充階段類似于計算密集型的批處理作業，而解碼階段類似于內存密集型、延遲關鍵的任務。靜態批處理將這兩種運算放到相同的 GPUs 上進行處理并不是實現高有效吞吐量的最佳策略。

注：此處說的是靜態批處理中的對prefill和decoding進行處理，并非是我們接下來會講到的融合方案。

階段特點

前面章節從技術和流程特點分析了prefill和decode之間的差異，我們這里結合指標和資源利用來繼續分析。

Prefill 和 Decode 請求有著不同的計算特性，并且需要滿足不同的 SLO。

• Prefill。主要滿足 TTFT 指標。這是一種Compute Bound 的運算，對計算資源的需求量極大。其計算量隨輸入提示長度的平方增加。耗時超線性增加（吞吐量下降，如果耗時線性增加，那么平均吞吐量不變），哪怕是小批次的預填充任務，甚至單個較長的預填充任務，都可能使 GPU 的計算能力達到飽和，打滿計算資源。一旦加速器達到飽和狀態，預填充階段的吞吐量就會保持不變（我們將其命名為加速器飽和閾值）。

• Decode：主要滿足 TPOT 指標。這是一種 Memory Bound 的運算，更容易受到 GPU 內存帶寬限制的影響。對延遲敏感的任務，其資源使用量隨生成的token長度呈次線性增長。Decode 計算需要大 batch 請求提高計算強度，充分利用計算資源。隨著bs的增大，耗時略有增加（主要是因為數據增多導致傳輸耗時增多，計算耗時基本認為不變），吞吐量大概是線性增加。隨著批量大小增加，解碼階段的吞吐量繼續增加，但一旦內存帶寬達到飽和狀態，吞吐量就會達到峰值。

互相影響

Prefill 和 Decode 放在相同的 GPUs 上運行會導致它們之間相互干擾。我們看看當 GPU 在同時處理 Prefill 請求和 Decode 請求時，Prefill 和 Decode 如何互相影響性能。

• 對于 Prefill 請求，由于同一個 batch 有其它 Decode 請求，所以在 GPU 中計算耗時對比沒有 Decode 請求時會略微增加，從而增加 Prefill 請求的 TTFT。
• 對于 Decode 請求，由于同一個batch 有 Prefill 請求，并且 Prefill 請求計算耗時通常會比 Decode 請求耗時大很多（一般是好幾倍甚至幾十倍的關系）。批處理中的解碼任務必須等待更長的預填充作業完成，會造成 Decode 請求被 Prefill 請求拖慢，必須等待 Prefill 請求完成計算，Decode 請求才會開始下一個 token 的計算，從而大大增加了 Decode 請求的 TPOT。

這種干擾是經典的系統問題。

• 混合運行預填充請求會導致嚴重的減速，因為我們在已經飽和的硬件上繼續添加計算密集型作業。
• 混合運行預填充和解碼請求會對兩者都產生負面影響，因為我們同時運行批處理和延遲關鍵的作業。
• 混合解碼請求會導致吞吐量下降，因為我們不了解內存帶寬和容量使用情況，從而導致爭用和排隊阻塞。

優先安排一個階段可能會導致無法滿足另一個階段的延遲要求。比如，較高的到達率（req/s）會生成更多的預填充作業，如果優先考慮預填充作業，則需要更多的GPU時間來滿足TTFT要求，這反過來會對TPOT產生不利影響。

合并預填充和解碼也會耦合它們的GPU資源分配策略和并行策略，并阻止實現更適合滿足每個階段特定延遲要求的不同并行性策略。然而，每個階段都具有其獨特的計算特性和延遲要求，需要更多的異構資源分配。

• GPU 分配策略。如果希望同時滿足 TTFT 和 TPOT 的 SLO，系統必須過度配置資源以滿足時延目標，尤其是在 SLO 要求嚴格的情況下。這兩種運算對 GPU 資源的需求不一樣。相同數量的 request 輸入，Prefill 需要更多的 GPU 滿足 TTFT，而 Decode 可以需要更少的 GPU 滿足 TPOT。但原有的服務架構中，Prefill 階段和 Decode 階段會共享 GPU 分配策略，因此，為了滿足 Prefill 運算的 TTFT，會配置更多的 GPU，這對于 Decode 運算是超額的；為了支持一定總量的吞吐，需要降低每個 GPU 處理的 requests 數量，這樣也需要超額配置 GPU 資源；隨之而來的就是部署成本的上升，這種“充值”的方式并不是長久之策。
• 模型并行策略。不同的場景下，由于它們的計算模式和時延目標截然不同，那么 Prefill 和 Decode 的最優模型并行策略可能不同。假如服務對 TTFT 較嚴格，對 TPOT 需求較低，那么 Prefill 階段更傾向使用 Tensor Parallelism（簡稱 TP）來滿足嚴格的時延目標，而 Decode 階段更傾向使用數據或流水線并行提高吞吐量。但現有的服務架構中，預填充和解碼計算的并行策略（張量、流水線或數據并行）本質上是耦合的，Prefill 階段和 Decode 階段會共享模型并行策略。因此，無法針對不同階段采用同時兼顧吞吐與SLOs更優的配置。

將預填充和解碼作業解除批處理并按順序調度并不能減輕干擾。由于解碼作業需要等待GPU上正在進行的預填充作業，解碼作業可能會經歷更長的排隊延遲。此外，專門用于解碼的批次通常會導致GPU利用率降低。優先處理某個階段的任務會對另一個階段的延遲產生不利影響，從而使優先調度失效。

2.4 挑戰

2.4.1 從宏觀角度來應對

從宏觀角度看。既然問題出自“基于Transformer的生成模型的多迭代特性”和“以請求為粒度的調度策略”之間的沖突，有研究人員就提出了采用迭代級別調度來促進各個階段的并發處理。OSDI 2022 上發表的 Orca 是第一篇解決這個問題的論文。

對于靜態批處理，其批次大小是在批次啟動時候就確定的，批處理的大小在批次中有所請求的推理完成之前保持不變。而對于迭代級別調度，批次的大小是在每個迭代中動態確定的，而不是在推理過程的開始時就固定下來。這意味著一旦批次中的某個序列完成生成，就可以立即插入一個新的序列以繼續利用GPU進行計算。這會讓GPU不會因為等待批次中的所有序列完成生成而被浪費，讓計算資源得到更高效的利用，從而可以實現更高的推理吞吐量，加快整個推理過程的速度。

2.4.2 從微觀角度來應對

從微觀角度看，目前的挑戰就是如何權衡TTFT 和TPOT之間的balance。假設請求池里面是來自不同用戶的不同狀態的請求，有的請求等待做pefill，有的請求等待做decode，但是只有一個GPU。

問題1的挑戰

針對問題1：如何處理提前完成的請求。

因為當前整個batch都沒有完成，用戶只能等待。因此有研究人員想到：是否可以讓batch中已經做完推理的請求釋放資源，讓新來的請求執行prefill操作。原有的請求繼續執行目前的操作（prefill或者decode）。

問題2的挑戰

針對問題2：如何處理延遲加入的請求。

由于預填充和解碼階段具有不同的特征，MaaS提供商設置不同的指標來衡量它們對應的服務級別目標（SLOs）。具體來說，預填充階段主要關注請求到達和生成第一個標記之間的延遲，即時間到第一個標記（TTFT）。另一方面，解碼階段關注同一請求連續標記生成之間的延遲，稱為標記之間的時間（TBT）。

如何滿足以上指標？有不同的思路。

• 當使用流水線并行時，有研究人員想到是否可以在氣泡處插入新請求做prefill或者decode操作？這樣可以在不影響當前請求的情況下填滿氣泡。
• 當不使用流水線并行時，接下來是處理等待prefill的請求？還是處理等待decode的請求？其實取決于在TTFT和TPOT兩個指標中犧牲哪個，保全哪個。假設請求池有若干的prefill請求，如果優先處理等待prefill的請求，就可以快速給用戶返回第一個token。但是需要連續處理好幾個處理等待prefill的請求，無法做decode。這是在犧牲TPOT保全TTFT；如果為了讓用戶快速獲取后續的token，讓系統連續處理幾個等待decode的請求，prefill請求就要等待，這樣是在犧牲TTFT好全TPOT。

似乎無論調度策略怎么調整，TTFT和TPOT這兩個指標都存在強耦合關系。是prefill和decode這兩個階段在本質上就無法共存，必須分開獨立發展？還是可以做融合，共同進步？

2.5 派系

針對上面的疑問，處理Prefill和Decode有融合派和分離派兩大流派。此處引用知乎大神方佳瑞的論斷。

2.5.1 融合派

融合派的思路是將prefill放到Decode的間隙進行處理，或者把prefill和某個decode一起進行推理。開山之作是發布在2022年OSDI的Orca，其將Prefill和Decode融合，作為一個批次迭代（batching step）做前向傳播。在這之后，以Sarathi（https://arxiv.org/pdf/2401.08671）和FastGen（https://arxiv.org/pdf/2403.02310）為代表，將prefill序列拆成chunk，可以把每個chunk的prefill的處理插入到decode階段之中，甚至和decode階段進行融合。

融合派的優點是，prefill chunk和decode融合之后的處理時間和只做decode處理時間可能類似，這相當于prefill把decode沒有用滿的算力再利用起來。而且因為decode和prefill階段都需要讀一些固定的數據（比如模型權重），所以decode也可以搭上prefill的便車，把多次的數據讀取合并起來，一次讀取，多次使用。

融合派的缺點是，Prefill和Decode在相同設備上，二者并行度被綁定。如果prompt長度有限，則在請求的處理時間中，prefill階段占比較小，尚可利用好算力。對于長序列，則prefill的占比提高，prefill并行度不夠靈活的缺陷就暴露出來。

2.5.2 分離派

分離派：考慮Prefill/Decode性質差異，人們開始嘗試把Prefill和Decode放到不同的設備來分別處理。比如，Splitwise（https://arxiv.org/abs/2311.18677）和DistServe（https://arxiv.org/pdf/2401.09670）就是其中典型代表。

分離預填充和解碼自然地解決了兩個階段之間的干擾，并使每個階段都能專注于其優化目標 - TTFT或TPOT。每種類型的實例可以使用不同的資源和并行性策略來滿足各種延遲要求。通過調整為兩種類型的實例提供的GPU數量和并行性，我們可以最大化整個系統的每個設備的吞吐量，避免過度配置資源，最終實現符合服務質量的每個查詢成本的降低。

分離派可以給Prefill和Decode設置不同的并行度，二者有更大的靈活性，使得整個系統朝著更好利用算力和更好利用帶寬這兩個不同方向對立發展，可以更好的對硬件進行優化。比如對于prefill和decode采用不同類型的GPU卡，或者可以針對長序列來分配多一些GPU進行處理，進而降低長序列的TTFT。

分離派遇到的最大挑戰是如何在不同設備之間傳輸KVCache，這導致網絡成本很大，從而需要比較高規格的網絡硬件來在集群中進行互聯。而如何處理長序列的并行，也是一個棘手問題。

接下來會看看兩種派別如何處理這些問題。

0x03 融合派

融合派的思路是搭便車。盡管很多推理引擎采用了最先進的批處理和調度技術，但是decode階段沒有充分利用到GPU的計算資源，因此可以將prefill放到Decode的間隙進行處理，或者把prefill和某個decode一起進行推理。典型代表是：

• ORCA：把調度粒度從請求級別調整為迭代級別，并結合選擇性批處理（selective batching）來進行優化。
• Sarathi-Serve：利用Chunked Prefill策略通過將不同長度的prompts拆分成長度一致的chunks來進行prefill，同時利用這些chunks間隙進行decode 操作。

注意：目前業界把依據ORCA思想實現的方案叫做Continuous Batching（連續批處理）。連續批處理是一種優化技術，它允許在生成過程中動態地調整批處理的大小。具體來說，一旦一個序列在批處理中完成生成，就可以立即用新的序列替代它，從而提高了GPU的利用率。這種方法的關鍵在于實時地適應當前的生成狀態，而不是等待整個批次的序列都完成。

與靜態批處理不同，連續批處理采用了迭代級別的調度。它并不等待每個序列在批次中完成生成后再進行下一個序列的處理。相反，調度程序在每個迭代中根據需要確定批次的大小。這意味著在每次迭代之前，調度程序檢查所有請求的狀態。一旦某個序列在批次中完成生成，就可以立即將一個新的序列插入到相同位置，同時刪除已完成的請求。

3.1 ORCA

針對如何處理“提前完成和延遲加入的請求”這個挑戰，ORCA給出的解決方案是用迭代級調度減少空閑時間，即以迭代為粒度（iteration-level）控制執行，而不是請求級粒度（request-level），并結合選擇性批處理（selective batching）來進行優化。

3.1.1 迭代級調度

目標

迭代級調度的目標是：及時檢測出推理完畢的請求，將其從batch中移出，以便新請求可以填補到舊請求的位置上，這樣新請求和舊請求能接連不斷組成新的batch。

調度

前文提到，Orca是第一篇提到迭代級別調度（Iteration-Level Schedule）的論文。具體來說就是：一個batch中的所有請求每做完1次iteration（prefill或者decode），scheduler就和engine交互一次，去檢查batch中是否有做完推理的請求，以此決定是否要更新batch。這樣就可以在每次GPU推理的空隙，可以插入調度操作，實現Batch樣本的增刪和顯存的動態分配釋放。

下圖給出了請求粒度的調度和迭代粒度調度的區別。前者需在整批請求全部完成前對調度批次進行多次迭代，而對于ORCA，服務系統在調度任務時，每次只向 Execution Engine 提交一次迭代的計算，而非等到完成整個 Request才能處理。這樣 ORCA 就可以在每個迭代都動態更改要處理的請求，新請求只需等待單次迭代即可被處理，從而避免early-finish的請求等待其他請求的結束。通過迭代級調度，調度器能夠完全控制每個迭代中處理哪些請求以及處理數量。

方案

下圖展示了采用迭代級調度的ORCA系統架構和整體工作流程。ORCA系統包括如下模塊：

• Endpoint（端點）。用于接收推理請求并發送響應。
• Request Pool（請求池）。新到達的請求被放入請求池中，該組件負責管理系統中所有請求的生命周期。
• Scheduler（調度器）。調度器監控請求池，負責以下任務：從池中選擇一組請求，調度執行引擎對這些請求執行模型迭代；接收執行引擎返回的執行結果（即輸出token），并將每個輸出token追加到對應請求中來更新請求池。
• Execution Engine（執行引擎）。執行引擎是執行實際張量操作的抽象層，可在跨多個機器分布的多個GPU上并行化。

我們接下來看看下圖中的工作流程，其中，虛線表示組件之間的交互，交互發生在執行引擎的每次迭代中。$x_{ij}$是第i個請求的第j個token。陰影token表示從客戶端接收到的輸入token，而非陰影token由ORCA生成。例如，請求$x_1$最初帶有兩個輸入標記（$x_{11}$，$x_{12}$），到目前為止已經運行了兩次迭代，其中第一次和第二次迭代分別生成了$x_{13}$和$x_{14}$。另一方面，請求$x_3$只包含輸入標記$x_{31}$,$x_{32}$，請求$x_4$包括$x_{41}$，$x_{42}$，$x_{43}$，因為它們還沒有運行任何迭代。

工作流程分為如下幾步：

• 調度器與請求池交互，以決定下一步運行哪些請求。對應下圖標號?。
• 調度器調用引擎為所選定的四個請求（$x_1,x_2,x_3,x_4$）執行一次迭代。此時，因為$x_3$和$x_4$還沒有運行任何迭代，因此調度器為$x_3$移交$x_{31}$，$x_{32}$給執行引擎，為$x_4$移交$x_{41}$，$x_{42}$，$x_{43}$給執行引擎。對應下圖標號?。
• 引擎對四個請求運行模型迭代，對應下圖標號?。
• 引擎把生成的輸出token（$x_{15}$, $x_{23}$, $x_{33}$, $x_{44}$）返回給調度器，對應下圖標號?。調度器在每次引擎返回，接收該迭代的執行結果之后會檢查請求是否完成。如果請求完成，請求池就會刪除已完成的請求，并通知端點發送響應，返回給客戶端。
• 對于新到達的請求，在當前迭代執行完畢后，它有機會開始處理（即調度器可能選擇新請求作為下一個執行對象）。因為新到達的請求只需等待一次迭代，從而顯著減少了排隊延遲。

ORCA對于中止請求（Canceled Requests）并沒有進行處理，實際上應該把這些請求會被及時從Batch中剔除并釋放相應顯存。

算法

下圖詳細描述如何在每次迭代中選擇請求的算法。此算法中對KV Cache釋放時機控制得不是很理想。在請求生成結束時就立即釋放其K/V Cache。在多輪對話場景中，這個機制會導致冗余計算，即“上一輪對話生成K/V Cache → 釋放K/V Cache顯存 → 通過本輪對話的Prompt生成 之前的K/V Cache”。這樣會惡化后續幾輪對話的First Token Time（產生第一個Token的時延）指標。

3.1.2 選擇性批處理

ORCA時代還沒有Paged Attention，因此需要Selective Batching來將注意力計算從 Batching 中解耦。即，為了提高計算效率，需要想辦法讓引擎能夠以批處理方式處理任何選定的請求集。

問題

在前面分析中，我們其實做了一個簡化的假設，即所有請求序列具有相同的長度。這是因為GPU的特殊性，如果想批量執行多個請求，每個請求的執行應該包含相同的操作，且消耗形狀相同的輸入張量。然而，在現實中，請求序列的長度是不同的。誠然Padding+Masking的方法可以解決，但嚴重浪費算力和顯存，對于算力和顯存均有限的推理GPU是不利的。

當使用迭代級別調度時，上述挑戰會愈發加劇。因為：

• 請求池中的請求可能具有不同的特征。
• prefill和decode的計算方式不同。
  • prefill過程是長序列并行計算的，decode過程是token by token的。
  • prefill過程不需要讀取KV cache，decode過程需要讀取KV cache。
  • 對于prefill，各個請求的prompt長度是不一致的。
  • 對于decode，不同請求的decode token的index不一樣，意味著它們計算attention的mask矩陣也不一樣。
• 迭代級調度方法可能導致同一個批處理中的不同請求的處理進度不一樣，即輸入張量的形狀會因為已處理的token數量不同而不一致。

我們用上面架構圖作為例子來進行分析，來看看即使對于一對請求$(x_i,x_j)$，也不能保證它們的下一次迭代可以合并、替換為批處理版本。有三種情況導致請求對不能合并批處理：

• 兩個請求都處于初始化階段，但輸入token數量不同（如下圖中的$x_3$和$x_4$）或者說輸入張量的“長度”維度不相等，因此無法將兩個請求進行批處理。

• 兩個請求都處于增量階段，但各自處理的token索引不同（如$x_1$和$x_2$）。由于每個請求處理的token索引不同，導致注意力鍵和值的張量形狀不同，因此也不能合并批處理。

• 請求處于不同階段：有的處于初始化階段，有的處于增量階段（如$x_1$和$x_3$）。由于不同階段的迭代輸入token數量不同（初始化階段迭代并行處理所有輸入token以提高效率，而增量階段每次迭代僅處理一個token），因此無法合并批處理。

對于第二種情況，在典型的批處理機制中，每次迭代時，Transformer層都會接收一個由批中多個[L,H]請求輸入張量拼接而成的三維輸入張量[B,L,H]，其中B是批大小，L是一起處理的token數，H是模型的隱藏大小。

但是在下圖中，“iter 1”（初始化階段）接收形狀為[2,2,H]的輸入張量，而“iter 2”（增量階段）接收形狀為[2,1,H]的張量。然而，在上圖中，當調度器決定對批(x1, x2, x3, x4)執行迭代時，初始化階段請求(x3: [2,H]和x4: [3,H])的輸入無法合并成一個形狀為[B,L,H]的單一張量，因為x3和x4的輸入token數不同，分別為2和3。

上述關于批處理的主要問題在于，前述三種情況對應于形狀不規則的輸入（或狀態）張量，這些張量無法合并成一個大的張量并輸入到批處理操作中。因此，并非所有的請求都能在任意Iteration被Batching到一起。僅當兩個選定請求處于同一階段，且（在初始化階段）具有相同數量的輸入token或（在增量階段）具有相同的token索引時，批處理才適用。這一限制大大降低了在實際工作負載中執行批處理的可能性，因為調度器需要同時找到兩個符合批處理條件的請求。

思路

解決這些問題的一個好思路是：盡量找到這些請求計算時的共同之處，使得計算能最大化合并。對于有差異的部分再單獨處理。我們先以一個transformer decode block為例，回顧一下序列要經過哪些計算。下圖是decoder block的各種計算類型。可以看到，Transformer decoder block 在計算上可以看做六個操作的總和：pre-proj，attn，post-proj，ffn_ln1，ffn_ln2，others（比如 layer normalization，activation functions，residual connection）。Transformer 輸出一個形狀為 [B, L, H] 的張量。其中 B 是 batch size，L 是 input tokens length，H 是模型的 embedding size。每個 token 的 KV Cache 大小均為 [1, H]。

我們把上面的介紹稍作提煉，得到如下重要信息：Transformer 層中的操作可以分為兩種類型: Attention 和 non-Attention，這兩種模塊的算子特點不同。

• preproj/postproj/FFN1/FFN2：這幾個模塊中主要是Add、Linear、GeLU等算子，這些算子的特點是：
  • 不需要區分 token 來自于哪個請求。因此，雖然它們是 token-wise 的，但可以使用批處理實現。
  • 和輸入序列長度無關。這意味著我們可以把一個batch中所有的tokens都展平成一行進行計算（維護好各自的位置向量就好），這樣不同長度的輸入也可以組成batch，從而進行計算。例如，上述x3和x4的輸入張量可以組合成一個二維張量[ΣL,H] = [5,H]，而不需要明確的批處理維度。
  • 需要從顯存讀取模型權重。讀取模型權重意味著我們應該盡量增大batch size，使得一次讀取能就可以造福更多請求，以此減少IO次數。
• attn： 該模塊的特點是：
  • 由于計算受各個序列的差異性影響（例如不同序列的mask矩陣不同、是否需要讀取KV cache），因此需要將序列拆分開獨立處理，即batch維度是重要的。而由于attn部分本身不涉及到權重讀取，因此你把序列拆分開處理，也不會在這一方面上帶來額外的IO開銷。
  • 對于注意力操作, 無論是 token-wise 還是 request-wise 的 batching 都無法執行，這是因為 Attention 操作需要一個請求的概念來計算同一請求的 token 之間的注意力。
  • 不對Attention層進行批處理對效率的影響很小，因為Attention層的操作不涉及到模型參數的重復使用，無法通過批處理來減少GPU內存讀取。

方案

總結上述思路：Transformer Layer里，并非所有的算子都要求組成批次的輸入具有相同的形狀。基于上述思路，Orca 提出了第二點核心技術： Selective batching（選擇性批處理），它不是對構成模型的所有張量操作（注意力和非注意力）都進行批處理, 而是有選擇地將批處理僅應用于少數非注意力操作，即對于不同類型的請求應用于不同類型的操作來解決問題，具體如下：

• 單獨處理每個注意力操作。即對于必須有相同Shape才能Batching的算子（例如Attention）會對不同的輸入進行單獨的計算。
• 對其他層（例如MLP層）則進行批處理。

上圖展示了處理圖4中描述的請求批(x1, x2, x3, x4)的選擇性批處理機制。示意圖描述的是，系統正在生成請求x3、x4的第1個Token（論文稱為Initiation Phase），同時生成請求x1、x2的后續Token（論文稱為Increment Phase）。因為產生第1個Token后才會有K/V Cache，所以從Attention K/V Manager那里看有沒有request對應的K/V cache，就能分辨出每個request處于哪個phase。具體機制如下：

• 應用批處理。
  • 該批有7個輸入token要處理，因此我們將輸入張量的形狀設置為[7,H]，[7,H]的意義是：給定形狀為 [(s1, h), (s2, h), ...] 的輸入，我們將它們堆疊成一個形狀為 (sum(si), h) 的大矩陣。
  • 對這個堆疊矩陣應用Linear這個非注意力操作。Linear的計算不涉及Token之間的交互，因此將輸入的Batch維度和Seq維度Reshape成1個維度，即可完成Linear的Batch計算。
• 將密集層的結果分割回 [(s1, h), (s2, h), ...]。在注意力操作之前，我們插入一個拆分（split）操作，目的是把batch中的每個請求的張量區分出來。（注：有的讀者可能會疑惑為什么Linear的輸入和輸出維度分別是[7,H]和[7,3H]，為什么升維了？這是因為通過矩陣拼接可以實現用一個Linear同時計算出QKV，即輸出包含了有H維的Q、H維的K和H維的V，合計3H維）。
• 對拆分后的張量中的每個請求分別運行注意力操作。Attention的計算涉及Token之間的交互，OCRA的操作是：將輸入按Batch維Split，每個樣本分別計算Attention。
• 然后，通過合并操作將注意力操作的輸出合并回形狀為[7,H]的張量，以恢復對其他操作的批處理功能。

雖然ORCA是把Prefill和decode階段合并在一起處理，但其實它們的計算模式差很多，分開做能夠更好地做優化（例如計算Linear時，Initiation Phase的計算Kernel更接近GEMM，Increment Phase的計算Kernel更接近GEMV）

后續有的工作就是將Initiation Phase的請求批量做完計算后，再這些請求與已在decode階段的請求合并做批量計算，有的框架通過超參數來管理這個問題：等待已服務比和等待結束序列標記的請求比（waiting_served_ratio）。

3.1.3 流水線并行

另外，ORCA的調度器實現了引擎中工作線程對多個批次的流水線執行。

ORCA的調度器使引擎中worker的執行能夠跨多個批次進行流水線處理。若當前已調度批次數n_scheduled達到工作線程數n_workers（算法1的第9-10行）時，調度器就不再等待。通過這種方式，調度器保持引擎中并發運行的批次數為n_workers，這意味著引擎中的每個工作線程都在處理一個批次，而不會處于空閑狀態。

下圖展示了3個ORCA工作線程的執行流水線，最大批次大小為2。假設請求A先于B到達，B先于C到達，以此類推。首先，調度器根據到達時間選擇請求A和B，并調度引擎來處理包含A和B的批次（我們稱之為AB批次），其中Worker1、Worker2和Worker3依次處理該批次。只有在調度器注入另外兩個批次CD和EF之后，調度器才等待AB批次返回。一旦AB批次返回，請求A和B會再次被選中并調度，因為它們是請求池中最早到達的請求。

3.2 Sarathi-Serve

3.2.1 挑戰

ORCA雖然很優秀，但是依然存在兩個問題：GPU利用率不高，流水線依然可能導致氣泡問題。

GPU利用率

我們來看sarathi-serve做的一個實驗。左右兩圖分別刻畫了在不同的batch size下，prefill和decode階段的處理時間和計算強度。可以觀察到如下：

• prefill：提升batch size時，速度的變化不明顯。即使在批處理大小為1的情況下也會使GPU計算飽和，并導致跨批處理大小的每個token時間幾乎恒定。
• decode：提升batch size時，處理速度降低的線性趨勢非常明顯。這是因為decode是memory-bound的，decode階段的算力嚴重打不滿，所以當增大batch size時，不僅能多利用算力，也能把多次讀取合并成一次讀取，降低處理速度。當batch size達到一定閾值時，速度的降低幅度也達到瓶頸。著批大小的增加，用于解碼的線性算子的增量成本幾乎為零。注意力成本不會從批處理大小中受益，因為它是內存受限的。

流水線

雖然Orca已經能在一定程度上改善pp氣泡問題了，但是它仍然可能導致氣泡問題，我們以下圖為例：

觀察到圖中一共有3種類型的bubble：

• PB1: 因為兩個連續的微批中，prefill序列長度不一致而產生的bubble。
• PB2: 因為prefill和decode階段計算時間的差異而產生的bubble。
• PB3: 不同微批的decode計算時間差異而產生的bubble，這是因為不同micro-batch在做decode時，要讀取的KV cache的長度不一致，這也導致了在讀取數據上所花費的時間不一致。

3.2.2 原因分析

產生以上問題的原因在于：ORCA組裝batching的過程比較隨機，一個batch中做prefill和做decode的請求有多少條是不確定的，只是大體按照先來后到的原則做動態組裝。這就造成了一些問題。

• 如果一個batch中prefill的請求非常多，或者遇到非常長的 system prompts，那么prefill tokens會占據大量計算資源，使得整個batch變成compute-bound。
• 如果一個batch中做decode的請求非常多（比如當所有的請求都沒做完推理時，或者請求隊列中沒有新序列可以調度時），這個batch就可能變成memory-bound的。

因此，我們需要一個機制來保證每個batch中prefill和decode請求的均衡程度。或者說，假設一個batch中能夠容納的token（即達到GPU計算能力的上限）是確定的，我們要找到一個方案來按照一定比例去分配做prefill的tokens和做decode的tokens，通過重疊這兩種請求來讓這個batch做到性能最大化，解決吞吐量和延遲之間的權衡。在這種思路下，prefill的序列是必定要拆開的，這就是Sarathi提出來的Chunked-prefills（分塊預填充）。

3.2.3 Chunked-prefills

Chunked-prefills方案的核心思想是：將長序列的預填充請求拆分為幾乎相等大小的小塊，然后構建了由prefill小塊和decode組成的混合batch。或者說，Chunked Prefill策略通過將不同長度的prompts拆分成長度一致的chunks來進行prefill，以避免長prompt阻塞其他請求，同時利用這些chunks的間隙進行decode的插入/捎帶（piggyback）操作，從而減少延遲并提高整體的吞吐。

decode 階段的開銷不僅來自從 GPU 內存中獲取 KV Cache，還包括提取模型參數。而通過這種 piggyback 方法，decode 階段能夠重用 prefill 時已提取的模型參數，幾乎將 decode 階段從一個以內存為主的操作轉變為一個計算為主的操作。因此，這樣構建的混合批次具有近乎均勻的計算需求（而且增加了計算密集性），使我們能夠創建平衡的微批處理調度，緩解了迭代之間的不平衡，導致GPU的管道氣泡最小化，提高了GPU的利用率。也最小化了計算新預填充對正在進行的解碼的 TBT（Time-Between-Tokens）的影響，從而實現了高吞吐量和低 TBT 延遲。

對比

下圖給出了傳統方案，ORCA方案和Sarathi方案的時間線。

• 默認機制（下圖(a)）僅在請求級別進行批處理。在這種情況下，就緒請求被分批處理，但是只有當正在處理的批次的所有請求都完成之前，調度器才會接收新的批次。由于請求可能具有較長的token生成階段，這可能會導致請求在其間到達的等待時間較長，從而導致高TTFT和高E2E延遲。
• 連續批處理（下圖(b)）是在默認機制上的優化。在這種情況下，調度決策是在模型的每次前向傳遞之前做出的。然而，任何給定的批處理要么只包含處于prompt（提示）階段的請求，要么只包含decode階段的請求。prompt階段被認為更重要，因為它會影響TTFT。因此，等待的prompt請求可以搶占decode階段。雖然這會導致TTFT縮短，但它會大大增加TBT，從而增加E2E延遲。
• 混合批處理（下圖?）就是Sarathi方案。通過這種批處理，在每次前向傳播時都會做出調度決策，prefill和decode階段可以一起運行。這減少了對TBT的影響，但并沒有完全消除這種影響，因為與prompt階段一起調度的decode階段將經歷更長的運行時間。

注意：這里的ORCA指的是后續的類ORCA優化方案。

實現

要使用預填充來附帶解碼，我們需要注意兩件事。

• 首先，我們需要確定可以攜帶的解碼的最大可能批量大小，并確定構成預填充塊的預填充token的數量。
• 其次，為了真正利用混合批的GPU飽和預填充計算來提高解碼效率，我們需要將預填充塊和批解碼的線性運算計算融合到一個操作中。

另外，動態分割的關鍵是將較長的預填充分成更小的塊（chunk），從而通過將預填充塊與多個解碼任務組合形成批處理，并充分調動 GPU，這個過程稱為捎帶確認（piggybacking）。

要使用預填充來附帶解碼，我們需要注意兩件事。首先，我們需要確定可以攜帶的解碼的最大可能批量大小，并確定構成預填充塊的預填充token的數量。其次，為了真正利用混合批的GPU飽和預填充計算來提高解碼效率，我們需要將預填充塊和批解碼的線性運算計算融合到一個操作中。

該實現中很重要的一點就是如何確定chunk的大小，Sarathi提供了“固定”和“動態”兩種chunk size策略。

• 固定策略：該策略會依據硬件和profilling實驗計算出來一個可以最大限度把GPU利用起來的單batch中的tokens數量。這個是batch的token總配額，其在運行過程中會盡量保持不變，而prefill tokens數量會隨著decode tokens的增減而變化，但是因為decode tokens數量一般也不多，所以prefill tokens數量和整體batch tokens配額也不會相差很多。
• 動態策略：該策略希望對于一個請求，其prefill tokens的數量能隨著迭代次數的增加而減少。這是因為如果一個prompt特別長，它在每次迭代中都會占據很多計算資源，從而歷史累積的decode序列和新來的請求受到影響。因此對于這種新進入batch的長序列請求，Sarathi會在開始多配置一些prefill tokens額度，后續隨著迭代次數的增加，遞減這個配額，降低它對其它迭代的影響。

兩階段流水線的chunked-prefills運作流程如下：

具體步驟是：

• 依據GPU的性能來確定每個batch中最多能處理的tokens數量。對應上圖標號1。圖上有4個請求（A，B，C，D），被分別拆分成小塊（chunk）。

• 當整個系統剛啟動時，batch中只有做prefill的序列。系統會依據整個prefill的長度預先分配好KV cache空間，這確保在這條prefill的后續迭代中，一定有充足的KV Cache空間。對應上圖標號2。

• 往batch中添加需要做decode的序列，直到KV Cache空間不足（因為decode操作需要有對應的KV Cache）。對應上圖標號3.1。同時也要根據這個batch中剩余的tokens預算，對需要做prefill的序列做chunk切割，把對應的prefill tokens添加進batch中。對應上圖標號3.2。

  比如上圖中產生了$A_{d1}$這個需要做decode的迭代。為了進一步處理$A_{d1}$，此時需要在A所在batch中分配1 token的配額給$A_{d1}$。同時也要去等待隊列中按FCFS（先到先服務）的原則找出請求C，依據配額比例把C切分成$C_{p1}$和$C_{p2}$，然后把$C_{p1}$放到batch中。

• 推理的每一步后，scheduler都會重新組建batch。因為Sarathi-Serve依然采用的是iteration-level schedules。

我們針對上圖，再次看看當A做完prefill之后，Orca和Sarathi的區別。

• Orca在硬件資源允許的情況下，會讓CD做prefill，AB繼續做decode。但由于decode和prefill的完整序列綁定，也使得整個decode的計算時間變長了。所以這其實也算是一種decode暫停。
• sarathi-serve也允許decode和prefill一起做，但是它通過合理控制每個batch中prefill tokens的數量，使得decode階段幾乎沒有延遲。這樣即保了延遲，又保了吞吐。

分析

在 Linear 層，Decode 和 Prefil可以一起執行，在 Attention 階段分開執行，不過在很多場景下，Attention 時間占比比較小。這樣就可以將 Decode 和 Prefill 一起變成計算密集型任務。此時 Decode 應該是湊批越大越好。但是同時帶來的問題就是，Decode 的時延會比較高。所以需要在吞吐和時延之間有個平衡。在 Prefill 上，達到 GPU 計算瓶頸之后，更長的序列會使得 Prefill 時延陡增，所以在遇到這個拐點之后，Perfill 吞吐不再增加。這個值就是我們需要找到的 Chunk Size 值，Decode + Prefill 的總 Token Size 應該小于等于這個 Chunk Size，但是這個 Chunk Size 不一定很好尋找。而且 Decode 和 Prefill 的請求密度，不一定能達到最完美的比例。最完美的比例就是：在 Prefill 的每一輪執行中，Decode 都能湊批執行。如果 Decode 的密度高，那么可能 Decode 會單獨運行。如果 Prefill 密度高，那么 Prefill 會單獨運行。

另外，做了 chunked prefill 后，prefill 的開銷會略微增大。因為在計算后續 chunk 時，需要把這些 chunk’ 對應的 KV 不斷從GPU memory 中讀到 kernel 中。不做 chunked prefill 時，這些 KV 始終在 kernel中。做chunked prefill的好處是，采用 piggyback 的方式捎帶 decode 到 chunk 的 bubble 后，可以直接復用 prefill 階段加載的 模型參數。這幾乎可以讓 decode 從一個 memory bound 操作轉換為一個 compute bound 操作。

3.2.4 chunked-prefills VS 分離式推理架構

我們可以看到，在使用chunked-prefills的策略下，通過合理劃分prefill tokens和decode tokens比例，似乎也能同時保全TTFT和TPOT/TBT，利用好GPU。那么在這樣的前提下，分離式推理架構還有什么優勢呢？

其實，由于 Chunked Prefill 的提出，Prefill 和 Decode 節點的邊界已經模糊了起來。Mooncake 論文就指出，設計一個單獨的彈性預填充池的必要性和最佳實踐仍然存在爭議。隨著分塊預填充的引入，這種分離是否仍然必要？因為 Chunked Prefill 有兩個明顯的好處：1）沒有分離，所有節點都被平等對待，使調度更容易；2） 在解碼批中嵌入分塊預填充可以提高解碼批的計算強度，從而獲得更好的MFU。

Mooncake繼續保留了分離架構（使用獨立的 Prefill 節點）。只有當請求可以在不分塊（比如特別短的 prompt，可以直接一次性加入到 decode 的 continues batch 里提升 MFU）、且不影響TBT SLO的情況下轉發時，請求的預填充才會內聯到解碼批中。這一決定有兩個主要原因：1）預填充節點需要不同的跨節點并行設置來處理長上下文。2）它提供了一個保存VRAM的獨特機會。

另外也許還有如下幾點原因：

• 分塊預填充會導致預填充的計算開銷，因此選擇明顯低于 GPU 飽和點的塊大小會延長預填充任務的執行時間。
• 仍然存在prefill階段無法最大化MFU的可能。因為在chunk-prefill中，我們只是用profiling估算出在特定設備上一個batch的最大tokens配額，這些tokens包括prefill和decode。這個size是對整體的，而不是單獨對prefill或decode的。而且如果序列長度除不盡tiles尺寸，則又會產生額外的計算開銷。
• 即使塊大小被優化到幾乎最大化 GPU 使用率，分塊預填充也會顯著增加預填充任務的內存訪問量，因為需要將 KV Cache 從 GPU 的 HBM 加載到 SRAM 以用于每個后續塊。而且長序列可能會持久地占據著KV cache的存儲空間以及gpu的計算資源。

至于 TPOT，將預填充和解碼在批次中合并實際上會降低所有這些解碼任務的速度。

總之，分塊預填充可能有助于最大化整體吞吐量，但由于動態分割無法完全解耦預填充和解碼操作，會導致資源爭用以及 TTFT 與 TPOT 之間的妥協。當應用程序無法在 TTFT 和 TPOT 之間進行權衡，而是要同時遵守這兩者時，解耦就成為更好的選擇。

所以，基于以上這些對chunked-prefills策略缺陷的猜想，或許使用分離式架構，對prefill階段獨立開發一套策略，可能可以更加針對性地解決以上問題。當然，這也取決于各策略的具體實現、業務場景和真實的實驗效果。

我們接下來就看看分離式推理架構。

注：筆者并非說明分離式推理架構就一定比融合式更好，而是各有優劣。比如，論文“Injecting Adrenaline into LLM Serving: Boosting Resource Utilization and Throughput via Attention Disaggregation”就指出了LLM服務系統中的PD分離導致了嚴重的GPU資源浪費的問題。具體來說，運行計算密集型Prefill階段的GPU的HBM容量和帶寬的利用率較低。內存密集型的Decode階段則面臨算力資源利用率低的問題。該論文也給出了改進方案Adrenaline。

0x04 分離方式

4.1 總體思路

分離方案背后的邏輯很簡單：將預填充（compute-bound）和解碼（memory-bound）解耦到不同的 GPU 中，使得不管是在硬件分配還是在并行策略上，這兩者都能朝著獨立的方向優化，分別處理TTFT和TPOT/TBT，同時提升吞吐和降低延遲。而無需再像合并式推理架構那樣，總是在這兩者之間做trade off，這自然解決了上述兩個問題：

• 預填充和解碼之間沒有干擾，使得兩個階段都能更快地達到各自的 SLO。Prefill 和 Decode 不會互相影響 TTFT 和 TPOT。
• 資源分配和并行策略解耦，從而通過為預填充和解碼量身定制資源分配（GPUs 數量）、并行性策略和優化策略來獨立地擴展每個階段，借此確保最大化每個GPU的吞吐量。比如，對于長序列可以多配置些GPU，短序列少配置一些GPU。

我們以分離式架構為引子，討論了解耦prefill和decode過程帶來的好處：

4.1.1 框架

那分離式的框架長什么樣呢？我們直接來看DistServe提供的一張架構簡圖，該圖展示了請求在這種解耦系統中被處理的方式。在這種架構中，prefill和decode不再共享一塊/一群gpu，而是被分布在不同的GPU上，這些GPU分別屬于不同的實例（Prefill instance和Decoding instance）。一個prefill isntance中包含一個完整的模型，它可能占據1或多塊gpu。decode instance也是類似配置，只是它與prefill實例不再共享gpu。這種服務架構的處理流程主要分為以下三步：

• Prefill：當一個請求進入系統時，它首先被發送到一個prefill isntance。Prefill isntance對新到達的請求執行 Prefill 運算，得到第一個 tokens 以及 prompt 對應的 KVCache；
• Migrate：將 Prefill 后的請求、對應的 KVCache 遷移至 Decode isntance；
• Decode： Decode isntance對遷移的請求循環執行 Decode 運算。一旦完成生成，請求就離開了系統。

三個階段綜合起來是一個流水線并行。把不同時期的 GPU 放在不同的分區去用，然后數據從prefill到KV Cache，再到decode。KV Cache用來傳遞中間結果。如果從大一統的角度看（把 KVCache 當做內存管理），這就是一個分布式系統。

從直覺上看，分離式架構相比于合并式架構，多加載了模型副本（耗顯存），同時還涉及到gpu間的KV Cache傳輸（耗時間），似乎比合并式架構更差。那么為何人們還要采用這種架構呢？接下來，我們就來回答這個問題。

4.1.2 論證

我們接下來從不同角度來看看研究人員是如何論證的。我們首先看看DistServe和TetriInfer的觀察和洞見，然后再從KV Cache角度進行分析。

DistServe

下圖說明了在使用現有系統提供服務時，隨著請求率的增加，P90（滿足90%的SLO達成率） TTFT和TPOT如何變化。具體實驗用1張A100 80G的gpu運行13B LLM，input_length = 512, output_length = 64。圖中：

• 橫軸表示每秒到達這塊gpu上的請求數量，記為rps(requests per second)
• 藍線采用PD合并的架構做實驗。如果我們同時達到設定好的TTFT SLO和TPOT SLO，我們的最大rps = 1.6，我們也記這個指標為goodput。
• 黃線：只讓GPU處理prefill的請求時，它的goodput=5.6。
• 綠線：只讓GPU處理decode的請求時，它的goodput=10。

可以發現，一張卡只做prefill或者只做decode的goodput，都要比它既做prefill又做decode的goodput要高。所以我們可以估算：如果為預填充分配2個GPU和解碼分配1個GPU，我們可以有效地提供模型的總體吞吐量（goodput ）為10 rps，或者每個GPU的平均goodput 為3.3 rps，比合并式goodput（值為1.6）高2.1倍。

再換一個角度粗糙地看：分離式架構中，我們三張卡承受起了10 reqs/s的流量；而合并式架構中由于單卡流量承受為1.6 reqs/s，所以我們需要6張卡才能承受起10 reqs/s的流量。這說明為了滿足延遲要求，合并式架構必須過度配置計算資源。

以上實驗說明了把prefill和decode放在一起會導致強烈的互相干擾，而且會耦合它們的資源分配，并阻止實現更適合滿足每個階段特定延遲要求的不同并行性策略。

TetriInfer

與LLM交互的方式有很多種，從簡單的聊天到更復雜的下游任務，如文檔摘要、內容創建等。這些交互方式的性質往往各不相同。比如摘要任務具有較長的輸入提示和較短的生成token，而上下文創建任務則相反。不同下游任務的token長度可以相差兩個數量級以上。為了研究這些推理請求在同時運行時的性能如何，TetriInfer作者將推理請求分為兩個維度（預填充和解碼長度）和一個屬性（輕量級或重量級），從而得到四種不同的請求類型：重量級預填充、輕量級預填充、重量級解碼和輕量級解碼。這里，重量級預填充指的是較長的token長度，而輕量級預填充指的是較短的token長度。輕量級解碼是指生成少量token的解碼請求，例如少于100個。重量級解碼是指生成大量token的解碼請求，例如大于512個token。

該論文混合了不同長度的 prefill 和 decode 請求，進行了廣泛的測試，觀察到所有組合都存在嚴重的互相干擾。下圖展示了預填充和預填充，預填充和解碼，解碼和解碼這幾種組合的效果。

• 混合預填充請求。當批次中的token總數大于加速器飽和閾值時，再運行預填充請求會導致嚴重的減速，因為我們在已經飽和的硬件上繼續添加計算密集型作業。

• 預填充和解碼。混合運行預填充和解碼請求會對兩者都產生負面影響，因為我們同時運行批處理和延遲關鍵的作業。即使只有一個其他重量級預填充請求在同一連續批次中，其解碼延遲也會增加5倍！一旦共運行的輕量級解碼請求超過7個，prefill的延遲就會增加。兩者的延遲大約都增加了2.5倍。

• 解碼和解碼（混合解碼請求）。混合解碼請求會導致吞吐量下降，因為我們不了解內存帶寬和容量使用情況，從而導致爭用和排隊阻塞。而且，與全部輕量級解碼請求的批次相比，增加重量級解碼請求可能嚴重影響吞吐量和延遲。

根本原因很簡單：當前的LLM系統對LLM預填充和解碼階段所展示的不同特征一無所知。預填充階段類似于計算密集型的批處理作業，而解碼階段類似于內存密集型、延遲關鍵的任務。把兩者混合推理，會讓彼此干擾。

KV Cache

讀者可能會有疑問：因為解耦，所以需要在預填充和解碼 GPU 之間傳輸中間狀態（即 KV Cache），這看起來是一種瓶頸，因為當模型過大的時候，可能單機無法完全放下一個 Prefill Worker 和 Decode Worker。對于這種模型，將KV Cache從預填充轉移到解碼實例會產生顯著的開銷，遷移耗時則會大幅增加。

我們接下來利用DistServe進行分析。針對KV Cache傳輸問題，為了避免遷移開銷上升，需要合理放置 Worker。DistServe作者通過精心放置預填充和解碼工作器來利用高帶寬網絡來有效地隱藏 KV Cache 傳輸開銷。具體策略是：因為KVCache 按層存儲，所以可以把模型按層劃分多段，每段放到不同機器，不同段的模型采用 PP，然后再對同一段的模型進行單機的模型并行策略的搜索。這樣可以保證：

• 跨機傳輸僅出現在 PP 層間。
• Prefill 和 Decode Worker 相同層的 KVCache 在同一個機器內，Prefill 和 Decode Worker之間可以使用節點內NVLINK帶寬進行傳輸。

從而顯著減少傳輸延遲。并且當模型越大、序列越長、多卡通信設備帶寬越高，KVCache 遷移開銷占比越低。

下圖左：使用DistServe在ShareGPT數據集上提供OPT-175B時的延遲細分。右：三種OPT模型KV Cache傳輸時間的CDF效果。可以看到傳輸問題已經得到了極大的緩解。通過適當的放置（placement），KV Cache 傳輸開銷可以有效地最小化，甚至低于一個解碼步數的時間。

分別配置

Prefill 和 Decode 的特性很不同，所以參數配置上很多都是不同的。例如湊批策略，量化類型，TP 大小，PP 大小，任務超時時間，重試時間，顯存申請策略，RDMA 軟硬件隊列大小，是否可回退執行等。我們用阿里RTP-LLM的經驗來看看如何配置。

• batch size。分離之后，Prefill 本身的 Batch Size 仍然不會很大。為了充分發揮 Decode 的能力，我們應該配置更多的 Prefill，減少 Decode 實例個數，使得 Decode 能進行大 Batch 推理。
• 資源分配：Prefill 應該選擇計算能力強的卡型，Decode 應該著重選擇顯存大的卡型。
• 量化方案：可以在 Prefill 機器部署 FP16 / W8A8 的模型文件，來獲得相對不錯的 Prefill 性能；而 Decode 機器則可以自由地選擇 W4A16 / W4A8 的方案，來獲得整體的更優性能。

我們接下來利用幾篇論文來繼續學習。

• DistServe：將預填充和解碼計算分配至不同的 GPU，從而消除了二者之間的干擾。針對應用程序的總時延和每個 Token 運行時間的需求，DistServe 為每個階段量身定制了資源分配和并行策略的優化方案。此外，DistServe 還根據服務集群的帶寬來確定如何部署這兩個階段，以最小化由任務分解引起的通信。
• SplitWise：特點是分布式調度策略、PD分離、分層KV Cache傳輸，并且增設了第三個機器池，專門用于處理 Prefill 和 Decode 階段的混合批處理，并且能夠根據實時計算需求靈活調整其規模。
• MemServe：用統一的分布式視角對KV Cache進行管理。
• TetriInfer：結合了分塊預填充和兩階段分離，以及預測性的兩階段調度算法來優化資源利用率，且僅針對 prefill 進行分塊處理。SplitWise、DistServe 和 TetriInfer的靜態并行性和分區策略并不靈活，無法處理動態工作負載。
• Mooncake：構建了以 KVCache 存儲為中心、基于RDMA的 P-D 分離調度集群和推理架構，形成了 Prefill/Decode Pool 以及分布式異構介質 KVCache Pool；融合了緩存感知、負載均衡和以服務水平目標（SLO）為導向的決策機制。

4.2 DistServe

DistServe有如下幾個特色：

• 將預填充和解碼計算分配給不同的GPU，從而消除了二者之間的干擾。
• 根據應用程序的TTFT和TPOT要求，DistServe為每個階段分別定制了資源分配和并行策略的共同優化策略。
• 根據服務集群的帶寬來確定如何部署這兩個階段，以最小化由任務分解引起的通信。

4.2.1 分析

我們首先看看預填充和解碼對于并行策略的述求。

并行劃分

常見的并行技術如下：

• 數據并行：權重復制，輸入數據分布在不同機器上
• 模型并行：將權重進行切分，根據權重切分的方式，輸入數據有時候也要一起進行切分
• 流水線并行：把不同的子圖放在不同的設備上運行，一起形成流水線，這需要對輸入數據切分batch

Alpa提出了另外一個劃分維度，即算子間、算子內并行劃分方法，通過“是否切分了張量的維度”來區分不同的并行。

• 算子內并行：切分了tensor維度的并行方式，包括數據并行和算子并行。
• 算子間并行：不切分tensor，只是把子圖進行不同的擺放分布，包括流水線并行。

prefill實例

針對prefill實例，我們的優化目標是使用最少的資源滿足服務的TTFT延遲要求。

預填充步驟通常是計算密集型的。一旦GPU變為計算限制，批次中加入再多請求也不會提高GPU使用率，而是成比例地延長批處理的總處理時間，無意中延遲了所有的請求。對于預填充實例，需要事先對LLM和GPU進行分析，以確定關鍵的輸入長度閾值。超過該閾值，預填充階段將變為計算限制。只有當計劃請求的輸入長度低于閾值時，才應考慮在批處理中添加更多的請求。

因此，適合prefill實例的并行策略是：在較低的rate(req/s)下，執行時間是主要因素，因此，算子內并行（intra-op）更高效，而隨著到達率的增加，排隊延遲變得更加顯著，此時算子間并行（inter-op）更優越。

decode實例

對于解碼實例，我們的優化目標是以最小的計算資源滿足應用的TPOT需求。

由于單個解碼作業嚴重受帶寬限制，批處理是避免低GPU利用率（因此高每個GPU的吞吐量）的關鍵。disaggregation（解聚）可以將多個預填充實例分配給單個解碼實例。這種方法允許在專用GPU上積累更大的解碼階段批處理大小，而不會犧牲TPOT。在disaggregation之后，解碼的批處理大小可能受到GPU內存容量的限制，因為需要為所有活動請求維護KV Cache。

intra-op減少了延遲，但隨著通信和分區后利用率的降低而遞減，而inter-op幾乎可以線性地擴展吞吐量。因此適合decode實例的并行策略是：當TPOT SLO要求嚴格時，intra-op對于減少TPOT以滿足延遲目標至關重要。而inter-op更適合線性增加吞吐量。

4.2.2 優化方向

因為prefill和decode的不同特點，使得在分離式框架下，有不同的優化方向。下表第一列是prefill的特點，第二列是decode的特點，第三列是優化方向。比如針對存儲，我們可以使用不同型號的硬件；針對batch策略，我們對prefill和decode使用不同的策略，分別優化。

prefill	decode	優化方向
算完KV Cache，發給decode階段后，可以使用策略來清除KV Cache。	逐token生成過程中頻繁讀取KV Cache，需要盡可能保存KV Cache。	計算和存儲獨立優化
因為prefill階段是compute-bound，故隨著batch size的增加，吞吐量的增長趨勢趨于平緩。	因為decode階段是memory-bound，故隨著batch size的增加，吞吐量的增長趨勢越來越顯著。提升batch size就能提升計算強度，進而提升吞吐量。	batch策略獨立優化
在不同條件下，對并行方式有傾向性：rate (req/s)較小時，適合TP；rate較大時，適合PP。	隨GPU數量的增加，PP方式可以產生更高吞吐量；TP方式可以產生更低延遲（處理單個請求更快）。	并行策略獨立優化

因此，在給定模型、工作負載特性、延遲要求和SLO達成目標下，DistServe需要確定：

• 預填充和解碼實例的并行性策略
• 部署每個實例類型的數量
• 如何將它們放置在物理集群上

終極目標是找到最大化每個GPU吞吐量的配置策略。

4.2.3 算法

接下來我們看看DistServe如何搜索 Prefill、Decode 最優模型并行配置。

在給定不同的集群設置的情況下，一個關鍵的設計考慮因素是管理分解的預填充和解碼階段之間的通信。DistServe使用了兩種算法：一種用于具有高速跨節點網絡的集群，另一種用于缺乏此類基礎設施的環境，后者引入了額外的約束。總體思路是：在約束條件基礎上，對于給定現有的 GPU 數量等硬件資源，羅列出TP和PP的可能性。然后用模型器來對每個可能性進行計算，得出該可能性下每張卡的 Goodput，選擇可以讓每個 Worker 最大化 Goodput 的 TP、PP 配置。

高節點親和力的集群上如何放置

在具有高節點親和力的集群上通常裝有Infiniband，KV Cache在節點之間的傳輸開銷可以忽略不計，DistServe可以在沒有約束的情況下有效地部署預填充和解碼實例。我們提出了一個雙層放置算法用于這種情況：首先優化預填充和解碼實例的并行性配置，以實現階段級別的最大每個GPU吞吐量；然后，我們使用復制來匹配整體流量速率。

算法1概述了該過程。我們列舉了在集群容量限制下的，預填充和解碼實例的所有可行并行配置例如，對于特定的預填充階段配置，我們使用simu_prefill來模擬并找到它們的最大吞吐量（類似于使用simu_decode進行解碼）。在確定了預填充和解碼實例的最佳并行配置后，我們根據它們的吞吐量來復制它們，以實現用戶所需的總體流量率。

低節點親和力集群的放置

上面的算法假設我們可以將預填充和解碼放置在集群的任意兩個節點之間，并且KV Cache傳輸利用了高帶寬。然而，在許多真實的集群中，節點內的GPU通過高帶寬NVLINK進行訪問，而GPU的帶寬又很有限。接下來，論文將開發一種算法來解決這個約束。

算法的關鍵是：中間狀態的傳輸僅發生在預填充和解碼實例的相應層之間。算法具體如下：

• 我們首先枚舉層間并行性來獲取所有可能的實例段。利用層間并行性，我們將層分組為階段，并將每個實例劃分為段，稱為實例段，每個段維護一個特定的層間階段。通過將同一階段的預填充和解碼段放置在同一節點上，我們強制中間狀態的傳輸只通過NVLINK進行。但是對于大型模型，在一個8-GPU節點上可能無法承載甚至一個預填充和解碼實例對。我們將這作為額外的放置約束，并與模型并行性一起進行優化。

• 鑒于每個節點通常只有8個GPU的限制，對于每個段，我們通過調用get_intra_node_configs來枚舉所有可能的節點內配置。

• 然后，我們使用模擬器找到最佳配置并復制它以滿足目標流量率。

4.3 SplitWise

SplitWise和DistServe思路類似，具體如下。

• 階段識別：將LLM推理請求分為兩個階段，prompt計算和token生成。不同階段放在不同機器上執行。

• 硬件適配：為每個階段選擇最適合的硬件，優化資源使用。

• 集群設計：設計了同質和異質集群，針對吞吐量、成本和功耗進行優化。

• 調度策略：使用兩級調度系統，包括集群級調度器（CLS）和機器級調度器（MLS），以優化請求分配和機器資源管理。

因此我們只看看其特殊之處，或者說只看我們在DistServe中沒有介紹的地方，藉此補齊。

4.3.1 調度方案

當到達請求的 prompt 長度有差異性的時候，預填充和解碼就會出現壓力的不均衡問題。因為整體的吞吐取決于 P 和 D 的全局資源利用，當 P 過載但 D 閑置，或者 P 閑置但 D 過載的時候，成本和性能都不是最優的。所以就需要考慮在 P 和 D 之間做負載均衡，要么從整個節點層面直接切換 P 和 D 的角色，要么 P 和 D 節點能夠承擔一些混雜的請求，比如通過 chunked prefill。

首先我們來看一下作者給出了什么樣的解決方案。下圖展現了SplitWise的整體架構圖，

• 首先，架構圖上有三個資源池，分別是Prompt Pool，Token Pool和Mixed Pool。在Prompt Pool中只做Prefill階段，Token Pool中只做Decode階段，實現了Prefill和Decode的分離部署。在Mixed Pool中以合并部署的方式執行Prefill和Decode混合的Batch，當請求負載變化劇烈、實例來不及伸縮或角色轉換時，由Mixed Pool來執行混合的Batch。
• 其次，架構圖上還有兩個級別的調度器，分別是cluster-level scheduler (CLS)和machine-level scheduler (MLS) 。集群級調度程序（圖上標號1）負責將傳入的請求路由到特定的機器并重新分配機器。機器級調度程序（圖上標號2）維護待處理隊列并管理每個機器上的請求批處理。

集群級調度

集群級調度程序（圖上標號1）負責將傳入的請求路由到特定的機器并重新分配機器。主要功能如下：

• 請求路由。CLS針對請求采用“加入最短隊列”的方法（Join the Shortest Queue (JSQ)）的路由策略進行調度。在請求來的時候，CLS就決定后續要執行Prefill和Decode的實例。每臺機器向CLS報告其內存容量或待處理隊列的任何變化。
• 機器管理。CLS主要維護兩個機器池：即時機器池和token機器池。Splitwise根據輸入/輸出token分布和預期負載（即每秒請求數）將機器初始分配給一個池。如果這些值與初始假設相差較大，Splitwise會對機器進行粗粒度重新分配，并在即時機器池和token機器池之間移動機器。
• 混合機器池。為了滿足SLO并避免在負載較高時出現性能斷崖，Splitwise維護一個特殊的混合池，混合池是一個中間緩沖。提示池或token池中的機器可以動態地移入和移出混合池，混合池根據請求速率和輸入/輸出token的分布動態增長和收縮，而且沒有明顯的池切換延遲。如果CLS嘗試使用JSQ為請求分配即時機器和token機器，并發現所選機器的隊列超過閾值，則會在混合池中尋找目標機器。混合池中的機器與非Splitwise機器完全相同，采用混合批處理方式運行。一旦混合請求隊列被處理完畢，CLS會將機器過渡回其原始池。例如，當隊列過長時，我們可以將即時機器移入混合池以運行token；一旦機器完成運行token，我們將機器過渡回即時池。

機器級調度（Cluster-level scheduling）

機器級調度程序（圖上標號2）在每臺機器上運行，負責跟蹤GPU內存利用率，維護待處理隊列，選擇每次迭代的批量大小和批量請求，并向CLS報告相關狀態。

• Prompt 機器（Prompt machines）。MLS使用先到先服務（FCFS）調度算法來調度請求。因為token數目過了一定閾值之后，吞吐量開始下降。因此，MLS限制了將多個即時請求批量處理到總共2048個token。這是一個可配置的值，對于不同的模型或硬件可能會有所不同。

• token機器（Token machines）。MLS使用FCFS調度算法盡可能地調度token和批量處理。token生成吞吐量隨著批量大小的增加而增加，直到機器的內存耗盡。因此，當機器接近內存耗盡時，MLS會開始把token放到隊列中。

• 混合機器（Mixed machines）。為了滿足TTFT的SLO，MLS需要優先運行即時請求，并立即調度任何新的即時請求進入待處理隊列。如果機器正在運行decode階段并且沒有容量，MLS將搶占token。為了避免由于搶占而導致decode階段饑餓，我們增加較老token的優先級，并限制每個token可以進行的搶占次數。

4.3.2 KV Cache傳輸

在Splitwise中，我們需要將KV Cache從prompt機傳輸到token機以完成推理。KV Cache是在請求的提示階段生成的，并在token生成階段不斷增長。因此作者提出，他們給出的這種分離架構的方案引入的最大開銷就是把KV Cache從Prefill階段遷移到Decode階段的開銷。這種傳輸延遲是與Splitwise相關的主要開銷。在本節中，我們將討論KVcache傳輸的影響以及如何優化它。

將活躍KV Cache從預填到解碼實例的傳輸有兩種方式：按層或按請求。

• 按請求方法是在預填階段完成后傳輸KV Cache。上圖左側使用的是未經優化的以串行方式運行的線性KV Cache傳輸，單個請求批次所經歷的時間就包括Prompt時間，等待KV Cache傳輸的時間和生成Token的時間。KV Cache傳輸僅在即時階段完成且生成第一個token后才開始。此外，它需要在下一個輸出token可以生成之前完成。這直接影響了推理的最大TBT和端到端延遲。傳輸所需的時間取決于KV Cache的大小（與即時token數量成正比）和即時機器與token機器之間的互連帶寬。即使在機器之間使用快速的InfiniBand連接，對于大型即時大小，KV Cache的開銷很容易成為TBT的顯著部分。

• 按層方法在一層完成計算后傳輸KV Cache。右側是SplitWise優化后的異步KV Cache傳輸。由于預填充是逐層處理的，在完成一層的計算以后，就可以將這一層的 KV Cache 發送出去。因此可以將KVCache的傳輸和轉儲與計算重疊。在每一層的注意力計算開始之前，模型等待該層的KVCache的異步加載完成，并觸發下一層的異步KV Cache加載。在完成關注度計算后，啟動該層的KVCache的異步存儲。傳輸重疊允許預填充實例的執行時間大致等于KVCache加載時間或標準預填充時間，具體取決于前綴緩存的相對比例。Splitwise發現按層方式優于按請求方式，因為按層方式重疊了計算和通信。

由于批次中的token數量在計算開始時已知，Splitwise選擇最佳的KV Cache傳輸技術。

• 一方面，對于較大的prompt，Splitwise使用逐層傳輸。逐層KV Cache傳輸與下一層的快速計算并行進行。這需要每層細粒度的同步以確保正確性。因此，可能會產生性能干擾并增加TTFT。
• 因此，對于小的prompt，Splitwise使用序列化的KV Cache傳輸。這是因為小的prompt的KV Cache很小，不需要支付由每層傳輸所需的細粒度層同步的開銷。

4.4 MemServe

MemServe通過其MemPool系統跨服務實例管理分布式CPU DRAM和GPU HBM資源。通過一套全面的分布式內存池API來管理活躍KV Cache和歷史KV Cache。MemServe為歷史KV Cache檢索提供了一個基于token的快速索引層，抽象出硬件異構性的跨實例數據交換機制，以及一個實現基于提示樹的本地化策略的全局調度器，以增強緩存重用。這些策略相輔相成，共同顯著提高了作業完成時間和首次token性能。

因為上下文緩存和分布式推理等技術優化擴展了KV Cache的壽命和領域，所以LLM服務已經從無狀態轉變為有狀態系統，這迫切需要一種新的架構方法。針對這個問題，論文作者提出了MemServe（Memory-enhanced model Servin），借此在統一系統內處理請求間和請求內的優化。上下文緩存和分布式推理主要是從流水線角度進行優化，MemServe 方案是把 KVCache 當做分布式內存進行管理，這是統一的分布式系統視角，decode 機器和 prefill 機器都是和pool進行交互，不再是個單純的流水線的邏輯了。

MemServe特點如下：

• 為了解決在分布式實例之間管理KV Cache的挑戰，MemServe引入了一個跨服務的管理分布式內存和KV Cache的彈性內存池，或者稱為MemPool，它是管理所有集群內存的基座，包括CPU DRAM和GPU HBM。或者說，MemServe將MemPool組件抽象出來，然后將分布式推理作為MemPool的一個用例來構建。
• MemPool提供了一套豐富的API用于管理分布式內存和KV Cache。利用這些API，MemServe將上下文緩存與分布式推理結合起來。
• 為了最大化KV Cache的重用，MemServe采用了一個全局調度器，該調度器使用新穎的全局提示樹的局部感知策略來增強緩存重用。

4.4.1 動機

利用LLM中依賴關系來對KV Cache進行優化的技術可以分為兩種類型：跨請求和請求內。

• 請求內優化：這種優化類型利用單個請求內的依賴關系以增強性能。兩個顯著的例子是分布式推理（將一個請求分成兩個子請求以獲得更好的調度）和序列并行（將一個請求分成多個子請求以分發負載）。
• 請求間優化：這種優化類型利用請求之間的依賴關系以獲得更好的性能。上下文緩存（Context caching）是該類別中唯一已知的技術。為了構建上下文緩存，模型存儲并重用自注意機制中的KV Cache，以避免在類似或重復請求之間進行冗余計算。這在多個請求共享公共前綴或上下文的情況下非常有用。

這些依賴利用技術中的一個共同主題是，它們需要新穎的邏輯來管理和傳輸KV Cache。

• 在請求內方法中，需要將KV Cache的作用域從單個實例擴展到分布式實例。這就需要有效的機制來在實例之間傳輸KV Cache，分布式推理和序列并行都是如此。
• 請求間方法需要將KV Cache的生命周期從單個請求延長到潛在無限，即跨請求。這依賴兩種機制：首先，需要一個索引來找到請求之間的依賴關系，從而找到保留的KV Cache。其次，需要修改的推理引擎和注意力核心來重用歷史KV Cache。

由于目前方案缺少機制來管理LLM的中間KV Cache數據，因此存在兩個關鍵問題。

• LLM服務系統無法同時應用任何現有的請求間和請求內依賴性利用優化。當前的上下文緩存（請求間）方法設計時沒有考慮到請求內場景。因此，分離的推理（請求內）無法從上下文緩存中受益，因為它缺乏利用KV Cache從解碼返回到預填充實例以供將來重用的機制。同樣，序列并行性將KV Cache分布到多個實例，但是缺乏保留和重用它所需的機制和算法。這個問題是由于請求內技術將一個緊密耦合的請求分解為多個松散耦合的子請求，使得在分布式環境中復雜化了KV Cache管理。
• LLM服務系統缺乏一個全面的、自上而下的設計來有效地利用現有的請求間技術。上下文緩存通過在同一服務實例中運行共享相同前綴的請求來重復使用歷史KV Cache。然而，當前的LLM服務系統根據負載或會話ID跨多個服務實例調度請求，這未能最大化跨會話的KV Cache重用。

這些問題的出現是因為現有的LLM服務系統是建立在KV Cache僅僅是單個請求在單個實例上的中間數據的假設之上的。隨著新興的依賴利用技術的出現，KV Cache的生命周期已經延長，其管理范圍也擴展到了分布式設置。這種范式轉變需要對LLM服務架構進行根本性的重新思考。

4.4.2 方案

MemServe被設計為一個大規模的、可以高效處理請求間和請求內優化的LLM服務系統。它包括三個主要組件：全局調度器、多種類型推理實例和一個彈性內存池（MemPool），MemServe是這么介紹自己的：we propose Memory-enhanced model Serving, or MemServe, to handle inter-request and intra-request optimizations within a uniffed system.

這里有幾個關鍵點：

• 統一系統。這里涉及了幾種 instances 類型，P-only、D-only、PD-colocated 機器類型（PD 就是 prefill-decode 的意思），同時在運算時就分為 (1) PD-colocated（只用第三種 instances 的）, (2) PD-colocated with caching, (3) PD-disaggregated （用前兩種 instances、或者混著一起用）, (4) PD-disaggregated with caching 。
• 內存增強（Memory-enhanced）。在分散推理的 setting 下引入了一個彈性內存池 MemPool，用于管理分布式內存和跨服務實例的 KVCache。

彈性內存池

MemPool是MemServe的核心組件，提供三種類型的API：內存、索引和分布式數據傳輸。它在每個推理實例內運行，使用固定大小的內存分配器管理所有本地內存。

內存池管理推理集群中的所有內存，包括CPU DRAM和GPU HBM。內存池在每個推理實例中運行，共同提供一組分布式內存池API。它管理正在進行的請求使用的活躍KV Cache以及請求完成后保留的歷史KV Cache。

下圖展示了MemPool啟用的用例。圓1是上下文緩存。圓2是分解推理。圓3是序列并行性。灰色實線表示MemPool索引API調用。實線表示MemPool分布式API。MemPool在一個平臺上支持所有用例。

索引

當 KVCache 越來越大的時候，內存池中的布局就是一個非常大的問題，所以MemPool使用內部索引將提示標記映射到KV歷史緩存，管理正在進行的請求的活躍KV Cache以及請求完成后保留的歷史KV Cache，確保高效檢索緩存數據。

有了索引，內存的檢索就靠索引了，其實也就是給 pagedattention 中的 page 加一個 id，可以通過其他的檢索機制來優化這個檢索過程。論文中給出的就是 global prompt trees 的結構機制來進行優化。

每當引擎調用插入、匹配、刪除等操作時，MemPool都會遍歷該索引。LLM服務世界有三種索引方法：標記、會話和文檔ID。基于標記的索引通用性強，因為它適用于任何共享提示前綴情況。會話和文檔ID索引更簡單，但只能在聊天會話內或跨會話使用相同文檔時重用共享提示。論文采用基于標記的索引方法，以實現廣泛適用性。為了實現這個索引，MemPool利用了SGLang提出的基數樹。

調度

MemServe 的全局調度器一方面負責整個框架和調度，另一方面還維護了 global prompt trees 全局提示樹 和 locality-aware scheduling policy 本地感知調度等策略來優化內存管理和 KVCache 管理。每個節點上會有對全局樹型緩存進行分布式維護。調度的時候，請求的提示詞會通過對所有類型的樹查詢完成對全局的查詢，就可以通過策略模塊，針對分布式負載情況選擇具有最長公共前綴（即最大保留歷史KVCache）的實例，以達到最優的檢索和訪問效率了。

傳輸API

MemPool提供了一個簡單的數據傳輸API，抽象了三種異構性：并行性、網絡和內存介質。MemServe通過使用MemPool API在四個步驟中彌合了上下文緩存（請求間）和分布式推理（請求內）之間的差距：(a)首先使用分布式API來復制分布式推理，(b)然后使用索引API向僅預填充實例添加緩存，?將相同的緩存應用于僅解碼實例，(d)最后我們啟用解碼到預填充數據傳輸。下圖展示了使用MemPool API通過上下文緩存增強分解推理。引擎箱（engine box）是指一個經過調整的推理引擎，如vLLM。圓圈數字表示構建解決方案所采取的步驟。A-KV是活躍KV Cache。H-KV是歷史KV Cache。

我們接下來看看如何在四個不同 level 的緩存設計機制中逐步構建一個完整的設計。

• PD-Basic。這是naive/vanilla 的機制，其實就是 DistServe 和 Splitwise提出的基本的解聚推理架構。作為優化的 baseline 。

• PD-Caching-1 只在 P-only instances 啟用緩存。將活躍KV Cache退休為歷史KV Cache，以便未來的推理可以利用保存的數據減少重計算（上圖中的第2步）。這種緩存設計只保留了預填階段產生的歷史KV Cache，而沒有保留解碼階段的緩存，因此適用于共享長公共前綴提示的工作負載，例如系統前綴。這種設計的主要缺點是，在多輪對話場景中（例如文檔QA），預填實例需要重復地將相同的活躍KV Cache集轉發給解碼實例，浪費帶寬并影響到第二個標記的時間。因此，我們提出下一個設計方案來解決這個問題。

• PD-Caching-2 在 D-only 也就是 decoder 節點進行緩存。這個設計在解碼實例中啟用緩存，以減少重復的數據移動。我們在PD-Caching-1的基礎上進行了兩個關鍵改變。首先，預填實例現在調用transfer_with_insert而不是transfer，這樣解碼實例將預填階段產生的傳輸的KV Cache插入到其本地索引中。其次，在請求完成后，解碼實例調用insert將解碼階段產生的KV Cache保留到其本地索引中。借助于局部感知調度，預填實例現在只需要遞增地傳輸新的KV Cache數據。盡管這種設計減少了從預填實例到解碼實例的數據移動，但它并沒有改善預填實例的上下文緩存，因為它缺乏來自解碼階段的歷史KV Cache。因此，在多輪對話場景中，隨著提示增加，上下文緩存的好處保持不變。因此，我們提出下一個設計方案來解決這個問題。

• PD-Caching-3 允許P-only 和 D-only 的機器可以互相傳輸并維護兩邊的索引。這個設計實現了解聚推理架構的全面上下文緩存。我們在PD-Caching-2的基礎上進行了一個改變：在請求完成后，解碼實例調用transfer_with_insert將解碼階段產生的KV Cache傳輸到預填實例（上圖中的第5步）。因此，預填實例保留的歷史KV Cache增長，上下文緩存的好處隨著輪數的增加線性增加。我們其實可以把這里的 Caching Pool 理解為一個不完全的同步所有數據的異步共享機制，各種類型的節點間即要維護自己的緩存也要按需推送和拉取部分遠程數據。

傳輸操作

將活躍KV Cache從預填到解碼實例的傳輸有兩種方式：按層或按請求。按層方法在一層完成計算后傳輸KV Cache。按請求方法在預填階段完成后傳輸KV Cache。Splitwise發現按層優于按請求，因為它重疊了計算和通信，從而加快了時間到第二個token（或TTST）。

當負載較低時，MemServe作者也觀察到了同樣現象。然而，隨著負載的增加，兩者都會產生非常大的開銷，根本原因是(1)內存布局離散，(2)網絡原語缺失。

• 基于分頁的動態內存管理，如PagedAttention，是現在LLM服務系統中的事實標準。無論分頁機制是在引擎 還是在驅動程序中實現，KV Cache都被分區并存儲在固定大小的內存塊中。塊大小是可配置的，通常以token數量為單位。而現有的引擎以細粒度的方式管理KV Cache。例如，vLLM為每個LLM層分配兩個塊。對于具有L層和每個塊8個token的LLM，引擎需要2? L塊來存儲8個token的KV Cache。盡管分頁提高了利用率，但離散的內存布局在使用現有AI網絡堆棧實現分布式推理時面臨巨大挑戰。

• AI中的網絡堆棧技術大多是諸如NCCL之類的集合通信庫。這些庫最適合使用張量或流水線并行處理的典型AI工作負載，但在支持LLM服務的請求內優化方面，如分布式推理或序列并行，集合通信庫表現不佳。LLM服務的這些新模式需要在HBM或DRAM之間進行有效的點對點、收集和分散原語。但是由于KV Cache是離散的，因此無論使用按層還是按請求方法，網絡API調用的數量等于離散內存塊的數量，這就是為什么兩者隨著負載增加都會產生開銷的根本原因。

為了解決分頁和網絡原語不足引起的挑戰，我們提出通過將較小的KV塊聚合成大塊來減少碎片化，類似于使用大頁。具體而言，我們將每層的兩個塊聚合成一個塊：新的塊大小等于 2? L 個較小的塊。這有效地將網絡API調用的數量減少了 2? L 倍。這種優化僅適用于按請求方法，因為按層方法不可避免地需要至少調用網絡API 次。

上圖比較了按層、按請求和按請求聚合（論文提出的優化）的跨內存布局和傳輸時間線。論文的測試表明，在低負載下，按層實現了最低的JCT(Job Completion Time)，但在高負載下，因為減少了網絡調用，通過層聚合（bylayer-agg）優于按層。

4.5 TetriInfer

論文”Inference without Interference: Disaggregate LLM Inference for Mixed Downstream Workloads“提出了TetriInfer，TetriInfer的意思是希望能像俄羅斯方塊一樣堆疊好空間，不要亂放出空當。

4.5.1 動機

避免干擾的一個簡單解決方案是為每個下游任務靜態地分配資源。鑒于LLM服務基礎設施的高成本，這個解決方案是不實際的。因此，TetriInfer構建了一個分布式的LLM推理服務系統，通過根據請求特征仔細安排和分組來解決這些問題。TetriInfer作者的思路如下。

• 分離 prefill 與 decode 到不同的硬件。TetriInfer具有專用的預填充和解碼實例，這兩個是虛擬概念，每個實例可以獨立擴展并在負載變化時翻轉角色。TetriInfer將預填充請求僅調度到預填充實例，解碼請求也是如此。預填充實例將預填充的KV Cache傳輸給解碼實例。
• 限制 prefill 的 size，就是找到使得 batch 剛好處在 compute-bound 和 memory-bound 的臨界點上。
  • 為了避免預填充階段的干擾，應該限制在單個預填充迭代中處理的token數量，以便充分利用硬件而不產生額外的懲罰。于是TetriInfer將輸入提示分割并填充為固定大小的塊，使預填充在一個固定大小的計算單元中運行，以便加速器始終接近其計算飽和限制。
  • Sarathi也提出了用于相同目的的分塊預填充（Chunked-prefills），但是Sarathi是將預填充-解碼混合在一起。相比之下，TetriInferj僅涉及運行預填充塊，因為TetriInfer已經將LLM的預填充和解碼分解為單獨的實例。
• 使用基于輸出的長度預測來調度 decode 階段任務，以避免調度熱點。TetriInfer利用基于LLM的長度預測模型來推測解碼請求的生成token數量，然后相應地進行調度。
  • 在每個預填充實例上運行長度預測器，因此預填充實例可以根據解碼實例的負載情況做出明智的決策，以確保某些解碼請求具有足夠的資源來運行。長度預測器使用小型LLM模型進行預測，并繼續使用固定大小的批次，而不是分塊預填充。這是由于該模型的小尺寸，其沒有像較大模型那樣明顯的計算飽和閾值。
  • 使用智能的兩級調度算法，并結合預測的資源使用情況，以避免解碼調度熱點。

4.5.2 實現

TetriInfer具體分為四個主要模塊。集中式控制平面、預填充實例、解碼實例和長度預測模型。

• 集中式控制平面。它由全局調度器和集群監視器組成。全局調度器根據負載向預填充實例發送請求，并接收來自解碼實例的流式輸出。集群監視器從預填充和解碼實例收集統計信息，并定期向預填充實例廣播負載信息。它添加、刪除和翻轉預填充或解碼實例。

• 預填充實例。它們只運行LLM推理請求的預填充階段。每個預填充實例都有一個本地調度器、一個長度預測器、主LLM引擎和一個調度器。

  • 為了避免預填充請求之間的干擾，我們使用預填充調度器和分塊預填充來對所有提示進行排序和分區為固定大小的塊。
  • 預填充實例的調度器對提高預填充階段的延遲和吞吐量至關重要。調度器維護一個原始請求隊列，用于存儲全局調度器的請求，以及一個已調度隊列，用于存儲排序后的請求。在這項工作中，我們設計并實現了三種調度策略：先到先服務（FCFS）、最短作業優先（SJF）和最長作業優先（LJF）。我們之所以可以使用后兩種策略，是因為我們可以根據提示中的標記數量準確估計請求的預填充時間。

• 解碼實例。它們在虛擬上與預填充實例分離，并且只運行LLM推理請求的解碼階段。每個解碼實例可以接收來自任何預填充實例的請求。它運行一個本地調度器，具有三個預定義的策略，用于選擇要在主LLM引擎中運行的解碼請求。

• 長度預測模型。預測模型是一個經過離線微調的小型LLM模型，用于預測LLM推理請求的生成長度。TetriInfer的預填充調度器和解碼實例的本地調度器利用推測的信息來調度解碼實例，避免了在§2.2.3中測量到的熱點問題。

4.6 Mooncake

Mooncake的核心理念是用存儲資源換取計算效率，其采用了以KVCache為中心的分離架構，將預取和解碼集群分開（將LLM服務的不同部分分解為專用資源池），并利用RDMA構建跨節點共享緩存。它還利用GPU集群的CPU、DRAM和SSD資源來實現KVCache的分層緩存。

4.6.1 洞見

論文的核心目標和主要洞見如下：

• 因為長上下文的需求始終存在，所以KVCache 的容量會長期保持高位。

• 盡可能多地重用 KVCache 以減少所需的計算資源和在節點間轉移數據的耗時。

• 最大化每個批次中的token數量以提高模型 FLOPs 利用率。

• 遠程調用 KVCache 會延長 TTFT，而大的batch size會導致更大的 TBT 。

• 存儲在較低層級存儲上的KVCache 會帶來更多問題。

4.6.2 實現

分離式架構

如下圖所示，Mooncake采用了分離式架構。這里分離有兩層含義：不僅將預填充與解碼節點分開，還將GPU集群的CPU、DRAM、SSD和RDMA資源分組，以實現分離的KVCache。這種分離的緩存利用了未充分利用的資源，提供了充足的緩存容量和傳輸帶寬。

• 將Prefill與Decode計算資源分開。Prefill階段優化目標是利用request間存在共同前綴的機會，盡可能復用KVCache，同時滿足TTFT，最大化MFU，讓KVCache小于CPU內存限制。Decode優化目標為最大化吞吐，滿足TBT，讓KVCache小于GPU顯存限制。

• 將KVCache和計算分離開。Mooncake 將GPU集群的CPU、DRAM、SSD和RDMA資源分組組成Distributed KVCache Pool，通過共享緩存提升命中率，減少重復計算。KVCache也是分塊以Paged方式管理。

所以，Mooncake 將單個同構 GPU 集群的資源打散并重新組織成三個可以獨立彈性伸縮的資源池。

• prefill pool。 Prefill Pool 處理用戶輸入，主要對TTFT負責。因為 Prefill 相對計算密集，這一部分也承擔著抬高整體資源利用率的任務。
• decoding pool。Prefill 處理完之后對應的 KVCache 會被送到 Decode Pool 進行自回歸式的流式輸出。雖然我們希望盡可能在一個 batch 中聚合更多 token 以提升 MFU，但這一部分主要需要對 TBT 負責。
• kv cache pool：利用每臺 HGX 機器上組成了一個 KVCache Pool 來進行全局的 Prefix Cache。全局 Prefix Cache 通過全局的調度能夠大幅度提升復用率從而提升總吞吐。由此帶來了一系列如何調度，分配，復制 KVCache 的問題，而且必須和 Prefill/Decode 的調度共同設計，因此我們把 Mooncake 的架構稱之為以 KVCache 為中心的架構。

我們接下來看看幾種pool的設計思路。

prefill pool

為了設計一個單獨的預填充節點池來無縫處理上下文長度的動態分配，Mooncake采用分塊流水線并行機制（CPP/chunked pipeline parallelism）來擴展單個請求的處理跨多個節點。與基于傳統序列并行的解決方案相比，CPP減少了網絡消耗，并簡化了對頻繁彈性擴展的依賴，而且通過進一步逐層預填充（Layer-wise Prefill）補充后，KVCache的流傳輸可以重疊延遲。這對于降低長上下文輸入的TTFT是非常必要的。

Layer-wise Prefill

分塊/層完成 prefill，即 block/layer-wise 的設計，可以理解為對 pagedattention 中 page 的一種細化，也是使用一個統一視角的 pool 來統一維護 KVCache 的存儲。

預填充階段的主要目標是盡可能多地重用KVCache，以避免冗余計算。因此，使用最頻繁的kvcache block應該復制到多個節點以避免獲取擁塞，而不怎么使用的block應該被換出以降低預留成本。

由于預加載是逐層進行且受計算限制，因此可以重疊 KVCache 的傳輸和轉儲與計算，進一步降低其占用成本，從而有效降低長上下文請求的延遲。在 Mooncake 中，KVCache 的加載和存儲通過啟動和等待操作異步執行。在每個層的注意力計算開始之前，模型會等待該層的 KVCache 的異步加載完成，并觸發下一層的異步 KVCache 加載。在注意力計算完成后，會啟動該層 KVCache 的異步存儲。一旦所有層的計算完成，進程會等待所有異步存儲操作完成。傳輸重疊使預加載實例的執行時間大致等同于 KVCache 加載時間或標準預加載時間，具體取決于前綴緩存比例相對于輸入長度的大小。

Multi-node Prefill

由于長上下文預填充中存在豐富的并行性，人們提出了很多并行策略。

• 數據并行：長序列Prefill的batch size是1，沒法數據并行。

• 張量并行：雖然用多個GPU并行處理長上下文請求是可取的。然而，將張量并行性擴展到超過一個節點需要每層做兩個昂貴的、基于RDMA的全局歸約操作，這顯著降低了預填充節點的MFU。所以張量并行難以拓展出節點。

• 序列并行。序列并行（SP）將請求的輸入序列分區到不同的節點上，以實現加速。這些SP方法利用了注意力操作符的關聯屬性，在Ring Attention 或Striped Attention的實現過程中至少每層需要跨節點通信一次。這大大降低了網絡消耗并改善了MFU。然而，采用SP仍然導致MFU較使用僅單節點TP時更差。而且，SP推理需要在每個卡上復制模型參數，對大模型來說是沉重的負擔。而且，SP每層都要通信，占用寶貴的網絡帶寬，這些帶寬更應該留給KV Cache的傳輸。

理想的部署應該將預填充節點組織成兩組：一組僅使用TP，另一組使用SP。只有在需要滿足TTFT SLO時才將請求分派給SP組。但是這種進一步的分解導致在動態調整每組節點數量時出現問題，因為靜態并行設置可能導致整個集群的利用率降低。

為了解決這個問題，Mooncake利用了僅解碼器模型的自回歸屬性，并為長上下文預填充實現了分塊流水線并行（CPP）。Mooncake將預填充集群中的每X個節點分組成一個流水線預填充節點組。對于每個請求，其輸入標記被分成塊，每個塊不超過預填充塊。同一請求的不同塊可以由不同節點同時處理，從而實現處理的并行化并減少TTFT。CPP提供了兩個主要好處：1）類似于訓練中的流水線并行，它僅在每個管道階段的邊界處需要跨節點通信，這可以很容易地與計算重疊。這導致更好的MFU和與KVCache傳輸的網絡資源爭用更少。2）它自然適用于短和長上下文，對于短上下文預填充沒有明顯的開銷，并避免頻繁地動態調整節點分區。

這種基于流水線的加速方法已經在訓練系統中進行了探索，但Mooncake首次在推理階段應用。

4.6.3 Decode Pool

Decode Pool 在業界的實踐是最多的，此處略過。

4.6.4 KVCache pool

分布式KVCache設計支持多對話會話共享緩存，避免重復存儲和計算；也提供分層存儲管理：高頻KVCache駐留DRAM，低頻數據下沉至SSD。

PD 直連就是預填充節點直接將 KV Cache 發送給解碼節點，它的好處是延遲低，但是會將P、D 節點耦合在一起。一旦解碼節點出現了問題，重新調度時無法僅調度 D 節點，需要重新進行整個預填充、解碼過程，導致代價太高。使用 KV Cache Store/Pool 是在 P 和 D 之間增加了一個中間存儲，預填充節點先將 KV Cache 寫到中間存儲，解碼節點從中間存儲讀。這樣做數據會多傳輸一次，增加了延遲和實現復雜度。好處是將P和D節點進行解耦合，容錯性更好。預填充階段也可以利用這個中間存儲做 Prefix Caching。

根據請求模式，KVCache pool可以使用緩存驅逐算法，如LRU（最近最少使用）、LFU（最不頻繁使用）或基于請求特征的緩存淘汰算法來進行調整。這些KVCache塊在CPU和GPU之間的傳輸由一個單獨的（GPUDirect）基于RDMA的組件Messenger處理。這種架構還使我們能夠為外部用戶提供上下文緩存API，以實現對KVCache的更高重用。

下圖顯示了CPU內存中的KVCache池，以及KVCache塊的存儲和傳輸邏輯。每個塊都附有一個哈希值，該哈希值由其自身的哈希值和用于重復數據刪除的prefix決定。

4.6.5 工作流

為了調度所有這些分散的組件，Mooncake實現了一個名為Conductor的全局調度程序。Conductor負責根據當前的KVCache分布和工作負載分派請求。如果對未來的推理有益，它還會復制或交換某些KVCache塊。

因為Mooncake 利用了 GPU 集群中未充分利用的 CPU、DRAM 和 SSD 資源希望可以得到充分利用。再簡單說就是GPU和CPU、內存等資源盡量打滿。那么就需要判斷當前推理瓶頸在什么地方？若是計算瓶頸，但所有的GPU打滿了，那么是不是可以放到CPU進行計算。若是存儲瓶頸，但是顯存已經占滿了，那么是不是可以放到內存中。若是傳輸瓶頸，是不是可以進行量化或者其他方式。

對于每個新請求，Conductor根據請求長度和prefix_len（因實例而異）估計相應的執行時間。然后將該請求的估計等待時間相加，以獲取該實例上的TTFT。Conductor將請求分配給TTFT最短的實例，并相應地更新該實例的緩存和隊列時間。為了預測請求的預填充階段的計算時間，Mooncake使用了從離線測試數據派生的預測模型。該模型基于請求的長度和前綴緩存命中長度來估計預填充計算時間。請求的排隊時間是通過匯總所有排隊請求的預填滿時間來計算的。

由于較高的實例負載，請求可能不總是被定向到具有最長前綴高速緩存長度的預填充實例。在這種情況下，如果估計的額外prefilling時間短于傳輸時間，則Conductor將高速緩存的位置和請求轉發到替代實例。如果最佳遠程前綴匹配長度不大于當前本地可重復使用的前綴乘以閾值，則我們更愿意直接計算輸入token。

上圖展示了請求的典型工作流程。當一個請求到達之后，Conductor選擇一對預處理節點和一個解碼節點，并啟動包含四個步驟的工作流程：

• KVCache重用。這是一種以KVCache為中心的調度。所選的預填充節點（組）接收到一個包含三個部分的請求：原始輸入、可重用的前綴緩存塊ID和分配給請求的完整緩存塊ID。Conductor根據前綴緩存塊ID將前綴緩存從遠程CPU內存加載到GPU內存，以啟動請求。如果不存在前綴緩存，則跳過此步驟。這種選擇平衡了三個目標：盡可能多地重用KVCache、平衡不同預填充節點的工作負載，并保證TTFT SLO。
• 增量預填充。預填充節點（組）使用前綴緩存完成預填充階段，并將新生成的增量KVCache存儲回CPU內存。如果未緩存的輸入token數量超過一定閾值（prefill_chunk），則將預填充階段分成多個塊并以流水線方式執行。此閾值被設定為一個可以充分利用相應GPU的計算能力的數值，通常大于1000個token。
• KVCache傳輸。每個節點上都部署一個Messenger服務，以管理和傳輸這些緩存。每個Messenger作為其相應推理實例中的獨立進程運行。Messanger將每個模型層生成的KVCache流式傳輸到目標解碼節點的CPU內存，此KV Cache的傳輸是異步執行的，并與上述增量預填充步驟重疊，這樣可以減少等待時間。
• 解碼。當解碼節點的CPU DRAM中接收到所有KVCache后，Mooncake會以連續批處理的方式加入下一批請求。Conductor根據解碼節點當前的負載預先選擇解碼節點，以確保不違反TBT SLO。然而，本地調度器（local scheduler）會再次檢查這個SLO，因為在預填階段之后，預期的負載可能已經發生變化。這種雙重檢查可能會導致請求被拒絕，這種情況下對應的預填成本就浪費了。

4.6.6 調度

Mooncake實現了基于KVCache的、平衡了實例負載和用戶體驗的請求調度算法。

下面的算法1詳細介紹了緩存感知預填充調度機制。對于每個新請求，其輸入token被劃分為幾個塊，并為每個塊計算哈希鍵。這涉及在一個塊中生成一個與前一個塊的哈希鍵連接的token哈希鍵（如果可用）。然后將請求的塊鍵逐一與每個預填充實例的緩存鍵進行比較，以識別前綴匹配長度（pref ix_len）。

為何可以進行全局調度？這是因為對于每 X byte 的 KVCache，其重新生成所需的算力正比于 X*hd 再乘以一個較大的常數，hd 是模型的 hidden dimension。因此只要每卡算力和每卡通訊帶寬的比值小于 hd 乘以這個常數，那么相比原地重算，從遠端傳輸 KVCache 不僅僅減少了計算量，還減少了 TTFT。

另外，下面還涉及到一個早期拒絕策略。Mooncake估計了兩階段的負載，在開始時處理前對于可能超出負載的請求直接退回或延遲處理，優化服務器超載問題。早期拒絕策略減少了過載場景中計算資源的浪費。如果在預填充階段之后沒有可用的解碼時隙，我們需要提前拒絕某些請求，以節省浪費的計算資源。

0xEE 個人信息

★★★★★★關于生活和技術的思考★★★★★★

微信公眾賬號：羅西的思考

如果您想及時得到個人撰寫文章的消息推送，或者想看看個人推薦的技術資料，敬請關注。

0xFF 參考

Orca: A Distributed Serving System for Transformer-Based Generative Models

SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills

手抓餅熊：Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving

手抓餅熊：MemServe: Context Caching for Disaggregated LLM Serving with Elastic Memory Pool

手抓餅熊：CachedAttention(原AttentionStore)

手抓餅熊：大模型Prefix場景Attention優化（三）

聊聊大模型推理服務中的優化問題 刀刀寧

大模型推理新突破：分布式推理技術探索與實踐 大模型推理新突破：分布式推理技術探索與實踐 極客邦科技InfoQ

大模型推理分離架構五虎上將 趙尚春

關于Deepseek采用EP推理方式的一些思考 楊鵬程

基于 chunked prefill 理解 prefill 和 decode 的計算特性 Chayenne Zhao

https://zhuanlan.zhihu.com/p/718715866

LLM PD 分離背后的架構問題 極客博哥

https://zhuanlan.zhihu.com/p/27836625742

Splitwise: Efficient generative llm inference using phase splitting

DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving

https://hao-ai-lab.github.io/blogs/distserve/

akaihaoshuai：LLM推理加速調研

Inference without interference: Disaggregate llm inference for mixed downstream workloads

Mooncake閱讀筆記：深入學習以Cache為中心的調度思想，譜寫LLM服務降本增效新篇章 方佳瑞

Mooncake：基于KVCache的LLM服務架構設計與分析 常華Andy

vLLM 最近有哪些更新? OLDPAN

大模型推理新突破：分布式推理技術探索與實踐 QCon [InfoQ]

Injecting Adrenaline into LLM Serving: Boosting Resource Utilization and Throughput via Attention Disaggregation Yunkai Liang, Zhangyu Chen, Pengfei Zuo, Zhi Zhou, Xu Chen, Zhou Yu