KVNAND: Efficient On-Device Large Language Model Inference Using DRAM-Free In-Flash Computing

TL;DR

KVNAND tackles the memory bottleneck of long-context on-device LLM inference by storing both weights and KV cache in compute-enabled 3D NAND flash and executing memory-bound operations entirely within IFC. It introduces discrete and compact variants, a design-space exploration framework, and a page-level KV cache mapping to maximize flash locality and reliability. Through detailed evaluation on multiple models and context lengths, KVNAND achieves near-dramatic speedups and energy efficiency gains over DRAM-based IFC baselines, while mitigating out-of-memory at 100K contexts. The approach paves the way for practical long-context, privacy-preserving edge LLM deployment with low cost and robust reliability.

Abstract

Deploying large language models (LLMs) on edge devices enables personalized agents with strong privacy and low cost. However, with tens to hundreds of billions of parameters, single-batch autoregressive inference suffers from extremely low arithmetic intensity, creating severe weight-loading and bandwidth pressures on resource-constrained platforms. Recent in-flash computing (IFC) solutions alleviate this bottleneck by co-locating weight-related linear computations in the decode phase with flash, yet still rely on DRAM for the key-value (KV) cache. As context length grows, the KV cache can exceed model weights in size, imposing prohibitive DRAM cost and capacity requirements. Attempts to offload KV cache to flash suffer from severe performance penalties. We propose KVNAND, the first DRAM-free, IFC-based architecture that stores both model weights and KV cache entirely in compute-enabled 3D NAND flash. KVNAND addresses the fundamental performance challenges of flash under intensive KV cache access by leveraging IFC for all memory-bound operations to reduce data transfer overhead, introducing head-group parallelism to boost throughput, and employing page-level KV cache mapping to align token access patterns with flash organization. In addition, we propose a design space exploration framework that evaluates discrete and compact KVNAND variants to balance weight and KV placement, automatically identifying the optimal design trade-off. These techniques mitigate latency, energy, and reliability concerns, turning flash into a practical medium for long-context KV storage. Evaluations on MHA 7B and GQA 70B LLMs show that KVNAND achieves 1.98\(\times\)/1.94\(\times\)/2.05\(\times\) geomean speedup at 128/1K/10K-token contexts compared to DRAM-equipped IFC designs and addresses out-of-memory failures at 100K context length.

Figures (16)

• Figure 1: Different on-device LLM inference architectures.
• Figure 2: General LLM architecture and inference flow.
• Figure 3: (a) KV cache size of various models as the context length increases. (b) Number of NAND flash pages occupied by KV cache and model weights.
• Figure 4: (a) Roofline model analysis of IFC and smartphone NPU. (b) Comparison of attention computation workflows when KV cache is stored in DRAM and in NAND with IFC.
• Figure 5: (a) Page-read (PGRD) counts for LLaMA3.1-8B KV cache with a 50K-token context (1K input and 1K generated tokens per request). (b) Acceptable BER in model weights and KV cache resulting in a 10% accuracy/perplexity degradation.