在NVIDIA TensorRT-LLM中引入新型KV缓存重用优化策略

创作时间:

作者:

@小白创作中心

在NVIDIA TensorRT-LLM中引入新型KV缓存重用优化策略

引用

nvidia

https://developer.nvidia.com/zh-cn/blog/introducing-new-kv-cache-reuse-optimizations-in-nvidia-tensorrt-llm/

NVIDIA TensorRT-LLM最近引入了两个新的KV缓存优化功能：基于优先级的KV缓存移除和KV缓存事件API。这些优化可以更好地控制KV缓存的使用，提高缓存命中率，并优化资源利用。

语言模型通过预测下一个令牌（给定所有先前的令牌，包括输入文本令牌）来生成文本。在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）。

图1、借助NVIDIA TensorRT-LLM中基于优先级的驱逐API，利用对工作负载的了解更好地控制KV缓存重用机会

对于每个请求，您可以在输入上下文中为标记的离散范围指定优先级和持续时间值，并为解码阶段分配的块指定优先级和持续时间。标记范围的优先级应用，直到在未重复使用一段时间后持续时间结束，或直到与这些范围对应的块被移除。

选择要移除的块时，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）。

图2、在NVIDIA TensorRT-LLM中使用KV缓存事件API，通过事件驱动型KV感知请求路由优化KV缓存重用机会