Abstract: Certain aspects of the present disclosure provide techniques and apparatus for generating a response to a query input in a generative artificial intelligence model. An example method generally includes receiving a plurality of sets of tokens generated based on an input prompt and a first generative artificial intelligence model, each set of tokens in the plurality of sets of tokens corresponding to a candidate response to the input prompt; selecting, using a second generative artificial intelligence model and recursive adjustment of a target distribution associated with the received plurality of sets of tokens, a set of tokens from the plurality of sets of tokens; and outputting the selected set of tokens as a response to the input prompt.

Abstract: A processing-in-memory (PIM) device implements block quantization techniques for matrix-vector operations. The PIM device performs matrix-vector operations between portions of a weight matrix and an input vector, and copies results to a register. A read operation retrieves the copied results while additional matrix-vector operations are performed in parallel. The device may apply scaling factors to the results using multipliers within the PIM device. In some implementations, the weight matrix includes data columns and scaling factor columns interspersed at regular intervals. The scaling factors may be applied to accumulated results using parallel multiplication operations. Disclosed techniques enable efficient implementation of block quantization for applications such as Large Language Models while managing computational resources within the PIM architecture.

Abstract: Techniques and apparatus for generating a response to an input prompt using efficient self-speculative decoding in a generative artificial intelligence model. An example method generally includes receiving an input prompt for processing. A forecast embedding representing one or more forecasted tokens responsive to the input prompt is generated. Generally, the one or more forecasted tokens include tokens speculatively decoded by a generative artificial intelligence model based on generation of an initial response token in response to the input prompt. A bias parameter for the input prompt is determined. Generally, the bias parameter includes an embedding representation representing an error metric between the one or more forecasted tokens and an accepted set of tokens responsive to the input prompt. Using the generative artificial intelligence model, a response to the input prompt is generated based on the input prompt, the forecast embedding, and the bias parameter, and the generated response is output.

Abstract: Certain aspects of the present disclosure provide techniques and apparatus for machine learning. In an example method, a memory associated with a generative machine learning model is accessed while data is processed using the generative machine learning model. The memory stores, for each token of a set of tokens for, or generated by, the generative machine learning model, a respective key tensor and a respective value tensor corresponding to the token. An average key tensor of the key tensors corresponding to the set of tokens is determined. A respective cosine similarity metric is generated for each key tensor, based on the key tensor and the average key tensor. At least one key tensor corresponding to at least one token of the set of tokens is evicted from the memory in response to determining that the respective cosine similarity metric for the at least one key tensor satisfies a condition.

Abstract: Certain aspects of the present disclosure provide techniques and apparatus for generating a response to a query input into a generative artificial intelligence model. The method generally includes generating, based on an input query and a first generative model, a plurality of sets of tokens, each set of tokens in the plurality of sets of tokens corresponding to a candidate response to the input query; outputting, to a second generative model, the plurality of sets of tokens for verification; receiving, from the second generative model, an indication of a selected set of tokens from the plurality of sets of tokens based on the input query and the plurality of sets of tokens; and outputting the selected set of tokens as a response to the input query.

Abstract: Disclosed are systems, apparatuses, processes, and computer-readable media for model training. A device may process, using a linear layer, an embedding generated from a first output token and input features to generate first features, wherein the first output token is generated by a previous iteration of a token predictor and wherein the input features are generated by a previous iteration of a decoding layer. A device may process, using the decoding layer, the first features to generate second features having first dimensions. A device may process, using a down-projection layer, the second features to generate third features having second dimensions smaller than the first dimensions. A device may generate, using the token predictor and the third features, a second output token.

Abstract: Disclosed are systems, apparatuses, processes, and computer-readable media for model training. A device may process, using a linear layer, an embedding generated from a first output token and input features to generate first features, wherein the first output token is generated by a previous iteration of a token predictor and wherein the input features are generated by a previous iteration of a decoding layer. A device may process, using the decoding layer, the first features to generate second features having first dimensions. A device may process, using a down-projection layer, the second features to generate third features having second dimensions smaller than the first dimensions. A device may generate, using the token predictor and the third features, a second output token.

Abstract: Certain aspects provide techniques and apparatus for executing queries in a computing system using machine learning models. An example method generally includes receiving a plan to satisfy a request in the computing system and event log data associated with execution of the plan. The plan generally specifies a first plurality of function calls at a first level of granularity. Using a plan refinement machine learning model, a refined plan is generated when the event log data indicates that execution of the generated plan results in one or more execution errors and the one or more execution errors are solvable. Generally, the refined plan specifies a second plurality of function calls at a second level of granularity, the second level of granularity being finer than the first level of granularity.

Abstract: Certain aspects of the present disclosure provide techniques and apparatus for machine learning. In an example method, an input prompt comprising a set of tokens is accessed as input to a generative machine learning model. A first key tensor and a first value tensor are generated for a first token of the set of tokens, and the first key tensor and the first value tensor are stored in a memory. A first retention score is generated, for the first token, based on the first key tensor, the first value tensor, and a second token of the set of tokens. The first key tensor and the first value tensor are evicted from the memory in response to determining that the first retention score is a lowest retention score of the memory.

Abstract: Certain aspects of the present disclosure provide techniques and apparatus for generating a response to a query input in a generative artificial intelligence model. An example method generally includes receiving an input prompt for processing; generating a set of forecasted parameters for the input prompt using a parameter prediction model; generating, using a generative artificial intelligence model, a response to the input prompt based on the input prompt and the set of forecasted parameters; and outputting the generated response.

Abstract: Certain aspects of the present disclosure provide techniques and apparatus for efficient inferencing using a machine learning model. An example method generally includes receiving an input including a set of tokens for processing by a transformer neural network. The set of tokens for processing by the transformer neural network is partitioned into a first set of tokens and a second set of tokens. Using at least one state space model, at least one compressed token representing the first set of tokens is generated. An output token is generated, using the transformer neural network, based on the compressed token and the second set of tokens. A response to the input is generated based on the output token.

Abstract: Certain aspects of the present disclosure provide techniques and apparatus for machine learning. In an example method, a machine learning model comprising a plurality of layers, and a set of input data for the machine learning model, are accessed. A combination of hyperparameters for the machine learning model is selected based on the set of input data, comprising selecting, for each respective layer of the plurality of layers, a respective cache size based on the input data. The machine learning model is deployed according to the combination of hyperparameters.

Abstract: Certain aspects of the present disclosure provide techniques and apparatus for machine learning. In an example method, an input prompt comprising a set of tokens is accessed as input to a generative machine learning model. A first key tensor and a first value tensor are generated for a first token of the set of tokens, and the first key tensor and the first value tensor are stored in a memory. A first retention score is generated, for the first token, based on the first key tensor, the first value tensor, and a second token of the set of tokens. The first key tensor and the first value tensor are evicted from the memory in response to determining that the first retention score is a lowest retention score of the memory.

Abstract: Certain aspects of the present disclosure provide techniques and apparatus for machine learning. In an example method, an input prompt comprising a set of tokens is accessed as input to a generative machine learning model. A first key tensor and a first value tensor are generated for a first token of the set of tokens, and the first key tensor and the first value tensor are stored in a memory. A first retention score is generated, for the first token, based on the first key tensor, the first value tensor, and a second token of the set of tokens. The first key tensor and the first value tensor are evicted from the memory in response to determining that the first retention score is a lowest retention score of the memory.

Abstract: Techniques and apparatus for efficiently adapting a machine learning model to perform a variety of tasks using different adapters are provided. An example method generally includes receiving an input including a sequence of tokens associated with at least an input prompt into a neural network. The sequence of tokens is generated by a transformer block and a first set of adapters associated with the transformer block. A second set of adapters associated with the transformer block is loaded. An output of the transformer block is generated based on a key-value cache associated with the input and on weights associated with the transformer block. An output of the second set of adapters associated with the transformer block is generated based on the key-value cache associated with the input and on adapter weights associated with the second set of adapters.

Abstract: Techniques and apparatus for efficiently adapting a machine learning model to perform tasks using adapters are provided. An example method generally includes receiving an input for processing by a transformer block in a neural network. An output of the transformer block is generated based on the received input and weights associated with the transformer block. An output of an adapter associated with the transformer block is generated based on a copy of the received input and adapter weights associated with the adapter. Key-value data associated with the output of the transformer block and key-value data associated with a combination of the output of the transformer block and the output of the adapter are stored in a cache for subsequent inferencing rounds. A response to the input is generated based on the combination of the output of the transformer block and the output of the adapter.

Abstract: Certain aspects provide techniques and apparatus for executing queries in a computing system using machine learning models. An example method generally includes receiving a plan to satisfy a request in the computing system and event log data associated with execution of the plan. The plan generally specifies a first plurality of actions to be performed by the computing system at a first level of granularity. Using a plan refinement machine learning model, a refined plan is generated when the event log data indicates that execution of the generated plan results in one or more execution errors and the one or more execution errors are solvable. Generally, the refined plan specifies a second plurality of actions to be performed by the computing system at a second level of granularity, the second level of granularity being finer than the first level of granularity.

Abstract: A method for optimizing the compilation of a machine learning model to be executed on target edge devices is provided. Compute nodes of a plurality of compute nodes are allocated to a compiler optimization process for a compiler of said machine learning model. The machine learning model has a compute graph representation having nodes that are kernel operators necessary to execute the machine learning model and edges that connect said kernel operators to define precedence constraints. A round of optimization is scheduled for the process amongst the allocated compute nodes. At each allocated compute node a sequencing and scheduling solution is applied per round to obtain a performance metric for the machine learning model. From each compute node the performance metric is received and a solution that has the best performance metric is identified and implemented for execution of the machine learning model on the target edge devices.

Abstract: Certain aspects of the present disclosure provide techniques and apparatus for improved machine learning. In an example method, a sequence of tokens is accessed as input to an attention operation. For a first token, an attention output is generated based on a window of tokens relative to the first token, comprising generating a first positional embedding for an influential token, generating a second positional embedding for the first token, and generating the attention output based on the first and second positional embeddings. For a second token, an attention output is generated based on a window of tokens relative to the second token, where the second window of tokens includes the first token, comprising generating a third positional embedding for the influential token, generating a fourth positional embedding for the second token, and generating the attention output based on the second, third, and fourth positional embeddings.

Abstract: A method performed by a computing device includes determining a partition for depth-first processing by a multi-layer artificial neural network (ANN) of the computing device. The computing device comprising a processor, on-chip memory, and off-chip memory. The first partition determined based on an amount of on-chip memory used by the first partition, an available amount of on-chip memory, and a size of a write back to the off-chip memory. The method also includes processing, at the device via the multi-layer ANN, an input, using the depth-first processing in accordance with the partition.