feat: add triton_op for AOT QDP plugins from Triton kernels

docsrc/py_api/kernels.rst

@@ -17,19 +17,22 @@ torch_tensorrt.kernels
 Overview
 --------
 
-The ``kernels`` module registers NVRTC-compiled CUDA C++ kernels as
-TensorRT Quick Deployable Plugins. Tensor-only declarative kernels use
-Ahead-of-Time (AOT) plugin launches when available; kernels with
-``ScalarInput`` compile through TensorRT's QDP JIT path because QDP AOT
-extra arguments currently support symbolic integer expressions, not
-arbitrary runtime floats.
-
-A single function — :func:`cuda_kernel_op` — handles both the declarative
-case (drive everything from a :class:`KernelSpec` dataclass) and the
-override case (supply ``meta_fn`` / ``eager_fn`` / ``aot_fn`` / ``schema``
-keyword arguments when the declarative DSL doesn't cover your kernel).
-:func:`ptx_op` is a parallel entry point for kernels that are already
-compiled to PTX bytes.
+The ``kernels`` module registers custom kernels — CUDA C++ compiled with
+NVRTC, or Triton — as TensorRT Quick Deployable Plugins. Tensor-only
+declarative kernels use Ahead-of-Time (AOT) plugin launches when
+available; kernels with ``ScalarInput`` compile through TensorRT's QDP JIT
+path because QDP AOT extra arguments currently support symbolic integer
+expressions, not arbitrary runtime floats.
+
+Three entry points share one registration funnel:
+
+* :func:`cuda_kernel_op` handles both the declarative case (drive
+  everything from a :class:`KernelSpec` dataclass) and the override case
+  (supply ``meta_fn`` / ``eager_fn`` / ``aot_fn`` / ``schema`` keyword
+  arguments when the declarative DSL doesn't cover your kernel).
+* :func:`ptx_op` registers kernels that are already compiled to PTX bytes.
+* :func:`triton_op` registers a ``@triton.jit`` kernel, compiling it to
+  PTX and deriving the AOT launch for you.
 
 Entry points
 ------------
@@ -117,18 +120,97 @@ Pre-compiled PTX entry point
 
 .. autofunction:: ptx_op
 
+Triton entry point
+------------------
+
+.. autofunction:: triton_op
+
+:func:`triton_op` is the Triton analogue of :func:`cuda_kernel_op`. It
+compiles the kernel once with ``triton.compile``, then registers the
+PyTorch custom op, the TRT plugin descriptor, the AOT impl embedding the
+PTX, and the Torch-TensorRT converter — replacing the hand-written
+``@trtp.aot_impl`` boilerplate in the :ref:`aot_plugin` example::
+
+    import tensorrt.plugin as trtp
+    import torch
+    import triton
+    import triton.language as tl
+    import torch_tensorrt.kernels as ttk
+
+    @triton.jit
+    def add_one_kernel(x_ptr, n_elements, y_ptr, BLOCK_SIZE: tl.constexpr):
+        pid = tl.program_id(0)
+        offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_elements
+        tl.store(y_ptr + offsets, tl.load(x_ptr + offsets, mask=mask) + 1, mask=mask)
+
+    def add_one_meta(X: torch.Tensor) -> torch.Tensor:
+        return torch.empty_like(X)
+
+    ttk.triton_op(
+        "my::add_one",
+        kernel=add_one_kernel,
+        signature={"x_ptr": "*fp32", "n_elements": "i32", "y_ptr": "*fp32"},
+        constexprs={"BLOCK_SIZE": 256},
+        grid=lambda inputs, outputs: (trtp.cdiv(inputs[0].shape_expr.numel(), 256),),
+        meta_fn=add_one_meta,
+        extra_args_fn=lambda inputs, outputs: [
+            trtp.SymInt32(inputs[0].shape_expr.numel())
+        ],
+    )
+
+``signature`` lists the kernel's non-constexpr parameters in declaration
+order (pointers as ``*<dtype>``, scalars as the bare dtype) and
+``constexprs`` supplies the ``tl.constexpr`` values baked into the PTX.
+``grid`` and ``extra_args_fn`` receive ``trtp.TensorDesc`` objects, so use
+``.shape_expr`` to stay symbolic and keep one engine valid across shapes.
+
+Because the launch is built from symbolic shape expressions, ``triton_op``
+supports dynamic shapes by default.  Pass ``eager_fn`` to also give the op
+a CUDA implementation outside TensorRT, or ``aot_fn`` to replace the
+derived launch entirely.  ``num_warps`` and ``num_stages`` are forwarded to
+``triton.compile``; ``num_warps`` also sets the launch's threads-per-block.
+
+The calling convention is enforced at registration time, because violating
+it does not fail loudly — the kernel would read whatever TensorRT happened
+to place in those argument slots.  :func:`triton_op` raises
+:class:`ValueError` if ``signature`` begins or ends with a scalar,
+interleaves scalars between pointers, disagrees with ``meta_fn``'s arity,
+or declares scalars without an ``extra_args_fn`` to supply them.
+
+.. note::
+
+   A ``triton_op`` registration compiles a single PTX for the given
+   ``signature`` and ``constexprs``.  Inputs or outputs whose dtypes differ
+   from the compiled ones are detected during conversion: the op is left
+   out of the engine and runs in PyTorch, with a warning naming the
+   mismatch.  Register a second op if you need a second dtype —
+   multi-config autotuning and dtype specialization are not yet supported.
+
+   Triton emits PTX at the ISA version of its own bundled ``ptxas``, which
+   can be newer than the installed CUDA driver accepts.  ``triton_op``
+   detects the driver's maximum ISA and recompiles at that version when
+   needed, so the embedded PTX always loads.  It also strips Triton's
+   trailing zero-sized scratch parameters, which TensorRT's AOT launcher
+   does not supply; a kernel needing non-zero scratch cannot use the AOT
+   QDP path and raises at registration time.
+
 Kernel signature convention
 ---------------------------
 
-All entry points assume the ``__global__`` kernel takes its arguments in
-the fixed order::
+All entry points assume the kernel takes its arguments in the fixed
+order::
 
     (input_ptrs..., extras..., output_ptrs...)
 
-Pointers are ``void*`` cast to the appropriate element type.  Extras
-follow the order declared in :attr:`KernelSpec.extras` for the
-declarative path, or the order your ``aot_fn`` builds for the override
-path.
+This matches the order TensorRT passes tensor pointers and AOT extra
+arguments, so no PTX rewriting is needed.  In a CUDA C++ ``__global__``
+kernel, pointers are ``void*`` cast to the appropriate element type; in a
+Triton kernel they are the ``*<dtype>`` parameters declared in
+``signature``.  Extras follow the order declared in
+:attr:`KernelSpec.extras` for the declarative path, the order
+``extra_args_fn`` returns for :func:`triton_op`, or the order your
+``aot_fn`` builds for the override path.
 
 Error behavior
 --------------

docsrc/tutorials/extensibility/plugins/index.rst

@@ -15,3 +15,4 @@ in serialized engines.
    Example: Custom Kernels with NVRTC in TensorRT AOT Plugins <../../_rendered_examples/dynamo/nvrtc_aot_plugin>
    Example: Auto-derived CUDA Kernel Plugins via cuda_kernel_op <../../_rendered_examples/dynamo/cuda_kernel_op>
    Example: Pre-compiled PTX Kernels via ptx_op <../../_rendered_examples/dynamo/ptx_op>
+   Example: Triton Kernel AOT Plugins via triton_op <../../_rendered_examples/dynamo/triton_op>

docsrc/tutorials/extensibility/plugins/plugins.rst

@@ -23,6 +23,10 @@ depending on your kernel language and performance requirements:
      - Triton
      - Pre-compiled PTX embedded in engine
      - :ref:`aot_plugin`
+   * - QDP declarative (AOT)
+     - Triton
+     - Pre-compiled PTX embedded in engine
+     - :ref:`triton_op`
    * - QDP auto-generate (AOT)
      - CUDA C++ via NVRTC
      - Pre-compiled PTX embedded in engine
@@ -119,6 +123,7 @@ For complete end-to-end examples see:
 
 * :ref:`auto_generate_plugins` — Triton kernel, QDP JIT plugin
 * :ref:`aot_plugin` — Triton kernel, QDP AOT plugin (pre-compiled PTX, no Python overhead at runtime)
+* :ref:`triton_op` — the same Triton AOT plugin registered in one ``triton_op`` call
 * :ref:`nvrtc_aot_plugin` — CUDA C++ kernel compiled with NVRTC, QDP AOT plugin
 * :ref:`custom_kernel_plugins` — manual plugin + converter registration (legacy approach)
 

docsrc/user_guide/compilation/unsupported_ops.rst

@@ -271,9 +271,10 @@ FAQ
 
 **"I have a Triton kernel. Can I use it in a serialized TRT engine?"**
 
-    Triton kernels can be wrapped as TRT plugins via the AOT plugin path. See
-    :ref:`aot_plugin` for an end-to-end example of compiling a Triton kernel as a
-    TRT plugin for use in a serialized engine.
+    Triton kernels can be wrapped as TRT plugins via the AOT plugin path. The
+    shortest route is ``torch_tensorrt.kernels.triton_op``, which compiles the
+    kernel and registers the op, plugin, and converter in one call — see
+    :ref:`triton_op`. :ref:`aot_plugin` shows the same thing built by hand.
 
 **"Operator X is listed as supported but my model still falls back"**
 

examples/dynamo/triton_op.py

@@ -0,0 +1,126 @@
+"""
+.. _triton_op:
+
+Register a Triton kernel as an AOT QDP plugin via ``torch_tensorrt.kernels.triton_op``
+======================================================================================
+
+This is the productized form of :ref:`aot_plugin`. That example shows the raw
+mechanism — a ``@triton.jit`` kernel wired to a TensorRT AOT Quick Deployable
+Plugin by hand: you write ``@torch.library.custom_op`` + ``register_fake``,
+``@trtp.register``, and a ~40-line ``@trtp.aot_impl`` that builds the
+``ASTSource`` signature, calls ``triton.compile``, and assembles the launch
+parameters, then ``generate_plugin_converter``.
+
+``triton_op`` collapses all of that into a single call. You provide the Triton
+``signature`` / ``constexprs`` / ``grid`` and a ``meta_fn``; the library
+compiles the kernel to PTX, derives the AOT launch, registers the PyTorch op,
+the plugin descriptor + AOT impl, and the Torch-TensorRT converter.
+
+Calling convention: the kernel's non-constexpr parameters must be declared as
+``(input_ptrs..., extra_scalars..., output_ptrs...)`` so no PTX rewriting is
+needed.
+"""
+
+import argparse
+
+import tensorrt.plugin as trtp
+import torch
+import triton
+import triton.language as tl
+
+import torch_tensorrt
+import torch_tensorrt.kernels as ttk
+
+# %%
+# Step 1: Define the Triton kernel (pure Triton, unchanged from aot_plugin.py)
+# ----------------------------------------------------------------------------
+
+
+@triton.jit
+def add_one_kernel(x_ptr, n_elements, y_ptr, BLOCK_SIZE: tl.constexpr):
+    pid = tl.program_id(0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    x = tl.load(x_ptr + offsets, mask=mask)
+    output = x + 1
+    tl.store(y_ptr + offsets, output, mask=mask)
+
+
+# %%
+# Step 2: Describe the op and register it with a single ``triton_op`` call
+# ------------------------------------------------------------------------
+#
+# * ``meta_fn`` — shape/dtype inference for FakeTensors (the schema is inferred
+#   from its type hints).
+# * ``signature`` / ``constexprs`` — exactly what ``triton.compile`` needs.
+# * ``grid`` / ``extra_args_fn`` — computed from ``trtp.TensorDesc`` inputs,
+#   using symbolic shapes so one engine works across sizes.
+# * ``eager_fn`` — optional; lets ``torch.ops.my.add_one`` also run in eager.
+
+BLOCK_SIZE = 256
+
+
+def add_one_meta(X: torch.Tensor) -> torch.Tensor:
+    return torch.empty_like(X)
+
+
+def add_one_eager(X: torch.Tensor) -> torch.Tensor:
+    Y = torch.empty_like(X)
+    grid = lambda meta: (triton.cdiv(X.numel(), meta["BLOCK_SIZE"]),)
+    add_one_kernel[grid](X, X.numel(), Y, BLOCK_SIZE=BLOCK_SIZE)
+    return Y
+
+
+ttk.triton_op(
+    "my::add_one",
+    kernel=add_one_kernel,
+    signature={"x_ptr": "*fp32", "n_elements": "i32", "y_ptr": "*fp32"},
+    constexprs={"BLOCK_SIZE": BLOCK_SIZE},
+    grid=lambda inputs, outputs: (trtp.cdiv(inputs[0].shape_expr.numel(), BLOCK_SIZE),),
+    meta_fn=add_one_meta,
+    extra_args_fn=lambda inputs, outputs: [trtp.SymInt32(inputs[0].shape_expr.numel())],
+    eager_fn=add_one_eager,
+    supports_dynamic_shapes=True,
+)
+
+
+# %%
+# Step 3: Use it — the op lowers to the AOT QDP plugin inside the engine
+# ----------------------------------------------------------------------
+
+
+class AddOne(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.ops.my.add_one(x)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--min_block_size", type=int, default=1)
+    args = parser.parse_args()
+
+    x = torch.randn(4, 256, device="cuda", dtype=torch.float32)
+    model = AddOne().cuda().eval()
+
+    ref = x + 1
+    eager_out = model(x)
+    assert torch.allclose(eager_out, ref), "eager path mismatch"
+    print("eager path OK")
+
+    trt_model = torch_tensorrt.compile(
+        model,
+        inputs=[x],
+        min_block_size=args.min_block_size,
+    )
+    print("engine compiled with the AOT QDP plugin")
+
+    trt_out = trt_model(x)
+    if torch.allclose(trt_out, ref, atol=1e-5):
+        print("triton_op AOT QDP plugin ran correctly under Torch-TensorRT")
+    else:
+        print(
+            "WARNING: TRT output did not match. Check that the CUDA driver "
+            "supports the kernel's PTX ISA (toolkit newer than driver can "
+            "require capping); registration/compile/build succeeded."
+        )

py/torch_tensorrt/kernels/__init__.py

@@ -22,6 +22,11 @@
     Useful when the PTX comes from an external compiler (Triton, a cached
     NVRTC output, etc.).
 
+``triton_op`` — register a ``@triton.jit`` kernel. Compiles the kernel to PTX
+    for you and derives the AOT launch, so you only supply the Triton
+    ``signature`` / ``constexprs`` / ``grid`` and a ``meta_fn`` — no
+    hand-written ``@trtp.aot_impl`` compile boilerplate.
+
 Minimal example — declarative ``cuda_kernel_op``::
 
     import torch, torch_tensorrt
@@ -65,7 +70,7 @@
     SameAs,
     ScalarInput,
 )
-from torch_tensorrt.kernels._ops import cuda_kernel_op, ptx_op
+from torch_tensorrt.kernels._ops import cuda_kernel_op, ptx_op, triton_op
 
 __all__ = [
     "Custom",
@@ -81,4 +86,5 @@
     "ScalarInput",
     "cuda_kernel_op",
     "ptx_op",
+    "triton_op",
 ]