Source-linked AI summary
S$^{3}$: Increasing GPU Utilization during Generative Inference for Higher Throughput
Yunho Jin, Chun-Feng Wu, David Brooks, Gu-Yeon Wei
TL;DR
Transformer-based LLM serving is constrained by memory used for growing KV caches, and maximum-length reservations can restrict batch size and throughput. S3 predicts output lengths, allocates memory and schedules queries accordingly, and handles mispredictions; it generates up to 6.49× more sequences under the same latency SLO.
Problem
Growing KV caches constrain batch size and GPU utilization, while reserving maximum sequence lengths wastes memory when actual outputs are shorter.
Source: §1 · p. 1Method
S3 predicts output sequence lengths, allocates corresponding memory, schedules queries, and adjusts allocations when predictions are wrong.
Source: §1 · pp. 2–3, §3 · p. 4Results
6.49× more sequences are generated under the same average-reading-speed latency SLO.
Source: §1 · p. 2Takeaways & Limitations
S3 expands the latency-throughput trade-off frontier by allocating varying memory sizes to different inputs.
Source: §6 · p. 9Evaluation assumes request traces resemble public question-and-answer datasets and uses average English reading speed as the latency SLO because text-generation traces and standardized SLOs are unavailable.
Source: §4.1 · p. 7, §6 · p. 9
Abstract
from arXiv · showhide
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.
1 Introduction
Transformer-based LLM generation is memory-bound: growing KV caches constrain batch size and GPU utilization. S3 predicts output lengths, allocates memory accordingly, schedules queries, and handles prediction errors to improve throughput.
1 Introduction
- Autoregressive generation makes the KV cache grow with batch size and sequence length, limiting maximum batch size and reducing GPU utilization and throughput.Source: §1 · pp. 1–2
- S3 predicts each query’s output length, allocates corresponding memory, schedules sequences, and adjusts allocations when predictions are wrong.Source: §1 · p. 2
- S3 targets higher throughput under latency SLOs in both online and offline settings, which differ in how tightly they constrain latency.Source: §1 · p. 2
- 6.49× more sequences are generated under the same average-reading-speed latency SLO.Source: §1 · p. 2
- With 6 GPUs, S3 provides almost identical throughput to a vanilla system with 10 GPUs.Source: §1 · p. 3
- S3 preempts sequences that exceed allocated memory and retrains its predictor using those mispredictions.Source: §1 · p. 3
2 Background and Motivation
Autoregressive generation retains past-token information in a growing KV cache, so memory limits constrain batch size and GPU utilization. Existing allocation strategies trade frequent memory allocation for wasted reservations; knowing output lengths could address both problems.
2.1 Generative AI Models
- Autoregressive models generate one token per iteration, using self-attention to attend to the current token and previously generated tokens.Source: §2.1 · p. 3
- A KV cache stores prior-token keys and values to avoid recomputing them; in GPT-NEOX, it uses 1MB per token.Source: §2.1 · p. 3
2.2 KV Cache Management on GPUs
- HF-Transformers allocates cache memory as sequences grow, while FasterTransformer reserves the maximum sequence length in advance.Source: §2.2 · p. 3
- For GPT-NEOX on an 80GB A100, FasterTransformer’s 2.2GB-per-sequence reservation limits batch size to fewer than 20, even when sequences are short.Source: §2.2 · pp. 3–4
2.3 Observation
- GPT-J reaches only 12.19% GPU utilization at sequence length 1024 with a maximum batch of 8 sequences; GPT-NEOX also hits memory limits at smaller batches and shorter sequences.Source: §2.3 · p. 4, Figure 2 · p. 4
- Larger batches can activate idle GPU compute, but the memory cliff prevents further utilization gains.Source: §2.3 · p. 4
- Knowing generated sequence lengths would enable exact cache allocation, avoiding both repeated allocations and unnecessary reservations.Source: §2.3 · p. 4
3 S3 Design
S3 combines output-length prediction, HBM-aware scheduling, and supervision to maximize GPU utilization while handling prediction errors.
3 S3 Design
- S3 predicts output lengths, schedules requests into GPU batches using those predictions, and uses a supervisor to monitor HBM and handle mispredictions.The predictor, scheduler, and supervisor form the framework's three components.Source: §3 · pp. 4, 6
- DistilBERT classifies output lengths into 10 buckets; bucket accuracy is 98.61% on Alpaca, 77.13% on Natural-Question, and 65.6% on Pile.The predictor is fine-tuned using questions and answer lengths from Alpaca, with separate models evaluated on Natural-Question and Pile.Source: §3 · p. 5
- The scheduler treats available HBM as a single bin and sequences' KV-cache sizes as item weights, then applies decreasing first fit [22].It orders requests by decreasing size and adds each sequence when its KV cache fits in the remaining HBM.Source: §3 · p. 5
- For unequal-length batches, S3 follows ORCA's selective batching and iteration-level scheduling, avoiding padding and waiting for every sequence to finish.Selective batching avoids padding short sequences, while iteration-level scheduling checks for finished sequences at each iteration.Source: §3 · pp. 5–6
- The supervisor admits requests when HBM permits and repairs short predictions by evicting overrun sequences, compacting KV-cache rows, and doubling their reserved memory.It moves evicted sequence state back to the request pool before freeing GPU memory.Source: §3 · p. 6
- The supervisor retrains the predictor in the background on mispredicted sequences; each iteration averages 11ms, and training for 10 epochs adds less than 10% overhead.The paper compares retraining time with sequence generation, which takes several seconds or more.Source: §3 · p. 6
4 Evaluation
S3’s evaluation reports higher throughput than ORCA in offline and SLO-constrained online settings, compares GPU-count cost, and measures runtime overhead.
4.1 Throughput Analysis
- Offline, S3 outperforms ORCA by 1.13× and up to 6.49× across models, with differences from Oracle of 9.34% to 40.52%.Figure 4(a) reports maximum throughput in sequences per second for different models.Source: Figure 4 · p. 7, §4.1 · p. 7
4 Evaluation
- Evaluation uses Alpaca question-answering prompts on 6B–175B models on an 80GB A100, comparing FasterTransformer-based S3 with ORCA and a perfect-predictor Oracle.Source: §4 · pp. 6–7, §4.1 · p. 7
4.1 Throughput Analysis
- Feed-forward layers benefit from shared-weight batching, whereas sequence-specific self-attention does not; gains are larger when parallelized work dominates serialized work.Source: §4.1 · p. 7
4.2 Cost Analysis
- 6-GPU S3 delivers maximum throughput similar to 10-GPU ORCA, which is 0.92% higher, despite S3’s 2.33× larger batch.The higher latency offsets S3’s batch-size advantage; pipelined execution schedules batches as GPUs process them concurrently.Source: Figure 5 · p. 8, §4.2 · p. 8
4.3 Overhead: Latency Breakdown
- Penalty plus predictor, scheduler, and supervisor overhead averages 11% compared with generation; lower prediction accuracy raises transfer penalty, while heavier predictors raise overhead.With a decent predictor, overhead is amortized as the model generates longer sequences.Source: Figure 6 · p. 8, §4.3 · pp. 8–9
4.1 Throughput Analysis
- Online, S3 generates 1.13× to 6.49× more sequences than ORCA while meeting the latency SLO in most cases.Oracle’s batch size was limited for GPT-J, LLAMA 13B, and GPT-NEOX when its maximum batch size exceeded the SLO; the SLO was 0.1875 seconds per token.Source: §4.1 · pp. 7–8
5 Related Works
Prior work addresses memory use and sequence-length prediction, but throughput-oriented serving systems do not preempt sequences based on predictions. S3 combines prediction-based scheduling with existing KV-cache reduction approaches.
Machine learning serving systems
- Throughput-oriented serving systems allocate the same memory to every sequence and do not preempt sequences based on predictions, while S3 can reduce their memory requirements.The paper says this could remove the need for larger, slower memories such as SSDs and reduce memory cost.Source: §5 · p. 9
Reducing the KV cache overhead
- KV-cache reduction methods—including low-rank approximation, sparsity, multi-query attention, compression, and quantization—are complementary to S3 and can reduce the penalty from mispredictions.The paper says these methods could allow use of a less accurate but lighter predictor.Source: §5 · p. 9
Sequence length prediction
- Yan et al. predict output length for machine translation with a convolution-based network, whereas S3 uses a small Transformer predictor for text-generation workloads.The paper describes this predictor as estimating output sequence lengths for text generation.Source: §5 · p. 9
6 Limitations and Conclusion
The evaluation assumes question-and-answering-like request traces and an average-reading-speed latency SLO. S3 predicts output lengths, schedules sequences accordingly, and handles prediction errors to pursue high throughput in generative-model serving.
Limitations
- The evaluation assumes request traces mimic public question-and-answering datasets because public text-generation traces are unavailable.The authors say analyzing actual traces would facilitate deployment in commercial workloads.Source: §6 · p. 9
- Because text generation lacks standardized latency SLOs, the evaluation uses the average reading speed of an English reader as its SLO.The authors say broader evaluation depends on organizations releasing different SLOs for text-generation applications.Source: §6 · p. 9
Conclusion
- S3 estimates output lengths, schedules sequences accordingly, and handles prediction errors while allocating different memory sizes to different inputs.The conclusion says this approach expands the conventional latency-throughput trade-off frontier.Source: §6 · p. 9