注意
前往末尾 下载完整的示例代码
自动为自定义 Kernel 生成 Plugin¶
我们将演示如何使用 Torch-TensorRT 以及 TensorRT 10.7 中基于 Python 的新 plugin 系统,自动为自定义 Kernel 生成 plugin。
如果 Torch-TensorRT 不知道如何在 TensorRT 中编译某些操作,它支持回退到 PyTorch 的实现。然而,这会带来图中断 (graph break) 的开销,并降低模型的性能。解决操作支持不足的最简单方法是添加分解 (decomposition)(参见:为 Dynamo 前端编写 Lowering Pass),它使用 Torch-TensorRT 支持的 PyTorch 操作来定义该操作;或者添加 Converter(参见:为 Dynamo 前端编写 Converter),它使用 TensorRT Operator 来定义该操作。
在某些情况下,这两种方法都不是最佳选择,例如该 Operator 是一个不属于标准 PyTorch 的自定义 Kernel,或者 TensorRT 本身无法原生支持它。
对于这些情况,可以使用 TensorRT Plugin 在 TensorRT Engine 内部替换该 Operator,从而避免图中断带来的性能和资源开销。
以前,这涉及到一个复杂的过程,不仅需要构建高性能 Kernel,还需要设置它以便在 TensorRT 中运行(参见:在 Torch-TensorRT 中使用 TensorRT Engine 内的自定义 Kernel)。从 TensorRT 10.7 开始,引入了一种新的 Python 原生 plugin 系统,这大大简化了此过程。该 plugin 系统还允许 Torch-TensorRT 自动生成必要的转换代码,以便将 PyTorch 中的操作转换为 TensorRT。
在 PyTorch 中编写自定义 Operator¶
先前的教程已涵盖如何在 PyTorch 中创建自定义 Operator,这些 Operator 随后可与 Torch-TensorRT 一起使用。
在此,我们在 Triton 中定义一个简单的逐元素乘法 Operator。然后,该 Operator 连同其主机启动代码以及一个“meta-kernel”一起在 PyTorch 中注册为自定义 op。Meta-kernel 是一个描述该 Operator 将执行的形状和数据类型转换的函数。Dynamo 和 Torch-TensorRT 会使用此 meta-kernel,因此必须对其进行定义。
from typing import Tuple
import tensorrt_bindings.plugin as trtp
import torch
import torch_tensorrt
import triton
import triton.language as tl
@triton.jit
def elementwise_scale_mul_kernel(X, Y, Z, a, b, BLOCK_SIZE: tl.constexpr):
pid = tl.program_id(0)
# Compute the range of elements that this thread block will work on
block_start = pid * BLOCK_SIZE
# Range of indices this thread will handle
offsets = block_start + tl.arange(0, BLOCK_SIZE)
# Load elements from the X and Y tensors
x_vals = tl.load(X + offsets)
y_vals = tl.load(Y + offsets)
# Perform the element-wise multiplication
z_vals = x_vals * y_vals * a + b
# Store the result in Z
tl.store(Z + offsets, z_vals)
@torch.library.custom_op("torchtrt_ex::elementwise_scale_mul", mutates_args=()) # type: ignore[misc]
def elementwise_scale_mul(
X: torch.Tensor, Y: torch.Tensor, b: float = 0.2, a: int = 2
) -> torch.Tensor:
# Ensure the tensors are on the GPU
assert X.is_cuda and Y.is_cuda, "Tensors must be on CUDA device."
assert X.shape == Y.shape, "Tensors must have the same shape."
# Create output tensor
Z = torch.empty_like(X)
# Define block size
BLOCK_SIZE = 1024
# Grid of programs
grid = lambda meta: (X.numel() // meta["BLOCK_SIZE"],)
# Launch the kernel with parameters a and b
elementwise_scale_mul_kernel[grid](X, Y, Z, a, b, BLOCK_SIZE=BLOCK_SIZE)
return Z
逐元素操作的 meta-kernel 就是其中一个输入的形状和 dtype,因为在操作过程中我们不会改变形状。
@torch.library.register_fake("torchtrt_ex::elementwise_scale_mul")
def _(x: torch.Tensor, y: torch.Tensor, b: float = 0.2, a: int = 2) -> torch.Tensor:
return x
在此,我们使用 Torch-TensorRT 中的自动 plugin 创建功能,该功能支持使用 TensorRT QDP API 进行 plugin 注册。
torch_tensorrt.dynamo.conversion.plugins.generate_plugin(
"torchtrt_ex::elementwise_scale_mul"
)
# # %%
# # Generating the Converter
# # -------------------------------------------------------------------
# # Given that we have defined the custom operator in PyTorch and TensorRT, we can now generate the converter for the operation.
# # As long as the namespace and names match, the following function will automatically generate the converter for the operation.
# # If plugins require an output allocator to dynamically allocate output buffers, like data dependent operators, please set requires_output_allocator to True.
torch_tensorrt.dynamo.conversion.plugins.generate_plugin_converter(
"torchtrt_ex::elementwise_scale_mul",
supports_dynamic_shapes=True,
requires_output_allocator=False,
)
# # %%
# # Above two commands can be replaced with the following single one line:
# torch_tensorrt.dynamo.conversion.plugins.custom_op("torchtrt_ex::elementwise_scale_mul", supports_dynamic_shapes=True, requires_output_allocator=False)
将我们的 Converter 与模型一起使用¶
现在,我们可以在模型中使用自定义 Operator 并用 Torch-TensorRT 编译它。可以看到,自定义 Operator 被用作模型前向传播中的一个操作。此时编译模型的过程与标准 Torch-TensorRT 用法相同。
class MyModel(torch.nn.Module): # type: ignore[misc]
def __init__(self):
super().__init__()
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
z = torch.add(x, y)
res = torch.ops.torchtrt_ex.elementwise_scale_mul.default(x, z, b=0.5)
return res
my_model = MyModel().to("cuda")
m = torch.randint(0, 5, (64, 64), device="cuda", dtype=torch.float)
n = torch.randint(0, 5, (64, 64), device="cuda", dtype=torch.float)
with torch_tensorrt.logging.errors():
model_trt = torch_tensorrt.compile(
my_model, inputs=[m, n], debug=True, min_block_size=1
)
for i in range(300):
res = model_trt(m, n)
assert torch.allclose(res, my_model(m, n))
print("Ran with custom plugin!")
脚本总运行时间: ( 0 分钟 0.000 秒)