TileLang FP16 矩阵乘法自动调优基准测试搜索空间设计、内核剖析与 TFLOPs 结果复现【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang本篇基于benchmark/matmul/README.md完整讲解 TileLang 仓库中 FP16 大矩阵乘法Matmul基准测试的设计与复现方法从benchmark_matmul.py的自动调优autotune搜索空间、块级 GEMM 内核结构到autotunejit装饰器的底层执行机制最终给出与 README 一致的 TFLOPs 数据表及逐列解读。读完本文你可以独立复现该基准、修改搜索空间并理解各调优参数block 尺寸、流水线级数、rasterization 开关对内核性能的影响路径。一、基准测试目标与环境benchmark/matmul/README.md记录了benchmark_matmul.py在M N 8192、不同K维度下使用默认自动调优搜索空间测得的 FP16 矩阵乘法吞吐。测试环境为仓库提交17bd0a6c651f599bec1397e0b91830c3ddc93076GPUNVIDIA H800 SXM驱动560.35.05计算口径为C A B.T其中 A、B 均为float16累加使用float32以保证数值精度。吞吐量按2*M*N*K / latency计算浮点运算量为两次乘加各一次故系数为 2。二、内核定义benchmark_matmul.py 中的 matmul基准内核位于 benchmark_matmul.py核心是一个标准的“分块 流水线 TensorCore”GEMM 结构。关键代码如下源码注释已保留autotune( configsget_configs, warmup3, rep20, ) jit( out_idx[2], ) def matmul( M, N, K, with_roller, block_MNone, block_NNone, block_KNone, num_stagesNone, thread_numNone, policyNone, enable_rasterationNone, ): dtype T.float16 accum_dtype T.float32 T.prim_func def main( A: T.Tensor((M, K), dtype), B: T.Tensor((N, K), dtype), C: T.Tensor((M, N), dtype), ): # x 维绑定 N 方向块索引y 维绑定 M 方向块索引 with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threadsthread_num) as (bx, by): A_shared T.alloc_shared((block_M, block_K), dtype) B_shared T.alloc_shared((block_N, block_K), dtype) C_local T.alloc_fragment((block_M, block_N), accum_dtype) # 寄存器片段fp32 累加 C_shared T.alloc_shared((block_M, block_N), dtype) T.use_swizzle(panel_size10, enableenable_rasteration) # rasterization 开关 T.clear(C_local) # K 维分块循环由 num_stages 控制软件流水线级数 for k in T.Pipelined(T.ceildiv(K, block_K), num_stagesnum_stages): T.copy(A[by * block_M, k * block_K], A_shared) T.copy(B[bx * block_N, k * block_K], B_shared) T.gemm(A_shared, B_shared, C_local, transpose_BTrue, policypolicy) # 结果经 shared memory 中转写回全局内存 T.copy(C_local, C_shared) T.copy(C_shared, C[by * block_M, bx * block_N]) return main各要素的仓库内实现位置T.Kernel线程块网格划分网格尺寸为ceildiv(N, block_N) × ceildiv(M, block_M)threads即每个块的线程数T.copy全局内存到 shared memory 的块拷贝定义于 copy_op.pyT.PipelinedK 维分块循环并插入num_stages级流水线定义于 loop.pynum_stages决定异步预取重叠程度T.gemm块级矩阵乘transpose_BTrue表明 B 以(N, K)布局参与计算实现见 gemm_op.pyT.use_swizzle(panel_size10, ...)控制块网格的 rasterization/swizzle 排列以改善 L2 命中率定义于 annotations.py。注意配置字典中的键名是enable_rasteration源码即此拼写。2.1 默认搜索空间with_rollerFalseget_configs在不开启 Roller 时用itertools.product枚举如下笛卡尔积见 benchmark_matmul.py L112-L121参数候选值含义block_M64, 128, 256M 方向块尺寸block_N64, 128, 256N 方向块尺寸block_K32, 64K 方向块尺寸num_stages0, 1, 2, 3流水线级数0 表示不流水线thread_num128, 256每线程块线程数policyT.GemmWarpPolicy.Squarewarp 在块内的划分策略enable_rasterationTrue, False是否启用 rasterization/swizzle合计 3×3×2×4×2×1×2 288 个候选配置每个配置都会被完整编译、跑 warmup 与计时最后选出延迟最低的best_config。2.2 Roller 模式with_rollerTrue开启--with_roller后get_configs不再手搓笛卡尔积而是调用 TileLang 的 carver 模块按架构推导调度提示见 benchmark_matmul.py L57-L110以MatmulTemplate(M, N, K, in_dtypeT.float16, out_dtypeT.float16, accum_dtypeT.float32)描述问题模板位于 matmul.py并对 CUDA 目标使用CUDA(cuda)对 HIP/RDNA 目标用auto_infer_current_arch()调用carve_template.recommend_hints(topk10)得到最多 10 条调度提示每条 hint 被转成一个配置字典hint.block→block_M/block_Nhint.rstep[0]→block_Khint.pipeline_stage→num_stages若 hint 带warp划分则按block_rows block_M // warp_m、block_cols block_N // warp_n计算thread_num并生成对应的T.GemmWarpPolicy否则回退为Square策略hint.rasterization_plan is not NoRasterization决定enable_rasteration。因此两种模式给出的是不同粒度的搜索空间默认模式是固定的 288 组粗粒度枚举Roller 模式是面向当前架构的 top-10 推荐配置。三、自动调优执行机制autotune jit 装饰器链matmul函数由两层装饰器组成autotune(configsget_configs, warmup3, rep20) jit(out_idx[2])jit(out_idx[2])声明第 3 个张量C索引 2为输出其余A、B由调优器自动供给随机输入autotune(configs..., warmup3, rep20)基准脚本将 warmup 覆盖为 3、rep 覆盖为 20autotune的公开默认值是warmup25, rep100, timeout100见 tuner.py L1318-L1338即每个候选配置计时前跑 3 次预热、取 20 次重复的延迟。调用matmul(M, N, K, with_roller)时实际发生的事实现位于 tuner.py 的AutoTuneImpl.__call__与AutoTuner.run缓存检查generate_cache_key以 tilelang 版本、函数源码、闭包自由变量、配置列表、编译参数CompileArgs的哈希与计时参数ProfileArgs的哈希组合生成 SHA-256 键tuner.py L321-L359。若命中内存或磁盘缓存缓存目录位于 KernelCache 命名空间下的autotuner子目录落盘结构见 param.py L406-L495则直接返回上次结果避免重复调优注意当提供ref_prog/supply_prog等回调时缓存会被禁用并行编译配置被拆成编译单元提交到ThreadPoolExecutor工作线程数由TILELANG_AUTO_TUNING_CPU_UTILITIES/TILELANG_AUTO_TUNING_CPU_COUNTS/TILELANG_AUTO_TUNING_MAX_CPU_COUNT等环境变量控制tuner.py L428-L447CUDA tvm_ffi后端下还支持enable_grouped_compile分组编译基准计时编译完成的内核进入队列由 benchmark worker 通过 profiler 的do_bench(n_warmupwarmup, n_repeatrep, backendevent)计时CUDA event 后端timeout超时的配置会被记录并跳过返回结果最终返回AutotuneResult字段包含latency最优延迟、config最优配置字典、ref_latency参考实现延迟若提供了ref_prog、func与kernel定义见 param.py L154-L172。一个实用技巧源自autotune的 docstring如果不想每次调优可以直接把调优参数显式传入来“锁定”某组配置例如matmul(M, N, K, False, block_M128, block_N128, num_stages2)此时调优器检测到可调参数已被指定会跳过搜索直接 JIT 编译tuner.py L935-L956。本基准脚本的matmul未传入ref_prog脚本顶层虽有ref_program函数但默认调优路径未挂载因此ref_latency为Noneif __name__ __main__中的 “Reference TFlops” 一行只在存在参考延迟时才打印。四、复现步骤4.1 按 README 原样复现M N 8192多组 K在仓库根目录执行这是 benchmark/matmul/README.md 给出的原始命令cd benchmark/matmul python - PY from benchmark_matmul import matmul M 8192 N 8192 for K in [256, 512, 1024, 2048, 4096, 8192, 16384]: res matmul(M, N, K, False) tflops 2 * M * N * K / res.latency * 1e-12 print(fK{K:5d} latency{res.latency:.6f}s TFlops{tflops:.3f}) PY其中2 * M * N * K是浮点运算量res.latency是调优后的最优延迟秒乘1e-12换算为 TFLOPs。首次运行会经历完整的 288 配置编译与计时过程每组 K 各一次共 7 次调优耗时较长后续同机同版本重跑可命中磁盘缓存显著加速。4.2 命令行方式benchmark_matmul.py自带 CLI支持任意维度与 Roller 开关python benchmark_matmul/benchmark_matmul.py --m 8192 --n 8192 --k 8192 # 默认 MNK16384 python benchmark_matmul/benchmark_matmul.py --m 8192 --n 8192 --k 16384 --with_roller参数--m / --n / --k设置矩阵维度默认均为 16384--with_roller启用 BitBLAS Roller 推导搜索空间。运行输出包括Best latency (s): ... Best TFlops: ... Best config: {...} # 最优 block_M/block_N/block_K/num_stages/thread_num/policy/enable_rasteration五、结果表与解读README 记录的原始数据NVIDIA H800 SXMFP16M N 8192KLatency (s)Throughput (TFLOPs)2560.0890563865120.13206452010240.21881662820480.39011270540960.74675273681921.449888758163842.871168766可以读出三点规律小 K 受写回带宽限制K256 时 M×N 结果矩阵8192×8192×2B ≈ 128 MiB的写回量与计算量相比过大算术强度低吞吐仅 386 TFLOPs此时内核行为接近带宽受限而非算力受限。吞吐随 K 单调爬升并趋于饱和从 K256 到 K16384TFLOPs 从 386 升至 766且 K≥4096 后增量收窄736 → 758 → 766说明大 K 下内核逐渐进入算力主导区间吞吐逼近该硬件上该实现的稳定上限。延迟近似随 K 线性增长K 从 256 翻倍到 16384延迟从 0.089 s 增至 2.87 s符合 GEMM 计算量 ∝ K 的预期。需要强调的是表中数字绑定具体硬件H800 SXM 驱动 560.35.05与具体提交17bd0a6c在你的 GPU 上复现时应以本机实测为准跨硬件对比没有意义。六、延伸阅读同目录的稀疏 Matmul 基准benchmark/matmul目录下另有两个稀疏基准脚本可视为本文方法的扩展benchmark_matmul_sp.py结构-稀疏2:4 式压缩 A 为A_sparse: (M, K//2) 掩码元数据E: (M, K//e_factor)使用T.gemm_sp(A_shared, E_shared, B_shared, C_local, ...)调用稀疏 TensorCore 指令支持--accum_dtypefloat/float16、--e_dtypeint16/int8/int32以及--bench_torch_sparse cutlass|cusparselt与 PyTorch 稀疏实现对照仅限 sm80benchmark_matmul_sp_compress.py稀疏压缩格式的对应基准。两者与稠密版共用同一套autotune jit模式与参数命名阅读顺序建议先本文的稠密基准再看稀疏版中e_factor与T.gemm_sp的差异。七、关键文件索引内容路径基准结果记录本文主体文档benchmark/matmul/README.mdFP16 稠密 Matmul 基准内核benchmark_matmul.py稀疏 Matmul 基准benchmark_matmul_sp.py自动调优器实现AutoTuner / AutoTuneImpl / autotune 装饰器tilelang/autotuner/tuner.py调优结果与磁盘缓存结构AutotuneResulttilelang/autotuner/param.pyMatmul 调度模板Roller 模式tilelang/carver/template/matmul.py官方自动调优编程指南docs/programming_guides/autotuning.mdGEMM 示例与更复杂的调优用法examples/gemm/【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 TileLang FP16 矩阵乘法自动调优基准测试：搜索空间设计、内核剖析与 TFLOPs 结果复现
TileLang FP16 矩阵乘法自动调优基准测试搜索空间设计、内核剖析与 TFLOPs 结果复现【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang本篇基于benchmark/matmul/README.md完整讲解 TileLang 仓库中 FP16 大矩阵乘法Matmul基准测试的设计与复现方法从benchmark_matmul.py的自动调优autotune搜索空间、块级 GEMM 内核结构到autotunejit装饰器的底层执行机制最终给出与 README 一致的 TFLOPs 数据表及逐列解读。读完本文你可以独立复现该基准、修改搜索空间并理解各调优参数block 尺寸、流水线级数、rasterization 开关对内核性能的影响路径。一、基准测试目标与环境benchmark/matmul/README.md记录了benchmark_matmul.py在M N 8192、不同K维度下使用默认自动调优搜索空间测得的 FP16 矩阵乘法吞吐。测试环境为仓库提交17bd0a6c651f599bec1397e0b91830c3ddc93076GPUNVIDIA H800 SXM驱动560.35.05计算口径为C A B.T其中 A、B 均为float16累加使用float32以保证数值精度。吞吐量按2*M*N*K / latency计算浮点运算量为两次乘加各一次故系数为 2。二、内核定义benchmark_matmul.py 中的 matmul基准内核位于 benchmark_matmul.py核心是一个标准的“分块 流水线 TensorCore”GEMM 结构。关键代码如下源码注释已保留autotune( configsget_configs, warmup3, rep20, ) jit( out_idx[2], ) def matmul( M, N, K, with_roller, block_MNone, block_NNone, block_KNone, num_stagesNone, thread_numNone, policyNone, enable_rasterationNone, ): dtype T.float16 accum_dtype T.float32 T.prim_func def main( A: T.Tensor((M, K), dtype), B: T.Tensor((N, K), dtype), C: T.Tensor((M, N), dtype), ): # x 维绑定 N 方向块索引y 维绑定 M 方向块索引 with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threadsthread_num) as (bx, by): A_shared T.alloc_shared((block_M, block_K), dtype) B_shared T.alloc_shared((block_N, block_K), dtype) C_local T.alloc_fragment((block_M, block_N), accum_dtype) # 寄存器片段fp32 累加 C_shared T.alloc_shared((block_M, block_N), dtype) T.use_swizzle(panel_size10, enableenable_rasteration) # rasterization 开关 T.clear(C_local) # K 维分块循环由 num_stages 控制软件流水线级数 for k in T.Pipelined(T.ceildiv(K, block_K), num_stagesnum_stages): T.copy(A[by * block_M, k * block_K], A_shared) T.copy(B[bx * block_N, k * block_K], B_shared) T.gemm(A_shared, B_shared, C_local, transpose_BTrue, policypolicy) # 结果经 shared memory 中转写回全局内存 T.copy(C_local, C_shared) T.copy(C_shared, C[by * block_M, bx * block_N]) return main各要素的仓库内实现位置T.Kernel线程块网格划分网格尺寸为ceildiv(N, block_N) × ceildiv(M, block_M)threads即每个块的线程数T.copy全局内存到 shared memory 的块拷贝定义于 copy_op.pyT.PipelinedK 维分块循环并插入num_stages级流水线定义于 loop.pynum_stages决定异步预取重叠程度T.gemm块级矩阵乘transpose_BTrue表明 B 以(N, K)布局参与计算实现见 gemm_op.pyT.use_swizzle(panel_size10, ...)控制块网格的 rasterization/swizzle 排列以改善 L2 命中率定义于 annotations.py。注意配置字典中的键名是enable_rasteration源码即此拼写。2.1 默认搜索空间with_rollerFalseget_configs在不开启 Roller 时用itertools.product枚举如下笛卡尔积见 benchmark_matmul.py L112-L121参数候选值含义block_M64, 128, 256M 方向块尺寸block_N64, 128, 256N 方向块尺寸block_K32, 64K 方向块尺寸num_stages0, 1, 2, 3流水线级数0 表示不流水线thread_num128, 256每线程块线程数policyT.GemmWarpPolicy.Squarewarp 在块内的划分策略enable_rasterationTrue, False是否启用 rasterization/swizzle合计 3×3×2×4×2×1×2 288 个候选配置每个配置都会被完整编译、跑 warmup 与计时最后选出延迟最低的best_config。2.2 Roller 模式with_rollerTrue开启--with_roller后get_configs不再手搓笛卡尔积而是调用 TileLang 的 carver 模块按架构推导调度提示见 benchmark_matmul.py L57-L110以MatmulTemplate(M, N, K, in_dtypeT.float16, out_dtypeT.float16, accum_dtypeT.float32)描述问题模板位于 matmul.py并对 CUDA 目标使用CUDA(cuda)对 HIP/RDNA 目标用auto_infer_current_arch()调用carve_template.recommend_hints(topk10)得到最多 10 条调度提示每条 hint 被转成一个配置字典hint.block→block_M/block_Nhint.rstep[0]→block_Khint.pipeline_stage→num_stages若 hint 带warp划分则按block_rows block_M // warp_m、block_cols block_N // warp_n计算thread_num并生成对应的T.GemmWarpPolicy否则回退为Square策略hint.rasterization_plan is not NoRasterization决定enable_rasteration。因此两种模式给出的是不同粒度的搜索空间默认模式是固定的 288 组粗粒度枚举Roller 模式是面向当前架构的 top-10 推荐配置。三、自动调优执行机制autotune jit 装饰器链matmul函数由两层装饰器组成autotune(configsget_configs, warmup3, rep20) jit(out_idx[2])jit(out_idx[2])声明第 3 个张量C索引 2为输出其余A、B由调优器自动供给随机输入autotune(configs..., warmup3, rep20)基准脚本将 warmup 覆盖为 3、rep 覆盖为 20autotune的公开默认值是warmup25, rep100, timeout100见 tuner.py L1318-L1338即每个候选配置计时前跑 3 次预热、取 20 次重复的延迟。调用matmul(M, N, K, with_roller)时实际发生的事实现位于 tuner.py 的AutoTuneImpl.__call__与AutoTuner.run缓存检查generate_cache_key以 tilelang 版本、函数源码、闭包自由变量、配置列表、编译参数CompileArgs的哈希与计时参数ProfileArgs的哈希组合生成 SHA-256 键tuner.py L321-L359。若命中内存或磁盘缓存缓存目录位于 KernelCache 命名空间下的autotuner子目录落盘结构见 param.py L406-L495则直接返回上次结果避免重复调优注意当提供ref_prog/supply_prog等回调时缓存会被禁用并行编译配置被拆成编译单元提交到ThreadPoolExecutor工作线程数由TILELANG_AUTO_TUNING_CPU_UTILITIES/TILELANG_AUTO_TUNING_CPU_COUNTS/TILELANG_AUTO_TUNING_MAX_CPU_COUNT等环境变量控制tuner.py L428-L447CUDA tvm_ffi后端下还支持enable_grouped_compile分组编译基准计时编译完成的内核进入队列由 benchmark worker 通过 profiler 的do_bench(n_warmupwarmup, n_repeatrep, backendevent)计时CUDA event 后端timeout超时的配置会被记录并跳过返回结果最终返回AutotuneResult字段包含latency最优延迟、config最优配置字典、ref_latency参考实现延迟若提供了ref_prog、func与kernel定义见 param.py L154-L172。一个实用技巧源自autotune的 docstring如果不想每次调优可以直接把调优参数显式传入来“锁定”某组配置例如matmul(M, N, K, False, block_M128, block_N128, num_stages2)此时调优器检测到可调参数已被指定会跳过搜索直接 JIT 编译tuner.py L935-L956。本基准脚本的matmul未传入ref_prog脚本顶层虽有ref_program函数但默认调优路径未挂载因此ref_latency为Noneif __name__ __main__中的 “Reference TFlops” 一行只在存在参考延迟时才打印。四、复现步骤4.1 按 README 原样复现M N 8192多组 K在仓库根目录执行这是 benchmark/matmul/README.md 给出的原始命令cd benchmark/matmul python - PY from benchmark_matmul import matmul M 8192 N 8192 for K in [256, 512, 1024, 2048, 4096, 8192, 16384]: res matmul(M, N, K, False) tflops 2 * M * N * K / res.latency * 1e-12 print(fK{K:5d} latency{res.latency:.6f}s TFlops{tflops:.3f}) PY其中2 * M * N * K是浮点运算量res.latency是调优后的最优延迟秒乘1e-12换算为 TFLOPs。首次运行会经历完整的 288 配置编译与计时过程每组 K 各一次共 7 次调优耗时较长后续同机同版本重跑可命中磁盘缓存显著加速。4.2 命令行方式benchmark_matmul.py自带 CLI支持任意维度与 Roller 开关python benchmark_matmul/benchmark_matmul.py --m 8192 --n 8192 --k 8192 # 默认 MNK16384 python benchmark_matmul/benchmark_matmul.py --m 8192 --n 8192 --k 16384 --with_roller参数--m / --n / --k设置矩阵维度默认均为 16384--with_roller启用 BitBLAS Roller 推导搜索空间。运行输出包括Best latency (s): ... Best TFlops: ... Best config: {...} # 最优 block_M/block_N/block_K/num_stages/thread_num/policy/enable_rasteration五、结果表与解读README 记录的原始数据NVIDIA H800 SXMFP16M N 8192KLatency (s)Throughput (TFLOPs)2560.0890563865120.13206452010240.21881662820480.39011270540960.74675273681921.449888758163842.871168766可以读出三点规律小 K 受写回带宽限制K256 时 M×N 结果矩阵8192×8192×2B ≈ 128 MiB的写回量与计算量相比过大算术强度低吞吐仅 386 TFLOPs此时内核行为接近带宽受限而非算力受限。吞吐随 K 单调爬升并趋于饱和从 K256 到 K16384TFLOPs 从 386 升至 766且 K≥4096 后增量收窄736 → 758 → 766说明大 K 下内核逐渐进入算力主导区间吞吐逼近该硬件上该实现的稳定上限。延迟近似随 K 线性增长K 从 256 翻倍到 16384延迟从 0.089 s 增至 2.87 s符合 GEMM 计算量 ∝ K 的预期。需要强调的是表中数字绑定具体硬件H800 SXM 驱动 560.35.05与具体提交17bd0a6c在你的 GPU 上复现时应以本机实测为准跨硬件对比没有意义。六、延伸阅读同目录的稀疏 Matmul 基准benchmark/matmul目录下另有两个稀疏基准脚本可视为本文方法的扩展benchmark_matmul_sp.py结构-稀疏2:4 式压缩 A 为A_sparse: (M, K//2) 掩码元数据E: (M, K//e_factor)使用T.gemm_sp(A_shared, E_shared, B_shared, C_local, ...)调用稀疏 TensorCore 指令支持--accum_dtypefloat/float16、--e_dtypeint16/int8/int32以及--bench_torch_sparse cutlass|cusparselt与 PyTorch 稀疏实现对照仅限 sm80benchmark_matmul_sp_compress.py稀疏压缩格式的对应基准。两者与稠密版共用同一套autotune jit模式与参数命名阅读顺序建议先本文的稠密基准再看稀疏版中e_factor与T.gemm_sp的差异。七、关键文件索引内容路径基准结果记录本文主体文档benchmark/matmul/README.mdFP16 稠密 Matmul 基准内核benchmark_matmul.py稀疏 Matmul 基准benchmark_matmul_sp.py自动调优器实现AutoTuner / AutoTuneImpl / autotune 装饰器tilelang/autotuner/tuner.py调优结果与磁盘缓存结构AutotuneResulttilelang/autotuner/param.pyMatmul 调度模板Roller 模式tilelang/carver/template/matmul.py官方自动调优编程指南docs/programming_guides/autotuning.mdGEMM 示例与更复杂的调优用法examples/gemm/【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考