[Bug已解决] AOTInductor 下 bmm 大 batch 动态形状 CUDA 失效DISABLED test解决方案一、现象长什么样PyTorch CI 里有一条DISABLED test_bmm_large_batch_dynamic_cuda (__main__.AOTInductorTestDualWrapper)它测试的是AOTInductor提前编译部署见 26 节在 CUDA 上对大 batch 的bmm批量矩阵乘且动态形状的编译支持。被DISABLED说明这个组合下 AOTInductor 生成的代码有 bug——可能大 batch 的 bmm 在动态形状下越界、形状推导错、或结果不对。bmm(input, mat2)是「批量矩阵乘」input 形状[b, n, m]mat2 形状[b, m, p]输出[b, n, p]对每个 batch 独立做矩阵乘。当bbatch很大且是动态时AOTInductor 的编译容易出问题。本文讲清楚 bmm、大 batch 动态形状的坑、以及如何规避。二、bmm 是什么torch.bmm做批量矩阵乘batch 维必须匹配import torch a torch.randn(4, 3, 5) # [b4, n3, m5] b torch.randn(4, 5, 2) # [b4, m5, p2] c torch.bmm(a, b) # [4, 3, 2] print(c.shape)对比torch.matmulmatmul会「广播」batch 维如[4,3,5] [5,2]也能乘而bmm要求 batch 维严格相等不做广播。这使它更适合「对齐的批量算子」。三、为什么「大 batch 动态形状」在 AOTInductor 下会踩坑AOTInductor 在编译期要确定所有形状 / 循环边界见 26 节 int64 上界问题。当涉及大 batchb 很大循环 / 网格划分可能用 int64 索引AOTInductor 生成的 kernel 在 int64 边界上可能算错同源 26 节动态 batchb 在编译期未知AOTInductor 要为「符号 batch 维」生成能适应任意 b 的代码符号推导若出错形状错位两者叠加大且动态的 batch让 bmm 的「每 batch 偏移计算」在 AOT 生成的代码里越界或错乱。于是 CI 用「大 batch 动态」的结果一致性断言守护发现不对就DISABLED。四、可运行复现 bmm 大 batch 动态 AOT 流程下面脚本演示 bmm 动态形状 AOT 导出无 GPU / 无 AOT 工具时优雅跳过import torch from torch.export import Dim def demo(): model torch.nn.Linear(8, 8) # 用 Linear 模拟每 batch 的小计算 # 构造 bmmbatch 维动态 batch Dim(batch, min1, max1024) def f(a, b): return torch.bmm(a, b) a torch.randn(4, 3, 8) b torch.randn(4, 8, 2) try: ep torch.export.export( f, (a, b), dynamic_shapes{a: {0: batch}, b: {0: batch}}, ) print(bmm 动态导出成功) # AOT 编译需要工具链 # torch._inductor.aoti_compile_and_package(ep, bmm.pt2) except Exception as e: print(bmm 动态导出/AOT 失败, type(e).__name__, str(e)[:120]) if __name__ __main__: demo()如果你在 AOT 编译大 batch 动态 bmm 时失败 / 结果错就复现了该问题。五、解决方案一限制动态 batch 上界Constraint和 26 节同理给 batch 维设明确上界让 AOTInductor 生成正确的循环边界避免大 batch 的 int64 越界import torch from torch.export import Dim batch Dim(batch, min1, max512) # 明确上界 ep torch.export.export( f, (a, b), dynamic_shapes{a: {0: batch}, b: {0: batch}}, )上界不要太离谱如 1e9否则 AOT 仍可能生成溢出边界。设成你业务里 batch 的真实最大值。六、解决方案二大 batch 拆成小 chunk避免单次超大 bmm如果 batch 极大如几万把 bmm 拆成多个小 batch 的 bmm避免单次超大 bmm 触发 AOT bug也降显存import torch def chunked_bmm(a, b, chunk1024): # a: [B, n, m], b: [B, m, p] B a.size(0) outs [] for s in range(0, B, chunk): e min(s chunk, B) outs.append(torch.bmm(a[s:e], b[s:e])) return torch.cat(outs, dim0) a torch.randn(5000, 3, 8) b torch.randn(5000, 8, 2) out chunked_bmm(a, b) print(分块 bmm 输出, out.shape)分块后每次 bmm 的 batch 都在安全范围AOT 编译更稳。七、解决方案三先 torch.compile 验证再 AOT 部署在投入 AOTInductor 部署前先用torch.compileJIT验证 bmm 大 batch 动态在 CUDA 上正确compiled torch.compile(lambda a, b: torch.bmm(a, b)) out compiled(a, b) print(torch.compile bmm 正常)如果torch.compile正常但 AOT 失败说明是 AOT 专属的生成 bug可暂时用torch.compile部署牺牲 AOT 的无 Python 部署优势等 PyTorch 修复。八、解决方案四用 aot_eager 隔离「导出」与「最终编译」aot_eager后端做 AOT 的图捕获但不做激进 kernel 生成用来验证 bmm 能否被 AOT 流程接受compiled torch.compile(lambda a, b: torch.bmm(a, b), backendaot_eager) out compiled(a, b)若aot_eager正常但完整 AOTInductor 错说明「最终 kernel 生成大 batch 动态」有 bug而非导出逻辑。九、解决方案五升级 PyTorchtest_bmm_large_batch_dynamic_cuda是 AOTInductor 对大 batch 动态 bmm 的 Known Issue。新版本会逐步修复 int64 边界 / 动态 batch 推导。查看并升级import torch print(PyTorch, torch.__version__)十、如何判断你踩的是同一条你用 AOTInductor 部署且模型含bmmbatch 很大且是动态形状现象是「AOT 编译后结果错 / 越界 / 崩溃」限制 batch 上界 / 分块 / 退回 torch.compile 后恢复。命中即说明踩中该 AOT bmm 大 batch 动态 bug。十一、小结DISABLED test_bmm_large_batch_dynamic_cuda揭示AOTInductor 在大 batch 动态 bmm 的 CUDA 代码生成有 bugint64 边界 / 动态 batch 推导。应对给 batch 维设明确上界第五节避免大 batch 溢出超大 batch分块 bmm第六节单次规模可控部署前先torch.compile验证逻辑第七节确认是 AOT 专属问题用aot_eager隔离导出与最终编译第八节升级到修复该 bug 的 PyTorch第九节。bmm 是「批量算子」的主力batch 维的动态 巨大正是 AOTInductor 编译期确定性的软肋。把 batch 上界钉死、把超大 batch 拆小让编译期能生成正确的循环边界——AOT 部署才能既快又对。