SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$\mathbf{S}^{3}$ : Increasing GPU Utilization during Generative Inference for Higher Throughput
Yunho Jin
Chun-Feng Wu
National Yang Ming
Chiao Tung University
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.
Text generation has become increasingly popular in various services, leading to a surge in the use of large language models (LLMs) known for their high accuracy. LLMs have distinct features that differentiate them from non-Transformer-based models, including a requirement for massive amounts of compute and memory resources. Specifically, we observe that Transformer-based LLMs are often limited by memory capacity and bandwidth, resulting in significant underutilization of compute resources. When serving GPT-J on an NVIDIA A100 GPU, the utilization of GPU compute resources can be as low as  $0.4\%$ . This highlights the memory-bound nature of the LLMs and the need for efficient memory utilization to increase the utilization of GPU compute resources.
A common approach to boost GPU utilization and enhance throughput is to increase batch size. This is due to the fact that inputs within a batch share the same model weights, thus the GPU only needs to load the model weight from its high bandwidth memory (HBM) to the on-chip SRAM once and reuse it for all inputs within the batch. The GPU uses more of its compute resources when processing the same model weights. Increasing batch size is a simple optimization technique and is, therefore, commonly used in serving convolutional and fully-connected neural network models [1,2].
However, the self-attention layer in Transformer-based text generation LLMs presents a challenge to this simple optimization due to its autoregressive nature. Specifically, when generating a new token in a sequence, the model needs to attend to all previous tokens in the sequence, requiring the model to retain all information from previous tokens and store them in HBM. We call this region in the HBM holding the information key/value cache (KV cache). The size of the KV cache grows with larger batch sizes and longer sequences, which limits the maximum batch size, thereby lowering GPU utilization and ultimately reducing throughput. To support the growing KV cache size, Huggingface's Transformers library [3] constantly allocates new memory at each token generation
Figure 1: Latency versus throughput trade-off among different models (left, online scenario with single GPU running models within the size confines of the GPU) and the number of GPUs (right, offline scenario distributing GPT-3 175B to 6, 8, and 10 GPUs) when generating 60 tokens, inspired by FlexGen [6]. 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 in vanilla systems and dotted lines show how much  $\mathbf{S}^3$  can expand the trade-off. The numbers represent the maximum batch sizes for  $\mathbf{S}^3$  and vanilla systems. The vertical line in the left figure denotes the latency SLO for reading a 60-token-long sequence.
and incurs latency overhead associated with memory allocation. This improves usability since users do not have to know the output sequence length but suffers long inference latency. Alternatively, NVIDIA's FasterTransformer library [4] pre-allocates memory for the maximum sequence length, which ends up wasting memory by allocating more than is necessary. This approach limits batch size and prioritizes latency over throughput. The trend towards long maximum sequence lengths in recent models (e.g., 8K or 32K) [5] amplifies the KV cache overhead, demanding more efficient memory utilization in serving Transformer-based text generation models.
In light of these challenges, we propose  $S^3$ , **scheduling** sequences with speculation, a framework that maximizes throughput via predicting the output sequence length and reducing memory waste. The frequent memory allocations in Huggingface's Transformers and the limited batch size in FasterTransformer stem from the fact that we lack prior knowledge of the output sequence length.  $S^3$  addresses these issues by predicting the expected output sequence length and allocating the corresponding amount of memory for each query. It also schedules sequences to increase the GPU utilization. Finally,  $S^3$  runs a supervisor in the background that detects mispredictions and adjusts the size of allocated memory to be more accurate. By integrating these components together,  $S^3$  optimizes memory usage and scheduling to maximize throughput during deployment of Transformer-based LLMs for text generation on GPUs.
There are two types of LLM deployment scenarios: online and offline. Online scenarios such as chatbots [7, 8] require service providers to generate a sequence within a tight latency service level objective (SLO) constraint. Offline scenarios include applications such as scoring [9] or data wrangling [10], and have loose latency SLOs, emphasizing throughput over end-to-end latency. In contrast, FaterTransformers [4] and xFormers [11] prioritize reducing latency. We argue that ensuring the latency remains below the SLO constraint renders it unnecessary to prioritize further latency reduction. To this end, we design  $\mathbf{S}^3$  to achieve higher throughput under those latency SLOs. Figure 1 highlights how much  $\mathbf{S}^3$  can improve throughput when trading off latency. For online scenarios, we assume a latency SLO set to the average reading speed of English readers, 4 words per second [12] and 0.1875 second per token [13]. The models on the left figure that are smaller than 100 billion parameters satisfy this SLO for all possible batch sizes. This SLO offers service providers to improve throughput over all models. In fact, for LLAMA-33B [14], there are opportunities to get throughput benefits with no latency penalties. The right figure, sweeping over different number of GPUs, shows opportunities to maximize throughput in offline scenarios that have loose latency SLO.
We evaluate  $S^3$ , assessing both its throughput and cost-efficiency. Our analysis includes both offline and online scenarios. In online scenarios under the average reading speed latency SLO constraint,
we find that  $\mathrm{S}^3$  can generate up to  $6.49\times$  more sequences while adhering to the same SLO constraint. In offline scenarios, we observe that  $\mathrm{S}^3$  achieves a speedup up to  $6.49\times$  for different models.  $\mathrm{S}^3$  does not affect the models' perplexity as it does not change the models' architectures. Furthermore, we evaluate the cost-efficiency of  $\mathrm{S}^3$  and find that using 6 GPUs,  $\mathrm{S}^3$  provides almost identical throughput compared to a vanilla system with 10 GPUs.
- We increase the achievable batch size under longer latency SLO and allow service providers to serve with higher throughput in both online and offline scenarios by using larger batch sizes.
- We fine-tune a Distillbert model to predict output sequence lengths given an input prompt.
- We provide a mechanism to recover from mispredictions.  $S^3$  preempts sequences that exceed their allocated memory and retrain the predictor to learn from its mistakes.
2.1 Generative AI Models
A Transformer-based generative model is autoregressive. It predicts the most probable token based on past tokens. Since the model generates one token at a time, it has to iterate over itself  $n$  times to generate a sequence that is  $n$ -tokens long. One iteration involves an input token traversing through the model which is a stack of transformer layers containing one attention, two layer norm, and two feed-forward layers. Especially, the self-attention layer uses information on the past tokens to generate the next token.
For example, the model at the  $i^{th}$  iteration attends the current token  $(t_i)$  with every token it already generated  $(t_0, \dots, t_{i-1})$  in the self-attention layer. We can express the self-attention layer as:
$$
h _ {o u t} = s o f t m a x (\frac {q _ {i} \cdot K ^ {T}}{\sqrt {d _ {h}}}) \cdot V
$$
where  $d_h$  is hidden dimension of the model,  $h_{out}, q_i \in \mathbb{R}^{d_h}$  are output hidden vector, current query vector, respectively, and  $K, V \in \mathbb{R}^{i \times d_h}$  are key and value matrices. The  $j_{th}$  rows in the  $K$  and  $V$  matrices represent key and value vectors of  $t_j$ , respectively. The two dot products attend  $t_i$  to all the key and value vectors in the current sequence.
The model stores  $K$  and  $V$  matrices as the key/value (KV) cache to avoid having to generate key and value vectors at each iteration. Otherwise, it has to store hidden states of every previous token and multiply it with a weight matrix  $W_{QKV} \in \mathbb{R}^{d_h \times 2d_h}$  at every transformer layer. This would require  $2(i - 1)d_h$  FLOPs per layer almost identical to  $2id_h$  FLOPs for the self-attention layer at  $t_i$ . The size of the KV cache is  $4ld_h$  bytes per token when using half-precision numbers. The cache uses 2 bytes for every number for both the key and value cache where  $l$  is the number of transformer layers in a model. For example, GPT-NEOX [15], a 20 billion parameter model, has 44 layers and 6144 hidden dimensions and thus uses 1MB per KV cache per token.
2.2 KV Cache Management on GPUs
The KV cache is relatively small (e.g., several MBs) and can be easily stored in the GPU HBM (high-bandwidth memory) when the sequence is short as the cache stores information about the previous tokens in the sequence. It grows as the model generates more tokens. Using this dynamic nature, Huggingface's Transformers [3] (HF-Transformers) constantly allocates more memory to the KV cache and stalls the computation until it is complete. This approach allows the library to allocate the exact amount of memory for each cache at the cost of frequent memory accesses.
To mitigate this, NVIDIA's FasterTransformer library reserves the maximum sequence length of memory for every sequence [4, 16]. It removes redundant memory accesses by simply filling in the reserved memory in an append-only fashion. However, this approach comes with its own drawback as it reserves more than strictly-necessary memory for sequences. For GPT-NEOX with a maximum sequence length of 2048 tokens, FasterTransformer reserves 2048 tokens even for sequences that
Figure 2: (a) GPT-J's and (b)GPT-NEOX's NVIDIA A100 compute utilization.
end up being 50 tokens long. The maximum batch size that FasterTransformer can use for this model with a 80GB A100 GPU is less than 20 with the 40GB model size and 2.2GB per sequence. The small batch size underutilizes the massive compute resources in the GPU. The rationale for reserving the maximum sequence length amount of memory even with the underutilization problem is to ensure that it generates full sequences and enhances user experience.
2.3 Observation
Language models are memory bound We demonstrate the extent of GPU resource underutilization when we run GPT-J with 6B parameters on an A100 GPU. The relatively smaller model shows a wider spectrum of different batch sizes and sequence lengths since we can fit larger batches with longer sequences in the GPU. Fig. 2 (a) shows the GPU utilization swept over different batch sizes and sequence lengths. As the figure denotes, increasing the batch size achieves a higher utilization but eventually faces a memory cliff, where an out-of-memory (OOM) error kills the process. Take the batches with 1024 sequence length for example. 32 sequences are the maximum batch size and thus  $12.56\%$  is the maximum utilization that we can achieve with this sequence length. Fig. 2 (b) shows similar underutilization in the larger GPT-NEOX model due to the memory cliff problem. This model faces the memory cliff with smaller batch sizes and shorter sequences since the model consumes more of the GPU memory. Please note that HF-Transformers still needs to know the output sequence length before batching inputs to avoid the memory cliff.
As the figures illustrate, increasing the batch size can enhance throughput in neural networks. This approach does not require intricate optimizations and enables the GPU to load the model weight from its HBM to the on-chip SRAM only once and share it among a larger number of inputs. By doing so, the GPU can activate its idle compute resources and concurrently handle multiple inputs. Nevertheless, the memory cliff poses a challenge, limiting the utilization of additional resources. However, as we elaborate, this issue can be resolved.
Reasons behind the inefficiencies Both the frequent memory allocations in HF-Transformers and the limited batch size in FasterTransformer come from the fact that we are not aware of the generated sequence length. If we know the precise length of the generated sequence, we can allocate exact memory to each sequence and resolve the repetitive memory reservation and the unnecessary memory allocation problems in HF-Transformers and FasterTransformer, respectively.
3 S $^3$  Design
$S^3$  is a system-algorithm co-design framework that maximizes GPU utilization with sequence length prediction to achieve higher throughput.  $S^3$  has three components as shown in Fig. 3: 1) predictor, 2) scheduler, and 3) supervisor. A text generation query arrives in a request pool in the host DRAM. The predictor then predicts its output sequence length which the scheduler uses to batch requests. The scheduler dispatches the batch to the GPU and the text generator model generates texts in the
Figure 3: Overview of  $\mathbf{S}^3$ . The boxes in yellow denote new components proposed by  $\mathbf{S}^3$ .
batch. The supervisor oversees the GPU utilization and handles mispredictions. We describe each component in detail and how they interact with each other in this section.
Output sequence length predictor We use a predictor to predict the output sequence length and resolve the frequent and redundant memory allocation problems. Specifically, we fine-tune a Distilbert [17] model that was trained for sequence classification to classify which length bucket the output sequence length falls into. We bucketize the sequence lengths since it is known that machine learning models are not as capable of the "last mile" search compared to narrowing the candidates down to the last mile. Each bucket is allocated the range of  $\frac{\text{max sequence length}}{\text{number of buckets}}$  and we use 10 buckets.
To this end, we fine-tune the model on the Alpaca dataset [18], one of the representative question-and-answering datasets, and use the questions as inputs and the lengths of the answers as labels. We observe that this predictor predicts the correct bucket with  $98.61\%$  accuracy. The average distance between the wrong predictions and the correct bucket is 1.03 meaning that the error converges to 0 when we double the bucket size. We also evaluate the predictor on a model fine-tuned with Google Natural-Question dataset [19] and observe an accuracy of  $77.13\%$ . It makes smaller mistakes more often than larger ones. For completeness, we fine-tune a model on the Pile dataset [20], a non-question-and-answering dataset, and see  $65.6\%$  accuracy. The predictor shows surprisingly high accuracy compared to randomly guessing the bins as the latter is correct only  $10\%$  of the time.
We choose Distilbert, a 66 million parameter model for its small size and fast inference. The model size is negligible since it is smaller than even a single transformer layer in the billion-scale models (e.g., 214 million for 6 billion GPT-J [21] model). The latency is also negligible since the predictor model runs only once when a request arrives at the server while the text generation model runs  $n$  times to generate an  $n$ -token long output sequence. We measure that Distilbert takes 3.7ms to run compared to several seconds for the text generation models on an NVIDIA A100 GPU.
Length-aware sequence scheduler The scheduler batches and schedules sequences based on the predicted results to maximize the GPU utilization without exceeding the GPU HBM capacity. We can formulate this problem as a variant of the bin packing problem with a single bin. The capacity bin is the HBM size and the item weight is the size of the KV cache for each sequence.
We use the decreasing first fit algorithm [22] as the solution to the bin packing problem for its simplicity. The scheduler queues the lengthiest sequences first, reserving room for shorter sequences within the GPU's available HBM. It orders sequences in the request pool by length in decreasing order and iterates through the pool to check if the KV cache of the current sequence does not exceed the available HBM. If so, it includes the sequence in the current batch and reduces the available HBM by the size of the KV cache. The scheduler continues this process until either there is no available HBM or it has iterated through the entire request pool. This approach has been consciously adopted due to its minimal associated overhead, while still maintaining an approximately-optimal resolution for the problem at hand.
The resulting batch is irregularly shaped where some sequences are longer than others. Unfortunately, current frameworks either do not support irregularly shaped batch [3, 4] or support it with limited performance [23]. Those that do not support this functionality pad the short sequences with padding tokens to match the length of sequences in the same batch. The padding tokens waste both computation and memory since they do not hold any useful information. ORCA [16] introduces an interesting solution to this problem termed selective batching. The authors of the work identify
Table 1: Model architecture used in the evaluations
Model	Num Params	Num layers	Model dim	Num heads
GPT-J [21]	6B	28	5120	16
LLAMA 13B [14]	13B	40	4096	40
GPT-NEOX [15]	20B	44	6144	64
LLAMA 33B [14]	30B	60	6656	52
GPT3 175B [26]	175B	96	12288	96
that inputs to certain layers (e.g., feed-forward) share identical weights, in contrast to inputs to other layers (i.e., self-attention) which do not share weights. Consequently, this enables a streamlined batch processing flow, wherein layers with shared weights are batch-processed and the batch is momentarily unpacked to process each input serially through the attention layers. ORCA shows that this has a negligible impact on the latency since the inputs to the self-attention layers do not share weights and do not benefit from batching. As such, we follow ORCA and use its selective batching technique.
Also borrowing from ORCA the iteration-level scheduling technique,  $S^3$  does not wait until all sequences in the current batch finish generation. Instead, it checks if any sequence in the batch has finished generation at every iteration. This grants  $S^3$  higher scheduling flexibility and removes any redundant waiting time. Finally, if a model cannot fit in one GPU,  $S^3$  uses pipeline parallelism and shard the models in Transformer layer granularity.
Supervisor The supervisor is in charge of supervising the GPU utilization and handling mispredictions. The supervisor runs in the background to check for the available space in the HBM and passes the information to the scheduler. The scheduler then appends a sequence to the running batch if the available memory is large enough for the sequence.
The supervisor is also responsible for handling mispredictions. In the case of short predictions, the supervisor preempts those sequences that exceed their reserved memory. It monitors the length of the current output sequences and evicts them if they are not finished but used up its reserved memory. It asynchronously moves the current state of those sequences including the KV cache and the generated tokens to the request pool and frees up the GPU memory. Now the K and V matrices are fragmented with blank rows where the evicted KV cache was originally stored in. The supervisor shifts the rows below the blank one so that all rows are stored contiguously. This memory format is required by current libraries [24,25] and also resolves the fragmentation issue. Finally, the supervisor doubles the assigned memory for the evicted sequences to fix the short misprediction.
Finally, the supervisor constantly trains the predictor in the background. It uses the sequences that the predictor mistook to train the predictor so that it can learn from its mistakes. This training time is relatively short and our measurement shows that each training iteration takes 11ms on average while sequence generation takes several seconds or more. This implies that the retraining overhead is less than  $10\%$  even if we train the predictor for 10 epochs.
Putting it all together We summarize this section with an explanation of how  $\mathrm{S}^3$  uses each component to serve a sequence-generation request. First, text-generation requests arrive at the request pool. The predictor predicts the output sequence length of the sequences in the pool. The supervisor runs in the background and checks the current HBM usage. Next, the scheduler uses both the predictions and the available HBM to batch requests for maximum GPU utilization. It finishes its job by scheduling that batch to the GPU which generates the scheduled sequences.
We show that  $S^3$  achieves higher throughput by predicting the output sequence length. It does so by using larger batch sizes hence smaller numbers of iterations. One iteration refers to processing a Transformer model once to generate a token. We also show that  $S^3$  can reduce the cost of serving models by using fewer GPUs.
(a) Maximum throughput (Alpaca [18])
(b) Maximum throughput (Google-NQ [19])
(c) Maximum throughput (The Pile [20])
(d) Generated sequences under latency SLO
Figure 4: Latency and throughput of different models and datasets.
Table 2: Maximum throughput of the three systems measured in tokens/s
Baseline	S3	Oracle
GPT-J	2061.16	2349.67	2569.09
LLAMA-13B	1018.15	1641.55	1907.09
GPT-NEOX	490.15	1344.94	1530.05
LLAMA-30B	91.46	593.73	834.30
Environment We run our evaluation on an NVIDIA 80GB A100 GPU connected to the host DRAM via PCIe  $4.0 \times 8$  in a Lenovo ThinkSystem SD650-N V2 Server [27].
Baselines We compare with ORCA [16], a Transformer serving system that increases throughput by iteration level scheduling and selective batching. ORCA has to allocate the maximum sequence length of memory when it is not aware of the output sequence length to guarantee full sequence generation.  $S^3$  predicts the output sequence length and allocates memory based on the prediction. We implement the systems on top of FasterTransformer [4] since this library is faster than HF- Transformers [3] due to more optimizations. We also compare  $S^3$  with an ideal system with a perfect predictor which we term Oracle.
Models We use models ranging from 6 billion parameters to 175 billion parameters for the evaluation. The specifics of these models are explained in table 1.
4.1 Throughput Analysis
We evaluate  $S^3$ 's throughput using Alpaca [18], Google Natural Questions (Google-NQ) [19], and The Pile [20] datasets. Specifically, we query  $S^3$  with questions and ask it to generate the answers.