Compute management engine in an artificial intelligence system
Methods, systems, and computer storage media for providing compute management using a compute management engine in an artificial intelligence (AI) system. A compute management engine supports dynamically switching between two modes of operation for an inference phase of a generative artificial AI model. The compute management engine employs a bypass engine that causes prompt stage operations to be executed without an in-memory compute engine and causes auto-regression stage operations to be executed with the in-memory compute engine. In operation, an inference phase operation is accessed. When the inference phase operation is a prompt stage operation, the inference phase operation is executed without an in-memory compute engine. When the inference phase operation is an auto-regressive stage operation, the inference phase operation is executed with the in-memory compute engine. Memory output is generated for the inference phase operation to cause a processor to output a processor output for the inference phase operation.
1 . A computerized system comprising:
one or more computer processors; and
computer memory storing computer-useable instructions that, when used by the one or more computer processors, cause the one or more computer processors to perform operations, the operations comprising:
accessing a prompt stage operation associated with a generative artificial intelligence (AI) model, the generative AI model is associated with memory that supports executing generative AI model operations associated with the generative AI model;
wherein the prompt stage operation is associated with a prompt stage of an inference phase of the generative AI model;
based on the prompt stage operation, dynamically causing the prompt stage operation to bypass an in-memory compute engine of the memory;
executing the prompt stage operation without using the in-memory compute engine;
accessing an auto-regressive stage operation associated with the generative AI model,
wherein the auto-regression stage operation is associated with an auto-regression stage of the inference phase of the generative AI model;
based on the auto-regressive stage operation, dynamically causing the auto-regressive stage operation to be executed using the in-memory compute engine of the memory; and
executing the auto-regressive stage operation using the in-memory compute engine.
2 . The system of claim 1 , wherein the generative AI model is a large language model (“LLM”) associated with a plurality of operations comprising vector-matrix operations, wherein a vector or a matrix of the plurality of operations is associated with a word embedding, a model parameter, or an intermediate representation of text.
3 . The system of claim 1 , wherein the prompt stage operation or the auto-regressive stage operation is associated with a dot product operation.
4 . The system of claim 1 , wherein the prompt stage operation supports processing input as a batch of new tokens, and wherein the auto-regressive stage operation supports processing individual tokens generated via the generative AI model.
5 . The system of claim 1 , wherein bypassing the in-memory compute engine for the prompt stage operation is based on a bypass engine that implements one or more of: a dedicated signal, a control register, a pre-defined address map, or two types of read commands that support identifying prompt stage operations and auto-regressive stage operations.
6 . The system of claim 1 , wherein executing the prompt stage operation without using the in-memory compute engine comprises executing memory readouts from the memory without using the in-memory compute engine and communicating data to a processor associated with a processor compute engine.
7 . The system of claim 1 , wherein causing the auto-regressive stage operation to be executed using the in-memory compute engine comprises communicating the auto-regressive operation to the in-memory compute engine.
8 . The system of claim 1 , wherein executing the auto-regressive stage operation using the in-memory compute engine comprises computing a dot product result of a first vector and a second vector.
9 . The system of claim 1 , the operation further comprising:
communicating an inference phase operation associated with the generative artificial intelligence (AI) model, the inference phase operation is communicated to the memory;
based on the communicating the inference phase operation, receiving a memory output associated with the inference phase operation, wherein the memory output is generated without the in-memory compute engine of the memory when the inference phase operation is the prompt stage operation, and wherein the memory output is generated with the in-memory compute engine of the memory when the inference phase operation is the auto-regressive stage operation; and
using the memory output, generating a processor output associated with the inference phase operation.
10 . The system of claim 1 , further comprising a bypass engine configured to execute the operations comprising:
accessing a first inference phase operation associated with the generative artificial intelligence model;
determining that the first inference phase operation is the prompt stage operation;
based on determining that the first inference phase operation is the prompt stage operation, causing the prompt stage operation to bypass the in-memory compute engine;
accessing a second inference phase operation associated with the generative AI model;
determining that the second inference phase operation is the auto-regressive stage operation; and
based on determining that the second inference phase operation is the auto-regressive stage operation, causing the auto-regressive stage operation to be executed with the in-memory compute engine.
11 . One or more computer-storage media having computer-executable instructions embodied thereon that, when executed by a computing system having a processor and memory, cause the processor to perform operations, the operations comprising:
communicating an inference phase operation associated with a generative artificial intelligence (AI) model, the inference phase operation is communicated to memory that supports executing generative AI model operations associated with the generative AI model;
based on the communicating the inference phase operation, receiving a memory output associated with the inference phase operation, wherein the memory output is generated without an in-memory compute engine of the memory when the inference phase operation is a prompt stage operation,
wherein the prompt stage operation is associated with a prompt stage of an inference phase of the generative AI model,
and wherein the memory output is generated with the in-memory compute engine of the memory when the inference phase operation is an auto-regressive stage operation
wherein the auto-regression stage operation is associated with an auto-regression stage of the inference phase of the generative AI model; and
using the memory output, generating a processor output associated with the inference phase operation.
12 . The media of claim 11 , wherein the inference phase operation is a processor compute engine that executes the inference phase operation using the memory output generated without the in-memory compute engine.
13 . The media of claim 11 , wherein, when the inference phase operation is the prompt stage operation, the memory output is a memory readout that is not associated with the in-memory compute engine.
14 . The media of claim 11 , wherein, when the inference phase operation is the auto-regressive stage operation, the memory output is a dot product of a first vector and a second vector, wherein the dot product is executed using the in-memory compute engine.
15 . The media of claim 11 , the operations further comprising:
accessing a first inference phase operation associated with the generative artificial intelligence model;
determining that the first inference phase operation is the prompt stage operation;
based on determining that the first inference phase operation is the prompt stage operation, causing the prompt stage operation to bypass the in-memory compute engine;
accessing a second inference phase operation associated with the generative AI model;
determining that the second inference phase operation is the auto-regressive stage operation; and
based on determining that the second inference phase operation is the auto-regressive stage operation, causing the auto-regressive stage operation to be executed with the in-memory compute engine.
16 . A computer-implemented method, the method comprising:
accessing a first inference phase operation associated with a generative artificial intelligence (AI) model, the generative AI model is associated with memory that support executing operations associated with the generative AI model;
determining that the first inference phase operation is a prompt stage operation, wherein the prompt stage operation is associated with a prompt stage of an inference phase of the generative AI model;
based on determining that the first inference phase operation is the prompt stage operation, dynamically causing the prompt stage operation to bypass an in-memory compute engine;
accessing a second phase inference operation associated with generative AI model;
determining that the second inference phase operation is an auto-regressive stage operation, wherein the auto-regression stage operation is associated with an auto-regression stage of the inference phase of the generative AI model; and
based on determining that the second inference phase operation is the auto-regressive stage operation, dynamically causing the auto-regressive stage operation to be executed with the in-memory compute engine.
17 . The method of claim 16 , wherein bypassing the in-memory compute engine for the prompt stage operation is based on a bypass engine that implements one or more of: a dedicated signal, a control register, a pre-defined address map, or two types of read commands corresponding to prompt stage operations and auto-regressive stage operations.
18 . The method of claim 16 , the method further comprising executing the first inference phase operation without using the in-memory compute engine.
19 . The method of claim 16 , the method further comprising executing the second inference phase operation using the in-memory compute engine.
20 . The method of claim 16 , wherein the prompt stage operation supports processing input as a batch of new tokens, and wherein the auto-regressive stage operation supports processing individual tokens generated via the generative AI model.