语言模型通过预测下一个令牌 (给定所有先前的令牌,包括输入文本令牌) 来生成文本。在 LLM 服务中,先前令牌的键和值元素用作历史语境,用于生成下一组令牌。从先前的 token 中缓存这些关键和值元素可以避免昂贵的重新计算,并有效地提高吞吐量。但是,键值 (KV) 缓存会随着语言模型的大小、批处理请求的数量和序列上下文长度呈线性增长,从而导致内存需求不断增长。
NVIDIA TensorRT-LLM 可提供多种 KV 缓存优化,以在显存大小增长与避免昂贵的重新计算之间实现具有挑战性的平衡。TensorRT-LLM 是一个开源库,可为 NVIDIA GPUs 上的众多热门大语言模型 ( LLMs ) 提供先进的推理支持。TensorRT-LLM KV 缓存包括多项优化,例如支持分页 KV 缓存、量化 KV 缓存、循环缓冲区 KV 缓存和 KV 缓存重复使用 。
在本文中,我们将深入探讨已引入 TensorRT-LLM 的两个新的高级功能。这些功能可实现对 KV 缓存的更精细控制,并提供 TensorRT-LLM KV 缓存的可见性,以便在 KV 缓存感知路由等上游应用中使用。
基于优先级的 KV 缓存移除
完成 LLM 请求后,系统会存储与这些请求关联的 KV 缓存块。鉴于 KV 缓存的有限大小,可能需要移除一些缓存块,以便为新序列腾出空间。默认情况下,驱逐遵循最近使用最少(LRU)的策略。
基于优先级的驱逐是 TensorRT-LLM Executor API 的一项新功能,使用户能够影响选择驱逐块的方式。用户可以指定两个用于指导块驱逐的属性:优先级和持续时间。优先级值用于设置相对保留优先级(在缓存中保留该块的重要性),持续时间值用于设置此优先级级别应应用的时长。
struct TokenRangeRetentionConfig {
# The beginning of this range
start: int
# The end of the range. Set to null to extend to end of sequence
end: optional<int>
# The priority level assigned to the range. 0->100
priority: int
# The duration this priority should apply for
duration: optional<int>
}
# Optional parameter to executor request
struct KvCacheRetentionConfig {
# List of priority assignments in context
ranges: list<TokenRangeRetentionConfig>
# Priority assigned to decode tokens
decode_priority: optional<int>
# Duration the decode priority applies for
decode_duration: optional<int>
}
借助 Priority Based – Eviction API,LLM 部署者能够利用有关其工作负载的知识,通过保留可能被重复使用的块来改善重复使用机会。例如,部署器可能希望与系统提示对应的块尽可能长时间地保留在缓存中,或者延迟关键型请求中可能涉及的块应具有高于其他块的优先级 (图 1)。

对于每个请求,您可以在输入上下文中为标记的离散范围指定优先级和持续时间值,并为解码阶段分配的块指定优先级和持续时间。标记范围的优先级应用,直到在未重复使用一段时间后持续时间结束,或直到与这些范围对应的块被移除。
选择要移除的块时,TensorRT-LLM 会考虑块内令牌的优先级。例如,带有 500- token 系统提示的请求可以将 token 范围 [0, 500) 设置为最高优先级。这样,与这些 token 对应的缓存块将仅在绝对必要时被移除。或者,如果您知道块永远不会被重复使用,您可以将此请求的块设置为最低优先级,以确保在其他块之前先将它们逐出。
这种新实现还会偏置离根更远的块,即使未设置优先级别,也会导致性能略有提升。我们的内部基准测试表明,基于优先级的驱逐可将缓存命中率提高约 20%,并且会因工作负载而异。
# Priority-based eviction usage examples
#Example 1: One-off request
KvCacheRetentionConfig(
[TokenRangeRetentionConfig(start=0, end=null, priority=0)],
decode_priority=0
)
#Example 2: High Priority system prompt
KvCacheRetentionConfig(
[TokenRangeRetentionConfig(start=0, end=1000, priority=100)]
)
#Example 3: Retain context blocks for 30 seconds, and decode blocks for 10 seconds
KvCacheRetentionConfig(
[TokenRangeRetentionConfig(start=0, end=null, priority=100, duration=30s)],
decode_priority=100, decode_duration=10s)
KV 缓存事件 API
在由 LLM 提供支持的大规模应用中,部署者通常会调配模型的多个服务实例来分发传入请求。这就提出了一个问题,哪个实例应该处理新请求?请求通常被路由到平衡负载,以确保高效利用和快速处理任何请求。任何实例上的 KV 缓存大小均代表增长和接受新工作的能力。
但是,基于负载的路由可能并不理想。如果中等加载的实例已经计算并缓存了新请求的键和值,可能仍首选将请求路由到此实例,以优化缓存重用。借助 KV 缓存事件 API,请求路由系统能够跟踪哪些实例已缓存或已移除块,从而实现更智能的重复使用和更高的性能。
TensorRT-LLM 执行程序 API 现在提供了一种跟踪 KV 缓存更新的方法。
struct KVCacheEvent {
event_id: long // Auto-incrementing event id
data: variant<CreatedData, StoredData, RemovedData, UpdatedData>
}
struct StoredBlockData {
blockHash: id // Unique identifier for the block.
tokens: list<Token>
loraId: id
cacheLevel: int // The cache level of the block (0 or 1, primary or secondary)
priority: int // The priority level of this block
}
struct StoredData {
parentHash: optional<id> // The parent of the sequence of blocks that was stored.
blocks: list<StoredBlockData> // The list of stored blocks
}
struct RemovedData {
blockHashes: list<id> // The hashes of blocks that were removed
}
# Set the max size of the internal event buffer. Defaults to 0 (no events)
kv_cache_config = KvCacheConfig(event_buffer_max_size=16384)
executor_config = ExecutorConfig(kv_cache_config)
executor = Executor(executor_config)
# Get an event manager
eventManager = executor.getKvCacheEventManager()
# Wait for new events. Once it returns, it implicitly clears the internal queue of events. Optionally provide a timeout value. If there's no events within this timeout, it returns an empty list.
events = eventManager.getLatestEvents()
存储缓存块以供重复使用、删除或更新时,系统会触发一个事件。应用可以实时使用这些事件,以最终获得 TensorRT-LLM KV 缓存当前状态的一致视图。这对于跟踪 KV 缓存重用机会特别有用。它可以在单个执行程序的规模上使用,以预测哪些请求将得到更多重复使用,也可以在多个执行程序中进行聚合,以做出 KV 感知型路由和调度决策(图 2)。

通过引入用于 KV 缓存重复使用和管理的基于优先级的驱逐和事件感知路由,TensorRT-LLM 为您提供了实现 KV 缓存精细控制的杠杆,以便您可以利用工作负载方面的知识来优化 KV 缓存管理。
总结
NVIDIA TensorRT-LLM 提供多项优化,可随时随地在 NVIDIA 加速的基础设施 (包括云、数据中心和工作站) 中高效部署生成式 AI 应用。这些优化可在同一硬件上实现显著加速和更好的缓存重复利用。这最终能够减少资源来处理相同的工作负载,降低能源成本并提高总拥有成本。