Selective Batching for Inference System for Transformer-Based Generation Tasks

A method of executing a batch of requests by applying a machine-learning transformer model, comprising:

• receiving a batch of requests including one or more token sequences, wherein a length of a first token sequence for a first request in the batch is different from a length of a second token sequence for a second request in the batch;
• accessing a machine-learning transformer model including at least a set of decoders;
• for one or more iterations, repeatedly performing the steps of:
• generating one or more output tokens for the requests by applying the set of decoders to one or more inputs for the requests, the generating further comprising, for at least one decoder in the set:
• generating one or more queries, one or more keys, and one or more values for the requests by applying a QKV weight tensor to one or more input representations, the one or more queries, the one or more keys, and the one or more values generated by a batch operation,
• splitting a first query for the first request from the one or more queries, a first key for the first request from the one or more keys, and a first value for the first request from the one or more values, and splitting a second query for the second request from the one or more queries, a second key for the second request from the one or more keys, and a second value for the second request from the one or more values,
• generating a first attention output for the first request by at least combining the first query, the first key, and the first value,
• separately generating a second attention output for the second request by at least combining the second query, the second key, and the second value for the second request,
• concatenating at least the first attention output and the second attention output into a concatenated tensor, and
• generating one or more output representations by at least applying one or more weight tensors to the concatenated tensor, the one or more output representations generated by a batch operation;
• setting the one or more output tokens as the one or more inputs to the set of decoders for a next iteration;
• and
• providing output tokens generated for at least one request to a client device as a response to the at least one request.