调试 Triton¶
本教程提供了调试 Triton 程序的指导。它主要为 Triton 用户编写。对探索 Triton 后端(包括 MLIR 代码转换和 LLVM 代码生成)感兴趣的开发者可以参考此章节,以探索调试选项。
使用 Triton 的调试操作¶
Triton 包含四个调试操作符,允许用户检查张量值
static_print
和static_assert
用于编译时调试。device_print
和device_assert
用于运行时调试。
device_assert
仅在 TRITON_DEBUG
设置为 1
时执行。其他调试操作符无论 TRITON_DEBUG
的值如何都会执行。
使用解释器¶
解释器是一个简单且有用的调试 Triton 程序的工具。它允许 Triton 用户在 CPU 上运行 Triton 程序并检查每个操作的中间结果。要启用解释器模式,请将环境变量 TRITON_INTERPRET
设置为 1
。此设置会导致所有 Triton 内核绕过编译,并通过解释器使用 Triton 操作的 numpy 等效项进行模拟。解释器顺序处理每个 Triton 程序实例,一次执行一个操作。
使用解释器主要有三种方式
使用 Python
print
函数打印每个操作的中间结果。要检查整个张量,请使用print(tensor)
。要检查idx
处的单个张量值,请使用print(tensor.handle.data[idx])
。附加
pdb
以对 Triton 程序进行逐步调试TRITON_INTERPRET=1 pdb main.py b main.py:<line number> r
导入
pdb
包并在 Triton 程序中设置断点import triton import triton.language as tl import pdb @triton.jit def kernel(x_ptr, y_ptr, BLOCK_SIZE: tl.constexpr): pdb.set_trace() offs = tl.arange(0, BLOCK_SIZE) x = tl.load(x_ptr + offs) tl.store(y_ptr + offs, x)
局限性¶
解释器有几个已知的局限性
它不支持
bfloat16
数字类型的操作。要对bfloat16
张量执行操作,请使用tl.cast(tensor)
将张量转换为float32
。它不支持间接内存访问模式,例如
ptr = tl.load(ptr) x = tl.load(ptr)
使用第三方工具¶
对于 NVIDIA GPU 上的调试,compute-sanitizer 是一个检查数据竞争和内存访问问题的有效工具。要使用它,请在运行 Triton 程序的命令前加上 compute-sanitizer
。
对于 AMD GPU 上的调试,您可能想尝试 ROCm 的 LLVM AddressSanitizer。
要详细可视化 Triton 程序中的内存访问,请考虑使用 triton-viz 工具,它与底层 GPU 无关。