TorchInductor的代码生成是torch.compile后端的灵魂。前端把FX图lower成IR节点调度器再把IR节点聚合成内核组最后落到具体设备上生成可执行代码——CPU上是C/OpenMP/AVXGPU上是Triton。这两套内核生成器分别占据torch/_inductor/codegen/cpp.py和torch/_inductor/codegen/triton.py两个大文件也是绝大多数性能问题的汇聚地。这篇文章是“TorchInductor代码生成”系列的第二篇聚焦CppKernel与TritonKernel的内核生成过程。适合已经跑过torch.compile、想弄明白“编译到底干了什么”的工程同学也适合准备二次开发Inductor、或想自己写一个代码生成后端的同学。我会从调度器如何分发设备、C内核如何向量化、Triton内核如何组织grid与mask一路讲到调试排障和自定义扩展尽量做到读完之后你能自己上手改代码。1. TorchInductor为什么需要两套内核生成器1.1 一条计算图从FX到机器码的完整路径先用一条链路把位置讲清楚。你调用torch.compile(backendinductor)的时候入口是torch._inductor.compile_fx它做三件事先是GraphLowering把FX IR降到Inductor自己的IR节点比如Pointwise、Reduction、TemplateBuffer然后是Scheduler把这些IR节点按依赖关系聚合成一串SchedulerNode并决定哪些节点能融合进同一个内核最后才是真正调用codegen把每个SchedulerNode变成一段可执行代码。这个“内核生成”阶段就是本文真正的主角。Inductor内部通过Scheduling类做设备分派CPU设备走CppScheduling生成C代码CUDA设备走TritonScheduling生成Triton语言代码。后续所有优化——循环融合、向量化、并行、mask处理——都发生在这一层。我见过不少同学把注意力全放在IR层研究Pointwise怎么split、Reduction怎么切分但一落到codegen就发怵。其实codegen层没那么神秘它就是一台“代码打印机”输入是结构化的IR节点输出是两端风格迥异的源代码。理解这层的关键在于知道它的输入有哪些可打印的语义信息以及它在什么约束下决定“怎么循环、怎么寻址、怎么并行”。1.2 为什么不能只做一套后端问得最多的问题是为什么不统一用一套代码生成后端答案很简单CPU和GPU的并行模型完全不同。CPU上你有几十个物理核每个核里有AVX向量寄存器最有效的表达是“外层OpenMP并行打满多核、内层循环步长等于向量宽度、连续内存向量化”这套东西用C写最自然。GPU上是几千个线程块加几万线程最有效的表达是“每个线程块处理一块数据、线程间通过shared memory协作、块间用global atomic或二次加载规约”这套东西直接用Triton写最自然。如果强行用一套抽象最终只会得到一个两头都不讨好的中间层。Inductor的选择很务实IR阶段共享一套代码生成阶段按设备分裂。这样CPU和GPU可以各自演进互不拖累。后面你会看到虽然两套代码生成器面对的是同一个IR结构但生成的代码风格、循环组织、编译入口几乎没有任何共同点这正是刻意为之的结果。2. CppKernel内核生成拆解2.1 CppKernel的核心组件与继承关系torch/_inductor/codegen/cpp.py里的类家族核心是CppKernel。它继承自codegen/common.py里的Kernel基类这个基类维护了bufs、vars、body这些编译期状态以及最关键的ops抽象。CppKernel在CPU后端里承担的角色是把一组IR节点翻译成C表达式再拼装成完整函数。CppKernel下面还有几个重要子类CppVecKernel负责向量化代码生成会在循环里插入at::vec::VectorizedT类型的向量临时变量CppOuterKernel负责处理reduction和多层循环CppMatmulKernel负责处理CPU上的矩阵乘模板。CppScheduling则负责在创建kernel时决定到底实例化哪个子类并通过does_require_kernel判断当前节点是否真的需要单独生成一个kernel。这一层分派逻辑非常关键因为很多很小的节点会被融合到别人的kernel里根本不需要自己生成代码。实际看代码你会发现CppKernel里大量工作是在管理“表达式字符串”——而不是直接操作Tensor。比如一个add操作经过CppOverrides里的override处理后会变成tmp0 tmp1这种纯C片段缓存到body里最终统一拼接。这个设计参考了TVM的codegen风格优点是可以很方便地做常量折叠和公共子表达式消除缺点是调试字符串非常痛苦——后面排障章节再细说。2.2 一个Pointwise算子的C代码生成全过程为了说清楚我拿一个最简单场景演示y x * 0.5 1.0假设x是连续float张量长度1024。这个Pointwise节点被Scheduler判定为可独立生成kernel后CppScheduling会创建一个CppVecKernel然后调用codegen_body。codegen_body会遍历这个节点对应的LoopsIR把循环维度提取出来。这里的关键是步长向量化宽度。在AVX2机器上float是8个元素一个向量循环步长就是8double是4如果机器只支持SSEfloat又是4。Inductor通过DTYPE_TO_VECTOR_WIDTH表查询这个宽度并据此生成步长。生成的C代码大致长这样#include ATen/ATen.h #include ATen/core/Tensor.h #include omp.h extern C void kernel_cpp_0( float* __restrict__ in_ptr0, float* __restrict__ out_ptr0, long ks0 ) { #pragma omp parallel for for (long i0 0; i0 ks0; i0 8) { auto tmp0 at::vec::Vectorizedfloat(in_ptr0 i0); auto tmp1 at::vec::Vectorizedfloat(0.5f); auto tmp2 tmp0 * tmp1; auto tmp3 at::vec::Vectorizedfloat(1.0f); auto tmp4 tmp2 tmp3; tmp4.store(out_ptr0 i0); } }注意几个细节。第一ks0是元素总数循环条件是i0 ks0步长8这意味着如果长度不是8的倍数尾部会漏掉。Inductor处理尾部有两条路一是大多数生成的内部场景里张量被pad到对齐长度二是非对齐场景会生成一个标量尾部循环走非向量化路径。我实际调试中遇到的大部分“结果不对”或“越界写”问题都出在这个尾部处理上排查时先看kernel参数里有没有pad字段。第二连续内存假设。上面代码里in_ptr0 i0直接用连续地址偏移前提是Scheduler判定这个输入tensor是contiguous的。如果stride不为1Inductor会退化为带stride的循环或者干脆放弃向量化。你不用手动处理这个但要知道这决定了你能否吃到AVX红利。第三临时变量tmp0到tmp4完全是寄存器级别的内存临时值C编译器会优化成YMM寄存器操作。Inductor生成C时故意把每个操作拆成一行方便在生成的代码里做diff和定位。可读性在这里不只是给开发者看的它本身也让编译器能更清楚地做指令调度。2.3 向量化、OpenMP与reduction的协同细节再看两个真正影响性能的细节。OpenMP并行化通常发生在最外层循环。Inductor在正确的位置插入#pragma omp parallel for默认用static调度而不是dynamic调度原因是负载均衡在大多数静态shape下不需要dynamic。如果你发现多核利用率低可以先用环境变量OMP_NUM_THREADS核对线程数再去看生成的循环里是不是每层都正确加了pragma。生成代码是展开的如果你看到多个并行for嵌套在一起这通常是性能陷阱——线程数被外层抻满内层并行就没有意义了。reduction节点的情况要复杂得多。CPU上的规约不能简单地用OpenMP parallel for包一个sum因为跨线程的sum需要归并。Inductor生成的C规约核心是“先在本线程内做向量化局部规约再把局部结果通过reduction逻辑combine”。具体来说CppVecKernel会把规约维度切成向量块每个线程处理若干块后得到部分和最后在函数出口把这些部分和累加。这个模式跟手写OpenMP reduction几乎等价但生成的代码可读性差不少初学者容易看到一大坨局部临时变量就懵了。你不用逐行读重点看有没有并行规约标记以及最终结果的累加顺序是否稳定——浮点累加顺序不同会让结果有微小差异这在数值敏感场景值得注意。2.4 CppKernel的边界与已知短板跑了几个月CPU编译后我对CppKernel的边界有几点体会。它能稳定产出的高性能场景是elementwise融合、channel-last卷积附近的重排、以及规律shape的reduction。复杂场景比如稀疏形状、动态shape、复杂数据依赖代码生成器会频繁退化到标量循环性能掉得厉害。另一个痛点是编译速度生成C后需要调用gcc/clang实际编译还要做caching首次编译一个小模型也要几秒动态shape场景会把编译时间放大到不可接受。官方也一直在推precompile和persistent cache但在动态shape下cache miss率很高这是CPU后端的真实瓶颈之一。所以你想在生产里用CPU编译推理预算时间的时候一定要把“首次编译开销”算进去别让模型第一个batch卡在编译上。3. TritonKernel内核生成拆解3.1 TritonKernel的对象模型与构造流程GPU侧的代码生成主文件是torch/_inductor/codegen/triton.py。这里的核心类是TritonKernel它与CppKernel在同级但干的事完全不同。它不再生成C字符串而是生成一段Python代码这段Python代码最终会被triton.jit编译成GPU kernel。TritonScheduling会在构造时收集很多元信息xnumel问题大小、rnumel规约轴大小、是否有mask、是否需要persistent reduction、要用多少个XBLOCK/RBLOCK。收集完这些codegen流程才开始。整个过程可以理解成先把所有IR节点转成Triton语言表达式然后用triton.jit装饰器包成函数体最后生成grid的lambda表达式。和Cpp一样这里也是先做表达式拼接再做内核组装但Triton代码里有明显的Python级别语法比如tl.arange、tl.load、tl.store这让生成的代码比C更好读。3.2 Pointwise内核的Triton代码生成流程还是用y x * 0.5 1.0这个例子GPU上的生成结果大致像这样import triton import triton.language as tl from torch._inductor.triton_heuristics import grid, start_graph, end_graph triton.jit def triton_poi_fused_add_mul_0(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr): xoffset tl.program_id(0) * XBLOCK xindex xoffset tl.arange(0, XBLOCK)[:] xmask xindex xnumel x0 xindex tmp0 tl.load(in_ptr0 x0, maskxmask) tmp1 tmp0 * 0.5 tmp2 tmp1 1.0 tl.store(out_ptr0 x0, tmp2, maskxmask) grid lambda META: (triton.cdiv(xnumel, META[XBLOCK]),)真正的Inductor会额外处理常量绑定和tuning参数但主体就是这样。有几个细节值得展开。第一XBLOCK是tl.constexpr类型这表示它在编译期是常量Triton编译器会用它做循环展开和寄存器分配。Inductor的autotune会在运行时尝试不同的XBLOCK值比如16、32、64、128每个值都编译出一个kernel变体然后benchmark挑最快。所以如果你在日志里看到同一个kernel被编译了几十次那不是bug是autotune在搜索。第二xmask的作用是处理边界。GPU的grid大小是cdiv(xnumel, XBLOCK)当元素总数不是XBLOCK的整数倍时最后一个线程块会越界所以tl.load/tl.store都要带mask。Inductor只在必要的时候生成mask如果Scheduler判定xnumel能被XBLOCK整除它会省略mask省掉一部分比较和分支开销。实测中是否生成mask对小型kernel影响很大能省则省是这里的优化哲学。第三tl.program_id(0)对应CUDA的blockIdx.xtl.arange(0, XBLOCK)[:]对应blockDim.x。它没有直接暴露threadIdx因为Triton的抽象级别更高程序员写的是“block内一组连续线程处理一段连续数据”具体映射交给编译器。这跟CppKernel里直接写OpenMP和AVX是两种思路一个把线程模型抽象化一个把硬件特性显式化。3.3 Reduction与Persistent Reduction的分支GPU上的规约比CPU复杂因为跨线程块规约需要显式处理。Inductor对Reduction节点有两种策略普通reduction和persistent reduction。普通reduction会把规约轴切到RBLOCK大小的块里每个线程块加载自己的分片用tl.sum沿axis0做块内规约再通过global memory上的atomic操作把部分结果合并。这个方案实现简单但在rnumel很小的时候atomic竞争会成为瓶颈。persistent reduction的思路是让足够多的线程块覆盖整个规约轴每个线程块保持“常驻”直到完成自己的分片规约配合tl.sum和finalize逻辑。Inductor对persistent reduction有专门的代码路径生成的网格维度更高代码更长但性能在大规约维度下更好。选择哪条路径由启发式决定内部参数主要在rnumel和XBLOCK/RBLOCK的匹配度。调优时如果看到某个reduction kernel慢先确认它走的是哪条路径再考虑调RBLOCK。3.4 GEMM与卷积的Template Kernel路径除了常规pointwise/reductionTriton后端还有一条Template Kernel路径主要用于GEMM和卷积。模板是Inductor里预先写好的高性能计算模板类似CUTLASS的Tile描述由Triton代码片段加可替换的算子片段组成。Cpp后端对应的是CppMatmulKernelTriton后端对应的是TritonTemplateKernel。模板机制的意义在于对GEMM这种结构极其规整、调参与手写同样重要的算子用通用pointwise codegen去生成几乎不可能达到手写tile GEMM的性能。Inductor的处理是先让调度器识别出Matmul节点然后直接套模板把A、B、bias等传入模板利用tl.dot生成基于TensorCore的tile GEMM。你生成的代码里会看到tl.dot、tl.max_contiguous这些提示这些是给Triton编译器做register tiling的线索。后续想改GEMM性能优先看模板而不是看通用codegen。4. CppKernel与TritonKernel的对比与选型4.1 生成机制对比字符串、编译入口与调试体验把这两个后端放在一起比较最有意思的不是性能而是设计哲学。CppKernel是“直接生成接近手写的C”之后系统编译器接管TritonKernel是“生成一种高级GPU语言”之后Triton编译器接管。两者都避开了直接操作低级IR或PTX因为它们都清楚在当下的工程现实里让成熟编译器去优化已经很成熟的表达比从头写编译器更可靠。我用一张表总结对比维度CppKernelTritonKernel目标设备CPUCUDA GPU生成语言CTriton语言编译入口gcc/clang omptriton.jit PTX并行模型OpenMP多核 AVX向量化线程块 线程组主要优化手段向量化、循环融合、OpenMP调度XBLOCK自动调优、mask消除、TensorCore调参维度OMP_NUM_THREADS、向量宽度XBLOCK/RBLOCK、num_warps、num_stages调试手段看生成C、gdb看生成Triton、cuda-gdb、dump PTX主要瓶颈动态shape编译慢、源码缓存autotune时间长、Triton版本兼容从调试体验看CppKernel生成的C我基本能一行行读懂问题定位快TritonKernel生成的Triton也能读但真到了数据竞争和bank conflict你还是需要dump PTX或切回CUDA写个等价核做对比。4.2 性能差异与适用场景性能上两者的天花板都足够高。CppKernel在规整的CPU算子融合上能追平手写OpenMPAVX的版本TritonKernel在GEMM和规约上因为可以直接利用TensorCore性能接近CUTLASS。但两者都很吃shape规律性静态shape下融合策略稳定性能好动态shape下代码生成器会为了通用性牺牲掉很多优化性能掉得厉害。适用场景可以这样分CPU侧如果你的模型以连续shape、规则内存布局为主比如NCHW卷积加BN加ReLU融合Cpp后端收益明显如果遇到复杂分支、稀疏访问、动态shape别对性能抱太高预期。GPU侧大规模GEMM、Matmul加pointwise融合、规约算子Triton后端表现很好而小算子单kernel启动都不到10微秒则是启动开销主导生成任何代码都救不了要走上层kernel合并和CUDA Graph。4.3 维护成本与生态位维护成本上CppKernel和TritonKernel并不对等。Triton后端是Inductor当前最活跃的部分社区PR多、代码更新快跟着版本走就有新优化可用Cpp后端因为使用面相对窄活跃度低很多优化停留在较旧版本。如果你的团队想长期维护一个fork我建议把主要精力放在ir.py和scheduler.py上这两个文件是两套后端的公共层codegen层尽量保持最少修改因为跨版本跟上游合并时codegen层冲突最多。5. 实战排障内核生成常见问题5.1 调试内核生成的通用三板斧排障内核生成问题我的三板斧是看日志、看代码、看产物。看日志用环境变量TORCH_LOGSinductor或者TORCH_LOGSoutput_code。前者会输出Inductor整个编译过程包括每个kernel被创建、融合、编译的信息后者会把最终生成的源代码打印出来。我几乎每次都是先看output_code因为生成代码本身就能暴露80%的问题——比如循环步长错了、mask没了、参数顺序对不上。看代码就直接打开生成的源文件。TORCH_COMPILE_DEBUG1会把所有编译中间产物和缓存路径打印出来你可以找到生成代码的物理路径用文本编辑器打开看。Triton路径下还能找到cache目录里的.ptx文件那已经是最底层的产物了。看产物实际上是指与期望结果做差分。我已经数不清多少次靠torch.testing.assert_close来定位代码生成bug用一个极小的输入跑一遍torch.compile后的模型和原始模型对比输出。如果数值不对就缩小输入比如固定到单个block的大小看边界是否正确处理。5.2 CppKernel编译和运行问题CppKernel最常见的编译问题是本机编译器版本不够。Inductor生成的C会用到一些较新的C17特性以及at::vec头文件老版本gcc可能编译失败。日志里会直接看到compile command和报错处理办法很简单升级gcc或者显式指定CC/CXX环境变量让Inductor使用新版编译器。运行时问题里最折磨人的是“结果错但没崩”。优先级最高的排查项是contiguous假设。生成代码里如果用的是in_ptr0 i0这种连续地址而tensor实际不连续结果必然错。Inductor通常不会犯这个错但当你手写了自定义op并绕过了layout假设时就容易踩中。确认方式是在生成代码里搜stride变量看是否被正确使用。还有一个高频坑是OpenMP线程数。模型跑着一会儿快一会儿慢往往不是代码生成问题而是OMP_NUM_THREADS没设置好导致多实例并发时线程超订。设置OMP_NUM_THREADS到物理核心数比折腾生成代码更有效。5.3 TritonKernel的版本与缓存问题Triton侧的问题完全可以另开一篇博客这里说三个最常见的。一是Triton版本和PyTorch不匹配。不同PyTorch版本依赖的Triton接口变化很大常见症状是import triton报错或者triton.compile的签名对不上。解决办法是严格使用该torch版本锁定的triton。装错版本后Inductor会在编译阶段莫名其妙地crash或者生成无效的kernel。二是stale kernel cache。Triton在~/.triton/cache下缓存编译产物如果你改了环境或代码但缓存没失效新逻辑不生效。排查时先删掉整个cache目录再跑排除“旧内核”干扰。PyTorch自己的inductor cache也有类似问题TORCHINDUCTOR_CACHE_DIR可以手动指定位置。三是autotune导致首次运行特别慢。这是Triton后端的固有代价每个kernel都要尝试多个XBLOCK、num_warps组合并benchmark。如果你不在乎那一点性能可以关掉autotune或者减小搜索空间比如用torch._inductor.config.triton.autotune_at_compile_time控制编译期的搜索强度。线上服务坚决建议预热。5.4 性能反直觉时的归因思路还有一个经常遇到的现象代码生成没报错编译也正常但性能就是比手写CUDA差一截。先不要怀疑codegen能力用nvprof或者torch.profiler看kernel占用率。很多时候瓶颈在内存访问模式而不在计算比如Triton生成了巨大的num_warps但grid很小导致一个block里线程都在等内存。此时不是改代码生成器而是调kernel的num_warps/num_stages参数。另一个归因方向是看是不是落到了某个低效但能用的fallback路径。Inductor很多算子有数个实现路径某些情况下会走最保守的路径性能不好是正常的。用TORCH_LOGSinductor搜一下路径选择日志往往直接能看到原因。别一上来就怀疑生成的代码有问题很多“性能不佳”其实是策略选择导致的与代码生成器的执行质量无关。6. 从内核生成再往下走6.1 如何扩展一个新的后端或新的内核类型如果看完这篇你想动代码最现实的入口不是直接改cpp.py或triton.py而是改两个地方IR节点和Scheduler的分派逻辑。IR层决定了codegen能拿到什么语义信息Scheduler层决定了哪些节点被融合。这两个层面对齐后codegen层只是“按格式打印”。举个例子你想给某个设备新增一个codegen后端核心工作是实现一个继承自Kernel的类并实现codegen_body和codegen_functions方法再把设备类型映射到你的Scheduling类上。你不需要从零生成所有专有代码很多基础设施比如缓存、autotune和前缀文件管理都来自common.py和graph.py。真正难的是想清楚你的后端要暴露哪些ops因为IR层调用的ops接口决定了kernel能表达哪些算子。6.2 我个人最看重的一个小技巧最后分享一个小技巧无论改CppKernel还是TritonKernel都先在生成的代码文件上做最小改动再单独编译。CPU侧直接把生成的C存成.cpp文件写个main函数手动编译运行GPU侧把生成的Triton文件存成.py单独跑triton.compile验证语法。这会让你的调试循环从分钟级缩到秒级比在torch.compile整个流程里反复试错高效得多。这是我从一开始接触Inductor就养成的习惯现在也推荐给每一个要折腾内核生成的读者。