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

一个看起来有点反常的现象，最近我在做国产芯片算子适配的时候，用Triton写了一个GEMM kernel，跑出来的性能不仅比芯片厂商提供的原生算子库高，而且在跨芯片迁移后依然能保持稳定反超。这个项目的代号就叫FlagOS，核心思路其实不复杂：用Triton作为统一的算子表达层，再针对不同硬件后端做调度级优化。很多做AI Infra的朋友听说Triton能跨芯片跑，第一反应是“能跑通就不错了”，但实测下来，性能和原生CANN算子库掰手腕是完全可能的。

这篇就围绕这个项目，把我拆解GEMM、搭跨芯片适配层、踩坑调优的完整过程写清楚。内容会比较长，涉及的东西包括Triton的安装与工程化配置、GEMM tile参数怎么选、调度策略怎么针对不同芯片调整、以及最后怎么做到性能反超。适合正在做算子开发、推理引擎适配、或者想把手动CUDA/昇腾算子迁移到Triton的工程师参考。

1. 为什么要做 FlagOS：Triton 写算子，跨芯片跑的困境

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), dtype=tl.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块数据的时候，提前把第i+1块数据搬进片上，用数据搬运时间把计算时间藏起来。这五个参数之间是强耦合的，比如你加大BLOCK_M，寄存器不够用，就得靠更小的BLOCK_K或者调整num_warps来平衡。

在FlagOS里，我们把这五个参数的搜索放到了编译期自动完成。写kernel的时候不需要指定具体的block尺寸，而是给一组候选值，编译阶段对不同组合做benchmark，直接挑一个在当前芯片上最快的结果埋进去。这个步骤虽然听起来简单，实际在跨芯片上特别有效，因为N卡上最好的一组参数可能在昇腾上变成灾难。

2.3 为什么原生库跑不满理论峰值

再说一个很多人没意识到的事：原生库并不总是快。厂商给你的GEMM接口为了兼顾各种形状、各种维度对齐条件、各种batch情况下的鲁棒性，内核里会有大量分支判断、边界处理和数据预处理逻辑。当你的输入恰好是规则的、对齐良好的矩阵时，这些通用逻辑反而成为一种浪费。

另外，原生库的GEMM往往按“不假设数据布局”来做。默认情况下它拿到的是row-major的连续数据，为了做高效的矩阵分块，库里会在内部做一次数据布局转换，把数据重排成更适合分块计算的格式。这个过程是有开销的，有时候还躲不掉。而你自己写Triton kernel的时候，可以直接指定你要的数据排布，省掉这一步骤。

我实测过一个场景：一张非标准的GEMM，M=1024，N=1024，K=2048，float16输入，也试过带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，原生库跑一遍GEMM+bias+ReLU要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_COMPILE=1这种环境变量强制每次启动都做编译，方便调试。这个变量在调试阶段极其有用，不然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( base=a_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( base=b_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( base=c_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), dtype=tl.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大小，而不是去折腾寄存器。

我们当时做了一轮针对性的对比实验，在某个目标芯片上，原生库的GEMM（M=2048，N=2048，K=2048，FP16）耗时约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_COMPILE=1，导致修改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、不同数据布局下的最优解是天差地别的。老老实实跑一轮自动调参，比什么都靠谱。