PagedAttention

大型語言模型（LLM）在生成過程中常受限於記憶體瓶頸。在生成的解碼階段，先前 Token 所產生的所有注意力鍵（Keys）與值（Values）都會被儲存在 GPU 記憶體中以便重複使用。這被稱為 KV 快取（KV cache），對於大型模型與長序列而言，它可能會佔用大量的記憶體。

PagedAttention 試圖透過將 KV 快取劃分為區塊，並透過查詢表（lookup table）來存取，進而優化記憶體的使用。如此一來，KV 快取就不需要存放在連續的記憶體空間中，而是按需求分配區塊。這種記憶體效率的提升，能夠增強記憶體密集型工作負載下的 GPU 利用率，從而支援更多的推論批次。

利用查詢表來存取記憶體區塊，也有助於在多次生成之間共享 KV。這對於諸如 平行採樣（parallel sampling） 等技術非常有幫助，因為該技術會針對同一個提示詞同時生成多個輸出。在這種情況下，快取中的 KV 區塊可以在這些生成過程間共享。

TGI 的 PagedAttention 實作運用了由 vLLM 專案 所開發的自定義 CUDA 核心。您可以透過 該專案頁面 進一步了解這項技術。

在 GitHub 上更新