📅 创建时间:2026-06-03 🏷️ 标签:#AutoScheduling #AutoTVM #Ansor #CostModel #Tuning #MeasurementBased #Tensorization 📚 前置知识:[[/04-ai/01-llm-engineering/07-llm-evolution]](LLM 发展脉络)[[19-tvm-te]](TVM TE)[[17-kernel-primer]](Kernel 开发入门) 📚 相关知识:[[20-triton]](Triton)[[18-cutlass]](CUTLASS)
基于成本模型与实测搜索的自动调度 / Automatic Scheduling with Cost Models and Measurement-Based Search
┌──────────────────────────────────────────────────────────────────────────────┐ │ 📖 场景:AutoScheduler 调优后上线性能抖动 │ ├──────────────────────────────────────────────────────────────────────────────┤ │ 你用 TVM 的 AutoScheduler(Ansor)给一个自定义算子调优: │ │ - 跑了 8 小时,2000 个配置,找到了一个比默认快 5 倍的 kernel │ │ - 但上线后发现性能抖动很大:有时候快,有时候慢 │ │ - 你怀疑 Ansor 的成本模型有偏差,在特定输入大小上失真了 │ └──────────────────────────────────────────────────────────────────────────────┘
第1节 为什么需要自动调度——手工 schedule 参数空间太大
1.1 Schedule 参数空间有多大
TVM TE 的 schedule 原语看似简单,但组合起来参数空间是爆炸的:
# 一个 GEMM 的 schedule 需要选择:
#
# 1. Tile 大小:
# - block_m: 16, 32, 64, 128, 256 (5个选项)
# - block_n: 16, 32, 64, 128, 256 (5个选项)
# - block_k: 16, 32, 64 (3个选项)
#
# 2. Thread 绑定:
# - thread_m: 4, 8, 16, 32 (4个选项)
# - thread_n: 4, 8, 16, 32 (4个选项)
#
# 3. Loop order (每个 tile 后的 reorder):
# - 可能的排列数: 5! = 120 (但很多是非法的)
#
# 4. Unroll 策略:
# - 哪些 axis unroll: 2^5 = 32
# - unroll 阈值: 4, 8, 16, 32, 64
#
# 5. Vectorize:
# - 哪些 axis vectorize: 2^5 = 32
# - vector size: 2, 4, 8, 16
#
# 总参数空间: 5×5×3×4×4×32×5×32×32 = 估计 > 10^9 种配置!
# 手工测试?不可能
# Ansor 的价值:在巨大参数空间中自动找到最优配置1.2 手工调优 vs 自动调度
┌─────────────────────────────────────────────────────────────────────────┐
│ 手工调优的局限 │
│ │
│ 1. 依赖专家经验:需要理解硬件细节 + schedule 原语 │
│ 2. 难以穷举:10^9 种配置,经验只能覆盖一小部分 │
│ 3. 难以迁移:不同硬件需要重新调优 │
│ 4. 难以维护:手动找到的参数难以复现和分析 │
│ │
│ 自动调度的价值 │
│ │
│ 1. 系统化搜索:可以探索手工难以触及的配置 │
│ 2. 可复现:所有配置和结果都有记录 │
│ 3. 可迁移:相同算子在不同硬件上可以自动重新搜索 │
│ 4. 可积累:历史调优数据可以用于 cost model 训练 │
└─────────────────────────────────────────────────────────────────────────┘第2节 AutoTVM 架构——调度空间、代价模型、剪枝、测量
2.1 AutoTVM 的四个核心组件
┌─────────────────────────────────────────────────────────────────────┐
│ AutoTVM 系统架构 │
│ │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ 1. Search Space (搜索空间) │ │
│ │ │ │
│ │ 你用 TVM TE 定义算子 │ │
│ │ AutoTVM 提供 ConfigSpace 描述所有可能的 schedule 配置 │ │
│ │ 例如: tile_sizes ∈ {16, 32, 64}, unroll ∈ {0, 1} │ │
│ └─────────────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ 2. Cost Model (代价模型) │ │
│ │ │ │
│ │ 给定一个 schedule 配置,预测执行时间 │ │
│ │ 两种方式: │ │
│ │ - 静态模型: Roofline model, analytical model │ │
│ │ - ML 模型: GBDT, Neural Network │ │
│ │ │ │
│ │ CostModel.predict(config) → 预测时间 (ms) │ │
│ └─────────────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ 3. Search Policy (搜索策略) │ │
│ │ │ │
│ │ 如何在巨大空间中高效搜索? │ │
│ │ - Random Search: 随机采样 │ │
│ │ - Evolutionary: 遗传算法,变异+交叉 │ │
│ │ - RL-based: 强化学习(PPO 等) │ │
│ │ - UCB: Upper Confidence Bound,平衡探索与利用 │ │
│ │ │ │
│ │ policy.select() → 选择下一个要评估的配置 │ │
│ └─────────────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ 4. Measurement (实际测量) │ │
│ │ │ │
│ │ Cost Model 有偏差 → 需要真实测量 │ │
│ │ 在目标硬件上实际运行 kernel,记录执行时间 │ │
│ │ 测量结果反馈给 Cost Model,用于下次预测 │ │
│ │ │ │
│ │ measure(builder, runner, config) → 实测时间 (ms) │ │
│ └─────────────────────────────────────────────────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────┘2.2 Search Space 定义
import tvm
from tvm import te, autotvm
# 定义算子 (同 TVM TE)
N, F, K = 1024, 512, 3
# placeholder
data = te.placeholder((N, F), name="data", dtype="float32")
weight = te.placeholder((F, F), name="weight", dtype="float32")
bias = te.placeholder((F,), name="bias", dtype="float32")
# Conv2d 简化为 GEMM
# 实际使用 Conv2d 更复杂,这里用全连接层演示
k = te.reduce_axis((0, F), name="k")
out = te.compute(
(N, F),
lambda i, j: te.sum(data[i, k] * weight[k, j], axis=k) + bias[j],
name="out"
)
# ============================================================
# 定义 Search Space (ConfigSpace)
# ============================================================
# 使用 @autotvm.config 不同参数化的装饰器
@autotvm.config 不同参数化的装饰器
@autotvm.tuner "grid" # 使用网格搜索调优器
def tune_gemm(cfg, M, N, K):
# cfg: 配置对象,调优器会修改这个对象
# cfg.space: 搜索空间定义
s = te.create_schedule(out.op)
# ============================================================
# 定义可调参数
# ============================================================
# 方法 1: cfg.space 定义离散取值
# cfg.define_split("tile_x", N, num_outputs=2)
# cfg.define_knob("tile_y", [16, 32, 64])
# 方法 2: cfg.define_float_range / cfg.define_int_range
# cfg.define_knob("unroll_k", [0, 1]) # 0 = 不 unroll, 1 = unroll
# 在 schedule 中使用配置
i, j = s[out].op.axis
k = s[out].op.reduce_axis[0]
# cfg.define_split 生成一个可调参数
# 第二个参数是 axis 的长度
# num_outputs=2 表示拆成 outer 和 inner
ko, ki = s[out].split(k, cfg="tile_k") # 使用名为 "tile_k" 的配置
# cfg.define_knob 生成一个离散的可调参数
# 可以直接用 cfg["xxx"] 获取当前值
if cfg["unroll_k"].value == 1: # 访问 knob 的当前值
s[out].unroll(ki)
# 定义 thread binding
i_outer, i_inner = s[out].split(i, factor=32)
j_outer, j_inner = s[out].split(j, factor=32)
s[out].reorder(i_outer, j_outer, ko, i_inner, j_inner, ki)
# 绑定到 GPU
s[out].bind(i_outer, te.thread_axis("blockIdx.x"))
s[out].bind(j_outer, te.thread_axis("blockIdx.y"))
s[out].bind(i_inner, te.thread_axis("threadIdx.x"))
s[out].bind(j_inner, te.thread_axis("threadIdx.y"))
return s, [data, weight, bias, out]2.3 Cost Model
# AutoTVM 支持的 Cost Model 类型
from tvm import autotvm
# 1. 静态模型 (Roofline-based)
# 适用于: 结构简单的算子,算术强度容易估计
model_static = autotvm.tensorize_cost_model.RooflineModel()
# 2. GBDT 模型 (Gradient Boosting)
# 适用于: 复杂算子,需要从大量数据中学习模式
model_gbdt = autotvm.model.GBDTTuner(
feature_type="knob", # 使用 schedule knob 作为特征
loss_type="rank:listMLE", # 排序损失
)
# 3. Neural Network 模型
# 适用于: 大规模搜索,需要更好的泛化
model_nn = autotvm.model.NNTTuner(
model="mlp", # 多层感知机
loss_type="rank:cosine", # 余弦相似度排序
)
# Cost Model 训练数据
# (schedule_config, measured_latency) pairs
training_data = [
(config_1, 1.2), # config_1 实测 1.2ms
(config_2, 0.8), # config_2 实测 0.8ms
(config_3, 2.1), # config_3 实测 2.1ms
# ...
]
# 训练
model_static.fit(training_data)
model_gbdt.fit(training_data, valid_data=validation_data)
# 预测
predicted_latency = model_gbdt.predict(unseen_config)
# 返回: 预测的 latency (ms)2.4 Search Policy
# AutoTVM 的搜索策略
from tvm import autotvm
# 1. Random Search
# 最简单,但效率低
tuner_random = autotvm.tuner.RandomTuner(task)
# 2. Grid Search
# 适用于: 搜索空间小,参数离散
tuner_grid = autotvm.tuner.GridSearchTuner(task)
# 3. GBDT-based Tuner
# 使用 GBDT 模型指导搜索
tuner_gbdt = autotvm.tuner.GBDTTuner(
task,
loss_type="rank:pairwise", # pairwise ranking loss
num_feature_type="knob",
)
# 4. XGBoost Tuner
# 效果通常比 GBDT 好
tuner_xgb = autotvm.tuner.XGBTuner(
task,
plan_size=64, # UCB 中 exploration 的数量
feature_type="knob",
)
# 5. RL-based Tuner
# 使用强化学习策略
tuner_rl = autotvm.tuner.RLTuner(task, num_agents=4)
# 运行调优
tuner = tuner_xgb
tuner.tune(
n_trial=1000, # 最大试验次数
early_stopping=100, # 提前停止(100次无改进后)
measure_option=autotvm.measure_option(
builder=autotvm.LocalBuilder(build_func="cuda"), # 本地编译
runner=autotvm.LocalRunner( # 本地运行
number=10, # 每个配置跑10次取平均
repeat=3, # 重复3轮
min_repeat_ms=100, # 最少运行100ms
)
),
)第3节 Ansor(AutoScheduler)——自动程序合成
3.1 Ansor vs AutoTVM 的区别
┌────────────────────────────────────────────────────────────────────────┐
│ AutoTVM vs Ansor │
│ │
│ AutoTVM: │
│ - 你手动定义 search space (define_split, define_knob) │
│ - 搜索空间有限(你定义了什么,就搜什么) │
│ - 优点:精确控制,适合专家 │
│ - 缺点:容易遗漏好配置,搜索空间爆炸时难定义 │
│ │
│ Ansor (AutoScheduler): │
│ - 自动推导 search space(从 compute graph) │
│ - 搜索空间更大、更完整 │
│ - 程序合成器 (Program Synthesizer) 可以探索程序变体 │
│ - 更好的覆盖率 │
│ │
└────────────────────────────────────────────────────────────────────────┘3.2 Ansor 的核心概念
# Ansor 使用流程
from tvm import auto_scheduler
# ============================================================
# 步骤 1: 定义算子 (同 TVM TE)
# ============================================================
N, F = 1024, 512
data = te.placeholder((N, F), name="data", dtype="float32")
weight = te.placeholder((F, F), name="weight", dtype="float32")
bias = te.placeholder((F,), name="bias", dtype="float32")
k = te.reduce_axis((0, F), name="k")
out = te.compute(
(N, F),
lambda i, j: te.sum(data[i, k] * weight[k, j], axis=k) + bias[j],
name="out"
)
# ============================================================
# 步骤 2: 创建 Task
# ============================================================
# Task: 一个需要调优的算子
# Ansor 会自动为这个 task 生成 search space
task = auto_scheduler.create_task(
func=lambda: (out.op, [data, weight, bias, out]), # compute + I/O tensors
target="cuda", # 目标硬件
)
# 打印 task 信息
print(task)
# ============================================================
# 步骤 3: 配置 Search Policy
# ============================================================
# AutoScheduler 使用的搜索策略
# 默认使用 evolutionary search + cost model
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=2000, # 最大测量次数 (耗时长时减少)
num_measures_per_round=64, # 每轮测量多少个配置
early_stopping=100, # 提前停止
num_search_iterations=100, # 搜索迭代次数
min_sample_interval_ms=50, # 两次采样之间的最小间隔
# Builder: 在目标硬件上编译 kernel
builder=auto_scheduler.LocalBuilder(build_func="cuda"),
# Runner: 在目标硬件上运行 kernel 并测量时间
runner=auto_scheduler.LocalRunner(
number=10,
repeat=3,
min_repeat_ms=100,
),
)
# ============================================================
# 步骤 4: 运行调优
# ============================================================
# 方法 1: 自动调优
policy = auto_scheduler.SketchPolicy(
task,
verbose=1,
search_policy="evolutionary", # 使用 evolutionary search
)
sch, args = auto_scheduler.auto_schedule(task, policy=policy, tuning_options=tune_option)
# 方法 2: 手动调优循环
# 给你更多控制
policy = auto_scheduler.SketchPolicy(task, verbose=1)
tune_option.runner = auto_scheduler.LocalRunner(number=10, repeat=3)
# 迭代搜索
for trial in range(1000):
# 采样一个配置
configs = policy.sample_valid_config(num=64)
# 测量配置
results = []
for config in configs:
latency = auto_scheduler.measure_one_config(task, config, tune_option)
results.append((config, latency))
# 更新 cost model
policy.update(results)
# 检查是否收敛
if trial > 100 and best_latency_stable:
break
# ============================================================
# 步骤 5: 提取最优 schedule
# ============================================================
sch, args = auto_scheduler.auto_schedule(task)
print(sch)3.3 Program Synthesizer(程序合成器)
Ansor 的 Program Synthesizer:
AutoTVM 的问题: 你需要手动定义 search space
- 你定义了 tile_x ∈ {16, 32, 64}
- 但如果最优值是 48 呢?你错过了
Ansor 的解决方案: Program Synthesizer
- 自动从 compute graph 推导出所有可能的程序变体
- 不需要手动定义 search space
- 能够发现手工难以想到的程序结构
程序合成的工作原理:
1. 从 compute graph 提取 "sketch"
sketch = 程序骨架(不完整,有空白需要填充)
2. 对 sketch 的空白进行填充
例如: tile factor, loop order, inline/compute_at 决策
3. 通过 mutation (变异) 探索程序变体
- Change tile size
- Change loop order
- Change compute_at position
- Fuse/split axis第4节 Cost Model——预测配置性能
4.1 Cost Model 的类型
| 模型类型 | 原理 | 优点 | 缺点 |
|---|---|---|---|
| 静态 (Roofline) | 基于算术强度和硬件参数计算 | 可解释,不需要训练数据 | 精度低,不考虑具体硬件特性 |
| GBDT/XGBoost | 梯度提升树 | 训练快,能处理离散特征 | 容易过拟合,泛化差 |
| Neural Network | MLP/Transformer 预测 latency | 泛化能力强 | 需要更多训练数据 |
| Hybrid | 结合静态 + ML 模型 | 平衡可解释性和精度 | 实现复杂 |
4.2 静态 Cost Model(Roofline-based)
# Roofline Cost Model 的实现思路
class RooflineCostModel:
"""
基于 Roofline Model 的静态代价模型
步骤:
1. 从 schedule 推导出算术强度 (Flops/Bytes)
2. 使用 Roofline 公式计算理论上限
3. 结合 occupancy 估算实际性能
"""
def __init__(self, hardware_params):
self.bandwidth = hardware_params["bandwidth_gbps"]
self.peak_tflops = hardware_params["peak_tflops"]
def predict(self, schedule, compute):
"""
预测给定 schedule 的性能
"""
# 步骤 1: 分析 schedule 的内存访问模式
memory_access = self.analyze_memory_access(schedule)
flops = self.count_flops(compute)
# 步骤 2: 计算算术强度
ai = flops / memory_access
# 步骤 3: Roofline 上限
roofline_peak = min(
self.bandwidth * ai, # Memory-bound 上限
self.peak_tflops # Compute-bound 上限
)
# 步骤 4: 考虑 occupancy 和其他因素
estimated_occupancy = self.estimate_occupancy(schedule)
predicted_latency = roofline_peak / estimated_occupancy
return predicted_latency
def analyze_memory_access(self, schedule):
"""
分析 schedule 的内存访问次数
"""
# 从 schedule IR 分析:
# - global memory loads/stores
# - shared memory 使用
# - register pressure
pass4.3 ML-based Cost Model
# GBDT Cost Model 的特征工程
class GBDTCostModel:
"""
使用 GBDT 的代价模型
特征类型:
1. Knob Features: 直接从 schedule config 提取
- tile sizes
- unroll flags
- vectorize factors
2. Schedule Features: 从 schedule IR 提取
- loop depth
- memory access pattern
- shared memory 使用量
3. Compute Features: 从 compute 定义提取
- 算子类型
- 数据大小
- reduce axis
"""
def extract_features(self, schedule, compute):
features = {}
# Knob Features
features["tile_m"] = schedule.config["tile_m"]
features["tile_n"] = schedule.config["tile_n"]
features["tile_k"] = schedule.config["tile_k"]
features["unroll_k"] = schedule.config["unroll_k"]
# Schedule Features
features["num_loops"] = self.count_loops(schedule)
features["has_shared_memory"] = self.check_smem(schedule)
features["has_vectorize"] = self.check_vectorize(schedule)
# Compute Features
features["flops"] = self.count_flops(compute)
features["memory_bytes"] = self.count_bytes(compute)
features["arithmetic_intensity"] = features["flops"] / max(features["memory_bytes"], 1)
return features
def train(self, X, y):
"""
X: 特征矩阵 (N × num_features)
y: 实测 latency (N,)
"""
import xgboost as xgb
self.model = xgb.XGBRegressor(
n_estimators=100,
max_depth=6,
learning_rate=0.1,
)
self.model.fit(X, y)
def predict(self, features):
return self.model.predict([features])[0]第5节 测量-based Tuning 的陷阱
5.1 硬件干扰
硬件干扰是测量结果抖动的主要原因:
1. Thermal Throttling (热节流)
- GPU 温度过高时自动降频
- 表现为: 长时间运行后性能下降
- 解决: 测量前让 GPU 冷却,监控温度
2. Background Processes (后台进程)
- 其他进程占用 GPU
- 表现为: 随机的高 latency
- 解决: 关闭其他进程,使用独占 GPU
3. Power Capping (功耗限制)
- GPU 功耗达到上限时降频
- 表现为: 多个 kernel 同时运行时性能下降
- 解决: 使用 nvidia-smi 监控功耗
4. Memory Allocation Overhead (内存分配开销)
- 第一次运行需要分配显存
- 表现为: 第一次测量比后续慢很多
- 解决: warm-up run,忽略第一次测量
5. JIT Compilation Overhead (JIT 编译开销)
- TVM/Ansor 首次编译需要时间
- 表现为: 第一次运行慢,后续快
- 解决: 预热编译,结果中排除首次5.2 输入数据差异
# 测量时输入数据的差异导致性能差异
# 问题 1: 不同 batch size
batch_sizes = [1, 4, 16, 64]
for batch in batch_sizes:
# 不同 batch 的算术强度不同
# 小 batch: 算术强度低,memory-bound
# 大 batch: 算术强度高,compute-bound
latency = measure(batch=batch)
print(f"batch={batch}, latency={latency}ms")
# 问题 2: 不同 tensor shape
shapes = [(1024, 1024), (2048, 2048), (4096, 4096)]
for shape in shapes:
# 不同 shape 的 tiling 效果不同
# 小的 shape 可能无法充分利用 GPU
latency = measure(M=shape[0], N=shape[1], K=shape[1])
print(f"shape={shape}, latency={latency}ms")
# 问题 3: 不同 data layout
layouts = ["NCHW", "NHWC", "CHWN"]
for layout in layouts:
# 不同 layout 的内存访问模式不同
# NCHW: 连续的 channel
# NHWC: 连续的 spatial
latency = measure(layout=layout)
print(f"layout={layout}, latency={latency}ms")
# 最佳实践:
# 1. 测量时要固定输入 shape 和 layout
# 2. 在目标部署场景的 shape 上进行测量
# 3. 使用与生产环境相同的数据进行测量5.3 JIT 编译开销
# JIT 编译开销的处理
# 问题: TVM 会动态编译 kernel
# 第一次运行慢,后续运行快
# 方案 1: Warm-up (预热)
for _ in range(5):
# 运行一个 dummy kernel,让 TVM 完成编译
dummy_kernel(data)
# 不记录这次的时间
# 方案 2: 记录编译时间
import time
start_compile = time.time()
kernel = tvm.build(s, args, target="cuda")
compile_time = time.time() - start_compile
# 测量时排除编译时间
for i in range(10):
start_run = time.time()
kernel(*args)
run_time = time.time() - start_run
# 如果 compile_time 被计入,需要减去
corrected_time = run_time - compile_time if i == 0 else run_time
record(corrected_time)
# 方案 3: 使用 AOT (Ahead-of-Time) 编译
# 在测量前预先编译所有可能的配置
configs = [...] # 所有候选配置
compiled_kernels = {}
for config in configs:
s = apply_config(original_compute, config)
compiled_kernels[config] = tvm.build(s, args, target="cuda")
# 测量时直接使用预编译的 kernel
for config in configs:
kernel = compiled_kernels[config]
latency = benchmark_kernel(kernel, args)第6节 Tensorization——映射到硬件指令
6.1 什么是 Tensorization
Tensorization = 把特定计算 pattern 映射到硬件的矩阵乘指令
传统 GEMM:
for i in range(M):
for j in range(N):
for k in range(K):
C[i,j] += A[i,k] * B[k,j]
每步: 1次乘 + 1次加 = 2 FLOPs
效率低,没有利用 SIMD/Tensor Core
Tensorized GEMM:
使用 WMMA/Tensor Core 指令:
- WMMA: Warp-level Matrix Multiply Accumulate
- 每条指令: 16×16×16 矩阵乘 = 8192 FLOPs
效率 = 8192 FLOPs / 1 instruction = 高得多6.2 在 TVM 中使用 Tensorization
# TVM TE 中的 tensorization
from tvm import te
from tvm import topi
# 定义 GEMM (同之前)
A = te.placeholder((1024, 1024), name="A")
B = te.placeholder((1024, 1024), name="B")
k = te.reduce_axis((0, 1024), name="k")
C = te.compute((1024, 1024), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k))
# 创建 schedule
s = te.create_schedule(C.op)
# 定义 tiling
i, j = s[C].op.axis
k = s[C].op.reduce_axis[0]
# 使用 tensorization 需要指定 tensor intrinsic
# tensor intrinsic 定义了如何使用硬件指令
# 例如: 使用 wmma tensorization (Ampere Tensor Core)
# 定义 tensor intrinsic
from tvm.topi import tensorize
# 设置 tensorization 引用的 intrin
# 需要定义:
# 1. input shape: 16×16×16
# 2. output shape: 16×16
# 3. compute: C = A @ B + C
# 这部分需要定义 TensorIntrin 对象
# TVM 提供了预定义的 tensorization intrinsic
# 例如: topi.nn.mma_intrin_group
# 配置 schedule 使用 tensorization
# s[C].tensorize(j, wmma_intrin_group)
# 注意事项:
# 1. Tile size 必须匹配 tensor intrinsic 的要求 (通常是 16 的倍数)
# 2. 需要正确设置 memory layout
# 3. 累加类型和输入类型可能不同 (fp16 input + fp32 accumulator)第7节 实际调优工作流
7.1 从定义 Search Space 到 Production
完整调优工作流:
┌─────────────────────────────────────────────────────────────────────┐
│ Phase 1: 准备 │
│ │
│ 1. 确定目标硬件 │
│ - GPU: A100, H100, RTX 4090... │
│ - 指定 target="cuda" 或 target="opencl" │
│ │
│ 2. 确定目标 workload │
│ - 具体的 shape: (B, H, S, D) │
│ - 具体的 dtype: float16, bfloat16, int8... │
│ - 实际的 batch size 和序列长度 │
│ │
│ 3. 准备测量环境 │
│ - 关闭其他进程 │
│ - 设置 GPU 为独占模式 │
│ - 确保散热良好 │
└─────────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────────┐
│ Phase 2: 调优 │
│ │
│ 1. 定义算子和 search space │
│ - 使用 TVM TE 定义 compute │
│ - 使用 AutoTVM 或 Ansor 定义 search space │
│ │
│ 2. 选择调优器 │
│ - 小搜索空间: GridSearch │
│ - 中等搜索空间: XGBoost / GBDT │
│ - 大搜索空间: RL-based / Evolutionary │
│ │
│ 3. 运行调优 │
│ - 设置 num_measure_trials (取决于时间预算) │
│ - 监控测量结果 │
│ - 注意排除 warm-up 和 outlier │
│ │
│ 4. 验证结果 │
│ - 在多个输入 shape 上验证 │
│ - 检查性能稳定性 (多次运行的方差) │
│ - 与 baseline (cuBLAS 等) 对比 │
└─────────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────────┐
│ Phase 3: 部署 │
│ │
│ 1. 导出最优配置 │
│ - 保存 schedule 和 config 到文件 │
│ - 包括: schedule JSON, tuning records │
│ │
│ 2. 集成到生产代码 │
│ - 加载预调优的 schedule │
│ - 或使用 compiled runtime │
│ │
│ 3. 监控和回归测试 │
│ - 上线后监控实际性能 │
│ - 定期重新调优 (硬件变化、驱动更新) │
└─────────────────────────────────────────────────────────────────────┘7.2 AutoTVM vs Ansor vs RL Schedulers 对比
| 维度 | AutoTVM | Ansor (AutoScheduler) | RL-based |
|---|---|---|---|
| 搜索空间定义 | 手动定义 (define_split, define_knob) | 自动推导 | 手动定义 |
| 覆盖率 | 受限于手动定义 | 高(程序变体搜索) | 受限于定义 |
| 需要专家知识 | 高 | 低 | 高 |
| 调优时间 | 短(搜索空间小) | 长(搜索空间大) | 中等 |
| 泛化能力 | 低(只覆盖定义的配置) | 高(能发现新变体) | 中等 |
| 成本模型 | 可选 | 内置 | 可选 |
| Tensorization 支持 | ✅ 手动指定 | ✅ 自动推断 | ❌ 不支持 |
| 适用场景 | 专家调优、小搜索空间 | 大搜索空间、自动搜索 | 研究、复杂策略 |
7.3 调优脚本示例
#!/usr/bin/env python3
"""
AutoTVM / Ansor 调优脚本模板
"""
import tvm
from tvm import te, autotvm, auto_scheduler
import numpy as np
def tune_with_autotvm(task, num_trials=500):
"""使用 AutoTVM 调优"""
# 创建调优器
tuner = autotvm.tuner.XGBTuner(task)
# 配置测量选项
measure_option = autotvm.measure_option(
builder=autotvm.LocalBuilder(build_func="cuda"),
runner=autotvm.LocalRunner(
number=10,
repeat=3,
min_repeat_ms=100,
enable_cpu_affinity=False,
),
)
# 运行调优
tuner.tune(
n_trial=num_trials,
early_stopping=50,
measure_option=measure_option,
callbacks=[
autotvm.callback.progress_bar(num_trials),
autotvm.callback.log_to_file("tuning_records.log"),
],
)
# 加载最优配置
with autotvm.apply_history_best("tuning_records.log"):
with tvm.target.Target("cuda"):
sch, args = tvm.build(task.compute_dag, task.config_space)
return sch, args
def tune_with_ansor(task, num_trials=1000):
"""使用 Ansor 调优"""
# 配置搜索策略
policy = auto_scheduler.SketchPolicy(
task,
verbose=1,
search_policy="evolutionary",
)
# 配置测量选项
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=num_trials,
num_measures_per_round=64,
early_stopping=100,
builder=auto_scheduler.LocalBuilder(build_func="cuda"),
runner=auto_scheduler.LocalRunner(
number=10,
repeat=3,
min_repeat_ms=100,
),
)
# 运行调优
sch, args = auto_scheduler.auto_schedule(task, policy=policy, tuning_options=tune_option)
return sch, args
def main():
# 定义算子
N, F = 1024, 512
data = te.placeholder((N, F), name="data", dtype="float32")
weight = te.placeholder((F, F), name="weight", dtype="float32")
k = te.reduce_axis((0, F), name="k")
out = te.compute(
(N, F),
lambda i, j: te.sum(data[i, k] * weight[k, j], axis=k),
name="out"
)
# 创建 task
task = tvm.auto_scheduler.create_task(
func=lambda: (out.op, [data, weight, out]),
target="cuda",
)
# 选择调优方法
use_ansor = True # 设置为 False 使用 AutoTVM
if use_ansor:
sch, args = tune_with_ansor(task, num_trials=1000)
else:
sch, args = tune_with_autotvm(task, num_trials=500)
# 编译
with tvm.target.Target("cuda"):
module = tvm.build(sch, args, "cuda")
# 保存编译结果
module.export_library("tuned_gemm.tar")
print("调优完成!结果已保存。")
if __name__ == "__main__":
main()升华
┌────────────────────────────────────────────────────────────────────────────┐
│ 🚀 自动调度核心原则 │
├────────────────────────────────────────────────────────────────────────────┤
│ │
│ ① 先测量再搜索:代价模型有偏差,测量验证是必须的 │
│ 搜索的结果只是候选,必须在真实硬件上验证 │
│ │
│ ② 搜索空间决定上限:Ansor 自动推导优于手动定义 │
│ 手动定义的搜索空间可能遗漏最优配置 │
│ │
│ ③ 测量环境要干净:硬件干扰是抖动的主要原因 │
│ 热节流、后台进程、JIT 开销都会影响测量准确性 │
│ │
│ ④ 测量要有代表性:只在目标 shape 上调优 │
│ 小 batch 调优的结果可能不适用于大 batch │
│ │
│ ⑤ 调优不是一次性:上线后要监控,定期重新调优 │
│ 驱动更新、硬件老化都可能导致性能变化 │
│ │
└────────────────────────────────────────────────────────────────────────────┘"AI 可查 vs 必须理解"清单
必须理解(不理解就等于不会):
- 🔴 为什么需要自动调度:手工 schedule 参数空间爆炸问题
- 🔴 AutoTVM vs Ansor 的区别:搜索空间是如何定义的
- 🔴 Cost Model 的作用和局限:为什么 cost model 预测不等于真实性能
- 🔴 测量-based tuning 的陷阱:热节流、JIT 开销、数据差异如何影响结果
- 🔴 测量环境的重要性:为什么要在与生产相同的环境中测量
AI 可查(知道去哪查就行):
- ✅ AutoTVM 的具体 API(tuner.tune、measure_option 等)
- ✅ Ansor 的具体参数(SketchPolicy、TuningOptions)
- ✅ XGBoost/GBDT 的特征工程方法(特征类型和提取代码)
- ✅ 特定硬件的 optimal 参数(benchmark 数据)
- ✅ TVM tensorization intrinsics 的具体定义(TVM 源码)
学习状态:🟡 开始学习