企业官网建设流程全解析

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

【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang

GEMV（矩阵-向量乘）是 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, mask=mask, other=0.0) x_vals = tl.load(x_ptr + offs_k, mask=mask, other=0.0) dot = tl.sum(a_vals * x_vals, axis=0) 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_dealloc<24>(); 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_alloc<240>(); C_reg[0] = 0.000000e+00f; 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实现中有明确文档：dim=0 对应threadIdx.x，dim=1 对应threadIdx.y；
• T.atomic_add：TileLang 语言级原子操作，实现位于 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 = 128（128 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), dtype="handle"), ): T.evaluate( T.tvm_thread_allreduce( T.uint32(1), C_accum[0], True, C_reduced[0], tk, dtype="handle", )) 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.cc：ThreadAllreduceBuilder会读取目标设备的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.000000e+00f; 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( configs=get_configs(), warmup=3, rep=20, ) @jit( out_idx=[-1], supply_type=tl.TensorSupplyType.Integer, ref_prog=ref_program, skip_check=False, target="auto", ) def kernel( BLOCK_N=None, reduce_threads=None, ): 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), dtype="handle"), ): T.evaluate( T.tvm_thread_allreduce( T.uint32(1), C_accum[0], True, C_reduced[0], tk, dtype="handle", )) C[bn * BLOCK_N + tn] = C_reduced[0] return main return kernel()

工作流说明：kernel(BLOCK_N=None, reduce_threads=None)中为None的参数即待搜索变量，autotune按get_configs()枚举的笛卡尔积逐一编译、以ref_prog（即A @ B.T）校验数值并计时，最终返回最优配置对应的编译结果；warmup=3, rep=20控制计时轮次。调优后内核延迟约0.0067 ms。

基准测试：Hopper GPU 上的完整优化曲线

文档在 Hopper GPU 上给出的各阶段延迟（float16 输入、float 累加）：

优化路径的贡献可以清晰拆分：K 维并行带来约 7 倍提升（0.166 → 0.024 ms），向量化再带来约 3 倍提升（0.024 → 0.008 ms），归约原语替换贡献最后约 16%（0.008 → 0.0067 ms）。

实操与延伸

运行完整示例（含正确性校验与逐版本计时，默认N=K=1024，可用--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, atol=1e-2, rtol=1e-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