MemLLM speeds up on-device language models by cutting storage delays
MemLLM: End-to-end In-Storage Acceleration of On-device Retrieval-Augmented Generation using Precomputed KV Caches
Hardware Architecture
Summary
Running big language models on devices like phones can be slow because of how data is handled. The authors created MemLLM, which puts some of the computing work directly inside the storage device. This change cuts down the slow data transfers between storage and the main processor. Their system makes these language tasks about four times faster on real devices compared to using a GPU.
What this means in practice
- •For mobile device developers: Accelerate personalized language model inference on edge devices by integrating MemLLM's in-storage processing architecture to reduce latency and power consumption.
- •For embedded system engineers: Build efficient on-device AI applications by minimizing PCIe data transfers and eliminating DRAM bottlenecks using MemLLM’s fused compute-storage design.
- •For cloud storage providers: Improve server-side inference throughput by offloading KV cache computations to storage units adopting MemLLM principles to reduce CPU-GPU data movement.$Commercial implications: Enables faster, more scalable AI services by embedding compute in storage hardware, creating products appealing to data center operators.
Authors
Seungkwan Kang, Hyungseok Ko, Heemin Kim, Myoungsoo Jung
Abstract
Retrieval-augmented generation (RAG) on edge devices enables personalized large language model (LLM) inference while preserving privacy. However, RAG's inherent length asymmetry makes the prefill phase the dominant bottleneck. While precomputing the KV caches offline alleviates the prefill cost on server systems, the multi-gigabyte caches over the narrow interconnect create a severe I/O wall for on-device systems. To mitigate this, we propose MemLLM, an end-to-end in-storage processing (ISP) architecture that executes the complete attention pipeline based on precomputed KV caches within the SSD. MemLLM utilizes per-channel accelerators to fully utilize the aggregated NAND channel bandwidth and accelerates all KV cache-related computations on the storage-side to minimize PCIe transfers. To efficiently handle the sparse updates required by RAG, MemLLM introduces a fused recomputation-query kernel paired with a hardware-based filter logic to halve the NAND read volume and eliminate write-back penalties. Furthermore, MemLLM's global SRAM-based stateful reduction engine enables a completely DRAM-less dataflow, thereby removing the internal memory bottlenecks and bounding the PCIe egress to an $O(1)$ footprint. A full-stack software stack allows asynchronous, zero-copy transfers from the ISP to the host and provides transparent execution within modern ML frameworks. Evaluated across eight real-world RAG datasets and six on-device LLM models, MemLLM achieves an average end-to-end speedup of 4.06$\times$ over the GPU baseline.