在 NVIDIA TensorRT-LLM 中引入新型 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)。

Image shows TensorRT-LLM priority-based eviction API in action allowing more fine-grained control over KV cache reuse. It is differentiated from the default LRU eviction policy where developers cannot control a block’s persistence for later reuse. — *圖 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）。