[Paper] PIM-SHERPA：通过解决 PIM 内存属性和布局不一致性实现设备端 LLM 推理的软件方法

Source: arXiv - 2603.09216v1

Overview

论文 PIM‑SHERPA 解决了一个阻碍设备端大语言模型（LLM）推理的实际难题：在使用内存中处理（PIM）加速器时，prefill（计算密集）阶段和 decode（内存带宽密集）阶段的内存处理方式冲突。通过引入一种纯软件方案来调和这些内存属性和布局的不匹配，作者在真实硬件上展示了接近最优的性能，同时将内存使用量降低了近一半。

关键贡献

• 识别“内存属性不一致”：prefill 需要在可缓存区域预置权重以供复用，而 decode 需要在不可缓存区域以触发 PIM 操作。
• 暴露“权重布局不一致”：主机友好的权重顺序与为高效内存计算所需的交错（PIM‑感知）布局不同。
• 提出 PIM‑SHERPA，一个纯软件框架，通过两种互补技术解决上述不一致：
  1. DRAM 双缓冲 (DDB) – 在不可缓存的 DRAM 中保持一份 PIM‑就绪的权重副本，同时将下一层的交错权重预取到小的可缓存缓冲区。
  2. 在线权重重排 (OWR) – 在每次 GEMM 之前即时进行交错内存拷贝，消除对永久 PIM‑感知副本的需求。
• 在 Llama 3.2 上实现约 50 % 的 DRAM 容量节省，性能接近理论 PIM 最优值，仅相差几百分点。
• 展示该方案在产品级 PIM‑启用系统上可运行，而非仅在模拟环境中，标志着首个针对该类硬件的实用纯软件修复。

方法论

1. 系统特性化 – 作者对典型的设备端 LLM 推理流水线进行画像，将 prefill（大矩阵乘法）与 decode（重复的 token 生成）分离。他们测量每个阶段如何与具备 PIM 能力的 DRAM 模块的内存层次结构交互。
2. DDB 设计 – 构建了一个双缓冲方案，其中：
   • 活动 层的权重永久驻留在不可缓存区域（PIM 所需）。
   • 下一 层的权重提前流入一个小的可缓存缓冲区，使得 prefill 阶段能够复用这些权重而不违反 PIM 触发条件。
3. OWR 设计 – 与其维护两份拷贝，OWR 在每次 GEMM 调用前执行轻量级的 “swizzle” 拷贝。该拷贝按需将权重重新排列为 PIM 友好的布局，利用 CPU 缓存将开销保持在低水平。
4. 与现有 LLM 运行时的集成 – 两项技术被注入到标准推理栈（分词器 → 嵌入层 → Transformer 层 → 输出）中，对高级代码的改动最小。
5. 评估 – 实验在真实的 PIM‑enabled DRAM 平台上使用 Llama 3.2 模型（7 B 参数）进行。基线包括一种朴素的 PIM 仿真，即将每层权重同时存放在可缓存和不可缓存区域。捕获的指标有：DRAM 使用量、每个 token 的延迟以及整体吞吐量。

结果与发现

含义：DDB 和 OWR 都将内存消耗削减约一半，同时将延迟保持在理想 PIM 性能的几百分点以内。OWR 中即时复制的少量开销被无需维护双缓冲的简易性所抵消。

Practical Implications

• Edge devices can host larger LLMs: 通过释放 DRAM，制造商可以在相同的硅片上容纳参数更多的模型，从而在手机、可穿戴设备或物联网网关上实现更丰富的对话式 AI。
• Zero‑hardware‑change deployment: 由于 PIM‑SHERPA 完全是软件实现，现有支持 PIM 的 DRAM 模块只需通过固件或驱动更新即可升级，缩短上市时间。
• Simplified memory management for developers: 双缓冲和按需 swizzle 被抽象到应用层之外，开发者仍可使用熟悉的框架（如 PyTorch、TensorFlow），仅需一个轻量级运行时 shim。
• Cost‑effective scaling: 降低 DRAM 容量需求直接转化为更低的 BOM 成本和功耗——这对电池供电的设备尤为关键。
• Potential for other PIM workloads: 相同的属性布局对齐思路可以迁移到视觉 Transformer、推荐模型或任何在计算密集和内存密集阶段交替的工作负载。

局限性与未来工作

• 模型规模上限：虽然内存节省显著，但该方法仍需要整个模型（或至少是活动层）能够放入不可缓存区域；在当前边缘 DRAM 容量下，超大模型（>30 B）可能仍难以实现。
• 硬件特定性：这些技术利用了所评估 DRAM 的特定 PIM 触发语义；不同厂商可能提供不同的属性位，需要相应的适配。
• 高度并行工作负载的运行时开销：在存在大量并发推理流的场景中，实时的 swizzle 可能成为瓶颈；未来工作可以探索异步拷贝流水线或硬件辅助的 swizzling。
• 与量化和稀疏性的集成：本文聚焦于 FP16 权重；将该方法扩展到 4‑bit 量化或稀疏表示可能进一步压缩内存占用。

总体而言，PIM‑SHERPA 为在不重新设计硬件的前提下将强大的大语言模型带到边缘提供了一条务实的路径，并为新兴 PIM 生态系统中更广泛的软件驱动优化奠定了基础。

作者

• Sunjung Lee
• Sanghoon Cha
• Hyeonsu Kim
• Seungwoo Seo
• Yuhwan Ro
• Sukhan Lee
• Byeongho Kim
• Yongjun Park
• Kyomin Sohn
• Seungwon Lee
• Jaehoon Yu

论文信息

• arXiv ID: 2603.09216v1
• 分类: cs.DC
• 发表时间: 2026年3月10日
• PDF: 下载 PDF