INF^2: High-Throughput Generative Inference of Large Language Models using Near-Storage Processing

TL;DR

INF^2 introduces a near-storage processing framework that shifting KV-cache attention workloads to accelerators inside computational storage devices, reducing KV-cache I/O bottlenecks for large-scale LLM inference. The method combines Attention-Near Storage, delayed KV cache writeback, and cooperative X-cache to hide storage latency and utilize host memory, achieving up to 3.46× throughput gains over state-of-the-art baselines on real hardware. The work demonstrates a practical, PyTorch-ready system using off-the-shelf components, with a real accelerator on Samsung SmartSSDs and a scheduler enabling scalable deployment; it also provides extensive ablation and sensitivity analyses and discusses cost-effectiveness. By leveraging the internal bandwidth of CSDs and intelligent dataflow, INF^2 offers a viable path to high-throughput, long-context LLM inference without exclusively relying on more expensive GPUs. The authors plan to open-source INF^2 to facilitate broader adoption and further research in near-storage AI acceleration.

Abstract

The growing memory and computational demands of large language models (LLMs) for generative inference present significant challenges for practical deployment. One promising solution to address these challenges is offloading-based batched inference, which leverages host memory and disk as an extended memory hierarchy for GPUs. While the approach cost-effectively enables LLM inference, its performance is limited by substantial I/O overhead, primarily due to the large key-value (KV) cache sizes, which increase with batch size and LLM context window length. In this paper, we introduce INFerence-INFinity (INF^2), a framework that boosts generative inference throughput using computational storage devices (CSDs). The core of INF^2 is attention-near storage, which offloads memory-intensive self-attention operations to near-storage accelerators, significantly reducing traffic through the system interconnect. We also propose delayed KV cache writeback to hide storage write latency by delaying newly generated KV cache writes until the cache reaches sufficient size in system memory. Additionally, we introduce cooperative X-cache, a technique designed to further trade off the remaining memory capacity for storage bandwidth. Our methods effectively minimize idle time for computation, improving the overall throughput. To demonstrate the effectiveness of our approach, \thiswork has been implemented on PyTorch and evaluated on a real system. Our experiments show that INF^2 achieves up to 3.46$\times$ throughput improvement compared to state-of-the-art baselines. We will open-source INF^2 to facilitate broader adoption.

Figures (17)

• Figure 1: Overall procedure of offloading-based batched inference flexgendeepspeedinf. (a) At the prefill stage, the KV cache is generated for input prompt sequence and stored in the storage. (b) At the decoding stage, each MHA layer loads the stored KV cache from storage to compute attention scores.
• Figure 2: (a) An example of PCIe hierarchy with standard storage devices (b) An example of an environment with computational storage devices (e.g. SmartSSD smartssd).
• Figure 3: A motivational study. $s$ represents the LLM input context size, and $B$ denotes the model parameter count. (a) Memory consumption breakdown during offloading-based batched inference. (b) The main bottleneck in offloading-based LLM inference is KV cache traffic. (c) Scaling storage devices with RAID0 fails to improve performance due to the bandwidth limits of the shared PCIe interconnect.
• Figure 4: The overall procedure of the decoding stage for (a) the baseline and (b) Attention-near storage. In the baseline, the KV cache is loaded from storage to the host through the system interconnect. With Attention-near storage, the KV cache is loaded from storage directly to our custom accelerator via an internal PCIe switch.
• Figure 5: KV cache writeback procedure. (a) Naive decode stage with attention-near storage. (b) Decode stage with delayed KV cache writeback.