跳到內容

混合 KV 快取管理器

警告

本文件基於提交 458e74 編寫。此功能仍處於早期階段，可能會發生變化。

什麼是混合模型？

許多最近的“混合”LLM 在一個模型中結合了多種注意力型別。例如：

  1. 滑動視窗注意力 (sw) + 全注意力 (full)：gpt-oss, Gemma 2/3, Ministral, cohere 等。
  2. Mamba + full：Bamba, Jamba, Minimax 等。
  3. 區域性塊注意力 + full：Llama4

為了高效地服務這些模型，我們的 KVCacheManager 必須

  1. 為不同層型別分配不同的槽位，例如：
    • 全注意力層：為**所有** token 保留槽位。
    • 滑動視窗層：僅為最近的 **sliding_window_size** 個 token 保留槽位。
  2. 支援特定於層的 the prefix-cache 規則，例如：
    • 全注意力：快取命中的字首要求**所有** token 保留在 KV 快取中。
    • 滑動視窗：快取命中的字首僅要求最後 **sliding_window_size** 個 token 保留在 KV 快取中。

定義

  1. kv 隱藏大小：儲存單個 token 的 KV 快取對單個層所需的位元組數。
  2. 塊：為 kv 快取保留的記憶體被劃分為多個具有相同*頁大小*（定義如下）的*塊*。
  3. 塊大小：塊內的 token 數量。

  4. 頁大小：塊的物理記憶體大小，定義為

    \[ \text{num_layers} \times \text{block_size} \times \text{kv_hidden_size} \]

    num_layers 並不意味著模型中的總層數。確切數字取決於本文件的上下文。

    注意

    這與程式碼中的 KVCacheSpec.page_size_bytes 不同，後者定義為

    \[ \text{block_size} \times \text{kv_hidden_size} \]

分配

高層思路

我們為所有層型別使用單個記憶體池。記憶體池被劃分為多個具有相同頁大小的塊。 KVCacheManager 根據其注意力型別為不同層分配不同數量的塊。

核心挑戰是確保每種層型別使用相同的**頁大小**。對於僅全注意力模型，頁大小的定義很直接：

\[ \text{page_size} = \text{block_size} \times \text{num_hidden_layers} \times \text{kv_hidden_size} \]

然而，在混合模型中，num_hidden_layers 因注意力型別而異，這通常會產生不匹配的頁大小。下面的案例展示了我們如何統一它們。

案例 1：玩具模型

讓我們從一個玩具示例開始：一個模型有 1 個全注意力層和 3 個滑動視窗注意力層。所有層具有相同的 kv_hidden_size。

我們讓每個塊持有 block_size 個 token，用於一層，因此：

\[ \text{page_size} = \text{kv_hidden_size} \times \text{block_size} \]

KVCacheManager 為每個層分配不同數量的塊。

此案例僅為玩具示例。對於真實模型，請參閱以下案例。

案例 2：相同的 kv_hidden_size 和規則模式

當模型有更多層時，例如，20 個滑動視窗注意力層和 10 個全注意力層，且具有相同的 kv_hidden_size。為每層呼叫一次分配器（30 次呼叫）是可以的，但效率低下。作為一種解決方案，我們將需要相同數量塊的層分組，以減少呼叫次數。

分組是可行的，因為不同型別層之間的數量通常存在漂亮的比例。例如：

  • Gemma-2：1 sw : 1 full
  • Llama 4：3 local : 1 full

我們的示例可以看作是 2 sw : 1 full。我們可以像模型中有 2 個 sw 和 1 個 full 一樣分配塊，然後將結果重複 10 次來生成 30 個層的 block_ids。頁大小變為：

\[ 10 \times \text{kv_hidden_size} \times \text{block_size} \]

假設 block_size 為 16，滑動視窗大小為 32，請求長度為 112，那麼對於上述示例模型，我們需要分配 11 個塊（0-6 用於 full，7-8 用於 sw 組 1，9-10 用於 sw 組 2）。

Allocation Result

此處，“/”表示不需要塊（滑動視窗層不需要早期 token 的槽位）。

見下文正式定義。層被劃分為多個*KV 快取組*，以便：

  1. 每個組內具有相同的注意力型別：每個組僅包含具有相同注意力型別的層，因此對於給定的請求，它們需要相同數量的塊。這使得同一組中的層可以共享相同的塊 id，而不會造成記憶體浪費。
  2. 組之間的頁大小相同：因為我們的記憶體池只有一個頁大小。

我們的示例模型被劃分為 3 個 KV 快取組：

  • 組 0：10 個全注意力層（full.0 - full.9）
  • 組 1：10 個滑動視窗注意力層（sw.0 - sw.9）
  • 組 2：10 個滑動視窗注意力層（sw.10 - sw.19）

顯然，它滿足規則 1。對於規則 2，所有 3 個組都有

\[ 10 \times \text{kv_hidden_size} \times \text{block_size} \]

作為它們的頁大小。

案例 3：相同的 kv_hidden_size 和無規則模式

不幸的是，並非所有模型都具有如此優美的比例，並且案例 2 中的方法會產生太多的小組。例如，Gemma-3-27b 有 52 個滑動視窗注意力層和 10 個全注意力層。在案例 2 的約束下，它將是 26 個滑動視窗組和 5 個全注意力組，每個組包含 2 層。分配仍然效率低下。為了減少 KV 快取組的數量，我們使用所有注意力型別中最小的層數來分組層。例如，Gemma-3-27b 中每組為 min(52, 10) = 10 個層。然後分組結果為：

  • 組 0：10 個全注意力層（full.0 - full.9）
  • 組 1：10 個滑動視窗注意力層（sw.0 - sw.9）
  • 組 2：10 個滑動視窗注意力層（sw.10 - sw.19）
  • ...
  • 組 6：10 個滑動視窗注意力層（sw.40 - sw.49）
  • 組 7：2 個滑動視窗注意力層（sw.50 - sw.51）和 8 個填充層

如果這種啟發式方法在出現新模型時導致了糟糕的結果（例如，20 個 full + 30 個 sw，組大小應為 10 而不是 20），我們將更新此演算法。

這種情況發生在 Gemma-3 系列模型以及案例 2 中的模型，但帶有引入一個全注意力層的 eagle 投機解碼。該解決方案存在一些記憶體浪費，並不完美。請報告任何填充開銷無法接受的情況，以便我們改進演算法。

案例 4：不同的 kv_hidden_size（主要是混合 Mamba 模型）

一些架構（例如，Bamba、Jamba、Minimax）將標準注意力層與 Mamba 層交織在一起，其中每個 Mamba 層的每個 token 的狀態大小可能遠大於注意力層的 kv_hidden_size。因為我們只支援跨所有組的單個頁大小，所以我們必須協調這些不同的隱藏大小。

當前的演算法是：

  1. 增加註意力層的 block_size，直到 $$ \text{block_size} \times \text{kv_hidden_size}{\text{att}} \ge \text{state_size} $$}
  2. 填充每個層的 mamba 狀態為 $$ \text{block_size} \times \text{kv_hidden_size}_{\text{att}} $$
  3. 應用案例 3 中的分組策略。

注意

這可能導致注意力層的 block_size 超過 400，這太大了。另一種填充策略是增加 block_size，直到

\[ \text{block_size} \times \text{kv_hidden_size}_{\text{att}} \times \text{num_attn_layers} \ge \text{state_size}_{\text{mamba}} \]

此填充策略仍在進行中。

案例 5：KV 共享

KV 共享是指一個層使用另一個層的 KV 快取，例如 gemma-3n。在這些模型中，KVCacheManager 忽略所有具有 KV 共享的層，僅為需要 KV 快取的層分配 KV 快取，並在模型執行器中進行一些修補以將分配結果應用於 KV 共享層。

字首快取

為簡單起見，我們在本節中假設 block_size=1。

高層思路

塊池使用類似於 tuple(block_hash, group_id) -> block 的字典來捕獲完整的塊。這意味著不同組的相同 token 會獨立快取和逐出。

當有新請求進入時，我們檢查每個組的快取命中字首，並將這些組的交集作為請求的快取字首返回。有關檢查單個組的快取命中和執行交集的詳細演算法，請參閱下文。

案例 0：僅全注意力模型

對於全注意力層，為請求中的所有 token 分配塊。有關底層設計的詳細資訊，請參閱 字首快取。

要找到請求的最長快取命中字首，我們從左（第一個塊）到右（最後一個塊）列舉，檢查塊是否被快取，並在快取未命中時退出。例如，在下面的示例中，我們將返回前 7 個 token（0-6）作為快取命中字首（藍色塊是已快取的）。

Prefix Caching of Full Attention

案例 1：僅滑動視窗注意力模型

對於滑動視窗注意力層，記憶體分配的樸素實現是分配 sliding_window_size 個塊，並以迴圈方式填充塊。但這種樸素實現與字首快取不相容，因此我們沒有選擇這種設計。在 vLLM 中，我們為不同的 token 分配不同的塊，並釋放超出滑動視窗的塊。

對於新請求，快取命中字首僅需要快取最後 sliding_window_size - 1 個 token。假設 sliding_window_size = 4 且 block_size = 1，請求是一個 15 個 token 的提示（藍色塊是已快取的）。

Prefix Caching of Sliding Window Attention

有 3 種可能的快取命中字首：

  • 快取命中長度 5，使用 [2, 3, 4] 進行預填充計算 → [5, 6, …, 14]
  • 快取命中長度 6，使用 [3, 4, 5] 進行預填充計算 → [6, 7, …, 14]
  • 快取命中長度 14，使用 [11, 12, 13] 進行預填充計算 → [14]（最高效）

我們可以從右到左檢查快取命中，並在找到匹配時提前退出。這與全注意力模型相反，後者從左到右檢查並在匹配失敗時提前退出。一個潛在的缺點（與全注意力相比）是，當沒有匹配時，我們最終會遍歷整個 token 列表，而這通常是常見情況。這可能會導致不可忽略的開銷，但與 full + swa 結合使用效果很好，如下所述。

案例 2：滑動視窗注意力 + 全注意力模型

第一個問題是如何找到快取命中字首。我們需要透過以下方式“相交”全域性和滑動視窗注意力層的快取命中：

  1. 獲取全注意力模型的最長快取命中（從左到右掃描）。
  2. 獲取滑動視窗注意力模型的最長快取命中，該命中長度小於上述長度。透過從右到左開始檢查快取命中來實現，從全注意力模型的快取命中長度開始。

可以確保滑動視窗注意力層的最終快取命中也是全注意力層的快取命中。這比查詢每個組的所有可能字首並執行交集更有效，因為我們的方法可以在沒有快取命中時提前退出。

該演算法適用於具有完全兩種注意力型別的模型：全注意力 + X，其中 X 可以是任意高效的注意力演算法，如滑動視窗、Llama 4 區域性注意力、Mamba。它不支援沒有全注意力層的模型，也不支援超過 2 種注意力型別的模型。這足以滿足本文件編寫時的大多數混合模型。

第二個問題是快取逐出策略。目前，我們為所有 KV 快取組使用一個 LRU 佇列。當塊被釋放時（因為請求完成或塊超出滑動視窗），它們會被新增到 LRU 佇列。

案例 3：Mamba 模型

Mamba 模型的字首快取支援仍在開發中。一旦實現，具有 Mamba 層 + 全注意力層的模型可以透過案例 2 中的 full attention + X 演算法來支援。

實現

概述

Overview of Hybrid KV Cache Manager

KVCacheManager 被組織成 3 層：

  • KVCacheManager：排程程式和 KV 快取管理系統之間的介面。
  • KVCacheCoordinator：協調每個組的 SingleTypeKVCacheManagers 來生成請求的分配結果。根據模型的配置，會選擇以下協調器之一：
    • KVCacheCoordinatorNoPrefixCache：在停用字首快取時使用。
    • UnitaryKVCacheCoordinator：如果只有一個 KV 快取組。字首快取邏輯被簡化，無需交集。
    • HybridKVCacheCoordinator：處理恰好兩個 KV 快取組（必須包含一個全注意力組加上另一個高效注意力組）。其他情況尚未實現。您可以停用字首快取以使用 KVCacheCoordinatorNoPrefixCache。
  • SingleTypeKVCacheManager：每個例項管理一個 KV 快取組的分配和字首快取，實現注意力型別特定的邏輯（例如，全注意力、滑動視窗、Mamba）。

上圖中的藍色框顯示了 10 個全注意力層和 20 個滑動視窗注意力層的案例，因此：

  • 使用 HybridKVCacheCoordinator。
  • 為 3 個 KVCacheGroups 使用 1 個 FullAttentionManager 和 2 個 SlidingWindowManager。

記憶體佈局

對於一個具有 n 個 KVCacheGroups 的模型，每個組有 m 層，我們分配 m 個緩衝區。每個緩衝區由 n 層共享，每層來自一個組。

下圖顯示了一個具有 10 個全注意力層（full.0 - full.9）和 20 個滑動視窗注意力層（sw.0-sw.19）的模型。它遵循“分配”部分中的“案例 2”，並被劃分為 3 個組。

  • 組 0：10 個全注意力層（full.0 - full.9）
  • 組 1：10 個滑動視窗注意力層（sw.0 - sw.9）
  • 組 2：10 個滑動視窗注意力層（sw.10 - sw.19）

對於一個請求，我們分配 11 個塊，其中 block_id 0-6 分配給組 0，7-8 分配給組 1，9-10 分配給組 2。

在這樣一個示例中，物理記憶體被劃分為 10 個緩衝區（KVCacheTensor 0 - KVCacheTensor 9）。每個緩衝區由 3 層共享（例如，KVCacheTensor 0 由組 0 的 full.0、組 1 的 sw.0 和組 2 的 sw.10 共享），並被劃分為大小為 block_size * kv_hidden_size 的塊。這些 3 個注意力層的 KV 快取根據分配的 block_ids 被儲存在緩衝區的不同塊中。

Example Memory Layout

注意

一個邏輯“塊”對映到物理記憶體的 10 個緩衝區中的 10 個塊。