自動為自定義 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!")