LLM Inference in a Flash!
Abstract Overview
This paper studies how to adapt LLM decoding to flash-based compute-in-memory (CIM) architectures, which provide high storage capacity and near-data execution but suffer from limited write endurance, low off-chip bandwidth, and poor high-precision floating-point support. To address these hardware bottlenecks, the authors propose an end-to-end integer-only quantization pipeline that executes linear and nonlinear operators without floating-point offloading, alongside a dictionary-based KV-cache compression strategy. The compression scheme represents past keys and values as sparse codes over a static, flash-resident dictionary while maintaining a recent local window in dense form and applying query-aware hierarchical sparsity. Evaluated on Llama-3.1-8B and Qwen-2.5-7B, the approach preserves task accuracy, while an analytical system model estimates substantial decode latency and energy reductions on a CIM-SSD architecture compared to an NPU baseline.
Novelty
The core novelty is the joint algorithmic co-design tailored specifically to flash compute-in-memory constraints rather than GPU memory hierarchies. The method combines an integer-only execution pipeline across all operators with a projection-based sparse coding pursuit that replaces computationally heavy least-squares refitting to minimize write traffic and dynamic KV-cache data movement.
Results
On Llama-3.1-8B, the integer-only pipeline achieves a WikiText-2 perplexity of 7.5803 compared to 7.5454 for the FP16 baseline, eliminating floating-point operations from the critical decode path. The combined compression scheme delivers a 15× reduction in dynamic KV-cache traffic while retaining 99% of the FP16 full-KV LongBench score across both Llama-3.1-8B and Qwen-2.5-7B. Analytical system modeling projects that the default S2 CIM-SSD configuration achieves 3.1× latency and 2.7× energy savings at 1K context length, scaling to 4.4× latency and 6.8× energy savings at 256K context length relative to a DRAM-based NPU baseline.
Key Points
- An end-to-end integer-only LLM quantization pipeline replaces floating-point operations with integer approximations for Softmax, RMSNorm, and SiLU, leveraging activation smoothing, group-wise quantization, and INT16 requantization to preserve stability.
- A flash-resident dictionary KV-cache compression method uses projection-based pursuit and query-aware hierarchical refinement to represent past context as sparse codes, avoiding costly least-squares solving and minimizing flash write wear.
- Evaluation shows minimal accuracy loss across WikiText-2 and LongBench benchmarks, with analytical hardware modeling projecting up to 4.4× latency and 6.8× energy reductions for long-context inference on a CIM-SSD platform.