先把结论放在前面如果你用过torch.compile大概率听过 Inductor 这个名字也大概率知道它默认走的是 Triton 后端。但真正自己写过自定义后端、或者想把 Inductor 的循环代码生成和别的编译后端对接的人应该都体会过那种“官方文档不够用、源码又太绕”的卡壳感觉。这篇就来把“逐 range 编译”这条路彻底捋一遍重点放在 PiecewiseBackend 和 Inductor 的对接机制上。这里的“range”不是随手写个for i in range(3)那么浅层的概念它是 Inductor 调度器里的循环空间。所谓“逐 range 编译”本质上是把一个大循环按 range 边界拆成多个片段每个片段单独交给后端生成代码最后再拼接成完整的执行逻辑。这种做法的好处非常实际既能隔离编译单元、降低单次编译压力也能在不同硬件环境下用不同策略去优化不同区间可移植性和可控性都上来了。这篇文章适合三类人第一类是正在读 Inductor 源码、想知道代码生成链路怎么钩子的人第二类是想给torch.compile接一个自定义后端、却被各类接口绕晕了的进阶玩家第三类是吃瓜看热闹、但想真正理解 PyTorch 2.x 编译体系“为什么越改越复杂”的好奇型工程师。内容会偏底层一些但我尽量用大白话把每个环节的为什么讲清楚。1. 从 torch.compile 到 Inductor一条主线上的两个岔路口1.1 编译管线里到底发生了什么先把整条链路摊开。torch.compile启动之后第一步是 Dynamo它把 Python 字节码转成 FX Graph。这一步干的事可以理解成“把动态的 Python 执行过程固化成一张静态计算图”。接下来FX Graph 会交给后端而默认后端就是 Inductor。Inductor 在这里做的是“图到代码”的翻译工作接收 FX Graph经过调度器Scheduler做算子融合、内存规划、并行切分最终通过代码生成器输出可以直接运行的代码。默认情况下GPU 上输出的是一份 Triton 内核代码CPU 上输出的是 C 核函数。import torch def foo(x, y): return (x y) * 2 compiled torch.compile(foo, backendinductor)上面这段代码走的就是默认路径。你只需要知道一个关键点backend参数是可替换的。你可以传eager、aot_eager也可以传入一个自定义函数这个函数接收GraphModule和输入的example_inputs然后返回一个可调用对象。这就是 PiecewiseBackend 存在的入口位置。它不是一个被 PyTorch 强绑定死的类而是作为自定义后端注册进来的。它的角色是介于 FX Graph 和最终可执行代码之间的一张“分拣台”把一个完整的计算图拆成若干段每一段对应一个 range 循环然后逐段去调用真正干活的代码后端。1.2 为什么不是直接整图编译很多人会问为什么不直接把整个计算图丢给一个编译器非得切来切去这个问题的答案藏在编译优化和编译开销的矛盾里。一张大图直接给编译器编译器要处理的中间表示IR会非常庞大优化器会在各种循环嵌套、依赖分析之间做全局搜索编译时间会指数级上升而且有些优化对整图并不适用还会引入稳定性风险。逐 range 编译的思路本质上就是“分而治之”。先把计算图依据循环结构切开每个循环体的边界清楚、依赖简单编译后端处理起来更轻松。切成多段之后每段可以独立做向量化、展开、缓存分析甚至可以为不同段选择完全不同的底层代码生成策略。我在实际调试里遇到过一种非常典型的情况一个融合了四五个算子的 kernel直接生成 Triton 代码时某个中间循环的边界条件让 Triton 编译器生成了极其冗余的分支判断。后来我把这个循环按 range 切成两段头部一段用tl.multiple_of做对齐优化尾部一段走边界处理逻辑代码体积直接降了一半性能还提升了百分之十几。这就是“逐 range”带来的实打实收益。2. “逐 range 编译”到底编译的是什么2.1 循环结构和 range 的关系要理解逐 range 编译先得理解 Inductor 内部怎么看待循环。Inductor 的 Scheduler 接收 FX Graph 之后会构建一个依赖图然后按照依赖关系合并算子。合并后的结果是一个个“C/Triton Kernel”。每个 Kernel 内部由若干循环组成这些循环的空间维度比如张量长度、卷积的空间尺寸、reduction 的维度长度就是决定生成的代码要跑多少次的根本依据。在 Inductor 源码里循环通常被抽象成LoopBody、padded、size_hints等结构。调度最终会调用代码生成器按 loop nest 一层层生成 for 循环。在 Triton 后端里这个 for 循环会变成tl.range或for i in range(...)在 C 后端里会变成#pragma omp for包裹的for (long i ...)循环体。逐 range 编译所针对的正是这些“range 循环”本身。PiecewiseBackend 的典型做法是不直接让 Inductor 把整个 Kernel 编译出来而是把 Kernel 的循环空间拆成多个连续片段比如[0, 1000)、[1000, 2000)、[2000, N)每个片段分别编译成一个独立的子函数然后由一个统一的入口函数依次调用这些子函数。2.2 切分粒度怎么确定切分粒度是逐 range 编译里最值得花心思的设计点。切得太细片段数量多函数调用开销和编译时间都会增加切得太粗又回到了整图编译的老路上。从我自己的工程实践来看合理的粒度至少要同时考虑三个维度第一是硬件执行单元的宽度。GPU 上更适合按 block/thread 的整除关系来切尽量保证每个片段的元素数量是 block size 的整数倍CPU 上则要看 L1 Cache 和寄存器向量宽度片段最好能适配 SIMD 宽度。第二是边界条件的复杂度。如果一个循环尾部存在大量 if 分支比如if i n - tail那最好把前面对齐的部分和尾部残差部分拆开前者生成干净的向量化代码后者用标量逻辑兜底。第三是内存访问的局部性。片段太大访存跨越范围变大缓存命中率下降片段太小数据刚加载进缓存就要切走浪费带宽。我通常会用一个小脚本先统计各个算子的循环上界分布再结合硬件属性确定初始切分跑一轮微基准测试后调整。整个过程很依赖经验但核心原则是稳定的让每个片段尽量“同构”同构的循环体最容易生成高质量代码。2.3 与普通 range 的对比聊到这里可以把标题里的“range”和日常写 Python 的range做一个简单区分。日常写代码时for i in range(0, n, 2): ...这是 Python 语法层面的迭代器它控制的是 Python 解释器内部的循环流程。而 Inductor 里的 range是编译器中间表示层面上的循环上界和步长信息它描述的是“最终生成的目标代码应该怎么循环”而不是 Python 解释器怎么循环。所以“逐 range 编译”实际上是一个编译技术术语强调的是对编译目标代码中循环结构的逐段处理。这个点如果你在看课程或材料时没留意很容易被误当成简单的 Python 循环技巧那就跑偏了。3. PiecewiseBackend 与 Inductor 的对接机制3.1 对接入口register_backend 与 BackendCallerInductor 与自定义后端的对接官方给了一个相对明确的扩展点torch._dynamo.backends.register_backend。只要注册了后端名称就能在torch.compile里通过backendpiecewise来调用。from torch._dynamo import register_backend import torch register_backend(piecewise) def piecewise_backend(gm: torch.fx.GraphModule, example_inputs): # step 1: 把 GraphModule 交给 Inductor 的调度器做图优化 # step 2: 拦截 Inductor 生成的 range loop # step 3: 按 range 切分 loop pieces # step 4: 每个 piece 独立 codegen 编译 # step 5: 返回一个包装好的 callable ...这里要特别提醒一下register_backend注册的是后端工厂函数入参是 FX 图和样例输入返回值必须是可调用的通常是__call__方法优化过的 wrapper。如果你返回的是一个普通 class 实例要确保实例实现了__call__不然后面跑推理会直接报XXX object is not callable。3.2 Inductor 侧如何配合Inductor 本身并不直接暴露一个“按 range 切分”的官方接口但它的内部调度结构恰恰给这种做法留了空间。在torch._inductor.scheduler里每个被调度的 Kernel 都有明确的循环结构信息包括循环的依赖顺序、size 变量、range 上下界等。PiecewiseBackend 对接的常见写法是先调用 Inductor 的编译步骤但在生成最终代码之前做一次“改造”。比较轻量的方法是在 FX Graph 层面按节点分组把连续的、共享相同 loop range 的节点聚成 piece比较重的方法则是直接复用 Scheduler 的SchedulerNode利用它的get_loop_body去抽取每个节点的循环体。# 伪代码示意按节点依赖和 shared range 分组 def split_gm_into_pieces(gm): pieces [] current_piece [] last_range None for node in gm.graph.nodes: node_range extract_loop_range(node) if last_range is not None and node_range ! last_range: pieces.append(current_piece) current_piece [] current_piece.append(node) last_range node_range pieces.append(current_piece) return pieces这段伪代码不是 Inductor 源码里的原样逻辑但表达的核心思路和社区里常见的做法是一致的以循环范围为锚点把节点分组。实际落地时你会用到torch._inductor.ir里的LoopBlock、RangeConstraint之类的基础设施把它们一一代入边界条件。3.3 分段后的代码怎么拼回去分段编译最大的工程难点不在“切”而在“拼”。每个 piece 编译完之后产出的是独立的 kernel 调用函数之间需要交换中间结果这就涉及内存分配和中间张量的生命周期管理。局部切换两端的实现方式有两条路线。一条是“全局分配”在进入入口函数时把计算图中所有中间的 buffer 一次性分配好各 piece 的 kernel 直接读写固定地址。这样做省掉了 piece 之间的重复分配开销但内存占用会偏高。另一条是“流式传递”每个 piece 只接收上一个 piece 产出的 buffer并把当前 piece 的最终输出交给下一个。优点是内存峰值低缺点是 piece 之间耦合度高如果某个 piece 要做并行调度缓存同步会成为瓶颈。我在一个 CPU 推理优化项目里试过两种方案。全局分配的方式对内存带宽更友好适合单线程或少量线程的场景流式传递在多线程流水线里更有前途但实现复杂度高了一个档次。如果是第一次做对接建议先用全局分配跑通再考虑流式优化。注意无论选哪种方案都要处理好torch.fx图中的get_attr、placeholder和output节点。很多人在自定义后端里踩坑就是没有把输入 sample 和输出节点重新绑定到拼合后的新图上。4. 动手实现一个最小可用的 PiecewiseBackend4.1 建立工程骨架实际操作里我不会一上来就啃 Inductor 源码而是先写一个最小后端把链路跑通。下面是一个可运行的骨架基于 PyTorch 版本相对较新的接口。import torch from torch._dynamo import register_backend from functorch.compile import make_boxed_compiler import torch._inductor as inductor register_backend(piecewise_debug) def piecewise_debug_backend(gm: torch.fx.GraphModule, example_inputs): # 1. 先用 inductor 的 lowering 处理一遍图 from torch._inductor.compile_fx import compile_fx # 2. 为了调试方便先打印原始 graph print(--- Original GraphModule ---) gm.graph.print_tabular() # 3. 真正调用 inductor 的编译流程 compiled compile_fx(gm, example_inputs) # 4. 返回 boxed wrapper return make_boxed_compiler(compiled)这段代码注册了一个名叫piecewise_debug的后端它做的事其实还是调用默认的compile_fx并没有真正分段。但它有一个重要作用验证后端注册和调用链路是否打通。4.2 把 range 切分逻辑插进去真正实现分段需要在compile_fx的结果上再做一层包壳。思路是先用 Inductor 把整张图编译成若干个 kernel然后在运行时把 kernel 按照 range 策略依次执行。class PiecewiseCompiledFunction: def __init__(self, kernels, split_plans): self.kernels kernels self.split_plans split_plans def __call__(self, *args): results [] # 按 split plan 把输入切段 for i, plan in enumerate(self.split_plans): local_args [x[plan.slice_] if isinstance(x, torch.Tensor) else x for x in args] output self.kernels[i](*local_args) results.append(output) # 拼接输出 return torch.cat(results, dim0) def piecewise_compile(gm, example_inputs, split_sizes): generated inductor.compile_fx(gm, example_inputs) # 这里用 split_sizes 来定义 range 分片 kernels [] plans [] for i in range(len(split_sizes) - 1): start, end split_sizes[i], split_sizes[i 1] plan slice(start, end) plans.append(plan) kernels.append(generated) # 实际要按 plan 重新 codegen return PiecewiseCompiledFunction(kernels, plans)上面代码里的kernels.append(generated)是示意实际工作中每个 piece 需要用对应 slice 重新编译不能直接复用同一个生成结果。否则只是把一份 kernel 反复调用完全没有分段编译的意义。正确的做法是把需要被 split 的 Tensor 先用torch.narrow切片后作为输入写进新的 GraphModule 里再编译。也就是说每个 piece 都有一个独立的 subgraphsubgraph 的输入是切片后的张量输出是切片后的结果。4.3 一次实际跑通的例子我在本地环境PyTorch 2.3CPU 版跑过这样一个简单例子输入一个[1024, 512]的浮点张量计算x * 2 1。用split_sizes [0, 256, 512, 1024]分成三段编译。import torch torch.compile(backendpiecewise) def piecewise_relu(x): return torch.relu(x) * 2 1 x torch.randn(1024, 512) out piecewise_relu(x) print(out.shape)跑通的关键点在于注册后端的时候返回的 callable 需要对输入example_inputs的 dtype、device、shape 做一次适配否则后续真实输入形状变化时编译缓存会失效甚至直接报错。这个小 demo 性能上并没有优势但它验证了整条链路Dynamo 捕获计算图、自定义后端接管、内部调用 Inductor、按 range 分片、分片编译、调用拼接、返回结果。有了这个底座后面做任何高级优化都顺理成章了。5. 影响范围与应用场景这个机制能干什么5.1 自研硬件和异构计算的落地方案逐 range 编译最现实的使用场景是自研硬件和异构计算。很多 AI 加速芯片的编译器并不能直接吃 Inductor 生成的完整 Triton 内核但它们往往支持简单的循环代码。通过 PiecewiseBackend你可以把 Inductor 生成的复杂 kernel 拆成一段段相对简单的循环片段每一段再映射到硬件编译器能识别的 IR 上。这样就不需要为每一种新硬件从零搭建一套完整的图编译前端只要实现一个“range 到 target IR”的翻译器即可。从行业实践来看这个模式对应到很多加速卡的接入方案里都是先接torch.compile的自定义后端再用分段翻译的方式喂给自己底层的编译器工具链。成本比从算子库逐个手写要低而且能吃到 Inductor 在高层图优化上的红利。5.2 CPU 场景下的降级和调优另一个典型场景是 CPU。Triton 在 GPU 上很强但在 CPU 上默认的 Inductor 生成的是 C 代码走 OpenMP 路线。在某些老旧的 CPU 环境或者特殊指令集环境下生成的代码未必最优。逐 range 编译给了一个灵活的空间你可以把热循环的前半段用高性能方案比如显式 SIMD intrinsic后半段用兼容性更强的普通 C 代码两段拼接在一个 kernel 里。实测下来这种做法在 AVX2 支持不全的机器上特别有效能避免编译器生成非法指令导致的崩溃。5.3 调试、可观测性和 AOT 落盘分段编译还有一个隐藏好处它天然适合做调试和可观测性。整图编译时一旦某个算子出错你面对的是一个巨大、无法断点命中的编译产物。而把 range 切成段后每一段对应一小段百分百可控的代码出错时可以精确定位是哪个 loop piece 的问题。我做过一个内存对齐排查就是这么定位出来的生成代码里有段循环访问越界整图跑的时候要跑很久才触发崩坏。切成逐 range 的版本后很快锁定到第二个 piece 没有按 32 字节对齐做 pads最终在排布和分配逻辑里找到了根因。另外逐 range 编译也让 AOT提前编译落盘变得更容易。每一段可以在部署环境编译完序列化成独立的二进制或源码文件需要的时候再加载拼接。这样在部分需要国产化或离线部署的场景里可以绕过运行时编译器依赖只保留加载和执行组件。6. 常见问题与排查技巧实录6.1 注册后端后 torch.compile 报 “No such backend”这个是最常见的入门坑。通常原因有三类第一register_backend函数所在的模块没被导入第二后端名称拼写不一致第三在 Jupyter 里注册了但 kernel 状态被重启。排查办法很粗暴写一个单独的.py文件在最顶部导入注册模块然后再调用torch.compile。import my_compiler_backend # noqa: F401 import torch torch.compile(lambda x: x 1, backendmy_backend_name)如果还报错就在注册函数里加一个print标记确认它有没有被调用。很多时候是注册了但 Dynamo 走了缓存根本没走你新写的后端。6.2 中间张量 shape 对不上分段编译后每个 piece 的输入和输出张量 shape 必须和 GraphModule 里的预期严格一致。最容易出错的地方是torch.cat拼接阶段尤其是当某个 split 的 slice 大小不整除最后一个维度时。比如输入是[1000, 10]你切成[0, 256]、[256, 512]、[512, 1000]最后一段是488行而不是 256 的整数倍。如果编译时没有为最后一段做动态 shape 处理调用时会直接崩。解决方式是在编译每个 piece 时把 split 边界作为常量传入而不是动态计算。6.3 性能反而变差不是所有场景都适合逐 range 编译。小 kernel、短循环、算子简单的情况下分段引入的调用开销和内存拼接开销很可能超过它带来的编译优化收益。我建议做一个基准对比表用同一模型在不同 split 策略下的耗时和编译时间做对比策略编译时间推理耗时备注整图编译1.2s2.3ms默认 Inductor2 段切分1.5s2.1ms略有提升4 段切分2.1s2.6ms开销变大动态分段1.8s2.0ms需要 tuned 参数从表里可以明显看到段数翻倍并不一定带来正收益。只要某个 split 让内存访问模式变差整体表现就会被拖下来。不要盲目追求细粒度。注意分段后每个 piece 的 kernel 长度都不长时GPU 上的调度开销会成为瓶颈。GPU 场景建议每个 piece 至少包含几千个元素否则 kernel launch 的时间会盖过计算时间。6.4 与 torch.compile 缓存机制的冲突Dynamo 有自己的代码缓存机制当模型的输入 shape、dtype、device 不变时会直接复用编译产物。分段后端的动态分片逻辑如果依赖运行时输入 shape就会频繁击穿缓存。处理办法是分片策略要尽量在编译期确定。如果输入 shape 会变化在__call__里加一层检查按 shape 打标缓存不要每次都重新编译。否则你会发现看起来编译时间不长的程序真实运行里一小时能花掉四十分钟在编译上。6.5 调试技巧把生成的代码打印出来Inductor 生成的代码并不是完全黑盒。通过环境变量可以把它生成的 Triton 或 C 源码输出出来TORCH_LOGSoutput_code python my_script.py分段编译后你会看到每一段生成的源码。用文本对比工具比对不同段之间的差异经常能发现重复计算、无效加载、公共子表达式没有消除等问题。这个方法在我调试 piece 边界时帮了大忙。7. 从工程角度看这套机制值得投入吗如果只是图个性能提升逐 range 编译不一定是最优解。Inductor 默认的整图编译在绝大多数场景下已经足够好你花大力气拆段很可能得到的只是编译时间的小幅下降运行性能甚至原地踏步。但如果你是在做编译器适配、硬件接入、或者需要深度理解代码生成细节那 PiecewiseBackend 和 Inductor 对接这套东西就非常值得投入。它把最初看起来像一个黑盒的编译流程拆解成了你可以逐段验证和调优的组件。而且这种“按循环分段”的思路并不局限于 PyTorch 和 Inductor放到 TVM、MLIR 或者其他任意一个带循环结构的代码生成器里都能复用。我从第一次尝试跑通一个最小分段后端到真正把分段策略应用到一个真实推理项目里前前后后用了大概一周。最耗时间的不是写代码而是理解 Inductor 的调度器和循环表示之间那些隐含的约定。好在这些约定一旦摸清后面再调整分段逻辑就是改参数和边界条件的事情不会有推倒重来的风险。最后分享一个我的个人经验不要在 GPU 上刚起步的时候就去搞复杂的动态分片。先固定分片比例、跑通链路、打印生成代码、看性能数据等对负载特征有了足够的认知再引入动态调整。一步一步来比一开始就追求“全自动、最优切分”要靠谱得多。