Dual-group quantization of key-value tensors in transformer-based models

Abstract

KV tensors in a transformer model may be quantized in a dual-group manner. A key tensor or value tensor may be segmented into groups along both the token dimension and the channel dimension. A group size may be determined based on hardware constraint and model accuracy. The group size may indicate the total number of tokens or the total number of channels in each group. A part of the tensor may be segmented into groups having the group size, while the rest of the attention tensor may constitute an additional group having a larger size. The larger group may include keys or values corresponding to one or more recent tokens. Different groups within the tensor are quantized to variable precision levels with group-specific scale and zero-point values. The quantized tensor may be cached and used in a matrix multiplication operation in an attention module of the transformer model.

Description

DUAL-GROUP QUANTIZATION OF KEY-VALUE TENSORS IN TRANSFORMER-BASED MODELSCross-Reference to Related Application

[0001] This application claims the benefit of U.S. Provisional Patent Application No.63 / 650,476, filed May 22, 2024, and entitled "GROUPED KEY VALUE CACHE QUANTIZATION AND PROCESSING FOR ACCELERATED GENERATIVE INFERENCE IN MODELS," which is incorporated by reference in its entirety.Technical Field

[0002] This disclosure relates generally to neural network (also referred to as "deep neural network" or "DNN"), and more specifically, dual-group quantization of key-value (KV) tensors in transformer-based models.Background

[0003] DNNs are used extensively for a variety of artificial intelligence applications ranging from computer vision to speech recognition and natural language processing due to their ability to achieve high accuracy. However, the high accuracy comes at the expense of significant computation cost. DNNs have extremely high computing demands as there can be a large number of operations as well as a large amount of data to read and write.Therefore, techniques to improve efficiency of DNNs are needed.Brief Description of the Drawings

[0004] Embodiments will be readily understood by the following detailed description in conjunction with the accompanying drawings. To facilitate this description, like reference numerals designate like structural elements. Embodiments are illustrated by way of example, and not by way of limitation, in the figures of the accompanying drawings.

[0005] FIG. 1 illustrates an example transformer model, in accordance with various embodiments.

[0006] FIG. 2 illustrates an example embedding layer, in accordance with various embodiments.

[0007] FIG. 3 illustrates positional encoding, in accordance with various embodiments.

[0008] FIGS. 4A and 4B illustrate an example multi-head attention (MHA) layer, in accordance with various embodiments.

[0009] FIG. 5 illustrates an example linear classifier, in accordance with various embodiments.

[0010] FIG. 6 illustrates a first inference stage of a transformer model, in accordance with various embodiments.

[0011] FIG. 7 illustrates subsequent inference stages of the transformer model, in accordance with various embodiments.

[0012] FIG. 8 illustrates computations in a self-attention layer without KV caching, in accordance with various embodiments.

[0013] FIG. 9 illustrates computations in a self-attention layer with KV caching, in accordance with various embodiments.

[0014] FIG. 10 illustrates segmentation of an attention tensor for dual-group quantization, in accordance with various embodiments.

[0015] FIG. 11 illustrates segmentation of another attention tensor for dual-group quantization, in accordance with various embodiments

[0016] FIG. 12 illustrates a quantization flow in an MHA block of a transformer model during the prefill phase, in accordance with various embodiments.

[0017] FIGS. ISA and 13B illustrate a quantization flow in an MHA block with adaptive grouped precision attention, in accordance with various embodiments.

[0018] FIG. 14 illustrates a quantization flow in an MHA block with periodic quantization optimization, in accordance with various embodiments.

[0019] FIG. 15 illustrates mixed-precision, hybrid dual-group quantization of attention tensors, in accordance with various embodiments.

[0020] FIG. 16 illustrates a group MatMul layer implemented on a data processing unit and a digital signal processor, in accordance with various embodiments.

[0021] FIG. 17 is a block diagram of a DNN system, in accordance with various embodiments.

[0022] FIG. 18 is a block diagram of a DNN module, in accordance with various embodiments.

[0023] FIG. 19A illustrates an example sparse cell, in accordance with various embodiments.

[0024] FIG. 19B illustrates an example sparse cell array, in accordance with various embodiments.

[0025] FIG. 20 illustrates an example convolution, in accordance with various embodiments

[0026] FIG. 21 illustrates implement of group MatMul operations on a DNN accelerator in FIG. 17, in accordance with various embodiments

[0027] FIG. 22 illustrates a process of selecting optimal quantization configuration for KV cache of a transformer mode, in accordance with various embodiments.

[0028] FIG. 23 is a flowchart of a method for executing a transformer model, in accordance with various embodiments.

[0029] FIG. 24 is a flowchart of another method for executing a transformer model, in accordance with various embodiments.

[0030] FIG. 25 is a block diagram of an example computing device, in accordance with various embodiments.Detailed DescriptionOverview

[0031] The last decade has witnessed a rapid rise in artificial intelligence (Al) based data processing, particularly based on DNNs. DNNs are widely used in the domains of computer vision, speech recognition, image, and video processing mainly due to their ability to achieve beyond human-level accuracy. A DNN typically includes a sequence of layers. A DNN layer may include one or more deep learning operations (also referred to as "neural network operations"), such as convolution, matrix multiplication, layer normalization, batch normalization, SoftMax operation, pooling, elementwise operation, linear operation, nonlinear operation, and so on.

[0032] Input or output data of deep learning operations may be arranged in data structures called tensors. A tensor is a data structure having multiple elements across one or more dimensions. Examples of tensors include vector (which is one-dimensional (ID) tensor), matrix (which is two-dimensional (2D) tensor), three-dimensional (3D) tensors, fourdimensional (4D) tensors, and even higher dimensional tensors. A dimension of a tensor may correspond to an axis, e.g., an axis in a coordinate system. A dimension may be measured by the number of data points along the axis. The dimensions of a tensor may define the shape of the tensor. A DNN layer may receive one or more input tensors andcompute an output tensor from the one or more input tensors. Taking a convolutional layer for example, the input tensors include an activation tensor (also referred to as "input feature map (IFM)" or "input activation tensor") including one or more activations (also referred to as "input elements") and a weight tensor. The weight tensor may be a kernel (a 2D weight tensor), a filter (a 3D weight tensor), or a group of filters (a 4D weight tensor). A convolution may be performed on the input activation tensor and weight tensor to compute an output activation tensor in the convolutional layer.

[0033] Tensors in DNNs can be saved in X-major (e.g., XYZ or XZY format), Y-major formats (e.g., YXZ or YZX format), or Z-major formats (e.g., ZXY or ZYX format). The format of a tensor may define the order in which the data points in the tensor are stored, written, or read. The first character may represent the dimension in which data points are contiguous in memory. The second character may represent the dimension in which data points can be accessed after the contiguous data points are accessed in memory. The third character may represent the dimension in which data points are accessed after the data points in the dimension represented by the second character are exhausted. Taking the ZXY format for example, the access order first starts in the Z dimension, then moves to the X dimension, and finally moves to the Y dimension. Data points in the tensor are contiguous in memory in the Z dimension, meaning data points having the same (x, y) coordinates are contiguous in memory. Using tensor permutation, the tensor may be read from memory in a different format.

[0034] The significant improvements in DNN model size and accuracy coupled with the rapid increase in computing power of execution platforms have led to the adoption of DNN applications even within resource constrained mobile and edge devices that have limited energy availability. DNN models may be executed, e.g., for training or inference, by DNN accelerators. A DNN accelerator may be or include one or more data processing units. A data processing unit may also be referred to as a compute block or compute tile. A data processing unit may include processing elements (PEs) that can carry out neural network operations.

[0035] Large Language Models (LLMs), including transformer-based neural networks like GPT (Generative Pre-trained Transformer) and BERT (Bidirectional Encoder Representations from Transformers) are used for natural language processing (NLP) in various applications, from virtual assistants, document summarization, to coding assistants, language translationand speech recognition apps. Despite the remarkable capabilities of LLMs, deploying them and executing LLM inference on end user devices with DNN accelerators is not trivial due to their high computational intensity, expensive memory requirements, dynamic nature of input and output shapes, large sequence lengths, in addition to the inherent autoregressive nature of LLM generation stage.

[0036] The demand for computational resources during LLM inference depends on numerous factors, including LLM model architecture, number of parameters, and application. The term "large" in LLM signifies the scale of these models, encompassing both the extensive number of model parameters (weights), often reaching scales of billions of parameters, and the vast corpus of training data. In addition, the number of floating-point operations (FLOPs) is primarily determined by the input sequence length and the model size. The compute complexity of the MHA module in encoder and decoder scales quadratically with the input sequence length. Also, substantial portions of text may result in potentially thousands of words / tokens (e.g., in content summarization) leading to high sequence length, which leads to high compute demand. This challenge results in high inference latency and high energy consumption.

[0037] LLMs usually exhibit a substantial memory footprint, stemming from the need to store extensive pre-trained model parameters (in billions) and maintain transient states during autoregressive generation phase, particularly during decoder with KV cache inference ("cached inference") phase. These voluminous parameters not only surpass the capacity of standard device memory but also pose bottlenecks to the generation throughput. This challenge also impacts the energy consumption of the accelerator.

[0038] Low-precision representation of KV cache or token removal (eviction) can be used to limit the memory requirements. However, weight compression techniques like gradientbased ranking and other iterative approaches cannot be applied here. Since K and value tensors are usually activations that are updated with inference steps, they can be highly data-dependent. The optimization technique should be online. Other approaches can provide data-fusing strategies to improve memory access and throughput. However, these approaches cannot deal with the memory limitation issues. For static tensors like LLMs weights, compression methods including quantization and pruning have been explored. Approaches used for dynamically growing KV tensors include uniform quantization, outlier- aware dense-sparse quantization, group quantization, and KV cache reduction via tokendropping. However, these approaches can add additional compute overhead or suffer from significant accuracy drop on complex generative tasks like complex reasoning. Additionally, various compute overheads can delay the inference time token generation. Methods like clustering based non-uniform KV quantization often cannot be applied in a "plug and play" manner. These methods usually require additional calibration for the quantization to work and can affect the adaptability of the compression method on client end devices.Additionally, such methods can affect the real-time generation throughput. Simple quantization (such as uniform asymmetric quantization) requires computation of the scaling and zero-point values, which can be fast and has native hardware support for compression and decompression. However, these quantization methods are not useful to retain the model accuracy at low-precision storage.

[0039] Embodiments of the present disclosure may improve on at least some of the challenges and issues described above by providing dual-group quantization of KV tensors in transformer-based models. An example key tensor or value tensor may be segmented into groups along the token dimension and channel dimension of the tensor. Group-specific quantization parameters may be determined. The groups may be quantized separately using different quantization parameters to reach different data precisions in the quantized tensor. Keys or values corresponding to more recently generated tokens may have a higher precision than other keys or values.

[0040] An example transformer model includes an encoder, a decoder, and a head. An inference request may be made to request the transformer model to make a prediction. During inference, the encoder may receive input tokens and generate an encoder output from the input tokens. A token may be a representation (e.g., numerical representation) of a signal (e.g., a word, image, audio, video, etc.) provided to the transformer model. The decoder may receive inputs and generate decoder outputs from the dynamic inputs and the encoder output during a sequence of inference stages. The decoder inputs may be dynamic across inference stages as the token predicted in an inference stage is added to the decoder input for the next inference stage. KV tensors in the decoder may also be dynamic across inference stages. An inference stage may be a cycle of executing operations in the decoder and head to predict a new token. An inference stage may also be referred to as an inference phase or inference cycle. The first inference stage may populate the KV cache with keys and values corresponding to the input tokens received by the encoder. The decoder may use theKV tensors to generate an output, which is then used by the head to generate the first predicted token. In each subsequent inference stage, the KV cache may be further populated with new keys and values corresponding to the toke predicted in the previous inference stage. All the inference stages may constitute a whole inference process. A sequence of tokens may be obtained after all the inference stages are complete. The sequence of tokens ("output sequence") may be the prediction of the transformer model.

[0041] In various embodiments of the present disclosure, KV tensors may be quantized in a dual-group manner. For instance, an attention tensor (e.g., a key tensor or value tensor) may be segmented into groups along both the token dimension and the channel dimension. To segment the attention tensor, a group size may be determined. The group size may include a token dimension, which may indicate the total number of tokens in each group, and a channel dimension, which may indicate the total number of channels in each group. The group size may be determined based on one or more constraints of a DNN accelerator that runs the transformer model and estimated accuracy of the transformer model if the group size is used to quantize the attention tensor. The group size may be used to segment part of the attention tensor into groups. The rest of the attention tensor may constitute an additional group, which may have a larger channel dimension than the other groups. The larger group may include keys or values corresponding to the token predicted in the previous inference stage. Different groups within the attention tensor are quantized with group-specific scale and zero-point values. Different groups of the quantized tensor may get different data precisions, constituting a tensor with mixed precision. An attention block may have a hybrid dual-group quantization configuration, with which the key tensor may have a different dual-group quantization configuration from the value tensor. For instance, the group size for the key tensor may be different from the group size for the value tensor. Quantization of KV cache may happen in all inference stages or some of the inference stages during transformer inference. Periodic quantization may be performed by quantizing KV cache in inference stages selected based on a quantization interval.

[0042] The DNN accelerator may use the quantized attention tensor to perform a matrix multiplication (MatMul) operation in an attention layer (e.g., an MHA layer) in the decoder during the current inference stage. The DNN accelerator may be capable of MatMul operation on tensors of different precisions. For instance, the DNN accelerator may perform a MatMul operation on a key tensor with a low precision and a query tensor with a highprecision. A data processing unit (data processing unit) in the DNN accelerator may perform a sub-MatMul operation using each quantized group to compute partial sum and dequantize the partial sum using the corresponding scaling value and zero-point value. A data signal processor (digital signal processor) in the DNN accelerator may accumulate the dequantized partial sum to compute an output of the MatMul operation. The data processing unit may also use a query tensor having a higher precision to perform the MatMul operation. The mixed precision method can enable different precision levels within the same group, providing a sophisticated means to handle tensors with high outlier values without the need for complex segregation.

[0043] The present disclosure provides a dual-group quantization technique for maintaining a multi-precision KV cache within an LLM. It allows dynamic adjustment of data precisions within the KV cache for efficient attention computational to characterize an attention module within an LLM by adaptive grouped precision. Recent tokens are preserved in high precision while older tokens are compressed to low precision, which can optimize the quantization of a KV cache at regular intervals to balance precision and compression over time. The dual-group quantization technique can provide significant memory saving with minimum impact on the model accuracy. The hybrid dual-group quantization technique allows different dual-group quantization configurations for different tensor types and can further tailor the quantization to the unique characteristics of each tensor type. Moreover, the hardware-software co-design methodology for setting dual-group quantization configurations can maximize both model accuracy and hardware utilization, ensuring that the quantization strategies can be efficiently executed on specialized computational platforms.

[0044] The present disclosure provides a comprehensive overhaul of KV cache management in LLMs, driving forward the capabilities of language processing technology with an eye towards precision, efficiency, and hardware compatibility. With the quantization techniques in the present disclosure, LLMs can be efficiently run on resource-constrained devices (such as client devices) by leveraging the low-power inference capability of hardware accelerator.

[0045] For purposes of explanation, specific numbers, materials and configurations are set forth in order to provide a thorough understanding of the illustrative implementations. However, it will be apparent to one skilled in the art that the present disclosure may be practiced without the specific details or / and that the present disclosure may be practicedwith only some of the described aspects. In other instances, well known features are omitted or simplified in order not to obscure the illustrative implementations.

[0046] Further, references are made to the accompanying drawings that form a part hereof, and in which is shown, by way of illustration, embodiments that may be practiced. It is to be understood that other embodiments may be utilized, and structural or logical changes may be made without departing from the scope of the present disclosure. Therefore, the following detailed description is not to be taken in a limiting sense.

[0047] Various operations may be described as multiple discrete actions or operations in turn, in a manner that is most helpful in understanding the claimed subject matter. However, the order of description should not be construed as to imply that these operations are necessarily order dependent. In particular, these operations may not be performed in the order of presentation. Operations described may be performed in a different order from the described embodiment. Various additional operations may be performed or described operations may be omitted in additional embodiments.

[0048] For the purposes of the present disclosure, the phrase "A or B" or the phrase "A and / or B" means (A), (B), or (A and B). For the purposes of the present disclosure, the phrase "A, B, or C" or the phrase "A, B, and / or C" means (A), (B), (C), (A and B), (A and C), (B and C), or (A, B, and C). The term "between," when used with reference to measurement ranges, is inclusive of the ends of the measurement ranges.

[0049] The description uses the phrases "in an embodiment" or "in embodiments," which may each refer to one or more of the same or different embodiments. The terms "comprising," "including," "having," and the like, as used with respect to embodiments of the present disclosure, are synonymous. The disclosure may use perspective-based descriptions such as "above," "below," "top," "bottom," and "side" to explain various features of the drawings, but these terms are simply for ease of discussion, and do not imply a desired or required orientation. The accompanying drawings are not necessarily drawn to scale. Unless otherwise specified, the use of the ordinal adjectives "first," "second," and "third," etc., to describe a common object, merely indicates that different instances of like objects are being referred to and are not intended to imply that the objects so described must be in a given sequence, either temporally, spatially, in ranking or in any other manner.

[0050] In the following detailed description, various aspects of the illustrative implementations will be described using terms commonly employed by those skilled in the art to convey the substance of their work to others skilled in the art.

[0051] The terms "substantially," "close," "approximately," "near," and "about," generally refer to being within + / - 20% of a target value as described herein or as known in the art. Similarly, terms indicating orientation of various elements, e.g., "coplanar," "perpendicular," "orthogonal," "parallel," or any other angle between the elements, generally refer to being within + / - 5-20% of a target value as described herein or as known in the art.

[0052] In addition, the terms "comprise," "comprising," "include," "including," "have," "having" or any other variation thereof, are intended to cover a non-exclusive inclusion. For example, a method, process, device, or DNN accelerator that comprises a list of elements is not necessarily limited to only those elements but may include other elements not expressly listed or inherent to such method, process, device, or DNN accelerators. Also, the term "or" refers to an inclusive "or" and not to an exclusive "or."

[0053] The systems, methods and devices of this disclosure each have several innovative aspects, no single one of which is solely responsible for all desirable attributes disclosed herein. Details of one or more implementations of the subject matter described in this specification are set forth in the description below and the accompanying drawings.Example Transformer Model

[0054] FIG. 1 illustrates an example transformer model 100, in accordance with various embodiments. The transformer model 100 may transform input sequences into output sequences. In some embodiments, the transformer model 100 is a DNN that can learn context and meaning by tracking relationships in sequential data, such as sequential words in a sentence, sequential audio signals, sequential images, and so on. In an example, the transformer model 100 may be an LLM. The transformer model 100 includes an encoder block 110, a decoder block 120, and a head block 130. In other embodiment, different or additional components may be included in the transformer model 100. Further, functionality attributed to a component of the transformer model 100 may be accomplished by a different component included in the transformer model 100 or a different model or module.

[0055] The encoder block 110 receives input sequences and generates matrix representations of the input sequences. In the embodiments of FIG. 1, the encoder block110 receives an input 101 and generates an encoder output 102. The input 101 may be an input prompt. In some embodiments, the input 101 may include one or more input tokens, such as words, phrases, sentences, images, audio signals, other types of input tokens, or some combination thereof. In an example, the input 101 may include a prompt received from a user of the transformer model 100. The prompt may include a question or request made by the user. A word in the prompt may be an input token. The encoder output 102 may include one or more vectors that are contextualized representations of the input 101. Each vector in the encoder output 102 may represent a token in the input 101 with contextual understanding.

[0056] The encoder block 110 includes an embedding layer 113, a positional encoding layer 115, and a plurality of layers 140 (individually referred to as "layer 140"). In other embodiments, the encoder block 110 may have different, fewer, or more components. Also, the arrangement of the components in the encoder block 110 may be different from the arrangement shown in FIG. 1. For the purpose of illustration, the encoder block 110 has N layers in FIG. 1, where N is an integer. Each layer 140 may include one or more neural network operations. The layers 140 may transform a sequence of embeddings into a representation that encapsulates the learned information from the input 101. Different layers 140 may have different internal parameters, e.g., different weights, bias, or other types of internal parameters. In some embodiments, the layers 140 have identical components. The components in a layer 140 may be layers and may also be referred to as sub-layers of the layer 140. As shown in FIG. 1, a layer 140 includes four sub-layers: an MHA layer 141, an add & norm layer 142, a feed forward layer 143, and another add & norm layer 144.

[0057] The decoder block 120 iteratively generates outputs 103 using encoded representations generated by the encoder block 110. The decoder block 120 includes an embedding layer 123, a positional encoding layer 125, and a plurality of layers 150 (individually referred to as "layer 150"). For the purpose of illustration, the decoder block 120 has N layers in FIG. 1, where N is an integer. In the embodiments of FIG. 2, the number of layers 150 in the decoder block 120 is the same as the number of layers 140 in the encoder block 110. In other embodiments, the number of layers 150 in the decoder block 120 may be different from the number of layers 140 in the encoder block 110. Each layer 150 may include one or more neural network operations. Different layers 150 may havedifferent internal parameters. In some embodiments, the layers 150 may have identical components. The components in a layer 150 may be layers and may also be referred to as sub-layers of the layer 150. As shown in FIG. 1, a layer 150 includes six sub-layers: an MHA layer 151, an add & norm layer 152, an encoder-decoder attention layer 153, another add & norm layer 154, a feed forward layer 155, and another add & norm layer 156.

[0058] In some embodiments, a sequence of inference stages is performed in the decoder block 120 using encoder outputs, e.g., the encoder output 102. A matrix may be predicted through each inference stage. The outputs 103 may include a plurality of matrices. Each matrix may be further processed in the head block 130 to predict a token. The plurality of matrices may be used to predict a sequence of tokens. For the first inference stage, the decoder block 120 may receive one or more start tokens as input tokens and compute a first matrix from the input tokens and the output of the encoder block 110. The first matrix may be used by the head block 130 to predict a first token. The predicted token may be used as a new input token, in addition to the start token(s), in the second inference stage. Similarly, a second token may be predicted through the second inference stage and may be used in the third inference stage. This iteration may continue till all the inference stages are complete.

[0059] The head block 130 receives the output of the decoder block 120 and processes it in a linear layer 133 and a SoftMax layer 135. A linear operation may be performed on the output of the decoder block 120 in the linear layer 133. The linear operation may include a multiplication of the output of the decoder block 120 with a weight matrix. The output of the linear layer 133 may be a vector. In some embodiments, the head block 130 may function as a classifier. The number of data elements in the vector computed in the linear layer 133 may depend on the number of classes involved. In an example where there are M classes, where M is an integer, the vector computed in the linear layer 133 may have M data elements representing the prediction for the M classes, respectively.

[0060] The output of the linear layer 133 may be input into the SoftMax layer 135. A SoftMax function may be applied on the output of the linear layer 133 to compute probability scores. A probability score may have a value in the range from 0 to 1. In some embodiments, a probability value is computed for each data element in the vector computed in the linear layer 133. The highest one of the probability scores may be the key. The corresponding index of the key may point to the token that the transformer model 100 predicts as the next in the sequence. The final output of the transformer model 100 may bethe sequence of predicted tokens. In some embodiments, the head block 130 may be a language modeling head.

[0061] An embedding layer (e.g., the embedding layer 113 or the embedding layer 123) converts an input of the embedding layer (e.g., the input 101 or the outputs 103) into one or more embeddings. An embedding may be a vector, which is also referred to as an embedding vector or a vector embedding. The vector embedding may include a sequence of data elements. In some embodiments, the embedding layer 113 may generate a plurality of embeddings, each of which may be converted from a different input token in the input 101. The embeddings may capture the semantic meaning of the tokens in the input 101. The embeddings may be numerical representations that capture the relationships or meanings of words, phrases, or other data types. In an example where the input 101 is a prompt including a sequence of words, the embedding layer 113 may generate an embedding from each word in the input 101. The embedding layer 123 in the decoder block 120 may generate a plurality of embeddings from tokens received by the decoder block 120 in a similar manner as the embedding layer 113. Certain aspects of embedding layers are described below in conjunction with FIG. 2.

[0062] A positional encoding layer (e.g., the positional encoding layer 115 or the positional encoding layer 125) performs positional encoding on embeddings generated in the corresponding embedding layer. In some embodiments, the positional encoding layer may apply one or more positional encoding vectors (e.g., a positional encoding vector 104 or positional encoding vector 105) on vector embeddings from the corresponding embedding layer to generate new vector embeddings that represents the embeddings with positional context. The positional encoding vector may encode information about the position of the embedding in a sequence of embeddings. In some embodiments, the positional encoding layer performs an addition operation on a positional encoding vector and a vector embedding. The addition operation may be elementwise addition. The positional encoding layer may output an embedding matrix that includes the vector embeddings computed in the positional encoding layer. Certain aspects of positional encoding layers are described below in conjunction with FIG. 3.

[0063] An MHA layer (e.g., the MHA layer 141, the MHA layer 151, or the MHA layer 153) may implement a multi-head attention mechanism, which may be a multi-head selfattention mechanism or a multi-head cross-attention mechanism. In some embodiments,the MHA layer 141 or the MHA layer 151 may implement a self-attention mechanism. For self-attention, the queries, keys, and values may come from the same place. For instance, for the MHA layer 141, the queries, keys, and values may all come from the positional encoding layer 115. For the MHA layer 151, the queries, keys, and values may all come from the positional encoding layer 125. The self-attention mechanism may enable the transformer model 100 to relate each token with other tokens. The MHA layer may compute attention scores from embeddings generated in the corresponding positional encoding layer. In some embodiments, the MHA layer may receive one or more queries, one or more keys, and one or more values. In some embodiments, the MHA layer has a number of heads that receive different linearly projected versions of the queries, keys, and values and produce outputs in parallel that are then used to generate the final result.

[0064] In some embodiments, the queries, keys, and values input into the MHA layer 141 may be computed from vector embeddings generated by the positional encoding layer 115. The queries, keys, and values input into the MHA layer 151 may be computed from vector embeddings generated by the positional encoding layer 125. A query, key, or value may be a vector the represents a token in a sequence. In some embodiments, a query matrix Q E IRWx / lmay bgconnpU ed by multiply an embedding matrix X E IRWxd(e.g., an embedding matrix computed in a positional encoding layer) with a weight matrix WqE ]Rdx / l, where d is the dimension of a vector embedding, N is the number of vector embeddings in the embedding matrix, and h is the number of attention heads. Each row in the query matrix may be a query. A key matrix K E K. / Vx / ’ may be computed by multiple an embedding matrix X E HWxd(e.g., an embedding matrix computed in a positional encoding layer) with a weight matrix WkE lRdx / l. Each row in the key matrix may be a key. A value matrix V E IRiWx / lmay be computed by multiple an embedding matrix X E IRWxd(e.g., an embedding matrix computed in a positional encoding layer) with a weight matrix WvE ]R.dxh. Each row in the value matrix may be a value.

[0065] In some embodiments, the MHA layer 151 may implement masked multi-head selfattention. The MHA layer 151 may prevent positions from attending to subsequent positions. For instance, each token in the sequence may not be influenced by future tokens. This masking can ensure that the predictions of a particular position can depend on known outputs at positions before it and not depend on unknown outputs at positions after it.

[0066] In some embodiments, the MHA layer 153 may implement a cross-attention mechanism, such as encoder-decoder cross-attention. The MHA layer 153 may use outputs from the previous layer (i.e., the add & norm layer 152) as queries and use outputs from the encoder block 110 as keys and values. The cross-attention can align the encoder's input with the decoder's, empowering the decoder block 120 to identify and emphasize the most relevant parts of the encoder's input. Certain aspects of MHA layers are described below in conjunction with FIGS. 4A and 4B.

[0067] An add & norm layer in the transformer model 100, such as the add & norm layer 142, 144, 152, 154, and 156, has an addition operation followed by a layer normalization operation. The addition operation may be an addition of the output of the preceding layer and the input of the preceding layer. The preceding layer is a layer that is arranged right before the add & norm layer. For example, the preceding layer of the add & norm layer 142 is the MHA layer 141. As another example, the preceding layer of the add & norm layer 154 is the encoder-decoder attention layer 153.

[0068] Then the layer normalization operation is applied on the result of the addition operation, which may be denoted as LayerNorm(x + sublay er x ), where LayerNorm denotes layer normalization, x is the input of the preceding layer, and sublayer^x) denotes the output of the preceding layer. In some embodiments, the layer normalization operation may include a sequence of computations. In an example, the layer normalization operation may include a mean computation, which may be denoted as> where Axyzdenotes a data element in the input tensor, x may be the positional index of the data element in one of the spatial dimensions, y may be the positional index of the data element in the other one of the spatial dimensions, z may be the positional index of the data element in the channel dimension, and .xydenotes the output of the mean computation, which may be a 2D matrix. The mean computation may be channel-wise reduction operation. The layer normalization operation may convert i-txyto a 3D tensor ixyz, e.g., by replicating every data element over z output points.

[0069] The layer normalization operation may also include an elementwise subtraction, which may be denoted as Dxyz= Axyz— ixyz. The layer normalization operation may further include a variance computation denoted as a2xy=and a divisioncomputation denoted as Mxy=1Mxymay be a 2D tensor. The layernormalization operation may also convert Mxyto a 3D tensor Mxyz, e.g., by replicating every data element over z output points. Further, the layer normalization operation may have an element multip • lication denoted alayer normalization operation may further compute A"xyzX yz. LNxyzmay be the output of the layernormalization operation.

[0070] A feed forward layer (e.g., the feed forward layer 143 and the feed forward layer 155) may be a position-wise fully-connected feed forward network. In an example, the feed forward layer may include two linear layers with an activation function in between. An example of the activation function is Rectified Linear Unit (ReLU).

[0071] FIG. 2 illustrates an embedding operation in an embedding layer 200, in accordance with various embodiments. The embedding layer 200 may be an example of the embedding layer 113 or the embedding layer 123 in FIG. 1. As shown in FIG. 2, the embedding layer 200 receives an input sequence 201, which includes three words 202, 203, and 204. Each word may be a token. The embedding layer 200 generates a vector embedding 205 from the word 202. The embedding layer 200 also generates a vector embedding 206 from the word 203. The embedding layer 200 further generates a vector embedding 207 from the word 204. In the embodiments of FIG. 2, the vector embeddings 205, 206, and 207 have the same dimension, i.e., they each have five data elements. In other embodiments, the vector embedding 205, 206, or 207 may have a different dimension. Also, the input to the embedding layer 200 may be data of a type other than words, such as audio signals, images, and so on.

[0072] In some embodiments where the embedding layer 200 is in an encoder (e.g., the encoder block 110), the input sequence 201 may be an input received by the encoder, such as a prompt made by a user. The input sequence 201 may remain the same during inference of the encoder. In some embodiments where the embedding layer 200 is in a decoder (e.g., the decoder block 120), the input sequence 201 may change and the dimension of the input sequence 201 may be dynamic during inference of the decoder. In an example, the decoderinference may include a sequence of phases. Each inference stage may be conducted for predicting a token. For the first inference stage, the input sequence 201 may include one or more start tokens. For each subsequent inference stage (e.g., the second inference stage, the third inference stage, etc.), the input sequence 201 may include tokens predicted in the previous inference stages. The dimension of the input sequence may be increased by one after each inference stage.

[0073] FIG. 3 illustrates a positional encoding operation in a positional encoding layer, in accordance with various embodiments. The positional encoding layer may be an example of the positional encoding layer 115 or the positional encoding layer 125 in FIG. 1. The positional encoding operation includes an addition of a vector embedding 310 and a positional encoding vector 320. The vector embedding 310 may be generated by an embedding layer. The positional encoding vector 320 may encode information of the position of the token represented by the vector embedding 310 in a sequence of tokens. The positional encoding operation computes a vector embedding 330, which represents the token with positional context. In some embodiments, the positional encoding operation may be an elementwise addition operation. A data element in the vector embedding 330 may equal the sum of a data element in the vector embedding 310 and a data element in the positional encoding vector 320. In the embodiments of FIG. 3, the vector embedding 310, positional encoding vector 320, and vector embedding 330 have the same dimension, i.e., they each have five data elements. In other embodiments, the vector embedding 310, positional encoding vector 320, or vector embedding 330 may have a different dimension.

[0074] FIGS. 4A and 4B illustrate an example MHA layer 400, in accordance with various embodiments. The MHA layer 400 may be an example of the MHA layer 141 or the MHA layer 151 in FIG. 1. As shown in FIG. 4A, the MHA layer 400 includes linear layers 410, 420, and 430, a MatMul layer 440, a scale layer 450, a SoftMax layer 460, another MatMul layer 470, a concatenation layer 480, and another linear layer 490. In other embodiments, the MHA layer 400 may include fewer, more, or different layers.

[0075] The MHA layer 400 receive three input matrices: a query matrix 401, a key matrix 402, and a value matrix 403, which are inputs of the linear layers 410, 420, and 430, respectively. The linear layers 410, 420, and 430 are in a linear block 415 of the MHA layer 400. In some embodiments, the MHA layer 400 includes a plurality of linear blocks that includes the linear block 415. For the purpose of illustration, the MHA layer 400 includes hY1linear blocks in FIG. 4A, where h is an integer. Each of the linear blocks may have the same layers as the linear block 415. Each linear block may compute three parameter matrices from the query matrix 401, key matrix 402, and value matrix 403, respectively. For instance, the linear layer 410 may perform a multiplication of the query matrix 401 with a weight matrix to compute a parameter matrix 404 shown in FIG. 4B. The linear layer 420 may perform a multiplication of the key matrix 402 with a weight matrix to compute a parameter matrix 405 shown in FIG. 4B. The linear layer 430 may perform a multiplication of the value matrix 403 with a weight matrix to compute a parameter matrix 406 shown in FIG. 4B.

[0076] The parameter matrix 404 may be denoted as QW , where Q is the query matrix 401 andmatrix. The parameter matrix 405 may be denoted as KW- ', where K is the key matrix 401 andG ^-modei^k jsthe weight matrix. The parameter matrix 406 may be denoted as W , where V is the value matrix 401 and W / / G ^modeixd / jmatrix, i may indicate the index of the head. dqis the dimension of a query vector. dkis the dimension of a key vector. dvis the dimension of a value vector. In some embodiments, dq= dk= dv= dmodel / h.

[0077] The MatMul layer 440, scale layer 450, mask layer 455, SoftMax layer 460, and MatMul layer 470 are in an attention block 425 of the MHA layer. The attention block 425 may implement a scaled dot-product attention mechanism. In some embodiments, the MHA layer 400 includes a plurality of attention blocks that includes the attention block 425. For the purpose of illustration, the MHA layer 400 includes h attention blocks in FIG. 4A. Each of the attention blocks may have the same layers as the attention block 425. The linear block 415 and attention block 425 may constitute a head of the MHA layer 400. As the MHA layer 400 has h linear blocks and h attention blocks, the MHA layer 400 has h heads. A head may be denoted as headt= Attention (Q W , KW , VW / / ).

[0078] A matrix multiplication operation may be performed on the parameter matrices 404 and 405 in the MatMul layer 440, which computes a score matrix 406. In some embodiments, the score matrix 406 may establish the degree of emphasis each token should place on other tokens. The 406 may include a plurality of scores. Each token may be assigned a score in relation to other tokens within the same time step. A higher score may indicate a higherfocus or emphasis. The score matrix 406 may be scaled in the scale layer 450. In some embodiments, the score matrix 406 is scaled down in the scale layer 450 bydividing the scores in the score matrix by the square root of the dimension of the query vector and the key vector, which may be denoted as fd^. The output of the scale layer 450 may be a scaled matrix 408, which includes adjusted scores. The mask layer 455 may be optional in some embodiments. The mask layer 455 may add an attention mask (which may be an input to the attention block 425) to the output of the scale layer 450 to mask out some elements in the output of the scale layer 450. The positions of the masked out elements may be defined by the attention mask. A SoftMax function may be applied on the scaled matrix 408 in the SoftMax layer 460 to compute an attention weight matrix 409. The attention weight matrix 409 includes attention weights. The attention weights may be probability values ranging from 0 to 1. The SoftMax function may emphasize high scores while diminishing low scores, which can enhance the model's ability to determine which tokens should get more attention.

[0079] In the MatMul layer 470, a matrix multiplication operation is performed on the attention weight matrix 409 computed in the SoftMax layer 460 and the parameter matrix 407 computed from value matrix 403 in the linear layer 430. The result of the matrix multiplication operation is a single-head output matrix 411, which is an output of the attention block 425.

[0080] As the MHA layer 400 has h attention blocks, there will be h single-head output matrices. The single-head output matrices are concatenated in the concatenation layer 480 to form a concatenated matrix. A linear operation (also referred to as "linear transformation") is performed on the concatenated matrix using a weight matrix in the linear layer 490. In some embodiments, the MHA may be denoted as MultiHead Q, K, V') = Concat head1, head2, ... , head^W0, where Concat denotes concatenation, and W° £ ]hdvxdmodei jsthe weight matrix in the linear layer 490.

[0081] FIG. 5 illustrates an example linear classifier 500, in accordance with various embodiments. The linear classifier 500 may be used in transformer models. In some embodiments, the linear classifier 500 may generate tokens based on outputs of decoders. The linear classifier 500 may be an example of the linear block 130 in FIG. 1. As shown in FIG. 5, the linear classifier 500 includes a linear layer 510 and a SoftMax layer 520. In other embodiments, the linear classifier 500 may include fewer, more, or different components.

[0082] The linear layer 510 is provided with a matrix 501. The matrix 501 may be an output of a decoder, e.g., the decoder block 120. A linear transformation may be performed on the matrix 501 and a weight matrix in the linear layer 510. The weight matrix may include weights, which are internal parameters of the linear layer 510. The linear layer outputs a vector 502. In some embodiments, the dimension of the vector 502 (e.g., the total number of elements in the vector 502) may be equal to the total number of classes associated with the Al task being performed by the transformer model. The vector 502 is provided to the SoftMax layer 520. The SoftMax layer 520 generates a vector 503 from the vector 502. In some embodiments, the dimension of the vector 503 may equal the dimension of the vector 502. Each element in the vector 503 may correspond to a predicted token and may indicate a probability score of the predicted token. The probability score may indicate the probability that the prediction is correct. A predicted token 504 having the highest probability score may be selected and output from the linear classifier 500.

[0083] The output of the linear classifier 500 may be the output of the transformer model. The execution of the linear classifier 500 may be performed multiple times during inference of the transformer model. For instance, the transformer model may have multiple inference stages, and the linear classifier 500 may be executed at least once in each inference stage. The dimensions of the vectors and matrices shown in FIGS. 2-5 are example dimensions used for purpose of illustration and simplicity. Any of the vectors and matrices used or computed by operations illustrated in FIGS. 2-5 may have different dimensions.

[0084] FIG. 6 illustrates a first inference stage of a transformer model 600, in accordance with various embodiments. The transformer model 600 includes an encoder 610, a decoder 620, and a head 630. An example of the transformer model 600 may be the transformer model 100 in FIG. 1. In the embodiments of FIG. 6, the encoder 610 receives an input tensor 601. The input tensor 601 may be a feature map extracted from one or more images, text documents, audio files, videos, other types of data, or some combination thereof. In some embodiments, the input tensor 601 may be generated by another neural network, e.g., a CNN. The encoder 610 generates an output tensor 602 from the input tensor 601. The shape of the output tensor 602 may be denoted as [batch size, SLencoder, dmodei], where SLencOdermaY be the dimension along the X axis (i.e., the width of the output tensor 602), and dmodeimay be the dimension along the Y axis (i.e., the height of the output tensor 602). The encoder 610 may include a plurality of layers arranged in a sequence, such as the layersinside the encoder block 110 in IFG. 1. The output tensor 602 is provided to the decoder 620.

[0085] The decoder 620 receives the output tensor 602 and an input sequence 603. The input sequence 603 may be a sequence of tokens. A token may be a numerical representation of an input signal, such as word, image, audio signal, video signal, etc. The dimension of the input sequence 603, which may be denoted as SLinput, may be the total number of tokens in the input sequence 603. For the purpose of illustration and simplicity, SLinputis 4. In other embodiments, the input sequence 603 may have a different shape. For instance, the input sequence 603 may be a 2D tensor. The dimension of the 2D tensor along the X axis may be SLinput, while the dimension of the 2D tensor along the Y axis may be a batch size indicating the number of batches in the input sequence 603.

[0086] The decoder 620 computes an output tensor 604, a self-attention key tensor 605, a self-attention value tensor 606, a cross-attention key tensor 607, and a cross-attention value tensor 608. In some embodiments, the shape of the output tensor 604 may be denoted as [batch size, SLinput, dmodelJ. The shape of the self-attention key tensor 605 or the shape of the self-attention value tensor 606 may be denoted as N X[batch size, h,SLinput, dhead], where N is the number of identical layers in the decoder (e.g., the number of layers 150 in the decoder block 120), h is the total number of heads in a MHA layer, and dheadis the dimension of a query vector, key vector, or value vector. In some embodiments, dmodei= h X dhead. The shape of the cross-attention key tensor 607 or the shape of the cross-attention value tensor 608 may be denoted as N X [batch size, h,SLencoder, dhead]-

[0087] The output tensor 604 may be provided to the head 630 and the head 630 outputs a predicted token 609. The shape of the token 609 may be denoted as [batch size, 1], For the purpose of illustration and simplicity, batch size is 1 in FIG. 6. In other embodiments, batch size may be a larger number. The predicted token 609 may be stored in a buffer. In some embodiments, the predicted token 609 may be used to update the input sequence 603. For instance, the predicted token 609 may be added to the right of the input sequence 603. The updated input sequence may be used as the input sequence in the second inference stage. In the second inference stage, the decoder 620 may receive the updated input sequence and the output tensor 602 for predicting another token. The output tensor602 may remain the same during inference of the decoder 620. Certain aspects of subsequent inference stages are described below in conjunction with FIG. 7.

[0088] In some embodiments, the self-attention key tensor 605 and the self-attention value tensor 606 may be provided to a self-attention layer in the decoder 620, an example of such a self-attention layer is the MHA layer 151. The self-attention key tensor 605 may be stored in a self-attention key cache. The self-attention key cache may have the same shape as the self-attention key tensor 605. The self-attention value tensor 606 may be stored in a selfattention value cache. The self-attention value cache may have the same shape as the selfattention value tensor 606.

[0089] In some embodiments, the decoder 620 computes the self-attention key tensor 605 and the self-attention value tensor 606 from the input sequence 603. The input sequence603 may be dynamic during inference of the decoder 620. For instance, a new token may be added to the input sequence 603 after each inference stage, as described above. As the input sequence 603 changes, the self-attention key tensor 605 and the self-attention value tensor 606 would also change. For instance, the dimension of the self-attention key tensor 605 or the self-attention value tensor 606 along the X axis may increase as SLinputincreases. The self-attention key cache and the self-attention value cache may change during all the inference stages of the decoder 620 to accommodate the changes in the selfattention key tensor 605 and the self-attention value tensor 606.

[0090] In some embodiments, the cross-attention key tensor 607 and the cross-attention value tensor 606 may be provided to a cross-attention layer in the decoder 620, an example of such a cross-attention layer is the MHA layer 153. The cross-attention key tensor 607 may be stored in a cross-attention key cache. The cross-attention key cache may have the same shape as the cross-attention key tensor 607. The cross-attention value tensor 608 may be stored in a cross-attention value cache. The cross-attention value cache may have the same shape as the cross-attention value tensor 608. In some embodiments, the decoder 620 computes the cross-attention key tensor 607 and the cross-attention value tensor 606 from the output tensor 602 generated in the encoder 610. As the output tensor 602 does not change during inference of the decoder 620, the cross-attention key tensor 607 and the cross-attention value tensor 606 may remain the same during all the inference stages of the decoder 620. The cross-attention key cache and the cross-attention value cache may remain the same during all the inference stages of the decoder 620.

[0091] FIG. 7 illustrates subsequent inference stages of the transformer model, in accordance with various embodiments. In the second inference stage, the decoder 620 may reuse the self-attention key tensor 605, self-attention value tensor 606, cross-attention key tensor 607, and cross-attention value tensor 608. The decoder 620 also receives the predicted token 609. The decoder 620 may compute self-attention key vectors from the predicted token 609 and concatenate the self-attention key vectors with the self-attention key tensor 605 to generate a new self-attention key tensor 615. For instance, a selfattention key vector for each head may be added to the right of a self-attention key matrix in the self-attention key tensor 605, and the self-attention key vector and the self-attention key matrix may correspond to the same head. The elements highlighted with a dot pattern in the self-attention key tensor 615 are the self-attention key vectors generated from the predicted token 609.

[0092] Similarly, the decoder 620 may compute self-attention value vectors from the predicted token 609 and concatenate the self-attention value vectors with the self-attention value tensor 606 to generate a new self-attention value tensor 616. For instance, a selfattention value vector for each head may be added to the right of a self-attention value matrix in the self-attention value tensor 606, and the self-attention value vector and the self-attention value matrix may correspond to the same head. The elements highlighted with a dot pattern in the self-attention value tensor 616 are the self-attention value vectors generated from the predicted token 609.

[0093] The decoder 620 also generates an output tensor 614. The decoder 620 may generate the output tensor 614 using the new self-attention key tensor 615 and new selfattention value tensor 616. The output tensor 614 is used by the head 630 to generate another predicted token 619. The predicted token 619 is the output of the transformer model 600 in the second inference stage.

[0094] One or more other subsequent inference stages may be conducted. In each subsequent inference stage, the decoder 620 receives a token predicted in the previous inference stage, a self-attention key tensor generated in the previous inference stage, a selfattention value tensor generated in the previous inference stage, the cross-attention key tensor 607, and the cross-attention value tensor 608. The decoder 620 may, in the subsequent inference stage, generate a larger self-attention key tensor and a larger self-attention value tensor, in addition to an output tensor which can be used by the head 630 to predict a new token.

[0095] In embodiments where the total number of inference stages is N, the input sequence 603 is updated to an input sequence 613 after N — 1 inference stages. In the last inference stage (i.e., the Nth inference stage), the decoder 620 may receive the predicted token generated in the (N — th inference stage, the self-attention key tensor generated in the (N — th inference stage, the self-attention value tensor generated in the (N — l^th inference stage, the cross-attention key tensor 607, and the cross-attention value tensor 608. The decoder 620 may generate a self-attention key tensor 625 and a selfattention value tensor 626 using the predicted token generated in the (N — l)th inference stage, the self-attention key tensor generated in the (TV — l)t / r inference stage, and the self-attention value tensor generated in the ( / — l)t / i inference stage. The dimensions of the self-attention key tensor 625 or self-attention value tensor 626 along the X axis is SLinput+ N. The decoder 620 also generates an output tensor 624, which is used by the head 630 to generate the last predicted token 629. The N tokens predicted by the transformer model in the N inference stages may constitute an output tensor 639, which may be the final output of the transformer model.

[0096] FIG. 8 illustrates computations in a self-attention layer without KV caching, in accordance with various embodiments. The self-attention layer may be a multi-head selfattention layer. In some embodiments, the self-attention layer is in a decoder of a transformer. The computations in the self-attention layer may include multiplication of a query matrix 810 and a key matrix 820, which results in an attention weight matrix 830. In some embodiments, the self-attention layer may be a masked self-attention layer. One or more elements in the attention weight matrix 830 may be masked. For instance, the elements highlighted with a dotted pattern in FIG. 8 may be masked. The computations in the self-attention layer also includes multiplication of the attention weight matrix 830 and a value matrix 840, which results in an output matrix 850 encoding new tokens. In other embodiments, the computations in the self-attention layer may include other computations, such as computations with a scaling function, SoftMax function, and so on. For the purpose of simplicity and illustration, these computations are not shown in FIG. 8.

[0097] Each of the query matrix 810, key matrix 820, and value matrix 840 may include a vector for each of the tokens in the input sequence. For the purpose of illustration andsimplicity, the input sequence has four tokens: tokens 1-4. In the embodiments of FIG. 8, as the decoder does not implement KV caching, computations on all the key tokens in the key matrix 820 and all the value tokens in the value matrix 840 need to be conducted. Some of the computations have already been conducted in the previous inference stage, e.g., computations on the key tokens 1-3 and computations on the value tokens 1-3. The duplication of these computations can be a waste of computational resources, such as power, time, and so on.

[0098] FIG. 9 illustrates computations in a self-attention layer with KV caching, in accordance with various embodiments. Forthe purpose of illustration and simplicity, the self-attention in FIG. 9 may have the same query matrix, key matrix, and value matrix as the self-attention in FIG. 8. Different from the embodiments of FIG. 8, the decoder implements KV caching in the embodiment of FIG. 9. With the KV caching, the keys and values used in the previous inference stage(s) as well as data computed from the keys and values in the previous inference stage(s) are cached and can be reused in the current inference stage. The KV caching can reduce the amount of computations in the self-attention layer. Data that can be retrieved from cache is highlighted with a dotted pattern in FIG. 9. The amount of multiplication is reduced. Therefore, computational resources can be saved. The performance and efficiency of the transformer model can be improved. In some embodiments, the computations in FIG. 9 are computations in the fourth inference stage of a decoder, which is carried out after the generation of three tokens in three inference stages that were previously carried out.Example Dual-Group Quantization of KV Tensors

[0099] FIG. 10 illustrates segmentation of an attention tensor 1000 for dual-group quantization, in accordance with various embodiments. The attention tensor 1000 may be a key tensor or value tensor of an attention layer in a transformer model. The attention tensor 1000 may be an example of key tensors or value tensors in the transformer model 100 described above in conjunction with FIG. 1. For the purpose of illustration and simplicity, the attention tensor 1000 is a 2D tensor. In other embodiments, the attention tensor 1000 may be a 3D tensor.

[0100] As shown in FIG. 10, the attention tensor 1000 has a spatial shape IseqX dmodei. Iseqdenotes the token dimension of the attention tensor 1000, which may indicate the number of tokens in the attention tensor 1000. dmodeldenotes the channel dimension of theattention tensor 1000, which may indicate the number of channels in the attention tensor 1000. In some embodiments, dmodei= nheadsX dhead, where nheadsdenotes the number of heads and d / ieaddenotes the dimension of each head. The attention tensor 1000 is segmented into groups 1010A (individually referred to as "group 1010A") and another group 1010B. The groups 1010A and 1010B are collectively referred to as "group 1010" or "groups 1010." Each group 1010A has a spatial shape TgX Cg. The group 1010B has a spatial shape THP dmodei.

[0101] In some embodiments, the attention tensor 1000 may be quantized using dual-group quantization, as the segment of the attention tensor 1000 is along both the token dimension and the channel dimension. Taking the groups 1010A for example, there are three groups 1010A along the channel dimension and two groups 1010A along the token dimension. The groups 1010 are quantized separately. For instance, each of the groups 1010 may have a set of group-specific quantization parameters so that the groups 1010 have variable data precisions. In an example, the group 1010B may have a higher precision than the groups 1010A. All elements in the same group 1010A or 1010B share identical scale and zero-point values. In the embodiments of FIG. 10, dmodel— CgX 3, Iseq— TLP+ THP, and TLP— TgX 2, where LP stands for low precision and HP stands for high precision.

[0102] In the embodiments of FIG. 10, the dual-group quantization has a token-heavy configuration. The grouping of elements is biased towards the token dimension. Each group 1010A has a larger number of tokens Tg) compared to the number of channels (Cfl), i.e., Tg> Cg. As an example, Tg= 6, Cg= 2 in FIG. 10. This variant is designed to exploit the statistical similarities among tokens within the same group. The token-heavy configuration may be advantageous as tokens within certain segments of the input sequence might exhibit correlated features or similar patterns of activity. By grouping these tokens together and applying a common quantization scale and zero-point, their collective information can be more accurately represented with fewer bits, leading to a more compact and memoryefficient tensor. Token-heavy quantization can provide several advantages. For example, it can lead to a reduction in the overall memory footprint of the model, as fewer unique quantization parameters are needed across the grouped elements. As another example, it can improve computational efficiency during inference, as the shared scale and zero-point values allow for batched processing of quantized tokens. As yet another example, thisapproach can maintain the accuracy of the model by preserving the fidelity of token-level information that is critical for the model's performance.

[0103] FIG. 11 illustrates segmentation of another attention tensor 1100 for dual-group quantization, in accordance with various embodiments. The attention tensor 1100 may be a key tensor or value tensor of an attention layer in a transformer model. The attention tensor 1100 may be an example of key tensors or value tensors in the transformer model 100 described above in conjunction with FIG. 1. For the purpose of illustration and simplicity, the attention tensor 1100 is a 2D tensor. In other embodiments, the attention tensor 1100 may be a 3D tensor.

[0104] As shown in FIG. 11, the attention tensor 1100 has a spatial shape lseqx dmodel. lseqdenotes the token dimension of the attention tensor 1100, which may indicate the number of tokens in the attention tensor 1100. dmodeidenotes the channel dimension of the attention tensor 1100, which may indicate the number of channels in the attention tensor 1100. In some embodiments, dmodel= nheadsX dhead, where nheadsdenotes the number of heads and dheaddenotes the dimension of each head. The attention tensor 1100 is segmented into groups 1110A (individually referred to as "group 1110A") and another group 1110B. The groups 1110A and 1110B are collectively referred to as "group 1110" or "groups 1110." Each group 1110A has a spatial shape TgX Cg. The group 1110B has a spatial shape THP dmo ei.

[0105] In some embodiments, the attention tensor 1100 may be quantized using dual-group quantization, as the segment of the attention tensor 1100 is along both the token dimension and the channel dimension. Taking the groups 1110A for example, there are two groups 1110A along the channel dimension and three groups 1110A along the token dimension.The groups 1110 are quantized separately. For instance, each of the groups 1110 may have a set of group-specific quantization parameters so that the groups 1110 have variable data precisions. In an example, the group 1110B may have a higher precision than the groups 1110A. All elements in the same group 1110A or 1110B share identical scale and zero-point values. In the embodiments of FIG. 11, dmodei= CgX 2, Iseq= TLP+ THP, and TLP= TgX 3, where LP stands for low precision and HP stands for high precision.

[0106] FIG. 11 shows another variant of dual-group quantization that emphasizes the channel dimension, catering to the feature-level intricacies of the attention tensor 1100. Inthis scheme, each group includes a larger number of channels than tokens, i.e., Cg> Tgwith all elements within a group sharing the same scale and zero-point values, e.g., Tg= 2, Cg= 3 as shown in FIG. 11. In some embodiments, the channel-heavy configuration can be more advantageous as channels, representing feature dimensions, can exhibit similar statistical properties and thus benefit from shared quantization parameters. The channel dimension of K and value tensors may encapsulate the learned features that represent different aspects of the input data. Channel-heavy quantization can group these features based on their statistical behavior, allowing for a more efficient representation. By quantizing similar features together, the model can maintain the integrity of the information conveyed by the channels. This variant can offer distinct advantages. For instance, it can lead to a more streamlined model by reducing the number of unique quantization parameters required for the channels in KV cache, which can be particularly beneficial for models with a large number of features. This reduction in parameters can translate to lower memory requirements and faster processing times. Additionally, by preserving the relationships between features within each group, channel-heavy quantization can help maintain the model's discriminative power and accuracy.

[0107] In some embodiments, per-channel group quantization or per-token group quantization may be used within the dual-group quantization framework. Per-channel group quantization is a special case of token-heavy dual-group quantization where the number of channels in each group is 1. This approach can be effective for KV tensors where the channel-wise distribution of values varies significantly. Per-token group quantization is special case of channel-heavy dual-group quantization, with 1 token in each group. This can be advantageous when the sequence of tokens exhibits diverse value ranges. Both strategies may be derived from the principles of dual-group quantization and are optimized for different tensor characteristics. These special cases may represent the extremes of the dual-group quantization spectrum, where the quantization is heavily skewed towards one dimension.

[0108] For token-heavy and channel-heavy quantization, a critical parameter is the value of Cg. To align with the architecture of the attention mechanism, where each head processes different aspects of the input, Cgmay be selected to ensure that group boundaries eitherfall within a single head or coincide with the boundaries of individual attention heads. In some embodiments, the following constraint on Cgmay be applied:Cg= 2n, n = 0, 1, 2, ... ; Cg< dhead.

[0109] This design choice can simplify the quantization process by preventing groups from spanning across multiple heads, thereby avoiding added complexity in the computation of attention scores on the hardware, where MatMul corresponding to each head may be computed separately. This can facilitate independent and efficient processing of the quantized KV tensors. In addition, this approach can not only facilitate a more straightforward implementation of the quantization strategy but also enhance the model's ability to represent diverse features without the interference that could arise from group overlap between heads. In some embodiments, additional constraints may be applied on Cgbased on hardware utilization, which is described below.

[0110] The dual-group quantization variants described above may depend on the segmentation of KV tensors into contiguous blocks within the tensor structure. Unlike existing quantization techniques such as outlier quantization that rearranges outlier values in KV tensors into separate groups bringing additional complexity during inference, the dual- group quantization techniques do not invoke any rearrangement and performs grouping on the original KV tensors. By maintaining contiguous blocks of KV groups, the computing system (e.g., DNN accelerator) running LLM inference can leverage spatial locality, resulting in faster access times and reduced cache misses. The enhanced cache coherence can be very advantageous for LLMs that require rapid access to large volumes of KV cache.

[0111] In some embodiments, hybrid dual-group quantization for KV cache may be used. This novel approach involves quantizing keys and values using different variants of dual- group quantization. The dual-group quantization variants may be tailored to the unique characteristics and roles of each tensor. Hybrid quantization allows for the application of, for example, a channel-heavy dual-group quantization to the key tensor and a token-heavy dual-group quantization to the value tensor, or vice versa. The hybrid configuration may be set based on the empirical analysis of the model's performance (e.g., accuracy) on the downstream task and the hardware's computational efficiency. The hybrid quantization method can capitalize on the strengths of each dual-group quantization variant, optimizing the quantization process for both K and V in a manner that enhances the overall accuracy ofthe LLM. By selecting the most suitable dual-group quantization variant for each tensor, the precision and granularity of quantization can be aligned with the accuracy demands of the model and hardware utilization.

[0112] The dual-group quantization approaches described above are more advantageous than currently available quantization approaches such as per-tensor, per-head, or perchannel quantization approaches. These currently available quantization approaches usually apply a one-size-fits-all method to quantize KV tensors. The dual-group quantization approaches can maintain the integrity of the original tensor structure, avoiding the complexity and overhead associated with rearranging values, as seen in outlier quantization techniques. The result is a quantization strategy that can not only meet the model's accuracy demand but also streamline the inference process, enabling LLMs to operate with greater speed and reduced computational load. Thus, the flexibility of choosing among all the dual-group quantization approaches can provide a customizable solution that can be optimized for specific use cases and hardware configurations.

[0113] FIG. 12 illustrates a quantization flow in an MHA block of a transformer model during the prefill phase, in accordance with various embodiments. The attention block 1200 may be an MHA block, examples of which may include the MHA layer 151 and MHA layer 153 in FIG. 1. In the embodiments of FIG. 12, the attention block 1200 includes a MatMul layer 1210, a mask layer 1220, a SoftMax layer 1230, and another MatMul layer 1240. An example of the MatMul layer 1210 may be the MatMul layer 440 in FIG. 4A. An example of the mask layer 1220 may be the mask layer 455 in FIG. 4A. An example of the SoftMax layer 1230 may be the SoftMax layer 460 in FIG. 4A. An example of the MatMul layer 1240 may be the MatMul layer 470 in FIG. 4A. In other embodiments, the attention block 1200 may include fewer, more, or different layers.

[0114] The prefill phase may be the first inference phase, for which a KV cache is prefilled with keys and values corresponding to the input prompt. In the embodiments of FIG. 12, the KV cache includes a key tensor 1201 and a value tensor 1202, denoted as KHPand VHP, respectively. The key tensor 1201 and value tensor 1202 may be computed in seperate linear layers of the transformer model, such as the linear layer 420 and linear layer 430 in FIG. 4A. The key tensor 1201 and value tensor 1202 are input into the attention block 1200. Also, a query tensor 1203, denoted as QHP, is input into the attention block 1200. The MatMul layer 1210 may include a MatMul operation on the key tensor 1201 and querytensor 1203. The result of the MatMul operation is further processed in the mask layer 1220 and SoftMax layer 1230. The MatMul layer 1240 includes a MatMul operation on the output of the SoftMax layer 1230 and the value tensor 1202. In some embodiments, the key tensor 1201, value tensor 1202, and query tensor 1203 may all have high precision.

[0115] The key tensor 1201 is provided to a quantization module 1250. The quantization module 1250 converts the key tensor 1201 to a quantized key tensor denoted as KLPin FIG. 12. The quantized key tensor is then populated to a compressed KV cache 1270. Similarly, the value tensor 1202 is provided to the quantization module 1250. The quantization module 1250 converts the value tensor 1202 to a quantized value tensor denoted as VLPin FIG. 12. The quantized value tensor is also populated to the compressed KV cache 1270. The compressed KV cache 1270, compared with the original KV cache that includes the key tensor 1201 and value tensor 1202, has a smaller storage size, as the key tensor 1201 and value tensor 1202 have a high precision while the quantized key tensor and quantized value tensor have a low precision. By caching the quantized key tensor and quantized value tensor instead of the key tensor 1201 and value tensor 1202, the memory footprint and memory bandwidth usage can be reduced. The KV cache 1207 may be used in the next inference stage, e.g., the second inference stage. With the low KV precision, energy and time needed for computations in the second inference stage can also be reduced.

[0116] In some embodiments, the quantization module 1250 may quantize the key tensor 1201 and value tensor 1202 with different quantization functions. For instance, the quantization module 1250 may use different variants of dual-group quantization for the key tensor 1201 and value tensor 1202. In other embodiments, the quantization module 1250 may quantize the key tensor 1201 and value tensor 1202 with the same quantization function. The quantize function, utilizing predetermined DGQ configurations and hyperparameters (Tg, Cg), may group-quantize the key tensor 1201 and value tensor 1202 individually into the lower-precision tensors KLPand VLP. The quantization module 1250 may partition the key tensor 1201 or value tensor 1202 into groups and perform groupspecific quantization. In some embodiments, the total number of groups 1010A in the key tensor 1201 or value tensor 1202 may be determined by the equation:where linputis the length of the initial input prompt or sequence. The number of recent tokens generated, denoted as TKrec, is zero for the prefill inference phase.

[0117] These quantized tensors may be written directly to the main memory (e.g., a dynamic random-access memories (DRAM)), forming the compressed KV cache 1270. Additionally, the scale and zero-point values for each group are stored in memory in a high- precision format (e.g., FP16). As shown in FIG. 12, the unquantized KHP, VHP, and QHPare directly utilized in the attention computations in the prefill inference phase. The attention computations are performed in high precision. This approach eliminates the need for dequantization, thereby preserving the accuracy of the attention computations and maintaining computational efficiency. In some embodiments, high precision may be a floating-point precision, such as FP32, FP16, BP16, FP8, etc.; low precision may be an integer precision, such as INT8, INT4, INT2, INT2, FP4, FP6, MXFP4, MXFP6, etc. In some embodiments, the low precision is lower than FP16. In some embodiments, the key precision and value precision may be different in a single layer or in each head of a layer. The key precision and value precision may be different across layers.

[0118] During the memory bound cached inference phase of LLMs that runs for multiple inferences, quantization intervals, denoted as qint, may be determined. To facilitate high model accuracy and low quantization / dequantization overhead, a multi-precision KV cache maintenance mechanism may be used for cached inference stages, such as the second inference phase and subsequent inference phases. In autoregressive LLMs, the generation of each output token may be intricately dependent on the sequence of tokens that precede it, with the most recent tokens being particularly vital as they carry the immediate context necessary for accurate predictions. Leveraging this characteristic of many transformer models, multi-precision KV cache may be used to maintain both low-precision and high- precision keys, denoted as KLP, KHP, respectively, as well as both low-precision and high- precision values, denoted as VLPVHP, respectively.

[0119] As shown in FIG. 10, the attention tensor 1000 can be partitioned into two subtensors based on data precision: a low-precision subtensor including the groups 1010A and a high-precision subtensor that is the group 1010B. Similarly, the attention tensor 1100 in FIG. 11 can be partitioned into a low-precision subtensor including the groups 1110A and a high-precision subtensor that is the group 1110B. TLPand THPindicate the token dimension for the low-precision subtensor and the high-precision subtensor, respectively, in eachattention tensor. This can ensure that the keys or values corresponding to the recent tokens are kept at high precision across all quantization configurations, including token-heavy dualgroup quantization, channel-heavy dual-group quantization, and their extreme variants, perchannel group quantization and per-token group quantization. In some embodiments, THPindicates the number of tokens in high precision, may vary across models, datasets, or downstream applications. The value of THPmay be a user input.

[0120] In some embodiments, the presence of scale and zero-point parameters for THPmay be precision-dependent. For instance, FP32 and FP16 do not require these parameters, whereas INT8 does. The aggregate number of tokens at any inference stage may be denotedasIseq=TLP+ THP. THPmay continuously grow as the model keeps on generating new tokens. TLPaccounts for the total number of tokens across all quantized groups post- dual- group quantization, which may be denoted as TLP= linput+ n x Tg, where n represents the number of groups established post-prefill phase. The quantization module 1250 may set the hyperparameter qintsuch that qint> Tg, ensuring that the most contextually relevant information is preserved in the highest fidelity for the attention mechanism's computations, which may be crucial for the model to make accurate predictions. This can minimize the risk of error propagation that could arise from lower-precision representations, particularly for tokens that directly influence subsequent outputs. This can also simplify the quantization process by deferring the computation of scale and zero-point parameters for new tokens, which can be challenging given the continuously growing size of KV cache. Further, this can eliminate dequantization of the entire KV cache that needs to be concatenated with the current K and value tensor in each cached inference.

[0121] In addition or alternative to the multi-precision KV cache maintenance mechanism, a group-based attention mechanism that is adaptive to the input precision of the KV cache and query may be used in transformer model inference. The adaptive grouped precision attention mechanism can optimize computational efficiency and memory usage by performing mixed-precision MatMul operation using compressed low-precision K and value tensors and high-precision MatMul operation using high-precision K and value tensors. This ensures that the most contextually relevant tokens are processed with high precision to maintain the accuracy of the model's predictions. This adaptive approach allows thetransformer model to balance performance with hardware constraints, making it well-suited for the demanding requirements of real-time language processing tasks.

[0122] FIGS. 13A and 13B illustrate a quantization flow in an MHA block with adaptive grouped precision attention, in accordance with various embodiments. In some embodiments, the adaptive attention block 1310 may be used in lieu of the attention block 1200 in FIG. 12. The adaptive attention block 1310 facilitates the adaptive grouped precision attention mechanism. As shown in FIG. 13B, the adaptive attention block 1310 includes a group MatMul layer 1311, a MatMul layer 1312, a concatenation layer 1313, an attention mask 1314, a SoftMax layer 1315, another group MatMul layer 1316, another MatMul layer 1317, and an addition layer 1318. In other embodiments, the adaptive attention block 1310 may include fewer, more, or different components.

[0123] As shown in FIG. 13A, the adaptive attention block 1310 is associated with a KV cache 1320 including a LP cache 1330 and a HP cache 1340, a concatenator 1350, and another concatenator 1360. The LP cache 1330 may store a low-precision key tensor denoted as KLP. the shape of the key tensor is represented by [TLP, dmodei], which indicates the token dimension and the channel dimension of the key tensor. The key tensor is retrieved from the LP cache 1330 and input into the group MatMul layer 1311. The group MatMul layer 1311 also receives a high-precision query tensor denoted as QHP. The group MatMul layer 1311 may include group MatMul operations. A group MatMul operation is a MatMul operation on QHPand a group within KLP. In some embodiments, each group MatMul operation computes a partial sum. The group MatMul layer 1311 may include an accumulation operation in which the partial sums are accumulated to compute the output of the group MatMul layer 1311, the shape of which is dentoed as [1, TLP], The scale and zero-point for each group may be consumed inside the group MatMul layer 1311 to produce the full precision tensor. The granularity at which the group MatMul operations are executed depends on Tgand Cg. The dequantization of KLPmay be fused into the group MatMul layer 1311 to reduce overhead. In some embodiments, the partial sums or the accumulation of the partial sums may be dequantized in the group MatMul layer 1311. When the key tensor and Q tensor have the same low precision (INT8), the dequantization step can be eliminated altogether and integer MatMul operation can be performed. In this scenario, s and z may still be consumed for ensuring correct arithmetic.

[0124] The high-precision key tensor KHPwith a size [THP,dmodei] is read from the HP cache 1340 and concatenated with a current key tensor 1301 by the concatenator 1350. The current key tensor 1301 is a high-precision key tensor KHPwith a size [1, dmodei]. The current key tensor 1301 may be produced by a preceding linear projection layer. The concatenation operation yields a concatenated high-precision key tensor KHPwith a size [THP+ 1, dmodel]. Also, the high-precision value tensor VHPwith a size [THP, dmodel] is read from the HP cache 1340 and concatenated with a current value tensor 1302 by the concatenator 1360. The current value tensor 1302 is a high-precision value tensor VHPwith a size [1, dmodei]. The current value tensor 1302 may be produced by a preceding linear projection layer. The concatenation operation yields a concatenated high-precision value tensor VHPwith size [THP+ l, dmodel],

[0125] The concatenated high-precision key tensor and the query tensor 1303 are input into the MatMul layer 1312. In the MatMul layer 1312, the query tensor 1303 undergoes a high- precision MatMul operation with the concatenated key tensor, producing an output tensor of size [1, THP+ 1], In some embodiments, the concatenated key tensor and the query tensor 1303 has the same precision. In other embodiments, the query tensor 1303 may have a lower precision than the concatenated key tensor and may be upscaled to match the K precision. The MatMul layer 1312 may be the same or similar as the MatMul layer 1210 in the attention block 1200.

[0126] The outputs of the group MatMul layer 1311 and MatMul layer 1312 are concatenated along the token dimension in the concatenation layer 1313. The result of the concatenation is a concatenated tensor of size [1, TLP+ THP+ 1], This concatenated tensor undergoes attention operations, including masking and SoftMax, in the attention mask 1314 and SoftMax layer 1315, respectively. In some embodiments, the attention mask 1314 may be the same or similar as the mask layer 1220 in the attention block 1200. The SoftMax layer 1315 may be the same or similar as the SoftMax layer 1230 in the attention block 1200. Although the concatenation layer 1313 may introduce an additional overhead, its computational complexity would be minimal. The output of the SoftMax layer 1315 is sliced into two parts in the token dimension. The sizes of the two parts are [1, TLP] and [1, THP+ 1], respectively. The part of size [1, TLP] may be attention scores.

[0127] The compressed value tensor VLPof size (TLP, dmodei) is loaded from the LP cache 1330 and is input into the group MatMul layer 1316, in which the compressed value tensoris group-matrix-multiplied with the first slice of the SoftMax output, resulting in an output of size [1, dmodei]. Dequantization may be fused into the group MatMul layer 1316. The second SoftMax output slice of size [1, THP+ 1] is input into the MatMul layer 1317, in which a MatMul operation is performed on the second SoftMax output slice and the concatenated high-precision value tensor regular matrix-multiplied with the concatenated high-precision value tensor VHPof size [THP+ 1, dmodel], resulting in an output of size [1, dmodei]. The MatMul layer 1317 may be the same or similar as the MatMul layer 1240 in FIG. 12.

[0128] The addition layer 1318 receives the output of the group MatMul layer 1316 and the output of the MatMul layer 1317. The addition layer 1318 may include an elementwise accumulation, which may be done along the channel dimension dmodel. The addition layer 1318 computes an output of size [1, dmodei], which may be the output of the adaptive attention block 1310. This output may then be processed in a subsequent linear layer. The usage of mixed-precision attention can allow for efficient utilization of cache and memory resources while maintaining the accuracy of the attention mechanism. The process seamlessly integrates high and low precision computations, ensuring that the model's performance is optimized without incurring the overhead of dequantization.

[0129] In addition or alternative to the multi-precision KV cache maintenance mechanism described above or the adaptive grouped precision attention mechanism described above, a periodic quantization optimization for KV cache may be used in transformer model inference. FIG. 14 illustrates a quantization flow in an MHA block with periodic quantization optimization, in accordance with various embodiments. The MHA block includes an adaptive attention block 1410, which may be an example of the adaptive attention block 1310 in FIGS. 13A and 13B.

[0130] In some embodiments, periodic quantization of the high-precision section of the KV cache may be performed at designated quantization intervals, qint, during cached inference phases. In some embodiments, qintTkeyand T™luemay each be set to multiple of 2 to ensure a simplified quantized inference flow in the MHA block without the need to quantize K and V at different intervals. In some embodiments, qtntkey= Tkey= 2 and qintvalue= Tgalue= 6. The periodic quantization flow may occur at instances when infid=int ~ 1- infid denotes the inference stage index. For instance, the index of theprefill inference stage is denoted as infid= 0, while the index of the second inference stage is denoted as infid= 1. The inference process continues till the last inference phase is complete. In some embodiments, the inference process may end when a stopping criterion is met, e.g., when the generated token is the end of a sentence.

[0131] As shown in FIG. 14, the adaptive attention block 1410 is associated with a KV cache 1420 a LP cache 1430 and a HP cache 1440, a concatenator 1450, another concatenator 1460, a quantization module 1470, and another quantization module 1480. The LP cache 1430 may store a low-precision key tensor denoted as KLP. the shape of the key tensor is represented by [TLP,dmodei], which indicates the token dimension and the channel dimension of the key tensor. The operation of the adaptive attention block 1410 and the loading of the LP cache 1430 from memory may proceed in the same or similar manner as the operation of adaptive attention block 1310 in FIGS. 13A and 13B. The adaptive attention block 1410 in FIG. 14 receives QHP> KLP, LP> and concatenated KHPand VHP. Change may occur with the HP cache 1440. During the selected inference stage of index infid= qint— 1, the concatenated KHPand VHPof size [THP+ 1, dmodel] may pass through their quantization modules 1470 and 1480, respectively. The quantization modules 1470 and 1480 may apply the appropriate dual-group quantization configuration for K and value tensors separately. In some embodiments, the dual-group quantization configuration may remain the same across the prefill and cached inference phase. Subsequently, the resulting compressed K and value tensors are saved to the LP cache 1430, unlike the other inference stages. This process mirrors the quantization approach used during the prefill phase, like the quantization approach described above in conjunction with FIG. 12. During an unselected inference stage, the quantization modules 1470 and 1480 may be bypassed and no quantization on KV tensors may be performed.

[0132] The advantage of the periodic approach is that it allows for the periodic compression of the high-precision KV cache, reducing its size and thus optimizing memory usage. By directly quantizing the concatenated high-precision K and value tensor, there can be a streamlined workflow that avoids the need for dequantization, as the high-precision tensors are still utilized in high-precision MatMul operations in the adaptive attention block 1410. This can ensure that the attention mechanism continues to operate with high accuracy, leveraging the most recent and relevant context without incurring additional computational overhead. In essence, the quantization intervals serve as checkpoints where the modeladaptively manages the KV cache precision, ensuring that the adaptive attention block 1410 operates optimally.

[0133] Table 1 below shows the evolution of the mixed-precision KV cache across a series of inference stages, demonstrating the dynamic management of cache sizes and the efficacy of our quantization strategy. Organized by inference phase (infd), Table 1 provides a snapshot of the cache at each step of the model's autoregressive generation, with columns indicating the sizes of the low and high-precision K and value tensors post-inference. For the purpose of illustration and simplicity, Tgey= 2, Tgalue= 6, qint= max(2,6) = 6 in the embodiments of Table 1. Initially, at infid= 0, the low-precision KLPand VLPare established. No high-precision KV tensors may be established at infid= 0. As inferences progress (infid= 1 to 5), the high-precision KHPand VHPtensors grow, capturing the most recent tokens in memory. At the quantization interval (infid= n x qint= 6), where n = 1, 2, ... is an integer, the value tensor reaches a group boundary, allowing for full quantization and storage in memory. Simultaneously, accumulated groups in the key tensor are quantized and saved, exemplifying the strategic balance between high precision for recent context and efficient compression of older data.Table 1

[0134] Table 1 above shows that the subsequent inferences after quantization (infid= n x qint+ 1) may proceed without high-precision KV cache, streamlining the adaptive attention block by eliminating the need for concatenation operations when THP= 0. Quantization may occur for K / value tensors with a smaller Tgwhen sufficient high-precisiontokens align with a group boundary, further optimizing the process. In such cases, the adaptive attention block may adapt to handle varying group sizes, maintaining the model's performance and resource efficiency.

[0135] FIG. 15 illustrates mixed-precision, hybrid dual-group quantization of attention tensors, in accordance with various embodiments. FIG. 15 shows two attention tensors: a key tensor 1500 and a value tensor 1550, which may be used in the same MHA block. For the purpose of illustration and simplicity, each of the key tensor 1500 and value tensor 1550 is a 2D tensor and has a size IseqX dmodel. In other embodiments, the key tensor 1500 or value tensor 1550 may be a 3D tensor.

[0136] The key tensor 1500 and value tensor 1550 are quantized using a hybrid dual-group quantization method as the dual-group quantization configurations for the key tensor 1500 and value tensor 1550 are different. As shown in FIG. 15, the key tensor 1500 is segmented into groups 1510A (individually referred to as "group 1510A") and another group 1510B. The groups 1510A and 1510B are collectively referred to as "group 1510" or "groups 1510." Each group 1510A has a spatial shape Tgx Cg, where Tg= 3 and Cg= 2. The group 1510B has a spatial shape THPX dmodel. The value tensor 1550 is segmented into groups 1560A (individually referred to as "group 1560A") and another group 1560B. The groups 1560A and 1560B are collectively referred to as "group 1560" or "groups 1560." Each group 1560A has a spatial shape TgX Cg, where Tg= 2 and Cg= 3. The group 1510B has a spatial shape THp x dmodei. The key tensor 1500 is quantized using a token-heavy method, while the value tensor 1550 is quantized using a channel-heavy method.

[0137] In some embodiments of hybrid dual-group quantization, all the low-precision groups in an attention tensor may be quantized to the same precision. For instance, the groups 1510A (or the groups 1560A) may have the same precision, which is lower than the precision of the group 1510B (or the group 1560B), even though the groups 1510A (or the groups 1560A) may have different s and z parameters. In other embodiments (e.g., embodiments of FIG. 15), the low-precision groups may have different precisions, i.e., mixed-precision. The groups 1510A constitute a mixed-precision subtensor of size TMPX dmodeiwithin the key tensor 1500. The groups 1560A constitute a mixed-precision subtensor of size TMPX dmodeiwithin the value tensor 1550.

[0138] The method of hybrid dual-group quantization with mixed-precision can enhance the flexibility and efficiency of quantization by allowing different groups within the KV cache to be quantized at varying precision levels, such as FP16, FP8, FP4, INT8, INT4, INT2, and so on. Each group may be assigned its own scale and zero-point values, which are determined by a precision-based quantization function tailored to the specific precision level of the group. With mixed-precision, the KV tensors can be quantized in a more granular manner. The precision of each group may be selected based on the group's data distribution and the model's performance requirements, leading to a more customized and potentially more efficient quantization scheme. This approach can ensure that when KV tensors exhibit a high number of outliers, the majority can be quantized at lower precision, reserving higher precision for outlier groups. Consequently, this maintains LLM accuracy even with high degree of KV cache compression and eliminates the need for complex outlier segregation and tensor rearrangement in outlier quantization.

[0139] The quantization flows shown in FIGS. 12-14 may be applicable to the mixed- precision, hybrid dual-group quantization method. The mixed-precision, hybrid dual-group quantization method can facilitate a more nuanced control over the trade-off between model accuracy and computational resource usage. By applying different precision levels to different groups, the model's performance can be optimized for a variety of hardware configurations and operational constraints, ensuring that the LLM remains both powerful and practical for real-world applications.

[0140] FIG. 16 illustrates a group MatMul layer implemented on a data processing unit 1610 and a digital signal processor 1620, in accordance with various embodiments. An example of the data processing unit 1610 may be the data processing unit 1730 in FIG. 17. An example of the digital signal processor 1620 may be the digital signal processor 1735 in FIG. 17. For the purpose of illustration, the group MatMul layer receives an attention tensor 1601 including three groups 1602A-1602C as an input. In some embodiments, the groups 1602A- 1602C may have different quantization parameters. The group MatMul layer may include group MatMul operations, dequantization, and accumulation, which may be allocated to the data processing unit 1610 and digital signal processor 1620.

[0141] FIG. 16 shows a cycle of operations inside the data processing unit 1610. First, the group 1602A is loaded, e.g., from a local memory of the data processing unit 1610 to a processing engine in the data processing unit 1610. Second, a group MatMul operation isperformed on the group 1602A to compute a partial sum, e.g., by the processing engine. Third, a scale and bias are applied on the partial sum to dequantize the partial size, e.g., by a post-processing engine in the data processing unit 1610. The dequantized partial sum is drained, e.g., from the post-processing engine to the local memory. Fourth, the partial sums are stored in the local memory. The cycle may be repeated for the group 1602B, then the group 1602C, and lastly the group 1602D. The partial sums computed from the groups 1602A-1602C may be read from the local memory into the digital signal processorl620. The digital signal processor 1620 may accumulate the partial sums and generate an output of the group MatMul layer, as shown in FIG. 16.Example DNN System

[0142] FIG. 17 is a block diagram of a DNN system 1700, in accordance with various embodiments. The whole DNN system 1700 or a part of the DNN system 1700 may be implemented in one or more computing devices, such as the computing device 2500 in FIG. 25. The DNN system 1700 can generate and execute DNNs, such as transformer-based models (e.g., the transformer models described above), convolution-based models, and so on. As shown in FIG. 17, the DNN system 1700 includes a DNN module 1701 and a DNN accelerator 1702. In other embodiments, alternative configurations, different or additional components may be included in the DNN system 1700. For instance, the DNN system 1700 may include multiple DNN modules or multiple DNN accelerators. Further, functionality attributed to a component of the DNN system 1700 may be accomplished by a different component included in the DNN system 1700 or a different system. In some embodiments, the DNN module 1701 and DNN accelerator 1702 may include different types of processing units. In an example, the DNN module 1701 may be implemented by one or more central processing units (CPUs) or graphics processing units (GPUs). The DNN accelerator 1702 may also be referred to as a neural processing unit, Al accelerator, or Al processor. The DNN module 1701 and DNN accelerator 1702 may be implemented in the same chip or separate chips.

[0143] The DNN module 1701 facilitates generation and deployment of DNNs. In some embodiments, the DNN module 1701 may generate and train DNNs. For instance, the DNN module 1701 can define the layered architecture of a DNN. The DNN module 1701 can also determine the internal parameters of the DNN through a DNN training process. The DNN module 1701 may also determine one or more hyperparameters that define how the DNN istrained. An example hyperparameter is a sparsity ratio that defines the sparsity level of one or more deep learning tensors for the DNN.

[0144] The DNN module 1701 may also compress DNNs, e.g., during or after training. In some embodiments, the DNN module 1701 may compress KV tensors in a transformer model and generate compressed KV cache. For instance, the DNN module 1701 may use various dual-group quantization methods to quantize KV tensors. In some embodiments, the DNN module 1701 may prune weights in one or more layers of a DNN by changing nonzero valued weight to zeros. The DNN module 1701 may prune weights based on a target weight sparsity ratio. A weight sparsity ratio may be the ratio of the number of zero-valued weights to the total number of weights. In an example where the DNN module 1701 prunes weight during DNN training, the DNN module 1701 may prune weight of a layer to achieve a target sparsity ratio after one or more epochs. The DNN module 1701 may prevent the pruned weights from changing values during the rest of the training process. Alternatively, the DNN module 1701 may allow the pruned weights to change values so that a pruned, zero-valued weight may have a nonzero value after further training. The DNN module 1701 may prune weights of the layer again after one or more additional epochs.

[0145] The DNN module 1701 may deploy trained, compressed, or validated DNNs for use in deep learning applications. In some embodiments, the DNN module 1701 may distribute trained, compressed, or validated DNNs to devices or systems which may use the DNNs to perform tasks (e.g., image classification, motion planning, etc.) for which the DNNs were trained. In other embodiments, the DNN module 1701 may facilitate deployment of the DNNs using the DNN accelerator 1702. For instance, the DNN module 1701 may receive data from a device or system coupled with the DNN system 1700 and input the received data (or data generated by the DNN module 1701, e.g., based on the received data) into a DNN. The DNN module 1701 may generate instructions (e.g., configuration files) that control the operation of the DNN accelerator 1702 during the DNN execution. The DNN module 1701 may receive an output of the DNN from the DNN accelerator 1702. The DNN module 1701 may transmit the output of the DNN (or a result of processing the output of the DNN by the DNN module 1701) to the device or system. In some embodiments, the DNN module 1701 may control execution processes of trained, compressed, or validated DNNs. The DNN module 1701 may function as a compiler for DNNs executed by the DNN accelerator 1702.The DNN module 1701 may perform compilation of DNNs and generate compilation descriptors, based on which the DNNs may be executed.

[0146] The DNN module 1701 may generate executable transformer models. The DNN module 1701 may also facilitate execution of transformer models by the DNN accelerator 1702. For instance, the DNN module 1701 may be the host for the execution of neural network operations in transformer models by the DNN accelerator 1702, e.g., the host for transformer model inference. In some embodiments, the DNN module 1701 may receive an inference request, which may be a request to have a transformer model to make a prediction based on input data. The DNN module 1701 may facilitate cached inference of the transformer model, in which attention tensors (e.g., key tensors, value tensor, etc.) may be cached and reused in the inference of the transformer model. The inference for making the prediction may include a sequence of inference stages, which generates a sequence of predicted tokens. The sequence of predicted tokens may be the prediction of the transformer model. In some embodiments, the DNN module 1701 may facilitate generation and execution of transformer models with static shaped caches and inputs. Certain aspects of the DNN module 1701 are provided below in conjunction with FIG. 18.

[0147] The DNN accelerator 1702 executes DNNs provided by the DNN module 1701. For instance, the DNN accelerator 1702 can execute a DNN by running deep learning operations in the DNN. The process of carrying out a deep learning operation is also referred to as a process of executing the deep learning operation or a process of performing the deep learning operation. The execution of the DNN may be for training the DNN or for using the DNN to perform Al tasks. As shown in FIG. 17, the DNN accelerator 1702 includes a memory 1710, a DMA (direct memory access) engine 1720, data processing units 1730 (individually referred to as "data processing unit 1730"), and a digital signal processor 1735. In other embodiments, alternative configurations, different or additional components may be included in the DNN accelerator 1702. For example, the DNN accelerator 1702 may include more than one memory 1710, DMA engine 1720, or digital signal processor 1735. As another example, the DNN accelerator 1702 may include a single data processing unit 1730. Further, functionality attributed to a component of the DNN accelerator 1702 may be accomplished by a different component included in the DNN accelerator 1702 or by a different system. A component of the DNN accelerator 1702 may be implemented in hardware, software, firmware, or some combination thereof.

[0148] The memory 1710 stores data associated with deep learning operations performed by the DNN accelerator 1702. The memory 1710 may be a system memory. In some embodiments, the memory 1710 includes a DRAM. When the DNN accelerator 1702 executes operations in transformer models, at least part of the memory 1710 may be used to implement KV caches, such as self-attention KV caches and cross-attention KV caches, in the transformer models. The KV caches may be updated during inference of the transformation model. Layout of data in the KV caches may be determined to optimize the efficiency of the DNN accelerator 1701. Some KV caches may include compressed KV tensors, such as KV tensors generated by quantizing keys and values with dual-group quantization methods.

[0149] In some embodiments, the memory 1710 may store data to be used by the data processing units 1730 for DNN execution. The memory 1710 may store weights, such as weights of convolutional layers, which are determined by training DNNs. The memory 1710 may further store inputs to DNN layers or outputs of DNN layers, such as data generated by the data processing units 1730 from performing deep learning operations in DNNs. Example deep learning operations include convolutions (also referred to as "convolutional operations"), layer normalization operations, SoftMax operations, matrix multiplication operations, pooling operations, elementwise operations, activation functions, other types of deep learning operations, or some combination thereof. The memory 1710 may be a main memory of the DNN accelerator 1702. In some embodiments, the memory 1710 includes one or more DRAMs.

[0150] The DMA engine 1720 facilitates data transfer between the memory 1710 and local memories of the data processing units 1730. For example, the DMA engine 1720 can read data from the memory 1710 and write data into a local memory of a data processing unit 1730. As another example, the DMA engine 1720 can read data from a local memory of a data processing unit 1730and write data into the memory 1710. The DMA engine 1720 provides a DMA feature that allows the data processing unit 1730 to initiate data transfer between the memory 1710 and the local memories of the data processing units 1730 and to perform other operations while the data transfer is being conducted. In some embodiments, the DMA engine 1720 may read tensors from the memory 1710, modify the tensors in a way that is optimized for the data processing unit 1730 before it writes the tensors into the local memories of the data processing units 1730.

[0151] The data processing units 1730 perform deep learning operations in DNNs. For instance, a data processing unit 1730 may execute a DNN layer by running one or more deep learning operations in the DNN layer. A data processing unit 1730 may execute a layer, or a portion of a layer, at a time. In some embodiments, the operations of the DNN layers may be run by multiple data processing units 1730 in parallel. For instance, multiple data processing units 1730 may each perform a portion of a workload for a deep learning operation. Data may be shared between the data processing units 1730. A data processing unit 1730 may also be referred to as a neural processing unit, a compute block, or a compute tile.

[0152] The data processing units 1730 may be capable of running various types of deep learning operations, such as convolution, layer normalization, SoftMax operation, pooling, elementwise operation, linear operation, nonlinear operation, and so on. Deep learning operations performed by the data processing units 1730 include tensor operations, i.e., operations whose inputs are tensors or operations whose outputs are tensors. In an example, the data processing unit 1730 receives an input tensor and one or more convolutional kernels and performs a convolution with the input tensor and convolutional kernels. The result of the convolution may be an output tensor, which can be further computed, e.g., by the data processing unit 1730 or another data processing unit 1730.

[0153] In the embodiments of FIG. 17, each data processing unit 1730 includes a local memory 1740, a sparsity mode module 1750, a load module 1760, a processing engine 1770, a post-processing engine 1780, and a drain module 1790. Some or all the components of the data processing unit 1730 can be implemented on the same chip. In other embodiments, alternative configurations, different or additional components may be included in the data processing unit 1730. Further, functionality attributed to a component of the data processing unit 1730 may be accomplished by a different component included in the data processing unit 1730, a different data processing unit 1730, another component of the DNN accelerator 1702, or a different system. A component of the data processing unit 1730 may be implemented in hardware, software, firmware, or some combination thereof.

[0154] The local memory 1740 is local to the corresponding data processing unit 1730. In the embodiments of FIG. 17, the local memory 1740 is inside the data processing unit 1730. In other embodiments, the local memory 1740 may be outside the data processing unit 1730. Data in the local memory 1740 may be transferred to or from the memory 1710, e.g.,through the DMA engine 1720. For instance, KV caches may be copied from the memory 1710 to the local memory 1740. In some embodiments, data in the local memory 1740 may be transferred to or from the local memory of another data processing unit 1730. The local memory 1740 may store data received, used, or generated by the sparsity mode module 1750, the load module 1760, the processing engine 1770, the post-processing engine 1780, or the drain module 1790. Examples of the data may include input activations, weights, output activations, sparsity bitmaps, and so on.

[0155] In some embodiments, the local memory 1740 may store tensors to be processed by the processing engine 1770 or the post-processing engine 1780. The tensors may be input tensors of deep learning operations. The local memory 1740 may also store tensors generated by the processing engine 1770 or the post-processing engine 1780. The tensors may be output tensors of deep learning operations. In some embodiments, the local memory 1740 may store dense tensors (e.g., dense activation tensors, dense weight tensors, etc.), sparse tensors (e.g., sparse activation tensors, sparse weight tensors, etc.), and so on. A dense tensor may be a tensor from which zero-valued elements (if any) are not removed. A dense tensor may be converted to a sparse tensor by removing one or more zero-valued elements in the dense tensor. A sparse tensor may also be referred to as a compressed tensor or packed tensor. The process of converting a dense tensor to a sparse tensor may be referred to as sparsity encoding. Sparsity encoding may also generate a sparsity tensor. Each element in the sparsity tensor may correspond to a different element in the dense tensor and indicate whether the element in the dense tensor is zero or not. The sparsity tensor may indicate positions of elements of the sparse tensor in the dense tensor. The sparsity tensor may be a sparsity bitmap, each element of which is a bit. A sparse tensor may be converted to a dense tensor through a densifying process, in which one or more zeros may be added to the sparse tensor based on the sparsity tensor.

[0156] In some embodiments, the local memory 1740 includes one or more SRAMs. The local memory 1740 may be byte-addressable, and each memory address identifies a single byte (eight bits) of storage. In some embodiments, the local memory 1740 may include memory banks. The number of data banks in the local memory 1740 may be 16, 64, 128, 1756, 512, 1024, 2048, or other numbers. A memory bank may include a plurality of storage units. In an example, a data bank may include 8, 16, 64, or a different number of storage units. A memory bank or a storage unit in a memory bank may have a memoryaddress. In an example, a storage unit may store a single byte, and data larger than a single byte may be stored in storage units with consecutive memory addresses, i.e., adjacent storage units. For instance, a storage unit can store an integer number in the INT8 format, versus two storage units may be needed to store a number in the FP16 or BF16 format, which has 16 bits. In some embodiments, 16 bits can be transferred from the local memory 1740 in a single read cycle. In other embodiments, 16 bits can be transferred from the local memory 1740 in multiple read cycles, such as two cycles.

[0157] The sparsity mode module 1750 determines sparsity modes in which the data processing unit 1730 operates to execute DNN layers. For instance, the sparsity mode module 1750 may determine whether to accelerate a layer based on weight sparsity, activation sparsity, or both. The sparsity mode module 1750 select the sparsity mode for a layer from a group of sparsity modes that includes, for example, combined sparsity mode in which the layer is accelerated based on both weight sparsity and activation sparsity, activation sparsity mode in which the layer is accelerated based on activation sparsity but not based on weight sparsity, weight sparsity mode in which the layer is accelerated based on weight sparsity but not based on activation sparsity, and a dense mode in which the layer is not accelerated based on sparsity. In some embodiments (e.g., embodiments where a layer is executed by multiple data processing units 1730), the sparsity mode module 1750 may determine the sparsity mode for all the data processing units 1730 that executes the layer. In some embodiments, the sparsity mode module 1750 may receive configuration parameters from the DNN module 1701. A configuration parameter may correspond to a layer and indicate whether to accelerate the layer based on weight sparsity. The sparsity mode module 1750 may determine the sparsity mode of the layer based on the configuration parameter.

[0158] The load module 1760 loads data from the local memory 1740 to the processing engine 1770 or to the post-processing engine 1780. The load module 1760 may read tensors from the local memory 1740. The tensors may include sparse activation tensors, sparse weight tensors, activation sparsity tensors, weight sparsity tensors, and so on. In some embodiments, the load module 1760 may load data based on the sparsity mode determined by the sparsity mode module 1750. The load module 1760 may select different data to transmit to the processing engine 1770 in different sparsity modes. For instance, the load module 1760 may transmit an activation sparsity tensor and a weight sparsity tensor of alayer to the processing engine 1770 in the combined sparsity mode, while transmit the activation sparsity tensor but not the weight sparsity tensor to the processing engine 1770 in the activation sparsity mode and transmit the weight sparsity tensor but not the activation sparsity tensor to the processing engine 1770 in the weight sparsity mode. In the dense mode, the load module 1760 does not transmit either the activation sparsity tensor or the weight sparsity tensor to the processing engine 1770.

[0159] In some embodiments, the load module 1760 may process (e.g., densify) data stored in the local memory 1740 before providing the data to the processing engine 1770. In an example, the load module 1760, while operating in the weight sparsity mode, may densify sparse activation tensors to generate dense activation tensors based on corresponding activation sparsity tensors. For instance, the load module 1760 may add one or more zeros into a sparse activation tensor based on an activation sparsity tensor associated with the sparse activation tensor to generate the dense activation tensor. The dense activation tensor includes one or more elements than the sparse activation tensor. The additional element(s) are zero-valued. The load module 1760 may identify one or more elements in the activation sparsity tensor that correspond to the zero-valued element(s), determine the position of each of the zero-valued element(s) in the dense activation tensor, and insert the zero-valued element(s) into the sparse activation tensor based on the determined positions. After the densification, the load module 1760 may transmit the dense activation tensors to the processing engine 1770. The load module 1760 may also transmit corresponding sparse weight tensors and weight sparsity tensors to the processing engine 1770. Activation sparsity tensor of the dense activation tensors may not be loaded to the processing engine 1770.

[0160] In another example, the load module 1760, while operating in the activation sparsity mode, may densify sparse weight tensors to generate dense weight tensors based on corresponding weight sparsity tensors by inserting zeros into sparse weight tensors. The densification of sparse weight tensors may be similar to the densification of sparse activation tensors described above. After the densification, the load module 1760 may transmit the dense weight tensors to the processing engine 1770. The load module 1760 may also transmit corresponding sparse activation tensors and activation sparsity tensors to the processing engine 1770. Weight sparsity tensor of the dense weight tensors may not be loaded to the processing engine 1770. In yet another example, the load module 1760, whileoperating in the dense mode, may densify both sparse weight tensors and sparse activation tensors. The load module 1760 may generate the input tensor and weight tensor of the layer and transmit the tensors to the processing engine 1770 for executing the layer without sparsity acceleration.

[0161] The processing engine 1770 performs operations in DNNs. The processing engine 1770 may accelerate neural network operations based on sparsity in data. In some embodiments, the processing engine 1770 may operate in a dense mode in which sparsity acceleration is not performed. The processing engine 1770 may include one or more processing cells. In some embodiments, the processing cells may be arranged in one or more rows and one or more columns in the processing engine 1770. Each processing cell may include PEs that may be arranged in an array that includes rows and columns. All the PEs in the processing engine 1770 may constitute a bigger array that includes more rows and columns.

[0162] An example PE may be or may include one or more multiply-accumulate (MAC) units that can perform MAC operations. In some embodiments (e.g., embodiments where the data processing unit 1730 executes a convolutional layer), a computation in an MAC unit may be an MAC operation on an activation operand and a weight operand. The activation operand may be an activation tensor that may include one or more activations in the input tensor of the convolution. Different activations may be in different input channels. The weight operand may be a weight tensor that may include one or more weights in the filter of the convolution. The values of the weights are determined through training the DNN. The weights in the weight operand may be in different input channels.

[0163] In some embodiments, an MAC unit includes one or more multipliers for performing multiplications. An MAC unit may also include one or more accumulators ("adders") for performing accumulations. A column of MAC units is referred to as an MAC column. An MAC column may be associated with one or more MAC lanes. A MAC lane is a path for loading data e.g., by the load module 1760, into an MAC column. A MAC lane may be also referred to as a data transmission lane or data loading lane. An MAC column may have multiple MAC lanes. The loading bandwidth of the MAC column is an aggregation of the loading bandwidths of all the MAC lanes associated with the MAC column. With a certain number of MAC lanes, data can be fed into the same number of independent MAC units simultaneously. In some embodiments where an MAC column has four MAC lanes forfeeding activations or weights into the MAC column and each MAC lane may have a bandwidth of 16 bytes, the four MAC lanes can have a total loading bandwidth of 64 bytes.

[0164] In some embodiments, the processing engine 1770 may be capable of depthwise convolution, standard convolution, or both. In a depthwise convolution, an MAC unit may perform an MAC operation that includes a sequence of multiplications for an input operand and a weight operand. Each multiplication in the sequence (also referred to as a cycle) is a multiplication of a different activation in the input operand with a different weight in the weight operand. The activation and weight in the same cycle may correspond to the same channel. The sequence of multiplication produces a product operand that includes a sequence of products. The MAC operation may also include accumulations in which multiple product operands are accumulated to produce an output operand of the MAC unit. The processing engine 1770 may output multiple output operands at a time, each of which is generated by a different MAC unit. In a standard convolution, MAC operations may include accumulations across the channels. For instance, as opposed to generating an output operand, a MAC unit may accumulate products across different channels to generate a single output point.

[0165] In some embodiments, the processing engine 1770 may perform MAC operations in quantized deep learning operations, such as MAC operations in a quantized convolution. In some embodiments, an MAC unit in the processing engine 1770 may receive quantized activation and quantized weights and compute a quantized MAC result. The quantized MAC result may be a quantized value in an integer format and may be the output of the MAC unit. In some embodiments, the MAC unit may also include a quantization multiplier that can multiply a quantization scale with the quantized MAC result, and the output of the MAC unit may be a real value in a floating-point format. The MAC unit may include no quantization subtractors as zero-point offsetting is not needed for the MAC operations in quantized deep learning operations.

[0166] In some embodiments, the processing engine 1770 may include sparsity acceleration logic for facilitating sparsity acceleration. For instance, each processing cell in the processing engine 1770 may include one or more sparsity modules. In an example, each MAC column or each MAC row may have a corresponding sparsity module that accelerates MAC operations in the MAC column or MAC row. In some embodiments, a sparsity module accelerates computations in the processing engine 1770 based on sparsity in activations,sparsity in weights, or both. The sparsity module may include a storage unit that stores a sparsity tensor, which may be loaded to the storage unit by the load module 1760. The sparsity tensor may be an activation sparsity tensor, a weight sparsity tensor, or a combined sparsity tensor.

[0167] An activation sparsity tensor may be the sparsity tensor of an activation tensor and has the same number of elements as the activation tensor. An element in the activation sparsity tensor may indicate whether the corresponding element in the activation tensor is zero or not. For instance, a zero-valued in the activation sparsity tensor may indicate that the corresponding element in the activation tensor is zero. A one-valued in the activation sparsity tensor may indicate that the corresponding element in the activation tensor is nonzero. A weight sparsity tensor may be the sparsity tensor of a weight tensor and has the same number of elements as the weight tensor. An element in the weight sparsity tensor may indicate whether the corresponding element in the weight tensor is zero or not. For instance, a zero-valued in the weight sparsity tensor may indicate that the corresponding element in the weight tensor is zero. A one-valued in the weight sparsity tensor may indicate that the corresponding element in the weight tensor is nonzero. The sparsity module may generate a combined sparsity tensor using an activation sparsity tensor and a weight sparsity tensor. For instance, the sparsity module may multiply an element of the activation sparsity tensor with a corresponding element of the weight sparsity tensor to compute an element of the combined sparsity tensor. The positions of the three elements in their corresponding sparsity tensors may match. In some embodiments, each element in a sparsity tensor may be a bit, and the sparsity tensor may be referred to as a sparsity bitmap.

[0168] The sparsity module may use the sparsity tensor to identify activations and weights to be used in MAC operations by the MAC units. In an embodiment where the processing engine 1770 operates in the combined sparsity mode, the sparsity module may identify activations and weights that correspond to nonzero valued elements of a combined sparsity tensor. In an embodiment where the processing engine 1770 operates in the activation sparsity mode, the sparsity module may identify activations and weights that correspond to nonzero valued elements of an activation sparsity tensor. In an embodiment where the processing engine 1770 operates in the weight sparsity mode, the sparsity module may identify activations and weights that correspond to nonzero valued elements of a weightsparsity tensor. The sparsity module may be bypassed in the dense mode as no sparsity acceleration would be conducted.

[0169] The post-processing engine 1780 processes outputs of the processing engine 1770. The post-processing engine 1780 may include one or more post-processing elements. In some embodiments, the post-processing elements in the post-processing engine 1780 may be arranged in an array that has rows and columns. In some embodiments, the postprocessing engine 1780 computes activation functions. The post-processing engine 1780 may receive outputs of the processing engine 1770 as inputs to the activation functions. In addition or alternative to activation functions, the post-processing engine 1780 may perform other types of post processing on outputs of the processing engine 1770. For instance, the post-processing engine 1780 may apply a bias on an output of the processing engine 1770. In some embodiments, the post-processing engine 1780 may be bypassed for certain neural network operations.

[0170] The drain module 1790 drains data from the processing engine 1770 or from the post-processing engine 1780. The drain module may write the data to the local memory 1740. The drained data may be tensors, such as output tensors of neural network operations. In some embodiments, the drain module 1790 may drain data on a cell level. For each processing cell, the drain module 1790 may drain outputs of PEs in the processing cell based on a row index or column index of each PE. For instance, the drain module 1790 may use a sequence of cycles to drain data from a processing cell. The drain module 1790 may drain the output of some of the PE s in each cycle. The sequence of the cycles may be configured based on a configuration parameter indicating the operation mode of the load module 1760.

[0171] In some embodiments, the drain module 1790 includes sparsity encoding logic that can convert outputs of the processing engine 1770 from a dense format to a sparse format. For instance, the drain module 1790 may be implemented with one or more sparsity encoders. A sparsity encoder converts dense data to compressed data based on sparsity in the dense data. For instance, the sparsity encoder may remove zeros in an activation tensor computed by the processing engine 1770 to convert the activation tensor to a compressed activation tensor. The sparsity encoder may also generate sparsity tensors, including activation sparsity tensors.

[0172] In some embodiments, the data drained from the processing engine 1770 may be at least part of an output tensor (e.g., the output tensor 2030 in FIG. 20) of a deep learning operation. The sparsity encoder may generate a compressed version of the output tensor. The sparsity encoder may identify every zero-valued activation in the output tensor and remove these activations from the output tensor to generate a compressed activation tensor (aka "sparse activation tensor"). The sparsity encoder may also generate one or more sparsity tensors for the output tensor. A sparsity tensor may correspond to a portion of the output tensor (e.g., the vector 2035 in FIG. 20). The sparsity tensor may include sparsity elements (e.g., bits), each of which corresponds to a different activation in the vector and indicates whether the corresponding activation is zeroed or not.

[0173] The drain module 1790 may write the compressed activation tensor and the one or more sparsity tensors into the local memory 1740. The sparse activation tensor and the one or more sparsity tensors may be further loaded to the memory 1710, e.g., through the DMA engine 1720. Additionally or alternatively, the sparse activation tensor and the one or more sparsity tensors may be loaded by the load module 1760 to the processing engine 1770 for further computation, e.g., for performing a deep learning operation in the next layer.

[0174] The digital signal processor 1735 is another compute unit in the DNN accelerator 1702 in addition to the data processing units 1730. The digital signal processor 1735 may have an architecture optimized for the operational needs of digital signal processing. The digital signal processor 1735 may be implemented on a microprocessor chip, which may be separate from a chip implementing at least one of the data processing units 1730. In some embodiments, the digital signal processor 1735 may be a Streaming Hybrid Architecture Vector Engine (SHAVE) DSP. In some embodiments, the digital signal processor 1735 can be pipelined along with the inference pipeline and the DMA engine 1720, thereby enabling parallel computations to improve overall performance. In some embodiments, the digital signal processor 1735 may perform some computations in a neural network operation, while other computations in the neural network operation may be performed by the processing engine 1770 or the post-processing engine 1780. For instance, the digital signal processor 1735 may accumulate partial sums computed from group MatMul operations performed by the data processing units 1730.

[0175] FIG. 18 is a block diagram of a DNN module 1800, in accordance with various embodiments. The DNN module 1800 may be an embodiment of the DNN module 1701 inFIG. 17. As shown in FIG. 18, the DNN module 1800 includes an interface module 1810, a training module 1820, a compressing module 1830, a compiler 1840, a transformer execution module 1850, and a datastore 1860. In other embodiments, alternative configurations, different or additional components may be included in the DNN module 1800. Further, functionality attributed to a component of the DNN module 1800 may be accomplished by a different component included in the DNN module 1800 or a different module or system.

[0176] The interface module 1810 facilitates communications of the DNN module 1800 with other modules or systems. For example, the interface module 1810 establishes communications between the DNN module 1800 with an external database to receive data that can be used to train DNNs or input into DNNs to perform tasks. As another example, the interface module 1810 supports the DNN module 1800 to distribute DNNs to other systems, e.g., computing devices configured to apply DNNs to perform tasks. The interface module 1810 may receive inference requests from users of transformer models.

[0177] The training module 1820 trains DNNs by using a training dataset. The training module 1820 forms the training dataset. In an example where the training module 1820 trains an DNN to recognize objects in images, the training dataset includes training images and training labels. The training labels describe ground-truth classifications of objects in the training images. In some embodiments, each label in the training dataset corresponds to an object in a training image. In some embodiments, a part of the training dataset may be used to initially train the DNN, and the rest of the training dataset may be held back as a validation subset used by the training module 1820 to validate performance of a trained DNN. The portion of the training dataset not including the tuning subset and the validation subset may be used to train the DNN.

[0178] The training module 1820 also determines hyperparameters for training the DNN. Hyperparameters are variables specifying the DNN training process. Hyperparameters are different from parameters inside the DNN (e.g., weights of filters). In some embodiments, hyperparameters include variables determining the architecture of the DNN, such as number of hidden layers, etc. Hyperparameters also include variables which determine how the DNN is trained, such as batch size, number of epochs, etc. A batch size defines the number of training samples to work through before updating the parameters of the DNN. The batch size is the same as or smallerthan the number of samples in the training dataset.The training dataset can be divided into one or more batches. The number of epochs defines how many times the entire training dataset is passed forward and backwards through the entire network. The number of epochs defines the number of times that the deep learning algorithm works through the entire training dataset. One epoch means that each training sample in the training dataset has had an opportunity to update the parameters inside the DNN. An epoch may include one or more batches. The number of epochs may be 1, 5, 10, 50, 100, 500, 1000, or even larger.

[0179] The training module 1820 defines the architecture of the DNN, e.g., based on some of the hyperparameters. The architecture of the DNN includes an input layer, an output layer, and a plurality of hidden layers. The input layer of an DNN may include tensors (e.g., a multidimensional array) specifying attributes of the input image, such as the height of the input image, the width of the input image, and the depth of the input image (e.g., the number of bits specifying the color of a pixel in the input image). The output layer includes labels of objects in the input layer. The hidden layers are layers between the input layer and output layer. The hidden layers include one or more convolutional layers and one or more other types of layers, such as pooling layers, fully-connected layers, normalization layers, SoftMax or logistic layers, and so on. The convolutional layers of the DNN abstract the input image to a feature map that is represented by a tensor specifying the feature map height, the feature map width, and the feature map channels (e.g., red, green, blue images include 3 channels). A pooling layer is used to reduce the spatial volume of input image after convolution. It is used between two convolution layers. A fully-connected layer involves weights, biases, and neurons. It connects neurons in one layer to neurons in another layer. It is used to classify images between different categories by training.

[0180] In the process of defining the architecture of the DNN, the training module 1820 also adds an activation function to a hidden layer or the output layer. An activation function of a layer transforms the weighted sum of the input of the layer to an output of the layer. The activation function may be, for example, a ReLU activation function, a tangent activation function, or other types of activation functions.

[0181] After the training module 1820 defines the architecture of the DNN, the training module 1820 inputs a training dataset into the DNN. The training dataset includes a plurality of training samples. An example of a training sample includes an object in an image and a ground-truth label of the object. The training module 1820 modifies the parameters insidethe DNN ("internal parameters of the DNN") to minimize the error between labels of the training objects that are generated by the DNN and the ground-truth labels of the objects. The internal parameters include weights of filters in the convolutional layers of the DNN. In some embodiments, the training module 1820 uses a cost function to minimize the error.

[0182] The training module 1820 may train the DNN for a predetermined number of epochs. The number of epochs is a hyperparameter that defines the number of times that the deep learning algorithm will work through the entire training dataset. One epoch means that each sample in the training dataset has had an opportunity to update internal parameters of the DNN. After the training module 1820 finishes the predetermined number of epochs, the training module 1820 may stop updating the parameters in the DNN. The DNN having the updated parameters is referred to as a trained DNN.

[0183] The training module 1820 may also verify accuracy of trained or compressed DNNs. In some embodiments, the training module 1820 inputs samples in a validation dataset into a trained DNN and uses the outputs of the DNN to determine the model accuracy. In some embodiments, a validation dataset may be formed of some or all the samples in the training dataset. Additionally or alternatively, the validation dataset includes additional samples, other than those in the training sets. In some embodiments, the training module 1820 may determine an accuracy score measuring the precision, recall, or a combination of precision and recall of the DNN. The training module 1820 may use the following metrics to determine the accuracy score: Precision = TP / (TP + FP) and Recall = TP / (TP + FN), where precision may be how many the DNN correctly predicted (TP or true positives) out of the total it predicted (TP + FP or false positives), and recall may be how many the DNN correctly predicted (TP) out of the total number of objects that did have the property in question (TP + FN or false negatives). The F-score (F-score = 2 * PR / (P + R)) unifies precision and recall into a single measure.

[0184] The training module 1820 may compare the accuracy score with a threshold score. In an example where the training module 1820 determines that the accuracy score of the DNN is less than the threshold score, the training module 1820 may re-train the DNN. In one embodiment, the training module 1820 may iteratively re-train the DNN until the occurrence of a stopping condition, such as the accuracy measurement indication that the DNN may be sufficiently accurate, or a number of training rounds having taken place.

[0185] The compressing module 1830 compresses DNNs. For instance, the compressing module 1830 may quantize attention tensors in transformer models. In some embodiments, the compressing module 1830 may use dual-group quantization methods to quantize attention tensors. The compressing module 1830 may determine a quantization configuration for an attention tensor (e.g., a key tensor or value tensor). The quantization configuration may be represented by quantization configuration parameters that indicate the token dimension and channel dimension of a group, i.e., group token dimension and group channel dimension. In some embodiments, the compressing module 1830 may use a hardware-software co-design methodology to select an optimal quantization configuration for the attention tensor. For instance, the compressing module 1830 may determine the quantization configuration parameters based on a constraint of the DNN accelerator 1702 which is to execute the attention block. The compressing module 1830 may also estimate accuracy of the transformer model if the quantization configuration parameters are used for compressing the attention tensor and when the estimated accuracy does not meet a criterion, the compressing module 1830 may adjust the quantization configuration parameters.

[0186] The compressing module 1830 may partition the attention tensor along both the token dimension and channel dimension using the quantization configuration parameters to generate a plurality of groups. The compressing module 1830 may generate a plurality of groups having the same dimension and one or more other groups that have different size from the plurality of groups. For instance, the compressing module 1830 may generate a plurality of smaller groups and a larger group by partitioning the attention tensor based on the quantization configuration parameter. The larger group may be larger along the channel dimension than the smaller groups. In some embodiments, the larger group may include keys or values corresponding to more recently generated tokens (e.g., tokens predicted in more recent inference stages) than the keys or values in the smaller groups. The compressing module 1830 may also determine group-specific quantization parameters, such as scaling factor or zero-point factor. Different groups may have different scaling values or zero-point values. In some embodiments, the compressing module 1830 may determine scaling values and zero-point values for the smaller groups and quantize the smaller groups separately. In some embodiments, the compressing module 1830 may not quantize the larger group so that the larger group.

[0187] In some embodiments, the compressing module 1830 may use a mixed precision approach by assigning different precisions to the groups and converting the original precision of the attention tensor to the different precisions for each group. Different groups may have different precision. In some embodiments (e.g., embodiments in which the attention tensor is segmented into a plurality of smaller groups and a larger group), the larger group has a higher precision than the smaller groups. In an example, the precision of the larger group may be the same as the original precision of the attention tensor, while the smaller groups has one or more lower data precisions. In some embodiments, all the smaller groups may have the same precision. In other embodiments, different ones of the smaller groups may have different precisions.

[0188] In some embodiments, the compressing module 1830 may use a hybrid dual-group quantization approach by using different quantization configurations for the key tensor and the value tensor of an attention block. For instance, the compressing module 1830 may apply token-heavy dual-group quantization for the key tensor so that the group token dimension is larger than the group channel dimension but apply channel-heavy dual-group quantization for the value tensor so that the group channel dimension is larger than the group token dimension, or vice versa. In some embodiments, the compressing module 1830 may use periodic quantization. For instance, the compressing module 1830 may not compress the KV cache in every inference stage. Rather, the compressing module 1830 sets up a quantization interval and compresses the KV cache in inference instances selected based on the quantization interval. The quantization interval may vary for key and value. The compressing module 1830 may use various combinations of these quantization approaches.

[0189] After the compressing module 1830 generates a compressed / quantized attention tensor, the compressing module 1830 may store the compressed / quantized attention tensor as a KV cache in a memory. An example of the memory may be the memory 1710 in FIG. 17. The KV cache, which is a compressed KV cache, may be read from the memory and used in the next inference stage.

[0190] The compiler 1840 compiles information of DNNs to executable instructions that can be executed, e.g., by the DNN accelerator 1702, to carry out neural network operations in DNNs. The compiler 1840 may generate executable DNNs, e.g., DNNs that can be at least partially executed by the DNN accelerator 1701. In some embodiments, the compiler 1805may generate a graph representing a DNN. The graph may include nodes and edges. A node may represent a specific neural network operation in the DNN. An edge may connect two nodes and represent a connection between the two corresponding neural network operations. In an example, an edge may encode a tensor that flows from one of the neural network operations to the other neural network operation. The tensor may be an output tensor of the first neural network operation and an input tensor of the second neural network operation. The edge may encode one or more attributes of the tensor, such as size, shape, storage format, and so on. The compiler 1840 may use the graph to generate instructions (e.g., compilation descriptors). The instructions would be executed by components of the DNN accelerator 1702 to execute the DNN. Functionality attributed to the compiler 1840 may be accomplished by the compressing module 1830 or transformer execution module 1850, or vice versa.

[0191] The transformer execution module 1850 facilitates execution of transformer models, including transformer models with compressed KV cache. The transformer execution module 1850 may map operations in transformer models to the DNN accelerator 1710, e.g., to the data processing units 1730 and digital signal processor 1735 in the DNN accelerator 1710. For instance, the transformer execution module 1850 may map MatMul operations and group MatMul operations in a transformer model to the data processing unit 1730 and digital signal processor 1735. The transformer execution module 1850 may modify transformer models based on attributes of hardware resources (e.g., memories, compute elements, etc.) of the DNN accelerator 1710. The transformer execution module 1850 may also determine parameters that will be used in inference of transformer models. In some embodiments, the transformer execution module 1850 may determine optimal data layouts in KV caches in transformer models. Memory footprint can be a challenge to both inference latency and energy consumption of transformer model inference on DNN accelerators.During cached inference, self-attention KV cache and cross-attention KV cache may undergo transpose operations before getting fed to MatMul operations inside MHA layers. For instance, the transformer execution module 1850 may transpose a key tensor before the key tensor is multiplied with a query tensor.

[0192] The datastore 1860 stores data received, generated, used, or otherwise associated with the DNN module 1800. For example, the datastore 1860 stores the datasets used by the training module 1820 and training module 1820. The datastore 1860 may also store datagenerated by the training module 1820, such as the hyperparameters for training DNNs, internal parameters of trained DNNs (e.g., weights, etc.), data for sparsity acceleration (e.g., sparsity bitmap, etc.), and so on. The datastore 1860 may also store compressed KV tensors generated by the compressing module 1830. The datastore 1860 may store graphs, configuration parameters, compilation descriptors, or other data generated by the compiler 1840 or the transformer execution module 1850. The datastore 1860 may include one or more memories. In the embodiment of FIG. 18, the datastore 1860 is a component of the DNN module 1800. In other embodiments, the datastore 1860 may be external to the DNN module 1800 and communicate with the DNN module 1800 through a network. In some embodiments, the datastore 1860 may be implemented in the memory 1710.

[0193] FIG. 19A illustrates an example sparse cell 1900, in accordance with various embodiments. The sparse cell 1900 may be a processing cell in a processing engine, e.g., the processing engine 1770 in FIG. 17. The sparse cell 1900 includes 16 MAC units 1910 (individually referred to as "MAC unit 1910"), which constitutes a MAC array having four rows and four columns. The MAC array has a spatial shape of 19x4, meaning the height of the MAC array is four and the width of the MAC array is also 19. The sparse cell 1900 also includes 16 weight register files 1920 (individually referred to as "weight register file 1920"), 16 activation register files 1930 (individually referred to as "activation register file 1930"), four row buffers 1940 (individually referred to as "row buffer 1940"), and sparsity modules 1966 (individually referred to as "sparsity module 1966"). In other embodiments, the sparse cell 1900 may include fewer, more, or different components. For example, the sparse cell 1900 may include a different number of MAC units 1910, weight register files 1920, activation register files 1930, row buffers 1940, or sparsity modules 1966. As another example, the sparse cell 1900 may include column buffers in lieu of or in addition to the row buffers 1940. Also, the shape (e.g., the height or width) of the MAC array may be different.

[0194] The MAC units 1910 are configured to perform MAC operations. Each MAC unit 1910 may include one or more multipliers and one or more adders. A multiplier may multiply an activation with a weight at a time to compute a product. In some embodiments (e.g., embodiments where the MAC unit 1910 includes multiple multipliers), the multipliers may operate simultaneously to process multiple activation-weight pairs and compute multiple products in one cycle. An adder may accumulate products computed by the multipliers. Even though not shown in FIG. 19A, the sparse cell may include an adder tree including aplurality of adder tiers. The first tier may receive outputs of a plurality of MAC units 1910. The number of adders in the first tier may be half of the number of the MAC units 1910, and each adder may accumulate the outputs of two MAC units 1910. The second tier may receive outputs of adders in the first tier. The number of adders in the second tier may be half of the number of adders in the first tier, and each adder in the second tier may accumulate the outputs of two adders in the first tier. The adder tree may include one or more other tiers. The last tier may include a single adder that accumulates outputs of adders in the second last tier to compute a partial sum of the sparse cell 1900.

[0195] The weight register files 1920 store weights to be processed in MAC operations. In the embodiments of FIG. 19A, four weight register files 1920 are grouped into a storage set that stores data to be used by a column of MAC units 1910. There are four storage sets corresponding to the four columns of MAC units 1910. In some embodiments, a weight register file 1920 may correspond to a MAC unit 1910 and store data to be processed by the MAC unit. In some embodiments, all the 16 weight register files 1920 constitute a weight storage unit.

[0196] The activation register files 1930 stores activations to be processed in MAC operations. In the embodiments of FIG. 19A, four activation register files 1930 are grouped into a storage set that stores data to be used by a row of MAC units 1910. There are four storage sets corresponding to the four rows of MAC units 1910. In some embodiments, an activation register file 1930 may correspond to a MAC unit 1910 and store data to be processed by the MAC unit. In some embodiments, all the 16 activation register files 1930 constitute an activation storage unit. The row buffers 1940 store outputs of the MAC units 1910. Each row buffer 1940 may drain outputs of a single row of MAC units 1910.

[0197] The sparsity module 1966 facilitates dynamic sparsity-based acceleration in the sparse cell 1900. In the embodiments of FIG. 19A, each sparsity module 1966 includes a sparsity tensor storage unit 1965 and a control logic 1967. The sparsity tensor storage unit 1965 stores combined sparsity tensors. A combined sparsity tensor stored in the sparsity tensor storage unit 1965 may correspond to an activation tensor and a weight tensor. A nonzero element in the combined sparsity tensor may correspond to a nonzero activationweight pair that includes a nonzero activation and a nonzero weight. The position of the nonzero activation in the activation tensor may match the position of the nonzero weight inthe weight tensor. The product of the nonzero activation and nonzero weight would be nonzero.

[0198] The control logic 1967 may control transmission of activations and weights stored from the weight register files 1920 and the activation register files 1930 to the MAC units 1910 based on sparsity tensors. For instance, the control logic 1967 may select a subset of the weights stored in the weight register files 1920 and select a subset of activations stored in the activation register files 1930 based on a combined sparsity tensor. The selected weights and activations constitute nonzero activation-weight pairs. The control logic 1967 may transmit the selected weights and activations to the MAC units 1910 for performing MAC operations. The other weights stored in the weight register files 1920 and the other activations stored in the activation register files 1930 are skipped from computation. In the embodiments of FIG. 19A, each sparsity module 1966 controls sparsity acceleration in a respective MAC unit 1910. As the sparsity acceleration is either based on both weight sparsity and activation sparsity, 16 sparsity modules 1966 are used for acceleration computations in the 16 MAC units 1910.

[0199] As shown in FIG. 19A, the sparse cell 1900 is associated with multiplexers (MUXs) 1903, 1904, 1905, and 1906. In other embodiments, the sparse cell 1900 may be associated with a different number of MUXs or other devices. The MUX 1903 facilitates loading weights, e.g., from the local memory 340, into the weight register files 1920. The MUX 1904 facilitates loading activations, e.g., from the local memory 340, into the activation register files 1930. The MUX 1905 facilitates loading sparsity tensors into the sparsity tensor storage unit 1965. The MUX 1906 may be a drain MUX that can facilitate draining outputs of the MAC units 1910, e.g., to the local memory 340.

[0200] In some embodiments, the sparse cell 1900 may also execute matrix multiplications converted from Fourier transform operations. For an example Fourier transform operation, the MAC units 1910 may perform MAC operations in the two sequences of matrix multiplications converted from the Fourier transform operation. The weight register files 1920 may be used to store data points in transformation tensor of the Fourier transform operation. The activation register file 1930 may be used to store data points in the input tensor of the Fourier transform operation. The row buffers 1940 may store data points in the output tensor of the Fourier transform operation.

[0201] FIG. 19B illustrates a sparse cell array 1950, in accordance with various embodiments. The sparse cell array 1950 may be an example of the processing engine 1770 in FIG. 17. In FIG. 19B, the sparse cell array 1950 includes sparse cells 1960 (individually referred to as "sparse cell 1960") arranged in four columns and four rows, an activation memory 1970, and a weight memory 1980. The sparse cell array 1950 may also be referred to as a data processing unit. In other embodiments, the sparse cell array 1950 may include fewer, more, or different components. For instance, the sparse cell array 1950 may include a different number of columns, rows, or sparse cells 1960.

[0202] Each sparse cell 1960 may perform sparsity accelerated MAC operations. The sparse cells 1960 may facilitate dynamic sparsity mode. For instance, the sparsity modes of a sparse cell 1960 may be dynamically changed between a combined sparsity mode, an activation sparsity mode, a weight sparsity mode, and a dense mode. An embodiment of a sparse cell 1960 may be the sparse cell 1900 in FIG. 19A. The activation memory 1970 stores activations, such as activations in input tensors of deep learning operations. Activations may be loaded from the activation memory 1970 to sparse cells 1960. The weight memory 1980 stores weights, such as weights in filters of deep learning operations. Weights may be loaded from the weight memory 1980 to sparse cells 1960. The activation memory 1970 or weight memory 1980 may be a buffer. In other embodiments, the sparse cell array 1950 may include a dense data memory and a sparse data memory in lieu of the activation memory 1970 and weight memory 1980. The dense data memory may store dense tensors, e.g., dense tensors generated by the load module 1760. The sparse data memory may store sparse tensors.

[0203] The sparse cell array 1950 may also execute matrix multiplications in Fourier transform operations. The activation memory 1970 may be used to store input tensors of the Fourier transform operations. The weight memory 1980 may be used to store transformation matrices of the Fourier transform operations.

[0204] FIG. 20 illustrates an example convolution, in accordance with various embodiments. The convolution may be a deep learning operation in a convolutional layer of a DNN. The convolution can be executed on an activation tensor 2010 and filters 2020 (individually referred to as "filter 2020"). The filters may constitute a weight tensor of the convolution. The result of the convolution is an output tensor 2030. In some embodiments, theconvolution is performed by a DNN accelerator. The convolution may be executed by a processing engine, such as the processing engine 1770 in FIG. 17.

[0205] The activation tensor 2010 may be computed in a previous layer of the DNN. In some embodiments (e.g., embodiments where the convolutional layer is the first layer of the DNN), the activation tensor 2010 may be an image. In the embodiments of FIG. 20, the activation tensor 2010 includes activations (also referred to as "input activations," "elements," or "input elements") arranged in a 3D matrix. The activation tensor 2010 may also be referred to as an input tensor of the convolution. An input element is a data point in the activation tensor 2010. The activation tensor 2010 has a spatial size HinX WinX Cin, where Hinis the height of the 3D matrix (i.e., the length along the Y axis, which indicates the number of activations in a column in the 3D matrix of each input channel), Winis the width of the 3D matrix (i.e., the length along the X axis, which indicates the number of activations in a row in the 20D matrix of each input channel), and Cinis the depth of the 3D matrix (i.e., the length along the Z axis, which indicates the number of input channels). For the purpose of simplicity and illustration, the activation tensor 2010 has a spatial size of 7x7x3, i.e., the activation tensor 2010 includes three input channels and each input channel has a 7x7 20D matrix. Each input element in the activation tensor 2010 may be represented by a (X, Y, Z) coordinate. In other embodiments, the height, width, or depth of the activation tensor 2010 may be different.

[0206] Each filter 2020 includes weights arranged in a 3D matrix. The values of the weights may be determined through training the DNN. A filter 2020 has a spatial sizeCf, where Hf is the height of the filter (i.e., the length along the Y axis, which indicates the number of weights in a column in each kernel), Wf is the width of the filter (i.e., the length along the X axis, which indicates the number of weights in a row in each kernel), and Cf is the depth of the filter (i.e., the length along the Z axis, which indicates the number of channels). In some embodiments, C equals Cin. For purpose of simplicity and illustration, each filter 2020 in FIG. 20 has a spatial size of 3x3x3, i.e., the filter 2020 includes 20 convolutional kernels with a spatial size of 3x3. In other embodiments, the height, width, or depth of the filter 2020 may be different. The spatial size of the convolutional kernels is smaller than the spatial size of the 20D matrix of each input channel in the activation tensor2010.

[0207] An activation or weight may take one or more bytes in a memory. The number of bytes for an activation or weight may depend on the data format. For example, when the activation or weight has an INT8 format, the activation takes one byte. When the activation or weight has a FP16 format, the activation or weight takes two bytes. Other data formats may be used for activations or weights.

[0208] In the convolution, each filter 2020 slides across the activation tensor 2010 and generates a 20D matrix for an output channel in the output tensor 2030. In the embodiments of FIG. 20, the 20D matrix has a spatial size of 5x5. The output tensor 2030 includes activations (also referred to as "output activations," "elements," or "output element") arranged in a 3D matrix. An output activation is a data point in the output tensor 2030. The output tensor 2030 has a spatial size HoutX WouLX Cout, where Houtis the height of the 3D matrix (i.e., the length along the Y axis, which indicates the number of output activations in a column in the 20D matrix of each output channel), Woutis the width of the 3D matrix (i.e., the length along the X axis, which indicates the number of output activations in a row in the 20D matrix of each output channel), and Coutis the depth of the 3D matrix (i.e., the length along the Z axis, which indicates the number of output channels). Coutmay equal the number of filters 2020 in the convolution. Houtand Woutmay depend on the heights and weights of the activation tensor 2010 and each filter 2020. In an example where the kernel size is lxl, Houtand Woutmay equal to Hinand Win, respectively.

[0209] As a part of the convolution, multiply-accumulate (MAC) operations can be performed on a 3x3x3 subtensor 2015 (which is highlighted with a dotted pattern in FIG. 20) in the activation tensor 2010 and each filter 2020. The result of the MAC operations on the subtensor 2015 and one filter 2020 is an output activation. In some embodiments (e.g., embodiments where the convolution is an integral convolution), an output activation may include 8 bits, e.g., one byte. In other embodiments (e.g., embodiments where the convolution is a floating-point convolution), an output activation may include more than one byte. For instance, an output element may include two bytes.

[0210] After the MAC operations on the subtensor 2015 and all the filters 2020 are finished, a vector 2035 is produced. The vector 2035 is highlighted with a dotted pattern in FIG. 20. The vector 2035 includes a sequence of output activations, which are arranged along the Z axis. The output activations in the vector 2035 have the same (x, y) coordinate, but the output activations correspond to different output channels and have different Zcoordinates. The dimension of the vector 2035 along the Z axis may equal the total number of output channels in the output tensor 2030. After the vector 2035 is produced, further MAC operations are performed to produce additional vectors till the output tensor 2030 is produced. In the embodiments of FIG. 20, the output tensor 2030 is computed in a Z-major format. When the output tensor 2030 is computed in the ZXY format, the vector that is adjacent to the vector 2035 along the X axis may be computed right after the vector 2035. When the output tensor 2030 is computed in the ZYX format, the vector that is adjacent to the vector 2035 along the Y axis may be computed right after the vector 2035. The output tensor 2030 may be permuted, e.g., by the drain module 1790, and stored in a memory (e.g., the local memory 340) in an X-major format or Y-major format.

[0211] In some embodiments, the MAC operations on a 3x3x3 subtensor (e.g., the subtensor 2015) and a filter 2020 may be performed by a plurality of MAC units. One or more MAC units may receive an input operand (e.g., an activation operand 2017 shown in FIG. 20) and a weight operand (e.g., the weight operand 2027 shown in FIG. 20). The activation operand 2017 includes a sequence of activations having the same (x, y) coordinate but different z coordinates. The activation operand 2017 includes an activation from each of the input channels in the activation tensor 2010. The weight operand 2027 includes a sequence of weights having the same (x, y) coordinate but different z coordinates. The weight operand 2027 includes a weight from each of the channels in the filter 2020. Activations in the activation operand 2017 and weights in the weight operand 2027 may be sequentially fed into a MAC unit. The MAC unit may receive an activation and a weight ("an activation-weight pair") at a time and multiple the activation and the weight. The position of the activation in the activation operand 2017 may match the position of the weight in the weight operand 2027. The activation and weight may correspond to the same channel.

[0212] Activations or weights may be floating-point numbers. Floating-point numbers may have various data formats, such as FP32, FP16, BF16, and so on. A floating-point number may be a positive or negative number with a decimal point. A floating-point number may be represented by a sequence of bits that includes one or more bits representing the sign of the floating-point number (e.g., positive or negative), bits representing an exponent of the floating-point number, and bits representing a mantissa of the floating-point number. The mantissa is the part of a floating-point number that represents the significant digits of thatnumber. The mantissa is multiplied by the base raised to the exponent to give the actual value of the floating-point number.

[0213] In some embodiments, the output activations in the output tensor 2030 may be further processed based on one or more activation functions before they are written into the memory or inputted into the next layer of the DNN. The processing based on the one or more activation functions may be at least part of the post processing of the convolution. In some embodiments, the post processing may include one or more other computations, such as offset computation, bias computation, and so on. The results of the post processing may be stored in a local memory of the compute block and be used as input to the next DNN layer. In some embodiments, the input activations in the activation tensor 2010 may be results of post processing of the previous DNN layer.

[0214] FIG. 21 illustrates implement of group MatMul operations on the DNN accelerator 1700 in FIG. 17, in accordance with various embodiments. For the purpose of illustration and simplicity, FIG. 21 shows part of the DNN system 1700. In some embodiments, group MatMul operations may be mapped to convolutions. For instance, for a group MatMul operations on queries and keys, the queries may be treated (e.g., saved, read, or loaded) as activations and the keys may be treated (e.g., saved, read, or loaded) as convolution weights. For a group MatMul operations on attention weights / scores and values, the attention weights / scores may be treated as activations and the values may be treated as convolution weights.

[0215] In the embodiments of FIG. 21, the load module 1760 includes an activation reader 2110 and a weight reader 2130. The activation reader 2110 may read either queries or attention scores from the local memory 1740 and load the queries or attention scores to the processing engine 1770. The weight reader 2130 may read keys or values from the local memory 1740 and load the keys or values to the processing engine 1770. The keys or values may be in quantized / compressed KV tensors. The processing engine 1770 may perform MatMul operations on data from the load module 1760. In embodiments where the processing engine 1770 performs a sub-MatMul operation (e.g., a group MatMul operation in a group MatMul layer), the processing engine 1770 may compute a partial sum from each sub-MatMul operation. The post-processing engine 1780 is an array of PPEs 2135 in FIG. 21. Each of the PPEs 2135 may include a scaling module for scaling data and a bias module for applying a bias on data. The PPEs 2135 may dequantize a partial sum using the scale valueand zero-point value of the corresponding group as scale and bias, respectively. The drain module 1790 may drain the dequantized partial sums from the post-processing engine 1780 and write the dequantized partial sums into the local memory 1740. In some embodiments, the dequantized partial sums may have a floating-point format, such as FP32. The digital signal processor 1735 may read the dequantized partial sums from the local memory 1740 and accumulate the dequantized partial sums to generate an output of the group MatMul layer.Example Methods for Transformer Model Execution

[0216] FIG. 22 illustrates a process 2200 of selecting optimal quantization configuration for KV cache of a transformer model, in accordance with various embodiments. The process 2200 may be performed by a DNN module, such as the DNN module 1701 in FIG. 17 or the DNN module 1800 in FIG. 18. Although the process 2200 is described with reference to the flowchart illustrated in FIG. 22, many other processes for selecting optimal quantization configuration for KV cache may alternatively be used. For example, the order of execution of the steps in FIG. 22 may be changed. As another example, some of the steps may be changed, eliminated, or combined. The process may use a hardware-software co-design methodology to select optimal quantization configurations for KV tensors. The methodology can ensure that the chosen quantization configuration and associated hyperparameters not only meets the accuracy requirements of the model on the software side by also maximizes compute efficiency on the hardware accelerator side. The process may be executed on various types of processing units (e.g., CPU, GPU, DNN accelerators, etc.) before deployment of the LLM on the DNN accelerator.

[0217] In some embodiments, the process may begin with the selection of type of dualgroup quantization by an assessment of compute efficiency on the hardware accelerator, followed by the evaluation of model accuracy. It then guides the user through the process of maximizing hardware utilization while satisfying accuracy constraints. In the embodiments of FIG. 22, attention dimension is found in Step 2210. The attention dimension found in Step 2210 include dimensions of the key tensor and dimensions of the value tensor. In some embodiments, the group matrix multiplication of Q and the transpose of K (Q x KT) may be mapped to a convolutional operation on the data processing unit. Each group's query tensor Q, with dimensions (1, Cg ), may be treated as activations, while the corresponding keytensor, with dimensions (Tg , Cg ), may be treated as convolution weights. This results in an output of size (1, Tg }. The channel dimension Cg may be designated as the input channel (IC), and the token dimension Tg may be designated as the output channel (OC), establishing channel as the accumulation dimension for K. For the group MatMul following the SoftMax operation, which computes the masked attention weights multiplied by V for each group, the value tensor, with dimensions (Tg , Cg ), may be treated as convolution weights, yielding an output of size (1, Cg ). The token dimension Tg is the IC, and the channel dimension Cg is the OC, indicating that the token dimension is the accumulation dimension for V.

[0218] In Step 2220, quantization configuration is set. For instance, it may determine whether channel-heavy dual-group quantization or token-heavy dual-group quantization should be used for the key tensor and value tensor. In some embodiments, hybrid quantization configuration may be selected. For instance, channel-heavy dual-group quantization may be selected for the key tensor, and token-heavy dual-group quantization may be selected for the value tensor. In some embodiments, the accelerator may support per-OC scale and zero-point values, which aligns with the non-accumulation direction of both convolutions. The most efficient hardware utilization may be achieved when fine granular grouping is performed along this direction. Consequently, a channel-heavy dual- group quantization or its extreme variant, per-token group quantization, may be the most suitable strategy for K. In per-token group quantization, this alignment can ensure that each token can be quantized with its own s and z, allowing us to exploit the data processing unit architecture efficiently during the convolution operations. In contrast, a token-heavy dual- group quantization or per-channel group quantization is the preferred approach for V. In per-channel group quantization, this strategy can allow the model to leverage the hardware's capability to apply a unique s and z for each channel, enhancing the efficiency of the convolution operations involving the value tensor.

[0219] In Step 2230, accumulation granularity on hardware is found for both K and V. In some embodiments, the accumulation granularity is determined based on one or more hardware constraints, such as the accelerator's constraint on the number of ICs that can be accumulated in one cycle during a convolution operation on the data processing unit. In an example where the accelerator is designed to perform MAC operation (that includesaccumulation) on 16 ICs simultaneously, the IC group dimension needs to be set to 16 to achieve 100% hardware utilization. In this example, the accumulation granularity may be set to: Cgmn= 16 for the key tensor and= 16 for the value tensor. This constraint may further refine the quantization strategy, ensuring that the hardware's parallel processing capabilities can be fully leveraged.

[0220] In Step 2240, group dimensions are set to match hardware constraint. In some embodiment, the channel group dimension of the key tensor may be denoted where n may be an integer. The token group dimension of the key tensor may where m may be an integer. The channel group dimension of the value tensor2n. The token group dimension of the value tensor may be Tg = 2m. In some embodiments, a user input 2201 is received for starting the search of optimal values of n and m. For instance, the user input may include an initial value of n and an initial value of m, which corresponding to initial group dimension of the KV tensors. The group dimensions for both the K and value tensors may be adjusted by adjusting n and m to conform to the accelerator constraints, ensuring that the quantization strategy is both feasible and efficient on the intended computational platform. For instance, n and m may be adjusted so that the following criteria can be met:< Cg < dheadwith dheadrepresenting the dimension of each attention head, which may be the total channels in the tensor; Tg < linputfor the prefill phase quantization with linputrepresenting the length of the initial input prompt or sequence, which may be the total number of tokens in the tensor; Cg < dhead; and T™inlinput. By setting these group dimensions based on hardware constraints in the above manner, the hardware's native support for power-of-two sizes can be leveraged, which can lead to more efficient memory access patterns and computational throughput. This deliberate alignment with hardware constraints can be a key factor in achieving high performance and scalability for attention computations. Quantization configuration parameters 2202 may be stored in a memory.

[0221] In Step 2250, hybrid dual-group quantization may be applied to the KV cache using the quantization configuration parameters 2202. The KV cache may be generated after the prefill inference phase. One or more cached inference stages may be run with periodic quantization or adaptive grouped precision attention. The accuracy of the transformermodel may be evaluated against a subset of downstream task dataset to confirm the quantization maintains model performance.

[0222] In Step 2260, it is determined whether the accuracy is adequate. For instance, it may be determined whether the accuracy meets a threshold accuracy. In embodiments where the accuracy is adequate, Step 2270 is performed, in which the quantization configuration parameters are treated as the final quantization configuration parameters and are saved. In embodiments where the accuracy is inadequate, Step 2280 is performed. In some embodiments, Steps 2280, 2285, and 2290 may be performed separately for the key tensor and value tensor.

[0223] For the key tensor, it is determined whether Tg = 1 in Step 2280. In embodiments where Tg =# 1, Step 2290 is performed, in which Tg is decreased (e.g., halved) to reduce the number of tokens / output channels, which allows for multiple s and z within the same workload. For instance, with linput= 128, Tg = 32, Cg = 64, dmodei= 12 X 128, where 12 is nheads> we have 96 groups, each with distinct s and z values, but only 24 workloads, each containing 4 groups. In embodiments where Tg = 1, Step 2285 is performed, in which it is determined whether Cg = 1. In embodiments where Cg = 1, Step 2270 is performed. In embodiments where Cg = 1, Step 2290 is performed, in which Cg is decreased to increase the number of workloads as each IC group is processed separately. In some embodiments, Cg maybe halved in Step 2290. This adjustment may have no impact on hardware utilization as long as Cg remains above the minimum threshold (C™"1). Although more workloads may slightly increase computation time, this can generally be offset by the memory-bound nature of cached inference, where the time spent on memory transactions far exceeds that of computation. Thus, the minor increase in compute time can be negligible compared to the latency and energy savings from KV cache compression.

[0224] For the value tensor, it is determined whether Cg = 1 in Step 2280. In embodiments where Cg = 1, Step 2290 is performed, in which Cg = 1 is decreased (e.g., halved). In some embodiments, Cg may be adjusted at the same time Tg is adjusted. The adjusted Tg and Cg may be stored as the quantization configuration parameters 2202. Step 2250 and the following step are performed using the quantization configuration parameters 2202. For mixed-precision hybrid DGQ, each group may be a separate workload.

[0225] In embodiments where C = 1, Step 2285 is performed, in which it is determined whether T = 1. In embodiments where T = 1, Step 2270 is performed. In embodiments where T 1, Step 2290 is performed, in which Tg may be reduced to scale down tokens (i.e., ICs) per group for value tensor. In some embodiments, Tg maybe halved in Step 2290. In some embodiments, the adjustment of Tg and Cg may start simultaneously. Every time the quantization configuration parameters are adjusted, the new quantization configuration parameter may be saved as the quantization configuration parameters 2202, which may be used to estimate the model accuracy in Step 2250. The steps following Step 2250 may also be performed. This iterative process of reducing Cg and Tg may continue until the desired accuracy is achieved or group dimensions have reached their minimum values, e.g., Cg = Tg = 1. At either of these junctures, Step 2270 would be performed. Throughout the process in FIG. 22, accuracy may be the primary constraint, guiding the necessary reductions in different hybrid dual-group quantization configuration parameters to refine the quantization approach and ensure the performance of the transformer model.

[0226] FIG. 23 is a flowchart of a method 2300 for executing a transformer model, in accordance with various embodiments. The method 2300 may be performed by the DNN system 1700 in FIG. 17. Although the method 2300 is described with reference to the flowchart illustrated in FIG. 23, many other methods for executing transformer models may alternatively be used. For example, the order of execution of the steps in FIG. 23 may be changed. As another example, some of the steps may be changed, eliminated, or combined.

[0227] The DNN system 1700 segments 2310 an attention tensor of a transformer model into groups. The attention tensor comprises data elements that are keys or values. In some embodiments, the attention tensor is generated in an inference stage of the transformer model. In some embodiments, the attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels. In some embodiments, the DNN system 1700 segments the attention tensor into the groups along both the token dimension and the channel dimension. In some embodiments, a size of a group along the channel dimension is larger than a size of one or more other groups along the channel dimension. In some embodiments, a data precision of the group in the quantized attention tensor is higher than a data precision of the one or more other groups in the quantized attention tensor. In some embodiments, the size of the group along thechannel dimension is the same as a size of the attention tensor along the channel dimension. In some embodiments, the group includes a key or value corresponding to a token predicted by the transformer model in the inference stage.

[0228] In some embodiments, the DNN system 1700 determines a group size based on a constraint of an apparatus executing the subsequent inference stage of the transformer model. The group size indicates a number of tokens or a number of channels in each group. In some embodiments, the DNN system 1700 determines a group size based on an estimated accuracy of the transformer model when the attention tensor is quantized using the group size.

[0229] The DNN system 1700 determines 2320 quantization parameters for the groups. Different groups have different quantization parameters. Data elements within the same group have the same quantization parameter. In some embodiments, a group has two quantization parameters, such as a scale factor and a zero-point factor. Different groups have different scale or zero-point values.

[0230] The DNN system 1700 quantizes 2330 the data elements based on the quantization parameters. A data element is quantized based on a quantization parameter of a group that includes the data element.

[0231] The DNN system 1700 generates 2340 a quantized attention tensor comprising the quantized data elements. In some embodiments, the groups in the quantized attention tensor have different data precisions.

[0232] The DNN system 1700 caches 2350 the quantized attention tensor in a storage unit. In some embodiments, the storage unit is a memory, such as a DRAM. In some embodiments, the quantized attention tensor is to be used in a subsequent inference stage of the transformer model. In some embodiments, the subsequent inference stage of the transformer model is performed using the quantized attention tensor and a query tensor. The query tensor comprises data elements that have a higher data precision than data elements in the quantized attention tensor. In some embodiments, the subsequent inference stage of the transformer model comprises a MatMul operation that is partitioned into a plurality of sub-MatMul operations. The MatMul operation is performed by performing each sub-MatMul operation using a different group and accumulating results of the plurality of sub-MatMul operations.

[0233] FIG. 24 is a flowchart of another method 2400 for executing a transformer model, in accordance with various embodiments. The method 2400 may be performed by the DNN system 1700 in FIG. 17. Although the method 2400 is described with reference to the flowchart illustrated in FIG. 24, many other methods for executing transformer models may alternatively be used. For example, the order of execution of the steps in FIG. 24 may be changed. As another example, some of the steps may be changed, eliminated, or combined.

[0234] The DNN system 1700 caches 2410 a quantized attention tensor generated by quantizing an attention tensor of the transformer model. The quantized attention tensor comprises groups with different data precisions. A group comprises a subset of keys or values in the quantized attention tensor. In some embodiments, the attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels.

[0235] In some embodiments, the quantized attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels. A size of a group along the channel dimension is larger than a size of one or more other groups along the channel dimension. In some embodiments, a data precision of the group is higher than a data precision of the one or more other groups. In some embodiments, the size of the group along the channel dimension is the same as a size of the quantized attention tensor along the channel dimension.

[0236] In some embodiments, the DNN system 1700 segments the attention tensor into groups along both the token dimension and the channel dimension. The DNN system 1700 determines different quantization parameters for different groups within the attention tensor. The DNN system 1700 quantizes the attention tensor using the different quantization parameters. In some embodiments, the DNN system 1700 determines a group size based on a constraint of an apparatus performing the plurality of matrix multiplication operations. The group size indicates a number of tokens or a number of channels in each group. In some embodiments, the DNN system 1700 determines a group size based on an estimated accuracy of the transformer model when the attention tensor is quantized using the group size.

[0237] The DNN system 1700 computes 2420 partial sums by performing a plurality of matrix multiplication operations using the quantized attention tensor. A matrix multiplication operation is performed using a different group and outputs a partial sum. Insome embodiments, a group, which has a size larger than a size of one or more other groups, is generated in a first inference stage of the transformer model. The plurality of matrix multiplication operations performed using the quantized attention tensor is in a second inference stage of the transformer model. The second inference stage is after the first inference stage. In some embodiments, the plurality of matrix multiplication operations is performed further using a query tensor. The query tensor comprising data elements that have a higher data precision than data elements in the quantized attention tensor

[0238] The DNN system 1700 dequantizes 2430 the partial sums using different quantization parameters. In some embodiments, for each partial sum, the DNN system 1700 identifies a scale value and a zero-point value for the partial sum based on a group from which the partial sum is computed. The DNN system 1700 dequantizes the partial sum using the scale value and zero-point value.

[0239] The DNN system 1700 computes 2440 an output of an attention layer in the transformer model by accumulating the dequantized partial sums. In some embodiments, the dequantized partial sums are computed by a data processing unit, while the accumulation of the dequantized partial sums is performed by a digital signal processor. In some embodiments, the output of the attention layer is used to predict a new token. A new attention tensor is generated by concatenating the attention tensor with new keys or values that correspond to the new token. The new attention tensor is segmented into new groups and each of the new groups is quantized with group-specific quantization parameters. The quantized attention tensor is cached and is to be used in the next inference stage of the transformer model.Example Computing Device

[0240] FIG. 25 is a block diagram of an example computing device 2500, in accordance with various embodiments. In some embodiments, the computing device 2500 can be used as at least part of the DNN system 1700. A number of components are illustrated in FIG. 25 as included in the computing device 2500, but any one or more of these components may be omitted or duplicated, as suitable for the application. In some embodiments, some or all of the components included in the computing device 2500 may be attached to one or more motherboards. In some embodiments, some or all of these components are fabricated onto a single system on a chip (SoC) die. Additionally, in various embodiments, the computing device 2500 may not include one or more of the components illustrated in FIG. 25, but thecomputing device 2500 may include interface circuitry for coupling to the one or more components. For example, the computing device 2500 may not include a display device 2506, but may include display device interface circuitry (e.g., a connector and driver circuitry) to which a display device 2506 may be coupled. In another set of examples, the computing device 2500 may not include an audio input device 2518 or an audio output device 2508 but may include audio input or output device interface circuitry (e.g., connectors and supporting circuitry) to which an audio input device 2518 or audio output device 2508 may be coupled.

[0241] The computing device 2500 may include a processing device 2502 (e.g., one or more processing devices). The processing device 2502 processes electronic data from registers and / or memory to transform that electronic data into other electronic data that may be stored in registers and / or memory. The computing device 2500 may include a memory 2504, which may itself include one or more memory devices such as volatile memory (e.g., DRAM), nonvolatile memory (e.g., read-only memory (ROM)), high bandwidth memory (HBM), flash memory, solid state memory, and / or a hard drive. In some embodiments, the memory 2504 may include memory that shares a die with the processing device 2502. In some embodiments, the memory 2504 includes one or more non-transitory computer- readable media storing instructions executable to perform operations for executing transformer models (e.g., the method 2300 described in conjunction with FIG. 23 or the method 2400 described in conjunction with FIG. 24) or some operations performed by one or more components of the DNN system 1700. The instructions stored in the one or more non-transitory computer-readable media may be executed by the processing device 2502.

[0242] In some embodiments, the computing device 2500 may include a communication chip 2512 (e.g., one or more communication chips). For example, the communication chip 2512 may be configured for managing wireless communications for the transfer of data to and from the computing device 2500. The term "wireless" and its derivatives may be used to describe circuits, devices, systems, methods, techniques, communications channels, etc., that may communicate data through the use of modulated electromagnetic radiation through a nonsolid medium. The term does not imply that the associated devices do not contain any wires, although in some embodiments they might not.

[0243] The communication chip 2512 may implement any of a number of wireless standards or protocols, including but not limited to Institute for Electrical and Electronic Engineers(IEEE) standards including Wi-Fi (IEEE 802.10 family), IEEE 802.16 standards (e.g., IEEE 802.16-2005 Amendment), Long-Term Evolution (LTE) project along with any amendments, updates, and / or revisions (e.g., advanced LTE project, ultramobile broadband (UMB) project (also referred to as "3GPP2"), etc.). IEEE 802.16 compatible Broadband Wireless Access (BWA) networks are generally referred to as WiMAX networks, an acronym that stands for worldwide interoperability for microwave access, which is a certification mark for products that pass conformity and interoperability tests for the IEEE 802.16 standards. The communication chip 2512 may operate in accordance with a Global System for Mobile Communication (GSM), General Packet Radio Service (GPRS), Universal Mobile Telecommunications System (UMTS), High Speed Packet Access (HSPA), Evolved HSPA (E- HSPA), or LTE network. The communication chip 2512 may operate in accordance with Enhanced Data for GSM Evolution (EDGE), GSM EDGE Radio Access Network (GERAN), Universal Terrestrial Radio Access Network (UTRAN), or Evolved UTRAN (E-UTRAN). The communication chip 2512 may operate in accordance with Code-division Multiple Access (CDMA), Time Division Multiple Access (TDMA), Digital Enhanced Cordless Telecommunications (DECT), Evolution-Data Optimized (EV-DO), and derivatives thereof, as well as any other wireless protocols that are designated as 3G, 4G, 5G, and beyond. The communication chip 2512 may operate in accordance with other wireless protocols in other embodiments. The computing device 2500 may include an antenna 2522 to facilitate wireless communications and / or to receive other wireless communications (such as AM or FM radio transmissions).

[0244] In some embodiments, the communication chip 2512 may manage wired communications, such as electrical, optical, or any other suitable communication protocols (e.g., the Ethernet). As noted above, the communication chip 2512 may include multiple communication chips. For instance, a first communication chip 2512 may be dedicated to shorter-range wireless communications such as Wi-Fi or Bluetooth, and a second communication chip 2512 may be dedicated to longer-range wireless communications such as global positioning system (GPS), EDGE, GPRS, CDMA, WiMAX, LTE, EV-DO, or others. In some embodiments, a first communication chip 2512 may be dedicated to wireless communications, and a second communication chip 2512 may be dedicated to wired communications.

[0245] The computing device 2500 may include battery / power circuitry 2514. The battery / power circuitry 2514 may include one or more energy storage devices (e.g., batteries or capacitors) and / or circuitry for coupling components of the computing device 2500 to an energy source separate from the computing device 2500 (e.g., AC line power).

[0246] The computing device 2500 may include a display device 2506 (or corresponding interface circuitry, as discussed above). The display device 2506 may include any visual indicators, such as a heads-up display, a computer monitor, a projector, a touchscreen display, a liquid crystal display (LCD), a light-emitting diode display, or a flat panel display, for example.

[0247] The computing device 2500 may include an audio output device 2508 (or corresponding interface circuitry, as discussed above). The audio output device 2508 may include any device that generates an audible indicator, such as speakers, headsets, or earbuds, for example.

[0248] The computing device 2500 may include an audio input device 2518 (or corresponding interface circuitry, as discussed above). The audio input device 2518 may include any device that generates a signal representative of a sound, such as microphones, microphone arrays, or digital instruments (e.g., instruments having a musical instrument digital interface (MIDI) output).

[0249] The computing device 2500 may include a GPS device 2516 (or corresponding interface circuitry, as discussed above). The GPS device 2516 may be in communication with a satellite-based system and may receive a location of the computing device 2500, as known in the art.

[0250] The computing device 2500 may include another output device 2510 (or corresponding interface circuitry, as discussed above). Examples of the other output device 2510 may include an audio codec, a video codec, a printer, a wired or wireless transmitter for providing information to other devices, or an additional storage device.

[0251] The computing device 2500 may include another input device 2520 (or corresponding interface circuitry, as discussed above). Examples of the other input device 2520 may include an accelerometer, a gyroscope, a compass, an image capture device, a keyboard, a cursor control device such as a mouse, a stylus, a touchpad, a bar code reader, a Quick Response (QR) code reader, any sensor, or a radio frequency identification (RFID) reader.

[0252] The computing device 2500 may have any desired form factor, such as a handheld or mobile computer system (e.g., a cell phone, a smart phone, a mobile internet device, a music player, a tablet computer, a laptop computer, a netbook computer, an ultrabook computer, a personal digital assistant (PDA), an ultramobile personal computer, etc.), a desktop computer system, a server or other networked computing component, a printer, a scanner, a monitor, a set-top box, an entertainment control unit, a vehicle control unit, a digital camera, a digital video recorder, or a wearable computer system. In some embodiments, the computing device 2500 may be any other electronic device that processes data.Select Examples

[0253] The following paragraphs provide various examples of the embodiments disclosed herein.

[0254] Example 1 provides a method of executing a transformer model, including segmenting an attention tensor of a transformer model into groups, the attention tensor including data elements that are keys or values; determining quantization parameters for the groups, different groups having different quantization parameters, data elements within a same group having a same quantization parameter; quantizing the data elements based on the quantization parameters, a data element quantized based on a quantization parameter of a group that includes the data element; generating a quantized attention tensor comprising the quantized data elements; and caching the quantized attention tensor in a storage unit.

[0255] Example 2 provides the method of example 1, in which the attention tensor is generated in an inference stage of the transformer model, and the quantized attention tensor is used in a subsequent inference stage of the transformer model.

[0256] Example 3 provides the method of example 1 or 2, in which the attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels, and segmenting the attention tensor includes segmenting the attention tensor into the groups along both the token dimension and the channel dimension, in which a data precision of a group in the quantized attention tensor is higher than a data precision of one or more other groups in the quantized attention tensor.

[0257] Example 4 provides the method of example 3, in which a size of the group along the channel dimension is larger than a size of the one or more other roups alon the channel dimension.

[0258] Example 5 provides the method of example 3 or 4, in which the size of the group along the channel dimension is the same as a size of the attention tensor along the channel dimension.

[0259] Example 6 provides the method of any one of examples 3-5, in which the group includes a key or value corresponding to a token predicted by the transformer model in the inference stage.

[0260] Example 7 provides the method of any one of examples 2-6, in which segmenting the attention tensor includes determining a group size based on a constraint of an apparatus executing the subsequent inference stage of the transformer model or based on an estimated accuracy of the transformer model when the attention tensor is quantized using the group size, the group size indicating a number of tokens or a number of channels in each group.

[0261] Example 8 provides the method of any one of examples 1-7, in which the attention tensor is a key tensor, and the method further includes segmenting a value tensor into groups, in which a shape of a group in the value tensor is different from a shape of a group in the key tensor.

[0262] Example 9 provides the method of any one of examples 2-8, in which the subsequent inference stage of the transformer model is performed using the quantized attention tensor and a query tensor, the query tensor including data elements that have a higher data precision than data elements in the quantized attention tensor.

[0263] Example 10 provides the method of any one of examples 2-9, in which the subsequent inference stage of the transformer model includes a MatMul operation that is partitioned into a plurality of sub-MatMul operations and is performed by: performing each sub-MatMul operation using a different group; and accumulating results of the plurality of sub-MatMul operations.

[0264] Example 11 provides one or more non-transitory computer-readable media storing instructions executable to perform operations for executing a transformer model, the operations including segmenting an attention tensor of the transformer model into groups, the attention tensor including data elements that are keys or values; determiningquantization parameters for the groups, different groups having different quantization parameters, data elements within a same group having a same quantization parameter; quantizing the data elements based on the quantization parameters, a data element quantized based on a quantization parameter of a group that includes the data element; generating a quantized attention tensor comprising the quantized data elements; and caching the quantized attention tensor in a storage unit.

[0265] Example 12 provides the one or more non-transitory computer-readable media of example 11, in which the attention tensor is generated in an inference stage of the transformer model, and the quantized attention tensor is used in a subsequent inference stage of the transformer model.

[0266] Example 13 provides the one or more non-transitory computer-readable media of example 11 or 12, in which the attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels, and segmenting the attention tensor includes segmenting the attention tensor into the groups along both the token dimension and the channel dimension, in which a data precision of a group in the quantized attention tensor is higher than a data precision of one or more other groups in the quantized attention tensor.

[0267] Example 14 provides the one or more non-transitory computer-readable media of example 13, in which a size of the group along the channel dimension is larger than a size of the one or more other groups along the channel dimension.

[0268] Example 15 provides the one or more non-transitory computer-readable media of example 13 or 14, in which the size of the group along the channel dimension is the same as a size of the attention tensor along the channel dimension.

[0269] Example 16 provides the one or more non-transitory computer-readable media of any one of examples 13-15, in which the group includes a key or value corresponding to a token predicted by the transformer model in the inference stage.

[0270] Example 17 provides the one or more non-transitory computer-readable media of any one of examples 12-16, in which segmenting the attention tensor includes determining a group size based on a constraint of an apparatus executing the subsequent inference stage of the transformer model or based on an estimated accuracy of the transformer model when the attention tensor is quantized using the group size, the group size indicating a number of tokens or a number of channels in each group.

[0271] Example 18 provides the one or more non-transitory computer-readable media of any one of examples 11-17, in which the attention tensor is a key tensor, and the operations further include segmenting a value tensor into groups, in which a shape of a group in the value tensor is different from a shape of a group in the key tensor.

[0272] Example 19 provides the one or more non-transitory computer-readable media of any one of examples 12-18, in which the subsequent inference stage of the transformer model is performed using the quantized attention tensor and a query tensor, the query tensor including data elements that have a higher data precision than data elements in the quantized attention tensor.

[0273] Example 20 provides the one or more non-transitory computer-readable media of any one of examples 12-19, in which the subsequent inference stage of the transformer model includes a MatMul operation that is partitioned into a plurality of sub-MatMul operations and is performed by: performing each sub-MatMul operation using a different group; and accumulating results of the plurality of sub-MatMul operations.

[0274] Example 21 provides an apparatus, including a computer processor for executing computer program instructions; and a non-transitory computer-readable memory storing computer program instructions executable by the computer processor to perform operations for executing a transformer model, the operations including segmenting an attention tensor of a transformer model into groups, the attention tensor including data elements that are keys or values, determining quantization parameters for the groups, different groups having different quantization parameters, data elements within a same group having a same quantization parameter, quantizing the data elements based on the quantization parameters, a data element quantized based on a quantization parameter of a group that includes the data element, generating a quantized attention tensor comprising the quantized data elements, and caching the quantized attention tensor in a storage unit.

[0275] Example 22 provides the apparatus of example 21, in which the attention tensor is generated in an inference stage of the transformer model, and the quantized attention tensor is used in a subsequent inference stage of the transformer model.

[0276] Example 23 provides the apparatus of example 21 or 22, in which the attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels, and segmenting the attention tensorincludes segmenting the attention tensor into the groups along both the token dimension and the channel dimension.

[0277] Example 24 provides the apparatus of example 23, in which: a size of a group along the channel dimension is larger than a size of one or more other groups along the channel dimension, a data precision of the group in the quantized attention tensor is higher than a data precision of the one or more other groups in the quantized attention tensor, and the group includes a key or value corresponding to a token predicted by the transformer model in the inference stage.

[0278] Example 25 provides the apparatus of any one of examples 22-23, in which segmenting the attention tensor includes determining a group size based on a constraint of an apparatus executing the subsequent inference stage of the transformer model and an estimated accuracy of the transformer model when the attention tensor is quantized using the group size, the group size indicating a number of tokens or a number of channels in each group.Additional Examples

[0279] The following paragraphs provide additional examples of the embodiments disclosed herein.

[0280] Example 1 provides an apparatus, including a memory configured to store a quantized attention tensor generated by quantizing an attention tensor of a transformer model, the quantized attention tensor including groups with different data precisions, a group including a subset of keys or values in the quantized attention tensor; a data processing unit configured to: compute partial sums by performing a plurality of matrix multiplication operations using the quantized attention tensor, in which a matrix multiplication operation is performed using a different group and outputs a partial sum, and dequantize the partial sums using different quantization parameters; and a digital signal processor configured to compute an output of an attention layer in the transformer model by accumulating the dequantized partial sums .

[0281] Example 2 provides the apparatus of example 1, in which the memory is further configured to store the partial sums, and the digital signal processor is further configured to receive the partial sums from the memory.

[0282] Example 3 provides the apparatus of example 1 or 2, in which the attention tensor has a token dimension corresponding to a number of tokens and a channel dimensioncorresponding to a number of channels, and quantizing the attention tensor includes segmenting the attention tensor into groups along both the token dimension and the channel dimension; determining different quantization parameters for different groups within the attention tensor; and quantizing the attention tensor using the different quantization parameters.

[0283] Example 4 provides the apparatus of example 3, in which segmenting the attention tensor includes determining a group size based on a constraint of the apparatus, the group size indicating a number of tokens or a number of channels in each group.

[0284] Example 5 provides the apparatus of example 3 or 4, in which segmenting the attention tensor includes determining a group size based on an estimated accuracy of the transformer model when the attention tensor is quantized using the group size, the group size indicating a number of tokens or a number of channels in each group.

[0285] Example 6 provides the apparatus of any one of examples 1-5, in which the quantized attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels, and a size of a group along the channel dimension is larger than a size of one or more other groups along the channel dimension.

[0286] Example 7 provides the apparatus of example 6, in which a data precision of the group is higher than a data precision of the one or more other groups.

[0287] Example 8 provides the apparatus of example 6 or 7, in which the size of the group along the channel dimension is the same as a size of the quantized attention tensor along the channel dimension.

[0288] Example 9 provides the apparatus of any one of examples 6-8, in which the group includes a key or value corresponding to a token predicted by the transformer model in the inference stage.

[0289] Example 10 provides the apparatus of any one of examples 1-9, in which the data processing unit configured to perform the plurality of matrix multiplication operations further using a query tensor, the query tensor including data elements that have a higher data precision than data elements in the quantized attention tensor.

[0290] Example 11 provides a method of executing a transformer model, the method including caching a quantized attention tensor generated by quantizing an attention tensor of the transformer model, the quantized attention tensor including groups with differentdata precisions, a group including a subset of keys or values in the quantized attention tensor; computing partial sums by performing a plurality of matrix multiplication operations using the quantized attention tensor, in which a matrix multiplication operation is performed using a different group and outputs a partial sum; dequantizing the partial sums using different quantization parameters; and computing an output of an attention layer in the transformer model by accumulating the dequantized partial sums.

[0291] Example 12 provides the method of example 11, in which dequantizing the partial sums includes for each partial sum: identifying a scale value and a zero-point valuefor the partial sum based on a group from which the partial sum is computed; and dequantizing the partial sum using the scale value and zero-point value.

[0292] Example 13 provides the method of example 11 or 12, in which the attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels, and quantizing the attention tensor includes segmenting the attention tensor into groups along both the token dimension and the channel dimension; determining different quantization parameters for different groups within the attention tensor; and quantizing the attention tensor using the different quantization parameters.

[0293] Example 14 provides the method of example 13, in which segmenting the attention tensor includes determining a group size based on a constraint of an apparatus performing the plurality of matrix multiplication operations, the group size indicating a number of tokens or a number of channels in each group.

[0294] Example 15 provides the method of example 13 or 14, in which segmenting the attention tensor includes determining a group size based on an estimated accuracy of the transformer model when the attention tensor is quantized using the group size, the group size indicating a number of tokens or a number of channels in each group.

[0295] Example 16 provides the method of any one of examples 11-15, in which the quantized attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels, and a size of a group along the channel dimension is larger than a size of one or more other groups along the channel dimension.

[0296] Example 17 provides the method of example 16, in which a data precision of the group is higher than a data precision of the one or more other groups.

[0297] Example 18 provides the method of example 16 or 17, in which the size of the group along the channel dimension is the same as a size of the quantized attention tensor along the channel dimension.

[0298] Example 19 provides the method of any one of examples 16-18, in which the group is generated in a first inference stage of the transformer model, the plurality of matrix multiplication operations performed using the quantized attention tensor is in a second inference stage of the transformer model, and the second inference stage is after the first inference stage.

[0299] Example 20 provides the method of any one of examples 11-19, in which the plurality of matrix multiplication operations are performed further using a query tensor, the query tensor including data elements that have a higher data precision than data elements in the quantized attention tensor.

[0300] Example 21 provides one or more non-transitory computer-readable media storing instructions executable to perform operations for executing a transformer model, the operations including caching a quantized attention tensor generated by quantizing an attention tensor of the transformer model, the quantized attention tensor including groups with different data precisions, a group including a subset of keys or values in the quantized attention tensor; computing partial sums by performing a plurality of matrix multiplication operations using the quantized attention tensor, in which a matrix multiplication operation is performed using a different group and outputs a partial sum; dequantizing the partial sums using different quantization parameters; and computing an output of an attention layer in the transformer model by accumulating the dequantized partial sums.

[0301] Example 22 provides the one or more non-transitory computer-readable media of example 21, in which dequantizing the partial sums includes for each partial sum: identifying a scale value and a zero-point value forthe partial sum based on a group from which the partial sum is computed; and dequantizing the partial sum using the scale value and zeropoint value.

[0302] Example 23 provides the one or more non-transitory computer-readable media of example 21 or 12, in which the attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels, and quantizing the attention tensor includes segmenting the attention tensor into groups along both the token dimension and the channel dimension; determining different quantizationparameters for different groups within the attention tensor; and quantizing the attention tensor using the different quantization parameters.

[0303] Example 24 provides the one or more non-transitory computer-readable media of any one of examples 21-23, in which the quantized attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels, a size of a group along the channel dimension is larger than a size of one or more other groups along the channel dimension, and a data precision of the group is higher than a data precision of the one or more other groups.

[0304] Example 25 provides the one or more non-transitory computer-readable media of any one of examples 21-24, in which the plurality of matrix multiplication operations are performed further using a query tensor, the query tensor including data elements that have a higher data precision than data elements in the quantized attention tensor.

[0305] The above description of illustrated implementations of the disclosure, including what is described in the Abstract, is not intended to be exhaustive or to limit the disclosure to the precise forms disclosed. While specific implementations of, and examples for, the disclosure are described herein for illustrative purposes, various equivalent modifications are possible within the scope of the disclosure, as those skilled in the relevant art will recognize. These modifications may be made to the disclosure in light of the above detailed description.

Claims

1. Claims1. A method of executing a transformer model, comprising: segmenting an attention tensor of the transformer model into groups, the attention tensor comprising data elements that are keys or values; determining quantization parameters for the groups, different groups having different quantization parameters, data elements within a same group having a same quantization parameter; quantizing the data elements based on the quantization parameters, a data element quantized based on a quantization parameter of a group that includes the data element; generating a quantized attention tensor comprising the quantized data elements; and caching the quantized attention tensor in a storage unit.

2. The method of claim 1, wherein the attention tensor is generated in an inference stage of the transformer model, and the quantized attention tensor is used in a subsequent inference stage of the transformer model.

3. The method of claim 1 or 2, wherein the attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels, and segmenting the attention tensor comprises: segmenting the attention tensor into the groups along both the token dimension and the channel dimension, wherein a data precision of a group in the quantized attention tensor is higher than a data precision of one or more other groups in the quantized attention tensor.

4. The method of claim 3, wherein a size of the group along the channel dimension is larger than a size of the one or more other groups along the channel dimension.

5. The method of claim 4, wherein the size of the group along the channel dimension is the same as a size of the attention tensor along the channel dimension.

6. The method of any one of claims 3-5, wherein the group includes a key or value corresponding to a token predicted by the transformer model in the inference stage.

7. The method of any one of claims 1-6, wherein the attention tensor is a key tensor, and the method further comprises: segmenting a value tensor into groups, wherein a shape of a group in the value tensor is different from a shape of a group in the key tensor.

8. The method of any one of claims 2-7, wherein segmenting the attention tensor comprises: determining a group size based on a constraint of an apparatus executing the subsequent inference stage of the transformer model or based on an estimated accuracy of the transformer model when the attention tensor is quantized using the group size, the group size indicating a number of tokens or a number of channels in each group.

9. The method of any one of claims 2-8, wherein the subsequent inference stage of the transformer model is performed using the quantized attention tensor and a query tensor, the query tensor comprising data elements that have a higher data precision than data elements in the quantized attention tensor.

10. The method of any one of claims 2-9, wherein the subsequent inference stage of the transformer model comprises a matrix multiplication (MatMul) operation that is partitioned into a plurality of sub-MatMul operations and is performed by: performing each sub-MatMul operation using a different group; and accumulating results of the plurality of sub-MatMul operations.

11. One or more non-transitory computer-readable media storing instructions executable to perform operations for executing a transformer model, the operations comprising: segmenting an attention tensor of the transformer model into groups, the attention tensor comprising data elements that are keys or values;determining quantization parameters for the groups, different groups having different quantization parameters, data elements within a same group having a same quantization parameter; quantizing the data elements based on the quantization parameters, a data element quantized based on a quantization parameter of a group that includes the data element; generating a quantized attention tensor comprising the quantized data elements; and caching the quantized attention tensor in a storage unit.

12. The one or more non-transitory computer-readable media of claim 11, wherein the attention tensor is generated in an inference stage of the transformer model, and the quantized attention tensor is used in a subsequent inference stage of the transformer model.

13. The one or more non-transitory computer-readable media of claim 11 or 12, wherein the attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels, and segmenting the attention tensor comprises: segmenting the attention tensor into the groups along both the token dimension and the channel dimension, wherein a data precision of a group in the quantized attention tensor is higher than a data precision of one or more other groups in the quantized attention tensor.

14. The one or more non-transitory computer-readable media of claim 13, wherein a size of the group along the channel dimension is larger than a size of the one or more other groups along the channel dimension.

15. The one or more non-transitory computer-readable media of claim 14, wherein the size of the group along the channel dimension is the same as a size of the attention tensor along the channel dimension.

16. The one or more non-transitory computer-readable media of any one of claims 13-15, wherein the group includes a key or value corresponding to a token predicted by the transformer model in the inference stage.

17. The one or more non-transitory computer-readable media of any one of claims 11-16, wherein the attention tensor is a key tensor, and the operations further comprise: segmenting a value tensor into groups, wherein a shape of a group in the value tensor is different from a shape of a group in the key tensor.

18. The one or more non-transitory computer-readable media of any one of claims 12-16, wherein segmenting the attention tensor comprises: determining a group size based on a constraint of an apparatus executing the subsequent inference stage of the transformer model or based on an estimated accuracy of the transformer model when the attention tensor is quantized using the group size, the group size indicating a number of tokens or a number of channels in each group.

19. The one or more non-transitory computer-readable media of any one of claims 12-18, wherein the subsequent inference stage of the transformer model is performed using the quantized attention tensor and a query tensor, the query tensor comprising data elements that have a higher data precision than data elements in the quantized attention tensor.

20. The one or more non-transitory computer-readable media of any one of claims 12-19, wherein the subsequent inference stage of the transformer model comprises a matrix multiplication (MatMul) operation that is partitioned into a plurality of sub-MatMul operations and is performed by: performing each sub-MatMul operation using a different group; and accumulating results of the plurality of sub-MatMul operations.

21. An apparatus, comprising: a computer processor for executing computer program instructions; anda non-transitory computer-readable memory storing computer program instructions executable by the computer processor to perform operations for executing a transformer model, the operations comprising: segmenting an attention tensor of a transformer model into groups, the attention tensor comprising data elements that are keys or values, determining quantization parameters for the groups, different groups having different quantization parameters, data elements within a same group having a same quantization parameter, quantizing the data elements based on the quantization parameters, a data element quantized based on a quantization parameter of a group that includes the data element, generating a quantized attention tensor comprising the quantized data elements, and caching the quantized attention tensor in a storage unit.

22. The apparatus of claim 21, wherein the attention tensor is generated in an inference stage of the transformer model, and the quantized attention tensor is used in a subsequent inference stage of the transformer model.

23. The apparatus of claim 21 or 22, wherein the attention tensor has a token dimension corresponding to a number of tokens and a channel dimension corresponding to a number of channels, and segmenting the attention tensor comprises: segmenting the attention tensor into the groups along both the token dimension and the channel dimension.

24. The apparatus of claim 23, wherein: a size of a group along the channel dimension is larger than a size of one or more other groups along the channel dimension, a data precision of the group in the quantized attention tensor is higher than a data precision of the one or more other groups in the quantized attention tensor, and the group includes a key or value corresponding to a token predicted by the transformer model in the inference stage.

25. The apparatus of any one of claims 22-24, wherein segmenting the attention tensor comprises: determining a group size based on a constraint of an apparatus executing the subsequent inference stage of the transformer model and an estimated accuracy of the transformer model when the attention tensor is quantized using the group size, the group size indicating a number of tokens or a number of channels in each group.

Images