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.

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 using variable draft length. An example method generally includes determining (e.g., measuring or accessing) one or more operational properties of a device on which inferencing operations using a machine learning model are performed. A first draft set of tokens is generated using the machine learning model. A number of tokens included in the first draft set of tokens is based on the one or more operational properties of the device and a defined scheduling function for the machine learning model. The first draft set of tokens are output for verification.

Abstract: Certain aspects of the present disclosure provide techniques and apparatus for efficiently generating a response to a query input in a generative artificial intelligence model. An example method generally includes generating, based on an input prompt and using a first machine learning model, a set of tokens including one or more subsets of tokens. Each respective subset of the one or more subsets corresponds to a respective portion of a response to the input prompt and includes a fixed number of tokens corresponding to a beam width for a beam search through the set of tokens. The set of tokens is output to a second machine learning model for verification, and information identifying a selected sequence of tokens from the generated set of tokens is received from the second machine learning model. The selected sequence of tokens is output as the response to the input prompt.

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: 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.