LLM Inference Optimizations — Continuous Batching (Dynamic Batching)
Batching for large language models (LLMs) is challenging due to the iterative nature of their inference. The difficulty arises because requests within a batch can complete at different times, making it complex to release their resources and incorporate new requests that are at varying stages of completion. As a result, the GPU may be underutilized, especially when the generation lengths of sequences in a batch vary significantly from the longest sequence. This inefficiency is depicted in the figure below by the white squares following the end-of-sequence tokens for sequences 1, 3, and 4, indicating unused GPU resources.
Continuous Batching — ORCA: a distributed serving system for Transformer-based generative models
ORCA, which introduces the concept of Continuous Batching, features iteration-level scheduling and selective batching to effectively address the challenges associated with batching in large language models.
Feature #1 — Iteration-level scheduling
The system architecture and workflow of ORCA using iteration-level scheduling are illustrated in the figure below. ORCA exposes an endpoint (e.g., HTTPS or gRPC) for handling inference requests. These requests are placed in a request pool, which is monitored by a scheduler. The scheduler’s role is to select requests from the pool, schedule the execution engine to run a model iteration on them, receive output tokens from the engine, and update the pool accordingly. The engine performs tensor operations, potentially parallelized across multiple GPUs and machines. The scheduler dynamically decides which requests to process in each iteration, allowing for flexibility in handling incoming requests. Once a request is fully processed, it is removed from the pool, and a response is sent to the endpoint. Unlike traditional methods that process a batch of requests through multiple iterations until all are completed, ORCA’s scheduler can adjust which requests are processed in each iteration.
Feature #2 — Selective batching
Selective batching enhances the flexibility of batched execution in model processing by applying batching selectively, rather than universally across all tensor operations. This technique addresses the difficulty of handling irregularly shaped input tensors, which cannot be easily combined into a single large tensor for batch processing. In operations like non-Attention matrix multiplication and layer normalization, selective batching flattens these irregular tensors into a 2D structure without a batch dimension, enabling their joint processing as a single large 2D matrix. This approach is feasible because these operations function on a per-token basis, allowing tokens from different requests to be treated as if they belong to the same request. However, for the Attention operation, which requires unique key-value pairs for each request, batching is split, and each request is processed individually. Once the Attention operation completes, the outputs are merged back into a unified tensor for the subsequent operations. Additionally, ORCA maintains the keys and values generated by prior Attention operations within a K/V manager, allowing incremental requests to leverage them for computing attention across multiple iterations.
Distributed Architecture
ORCA combines two model parallelization techniques for Transformer models: intra-layer and inter-layer parallelism. Intra-layer parallelism, which is also used by FasterTransformer, divides matrix multiplications (such as Linear and Attention operations) and their parameters across multiple GPUs. Inter-layer parallelism, on the other hand, splits Transformer layers among multiple GPUs. ORCA assigns an equal number of Transformer layers to each GPU. For example, in a 4-layer GPT model, the layers can be divided into 2 inter-layer partitions, with each partition further divided into 3 intra-layer partitions, resulting in a total of 6 GPUs being used.
FasterTransformer is an open-source library developed by NVIDIA that optimizes and accelerates transformer models, supporting architectures like BERT, GPT-2, GPT-J, and T5. The development of FasterTransformer has now moved to TensorRT-LLM.
The ORCA execution engine facilitates distributed execution by assigning each worker process to a different machine, with each worker responsible for managing one or more CPU threads that control GPUs. Each worker handles an inter-layer partition of the model, and the number of threads is determined by the degree of intra-layer parallelism. When a model iteration is scheduled for a batch of requests, the engine master provides the first worker (Worker1) with the necessary information, including tokens and control messages. Worker1 then relays this information to its GPU-controlling threads, which issue the appropriate GPU kernels. Concurrently, Worker1 forwards the control message to the next worker (Worker2) without waiting for its own GPU tasks to complete. The final worker ensures that its GPU tasks are finished before collecting the output tokens and sending them back to the engine master. This approach enables efficient distributed execution across multiple machines and GPUs.
Scheduling Algorithm
The ORCA scheduler is designed to efficiently manage request processing by selecting a limited number of requests (up to a maximum batch size) based on their arrival time. When scheduling a request for the first time, it reserves memory slots for storing the necessary keys and values. The scheduling process involves selecting a batch of requests from a pool, where the Select function ensures that the selected requests are prioritized by arrival time. The scheduler also checks if sufficient memory slots are available for new requests, based on the maximum number of tokens required. If memory is available, it guarantees that the necessary buffers for keys and values can be allocated until the request is completed. The entire process is managed concurrently, with threads adding new requests and removing completed ones from the request pool.
The ORCA scheduler optimizes the execution of workers by pipelining their tasks across multiple batches. Instead of waiting for each batch to complete before scheduling the next, the scheduler ensures that the number of concurrently running batches equals the number of workers. This approach keeps all workers actively processing batches without any idle time. For example, with three workers and a maximum batch size of two, the scheduler processes requests in batches (e.g., AB, CD, EF). After scheduling the first batch (AB), the scheduler immediately schedules the next two batches (CD and EF) before waiting for the return of AB. Once batch AB is completed, those requests are re-scheduled if they are still the earliest in the request pool. This strategy maximizes worker utilization and throughput.
Evaluation
Thanks to its iteration-level scheduling and selective batching, ORCA achieves up to 37x higher throughput compared to FasterTransformer while maintaining the same level of latency.
References
- ORCA: A Distributed Serving System for Transformer-Based Generative Models
- https://www.anyscale.com/blog/continuous-batching-llm-inference
- https://discuss.pytorch.org/t/batched-matrix-vector-dot-product-attention/23461
- https://jaketae.github.io/study/seq2seq-attention/
