备注
点击 此处 下载完整示例代码
在 torch.compile 中使用用户定义的 Triton 内核¶
Created On: Apr 19, 2024 | Last Updated: May 02, 2025 | Last Verified: Nov 05, 2024
作者： 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}")
高级用法¶
Triton 的 autotune 特性是一个强大的工具，可以自动优化 Triton 内核的配置参数。它会探索一系列可能的配置，并选择在您的特定用例中性能最佳的那个。
与 torch.compile 配合使用时，triton.autotune 可以确保您的 PyTorch 模型运行尽可能高效。以下是使用 torch.compile 和 triton.autotune 的示例。
备注
torch.compile 仅支持传递给 triton.autotune 的 config 和关键参数。
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}")
可组合性¶
用户定义的 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 内核（无显式 |
|
|
|
|---|---|---|---|
支持推理 |
是 |
是 |
是 |
支持训练 |
大多数情况下 |
是 |
是 |
支持 |
是 |
是 |
是 |
支持 |
大多数情况下 |
大多数情况下 |
所有情况都支持 |
|
是 |
是 |
否 |
支持 AOTInductor |
是 |
是 |
否 |
支持 PyTorch 子系统，如 FlopCounterMode、CPU 回退、Tensor 子类 |
否 |
是 |
是 |
使用 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 组合使用时存在某些隐藏问题）。
def backward(ctx, grad):
x, = ctx.saved_tensors
return grad * x.cos()
def setup_context(ctx, inputs, output):
x, = inputs
ctx.save_for_backward(x)
mysin.register_autograd(backward, setup_context=setup_context)
请注意，反向传播必须由 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（或其他设备）回退：
@mysin.register_kernel("cpu")
def _(x):
return torch.sin(x)
x = torch.randn(3)
y = mysin(x)
assert torch.allclose(y, x.sin())
回退必须由 PyTorch 操作符组成。
添加 FlopCounter 公式¶
要在 PyTorch 的 flop 计数器中指定 Triton 内核报告多少浮点操作数（flops），请使用 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 进行安装。
>>> with FlopCounterMode() as flop_counter:
>>> y = mysin(x)
限制¶
从 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。
结论¶
在本教程中，我们探讨了如何在 torch.compile 中使用用户定义的 Triton 内核。我们深入学习了一个简单的向量加法内核的基本使用，以及涉及 Triton 的 autotune 功能的高级用法。我们还讨论了用户定义的 Triton 内核与其他 PyTorch 功能的可组合性，并指出了一些当前的限制。
