End-to-End Acceleration of Generative Models With Runtime Regularized KV Cache Management

Abstract

Despite their remarkable success in achieving high performance, Transformer-based models impose substantial computational and memory bandwidth requirements, posing significant challenges for hardware deployment. A key contributor to these challenges is the large KV cache, which increases data movement costs in addition to the model parameters. While various token pruning techniques have been proposed to reduce the computational complexity and storage requirements of the attention mechanism by eliminating redundant tokens, these methods often introduce irregularities in the sparsity patterns that complicate hardware implementation. To address these challenges, we propose a hardware and algorithm co-design approach. Our solution features a Runtime Cache Eviction (RCE) algorithm that removes the least relevant tokens and replaces them with newly generated ones, maintaining a constant KV cache size across blocks and inputs. To support this algorithm, we design an accelerator equipped with a KV Memory Management Unit (KV-MMU), which efficiently manages active tokens through eviction and replacement, thereby optimizing DRAM storage and access. Additionally, our design integrates batch processing and an optimized processing pipeline to improve end-to-end throughput, effectively meeting the requirements of both pre-filling and generation stages. The proposed system achieves up to $8\times $ KV cache size reduction with minimal accuracy degradation. In a 65 nm process, the proposed accelerator demonstrates $1.52\times $ energy savings and $3.62\times $ delay reductions when processing a batch size of 16, with only a 1.11% energy overhead attributed to the specialized KV-MMU.