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.

翻译：处理中内存（PIM）架构为数据密集型机器学习中的内存瓶颈提供了有前景的解决方案，但往往忽视了日益增长的激活内存占用问题。传统PIM方法难以应对Transformer模型在长上下文场景中生成的庞大KV缓存规模，这一规模常超出PIM的有限内存容量，而稀疏注意力等技术可能与PIM对数据局部性的需求产生冲突。现有PIM方法与量化技术往往不足以利用激活值的独特特性，或难以适应该特性。本文识别出面向PIM的专用激活量化在提升带宽和计算效率方面的机遇。我们探索了与激活特性及PIM内部带宽能力高度契合的聚类向量量化方法。在此基础上，提出AQPIM——一种基于乘积量化（PQ）的新型PIM感知激活量化框架，针对现代大型语言模型（LLM）进行优化。通过在内存内部直接执行量化，AQPIM充分利用PIM的高内部带宽，并实现对压缩数据的直接计算，显著降低注意力计算的内存占用与计算开销。针对PQ的精度挑战，AQPIM引入多项算法优化。评估表明，AQPIM实现了显著的性能提升，大幅减少GPU-CPU通信（该通信可占解码延迟的90%至98.5%），且相比最先进的PIM方法获得3.4倍加速。