当前位置：首页 > news >正文

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

news 2026/5/29 1:52:54

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类型	工艺	交叉阵列尺寸	存储密度	每芯片粒内存	典型功耗
Standard	ReRAM	128x128	2bit/cell	9568 Kb	1.2W
Shared ADC	SRAM	768x768	1bit/cell	9792 Kb	2.8W
Accumulator	ReRAM	256x256	2bit/cell	19200 Kb	1.5W
ADC-less	SRAM	128x128	1bit/cell	2416 Kb	0.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, key=lambda 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})^2

3.2 多目标强化学习训练

训练参数配置体现了我们在调参过程中的经验：

训练环境: 并行实例: 3（对应不同偏好向量） 每周期样本: 30,000 总训练步数: 25M 优化器: 算法: 改进PPO 学习率: 5e-4 折扣因子: 0.95 裁剪阈值: 0.1 硬件平台: AMD Ryzen Threadripper PRO 7985WX 训练时间: 5.5小时

关键训练技巧：

异步奖励处理：允许温度惩罚延迟反馈
课程学习：从简单负载逐步过渡到混合工作负载
目标向量插值：增强策略在未见偏好下的泛化能力

4. 实测性能与优化建议

4.1 不同NoI拓扑下的表现

我们在四种主流片上网络拓扑上验证THERMOS的适应性：

指标	Mesh	Floret	Hexamesh	Kite
最大吞吐提升	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] # 平衡模式

常见问题排查：