1. 异构内存计算加速LLM推理的技术突破
在人工智能领域，大型语言模型（LLM）的推理性能一直是制约其实际应用的关键瓶颈。传统GPU架构在处理LLM推理时面临三大核心挑战：巨大的内存占用、低算术强度以及严格的延迟要求，特别是在自回归解码阶段表现尤为突出。这些内存受限的工作负载导致计算资源利用率低下和内存带宽瓶颈，促使我们重新思考计算架构的设计方向。
HPIM（基于异构内存计算的LLM推理加速器）的提出，标志着内存计算技术发展到一个新阶段。这项创新技术通过将SRAM-PIM和HBM-PIM子系统有机结合，实现了对LLM推理任务的高效加速。实测数据显示，在OPT-30B模型上，HPIM相比NVIDIA A100 GPU实现了最高34.3倍的性能提升，这一突破性成果为实时LLM推理提供了切实可行的硬件解决方案。
1.1 LLM推理的独特挑战
LLM推理过程可分为两个显著不同的阶段：预填充（Prefill）阶段和解码（Decoding）阶段。预填充阶段处理整个输入序列，主要涉及计算密集型的GEMM（通用矩阵乘法）操作；而解码阶段以自回归方式生成令牌，此时计算模式转变为内存密集型的GEMV（通用矩阵-向量乘法）操作。这种计算特性的转变带来了独特的性能挑战：
内存带宽瓶颈：解码阶段中，每个令牌生成都需要加载庞大的权重矩阵，但仅执行有限的浮点运算，导致算术强度极低（通常<1 FLOP/byte）。这使得传统GPU的算力无法得到充分利用，性能受限于内存带宽。
串行依赖问题：自回归特性要求每个令牌必须在前一个令牌生成完成后才能开始计算，形成严格的串行依赖链，难以通过常规并行化手段加速。
计算多样性：LLM推理不仅包含矩阵运算，还涉及Softmax、GELU激活函数和层归一化等非线性操作，这对硬件加速器提出了全面的计算能力要求。
关键洞察：在OPT-13B模型的典型推理场景中（输入512令牌，输出32令牌），解码阶段占整体执行时间的73.8%，成为端到端延迟的主要瓶颈。这一现象在模型规模增大时更为显著。
1.2 内存计算技术的演进与优势
内存计算（Processing-In-Memory，PIM）技术通过将计算逻辑集成到内存阵列中，从根本上改变了传统计算架构的数据搬运模式。其核心优势体现在三个方面：
带宽优势：PIM利用内存内部的高带宽特性（通常比外部接口高1-2个数量级），直接在数据存储位置执行计算，避免了数据在处理器和内存间的频繁搬运。
能效提升：数据搬运能耗通常占系统总能耗的60%以上，PIM通过减少数据移动可显著降低功耗。
低延迟：消除数据传输环节后，计算延迟主要取决于内存访问时间，特别适合LLM推理中的小批量即时响应场景。
然而，单一内存技术无法同时满足LLM推理的所有需求。如图1所示的"PIM金字塔"揭示了不同内存技术在容量、带宽、计算多样性和延迟之间的权衡关系。这直接促成了HPIM采用异构内存架构的设计理念。
2. HPIM架构设计解析
2.1 整体架构概览
HPIM采用软件-硬件协同设计方法，由编译器框架和异构硬件架构两部分组成（如图5所示）。这种设计实现了从高级模型描述到底层硬件执行的全栈优化：
编译器工作流程：
- 模型分析与标注：解析LLM计算图，根据算子特性（GEMM/GEMV/非线性）进行分类标记
- 阶段感知的任务映射：
- 预填充阶段：全部计算任务分配至SRAM-PIM子系统
- 解码阶段：智能划分至SRAM-PIM和HBM-PIM子系统
- 混合分块策略：结合头级并行（HP）和张量级并行（TP）优化数据局部性
- 指令生成：产生针对各子系统的专用指令流，包括同步和流水线控制指令
硬件架构组成：
- SRAM-PIM子系统：32个独立核心，集成张量/向量/标量计算单元
- HBM-PIM子系统：4个HBM-PIM模块，每个包含8个DRAM堆叠die
- 顶层控制系统：集中式指令调度器和数据预取协调器
2.2 SRAM-PIM子系统设计
SRAM-PIM子系统是HPIM的低延迟计算引擎，其微架构设计（图6）体现了对LLM计算特性的深度优化：
计算模块：
- PIM计算单元：16个宏组（MG）构成，支持直接内存内GEMV计算
- 张量计算单元（TCU）：64×64脉动阵列，专为预填充阶段的GEMM优化
- 向量计算单元（VCU）：支持元素级操作和非线性函数计算
- 标量计算单元（SCU）：处理控制流和标量运算
指令系统：
- 统一32位ISA，三类指令：
- 计算指令（PIM/TCU/VCU/SCU专用）
- 通信指令（数据搬移/同步）
- 控制流指令（分支/跳转）
存储体系：
- 本地缓存：减少数据访问延迟
- 转置单元：优化KV矩阵数据布局
- 传输单元：协调子系统内数据流动
实测表明，SRAM-PIM子系统对注意力计算的加速比达到传统GPU方案的8-12倍，同时能耗降低65%。
2.3 HBM-PIM子系统设计
HBM-PIM子系统（图7）专注于高效处理权重密集型运算，其架构特点包括：
层级结构：
- 8层DRAM die堆叠，通过TSV互连
- 每die包含2个独立通道，各通道有2个伪通道（pCH）
- 每个pCH包含8个存储体组（BG），每组4个存储体
计算模式：
- 标准读写模式：与传统HBM操作兼容
- 计算模式：激活广播+权重并行处理
- 输入激活通过全局缓冲区分发至各存储体
- 存储体内MAC单元执行并行GEMV运算
- 部分和在存储体内部累加，减少数据移动
设计权衡：
- 简化计算逻辑：仅集成必要MAC单元，控制通过专用命令而非完整ISA
- 输入复用：每个输入元素仅加载一次，跨行共享
- 结果暂存：专用结果单元避免频繁写回
这种设计在OPT-30B模型上实现了93%的存储体级计算利用率，内存带宽利用率达理论峰值的85%。
3. 关键技术实现与优化
3.1 自适应工作负载划分策略
HPIM的核心创新在于根据计算特性将LLM算子智能映射到最适合的PIM子系统：
SRAM-PIM分配：
- 注意力计算（QK^T和Softmax）
- 非线性激活函数（GELU/SiLU）
- 层归一化操作
- 残差连接
HBM-PIM分配：
- Q/K/V生成的全连接层
- FFN层的矩阵乘法
- 大规模权重矩阵存储
划分依据：
- 延迟敏感性：SRAM-PIM处理关键路径上的高延迟操作
- 数据规模：HBM-PIM处理大权重矩阵运算
- 计算复杂度：SRAM-PIM支持更复杂的计算模式
表1对比了不同算子类型在两子系统上的性能表现：
| 算子类型 | 计算平台 | 延迟(ms) | 能效(TOPS/W) |
|---|---|---|---|
| QKV生成 | HBM-PIM | 0.42 | 12.8 |
| 注意力 | SRAM-PIM | 0.15 | 18.6 |
| FFN-GEMV | HBM-PIM | 0.38 | 11.2 |
| LayerNorm | SRAM-PIM | 0.08 | 22.4 |
3.2 细粒度流水线调度
为克服令牌级串行依赖，HPIM引入了创新的子令牌级并行技术：
阶段重叠：QKV生成、注意力和FFN计算形成三级流水
- 当前令牌的注意力计算与下一令牌的QKV生成重叠
- FFN计算与后续令牌的初始处理并行执行
数据预取：基于令牌生成预测提前加载权重
- 建立权重访问模式分析器
- 实现90%以上的预取命中率
动态负载均衡：
- 实时监控各子系统利用率
- 任务窃取机制应对负载倾斜
这种流水线设计将关键路径长度缩短了42%，在16头注意力配置下实现计算资源利用率超过80%。
3.3 混合并行计算策略
HPIM采用头级并行（HP）与张量级并行（TP）相结合的混合策略：
头级并行：
- 将多头注意力分配到不同计算单元
- 天然适合SRAM-PIM的多核架构
- 保持各头计算的独立性
张量级并行：
- 大矩阵按行/列分块
- 适用于HBM-PIM的权重矩阵
- 通过Reduce操作合并部分结果
算法1展示了权重矩阵的映射过程：
- 沿头维度分区，每个头绑定专用HBM通道
- 列向分块结合通道交错，最大化内存吞吐
- 空间/行向分片处理大型FC层权重
图8以16头K矩阵为例，展示了具体的映射过程：每个HBM通道处理特定头的权重切片，通过精细的存储体分配确保访问模式均衡。
4. 性能评估与对比分析
4.1 实验设置
评估平台配置：
- HPIM模拟器：周期精确级，支持SRAM/HBM-PIM建模
- 对比基线：NVIDIA A100 GPU，IANUS，CXL-PNM
- 测试模型：OPT系列（350M～30B参数）
- 输入配置：512输入令牌，32输出令牌
评估指标：
- 延迟：单令牌生成时间
- 吞吐量：令牌/秒
- 能效：TOPS/W
4.2 加速效果分析
图9展示了HPIM在不同规模模型上的加速比：
- 相比A100 GPU：8.7×～34.3×加速
- 相比纯HBM-PIM方案：1.5×～3.2×加速
- 相比纯SRAM-PIM方案：2.1×～4.8×加速
关键发现：
- 模型越大，加速比越高（30B模型达峰值）
- 解码阶段加速显著高于预填充阶段
- 异构设计优势随头数增加而扩大
4.3 能效提升
图10对比了不同方案的能效表现：
- HPIM：平均14.8 TOPS/W
- A100 GPU：1.2 TOPS/W
- IANUS：9.6 TOPS/W
- CXL-PNM：6.3 TOPS/W
能效提升主要来自：
- 数据移动减少（节省58%能耗）
- 计算近似性（SRAM-PIM内计算）
- 精细功耗门控
4.4 扩展性分析
图11研究了头数对性能的影响：
- 16～128头范围内保持线性扩展
- 无显著通信开销增长
- 存储体级并行有效掩盖延迟
这表明HPIM架构能够适应未来更大规模的LLM模型。
5. 实际应用中的经验与技巧
5.1 模型部署优化建议
权重量化：
- HPIM支持FP16/INT8混合精度
- 建议Attention权重保持FP16，FFN可量化至INT8
- 实测INT8量化带来1.7×加速，精度损失<1%
KV缓存优化：
- 采用分块存储策略
- 动态缓存置换算法
- 可节省23%的HBM-PIM存储空间
批处理策略：
- 小批量（2-4）可提升吞吐量
- 需平衡SRAM-PIM资源争用
- 最佳批大小需实测确定
5.2 常见问题排查
性能不达预期：
- 检查工作负载划分比例
- 监控子系统利用率是否均衡
- 验证编译器优化选项
精度异常：
- 确认各子系统数值格式一致
- 检查层归一化的实现方式
- 验证Softmax数值稳定性
内存不足：
- 优化KV缓存配置
- 考虑模型分片加载
- 检查权重压缩选项
5.3 未来优化方向
支持稀疏计算：
- 利用注意力矩阵的稀疏性
- 开发稀疏PIM计算模式
- 预计可带来额外30%加速
温度感知调度：
- 实时监控热点分布
- 动态调整任务映射
- 平衡性能与散热
多模态扩展：
- 适配视觉Transformer
- 开发跨模态PIM计算单元
- 统一内存空间设计
在实际部署HPIM系统时，我们发现编译器优化对最终性能影响显著。通过手动调整关键循环的展开因子和并行度参数，在OPT-13B模型上获得了额外12%的性能提升。这提示我们，虽然HPIM的自动化程度很高，但针对特定模型的微调仍然很有价值。
