Striking the Right Balance between Compute and Copy: Improving LLM Inferencing Under Speculative Decoding

TL;DR

This paper tackles the KV-cache bottleneck in LLM autoregressive decoding by introducing Balancing Memory and Compute (BMC), a scheme that reallocates KV caches only every $r$ iterations to trade occasional redundant computation for in-place KV updates and reduced memory traffic. It further demonstrates that the redundant padding can be repurposed for Speculative Decoding (SD), boosting token-generation efficiency. An analytical model predicts the optimal number of allocations $T$ as proportional to $\sqrt{N}$, which is validated across multiple models, backends, and even GPUs. Empirically, BMC yields up to 3.2× throughput gains over baseline systems, with additional gains when combined with SD, and maintains accuracy while being hardware-agnostic. The approach outperforms state-of-the-art inference servers on CPUs and GPUs while offering broad compatibility and scalability across sequence lengths and backends.

Abstract

With the skyrocketing costs of GPUs and their virtual instances in the cloud, there is a significant desire to use CPUs for large language model (LLM) inference. KV cache update, often implemented as allocation, copying, and in-place strided update for each generated token, incurs significant overhead. As the sequence length increases, the allocation and copy overheads dominate the performance. Alternate approaches may allocate large KV tensors upfront to enable in-place updates, but these matrices (with zero-padded rows) cause redundant computations. In this work, we propose a new KV cache allocation mechanism called Balancing Memory and Compute (BMC). BMC allocates, once every r iterations, KV tensors with r redundant rows, allowing in-place update without copy overhead for those iterations, but at the expense of a small amount of redundant computation. Second, we make an interesting observation that the extra rows allocated in the KV tensors and the resulting redundant computation can be repurposed for Speculative Decoding (SD) that improves token generation efficiency. Last, BMC represents a spectrum of design points with different values of r. To identify the best-performing design point(s), we derive a simple analytical model for BMC. The proposed BMC method achieves an average throughput acceleration of up to 3.2x over baseline HuggingFace (without SD). Importantly when we apply BMC with SD, it results in an additional speedup of up to 1.39x, over and above the speedup offered by SD. Further, BMC achieves a throughput acceleration of up to 1.36x and 2.29x over state-of-the-art inference servers vLLM and DeepSpeed, respectively. Although the BMC technique is evaluated extensively across different classes of CPUs (desktop and server class), we also evaluate the scheme with GPUs and demonstrate that it works well for GPUs.

Figures (16)

• Figure 1: Attention Computation with Iterative Allocation.
• Figure 2: Execution Time Contribution for OPT Models.
• Figure 3: KV Cache update and SDPA Computation Time under Iterative Allocation.
• Figure 4: Total Time taken for Attention Layer -- Upfront versus. Iterative Allocation.
• Figure 5: Speedup in Attention Block Latency with BMC.