Institute of Science Tokyo's AQPIM Squeezes 10x Faster AI Memory Without Extra Chips
A new HPCA paper redesigns processing-in-memory around product quantization, cutting KV cache footprint 6.5x and delivering 3.4x speedup over the previous SOTA PIM accelerator.
PRO- AQPIM performs Product Quantization inside HBM-PIM to compress KV cache 6.53x without breaking capacity limits (paper).
- Attention runs directly on compressed data as lookups plus sums, reusing existing FP16 MAC units with no dequantization.
- Page-aware windowed clustering keeps lookup tables inside a single DRAM row, turning random accesses into row-buffer hits.
- Delivers 3.4x speedup over AttAcc! and up to 14.29x better energy efficiency per decoding step.
- Adds only 0.43% area overhead to BankPEs; importance-weighted k-means preserves accuracy on LongBench.
- Evaluated on Mistral-7B and Llama-3.2-3B with a customized AttAcc simulator based on Ramulator 2.0.
Processing-in-Memory sounds like the obvious fix for the memory wall stalling long-context LLM inference: put compute next to the DRAM, skip the trip across the bus, and attention becomes bandwidth-rich instead of bandwidth-starved. The catch is that PIM has its own wall, and it concerns capacity rather than bandwidth. A new paper from a group at Institute of Science Tokyo, accepted to HPCA, proposes AQPIM, a co-design that puts product quantization directly inside HBM-PIM banks and computes attention on the compressed data without ever decoding it.
The capacity wall nobody talks about
Prior PIM accelerators for LLMs like AttAcc! and NeuPIMs treat attention as a bandwidth problem, using PIM's high internal bandwidth to speed up the GEMV between a fresh query and the KV cache. Existing PIM architectures buckle under the sheer volume of KV caches generated in long-context scenarios, which can reach hundreds of GBs. A SOTA accelerator like AttAcc! already requires as many as 40 HBM-PIM devices to support even short contexts, which becomes economically unviable at 128K or 1M token contexts.
The usual escape hatches all fight PIM's architecture. Offloading to host memory hits PCIe bandwidth, which is roughly 26x slower than HBM. Sparse attention produces scattered access patterns that clash with PIM's need for contiguous locality. Bolting standard integer quantization onto bank-level compute is prohibitive: simply adding integer MAC units could increase the logic area from 50% (FP16 only) to as much as 126% (FP16+INT32), severely compromising memory density.
Turning a weakness of clustering into a PIM strength
Clustering-based vector quantization, specifically Product Quantization (PQ), compresses activations well because their distribution is highly context-dependent, exhibiting significant locality and similarity, forming tight clusters rather than the evenly-spread distributions weights tend to have. The problem is that clustering is bandwidth-hungry: the iterative centroid recomputation has kept it out of online use on GPUs.
This story is for Pro members
You've reached the end of the free preview. Upgrade to AlphaSignal Pro to read the full article - and everything else behind the paywall.
Comments
