跳到內容

自動字首快取

字首快取 kv-cache 塊是 LLM 推理中一種流行的最佳化技術，可以避免重複的提示計算。核心思想很簡單——我們快取已處理請求的 kv-cache 塊，並在新請求帶有與先前請求相同的（或部分相同的）字首時重用這些塊。由於字首快取幾乎是免費的午餐，並且不會改變模型輸出，因此已被許多公共端點（例如 OpenAI、Anthropic 等）和大多數開源 LLM 推理框架（例如 SGLang）廣泛使用。

雖然實現字首快取的方法有很多，但 vLLM 選擇了一種基於雜湊的方法。具體來說，我們透過塊中的 token 以及塊之前的字首中的 token 來雜湊每個 kv-cache 塊。

                    Block 1                  Block 2                  Block 3
         [A gentle breeze stirred] [the leaves as children] [laughed in the distance]
Block 1: |<--- block tokens ---->|
Block 2: |<------- prefix ------>| |<--- block tokens --->|
Block 3: |<------------------ prefix -------------------->| |<--- block tokens ---->|

在上面的示例中，第一個塊中的 KV 快取可以透過 token “A gentle breeze stirred” 唯一標識。第三個塊可以透過塊中的 token “laughed in the distance” 以及字首 token “A gentle breeze stirred the leaves as children” 來唯一標識。因此，我們可以構建 hash(tuple[components]) 的塊雜湊，其中 components 是：

  • 父雜湊值：父雜湊塊的雜湊值。
  • 塊 token：此塊中 token 的元組。包含確切 token 的原因是為了減少潛在的雜湊值衝突。
  • 額外雜湊：使此塊唯一的其他值，例如 LoRA ID、多模態輸入雜湊（見下例）以及用於在多租戶環境中隔離快取的快取鹽。

注意 1

我們只快取完整的塊。

注意 2

上述雜湊鍵結構並非 100% 無衝突。理論上，不同的字首 token 仍然可能具有相同的雜湊值。為避免任何雜湊衝突，**在多租戶設定中，我們使用 SHA256** 作為雜湊函式，而不是內建雜湊。SHA256 自 vLLM v0.8.3 起支援，自 v0.10.2 起為預設值。它帶來的效能影響微乎其微，每 token 約 75ns（對於 50k token 的上下文，不到 4ms）。

多模態輸入雜湊示例
在此示例中，我們說明了字首快取如何處理多模態輸入（例如影像）。假設我們有一個包含以下訊息的請求：

messages = [
    {"role": "user",
     "content": [
         {"type": "text",
          "text": "What's in this image?"
         },
         {"type": "image_url",
          "image_url": {"url": image_url},
         },
    ]},
]

它將變成以下提示：

Prompt:
    <s>[INST]What's in this image?\n[IMG][/INST]

Tokenized prompt:
    [1, 3, 7493, 1681, 1294, 1593, 3937, 9551, 10, 4]

Prompt with placeholders (<P>):
    [1, 3, 7493, 1681, 1294, 1593, 3937, 9551, <P>, <P>, ..., <P>, 4]

正如我們所見，在 token 化後，[IMG] 將被一系列佔位符 token 替換，並且這些佔位符將在預填充期間被影像嵌入替換。字首快取支援這種情況的挑戰在於我們需要區分影像和佔位符。為了解決這個問題，我們對前端影像處理器生成的影像雜湊進行編碼。例如，上述提示中塊的雜湊將是（假設塊大小為 16，我們有 41 個佔位符 token）：

Block 0
    Parent hash: None
    Token IDs: 1, 3, 7493, 1681, 1294, 1593, 3937, 9551, <p>, ..., <p>
    Extra hash: <image hash>
Block 1
    Parent hash: Block 0 hash
    Token IDs: <p>, ..., <p>
    Extra hash: <image hash>
Block 2
    Parent hash: Block 1 hash
    Token IDs: <p>, ..., <p>
    Extra hash: <image hash>
Block 3
    Parent hash: Block 2 hash
    Token IDs: <p>, ..., <p>, 4
    Extra hash: <image hash>

在本文件的其餘部分，我們將首先介紹 vLLM v1 中用於字首快取的資料結構，然後介紹主要 KV 快取運算子（例如，分配、追加、釋放、淘汰）的字首快取工作流。最後，我們將使用一個示例來說明端到端的字首快取工作流。

安全快取隔離為提高共享環境中的隱私性，vLLM 支援透過可選的每個請求加鹽來隔離字首快取重用。透過在請求中包含 cache_salt，該值會被注入到第一個塊的雜湊中，確保只有具有相同鹽值的請求才能重用快取的 KV 塊。這可以防止基於時間的攻擊，攻擊者可以透過觀察延遲差異來推斷快取的內容。這提供了保護，同時又不犧牲效能。

{
  "messages": [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Here is a document with details about the world series: ..."},
    {"role": "user", "content": "Who won the world series in 2020?"}
  ],
  "cache_salt": "your-cache-salt"
}

透過此設定，快取共享僅限於明確同意共同鹽值的使用者或請求，從而在信任組內實現快取重用，同時隔離其他使用者。

資料結構

vLLM v1 中的字首快取實現在 KV 快取管理器中。基本構建塊是“Block”資料類（簡化版）：

class KVCacheBlock:
    # The block ID (immutable)
    block_id: int
    # The block hash (will be assigned when the block is full,
    # and will be reset when the block is evicted).
    block_hash: BlockHash
    # The number of requests using this block now.
    ref_cnt: int

    # The pointers to form a doubly linked list for the free queue.
    prev_free_block: "KVCacheBlock | None" = None
    next_free_block: "KVCacheBlock | None" = None

有兩個設計要點值得強調：

  1. 我們在初始化 KV 快取管理器時分配所有 KVCacheBlock，作為塊池。這避免了 Python 物件建立的開銷，並可以輕鬆地一直跟蹤所有塊。
  2. 我們將雙向連結串列指標直接引入 KVCacheBlock，這樣我們就可以直接構建一個空閒佇列。這給我們帶來了兩個好處：
    1. 我們可以在 O(1) 時間複雜度內將中間的元素移動到末尾。
    2. 我們可以避免引入另一個包裝了元素的 Python 佇列（例如 deque）。

因此，當 KV 快取管理器初始化時，我們將擁有以下元件：

Component Overview

  • 塊池：KVCacheBlock 的列表。
  • 空閒塊佇列：僅儲存頭尾塊的指標以進行操作。
  • 快取塊：從雜湊鍵到塊 ID 的對映。
  • 請求塊：從請求 ID 到已分配塊 ID 的對映。

操作

塊分配

新請求：排程器安排新請求並分配 KV 快取塊的工作流。

  1. 排程器呼叫 kv_cache_manager.get_computed_blocks() 來獲取已計算的塊序列。這透過雜湊請求中的提示 token 並查詢快取塊來完成。
  2. 排程器呼叫 kv_cache_manager.allocate_slots()。它執行以下步驟：
    1. 計算所需新塊的數量，如果塊不足則返回。
    2. “觸碰”已計算的塊。它將已計算塊的引用計數加一，如果塊未被其他請求使用，則將其從空閒佇列中移除。這是為了避免這些已計算的塊被淘汰。有關說明，請參閱下一節的示例。
    3. 透過彈出空閒佇列的頭部來分配新塊。如果頭部塊是快取塊，這還將“淘汰”該塊，以便其他請求從現在開始無法重用它。
    4. 如果分配的塊已裝滿 token，我們會立即將其新增到快取塊中，以便同一批中的其他請求可以重用它。

正在執行的請求：排程器安排正在執行的請求並分配 KV 快取塊的工作流。

  1. 排程器呼叫 kv_cache_manager.allocate_slots()。它執行以下步驟：
    1. 計算所需新塊的數量，如果塊不足則返回。
    2. 透過彈出空閒佇列的頭部來分配新塊。如果頭部塊是快取塊，這還將“淘汰”該塊，以便其他請求從現在開始無法重用它。
    3. 將 token ID 追加到現有塊以及新塊的槽中。如果一個塊已滿，我們將其新增到快取塊中進行快取。

重複塊
假設塊大小為 4，並且您傳送了一個帶有提示 ABCDEF 和解碼長度 3 的請求（請求 1）。

Prompt: [A, B, C, D, E, F]
Output: [G, H, I]

Time 0:
  Tokens: [A, B, C, D, E, F, G]
  Block Table: [0 (ABCD), 1 (EFG)]
  Cache Blocks: 0
Time 1:
  Tokens: [A, B, C, D, E, F, G, H]
  Block Table: [0 (ABCD), 1 (EFGH)]
  Cache Blocks: 0, 1
Time 2:
  Tokens: [A, B, C, D, E, F, G, H, I]
  Block Table: [0 (ABCD), 1 (EFGH), 2 (I)]
  Cache Blocks: 0, 1

現在塊 0 和塊 1 已被快取，並且我們再次傳送相同的請求（請求 2）並使用貪婪取樣，因此它將產生與請求 1 完全相同的輸出。

Prompt: [A, B, C, D, E, F]
Output: [G, H, I]

Time 0:
  Tokens: [A, B, C, D, E, F, G]
  Block Table: [0 (ABCD), 3 (EFG)]
  Cache Blocks: 0, 1
Time 1:
  Tokens: [A, B, C, D, E, F, G, H]
  Block Table: [0 (ABCD), 3 (EFGH)]
  Cache Blocks: 0, 1, 3

可以看到，塊 3 是一個新的完整塊並被快取。然而，它相對於塊 1 是冗餘的，這意味著我們快取了兩次相同的塊。在 v0 中，當檢測到塊 3 重複時，我們會釋放塊 3 並讓請求 2 使用塊 1，因此在時間 1 它的塊表變為 [0, 1]。然而，vLLM v1 中的塊表是追加式的，這意味著不允許將塊表從 [0, 3] 更改為 [0, 1]。因此，我們將為雜湊鍵 E-H 擁有重複的塊。當請求被釋放時，這種重複將被消除。

釋放

當請求完成時，如果沒有任何其他請求使用它們的塊（引用計數 = 0），我們將釋放所有塊。在此示例中，我們釋放請求 1 及其關聯的塊 2、3、4、8。我們可以看到釋放的塊以**反向**順序新增到空閒佇列的末尾。這是因為請求的最後一個塊必須雜湊更多的 token，並且不太可能被其他請求重用。因此，它應該首先被淘汰。

Free queue after a request us freed

淘汰 (LRU)

當空閒佇列的頭部塊（最近最少使用的塊）被快取時，我們必須淘汰該塊，以防止它被其他請求使用。具體來說，淘汰涉及以下步驟：

  1. 從空閒佇列頭部彈出塊。這是要淘汰的 LRU 塊。
  2. 從快取塊中刪除塊 ID。
  3. 移除塊雜湊。

示例

在此示例中，我們假設塊大小為 4（每個塊可以快取 4 個 token），並且 KV 快取管理器中總共有 10 個塊。

時間 1：快取為空，新請求進入。我們分配了 4 個塊。其中 3 個已滿並被快取。第四個塊部分填充，有 3 個 token（滿 4 個）。

Example Time 1

時間 2：請求 0 使塊 3 變滿，並要求一個新塊以繼續解碼。我們快取塊 3 並分配塊 4。

Example Time 2

時間 3：請求 1 進入，帶有 14 個提示 token，其中前 10 個 token 與請求 0 相同。我們可以看到只有前 2 個塊（8 個 token）命中快取，因為第三個塊只匹配了 4 個 token 中的 2 個。

Example Time 3

時間 4：請求 0 完成並釋放。塊 2、3 和 4 以反向順序新增到空閒佇列（但塊 2 和 3 仍被快取）。塊 0 和 1 未新增到空閒佇列，因為它們正在被請求 1 使用。

Example Time 4

時間 5：請求 1 完成並釋放。

Example Time 5

時間 6：請求 2 進入，帶有 29 個提示 token，其中前 12 個 token 與請求 0 相同。請注意，即使空閒佇列中的塊順序是 7 - 8 - 9 - 4 - 3 - 2 - 6 - 5 - 1 - 0，快取命中的塊（即 0、1、2）在分配前被觸碰並從佇列中移除，因此空閒佇列變為 7 - 8 - 9 - 4 - 3 - 6 - 5。結果，分配的塊是 0（快取）、1（快取）、2（快取）、7、8、9、4、3（淘汰）。

Example Time 6