PaperlayerView exact revision on arXiv

Source-linked AI summary

LLM Inference in a Flash!

Sebastian Zhao, Minseo Kim, Coleman Hooper, Luca Manolache, Michael W. Mahoney, Yakun Sophia Shao, Kurt Keutzer, Amir Gholami

arXiv:2609.16161v1cs.LG

TL;DR

Longer-context LLM inference increases memory pressure, while Flash CIM is constrained by limited floating-point support and write endurance. The paper combines integer-only inference with dictionary-based KV-cache compression. On LongBench, the combined method retains 99% of the FP16 full-KV baseline score for both tested models, with nearly 15× KV compression.

  • Problem

    Longer-context LLM decoding is memory-bandwidth-bound as each token reads the growing KV cache, while Flash's limited write endurance and floating-point support constrain its use for inference.

    Source: §1 · p. 1

  • Method

    The paper combines end-to-end integer-only inference with sparse dictionary coding that represents KV vectors using a static dictionary.

    Source: §1 · p. 2, §3.2.3 · p. 6

  • Results

    99% of the FP16 full-KV baseline score is retained on LongBench for both Llama-3.1-8B and Qwen-2.5-7B, with less than 1 point average degradation and nearly 15× KV compression.

    Source: §5.4 · p. 12

  • Takeaways & Limitations

    The design maps both weights and KV-cache operations to Compute-in-Flash while minimizing off-chip data transfers and dynamic memory operations.

    Source: §1 · p. 2

    The study focuses on generation; efficiently mapping the compute-bound prefill phase to compute-in-memory platforms remains future work.

    Source: §7 · p. 13

Abstract

from arXiv · show

Large Language Models (LLMs) have shown impressive capabilities across a range of natural language processing tasks, and LLM inference has emerged as a critical workload for enabling downstream applications. The demands of serving LLM inference are becoming increasingly challenging as requests shift toward longer sequences and heavier inference, driven by retrieval-augmented generation, inference-time compute scaling, and long-context applications. Additionally, these challenges are compounded by hardware trends, as memory capacity and communication bandwidth are not scaling as fast as increases in workload complexity. Compute-in-Flash is a promising solution to address memory bandwidth limitations by moving computation close to memory, and to exploit the large capacity of SSD technologies. However, it is challenging to deploy LLMs on these systems as they lack support for high-precision floating point operations and have limited write endurance. In our work, we aim to address these challenges by designing inference algorithms to enable LLM inference on Flash compute-in-memory devices. We present an end-to-end integer-only quantization approach to eliminate expensive floating-point computations. To address the limited write endurance, we design a dictionary-based KV cache compression strategy based on sparse dictionary coding that represents each KV vector as a linear combination of static dictionary vectors. These algorithmic improvements enable us to exploit the benefits of Compute-in-Flash for both model weights and KV cache, and to minimize expensive data transfer operations. Across Llama-3.1-8B and Qwen-2.5-7B, our combined method exhibits limited accuracy degradation while reducing dynamic KV cache traffic by 15$\times$.

1 Introduction

Longer-context LLM inference is memory-bandwidth-bound, while Flash offers capacity but is constrained by write endurance and limited floating-point support. The paper pairs integer-only inference with dictionary-based KV-cache compression and evaluates the combined design.

1 Introduction

  • LLM decoding is memory-bandwidth-bound because each token reads the growing KV cache, and longer contexts intensify this pressure.The paper identifies retrieval-augmented generation and inference-time compute scaling as workload trends that increase pressure.Source: §1 · p. 1
  • Flash offers greater capacity, but limited write endurance, low off-chip bandwidth, and lack of high-precision floating-point support constrain its use for LLM inference.Compute-in-memory moves computation close to memory to address bandwidth limits and can reduce transfer energy and latency.Source: §1 · pp. 1–2
  • The design enables LLM inference on Compute-in-Flash by combining end-to-end integer quantization, including nonlinear operations, with KV-cache compression using a large read-only dictionary to reduce dynamic state and write traffic.Source: §1 · p. 2
  • End-to-end integer-only quantization covers attention and nonlinear operations, avoiding floating-point computation and FP offloading on Compute-in-Flash devices.The paper also presents a system-level analytical model of latency and energy benefits.Source: §1 · p. 2
  • Sparse dictionary coding represents each KV vector using a static, flash-resident dictionary to minimize write traffic, while only lightweight metadata is handled by the host.The design exploits flash capacity and internal read bandwidth.Source: §1 · p. 2
  • 3.1×/2.7× and 4.4×/6.8× are the modeled latency/energy savings for short and long contexts, respectively, under the S2 configuration.The savings are quantified by a system-level analytical model under realistic Compute-in-Flash constraints (Section 5.5).Source: §1 · p. 2

2 Related Work

Prior work addresses Flash-based LLM inference through hardware and system strategies, while quantization and KV-cache methods often target other hardware constraints. This section positions algorithmic co-design for Flash CIM as a response to its limited floating-point support and write endurance.

2.1 Compute-in-Memory for LLM Inference

  • Prior Flash-CIM work targets hardware support or system scheduling and placement, leaving limited work on reshaping LLM inference algorithms for floating-point and write-endurance constraints.KVNAND [8] optimizes KV placement for device lifetime; reported write-volume estimates illustrate the endurance concern.Source: §2.1 · p. 3

2.2 Quantization

  • GPU-oriented quantization methods reduce precision or approximate nonlinearities, but many retain floating-point operations or dequantization that poorly match Flash/PIM controllers.The paper targets integer execution for modern LLM decoding, including nonlinear and reduction operators.Source: §2.2 · p. 3

2.3 KV Cache Compression

  • KV-cache research includes eviction, sparse attention, quantization, and rematerialization; Lexico [26] uses a static dictionary, while this work adapts dictionary compression and hierarchical sparsity to CIM constraints.The cited approaches address memory footprint or bandwidth, often in GPU settings.Source: §2.3 · p. 4

3.1 Integer-Only LLM Quantization

Flash/PIM controllers make floating-point nonlinear and reduction operators a deployment challenge. The quantization pipeline combines integer execution with adaptations for activation range and nonlinear operations.

3.1.1 Quantization Challenges for PIM

  • Integer-only execution targets Flash/PIM constraints by quantizing linear operators to W8A8 and implementing key nonlinear and reduction operators with integer approximations.The pipeline aims to minimize dequantization and host round-trips, while adapting prior integer-only formulations to modern LLMs and long contexts.Source: §3.1.1 · p. 4

3.1.2 Activation Smoothing

  • Activation smoothing rescales activations and compensates in weights, while group-wise scales address the sensitivity of nonlinear-operation inputs and outputs.Naive INT8 quantization can incur substantial error from activation outliers and high dynamic range.Source: §3.1.2 · p. 4

3.1.3 Nonlinear Operations

  • Tokenwise INT16 requantization makes Softmax and RMSNorm reductions valid across hidden dimensions without full-precision dequantization.Second-order polynomial approximations handle nonlinearities; overflow is managed by shifting inputs and outputs, and RMSNorm uses an integer square-root routine and fixed-point reciprocal.Source: §3.1.3 · p. 5

3.2 Dictionary-Based KV Cache Compression

To fit flash’s read-heavy, write-limited profile, the method stores a static dictionary and sparse KV codes, then attends over compressed context without materializing KV.

3.2.1 Hardware Constraints and Design Rationale

  • A static flash-resident dictionary and sparse KV codes replace frequent writes from appending each decoding step’s keys and values.A small recent window remains in full precision; older tokens are encoded and moved into the compressed context.Source: §3.2.1 · p. 5, §3.2.4 · p. 7

3.2.2 Learning the Overcomplete Dictionary

  • The hybrid dictionary-learning procedure trains only the decoder as the dictionary, using projection-based sparse codes to reconstruct buffered key and value activations.It extracts activations per attention layer and key/value stream, buffers them in a shuffled pool, and applies per-tensor quantization to the learned dictionary for flash CIM.Source: §3.2.2 · pp. 5–6, Figure 2 · p. 6

3.2.3 Projection-Based Sparse Coding

  • Projection-based greedy pursuit replaces OMP-style least-squares refitting with correlations, direct coefficients, and residual updates.The same compress() operator generates training codes and encodes inference-time KV states, avoiding train–inference mismatch.Source: §3.2.3 · p. 6

3.2.4 Overall Pipeline.

  • At runtime, older KV tokens are batch-compressed into sparse codes while a recent local window stays dense; attention aggregates dictionary values without materializing compressed keys and values.Compressed-context and local-window contributions are combined under one softmax normalization.Source: §3.2.4 · p. 7, Figure 3 · p. 7

3.2.5 Query-Aware Hierarchical Sparsity

  • Query-aware hierarchical sparsity first reconstructs all tokens with a low budget, then selectively raises the budget for relevant tokens to reduce compute and read I/O.The coarse key approximation uses the first Klow dictionary atoms to form a proxy attention logit; selected tokens are refined from Klow to Khigh.Source: §3.2.5 · p. 8

3.2.6 Why Flash CIM?

  • Dictionary-based KV coding fits flash CIM better than conventional accelerators, where repeated dictionary streaming can add memory traffic and offset compression gains.On conventional accelerators, the stated use case is when the KV cache no longer fits device memory, reducing host–device offload bandwidth.Source: §3.2.6 · p. 8

4 System Modeling

The section models decoding on two NPU–DRAM baselines and two CIM-SSD variants, then estimates latency and per-token energy by representing system components, data movement, and execution overlap.

Overview and metrics

  • The analytical model targets batch-size-1 autoregressive decoding of Llama 3.1–8B and reports latency, throughput, and total per-token energy, including KV read and write costs.The model is intended to capture trends with context length, controller SRAM capacity, and pursuit placement.Source: §4 · p. 9

Systems and workload

  • Four configurations compare an NPU with full-FP16 KV storage (B1), the same NPU with dictionary compression (B2), and two CIM-SSD variants (S1 and S2).B2 performs compression and reconstruction on the NPU; S2 places compressed codes in DRAM and fetches them to the CPU during reconstruction.Source: §4 · p. 9

Cost primitives and SRAM budget

  • The model combines compute and transfer costs, permits overlap only for plausible concurrent work, and treats decode stages and pursuit iterations as sequential.S1 SRAM must hold the local window, transient pursuit state, and growing compressed-history codes; S2 stores coefficients in DRAM.Source: §4 · pp. 9–10

Layer modeling

  • Linear projections are modeled as GEMV: B1 streams FP16 weights from DRAM, while CIM-SSD keeps INT8 weights stationary in flash and returns INT32 partial sums for INT8 requantization.CIM computation and internal-bus transfers overlap within the linear stage.Source: §4 · p. 9
  • On CIM-SSD, nonlinear operators run on the SSD controller using integer-only approximations, whereas the NPU baseline runs them in FP32.Examples include polynomial/shift-based exponentiation for IntSoftmax and integer square root and division for IntRMSNorm.Source: §4 · p. 10

KV cache write path

  • New INT8 KV vectors enter a local SRAM window and are periodically compressed by pursuit on the CIM array into sparse index–coefficient codes; B1 instead writes FP16 vectors to DRAM.B2 performs compression on the NPU with dictionaries and codes in DRAM, adding DRAM transfers during encoding.Source: §4 · p. 10

KV cache read path

  • For compressed context, S1 computes attention in the dictionary-score domain without materializing full KV, while S2 stores codes in DRAM and fetches them to the CPU for reconstruction.B1 reads the full FP16 cache; B2 reads compressed codes and dictionaries from DRAM and reconstructs KV on the NPU, incurring repeated dictionary transfers.Source: §4 · p. 10

5 Results

The combined pipeline is evaluated for quality, KV-cache compression, and projected system performance across model configurations and context lengths. Results include near-15× KV compression with limited LongBench degradation and modeled latency and energy savings for CIM-SSD.

5.1 Experimental Setup

  • Dictionary training uses Pile calibration activations, 32,768 atoms per layer and key/value, a 128-token window, and compression of the oldest 32 tokens when full.The hierarchical policy uses Khigh=16 and Klow=3; quantization uses group size 128.Source: §5.1 · p. 10

5.2 Integer-Only LLM Inference

  • 7.5803 WikiText-2 perplexity is within 0.0349 of the FP16 baseline while the final pipeline eliminates floating-point computation from the critical path.INT16 requantization at Softmax and RMSNorm inputs mitigates degradation from integer polynomial approximations.Source: §5.2 · p. 11, Table 1 · p. 11

5.3 Dictionary-Based KV Cache Compression

  • 15× lower effective KV-cache footprint is achieved without degrading LongBench score; projection-based coding avoids costly least-squares operations.The pipeline combines INT8 dictionary quantization with query-aware hierarchical sparsity, which uses Klow=3 for coarse scoring and refines the top 10% to Khigh=16.Source: §5.3 · p. 11, Table 2 · p. 11

5.4 End-to-End Results

  • 99% of the FP16 full-KV baseline score is retained for both models, with less than 1 point average degradation and nearly 15× KV compression.The end-to-end LongBench evaluation covers Llama-3.1-8B and Qwen-2.5-7B.Source: Table 3 · p. 12, §5.4 · p. 12

5.5 System-Level Modeling

  • Against B1, S1 is 3.4× faster and 2.7× more energy-efficient at 1K; B2 is 2.9× faster and 2.6× more energy-efficient at 256K.At 1K, B2 is slightly slower and less energy-efficient than B1 because loading its dictionary from DRAM offsets KV-transfer savings.Source: §5.5 · pp. 12–13
  • 3.1×/2.7× latency/energy savings at 1K and 4.4×/6.8× at 256K are projected for the default S2 configuration.S2 is the configuration treated as the default serving option because it supports both short and long contexts.Source: §5.5 · p. 12, Table 4 · p. 13
  • At 256K, S1 and B2 have comparable latency, but S1 uses 3.3× less energy by avoiding dictionary streaming from DRAM.The CIM-SSD architecture keeps dictionaries stationary in flash and accesses them through high-bandwidth internal channels.Source: §5.5 · p. 13
  • S1 is slightly faster at 1K, but S2 has a 1.5× latency advantage at 256K; S1 remains more energy-efficient across context lengths.S1 keeps codes in controller SRAM, while S2 places them in host DRAM; S2 is preferred when SRAM capacity is exceeded.Source: §5.5 · p. 13

6 Conclusion

The paper co-designs LLM inference for Compute-in-Flash, targeting memory bottlenecks with integer computation and read-only information for KV-cache management. Its modeled results show latency and energy savings at both short and long contexts.

6 Conclusion

  • The approach adapts inference to use efficient integer computation and KV-cache management based on read-only information to exploit internal read bandwidth.The paper presents algorithmic solutions for Compute-in-Flash constraints.Source: §6 · p. 13
  • 3.1×/2.7× latency/energy savings are reported for short contexts, and 4.4×/6.8× for long contexts.These are the reported savings for the approach described in the conclusion.Source: §6 · p. 13
  • The paper concludes that its adaptations enable efficient inference and reduce the core memory bottleneck for LLM inference.Source: §6 · p. 13

7 Limitations

The study focuses on generation, while efficient mapping of prefill remains future work. Its Value-dictionary design also assumes storing two copies to support different operations.

7 Limitations

  • The study evaluates the generation phase, not both prefill and decode.The authors identify mapping both phases efficiently as future work.Source: §7 · pp. 13–14
  • The compute-bound prefill phase is described as more challenging to map efficiently to compute-in-memory platforms, which typically have lower peak compute performance than CPU/GPU platforms.Source: §7 · p. 13
  • The design assumes two Value-dictionary copies: one transposed for pursuit and one non-transposed for attention.Source: §7 · p. 14

A System Modeling Parameters

The analytical system model specifies compute and memory energy, interface bandwidth, storage capacity, and peak-throughput assumptions.

A System Modeling Parameters

B Nonlinear Operation Offloading Ablations

The ablation compares nonlinear-operation offloading to the NPU with integer approximations in the controller. The integer approach uses less energy, achieves speedup, and scales better as context length increases.

B Nonlinear Operation Offloading Ablations

C Full System Results

The full-system results table reports decode latency and energy across three context lengths for two NPU-plus-DRAM baselines and two CIM-SSD variants.

C Full System Results

Loading 2609.16161v1…