Systems and methods for key-value (kv) cache pruning

A system for managing cache usage on a graphic processing unit (GPU) for a Transformer-based neural network model, the system comprising:

• a memory that stores a Transformer-based neural network model and a plurality of processor-executable instructions;
• a communication interface that receives an input sequence of tokens;
• and
• one or more hardware processors that read and execute the plurality of processor-executable instructions from the memory to perform operations comprising:
• transforming, by the Transformer-based neural network model, the input sequence of tokens into intermediate variables including at least a key matrix, a value matrix, and a query matrix;
• allocating, in one or more processor memories, a key cache for storing the key matrix and a value cache for storing the value matrix;
• generating a plurality of scores indicating magnitudes associated with a plurality of rows of the key matrix stored at the key cache, respectively;
• re-allocating at least a portion of the key cache by removing at least one or more rows of the key matrix having associated scores below a score threshold, each row of the key matrix corresponding to a channel of the key cache;
• and
• operating the Transformer-based neural network model on the one or more processors with a reduced key cache after the re-allocating.

A non-transitory machine-readable medium comprising a plurality of machine-executable instructions which, when executed by one or more processors, are adapted to cause the one or more processors to perform operations comprising:

• transforming, by a Transformer-based neural network model, an input sequence of tokens into intermediate variables including at least a key matrix, a value matrix, and a query matrix;
• allocating, in one or more processor memories, a key cache for storing the key matrix and a value cache for storing the value matrix;
• generating a plurality of scores indicating magnitudes associated with a plurality of rows of the key matrix stored at the key cache, respectively;
• re-allocating at least a portion of the key cache by removing at least one or more rows of the key matrix having associated scores below a score threshold, each row of the key matrix corresponding to a channel of the key cache;
• and
• operating the Transformer-based neural network model with a reduced key cache after the re-allocating.