Systems and methods for managing multi-tier key-value caching for generative ai models

Abstract

Systems and methods for managing multi-tier key-value (KV) caching for a generative artificial intelligence (AI) model are described. A system includes a lower tier memory, a higher tier memory, and a KV cache manager to: (1) determine whether a key-value (KV) cache corresponding to a user associated with a session exists in the lower tier memory, and (2) if the KV cache, or a portion thereof, does not exist in the lower tier memory, upon determining that at least a portion of the KV cache exists in the higher tier memory, load at least the portion of the KV cache into the lower tier memory. The system includes a KV cache engine to execute: (1) an extended prefill phase to generate keys and values for extending the KV cache in the lower tier memory, and (2) a decode phase to generate output tokens until a stopping criteria is met.