MLKAPS框架：基于自适应采样与决策树的HPC内核自动调优实践

1. 项目概述与核心挑战

在咱们高性能计算（HPC）这个行当里，性能调优一直是个既关键又头疼的活儿。特别是那些底层的计算内核，比如线性代数库里的矩阵分解函数，它们的性能对上层应用的影响是决定性的。传统上，这事儿得靠经验丰富的工程师手动去试，根据硬件特性、问题规模，一点点调整线程数、分块大小、循环展开因子这些参数。但说实话，人脑在面对动辄几十上百维的参数空间时，局限性太大了。你很难保证找到的就是全局最优，而且换个硬件平台或者问题规模，之前那套经验可能就不好使了，又得从头再来。

这就是自动调优（Auto-tuning）技术要解决的问题。它的核心思路很直观：与其靠人海战术或专家经验去“猜”，不如让算法自己去“学”和“找”。具体来说，就是构建一个代理模型，让它去学习“输入问题特征 + 参数配置”到“执行性能”之间的映射关系。然后，基于这个模型，用优化算法在庞大的参数空间里搜索最优配置。理想情况下，这个模型学得越好，搜索效率越高，最终找到的配置性能就越强。

然而，理想丰满，现实骨感。我在实际工作中发现，现有的自动调优框架普遍面临几个棘手的挑战：

1. 采样效率与成本：构建代理模型需要大量的性能采样点（即用不同的参数配置实际跑一下程序，记录耗时）。采样点太少，模型不准；采样点太多，调优过程本身耗时过长，失去实用价值。如何在有限的采样预算内，采到“信息量最大”的点，是个关键。
2. “全局准确”与“局部精准”的矛盾：我们希望代理模型对整个参数空间的性能都有个大致准确的把握（全局准确），但更希望它在可能的最优解附近预测得特别准（局部精准）。很多采样策略，比如随机采样或旨在提升全局均匀性的拉丁超立方采样，往往在提升全局准确性的同时，牺牲了最优解附近的精度。
3. 模型的可扩展性与部署：调优完成后，如何把找到的“最优配置规则”应用到实际程序中？是生成一个巨大的查找表，还是一个复杂的模型？前者可能内存爆炸，后者可能推理开销太大。我们需要一种既轻量又能快速决策的机制。
4. 与现有专家知识的融合：像 Intel MKL 这样的工业级库，已经包含了大量专家手工优化的成果。一个鲁棒的自动调优框架，不应该完全推倒重来，而应该能识别出哪些地方手工调优已经做得很好（避免性能回退），哪些地方还有优化潜力，并与之结合。

MLKAPS 这个框架，正是针对上述痛点的一次系统性工程实践。它不仅仅是一个算法玩具，而是瞄准了生产环境落地，目标是为 HPC 内核生成可以直接嵌入、用于运行时动态决策的轻量级决策树。下面，我就结合自己的理解，拆解一下它的核心设计思路和那些值得关注的实操细节。

2. MLKAPS 核心设计思路拆解

MLKAPS 的全称是 Machine Learning and Adaptive Sampling for HPC Kernel Auto-tuning，名字就点明了它的两大支柱：机器学习和自适应采样。它的工作流程可以概括为“两阶段解耦”模式，这与许多在线学习、边采样边优化的框架（如基于贝叶斯优化的方法）有显著区别。

2.1 两阶段解耦：采样建模与决策生成

第一阶段是“采样与建模”。在这个阶段，MLKAPS 的核心任务是利用有限的采样预算，尽可能高效地构建一个高质量的代理模型。这里的关键在于采样策略。MLKAPS 对比了多种策略：

• 随机采样：最简单，但效率低下，容易在非关键区域浪费采样点。
• 拉丁超立方采样：能保证参数空间投影的均匀性，提升了全局探索能力。
• HVS：一种旨在最大化每次采样信息增益的策略，目标是快速提升模型的全局准确性。
• GA-Adaptive：这是 MLKAPS 重点提出的策略。它的核心思想是偏向于最优解区域进行采样。通过结合遗传算法的搜索能力，它更关注于找到并精细刻画那些高性能区域，从而提升模型在最优解附近的局部精度。

为什么局部精度如此重要？想象一下，你要在一片山区里找最高峰。一张粗略的等高线地图（高全局精度）能告诉你山脉的大致走向和几个可能的高点区域。但如果你有一张在某个疑似最高点附近精度极高的地图（高局部精度），你就能更准确地判断它是否真的是最高点，以及攀登它的最佳路径。对于调优而言，最终我们只关心那个“最高点”（最优配置）及其精确性能。因此，GA-Adaptive 这种“好钢用在刀刃上”的思路，在实践中被证明更有效。原文中的实验数据也清晰显示，尽管 GA-Adaptive 构建的模型全局误差可能更大，但在最优解附近的预测误差（MAE）显著更低，最终带来的调优效果（平均加速比）也最好。

第二阶段是“优化与决策树生成”。当代理模型（MLKAPS 使用梯度提升决策树 GBDT，如 LightGBM）训练好后，我们就不再需要进行昂贵的真实性能评估了。我们可以在这个“模拟器”上，为任意给定的输入问题（如矩阵的维度 m, n），快速评估海量的参数配置，并找出预测性能最好的那一个。但这里又有一个问题：我们不可能为所有可能的输入组合都预计算并存储一个最优配置，那需要巨大的存储空间。

MLKAPS 的解决方案是生成决策树。决策树的妙处在于，它通过一系列简单的“if-else”规则，就能将连续的输入空间划分成不同的区域，并为每个区域分配一个固定的最优配置。这种规则化的表达非常紧凑，推理速度极快（几乎无开销），并且可以直接以源代码的形式嵌入到目标内核中，实现真正的运行时动态调优。例如，生成的 C 代码可能长这样：

int get_optimal_block_size(int m, int n) { if (m <= 1024) { if (n <= 512) return 64; else return 128; } else { if (m <= 2048) return 256; else return 512; } }

2.2 与主流框架的定位差异

理解 MLKAPS 的定位，有助于我们在实际项目中做出正确选择。原文中它主要与 Optuna 和 GPTune 进行了对比。

• vs. Optuna：Optuna 是一个强大的超参数优化框架，但它本质上是为“单个任务”寻找一组最优参数。对于 HPC 内核调优，这意味着如果你有 1000 种不同的（m, n）输入，你需要为每一种输入单独运行一次 Optuna 优化，成本不可接受。MLKAPS 通过构建一个统一的、能够泛化到不同输入的代理模型，一次性解决了整个输入空间的问题，实现了高效的迁移学习。在针对大规模输入空间时，MLKAPS 的优势明显。
• vs. GPTune：GPTune 是一种基于多任务贝叶斯优化的先进框架，它也能处理多个输入点。但其核心模型（如线性协同区域化模型 LMC）的复杂度随着任务数（输入点数量）和样本数的增加而非线性增长，这导致了严重的内存和计算可扩展性问题。原文实验显示，在样本数增加到几千时，GPTune 就会因内存耗尽而崩溃。而 MLKAPS 的 GBDT 模型和后续的决策树生成，其开销与样本数基本呈线性关系，能够轻松处理数万乃至更多样本，更适合需要大量采样才能收敛的复杂调优场景。

简单来说，MLKAPS 更适合于“大规模采样、大规模输入空间、需生成可部署决策树”的场景；而 GPTune 和 Optuna 则在“采样预算极其有限、输入点很少、参数空间有复杂约束”的场景下可能更有优势。

3. 关键组件深度解析与实操要点

3.1 GA-Adaptive 采样策略的实现心法

GA-Adaptive 是 MLKAPS 提升调优效果的关键。它的目标不是让模型对所有点都预测得准，而是让模型对“好”的点预测得特别准。在实操中，实现这样一个采样器需要注意以下几点：

1. 初始种群与探索：不能一开始就只盯着可能的好区域，否则会陷入局部最优。通常，我们会用一部分采样预算（例如前 20%）进行随机或拉丁超立方采样，以建立对全局性能分布的初步认知。这个初始模型虽然粗糙，但能为遗传算法提供初步的适应度评估依据。
2. 适应度函数设计：这是遗传算法的指挥棒。我们不能直接用代理模型预测的绝对性能作为适应度，因为模型初期可能不准。一个更稳健的做法是使用预期改进或置信上界。例如，可以采用高斯过程回归中常用的“期望提升”思想，或者简单地将“模型预测值 + 不确定性系数 * 模型预测标准差”作为适应度，鼓励算法去探索那些预测值高或者不确定性大的区域（兼顾利用与探索）。
3. 遗传操作与精英保留：交叉和变异操作需要根据参数类型（连续值、离散值、类别值）精心设计。对于连续参数，可以使用模拟二进制交叉和多项式变异；对于离散参数，则需要特殊的处理。同时，必须采用精英保留策略，确保每一代的最佳个体不会丢失。
4. 与代理模型的交互：GA-Adaptive 是一个迭代过程。每一代遗传算法产生的新候选点，都需要用当前的代理模型进行评估（预测适应度）。然后，从这些候选点中选出真正有潜力的几个（比如适应度最高的前 10%），进行真实评估（即实际运行内核），并将这个真实的（参数，性能）数据点加入到训练集中，更新代理模型。这样，模型就在我们关心的区域变得越来越精确。
5. 平衡参数：需要控制好两个比例：一是每一代中，用模型预测筛选的点与最终真实评估的点的比例；二是探索（ Exploration ）与利用（ Exploitation ）的平衡。这通常需要通过一些初步实验来确定。

注意：GA-Adaptive 的效能高度依赖于初始采样和代理模型初版的质量。如果初始模型完全偏离实际，遗传算法可能会被误导到错误的区域。因此，足够的初始随机采样至关重要。

3.2 从代理模型到可嵌入决策树

训练好的 GBDT 模型本身已经可以用于预测，但直接集成 GBDT 模型到 HPC 内核中是不现实的，因为它的推理过程涉及多棵树的遍历和求和，开销相对较大。决策树才是最终交付的“产品”。

1. 决策树生成流程：

   • 网格化输入空间：首先，将关心的输入空间（如 m 从 100 到 5000， n 从 100 到 5000）离散化为一个网格。网格密度是一个权衡，太密则后续优化计算量大，太疏则可能错过最优边界。原文使用了 46x46 的网格。
   • 网格点优化：对于网格上的每一个点 (m_i, n_i)，使用训练好的代理模型，在其参数配置空间内进行优化（例如使用局部搜索、单纯形法等轻量级优化器），找到针对该特定输入点的“模型预测最优配置”。
   • 决策树构建：现在，我们有了一个数据集：(m_i, n_i) -> optimal_parameters。利用这个数据集，使用决策树算法（如 CART）进行训练。决策树的学习目标，就是根据 (m, n) 预测出对应的 optimal_parameters。这里的一个关键技巧是，每个需要调优的参数（如分块大小、线程数）通常需要单独训练一棵决策树。

2. 决策树的优势与陷阱：

   • 优势：推理速度快，代码简洁，可嵌入，可解释性强（你可以看到清晰的规则）。
   • 陷阱：决策树会产生“分块效应”。由于它使用轴平行的分割，会在输入空间形成规则的矩形区域（如原文图10中的“棋盘格”状图案）。这意味着在区域边界附近，输入参数的微小变化可能导致配置的跳跃性改变，而实际性能可能是连续变化的。这可能会引入不必要的性能抖动。缓解方法包括使用更密的网格进行训练，或者使用更复杂的模型（如梯度提升树的叶子节点索引）作为标签，但会牺牲一些可解释性。

3.3 专家知识注入：避免性能回退的保险丝

这是 MLKAPS 设计中非常务实和工业化的一环。自动调优的结果再漂亮，如果在某些常见场景下性能反而比不过高度优化的手工版本（如 MKL），那也是无法接受的。MLKAPS 提供了一种优雅的集成方式。

具体做法是：在生成最终的决策树后，并不直接替换原有的手工调优逻辑。而是构建一个“专家树”或一个简单的选择逻辑。对于任何一个输入 (m, n)：

1. 分别查询MLKAPS 决策树和手工调优规则（或基准配置）给出的参数。
2. 在一个小的、有代表性的测试集上，快速评估这两种配置的实际性能（或使用代理模型预测性能）。
3. 选择性能更好的那一个作为最终配置。

这个过程可以离线进行，其结果就是一棵新的、融合了双方优点的决策树。这棵“专家树”保证了性能“永不回退”：在任何输入下，其性能至少不低于手工调优版本。原文图12的结果完美展示了这一点，在融合了 MKL 的专家知识后，所有性能回归点都消失了，同时又在其他区域保留了 MLKAPS 带来的加速收益。

在实际部署中，这相当于一个安全的“双保险”机制，极大地提升了自动调优方案的可接受度。

4. 实战：以 Intel MKL dgetrf 内核为例的调优全流程

假设我们现在要使用 MLKAPS 对 Intel MKL 中的dgetrf（双精度 LU 分解）函数进行调优。以下是基于原文思路梳理的一个可操作流程。

4.1 环境准备与参数空间定义

首先，需要明确调优的目标和范围。

1. 确定调优参数：对于dgetrf，我们需要分析其可能的性能相关参数。这些参数可能包括（具体需要查阅 MKL 文档或通过性能分析工具推测）：

   • BLOCK_SIZE: 算法分块的大小。
   • NUM_THREADS: 使用的线程数。
   • INTERNAL_THREAD_AFFINITY: 内部线程亲和性设置。
   • LOOP_UNROLL_FACTOR: 循环展开因子。
   • 等等。假设我们最终确定了 8 个关键的可调参数P1, P2, ..., P8。每个参数都有其取值范围（离散或连续）。

2. 定义输入空间：dgetrf的输入是矩阵维度m和n。我们需要定义一个有代表性的输入空间，例如m, n ∈ [100, 5000]。这个范围应覆盖目标应用场景的典型问题规模。

3. 搭建实验框架：

   • 编写一个封装程序，它能接受(m, n, P1, P2,..., P8)作为输入，调用对应配置的 MKLdgetrf，并精确测量其运行时间（需多次预热和运行取中位数或最小值以减少噪声）。
   • 准备一个采样调度器，能够根据 MLKAPS 核心的指令，执行指定的参数配置并返回性能数据。

4.2 分阶段采样与模型训练

我们设���总采样预算为 30,000 次真实评估。

1. 阶段一：初始探索（约 5,000 次）

   • 策略：采用拉丁超立方采样。
   • 操作：在(m, n, P1,..., P8)的联合空间内进行采样。确保每个参数的边缘分布都是均匀的。这有助于构建一个具有良好全局特性的初始代理模型Model_v0。
   • 目的：避免 GA-Adaptive 初期因模型太差而误入歧途。

2. 阶段二：自适应精细采样（约 25,000 次）

   • 策略：启动GA-Adaptive采样器。
   • 操作： a. 以Model_v0作为初始代理模型。 b. 设置遗传算法参数：种群大小（如 100）、迭代代数、交叉变异概率等。 c. 在每一代中，遗传算法基于当前代理模型预测的适应度（如“预测执行时间的倒数”）生成大量候选点。 d. 从候选点中，选择一部分（如预测最好的一批，加上一些随机扰动以保持多样性）进行真实评估。 e. 将新评估的数据加入训练集，更新代理模型（例如，继续训练 LightGBM 模型）。 f. 重复 c-e，直到耗尽剩余的采样预算。
   • 监控：在此过程中，可以定期（如每 1000 次采样）在独立的验证集上评估代理模型的性能，观察其局部精度（在历史最优解附近的表现）是否在提升。

4.3 决策树生成与专家融合

采样完成后，我们拥有一个包含 30,000 个样本的数据集(m, n, P1,..., P8, time)。

1. 构建最终代理模型：使用全部数据，重新训练一个最终的 LightGBM 模型Model_final。这个模型是我们对性能曲面的最佳估计。

2. 网格点优化：

   • 在(m, n)空间创建 46x46 的均匀网格。
   • 对于每个网格点(m_i, n_i)，固定输入，将(P1,..., P8)作为优化变量，以Model_final预测的执行时间为目标函数，运行一个快速的局部优化算法（如 Nelder-Mead 单纯形法或随机局部搜索），找到针对该点的最优参数组合P_opt_i。
   • 这样就得到了一个映射数据集：Grid: (m_i, n_i) -> P_opt_i。

3. 训练决策树：

   • 将P_opt_i中的每个参数单独作为目标变量，训练一棵决策树DT_Pk(k=1..8)。特征都是(m, n)。
   • 关键参数：决策树的最大深度、叶子节点最小样本数等需要调优，以在模型精度和树的大小之间取得平衡。树太深会导致规则复杂、过拟合；树太浅则拟合能力不足。可以通过交叉验证来选择。

4. 专家知识注入：

   • 准备 Intel MKL 默认的dgetrf性能作为基准。
   • 在同一个验证网格上，对于每个点，比较MLKAPS决策树推荐的配置和MKL默认配置的实际性能（可以用Model_final快速预测，也可以用小规模真实运行验证）。
   • 生成一个“选择器”逻辑：对于每个区域，选择性能更优的配置来源。这个逻辑本身也可以表达为一棵简单的决策树或一系列条件语句。
   • 将“选择器”逻辑与 MLKAPS 的决策树、MKL 默认配置逻辑合并，生成最终的、无性能回退的“专家决策树”代码。

4.4 集成与测试

将生成的专家决策树代码（一组if-else或switch-case函数）集成到你的应用程序中，替代原来直接调用 MKL 默认接口的方式。在调用dgetrf之前，先根据当前的m, n通过决策树函数计算出推荐的参数，并应用这些参数（可能需要通过 MKL 的环境变量或函数接口设置）。

最后，必须进行全面的集成测试。使用一系列未见过的、覆盖整个输入空间的测试用例，对比使用专家决策树调参后的性能与纯 MKL 默认性能，验证平均加速效果，并确保没有任何案例出现显著的性能下降。

5. 常见陷阱、问题排查与效能分析

在实际操作 MLKAPS 或类似流程时，你肯定会遇到各种问题。下面是我总结的一些常见坑点和排查思路。

5.1 采样阶段的问题

• 问题1：代理模型预测误差始终很大，优化找不到方向。

  • 可能原因1：采样噪声过大。性能测量本身不稳定（缓存、系统负载等干扰）。
    • 排查：检查单个配置多次运行的性能方差。确保测量程序做了充分的预热（运行多次丢弃前几次），并使用稳定的计时函数（如clock_gettime）。
    • 解决：增加每次配置的重复运行次数，取中位数或最佳值作为代表。但这会增加采样成本。
  • 可能原因2：参数空间存在无效配置。某些参数组合可能导致程序错误或异常缓慢，这些异常点污染了训练数据。
    • 排查：检查采样数据中是否存在极端异常值（如执行时间比其他点高几个数量级）。
    • 解决：在采样框架中加入超时机制和错误捕获。对于运行失败或超时的配置，赋予一个极差的性能值（如一个非常大的数），让模型学会避开这些区域。
  • 可能原因3：模型容量不足或过拟合。
    • 排查：观察训练误差和验证误差的差距。如果训练误差很低但验证误差很高，可能是过拟合；如果两者都高，可能是欠拟合。
    • 解决：调整 GBDT 模型参数（如树的数量、深度、学习率）。对于复杂关系，可能需要更深的树或更多树。同时，确保有足够大的独立验证集。

• 问题2：GA-Adaptive 采样很快陷入局部最优，后期改进缓慢。

  • 可能原因：探索与利用的平衡没做好。
    • 排查：观察遗传算法种群中个体的多样性。是否很快所有个体都聚集在一个狭小的区域？
    • 解决：提高变异算子的概率，或在选择策略中引入更多随机性（如锦标赛选择）。也可以定期注入一些完全随机的个体到种群中，打破僵局。

5.2 决策树生成与部署阶段的问题

• 问题3：生成的决策树在边界处性能不稳定。

  • 现象：当输入(m, n)在决策树某个分割边界附近微小变化时，性能可能出现跳变。
  • 原因：这是决策树模型的固有特性，称为“不连续性”。
  • 缓解：
    1. 增加训练网格密度：使用更细的网格进行优化，让决策树学习更精细的边界。
    2. 后处理平滑：对于边界附近的值，可以取相邻两个区域配置的性能预测值，进行加权平均或直接选择性能更好的一个，作为过渡区域的配置。
    3. 接受权衡：在大多数情况下，这种性能跳变带来的损失相对于整体加速收益是微小的，可以接受。

• 问题4：决策树推理引入了可观测的开销。

  • 排查：在极小的矩阵上（如 10x10），决策树调参带来的加速可能无法覆盖其本身判断逻辑的开销。
  • 解决：这是动态调优的普遍问题。通常的解决方案是设置一个阈值。当问题规模小于某个阈值时，直接使用一个固定的、针对小规模问题优化过的保守配置（例如，单线程、小分块），跳过决策树查询。这个阈值可以通过实验确定。

5.3 效能分析与对比

为了让你对 MLKAPS 的收益和成本有更具体的认识，我们可以量化分析一下：

核心价值判断：调优过程（数小时）是一次性的离线成本。而生成的决策树，可以在该硬件平台上的所有后续运行中持续产生收益。对于长期运行的科学计算任务或云服务提供商，这种前期投入的回报率非常高。

6. 扩展思考与未来方向

MLKAPS 框架已经展示出了强大的潜力，但从生产落地角度看，还有不少可以深化和扩展的方向。

1. 约束处理的自动化：原文提到 MLKAPS 目前需要用户手动处理参数间的约束关系（例如，分块大小不能超过矩阵维度）。一个更友好的框架应该能自动识别并处理这些约束，在采样和优化过程中排除无效配置。

2. 动态优化网格：目前使用固定网格进行决策树训练，可能无法精确捕捉性能曲面的复杂边界。未来可以探索自适应网格，在性能变化剧烈的区域使用更密的网格，在平坦区域使用更疏的网格，从而在固定计算预算下提升决策树的精度。

3. 多目标优化：目前 MLKAPS 只优化执行时间。在实际场景中，我们可能还需要考虑能耗、内存带宽占用、数值稳定性等多重目标。如何构建多目标代理模型，并生成满足帕累托最优的决策树集合，是一个有挑战性的方向。

4. 跨平台迁移学习：在一个硬件平台（如 Intel SPR）上训练得到的模型或决策树，能否经过少量校准，快速迁移到另一个相似平台（如 AMD Zen4）？这可以极大降低为新硬件调优的成本。MLKAPS 的代理模型结构为此提供了可能，但需要研究如何编码硬件特征。

5. 与编译期调优的结合：MLKAPS 侧重于运行时参数调优。而像 Tensor Comprehensions、TVM、MLIR 等框架则专注于编译期的循环变换、张量化等优化。两者并非互斥，未来可以探索将 MLKAPS 作为上层调优器，为底层代码生成器选择最优的编译参数，形成端到端的优化链路。

从我个人的实践经验来看，自动调优的成功，三分之一靠算法，三分之一靠工程实现（稳定的测量、高效的调度），还有三分之一靠对问题本身和硬件架构的深刻理解。MLKAPS 提供了一个优秀的算法和系统框架，但将其应用到具体的生产内核时，仍然需要工程师仔细地定义参数空间、设计测量实验、并解读结果。它不是一个“黑箱”魔法，而是一个强大的“杠杆”，帮你把专业经验更高效、更系统地转化为性能收益。