Optimizing key value cache for large language model inference
An input sequence is received from a client device. Large language model inference is performed 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. The one or more generated tokens are provided to the client device.
1 . A method, comprising:
receiving an input sequence from a client device;
performing large language model inference by processing the input sequence including by tokenizing the input sequence to create a plurality of tokens, converting the plurality of tokens into embeddings, and passing the embeddings 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 hybrid attention includes:
enabling local attention for a plurality of consecutive layers of the series of transformer layers; and
injecting global attention at regular intervals in between a first plurality of local attention enabled layers of the series of transformer layers and a second plurality of local attention enabled layers of the series of transformer layers;
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.
2 . The method of claim 1 , wherein the input sequence includes a question, a prompt, or a command.
3 . The method of claim 1 , further comprising adding corresponding positional embeddings to the embeddings associated with the plurality of tokens.
4 . The method of claim 1 , wherein the first plurality of local attention enabled layers and the second plurality of local attention enabled layers include five layers.
5 . The method of claim 1 , wherein the corresponding key vectors for each token associated with the input sequence and the corresponding value vectors for each token associated with the input sequence are shared across two or more transformer layers of the transformer layers.
6 . The method of claim 1 , wherein weights used to perform the multi-query attention are represented using Int8 precision.
7 . The method of claim 1 , wherein the one or more tokens are predicted based on an output vector associated with a last token in the input sequence.
8 . The method of claim 7 , wherein the one or more predicted tokens are selected based on a highest probability, a probability distribution, or a beam search method.
9 . The method of claim 8 , wherein the one or more selected tokens are appended to the input sequence.
10 . The method of claim 9 , wherein the one or more selected tokens are appended to the input sequence until an end condition has been satisfied.
11 . The method of claim 10 , wherein the end condition is a rule, an end-of-sequence token, a maximum token limit, or a user interrupt.
12 . The method of claim 1 , wherein the one or more generated tokens are provided to the client device as the one or more generated tokens are being selected.
13 . The method of claim 1 , wherein the one or more generated tokens are provided to the client device after end condition has been satisfied.
14 . 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 including by tokenizing the input sequence to create a plurality of tokens, converting the plurality of tokens into embeddings, and passing the embeddings 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 hybrid attention, the compute device is configured to:
enable local attention for a plurality of consecutive layers; and
inject global attention at regular intervals in between a first plurality of local attention enabled layers and a second plurality of local attention enabled layers;
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.
15 . The system of claim 14 , wherein the input sequence includes a question, a prompt, or a command.
16 . The system of claim 14 , wherein the corresponding key vectors for each token associated with the input sequence and the corresponding value vectors for each token associated with the input sequence are shared across two or more transformer layers of the transformer layers.
17 . The system of claim 14 , wherein weights used to perform the multi-query attention are represented using Int8 precision.
18 . The system of claim 14 , wherein the one or more tokens are predicted based on an output vector associated with a last token in the input sequence.
19 . The system of claim 14 , wherein the one or more selected tokens are appended to the input sequence.
20 . 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 including by tokenizing the input sequence to create a plurality of tokens, converting the plurality of tokens into embeddings, and passing the embeddings 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 hybrid attention includes:
enabling local attention for a plurality of consecutive layers; and
injecting global attention at regular intervals in between a first plurality of local attention enabled layers and a second plurality of local attention enabled layers;
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.