异构PIM架构热管理挑战与THERMOS解决方案

1. 异构PIM架构的热管理挑战与THERMOS解决方案在AI计算领域处理内存PIM架构正在彻底改变传统计算范式。作为一名长期从事异构计算研究的工程师我见证了PIM技术如何通过减少数据搬运来突破内存墙限制。但当我们把数十个不同特性的PIM芯片粒chiplet集成在2.5D封装中时新的挑战出现了——如何协调这些异构资源如何避免局部过热导致系统降频这正是THERMOS框架要解决的核心问题。当前主流方案如Simba调度器只关注通信成本Big-Little策略则局限于同构芯片粒的尺寸差异。这些方法都忽略了三个关键维度热动态特性不同PIM类型如ReRAM和SRAM的功率密度差异可达3倍目标冲突延迟优化往往需要集中计算而能耗优化倾向于分散负载拓扑影响Mesh/Floret等片上网络(NoI)会显著改变通信能耗分布THERMOS的创新在于将多目标强化学习MORL与物理感知调度相结合。其框架在实测中实现了4.59 DNN/s的吞吐量比Simba提升24%35%的平均延迟降低温度违规减少80%以上2. THERMOS架构设计解析2.1 异构PIM芯片粒的硬件特性我们首先分析实验采用的四种PIM芯片粒配置这些参数直接影响调度策略PIM类型工艺交叉阵列尺寸存储密度每芯片粒内存典型功耗StandardReRAM128x1282bit/cell9568 Kb1.2WShared ADCSRAM768x7681bit/cell9792 Kb2.8WAccumulatorReRAM256x2562bit/cell19200 Kb1.5WADC-lessSRAM128x1281bit/cell2416 Kb0.8W关键观察内存密度Accumulator类型适合存储大型权重矩阵功耗特性Shared ADC虽然计算能力强但容易成为热点混合精度ReRAM芯片支持2bit/cell适合低精度层2.2 分层调度机制THERMOS采用两级决策架构这种设计源自我们在实际部署中的经验教训——扁平化调度在80芯片粒系统中会导致决策延迟激增。第一级集群选择MORL策略class DDT_Policy: def __init__(self): self.tree_depth 5 # 可微分决策树深度 self.feature_dim 64 # 状态特征维度 def forward(self, state, preference): # state包含负载特征、温度分布、资源利用率 # preference∈{[1,0],[0,1],[0.5,0.5]}对应延迟/能耗/平衡 cluster_probs self.decision_tree(state, preference) return cluster_probs决策树每层仅需6μs的推理时间这对实时调度至关重要。我们通过三个并行的训练环境对应不同优化目标来生成综合策略。第二级芯片粒映射邻近驱动算法def proximity_mapping(cluster, prev_layer_chiplets): available_chiplets filter_memory_capacity(cluster) sorted_by_distance sorted( available_chiplets, keylambda c: weighted_distance(c, prev_layer_chiplets) ) return allocate_weights(sorted_by_distance)该算法实测平均耗时49.3μs主要开销来自计算加权距离考虑NoI跳数和链路带宽内存容量验证避免超额分配3. 热感知调度的实现细节3.1 温度建模与约束我们采用MFIT热模型进行实时监测其实时性通过两个技术保证离散状态空间模型将2.5D封装划分为580个热节点活性层2×2精细网格0.1°C精度被动层粗粒度建模100ms采样间隔平衡精度与开销仅增加0.015%时延温度约束通过双重奖励机制实现R_{total} \underbrace{R_{primary}}_{\text{即时奖励}} \gamma \underbrace{R_{secondary}}_{\text{温度惩罚}}其中温度惩罚项R_{secondary} \sum_{t} \max(0, T_i(t) - T_{threshold})^23.2 多目标强化学习训练训练参数配置体现了我们在调参过程中的经验训练环境: 并行实例: 3对应不同偏好向量 每周期样本: 30,000 总训练步数: 25M 优化器: 算法: 改进PPO 学习率: 5e-4 折扣因子: 0.95 裁剪阈值: 0.1 硬件平台: AMD Ryzen Threadripper PRO 7985WX 训练时间: 5.5小时关键训练技巧异步奖励处理允许温度惩罚延迟反馈课程学习从简单负载逐步过渡到混合工作负载目标向量插值增强策略在未见偏好下的泛化能力4. 实测性能与优化建议4.1 不同NoI拓扑下的表现我们在四种主流片上网络拓扑上验证THERMOS的适应性指标MeshFloretHexameshKite最大吞吐提升24%18%22%26%能耗降低8%22%9%2%EDP改进36%27%31%23%特别值得注意的是Floret拓扑在能耗方面的优势——其空间填充曲线特性天然适合数据流式负载。4.2 实际部署建议基于我们的部署经验给出以下实操建议硬件配置至少预留5%的芯片粒作为热备用thermal spare对Shared ADC类型芯片粒加强散热如微流道设计参数调优# 权衡延迟与能耗的偏好设置 def set_preference(workload_type): if workload_type latency_sensitive: return [0.8, 0.2] # 偏向延迟 elif workload_type energy_constrained: return [0.3, 0.7] # 偏向能耗 else: return [0.5, 0.5] # 平衡模式常见问题排查温度读数异常检查MFIT模型与物理布局的一致性验证功率传感器的校准我们曾因1Ω采样电阻偏差导致10°C误判调度延迟突增检查NoI拥塞情况使用内置性能计数器验证决策树缓存命中率应99%5. 扩展应用与未来方向虽然THERMOS当前针对AI负载优化但其框架可扩展至3D封装系统增加垂直维度热耦合考量存算一体芯片支持新型非易失存储器特性边缘设备适配资源受限场景需量化决策树我们在Jetson Xavier NX上的原型验证显示调度开销仅占0.14%执行时间每决策能耗44.73μJ对万级图像批处理可忽略一个有趣的发现是当系统负载70%时温度感知调度带来的收益会指数增长——这启发我们在高密度计算场景应更激进地采用热约束策略。