L3：面向可扩展长上下文 LLM 推理的 DIMM-PIM 集成架构与协同

大型语言模型（LLM）越来越需要处理长文本序列，但 GPU 内存限制迫使在内存容量和带宽之间进行艰难的权衡。虽然基于 HBM 的加速提供了高带宽，但其容量仍然受限。将数据卸载到主机端 DIMM 可提高容量，但引入了昂贵的数据交换开销。我们发现，关键的内存瓶颈完全在于多头注意力（MHA）的解码阶段，该阶段需要大容量来存储 KV 缓存，并需要高带宽进行注意力计算。我们的关键见解揭示了该操作与现代基于 DIMM 的存内处理（PIM）架构独特契合，后者同时提供容量和带宽的可扩展性。基于这一观察和见解，我们提出了 L3，一个集成 DIMM-PIM 和 GPU 设备的软硬件协同设计系统。L3 引入了三项创新：首先，硬件重新设计解决了 DIMM-PIM 中的数据布局不匹配和计算单元不匹配问题，提高了 LLM 推理的利用率。其次，通信优化实现了将数据传输开销与计算相隐藏。第三，自适应调度器协调 GPU-DIMM-PIM 操作，以最大化设备间的并行性。使用真实世界 trace 的评估表明，与 SOTA 的 HBM-PIM 解决方案相比，L3 实现了高达 6.1$\times$ 的加速，同时显著提高了批处理大小。