前阵子帮一个团队排查线上推理速度问题定位到某层算子在不同精度的实现差异非常大那趟排错让我又一次确认了一个观点想做 AI 工程绕不开“算子”这个概念。很多人写了不少模型代码但对算子始终是“用但不了解”一旦遇到性能瓶颈、显存爆炸、算子报错就很容易卡住。这篇文章我想把 AI 算子从概念到实践完整拆开从“是什么”讲到“怎么用”再讲到“怎么写、怎么优化、怎么排错”。无论你是刚入门的算法工程师、准备转做推理优化的开发者还是想了解框架内部机制的学生这篇文章应该都能给你一套完整且能上手的知识框架。1. 算子到底是什么一次把概念说清楚1.1 从一个最简单的数学操作说起“算子”这个词听起来挺唬人其实数学里早就有了。在线性代数中算子本质上是“从一个空间到另一个空间的映射”比如 y Wx b你输入一个向量或者矩阵它输出另一个向量或矩阵。深度学习里几乎所有计算都可以抽象成这种“张量进、张量出”的操作卷积、矩阵乘法、归一化、激活函数、损失计算全都是算子。真正让“算子”这个概念在工程里变得重要是因为深度学习框架普遍采用了计算图。你的模型在框架眼里就是一张由节点和边组成的图每个节点对应一个算子边对应节点之间流动的张量数据。你在 PyTorch 里调用torch.add、torch.nn.Conv2d表面上看是调了一个 Python 函数实际上就是往这张计算图里挂了一个算子节点再由后端去执行具体的计算。我见过不少同学第一次接触算子的反应是“这不是一个函数吗为什么还要单独研究”区别在于在用户视角它是一个 API在框架内部它是一个完整的执行单元包含了计算逻辑、内存分配、调度策略、硬件适配等多层内容。同一个加法在 CPU 上是一段循环加向量化指令在 GPU 上是一群线程并行执行在 NPU 上可能又是另一种指令结构。算子这个概念承担了“统一逻辑描述”和“分硬件实现”之间的桥梁作用。1.2 三类人看到的三种算子要理解算子先接受一个事实不同角色看到的算子完全不一样。用框架写模型的人看到的算子是一系列函数。比如torch.nn.functional.relu、torch.matmul、torch.nn.Conv2d他们只关心输入输出和参数。做框架开发的人看到的算子是一套需要注册、分发、适配的系统。以一个卷积为例PyTorch 内部从torch.nn.functional.conv2d出发会经历 ATen 层的统一接口、dispatch 机制的分发、然后落到某个后端实现。实现可能来自 cuDNN、MKL-DNNoneDNN也可能来自框架自带的高性能 kernel。开发者的任务就是保证这个分发正确、高效、可扩展。写硬件底层的人看到的算子则是线程、指令、寄存器、片上存储。同一个卷积在 GPU 上要怎么让线程块协作怎么复用共享内存什么时候做向量化加载这些统统决定性能。这部分工作通常由 CUDA、ROCm、OneAPI 或 NPU 工具链完成。一个经典问题是“模型跑得慢到底是谁的问题”如果没有三重视角你很难定位。比如你在 PyTorch 上调用F.conv2d慢的原因可能是没有启用cudnn.benchmark也可能是 NCHW 与 NHWC 内存布局不匹配导致框架内部发生了额外的转置也可能是卷积本身参数不适合 GPU 并行。这三类原因分别对应上面三类角色的领域。1.3 为什么算子决定了 AI 框架的天花板框架的“好用”和“强大”很大程度上取决于算子生态。一个深度学习框架能支持多少算子、每个算子在各种硬件上是否够快、是否容易扩展新硬件决定了它能在真实业务里走多远。PyTorch 能成为事实标准除了灵活的动态图一个很重要的原因是它背后有大量高质量的算子实现和成熟的扩展机制。更现实的角度是你在做模型训练和推理时大部分计算时间都花在算子执行上。一次 Transformer 训练主要时间消耗是矩阵乘法和注意力相关算子一次卷积网络的推理卷积算子可能占掉八成以上的耗时。模型的层数、参数量决定不了速度真正决定速度的是这些算子在目标硬件上跑得有多快。另外算子和显存、内存的关系极其密切。每个算子执行时都要分配输出张量如果算子切得太碎中间结果会大量占用显存算子的实现方式也决定了数据搬运次数搬来搬去可比算本身贵多了。这也是为什么业界一直强调“算子融合”。把多个算子合并成一个减少中间张量的读写往往能带来数倍加速。这些内容后面我会展开讲。2. 算子的“一生”从代码调用到硬件执行2.1 计算图里的一个节点先看宏观流程。你在 PyTorch 里写下x torch.relu(x)这一步发生在动态计算图中。动态图的好处是边执行边构建函数被调用的那一瞬间框架会创建一个用于自动微分的节点记录输入输出并立即执行底层算子。而在 TensorFlow 的静态图模式或者 TensorRT 的离线优化中算子会先被放置在一个完整的计算图里框架可以做全局优化比如合并相邻算子、重排执行顺序、消除公共子表达式然后再编译执行。这也是为什么静态图模式往往能跑得比动态图快一些因为它有机会从全局视角调配算子。这两种图模式的核心都是算子节点。不同之处在于动态图算子的“出生时间”是运行时静态图算子的“出生时间”是构建期。理解这一点之后你再看“PyTorch 为什么灵活”“TensorRT 为什么快”这类问题就会有非常清晰的感觉。2.2 一次卷积调用的完整旅程现在跟着一次 GPU 上的卷积走一遍完整流程。这有助于你把零散的知识点串起来。第一步Python 层。你调用torch.nn.Conv2d或F.conv2d得到一个输出张量。这时输入数据是四维的比如 NCHW 布局N 是 batchC 是通道H 和 W 是高和宽。第二步ATen 层。PyTorch 内部把这次调用转成 ATen 的conv2d接口。ATen 提供了一套统一张量操作接口不同后端可以各自实现同时可以通过 dispatcher 机制根据张量的 device、dtype、layout 参数选择具体实现。第三步dispatch。如果你的输入是 CUDA 张量框架会把调用分发到 CUDA 后端。这时它先查缓存看有没有适用于当前输入形状的卷积算法再调用 cuDNN 的cudnnConvolutionForward或者你自己实现的自定义 CUDA kernel。第四步硬件执行。cuDNN 会根据你的卷积参数通道数、kernel size、stride、dilation 等从几十种算法里选一种比如 implicit GEMM、Winograd、FFT 等在 GPU 上启动成千上万个线程并行计算。最后结果写回显存。这套链路每层都有性能陷阱。比如 dispatch 阶段的判断逻辑可能成为小算子调用的开销瓶颈cuDNN 算法选择如果每次形状都变可能一直触发基准测试数据布局如果和 kernel 实现不一致中间还可能插入 transpose。2.3 算子的好坏差距在哪里同样是实现一个算子不同人写出来的差距可以是数量级的。核心差异体现在几个方面。第一是访存效率。GPU 和 CPU 计算数据都有层次化存储寄存器、共享内存、缓存、显存、内存。一个算子如果能尽量把数据留在片上减少向低层存储搬运速度就会明显更快。第二是并行度。GPU 算得快靠的是大量线程同时干活。一个 kernel 如果只用了少量线程或者线程之间负载不均衡性能必然很差。矩阵乘法为什么难写就是因为要在分块、线程组织、数据复用之间找平衡。第三是计算密度。有些算子做一次加法要搬好几次数据这属于低计算密度有些算子一次加载可以做多次运算比如卷积结合 Winograd 降低乘法次数就是提升计算密度。第四是同步开销。多线程世界里同步是非常贵的。设计算子时最好让线程之间互不打扰各算各的块最后只做一次轻量同步。频繁同步的 kernel 基本跑不快。用个生活化类比解释算子就像一个外卖配送系统。如果你的策略是每个骑手只送一份餐取餐、送餐、回店循环那必然效率极低。好的策略是让一个骑手一次带一批订单规划路线甚至几个骑手分片区协作把往返次数降到最低。算子的访存优化就是做类似的事情。3. 动手实操从零实现一个算子的完整流程3.1 实操环境与工具选型讲概念讲太多容易飘这部分我们直接动手。实操需要一台带 NVIDIA GPU 的 Linux 机器装上 PyTorch 和 Triton。如果你没有 GPU也可以先把代码看懂逻辑部分在 CPU 上验证。为什么推荐先学 Triton 而不是直接上 CUDA因为 Triton 把很多底层细节托管了你不用手动管理线程块、共享内存只需要写每个线程块内部做什么剩下的交给编译器。这对于理解算子的“逻辑结构”特别友好初学阶段不容易被 CUDA 的__syncthreads、bank conflict 这些细节劝退。入门之后如果想深入再回头写 CUDA你会发现有了 Triton 的思维模型理解 CUDA 代码会顺畅很多。安装命令很简单一行搞定pip install triton torch3.2 第一个算子向量加法向量加法是最直觉的算子它把一个张量和另一个张量逐元素相加。我们用 Triton 写一个 GPU kernel。import torch import triton import triton.language as tl triton.jit def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): pid tl.program_id(axis0) block_start pid * BLOCK_SIZE offsets block_start tl.arange(0, BLOCK_SIZE) mask offsets n_elements x tl.load(x_ptr offsets, maskmask) y tl.load(y_ptr offsets, maskmask) tl.store(output_ptr offsets, x y, maskmask) def add_triton(x, y): output torch.empty_like(x) n_elements output.numel() BLOCK_SIZE 1024 grid (triton.cdiv(n_elements, BLOCK_SIZE),) add_kernel[grid](x, y, output, n_elements, BLOCK_SIZEBLOCK_SIZE) return output这段代码逻辑很短但它已经包含了一个算子最核心的骨架。看一下关键点。第一tl.program_id相当于 CUDA 里的blockIdx.x它告诉这个线程块处理哪一段数据。每个块处理BLOCK_SIZE个元素块的数量由grid决定。我们用的cdiv是向上取整除法保证元素数量不是块大小整数倍时也有足够的块来处理。第二tl.arange(0, BLOCK_SIZE)生成了块内的偏移。这里要注意第pid个块处理的元素范围是从pid * BLOCK_SIZE开始的这个偏移计算是新手最容易出错的地方。第三mask 的作用是防止越界。当总元素数不能被BLOCK_SIZE整除时最后一块会有一部分偏移超出张量长度mask保证这些位置不加载、不存储。很多新手写算子第一次跑出错误结果基本都是因为边界处理马虎。你可以用下面的代码验证结果x torch.randn(10000, devicecuda) y torch.randn(10000, devicecuda) result add_triton(x, y) assert torch.allclose(result, x y) print(passed)这个算子本身没有性能优势因为 PyTorch 的x y已经高度优化了。它最大的意义是让你看到算子的骨架取数据、算数据、存数据。所有更复杂的算子本质都是在这个骨架上增加更复杂的索引和计算逻辑。3.3 第二个算子一个迷你矩阵乘矩阵乘法是深度学习里最重要的算子也是我认为最值得亲手写一遍的算子。这里我们用 Triton 写一个基础版本。triton.jit def matmul_kernel( a_ptr, b_ptr, c_ptr, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, ): pid_m tl.program_id(axis0) pid_n tl.program_id(axis1) offs_m pid_m * BLOCK_SIZE_M tl.arange(0, BLOCK_SIZE_M) offs_n pid_n * BLOCK_SIZE_N tl.arange(0, BLOCK_SIZE_N) offs_k tl.arange(0, BLOCK_SIZE_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_SIZE_M, BLOCK_SIZE_N), dtypetl.float32) for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): a tl.load(a_ptrs) b tl.load(b_ptrs) acc tl.dot(a, b) a_ptrs BLOCK_SIZE_K * stride_ak b_ptrs BLOCK_SIZE_K * stride_bk offs_cm pid_m * BLOCK_SIZE_M tl.arange(0, BLOCK_SIZE_M) offs_cn pid_n * BLOCK_SIZE_N tl.arange(0, BLOCK_SIZE_N) c_ptrs c_ptr offs_cm[:, None] * stride_cm offs_cn[None, :] * stride_cn tl.store(c_ptrs, acc)这里我刻意省略了边界 mask为了让核心逻辑更清晰。你如果自己跑这段代码需要补上offs_m M和offs_n N的判断。这个实现展示了矩阵乘的核心策略把一个大的输出矩阵切成BLOCK_SIZE_M x BLOCK_SIZE_N的小块每个程序负责其中一个块的计算。内部循环沿着 K 维做累加每轮加载一小块 A 和一小块 B用tl.dot做乘加。这种切块方式使得数据可以在片上被反复使用避免每个输出元素都去显存重新读一遍完整输入。矩阵乘的优化空间非常大甚至可以说 GPU 性能优化的巅峰都在矩阵乘里。等你经历过BLOCK_SIZE调整带来的性能差异再回头看 cuBLAS 和 cuDNN就会知道那些库的价值有多大。3.4 生产级算子的隐藏步骤上面写的都是“裸算子”只在计算层面成立。真正放入生产框架里的算子通常还要穿上好几层外衣这里记录一下我踩过的坑。第一是 Shape 推导。框架在构建计算图时必须知道输入输出张量的形状。你不能只在 Python 端计算出结果就完事还要实现一个 shape inference 逻辑让框架在真正执行前就能分配内存、做静态图优化。第二是算子注册。以 PyTorch 为例你写了一个 C 实现的算子还需要通过TORCH_LIBRARY宏注册到框架里告诉框架“有这个算子它的 schema 是什么”。不注册框架就不知道它有自动微分规则、不知道它能不能参与图优化。很多初学自定义算子的人在这步翻车。第三是自动微分。训练框架里的算子往往要成对出现一个是 forward 计算一个是 backward 计算。backward 要能够根据输出的梯度算出对输入的梯度。以矩阵乘为例设 Y A B那么 dA dY * B^TdB A^T * dY。你需要自己推导并实现这些公式。第四是内存格式策略。同样的逻辑对 NCHW 和 NHWC 输入可能都要支持或者明确声明只支持一种。PyTorch 内部经常因为算子对内存格式的支持不同导致实际执行时插入显式的 transpose 操作白白浪费性能。我不建议你一上来就搞全套生产级算子但务必知道“能算”和“能用”之间还隔着 shape 推导、注册、反向实现、格式兼容这四道工序。这也是网上很多教程只讲 kernel 不讲工程落地的原因后者确实更枯燥但实际工作中绕不开。4. 算子性能优化从能跑到跑得快4.1 先判断算子属于哪一类访存密集型还是计算密集型优化算子前先判断它是哪种类型。这个方法极其简单又极其实用。访存密集型算子的特点是计算很简单但数据量大瓶颈在于数据搬运。典型如逐元素算子加法、乘法、ReLU、transpose、copy。这些算子几乎不做计算时间都花在把数据从内存搬到计算单元再搬回去。优化这类算子的核心是减少访存次数、提高访存效率比如利用向量化加载、保证合并访问。计算密集型算子的特点是计算量远大于访存量瓶颈在计算本身。典型如大矩阵乘、卷积、FFT。这类算子优化的核心是提升计算效率比如利用张量核心、降低无效计算、调整分块策略以提升数据复用。怎么判断算计算强度也就是“总计算量/总访问字节数”。如果计算强度低就是访存密集如果高就是计算密集。这个思路来自经典的 roofline model它可以告诉你一个算子在特定硬件上性能上限是访存限制还是计算限制。4.2 算子融合为什么是“免费的午餐”算子融合是我工作中收益最大、见效最快的优化手段。原理是把多个算子合并成一个 kernel 执行减少中间张量的分配和读写。举个例子Transformer 的注意力里常见的操作序列是QK^T之后接 softmax再接softmax_result V。标准的实现里QK^T得到一个中间矩阵写回显存softmax 读取它计算后再次写回最后一个矩阵乘再读一次。如果你把三者融合成一个算子让QK^T的结果留在片上寄存器或者共享内存里直接做 softmax再直接和 V 相乘省掉了两轮完整的数据搬运。大矩阵的读写往往比计算本身慢得多所以这种融合的加速比非常可观甚至可以达到两三倍。另一个典型案例是卷积和批量归一化的融合。推理阶段 BN 的均值、方差都是固定的它可以被数学上等价地折叠进卷积层的权重和偏置中。网上很多推理优化教程都提到过这个技巧本质就是算子融合。你算一次就发现训练和推理的差异远不止eval()那么简单融合后的网络计算量完全不变但因为少了好几轮中间张量读写速度提升非常明显。融合也不是越多越好。有两个问题一是融合后的算子可能过于复杂难维护、难调优二是当融合引入额外分支、padding 或特殊边界处理时可能降低计算密度。我的经验是优先融合访存密集型算子为主体的子图这类融合收益最明显。4.3 让机器帮你选参数自动调优同样的算子不同BLOCK_SIZE、不同循环展开策略、不同线程数性能差异常能达到数倍。手动一点一点试很累业界很早就开始做自动调优。Triton 自带简单的triton.autotune装饰器可以让你给关键参数准备一组候选值运行时自动选择最好的组合。cuDNN 也有 benchmark 模式PyTorch 里开启torch.backends.cudnn.benchmark True后它会针对当前输入形状去测试一批卷积算法选出最快的一个。TVM 的 Ansor、TensorRT 的自动 tuner 做得更深能在算子搜索空间中组合出接近手写优化的实现。我建议初学时不要过度依赖自动调优先把 BLOCK_SIZE、循环切分这些概念理解透因为自动调优的参数含义不清楚时你根本不知道它在试什么、结果是否可信。4.4 一个容易被忽视的问题小算子的调度开销算子不光有“计算时间”还有“调度时间”。当一个算子输入很小时比如一次add只处理几百个元素那么 kernel 启动的开销、框架 dispatch 的开销可能比计算本身还大。这种情况在推理场景大量出现尤其当模型结构很碎片化时。优化策略包括把小算子合并成大算子、用图优化把子图融合成一个 kernel、或者切换到推理引擎如 TensorRT、onnxruntime 来避开框架每步调度的开销。我自己遇到过 PyTorch 推理时一个简单的逐元素操作占掉大量时间融合进前一个算子后整体时间直接减半的案例。5. 踩坑实录我在算子开发与调试中遇到的典型问题5.1 三个印象最深的“低级”错误先介绍三个我自己实际遇到过、而且网上经常有人问的典型错误。第一个是索引算错。写向量加法 kernel 时最后一个 block 的偏移处理不对导致读取越界或写错位置。这类错误最坑的地方在于它不一定崩有时只是结果悄悄错几个数。后来我每次写 kernel 都会认真检查 mask 逻辑并且用torch.allclose对随机输入做验证。第二个是忽略了不同内存布局带来的转置开销。我在一个项目里用了 channels_last 格式的输入但某个自定义算子内部按 NCHW 逻辑实现结果每次调用前框架自动插入 transpose性能直接掉了一半。后来养成了习惯写算子前先确认输入输出布局并且在性能测试时对比是否出现了额外转换。第三个是一上来就过度融合。我有一段时间走火入魔试图把所有能融合的算子全塞进一个 kernel。结果 kernel 逻辑太复杂边界分支繁多最终性能不仅没提升反而比分离实现慢。教训是融合要顺应数据的访存模式和计算模式不是简单物理合并。5.2 数值一致性与正确性验证算子开发里写对功能只是第一步保证数值正确才是细致活。常见做法是对比同一个算子的不同实现比如你的 GPU 实现对比 PyTorch CPU 实现设置一定范围内的相对误差和绝对误差容限。如果涉及到变精度计算比如混合精度训练用了 FP16那么验证误差容限也需要相应放宽。PyTorch 里可以直接用torch.allclose设置rtol和atol。对自动微分相关的算子还可以用torch.autograd.gradcheck它会用数值微分法验证你写的 backward 是否正确这是自定义算子开发中非常实用的工具。5.3 性能排查工具怎么用排查算子的性能问题先分清是什么层面。如果是框架层 dispatch 开销可以用 PyTorch profiler 看每个 op 的时间占比如果是 kernel 内部慢可以用 NVIDIA Nsight Computencu看 kernel 的访存吞吐、计算吞吐、占用率等指标如果是数据搬运和同步问题可以用 Nsight Systems 看时间线中 kernel 之间的空隙。我最常用的流程是先用 profiler 定位耗时算子再对可疑 kernel 跑 ncu看它的 memory throughput 和 compute throughput 哪个更接近上限。如果 memory 接近上限说明访存优化空间不大看看能否减少数据量或融合如果 compute 接近上限看看能否用张量核心或降低精度。这套流程帮我快速定位过好几个问题值得养成习惯。下面用一个速查表总结常见现象、可能原因和入手方向。表现可能原因排查手段算子结果个别数值错误索引越界、mask 写错、没有考虑边界对比 CPU 实现检查边界偏移性能远低于预期未启用向量化、非合并访存、kernel 启动太频繁Nsight Compute 查看访存吞吐数据布局引起额外开销框架插入了 transpose 算子PyTorch profiler 看算子序列小算子调用开销占比高调度开销大于计算开销检查算子耗时分布考虑融合FP16 结果误差超标精度不足、计算顺序不同放宽 rtol/atol切换 FP32 验证显存占用过高中间张量过多、算子切分过碎检查算子图做算子融合5.4 调试时的几个好用技巧这里再分享几个小技巧。第一个是我习惯在自定义算子入口处用 Python 侧断言验证输入 shape 和 dtype。很多算子的 hidden bug 来自“输入其实是 int64”或者“某个维度的顺序错了”早点断言能节省大量排查时间。第二个是在开发初期先用小尺寸输入调正确性比如 1x16、2x32 这种规模方便手动核对结果跑通后再扩展到真实尺寸。我见过有人在 1024x1024 的大矩阵上调 bug每一步都要靠辅助打印非常痛苦。第三个是保存好基线结果。做一个算子优化之前先把当前正确的结果和性能记录保存下来。这样你每做一次修改都能快速对比是否变快、是否变慢、结果是否依然正确。没有基线优化很容易变成无头苍蝇。6. 走向进阶算子开发者的学习路径6.1 从框架层到硬件层的完整地图如果你看完前文决定深入算子方向我先给你画一张学习地图。第一个层次是框架层。你要熟悉 PyTorch 的 dispatch 机制、自定义 op 的写法、ATen 的作用。这一层让你知道算子怎么被组织进框架生态。第二个层次是编程模型也就是 Triton 或 CUDA。Triton 适合快速理解算子的数据并行结构CUDA 让你理解硬件细节比如线程层级、共享内存、寄存器分配、同步原语。我建议两个都学先用 Triton 找感觉再啃 CUDA 深化。第三个层次是高性能算子库包括 cutlass、cuDNN、oneDNN、FlashAttention 这些开源实现。不要当成黑盒而是当成最好的教材里面浓缩了几十年并行计算的优化经验。第四个层次是编译优化。TVM、XLA、Triton 编译器如何把高级算子描述转换为底层 kernel如何做自动调度。这个层次属于进阶中的进阶但能让你对算子的理解从“实现”跃迁到“生成”。6.2 由简到难的项目练手顺序听再多方法论都不如亲手写代码。练手顺序我自己比较推荐这样安排先写向量加法掌握最基本的 kernel 结构和边界处理然后写 reduce 类算子比如求和、softmax这能让你学到跨线程协作和归约的写法接着写矩阵乘这是终身受益的核心基本功再看 Attention 或 GELU 这类模型常用算子如何融合优化最后尝试做一个小型算子库把多个算子统一管理、统一测试。每一步都要做两件事验证正确性同时记录性能数据。实践里最有效的进步方式是“先抄袭再超越”把一个成熟实现完全读懂改成自己的版本再逐步替换部分逻辑观察性能变化。6.3 几个务实的工程建议给几个相对务实的建议。第一算子开发务必有自动化测试。算子的边界条件太多手动验证根本测不全。至少要覆盖不同 shape、不同 dtype、不同布局。第二性能测试要和正确性测试分开因为性能测试的 batch 往往更大逻辑也特殊。第三做算子不是写论文与其追求极端性能不如追求稳定的、可维护的、能融合进系统里的实现。一个性能不错但很难维护的算子长期来看是团队的负担。另外我特别建议多关注硬件特性。同样一个算子在不同 GPU 上表现可能完全不同Tensor Core 是否可用、L2 缓存大小、内存带宽高低都会影响你的优化方案。只看框架层很多问题永远解释不了。写在最后一点真实体会做算子开发这几年来我最大的体会是算子就是深度学习的“汇编语言”它离算法很远离硬件很近但它决定了你的想法最终能不能高效落地。最初学 CUDA 和 Triton 时我也被各种概念砸得头晕什么 grid、block、shared memory、memory coalescing每个词背后都是一大堆背景知识。后来慢慢发现只要抓住“数据怎么组织、访存怎么安排、并行怎么切分”这三件事绝大多数算子都是可以理解、可以优化的。如果你正在入门我的建议是不要停留在“会用框架调用算子”的层面抽一个周末用 Triton 写一遍向量加法再写一遍矩阵乘然后尝试把一个简单的两算子序列融合成一个 kernel。这个过程会逼迫你建立完整的、从框架到底层的思维模型。这个模型一旦建立起来等你再遇到性能优化、算子报错、显存瓶颈这些问题时就不再是瞎子摸象了。 AI算子从概念到实践：原理、实现与性能优化指南
前阵子帮一个团队排查线上推理速度问题定位到某层算子在不同精度的实现差异非常大那趟排错让我又一次确认了一个观点想做 AI 工程绕不开“算子”这个概念。很多人写了不少模型代码但对算子始终是“用但不了解”一旦遇到性能瓶颈、显存爆炸、算子报错就很容易卡住。这篇文章我想把 AI 算子从概念到实践完整拆开从“是什么”讲到“怎么用”再讲到“怎么写、怎么优化、怎么排错”。无论你是刚入门的算法工程师、准备转做推理优化的开发者还是想了解框架内部机制的学生这篇文章应该都能给你一套完整且能上手的知识框架。1. 算子到底是什么一次把概念说清楚1.1 从一个最简单的数学操作说起“算子”这个词听起来挺唬人其实数学里早就有了。在线性代数中算子本质上是“从一个空间到另一个空间的映射”比如 y Wx b你输入一个向量或者矩阵它输出另一个向量或矩阵。深度学习里几乎所有计算都可以抽象成这种“张量进、张量出”的操作卷积、矩阵乘法、归一化、激活函数、损失计算全都是算子。真正让“算子”这个概念在工程里变得重要是因为深度学习框架普遍采用了计算图。你的模型在框架眼里就是一张由节点和边组成的图每个节点对应一个算子边对应节点之间流动的张量数据。你在 PyTorch 里调用torch.add、torch.nn.Conv2d表面上看是调了一个 Python 函数实际上就是往这张计算图里挂了一个算子节点再由后端去执行具体的计算。我见过不少同学第一次接触算子的反应是“这不是一个函数吗为什么还要单独研究”区别在于在用户视角它是一个 API在框架内部它是一个完整的执行单元包含了计算逻辑、内存分配、调度策略、硬件适配等多层内容。同一个加法在 CPU 上是一段循环加向量化指令在 GPU 上是一群线程并行执行在 NPU 上可能又是另一种指令结构。算子这个概念承担了“统一逻辑描述”和“分硬件实现”之间的桥梁作用。1.2 三类人看到的三种算子要理解算子先接受一个事实不同角色看到的算子完全不一样。用框架写模型的人看到的算子是一系列函数。比如torch.nn.functional.relu、torch.matmul、torch.nn.Conv2d他们只关心输入输出和参数。做框架开发的人看到的算子是一套需要注册、分发、适配的系统。以一个卷积为例PyTorch 内部从torch.nn.functional.conv2d出发会经历 ATen 层的统一接口、dispatch 机制的分发、然后落到某个后端实现。实现可能来自 cuDNN、MKL-DNNoneDNN也可能来自框架自带的高性能 kernel。开发者的任务就是保证这个分发正确、高效、可扩展。写硬件底层的人看到的算子则是线程、指令、寄存器、片上存储。同一个卷积在 GPU 上要怎么让线程块协作怎么复用共享内存什么时候做向量化加载这些统统决定性能。这部分工作通常由 CUDA、ROCm、OneAPI 或 NPU 工具链完成。一个经典问题是“模型跑得慢到底是谁的问题”如果没有三重视角你很难定位。比如你在 PyTorch 上调用F.conv2d慢的原因可能是没有启用cudnn.benchmark也可能是 NCHW 与 NHWC 内存布局不匹配导致框架内部发生了额外的转置也可能是卷积本身参数不适合 GPU 并行。这三类原因分别对应上面三类角色的领域。1.3 为什么算子决定了 AI 框架的天花板框架的“好用”和“强大”很大程度上取决于算子生态。一个深度学习框架能支持多少算子、每个算子在各种硬件上是否够快、是否容易扩展新硬件决定了它能在真实业务里走多远。PyTorch 能成为事实标准除了灵活的动态图一个很重要的原因是它背后有大量高质量的算子实现和成熟的扩展机制。更现实的角度是你在做模型训练和推理时大部分计算时间都花在算子执行上。一次 Transformer 训练主要时间消耗是矩阵乘法和注意力相关算子一次卷积网络的推理卷积算子可能占掉八成以上的耗时。模型的层数、参数量决定不了速度真正决定速度的是这些算子在目标硬件上跑得有多快。另外算子和显存、内存的关系极其密切。每个算子执行时都要分配输出张量如果算子切得太碎中间结果会大量占用显存算子的实现方式也决定了数据搬运次数搬来搬去可比算本身贵多了。这也是为什么业界一直强调“算子融合”。把多个算子合并成一个减少中间张量的读写往往能带来数倍加速。这些内容后面我会展开讲。2. 算子的“一生”从代码调用到硬件执行2.1 计算图里的一个节点先看宏观流程。你在 PyTorch 里写下x torch.relu(x)这一步发生在动态计算图中。动态图的好处是边执行边构建函数被调用的那一瞬间框架会创建一个用于自动微分的节点记录输入输出并立即执行底层算子。而在 TensorFlow 的静态图模式或者 TensorRT 的离线优化中算子会先被放置在一个完整的计算图里框架可以做全局优化比如合并相邻算子、重排执行顺序、消除公共子表达式然后再编译执行。这也是为什么静态图模式往往能跑得比动态图快一些因为它有机会从全局视角调配算子。这两种图模式的核心都是算子节点。不同之处在于动态图算子的“出生时间”是运行时静态图算子的“出生时间”是构建期。理解这一点之后你再看“PyTorch 为什么灵活”“TensorRT 为什么快”这类问题就会有非常清晰的感觉。2.2 一次卷积调用的完整旅程现在跟着一次 GPU 上的卷积走一遍完整流程。这有助于你把零散的知识点串起来。第一步Python 层。你调用torch.nn.Conv2d或F.conv2d得到一个输出张量。这时输入数据是四维的比如 NCHW 布局N 是 batchC 是通道H 和 W 是高和宽。第二步ATen 层。PyTorch 内部把这次调用转成 ATen 的conv2d接口。ATen 提供了一套统一张量操作接口不同后端可以各自实现同时可以通过 dispatcher 机制根据张量的 device、dtype、layout 参数选择具体实现。第三步dispatch。如果你的输入是 CUDA 张量框架会把调用分发到 CUDA 后端。这时它先查缓存看有没有适用于当前输入形状的卷积算法再调用 cuDNN 的cudnnConvolutionForward或者你自己实现的自定义 CUDA kernel。第四步硬件执行。cuDNN 会根据你的卷积参数通道数、kernel size、stride、dilation 等从几十种算法里选一种比如 implicit GEMM、Winograd、FFT 等在 GPU 上启动成千上万个线程并行计算。最后结果写回显存。这套链路每层都有性能陷阱。比如 dispatch 阶段的判断逻辑可能成为小算子调用的开销瓶颈cuDNN 算法选择如果每次形状都变可能一直触发基准测试数据布局如果和 kernel 实现不一致中间还可能插入 transpose。2.3 算子的好坏差距在哪里同样是实现一个算子不同人写出来的差距可以是数量级的。核心差异体现在几个方面。第一是访存效率。GPU 和 CPU 计算数据都有层次化存储寄存器、共享内存、缓存、显存、内存。一个算子如果能尽量把数据留在片上减少向低层存储搬运速度就会明显更快。第二是并行度。GPU 算得快靠的是大量线程同时干活。一个 kernel 如果只用了少量线程或者线程之间负载不均衡性能必然很差。矩阵乘法为什么难写就是因为要在分块、线程组织、数据复用之间找平衡。第三是计算密度。有些算子做一次加法要搬好几次数据这属于低计算密度有些算子一次加载可以做多次运算比如卷积结合 Winograd 降低乘法次数就是提升计算密度。第四是同步开销。多线程世界里同步是非常贵的。设计算子时最好让线程之间互不打扰各算各的块最后只做一次轻量同步。频繁同步的 kernel 基本跑不快。用个生活化类比解释算子就像一个外卖配送系统。如果你的策略是每个骑手只送一份餐取餐、送餐、回店循环那必然效率极低。好的策略是让一个骑手一次带一批订单规划路线甚至几个骑手分片区协作把往返次数降到最低。算子的访存优化就是做类似的事情。3. 动手实操从零实现一个算子的完整流程3.1 实操环境与工具选型讲概念讲太多容易飘这部分我们直接动手。实操需要一台带 NVIDIA GPU 的 Linux 机器装上 PyTorch 和 Triton。如果你没有 GPU也可以先把代码看懂逻辑部分在 CPU 上验证。为什么推荐先学 Triton 而不是直接上 CUDA因为 Triton 把很多底层细节托管了你不用手动管理线程块、共享内存只需要写每个线程块内部做什么剩下的交给编译器。这对于理解算子的“逻辑结构”特别友好初学阶段不容易被 CUDA 的__syncthreads、bank conflict 这些细节劝退。入门之后如果想深入再回头写 CUDA你会发现有了 Triton 的思维模型理解 CUDA 代码会顺畅很多。安装命令很简单一行搞定pip install triton torch3.2 第一个算子向量加法向量加法是最直觉的算子它把一个张量和另一个张量逐元素相加。我们用 Triton 写一个 GPU kernel。import torch import triton import triton.language as tl triton.jit def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): pid tl.program_id(axis0) block_start pid * BLOCK_SIZE offsets block_start tl.arange(0, BLOCK_SIZE) mask offsets n_elements x tl.load(x_ptr offsets, maskmask) y tl.load(y_ptr offsets, maskmask) tl.store(output_ptr offsets, x y, maskmask) def add_triton(x, y): output torch.empty_like(x) n_elements output.numel() BLOCK_SIZE 1024 grid (triton.cdiv(n_elements, BLOCK_SIZE),) add_kernel[grid](x, y, output, n_elements, BLOCK_SIZEBLOCK_SIZE) return output这段代码逻辑很短但它已经包含了一个算子最核心的骨架。看一下关键点。第一tl.program_id相当于 CUDA 里的blockIdx.x它告诉这个线程块处理哪一段数据。每个块处理BLOCK_SIZE个元素块的数量由grid决定。我们用的cdiv是向上取整除法保证元素数量不是块大小整数倍时也有足够的块来处理。第二tl.arange(0, BLOCK_SIZE)生成了块内的偏移。这里要注意第pid个块处理的元素范围是从pid * BLOCK_SIZE开始的这个偏移计算是新手最容易出错的地方。第三mask 的作用是防止越界。当总元素数不能被BLOCK_SIZE整除时最后一块会有一部分偏移超出张量长度mask保证这些位置不加载、不存储。很多新手写算子第一次跑出错误结果基本都是因为边界处理马虎。你可以用下面的代码验证结果x torch.randn(10000, devicecuda) y torch.randn(10000, devicecuda) result add_triton(x, y) assert torch.allclose(result, x y) print(passed)这个算子本身没有性能优势因为 PyTorch 的x y已经高度优化了。它最大的意义是让你看到算子的骨架取数据、算数据、存数据。所有更复杂的算子本质都是在这个骨架上增加更复杂的索引和计算逻辑。3.3 第二个算子一个迷你矩阵乘矩阵乘法是深度学习里最重要的算子也是我认为最值得亲手写一遍的算子。这里我们用 Triton 写一个基础版本。triton.jit def matmul_kernel( a_ptr, b_ptr, c_ptr, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, ): pid_m tl.program_id(axis0) pid_n tl.program_id(axis1) offs_m pid_m * BLOCK_SIZE_M tl.arange(0, BLOCK_SIZE_M) offs_n pid_n * BLOCK_SIZE_N tl.arange(0, BLOCK_SIZE_N) offs_k tl.arange(0, BLOCK_SIZE_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_SIZE_M, BLOCK_SIZE_N), dtypetl.float32) for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): a tl.load(a_ptrs) b tl.load(b_ptrs) acc tl.dot(a, b) a_ptrs BLOCK_SIZE_K * stride_ak b_ptrs BLOCK_SIZE_K * stride_bk offs_cm pid_m * BLOCK_SIZE_M tl.arange(0, BLOCK_SIZE_M) offs_cn pid_n * BLOCK_SIZE_N tl.arange(0, BLOCK_SIZE_N) c_ptrs c_ptr offs_cm[:, None] * stride_cm offs_cn[None, :] * stride_cn tl.store(c_ptrs, acc)这里我刻意省略了边界 mask为了让核心逻辑更清晰。你如果自己跑这段代码需要补上offs_m M和offs_n N的判断。这个实现展示了矩阵乘的核心策略把一个大的输出矩阵切成BLOCK_SIZE_M x BLOCK_SIZE_N的小块每个程序负责其中一个块的计算。内部循环沿着 K 维做累加每轮加载一小块 A 和一小块 B用tl.dot做乘加。这种切块方式使得数据可以在片上被反复使用避免每个输出元素都去显存重新读一遍完整输入。矩阵乘的优化空间非常大甚至可以说 GPU 性能优化的巅峰都在矩阵乘里。等你经历过BLOCK_SIZE调整带来的性能差异再回头看 cuBLAS 和 cuDNN就会知道那些库的价值有多大。3.4 生产级算子的隐藏步骤上面写的都是“裸算子”只在计算层面成立。真正放入生产框架里的算子通常还要穿上好几层外衣这里记录一下我踩过的坑。第一是 Shape 推导。框架在构建计算图时必须知道输入输出张量的形状。你不能只在 Python 端计算出结果就完事还要实现一个 shape inference 逻辑让框架在真正执行前就能分配内存、做静态图优化。第二是算子注册。以 PyTorch 为例你写了一个 C 实现的算子还需要通过TORCH_LIBRARY宏注册到框架里告诉框架“有这个算子它的 schema 是什么”。不注册框架就不知道它有自动微分规则、不知道它能不能参与图优化。很多初学自定义算子的人在这步翻车。第三是自动微分。训练框架里的算子往往要成对出现一个是 forward 计算一个是 backward 计算。backward 要能够根据输出的梯度算出对输入的梯度。以矩阵乘为例设 Y A B那么 dA dY * B^TdB A^T * dY。你需要自己推导并实现这些公式。第四是内存格式策略。同样的逻辑对 NCHW 和 NHWC 输入可能都要支持或者明确声明只支持一种。PyTorch 内部经常因为算子对内存格式的支持不同导致实际执行时插入显式的 transpose 操作白白浪费性能。我不建议你一上来就搞全套生产级算子但务必知道“能算”和“能用”之间还隔着 shape 推导、注册、反向实现、格式兼容这四道工序。这也是网上很多教程只讲 kernel 不讲工程落地的原因后者确实更枯燥但实际工作中绕不开。4. 算子性能优化从能跑到跑得快4.1 先判断算子属于哪一类访存密集型还是计算密集型优化算子前先判断它是哪种类型。这个方法极其简单又极其实用。访存密集型算子的特点是计算很简单但数据量大瓶颈在于数据搬运。典型如逐元素算子加法、乘法、ReLU、transpose、copy。这些算子几乎不做计算时间都花在把数据从内存搬到计算单元再搬回去。优化这类算子的核心是减少访存次数、提高访存效率比如利用向量化加载、保证合并访问。计算密集型算子的特点是计算量远大于访存量瓶颈在计算本身。典型如大矩阵乘、卷积、FFT。这类算子优化的核心是提升计算效率比如利用张量核心、降低无效计算、调整分块策略以提升数据复用。怎么判断算计算强度也就是“总计算量/总访问字节数”。如果计算强度低就是访存密集如果高就是计算密集。这个思路来自经典的 roofline model它可以告诉你一个算子在特定硬件上性能上限是访存限制还是计算限制。4.2 算子融合为什么是“免费的午餐”算子融合是我工作中收益最大、见效最快的优化手段。原理是把多个算子合并成一个 kernel 执行减少中间张量的分配和读写。举个例子Transformer 的注意力里常见的操作序列是QK^T之后接 softmax再接softmax_result V。标准的实现里QK^T得到一个中间矩阵写回显存softmax 读取它计算后再次写回最后一个矩阵乘再读一次。如果你把三者融合成一个算子让QK^T的结果留在片上寄存器或者共享内存里直接做 softmax再直接和 V 相乘省掉了两轮完整的数据搬运。大矩阵的读写往往比计算本身慢得多所以这种融合的加速比非常可观甚至可以达到两三倍。另一个典型案例是卷积和批量归一化的融合。推理阶段 BN 的均值、方差都是固定的它可以被数学上等价地折叠进卷积层的权重和偏置中。网上很多推理优化教程都提到过这个技巧本质就是算子融合。你算一次就发现训练和推理的差异远不止eval()那么简单融合后的网络计算量完全不变但因为少了好几轮中间张量读写速度提升非常明显。融合也不是越多越好。有两个问题一是融合后的算子可能过于复杂难维护、难调优二是当融合引入额外分支、padding 或特殊边界处理时可能降低计算密度。我的经验是优先融合访存密集型算子为主体的子图这类融合收益最明显。4.3 让机器帮你选参数自动调优同样的算子不同BLOCK_SIZE、不同循环展开策略、不同线程数性能差异常能达到数倍。手动一点一点试很累业界很早就开始做自动调优。Triton 自带简单的triton.autotune装饰器可以让你给关键参数准备一组候选值运行时自动选择最好的组合。cuDNN 也有 benchmark 模式PyTorch 里开启torch.backends.cudnn.benchmark True后它会针对当前输入形状去测试一批卷积算法选出最快的一个。TVM 的 Ansor、TensorRT 的自动 tuner 做得更深能在算子搜索空间中组合出接近手写优化的实现。我建议初学时不要过度依赖自动调优先把 BLOCK_SIZE、循环切分这些概念理解透因为自动调优的参数含义不清楚时你根本不知道它在试什么、结果是否可信。4.4 一个容易被忽视的问题小算子的调度开销算子不光有“计算时间”还有“调度时间”。当一个算子输入很小时比如一次add只处理几百个元素那么 kernel 启动的开销、框架 dispatch 的开销可能比计算本身还大。这种情况在推理场景大量出现尤其当模型结构很碎片化时。优化策略包括把小算子合并成大算子、用图优化把子图融合成一个 kernel、或者切换到推理引擎如 TensorRT、onnxruntime 来避开框架每步调度的开销。我自己遇到过 PyTorch 推理时一个简单的逐元素操作占掉大量时间融合进前一个算子后整体时间直接减半的案例。5. 踩坑实录我在算子开发与调试中遇到的典型问题5.1 三个印象最深的“低级”错误先介绍三个我自己实际遇到过、而且网上经常有人问的典型错误。第一个是索引算错。写向量加法 kernel 时最后一个 block 的偏移处理不对导致读取越界或写错位置。这类错误最坑的地方在于它不一定崩有时只是结果悄悄错几个数。后来我每次写 kernel 都会认真检查 mask 逻辑并且用torch.allclose对随机输入做验证。第二个是忽略了不同内存布局带来的转置开销。我在一个项目里用了 channels_last 格式的输入但某个自定义算子内部按 NCHW 逻辑实现结果每次调用前框架自动插入 transpose性能直接掉了一半。后来养成了习惯写算子前先确认输入输出布局并且在性能测试时对比是否出现了额外转换。第三个是一上来就过度融合。我有一段时间走火入魔试图把所有能融合的算子全塞进一个 kernel。结果 kernel 逻辑太复杂边界分支繁多最终性能不仅没提升反而比分离实现慢。教训是融合要顺应数据的访存模式和计算模式不是简单物理合并。5.2 数值一致性与正确性验证算子开发里写对功能只是第一步保证数值正确才是细致活。常见做法是对比同一个算子的不同实现比如你的 GPU 实现对比 PyTorch CPU 实现设置一定范围内的相对误差和绝对误差容限。如果涉及到变精度计算比如混合精度训练用了 FP16那么验证误差容限也需要相应放宽。PyTorch 里可以直接用torch.allclose设置rtol和atol。对自动微分相关的算子还可以用torch.autograd.gradcheck它会用数值微分法验证你写的 backward 是否正确这是自定义算子开发中非常实用的工具。5.3 性能排查工具怎么用排查算子的性能问题先分清是什么层面。如果是框架层 dispatch 开销可以用 PyTorch profiler 看每个 op 的时间占比如果是 kernel 内部慢可以用 NVIDIA Nsight Computencu看 kernel 的访存吞吐、计算吞吐、占用率等指标如果是数据搬运和同步问题可以用 Nsight Systems 看时间线中 kernel 之间的空隙。我最常用的流程是先用 profiler 定位耗时算子再对可疑 kernel 跑 ncu看它的 memory throughput 和 compute throughput 哪个更接近上限。如果 memory 接近上限说明访存优化空间不大看看能否减少数据量或融合如果 compute 接近上限看看能否用张量核心或降低精度。这套流程帮我快速定位过好几个问题值得养成习惯。下面用一个速查表总结常见现象、可能原因和入手方向。表现可能原因排查手段算子结果个别数值错误索引越界、mask 写错、没有考虑边界对比 CPU 实现检查边界偏移性能远低于预期未启用向量化、非合并访存、kernel 启动太频繁Nsight Compute 查看访存吞吐数据布局引起额外开销框架插入了 transpose 算子PyTorch profiler 看算子序列小算子调用开销占比高调度开销大于计算开销检查算子耗时分布考虑融合FP16 结果误差超标精度不足、计算顺序不同放宽 rtol/atol切换 FP32 验证显存占用过高中间张量过多、算子切分过碎检查算子图做算子融合5.4 调试时的几个好用技巧这里再分享几个小技巧。第一个是我习惯在自定义算子入口处用 Python 侧断言验证输入 shape 和 dtype。很多算子的 hidden bug 来自“输入其实是 int64”或者“某个维度的顺序错了”早点断言能节省大量排查时间。第二个是在开发初期先用小尺寸输入调正确性比如 1x16、2x32 这种规模方便手动核对结果跑通后再扩展到真实尺寸。我见过有人在 1024x1024 的大矩阵上调 bug每一步都要靠辅助打印非常痛苦。第三个是保存好基线结果。做一个算子优化之前先把当前正确的结果和性能记录保存下来。这样你每做一次修改都能快速对比是否变快、是否变慢、结果是否依然正确。没有基线优化很容易变成无头苍蝇。6. 走向进阶算子开发者的学习路径6.1 从框架层到硬件层的完整地图如果你看完前文决定深入算子方向我先给你画一张学习地图。第一个层次是框架层。你要熟悉 PyTorch 的 dispatch 机制、自定义 op 的写法、ATen 的作用。这一层让你知道算子怎么被组织进框架生态。第二个层次是编程模型也就是 Triton 或 CUDA。Triton 适合快速理解算子的数据并行结构CUDA 让你理解硬件细节比如线程层级、共享内存、寄存器分配、同步原语。我建议两个都学先用 Triton 找感觉再啃 CUDA 深化。第三个层次是高性能算子库包括 cutlass、cuDNN、oneDNN、FlashAttention 这些开源实现。不要当成黑盒而是当成最好的教材里面浓缩了几十年并行计算的优化经验。第四个层次是编译优化。TVM、XLA、Triton 编译器如何把高级算子描述转换为底层 kernel如何做自动调度。这个层次属于进阶中的进阶但能让你对算子的理解从“实现”跃迁到“生成”。6.2 由简到难的项目练手顺序听再多方法论都不如亲手写代码。练手顺序我自己比较推荐这样安排先写向量加法掌握最基本的 kernel 结构和边界处理然后写 reduce 类算子比如求和、softmax这能让你学到跨线程协作和归约的写法接着写矩阵乘这是终身受益的核心基本功再看 Attention 或 GELU 这类模型常用算子如何融合优化最后尝试做一个小型算子库把多个算子统一管理、统一测试。每一步都要做两件事验证正确性同时记录性能数据。实践里最有效的进步方式是“先抄袭再超越”把一个成熟实现完全读懂改成自己的版本再逐步替换部分逻辑观察性能变化。6.3 几个务实的工程建议给几个相对务实的建议。第一算子开发务必有自动化测试。算子的边界条件太多手动验证根本测不全。至少要覆盖不同 shape、不同 dtype、不同布局。第二性能测试要和正确性测试分开因为性能测试的 batch 往往更大逻辑也特殊。第三做算子不是写论文与其追求极端性能不如追求稳定的、可维护的、能融合进系统里的实现。一个性能不错但很难维护的算子长期来看是团队的负担。另外我特别建议多关注硬件特性。同样一个算子在不同 GPU 上表现可能完全不同Tensor Core 是否可用、L2 缓存大小、内存带宽高低都会影响你的优化方案。只看框架层很多问题永远解释不了。写在最后一点真实体会做算子开发这几年来我最大的体会是算子就是深度学习的“汇编语言”它离算法很远离硬件很近但它决定了你的想法最终能不能高效落地。最初学 CUDA 和 Triton 时我也被各种概念砸得头晕什么 grid、block、shared memory、memory coalescing每个词背后都是一大堆背景知识。后来慢慢发现只要抓住“数据怎么组织、访存怎么安排、并行怎么切分”这三件事绝大多数算子都是可以理解、可以优化的。如果你正在入门我的建议是不要停留在“会用框架调用算子”的层面抽一个周末用 Triton 写一遍向量加法再写一遍矩阵乘然后尝试把一个简单的两算子序列融合成一个 kernel。这个过程会逼迫你建立完整的、从框架到底层的思维模型。这个模型一旦建立起来等你再遇到性能优化、算子报错、显存瓶颈这些问题时就不再是瞎子摸象了。