5x Faster Time to First Token with NVIDIA TensorRT-LLM KV Cache Early Reuse
In our previous
blog post
, we demonstrated how reusing the key-value (KV) cache by offloading it to CPU memory can accelerate time to first token (TTFT) by up to 14x on x86-based NVIDIA H100 Tensor Core GPUs and 28x on the NVIDIA GH200 Superchip. In this post, we shed light on KV cache reuse techniques and best practices that can drive even further TTFT speedups.
Introduction to KV cache
LLM models are rapidly being adopted for many tasks, including question-answering, and code generation. To generate a response, these models begin by converting the user’s prompt into tokens, which are then transformed into dense vectors. Extensive dot-product operations follow to mathematically model the relationships between the tokens and build a contextual understanding of the user input. The computational cost of generating this contextual understanding increases quadratically with the length of the input sequence.
This resource-intensive process generates keys and values, which are cached to avoid recomputation when generating subsequent tokens. Reusing the KV cache reduces the computational load and time needed to generate additional tokens—leading to a faster and more efficient user experience.
When reusing the KV cache, careful attention must be given to how long it remains in memory, which components to evict first when memory is full, and when it can be reused for new incoming prompts. Optimizing these factors can lead to incremental performance improvements in KV cache reuse. NVIDIA TensorRT-LLM offers three key features that specifically address these areas.
Early KV cache reuse
Traditional reuse algorithms require the entire KV cache computation to be completed before any portions of it can be reused with new user prompts. In scenarios such as enterprise chatbots, where system prompts—predefined instructions added to user queries—are essential to direct the LLM’s responses in line with enterprise guidelines, this method can be inefficient.
When a surge of users interacts with the chatbot simultaneously, each user would require a separate computation of the system prompt KV cache. With TensorRT-LLM, we can instead reuse the system prompt as it is being generated in real time, enabling it to be shared across all users during the burst, rather than recalculating it for each user. This can significantly accelerate inference for use cases requiring system prompts by up to 5x.
Figure 1. TensorRT-LLM KV cache reuse can speed up TTFT by up to 5x
Flexible KV cache block sizing
In reuse implementations, only entire cache memory blocks can be allocated for reuse. For example, if the cache memory block size is 64 tokens and KV cache is 80 tokens, only 64 tokens will be stored for reuse, while the remaining 16 tokens will need to be recomputed. However, if the memory block size is reduced to 16 tokens, all 64 tokens can be stored across five memory blocks, eliminating the need for re-computation.
This effect is most pronounced when the input sequences are short. For long input sequences, larger blocks can be more beneficial.  As is clear, the more granular the control you have over the KV cache, the better you can optimize it for your specific use case.
TensorRT-LLM provides fine-grained control over KV cache memory blocks, giving developers the ability to chop them into smaller blocks between 64 to 2 tokens. This optimizes the usage of allocated memory, increases reuse rates, and improves TTFT. When running LLAMA70B on NVIDIA H100 Tensor Core GPUs, we can speed up TTFT up to 7% in multi-user environments by reducing KV cache block size from 64 tokens to 8 tokens.
Figure 2. Impact of changing KV cache block size on inference speedup
Efficient KV cache eviction protocols
Partitioning the KV cache into smaller blocks and evicting unused ones can be effective for memory optimization, but it introduces dependency complexities. When a specific block is used to generate a response, and the result is stored as a new block, it can form a tree-like structure of dependencies.
Over time, the counters tracking the usage of the source blocks (the branches) may become stale as the dependent nodes (the leaves) are reused. Evicting the source block then requires the eviction of all dependent blocks, which would require recalculation of the KV cache for new user prompts, increasing TTFT.
To address this challenge, TensorRT-LLM includes intelligent eviction algorithms that can trace the dependent nodes from their source nodes and evict dependent nodes first, even if they have more recent reuse counters. This ensures more efficient memory management while preventing unnecessary evictions of dependent blocks.
Figure 3. A logical representation of KV cache eviction algorithm show how it can reduce the number of evicted blocks, increasing the likelihood of reuse
Getting started with TensorRT-LLM KV cache reuse
Generating KV cache during inference requires a lot of compute and memory resources. Using it efficiently is critical to improving model response, accelerating inference, and increasing system throughput. TensorRT-LLM provides advanced reuse features for developers looking to further optimize TTFT response times for peak performance.
To start using TensorRT-LLM KV cache reuse check out our
GitHub documentation
.
NVIDIA TensorRT-LLM KV 缓存早期重用实现首个令牌速度 5 倍提升
在上一篇
博客文章
中，我们展示了通过将键值 (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 倍。
图 1.TensorRT-LLM KV 缓存重复利用可将 TTFT 速度提升高达 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%。
图 2. 更改 KV 缓存块大小对推理加速的影响
高效的 KV 缓存回收协议
将 KV 缓存分割成较小的块并移除未使用的块对于内存优化可能有效，但这会引入依赖项复杂性。当使用特定块生成响应并将结果存储为新块时，它会形成树状的依赖结构。
随着时间的推移，跟踪源块（分支）使用情况的计数器可能会随着依赖节点（叶子）的重复使用而过时。然后，移除源块需要移除所有依赖块，这将需要重新计算 KV 缓存以用于新用户提示，从而增加 TTFT。
为应对这一挑战，TensorRT-LLM 包含智能驱逐算法，该算法可以追踪源节点中的依赖节点，并首先驱逐依赖节点，即使这些节点具有较新的重复使用计数器。这可以确保更高效的内存管理，同时防止不必要的依赖块驱逐。
图 3.KV 缓存驱逐算法的逻辑表示展示了如何减少被驱逐的块的数量，从而增加重复使用的可能性。
TensorRT-LLM KV 缓存重用入门
在推理过程中生成 KV 缓存需要大量的计算和内存资源。高效使用 KV 缓存对于改善模型响应、加速推理和提高系统吞吐量至关重要。TensorRT-LLM 为希望进一步优化 TTFT 响应时间以实现峰值性能的开发者提供了高级复用功能。
要开始使用 TensorRT-LLM KV 缓存，请参阅
我们的 GitHub 文档
。