AI算子详解：从张量计算图到自定义算子性能优化实践

1. 从一张数字流水线说起AI算子到底是什么AI算子这个词最近在技术社区里讨论热度肉眼可见地上升。不管你是做模型训练、推理加速还是自己捣鼓深度学习框架最后都会撞上算子这个概念。我第一次真正意识到它的重要性是在一个项目里被模型性能瓶颈卡了整整三周——查来查去问题不是出在模型结构上而是出在一个很不起眼的元素级算子上。那到底什么是AI算子你可以把整个深度学习模型想象成一条数字流水线。原始图像、文本、音频数据从一端进来经过无数道加工工序最终从另一端输出预测结果。每一道工序就是一个算子。有的负责把数据挪个位置比如转置、切片有的负责做数学运算比如矩阵乘、卷积、归一化有的负责改变数据的形状和分布比如reshape、softmax。算子就是这些工序被抽象出来的统一名字在底层对应着一小段可执行的数学逻辑和计算代码。为什么现在单独谈算子特别重要因为AI框架再花哨最终都要落到算子在硬件上能高效跑起来。如果你想让模型跑得更快、显存占用更小、部署在边缘设备上还能保持实时性那么对算子的理解就是绕不开的一课。这篇内容适合三类人看一是刚转型做AI开发、被各种概念绕晕的新手二是已经在用PyTorch、TensorFlow但想知道底层到底发生什么的进阶玩家三是准备自己手写算子做性能优化的工程同学。我会从概念讲到实践最后用一个手写加算子的完整过程收尾尽量把关键为什么都拆开讲清楚。2. 先把基础概念钉牢张量、计算图与算子的关系2.1 张量是数据容器算子是加工动作很多人第一次接触算子的时候会被张量Tensor和算子之间的关系搞混。一句话总结张量是数据算子是动作。张量你可以理解成一个多维数组它可以是0维的标量、1维的向量、2维的矩阵也可以是一张图片的形状 [3, 224, 224]更可以是一个批量文本序列的形状 [32, 128, 768]。算子则是作用在张量上的函数输入一个或多个张量经过内部计算输出新的张量。这套设计思路其实就是函数式编程的翻版。把数据和处理逻辑完全分离好处是极大的灵活性同一个张量可以由不同的算子以不同顺序来处理同一个算子也可以被复用到不同位置。框架层面只需要维护好数据流向和依赖关系剩下的事情就变得很模块化。在实际工程里有一类问题是新手最容易踩的维度的隐式变化。比如你用了一个reduce算子对张量求均值原来维度是 [batch, seq_len, hidden]结果你把某个轴给约掉了后面再接的其他算子突然报维度不匹配。这个问题表面上看是报错本质上是你在心里没想清楚算子的输入输出签名。养成一个习惯每写一个算子先在纸上写出输入shape到输出shape的映射关系能省掉大量调试时间。2.2 计算图是整个模型的施工蓝图当几百个算子组合在一起的时候它们之间的依赖关系和图数据结构便自然浮现。一个卷积算子的输出要喂给一个激活函数算子激活函数的输出又喂给池化算子或者归一化算子——这些连接关系画出来就是一张有向无环图也就是计算图。框架用计算图是为了做三件事分析、优化和调度。分析是指静态地看一下整张图里哪些节点可以并行、哪些节点有依赖优化是指做算子融合、常量折叠、死代码消除这类变换调度是把图中节点映射到实际硬件上比如哪几个算子放在CPU上跑、哪几个挪到GPU上。我在实际接触框架源码之前一直把计算图当成一个抽象概念。直到有一次在排查一个显存占用异常问题的时候打开profile工具看到整张图的节点执行顺序和内存分配情况才真正意识到计算图这层抽象对工程实践的意义。你不会直接跟图交互但底层全是在图上做文章。3. 算子的类型拆解从简单到复杂每个都是独立的知识点3.1 形状操作类算子最容易实现也最容易埋坑形状操作是所有算子里面最没技术含量但又最不能忽视的一类。reshape、transpose、concat、split、expand、squeeze、unsqueeze都属于这个范畴。它们的共同点是不改变张量里的元素数值只改变元素的排列方式和逻辑视图。举个例子一个形状为 [2, 3] 的张量底层内存是一段连续的6个浮点数逻辑上被组织成两行三列。reshape成 [3, 2] 之后内存不变只是被解释成三行两列。但这里有一个关键坑reshape要求张量在内存中是连续存储的。如果你先做了transpose得到的是一个视图不连续的张量此时直接reshape会报错或者触发隐式的数据拷贝。很多框架提供contiguous()方法来把数据物理上拷贝成连续排列然后再reshape。这个隐式拷贝很消耗时间尤其在跑大模型的时候尽量不要在关键路径上做transpose后立刻reshape这类操作。concat和split反而是真的会发生数据搬移的。concat会把多个张量按指定维度拼在一起如果维度匹配不对、或者张量都没有对齐轻则报错重则产生非预期结果。我习惯在做concat之前手写一行断言检查维度确保维度匹配无误。3.2 元素级运算算子性能调优的主战场所谓元素级运算就是输出张量中的每一个位置只取决于输入张量对应位置的值。最常见的包括加法、减法、乘法、除法、ReLU、Sigmoid、指数、对数等等。这种算子的特点是计算规律很简单但它占据了深度学习训练过程中的绝大多数时间。为什么因为模型结构里即便是矩阵乘法也会分成核心大运算和很多小的辅助运算。比如LayerNorm先要计算均值和方差然后做归一化最后再缩放和平移这里面每一步都可以拆成元素级的加减乘除和reduce操作。在NVIDIA GPU上这些元素级算子受限于内存带宽而不是计算能力。什么意思呢就是计算单元大部分时间是空闲的在等数据从显存里搬过来。这也是为什么现代推理引擎把很大精力放在算子融合上。举个例子把卷积BNReLU三个算子合成一个kernel让数据在寄存器里做完整条链路的计算不需要每算一步就写回显存再读出来。在工程实践里这种融合往往能带来数倍的速度提升。理解这一点你就理解了为什么广大AI框架里的算子数量比你自己想象的少得多——因为很多算子被提升融合了。3.3 规约类算子一个被忽视的性能瓶颈规约reduction算子是把一个张量按某个轴或多轴方向压缩成更小张量的操作。典型代表有sum、mean、max、min、argmax。在PyTorch里你几乎每天都会用到它们但很少人会想它们到底是怎么实现的。规约算子的难点在于跨维度的数据访问模式。拿二维矩阵在列方向求均值来举例输出的每个元素需要跨行访问同一列的数据而内存里同一行是连续的跨行就意味着内存跳跃访问。加上并行技术时多个线程同时更新一个输出位置还可能产生写冲突。要处理好这个问题需要根据张量大小选择不同的归约策略小张量用单线程串行中等张量用原子操作或者拆分多部分局部归约再合并大张量则要用树形规约来分摊开销。很多做自定义算子的人算法逻辑写得没问题性能却一塌糊涂大多就是栽在了没有考虑到规约算子的本质访问模式上。3.4 GEMM类算子深度学习的心脏矩阵乘法GEMM可以说是深度学习最核心的算子全连接层、卷积层、注意力机制里的QKV投影到最后都等价于矩阵乘法。GEMM的特点是计算密度高即每访问一个数据可以做很多次乘加运算所以完全可以用矩阵分块、寄存器重排、指令流水等技巧把硬件榨干。在工程实现上GEMM几乎不会直接用三重循环去写因为那样子缓存命中率极低。正确的做法是把大矩阵切成小块让子块能装进L1缓存和寄存器然后对每个子块做微内核计算。以CUDA为例通常会用tile大小为16x16或32x32配合共享内存的双缓冲机制来隐藏全局内存延迟。你能接触到的cuBLAS、cuDNN核心就是经过深度调优的GEMM实现。如果你对矩阵乘法的底层优化感兴趣强烈建议去看一个开源项目叫CUTLASS里面的分块策略、流水线设计是最佳的学习样例。4. 为什么算子实现如此讲究硬件特性和性能模型4.1 算子的执行效率和访存密度直接挂钩算子在硬件上的执行时间主要由两方面决定计算时间和访存时间。计算时间是你做算术操作需要花的时钟周期访存时间是把数据从内存/显存搬到处理器里需要花的时钟周期。这两者哪个是瓶颈取决于算子的性质。回到前面那个元素级运算的例子加法运算本身只需要一个FMA指令乘加指令周期但为了得到两个输入数据和一个输出数据你需要读3个值、写1个值。以NVMe SSD和CPU内存来类比计算就像在水龙头旁边接水访存就像从蓄水池打水过来如果你的手速远远快于打水速度那么你就是在等水。这也是为什么很多篇技术文章里反复提到访存密集型memory-bound和计算密集型compute-bound两个词。做算子优化的时候第一件事不是盯着汇编指令优化而是先判断这个算子属于哪一类。如果是访存密集型的你该思考的是怎么减少数据搬运量如果是计算密集型的你才需要考虑指令级并行、向量化和矩阵分块。方向错了后面全是白费力气。实际工作中我见过不少反例有人为了一两个周期的手写SIMD指令花了一整天结果发现整个算子大部分时间都在等待数据的到来优化指令根本没有意义。先分析特征再动手这是算子优化里最重要的一条原则。4.2 内存布局对算子性能的影响超乎想象内存布局通常指张量在物理内存中的排列方式。常见的布局有NCHW通道优先和NHWC像素点优先这决定了卷积算子在读取特征图时按什么顺序访问内存。如果访问模式和内存布局不匹配会导致大量高速缓存未命中。举一个实际案例在推理引擎里我们经常把一个视觉模型的输入从NCHW转成NHWC。因为对于深度卷积NHWC布局下同一像素点的各通道数据是连续的向量化加载和SIMD指令能一次性读取多个通道这对CPU端的推理优化特别有利。相反在NCHW布局下通道维被放到了最外层深度卷积需要跨大步长访问性能下降得厉害。TensorFlow的默认布局是NHWCPyTorch的默认布局是NCHW。这种差异导致从PyTorch导出的模型在转到TensorFlow部署时如果不做布局转换就会产生大量重排算子。有些推理引擎如TensorRT会自动做布局优化但如果你的场景是自研引擎就得格外关注这个问题。判断一个算子的内存访问是否高效可以看一个指标有效带宽利用率也就是实际访问的数据量除以总数据量。如果利用率不到70%大概率是内存布局或者访问模式出了问题。4.3 并行粒度与同步开销的权衡之道现代处理器无论CPU还是GPU都有大量的并行计算单元一个算子如果想跑得快必须把工作拆解成可以并行执行的任务。但并行粒度并不是越细越好因为创建任务、同步线程、等待结果都有额外开销。在GPU上每个线程处理一个元素是最自然的想法但未必是最高效的。比如一个用于2560x2560矩阵的纯粹元素级操作每个线程处理一个元素那么总共需要约650万个线程。GPU虽然能承载高并发但大量的线程调度也会有开销。更优的做法是让一个线程处理连续多个元素比如4个或8个使用向量化指令一次加载和计算多个数。这样既减少了线程总数又让内存访问变得更加规整。还有一种常见的优化策略是grid-stride loop。假设GPU有100个并行线程块而你这边有1000个独立任务那就让每个线程块循环处理多个任务每次跳跃的步长是总线程块数。这种方式的好处是任务分配均匀、可以处理任意规模的数据而且不需要频繁创建和销毁大量线程块。我在写自定义kernel时几乎总是使用这个模式总体稳定且容易调参。5. 工具选型与前后端环境准备从哪里开始动手5.1 初学自定义算子该选什么样的工具链写自定义算子的方式非常多从易到难分别是纯Python模拟、使用框架的自动微分API组合已有算子、使用框架提供的C扩展接口、手写内核并绑定到框架、用Triton编写GPU内核。对第一次想手写算子的朋友我建议先从PyTorch的C扩展切入因为它的基础设施最完善。只需要安装好PyTorch和一个支持C编译的环境用torch.utils.cpp_extension.load来加载C代码就能在Python里直接调用自己写的算子。如果你之后想上GPU再改用CUDA扩展来深入重写路径会很平滑。那Triton值不值得学非常值得。它有点像CUDA编程的高层语言你不用手动管理线程、共享内存这些底层细节只需要描述计算逻辑和内存访问模式编译器会帮你做优化。关键它生成的代码性能通常能接近手写CUDA的水平对入门者极其友好。我自己现在超过一半的算子原型都是用Triton写的只有遇到特别变态的需求才会降级回到CUDA手写。5.2 环境配置的一个建议先用容器或虚拟环境隔离我在过去的项目里吃过不少环境依赖的亏。PyTorch版本、CUDA版本、GCC版本、甚至Python版本任何一个不匹配都能让最简单的示例跑不起来。个人经验是用Docker镜像来做隔离把一套已确认可用的环境固化下来。如果是自己本地玩也强烈建议用虚拟环境工具加Python版本管理工具不要在系统Python里裸装各种包。另外建议在动手写算子之前先跑通框架自带的算子扩展示例。PyTorch官方文档里有一个线性层自定义示例麻雀虽小但五脏俱全把 forward、backward、autograd.Function 的接线方式讲得明明白白。把那个例子跑通一次你对整个流程的框架就有了体感后面写自己的算子会顺利很多。6. 实操全流程手写一个自定义加算子并从零验证6.1 需求描述与接口设计为了讲清楚整个流程这里选择一个绝对简单、但能覆盖主要踩坑点的例子实现一个支持广播机制的加算子。也就是说它接受两个任意形状的张量 A 和 B输出 A B。注意这不仅仅是形状完全一样的加法还包含 形状不同、按规则广播到相同形状后再做加法 的场景。比如 A 的维度是 [2, 3]B 的维度是 [3]输出应该是 [2, 3]。为什么选这个因为它既能让你理解元素级算子的基本实现又能让你遇到广播broadcasting这个非常经典、又特别容易把新手绕晕的问题。广播规则虽然看起来简单要自己实现的时候你得处理输入维度对齐、右对齐规则、维度为1的自动扩张等情况写完之后你会豁然开朗。采用PyTorch C扩展的方式来做。具体来说我们要写一个函数 torch_extension_add(A, B)内部调用我们自己用C实现的核心理算逻辑。6.2 第一步注册扩展并测试基础加法先建立一个文件 add_ops.cpp代码如下#include torch/extension.h #include iostream torch::Tensor custom_add(torch::Tensor a, torch::Tensor b) { TORCH_CHECK(a.device() b.device(), Input tensors must be on the same device); TORCH_CHECK(a.dtype() b.dtype(), Input tensors must have the same dtype); return a b; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def(custom_add, custom_add, Custom element-wise add); }再看Python侧加载代码import torch from torch.utils.cpp_extension import load_inline cpp_src #include torch/extension.h torch::Tensor custom_add(torch::Tensor a, torch::Tensor b) { TORCH_CHECK(a.device() b.device(), Device mismatch); return a b; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def(custom_add, custom_add); } ext load_inline( namecustom_add_ext, cpp_sourcescpp_src, functions[custom_add], verboseFalse, ) a torch.randn(2, 3) b torch.randn(2, 3) c ext.custom_add(a, b) print(c)第一次跑这个扩展记得预留几分钟让它编译。一旦编译完成你在Python里调用它和调用PyTorch自带算子体验没什么区别。注意这里用 load_inline 是为了方便讲解在正式项目里我更推荐把代码写进独立文件再用 setup.py 安装。踩过的一个坑load_inline 里如果忘记在 functions 列表写上函数名编译能过但在调用时会出现找不到函数报错。这属于非常隐蔽的坑排查了很久才注意到是注册列表没写完整。6.3 第二步实现广播逻辑关键细节上面的实现只是照搬了PyTorch原生的加法并没有体现我们自己的实现细节。下面用C手写一套没有自动广播逻辑的加法然后再解开广播的结。先不谈广播我们来写纯粹的、只支持相同形状加法的C函数#include torch/extension.h #include vector torch::Tensor naive_add_same_shape(torch::Tensor a, torch::Tensor b) { TORCH_CHECK(a.sizes() b.sizes(), Shapes must match for naive add); auto a_contig a.contiguous(); auto b_contig b.contiguous(); auto out torch::empty_like(a_contig); const float* a_ptr a_contig.data_ptrfloat(); const float* b_ptr b_contig.data_ptrfloat(); float* out_ptr out.data_ptrfloat(); int64_t numel a_contig.numel(); for (int64_t i 0; i numel; i) { out_ptr[i] a_ptr[i] b_ptr[i]; } return out; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def(naive_add_same_shape, naive_add_same_shape); }这个实现再简单不过了。从左到右挨个把两个张量对应位置的数字相加写到一个新的输出张量中。这里最关键的一步是调用了 contiguous()原因是如果传入的张量是一个转置后的视图内存中的数据并不是连续排列的我们直接用 data_ptr 按顺序遍历就会得到错误的结果。调用 contiguous() 确保按行主序排列。再来看广播。假设输入A形状 [4, 1, 6]B形状 [5, 1]按照NumPy广播规则它们会先从右往左对齐A形状 [4, 1, 6]B形状 [5, 1]对齐后第一维A是4B没有这一维结果取4第二维A是1B是5取5第三维A是6B是1取6结果形状应为 [4, 5, 6]。广播的通俗理解是维度为1的会被拉伸成对方维度的大小相当于把这个维度上的数据复制多份。要在C里实现它最直观的方式是先把输入张量用 expand 方法扩展成目标形状但 expand 返回的是视图需要再连续化。这里直接实现计算逻辑不用框架的底层算子torch::Tensor custom_broadcast_add(torch::Tensor a, torch::Tensor b) { auto a_contig a.contiguous(); auto b_contig b.contiguous(); // 通过 PyTorch 的 broadcast 机制获取输出形状 auto out_size at::infer_size(a_contig.sizes(), b_contig.sizes()); auto out torch::empty(out_size, a_contig.options()); // 把输入广播到目标形状 auto a_b a_contig.expand(out_size); auto b_b b_contig.expand(out_size); // 对于广播后的视图使用 TensorIterator 来做逐元素加法可以处理好边界情况 // 但为了展示手写逻辑这里用最朴素的索引循环。 auto a_acc a_b.accessorfloat, 3(); auto b_acc b_b.accessorfloat, 3(); auto out_acc out.accessorfloat, 3(); for (int64_t i 0; i out_size[0]; i) for (int64_t j 0; j out_size[1]; j) for (int64_t k 0; k out_size[2]; k) out_acc[i][j][k] a_acc[i][j][k] b_acc[i][j][k]; return out; }这个版本只固定处理 3D 张量维度一多就得写多层模板去展开。实际生产环境里PyTorch 之所以性能好且通用就是因为底层有一套叫做 TensorIterator 的机制它把广播、类型提升、内存连续化这些事情全部统一处理了内核只需要关心最内层连续的一维循环。这也是所有自研算子库应该借鉴的架构设计。6.4 第三步验证结果和边界情形写完算子后的第一件事不是直接跑到大模型里去测而是做最小化的单元验证。测试矩阵大致是同形状加法形状 [2, 3] 与 [2, 3]结果应与 PyTorch 原生加法一致一维广播 形状 [2, 3] 与 [3]结果应与 torch.add 一致多维广播 形状 [4, 1, 6] 与 [5, 1]结果应与 torch.add 一致反向广播 大张量加小张量B的维度更大的顺序交换情况标量加法 形状 [2, 3] 加一个标量在写测试的时候别图省事只比较最终数值还要考虑到浮点误差。用 allclose 函数而不是直接判断相等相对误差放宽到 1e-5。我第一次写这个测试时用等号判断结果因浮点运算顺序不同报了一大堆错折腾了半天才意识到是判断方式的问题。对于广播这个功能你还需要检查输入维度为1时拷贝出来的数据是否对应正确的原始数据。比如 A 在某个维度大小为 1广播成大小为 5 后那5个位置的值全都来自同一个原始数据这是对的。但如果实现时遇到维度不匹配且某一方不为1应当报错而不是静默地错乱算。我会在函数开头补一个显式检查如果两个张量对齐后某维度两边都不相等且都不为1就抛出清晰异常信息。7. 工程陷阱排查自定义算子时最常见的五类问题7.1 设备不匹配与数据类型不匹配框架本身对 device 和 dtype 检查得非常严格但你自己写扩展的时候很容易直接把 CPU 张量和 GPU 张量混在一起或者在 float32 和 float64 张量上执行同一个内核最后得到奇怪的结果甚至段错误。解决方案是在代码入口统一做检查这比把希望寄托在调用方更有把握。我在TORCH_CHECK里除了检查 device 之外还一定会检查 dtype。混合精度运算如果在算子内部做一定要显式先转成目标类型否则隐式转换的规则在不同版本里可能不一致会导致结果不合预期。7.2 生命周期问题数据指针挂在悬空引用上这段是被提得很多但依然高频出问题的坑。你拿到一个data_ptr但在内核还没跑完之前生成这个张量的临时对象已经被垃圾回收/内存释放了那你的指针就是野指针。症状表现为偶发性的内存访问错误看起来十分随机难以复现。防御性做法是在进入内核之前把需要的张量用一个局部变量持有保证它的生命周期覆盖到内核执行结束。如果你用torch::from_blob创建张量包装外部内存一定要设置 deleter 回调否则框架会在张量析构时尝试释放该内存可能触发不可预测的行为。7.3 自动求导机制没为自定义算子接通如果你只实现了 forward但没有实现 backward那么在训练时一旦对算子的输出调用 backward就会直接报错。框架不知道你的算子的梯度公式是什么自然无法反传。需要实现torch::autograd::Function的子类重写 forward 和 backward。对于加算子backward 就是把上游传下来的梯度分别拷贝到两个输入的对应位置如果是广播场景还要做梯度的规约因为广播等于数据复制梯度在复制方向要加起来。这个逻辑单独写起来不复杂但它涉及维度追踪容易出错。稳妥的方法是在 backward 函数里调用a.sizes()与b.sizes()将累加后的梯度再 reshape 回原形状。7.4 内存泄漏与显存增长用 C 扩展最容易出现的问题是内存分配后没有释放。如果你在算子内部new了一块内存记得在函数结束时删除或者在函数里返回一个由torch::from_blob包装的张量时一定要把释放逻辑交给 deleter。显存方面更麻烦每次算子调用如果都触发一次大的临时张量分配训练吃显存就会特别快。排查思路是在循环里调用算子之前和之后各打印一次显存占用如果持续增长多半是算子内部临时张量生命周期不正确。此时用工具或API追踪张量引用计数准确定位在哪里泄漏。7.5 多线程环境下的数据竞争某些框架算子会尝试并行执行来加速如果你的内核没有做加锁或原子处理多个线程同时写一个输出位置就会产生数据竞争。尤其是规约类的算子累积变量是共享的竞争会带来不可预期的数值。建议先用单线程验证逻辑正确性再逐步开放并行。个人经验是最大限度利用局部变量先在各线程内做局部累加最后再做一次合并这比在每个累加步骤都加原子操作要高效得多。8. 从自定义算子到算子库设计一个更大的视角8.1 为什么说算子实现是框架性能的胜负手深度学习框架过去十年的迭代本质上都是在拼算子性能和调度效率。谁的卷积更快、谁的注意力算子更省显存、谁在端侧能更高效地执行Transformer谁就能在对应的应用场景里占据优势。这也是为什么英伟达、AMD、各大云厂商都在大力投资算子体系和自动调优工具链。如果你要在一个新的AI加速芯片上跑模型理论上能跑通但性能要逼近native实现几乎一定会经历一个过程先白盒分析最耗时算子然后逐一手写优化内核。接近80%的收益都集中在5%的关键算子上要么是GEMM要么是Attention要么是LayerNorm先集中火力搞定这些top算子。8.2 值得长期关注的两个方向算子编译和自动调优算子编译的思想在于不写死算子而是把算子的计算逻辑以特定中间表示输入给编译器编译器根据目标硬件生成最佳实现。这意味着你写一次算子逻辑就可以在CPU、GPU、甚至新出的AI专用芯片上分别得到比较不错的内核这其实是当前各框架投入大量资源的方向。自动调优更是工程实践里非常实用的手段。有些库会根据输入形状在启动时跑一组benchmark然后选择最优实现。比如GEMM在某种形状下用这种分块方式最快换一个形状又可能是另一种配置最好。把这些配置组合都测一遍把最快的方案缓存下来后续同一形状就直接复用。对于个人开发者来说你不一定要把这些都实现出来但一定要理解它们的存在否则很容易把某个形状上优化到极致的内核误当成全场景通用最优解。9. 关于算子性能分析的三个实用工具性能问题不能靠猜要用数据说话。以 PyTorch 为例有自带的torch.profiler可以展示每个算子消耗的时间、显存、调用次数。这个工具贵在精简能快速帮你找出瓶颈。看到某一个算子时间占比异常高再深入它内部去做进一步分析。如果到了 GPU kernel 级别NVIDIA 官方提供nsight compute和nsight systems。nsight systems看整体程序各阶段的时间分布nsight compute专门分析单个核函数的占用率、访存带宽、指令吞吐。有时候你会发现一个kernel跑不快的原因居然是共享内存bank conflict或者寄存器溢出这类问题用编辑器肉眼看代码是绝对看不出来的。CPU 端则可以用perf、vtune这类工具。关注的核心指标包括L1缓存命中率、分支预测错误率、SIMD指令占比、内存带宽利用率。可以看到优化算子的本质除了算法之外很多是跟硬件微观特征在打交道。10. 最后想说的话关于 AI 算子学习我自己的路径是先用了很长时间的框架对算子有了直观感受再倒逼自己去实现算子、研究底层原理现在反而是从算子视角回头审视框架设计很多原先似是而非的概念突然都通透了。给入门者一个具体的建议路径先手写一个最简单的加算子并跑通自动求导然后把维度从固定变成通用处理广播再尝试把一个元素级算子的性能用向量化指令提升一次接着尝试写一个reduce算子体会线程协作的复杂最后选一个真实的模型结构比如BERT或ResNet用profiler找出最耗时的几个算子逐个手动重写优化。走完这一圈你对算子的理解会超过大部分只会调包的开发者。在动手的过程中一定会遇到各种莫名其妙的编译错误、内存错误、以及结果不对但不知道怎么排查的窘境。这些都是必经之路但每次解决问题后你学到的东西比看十篇教程都更深刻。把算子的理论学习、工具实践、性能分析和工程落地串成一个闭环这才是最有价值的学习模式。我这些年最有收获的项目几乎都是从这个闭环中走出来并且真正运用于生产环境的。如果你也是自己折腾算子或者正在为项目做独立模型加速希望上面的内容能给你带来一些新视角。实践过程中遇到的问题欢迎交流讨论。