PyTorch AOTInductor:Triton 内核 “index out of bounds“ 断言故障排查完全指南 PyTorch AOTInductorTriton 内核 index out of bounds 断言故障排查完全指南【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch本篇指南面向使用 PyTorch AOTInductorAOTI做 AOT 编译部署时遇到的典型运行期故障Triton 内核中Assertion index out of bounds: 0 tmpN ksM failed。读完本文你将掌握一套从报错信息出发、定位wrapper.cpp中失败内核、沿动态形状变量回溯到模型输入、再映射回原始 PyTorch 代码的完整排查链路并能针对空张量边界这类常见根因给出可落地的修复方案。1. 识别故障什么错误适用本指南当你在运行 AOTI 编译出的模型时看到如下形式的报错即可适用本指南的排查流程/var/tmp/torchinductor_*/.../*.py:NN: unknown: block: [X,Y,Z], thread: [X,Y,Z] Assertion index out of bounds: 0 tmpN ksM failed.报错信息的关键字段字段取值含义文件路径/var/tmp/torchinductor_*/*.py运行时生成的 Triton 内核文件行号:NN断言失败的生成内核中的行号Block/Thread[X,Y,Z]CUDA block 与 thread 索引断言0 tmpN ksM索引tmpN必须落在[0, ksM)区间内读懂这条断言tmpNTriton 内核中某个计算出来的索引值tmp前缀是 Inductor 对公共子表达式变量自动命名的产物ksM一个动态内核尺寸参数runtime value对应某个动态形状维度断言在tmpN 0或tmpN ksM时失败。这条断言从哪里来源码佐证该报错文案并非 Triton 本身产生而是 Inductor 代码生成阶段的确定性产物torch/_inductor/codegen/common.py 中的indirect_assert方法负责为间接索引生成越界检查。当同时存在上下界时它拼接出条件(lower var) (var upper)并生成断言语句{assert_function}({cond}, index out of bounds: {cond_print})——这正是报错中index out of bounds: 0 tmpN ksM文案的直接来源。若索引表达式带 mask例如由布尔索引产生的受条件控制的访问条件还会退化为({cond}) | ~({mask})的形式即 mask 为假的位置不参与越界判断。对 Triton 后端assert_function被实现为tl.device_assert见 torch/_inductor/codegen/triton.py运行期真正执行断言的辅助函数device_assert_then位于 torch/_inductor/runtime/triton_helpers.py内部就是调用tl.device_assert(cond, msg)。因此可以确认这条断言是 Inductor 在代码生成时根据索引表达式可推导的上下界自动插入的运行时保护它触发意味着运行时数据越过了编译期认为不可能的取值范围。2. Step 1获取 AOTI 编译产物排查的前提是能拿到当次编译产出的 AOTI 包。它通常是一个.pt2包或一个包含wrapper.cpp文件的已解压归档。关键文件*.wrapper.cpp中包含所有 Triton 内核源代码以注释形式内嵌内核启动配置launch configuration输入/输出张量的映射动态形状变量定义int64_t sXXX ...形式。这一说法可以直接在 Inductor 的编译流程中得到印证torch/_inductor/codecache.py 在写出wrapper.cpp后会给独立写出的kernel.cpp头部追加注释// Triton kernels are embedded as comments in {wrapper_path}在config.aot_inductor.package_cpp_only模式下内核代码甚至会被直接拼接进 wrapper。所以wrapper.cpp是单文件排查 AOTI 运行时故障的核心入口。3. Step 2在 C Wrapper 中定位失败内核搜索断言模式从报错中提取断言特征例如tmp18 ks0然后在 wrapper 中搜索# 搜索特定断言 grep -n tmpN ksM /path/to/*.wrapper.cpp # 获取断言附近的上下文前 80 行、后 20 行 grep -n -B80 -A20 tmpN ksM /path/to/*.wrapper.cpp找到完整内核定义内核在 C wrapper 中是以 Python 文档字符串注释形式内嵌的典型长相如下/* async_compile.triton(triton_red_fused_..., import triton import triton.language as tl ... def triton_red_fused_...(in_ptr0, out_ptr1, ks0, xnumel, r0_numel, ...):triton_red_fused_...这类命名前缀triton_red/triton_per本身就携带信息red表示带规约维度的内核内核参数中的r0_numel等是规约维度规模ks0则是参与索引计算的动态形状变量。4. Step 3理解内核逻辑——空张量ks0 0模式分析通向断言的代码路径。造成 index out of bounds 的常见模式是动态尺寸为 0 的空张量。当动态形状ks0 0时tmp13 (-1) 0 -1索引回绕wrap-around逻辑算出-1断言0 -1 0失败示例内核模式tmp13 (-1) ks0 # 即 ks0 - 1 tmp14 tl.where(tmp12, tmp10, tmp13) # 条件成立取 tmp10否则取 ks0-1 tmp15 ks0 tmp16 tmp14 tmp15 # 负数索引的回绕加 ks0 tmp17 tmp14 0 tmp18 tl.where(tmp17, tmp16, tmp14) # 若为负则加 ks0 # 断言0 tmp18 ks0 tl.device_assert(((0 tmp18) (tmp18 ks0)), index out of bounds)这段逻辑正是第 1 节所述indirect_assert生成断言前的典型索引计算(-1) ks0对应取最后一个元素这类语义负数回绕对应xindex风格的下标归一化。关键问题在于回绕公式tmp16 tmp14 ks0只在-ks0 tmp14 0时才能把索引拉回[0, ks0)。当ks0 0时tmp13恒为-1回绕后仍是-1而合法区间[0, 0)是空集断言必然失败。这解释了为什么编译时样本里没有 size 0 的输入运行期一旦出现就必然触发断言——不是内核写错而是输入分布超出了编译期假设。5. Step 4识别动态形状变量找到内核的调用点grep -n call_triton_KERNEL_NAME /path/to/*.wrapper.cpp调用点示例call_triton_red_fused_...(arg1415_1, buf696, s607, 1L, s13, ...);参数映射参数取值含义in_ptr0arg1415_1输入张量out_ptr1buf696输出缓冲区ks0s607动态形状——即失败的边界找到形状变量的定义grep -n int64_t s607 /path/to/*.wrapper.cpp它会揭示是哪个输入张量的哪个维度定义了该形状int64_t s607 arg1416_1_size[0];至此你已把内核里失败的 ks0与具体某个输入的第 0 维精确关联起来排查焦点从内核层上升到输入层。6. Step 5回溯到模型输入找到输入编号输入按顺序编号。确认该参数对应第几个输入grep -n inputs_info_\[INDEX\].name argNNN_1 /path/to/*.wrapper.cpp检查输入约束grep -n argNNN_1_size\[0\] /path/to/*.wrapper.cpp注意寻找形如这样的守卫if (arg_size[0] 230400) { // 只有上界检查——没有下界常见问题AOTI 只生成上界检查upper bound check不会生成 1的下界检查。也就是说编译期样本给出的尺寸上限被固化成了校验而该维度可以为 0这一事实完全未被表达运行期一旦传入 size 0 的张量就会直接穿透到内核层的device_assert。7. Step 6映射回模型代码利用 Source Node 注释C wrapper 中包含注释标明每个内核由哪些 PyTorch 操作生成grep -n -B5 call_triton_KERNEL_NAME /path/to/*.wrapper.cpp | grep Source Nodes注释示例// Topologically Sorted Source Nodes: [slice_1, sub_89, cumsum, ge_231, where_2, index_copy]把算子映射到 Python 代码ATen 算子对应的 Python 代码模式cumsumtorch.cumsum(tensor, dim0)subidx - 1geidx 0wheretorch.where(condition, ...)index_copytensor.index_copy(0, indices, source)上表组合cumsumsubgewhereindex_copy典型对应jagged/变长输入按长度做前缀和、生成有效索引、掩膜后写回的稀疏事件处理模式——这类代码恰恰最容易在lengths为空时让某个维度变成 0与第 4 节ks0 0模式互相印证。8. 根因分析常见根因运行期出现空张量jagged/变长张量在运行期尺寸为 0但编译期样本从未覆盖这种情况缺少下界守卫AOTI 只生成上界检查不生成 1的下界检查size 0 的输入畅通无阻地进入内核导出样本未覆盖边界AOTI 导出时的样本输入从未包含该边界情形。三者通常叠加成立样本缺边界3导致编译期无法感知1而缺失的下界守卫2让故障以内核断言这种较难读懂的形式暴露出来。9. 修复建议方案一在 forward 中加守卫def forward(self, lengths: torch.Tensor, ...) - torch.Tensor: if lengths.numel() 0: device lengths.device return torch.empty(0, self.output_dim, devicedevice) # ... 其余方法体方案二修复具体算子在出问题的操作处增加空张量处理def process_events(self, lengths: torch.Tensor, ...): if lengths.numel() 0: return torch.empty(0, self.emb_dim, devicelengths.device) # ... 其余方法体方案三在 AOTI 导出时纳入边界样本确保 AOTI 导出阶段的样本输入包含空张量size 0最小尺寸张量size 1预期最大尺寸。三个方案中方案一/二在模型侧显式处理边界最稳妥方案三让编译期见过边界可配合显式守卫形成双重保险。10. 速查常用调试命令与环境变量AOTI Wrapper 搜索命令汇总# 按断言模式找内核 grep -n tmpN ksM *.wrapper.cpp # 获取完整内核上下文 grep -n -B80 -A20 ASSERTION_PATTERN *.wrapper.cpp # 找内核调用点 grep -n call_KERNEL_NAME *.wrapper.cpp # 找动态形状定义 grep -n int64_t SHAPE_VAR *.wrapper.cpp # 找输入映射 grep -n inputs_info_\[INDEX\].name *.wrapper.cpp # 找尺寸约束 grep -n SHAPE_VAR_size\[0\] *.wrapper.cpp调试用环境变量# 在 torch.compile 期间开启调试输出 export TORCH_COMPILE_DEBUG1 # 把生成的内核保存到持久目录 export TORCHINDUCTOR_CACHE_DIR/path/to/save/kernels # 开启 CUDA launch blocking获得精确堆栈 export CUDA_LAUNCH_BLOCKING1说明报错中出现的/var/tmp/torchinductor_*/.../*.py正是 Inductor 内核缓存的默认临时目录运行环境重启后可能被清理因此TORCHINDUCTOR_CACHE_DIR指向持久路径对复现排查很关键CUDA_LAUNCH_BLOCKING1让内核错误在 launch 处同步报出便于把断言与触发它的调用点一一对应会牺牲部分性能仅用于调试。小结AOTI 场景下的index out of bounds断言本质上是一次运行期输入超出编译期假设的信号。排查路径可以概括为一条链报错断言 → wrapper.cpp 中的内核定义Step 2→ 内核内索引计算与ks0的语义Step 3→int64_t sXXX argNNN_1_size[D]的形状定义Step 4→ 输入映射与缺失的下界守卫Step 5→ Source Nodes 注释映射回 Python 代码Step 6。其底层机制在 torch/_inductor/codegen/common.py 的indirect_assert、torch/_inductor/codegen/triton.py 的tl.device_assert实现以及 torch/_inductor/codecache.py 的 wrapper 生成逻辑中都有明确对应。按此链路走通一次绝大多数空张量/边界尺寸类 AOTI 运行时断言都能定位到具体的模型输入并给出修复。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考