# BATON: Enhancing Batch-wise Inference Efficiency for Large Language Models via Dynamic Re-batching

arXiv: https://arxiv.org/abs/2410.18701 · HTML: https://mlsanity.com/paper/2410.18701/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 44 references that are in the map.
- First posted on arXiv 24 Oct 2024.

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers](https://mlsanity.com/paper/2210.17323/) (2022)
- [DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving](https://mlsanity.com/paper/2401.09670/) (2024)
- [Splitwise: Efficient generative LLM inference using phase splitting](https://mlsanity.com/paper/2311.18677/) (2023)
- [Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve](https://mlsanity.com/paper/2403.02310/) (2024)
- [SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills](https://mlsanity.com/paper/2308.16369/) (2023)
- [Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training](https://mlsanity.com/paper/2110.14883/) (2021)
- [Response Length Perception and Sequence Scheduling: An LLM-Empowered LLM Inference Pipeline](https://mlsanity.com/paper/2305.13144/) (2023)
- [TensorFlow-Serving: Flexible, High-Performance ML Serving](https://mlsanity.com/paper/1712.06139/) (2017)
- [Batch Prompting: Efficient Inference with Large Language Model APIs](https://mlsanity.com/paper/2301.08721/) (2023)
- [S3: Increasing GPU Utilization during Generative Inference for Higher Throughput](https://mlsanity.com/paper/2306.06000/) (2023)
- [Full Stack Optimization of Transformer Inference: a Survey](https://mlsanity.com/paper/2302.14017/) (2023)
- [TurboTransformers: An Efficient GPU Serving System For Transformer Models](https://mlsanity.com/paper/2010.05680/) (2020)
- [D'ej`aVu: KV-cache Streaming for Fast, Fault-tolerant Generative LLM Serving](https://mlsanity.com/paper/2403.01876/) (2024)
- [A Survey of Techniques for Optimizing Transformer Inference](https://mlsanity.com/paper/2307.07982/) (2023)
- [DTLLM-VLT: Diverse Text Generation for Visual Language Tracking Based on LLM](https://mlsanity.com/paper/2405.12139/) (2024)
- [Understanding INT4 Quantization for Transformer Models: Latency Speedup, Composability, and Failure Cases](https://mlsanity.com/paper/2301.12017/) (2023)

## Built on by (most cited first, 1 of 1)

- [Polar Sparsity: High Throughput Batched LLM Inferencing with Scalable Contextual Sparsity](https://mlsanity.com/paper/2505.14884/) (2025)

## Abstract (arXiv)

The advanced capabilities of Large Language Models (LLMs) have inspired the development of various interactive web services or applications, such as ChatGPT, which offer query inference services for users. Unlike traditional DNN model, the inference of LLM entails different iterations of forward computation for different queries, which result in efficiency challenges for existing run-to-completion batch-wise inference. Hence, some methods refine batch-wise inference to iteration-level by duplicating all nonlinear layers of LLM. However, this approach not only increases resource usage but also introduces idle computations to the batch due to the prefilling of newly added queries. Therefore, we propose BATON, an efficient batch-wise LLM inference scheme by dynamically adjusting processing batch, which can achieve near-zero idle computations without incurring additional resource consumption. To do so, BATON 1) shapes the vectors involved in the inference of the newly inserted query and processing batch to align dimensions and generates a new attention mask based on vector shaping to ensure inference correctness, which enables query inserting without consuming additional resource; 2) embeds prefilled Keys and Values of the new query into the KV_Cache of the processing batch by leveraging the prefilling and decoding separation mechanism, eliminating idle computations to the batch introduced by the prefilling process of the new query. Experimental results show that compared to the state-of-the-art solution Orca, BATON improves query processing by up to 1.75 times.