[Bug已解决] torch.compile(modemax-autotune) 在 bmm 自动调优时因模板输入去重导致位置参数过多崩溃解决方案一、现象长什么样你用torch.compile开了最强优化档训练或推理一个含bmm批量矩阵乘的模型model MyModel() model torch.compile(model, modemax-autotune) out model(x)结果在编译阶段autotuning自动调优内核参数崩了报错类似TypeError: triton kernel() takes N positional arguments but M were given # 或 torch.compile(modemax-autotune) crashes with too many positional arguments during bmm autotuning when template inputs are deduplicated本 issuepytorch/pytorch#188069点出根因当 bmm 的 template Triton 模板内核输入被「去重deduplicated」后Inductor 在生成 autotuning 调用时仍按「未去重」的参数列表传参导致传给 Triton kernel 的位置参数比它声明的多触发「too many positional arguments」崩溃。 本文聚焦max-autotune 的 autotuning 是什么、template 输入去重怎么会多传参、怎么绕开这个编译期崩溃。二、背景max-autotune 与 bmm 的 autotuningtorch.compile(modemax-autotune)会做最激进优化其中重要一环是autotuning自动调优对像bmm/mm这类算子Inductor 生成多个 Triton 内核候选不同 BLOCK 大小、不同流水线策略然后用真实输入实际跑一遍选最快的。bmm的 Triton 模板内核签名形如triton.jit def bmm_kernel(inp1_ptr, inp2_ptr, out_ptr, M, N, K, stride_..., ...): ...关键点Triton kernel 的位置参数是固定的多一个少一个都会TypeError。 Inductor 在 autotuning 时会把「输入张量 形状/步长标量」打包成参数列表传给 kernel。问题出在当 bmm 的多个 template 输入其实是同一个张量例如bmm(a, a)或某优化把重复输入去重时Inductor 的「去重」逻辑改了参数列表但生成 autotuning 调用的代码没同步仍按原列表传参 → 位置参数数量不匹配 → 崩溃。三、为什么 template 输入去重会多传参Inductor 生成 autotuning 代码时大致流程收集 template 的「输入张量列表」如[inp1, inp2]为减少冗余做去重dedup若inp1 is inp2同一对象列表变成[inp1]但生成「调用 kernel」的代码时用了未去重前的索引/数量来构造参数于是仍尝试传inp1, inp2两个张量Triton kernel 声明只接收一个去重后的张量 标量 → 实际传了两个 →too many positional arguments。 或者反过来去重后 kernel 签名少了参数但调用处没减 → 多传。总之是Inductor 的「去重」与「autotuning 调用生成」两套逻辑没对齐属于代码生成 bug。 注意它只在modemax-autotune会真正跑 autotuning时触发普通modedefault不跑完整 autotuning往往不崩。四、最小可运行复现带守卫下面演示「bmm 相同输入」这种易触发去重的情形。实际崩在 autotuning需 GPU max-autotune用守卫包裹import torch def make_model(): # 一个含 bmm 的小模型 class M(torch.nn.Module): def forward(self, a, b): # bmm 批量矩阵乘 return torch.bmm(a, b) return M() def run_max_autotune(): if not torch.cuda.is_available(): print([skip] 无 GPU仅说明max-autotune 的 bmm autotuning 可能崩) return model make_model() try: compiled torch.compile(model, modemax-autotune) a torch.randn(2, 4, 8, devicecuda) b torch.randn(2, 8, 3, devicecuda) out compiled(a, b) print(max-autotune 输出形状:, out.shape) except TypeError as e: if positional arguments in str(e): print([确认] 命中 bmm autotuning 多传参 bug (#188069):, e) else: print(其它 TypeError:, e) except Exception as e: print(其它错误:, e) if __name__ __main__: run_max_autotune()要点modemax-autotune才会跑 autotuning 触发若把a、b设成同一张量torch.bmm(a, a)去重更易触发。五、解决方案一不用 max-autotune降级 mode最直接绕过别用max-autotune改用不跑完整 autotuning 的模式import torch # 默认模式不跑完整 bmm autotuning避开该 bug compiled torch.compile(model, modedefault) # 或 reduce-overheadCUDA graphs不一定触发该 autotuning 路径 compiled torch.compile(model, modereduce-overhead)代价少了一部分 autotuning 带来的峰值性能但编译稳定不崩。绝大多数场景default已经很快。六、解决方案二关闭 autotuning 相关开关如果一定要 max-autotune 的部分优化可单独关掉 autotuningimport torch._inductor.config as cfg # 关闭 bmm / mm 的 autotuning用固定内核不再多传参 cfg.max_autotune False cfg.max_autotune_gemm False cfg.autotune_local_cache False # 或只对当前 compile 关 compiled torch.compile( model, modemax-autotune, options{max_autotune: False}, # 关掉 autotune 部分 )这样保留max-autotune的其它优化如 layout 优化但不跑会崩的 bmm autotuning。七、解决方案三避免 bmm 输入是同一张量 / 强制不共享若你的模型里恰好bmm(x, x)输入相同去重更易触发。临时改写下让两个输入不是同一对象import torch def safe_bmm(a, b): # 若 a、b 是同对象复制一份避免 Inductor 去重触发多传参 if a is b: b b.clone() return torch.bmm(a, b) # 或在前向里显式分开 out torch.bmm(q, k) # q、k 本就不同 → 一般不触发 # 若必须 bmm(a,a)用 safe_bmm(a, a)注意这只是降低触发概率去重也可能因其它优化发生根本修复还是依赖 PyTorch 升级。八、解决方案四升级到修复版本#188069 是 Inductor 已知缺陷官方会修。升级pip install --upgrade torch --index-url https://download.pytorch.org/whl/cu124 python -c import torch; print(torch.__version__)判断修复同样模型torch.compile(modemax-autotune)含bmm不再报too many positional arguments编译通过。修复前用「降级 mode / 关 autotuning」绕过。九、排查清单报too many positional arguments且栈在triton/autotuning/bmm template→ 确认是 #188069。是否用了modemax-autotune是 → 先降级modedefault验证能编译确认是 autotuning 触发。关 autotuningcfg.max_autotuneFalse/cfg.max_autotune_gemmFalse保留其它优化。检查bmm(x, x)这类同输入临时.clone()分开降低触发概率。升级 PyTorch 到修复版本CI 里对触发模型暂时不用 max-autotune。若贡献修复定位 Inductor template 的「dedup 输入」与「autotuning 调用生成」对齐逻辑。十、小结torch.compile(modemax-autotune) crashes with too many positional arguments during bmm autotuning when template inputs are deduplicated#188069的本质是Inductor 在 max-autotune 的 bmm autotuning 阶段对 template 输入做了「去重」但生成 Triton kernel 调用参数的代码没同步仍按未去重列表传参导致传给 kernel 的位置参数比声明多触发 TypeError 崩溃。它只在跑完整 autotuning 时触发。 应对降级 mode用modedefault/reduce-overhead不跑完整 bmm autotuning编译稳定关 autotuningcfg.max_autotuneFalse/max_autotune_gemmFalse保留其它优化避开同输入bmm(x, x)临时.clone()分开降低去重触发概率升级修复版本后 max-autotune 可正常用于含 bmm 的模型。 记住崩溃在「编译期 autotuning」不是运行时只要不跑完整 autotuningbmm 模型照样能被torch.compile加速。