Optimizing key value cache for large language model inference

A system, comprising:

• a compute device configured to:
• receive an input sequence from a client device;
• perform large language model inference by processing the input sequence through a series of transformer layers to generate one or more tokens including by performing hybrid attention, multi-query attention, and cross-layer key value sharing,
• wherein to perform the multi-query attention at each transformer layer of the series of transformer layers, the compute device is configured to determine, for each token, corresponding key and value vectors among a plurality of attention heads within the same transformer layer, wherein performing cross-layer key value sharing includes reusing at least a portion of key vectors and value vectors generated by a preceding transformer layer for attention computations in one or more subsequent transformer layers so as to reduce a key-value cache size and memory bandwidth needed for inference, wherein the cross-layer key value sharing occurs between adjacent layers in which the local attention is performed and between non-adjacent layers in which the global attention is performed;
• and provide the one or more generated tokens to the client device;
• and
• a memory coupled to the compute device and configured to provide the compute device with instructions.

A computer program product embodied in a non-transitory computer readable medium and comprising computer instructions for:

• receiving an input sequence from a client device;
• performing large language model inference by processing the input sequence through a series of transformer layers to generate one or more tokens including by performing hybrid attention, multi-query attention, and cross-layer key value sharing,
• wherein performing the multi-query attention at each transformer layer of the series of transformer layers includes determining, for each token, corresponding key and value vectors among a plurality of attention heads within the same transformer layer, wherein performing cross-layer key value sharing includes reusing at least a portion of key vectors and value vectors generated by a preceding transformer layer for attention computations in one or more subsequent transformer layers so as to reduce a key-value cache size and memory bandwidth needed for inference, wherein the cross-layer key value sharing occurs between adjacent layers in which the local attention is performed and between non-adjacent layers in which the global attention is performed;
• and
• providing the one or more generated tokens to the client device.