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

Summary

This work builds an LLM inference platform, called S^3, around a sequence length predictor. The sequence length predictor is used to 1. batch LLM generations 2. pre-allocate kv cache (where all seq have similar predicted seq len) S^3 also has a method for handling seq len prediction errors. They do pipelined generation. They unpad shorter seq to ignore wasting compute.

Strengths

LLM inference is becoming a larger and larger part of the total compute used within large organizations. Decrease inference costs is an extremely impactful line of work. The analysis of overheads is good.

Weaknesses

For kv cache pre-allocation, the authors note that pre-allocation limits seq len, not pre-allocating kv cache slows generation. An obvious baseline is pre-allocating for S tokens at a time (eg S=64) and freeing kv cache mem of generations which have finished (this frees mem for the really long generations). This simple baseline would decrease the kv cache concat overhead by a factor of S and eliminates the overhead of the S^3 algo. The work uses a sequence length predictor to build an LLM inference platform. This is great systems research, I'm not sure it's positioned well for NeurIPS. The work doesn't show how much of the benefit comes from pipelining vs kv cache preallocation vs seq len batching vs unpading. I'm not sure how much improvement each method is or if we're introducing other overheads by including methods such as seq len prediction. There are great works on improving LLM inference. Why only compare to ORCA?

Questions

The authors note: "More GPUs shard the model into finer pieces and leave more space for storing the KV cache, allowing us to increase the batch size. " - line 280 The method also pipelines generations; kv caches for all generations must be stored in mem, this in turn lowers the possible batch size. Is it possible to quantify the trade off? Will the code be made public? People can easily benchmark the implementation helping the authors identify shortcomings and improve the overall performance over time.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

2 fair

Limitations

Requires training a seq len predictor which hopefully generalizes to different data distributions.

Summary

The paper proposes a scheme that increases the throughput during inference on Transformer large language models (LLM). Typically, LLMs require large amounts of memory, for model parameters and for the KV (key/value) cache. The KV cache size depends of the output sequence length, which is not known when inference starts. Some implementations allocate memory for the KV cache in small increments, causing large latency, while others preallocate for up to the maximum output sequence length, causing potential memory waste. The authors propose a method, called S3 (scheduling sequences with speculation), that predicts the output length, and allocates the KV cache accordingly. The predictor is a fine-tuned Distilbert model (66M parameters) with small size and fast prediction time. A scheduler batches requests according to a greedy strategy. A supervisor is in charge of checking GPU utilization and handling mispredictions, while at the same time training the predictor in the background. An experimental analysis shows the flexibility of S3 in offering a trade-off between latency and throughput.

Strengths

Transformers are memory and compute intensive. Designing schemes that make better use of memory, thus increasing the latency and/or throughput is an important topic. The proposed solution is simple, and the predictor ran as part of S3 adds negligible overhead. The experimental results show improvements, generating by up to 6.49 times more sequences compared to other existing system.

Weaknesses

The system will need more analysis in the future, based on real traces (which are not available for research community now), and considering realistic service level objectives.

Questions

1. Could you add to the experimental section results for the case when no KV caching is used?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

N/A

Summary

The paper tackles the problem of predicting the number of generated tokens for transformers in text generation tasks. This will help with better memory allocation and batch size management. The previous systems either used dynamic memory allocation (Hugging Face) which incurs inference overhead, or preallocated memory for over-estimated output length, limiting the batch size. This paper's prediction of output length achieves a larger batch size, thus increasing the throughput. Experiments show a 6.49x improvement in throughput. The predictor is get by training a Distilbert model.

Strengths

1. The idea is simple but effective. 2. The paper is clear, well-written, and easy to read. 3. The evaluation is clear and convincing. It contains various settings, including different models, different hardware setups, online/offline with clear information presented (latency breakdown, batch size).

Weaknesses

1. It has not been discussed that different models can generate outputs with different lengths. 2. The evaluation does not contain different request patterns. 3. Missing the evaluation of throughput in term of token/s.

Questions

1. Figure 1: is it get from generating with 60 token context length? This is not a long context. 2. line 71: "6.49× more sequences" is vague, because sequences can have different lengths. 3. line 212: "The supervisor shifts the rows below the blank one so that all rows are stored contiguously.". How much time overhead does this step cause? 4. Why are you measuring the throughput as sequence/s, not token/s? Sequences have different lengths, which makes sequence/s hard to interpret. I can understand that you want to count how many requests can be served, but maybe also add plots for token/s. 5. When you are using the Alpaca dataset, how did you setup the request distribution (coming time, etc)? 6. How does your method compare to this one "Fast Distributed Inference Serving for Large Language Models" that was released recently? minor: line 47: typo before "since most of the HBM is used to..." line 307: [31, 39] are not focus on throughput.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

To accurately predict the output length, the distribution of the test trace should be similar to the data used to train the predictor. Also, different models will generate answers with different lengths, which requires the training data to reflect a similar behavior to the model used in the generation.

Summary

This paper presents a new solution to the challenges of GPU underutilization and increasing batch size in the generation task of Large Language Models (LLMs), rooted in the memory-intensive requirement to retain the K and V values of prior tokens. To tackle these issues, the paper proposes an efficient framework named S^3, which uses a fine-tuned Distillbert model as a predictor to forecast the output sequence length based on an input prompt. This predictive model guides query generation scheduling and manages mispredictions by halting sequences exceeding allocated memory and doubling their assigned memory. The proposed approach increases the maximum configurable batch size in both online and offline scenarios and improves throughput, reporting a throughput increase of up to 6.49 times. Although this paper offers improvements in handling memory allocation for LLM inference, there are some concerns about evaluation and ablation of the proposed techniques.

Strengths

- This paper tackles the vital challenge of predicting output sequence length through a unique DistillBERT-based co-optimization of the system and algorithm. This approach effectively addresses the limitations of existing frameworks, namely HF-Transformer and FasterTransformer, leading to improved throughput while meeting SLO latency. - While larger models typically yield better performance, as the space for K, V caches reduce, effective memory allocation management becomes increasingly essential. The S^3 framework notably elevates the maximum throughput, particularly in larger models. - The S^3 framework optimizes the latency-throughput trade-off by adjusting memory allocation, leaving the model architecture intact and thereby maintaining model perplexity. Operating independently and focusing on performance enhancement, the framework doesn't impact the accuracy of existing LLM models.

Weaknesses

- While a major contribution of the paper is the prediction of output sequence length, the effectiveness of this proposed predictor isn't thoroughly evaluated. Tables 2 and 3 show a significant difference between the batch size predictions of S^3 and Oracle. The reasons behind this discrepancy need to be explained. - The effectiveness of the predictor appears to vary greatly with the fine-tuning dataset, with accuracies ranging from 65.6% to 98.6%. The paper proposes online learning, but this isn't evaluated. Additionally, the paper lacks an ablation study on the choice of predictor models and sizes. - The paper only compares the latency-throughput trade-off between vanilla systems and S^3 and doesn't compare average throughput and latency in specific tasks. This is crucial, as output sequence length prediction accuracy, which seems to be task-dependent, would likely impact this. - While S^3 meets the SLO limit, its implementation generally results in an increase in end-to-end latency. The paper doesn't adequately address how much each proposed technique contributes to this overhead. - The paper introduces a supervisor component to handle mispredictions, but the description of this technique is unclear. The supervisor appears to allocate double the memory when a previous allocation isn't sufficient. However, in typical scenarios with scarce memory, this could lead to problems if the predictor misestimates the batch size. More detailed explanations of various scenarios would be useful.

Questions

- The figure captions could use more detail. For instance, is Figure 1 (left) showing an online scenario, and how many GPUs were used? Similarly, does Figure 1 (right) represent an offline scenario, and how large is the model size? - The term "number of iterations" requires clarification. It's evaluated in the sections assessing performance, but isn't clearly defined. - In Section 4.1, it's mentioned that "Oracle exceeds the SLO for GPT-J, LLAMA 13B, and GPT-NEOX when it chooses the maximum batch size". It's unclear why Oracle would select a batch size that violates the SLO. - Both Figures 4 (a) and (b) claim a maximum performance improvement of 6.49x. How can this be, especially when Figure 4 (b) pertains to "Generated sequences under latency SLO"? Is the data in Figure 4 (b) based on a scenario of maximum throughput? - The application of S^3 seems to yield less throughput improvement as the model size decreases. Is there still a significant performance improvement when applying S^3 to smaller models like GPT-3 small, which has around 125M parameters? - Regarding throughput-related performance metrics, it's unclear how Oracle's (ideal predictor and scheduler) performance was measured. The scheduling might have been done under an entirely ideal scenario, or it might be based on the assumption that the predictor makes no errors in predicting input sequences.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

The authors properly stated the limitations of the proposed methods.

Summary

The paper proposes a simple yet effective systematic solution to increase GPU usage and throughput during inference. The authors first make an interesting observation that the existing large language models are bounded by memory and the inefficiency is caused by the lack of awareness of the sequence length. Accordingly, the authors first fine-tune a small model DistillBert to predict the sequence length. Based on it, the authors accordingly schedule the generation of queries and deal with mispredictions. The authors conduct experiments with models from 6 billion to 175 billion parameters to evaluate the effectiveness of the proposed system.

Strengths

++ The paper solves a practical problem, personally, I think improving the GPU utilization can benefit the community and the deployment of different applications. ++ The proposed solution is simple yet effective. Predicting the output sequence length first is a simple and intuitive way to schedule the network and better utilize the GPU. Although incorporating a new model introduces additional lags, the authors show that such overhead is negligible in Figure 6. ++ The results look promising; the proposed method allows for larger inference batch size and during the inference, e.g., 32 -> 256 for LLaMA-13B and 4 -> 32 for LLaMA-32B, and hence achieves higher throughput.

Weaknesses

-- Subsection Length-aware sequence scheduler in Section 3 is not very intuitive to the readers, especially the ones that are not directly working on the same topic. It would be better to illustrate more on how the bin packing problem is solved and the batching technique in ORCA (with an Algorithm or diagram or more descriptions). -- The authors only evaluate the proposed method on NVIDIA A100 GPU. However, it seems the design is not dependent on the architecture of A100 GPU. It would be better to also evaluate the proposed method on other GPU architectures. -- Minor: Distilbert in L169 and Distil-bert in L155

Questions

See weakness.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Rebuttal comments

Thank you for running the baseline; I stand corrected, having the predictor and sorting does seem to help substantially. It'd also improve the paper if the authors could show performance on MultiQuery and GroupedQuery (with 8 groups cuz TensorParallelism) attention given these are gaining a lot of popularity (ie they are starting to be used in most new popular networks) If performance on MultiQuery and GroupedQuery Attention is shown, I'd change my rating to a `6: Weak Accept` (assuming I can figure out the UI)

Thanks for the feedback and a follow-up question

I thank the authors for their feedback! I still have a follow-up question. Is the authors' implementation dependent on GPU types? If so, which feature of NVIDIA A100 GPU is necessary and can your implementation work on other popular GPU types, e.g., V100 or GeForce 3090 GPUs?

Thank you for the answers

I appreciate the authors' response to my questions, which resolved many of my concerns. But still, there are remaining concerns that the analysis of the effectiveness of the proposed predictor was mostly empirical. Thus, I would like to maintain my original rating.