Method for managing kv cache in transformer model based on reinforcement learning, and apparatus therefor

Abstract

The disclosure is intended to efficiently allocate the KV cache by predicting the length of an output token sequence (or the number of tokens) of the transformer model according to an input prompt (or input token sequence) through a neural network model, thereby efficiently utilizing the limited memory of a processor such as a GPU or the like. According to the disclosure, there is provided a method, performed in a computing device, for managing a KV cache for operation of a transformer model, and the method may include: training a KV cache manager comprising a neural network, based on a plurality of input prompts of the transformer model and an output sequence of the transformer model for each input prompt; predicting, by the trained KV cache manager, based on a first input prompt of the transformer model, the number of tokens to be included in a first output sequence of the transformer model for the first input prompt; and determining a size of a KV cache for the first input prompt, based on the predicted number of tokens.

Dependencies are inferred from the claim text.