PagedAttention

大语言模型（LLM）在生成过程中会遇到内存限制的问题。在生成的解码阶段，所有先前词元（token）生成的注意力键（key）和值（value）都存储在 GPU 内存中以供重用。这被称为 KV 缓存，对于大型模型和长序列，它可能会占用大量内存。

PagedAttention 尝试通过将 KV 缓存划分为块来优化内存使用，这些块通过一个查找表进行访问。因此，KV 缓存无需存储在连续的内存中，并且块是按需分配的。这种内存效率可以提高内存密集型工作负载的 GPU 利用率，从而支持更多的推理批次。

使用查找表访问内存块也有助于在多个生成任务之间共享 KV。这对于像并行采样这样的技术很有帮助，其中为同一提示同时生成多个输出。在这种情况下，缓存的 KV 块可以在这些生成任务之间共享。

TGI 的 PagedAttention 实现利用了 vLLM 项目 开发的自定义 CUDA 内核。您可以在该项目页面上了解更多关于这项技术的信息。

< > 在 GitHub 上更新