分块注意力

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

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

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

TGI 的分块注意力实现利用了 vLLM 项目 开发的自定义 cuda 内核。您可以在 项目的页面 中了解更多关于此技术的信息。

< > 在 GitHub 上更新