Triton跨芯片优化GEMM：反超原生库的实战与自动调优策略

1. 为什么Triton写出的GEMM，跨芯片时反而能反超原生库？

先交代一下背景。我最近在做一个跨芯算子迁移项目，目标很朴素：把一套基于CUDA生态写的算子在多款国产AI加速芯片上跑起来，性能不能比原卡差太多，最好还能反超。项目初期大家默认的方案是调用各家芯片SDK自带的GEMM库，毕竟原厂库是芯片厂商自己调的，适配性、性能那都是压箱底的功夫。

结果实测下来打脸了。在某些加速芯片上，用Triton重写的GEMM内核，配合众智FlagOS平台做跨芯片编译调度，性能竟然超过了芯片厂商SDK原生的GEMM算子，最高一档领先幅度在特定形状上到了20%上下。这个结果让我重新审视了一件事：原生库的"原"，未必等于"适合你的场景"。

这里先快速对齐几个概念。

• FlagOS：一个面向算子跨芯片迁移与调优的开放基础设施平台，核心是提供统一的算子开发编译链路、自动调优框架和可复用的调度策略，解决同一套算子在多种加速芯片上的移植适配问题。
• Triton：一个Python式的GPU编程语言和编译器，本质上是个面向GPU/AI加速器的DSL。它把CUDA里的grid、block、warp、shared memory、同步这些细节做了抽象，让开发者用类NumPy的语法描述算子，再由编译器自动映射到SIMT指令、张量核指令以及各种访存调度上。
• GEMM：通用矩阵乘法，也就是C = A × B + C。几乎所有深度学习模型的计算瓶颈都压在这里，所以GEMM也常被称为AI计算的"基准货币"。

那为什么跨芯片时，Triton能反超原生库？道理其实不复杂：原生库的优化是面向单一硬件深度定制的，库函数把方块切分、数据布局、指令发射、流水线深度全部固化了。当你换一张芯片，它的寄存器堆结构、访存带宽、张量核位宽、缓存层次跟原设计不完全匹配时，这个"固化"就变成了一种妥协。而Triton内核可读可改，配合FlagOS的自动调优搜索，可以直接针对目标芯片重新生成调度方案，相当于每张新卡都能"量体裁衣"。

这不是说原生库不行，而是说跨芯场景下，通用可移植内核+自动调优的组合，往往比抱着一个未经适配的原生库更靠谱。这一点，做异构计算的老兵应该都有共鸣：没有万能的轮子，只有针对特定路况调优的轮子。

下面我用实际项目的角度，把这次FlagOS跨芯算子优化的完整过程和盘托出。

2. 写一个能"跨芯"的Triton GEMM内核，先从理解GEMM的底层规律开始

2.1 GPU上GEMM的真实执行瓶颈：计算、访存、还是调度？

在贴代码之前，得先搞清楚GEMM在AI加速器上到底怎么跑。很多人一写GEMM就想着用张量核，以为把矩阵扔给mma指令就万事大吉，实际远没那么简单。

GEMM的执行过程可以拆成三个层次：宏观上是分块矩阵的调度，中观上是数据块在内存和计算单元之间的搬运，微观上才是乘加指令本身。一张AI芯片的计算峰值再高，如果数据喂不进计算单元，峰值就是一串永远够不到的数字。业界常用一个指标叫算术强度（arithmetic intensity），它等于计算量除以数据搬运量。GEMM这么做是有原因的：矩阵乘天然具备超高的计算复用率，一个数据元素可以被反复用于多次乘加运算，所以理论上算术强度可以做得很大，瓶颈通常不在计算，而在怎么把数据高效地调度进寄存器或者张量核。

我说个更直观的比喻：GEMM像一个大厨房里的备菜流水线。计算单元是灶台，内存是仓库，缓存是案板。如果能提前把接下来要炒的菜全部切好、码在案板上，灶台就能一刻不停地炒；如果每次要炒菜才去仓库取食材，灶台就得空转等菜。GEMM优化的本质，就是让"取菜—切菜—炒菜"这三个环节始终重叠起来，互不等待。

所以我写跨芯GEMM时，脑子里始终绷着三根弦：

• 数据能不能按时到达计算单元；
• 计算单元能不能持续满负荷工作；
• 中间结果能不能安全、高效地留在离计算单元最近的地方（寄存器/共享内存）。

三根弦只要有一根松了，性能立刻掉下来。Triton写的GEMM再花哨，也逃不出这三条铁律。

2.2 用Triton写GEMM内核时的可移植设计

Triton的哲学是"用Python描述计算，让编译器搞定调度"。它不直接暴露shared memory和线程同步，而是提供一组更高级的原语：tl.dot负责矩阵乘法的核心计算，tl.load/tl.store负责数据搬运，编译器自动决定如何在底层安排加载、同步和计算。这种抽象对跨芯非常有利——因为各家芯片的并行模型在细节上有差异，而Triton的编译器负责针对不同后端生成对应的中间表示并调度到目标硬件。

但抽象不等于万能。跨芯GEMM内核设计的关键，是把"计算逻辑"和"硬件调度"尽量分离。计算逻辑部分：分块方式、累加逻辑、精度处理，这些是芯片无关的，要写成清晰的数学表达。硬件相关部分：块大小、流水线级数、访存布局，这些是芯片强相关的，要写成可配参数，交给FlagOS的自动调优去搜索。

我最初迁移到FlagOS上的GEMM内核，核心代码大致是这个样子（简化版，演示跨芯可移植写法的骨架）：

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, GROUP_M: tl.constexpr, ): pid = tl.program_id(0) num_pid_m = tl.cdiv(M, BLOCK_M) num_pid_n = tl.cdiv(N, BLOCK_N) # 用grouped ordering做L2友好的调度 num_pid_in_group = GROUP_M * num_pid_n group_id = pid // 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 % group_size_m) pid_n = (pid % num_pid_in_group) // group_size_m offs_am = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % M offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N offs_k = tl.arange(0, BLOCK_K) a_ptrs = A_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) b_ptrs = B_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k in range(0, tl.cdiv(K, BLOCK_K)): a = tl.load(a_ptrs) b = tl.load(b_ptrs) # tl.dot是跨芯编译的关键抽象：编译到目标芯片的矩阵乘指令 accumulator = tl.dot(a, b, accumulator) a_ptrs += BLOCK_K * stride_ak b_ptrs += BLOCK_K * stride_bk offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) c_ptrs = C_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] tl.store(c_ptrs, accumulator)

这段代码最初照着Triton官方的GEMM示例改的，改动不大，但迁移到FlagOS以后有两点要注意：

1. tl.dot是硬性计算接口，跨芯片时编译器会把它映射到目标芯片的矩阵乘指令（比如N卡上的mma指令，或者其他芯片上对应的张量指令）。写内核时不要试图自己拆解tl.dot，否则一旦拆成FFMA或者element-wise乘加，就没法吃到芯片自带的高效矩阵指令了。
2. BLOCK_M、BLOCK_N、BLOCK_K这些参数不要写死。同一段代码，在A芯片最优可能是128x128x32，在B芯片最优可能是64x64x64。把它们拦成可搜索参数，是后面FlagOS自动调优能跑起来的前提。

提示：这个内核里有% M和% N的取模操作，是为了处理M和N不能被块大小整除的边缘情况。在自动调优时，尽量让取模分支在编译期就能被消除（即M、N刚好等于块大小的整数倍），否则这部分的地址计算会引入额外的整数运算开销，影响内核整体效率。

2.3 为什么说数据复用是GEMM调优的第一性原理

写Triton内核时可以偷懒把每个线程块的A、B块直接从显存里加载，然后tl.dot算完拉倒。这样代码最短，但性能会很难看。原因在于你不去管它的时候，编译器难以判断什么样的调度策略在当前芯片上最合理。

GEMM的调优本质上是在做"空间换时间"的生意：把一块数据从慢速存储搬到快速存储，多搬一次、多复用一次，计算就会快一分。所以所有高性能GEMM库，无论N卡还是国产卡，都会在kernel里做多级数据缓存：

• A的一小块和B的一小块先加载到共享内存；
• 再从共享内存加载进寄存器/张量核参与计算；
• 计算完的中间结果留在累加器里，不急着写回显存。

Triton里通过tl.load加载的数据，你能主动控制的是"从全局内存加载到片上缓冲"这一步。而"从片上缓冲到寄存器"在编译器自动安排的调度里完成，对你的代码来说是一个隐式过程。涉及到一个关键的编译参数叫num_stages，它控制流水线预取的深度——也就是提前多少个循环迭代预取数据到片上，用来覆盖计算延迟。

我在跨芯调优中第一个要搜的参数就是这个num_stages。在原生N卡上，它设置成3或4通常有提升；但换到另一款加速芯片时，由于它的片上内存带宽和延迟特性不一样，最优值可能变成2、4，甚至6。这个参数不试不知道，全靠实测说话。

3. FlagOS跨芯算子优化的实操路线：从编译适配到自动调优

3.1 平台侧做了什么：统一IR、代码生成与算子库对接

如果没有FlagOS这类平台，跨芯迁移的常规做法是什么？拿到一块新芯片SDK，读文档，看它的编程模型，把CUDA代码一行行翻译成新语言（比如CANN的Ascend C）。这个过程的痛点是：每次新芯片都要重写一遍内核，而且你根本不可能比原厂更了解自己的硬件。

FlagOS的思路不同。它把算子开发抽象成了"前端语言无关、后端芯片由编译器和调度框架适配"的中间表示层。Triton写的主机端代码进入FlagOS编译链路后，先被解析成统一的算子中间表示，再经由FlagOS面向不同芯片的后端，做指令选择、访存调度、管线编排，最后生成本地可执行代码。

这意味着两件事。第一，你要移植的算子的逻辑，只需要写一次；第二，算子性能的最终走向，由平台针对具体芯片生成的调度策略决定，而不是由开发者的手工指令优化决定。

具体到这次GEMM优化，FlagOS帮我们省力最多的是三块：

• 自动同步与流水线插入：在多芯片上做tl.load循环时，编译器要自动安排"数据加载完成后再开始计算"的依赖关系。手写CUDA时这是最烦的部分；Triton+FlagOS把这部分自动化了，但要命的是流水线深度（num_stages）对每张芯片的收益不一样，必须去搜。
• 高效访存布局生成：A、B矩阵的加载，在目标芯片上到底是按行连续拷贝更快、还是按列连续拷贝更快、还是用swizzled模式避开bank conflict更快，平台会在自动调优时生成不同变体来尝试。
• 与原生算子库的互操作接口：调试阶段可以拿平台生成的内核直接跟芯片SDK的GEMM对比，不用另外写benchmark框架，方便我们量化反超幅度。

3.2 自动调优搜什么参数？一张表说清关键搜索空间

跨芯优化最大的玄学就是参数组合太多。给了平台一个好内核，平台能否找到最优配置，直接决定了反超原生库的成败。我这边把搜索项整理成了这样一张表，供大家参考：

调优搜索的本质就是在这些参数里做组合爆炸。人肉试的话一次编译加benchmark几分钟到十几分钟，几十个组合跑下来一天就没了。FlagOS的自动调优器会自动跑网格搜索或者贝叶斯优化，根据实测吞吐反馈迭代，类似深度学习里的超参搜索。我们通常给它设定一个预算上限，比如最多跑30个配置，它把每张卡上最优的一小撮配置返回给你。

3.3 一个容易被忽略的跨芯调优杠杆：数据布局对齐

写跨芯GEMM内核，有一类坑是芯片无关的，但每张芯片踩的方式都不同——就是数据布局对齐。矩阵的leading dimension不被对齐因子整除时，tl.load在连续访存时会退化成非向量化加载，带宽立刻掉一个档次。

模型框架里的矩阵往往来自前一层算子的输出，其内存布局不是我们能控制的。所以FlagOS平台在生成内核时，有一个选项是让内核把非对齐数据先拷进对齐的共享内存缓冲区，再做对齐访存；但那相当于多了一轮拷贝开销。更聪明的做法是在系统层面调度时插入一个layout转换算子，把整个模型的GEMM输入统一成对齐布局，再连续跑多个GEMM，这样摊薄了转换成本。

这个细节我们在优化某个推理模型时效果很明显。原始输入矩阵行跨度是1012字节，不是对齐的整数倍，GEMM吞吐只有理论峰值的45%。在FlagOS里插入了一次类型为「layout_alignment」的预转换之后，GEMM整体吞吐提升到理论峰值的70%，虽然多了一次整矩阵拷贝，但整体端到端时间反而下降了接近三成。

提示：迁移跨芯GEMM时，第一步先检查输入张量的内存对齐情况。如果行跨度对芯片的自然对齐因子（通常是32字节或者128字节）取模不是0，很多芯片的加载单元都会走低效路径。先在框架侧把布局统一，再调内核参数，优先级更高。

4. 实测性能：Triton跨芯GEMM反超原生库，反超在哪里？

4.1 测量方法：不是口号式对比，而是同条件同精度

性能对比要严谨，否则容易被芯片差异、精度差异、实现方式差异误导。我的对照方法如下：

• 同一台机器上插两款目标芯片，测前用厂商自带的监控工具锁定峰值频率；
• 用同一份二进制输入数据，同一组矩阵形状，同一个输出缓存，同一套内存分配方式；
• 将原厂SDK的GEMM（记为baseline）与FlagOS+Triton生成的GEMM（记为flagged）分别跑100次，取中位数；
• 精度校验：两者的输出误差要满足1e-2以下的相对误差（fp32累加场景），否则对比无意义；
• 全部数据在两种芯片上都独立重复一遍，避免单卡个体差异。

我用一批典型的深度学习矩阵尺寸做了对比：M、N覆盖256到4096，K覆盖128到2048。挑其中代表性的几组，结果如下表（性能用相对值，以baseline在该形状下的吞吐为1.00）：

这两行K=128的小K形状意味深长：GEMM的算术强度和K成正比，K越小，计算复用率越低，整个算子的瓶颈越偏向访存而不是计算。在这种形状下，原生库久经考验的访存优化反而更有优势，而我们的Triton内核即便经过了调优，也难以在访存密集场景下追平原厂。这个结果让我认识到，所谓"反超原生"，是有适用条件的，不能无限外推。

4.2 反超背后的三个因素拆解：调度、流水线、L2亲和

为什么大K、大方块形状下Triton+FlagOS能反超原生？三个因素在起作用：

第一，线程块调度的L2亲和性。原生库的线程块调度策略是固定的，按某一种编号顺序发射线程块。而FlagOS的自动调优在搜索GROUP_M参数时，实际上是在搜索一种新的线程块发射顺序，让同一个L2缓存slice上相邻发射的线程块尽量共享A、B矩阵的数据块。实测下来GROUP_M从默认的1调成8以后，L2命中率明显上升，等效带宽提升。

第二，流水线深度与片上资源比例。原生库对这个参数有自己的偏好，但换芯片后有可能不是最优。我们的内核在调优时把num_stages从默认3调成对应芯片上的最优值，有些芯片最优值是5，多预取两级数据，延迟隐藏更充分。Triton编译器在展开循环时把加载、计算、累加完全overlap，原生库的流水线固定深度做不到这一点。

第三，对张量核的利用方式。某些芯片上，Triton编译器生成的矩阵乘指令顺序，恰好比原生库能更好地配对同地址的数据块——这其实有一点运气成分，但是自动调优某种程度上把"运气"变成了"可搜索的空间"。参数搜到了，性能就是你的。

提示：如果你的目标是"跨芯性能反超原生"，一定要先找芯片的"优势形状区间"。K中等偏大、块形状整齐的GEMM最容易突破原厂库；而小K、访存密集的形状不要恋战，直接调原生库或者接受略低的性能即可。

5. 跨芯GEMM踩坑实录：五个让内核变慢的隐藏陷阱

写Triton内核很爽，但跨芯片移植之后性能大幅回退的事我也遇到过不少次。在这里把五个最有共性的坑列出来，照方抓药的话至少能避开大半性能雷区。

5.1 精度模式的默认值坑：fp32变成tf32后误差失控

Triton在N卡上默认会把fp32的tl.dot转成tf32精度，因为N卡的tensor core对tf32有硬件加速。换到某款国产加速芯片时，它的"tf32支持"语义可能完全不同，甚至根本没有tf32模式，编译器会把tf32输入转成bf16再计算——精度直接掉两个数量级。

我在一次端到端模型测试中发现，FlagOS上生成的GEMM输出误差比baseline大了将近20倍，排查半天才发现是输入类型被自动转了bf16。

解决方式：在tl.dot之前显式用tl.inline_asm_elementwise或者数据类型转换，把输入固定成你想要的精度；如果是fp32的严格需求，可以在dot上调input_precision="ieee"（不同Triton版本的写法略有差异，可以参考当前版本文档）。总之不要相信跨芯片的"默认精度"，每张卡都手动指定。

5.2 共享内存的bank conflict：一个悄悄吃掉30%带宽的元凶

Triton隐藏了共享内存的细节，但不代表bank conflict消失了。编译器生成的共享内存访问模式，在不同芯片上bank数量不一样时，可能带病运行。某一次在芯片B上性能始终上不去，我用平台提供的分析工具看内核的访存流水线，发现load阶段的效率只有65%，明显bank conflict。

解决方式是调整数据在共享内存中的布局，Triton里通常通过tl.trans重新排布加载后的张量维度，或者给共享内存分配加上padding偏移。FlagOS平台在生成内核时可以设置swizzle参数，我在调优搜索里加上了这个变量，跑完一组搜索后，load效率提升到92%以上。

5.3num_stages过大引发的片上内存溢出与spill

这是最诱人的坑：调大num_stages确实能提升流水线深度，但超过了芯片片上共享内存容量之后，编译器会把多余的数据spill到全局内存，性能断崖式下跌，比不设流水线还慢。关键特征：原本吞吐随着num_stages递增，到某个值忽然掉到原先的六成甚至更低。

排查方式很简单：把num_stages从1跑到8，看每个档位的吞吐曲线。如果出现"先升后崩"的形状，说明芯片的片上容量就是这个算子的极限。记录崩溃点前一档的数值，作为该芯片最优值。不同芯片的最优num_stages差别很大，5.1节里提到的搜索空间设计，必须包含这个变量。

5.4 矩阵非方块的形状启发式调优失效

很多Triton教程用方阵举例，给的建议参数（比如BLOCK_M=BLOCK_N=128）也是围绕方阵调的。真实模型里M和N严重不等的情况非常多：比如M=256、N=4096，强行用固定块参数会导致大量线程块在行方向上空转，或者列方向上并行度不足。

我在调优时把BLOCK_M和BLOCK_N解耦，并且加入一个约束：num_pid_m * num_pid_n要足够大，保证芯片的并发执行能力被填满。对于M明显小于N的形状，把BLOCK_M设小一点（比如64），BLOCK_N放大（比如256），这样整体SM占用率更高。这一点在跨芯时特别重要，因为不同芯片的并发执行单元数量差异非常大，照抄主流卡的最优参数很容易翻车。

5.5 编译缓存未命中导致的调优效率灾难

跨芯片调优另一个烦心问题是编译时间。Triton的JIT编译一版内核要几秒钟到几十秒，自动调优要跑几十上百个配置，编译时间一旦叠加起来，一个晚上才调完一轮，迭代效率极低。

实践中有个有效技巧：在FlagOS里把编译缓存和调优知识库打开，让平台自动缓存"芯片+形状+参数组合"对应的编译产物和对应性能。再次调优相同或相似的矩阵形状时，直接命中缓存，编译全跳过，调优速度能提升一个数量级。跨芯迁移的算子往往会在不同shape上反复调，这个缓存机制省下的时间非常可观。

6. 现在的成绩和下一步：跨芯算子优化的边界与野心

用Triton+Triton在FlagOS上跑通GEMM反超原生后，我更关心的是：这能不能复制到其他算子上去？

我的结论是，至少有三类算子的路径是通的：

• 卷积类算子：如果能用tl.gather或tl.dot的模式把im2col折叠掉，很多卷积也能用Triton重写并在跨芯场景下追平乃至反超原生库；
• Attention类算子：flash attention的fwd内核，Triton社区已经有成熟的参考实现，跨芯移植后通过调优搜索，有希望在长序列场景下逼近甚至超过原厂SDK的融合算子；
• Elementwise与Reduce组合算子：这类算子访存密集，反超难度大，但能把跨芯一致性做好，避免每张卡各写一套的痛苦。

这次实践也让我重新思考一件事：当AI加速芯片不再只有一个生态霸主时，"算子生态"的核心竞争力，未必是每个算子都在原生库里做到极致，而是"同一份算子代码，在每张新芯片上都能自动变成最优形态"。FlagOS这类平台的价值，恰好就是把"手写调优"转变为"编译器+自动调优器去搜索"，把跨芯性能从玄学变成了可复现的流程。

如果后续能进一步把模型中所有关键算子的调优知识沉淀到平台的知识库里，让算子的形状、芯片型号、最优参数组合这些数据自动积累复用，那跨芯迁移这件事，有望从"每次从零开始"变成"一次适配、处处受益"。

最后想分享一个个人体会：在所有加速芯片的调优文档里，最不缺的就是"如何在我们的卡上写最优代码"，最缺的其实是"为什么同一段代码在不同卡上表现差异这么大"的解释。而以Triton为代表的算子语言和FlagOS这样的跨芯基础设施，某种程度上正是在重构这座桥梁：内核代码负责表达算法意图，而芯片相关的细节，则交给编译器和自动调优去解决。这条路还没走完，但至少GEMM这一步，已经足够让人看到方向了。