TileLang 自动调优实战AutoTuner 配置搜索、输入供给、基准验证与结果缓存全解析【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang本文围绕 TileLang 的内置自动调优器autotuner展开讲解如何在不改动 kernel 主体逻辑的前提下用装饰器或编程式 API 搜索block_M/block_N/block_K、num_stages、threads等配置空间自动完成并行编译、正确性校验、计时基准测试并选出最优 kernel。读完本文你将掌握tilelang.autotune装饰器与AutoTuner.from_kernel(...).run()两种工作流的完整用法、输入张量供给的三种方式、缓存键与磁盘缓存文件的构成以及控制并行度与超时的全部环境变量能够把调优结果稳定地复用于生产或 CI 流程。两种调优工作流装饰器与编程式 APITileLang 的自动调优器承担五项职责在给定配置空间中挑选候选、并行编译候选 kernel、校验正确性、基准测试计时并把最优结果缓存以供复用。其入口由 tilelang/autotuner/tuner.py 中的AutoTuner类与AutoTuneImpl装饰器封装实现公共接口在 tilelang/autotuner/init.py 中统一导出。官方指南文档见 autotuning.md。装饰器方式tilelang.autotune叠加在tilelang.jit之上用法是在tilelang.jit之上再叠一层tilelang.autotune(configs...)并把可调参数暴露为带默认值的函数实参。调优器会用配置空间中的值逐个覆盖这些实参进行编译与评测import tilelang import tilelang.language as T def matmul_configs(M, N, K): # 示例配置空间可按目标硬件裁剪 tiles [64, 128] stages [2, 3] threads [128, 256] return [ dict(block_MBM, block_NBN, block_KBK, num_stagesS, threadsTH) for BM in tiles for BN in tiles for BK in [32, 64] for S in stages for TH in threads ] tilelang.autotune(configsmatmul_configs, warmup25, rep100, timeout60) tilelang.jit(out_idx[-1]) def matmul(M: int, N: int, K: int, block_M: int 128, block_N: int 128, block_K: int 32, threads: int 128, num_stages: int 3, dtype: str float16, accum_dtype: str float32): T.prim_func def kernel(A: T.Tensor((M, K), dtype), B: T.Tensor((K, N), dtype), C: T.Tensor((M, N), dtype)): with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threadsthreads) as (bx, by): A_s T.alloc_shared((block_M, block_K), dtype) B_s T.alloc_shared((block_K, block_N), dtype) C_f T.alloc_fragment((block_M, block_N), accum_dtype) T.clear(C_f) for ko in T.Pipelined(T.ceildiv(K, block_K), num_stagesnum_stages): T.copy(A[by * block_M, ko * block_K], A_s) T.copy(B[ko * block_K, bx * block_N], B_s) T.gemm(A_s, B_s, C_f) T.copy(C_f, C[by * block_M, bx * block_N]) return kernel # 用法通过上下文提供固定输入推荐保证各配置在相同数据上评测 import torch M N K 1024 A torch.randn(M, K, devicecuda, dtypetorch.float16) B torch.randn(K, N, devicecuda, dtypetorch.float16) C torch.empty(M, N, devicecuda, dtypetorch.float16) from tilelang.autotuner import set_autotune_inputs with set_autotune_inputs(A, B, C): tuned_kernel matmul(M, N, K) # 编译、调优并返回最优 kernel tuned_kernel(A, B, C) # 运行最优 kernel装饰器行为的几个要点均可在源码中核对configs可以是字典列表也可以是可调用对象(args...) - list[dict]。当传入 callable 时调优器会以待调优函数的实参调用它生成配置列表见 tuner.py#L866-L869。每个配置字典的键必须与可调函数实参同名出现无法匹配的多余键会直接抛出Unused keys in config: ...错误tuner.py#L903-L918。唯一的保留键是pass_configs用于按配置单独覆盖编译 pass 选项。装饰器返回的 callable 对每一组实参元组执行一次调优并在进程内按实参键缓存结果tuner.py#L1286-L1301。装饰器默认参数为warmup25、rep100、timeout100秒单配置基准测试超时其余校验与供给参数见 tuner.py#L1318-L1339。跳过调优的技巧直接显式传入可调参数调优器检测到可调实参已被赋值时会跳过搜索、直接 JIT 编译tuner.py#L944-L956。装饰器 docstring 中也给出了这一写法示例。装饰器还支持do_not_specialize(...)列出的实参变化不会触发重新调优即不进入调优缓存键tuner.py#L1273-L1285适合把“与性能无关但会频繁变化”的参数与调优解耦。编程式方式AutoTuner类需要显式管理配置与编译/评测参数时使用AutoTunerfrom tilelang.autotuner import AutoTuner kernel_factory matmul # 上面的函数已被 tilelang.jit 装饰 tuner AutoTuner.from_kernel(kernel_factory(M, N, K), configsmatmul_configs(M, N, K)) tuner.set_profile_args( warmup25, rep100, timeout60, supply_typetilelang.TensorSupplyType.Auto, # 或提供 supply_prog / ref_prog ref_proglambda A, B, C: torch.allclose(C, (A B).to(C.dtype), rtol1e-2, atol1e-2), ) tuner.set_compile_args( targetauto, # 或 cuda / hip / metal execution_backendauto, # 按目标平台解析 out_idx[-1], # 多个输出时选择返回哪些 pass_configs{ # 可选的 TVM pass / 标志 # tilelang.PassConfigKey.EXAMPLE_KEY: value, }, ) artifact tuner.run() # 编译 运行 校验全部配置 best_kernel artifact.kernel # JITKernel best_latency artifact.latency best_config artifact.config # 复用最优 kernel best_kernel(A, B, C)编程式路径相比装饰器额外暴露了run()的高级开关tuner.py#L809-L838use_pipelineTrue编译与基准测试流水线化编译完成一个立即开跑减少总墙钟时间enable_grouped_compileTrue与group_compile_sizeN分组编译目前仅对 CUDA tvm_ffi 组合生效其余组合会告警并回退为逐配置编译tuner.py#L410-L426benchmark_multi_gpuTrue与benchmark_devices[...]将不同配置分散到多块 CUDA GPU 上并行基准测试无效设备号会被忽略并告警tuner.py#L744-L807early_stopTrue与early_stop_factor2.0若某配置的前段计时已超过当前最优延迟的倍数阈值则提前截断加速大空间搜索。编译参数与评测参数分别对应CompileArgs与ProfileArgs两个不可变 dataclass完整字段定义见 param.py#L47-L151。ProfileArgs的默认值为warmup25、rep100、timeout30、backendevent、rtol1e-2、atol1e-2、max_mismatched_ratio0.01、skip_checkFalse。其中backend可选eventCUDA event 计时默认、cupti、cudagraph。仓库中的真实示例以下示例可直接打开对比不同写法均确认存在于当前仓库examples/gdn/example_chunk_delta_h.py —— 用autotune扫描配置空间examples/deepseek_nsa/benchmark/benchmark_nsa_fwd.py —— 使用tilelang.autotune的 NSA 前向基准examples/quickstart.py —— 对调优后的 kernel 使用get_profiler复测examples/hadamard_transform/example_hadamard.py —— 自定义 warmup 的 profiler 用法examples/dynamic_shape/example_dynamic.py —— 动态形状下的调优与 profilerexamples/gemm/example_gemm_persistent.py —— 持久化persistent与普通 GEMM 的对比。输入张量供给三种方式与优先级调优器需要具体输入来完成编译、校验和计时。按优先级从高到低有三种供给方式上下文管理器固定输入推荐with set_autotune_inputs(A, B, C): tuned matmul(M, N, K)set_autotune_inputs接受可变参数或单个列表set_autotune_inputs(a, b, c)与set_autotune_inputs([a, b, c])均可其实现基于线程局部thread-local栈支持嵌套且避免工作线程间的相互干扰见 capture.py#L100-L124。当程序运行在该上下文中时set_profile_args会把捕获到的张量“冻结”为内部supply_prog按设备缓存并在需要时clone保证每个候选配置都在同一份数据上评测tuner.py#L246-L285此时若再显式传入supply_prog会被忽略并给出警告。自定义供给程序def supply_prog(signature): # signature 持有描述形状/数据类型的 KernelParam 列表 # 返回与 kernel 实参一一对应的 torch 张量列表 return [A, B, C] tuner.set_profile_args(supply_progsupply_prog)内置生成器supply_typetuner.set_profile_args(supply_typetilelang.TensorSupplyType.Normal)TensorSupplyType枚举定义了Integer、Uniform、Normal、Randn、Zero、One与Auto默认按 dtype 启发式选择均匀整数/浮点区间定义见 tensor.py#L32-L39。两点重要限制内置生成器要求静态形状。若 PrimFunc 使用符号维度T.dyn生成器会直接抛出“必须具有静态形状”的错误tensor.py#L51-L62此时必须通过方式 1 或 2 提供具体输入。标量输入不能被自动生成。若 kernel 中存在既非输出、也无法自动生成的标量参数调优器会抛出明确异常提示使用with set_autotune_inputs(...)提供具体输入tuner.py#L299-L319。若 kernel 涉及 Float8 类型依赖 PyTorch 2.1 的torch.float8_*支持。supply_prog生效时supply_type的设置无效两者不可兼用见 tuner.py#L287-L290。另有cache_input_tensors开关开启后各配置间复用同一份输入张量若检测到缓存张量与新配置的 dtype/形状不兼容调优器会告警并重新生成tuner.py#L679-L708。正确性校验与容差评测前先做正确性检查有三种方式ref_prog提供参照程序接收相同输入并对结果做检查可以返回布尔值也可在结果不匹配时抛异常manual_check_prog一个自定义可调用对象检查输出并在不匹配时抛异常可与ref_prog组合使用见 tuner.py#L710-L716skip_checkTrue跳过正确性检查更快但需谨慎使用。数值漂移由以下容差参数控制参数默认值含义rtol1e-2相对误差容限atol1e-2绝对误差容限max_mismatched_ratio0.011%允许的最大元素不匹配比例校验实现走 profiler 的assert_allclose/manual_assert_closetuner.py#L710-L716。需要特别注意一旦设置了ref_prog、supply_prog或manual_check_prog任意一个回调调优器会放弃磁盘缓存——generate_cache_key直接返回None因为任意回调缺少可靠的可持久化身份可能闭包持有不可 pickle 的运行时状态宁可不复用缓存也不冒“不同输入或校验行为”的错用风险tuner.py#L321-L359。如果你的调优流程固定使用ref_prog这是需要知晓的缓存语义。配置空间设计与最佳实践可调维度Tile 尺寸block_M、block_N、block_K软件流水线num_stages每 block 线程数threads也可为(x, y)元组可选dtype 变体、epilogue 策略、小型调度开关通过保留键pass_configs还可以按配置覆盖编译器 pass 选项见仓库测试 test_pass_configs_tuning.py 对这一机制的验证。实用建议来自指南与源码行为印证从一个能跑通的基线出发先调一个小而有意义的空间再逐步扩大尊重硬件上限共享内存字节数、每线程寄存器、每 block 最大线程数把不可能成立的配置在生成阶段就排除掉block 尺寸尽量取向量宽度与 warp 尺寸的整数倍保证访存向量化与 MMA 对齐用set_autotune_inputs保证每个配置在同一份数据上被测量避免不同数据导致的计时偏差把最终的最优配置记录并固化为函数默认值稳定后不再反复搜索。配置空间还可以写成可调用对象依据问题规模动态裁剪搜索空间既保持针对性又保证配置合法性def matmul_configs(M, N, K): large min(M, N, K) 1024 tiles [128] if large else [64, 128] for BM in tiles: for BN in tiles: for BK in [32, 64]: for S in [2, 3]: for TH in [128, 256]: yield dict(block_MBM, block_NBN, block_KBK, num_stagesS, threadsTH)并行编译、基准测试与超时机制调优器用线程池并行编译各配置并对每个配置施加独立超时。CUDA 目标下每个编译 worker 线程在入口处执行torch.cuda.set_device绑定到当前设备避免多上下文串扰tuner.py#L467-L472。worker 数量的决定逻辑在_resolve_num_compile_workerstuner.py#L428-L447TILELANG_AUTO_TUNING_CPU_COUNTS 0时直接取min(该值, 可用 CPU 数)否则按TILELANG_AUTO_TUNING_CPU_UTILITIES默认 0.9× 可用 CPU 数计算最后用TILELANG_AUTO_TUNING_MAX_CPU_COUNT默认 -1即不限封顶。关于超时的实现细节值得注意timeout是通过把每次基准调用放入一个守护子线程并join(timeout...)来可移植地强制的tuner.py#L582-L607。由于原生的 / CUDA 的调用无法被强行中断超时的守护线程会在底层调用返回前继续存在——这意味着超时只保证主流程不阻塞而不保证 GPU 侧工作立即停止。若配置timeout 0启动时会打印一条说明该机制的警告日志。所有调优日志同时写入工作目录下的autotuner.log文件与标准输出文件 handler 为 DEBUG 级、控制台为 INFO 级见 tuner.py#L66-L79。单配置失败编译失败、超时、评测错误只记日志并跳过不会中断整体搜索只有当没有任何配置成功时才会抛出RuntimeError: Auto-tuning failed: ...tuner.py#L1137-L1140。缓存机制缓存键、磁盘文件与环境变量自动调优器把最优产物缓存到两层进程内内存每进程一张_memory_cache表命中时会提示优先使用tilelang.autotune装饰器以便跨实参元组复用tuner.py#L873-L892以及磁盘目录$TILELANG_CACHE_DIR/autotuner之下按缓存键分目录。缓存键的构成generate_cache_key用 SHA-256 对以下内容取哈希tuner.py#L321-L359TileLang 版本号函数默认参数值、闭包中的自由变量如被外层函数捕获的M/N/K常量必须可序列化函数源码inspect.getsource(self.fn)完整配置列表configsCompileArgs哈希含out_idx、execution_backend、target、target_host、verbose、规范化后的pass_configsparam.py#L79-L95ProfileArgs哈希含warmup/rep/timeout/backend、supply_type、rtol/atol/max_mismatched_ratio/skip_check/cache_input_tensorsparam.py#L136-L151以及如前所述存在自定义回调时整体返回None禁用磁盘缓存。磁盘缓存文件清单每个缓存键目录下的文件命名与kernel_cache保持一致常量定义见 param.py#L31-L43最优配置与延迟best_config.json、latency.json含latency与ref_latency两个字段、out_idx.jsonkernel 源码与库device_kernel.cu、host_kernel.cu、kernel_lib.so具体文件名依后端而异见下函数与参数function.pklcloudpickle 序列化的函数、params.pklkernel 参数。不同执行后端的库产物文件名不同param.py#L566-L575执行后端库产物nvrtckernel.cubinkernel.pyPython launcher加载时必须两者齐备tvm_ffiexecutable.socutedslkernel.pyCuTeDSL 的“库”是 Python 源码另有可选的 launcher.so与kernel.cubin其他如cythonkernel_lib.sosave_to_disk采用“临时 staging 目录 fsync 原子 rename”的方式落盘避免并发读者看到半写入的缓存条目param.py#L406-L495load_from_disk会按后端校验必需文件齐备后才通过JITKernel.from_database重建 kernelparam.py#L497-L564。原子保存行为有专门的回归测试 test_tilelang_autotune_atomic_save.py缓存键行为见 test_autotune_cache_key.py。控制缓存与环境变量以下环境变量集中定义在 tilelang/env.pyAuto-tuning 段见 env.py#L405-L409环境变量默认值作用TILELANG_CACHE_DIR~/.tilelang/cacheTileLang 缓存根目录自动调优磁盘缓存位于其autotuner子目录env.py#L368TILELANG_DISABLE_CACHE0置1全局禁用所有 kernel 缓存优先级最高TILELANG_AUTO_TUNING_DISABLE_CACHE0仅禁用自动调优磁盘缓存TILELANG_AUTO_TUNING_CPU_UTILITIES0.9编译 worker 占用的 CPU 比例TILELANG_AUTO_TUNING_CPU_COUNTS-1显式指定 CPU 数-1表示自动TILELANG_AUTO_TUNING_MAX_CPU_COUNT-1worker 数上限-1表示不限后端相关说明NVRTC 后端会把.cubin与 Python launcher 一起持久化Torch 后端不会把编译产物写到磁盘——源码中会打印Torch backend does not support cache saving to disk.警告tuner.py#L1157-L1158此时只有进程内缓存生效。替代方案par_compile手工扫描如果希望完全自己掌控基准测试流程可以用JITImpl.par_compile批量编译一组配置然后自行计时。它在 tilelang/jit/init.py#L403-L450 中定义tilelang.jit def factory(M, N, K, block_M128, block_N128, block_K32): T.prim_func def k(A: T.Tensor((M, K), float16), B: T.Tensor((K, N), float16), C: T.Tensor((M, N), float16)): ... return k impl factory # JITImpl cfgs [ dict(block_M64, block_N128, block_K32), dict(block_M128, block_N128, block_K64), ] kernels impl.par_compile(cfgs, num_workers4) # 之后自行 benchmarkkernelsipar_compile支持ignore_errorTrue单配置编译失败时记警告并置None而不是抛异常num_workersNone时由系统决定并发度。相比AutoTuner.run()这条路没有自动校验与缓存适合需要自定义计时方案例如固定 CUDA graph、多形状混合负载的场景。保存与复用调优结果AutotuneResult编程式路径返回的AutotuneResult可以整体落盘并在之后重新加载适用于 CI、多主机工作流或“带着调好的配置发货”artifact tuner.run() # AutotuneResult # 保存到磁盘 from pathlib import Path save_dir Path(out/best/matmul_1024) artifact.save_to_disk(save_dir, verboseTrue) # 之后重载 from tilelang.autotuner.param import AutotuneResult, CompileArgs restored AutotuneResult.load_from_disk(save_dir, CompileArgs()) best restored.kernel best(A, B, C)AutotuneResult携带latency、config、ref_latency、libcodekernel 源码、funcPrimFunc与kernel已编译 kernel六个字段param.py#L154-L172保存的目录中同时含人类可读的 JSONbest config / latency与 kernel 源码便于审计。注意 Torch 执行后端不持久化编译二进制重载时需要重新编译或改用其他后端。设备与后端选择通过set_compile_args显式指定编译期选项targetauto | cuda | hip | metal归一化为 TVM Target也支持{kind: cuda, arch: sm_90}这样的目标字典execution_backendauto | tvm_ffi | cython | nvrtc | torchpass_configs{...}在实验性场景下切换 TileLang/TVM pass。这三个参数均支持环境变量缺省未显式传入时target读TILELANG_DEFAULT_TARGET默认auto、execution_backend读TILELANG_EXECUTION_BACKEND默认auto、verbose读TILELANG_VERBOSEtuner.py#L156-L209、env.py#L411-L415。在多 GPU 的 CUDA 机器上编译 worker 与基准 worker 都会逐线程绑定 CUDA 设备编译端torch.cuda.set_devicetuner.py#L467-L472基准端_benchmark_worker_loop进入时按worker_device设设备tuner.py#L544-L551避免上下文混用。常见问题排查“No configurations to tune”确认configs是非空列表或返回非空列表的 callable空配置空间会抛ValueErrortuner.py#L917-L918。超时增大timeout确认输入张量放得进设备显存检查参照校验本身是否成为瓶颈。动态形状用set_autotune_inputs或自定义supply_prog提供具体输入内置生成器不支持符号维度。磁盘缓存没生效检查TILELANG_DISABLE_CACHE/TILELANG_AUTO_TUNING_DISABLE_CACHE注意执行后端Torch 后端不落盘以及是否传入了ref_prog等回调回调会使缓存键为None。结果全部失败查看工作目录下的autotuner.log其中记录了每个失败配置的 DEBUG 级错误与 traceback。调优器相关的完整行为覆盖在 testing/python/autotune/ 测试目录中例如超时机制test_tilelang_autotune_timeout.py、输入捕获test_tilelang_autotune_with_inputs.py、标量输入校验test_tilelang_autotune_scalar_inputs.py与 eager 模式test_tilelang_autotune_eager_mode.py可作为各特性的可执行参考。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考