在上一篇 博客文章 中，我们展示了通过将键值 (KV) 缓存卸载到 CPU 内存来重复使用该缓存如何在基于 x86 的 NVIDIA H100 Tensor Core GPU 上将首次生成令牌 (Time To First Token, TTFT) 的速度提升高达 14 倍，在 NVIDIA GH200 Superchip 上可将其提升高达 28 倍。在本文中，我们介绍了可以进一步推动 TTFT 加速的 KV 缓存重复使用技术和最佳实践。
KV 缓存简介
许多任务（包括问答和代码生成）都迅速采用了 LLM 模型。为了生成响应，这些模型首先将用户的提示符转换为 tokens，然后将其转换为 dense vectors。随后会进行大量的 dot-product operations，以数学方式对 tokens 之间的关系建模，并构建对用户输入的上下文理解。生成这种上下文理解的计算成本会随着输入序列的长度而二次增加。
此资源密集型流程会生成键和值，这些键和值会被缓存，以避免在生成后续 token 时进行重新计算。重复使用 KV 缓存可减少生成额外 token 所需的计算负载和时间，实现更快、更高效的用户体验。
重复使用 KV 缓存时，必须仔细注意其在内存中的剩余时间、在内存已满时首先移除哪些组件，以及何时可重复用于新的传入提示。优化这些因素可能会带来 KV 缓存重复使用的性能增益。NVIDIA TensorRT-LLM 提供了三个关键功能，专门解决了这些问题。
早期 KV 缓存复用
传统的复用算法需要先完成整个 KV 缓存的计算，然后才能根据新用户的提示重复使用其任何部分。在企业聊天机器人等场景中，系统提示（添加到用户查询中的预定义指令）对引导 LLM 根据企业指南做出响应至关重要，这种方法可能效率低下。
当大量用户同时与聊天机器人交互时，每个用户都需要单独计算系统提示 KV 缓存。借助 TensorRT-LLM，我们可以重新使用实时生成的系统提示，使其在突发期间在所有用户之间共享，而不是为每个用户重新计算。这可以显著加快推理速度，最高可达 5 倍。
灵活调整 KV 缓存块大小
在重用实现中，只能分配整个缓存内存块以供重用。例如，如果缓存内存块大小为 64 个令牌，KV 缓存为 80 个令牌，则仅存储 64 个令牌以供重用，而其余 16 个令牌则需要重新计算。但是，如果将内存块大小减少到 16 个令牌，则可在五个内存块中存储全部 64 个令牌，从而消除重新计算的需求。
当输入序列较短时，这种效果最为明显。对于长输入序列，更大的块会更有好处。显然，您对 KV 缓存的控制越细致，就越能针对特定用例进行优化。
TensorRT-LLM 可对 KV 缓存内存块进行细粒度控制，使开发者能够将这些块分割成 64 到 2 个令牌之间的较小块。这可优化已分配内存的使用，提高复用率并改善 TTFT。在 NVIDIA H100 Tensor Core GPU 上运行 LLAMA70B 时，我们可以通过将 KV 缓存块的大小从 64 个令牌减少到 8 个令牌，在多用户环境中将 TTFT 速度提升高达 7%。
高效的 KV 缓存回收协议
将 KV 缓存分割成较小的块并移除未使用的块对于内存优化可能有效，但这会引入依赖项复杂性。当使用特定块生成响应并将结果存储为新块时，它会形成树状的依赖结构。
随着时间的推移，跟踪源块（分支）使用情况的计数器可能会随着依赖节点（叶子）的重复使用而过时。然后，移除源块需要移除所有依赖块，这将需要重新计算 KV 缓存以用于新用户提示，从而增加 TTFT。
为应对这一挑战，TensorRT-LLM 包含智能驱逐算法，该算法可以追踪源节点中的依赖节点，并首先驱逐依赖节点，即使这些节点具有较新的重复使用计数器。这可以确保更高效的内存管理，同时防止不必要的依赖块驱逐。
TensorRT-LLM KV 缓存重用入门
在推理过程中生成 KV 缓存需要大量的计算和内存资源。高效使用 KV 缓存对于改善模型响应、加速推理和提高系统吞吐量至关重要。TensorRT-LLM 为希望进一步优化 TTFT 响应时间以实现峰值性能的开发者提供了高级复用功能。
要开始使用 TensorRT-LLM KV 缓存，请参阅 我们的 GitHub 文档 。
