AQPIM在内存中量化激活值,突破大模型推理的内存瓶颈
AQPIM: Breaking the PIM Capacity Wall for LLMs with In-Memory Activation Quantization

- 在内存中直接对激活值进行分组量化,利用PIM高带宽特性
- 减少90%~98.5%的GPU-CPU通信,实现3.4倍于当前最优方案的速度提升
- 专为大语言模型设计,适合追求低延迟推理的部署场景
存内计算(PIM)架构能有效缓解数据密集型机器学习中的内存瓶颈,但常忽视激活值内存占用日益增长的问题。基于Transformer的大模型在长上下文场景下生成的键值缓存(KV cache)规模庞大,常超出PIM有限的内存容量,而稀疏注意力等技术又与PIM对数据局部性的需求冲突。现有PIM方法和量化策略往往不充分或不匹配激活特征。本文发现针对PIM特性的激活量化存在优化空间,提出基于乘积量化(Product Quantization, PQ)的AQPIM框架,直接在内存中完成量化,充分利用PIM内部高带宽,实现压缩数据上的直接计算,显著降低内存占用与注意力计算开销。通过多项算法优化解决PQ精度问题,实验表明:AQPIM可将解码延迟中占比90%~98.5%的GPU-CPU通信大幅减少,并相较当前最优PIM方案提速3.4倍。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。