AOTInductor Triton SASS 检查
内容提要
本文介绍如何检查AOTInductor生成的Triton内核SASS。AOTInductor是PyTorch后端编译器,可生成预优化推理引擎,其Triton内核以cubin文件存于pt2包中。解压后用nvdisasm或cuobjdump反汇编即可查看SASS。示例验证了128位向量化内存指令LDG.E.128和STG.E.128确实被生成。
延伸解读
为何需要检查 Triton 内核的 SASS
AOTInductor 作为 PyTorch 的提前编译后端,会将 Triton 内核编译为 cubin 文件。通过反汇编这些 cubin 得到 SASS,可以验证编译器是否按预期生成了特定指令,例如 128 位向量化内存访问。这对于性能调优和确认底层代码生成行为至关重要,尤其是在依赖向量化提升内存带宽的场景中。
128 位向量化指令的生成条件
Triton 编译器根据输入张量的基地址和静态偏移推断向量化指令。PyTorch 缓存分配器默认返回 256 字节对齐的张量,若步长偏移也满足 16 字节对齐,编译器便可生成 LDG.E.128 和 STG.E.128 指令。示例中通过 tl.multiple_of 提示对齐,最终在 SASS 中确认了这些指令的存在。
从 pt2 包提取并反汇编 cubin 的步骤
AOTInductor 生成的 pt2 包本质是 zip 归档,解压后可在 model/data/aotinductor 目录下找到 cubin 文件。使用 nvdisasm 或 cuobjdump --dump-sass 即可反汇编。示例中两种工具均成功输出 SASS,并定位到 LDG.E.128 和 STG.E.128 指令,验证了向量化的实现。
验证流程中的关键检查点
示例流程包括导出计算图、AOTI 编译打包、执行并验证按位一致,以及解包检查 cubin。按位一致确保编译未改变计算语义,而 SASS 检查则确认了底层指令符合预期。这种组合方法既保证了正确性,又提供了性能相关的底层洞察,适用于调试和优化自定义 Triton 内核。
Q&A
AOTInductor 生成的 Triton 内核保存在哪里?
AOTInductor 生成的 Triton 内核以 cubin 文件形式保存在 pt2 包(zip 归档)中。
如何查看 AOTInductor 生成的 Triton 内核的 SASS?
解压 pt2 包后,使用 nvdisasm 或 cuobjdump --dump-sass 对 cubin 文件进行反汇编即可查看 SASS。
Triton 编译器在什么条件下会生成 128 位向量化内存指令?
当输入张量基地址和步长偏移都满足 16 字节对齐时,Triton 编译器可以生成 128 位的 LDG.E.128 / STG.E.128 向量化内存指令。
为什么 PyTorch 张量默认适合生成 128 位向量化指令?
因为 PyTorch 缓存分配器分配的张量默认是 256 字节对齐的,这满足 16 字节对齐要求,从而有利于生成 128 位向量化指令。
在示例中如何验证生成的 SASS 包含 128 位向量化指令?
示例通过解压 pt2 包,对 cubin 文件运行 nvdisasm 和 cuobjdump,在输出中搜索 LDG.E.128 和 STG.E.128 指令来确认。
AOTInductor 是什么?
AOTInductor 是 PyTorch 的后端编译器,可以生成预优化的深度学习推理引擎,并兼容自定义 Triton 内核和性能自动调优。