调试 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 无关。