arXiv: 2604.18137 · PDF

作者: Kosuke Matsushima, Yasuyuki Okoshi, Masato Motomura, Daichi Fujiki

主分类: cs.AR · 全部: cs.AI, cs.AR, cs.LG

命中关键词: large language model, llm, rag, kv cache, quantization, attention, transformer, gpu, latency

TL;DR

AQPIM 在 PIM 内部用 Product Quantization 压缩 LLM 激活/KV cache，突破 PIM 容量墙，相比 SOTA PIM 方案加速 3.4 倍。

核心观点

  • 传统 PIM 聚焦权重瓶颈，却忽视了长上下文 Transformer 激活（尤其 KV cache）的内存压力。
  • 稀疏 attention 与 PIM 所需的数据局部性冲突，通用量化方法也不适配 PIM 特性。
  • 基于聚类的向量量化（VQ）与激活分布及 PIM 内部带宽高度契合。
  • 提出 AQPIM：PIM-aware 的 PQ 激活量化框架，在内存内直接量化并对压缩数据进行计算。

方法

  • 以 Product Quantization 为基础，将激活切分子向量并用码本索引替代原始值。
  • 量化过程直接发生在 PIM 内部，利用其高内部带宽避免数据搬运。
  • attention 计算在压缩域上执行，减少解码算力与访存开销。
  • 引入多项算法优化（针对 PQ 精度损失），以适配现代 LLM 的激活分布特性。

实验

摘要未给出具体数据集、模型规模或基线细节，仅提到与 SOTA PIM 方法对比，并测量 GPU-CPU 通信占 decoding latency 比例。

结果

  • GPU-CPU 通信（原占 decoding latency 的 90%~98.5%）被显著削减。
  • 相对 SOTA PIM 方案获得约 3.4× 端到端加速。
  • 精度影响通过 PQ 算法优化缓解，但摘要未披露具体精度数字。

为什么重要

对 LLM 推理基础设施从业者而言，AQPIM 把 PIM 的价值从"权重搬运"拓展到"激活 / KV cache 压缩+就地计算"，为长上下文推理提供了绕过 PIM 容量墙的新路径，并直接攻击 GPU-CPU 通信这一 decoding 主瓶颈。

与已有工作的关系

  • 延续 HBM-PIM、UPMEM、Samsung/SK Hynix PIM 等面向 ML 的 PIM 架构路线。
  • 借鉴 Product Quantization / 向量量化在 ANN 检索与模型压缩（如 PQ-KV、ScaNN）中的思路。
  • 与 KV cache 量化方案（KVQuant、KIVI、SmoothQuant 等）互补，但将量化下沉至内存侧。
  • 和稀疏 attention（H2O、StreamingLLM）形成对照，强调稠密但压缩的 PIM 友好方案。

尚未回答的问题

  • 在主流 LLM（LLaMA、Qwen 等）和长上下文 benchmark 上的 perplexity / 任务精度损失具体多少？
  • PQ 码本训练成本、在线更新策略以及对 prefill 阶段的影响？
  • 对不同 PIM 硬件（HBM-PIM vs. LPDDR-PIM）和 batch size 的通用性？
  • 与 KV 量化 + 稀疏 attention 组合后的收益上限在哪里？

论文图表

图 1: Figure 1 (extracted from PDF)

图 1

原始摘要

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.