update

Author: bowang007

examples/dynamo/aot_plugin.py

@@ -53,13 +53,8 @@
 
 @triton.jit
 def add_one_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
-    # Arg order matters for the AOT path: TRT launches the embedded PTX with
-    # arguments in (input_ptrs, output_ptrs, extra_args) order — inputs first,
-    # then outputs, then anything from ``extra_args`` in ``@trtp.aot_impl``.
-    # If this kernel declared ``(x_ptr, n_elements, y_ptr, ...)`` then TRT
-    # would feed the output pointer into ``n_elements`` and ``n_elements``
-    # into ``y_ptr`` at launch, which is a wild pointer dereference (engine
-    # builds fine, ``enqueueV3`` returns -1 and the process segfaults).
+    # AOT path requires (inputs, outputs, extra_args) order — swapping any
+    # two slots feeds the wrong value into the kernel and segfaults.
     pid = tl.program_id(0)
     block_start = pid * BLOCK_SIZE
     offsets = block_start + tl.arange(0, BLOCK_SIZE)
@@ -71,13 +66,8 @@ def add_one_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
 
 @triton.jit
 def add_one_inplace_kernel(x_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
-    # Distinct kernel for the aliased-I/O variant. The plugin descriptor
-    # declares its output as ``X.aliased()`` — at runtime TRT passes a
-    # *single* pointer for the aliased I/O pair (the shared buffer), not two.
-    # If we re-used ``add_one_kernel`` here, TRT would supply: pointer,
-    # n_elements, padding... and the kernel's ``y_ptr`` slot would absorb
-    # ``n_elements`` while ``n_elements`` would read the padding zero — the
-    # mask would be all-false and the kernel would do nothing.
+    # Aliased-I/O variant: TRT routes a single pointer for the shared
+    # input/output buffer, so the kernel takes one pointer, not two.
     pid = tl.program_id(0)
     block_start = pid * BLOCK_SIZE
     offsets = block_start + tl.arange(0, BLOCK_SIZE)
@@ -191,29 +181,13 @@ def add_plugin_aot_impl(
     launch_params.block_x = compiled_kernel.metadata.num_warps * 32  # threads per block
     launch_params.shared_mem = compiled_kernel.metadata.shared  # bytes of shared mem
 
-    # ``extra_args`` are scalar arguments appended to the kernel's argument list at
-    # launch. ``n_elements`` is passed as a 32-bit symbolic integer so TRT
-    # evaluates it from the actual tensor size at runtime.
-    #
-    # Triton >= 3.x always emits two trailing ``.param .u64 .ptr`` slots in
-    # the compiled PTX for ``global_scratch`` and ``profile_scratch`` — even
-    # when their sizes (``compiled_kernel.metadata.global_scratch_size`` /
-    # ``profile_scratch_size``) are 0. Triton's own runtime allocates
-    # zero-sized scratch buffers and passes those pointers at launch; TRT's
-    # AOT plugin path doesn't know about them and would otherwise leave the
-    # two trailing slots filled with stale register state — symptom:
-    # ``Failed to enqueue status -1`` and a segfault on the first call.
-    # We pad ``extra_args`` with four ``SymInt32(0)`` (two per u64 slot) so
-    # the kernel sees null pointers for both scratch params; since their
-    # sizes are 0 the kernel never dereferences them.
-    extra_args = trtp.SymIntExprs(1 + 4)
-    extra_args[0] = trtp.SymInt32(N)
-    for _i in range(1, 5):
-        extra_args[_i] = trtp.SymInt32(0)
+    extra_args = torch_tensorrt.dynamo.conversion.plugins.make_aot_extra_args(
+        [trtp.SymInt32(N)], compiled_kernel=compiled_kernel
+    )
 
     return (
-        compiled_kernel.metadata.name,  # kernel function name in PTX
-        compiled_kernel.asm["ptx"],  # PTX source — embedded in TRT engine
+        compiled_kernel.metadata.name,
+        compiled_kernel.asm["ptx"],
         launch_params,
         extra_args,
     )
@@ -244,35 +218,15 @@ def add_plugin_aot_impl(
 # In-place variant: aliased plugin I/O
 # -----------------------------------------
 #
-# This second registration shows the same kernel exposed as an *in-place* plugin —
-# the engine mutates the input buffer directly instead of allocating a separate
-# output. Useful for KV-cache updates and any pattern where only a small slice of
-# a large state changes per call.
-#
-# Three things change vs. ``my::add_one`` above:
-#
-# 1. ``mutates_args=("X",)`` on the torch op. This is the load-bearing signal —
-#    it tells the QDP descriptor in ``_generate_plugin.py`` that input ``X`` is
-#    a candidate for aliasing, and it also tells PyTorch's autograd and
-#    functionalization machinery that the op has side effects on ``X``.
-#
-# 2. The registered fake returns ``X`` by identity (``return X``). Combined with
-#    the schema's ``mutates_args``, this is what makes the descriptor emit
-#    ``X.aliased()`` (see ``_generate_plugin._generic_plugin_desc``) instead of
-#    building a fresh output ``TensorDesc``.
+# Same kernel exposed as an in-place plugin: the engine mutates the input
+# buffer directly via QDP ``TensorDesc.aliased()`` instead of allocating a
+# separate output. Useful for KV-cache updates and similar patterns.
 #
-# 3. The descriptor itself uses ``X.aliased()``. ``aliased()`` returns a
-#    ``TensorDesc`` that shares its data buffer with ``X`` — TRT will route the
-#    same pointer to both the input and output binding at runtime.
-#
-# The eager torch impl has to mutate ``X`` itself (so the semantics match what
-# the engine will do) and return ``X.clone()``. ``torch.library`` forbids
-# returning an input by identity from a custom op, hence the clone.
-#
-# Note on the AOT kernel: we re-use ``add_one_kernel`` unchanged. Its signature
-# takes two pointers (``x_ptr``, ``y_ptr``). With aliased I/O declared, TRT
-# passes the same buffer for both — the kernel reads from ``x_ptr`` and writes
-# to ``y_ptr``, which is the same memory, so the effect is in-place.
+# Two signals together declare aliasing to the framework:
+#  * ``mutates_args=("X",)`` on the torch op
+#  * the registered fake returns ``X`` by identity
+# The eager impl must mutate ``X`` itself and return a clone — ``torch.library``
+# rejects returning an input by identity.
 
 
 @torch.library.custom_op("my::add_one_inplace", mutates_args=("X",))  # type: ignore[misc]
@@ -281,26 +235,16 @@ def add_one_inplace(X: torch.Tensor) -> torch.Tensor:
     BLOCK_SIZE = 256
     grid = lambda meta: (triton.cdiv(X.numel(), meta["BLOCK_SIZE"]),)
     add_one_inplace_kernel[grid](X, X.numel(), BLOCK_SIZE=BLOCK_SIZE)
-    # Must not return X by identity — torch.library's no-alias constraint
-    # rejects that. The TRT path doesn't observe this clone (aliasing is
-    # declared at the descriptor level), it's purely for the eager impl.
     return X.clone()
 
 
 @torch.library.register_fake("my::add_one_inplace")
 def _(X: torch.Tensor) -> torch.Tensor:
-    # Identity return is the secondary signal the descriptor uses to detect
-    # aliasing. Combined with ``mutates_args=("X",)`` above, this is what
-    # makes ``_generic_plugin_desc`` emit ``X.aliased()`` for the output.
     return X
 
 
 @trtp.register("my::add_one_inplace")
 def add_one_inplace_desc(X: trtp.TensorDesc) -> Tuple[trtp.TensorDesc]:
-    # ``aliased()`` is the QDP API that declares output-shares-storage-with-input.
-    # Engine build will fail with "PreviewFeature::kALIASED_PLUGIN_IO_10_03 not
-    # enabled" unless the build config enables that preview feature; the
-    # converter wires this on for you when it sees a non-empty aliased_map.
     return X.aliased()
 
 
@@ -313,7 +257,6 @@ def add_one_inplace_aot_impl(
     type_str = "fp32" if X.dtype == trt.float32 else "fp16"
 
     block_size = 256
-    # Single-pointer signature — see the comment on ``add_one_inplace_kernel``.
     src = triton.compiler.ASTSource(
         fn=add_one_inplace_kernel,
         signature={
@@ -332,14 +275,9 @@ def add_one_inplace_aot_impl(
     launch_params.block_x = compiled_kernel.metadata.num_warps * 32
     launch_params.shared_mem = compiled_kernel.metadata.shared
 
-    # See the matching note on the non-in-place ``add_plugin_aot_impl``:
-    # Triton 3.x emits two trailing ``.param .u64 .ptr`` slots for the
-    # global/profile scratch buffers, and TRT's AOT path needs them zeroed
-    # explicitly via ``extra_args`` so the kernel doesn't read stale state.
-    extra_args = trtp.SymIntExprs(1 + 4)
-    extra_args[0] = trtp.SymInt32(N)
-    for _i in range(1, 5):
-        extra_args[_i] = trtp.SymInt32(0)
+    extra_args = torch_tensorrt.dynamo.conversion.plugins.make_aot_extra_args(
+        [trtp.SymInt32(N)], compiled_kernel=compiled_kernel
+    )
 
     return (
         compiled_kernel.metadata.name,
@@ -376,10 +314,6 @@ def forward(self, X: torch.Tensor) -> torch.Tensor:
 
 
 class MyInplaceModel(torch.nn.Module):
-    """Drives the in-place plugin. The op mutates ``X`` in place; the returned
-    tensor carries the post-mutation value (a clone, only to satisfy
-    torch.library's no-alias rule)."""
-
     def forward(self, X: torch.Tensor) -> torch.Tensor:
         return torch.ops.my.add_one_inplace.default(X)
 
@@ -416,18 +350,8 @@ def forward(self, X: torch.Tensor) -> torch.Tensor:
     # In-place plugin demo
     # ---------------------
     #
-    # The standard "compile once, run many times" comparison pattern doesn't
-    # work for an in-place op because each call mutates the input — running
-    # eager and TRT on the same buffer double-applies the mutation. We work
-    # off a base tensor and clone for each call instead.
-    #
-    # Three things to verify, beyond "it ran":
-    # 1. The compiled module contains a TRT engine (not a PyTorch fallback —
-    #    a regression here would silently pass on value because the eager
-    #    kernel mutates the input the same way).
-    # 2. The engine's return value matches the expected post-mutation tensor.
-    # 3. The user's input buffer was mutated in place by the engine — the
-    #    actual reason to use aliased plugin I/O in the first place.
+    # In-place ops mutate their input, so eager and TRT must run on separate
+    # cloned buffers; otherwise each comparison double-applies the mutation.
 
     print("\nIn-place plugin demo:")
     inplace_model = MyInplaceModel().to("cuda").eval()
@@ -452,22 +376,18 @@ def forward(self, X: torch.Tensor) -> torch.Tensor:
         if isinstance(m, (PythonTorchTensorRTModule, TorchTensorRTModule))
     ]
     assert engine_submodules, (
-        "Expected a TRT engine submodule for the in-place plugin path, but the"
-        " compiled module is pure PyTorch — check that the un-functionalize"
-        " lowering pass restored the mutating op before partitioning. Graph:\n"
-        f"{model_trt_inplace.graph}"
+        "Expected a TRT engine submodule for the in-place plugin path; got a"
+        f" pure-PyTorch fallback. Graph:\n{model_trt_inplace.graph}"
     )
     print(f"  TRT engine submodule(s) present: {len(engine_submodules)}")
 
     with torch.no_grad():
         trt_input = base.clone()
         trt_out = model_trt_inplace(trt_input)
         assert torch.allclose(trt_out, expected_post), "TRT output mismatch"
-        assert torch.allclose(trt_input, expected_post), (
-            "Engine did not mutate the input buffer — aliased plugin I/O is"
-            " not active. Check that PreviewFeature.ALIASED_PLUGIN_IO_10_03"
-            " was enabled and that the descriptor emitted X.aliased()."
-        )
+        assert torch.allclose(
+            trt_input, expected_post
+        ), "Engine did not mutate the input buffer — aliased plugin I/O is not active."
 
     print("  Output matches expected post-mutation value.")
     print("  Input buffer was mutated in place by the TRT engine.")

py/torch_tensorrt/dynamo/conversion/plugins/__init__.py

@@ -1,3 +1,4 @@
+from torch_tensorrt.dynamo.conversion.plugins._aot_utils import make_aot_extra_args
 from torch_tensorrt.dynamo.conversion.plugins._custom_op import custom_op
 from torch_tensorrt.dynamo.conversion.plugins._generate_plugin import generate_plugin
 from torch_tensorrt.dynamo.conversion.plugins._generate_plugin_converter import (

py/torch_tensorrt/dynamo/conversion/plugins/_aot_utils.py

@@ -0,0 +1,37 @@
+"""Helpers for writing AOT QDP plugins backed by Triton kernels."""
+
+from typing import Any, Sequence
+
+
+def _has_triton_scratch_params(compiled_kernel: Any) -> bool:
+    md = getattr(compiled_kernel, "metadata", None)
+    if md is None:
+        return False
+    return hasattr(md, "global_scratch_size") and hasattr(md, "profile_scratch_size")
+
+
+def make_aot_extra_args(
+    user_args: Sequence[Any],
+    *,
+    compiled_kernel: Any = None,
+) -> Any:
+    """Build a ``trtp.SymIntExprs`` for an AOT plugin's ``extra_args`` return.
+
+    When ``compiled_kernel`` is a Triton-compiled kernel, four trailing
+    ``SymInt32(0)`` are appended to cover the two ``.param .u64 .ptr`` slots
+    (``global_scratch``, ``profile_scratch``) that Triton >= 3.x always emits
+    in PTX even when their sizes are zero. TRT's AOT plugin path does not
+    plumb those slots through, so without padding ``enqueueV3`` reads stale
+    register state for them and segfaults on the first call.
+    """
+    import tensorrt.plugin as trtp
+
+    pad = 4 if _has_triton_scratch_params(compiled_kernel) else 0
+    total = len(user_args) + pad
+    out = trtp.SymIntExprs(total)
+    for i, arg in enumerate(user_args):
+        out[i] = arg
+    zero = trtp.SymInt32(0)
+    for i in range(len(user_args), total):
+        out[i] = zero
+    return out

py/torch_tensorrt/dynamo/conversion/plugins/_generate_plugin.py

@@ -71,31 +71,34 @@ def _generate_plugin(plugin_name: str) -> None:
 
     # retrieve the corresponding torch operation using the passed in string
     torch_op = getattr(getattr(torch.ops, namespace), name)
+    default_schema = torch_op.default._schema
 
     # Positional indices of tensor inputs marked as mutated in the schema
-    # (Tensor(aN!) ... -> ...). These are candidates for in-place QDP outputs
-    # via `TensorDesc.aliased()` — the actual alias map (output_idx -> input_idx)
-    # is determined at fake-run time by checking output-to-input identity, since
-    # torch.library does not stamp aliasing onto the return alias_info.
-    _schema = torch_op._schemas[""]
+    # (Tensor(aN!) ... -> ...) — candidates for in-place QDP outputs via
+    # `TensorDesc.aliased()`. The actual alias map (output_idx -> input_idx)
+    # is decided at fake-run time by checking output-to-input identity.
     _tensor_arg_positions = [
         i
-        for i, a in enumerate(_schema.arguments)
+        for i, a in enumerate(default_schema.arguments)
         if a.type.isSubtypeOf(torch._C.TensorType.get())
     ]
     _mutated_tensor_arg_positions = {
         _tensor_arg_positions.index(i)
-        for i, a in enumerate(_schema.arguments)
+        for i, a in enumerate(default_schema.arguments)
         if a.type.isSubtypeOf(torch._C.TensorType.get())
         and a.alias_info is not None
         and a.alias_info.is_write
     }
 
+    # Cached frozenset of aliased output indices, populated by the descriptor
+    # on first build-time call so the runtime impl skips the per-call probe.
+    _aliased_indices_cache: list[Any] = [None]
+
     # helper function that generates the required signature based on the torch operation
     def generate_signature(
         torch_op: Callable[[Any], Any],
     ) -> Tuple[str, str, str, dict[str, Any], dict[str, Any]]:
-        schema = torch_op._schemas[""]
+        schema = torch_op.default._schema
 
         arg_list = []
 
@@ -202,21 +205,19 @@ def _generic_plugin_desc(*args: Any, **kwargs: Any) -> Tuple[trtp.TensorDesc, ..
 
             output = torch_op(*fake_args, *non_tensor_args, **torch_kwargs)
 
-        # Normalize to a list of fake outputs. Multi-output torch ops return
-        # a tuple; single-output ops return a bare Tensor.
         outputs_list = list(output) if isinstance(output, (tuple, list)) else [output]
 
-        # Build output_idx -> input_idx alias map. An output is aliased to a
-        # mutated input when both (a) the input is declared mutates_args in the
-        # schema and (b) the fake kernel returned that exact tensor by identity.
-        # The schema gate prevents accidental aliasing when a fake kernel
-        # incidentally returns an input (e.g. an identity op without mutation).
+        # output_idx -> input_idx alias map: an output aliases a mutated input
+        # iff the schema marks the input as mutated AND the fake returns that
+        # tensor by identity. The schema gate prevents accidental aliasing on
+        # incidental identity returns from non-mutating ops.
         alias_map: dict[int, int] = {}
         for out_idx, fake_out in enumerate(outputs_list):
             for in_idx in _mutated_tensor_arg_positions:
                 if in_idx < len(fake_args) and fake_out is fake_args[in_idx]:
                     alias_map[out_idx] = in_idx
                     break
+        _aliased_indices_cache[0] = frozenset(alias_map.keys())
 
         input_node_expr = list(
             itertools.chain.from_iterable(
@@ -239,10 +240,14 @@ def _generic_plugin_desc(*args: Any, **kwargs: Any) -> Tuple[trtp.TensorDesc, ..
 
         out_descs = []
         for out_idx, fake_out in enumerate(outputs_list):
-            # In-place output: tell TRT this output shares its buffer with the
-            # mutated input. `aliased()` carries the input's shape/dtype so we
-            # don't rebuild the descriptor from the symbolic expressions.
             if out_idx in alias_map:
+                # Aliased output shares its buffer with a mutated input;
+                # `aliased()` carries the input's shape/dtype.
+                # Limitation: TRT preview-feature ALIASED_PLUGIN_IO_10_03
+                # inserts a defensive copy that breaks aliasing when a
+                # multi-output plugin's aliased output is consumed by
+                # another TRT layer in the same engine. Single-output
+                # plugins (the common KV-cache pattern) work end-to-end.
                 out_descs.append(tensor_args[alias_map[out_idx]].aliased())
                 continue
 
@@ -312,13 +317,15 @@ def _generic_plugin_impl(
 
         dest_tensors = [torch.as_tensor(o, device="cuda") for o in outputs]
 
-        # Outputs declared `.aliased()` share storage with their input — the
-        # mutation already lands in the destination buffer, so copying onto it
-        # would either be a no-op self-copy or, worse, clobber the in-place
-        # result if the eager kernel returned a fresh tensor.
-        aliased_outputs = {
-            i for i, o in enumerate(outputs) if o.get_aliased() is not None
-        }
+        # Skip copy_ for aliased outputs: storage is shared with the input
+        # the eager op already mutated, so copying would either no-op against
+        # itself or clobber the in-place result when the eager op returns a
+        # fresh tensor (as torch.library forces it to do).
+        aliased_outputs = _aliased_indices_cache[0]
+        if aliased_outputs is None:
+            aliased_outputs = frozenset(
+                i for i, o in enumerate(outputs) if o.get_aliased() is not None
+            )
 
         stream = torch.cuda.ExternalStream(stream)
         with torch.cuda.stream(stream):