跳到內容

自動字首快取

PagedAttention 的核心思想是將每個請求的 KV 快取劃分為 KV 塊。每個塊包含固定數量標記的注意力鍵和值。PagedAttention 演算法允許這些塊儲存在非連續的物理記憶體中，從而透過按需分配記憶體來消除記憶體碎片。

為了自動快取 KV 快取，我們利用以下關鍵觀察：每個 KV 塊可以透過塊內的標記以及塊字首中的標記唯一標識。

                    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 快取可以透過標記“A gentle breeze stirred”唯一標識。第三個塊可以透過塊中的標記“laughed in the distance”，以及字首標記“A gentle breeze stirred the leaves as children”唯一標識。因此，我們可以建立以下一對一對映

hash(prefix tokens + block tokens) <--> KV Block

透過此對映，我們可以在 vLLM 的 KV 快取管理中增加另一層間接性。以前，vLLM 中的每個序列都維護著從其邏輯 KV 塊到物理塊的對映。為了實現 KV 塊的自動快取，我們將邏輯 KV 塊對映到它們的雜湊值，並維護一個包含所有物理塊的全域性雜湊表。透過這種方式，所有共享相同雜湊值的 KV 塊（例如，兩個請求之間共享的字首塊）都可以對映到相同的物理塊並共享記憶體空間。

此設計實現了自動字首快取，而無需在 KV 塊之間維護樹結構。更具體地說，所有塊都是相互獨立的，可以自行分配和釋放，這使得我們能夠像作業系統中的普通快取一樣管理 KV 快取。

通用快取策略

將所有 KV 塊儲存在雜湊表中使 vLLM 能夠快取早期請求的 KV 塊，從而節省記憶體並加速未來請求的計算。例如，如果一個新請求與前一個請求共享系統提示，則共享提示的 KV 快取可以直接用於新請求而無需重新計算。然而，KV 快取的總空間是有限的，當快取已滿時，我們必須決定保留或驅逐哪些 KV 塊。

透過雜湊表管理 KV 快取允許我們實現靈活的快取策略。例如，在當前的 vLLM 中，我們實現了以下驅逐策略

  • 當沒有空閒塊時，我們將驅逐引用計數（即，當前使用該塊的請求數量）等於 0 的 KV 塊。
  • 如果有多個引用計數等於 0 的塊，我們優先驅逐最近最少使用的塊（LRU）。
  • 如果有多個塊的最後訪問時間相同，我們優先驅逐位於最長字首末尾的塊（即，其前面有最大數量塊的塊）。

請注意，當應用於全注意力模型時，此驅逐策略有效實現了與 RadixAttention 中完全相同的策略，該策略優先驅逐字首樹中引用計數為零和最近最少使用的葉節點。

然而，基於雜湊的 KV 快取管理為我們提供了處理更復雜服務場景和實現超越上述策略的更復雜驅逐策略的靈活性

  • 多 LoRA 服務。當為多個 LoRA 介面卡提供服務時，我們可以簡單地讓每個 KV 塊的雜湊值也包含請求查詢的 LoRA ID，以啟用所有介面卡的快取。透過這種方式，我們可以聯合管理不同介面卡的 KV 塊，這簡化了系統實現並提高了全域性快取命中率和效率。
  • 多模態模型。當用戶輸入不僅僅包含離散標記時，我們可以使用不同的雜湊方法來處理不同模態輸入的快取。例如，對影像使用感知雜湊來快取相似的輸入影像。