一句话:提出 GPU-PIM 混合系统 REPA,用可重构 ReRAM PIM 联合加速 LLM 推理中的 KV cache 卸载与处理。
大语言模型推理中的 KV cache 导致高内存占用与解码性能下降,已有工作通常只通过卸载或阶段拆分解决其一。本文提出 GPU-PIM 混合系统 REPA,联合加速 KV cache 卸载与处理:利用可重构 ReRAM PIM 实现快速持久化并平衡处理速度与容量。为释放并行性,REPA 在架构、数据映射和流水线层面优化,包括批量式内存指令与多级控制器、局部性感知映射、子批次流水线和传输重叠。实验显示,其相比 NVIDIA A100 推理快 1.5--6.5 倍、能效高 8--10 倍,长上下文下超过先进 DRAM PIM 最多 1.4 倍;集成现有卸载系统可提升卸载速度 1.4--2.0 倍、端到端加速 1.2--1.4 倍。
The use of KV cache in LLM inference leads to large memory footprint and sub-optimal decoding performance. Prior studies typically address one of these two limitations by either offloading or stage-split inference. In this paper, we explore and reveal the possibility of a joint solution, and propose REPA, a GPU-PIM hybrid system to prototype this idea. We leverage reconfigurable ReRAM PIM to achieve fast KV cache persistence, and balance the requirement of processing speed and memory capacity. To fully unleash the parallelization potential of REPA, we propose optimizations in (1) architecture, (2) data mapping and (3) pipelining: (1) We propose bulk-wise memory instructions and multi-level controllers to enable finer-grained parallelism in the PIM device. (2) We propose locality-aware data mapping to make the best of the aforementioned architectural optimization, and reduce long-range data transfer on chip. (3) We adopt sub-batch pipelining to reduce idleness in batches, and propose transfer overlapping to shadow the KV cache transfer by computation. Experimental results show that REPA exhibits high inference speed, energy efficiency and integratability. It is 1.5--6.5× faster, and 8--10× more efficient than NVIDIA A100. It also outperforms state-of-the-art DRAM PIM systems by up to 1.4× for long context inference. When integrated into existing offloading systems, REPA achieves 1.4--2.0× offloading speed, and 1.2--1.4× end-to-end speedup, showcasing its high potential for fast KV cache offloading and processing.
@inproceedings{HongY0026,
title = {REPA:
<u>Re</u>
configurable
<u>P</u>
IM for the Joint
<u>A</u>
cceleration of KV Cache Offloading and Processing},
author = {Yang Hong and Junlong Yang and Bo Peng and Jianguo Yao},
booktitle = {ASPLOS},
year = {2026},
doi = {10.1145/3779212.3790212}
}