Triton 自动调优写 3 组配置、指 1 个 key省掉手动调参的一整天【免费下载链接】tritonDevelopment repository for the Triton language and compiler项目地址: https://gitcode.com/GitHub_Trending/tri/tritonTriton 的 autotune自动调优能替你自动完成 GPU 内核参数调优给出几组候选配置它首次运行时逐组实测、挑出最快的一组并缓存复用省掉你手动改参数、重编译、量时间的反复。矩阵乘法里BLOCK_SIZE 到底填多少矩阵乘法的 BLOCK_SIZE 没有标准答案填 64、128 还是 256取决于矩阵尺寸和你用的 GPU人工几乎无法穷举。写一个 matmul 内核通常要定 BLOCK_SIZE_M、BLOCK_SIZE_N、BLOCK_SIZE_K 三个分块大小再配 num_warps每块用多少 warp1 warp32 线程和 num_stages循环软件流水级数两个编译选项。就算每个参数只给 3 个候选组合数也轻松破百。你不可能对每一种都手动改代码、重新编译、再逐组量一遍耗时。这块组合爆炸的活正好交给 autotune。原理速览autotune 就像一次试吃autotune 的原理一句话就能说清把每组配置都真跑一遍、量出耗时留下最快的那组之后复用。打个比方像厨师调一道菜的口味。他不会一次定死盐量而是先做 3 个小份——偏咸、偏淡、正好——让你尝记下哪份最合口味之后就一直按那个比例出菜。autotune 做的正是这件事对应关系是configs你列出的几份小样即不同 BLOCK_SIZE、num_warps 的组合首次调用autotune 逐份试吃每组配置在 GPU 上真实执行一次并计时记结果选出耗时最短的那组写进内存缓存再次调用同样的输入直接复用最优配置不再重复实测。上图就是 autotune 在试吃的东西同一块数据换一种分块BLOCK_SIZEGPU 上的计算量、访存方式都不同快慢自然也不同。它不靠猜全靠实测挑最快。✅ 最小上手示例一个能跑的求和内核给 triton.jit 函数套一层 triton.autotune 装饰器传 configs 和 key 两个参数就能自动调优。import triton import triton.language as tl configs [ triton.Config({BLOCK_SIZE: 128}, num_warps4), triton.Config({BLOCK_SIZE: 256}, num_warps8), triton.Config({BLOCK_SIZE: 512}, num_warps8), ] triton.autotune(configsconfigs, key[x_size]) triton.jit def sum_kernel(x_ptr, x_size, BLOCK_SIZE: tl.constexpr): pid tl.program_id(0) offs pid * BLOCK_SIZE tl.arange(0, BLOCK_SIZE) return tl.sum(tl.load(x_ptr offs, maskoffs x_size))两个关键参数configs一个 triton.Config 列表每个 Config 描述一组候选参数。字典里的键如 BLOCK_SIZE会作为常量在编译时固定下来并传进内核num_warps、num_stages 是编译选项。上面 3 组配置就是 autotune 要试吃的 3 份小样。key一组触发重新调优的参数名。key[x_size] 表示只有 x_size 变了才重跑 3 组配置x_size 不变就直接复用。矩阵乘法一般写 key[M, N, K]即矩阵尺寸变了才重新调优。 调优提速缓存、剪枝与看日志调优省时间有三个抓手按需取用。按输入尺寸自动缓存最优配置同一份 key 的输入autotune 只在内存里调一次后续调用直接复用。想让跨进程、跨运行也复用把装饰器参数改成 cache_resultsTrue或设环境变量 TRITON_CACHE_AUTOTUNING1各组耗时会写进磁盘缓存下次启动也能直接命中。用 top_k 剪掉明显不优的配置配置很多时先剪枝再实测更省。prune_configs_by{top_k: 0.3} 表示先用内置性能模型估算每组耗时只保留最快的 30% 去实测。triton.autotune( configsconfigs, key[M, N, K], prune_configs_by{top_k: 0.3}, )打开调优日志看清每组配置的真实耗时设 TRITON_PRINT_AUTOTUNING1 就能打印调优过程。运行前先 export TRITON_PRINT_AUTOTUNING1控制台会输出每个 key 的调优总耗时和选中的最优配置方便你核对到底哪组最快、调优花了多久。一句话收尾autotune 不神秘列出几组 configs、指一个 key它替你试吃、留最快的一组合、之后复用。想继续深入官方 API 见 docs/python-api/triton.rst矩阵乘法的完整调优示例见 python/tutorials/03-matrix-multiplication.py。【免费下载链接】tritonDevelopment repository for the Triton language and compiler项目地址: https://gitcode.com/GitHub_Trending/tri/triton创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考