ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀，每一条都经过真实交付验证。

首页 / 资讯中心 / 正文

TileLang GEMV 优化实战：从朴素内核到超越 cuBLAS 的线程级调优路径

2026/9/16 11:21:23 拓冰建站 浏览量
TileLang GEMV 优化实战：从朴素内核到超越 cuBLAS 的线程级调优路径 TileLang GEMV 优化实战从朴素内核到超越 cuBLAS 的线程级调优路径【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelangGEMV矩阵-向量乘是 LLM 推理阶段的核心算子本文以 TileLang 官方 GEMV 教程为主线完整复现从 Triton 风格朴素实现出发、逐步引入 Split-K 并行、向量化访存、tvm_thread_allreduce归约直至autotune自动调优的四步优化路径并结合仓库源码解释每一步生成的 CUDA 代码变化与底层 pass 机制读完后你可以独立编写并调优一个内存带宽受限型向量算子内核。问题定位为什么 GEMV 需要线程级优化GEMV 可以看作一个退化的 GEMM 特例(1, k) * (k, n)的矩阵乘法输出是长度为 N 的向量。与 GEMM 不同GEMV 的算术强度每字节数据对应的浮点运算数极低瓶颈从计算单元转向内存带宽——这也是后续所有优化的出发点要么提高并发度把带宽吃满要么用更宽的访存事务减少加载次数要么用更高效的归约路径替换共享内存原子加。文档原文标注此内容为实验性质experimental示例代码位于 examples/gemv/example_gemv.py配套正确性测试为 examples/gemv/test_example_gemv.py性能回归脚本为 examples/gemv/regression_example_gemv.py。第一步Triton 朴素实现作为参照基线用 Triton 这类块级 DSL 写 GEMV 十分直接每个 program 负责输出向量的一行沿 K 维分块加载求内积triton.jit def _gemv_naive( x_ptr, A_ptr, y_ptr, N, K, BLOCK_SIZE_K: tl.constexpr, ): n tl.program_id(0) offs_k tl.arange(0, BLOCK_SIZE_K) mask offs_k K a_ptrs A_ptr n * K offs_k a_vals tl.load(a_ptrs, maskmask, other0.0) x_vals tl.load(x_ptr offs_k, maskmask, other0.0) dot tl.sum(a_vals * x_vals, axis0) tl.store(y_ptr n, dot)这种写法胜在简洁但块级抽象限制了线程级控制——你无法决定每个线程加载多少元素、如何组织归约、是否向量化。这正是 TileLang 的定位暴露更底层的编程控制面。第二步TileLang 朴素实现——GEMM 分块策略的直接套用有 CUDA C 基础的话最自然的起点是复用 GEMM 的 tiling 思路一个 block 负责 N 维的一段块内每个线程负责 C 的一个输出元素K 维分块串行累加def naive_gemv( N: int, K: int, BLOCK_N: int, BLOCK_K: int, dtype: str float16, accum_dtype: str float, ): T.prim_func def main( A: T.Buffer((K,), dtype), B: T.Buffer((N, K), dtype), C: T.Buffer((N,), dtype), ): with T.Kernel(T.ceildiv(N, BLOCK_N)) as bn: tn T.get_thread_binding(0) # tn threadIdx.x A_shared T.alloc_shared((BLOCK_K,), dtype) B_shared T.alloc_shared((BLOCK_N, BLOCK_K), dtype) C_reg T.alloc_local((1,), accum_dtype) T.clear(C_reg) for bk in T.serial(T.ceildiv(K, BLOCK_K)): for tk in T.serial(BLOCK_K): A_shared[tk] A[bk * BLOCK_K tk] B_shared[tn, tk] B[bn * BLOCK_N tn, bk * BLOCK_K tk] for tk in T.serial(BLOCK_K): C_reg[0] A_shared[tk].astype(accum_dtype) * B_shared[tn, tk].astype(accum_dtype) C[bn * BLOCK_N tn] C_reg[0] return main编译后缓存在~/.tilelang/cache生成的 CUDA 内核大致如下extern C __global__ void __launch_bounds__(256, 1) main_kernel(half_t* __restrict__ A, half_t* __restrict__ B, half_t* __restrict__ C) { extern __shared__ __align__(1024) uchar buf_dyn_shmem[]; float C_reg[1]; __shared__ uint64_t _mbarrier[2]; if (((int)threadIdx.x) 0) { tl::mbarrier_init(_mbarrier[0], 128); tl::mbarrier_init(_mbarrier[1], 128); } __syncthreads(); if (128 ((int)threadIdx.x)) { tl::warpgroup_reg_dealloc24(); for (int bk 0; bk 8; bk) { tl::mbarrier_wait(_mbarrier[1], ((bk 1) ^ 1)); for (int tk 0; tk 128; tk) { ((half_t*)buf_dyn_shmem)[tk] A[((bk * 128) tk)]; ((half_t*)buf_dyn_shmem)[(((((int)threadIdx.x) * 128) tk) - 16256)] B[...]; } tl::fence_proxy_async(); tl::mbarrier_cp_async_arrive(_mbarrier[0]); tl::mbarrier_arrive(_mbarrier[0]); } } else { tl::warpgroup_reg_alloc240(); C_reg[0] 0.000000e00f; for (int bk_1 0; bk_1 8; bk_1) { tl::mbarrier_wait(_mbarrier[0], (bk_1 1)); for (int tk_1 0; tk_1 128; tk_1) { C_reg[0] (C_reg[0] (((float)((half_t*)buf_dyn_shmem)[tk_1]) * ((float)((half_t*)buf_dyn_shmem)[(((((int)threadIdx.x) * 128) tk_1) 128)]))); } tl::fence_proxy_async(); tl::mbarrier_arrive(_mbarrier[1]); } C[((((int)blockIdx.x) * 128) ((int)threadIdx.x))] ((half_t)C_reg[0]); } }从源码结构看编译器把 256 个线程自动划分成了两个角色前 128 个线程threadIdx.x 128的 else 分支之外充当数据生产者把 A、B 异步搬运到共享内存后 128 个线程是消费者等待 mbarrier 信号后从共享内存累加。warpgroup_reg_dealloc/alloc是寄存器再平衡tl::fence_proxy_async()保证异步拷贝对计算代理可见。这个版本的代价很高文档实测约0.17 ms而 torch/cuBLAS 只需约 0.008 ms慢了近 20 倍。原因很简单——每个 block 只有一个消费者维度在算K 维完全串行大量线程空转等待。第三步引入 K 维并行Split-K与原子加提高并发的关键手段是把 K 维也切给线程。每个线程只负责一部分 K 的局部累加再用原子加合并到共享内存def naive_splitk_gemv( N: int, K: int, BLOCK_N: int, BLOCK_K: int, dtype: str float16, accum_dtype: str float, ): T.prim_func def main( A: T.Buffer((K,), dtype), B: T.Buffer((N, K), dtype), C: T.Buffer((N,), dtype), ): with T.Kernel(T.ceildiv(N, BLOCK_N), threads(BLOCK_N, BLOCK_K)) as bn: tn T.get_thread_binding(0) tk T.get_thread_binding(1) A_local T.alloc_local((1,), dtype) B_local T.alloc_local((1,), dtype) C_accum T.alloc_local((1,), accum_dtype) C_shared T.alloc_shared((BLOCK_N,), accum_dtype) if tk 0: C_shared[tn] 0 T.clear(C_accum) for bk in T.serial(T.ceildiv(K, BLOCK_K)): A_local[0] A[bk * BLOCK_K tk] B_local[0] B[bn * BLOCK_N tn, bk * BLOCK_K tk] C_accum[0] A_local[0].astype(accum_dtype) * B_local[0].astype(accum_dtype) T.atomic_add(C_shared[tn], C_accum[0]) C[bn * BLOCK_N tn] C_shared[tn] return main要点解析T.Kernel(blocks, threads(BLOCK_N, BLOCK_K))把线程块组织成二维threadIdx (tn, tk)。get_thread_binding(0/1)分别取threadIdx.x/y对应关系在 tilelang/language/kernel.py 的get_thread_binding实现中有明确文档dim0 对应threadIdx.xdim1 对应threadIdx.yT.atomic_addTileLang 语言级原子操作实现位于 tilelang/language/atomic.py直接映射到 CUDA 原子指令每个线程每轮只处理 1 个元素访存事务宽度为 1 个float16带宽利用率仍然很低。这一步之后延迟降到约 0.024 ms提升明显但仍未追平 torch/cuBLAS。用reduce_threads细化 K 维切分粒度当 K 很大时可以让每个线程一次处理多个连续元素通过reduce_threads参数控制一个输出元素由多少个线程共同归约def splitk_gemv( N: int, K: int, BLOCK_N: int, BLOCK_K: int, reduce_threads: int, dtype: str float16, accum_dtype: str float, ): TILE_K T.ceildiv(BLOCK_K, reduce_threads) T.prim_func def main( A: T.Buffer((K,), dtype), B: T.Buffer((N, K), dtype), C: T.Buffer((N,), dtype), ): with T.Kernel(T.ceildiv(N, BLOCK_N), threads(BLOCK_N, reduce_threads)) as bn: tn T.get_thread_binding(0) tk T.get_thread_binding(1) A_local T.alloc_local((TILE_K,), dtype) B_local T.alloc_local((TILE_K,), dtype) C_shared T.alloc_shared((BLOCK_N,), accum_dtype) C_accum T.alloc_local((1,), accum_dtype) if tk 0: C_shared[tn] 0 T.clear(C_accum) for bk in T.serial(T.ceildiv(K, BLOCK_K)): for k in T.serial(TILE_K): A_local[k] A[bk * BLOCK_K tk * TILE_K k] B_local[k] B[bn * BLOCK_N tn, bk * BLOCK_K tk * TILE_K k] for k in T.serial(TILE_K): C_accum[0] A_local[k].astype(accum_dtype) * B_local[k].astype(accum_dtype) T.atomic_add(C_shared[tn], C_accum[0]) C[bn * BLOCK_N tn] C_shared[tn] return main此时线程内局部缓冲区扩为TILE_K个元素线程tk负责 K 块中[tk * TILE_K, tk * TILE_K TILE_K)这一段BLOCK_K reduce_threads * TILE_K。这给了用户一个连续的搜索维度BLOCK_N × reduce_threads组合决定了线程块形状与每线程工作量。第四步向量化访存——把带宽瓶颈打掉GEMV 是典型的内存带宽受限算子最有效的单一优化是加宽访存事务等效于 CUDA 中的float2/float4加载。TileLang 通过T.vectorized显式声明向量化循环def splitk_gemv_vectorized( N: int, K: int, BLOCK_N: int, reduce_threads: int, dtype: str float16, accum_dtype: str float, ): MAX_TRANSACTION_SIZE_IN_BITS 128 TILE_K MAX_TRANSACTION_SIZE_IN_BITS // DataType(dtype).bits BLOCK_K reduce_threads * TILE_K T.prim_func def main( A: T.Buffer((K,), dtype), B: T.Buffer((N, K), dtype), C: T.Buffer((N,), dtype), ): with T.Kernel(T.ceildiv(N, BLOCK_N), threads(BLOCK_N, reduce_threads)) as bn: tn T.get_thread_binding(0) tk T.get_thread_binding(1) A_local T.alloc_local((TILE_K,), dtype) B_local T.alloc_local((TILE_K,), dtype) C_shared T.alloc_shared((BLOCK_N,), accum_dtype) C_accum T.alloc_local((1,), accum_dtype) if tk 0: C_shared[tn] 0 T.clear(C_accum) for bk in T.serial(T.ceildiv(K, BLOCK_K)): for k in T.vectorized(TILE_K): A_local[k] A[bk * BLOCK_K tk * TILE_K k] B_local[k] B[bn * BLOCK_N tn, bk * BLOCK_K tk * TILE_K k] for k in T.serial(TILE_K): C_accum[0] A_local[k].astype(accum_dtype) * B_local[k].astype(accum_dtype) T.atomic_add(C_shared[tn], C_accum[0]) C[bn * BLOCK_N tn] C_shared[tn] return main关键设计TILE_K不再由用户随意指定而是由硬件最大事务宽度反推——MAX_TRANSACTION_SIZE_IN_BITS 128128 bit 16 字节对应half类型下的 8 个元素、一次uint4加载。T.vectorized的语义定义见 tilelang/language/loop.py它生成带 vectorize 标注的 For 帧由后端 pass 决定实际向量化宽度。这一步把延迟压到约 0.0084 ms已经贴近 cuBLAS。第五步用tvm_thread_allreduce替换共享内存原子加atomicAdd到共享内存存在竞争开销更优做法是走 TVM 的tvm_thread_allreduce原语——它会针对归约轴上的线程数生成优化的树形/层级归约代码def splitk_gemv_vectorized_tvm( N: int, K: int, BLOCK_N: int, reduce_threads: int, dtype: str float16, accum_dtype: str float, ): MAX_TRANSACTION_SIZE_IN_BITS 128 TILE_K MAX_TRANSACTION_SIZE_IN_BITS // DataType(dtype).bits BLOCK_K reduce_threads * TILE_K T.prim_func def main( A: T.Buffer((K,), dtype), B: T.Buffer((N, K), dtype), C: T.Buffer((N,), dtype), ): with T.Kernel(T.ceildiv(N, BLOCK_N), threads(BLOCK_N, reduce_threads)) as bn: tn T.get_thread_binding(0) tk T.get_thread_binding(1) A_local T.alloc_local((TILE_K,), dtype) B_local T.alloc_local((TILE_K,), dtype) C_accum T.alloc_local((1,), accum_dtype) T.clear(C_accum) for bk in T.serial(T.ceildiv(K, BLOCK_K)): for k in T.vectorized(TILE_K): A_local[k] A[bk * BLOCK_K tk * TILE_K k] B_local[k] B[bn * BLOCK_N tn, bk * BLOCK_K tk * TILE_K k] for k in T.serial(TILE_K): C_accum[0] A_local[k].astype(accum_dtype) * B_local[k].astype(accum_dtype) C_reduced T.alloc_local((1,), accum_dtype) with T.attr( T.comm_reducer(lambda x, y: x y, [T.cast(0, accum_dtype)]), reduce_scope, T.reinterpret(T.uint64(0), dtypehandle), ): T.evaluate( T.tvm_thread_allreduce( T.uint32(1), C_accum[0], True, C_reduced[0], tk, dtypehandle, )) C[bn * BLOCK_N tn] C_reduced[0] return main写法要点T.attr(T.comm_reducer(...), reduce_scope, ...)声明归约作用域与组合算子这里是求和tvm_thread_allreduce内置函数在语言层的入口为 tilelang/language/tir/ir.py参数依次为归约值个数1、源累加器C_accum[0]、is_init标志、目标C_reduced[0]、归约线程轴tk共享内存分配被完全移除归约缓冲由编译器按线程数自动分配。底层实现在 src/transform/lower_thread_allreduce.ccThreadAllreduceBuilder会读取目标设备的thread_warp_size属性识别reduce_scope属性语句中的组合子把tvm_thread_allreduce调用改写为warp 内 shuffle 归约 跨 warp 共享内存层级归约的展开代码。以reduce_threads 32、BLOCK_N 2的调优配置为例最终生成的 CUDA 正是逐层折半的树形归约extern C __global__ void __launch_bounds__(64, 1) main_kernel(half_t* __restrict__ A, half_t* __restrict__ B, half_t* __restrict__ C) { float C_accum[1]; half_t A_local[8]; half_t B_local[8]; __shared__ float red_buf0[64]; C_accum[0] 0.000000e00f; for (int bk 0; bk 4; bk) { *(uint4*)(A_local 0) *(uint4*)(A ((bk * 256) (((int)threadIdx.y) * 8))); *(uint4*)(B_local 0) *(uint4*)(B ((((((int)blockIdx.x) * 2048) (((int)threadIdx.x) * 1024)) (bk * 256)) (((int)threadIdx.y) * 8))); for (int k 0; k 8; k) { C_accum[0] (C_accum[0] (((float)A_local[k]) * ((float)B_local[k]))); } } tl::fence_proxy_async(); __syncthreads(); red_buf0[((((int)threadIdx.x) * 32) ((int)threadIdx.y))] C_accum[0]; __syncthreads(); if (((int)threadIdx.y) 16) { red_buf0[...] (red_buf0[...] red_buf0[... 16]); } __syncthreads(); if (((int)threadIdx.y) 8) { /* 再折半 */ } ... if (((int)threadIdx.y) 1) { /* 最后一次折半 */ } __syncthreads(); C[((((int)blockIdx.x) * 2) ((int)threadIdx.x))] ((half_t)red_buf0[(((int)threadIdx.x) * 32)]); }可以看到三处编译器自动插入的优化uint4的 128-bit 向量化加载、__shared__ float red_buf0[64]归约缓冲、以及 5 轮__syncthreads()包裹的 32→16→8→4→2→1 折半归约。生成的代码与 TileLang 源程序高度同构同步与底层优化全部由编译器补全。这一步把延迟从约 0.0084 ms 进一步降到约 0.0069 ms文档实测已快于 torch/cuBLAS。第六步autotune 自动搜索最优配置BLOCK_N与reduce_threads是核心超参仓库提供tilelang.autotune装饰器自动搜索实现见 tilelang/autotuner/tuner.py支持configs、warmup、rep、timeout、ref_prog、skip_check、early_stop等参数并在检查阶段以rtol/atol验证数值正确性def get_best_config(N, K): def get_configs(): BLOCK_N [2, 4, 8, 32, 64, 128] reduce_threads [4, 8, 32] _configs list(itertools.product( BLOCK_N, reduce_threads, )) configs [{ BLOCK_N: c[0], reduce_threads: c[1], } for c in _configs] return configs autotune( configsget_configs(), warmup3, rep20, ) jit( out_idx[-1], supply_typetl.TensorSupplyType.Integer, ref_progref_program, skip_checkFalse, targetauto, ) def kernel( BLOCK_NNone, reduce_threadsNone, ): dtype float16 accum_dtype float MAX_TRANSACTION_SIZE_IN_BITS 128 TILE_K MAX_TRANSACTION_SIZE_IN_BITS // DataType(dtype).bits BLOCK_K reduce_threads * TILE_K T.prim_func def main( A: T.Buffer((K,), dtype), B: T.Buffer((N, K), dtype), C: T.Buffer((N,), dtype), ): with T.Kernel(T.ceildiv(N, BLOCK_N), threads(BLOCK_N, reduce_threads)) as bn: tn T.get_thread_binding(0) tk T.get_thread_binding(1) A_local T.alloc_local((TILE_K,), dtype) B_local T.alloc_local((TILE_K,), dtype) C_accum T.alloc_local((1,), accum_dtype) T.clear(C_accum) for bk in T.serial(T.ceildiv(K, BLOCK_K)): for k in T.vectorized(TILE_K): A_local[k] A[bk * BLOCK_K tk * TILE_K k] B_local[k] B[bn * BLOCK_N tn, bk * BLOCK_K tk * TILE_K k] for k in T.serial(TILE_K): C_accum[0] A_local[k].astype(accum_dtype) * B_local[k].astype(accum_dtype) C_reduced T.alloc_local((1,), accum_dtype) with T.attr( T.comm_reducer(lambda x, y: x y, [T.cast(0, accum_dtype)]), reduce_scope, T.reinterpret(T.uint64(0), dtypehandle), ): T.evaluate( T.tvm_thread_allreduce( T.uint32(1), C_accum[0], True, C_reduced[0], tk, dtypehandle, )) C[bn * BLOCK_N tn] C_reduced[0] return main return kernel()工作流说明kernel(BLOCK_NNone, reduce_threadsNone)中为None的参数即待搜索变量autotune按get_configs()枚举的笛卡尔积逐一编译、以ref_prog即A B.T校验数值并计时最终返回最优配置对应的编译结果warmup3, rep20控制计时轮次。调优后内核延迟约0.0067 ms。基准测试Hopper GPU 上的完整优化曲线文档在 Hopper GPU 上给出的各阶段延迟float16 输入、float 累加Kernel NameLatencytorch/cuBLAS0.00784 msTriton0.00773 msnaive_gemv0.16607 mssplitk_gemv0.02419 mssplitk_gemv_vectorized0.00809 mssplitk_gemv_vectorized_tvm0.00675 ms优化路径的贡献可以清晰拆分K 维并行带来约 7 倍提升0.166 → 0.024 ms向量化再带来约 3 倍提升0.024 → 0.008 ms归约原语替换贡献最后约 16%0.008 → 0.0067 ms。实操与延伸运行完整示例含正确性校验与逐版本计时默认NK1024可用--n/--k调整维度python examples/gemv/example_gemv.py # 或仅跑正确性回归 python examples/gemv/test_example_gemv.py正确性验证逻辑在 examples/gemv/example_gemv.py 的check_correctness_and_bench用profiler.assert_allclose(lambda x, y: x y.T, atol1e-2, rtol1e-2)与 PyTorch 参考比对再do_bench计时。此外示例文件还额外提供了一个块级思路的实现gemv_alloc_reducer基于T.alloc_reducer与T.Pipelined的软件流水线版本可作为与本文线程级路径的对照组两者在main的 benchmark 部分分别以 SIMT Latency 与 BlockReduce Latency 输出对比。小结这篇教程的价值在于展示 TileLang 相对块级 DSL 的核心差异——同一个 GEMV 问题通过T.get_thread_binding控制线程映射、T.atomic_add/tvm_thread_allreduce选择归约策略、T.vectorized控制访存宽度把朴素实现从慢于 cuBLAS 20 倍优化到反超 cuBLAS且每一步的中间产物生成的 CUDA 代码都可读、可诊断。对于其他带宽受限算子GEMV 类、量化反量化、逐元素归约这套并行度—事务宽度—归约路径的三层调优方法论同样适用具体超参应以autotune在目标硬件上的实测结果为准不同 N/K 规模下最优的BLOCK_N × reduce_threads组合并不一致。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考