AILFM：基于主动模仿学习的3D堆叠CPU大模型推理热管理调度框架

主动模仿学习3D S-NUCA大语言模型推理

于 2026-05-28 03:18:32 修改

·本内容遵循CC 4.0 BY-SA版权协议

1. 项目概述与核心挑战

在数据中心和边缘计算场景中，大语言模型（LFM）的推理任务正从GPU向通用CPU迁移，这背后是成本、可用性和灵活性的综合考量。然而，当我们将这些计算和内存都极度密集的负载，部署到新兴的3D堆叠静态非均匀缓存架构（3D S-NUCA）多核处理器上时，一个复杂且棘手的问题便浮出水面：如何在保证芯片“不发烧”的前提下，榨取出最高的推理性能？

传统的热管理方案，比如动态电压频率调节（DVFS）或者简单的线程迁移，在面对3D S-NUCA时往往力不从心。原因在于，3D堆叠在带来高带宽和低延迟优势的同时，也引入了严重的“热堆积”问题——热量难以从堆叠的硅层中快速散出。更微妙的是，由于缓存物理上分布在不同的层和位置，不同核心访问最后一级缓存（LLC）的延迟存在显著差异，即“缓存异构性”。一个在角落的核心可能“凉快”但访问数据“慢”，而一个在中心的核心可能“快”但更容易“过热”。大模型推理本身又是由多个计算模式迥异的“核”（Kernel）组成的，例如注意力计算（Attention）对缓存延迟极其敏感，而前馈网络（FFN）则更偏向计算密集型。这就形成了一个三维的优化难题：空间维度上的热与缓存延迟分布不均、时间维度上的热量累积与消散、以及负载维度上不同计算核的异构行为。

现有的方法大多基于简化的分析模型，试图用一个“万能公式”去套所有场景，结果往往是为了控制温度而过度降频，牺牲了性能，或者因无法精准感知不同计算核的特性而导致调度失策。我们需要的，是一个能像经验丰富的系统管理员一样，实时“感知”芯片温度、缓存状态和当前正在运行何种计算任务，并能做出“近乎最优”调度决策的智能体。这正是AILFM框架要解决的问题：它不是一个硬编码的规则引擎，而是一个通过主动模仿学习训练出来的“调度大脑”，能够学会在复杂的3D S-NUCA环境中，为不同的大模型计算核找到性能和热安全的最佳平衡点。

2. 核心思路：为什么是主动模仿学习？

要理解AILFM的巧妙之处，我们得先拆解“模仿学习”和“主动学习”这两个概念在系统调度中的实际意义。

模仿学习 的核心思想是“站在巨人的肩膀上”。在这个场景中，“巨人”就是一个理论上完美的调度器，我们称之为 Oracle（先知）。这个Oracle知晓系统的一切内部状态（温度、缓存命中率、核心频率等），并能瞬间计算出所有可能的调度动作（将线程从核心A迁移到核心B，或将某个核心的电压频率调低一档）所带来的未来收益。显然，在真实系统中构建这样一个全知全能的Oracle是不可能的，其计算开销也无法承受。但是，我们可以离线地，利用高性能模拟器（如CoMeT）和精确的功耗、热模型，为大量不同的系统状态和工作负载组合，预先计算出这些“最优”或“接近最优”的调度决策，生成一个庞大的“专家示范”数据集。

然而，直接让一个神经网络去死记硬背这个数据集有两个问题：第一，数据收集成本高，要覆盖所有可能的系统状态几乎不可能；第二，当遇到训练数据中未见过的新状态时，模型的决策可能非常不可靠。这就引出了 主动学习 的用武之地。

主动学习 的精髓是“知之为知之，不知为不知”。我们赋予学习到的调度策略（称为 学习策略）一种评估自身决策“信心”的能力。当它对当前状态很有把握时，就自主做出决策，享受极低的运行时开销；当它感到“不确定”或“没遇到过这种情况”时，就主动去“请教”Oracle（尽管在线上是查询一个轻量级的替代模型或规则），并将这次请教的结果作为新的学习样本。这个过程就像一个新司机：在熟悉的道路上自己开，遇到复杂路口或陌生路段就打开导航（Oracle），同时记住这次导航的路线，下次再遇到类似情况就可能自己处理了。

AILFM将两者结合，形成了 主动模仿学习 的闭环：

离线准备：利用Oracle生成高质量的初始示范数据。
在线学习与决策：一个轻量级神经网络作为学习策略，在运行时根据系统状态（特征）做出迁移或DVFS决策。
不确定性评估：通过蒙特卡洛Dropout等技术，实时估计当前决策的置信度。
选择性请教：置信度低时，触发查询（模仿Oracle的决策）；置信度高时，自主行动。
持续精进：将请教的决策作为新的训练数据，持续微调学习策略，使其越来越“像”Oracle，同时越来越“独立”。

这种方法的优势显而易见：它最终部署的是一个开销极小的神经网络，却能获得接近Oracle的调度质量，并且具备处理未见情况的能力，完美适配大模型推理中动态、异构的计算特征。

3. 系统建模与问题定义

在深入AILFM的实现细节前，我们必须形式化地定义我们要优化的目标，以及我们所处的战场——3D S-NUCA系统——的特性。

3.1 目标系统：3D S-NUCA多核处理器

想象一个由多层硅片垂直堆叠而成的“计算立方体”。每一层都布满了处理核心（Core）或内存块（Memory Bank）。这些核心通过一个三维的片上网络（3D NoC）互联，共享一个逻辑上统一、但物理上分布在整个立方体中的最后一级缓存（LLC）。这就是3D S-NUCA。

缓存非均匀性：由于数据块被静态映射到特定的缓存块（Bank），一个核心访问不同数据块的延迟，取决于该数据块所在的物理位置与核心之间的“距离”。这个距离通常用 平均曼哈顿距离 来衡量，即一个核心到所有其他核心（或缓存块）的跳数平均值。如图2所示，位于立方体几何中心的核心AMD值低（访问“平均距离”近，延迟低），但被其他核心包围，散热困难；位于角落的核心AMD值高（访问“平均距离”远，延迟高），但散热条件好。这种性能与热特性的空间耦合，是调度的核心挑战。
热挑战：3D堆叠使得单位体积内的功耗密度激增，但散热表面积并未同比增加。热量容易在堆叠层内部积聚，形成局部热点。一旦温度超过安全阈值（如85°C），轻则触发降频（DVFS）导致性能下降，重则引发硬件错误或永久性损伤。

3.2 大模型推理的工作负载分解

一个大模型（如ViT、LLaMA）的推理过程并非一个 monolithic 的巨块，而是可以分解为一系列具有不同计算和访存特征的 计算核。

嵌入层：将输入令牌转换为向量，通常涉及查表操作，对缓存容量敏感。
自注意力层：计算Query, Key, Value矩阵及其交互，需要频繁访问大量的中间状态（K/V缓存），对缓存带宽和延迟极度敏感。
前馈网络层：通常是几个全连接层，计算密集，访存模式相对规整，对缓存延迟相对不敏感。
语言模型头：将隐藏状态映射回词汇表，计算量大，也有一定的缓存访问需求。

如表1所示，我们对ViT-base的各个核在不同AMD值的核心上进行了性能剖析。结果清晰表明：注意力核的性能对AMD变化最为敏感，IPS（每秒指令数）下降超过35%；而FFN核最不敏感，IPS下降不到6.5%。这直接启示我们：调度必须“核感知”——把对缓存敏感的核（如Attention）尽量放在AMD值低（“快”）的核心上，哪怕那里热一些；而把计算密集的核（如FFN）可以调度到AMD值高（“慢”）但更凉爽的核心上。

3.3 优化问题形式化

我们的目标是在保证芯片最高温度不超过安全阈值的前提下，最小化整个大模型推理任务的执行时间。可用的控制“旋钮”有两个：

线程迁移：将一个正在执行的线程（及其关联的缓存数据状态）从一个核心移动到另一个核心。
动态电压频率调节：直接降低某个核心的电压和频率，减少其功耗和发热，但也会降低其计算性能。

这两个旋钮都有代价：迁移会导致目标核心的缓存“冷启动”，带来性能惩罚；DVFS则直接损失了计算吞吐量。因此，调度策略必须在“迁移开销”、“降频损失”和“避免热紧急事件”之间做精细的权衡。

我们将此形式化为一个约束优化问题：寻找一个调度策略π，使得总推理时间M(π)加上迁移或DVFS带来的开销O的最小值最小化，同时约束峰值温度T_peak低于阈值T_th。AILFM就是要学习出这个接近最优的策略π。

4. AILFM框架深度解析

AILFM的整体架构如图3所示，它是一个在线学习与决策系统，由“Oracle策略”和“学习策略”两部分协同工作。

4.1 特征工程：调度决策的“眼睛”

要让学习模型做出好的决策，首先得给它看对的“信息”。我们筛选并定义了五个核心的系统状态特征，作为模型的输入：

指令每秒：当前核心上运行线程的实时性能指标，直接反映计算吞吐量。
LLC每千指令失效数：衡量缓存局部性的关键指标。高MPKI意味着频繁访问LLC未命中，需要去更远的内存取数据，对延迟更敏感。
平均曼哈顿距离：该核心的“地理位置”属性，决定了其访问缓存的平均延迟和散热潜力。
功率预算：由底层的3D-TTP等热管理框架分配给的、在当前热约束下该核心允许消耗的最大动态功率。这直接关联到可用的频率档位。
温度状态：当前系统的峰值温度，以及相对于安全阈值的余量。

这五个特征构成了一个多维状态向量，能够较为全面地刻画“谁在哪儿、干得怎么样、有多热、还能干多猛”。

4.2 Oracle策略：基于混合高斯过程回归的专家示范

Oracle是智慧的源泉。我们需要一个能够在离线阶段，为任意给定的系统状态和计算核类型，快速评估“如果将这个线程从核心A迁移到核心B，能带来多少性能收益（或损失）”的模型。我们为此设计了一个 混合高斯过程回归 模型。

为什么是高斯过程回归？ 因为GPR是一种贝叶斯非参数模型，它不仅能给出预测值（迁移效用），还能给出预测的不确定性（方差）。这对于我们后续的主动学习机制至关重要。GPR特别适合处理我们这种规模不大但需要精确插值和不确定性估计的数据集。

为什么是“混合”？ 因为不同的大模型计算核（如Attention和FFN）对迁移的响应模式截然不同。用一个统一的GPR模型去拟合所有核的数据，效果会很差。因此，我们为每一种计算核类型训练一个独立的GPR专家模型。在运行时，通过一个轻量级的、基于IPS和MPKI加权和的核类型分类器，判断当前活跃的核是哪种类型，然后激活对应的那个GPR专家进行效用预测。

每个GPR专家的输入是源核心状态向量和目标核心状态向量。输出是一个标量，代表预测的迁移效用——正的效用意味着迁移可能带来性能提升，负的效用则意味着可能得不偿失。Oracle的策略就是：遍历所有可能的目标核心，用对应的GPR专家预测迁移效用，并检查迁移后是否会导致温度超标，最后选择那个效用最高且 thermally-safe 的目标进行迁移。

4.3 学习策略：带有不确定性估计的轻量级神经网络

线上运行的主角是学习策略，它是一个轻量级的全连接神经网络。通过神经架构搜索，我们确定了一个三隐藏层的结构，在预测精度和计算开销间取得了平衡。

蒙特卡洛Dropout：给模型装上“信心仪表盘” 这是实现主动学习的关键技术。通常，Dropout只在训练时使用，用于防止过拟合。在MC Dropout中，我们在模型推理时也保持Dropout开启。对于同一个输入状态，我们让网络进行N次（例如50次）前向传播，由于Dropout的随机性，每次会得到略微不同的输出。这N个输出的均值就是最终的预测值（例如，迁移到某个核心的效用），而这N个输出的方差，就成为了模型对此次预测的不确定性的度量。方差越大，说明模型越“犹豫不决”；方差越小，说明模型越“自信”。

选择性查询机制 我们设定一个不确定度阈值τ。在每一个调度时刻（例如每500纳秒），学习策略观察当前系统状态，并对其最优决策进行预测，同时计算该预测的不确定度u(s)。

如果 u(s) ≤ τ，说明模型对当前情况很熟悉、很自信，那么就直接采用学习策略自己的决策，享受极低的延迟开销。
如果 u(s) > τ，说明模型遇到了陌生或复杂情况，信心不足。此时，系统会触发一次查询，转而采用Oracle策略（实际上线上是一个简化的、快速的替代规则或一个小型模型）的决策来执行。同时，这次“状态-最优决策”对会被记录下来，用于后续更新学习策略。

这个机制完美平衡了效率与效果：大部分时间由快速神经网络决策，只在必要时才付出较高开销去寻求专家指导。

4.4 训练与更新：让模型越来越聪明

学习策略的训练数据来自两部分：

专家示范数据：离线阶段由Oracle生成的高质量（状态，动作）对。
在线查询数据：运行时，当不确定度高而触发Oracle查询时，记录下的新数据。

损失函数由两部分组成：

监督损失：针对专家示范数据和在线查询数据，让模型的预测尽量接近Oracle给出的“标准答案”。这部分确保模型能模仿专家。
自模仿损失：针对模型自己做出且被证明是成功的决策（例如，自主决策后系统性能稳定且温度安全），给予奖励，强化这类决策。这部分鼓励模型在熟悉领域更加自信和高效。

通过这种方式，学习策略不仅能从Oracle那里学，还能从自己的成功经验中学，不断进化，减少对Oracle的依赖，最终实现以极低开销获得接近Oracle的性能。

5. 实验评估与结果分析

我们在最先进的3D系统模拟器CoMeT上构建了一个64核（4x4x4拓扑）的3D S-NUCA系统仿真环境，并使用了ViT、BERT、LLaMA、Gemma、DeepSeek等多种大模型工作负载进行验证。对比的基线包括：

3QTM：基于强化学习的核心-内存协同热管理方法。
NeuroTAP：基于分析模型的3D DRAM热感知数据映射方法。
3D-DNaPE：将任务迁移到最冷相邻空闲核心的调度技术。
DLFM：一个没有Oracle指导、直接进行端到端学习的基线模型。

5.1 性能表现：全面领先

如图5和图6所示，在不同的输入序列长度和不同的热安全阈值下，AILFM在所有测试模型上均取得了最佳的性能。随着输入长度增加（计算和热压力增大），AILFM相对于基线方法的优势更加明显。例如，在DeepSeek-2.4B模型上，当序列长度达到1024时，AILFM比第二名3D-DNaPE的性能高出约35%。这证明了AILFM框架在处理大规模、高强度推理负载时的强大鲁棒性。

5.2 热安全保证：温度控制更精准

如图7所示，在将热阈值设定为75°C和85°C的两种严格条件下，AILFM成功地将系统的峰值温度维持在了最低水平，平均分别约为72.5°C和73.2°C，显著优于其他基线方法。NeuroTAP等基于分析模型的方法，由于无法适应大模型核级行为的动态变化，温度控制波动较大，时常接近甚至偶尔超过阈值。AILFM通过核感知的精准调度，在性能与温度之间找到了更优的平衡点。

5.3 开销与可扩展性：实用性的关键

运行时开销：如图8所示，即使在全系统64个核心满载的极端情况下，AILFM的调度器运行时开销也低于5%。这主要归功于其主体是一个轻量级神经网络，且大部分决策由自信的模型自主做出，避免了频繁调用复杂Oracle。
可扩展性：如表2所示，我们从2D平面结构扩展到3D堆叠结构，再扩展到更大的216核（6x6x6）3D拓扑，AILFM均能有效工作。与2D平面基线相比，3D S-NUCA架构带来了最高1.64倍的性能加速，而AILFM成功管理了随之增加的热密度，证明了其良好的可扩展性。

5.4 消融实验：验证核心设计

我们通过一系列消融实验，验证了AILFM各个组件的必要性：

Oracle指导的价值：如表3所示，纯Oracle（MoGPR）虽然准确率略高，但其单次决策开销巨大，导致高达24.6%的运行时开销，不具备实用性。AILFM以仅3.2%的开销，达到了与Oracle相近的调度精度，实现了效率与效果的完美结合。
核感知建模的重要性：如表6所示，如果移除“核感知”设计（即用一个统一的模型处理所有核），调度准确率从96.7%骤降至90.8%，峰值温度也从75.1°C上升至79.9°C。这说明忽视工作负载内部异构性是行不通的。
主动查询机制的有效性：如表5所示，不确定性阈值τ的选择是一个权衡。τ太小（如0.05）会导致查询过于频繁，开销增加；τ太大（如0.30）则查询太少，模型在陌生领域容易犯错，导致温度违规率上升。我们通过验证集选择了τ=0.15作为最佳平衡点，在维持低违规率的同时保持了低开销。

6. 实现启示与避坑指南

将AILFM从论文思想落地到实际系统，会面临一系列工程挑战。以下是一些关键的实操心得和注意事项：

6.1 特征采集与核类型识别

性能计数器是关键：IPS和LLC MPKI等特征需要从处理器的硬件性能计数器中实时读取。确保你有权限且了解如何正确配置和读取这些计数器。不同架构（x86, ARM）的计数器名称和访问方式可能不同。
核类型识别的轻量化：线上核类型分类器必须极其轻量。我们采用IPS和MPKI的加权和作为判断依据，权重λ1和λ2是通过离线数据分析学习得到的。也可以考虑使用更简单的、基于程序计数器或特定指令序列模式的轻量级指纹识别。切忌使用复杂的模型来做在线识别，那会成为新的性能瓶颈。
采样频率的权衡：调度器多久做一次决策？太频繁（如每100纳秒）会引入过多开销；太稀疏（如每毫秒）则可能错过最佳调度时机。我们的实验表明，对于大模型推理这种相对稳定的负载，每500纳秒到一个微秒进行一次调度评估是合理的。

6.2 Oracle模型的构建与简化

离线模拟数据的质量：Oracle的GPR模型严重依赖于离线模拟数据的质量和覆盖面。需要精心设计实验矩阵，覆盖不同的AMD值、功率预算、初始温度以及各种大模型核的组合。数据噪声要小，否则GPR会学习到错误的规律。
线上Oracle的替代：论文中线上查询的“Oracle”不可能是完整的MoGPR模型。在实践中，我们需要为线上部署准备一个简化版本。例如，可以将MoGPR预测的结果预先制成一个多维查找表，或者训练一个极度轻量级的小型神经网络来近似Oracle的行为。核心原则是：查询开销必须远低于一次错误的调度决策带来的损失。

6.3 学习策略的部署与更新

神经网络的小型化：尽管NAS找到了一个三层的网络，但在部署前仍需进行剪枝、量化等操作，进一步压缩模型大小，以适应嵌入式控制器的有限资源。
在线更新的策略：是否允许学习策略在部署后继续在线更新？这是一个需要谨慎权衡的问题。在线更新可以让模型适应硬件老化、工作负载漂移等变化，但也带来了稳定性和安全性的风险（例如，模型被“教坏”）。一个折中的方案是定期（例如每天）在后台用收集到的新数据重新训练模型，经过严格验证后再进行热替换。
不确定性校准：MC Dropout给出的不确定性估计是否准确可靠？需要在实际系统上进行校准。可以通过观察在模型高置信度下决策的失败率，来调整阈值τ或改进不确定性估计方法。

6.4 与底层硬件管理单元的协同

AILFM是一个位于操作系统或运行时层面的调度框架，它需要与底层的硬件管理单元紧密协同：

与DVFS控制器交互：AILFM做出降频决策后，需要通过ACPI或特定MSR寄存器向处理器发出指令。需要确保指令的延迟和粒度是可接受的。
与线程迁移机制交互：迁移线程涉及操作系统调度器的深度介入。在Linux中，可能需要修改内核调度器或利用cgroups、sched_setaffinity等接口，并处理好缓存亲和性（cache affinity）的破坏与重建带来的开销评估。
与温度传感器集成：需要能够以低延迟、高精度地读取来自不同核心和区域的温度传感器数据。这些数据是热约束判断的直接依据。

7. 总结与展望

AILFM框架为3D堆叠异构多核系统上的大模型推理热管理问题，提供了一条新颖且高效的解决路径。它通过主动模仿学习，将复杂的、基于物理模型的Oracle知识，蒸馏到了一个轻量级、低开销的神经网络调度器中，实现了核感知的精准调度。实验证明，该方法在保证热安全的前提下，能显著提升推理性能，且具备良好的可扩展性和实用性。

这项工作也为我们打开了新的思路。未来，我们可以探索将AILFM的思想扩展到更广泛的场景：

多任务混合负载：当前工作主要针对单个大模型推理。在云数据中心，多个不同模型或任务的实例可能混合部署。如何在这种更复杂的干扰环境下进行协同调度和热管理，是一个更大的挑战。
异构计算单元：未来的处理器可能集成CPU、GPU、NPU等多种计算单元。AILFM的框架可以扩展为不仅调度线程在核心间的迁移，还能决策将某个计算核卸载到更合适的加速器上执行，形成跨异构单元的全局热-性能优化。
学习框架的进化：能否用更先进的元学习、课程学习方法来加速学习策略的训练？能否让模型学会自动发现和定义新的、更有效的系统状态特征？

从实验室的模拟器到真实数据中心的机架，还有大量的工程化工作要做。但AILFM无疑为我们点亮了一盏灯，指明了一条通过机器学习让复杂芯片系统自我优化、智能调度的可行之路。在实际部署中，我个人的体会是，最大的挑战往往不是算法本身，而是如何将算法与现有复杂的软硬件栈无缝、稳定地集成起来，并建立起可靠的监控和回滚机制。这需要架构师、系统软件工程师和机器学习工程师的紧密协作。每一次将这样的智能调度系统成功推向生产环境，看着它平稳运行并带来实实在在的能效提升，都是对这项工作价值的最好印证。