Continuous Batching: Efficiently Optimizing LLM Inference

The Problem of Static Batching

In the field of large language model (LLM) inference, static batching is often used to process multiple requests simultaneously. This method involves grouping requests into fixed-size batches, processing each batch together. However, this approach has significant limitations in terms of efficiency.

Let’s consider three distinct requests:

• A: "The capital of France is" which generates 6 additional tokens.
• B: "Today's weather is so" which generates 50 additional tokens.
• C: "In machine learning, a transformer is" which generates 300 additional tokens.

In this scenario, requests A and B finish well before request C. Yet, their slots remain occupied and inactive while the GPU continues processing the remaining tokens for C. This leads to unnecessary consumption of computing cycles.

During the decoding process, the situation is as follows:

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 …300

A <BOS> <The> <capital> <of> <France> <is> <the> <capital> <of> <the> <Republic> <EOS> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> …<PAD>

B <BOS> <Today’s> <weather> <is> <so> <cold> <that> <it> <is> <difficult> <to> <see> <the> <sun> <.> <But> <this> <isn’t> <like> <we> <are> <going> <…>

C <BOS> <In> <machine> <learning> <,> <a> <transformer> <is> <a> <type> <of> <machine> <learning> <algorithm> <that> <can> <be> <used> <for> <…>

Special tokens such as <BOS> (beginning of sentence), <EOS> (end of sentence), and <PAD> (padding) illustrate how request A is padded to fill up to the end of the 300 tokens in this batch. The GPU consumes computing cycles that yield no results with the padding tokens. Furthermore, the response to request A is likely not delivered until request C is completed.

Example of Static Batching Code

To illustrate how static batching works, let’s consider a GPT-2 model executing six requests of varying lengths with a batch size of 3. Each request is associated with a prompt and a maximum number of tokens to generate.

MODEL_ID = "[openai](/dossier/openai)-community/gpt2"
BATCH_SIZE = 3
requests = [
    ("The capital of France is", 6),
    ("Today's weather is so", 50),
    ("In machine learning, a transformer is", 300),
    ("Once upon a time in a faraway land,", 30),
    ("Quantum computing differs from classical computing because", 180),
    ("The history of the Roman Empire began", 45),
]

The static batching function processes requests in waves of a defined size, with each wave executing until the longest request is completed. Shorter requests remain inactive, and no new slots can be utilized until the entire wave is finished.

def static_batching(requests: list[tuple[str, int]], tokenizer, model) -> list[str]:
    if not requests:
        return []
    tokenizer.padding_side = "left"
    results: dict[int, str] = {}
    indexed = list(enumerate(requests))
    for wave_start in range(0, len(indexed), BATCH_SIZE):
        wave = indexed[wave_start: wave_start + BATCH_SIZE]
        wave_max = max(cap for _req_id, (_prompt, cap) in wave)
        prompts = [p for _, (p, _) in wave]
        inputs = tokenizer(
            prompts, return_tensors="pt", padding=True, truncation=True
        ).to(model.device)
        with torch.no_grad():
            output_ids = model.generate(
                **inputs,
                max_new_tokens=wave_max,
                pad_token_id=tokenizer.eos_token_id,
                do_sample=False,
            )
        width = inputs.input_ids.shape[1]
        for slot, ((req_id, (prompt, cap)), row) in enumerate(zip(wave, output_ids)):
            text = prompt + tokenizer.decode(row[width:width + cap], skip_special_tokens=True)
            results[req_id] = text
    return [results[k] for k in sorted(results)]

Running this code shows that shorter requests unnecessarily wait for the completion of the longest request, thus illustrating the inefficiencies of static batching.

Towards a More Efficient Solution

Continuous batching, with its dynamic scheduling and irregular batching, promises to improve these inefficiencies. By adapting batch sizes and optimizing request management, this approach could offer significant gains in performance and processing speed.

Conclusion

Static batching presents limitations in terms of efficiency, particularly due to the waiting of shorter requests. In the following sections, we will explore how continuous batching and dynamic scheduling can improve this situation.