PyTorch TunableOp 实战指南：GEMM 自动调优的完整机制、离线调优工作流与 C++/Python 接口

TunableOp 是 PyTorch 中针对 GEMM 等算子的运行时调优框架：它把标准算子调用替换为"Tunable 实现"，自动对注册的多套 BLAS 库候选方案进行基准测试并选出最快者，同时通过 CSV 文件持久化调优结果，支持多进程多卡场景与离线批量调优。本文基于仓库内 aten/src/ATen/cuda/tunable/README.md 的完整内容，并结合 aten/src/ATen/cuda/tunable/Tunable.h、aten/src/ATen/cuda/tunable/TunableOp.h、torch/cuda/tunable.py 等源码，系统讲解其工作原理、结果文件格式、离线调优流程，以及环境变量、Python、C++ 三套控制接口的完整用法。

一、TunableOp 要解决什么问题

某些算子（如 GEMM）可以用多个库或多种技术来实现。例如在 CUDA 或 ROCm 上，GEMM 既可以用 blas 库也可以用 blasLt 库实现；而在 ROCm 平台上，rocblas 和 hipblaslt 还允许用户查询所有可选算法再自行挑选。那么，如何知道哪个实现最快、应该被选用？这正是 TunableOp 提供的答案：

• 针对给定输入参数组合，为算子注册多个候选实现；
• 首次调用时对每个候选方案循环运行并取平均执行时间；
• 选出最快方案，并把结果持久化到 CSV 文件，后续调用直接命中缓存，不再重复调优。

从源码看，TunableOp 的核心类模板定义在 TunableOp.h 中：TunableOp<ParamsT>::operator() 是统一入口，内部按"查结果 → 未命中且允许调优则 FindFastest → 否则回退 Default"的顺序分派；每个候选实现是一个 Callable<ParamsT>，通过 Call(params) 执行。该接口最初源自 onnxruntime（见 Tunable.h 文件头的版权声明），由 AMD 团队移植适配到 PyTorch。

二、启用 TunableOp 与启用调优是分开的两件事

这是 TunableOp 设计中最重要的概念区分：

• 启用 TunableOp（PYTORCH_TUNABLEOP_ENABLED）意味着 PyTorch 把标准算子替换为其 Tunable 实现。每次调用 TunableOp 时，首先检查该算子在当前输入下是否已有调优结果：若有，立即调用已调优实现，即使调优设置处于打开状态也不会再触发调优；若没有调优结果且调优已启用，则对该算子的每个注册实现进行基准测试并选择最快者。
• 启用调优（PYTORCH_TUNABLEOP_TUNING）独立于上一步，只控制"是否允许现场寻找最快实现"。

这种分离使得典型部署模式成为可能：第一阶段带调优运行以生成结果文件，第二阶段关闭调优、直接加载结果文件运行生产负载。

TunableOp.h 中的 operator() 实现印证了这一逻辑：先 mgr.Lookup(op_sig, concrete_sig) 查询结果；未命中且 IsTuningEnabled() 为真时执行 FindFastest(params) 并 mgr.Add(...) 写回结果管理器；未命中且调优关闭时，则可选择记录未调优算子（离线调优的采集开关），或在开启 wildcard fallback 时执行 LookupWildcardFallback 按通配模式匹配持久化条目。若最终没有任何结果，返回 ResultEntry::Default()，即原始的未调优 PyTorch 实现——保证任何配置下都有正确结果兜底。

三、结果文件的输入与输出

3.1 文件读写时机与命名

• 任意 TunableOp 首次被调用时，内部调优数据库会尝试从指定文件读入结果。默认文件名是 tunableop_results.csv。
• 为支持"多 GPU × 多进程"的调优场景，GPU 设备序号会自动插入文件名，避免多个进程互相覆盖同一文件。
• 若调优启用、且运行过程中发现了新的调优结果，程序退出时会自动把全部结果（启动时读入的 + 运行中新发现的）写回同一文件。这使得可以复用同一个文件、跨多个工作负载逐步累积调优库。
• 输出文件在应用终止时自动创建。该行为只能通过 C++ 与 Python API 控制，不能通过环境变量控制。

3.2 CSV 文件格式

指定文件名后，最终会得到一个形如这样的 CSV 文件：

Validator,PT_VERSION,2.2.0
Validator,ROCM_VERSION,6.0.0.0-12969-1544e39
Validator,HIPBLASLT_VERSION,0.6.0-a9c5cc7
Validator,ROCBLAS_VERSION,4.0.0-72e57364-dirty
GemmTunableOp_float_NT,nt_25088_4096_64,1219,1.262
GemmTunableOp_float_NT,nt_4096_4096_64,1216,0.033

Validator 行：如果 PyTorch 版本、ROCm 版本或 BLAS 库版本发生变化，TunableOp 会检测到并拒绝加载该调优文件，因为旧的调优结果很可能受其他软件变更的影响。

调优结果行：每行 4 个逗号分隔字段——算子名、算子参数、方案名、平均执行时间（执行时间是可选字段）。对 GemmTunableOp：第 1 个字段标明数据类型与输入是否转置（T 或 N），第 2 个字段标明 M、N、K 输入形状。

CSV 文件可以手工编辑，但需谨慎：

• 把方案名（第 3 字段）改成 "Default"，即可回退到 PyTorch 原始未调优实现；
• 在 ROCm 的 hipBLAS / hipBLASLt 场景下，如果你知道具体方案索引，可以直接替换该值来覆盖 TunableOp 自动选出的方案；
• 算子名与参数（第 1、2 字段）是内部命名，不应修改。

另外提供 verbose 输出选项，仅推荐调试时使用——它会输出大量诊断信息，可用于确认 TunableOp 是否真的在工作；除此之外，除非出现警告或错误，TunableOp 除文件输出外完全静默。

从源码看，Validator 机制由 Tunable.h 中的 TuningResultsValidator 类实现：它注册"获取函数 + 校验函数"对，mandatory_keys 固定包含 PT_VERSION，ROCm 构建下还会注册 ROCM_VERSION、HIPBLASLT_VERSION、ROCBLAS_VERSION 等；ReadFile 加载结果前先 ValidateAll，不通过则整体拒绝。

四、调优行为、预热与缓存效应

调优一个算子的过程是：遍历已注册实现列表，对每个实现循环运行多次并取平均执行时间来建立性能档案。相关细节与源码实现如下：

• 预热阶段：调优前有可选的 warmup 阶段，帮助硬件达到稳定的功耗状态。在 TunableOp.h 的 FindFastest 中，warmup 迭代数默认取 0（不启用），当用户设置了 max_warmup_duration 或 max_warmup_iter 时按"两者取小"计算。
• 指令缓存冲刷：调优期间硬件缓存命中率会比实际工作负载中更高，因此提供了冲刷指令缓存（ICache flush）与轮转输入张量（rotating buffer）两个选项，以获得更接近"算子运行在更大工作负载内"而非"紧循环反复执行"的忠实画像。源码中 WarmUp / ProfileStats 在每次迭代前调用 at::cuda::flush_icache()（受 IsICacheFlushEnabled 控制），FindFastest 则按 GetRotatingBufferSize() 分配多份参数副本轮流使用。
• 默认调优预算：默认每个候选方案运行 100 次迭代，或运行到 30ms 为止，取更小者，并计算平均执行时间。在成功建立档案的所有候选中选出最快者。
• 档案可能失败的两种情况：候选方案未达到与默认实现相同的数值精度，或候选方案返回错误码。对应源码中，若启用数值检查，FindFastest 会先用 Default 实现计算参考值，再对每个候选调用 NumericalCheck 比对；任何一步失败该候选即被跳过。

源码中还存在两级"快速筛除"（screening pass）：先用约 3 次迭代的粗测估计耗时，超过当前最优 1.5 倍的候选直接跳过；再视预算加测 10 次迭代，超过 1.15 倍再跳过，从而把调优总开销控制在 PYTORCH_TUNABLEOP_MAX_TUNING_DURATION_MS / MAX_TUNING_ITERATIONS 设定的边界内。

五、当前可 Tunable 的算子

5.1 ROCm 上的 TunableGemm

目前完整实现的是 ROCm 上的 TunableGemm。启用后，任何对 at::cuda::blas::gemm() 或 ::bgemm() 的调用都会路由到 TunableOp；给定一组输入参数 (transa, transb, m, n, k)，会尝试在 rocblas 与 hipblaslt 两个库的所有可用实现中选出最快者。候选实现分别定义在 GemmRocblas.h 与 GemmHipblaslt.h。

5.2 CUDA 上的行为

CUDA 构建的 PyTorch 在启用 TunableOp 时也能正确工作，但可用方案仅有 Default（即原来的 cuBLAS 默认路径，只是改为经由 TunableOp 调用）；对已经走 cuBLASLt 的 GEMM 路径，TunableGemm 会注册 cuBLASLt 的 heuristic 候选（数量由 PYTORCH_TUNABLEOP_CUBLASLT_REQUESTED_ALGO_COUNT 控制，默认 8），复用同一套结果缓存与计时机制。CUDA 侧候选逻辑见 GemmCublaslt.h。

六、离线调优（Offline Tuning）

6.1 动机

• 高显存利用率场景：常规在线调优需要为多个候选同时持有输入副本，显存吃紧的工作负载可能因此 OOM；
• 计算密集型场景：把整个工作负载的 GEMM 一次性收集下来，之后反复用不同调优参数或不同库进行调优，资源效率更高。

6.2 两步工作流

第 1 步：收集未调优 GEMM。 设置环境变量运行工作负载，生成 tunableop_untuned0.csv：

export PYTORCH_TUNABLEOP_ENABLED=1
export PYTORCH_TUNABLEOP_TUNING=0
export PYTORCH_TUNABLEOP_RECORD_UNTUNED=1
# ... 运行工作负载 ...

第 2 步：离线调优。 运行一个读取 untuned 文件并生成 tunableop_results0.csv 的 Python 脚本：

import torch.cuda.tunable as tunable
import os

os.putenv('PYTORCH_TUNABLEOP_ENABLED', '1')
os.putenv('PYTORCH_TUNABLEOP_TUNING', '1')
os.putenv('PYTORCH_TUNABLEOP_RECORD_UNTUNED', '0')
tunable.tune_gemm_in_file("tunableop_untuned0.csv")

从 torch/cuda/tunable.py 源码看，tune_gemm_in_file 会逐行扫描文件，对以 Gemm / ScaledGemm 开头的行调用 _process_single_offline_gemm：按第 1 字段解析算子类型（GemmTunableOp、GemmStridedBatchedTunableOp、ScaledGemmTunableOp、GemmAndBiasTunableOp）与数据类型/转置布局，按第 2 字段解析 M/N/K 与 leading dimension（ld 标记），并据此重建张量后调用 torch.mm / torch.bmm / torch._scaled_mm / F.linear 触发一次真实的（带调优的）执行，从而把在线调优机制"回放"到离线收集的形状上。

6.3 多 GPU 并行调优

多个 untuned 文件（例如来自多卡各 rank 的 tunableop_untuned0.csv ~ tunableop_untuned7.csv）可以汇总后分发到单机多卡并行调优。整体流程：

1. 汇集所有 GEMM 行并去重；
2. 以 round-robin 方式把 GEMM 分配给各 GPU；
3. 全部调优完成后，把各 GPU 的结果汇聚为单个文件，基础文件名追加 _full0（如 tunableop_results_full0.csv）；
4. 为方便用户随后在 N 卡上运行调优后的负载，该汇总文件再复制 N 份、每 GPU 一份。

if __name__ == "__main__":
    num_gpus = 8  # 调优过程中将使用的 GPU 数量
    tunable.mgpu_tune_gemm_in_file("tunableop_untuned?.csv", num_gpus)

注意 mgpu_tune_gemm_in_file 与单卡版 tune_gemm_in_file 的用法差异：

• 调用它的 Python 脚本主体必须包裹在 main() 中（如上面示例），因为内部使用 concurrent futures / 多进程；
• 文件名参数必须包含通配符（? 或 *）以展开 untuned 文件列表；
• num_gpus 必须介于 1 与可用 GPU 总数之间，否则抛出断言错误。

源码中 mgpu_tune_gemm_in_file 使用 mp.get_context("spawn") 为每块 GPU 派生一个进程（每个进程绑定单一 GPU，因为 TunableOp 输出文件名按进程初始化一次），全部进程结束后调用 torch.cuda.synchronize() 并执行 _gather_tunableop_results() 完成结果汇聚与文件复制。

七、Tuning Context：三套控制接口

TunableOp 的行为可通过环境变量、C++ 接口 at::cuda::tunable::getTuningContext()、以及 torch.cuda.tunable Python 接口来操纵。环境变量优先于 C++ / Python API 设置的值。

7.1 环境变量接口

环境变量在首次读取时被缓存，此后无法再程序化修改；需要运行时动态调整时请使用 C++ 或 Python API。

环境变量	说明
PYTORCH_TUNABLEOP_ENABLED	默认 0。设为 1 启用。
PYTORCH_TUNABLEOP_TUNING	默认 1。设为 0 禁用。
PYTORCH_TUNABLEOP_RECORD_UNTUNED	默认 0。设为 1 启用。
PYTORCH_TUNABLEOP_UNTUNED_FILENAME	默认 tunableop_untuned.csv。
PYTORCH_TUNABLEOP_VERBOSE	默认 0。设为 1 启用基本日志，2 为基本调优状态，3 为全量跟踪。
PYTORCH_TUNABLEOP_VERBOSE_FILENAME	默认 err（stderr）。设为 out 输出到 stdout，或指定文件名捕获 verbose 日志。
PYTORCH_TUNABLEOP_FILENAME	默认 tunableop_results.csv。
PYTORCH_TUNABLEOP_NUMERICAL_CHECK	默认关闭。设为 atol_rtol 形式启用，例如 1e-5_1e-5。
PYTORCH_TUNABLEOP_ROCBLAS_ENABLED	默认 1。设为 0 排除 rocblas 参与调优。
PYTORCH_TUNABLEOP_HIPBLASLT_ENABLED	默认 1。设为 0 排除 hipblaslt 参与调优。
PYTORCH_TUNABLEOP_MAX_TUNING_DURATION_MS	默认 30，单位毫秒。
PYTORCH_TUNABLEOP_MAX_TUNING_ITERATIONS	默认 100。
PYTORCH_TUNABLEOP_CUBLASLT_REQUESTED_ALGO_COUNT	默认 8，仅 CUDA。请求的 cuBLASLt heuristic 候选数量；小于 1 的值会被钳制为 1。
PYTORCH_TUNABLEOP_MAX_WARMUP_DURATION_MS	默认 0，即不使用。单位毫秒。
PYTORCH_TUNABLEOP_MAX_WARMUP_ITERATIONS	默认 0，即不使用。
PYTORCH_TUNABLEOP_ICACHE_FLUSH_ENABLED	默认 1。设为 0 禁用。
PYTORCH_TUNABLEOP_ROTATING_BUFFER_SIZE	默认（或 < 0）为查询 L2 缓存大小。设为 0 禁用。否则设为 operator 参数池的 MiB 数。例如设为设备内存缓存大小可保证每次调优迭代都使用冷缓存。
PYTORCH_TUNABLEOP_BLAS_LOG	默认 0。设为 1 启用，把 BLAS 参数写入调优 CSV 文件。

这些变量由 Tunable.cpp 中 TuningContext 的构造函数一次性读取并缓存到成员变量（enable_、tuning_enable_、max_tuning_duration_ms_、rotating_buffer_size_ 等），与"缓存后不可再改"的语义一致。

7.2 Python 接口

所有 Python API 位于 torch.cuda.tunable 模块（源码 torch/cuda/tunable.py，每个函数最终都转发到 torch._C._cuda_tunableop_* 绑定）：

Python API	说明
enable(val: bool = True) -> None	TunableOp 总开关。
is_enabled() -> bool
tuning_enable(val: bool = True) -> None	默认 True。
tuning_is_enabled() -> bool
record_untuned_enable(val: bool = True) -> None	默认 True。
record_untuned_is_enabled() -> bool
set_max_tuning_duration(duration: int) -> None	时长与迭代数同时设置时取较小者生效，至少运行 1 次迭代。
get_max_tuning_duration() -> int
set_max_tuning_iterations(iterations: int) -> None
get_max_tuning_iterations() -> int
set_cublaslt_requested_algo_count(count: int) -> None	仅 CUDA。小于 1 的值钳制为 1。
get_cublaslt_requested_algo_count() -> int	仅 CUDA。
set_filename(filename: str, insert_device_ordinal: bool = False) -> None	insert_device_ordinal=True 时自动插入当前设备序号，适合一进程一卡的场景，保证各进程写入不同文件。
get_filename() -> str
set_numerical_check_tolerances(enable: bool, atol: float, rtol: float) -> None	启用/禁用数值检查；atol 与 rtol 默认 1e-5。
get_results() -> Tuple[str, str, str, float]
get_validators() -> Tuple[str, str]
read_file(filename: Optional[str] = None) -> None	未给文件名时调用 get_filename()。
tune_gemm_in_file(filename: str) -> None	读取 untuned 文件并调优其中的 GEMM。
mgpu_tune_gemm_in_file(filename_pattern: str, num_gpus: int) -> None	读取一个或多个 untuned 文件，在一个或多个 GPU 上调优所有去重后的 GEMM。

7.3 C++ 接口

示例：

#include <ATen/cuda/tunable/Tunable.h>

at::cuda::tunable::getTuningContext()->EnableTunableOp(true);

C++ API	说明
void EnableTunableOp(bool value);
bool IsTunableOpEnabled() const;
void EnableTuning(bool value);
bool IsTuningEnabled() const;
void SetMaxTuningDurationMs(int max_duration_ms);
int GetMaxTuningDurationMs() const;
void SetMaxTuningIterations(int max_iter);
int GetMaxTuningIterations() const;
TuningResults GetTuningResults();
void SetFilename(const std::string& filename, bool insert_device_ordinal=false);
std::string GetFilename() const;
void WriteFileOnExit(bool value);
bool ReadFile(const std::string& filename={});
bool WriteFile(const std::string& filename={});

以上签名与 Tunable.h 中 TuningContext 类的公有方法一一对应。TuningContext 是进程内单例，通过 getTuningContext() 获取（见 Tunable.cpp 中函数级 static 实例），同时持有 TuningResultsManager（结果缓存，线程安全，带互斥锁）与 TuningResultsValidator（版本校验器）。

八、源码结构速览：一次 Tunable GEMM 调用的完整路径

结合仓库源码，可以把 TunableOp 的运行时链路梳理为：

1. 入口路由：at::cuda::blas::gemm() / bgemm() 在启用 TunableOp 时路由到 TunableGemm（TunableGemm.h），其 operator() 继承自 TunableOp.h 的模板类。
2. 签名查询：以 typeid 反修饰得到的算子签名（如 GemmTunableOp_float_NT）加上参数签名（如 nt_25088_4096_64）在 TuningResultsManager 中 Lookup。
3. 调优或回退：命中则直接调用；未命中且调优开启则 FindFastest 执行"探测 → 两级筛查 → 数值检查 → warmup → 正式计时"的完整流程（计时使用 StreamTimer 的 CUDA/HIP stream 事件）；未命中且调优关闭则走 record-untuned / wildcard fallback / Default 三级处理。
4. 结果持久化：退出时把 ResultsMap 连同 Validator 行写出为 CSV；读取时先过 TuningResultsValidator::ValidateAll，版本不匹配整体拒载。

动态形状场景下，TunableOp 还支持把 M/N/K/BATCH 维按位掩码（DynamicDimsMask）通配为 *，使调优结果可被后续形状复用；运行时未命中具体形状时 LookupWildcardFallback 按 token 模式匹配持久化的通配条目，命中即复用，不命中则回退 Default——这保证了"绝不产生错误结果"的底线。该机制的端到端行为可在测试 test/inductor/test_dynamic_tunable_ops.py 中查看，覆盖了 addmm、mm、bmm、baddbmm、_scaled_mm 在 NN/NT/TN/TT 布局下的具体命中、通配命中与双未命中三种结果。

九、实操建议小结

• 在线调优：PYTORCH_TUNABLEOP_ENABLED=1 默认即带调优，适合单卡开发机快速验证；
• 生产部署：先跑一遍生成 tunableop_results.csv，之后设 PYTORCH_TUNABLEOP_TUNING=0 只加载不现场调优，把调优开销移出推理路径；
• 多机多卡：用 insert_device_ordinal 或依赖自动插入的设备序号避免文件竞争；大规模收集后优先用 mgpu_tune_gemm_in_file 的离线批量调优；
• 调优画像失真（紧循环偏快）时：调大 PYTORCH_TUNABLEOP_ROTATING_BUFFER_SIZE、保留 ICache 冲刷、必要时设置 warmup 参数让硬件进入稳定功耗状态；
• 怀疑精度问题：开启 PYTORCH_TUNABLEOP_NUMERICAL_CHECK=1e-5_1e-5，让 TunableOp 自动排除与默认实现精度不一致的候选；
• 排查是否生效：PYTORCH_TUNABLEOP_VERBOSE=2（调优状态）或 =3（全量跟踪），配合 PYTORCH_TUNABLEOP_VERBOSE_FILENAME 重定向日志，避免污染 stderr。

TunableOp 的边界当前很清晰：ROCm 上具备跨 rocblas/hipblaslt 的完整算法搜索能力，CUDA 上以 cuBLASLt heuristic 候选为主、Default 兜底。对于以 GEMM 为性能瓶颈的推理与训练负载，这套"结果缓存 + 在线/离线双模式调优 + 版本校验"的机制提供了在不改模型代码的前提下换取 BLAS 层性能的手段。