SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Offline scenario: Maximum throughput Fig. 4 (a) - (c) reports the maximum throughput in sequences per second for different models and different datasets. We measure the throughput using the maximum batch size of each configuration. It shows that  $S^3$  outperforms ORCA by  $1.13 \times$  and up to  $6.49 \times$  and closely matches Oracle, differing from  $9.34\%$  and up to  $40.52\%$ . The difference is magnified with larger models because the batch size is limited even for  $S^3$  since most of the HBM
Table 3: Average batch size and number of iterations for different models
Model	ORCA	S3	Oracle
Batch size	Num iter	Batch size	Num iter	Batch size	Num iter
GPT-J	69.94	7988	530	1054	564.88	989
LLAMA 13B	31.61	17675	274.66	2034	527.04	1060
GPT-NEOX	16.89	33069	157.59	3545	292.19	1912
LLAMA 33B	4	139790	42.47	13155	77.35	7223
Figure 5: Maximum throughput of GPT3 running on different numbers of GPUs.
Figure 6: Latency breakdown of  $\mathbf{S}^3$
is used to hold model weights. The batch size in  $S^3$  gets cut off before saturating the throughput as shown in Fig. 1. We also report the throughput in tokens/s in table 2 evaluated on the Alpaca dataset.
We can notice that the maximum throughput increases by  $6.49 \times$  while the batch size increases by nearly  $10 \times$  for every model. This comes from the unique behavior of Transformers. Feed-forward layers in the models benefit from batching layers while self-attention layers do not. This is because inputs in feed-forwards share the same weight while inputs in self-attentions attend to their own sequences.  $\mathbf{S}^{3}$ 's performance benefit is increased when the parallelized portion is larger than the serialized portion. However, increasing the batch size entails adding more serialized inputs thereby growing the serialized portion. GPT-J is a good example of this characteristic which shows a similar throughput jump of  $1.13 \times$  from ORCA to  $\mathbf{S}^{3}$  and  $1.09 \times$  from  $\mathbf{S}^{3}$  and Oracle while the batch size differs by  $8.69 \times$  and  $1.97 \times$ , respectively.
Online scenario: SLO-aware throughput We now consider a scenario with a latency SLO constraint. We set the SLO as 0.1875 seconds to generate one token given that average English readers can read at 4 words per second [12] or 0.25 second per word, and 0.75 words per token [13]. We next calculate the latency SLO for each sequence by multiplying 0.1875 by its sequence length (e.g., 11.25s for a sequence with 60 tokens).
Fig. 4d reports the number of sequences that each model generates using ORCA,  $S^3$ , and Oracle. Oracle exceeds the SLO for GPT-J, LLAMA 13B, and GPT-NEOX when it chooses the maximum batch size. So we limit the batch size for these models during this evaluation.  $S^3$  generates a similar number of sequences with the ideal case and  $1.13 \times 6.49 \times$  more sequences compared to ORCA. The throughput increase is similar to the offline scenario since  $S^3$  meets the SLO in most cases with its maximum batch size. However, the difference between  $S^3$  and Oracle reduces by  $10\%$  compared to the offline scenarios because we limit the batch size hence the throughput of Oracle.
4.2 Cost Analysis
We evaluate  $S^3$  with different numbers of GPUs. We partition GPT-3 into 6 and 8 GPUs in a pipeline-parallel manner, allocating 16 and 12 transformer layers per GPU for each configuration, respectively. We also evaluate on 10 GPU setting where we allocate 10 layers to 9 GPUs and 6 to the remaining GPU.  $S^3$  pipelines each batch and schedules the a batch whenever the GPU processing the first partition passes its result to the next batch. This reduces GPU idle time by having every GPU processing batches concurrently.
Table 4: Average batch size and number of iterations for different system configurations on Alpaca
Num GPUs	ORCA	S3	Oracle
Batch size	Num iter	Batch size	Num iter	Batch size	Num iter
6	11.95	46734	114.22	4891	209.47	2667
8	28.68	19482	247.85	2254	470.65	1187
10	48.99	11403	399.62	1398	564.88	989
Table 5: Accuracies and runtime of different predictors on different datasets
Model accuracy (%)	Model size	Datasets	Runtime (ms)
Alpaca [18]	Google [19]	The Pile [20]
MS-minibert	22M	98.06	77.99	60.1	2.3
Distilbert-base	66M	98.6	82.68	65.6	4.1
Bert-base	110M	99.54	85.08	68.2	7.6
Bert-large	340M	99.6	89.25	71.9	14.5
Fig. 5 reports the maximum throughput using the different numbers of GPUs. First, we can see that  $S^3$  achieves similar throughput using 6 GPUs compared to ORCA with 10 GPUs. ORCA shows  $0.92\%$  higher throughput than  $S^3$  to be specific. More GPUs shard the model into finer pieces and leave more space for storing the KV cache, allowing us to increase the batch size. Table 4 supports this claim by reporting larger batch sizes with more GPUs.  $S^3$  can achieve similar effect with fewer GPUs by optimizing memory allocation strategy.
Naturally, it leads to a similar question with the one in the throughput evaluation on why  $S^3$  with 6 GPUs shows similar throughput with ORCA with 10 GPUs even with  $2.33 \times$  the batch size. The answer is in the pipelined execution and the sequential nature of self-attention layers in Transformer-based models as explained in 4.1. The systems complete a batch at every  $\left\lceil \frac{l}{number \, of \, GPUs} \right\rceil$  instead of at every  $l$  layers. For example,  $S^3$  with 6 GPUs completes a batch at every 16 layers instead of 96 for GPT3. Similarly, ORCA with 10 GPUs completes at every 10 layers and thus processes more quickly. The increase in the latency negates the throughput benefit from  $S^3$  such that the two systems show almost identical throughput even with the  $2.33 \times$  larger batch size when using  $S^3$ .
4.3 Overhead: Latency Breakdown
We evaluate runtime latency of each component in  $S^3$ . We classify the runtime latency into three categories: generation, penalty, and overhead. Generation is the time  $S^3$  spent on processing the model to generate tokens. Penalty denotes the time it took for  $S^3$  to preempt the KV cache and hidden states from the GPU and load it back to the GPU. Overhead includes the time it took for predictor, scheduler, and supervisor, combined. Fig. 6 show that penalty and overhead combined are negligible (i.e.,  $11\%$  on average) compared to the generation. Of course, the penalty would increase if the predictor is less accurate and triggers more data traffic between the GPU and the host. In contrast, the overhead will increase if we employ a more accurate but heavier predictor, thus introducing a new trade-off.
4.4 Predictor Ablation Study
We vary the predictor size from 22M to 340M parameters and report their accuracies on different datasets in table 5. We can observe a similar trend in all three datasets where a larger predictor generates more accurate predictions. The accuracies differ among different datasets since the length distribution differs. Specifically, Alpaca [18] showed the smallest variance among the length distributions in the three datasets compared to The Pile [20] showing the greatest.
Machine learning serving systems The high interest in machine learning has sparked numerous research in its service platforms [1,4,6, 16, 28-38]. Especially, the surprising performance of Transformer-based language models has directed many researchers to develop Transformer-specific
serving systems [4, 6, 34, 37-39]. Most of the systems focus on reducing the latency without much concern for throughput with the exceptions of [6, 31]. The throughput-oriented systems use memory hierarchy to store parts of the model in slower memory and to increase the batch size. However, they all allocate the same memory to every sequence and do not consider preempting a sequence based on a prediction.  $S^3$  can improve these systems by reducing the required memory size, removing the need for larger but slower memories such as SSDs, and reducing the fiscal cost of memory overall.
FastServe [39] tackles the head-of-line blocking problem caused by query-level scheduling while  $S^3$  addresses the inefficient memory allocation issue. It proactively manages KV cache similar to  $S^3$ 's supervisor by migrating sequences between the host memory and the GPU HBM. However, FastServe does not have an output sequence length predictor and thus uses a skip-join Multi-Level Feedback Queue since it is unaware of the job execution time. We expect  $S^3$  to work with FastServe so that  $S^3$ 's predictor delivers more information to FastServe for better scheduling.
Reducing the KV cache overhead The issue of attention layers in Transformers requiring quadratic computation and memory with respect to the sequence length has been extensively studied. Various approaches have been proposed to address this problem. Low-rank approximation [40, 41] and exploiting sparsity [42-44] are among the methods that aim to mitigate the issue. Another approach, known as multi-query attention [37, 45], suggests reducing the cache size by utilizing one attention head per key-value (KV) cache instead of employing multiple attention heads, as illustrated in Table 1. Additionally, works focusing on model size reduction through compression [46] and quantization [6, 47, 48] can also contribute to reducing the size of the KV cache. These approaches are complementary to our work and can be employed to reduce the penalty caused by mispredictions, allowing for the use of a less accurate but lighter predictor.
Sequence length prediction While there are limited works that predict output sequence length based on an input sequence, Yan et al. [49] propose a convolution-based small network with embedding layers to forecast output sequence length in machine translation. In our approach, we employ a similar strategy but utilize a small Transformer-based predictor to estimate the output sequence lengths in text generation tasks. This predictor allows us to accurately predict the output sequence length for text generation workloads.
Limitations We make the following assumptions in this work. We assume that text generation request traces mimic the publicly available question-and-answering datasets since there are no publicly-available traces. Analyzing the actual trace would facilitate deploying  $S^3$  in commercial workloads. Similarly, text generation task does not have any standardized latency SLO constraint as in other machine learning workloads [50] since it is a relatively new service. So we assume average reading speed of an English reader as our SLO. We will be able to evaluate  $S^3$  in broader scenarios if organizations publicly release different SLOs for different text generation applications.
Conclusion In summary, we introduce  $\mathbf{S}^3$ , a framework designed to achieve high throughput in serving Transformer-based generative models.  $\mathbf{S}^3$  leverages a predictor to estimate the output length of generated sequences and schedules them accordingly to maximize throughput. Additionally,  $\mathbf{S}^3$  handles potential prediction errors to guarantee reliability. By allocating varying memory sizes to different inputs,  $\mathbf{S}^3$  acknowledges that not all sequences should be treated equally. This approach expands the conventional trade-off between latency and throughput frontier, paving the way for new possibilities in optimizing the latency-throughput trade-off.
We thank the anonymous reviewers for their thoughtful comments and suggestions. This material is based upon work supported by the National Science Foundation under Grant No. 1704834 and No. 2118985 and supported in part by the Application Driving Architectures (ADA) Research Center and the National Science and Technology Council, Taiwan, under Grant No. 112-2222-E-A49-002-MY2. Chun-Feng Wu acknowledges the support from the Yushan Young Scholar Program by the Ministry of Education (MOE) in Taiwan.