Reproduce “BATON: Enhancing Batch-wise Inference Efficiency for Large Language Models via Dynamic Re-batching” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Peizhuang Cong, Qizhi Chen, Haochen Zhao, Tong Yang
Baton, developed at Peking University, enhances batch-wise inference efficiency for Large Language Models by introducing dynamic re-batching and a prefilling and decoding decoupling strategy. This approach reduces idle computation and optimizes GPU memory utilization, leading to a substantial increase in query processing throughput compared to existing methods.
Key insights
- Dynamically manipulating `input_token`, `attention_mask`, and `KV_Cache` enables 'relay race' re-batching, allowing completed queries to exit and new queries to enter a batch.
- Decoupling prefilling and decoding phases, by pre-computing KV-caches and embedding them, eliminates 'redundant prefilling' and maintains all queries in a batch at single-token decoding, significantly reducing idle computation.
- A resource releasing mechanism proactively trims `KV_Cache` and `attention_mask` tensors, preventing indefinite growth and ensuring memory efficiency over time.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.