分页注意力

LLM 在生成过程中会受到内存限制的困扰。在生成的解码部分，为先前令牌生成的所有注意力键和值都存储在 GPU 内存中以供重用。这被称为 KV 缓存，对于大型模型和长序列，它可能会占用大量内存。

PagedAttention 尝试通过将 KV 缓存分区为通过查找表访问的块来优化内存使用。因此，KV 缓存不需要存储在连续内存中，并且块是根据需要分配的。内存效率可以提高内存受限工作负载上的 GPU 利用率，从而可以支持更多的推理批次。

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

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

< > 在 GitHub 上更新