PagedAttention

大語言模型（LLM）在生成過程中會遇到記憶體限制的問題。在生成的解碼階段，所有先前詞元（token）生成的注意力鍵（key）和值（value）都儲存在 GPU 記憶體中以供重用。這被稱為 KV 快取，對於大型模型和長序列，它可能會佔用大量記憶體。

PagedAttention 嘗試透過將 KV 快取劃分為塊來最佳化記憶體使用，這些塊透過一個查詢表進行訪問。因此，KV 快取無需儲存在連續的記憶體中，並且塊是按需分配的。這種記憶體效率可以提高記憶體密集型工作負載的 GPU 利用率，從而支援更多的推理批次。

使用查詢表訪問記憶體塊也有助於在多個生成任務之間共享 KV。這對於像並行取樣這樣的技術很有幫助，其中為同一提示同時生成多個輸出。在這種情況下，快取的 KV 塊可以在這些生成任務之間共享。

TGI 的 PagedAttention 實現利用了 vLLM 專案 開發的自定義 CUDA 核心。您可以在該專案頁面上了解更多關於這項技術的資訊。

< > 在 GitHub 上更新