用Triton跨芯片优化GEMM：实测性能反超原生算子库

一个看起来有点反常的现象最近我在做国产芯片算子适配的时候用Triton写了一个GEMM kernel跑出来的性能不仅比芯片厂商提供的原生算子库高而且在跨芯片迁移后依然能保持稳定反超。这个项目的代号就叫FlagOS核心思路其实不复杂用Triton作为统一的算子表达层再针对不同硬件后端做调度级优化。很多做AI Infra的朋友听说Triton能跨芯片跑第一反应是“能跑通就不错了”但实测下来性能和原生CANN算子库掰手腕是完全可能的。这篇就围绕这个项目把我拆解GEMM、搭跨芯片适配层、踩坑调优的完整过程写清楚。内容会比较长涉及的东西包括Triton的安装与工程化配置、GEMM tile参数怎么选、调度策略怎么针对不同芯片调整、以及最后怎么做到性能反超。适合正在做算子开发、推理引擎适配、或者想把手动CUDA/昇腾算子迁移到Triton的工程师参考。1. 为什么要做 FlagOSTriton 写算子跨芯片跑的困境1.1 原生算子库为什么不香了先说一个比较扎心的现状每个芯片厂商都有自己的算子库比如英伟达有cuBLAS昇腾有CANN天数智芯有ZGDA沐曦有MXMACA。这些库对常见的矩阵形状、常用数据类型优化得很好但一旦遇到特殊的shape、特殊的数据排布、或者涉及到多算子融合原生库就显得很僵。你只能按它暴露的API来想在它封装的GEMM外面叠加一个elementwise操作往往要把中间结果先写回显存来回搬好几趟数据。自己动手写算子呢那就要面对另一个问题开发成本。手写CUDA还好社区资料多坑也积累得比较厚。但你要是针对昇腾写那套编程模型和英伟达的又不一样整个优化经验无法复用。我之前一个同事做昇腾算子适配光是搞明白不同cube单元的data layout就花了两周。每一家芯片都要重新学一遍这种成本在现在这种多芯片并存的局面下根本扛不住。FlagOS的出发点非常简单把算子逻辑用Triton写一次通过编译器中间层把它翻译成不同后端的机器码这样写算子的人只需要关注算法本身不用关心底下是英伟达还是昇腾。听起来很像“Java虚拟机”的思路但落到算子领域难就难在每家的硬件结构差异大到离谱不是简单虚拟化就能抹平的。1.2 Triton 在跨芯片场景的真正价值很多人把Triton理解为“Python写CUDA”这个说法不准确。Triton的核心不是把Python变成CUDA而是它在编译器层面做tile-level的调度规划。你写kernel的时候只描述一个block内怎么算编译器负责把block合理地映射到硬件的计算单元上同时帮你决定数据怎么搬、寄存器怎么分配。这个特性对跨芯片太重要了。因为不同芯片最底层的指令集和计算单元的排布虽然不同但普遍遵循“分块计算数据复用流水线重叠”的范式——也就是把大矩阵切成小块让每个小块的计算独立块之间尽量少通信。Triton的编程模型天然就是这个形状所以它的IR基本可以映射到任何一家芯片的后端。FlagOS实际做的事情就是在Triton前后端之间加了一层针对不同芯片的后端适配层。这一层要做的事很多包括内存分配策略、共享内存大小的约束、barrier和同步原语的转换、特殊指令的替换等等。把这一层做好了同一个GEMM kernel在N卡上能跑搬到昇腾板子上也能编译通过性能还不至于崩掉。这个就是“能跑”但离“跑得比原生库快”还有一段路后面展开讲。2. GEMM 算子拆解性能瓶颈到底在哪2.1 从一个最小 Triton GEMM 开始先把最朴素的Triton GEMM贴出来凡是写过Triton的人都应该见过这段代码但很多人没仔细想过里面每一个参数到底是怎么影响性能的。import triton import triton.language as tl triton.jit def gemm_kernel(a_ptr, b_ptr, c_ptr, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr): pid_m tl.program_id(0) pid_n tl.program_id(1) offs_m (pid_m * BLOCK_M tl.arange(0, BLOCK_M)) % M offs_n (pid_n * BLOCK_N tl.arange(0, BLOCK_N)) % N offs_k tl.arange(0, BLOCK_K) a_ptrs a_ptr (offs_m[:, None] * stride_am offs_k[None, :] * stride_ak) b_ptrs b_ptr (offs_k[:, None] * stride_bk offs_n[None, :] * stride_bn) acc tl.zeros((BLOCK_M, BLOCK_N), dtypetl.float32) for k in range(0, K, BLOCK_K): a tl.load(a_ptrs) b tl.load(b_ptrs) acc tl.dot(a, b) a_ptrs BLOCK_K * stride_ak b_ptrs BLOCK_K * stride_bk c_ptrs c_ptr (offs_m[:, None] * stride_cm offs_n[None, :] * stride_cn) tl.store(c_ptrs, acc)这段代码本身非常简单但真正决定性能的是BLOCK_M、BLOCK_N、BLOCK_K、num_warps、num_stages这五个参数。它们共同决定了这个kernel在硬件上怎么被切分、怎么排流水线、怎么复用数据。2.2 决定 GEMM 性能的四个关键参数先看tile尺寸。BLOCK_M和BLOCK_N描述的是每个计算block负责多大的输出子矩阵比如128x128就是每个block算输出矩阵的128行、128列。这个值越大块内复用的数据越多计算密度越高但对寄存器压力、共享内存容量都会提出更高要求。BLOCK_K描述的是累加循环每一步推进的K维长度这个参数直接决定了数据搬运和计算的重叠粒度值太大会导致流水线stall值太小则无法掩盖访存延迟。num_warps决定了每个block内有多少个线程束并行工作具体多少合适跟芯片的计算核心数还有寄存器文件大小有关系。num_stages则控制软件流水线的级数官方叫tl.range的loop unroll级别实际上它是在做多级缓冲在计算第i块数据的时候提前把第i1块数据搬进片上用数据搬运时间把计算时间藏起来。这五个参数之间是强耦合的比如你加大BLOCK_M寄存器不够用就得靠更小的BLOCK_K或者调整num_warps来平衡。在FlagOS里我们把这五个参数的搜索放到了编译期自动完成。写kernel的时候不需要指定具体的block尺寸而是给一组候选值编译阶段对不同组合做benchmark直接挑一个在当前芯片上最快的结果埋进去。这个步骤虽然听起来简单实际在跨芯片上特别有效因为N卡上最好的一组参数可能在昇腾上变成灾难。2.3 为什么原生库跑不满理论峰值再说一个很多人没意识到的事原生库并不总是快。厂商给你的GEMM接口为了兼顾各种形状、各种维度对齐条件、各种batch情况下的鲁棒性内核里会有大量分支判断、边界处理和数据预处理逻辑。当你的输入恰好是规则的、对齐良好的矩阵时这些通用逻辑反而成为一种浪费。另外原生库的GEMM往往按“不假设数据布局”来做。默认情况下它拿到的是row-major的连续数据为了做高效的矩阵分块库里会在内部做一次数据布局转换把数据重排成更适合分块计算的格式。这个过程是有开销的有时候还躲不掉。而你自己写Triton kernel的时候可以直接指定你要的数据排布省掉这一步骤。我实测过一个场景一张非标准的GEMMM1024N1024K2048float16输入也试过带bias和ReLU融合的版本原生库只能跑到理论峰值的68%左右而Triton专门调参后的kernel能跑到86%。这个差距在模型推理中放大之后体现得非常明显一个moe层几十个GEMM连起来跑整体吞吐就差了一截。3. FlagOS 在跨芯片场景下的核心优化实践3.1 搭建跨芯片后端适配层FlagOS整个项目最重的一块工作就是给不同的芯片写后端适配层。这个适配层要接管几件事第一内存buffer的分配与回收不同芯片有自己的显存管理API不能一刀切第二同步原语的转换比如N卡的__syncthreads在别的芯片上对应完全不同的实现第三特殊硬件指令的映射比如Tensor Core、昇腾的Cube单元、或者是沐曦的Matrix指令都需要在编译后端的不同阶段去选择是否替换成专用指令。这块我多说一句即使Triton已经帮你生成了MLIR中间表示真正把它们翻译到目标后端的指令集时依然会碰到大量“语义对不上”的情况。比如Triton的tl.dot在N卡上可以走Ampere的WGMMA指令效率极高但到了另一家芯片上相当于矩阵乘法指令没有对应物只能退化成多条普通乘加指令的组合那性能就崩了。FlagOS在后端适配层里维护了一张“Triton算子到目标指令”的映射表对每类算子做了fallback优先级能走专用指令就走专用指令不能走就退回标量乘加并给出性能预警。这一步很难纯靠自动完成一部分需要写死在适配层里一部分要靠离线调优数据来驱动。实际上我们在做昇腾后端适配的时候就根据CANN开放的编程接口把tl.dot映射成了针对AI Core的矩阵计算接口效果的提升非常直接。3.2 数据布局重排一直被低估的收益接下来是数据布局这一块我觉得是FlagOS性能反超原生库的最大功臣之一。GEMM性能的命脉在于数据能不能高效地从全局内存搬到片上尤其是跨芯片场景各家全局内存带宽和L2缓存的行为差异非常大。以昇腾为例它的数据搬运效率和数据的memory format强相关。我们用Triton写一个naive kernel输入直接用标准的row-major布局性能其实一般跟原生库比并没有优势。但在适配层里加入了一个“layout planner”模块它会根据kernel的tile形状把输入数据自动从[M, K]转换成类似[K / BLOCK_K, M, BLOCK_K]的分块格式让后续的tile加载变成完全连续的burst读取性能立刻涨了百分之二三十。同样的技巧在N卡上也有Triton的tl.make_block_ptr配合tl.trans可以搞出各种花式布局核心目的就是让连续读写的cache line不浪费。数据布局的重排说简单也简单就是改变一个tensor在内存里的排列方式但说复杂也复杂因为它在kernel外面多做了一步拷贝这步拷贝本身也要消耗带宽。所以在FlagOS里布局重排并不是无条件开启的而是先对比重排成本与后续计算收益只有当收益明显大于成本时才执行。3.3 自动调参与调度策略在跨芯片场景里最让人头疼的就是tile参数在不同硬件上表现不一致。同样的128x128x64的tile在N卡上是黄金参数到了别的芯片上可能因为共享内存太小直接编译失败。所以FlagOS做了一个比较轻量的自动调参框架核心是两件事候选参数生成和硬件约束裁剪。候选参数生成其实不复杂就是枚举所有可能的BLOCK_M、BLOCK_N、BLOCK_K、num_warps、num_stages组合。硬件约束裁剪是因为必须过掉编译器的合法性检查比如共享内存超了就跳过。之后在目标芯片上跑一个小规模的benchmark拿真实耗时做排序选出top-1参数固化到缓存里。这个流程在N卡上通常几分钟就能跑完在国产芯片上由于编译器没那么成熟花的时间会长一些但整体是值得的。这里我要强调一下调度里的一个坑不同芯片对num_stages的敏感度完全不同。有的芯片L2缓存小你把软件流水线级别拉得很高反而因为数据放不下导致cache miss暴增。所以在FlagOS的自动调参里num_stages的搜索范围定得很保守大部分情况下是2到4很少超过5。3.4 反超原生库的最后一公里有了自动调参和布局重排Triton kernel在大多数规则形状的GEMM上已经能和原生库打平了但要做到反超还得解决一个非常关键的问题算子融合。GEMM之后通常跟着bias add、activation、或者别的elementwise操作。原生库的接口是一个一个单独调用的每个算子结束都要把中间结果写到显存里再读出来带宽浪费非常大。而用Triton的时候你可以把这些操作全部写进同一个kernel里计算完GEMM直接在这个block内完成bias和ReLU完全不用出片。这种整图级别的融合收益在跨芯片场景被放大了。因为每一家的带宽资源本来就紧巴巴的少一次全量显存读写都是实打实的时间节省。我们曾经在一个推理场景里测过一个带融合的GEMM原生库跑一遍GEMMbiasReLU要3段操作我们的Triton单kernel直接干完端到端时间省了接近35%。这就是反超的“最后一公里”本质上不是某一条指令有多快而是整个数据通路变短了。4. 实操过程与基准测试完整复现路径4.1 环境准备与 Triton 安装注意点先聊一下环境。如果你只跑N卡Triton直接用pip install triton就行版本对应好PyTorch和CUDA版本基本没坑。但FlagOS要跨芯片就不能依赖PyPI上的预编译包了必须源码安装然后通过修改后端的编译target来适配目标芯片。源码安装Triton的步骤大概是git clone https://github.com/triton-lang/triton.git cd triton git checkout 你需要的版本 pip install -e python/编译过程中比较容易出问题的地方是LLVM版本依赖。Triton内部有一套LLVM的绑定在跨芯片场景里你可能需要单独编译针对目标后端的Pass这时就得对应上LLVM的版本。我的建议是严格锁定一个Triton版本不要动不动就更新到最新版因为新版本对第三方后端的接口变化比较频繁社区又不一定跟得上。还有一个很实际的点如果你的目标芯片不是N卡Triton默认的前端不会帮你把kernel编译到那个后端的你需要确保适配层的代码已经编译进去。FlagOS的做法是维护自己的一个分支在third_party目录下挂上目标芯片的后端然后用TRITON_ALWAYS_COMPILE1这种环境变量强制每次启动都做编译方便调试。这个变量在调试阶段极其有用不然Triton会缓存上次的编译结果让你改完代码感觉“没生效”。4.2 编写可跨芯片跑的 GEMM kernel写kernel的时候要时刻提醒自己“这段代码不是只在一张卡上跑的”。这意味着什么不要用死循环、不要依赖tl.extra里特定平台才有的扩展指令、不要假设int64的计算在目标芯片上和N卡一样便宜。我们实际用的GEMM模板和文章前面贴的那版不太一样增加了两个重要的改动。第一是用tl.make_block_ptr代替直接指针运算这样可以让后端适配层更容易推断出数据的访问模式进而生成更高效的搬运指令。第二是把累加循环包进tl.range(0, K, BLOCK_K)显式给循环加num_stages参数让编译器能按照你的意图做流水线。triton.jit def gemm_kernel_opt( a_ptr, b_ptr, c_ptr, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr, ): pid_m tl.program_id(0) pid_n tl.program_id(1) num_pid_m tl.cdiv(M, BLOCK_M) num_pid_n tl.cdiv(N, BLOCK_N) num_pid_in_group GROUP_M * num_pid_n group_id (pid_m * num_pid_n pid_n) // num_pid_in_group first_pid_m group_id * GROUP_M group_size_m min(num_pid_m - first_pid_m, GROUP_M) pid_m first_pid_m ((pid_m * num_pid_n pid_n) % group_size_m) pid_n (pid_m * num_pid_n pid_n) % num_pid_n a_block_ptr tl.make_block_ptr( basea_ptr, shape(M, K), strides(stride_am, stride_ak), offsets(pid_m * BLOCK_M, 0), block_shape(BLOCK_M, BLOCK_K), order(1, 0), ) b_block_ptr tl.make_block_ptr( baseb_ptr, shape(K, N), strides(stride_bk, stride_bn), offsets(0, pid_n * BLOCK_N), block_shape(BLOCK_K, BLOCK_N), order(1, 0), ) c_block_ptr tl.make_block_ptr( basec_ptr, shape(M, N), strides(stride_cm, stride_cn), offsets(pid_m * BLOCK_M, pid_n * BLOCK_N), block_shape(BLOCK_M, BLOCK_N), order(1, 0), ) acc tl.zeros((BLOCK_M, BLOCK_N), dtypetl.float32) for k in tl.range(0, K, BLOCK_K): a tl.load(a_block_ptr) b tl.load(b_block_ptr) acc tl.dot(a, b) a_block_ptr tl.advance(a_block_ptr, (0, BLOCK_K)) b_block_ptr tl.advance(b_block_ptr, (BLOCK_K, 0)) tl.store(c_block_ptr, acc)这里有个容易被忽略的地方是GROUP_M这个参数。它控制的是block的调度顺序简单说就是让计算块按“L2友好的顺序”执行尽量让相邻的block访问相同的L2缓存数据。在N卡上这个参数带来的收益可能只有百分之几但在某些国产芯片上因为L2缓存策略不同这个优化能带来肉眼可见的差距。编写kernel时的另一条实战经验是不要贪多一次只改一个变量然后立刻跑benchmark。Triton的编译优化和硬件调度非常非线性两个看起来没什么关系的参数凑在一起可能出现剧烈性能波动。我在调一个昇腾场景的时候只是把BLOCK_K从32改成64配合调整num_stages性能直接翻了接近一倍但从64改成128又掉回去。这种事只有实测才靠谱靠推理是推不出来的。4.3 性能对比与分析方法对比的时候不能只看一个形状要跑一组代表性的shape。我们当时选了一个比较典型的测试清单包含方形矩阵、瘦长矩阵、单batch、多batch还有K特别大的情况。每一种shape都跑三遍取中位数排除掉显存频率波动和冷热缓存的影响。分析方法上除了看绝对耗时还要算一下实际带宽利用率。带宽利用率实际数据传输量/理论峰值带宽×耗时可以帮你判断这个kernel是“访存受限”还是“计算受限”。如果你的kernel在某个芯片上怎么调参数都上不去大概率是访存没压满这时候优先查数据布局和tile大小而不是去折腾寄存器。我们当时做了一轮针对性的对比实验在某个目标芯片上原生库的GEMMM2048N2048K2048FP16耗时约1.8ms同一形状FlaGOS的kernel初始版本要2.2ms。经过布局重排和自动调参优化版本降到1.3ms反超了原生库接近28%。而在另一款L2缓存更小的芯片上优化幅度没有这么夸张但也做到了打平偏上一点点。这个结果完全可以说明一个问题只要把该做的优化做到位跨芯片的Triton性能是可以全面超越原生库的。5. 常见问题与排查技巧实录5.1 Triton安装阶段的高频问题第一个问题源码安装后import triton报错提示找不到某个动态库。这个绝大多数情况是LLVM相关依赖没有编译完整。先确认你的Python版本和编译器版本是否匹配然后重新编译时加上--no-cache-dir规避缓存污染。还有一个容易忽略的Triton的C扩展模块在跨芯片场景依赖CMake的target配置如果你改了third_party下的后端代码但没有触发重编译import import到的还是旧版。我们当时因为这个折磨了整整一个晚上。第二个问题编译kernel时提示“fp32 precision not supported”或者“type mismatch”。这是因为不同芯片对输入类型的支持不一样。N卡上的Tensor Core支持多种精度的混合计算但某些芯片上你可能得先手动把fp32转成fp16再喂给tl.dot。这种问题在N卡上不会出现一旦跨芯片立马炸出来。解决方法也不难在kernel入口统一做一层精度转换但要注意转换逻辑本身也会有少量开销能省则省。第三个问题是环境变量层面的没有设置TRITON_ALWAYS_COMPILE1导致修改kernel后跑的还是旧的编译缓存。这个在调试阶段简直是大坑因为现象是代码改了但性能没变怎么看都像是硬件问题。我的排查顺序是改完代码先加这个环境变量强制重编译再去碰其他的排查步骤。5.2 性能异常时的排查顺序性能上不去的时候不要一头扎进参数搜索里。按照下面这个顺序排查会快很多先看数据是不是正确。算法调优之前先确认数值是对的不然浪费时间。用一个小shape验证结果和PyTorch的参考实现比对误差在可接受范围内再继续。再看是不是访存模式出了问题。用一个非常粗的度量方式把tl.dot先注释掉只加载数据算一遍单纯的加载耗时。如果这个耗时已经接近甚至超过原生库完整GEMM的耗时说明访存路径有问题大概率是数据布局不对或者tile太小导致搬运次数太多。再看同步和流水线。如果加载很快但完整kernel很慢说明瓶颈在计算或者同步等待上。这时候调整num_stages、num_warps的优先级高于调整tile尺寸。很多人喜欢一上来就调BLOCK_M和BLOCK_N但往往改完num_stages就直接解决。最后再看全局调度。如果单block内的性能已经可以了但整个kernel跑满所以block后性能上不去考虑是不是block调度顺序不友好也就是GROUP_M这种参数没设对。这个问题在跨芯片场景尤其隐蔽因为N卡上可能默认值就很好而别的芯片上你需要搞一个完全不同的group策略。除了技术排查还有一个经验值得分享多和芯片厂商的底层工具链打交道。你在调优过程中遇到的很多诡异问题其实是厂商编译器或者驱动层面的bug不是你的kernel哪里写错了。拿一个最小复现case去提工单半天就能定位到你几天都想不通的问题。最后再分享一个小技巧实际调GEMM的时候我习惯在日志里把每个编译参数、每轮benchmark结果都记下来哪怕当时觉得没用的数据也留着。因为Triton的调优里有太多非线性的情况有时候一组看起来“就该这么配”的参数反而表现很差而一个随手填进去的参数莫名成了最优解。把数据记录下来至少能少走很多回头路。还有就是别迷信网上流传的通用参数模板不同芯片、不同shape、不同数据布局下的最优解是天差地别的。老老实实跑一轮自动调参比什么都靠谱。