Selective KV Cache Protection for Noise-Resilient LLM Inference on Analog Compute-In-Memory Systems

2026-07-31

AI Summary

This work addresses the incompatibility between the dynamic key-value (KV) cache updates in large language model attention mechanisms and the static weight paradigm of analog-in-memory computing (AiMC) systems, which renders inference highly sensitive to hardware noise and severely degrades accuracy. The study presents the first systematic analysis of this vulnerability, revealing that initial and recent tokens are particularly noise-sensitive. To mitigate this, the authors propose a hierarchical protection strategy: critical tokens’ KV caches are processed via a high-precision digital path, while the rest are handled by analog units. A sliding-window scheduling scheme coupled with an ownership migration mechanism jointly controls digital overhead. Evaluated across nine LLMs, the approach reduces average perplexity from 33.91 to 11.95—approaching the noise-free baseline of 11.06—and achieves 91.2% utilization of dynamic KV programming rows, significantly enhancing noise robustness while preserving high energy efficiency.

Abstract

Analog compute-in-memory (CIM) arrays have emerged as a promising substrate for energy-efficient LLM inference, particularly for weight-stationary computations in linear layers. However, extending analog CIM to attention mechanisms introduces a fundamental challenge: KV cache operations demand repeated in-situ weight updates, and the resulting mismatch with the weight-stationary paradigm exposes dynamic computations to significant hardware noise, a critical problem that remains largely unexplored. In this paper, we present the first systematic study of dynamic attention computation on analog CIM arrays, revealing that initial and recent tokens exhibit disproportionate vulnerability to hardware noise. Motivated by this token-level insight, we propose a hierarchical token protection strategy that keeps sink tokens and a sliding recent-token window on a higher-precision digital path while processing the bulk KV cache on analog CIM. A co-designed scheduler combines analog programming, ownership transition, and bulk-MVM tile formation to bound digital overhead. Evaluations on nine LLMs show that our approach lowers average perplexity under analog noise from 33.91 to 11.95, close to the clean baseline of 11.06, while improving dynamic-KV programming-row utilization from 23.1\% to 91.2\%.

Problem

• analog compute-in-memory
• KV cache
• hardware noise
• attention mechanism
• LLM inference

Innovation

• analog compute-in-memory
• KV cache protection
• noise-resilient LLM inference
• hierarchical token scheduling
• dynamic attention computation