PIM架构如何突破LLM推理的能效瓶颈
1. PIM架构与LLM推理的效能瓶颈突破
在大型语言模型（LLM）推理任务中，传统计算架构面临的核心矛盾在于：模型参数规模呈指数级增长与内存带宽增长缓慢之间的不匹配。以GPT-3 175B模型为例，仅参数加载就需要350GB内存空间，而即使是高端GPU如NVIDIA A100，其HBM2e内存带宽也仅2TB/s。这种差距导致系统性能严重受限于数据搬运而非实际计算。
处理内存计算（Processing-In-Memory, PIM）技术通过将计算单元嵌入内存bank内部，从根本上改变了数据流动范式。如图7(a)所示，在传统PIM执行全连接层（FC）核时，DRAM访问能耗占比高达86.7%，而实际计算能耗仅占13.3%。这种能量分布揭示了传统架构的效率瓶颈——数据搬运消耗了绝大部分能量而非用于有效计算。
关键发现：当采用DRAM数据重用技术（data reuse level=64）时，单次DRAM访问的数据可复用64次计算，使DRAM访问能耗占比降至33.1%（图7b）。这意味着通过智能数据复用，理论上可获得约2.6倍的能效提升空间。
2. DRAM数据重用机制深度解析
2.1 数据重用的硬件实现基础
在现代DRAM架构中，每次行激活（row activation）会传输约8KB数据到行缓冲区。对于LLM中的矩阵运算，特别是FC层的权重矩阵计算，同一批权重参数往往需要与多个输入向量相乘。传统架构需要反复从DRAM读取相同权重，而PIM架构通过在bank内部保留激活的行数据，实现跨多次计算的参数复用。
具体实现依赖三个关键技术：
- Bank级计算单元：每个DRAM bank集成浮点运算单元（FPU），支持本地数据直接计算
- 行缓冲区扩展：延长行缓冲区保持时间，避免频繁预充电
- 计算调度优化：将计算密度高的操作集中调度到已激活的bank
2.2 数据重用与并行计算的协同设计
图7(c)揭示了数据重用级别与并行配置的权衡关系。实验测量了三种配置：
- 1P1B：每bank 1个FPU
- 2P1B：每bank 2个FPU
- 4P1B：每bank 4个FPU
当数据重用级别≥4时，4P1B配置的功耗降至116W以下（HBM3功率预算），此时相比无重用情况可获得：
- 计算吞吐提升：4倍于基础配置
- 能效比提升：2.8倍（单位能量完成的计算量）
3. 异构PIM架构设计
3.1 FC-PIM与Attn-PIM的差异化配置
针对LLM中FC层和注意力层的不同特性，我们采用异构PIM设计：
| 特性 | FC-PIM配置 | Attn-PIM配置 |
|---|---|---|
| 计算强度 | 高（4P1B） | 低（1P2B） |
| 数据重用率 | ≥64 | ≤4 |
| 带宽需求 | 中等（权重稳定） | 高（KV缓存动态） |
| 面积分配 | 40% bank用于计算 | 15% bank用于计算 |
FC-PIM设计要点：
- 采用4P1B配置最大化计算并行度
- 通过权重矩阵分块（blocking）实现bank间负载均衡
- 支持混合精度计算（FP16累加+FP32输出）
Attn-PIM设计创新：
- 采用1P2B配置优化内存容量
- 动态KV缓存管理机制
- 支持稀疏注意力计算（可跳过无效bank计算）
3.2 面积与功耗的协同优化
在22nm工艺下，关键面积参数为：
- 单个HBM bank面积：0.83mm²
- 单个FPU面积：0.1025mm²
- 最大HBM裸片面积：121mm²
通过公式推导，4P1B配置的最大bank数量为：
m(0.1025×4 + 0.83) ≤ 121 → m ≤ 97实际采用96 banks（3个bank groups）的设计，在满足面积约束的同时提供：
- 理论算力：1.2TFLOPS @666MHz
- 存储容量：12GB（FC-PIM）/16GB（Attn-PIM）
4. 系统集成与数据调度
4.1 层次化互连架构
如图5(a)所示，系统采用三级互连：
- NVLink级：连接FC-PIM与主机处理器（带宽900GB/s）
- PCIe/CXL级：连接Attn-PIM设备（支持≤4096设备扩展）
- NoC级：bank间网状网络（延迟<10ns）
4.2 数据分区策略
注意力层：
- 按attention head分区
- K^T矩阵：伪通道级列分割 + bank级行分割
- V矩阵：伪通道级行分割 + bank级列分割
FC层：
- 权重矩阵分块（block size=256×256）
- 输入向量广播（通过bank间快速通道）
5. 实测性能与优化效果
5.1 端到端性能对比
在LLaMA-65B模型上的测试结果显示（batch size=64, spec length=4）：
| 系统 | 相对性能 | 能效比 |
|---|---|---|
| A100+AttAcc | 1.0× | 1.0× |
| PAPI（本设计） | 1.8× | 3.4× |
| AttAcc-only | 0.16× | 0.9× |
关键优势体现在：
- 动态负载均衡：FC层任务在GPU与PIM间智能分配
- 流水线优化：计算与通信重叠率达73%
5.2 敏感度分析
图10展示不同并行配置下的性能变化：
- 批量大小（RLP）从4增加到128时，PAPI保持1.5-1.8×优势
- 推测长度（TLP）从1增加到8时，AttAcc-only性能下降40%而PAPI仅降12%
6. 工程实现中的关键挑战
6.1 热管理策略
在4P1B配置下，实测bank局部热点可达85°C。我们采用：
- 动态频率调节（666MHz→500MHz@75°C）
- 计算负载轮转（bank级负载均衡）
- 硅中介层热扩散设计
6.2 可靠性保障
针对DRAM计算引入的可靠性问题：
- 错误检测：每128次计算插入1次校验操作
- 银行隔离：故障bank可软件屏蔽
- 误差补偿：输出层添加轻量级校准网络
7. 前沿扩展：支持MoE架构
对于混合专家模型（Mixture of Experts），PAPI架构展现出独特优势：
- 稀疏计算：只有激活的expert权重需要加载
- 动态分配：专家权重可分布在不同bank
- 能效优化：非活跃FPU可时钟门控
实测在Switch Transformer模型上，PAPI可实现：
- 专家激活延迟：<50ns
- 稀疏计算能效：28TOPS/W
这种设计使得PAPI不仅能适应当前密集LLM，也为未来稀疏化模型提供了硬件基础。通过将计算范式从"数据搬运"转变为"计算移动"，PIM架构正在重塑LLM推理的硬件格局。