注意
转到末尾 下载完整的示例代码。
在 torch.compile 中使用用户定义的 Triton 内核#
创建日期：2024 年 4 月 19 日 | 最后更新：2025 年 5 月 2 日 | 最后验证：2024 年 11 月 5 日
作者： Oguz Ulgen
用户定义的 Triton 内核可用于优化模型计算的特定部分。这些内核使用 Triton 语言编写，旨在更容易地实现峰值硬件性能。通过将用户定义的 Triton 内核与 torch.compile 结合使用，您可以将这些优化的计算集成到 PyTorch 模型中，从而有可能显著提高性能。
此秘籍演示了如何将用户定义的 Triton 内核与 torch.compile 结合使用。
先决条件#
在开始此秘籍之前，请确保您已具备以下条件
对
torch.compile和 Triton 有基本了解。请参阅PyTorch 2.3 或更高版本
支持 Triton 的 GPU
import torch
from torch.utils._triton import has_triton
基本用法#
在此示例中，我们将使用 Triton 文档中简单的向量加法内核与 torch.compile。有关参考信息，请参阅 Triton 文档。
if not has_triton():
print("Skipping because triton is not supported on this device.")
else:
import triton
from triton import language as tl
@triton.jit
def add_kernel(
in_ptr0,
in_ptr1,
out_ptr,
n_elements,
BLOCK_SIZE: "tl.constexpr",
):
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
x = tl.load(in_ptr0 + offsets, mask=mask)
y = tl.load(in_ptr1 + offsets, mask=mask)
output = x + y
tl.store(out_ptr + offsets, output, mask=mask)
@torch.compile(fullgraph=True)
def add_fn(x, y):
output = torch.zeros_like(x)
n_elements = output.numel()
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=4)
return output
x = torch.randn(4, device="cuda")
y = torch.randn(4, device="cuda")
out = add_fn(x, y)
print(f"Vector addition of\nX:\t{x}\nY:\t{y}\nis equal to\n{out}")
Vector addition of
X: tensor([-1.1766, -2.1809, 0.4788, -1.8701], device='cuda:0')
Y: tensor([ 0.5981, -0.1185, -0.7021, -1.4897], device='cuda:0')
is equal to
tensor([-0.5785, -2.2994, -0.2233, -3.3598], device='cuda:0')
高级用法#
Triton 的自动调优功能是一个强大的工具，可以自动优化 Triton 内核的配置参数。它探索一系列可能的配置，并选择最能为您的特定用例提供最佳性能的配置。
与 torch.compile 结合使用时，triton.autotune 有助于确保您的 PyTorch 模型尽可能高效地运行。以下是使用 torch.compile 和 triton.autotune 的示例。
注意
torch.compile 仅支持 triton.autotune 的配置和关键参数。
if not has_triton():
print("Skipping because triton is not supported on this device.")
else:
import triton
from triton import language as tl
@triton.autotune(
configs=[
triton.Config({"BLOCK_SIZE": 4}, num_stages=3, num_warps=8),
triton.Config({"BLOCK_SIZE": 4}, num_stages=4, num_warps=4),
triton.Config({"BLOCK_SIZE": 2}, num_stages=3, num_warps=8),
triton.Config({"BLOCK_SIZE": 2}, num_stages=4, num_warps=4),
],
key=[],
)
@triton.jit
def add_kernel_autotuned(
in_ptr0,
in_ptr1,
out_ptr,
n_elements,
BLOCK_SIZE: "tl.constexpr",
):
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
x = tl.load(in_ptr0 + offsets, mask=mask)
y = tl.load(in_ptr1 + offsets, mask=mask)
output = x + y
tl.store(out_ptr + offsets, output, mask=mask)
@torch.compile(fullgraph=True)
def add_fn(x, y):
output = torch.zeros_like(x)
n_elements = output.numel()
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
add_kernel_autotuned[grid](x, y, output, n_elements)
return output
x = torch.randn(4, device="cuda")
y = torch.randn(4, device="cuda")
out = add_fn(x, y)
print(f"Vector addition of\nX:\t{x}\nY:\t{y}\nis equal to\n{out}")
Vector addition of
X: tensor([-0.0302, 0.2330, -0.8989, 0.2358], device='cuda:0')
Y: tensor([1.5251, 0.5872, 0.3729, 1.0297], device='cuda:0')
is equal to
tensor([ 1.4949, 0.8202, -0.5260, 1.2655], device='cuda:0')
可组合性#
用户定义的 Triton 内核不自动支持所有 PyTorch 子系统。这可以在以下用例中看到
添加 CPU 回退
添加
FlopCounter公式与张量子类组合
要与附加的 PyTorch 子系统组合，请使用 torch.library.triton_op。
triton_op 是一种结构化的自定义操作符定义方式，它由一个或多个 Triton 内核支持：与常规自定义操作符 (torch.library.custom_op) 一样，您可以通过 torch.library 指定与 PyTorch 子系统的交互。但是，与 torch.library.custom_op 不同，后者会创建与 torch.compile 相关的“不透明”可调用对象，torch.compile 会跟踪 triton_op 的实现以应用优化。
这是集成 Triton 内核与 PyTorch 时应使用的 API 图表。
Triton 内核（无显式 |
|
|
|
|---|---|---|---|
支持推理 |
是 |
是 |
是 |
支持训练 |
在大多数情况下 |
是 |
是 |
支持 |
是 |
是 |
是 |
支持 |
在大多数情况下 |
在大多数情况下 |
在所有情况下 |
torch.compile 是否跟踪到实现中？ |
是 |
是 |
否 |
支持 AOTInductor |
是 |
是 |
否 |
支持 PyTorch 子系统，如 FlopCounterMode、CPU 回退、张量子类 |
否 |
是 |
是 |
使用 triton_op 封装 Triton 内核#
使用 torch.library.triton_op 封装可能调用一个或多个 Triton 内核的函数。使用 torch.library.wrap_triton 封装对 Triton 内核的调用。
from torch.library import triton_op, wrap_triton
@triton_op("mylib::mysin", mutates_args={})
def mysin(x: torch.Tensor) -> torch.Tensor:
out = torch.empty_like(x)
n_elements = x.numel()
wrap_triton(sin_kernel)[(n_elements,)](x, out, n_elements, BLOCK_SIZE=4)
return out
@triton.jit
def sin_kernel(
in_ptr0,
out_ptr,
n_elements,
BLOCK_SIZE: "tl.constexpr",
):
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
x = tl.load(in_ptr0 + offsets, mask=mask)
output = tl.sin(x)
tl.store(out_ptr + offsets, output, mask=mask)
您可以通过以下两种方式之一调用 triton_op。
x = torch.randn(3, device="cuda")
y = mysin(x)
z = torch.ops.mylib.mysin.default(x)
assert torch.allclose(y, x.sin())
assert torch.allclose(z, x.sin())
结果 triton_op 可与 torch.compile 和 AOTInductor 配合使用。
y = torch.compile(mysin)(x)
assert torch.allclose(y, x.sin())
添加训练支持#
使用 register_autograd 为 triton_op 添加自动梯度公式。优先使用此方法，而不是使用 torch.autograd.Function（后者与 torch.compile 存在各种可组合性陷阱）。
请注意，反向传播必须是 PyTorch 理解的操作符的组合。如果希望反向传播调用 Triton 内核，则这些内核也必须封装在 triton_op 中。
@triton.jit
def cos_kernel(
in_ptr0,
out_ptr,
n_elements,
BLOCK_SIZE: "tl.constexpr",
):
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
x = tl.load(in_ptr0 + offsets, mask=mask)
output = tl.cos(x)
tl.store(out_ptr + offsets, output, mask=mask)
@triton_op("mylib::mycos", mutates_args={})
def mycos(x: torch.Tensor) -> torch.Tensor:
out = torch.empty_like(x)
n_elements = x.numel()
wrap_triton(cos_kernel)[(n_elements,)](x, out, n_elements, BLOCK_SIZE=4)
return out
def backward(ctx, grad):
x, = ctx.saved_tensors
return grad * mycos(x)
def setup_context(ctx, inputs, output):
x, = inputs
ctx.save_for_backward(x)
mysin.register_autograd(backward, setup_context=setup_context)
添加 CPU 回退#
Triton 内核无法在 CPU 上运行。使用 register_kernel 为 triton_op 添加 CPU（或任何其他设备）回退。
回退必须由 PyTorch 操作符组成。
添加 FlopCounter 公式#
要指定 Triton 内核在 PyTorch 的浮点运算计数器下报告的浮点运算次数，请使用 register_flop_formula。
from torch.utils.flop_counter import FlopCounterMode, register_flop_formula
@register_flop_formula(torch.ops.mylib.mysin)
def _(x_shape):
numel = 1
for s in x_shape:
numel *= s
return numel
x = torch.randn(3, device="cuda")
FlopCounterMode 需要 tabulate。在运行下面的代码之前，请确保已安装 tabulate，或者通过运行 pip install tabulate 进行安装。
限制#
截至 PyTorch 2.3，torch.compile 中对用户定义的 Triton 内核的支持包括动态形状、torch.autograd.Function、JIT inductor 和 AOT inductor。您可以将这些功能结合使用以构建复杂、高性能的模型。
PyTorch 2.6 添加了 torch.library.triton_op，它增加了对张量子类中用户定义的 Triton 内核和其他高级功能的支持。
但是，需要注意以下限制
Triton 功能： 尽管
triton.heuristics可以单独使用或在triton.autotune之前使用，但它不能在triton.autotune之后使用。这意味着如果triton.heuristics和triton.autotune要一起使用，则必须首先使用triton.heuristics。
结论#
在此秘籍中，我们探讨了如何将用户定义的 Triton 内核与 torch.compile 结合使用。我们深入研究了简单向量加法内核的基本用法以及涉及 Triton 自动调优功能的高级用法。我们还讨论了用户定义的 Triton 内核与其他 PyTorch 功能的可组合性，并强调了一些当前的限制。
另请参阅#
脚本总运行时间：（0 分钟 3.623 秒）
