OasisKV：以前瞻稀疏預取將解碼過程中的KV快取擴展至HBM之外

摘要

大型語言模型（LLM）的推論服務日益受到記憶體，而非運算能力的限制。隨著長上下文與長篇推理工作負載日益普遍，鍵值（KV）快取在 LLM 令牌生成（即解碼）期間主導了記憶體佔用與記憶體流量。尤其，HBM 容量已成為稀缺且昂貴的資源，嚴重限制了推論批次大小與系統吞吐量。本文提出 OasisKV，一種以記憶體為中心的 LLM 推論系統設計，透過在 LLM 解碼期間將完整 KV 快取儲存與 HBM 解耦，來緩解 HBM 容量壓力。由於解碼時的注意力天然具有稀疏性，OasisKV 僅將最相關令牌的 KV 條目保留在 HBM 中以供注意力計算使用。我們觀察到，使用推測解碼（SD）所草擬的前瞻令牌，可以提前準確預測未來重要的令牌。OasisKV 採用高效的注意力背景管線來識別重要的 KV 區塊。這些區塊隨後從容量較高的記憶體層級（例如主機或遠端記憶體）預取，並在下一步解碼使用之前暫存於 HBM 中。我們基於 vLLM 實作 OasisKV。前瞻預測的準確度足以在 2,048 個令牌的 KV 預算下，將準確度維持在與完整注意力相差 0.7 個百分點以內。這使得 OasisKV 能將稀疏性轉化為吞吐量提升：在推理工作負載上，相較於密集 vLLM 可達 1.69 倍，且僅損失 0.1 個百分點的準確度；在多 GPU 長上下文服務上，則可達 2.1 倍。在預填充-解碼分離的架構下，OasisKV 可達到約 2 倍的密集吞吐量，同時每個請求的 KV 用量減少 6.5 至 9.7 倍，且相較於完整 KV 傳輸，解碼節點主機記憶體占用減少 2.2 至 2.6 倍。

English

Large language model (LLM) inference serving is increasingly constrained by memory rather than compute. As long-context and long-form reasoning workloads become more prevalent, the key-value (KV) cache dominates both memory footprint and memory traffic during LLM token generation, i.e., decode. In particular, HBM capacity has become a scarce and costly resource that heavily limits inference batch size and system throughput. This paper presents OasisKV, a memory-centric LLM inference system design that alleviates HBM capacity pressure by decoupling full KV-cache storage from HBM during LLM decoding. Because decode-time attention is naturally sparse, OasisKV keeps only the KV entries of the most relevant tokens in HBMs for attention computation. We observe that future important tokens can be predicted accurately in advance using lookahead tokens drafted by speculative decoding (SD). OasisKV employs an efficient attention background pipeline to identify important KV blocks. They are then prefetched from higher-capacity memory tiers (e.g., host or remote memory) and staged in HBMs before being used in the next decode step. We implement OasisKV based on vLLM. The lookahead prediction is accurate enough to keep accuracy within 0.7 points of full attention under a 2,048-token KV budget. This lets OasisKV turn sparsity into throughput gain: 1.69times over dense vLLM on the reasoning workload at 0.1 points of accuracy loss, and up to 2.1times on multi-GPU long-context serving. Under prefill--decode disaggregation, OasisKV reaches about 2times dense throughput while admitting each request with 6.5--9.7times less KV and holding 2.2-2.6 less decode-node host memory than full KV transfer.