PaperlayerView exact revision on arXiv

Source-linked AI summary

MemLLM: End-to-end In-Storage Acceleration of On-device Retrieval-Augmented Generation using Precomputed KV Caches

Seungkwan Kang, Hyungseok Ko, Heemin Kim, Myoungsoo Jung

arXiv:2610.22579v1cs.AR

TL;DR

On-device RAG’s long retrieved contexts make prefill a bottleneck, while precomputed KV caches create an I/O wall across narrow PCIe. MemLLM runs the precomputed-KV attention pipeline inside the SSD, combining recomputation and query attention with per-channel accelerators and SRAM-based reduction. Across eight real-world datasets and six on-device LLM models, it achieves a 4.06× average end-to-end speedup over GPU-based systems.

  • Problem

    Long retrieved contexts and short answers make prefill a bottleneck in on-device RAG, while streaming precomputed KV caches over limited PCIe bandwidth creates an I/O wall.

    Source: §I · p. 1, A. The I/O Wall of Precomputed KV Caches · p. 4

  • Method

    MemLLM executes the precomputed-KV attention pipeline within the SSD, fusing selective recomputation with query attention and using per-channel accelerators and global SRAM reduction.

    Source: §I · p. 2

  • Results

    4.06× average end-to-end speedup over GPU-based systems across eight real-world RAG datasets and six on-device LLM models.

    Source: §IX · p. 13

  • Takeaways & Limitations

    The attention computation’s PCIe traffic is bounded to O(1), independent of context length.

    Source: §I · p. 2

    MemLLM targets long-context, short-answer on-device RAG with precomputed document KV caches and selective recomputation; its benefit narrows for long answers or memory-resident KV caches, and deployment assumes a computational storage device.

    Source: §VII · p. 13

Abstract

from arXiv · show

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.

I. INTRODUCTION

On-device RAG shifts inference toward prefill because retrieved contexts are long and answers are short, but moving precomputed KV caches across PCIe creates an I/O wall. MemLLM performs the attention pipeline inside the SSD, combining recomputation and query attention while using per-channel accelerators and DRAM-less reduction.

I. INTRODUCTION

  • Long retrieved contexts and short answers shift the on-device RAG bottleneck to prefill, where processing the context to build its KV cache creates a performance barrier.Retrieved contexts are often 100–1000× longer than generated answers.Source: §I · p. 1
  • Precomputing KV caches replaces server-side prefill computation with storage reads, but multi-gigabyte caches crossing narrow PCIe can take longer than the computation they replace.Edge-device memory limits also require these caches to remain in persistent storage.Source: §I · p. 1
  • Prior storage-side approaches do not handle selective recomputation: independently precomputed document caches require recomputing about 15% of tokens, then rereading the cache for query attention.This disjoint flow streams the multi-gigabyte cache across PCIe twice.Source: §I · pp. 1–2
  • MemLLM fuses recomputation with query attention in one NAND-read pass, uses per-channel Elastic-SIMD accelerators, and reduces partial outputs in global SRAM to avoid SSD DRAM traffic.The fused kernel halves NAND reads and eliminates write-back; PCIe traffic for attention is bounded to O(1), independent of context length.Source: §I · p. 2
  • 4.06× average end-to-end speedup over the GPU baseline across eight real-world datasets and six LLM models.The evaluation used an RTL implementation of Elastic-SIMD on a 16nm FPGA, GPU-kernel measurements on a Jetson Orin NX, and a cycle-accurate SSD simulator.Source: §I · p. 2

II. BACKGROUND

Decoder-only LLM inference uses prefill for prompt-wide computation and autoregressive decode with KV caches to avoid recomputing prior attention. In RAG, retrieval augments the prompt with relevant passages, making long-context prefill a major latency component.

A. Large Language Model (LLM) Inference

  • Grouped-Query Attention shares one K/V pair across each query-head group to compress the KV cache, unlike Multi-Head Attention’s separate K/V matrices per query head.Scaled dot-product attention weights value vectors using query-key scores.Source: A. Large Language Model (LLM) Inference · p. 2
  • Prefill processes the prompt in parallel with GEMMs, while autoregressive decode generates one token per iteration and uses stored K/V vectors to reduce prior-token attention to linear-time GEMV.KV-cache size grows with sequence length and can exhaust edge-device GPU memory.Source: A. Large Language Model (LLM) Inference · pp. 2–3
  • Online softmax supports attention beyond on-chip SRAM capacity by processing tiles and merging partial outputs using running maxima and sums of exponentials.The passage states that this enables long attentions within a fixed SRAM budget without accuracy loss.Source: A. Large Language Model (LLM) Inference · p. 3

B. Retrieval-Augmented Generation (RAG)

  • RAG retrieves top-k passages from a pre-indexed vector database and concatenates them with the query; with long retrieved contexts and short answers, prefill dominates end-to-end latency.The paper assumes retrieved chunks are finalized before attention begins.Source: B. Retrieval-Augmented Generation (RAG) · p. 3

C. Precomputed KV Caches for RAG

Independent offline KV caches save prefill computation but break cross-document attention, requiring selective recomputation; on-device storage and narrow PCIe bandwidth make cache movement a bottleneck.

C. Precomputed KV Caches for RAG

  • Selective recomputation restores cross-attention broken by independently precomputed document caches by recomputing roughly 15% of tokens flagged as HKVD.Those tokens require full attention recomputation using model weights; the others are consumed directly from storage.Source: C. Precomputed KV Caches for RAG · p. 3
  • The disjoint recomputation and query-prefill phases each read the entire KV cache, sending multi-gigabyte data across PCIe twice.Target tokens are scattered across sequence positions, so recomputation also uses an irregular attention mask.Source: C. Precomputed KV Caches for RAG · p. 3

III. CHALLENGE AND MOTIVATION

  • On-device RAG's long contexts create severe performance challenges for resource-constrained edge devices.The section analyzes these challenges to derive architectural requirements for an on-device RAG accelerator.Source: §III · p. 3

A. The I/O Wall of Precomputed KV Caches

  • The analysis uses datasets from finance, biomedical, legal, and general-knowledge domains; Wikipedia-scale and broad-web corpora exceed edge storage capacity and are evaluated using access-pattern traces.The large datasets are used to analyze generality and scalability, not as intended edge deployments.Source: A. The I/O Wall of Precomputed KV Caches · pp. 3–4
  • Contexts typically span thousands of tokens while answers are short; context-to-answer ratios reach approximately 1000:1 for TriviaQA, 90:1 for CuAD, and 440:1 for CovidQA.This length asymmetry inflates prefill time while limiting decode iterations.Source: A. The I/O Wall of Precomputed KV Caches · p. 4
  • 2.1× end-to-end speedup over the full-compute baseline on servers contrasts with below 0.58× on-device, where storage I/O dominates latency.On-device storage I/O reaches 80.2% of total latency.Source: A. The I/O Wall of Precomputed KV Caches · p. 4, Figure 4 · p. 4, Figure 5 · p. 4

B. Analysis of End-to-End On-Device RAG

With storage-resident KV caches, both decode and prefill are I/O-bound; decode-only in-storage processing shifts the bottleneck to prefill, motivating an end-to-end storage-side pipeline.

B. Analysis of End-to-End On-Device RAG

  • Decode streams the full multi-gigabyte KV cache over PCIe at every token step and accounts for over 64% of baseline latency.Even ideal decode-only ISP shifts the bottleneck to prefill, whose I/O exceeds 96% of latency regardless of context length.Source: B. Analysis of End-to-End On-Device RAG · pp. 4–5, Figure 6 · p. 5, Figure 7 · p. 5
  • Selective recomputation forces a PCIe round-trip during prefill: KV is transferred to the GPU, recomputed, written back, and reread for query prefill.The resulting disjoint execution motivates absorbing recomputation and query attention within storage.Source: B. Analysis of End-to-End On-Device RAG · p. 5
  • A 2D systolic array underutilizes its processing elements on decode GEMV and thin recomputation GEMMs, while irregular masks add boundary-padding overhead.Under a strict 3W budget, dedicating silicon to a large, mostly underused 2D array is architecturally prohibitive.Source: B. Analysis of End-to-End On-Device RAG · p. 5, Figure 8 · p. 5
  • SIMD handles GEMV and vector operations efficiently but lacks low-cost inter-PE forwarding for dense GEMMs, routing partial sums through shared memory.The resulting internal-memory bandwidth bottleneck limits the compute ceiling for projection operators.Source: B. Analysis of End-to-End On-Device RAG · p. 5, Figure 8 · p. 5
  • The analysis identifies elasticity between GEMV and GEMM execution and sparsity-aware filtering and compaction as requirements for an on-device RAG ISP.These requirements address heterogeneous operators and irregular attention masks without control or padding overheads.Source: B. Analysis of End-to-End On-Device RAG · p. 5

IV. OVERVIEW

MemLLM performs the full KV-cache attention computation within the SSD to avoid massive cache transfers, minimizing I/O through fused operations, channel-distributed compute, and SRAM-confined reductions.

IV. OVERVIEW

  • MemLLM applies three I/O-minimization principles: fuse operations to halve NAND reads, distribute compute across channels, and keep partial reductions in on-chip SRAM.The design aims to bypass controller DRAM and avoid massive KV-cache transfers.Source: §IV · p. 5

A. System Overview

  • The host GPU retains model weights and executes projection and FFN layers, while the SSD stores KV caches and performs attention over them.A vendor-specific NVMe command carries tensor metadata and DMA addresses for zero-copy GPU-to-ISP transfers over PCIe.Source: A. System Overview · p. 5
  • A 1 MB global SRAM holds a repeatedly broadcast query and accumulates partial attention outputs with in-place online softmax correction, excluding DRAM from the attention data path.Controller DRAM remains dedicated to FTL management.Source: A. System Overview · p. 5
  • Per-channel accelerators filter invalid tokens, switch between GEMV and GEMM modes, and double-buffer KV pages in local scratchpads.Data flows directly from NAND pages to each accelerator without traversing a shared internal SSD bus, dedicating aggregate channel bandwidth to computation.Source: A. System Overview · pp. 5–6

B. End-to-End Dataflow

MemLLM combines selective recomputation with query attention in one storage-side pass, keeping the bulk KV cache inside the SSD. Its design reduces end-to-end RAG latency while bounding PCIe egress independently of context length.

B. End-to-End Dataflow

  • The host identifies roughly 15% of high-deviation tokens for recomputation, while MEMLLM fuses that work with query attention in one NAND-read pass.The recomputation ratio is a static, user-determined hyperparameter (5–18%) that trades generation quality against compute cost.Source: B. End-to-End Dataflow · p. 6
  • Only compact query and output data, plus the host-recomputed KV subset, cross PCIe; egress remains O(1) regardless of context length.The bulk precomputed KV cache stays in the SSD, unlike storage-offloading baselines that repeatedly stream the full O(N) cache.Source: B. End-to-End Dataflow · p. 6

C. Performance Improvement Analysis

  • 7.1× end-to-end RAG latency reduction versus the storage-offloading baseline follows from the three I/O-minimization principles.Fusing recomputation and query attention reduces latency by 38%; per-channel compute and Stream Filter add 39%, and DRAM-less reduction adds 63%, reaching 86% of the aggregate-NAND-read minimum.Source: C. Performance Improvement Analysis · p. 6

V. DESIGN AND IMPLEMENTATION

  • The design implementation covers a fused kernel and internal dataflow, followed by Elastic-SIMD, NAND-aware tiling, DRAM-less dataflow, and software integration.The supplied figure identifies channel-parallel tiling and per-channel compute as elements of the internal dataflow.Source: §V · pp. 6–7, Figure 12 · p. 7

A. Fused Kernel and Internal Dataflow

The fused attention kernel handles selective recomputation and query attention together while the KV cache remains in flash. Per-channel Elastic-SIMD and NAND-aware tiling support the computation across prefill and decode.

A. Fused Kernel and Internal Dataflow

  • The fused kernel scores stacked recomputation-target and query vectors against precomputed, recomputed, and query keys using a jagged attention mask.The mask applies hybrid causal boundaries, and the tiled online-softmax computation matches software CacheBlend up to floating-point rounding.Source: A. Fused Kernel and Internal Dataflow · p. 7
  • The Stream Filter substitutes recomputed KV entries at selected positions, drops masked-out tokens, and merges them with the precomputed KV cache.New KV entries are appended out of place, preserving the original cache and avoiding read-modify-write; background writes overlap host projection and FFN work.Source: A. Fused Kernel and Internal Dataflow · p. 7

B. Elastic-SIMD Microarchitecture

  • Elastic-SIMD uses a 1D PE array and reconfigurable routing to support broadside GEMV for decode and cascaded thin-GEMMs for prefill.Its per-PE vector ALU performs online-softmax correction in registers without stalling the MAC pipeline.Source: B. Elastic-SIMD Microarchitecture · p. 7

C. NAND-Aware Tiling Algorithm

  • NAND-stationary tiling fixes KV tiles at 32 tokens per 16KB page and fetches each page once while query tiles stream through it.Round-robin striping across dies pipelines sensing, transfer, and compute to hide dominant sensing latency.Source: C. NAND-Aware Tiling Algorithm · p. 8

D. DRAM-less Elastic-SIMD Dataflow

MEMLLM computes attention through register-resident phases and global-SRAM reduction, avoiding intermediate SRAM writes and bypassing controller DRAM. Its software stack submits asynchronous, zero-copy transfers between GPU memory and the SSD.

D. DRAM-less Elastic-SIMD Dataflow

  • The DRAM-less dataflow computes scores, partial softmax, and partial outputs in registers before merging channel outputs through stateful log-sum-exp reduction in global SRAM.The design avoids writing intermediate score matrices to scratchpad; Qfused and Ofinal fit on-chip, leaving controller DRAM available for FTL operations.Source: D. DRAM-less Elastic-SIMD Dataflow · p. 8

E. MemLLM Software Stack

  • The software stack uses NVMe passthrough and a vendor-specific fused-attention command to submit operands without POSIX-layer tensor semantics.A lightweight handle carries layout information and DMA addresses; dma-buf maps GPU VRAM for peer-to-peer transfers, while io_uring enables asynchronous submission and overlap with host computation.Source: E. MemLLM Software Stack · p. 8

VI. EVALUATION

The evaluation combines FPGA-derived hardware metrics with cycle-level simulation of host, PCIe, and storage behavior. It compares MEMLLM against GPU and ISP baselines and evaluates incremental design ablations.

A. Methodology

  • Elastic-SIMD was prototyped on a 16nm FPGA, with cycle counts, area, and power measured against reference SIMD and systolic-array designs.The evaluation estimates area using CACTI and 45nm synthesis; the per-device flash-coordinator buffer accounts for more than 99% of its area at 1.6mm2.Source: A. Methodology · pp. 8–9
  • The trace-driven cycle-level simulator combines profiled Jetson Orin NX GPU kernels, prototype-derived ISP models, measured PCIe delays, and SimpleSSD storage dynamics.The RAG workload is assumed to have roughly 8GB of host memory available from the 16GB module.Source: A. Methodology · p. 9
  • The GPU baselines use full prefill or precomputed KV caches, while InstAttention offloads decode attention and AiF places model-weight computation in storage.InstAttention leaves prefill and recomputation on the host GPU; AiF can encounter out-of-memory errors for long contexts because its KV cache resides in GPU memory.Source: A. Methodology · p. 9
  • The ablations add fused recomputation-query attention, per-channel accelerators, and a global-SRAM reduction engine in successive configurations.The full MEMLLM configuration preserves an approximately 3W compute budget while enabling DRAM-less operation and eliminating PCIe egress bottlenecks.Source: A. Methodology · p. 9

B. Overall Performance

MemLLM outperforms the compared systems across datasets while reducing PCIe traffic by keeping KV-cache processing in storage. Its measured power stays within the stated 3W budget, and its accuracy difference from the baseline is within the margin of error.

Overall speedup

  • 4.06× average speedup over the GPU baseline and 4.18× over GPU-PreCompute, with MemLLM outperforming all compared systems across every dataset.GPU-PreCompute averages a 2.86% performance drop as context size increases.Source: B. Overall Performance · p. 10, Figure 15 · p. 10

Cumulative design gains

  • MemLLM achieves a 1.80× average speedup over baseline by running selective recomputation in the ISP; fusion adds 1.11×, per-channel accelerators add 1.82×, and global reduction adds 1.12×.Fusion halves NAND access volume, while global reduction prevents intermediate outputs from spilling into internal DRAM.Source: B. Overall Performance · p. 10, Figure 19 · p. 11

Power

Accuracy

  • MemLLM differs from the baseline by less than 0.3% on TriviaQA accuracy and F1-score, within the margin of error.The reported F1-score difference is 1.35× smaller than CacheBlend’s reported 0.002 deviation.Source: B. Overall Performance · p. 10, Figure 16 · p. 10

PCIe traffic

  • MemLLM limits PCIe traffic to the first layer’s precomputed KV cache, query vectors, and final attention outputs, making its footprint near-zero and independent of context length.The first-layer cache fetch is less than 5% of total cache volume.Source: B. Overall Performance · p. 10, Figure 17 · p. 10

Utilization

  • MemLLM eliminates PCIe movement and saturates aggregated NAND bandwidth, completing the task approximately 4× faster than GPU-Pre.During fused HKVD recomputation, GPU utilization reaches 100%.Source: Figure 18 · p. 10, B. Overall Performance · p. 11

C. Ablation Analysis

The ablations show that reducing NAND reads and distributing computation require a DRAM-less reduction path and pipelined NAND access to realize their benefits. Sensitivity and scalability tests examine recomputation, tile size, context and answer lengths, channel count, host memory, and retrieval and compression methods.

C. Ablation Analysis

  • The fused kernel halves NAND read time, per-channel accelerators reduce exposed NAND time, and global SRAM reduction avoids DRAM contention; pipelining raises effective throughput from nearly 11 GB/s to approximately 19 GB/s.Without DRAM-less reduction, partial-output traffic inflates computation latency and can exceed Base DRAM traffic on long contexts; double buffering hides NAND read latency.Source: C. Ablation Analysis · p. 11, Figure 20 · p. 11

D. Sensitivity and Scalability Analysis

  • Elastic-Filter fuses recomputed KV entries during NAND loading, eliminating the remaining SRAM staging bottleneck across recomputation ratios from 5% to 25%, including CacheBlend’s suggested 18%.The decoupled SIMD-plus-systolic baseline incurs substantial SRAM I/O, while Elastic reduces it through register-centric dataflows.Source: D. Sensitivity and Scalability Analysis · pp. 11–12
  • Latency improvement saturates beyond tK = 12, when total SRAM across channels reaches approximately 2MB.Larger tiles initially hide more exposed NAND read latency but increase SRAM requirements linearly.Source: D. Sensitivity and Scalability Analysis · p. 12
  • As context grows, baseline PCIe traffic scales linearly, while MemLLM maintains an O(1) PCIe footprint despite increasing compute time.The reported mechanism is keeping KV-cache processing within storage.Source: D. Sensitivity and Scalability Analysis · p. 12
  • MemLLM retains the highest end-to-end speedup across evaluated answer lengths by combining prefill savings with per-channel accelerators that meet decode throughput needs.InstAttention and AiF remove decode-phase PCIe transfers, but MemLLM continues to dominate in the reported comparisons.Source: D. Sensitivity and Scalability Analysis · p. 12
  • MemLLM scales near-linearly with NAND channel count by placing Elastic-SIMD units at channel endpoints and using aggregate NAND bandwidth without internal contention.The GPU baseline remains flat with storage parallelism, while InstAttention saturates at the SSD internal bus.Source: D. Sensitivity and Scalability Analysis · p. 12
  • At 16GB host DRAM, MemLLM maintains a 1.83× performance advantage over baseline, with latency remaining constant as host memory increases.The GPU-Full and GPU-Pre baselines improve with more host memory, but PCIe transfers remain unavoidable under the stated shared-memory budget for long contexts.Source: D. Sensitivity and Scalability Analysis · p. 12
  • Across the range both systems can serve, MemLLM retains its advantage over AiF while supporting contexts beyond AiF’s below-48k memory bound.With 100-token answers, decode dominates and AiF degrades more gracefully as context grows, but its KV cache must fit in GPU memory.Source: D. Sensitivity and Scalability Analysis · p. 12

VII. DISCUSSION

MemLLM targets on-device RAG with precomputed document KV caches, selective recomputation, long contexts, and short answers. Its deployment assumes computational storage with per-channel accelerators and central coordination logic.

Scope

  • MemLLM’s benefit narrows when long answers shift the bottleneck to decode or when KV caches fit in memory, though it still leads in the stated comparisons.The paper identifies the memory-constrained, long-context regime as the common case for practical on-device RAG.Source: §VII · p. 13

Deployment requirements

  • Deployment requires computational storage, and the custom NVMe command uses a vendor-specific opcode; the described software path relies on io_uring-based NVMe passthrough and GPU–device P2P DMA support.The paper notes that unified-memory SoCs can simplify the path relative to conventional systems.Source: §VII · p. 13

VIII. RELATED WORKS

Prior work addresses prompt-cache reuse and selected LLM operations in storage, but does not support the complete precomputed-KV RAG attention pipeline. MemLLM combines per-channel accelerators, fused recomputation-query processing, and DRAM-less reduction to execute that pipeline in the SSD.

VIII. RELATED WORKS

  • For RAG, independently precomputed document caches disrupt cross-attention, so CacheBlend and CacheClip [30] selectively recompute tokens to restore quality.These software optimizations speed high-bandwidth servers but encounter an I/O wall on resource-constrained edge devices.Source: §VIII · p. 13
  • Decode-phase ISPs,, AiF, and RAGX address other LLM inference or retrieval workloads, but the cited approaches do not support end-to-end RAG compute over precomputed KV caches.The passage attributes the mismatch to selective recomputation or long-context memory limits.Source: §VIII · p. 13

IX. CONCLUSION

  • MemLLM executes the complete RAG attention pipeline in the SSD, using per-channel Elastic-SIMD accelerators, a fused recomputation-query kernel, and DRAM-less global reduction.Across eight real-world RAG datasets and six on-device LLM models, it achieves an average 4.06× end-to-end speedup compared with GPU-based systems.Source: §IX · p. 13

Related papers

Loading 2610.22579v1…