Dynamic batching for inference system for transformer-based generation tasks
An inference system applies a machine-learning transformer model to a batch of requests with variable input length or variable target length or variable internal sate length by selectively batching a subset of operations in the transformer model but processing requests in the batch individually for a subset of operations in the transformer model. In one embodiment, the operation to be processed individually is an attention operation of an encoder or a decoder of the transformer model. By selective batching, the inference system can allow batching operations to be performed for a batch of requests with variable input or target length or internal state length to utilize the parallel computation capabilities of hardware accelerators while preventing unnecessary computations that occur for workarounds that restrain the data of a batch of requests to a same length.
1 . A non-transitory computer-readable storage medium storing computer program instructions executable to perform operations comprising:
receiving, by a serving system, one or more requests for execution, the serving system including one or more execution engines each coupled to access a machine-learning transformer model;
scheduling a batch of requests including the one or more requests for execution on an execution engine;
generating, by the execution engine, a first set of output tokens by applying the transformer model to a first set of inputs for the batch of requests, wherein applying the transformer model comprises applying at least one batch operation to one or more input tensors associated with the batch of requests;
obtaining a sequence of input tokens for a new request;
scheduling a second batch of requests for execution on the execution engine, the second batch of requests including the new request and at least one request in the batch of requests, wherein a length of an internal state for the new request is a length of keys or values for the new request of an attention operation of the machine-learning transformer model, wherein a length of another internal state for the at least one request is a length of keys or values for the at least one request of the attention operation, wherein the length of the internal state for the new request is different from the length of the another internal state for the at least one request; and
generating, by the execution engine, a second set of output tokens by applying the transformer model to a second set of inputs for the second batch, wherein a key or value is added to the internal state for the new request and another key or value is added to the another internal state for the at least one request.
2 . The non-transitory computer-readable storage medium of claim 1 , wherein an input for the at least one request in the second set of inputs comprises at least two tokens.
3 . The non-transitory computer-readable storage medium of claim 1 , wherein the at least two tokens for the at least one request comprises one or more output tokens from the first set of output tokens.
4 . The non-transitory computer-readable storage medium of claim 1 , the operations further comprising:
iteratively generating a third set of output tokens for the at least one request; and
responsive to determining that the third set of output tokens reaches a threshold number, providing the third set of output tokens.
5 . The non-transitory computer-readable storage medium of claim 1 , wherein the second batch of requests is scheduled responsive to determining that the execution engine has memory available to execute the second batch of requests.
6 . The non-transitory computer-readable storage medium of claim 1 , wherein the at least one request is associated with a cache in the execution engine for storing elements of the another internal state for the at least one request, and responsive to determining that the at least one request has been completed, freeing the cache for the at least one request in the execution engine.
7 . The non-transitory computer-readable storage medium of claim 1 ,
wherein the internal state includes a key cache storing keys generated for the new request or a value cache storing values generated for the new request,
wherein the another internal state includes another key cache storing keys generated for the at least one request or another value cache storing values generated for the at least one request, and
wherein a length of the key cache is different from a length of the another key cache or a length of the value cache is different from a length of the another value cache.
8 . The non-transitory computer-readable storage medium of claim 7 , wherein the length of the key cache is a number of keys for the new request and the length of the another key cache is a number of keys for the at least one request, and wherein the length of the value cache is a number of values for the new request and the length of the another value cache is a number of values for the at least one request.
9 . The non-transitory computer-readable storage medium of claim 1 , wherein the length of the sequence of input tokens for the new request is different from the length of the another sequence of input tokens for the at least one request.
10 . The non-transitory computer-readable storage medium of claim 9 , wherein the length of the sequence of input tokens is a number of input tokens in the sequence and the length of the another sequence of input tokens is a number of input tokens in the another sequence.
11 . The non-transitory computer-readable storage medium of claim 1 , wherein applying the at least one batch operation to the one or more input tensors associated with the batch of requests further comprises:
concatenating the one or more input tensors associated with the batch of requests into a concatenated input tensor; and
applying an operation for the batch operation to the concatenated input tensor to generate an output tensor.
12 . The non-transitory computer-readable storage medium of claim 11 , wherein concatenating the one or more input tensors comprises concatenating elements of the one or more input tensors along one dimension.
13 . The non-transitory computer-readable storage medium of claim 1 , wherein each token in the sequence of input tokens represents a text unit.
14 . The non-transitory computer-readable storage medium of claim 1 , wherein applying the transformer model to the first set of inputs further comprises applying an attention operation to the first set of inputs, and wherein the attention operation is configured to obtain queries, keys, values for the first set of inputs and generate attention outputs by combining the queries, the keys, and the values.
15 . The non-transitory computer-readable storage medium of claim 1 , wherein applying the transformer model to the second set of inputs further comprises applying an attention operation to the second set of inputs, and wherein the attention operation is configured to obtain queries, keys, values for the second set of inputs and generate attention outputs by combining the queries, the keys, and the values.
16 . The non-transitory computer-readable storage medium of claim 15 , wherein applying the attention operation to the second set of inputs further comprises:
generating a first attention output for the at least one request by combining a first query, a first key, and a first value for the at least one request, and
separately generating a second attention output for the new request by combining a second query, a second key, and a second value for the new request.
17 . The non-transitory computer-readable storage medium of claim 1 , wherein a length of the sequence of input tokens for the new request is different from a length of another sequence of input tokens for the at least one request.
18 . A method comprising:
receiving, by a serving system, one or more requests for execution, the serving system including one or more execution engines each coupled to access a machine-learning transformer model;
scheduling a batch of requests including the one or more requests for execution on an execution engine;
generating, by the execution engine, a first set of output tokens by applying the transformer model to a first set of inputs for the batch of requests, wherein applying the transformer model comprises applying at least one batch operation to one or more input tensors associated with the batch of requests;
obtaining a sequence of input tokens for a new request;
scheduling a second batch of requests for execution on the execution engine, the second batch of requests including the new request and at least one request in the batch of requests, wherein a length of an internal state for the new request is a length of keys or values for the new request of an attention operation of the machine-learning transformer model, wherein a length of another internal state for the at least one request is a length of keys or values for the at least one request of the attention operation, wherein the length of the internal state for the new request is different from the length of the another internal state for the at least one request; and
generating, by the execution engine, a second set of output tokens by applying the transformer model to a second set of inputs for the second batch, wherein a key or value is added to the internal state for the new request and another key or value is added to the another internal state for the at least one request.
19 . The method of claim 18 , wherein an input for the at least one request in the second set of inputs comprises at least two tokens.
20 . The method of claim 19 , wherein the at least two tokens for the at least one request comprises one or more output tokens from the first set of output tokens.
21 . The method of claim 18 , further comprising:
iteratively generating a third set of output tokens for the at least one request; and
responsive to determining that the third set of output tokens reaches a threshold number, providing the third set of output tokens.
22 . The method of claim 18 , wherein the second batch of requests is scheduled responsive to determining that the execution engine has memory available to execute the second batch of requests.
23 . The method of claim 18 , wherein the at least one request is associated with a cache in the execution engine for storing elements of the another internal state for the at least one request, and responsive to determining that the at least one request has been completed, freeing the cache for the at least one request in the execution engine.
24 . The method of claim 18 ,
wherein the internal state includes a key cache storing keys generated for the new request or a value cache storing values generated for the new request,
wherein the another internal state includes another key cache storing keys generated for the at least one request or another value cache storing values generated for the at least one request, and
wherein a length of the key cache is different from a length of the another key cache or a length of the value cache is different from a length of the another value cache.
25 . The method of claim 24 , wherein the length of the key cache is a number of keys for the new request and the length of the another key cache is a number of keys for the at least one request, and wherein the length of the value cache is a number of values for the new request and the length of the another value cache is a number of values for the at least one request.
26 . The method of claim 18 , wherein the length of the sequence of input tokens for the new request is different from the length of the another sequence of input tokens for the at least one request.
27 . The method of claim 26 , wherein the length of the sequence of input tokens is a number of input tokens in the sequence and the length of the another sequence of input tokens is a number of input tokens in the another sequence.
28 . The method of claim 18 , wherein applying the at least one batch operation to the one or more input tensors associated with the batch of requests further comprises:
concatenating the one or more input tensors associated with the batch of requests into a concatenated input tensor; and
applying an operation for the batch operation to the concatenated input tensor to generate an output tensor.
29 . The method of claim 28 , wherein concatenating the one or more input tensors comprises concatenating elements of the one or more input tensors along one dimension.
30 . The method of claim 18 , wherein each token in the sequence of input tokens represents a text unit.
31 . The method of claim 18 , wherein applying the transformer model to the first set of inputs further comprises applying an attention operation to the first set of inputs, and wherein the attention operation is configured to obtain queries, keys, values for the first set of inputs and generate attention outputs by combining the queries, the keys, and the values.
32 . The method of claim 18 , wherein applying the transformer model to the second set of inputs further comprises applying an attention operation to the second set of inputs, and wherein the attention operation is configured to obtain queries, keys, values for the second set of inputs and generate attention outputs by combining the queries, the keys, and the values.
33 . The method of claim 32 , wherein applying the attention operation to the second set of inputs further comprises:
generating a first attention output for the at least one request by combining a first query, a first key, and a first value for the at least one request, and separately generating a second attention output for the new request by combining a second query, a second key, and a second value for the new request.
34 . The method of claim 18 , wherein a length of the sequence of input tokens for the new request is different from a length of another sequence of input tokens for the at least one request.