AQPIM: Breaking the PIM Capacity Wall for LLMs with In-Memory Activation Quantization

2026-04-20

AI Summary

This work addresses the limitation of Processing-in-Memory (PIM) architectures in large language model inference with long contexts, where activation memory—particularly KV cache—exceeds PIM capacity. To overcome this challenge, the authors propose a PIM-aware activation quantization framework that, for the first time, jointly leverages PIM hardware characteristics and the statistical distribution of activations. Their approach introduces a clustering-based product quantization strategy that enables direct compression of activations and attention computation within PIM, substantially mitigating accuracy degradation. By offloading critical operations to PIM, the method dramatically reduces GPU–CPU communication overhead—which accounts for 90% to 98.5% of decoding latency—and achieves a 3.4× speedup over the state-of-the-art PIM-based solution.

Technology Category

• Machine Learning: Hardware-aware ML
• Natural Language Processing: (Large) Language Models
• Data Mining & Knowledge Management: Intelligent Query Processing

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Large language models for search

Abstract

Processing-in-Memory (PIM) architectures offer a promising solution to the memory bottlenecks in data-intensive machine learning, yet often overlook the growing challenge of activation memory footprint. Conventional PIM approaches struggle with massive KV cache sizes generated in long-context scenarios by Transformer-based models, frequently exceeding PIM's limited memory capacity, while techniques like sparse attention can conflict with PIM's need for data locality. Existing PIM approaches and quantization methods are often insufficient or poorly suited for leveraging the unique characteristics of activations. This work identifies an opportunity for PIM-specialized activation quantization to enhance bandwidth and compute efficiency. We explore clustering-based vector quantization approaches, which align well with activation characteristics and PIM's internal bandwidth capabilities. Building on this, we introduce AQPIM, a novel PIM-aware activation quantization framework based on Product Quantization (PQ), optimizing it for modern Large Language Models (LLMs). By performing quantization directly within memory, AQPIM leverages PIM's high internal bandwidth and enables direct computation on compressed data, significantly reducing both memory footprint and computational overhead for attention computation. AQPIM addresses PQ's accuracy challenges by introducing several algorithmic optimizations. Evaluations demonstrate that AQPIM achieves significant performance improvements, drastically reducing of GPU-CPU communication that can account for 90$\sim$98.5\% of decoding latency, together with 3.4$\times$ speedup over a SOTA PIM approach.

Problem

• Processing-in-Memory
• activation memory
• Large Language Models
• KV cache
• memory bottleneck

Innovation

• Processing-in-Memory
• Activation Quantization
• Product Quantization
• Large Language Models
• KV Cache