Dynamic quantization and memory management of key-value cache for serving large language models

One or more non-transitory computer-readable media storing instructions executable by a computing processor to perform operations for memory management, the operations comprising:

• determining an importance score of a key-value cache page based on one or more attention scores of one or more tokens, the key-value cache page having one or more key tensors and one or more value tensors calculated for one or more tokens by an attention head of a neural network;
• determining a quantization level based on the importance score;
• retrieving the key-value cache page at the determined quantization level from a memory;
• and
• storing the key-value cache page quantized at the quantization level in a key-value cache block of a key-value cache.

A method for memory management, comprising:

• determining an importance score of a key-value cache page based on one or more attention scores of one or more tokens, the key-value cache page having one or more key tensors and one or more value tensors calculated for one or more tokens by an attention head of a neural network;
• determining a quantization level based on the importance score;
• and
• storing the key-value cache page quantized at the quantization level in a key-value cache block of a key-value cache.