PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Peizhuang Cong, Qizhi Chen, Haochen Zhao, Tong Yang

arXiv:2410.18701v1cs.LG

TL;DR

Autoregressive LLM queries take differing numbers of iterations, making run-to-completion batches waste computation and complicating query insertion. BATON dynamically reshapes and embeds queries while decoupling prefilling from decoding. It achieves 1.29-1.75× higher query processing throughput than the state-of-the-art solution.

  • Problem

    Autoregressive queries require different numbers of iterations, creating idle computation in run-to-completion batches and making direct insertion difficult as vector dimensions change.

    Source: §1 · p. 2, §2.2 · p. 3

  • Method

    BATON dynamically shapes query vectors and masks for alignment, then embeds prefilled Keys and Values into the processing batch’s KV-Cache through prefilling-decoding separation.

    Source: §1 · p. 2, §4.2 · p. 6

  • Results

    1.29-1.75× higher query processing throughput than the state-of-the-art solution.

    Source: §7 · p. 9

  • Takeaways & Limitations

    BATON supports dynamically replenished batches with near-zero idle computations while avoiding additional resource and computational bubbles.

    Source: §1 · p. 2

    Validation on real query traces, including dynamic batch scaling and preemptive scheduling, and migration to different inference frameworks are future work.

    Source: §7 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

LLM serving commonly batches queries for GPU parallelism, but autoregressive generation gives queries different completion lengths, leaving completed queries doing idle work. BATON dynamically replenishes batches and reports higher query-processing throughput than the state-of-the-art solution.

  • Autoregressive decoding generates one token per iteration until EOS or the maximum length, so queries in a batch can finish at different times.Completed queries continue producing EOS while the batch runs, creating idle computation.Source: §1 · pp. 1–2
  • Orca refines scheduling to the iteration level but replicates self-attention modules, consuming resources and allowing new-query prefilling to stall batch decoding.Source: §1 · p. 2
  • BATON dynamically removes completed queries and inserts new ones into the processing batch with near-zero idle computations.The paper also describes query interruption and recovery for preemptive scheduling and flexible batch-size scaling.Source: §1 · p. 2
  • Vector padding aligns new and active queries with masks; P&D decoupling embeds new-query Keys and Values into the batch KV-Cache without redundant prefilling.Source: §1 · p. 2
  • 1.29-1.75× higher query processing throughput than the state-of-the-art solution Orca.Source: §1 · p. 2

2 Preliminary and Motivation

The preliminaries describe Transformer attention and KV-Cache inference, while the motivation identifies idle computation and insertion barriers in run-to-completion batches. BATON targets these challenges with query-level replenishment.

2.1 Preliminaries

  • Transformer self-attention transforms inputs into Q, K, and V, uses scaled softmax attention weights to combine V, and is arranged in layers with FFNs or MLPs.Source: §2.1 · p. 2
  • KV-Cache reuses previously computed K and V across autoregressive iterations, so after prefilling, decoding can process only the latest output token.Source: §2.1 · pp. 2–3
  • Batch-wise inference pads query vectors to compatible lengths and uses attention masks to distinguish original tokens from padding.Source: §2.1 · p. 3

2.2 Motivation

  • Because query lengths and decoding iterations differ, run-to-completion batches waste computation; inserting a new query also faces changing vector dimensions and redundant prefilling.Orca-style replication addresses dimension misalignment with extra parameters, while synchronization can leave existing queries idle during new-query prefilling.Source: §2.2 · pp. 3–4

3 Baton Overview

Baton seeks query-level continuous inference while limiting additional resource and computation overhead. It combines vector shaping with prefilling-decoding separation to support query insertion.

3.1 Desired properties

  • The desired properties are query-level continuous inference, allowing new queries to enter after any query completes, and little or no added resource or computation overhead.Source: §3.1 · p. 4

3.2 Baton solutions

  • Vector shaping pads new and active query vectors to align dimensions and creates attention masks that preserve correctness in subsequent inference iterations.Source: §3.2 · p. 4
  • Vector embedding decouples prefilling from decoding so a new query can join after prefilling without padding the existing queries; it also supports hybrid GPU-CPU KV-Cache storage.Source: §3.2 · p. 4

4 Baton Design

BATON dynamically inserts queries into an active batch using padding and attention masks, then embeds prefilled keys and values to avoid the extra prefill work that would otherwise burden existing queries. It also supports preemptive scheduling and batch-size adjustment.

4.1 Vector Shaping

  • Vector shaping aligns new and active queries with padding and attention masks, enabling continued inference while preserving query outputs.Completed queries can be removed and replaced; front placeholder padding can be released to limit cache and mask growth.Source: §4.1 · pp. 4–5

4.2 Vector Embedding

  • Prefilling the new query separately lets BATON embed its keys and values into the batch cache, avoiding extra prefill computation for active queries.When the new query’s KV length fits the existing cache, placeholders can be reused; a longer KV requires cache and attention-mask expansion.Source: §4.2 · p. 6

4.3 Additional functionalities

  • BATON can preempt the lowest-priority query, insert a high-priority one, and reinsert the interrupted query when a batch slot opens.It temporarily stores the interrupted query’s keys and values.Source: §4.3 · p. 7
  • When GPU use exceeds a predefined threshold, BATON can move batch queries to host memory and release their GPU resources, or add queries to scale the batch up.Resource use can be monitored in real time.Source: §4.3 · p. 7

5 Evaluation

The evaluation compares BATON with batch-wise and prefilling/decoding baselines using overall and detailed inference tests. BATON improves throughput over Orca, while detailed analyses track query completion, output-token rates, and KV_Cache use.

5.1 Experimental setup

  • The setup used deterministic Llama2-7b-chat-hf inference, two simulated query sets, and A6000 and 4090 GPUs for overall and detailed evaluations.Benchmark was the Transformers batch-wise strategy; PD grouped similar-length queries for prefilling and randomly combined batches for decoding, while Baton-PD added P&D decoupling to Baton.Source: §5.1 · p. 7

5.2.1 Target-1: overall evaluation.

  • Against Orca, BATON improves throughput by 1.29–1.75×; versus Benchmark (PD), it reduces processing time by 46.25% (45.98%) at batch 4 and 59.90% (59.64%) at batch 8.The overall tests used batch sizes 4 and 8, with sequences duplicated for the batch-4 experiment.Source: Table 1 · p. 7, §5.2.1 · pp. 7–8

5.2.2 Target-2: detailed analyses.

  • Figure 6 shows BATON’s cumulative completed-query curve is less step-like than Benchmark’s, which can be delayed by longer-iteration queries.BATON can update batch queries continuously without interruption.Source: Figure 6 · p. 8, §5.2.2 · p. 8
  • Figure 7’s output-token distribution shows Benchmark’s repeated slowdown phases, while BATON’s output rate declines later when no new queries are available.The Benchmark pattern is attributed to the previously described delay; BATON’s later decrease occurs when replenishment stops.Source: Figure 7 · p. 8, §5.2.2 · pp. 8–9
  • BATON speeds throughput over Benchmark by 1.94×, 1.89×, 1.66×, 1.61×, and 1.31× at batch sizes 2, 4, 6, 8, and 10, respectively, with gains declining as the batch empties.The authors attribute the decline to there being no additional queries to replenish the processing batch during later inference.Source: §5.2.2 · p. 8
  • Figure 8 shows BATON releases KV_Cache memory incrementally rather than synchronously at batch completion, with comparable peak usage and consistently higher utilization than Benchmark.The evaluation measured KV_Cache memory only, excluding model parameters.Source: Figure 8 · p. 9, §5.2.2 · p. 9

6 Related Work

LLM services make inference efficiency and operational overhead important research directions. The section surveys inference-optimization approaches and situates Baton as an orthogonal strategy for multi-cloud deployments.

6 Related Work

  • LLM services handle vast numbers of daily inferences, making service quality and inference overhead key research directions.Source: §6 · p. 9
  • Inference-efficiency work spans kernel customization, parallel computing, quantization, and batch-processing methods.One batch-processing example groups queries by input-vector length.Source: §6 · p. 9
  • Existing approaches are described as primarily targeting single-computing-engine scenarios, while Baton is designed for multi-cloud services and operates orthogonally to them.The section says Baton can be integrated with these strategies.Source: §6 · p. 9

7 Conclusion

Baton dynamically removes completed queries and inserts new ones into the processing batch with near-zero idle computation. Its prototype improves query-processing throughput over the state-of-the-art solution, while future work targets broader validation.

7 Conclusion

  • 1.29–1.75× improvement in query-processing throughput over the state-of-the-art solution.The paper reports building a prototype and conducting extensive experiments.Source: §7 · p. 9
  • Baton enables completed queries to be removed and new queries inserted dynamically, with near-zero idle computations.Source: §7 · p. 9
  • Future work will validate Baton with real query traces and explore dynamic batch-size scaling, preemptive scheduling, and migration to other inference frameworks.These efforts are intended to provide a more comprehensive experimental evaluation.Source: §7 · p. 9
Loading 2410.18701v1…