注意
跳转到末尾 下载完整的示例代码。
使用用户定义的 Triton 内核配合 torch.compile#
创建于: 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.2986, 0.5554, 2.0662, -0.0785], device='cuda:0')
Y: tensor([ 0.1402, 1.0969, -0.6538, 0.1946], device='cuda:0')
is equal to
tensor([-1.1584, 1.6523, 1.4124, 0.1160], 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.9884, 1.0321, 0.0053, 0.6395], device='cuda:0')
Y: tensor([-0.8745, 0.7483, -2.6247, -0.2872], device='cuda:0')
is equal to
tensor([ 0.1139, 1.7804, -2.6194, 0.3523], device='cuda:0')
可组合性#
用户定义的 Triton 内核并不自动支持所有 PyTorch 子系统。这在以下用例中可以看到：
添加 CPU 回退
添加
FlopCounter公式与 Tensor 子类组合
要与额外的 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 |
是 |
是 |
否 |
支持 FlopCounterMode、CPU 回退、Tensor 子类等 PyTorch 子系统 |
否 |
是 |
是 |
使用 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，它增加了对 Tensor 子类和其他高级功能中的用户定义 Triton 内核的支持。
但是，有一些限制需要注意：
Triton 功能： 虽然
triton.heuristics可以独立使用，或者在triton.autotune之前使用，但不能在triton.autotune之后使用。这意味着如果triton.heuristics和triton.autotune要一起使用，则必须先使用triton.heuristics。
结论#
在此示例中，我们探讨了如何将用户定义的 Triton 内核与 torch.compile 一起使用。我们深入研究了一个简单的向量加法内核的基本用法以及涉及 Triton 自动调优功能的高级用法。我们还讨论了用户定义的 Triton 内核与其他 PyTorch 功能的可组合性，并强调了一些当前的局限性。
另请参阅#
脚本总运行时间： (0 分 3.575 秒)
