S$^{3}$: Increasing GPU Utilization during Generative Inference for Higher Throughput

TL;DR

This work tackles the memory-bound nature of generative LLM inference caused by the growing KV cache, which constrains batch size and GPU throughput. It introduces S3, a system comprising an output-length predictor, a length-aware scheduler, and a supervisor that handles mispredictions, to allocate memory precisely and batch requests intelligently. Across multiple models and online/offline scenarios, S3 achieves up to 6.49x throughput improvements over worst-case baselines while maintaining latency SLOs, and can reduce required GPU counts for comparable throughput. The approach expands the latency-throughput trade-off frontier and offers practical cost benefits for deploying Transformer-based generative models at scale.

Abstract

Generating texts with a large language model (LLM) consumes massive amounts of memory. Apart from the already-large model parameters, the key/value (KV) cache that holds information about previous tokens in a sequence can grow to be even larger than the model itself. This problem is exacerbated in one of the current LLM serving frameworks which reserves the maximum sequence length of memory for the KV cache to guarantee generating a complete sequence as they do not know the output sequence length. This restricts us to use a smaller batch size leading to lower GPU utilization and above all, lower throughput. We argue that designing a system with a priori knowledge of the output sequence can mitigate this problem. To this end, we propose S$^{3}$, which predicts the output sequence length, schedules generation queries based on the prediction to increase device resource utilization and throughput, and handle mispredictions. Our proposed method achieves 6.49$\times$ throughput over those systems that assume the worst case for the output sequence length.

Figures (6)

• Figure 1: Latency versus throughput trade-off among different models (left) and the number of GPUs (right, distributing GPT-3 to 6, 8, and 10 GPUs) when generating 60 tokens, inspired by FlexGen flexgen. The markers in the lines represent batch sizes, from 1 to the maximum batch size that can be loaded on an A100 GPU, incrementing by the power of two. Allocating the exact amount of memory for each sequence expands the curve to higher throughput at the cost of higher latency. The solid lines show the trade-off that vanilla systems face and dotted lines show how much S3 can expand the trade-off. The numbers represent the maximum batch sizes for S3 and vanilla system. The vertical line in the left figure denotes the latency SLO for reading a 60-token long sequence.
• Figure 2: (a) GPT-J's and (b)GPT-NEOX's NVIDIA A100 compute utilization.
• Figure 3: Overview of S3. The boxes in yellow denote new components proposed by S3.
• Figure 4: Latency and throughput of different models.
• Figure 5: Maximum throughput of GPT3 running on different numbers of GPUs.