TensorRT 10.9 OSS Release. (#4381)

Signed-off-by: Leo Dong <leod@nvidia.com>

python/packaging/bindings_wheel/tensorrt/__init__.py

@@ -31,20 +31,25 @@ else:
 
 
 if not _libs_wheel_imported and sys.platform.startswith("win"):
+    log_found_dlls = bool(int(os.environ.get("TRT_LOG_FOUND_DLLS", 0)))
     # On Windows, we need to manually open the TensorRT libraries - otherwise we are unable to
     # load the bindings. If we imported the tensorrt_libs wheel, then that should have taken care of it for us.
     def find_lib(name):
         paths = os.environ["PATH"].split(os.path.pathsep)
+
         # Add ../##TENSORRT_MODULE##.libs to the search path.  This allows repackaging non-standalone TensorRT wheels as standalone
         # using delvewheel (with the --no-mangle-all flag set) to work properly.
         paths.append(os.path.join(os.path.dirname(__file__), os.pardir, "##TENSORRT_MODULE##.libs"))
+
         for path in paths:
             libpath = os.path.join(path, name)
             if os.path.isfile(libpath):
+                if log_found_dlls:
+                    print(f"Found {name} in path: {libpath}")
                 return libpath
 
-        if name.startswith("cudnn") or name.startswith("cublas"):
-            return ""
+        if name.startswith("nvinfer_builder_resource"):
+            return None
 
         raise FileNotFoundError(
             "Could not find: {:}. Is it on your PATH?\nNote: Paths searched were:\n{:}".format(name, paths)
@@ -54,11 +59,9 @@ if not _libs_wheel_imported and sys.platform.startswith("win"):
     LIBRARIES = {
         "tensorrt": [
             "nvinfer_##TENSORRT_MAJOR##.dll",
-            "cublas64_##CUDA_MAJOR##.dll",
-            "cublasLt64_##CUDA_MAJOR##.dll",
-            "cudnn64_##CUDNN_MAJOR##.dll",
             "nvinfer_plugin_##TENSORRT_MAJOR##.dll",
             "nvonnxparser_##TENSORRT_MAJOR##.dll",
+            "nvinfer_builder_resource_##TENSORRT_MAJOR##.dll",
         ],
         "tensorrt_dispatch": [
             "nvinfer_dispatch_##TENSORRT_MAJOR##.dll",
@@ -70,8 +73,10 @@ if not _libs_wheel_imported and sys.platform.startswith("win"):
 
     for lib in LIBRARIES:
         lib_path = find_lib(lib)
-        if lib_path != "":
-            ctypes.CDLL(lib_path)
+        if not lib_path:
+            continue
+        assert os.path.isfile(lib_path)
+        ctypes.CDLL(lib_path)
 
 del _libs_wheel_imported
 

python/packaging/bindings_wheel/tensorrt/plugin/_export.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,6 +15,7 @@
 # limitations under the License.
 #
 
+import tensorrt as trt
 from types import ModuleType
 import importlib
 
@@ -34,3 +35,5 @@ def public_api(module: ModuleType = None, symbol: str = None):
         return obj
 
     return export_impl
+
+IS_AOT_ENABLED = hasattr(trt, "QuickPluginCreationRequest")

python/packaging/bindings_wheel/tensorrt/plugin/_lib.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,13 +20,16 @@ import types
 import typing
 from typing import Callable, Tuple, List
 import numpy as np
-
-from ._plugin_class import _TemplatePlugin
+from ._plugin_class import _TemplateJITPlugin
+from ._export import IS_AOT_ENABLED
+if IS_AOT_ENABLED:
+    from ._plugin_class import _TemplateAOTPlugin
 from ._validate import (
     _parse_register_inputs,
     _parse_register_return,
     _validate_autotune,
     _validate_impl,
+    _validate_aot_impl,
     _validate_name_and_namespace,
 )
 from ._utils import (
@@ -91,11 +94,13 @@ class PluginDef:
         self.plugin_id = None  # includes namespace (format is ns::name)
         self.register_func = None
         self.impl_func = None
+        self.aot_impl_func = None
         self.autotune_func = None
         self.autotune_attr_names = None
         self.input_tensor_names = None
         self.input_attrs = None  # map name -> type
         self.impl_attr_names = None
+        self.aot_impl_attr_names = None
         self.num_outputs = None
         self.input_arg_schema = None
         self.expects_tactic = None
@@ -195,24 +200,26 @@ class PluginDef:
                     )
                 )
 
-        plg = plg_creator.create_plugin(
-            name,
-            namespace,
-            trt.PluginFieldCollection(fields),
-            trt.TensorRTPhase.BUILD,
-        )
-        plg.init(
-            self.register_func,
-            attrs,
-            self.impl_attr_names,
-            self.impl_func,
-            self.autotune_attr_names,
-            self.autotune_func,
-            self.expects_tactic,
-        )
+        def create_plugin_instance(quick_plugin_creation_request: "trt.QuickPluginCreationRequest" = None):
+            if quick_plugin_creation_request is None:
+                plg = plg_creator.create_plugin(
+                    name,
+                    namespace,
+                    trt.PluginFieldCollection(fields),
+                    trt.TensorRTPhase.BUILD
+                )
+            else:
+                plg = plg_creator.create_plugin(
+                    name,
+                    namespace,
+                    trt.PluginFieldCollection(fields),
+                    trt.TensorRTPhase.BUILD,
+                    quick_plugin_creation_request
+                )
 
-        return input_tensors, [], plg
+            return input_tensors, [], plg
 
+        return create_plugin_instance
 
 class _TemplatePluginCreator(trt.IPluginCreatorV3Quick):
     def __init__(self, name, namespace, attrs):
@@ -246,7 +253,7 @@ class _TemplatePluginCreator(trt.IPluginCreatorV3Quick):
 
         self.field_names = trt.PluginFieldCollection(field_names)
 
-    def create_plugin(self, name, namespace, fc, phase):
+    def create_plugin(self, name, namespace, fc, phase, qpcr: "trt.QuickPluginCreationRequest" = None):
         desc = QDP_REGISTRY[f"{namespace}::{name}"]
         name = name
         namespace = namespace
@@ -271,18 +278,83 @@ class _TemplatePluginCreator(trt.IPluginCreatorV3Quick):
                 else:
                     attrs[f.name] = attr_type_annot(f.data)
 
-        plg = _TemplatePlugin(name, namespace, desc.num_outputs)
-        plg.init(
-            desc.register_func,
-            attrs,
-            desc.impl_attr_names,
-            desc.impl_func,
-            desc.autotune_attr_names,
-            desc.autotune_func,
-            desc.expects_tactic,
-        )
-        return plg
+        jit_or_aot = None # True if JIT is to be created, False if AOT. Not None will be asserted before plugin creation.
 
+        if qpcr is None:
+            plg = _TemplateJITPlugin(name, namespace, desc.num_outputs)
+
+            plg.init(
+                desc.register_func,
+                attrs,
+                desc.impl_attr_names,
+                desc.impl_func,
+                desc.autotune_attr_names,
+                desc.autotune_func,
+                desc.expects_tactic,
+            )
+
+            return plg
+
+        # If there is a strict preference, that takes precedence
+        if qpcr == trt.QuickPluginCreationRequest.STRICT_AOT:
+            if desc.aot_impl_func is None:
+                raise ValueError(f"AOT implementation requested, but not defined for '{desc.plugin_id}'. Was @trt.plugin.aot_impl defined?")
+            jit_or_aot = False
+        elif qpcr == trt.QuickPluginCreationRequest.STRICT_JIT:
+            if desc.impl_func is None:
+                raise ValueError(f"JIT implementation requested, but not defined for '{desc.plugin_id}'. Was @trt.plugin.impl defined?")
+            jit_or_aot = True
+        else:
+            aot_defined = desc.aot_impl_func is not None
+            jit_defined = desc.impl_func is not None
+
+            # A preferemce must be indicated if both AOT and JIT implementations are defined
+            if aot_defined and jit_defined:
+                if qpcr == trt.QuickPluginCreationRequest.PREFER_AOT:
+                    jit_or_aot = False
+                elif qpcr == trt.QuickPluginCreationRequest.PREFER_JIT:
+                    jit_or_aot = True
+                else:
+                    raise ValueError(f"Plugin '{desc.plugin_id}' has both AOT and JIT implementations. NetworkDefinitionCreationFlag.PREFER_AOT_PYTHON_PLUGINS or NetworkDefinitionCreationFlag.PREFER_JIT_PYTHON_PLUGINS should be specified.")
+            else:
+                # If only one implementation is defined, use that.
+                # Any preference specified is ignored. If the preference is strong, a strict flag should have been specified.
+                if aot_defined:
+                    jit_or_aot = False
+                elif jit_defined:
+                    jit_or_aot = True
+                else:
+                    raise ValueError(f"Plugin '{desc.plugin_id}' does not have either a AOT or JIT implementation.")
+
+        assert jit_or_aot is not None
+
+        if jit_or_aot:
+            plg = _TemplateJITPlugin(name, namespace, desc.num_outputs)
+
+            plg.init(
+                desc.register_func,
+                attrs,
+                desc.impl_attr_names,
+                desc.impl_func,
+                desc.autotune_attr_names,
+                desc.autotune_func,
+                desc.expects_tactic,
+            )
+
+        else:
+            plg = _TemplateAOTPlugin(name, namespace, desc.num_outputs)
+
+            plg.init(
+                desc.register_func,
+                attrs,
+                desc.aot_impl_attr_names,
+                desc.aot_impl_func,
+                desc.autotune_attr_names,
+                desc.autotune_func
+            )
+
+        # the caller can determine if the created plugin is an AOT or JIT plugin by inspecting the interface info
+        return plg
 
 def _register_plugin_creator(name: str, namespace: str, attrs_types):
     plg_registry = trt.get_plugin_registry()
@@ -445,6 +517,102 @@ def impl(plugin_id: str) -> Callable:
 
     return decorator
 
+# Decorator for `tensorrt.plugin.aot_impl`
+@public_api()
+def aot_impl(plugin_id: str) -> Callable:
+    """
+    Wraps a function to define an Ahead-of-Time (AOT) implementation for a plugin already registered through `trt.plugin.register`.
+
+    This API is only intended to be used as a decorator. The decorated function is not required to have type hints for input arguments or return value;
+    however, any type hints specified will be validated against the `trt.plugin.register` signature for consistency.
+
+    The schema for the function is as follows:
+    .. code-block:: text
+
+        (inp0: TensorDesc, inp1: TensorDesc, ..., attr0: SupportedAttrType, attr1: SupportedAttrType, outputs: Tuple[TensorDesc], tactic: Optional[int]) -> Tuple[str, str, KernelLaunchParams, SymExprs]
+
+    * Input tensors are passed first, each described by a `TensorDesc`.
+    * Plugin attributes are declared next.
+       * Not all attributes included in `trt.plugin.register` must be specified here -- they could be a subset.
+       * NOTE: Plugin attributes are not serialized into the engine when using an AOT implementation.
+    * `tactic` is an optional argument. If the plugin is using custom tactics, it must be specified to receive the tactic value to use for the current execution of the plugin.
+
+    Args:
+        plugin_id: The ID for the plugin in the form "{namespace}::{name}", which must match that used during `trt.plugin.register`
+
+    :returns:
+        - kernel_name: The name of the kernel.
+        - compiled_kernel: Compiled form of the kernel. Presently, only PTX is supported.
+        - launch_params: The launch parameters for the kernel
+        - extra_args: Symbolic expressions for scalar inputs to the kernel, located after the tensor inputs and before the tensor outputs
+
+    .. code-block:: python
+        :linenos:
+        :caption: Implementation of an elementwise plugin with an OpenAI Triton kernel
+
+        import tensorrt.plugin as trtp
+        import triton
+        import triton.language as tl
+
+        @triton.jit
+        def add_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
+            x = tl.load(x_ptr + offsets, mask=mask)
+            tl.store(y_ptr + offsets, x + 1, mask=mask)
+
+        @trtp.register("my::add_plugin")
+        def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> Tuple[trtp.TensorDesc]:
+            return inp0.like()
+
+        @trtp.aot_impl("my::elemwise_add_plugin")
+        def add_plugin_aot_impl(
+            inp0: trtp.TensorDesc, block_size: int, single_tactic: bool, outputs: Tuple[trtp.TensorDesc], tactic: int
+        ) -> Tuple[Union[str, bytes], Union[str, bytes], trtp.KernelLaunchParams, trtp.SymExprs]:
+
+            type_str = "fp32" if inp0.dtype == trt.float32 else "fp16"
+
+            src = triton.compiler.ASTSource(
+                fn=add_kernel,
+                signature=f"*{type_str},i32,*{type_str}",
+                constants={
+                    "BLOCK_SIZE": block_size,
+                },
+            )
+
+            compiled_kernel = triton.compile(src)
+
+            N = inp0.shape_expr.numel()
+            launch_params = trtp.KernelLaunchParams()
+
+            # grid dims
+            launch_params.grid_x = trtp.cdiv(N, block_size)
+            # block dims
+            launch_params.block_x = compiled_kernel.metadata.num_warps * 32
+            # shared memory
+            launch_params.shared_mem = compiled_kernel.metadata.shared
+
+            extra_args = trtp.SymIntExprs(1)
+            extra_args[0] = trtp.SymInt32(N)
+
+            return compiled_kernel.metadata.name, compiled_kernel.asm["ptx"], launch_params, extra_args
+    """
+    def decorator(aot_impl_func: Callable):
+        if plugin_id not in QDP_REGISTRY:
+            raise ValueError(
+                f"Plugin {plugin_id} is not registered. Did you register it with tensorrt.plugin.register API?"
+            )
+
+        plugin_def = QDP_REGISTRY[plugin_id]
+        aot_impl_attr_names = _validate_aot_impl(aot_impl_func, plugin_def)
+
+        plugin_def.aot_impl_func = aot_impl_func
+        plugin_def.aot_impl_attr_names = aot_impl_attr_names
+        return aot_impl_func
+
+    return decorator
+
 
 # Decorator for `tensorrt.plugin.autotune`
 @public_api()

python/packaging/bindings_wheel/tensorrt/plugin/_plugin_class.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,25 +15,27 @@
 # limitations under the License.
 #
 import tensorrt as trt
-from typing import Tuple
+from typing import Tuple, Union
 
 import numpy as np
 from ._utils import _numpy_to_plugin_field_type, _built_in_to_plugin_field_type
-from ._tensor import TensorDesc, Tensor, Shape, ShapeExpr, ShapeExprs
+from ._tensor import TensorDesc, Tensor, Shape, ShapeExpr, ShapeExprs, SymIntExpr, SymExprs, SymInt32
+from ._export import IS_AOT_ENABLED
+if IS_AOT_ENABLED:
+    from ._tensor import KernelLaunchParams
 from ._autotune import _TypeFormatCombination
 
+from ._export import public_api
 
-class _TemplatePlugin(
+class _TemplatePluginBase(
     trt.IPluginV3,
     trt.IPluginV3QuickCore,
     trt.IPluginV3QuickBuild,
-    trt.IPluginV3QuickRuntime,
 ):
     def __init__(self, name, namespace, num_outputs):
         trt.IPluginV3.__init__(self)
         trt.IPluginV3QuickCore.__init__(self)
         trt.IPluginV3QuickBuild.__init__(self)
-        trt.IPluginV3QuickRuntime.__init__(self)
 
         self.plugin_version = "1"
         self.input_types = []
@@ -46,28 +48,6 @@ class _TemplatePlugin(
         self.autotune_combs = []
         self.supported_combs = {}
         self.curr_comb = None
-        self.expects_tactic = False
-
-    def init(
-        self,
-        register_function,
-        attrs,
-        impl_attr_names,
-        impl_function,
-        autotune_attr_names,
-        autotune_function,
-        expects_tactic,
-    ):
-        self.register_function = register_function
-        self.impl_function = impl_function
-        self.attrs = attrs
-        self.impl_attr_names = impl_attr_names
-        self.autotune_attr_names = autotune_attr_names
-        self.autotune_function = autotune_function
-        self.expects_tactic = expects_tactic
-
-    def get_capability_interface(self, type):
-        return self
 
     def get_num_outputs(self):
         return self.num_outputs
@@ -140,7 +120,7 @@ class _TemplatePlugin(
 
     def get_output_shapes(self, inputs, shape_inputs, exprBuilder):
         assert len(shape_inputs) == 0  # Shape inputs are not yet supported for QDPs
-        ShapeExpr._exprBuilder = exprBuilder
+        SymIntExpr._exprBuilder = exprBuilder
         self.input_descs = []
         for i in range(len(inputs)):
             desc = TensorDesc()
@@ -247,6 +227,45 @@ class _TemplatePlugin(
 
         return ret_supported_combs
 
+    def get_aliased_input(self, output_index: int):
+        return self.aliased_map[output_index]
+
+    def get_valid_tactics(self):
+        tactics = self.supported_combs.get(self.curr_comb)
+        assert tactics is not None
+        return list(tactics)
+
+    def set_tactic(self, tactic):
+        self._tactic = tactic
+
+class _TemplateJITPlugin(_TemplatePluginBase, trt.IPluginV3QuickRuntime):
+    def __init__(self, name, namespace, num_outputs):
+        super().__init__(name, namespace, num_outputs)
+        trt.IPluginV3QuickRuntime.__init__(self)
+
+        self.expects_tactic = False
+
+    def init(
+        self,
+        register_function,
+        attrs,
+        impl_attr_names,
+        impl_function,
+        autotune_attr_names,
+        autotune_function,
+        expects_tactic,
+    ):
+        self.register_function = register_function
+        self.impl_function = impl_function
+        self.attrs = attrs
+        self.impl_attr_names = impl_attr_names
+        self.autotune_attr_names = autotune_attr_names
+        self.autotune_function = autotune_function
+        self.expects_tactic = expects_tactic
+
+    def get_capability_interface(self, type):
+        return self
+
     def enqueue(
         self,
         input_desc,
@@ -305,20 +324,136 @@ class _TemplatePlugin(
         else:
             self.impl_function(*input_tensors, *val, output_tensors, stream=stream)
 
-    def get_aliased_input(self, output_index: int):
-        return self.aliased_map[output_index]
-
-    def get_valid_tactics(self):
-        tactics = self.supported_combs.get(self.curr_comb)
-        assert tactics is not None
-        return list(tactics)
-
-    def set_tactic(self, tactic):
-        self._tactic = tactic
-
     def clone(self):
-        cloned_plugin = _TemplatePlugin(
+        cloned_plugin = _TemplateJITPlugin(
             self.plugin_name, self.plugin_namespace, self.num_outputs
         )
         cloned_plugin.__dict__.update(self.__dict__)
         return cloned_plugin
+
+if IS_AOT_ENABLED:
+    class _TemplateAOTPlugin(
+        _TemplatePluginBase,
+        trt.IPluginV3QuickAOTBuild,
+    ):
+        def __init__(self, name, namespace, num_outputs):
+            _TemplatePluginBase.__init__(self, name, namespace, num_outputs)
+            trt.IPluginV3QuickAOTBuild.__init__(self)
+            self.kernel_map = {}
+
+        def set_tactic(self, tactic):
+            self._tactic = tactic
+
+        def init(
+            self,
+            register_function,
+            attrs,
+            aot_impl_attr_names,
+            aot_impl_function,
+            autotune_attr_names,
+            autotune_function
+        ):
+            self.register_function = register_function
+            self.aot_impl_function = aot_impl_function
+            self.attrs = attrs
+            self.aot_impl_attr_names = aot_impl_attr_names
+            self.autotune_attr_names = autotune_attr_names
+            self.autotune_function = autotune_function
+
+        def get_capability_interface(self, type):
+            return self
+
+        def get_kernel(self, inputDesc, outputDesc):
+            io_types = []
+            io_formats = []
+
+            for i, desc in enumerate(inputDesc):
+                io_types.append(desc.type)
+                io_formats.append(desc.format)
+
+            for i, desc in enumerate(outputDesc):
+                io_types.append(desc.type)
+                io_formats.append(desc.format)
+
+            key = (tuple(io_types), tuple(io_formats), self._tactic)
+
+            assert key in self.kernel_map, "key {} not in kernel_map".format(key)
+
+            kernel_name, ptx = self.kernel_map[key]
+
+            return kernel_name, ptx.encode() if isinstance(ptx, str) else ptx
+
+        def get_launch_params(self, inDimsExprs, in_out, num_inputs, launchParams, symExprSetter, exprBuilder):
+
+            SymIntExpr._exprBuilder = exprBuilder
+
+            if len(self.attrs) > 0:
+                _, val = zip(*self.attrs.items())
+            else:
+                val = ()
+
+            io_types = []
+            io_formats = []
+
+            for i, desc in enumerate(in_out):
+                if i < num_inputs:
+                    self.input_descs[i]._immutable = False
+                    self.input_descs[i].shape = Shape(desc)
+                    self.input_descs[i].dtype = desc.desc.type
+                    self.input_descs[i].format = desc.desc.format
+                    self.input_descs[i].scale = desc.desc.scale
+                    io_types.append(desc.desc.type)
+                    io_formats.append(desc.desc.format)
+                    self.input_descs[i]._immutable = True
+                else:
+                    self.output_descs[i - num_inputs]._immutable = False
+                    self.output_descs[i - num_inputs].shape = Shape(desc)
+                    self.output_descs[i - num_inputs].dtype = desc.desc.type
+                    self.output_descs[i - num_inputs].format = desc.desc.format
+                    self.output_descs[i - num_inputs].scale = desc.desc.scale
+                    io_types.append(desc.desc.type)
+                    io_formats.append(desc.desc.format)
+                    self.output_descs[i - num_inputs]._immutable = True
+
+            kernel_name, ptx, launch_params, extra_args = self.aot_impl_function(
+                *self.input_descs, *val, self.output_descs, self._tactic
+            )
+
+            if not isinstance(kernel_name, str) and not isinstance(kernel_name, bytes):
+                raise TypeError(f"Kernel name must be a 'str' or 'bytes'.  Got: {type(kernel_name)}.")
+
+            if not isinstance(ptx, str) and not isinstance(ptx, bytes):
+                raise TypeError(f"PTX/CUBIN must be a 'str' or 'bytes'.  Got: {type(ptx)}.")
+
+            if not isinstance(launch_params, KernelLaunchParams):
+                raise TypeError(f"Launch params must be a 'tensorrt.plugin.KernelLaunchParams'.  Got: {type(launch_params)}.")
+
+            if not isinstance(extra_args, SymExprs):
+                raise TypeError(f"Extra args must be a 'tensorrt.plugin.SymIntExprs'.  Got: {type(extra_args)}.")
+
+            launchParams.grid_x = launch_params.grid_x()
+            launchParams.grid_y = launch_params.grid_y()
+            launchParams.grid_z = launch_params.grid_z()
+            launchParams.block_x = launch_params.block_x()
+            launchParams.block_y = launch_params.block_y()
+            launchParams.block_z = launch_params.block_z()
+            launchParams.shared_mem = launch_params.shared_mem()
+
+            self.kernel_map[(tuple(io_types), tuple(io_formats), self._tactic)] = (kernel_name, ptx)
+
+            symExprSetter.nbSymExprs = len(extra_args)
+
+            for i, arg in enumerate(extra_args):
+                if not isinstance(arg, SymInt32):
+                    raise TypeError(f"Extra args must be a 'tensorrt.plugin.SymInt32'.  Got: {type(arg)}.")
+                symExprSetter[i] = arg()
+
+        def get_timing_cache_id(self):
+            return ""
+
+        def clone(self):
+            cloned_plugin = _TemplateAOTPlugin(
+                self.plugin_name, self.plugin_namespace, self.num_outputs
+            )
+            cloned_plugin.__dict__.update(self.__dict__)
+            return cloned_plugin

python/packaging/bindings_wheel/tensorrt/plugin/_tensor.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,67 +18,223 @@
 import tensorrt as trt
 from typing import Tuple, Union
 import numpy as np
-from ._export import public_api
+from ._export import public_api, IS_AOT_ENABLED
+from abc import ABC, abstractmethod
 
-# Symbolic expression for a given dimension of a tensor
-@public_api()
-class ShapeExpr:
+class SymExpr(ABC):
+    
+    @abstractmethod
+    def _op(self, op, other):
+        pass
+
+    @abstractmethod
+    def __add__(self, other):
+        pass
+
+    @abstractmethod
+    def __sub__(self, other):
+        pass
+
+    @abstractmethod
+    def __mul__(self, other):
+        pass
+
+    @abstractmethod
+    def __floordiv__(self, other):
+        pass
+
+    @abstractmethod
+    def __eq__(self, other):
+        pass
+
+    @abstractmethod
+    def __lt__(self, other):
+        pass
+
+    @abstractmethod
+    def __repr__(self):
+        pass
+
+    @property
+    @abstractmethod
+    def is_constant(self) -> bool:
+        pass
+
+    @property
+    @abstractmethod
+    def constant_value(self) -> int:
+        pass
+    
+    # Evaluate the underlying trt.IDimensionExpr, if so done lazily
+    @property
+    @abstractmethod
+    def _expr(self):
+        pass
+
+class SymIntExprMeta(type(SymExpr)):
+    pass
+
+class SymIntExpr(SymExpr, metaclass=SymIntExprMeta):
     """
-    Symbolic expression for single dimension of a tensor
+    Symbolic integer (scalar) expression
     """
     _exprBuilder = None  # trt.IExprBuilder instance. Populated when a shape-calculation context is entered.
 
-    def __init__(self, value: Union[int, trt.IDimensionExpr, "ShapeExpr"] = None):
+    def __init__(self, value: Union[int, trt.IDimensionExpr, "SymIntExpr"] = None):
         """
         Args:
-            value (Union[int, trt.IDimensionExpr, ShapeExpr], optional): Constant or another symbolic expression. Defaults to creating a fake shape expression.
+            value (Union[int, trt.IDimensionExpr, SymIntExpr], optional): Constant or another symbolic expression. Defaults to creating a fake shape expression.
         """
+        self._int_expr = None
+        if isinstance(value, int):
+            if SymIntExpr._exprBuilder is None:
+                self._int_expr = None
+            else:
+                self._int_expr = SymIntExpr._exprBuilder.constant(value)
+        elif isinstance(value, trt.IDimensionExpr):
+            self._int_expr = value
+        elif isinstance(value, SymIntExpr):
+            self._int_expr = value._int_expr
+
+    def _op(self, op: trt.DimensionOperation, other: Union[int, "SymIntExpr"]):
+        if isinstance(other, int):
+            other = SymIntExpr(other)
+        return SymIntExpr(SymIntExpr._exprBuilder.operation(op, self._expr, other._expr))
+
+    # Binary operations for +, -, *, //, ==. <
+    # Those for ceil_div, max and min are provided as top-level functions of tensorrt.plugin
+    def __add__(self, other: Union[int, "SymIntExpr"]):
+        return self._op(trt.DimensionOperation.SUM, other)
+
+    def __sub__(self, other: Union[int, "SymIntExpr"]):
+        return self._op(trt.DimensionOperation.SUB, other)
+
+    def __mul__(self, other: Union[int, "SymIntExpr"]):
+        return self._op(trt.DimensionOperation.PROD, other)
+
+    def __floordiv__(self, other: Union[int, "SymIntExpr"]):
+        return self._op(trt.DimensionOperation.FLOOR_DIV, other)
+
+    def __eq__(self, other: Union[int, "SymIntExpr"]):
+        return self._op(trt.DimensionOperation.EQUAL, other)
+
+    def __lt__(self, other: Union[int, "SymIntExpr"]):
+        return self._op(trt.DimensionOperation.LESS, other)
+
+    def __repr__(self):
+        if self._is_dummy:
+            return f"FakeSymIntExpr[id={id(self)}]"
+        elif not self.is_constant:
+            return f"SymIntExpr[id={id(self)}]"
+        return f"SymIntExpr[{self._expr.get_constant_value()}]"
+
+    @property
+    def is_constant(self) -> bool:
+        """
+        `True` if this integer expression is a build-time constant, `False` otherwise.
+
+        Raises:
+            RuntimeError: For fake :class:`SymIntExpr`\s. Check :attr:`is_fake` to determine accessibility.
+        """
+        if self._is_dummy:
+            raise RuntimeError(
+                "Not accessible for fake 'SymIntExpr's. Check is_fake to determine accessibility."
+            )
+        return self._expr.is_constant()
+
+    def constant_value(self) -> int:
+        """
+        Return value of the constant integer expression.
+
+        Raises:
+            RuntimeError: For non-constant integer expressions. Check :attr:`is_constant` to determine accessibility.
+        """
+        if not self.is_constant:
+            raise RuntimeError(
+                "Not accessible for non-constant integer expressions. Check is_constant to determine accessibility."
+            )
+        return self._expr.get_constant_value()
+    
+    # Evaluate the underlying trt.IDimensionExpr, if so done lazily
+    @property
+    def _expr(self):
+        return self._int_expr
+
+    def _clone(self):
+        return SymIntExpr(self + 0)
+
+class SymExprImpl(trt.ISymExpr):
+    def __init__(self, expr: SymIntExpr):
+        trt.ISymExpr.__init__(self)
+        self.type = trt.PluginArgType.INT
+        if isinstance(expr, SymInt32):
+            self.dtype = trt.PluginArgDataType.INT32
+        elif isinstance(expr, SymInt16):
+            self.dtype = trt.PluginArgDataType.INT16
+        elif isinstance(expr, SymInt8):
+            self.dtype = trt.PluginArgDataType.INT8
+        else:
+            raise ValueError(f"Unknown SymIntExpr type {type(expr)}")
+        
+        self.expr = expr._expr
+@public_api()
+class SymInt32(SymIntExpr):
+    """
+    Symbolic expression for a 32-bit integer
+    """
+    def __init__(self, value: Union[int, trt.IDimensionExpr, SymIntExpr] = None):
+        super().__init__(value)
+
+    def __call__(self):
+        return SymExprImpl(self)
+@public_api()
+class SymInt8(SymIntExpr):
+    """
+    Symbolic expression for an 8-bit integer
+    """
+    def __init__(self, value: Union[int, trt.IDimensionExpr, "SymIntExpr"] = None):
+        super().__init__(value)
+
+@public_api()
+class SymInt16(SymIntExpr):
+    """
+    Symbolic expression for a 16-bit integer
+    """
+    def __init__(self, value: Union[int, trt.IDimensionExpr, "SymIntExpr"] = None):
+        super().__init__(value)
+
+# Symbolic expression for a given dimension of a tensor
+@public_api()
+class ShapeExpr(SymInt32):
+    """
+    Symbolic expression for single dimension of a tensor
+    """
+    def __init__(self, value: Union[int, trt.IDimensionExpr, "ShapeExpr", SymIntExpr] = None):
+        """
+        Args:
+            value (Union[int, trt.IDimensionExpr, ShapeExpr, SymIntExpr], optional): Constant or another symbolic expression. Defaults to creating a fake shape expression.
+        """
+        super().__init__(value)
+        self._exprBuilder = SymIntExpr._exprBuilder
         self._is_dummy = False
-        self._dim_expr = None
         self._is_size_tensor = False
         if value is None:
             self._is_dummy = True
         elif isinstance(value, int):
             if self._exprBuilder is None:
-                self._dim_expr = None
                 self._is_dummy = True
-            else:
-                self._dim_expr = ShapeExpr._exprBuilder.constant(value)
-        elif isinstance(value, trt.IDimensionExpr):
-            self._dim_expr = value
         elif isinstance(value, ShapeExpr):
-            self._dim_expr = value._dim_expr
             self._is_dummy = value._is_dummy
             self._is_size_tensor = value._is_size_tensor
+        elif isinstance(value, SymIntExpr):
+            pass
 
     def _op(self, op: trt.DimensionOperation, other: Union[int, "ShapeExpr"]):
         if self._is_size_tensor:
             raise ValueError("It is not permitted to perform binary operations on size tensor expressions") # trt limitation
         if self._is_dummy:
             return ShapeExpr()
-        if isinstance(other, int):
-            other = ShapeExpr(other)
-        return ShapeExpr(ShapeExpr._exprBuilder.operation(op, self._expr, other._expr))
-
-    # Binary operations for +, -, *, //, ==. <
-    # Those for ceil_div, max and min are provided as top-level functions of tensorrt.plugin
-    def __add__(self, other: Union[int, "ShapeExpr"]):
-        return self._op(trt.DimensionOperation.SUM, other)
-
-    def __sub__(self, other: Union[int, "ShapeExpr"]):
-        return self._op(trt.DimensionOperation.SUB, other)
-
-    def __mul__(self, other: Union[int, "ShapeExpr"]):
-        return self._op(trt.DimensionOperation.PROD, other)
-
-    def __floordiv__(self, other: Union[int, "ShapeExpr"]):
-        return self._op(trt.DimensionOperation.FLOOR_DIV, other)
-
-    def __eq__(self, other: Union[int, "ShapeExpr"]):
-        return self._op(trt.DimensionOperation.EQUAL, other)
-
-    def __lt__(self, other: Union[int, "ShapeExpr"]):
-        return self._op(trt.DimensionOperation.LESS, other)
+        return ShapeExpr(super()._op(op, other))
 
     def __repr__(self):
         if self._is_dummy:
@@ -116,7 +272,7 @@ class ShapeExpr:
             raise RuntimeError(
                 "Not accessible for fake 'ShapeExpr's. Check is_fake to determine accessibility."
             )
-        return self._expr.is_constant()
+        return super().is_constant
 
     def constant_value(self) -> int:
         """
@@ -129,12 +285,7 @@ class ShapeExpr:
             raise RuntimeError(
                 "Not accessible for non-constant shape expressions. Check is_constant to determine accessibility."
             )
-        return self._expr.get_constant_value()
-
-    # Evaluate the underlying trt.IDimensionExpr, if so done lazily
-    @property
-    def _expr(self):
-        return self._dim_expr
+        return super().constant_value()
 
     def _clone(self):
         ret = ShapeExpr(self + 0)
@@ -172,73 +323,163 @@ class SizeTensorShapeExpr(ShapeExpr):
 
     @property
     def _expr(self):
-        if self._dim_expr is not None:
-            return self._dim_expr
-
-        self._dim_expr = super()._exprBuilder.declare_size_tensor(self._size_tensor_desc.index, self._size_tensor_desc.opt._expr, self._size_tensor_desc.upper_bound._expr)
-        return self._dim_expr
+        if self._int_expr is not None:
+            return self._int_expr
 
+        self._int_expr = super()._exprBuilder.declare_size_tensor(self._size_tensor_desc.index, self._size_tensor_desc.opt._expr, self._size_tensor_desc.upper_bound._expr)
+        return self._int_expr
+    
     def __repr__(self):
         return f"ShapeExpr[is_size_tensor = True, id={id(self)}]"
 
+def _from_scalar(s):
+    if isinstance(s, int):
+        return SymInt32(s)
+    elif isinstance(s, float):
+        raise ValueError("Float symbolic expressions are not supported")
+    else:
+        raise ValueError(f"Unsupported type: '{type(s)}'")
+
 # Iterable holding `ShapeExpr`s
 @public_api()
-class ShapeExprs:
-    def __init__(self, length: int, _is_dummy: bool = False):
+class SymExprs:
+    def __init__(self, length: int):
         """
-        Iterable holding :class:`ShapeExpr`\s
+        Iterable holding symbolic expressions
 
         Args:
             length (int): Number of dimensions of the tensor
         """
         self._length = length
-        self._is_dummy = _is_dummy
-        if _is_dummy:
-            self._shapes = [ShapeExpr()] * length
-        else:
-            self._shapes = [None] * length
+        self._exprs = [None] * length
 
     @classmethod
-    def from_tuple(cls, shape_exprs: Tuple[Union[ShapeExpr, int]]) -> "ShapeExprs":
+    def from_tuple(cls, shape_exprs: Tuple[Union[SymExpr, int]]) -> "SymExprs":
+        """
+        Args:
+            shape_exprs (Tuple[Union[SymExpr, int]]): Tuple to construct :class:`SymExprs` from
+        """
+
+        shape_exprs_ = tuple([e if isinstance(e, SymExpr) else _from_scalar(e) for e in shape_exprs])
+        inst = cls(len(shape_exprs_))
+        inst._exprs = list(shape_exprs_)
+        return inst
+
+    def __iter__(self):
+        return iter(self._exprs)
+
+    def __getitem__(self, index):
+        return self._exprs[index]
+
+    def __len__(self):
+        return self._length
+
+    def __setitem__(self, index, expr):
+        if index >= self._length:
+            raise IndexError("Index out of range")
+        
+        if not isinstance(expr, SymExpr):
+            expr = _from_scalar(expr)
+
+        self._exprs[index] = expr
+
+    def __repr__(self):
+        return f"SymExprs[{', '.join([s.__repr__() for s in self._exprs])}]"
+
+@public_api()
+class ShapeExprs(SymExprs):
+    def __init__(self, length, _is_dummy = False):
+        """
+        Iterable holding :class:`ShapeExpr`\s, representing a tensor shape
+
+        Args:
+            length (int): Number of dimensions of the tensor
+        """
+        if length > trt.Dims.MAX_DIMS:
+            raise ValueError(f"ShapeExprs can only support up to trt.Dims.MAX_DIMS = {trt.Dims.MAX_DIMS} dimensions. {length} given.")
+        
+        super().__init__(length)
+
+        self._is_dummy = _is_dummy
+        if _is_dummy:
+            self._exprs = [ShapeExpr()] * length
+
+    @classmethod
+    def from_tuple(cls, shape_exprs: Tuple[Union[ShapeExpr, int]]) -> "ShapeExpr":
         """
         Args:
             shape_exprs (Tuple[Union[ShapeExpr, int]]): Tuple to construct :class:`ShapeExprs` from
         """
+
         shape_exprs_ = tuple([e if isinstance(e, ShapeExpr) else ShapeExpr(e) for e in shape_exprs])
         inst = cls(len(shape_exprs_))
-        inst._shapes = list(shape_exprs_)
+        inst._exprs = list(shape_exprs_)
         return inst
-
+    
     def numel(self) -> ShapeExpr:
         """
         Returns a symbolic expression for the number of elements
         """
         ret = ShapeExpr(1)
-        for s in self._shapes:
+        for s in self._exprs:
             ret *= s
         return ret
 
-    def __iter__(self):
-        return iter(self._shapes)
-
-    def __getitem__(self, index):
-        return self._shapes[index]
-
-    def __len__(self):
-        return self._length
-
-    def __setitem__(self, index, shape):
+    def __setitem__(self, index, value):
         if index >= self._length:
             raise IndexError("Index out of range")
-        self._shapes[index] = shape
+        
+        if not isinstance(value, ShapeExpr):
+            if not isinstance(value, int):
+                raise ValueError(f"Value should be int or ShapeExpr. Got '{type(value)}'")
+            value = ShapeExpr(value)
 
+        self._exprs[index] = value
+    
     def __repr__(self):
-        return f"ShapeExprs[{', '.join([s.__repr__() for s in self._shapes])}]"
+        return f"ShapeExprs[{', '.join([s.__repr__() for s in self._exprs])}]"
 
     def _clone(self):
-        ret = ShapeExprs.from_tuple((e._clone() for e in self._shapes))
+        ret = ShapeExprs.from_tuple((e._clone() for e in self._exprs))
         ret._is_dummy = self._is_dummy
         return ret
+    
+@public_api()
+class SymIntExprs(SymExprs):
+    def __init__(self, length):
+        """
+        Iterable holding :class:`SymIntExpr`\s
+
+        Args:
+            length (int): Number of symbolic expressions in the iterable
+        """   
+        super().__init__(length)
+
+    @classmethod
+    def from_tuple(cls, shape_exprs: Tuple[Union[SymIntExpr, int]]) -> "SymIntExpr":
+        """
+        Args:
+            shape_exprs (Tuple[Union[SymIntExpr, int]]): Tuple to construct :class:`SymIntExprs` from
+        """
+
+        shape_exprs_ = tuple([e if isinstance(e, SymIntExpr) else SymIntExpr(e) for e in shape_exprs])
+        inst = cls(len(shape_exprs_))
+        inst._exprs = list(shape_exprs_)
+        return inst
+
+    def __setitem__(self, index, value):
+        if index >= self._length:
+            raise IndexError("Index out of range")
+        
+        if not isinstance(value, SymIntExpr):
+            if not isinstance(value, int):
+                raise ValueError(f"Value should be int or SymIntExpr. Got '{type(value)}'")
+            value = SymIntExpr(value)
+
+        self._exprs[index] = value
+    
+    def __repr__(self):
+        return f"SymIntExprs[{', '.join([s.__repr__() for s in self._exprs])}]"
 
 # Numerical representation of a tensor shape
 @public_api()
@@ -247,7 +488,7 @@ class Shape:
     Numerical representation of a tensor shape
     """
     def __init__(
-        self, tensor_desc: Union[Tuple[int], trt.DynamicPluginTensorDesc, trt.PluginTensorDesc]
+        self, tensor_desc: Union[Tuple[int], trt.DynamicPluginTensorDesc, trt.PluginTensorDesc] = None
     ):
         self._is_dynamic = None  # set lazily
         if isinstance(tensor_desc, trt.DynamicPluginTensorDesc):
@@ -353,6 +594,7 @@ class Shape:
         ret.__dict__.update(self.__dict__)
         return ret
 
+
 # Descriptor for a tensor
 # A `TensorDesc` never contains nor refers to any tensor data.
 @public_api()
@@ -848,3 +1090,39 @@ class Tensor:
 
         cloned._aliased_to = self
         return cloned
+
+if IS_AOT_ENABLED:
+    @public_api()
+    class KernelLaunchParams:
+        """
+        Args:
+            grid_x (Union[int, trt.IDimensionExpr, SymInt32], optional): The grid x dimension. Defaults to 1.
+            grid_y (Union[int, trt.IDimensionExpr, SymInt32], optional): The grid y dimension. Defaults to 1.
+            grid_z (Union[int, trt.IDimensionExpr, SymInt32], optional): The grid z dimension. Defaults to 1.
+            block_x (Union[int, trt.IDimensionExpr, SymInt32], optional): The x dimension of each thread block. Defaults to 1.
+            block_y (Union[int, trt.IDimensionExpr, SymInt32], optional): The y dimension of each thread block. Defaults to 1.
+            block_z (Union[int, trt.IDimensionExpr, SymInt32], optional): The z dimension of each thread block. Defaults to 1.
+            shared_mem (Union[int, trt.IDimensionExpr, SymInt32], optional): Shared-memory per thread block in bytes. Defaults to 0.
+        """
+        def __init__(self,
+                    grid_x: Union[int, trt.IDimensionExpr, SymInt32] = 1,
+                    grid_y: Union[int, trt.IDimensionExpr, SymInt32] = 1,
+                    grid_z: Union[int, trt.IDimensionExpr, SymInt32] = 1,
+                    block_x: Union[int, trt.IDimensionExpr, SymInt32] = 1,
+                    block_y: Union[int, trt.IDimensionExpr, SymInt32] = 1,
+                    block_z: Union[int, trt.IDimensionExpr, SymInt32] = 1,
+                    shared_mem: Union[int, trt.IDimensionExpr, SymInt32] = 0):
+            self.grid_x = SymInt32(grid_x)
+            self.grid_y = SymInt32(grid_y)
+            self.grid_z = SymInt32(grid_z)
+            self.block_x = SymInt32(block_x)
+            self.block_y = SymInt32(block_y)
+            self.block_z = SymInt32(block_z)
+            self.shared_mem = SymInt32(shared_mem)
+
+
+        def __setattr__(self, name, value):
+            if name in ["grid_x", "grid_y", "grid_z", "block_x", "block_y", "block_z", "shared_mem"]:
+                self.__dict__[name] = SymInt32(value)
+            else:
+                raise AttributeError(f"KernelLaunchParams object has no attribute '{name}'")

python/packaging/bindings_wheel/tensorrt/plugin/_validate.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,9 +18,13 @@
 import inspect
 import numpy as np
 import typing
+import types
 
 from ._utils import _is_numpy_array, _join_with, _infer_numpy_type, _is_npt_ndarray
-from ._tensor import TensorDesc, Tensor
+from ._tensor import TensorDesc, Tensor, SymExprs
+from ._export import IS_AOT_ENABLED
+if IS_AOT_ENABLED:
+    from ._tensor import KernelLaunchParams
 from ._autotune import AutoTuneCombination
 
 SERIALIZABLE_BUILTIN_TYPES = (int, float, bytes, bool, str)
@@ -93,7 +97,7 @@ def _parse_register_inputs(register_func, lazy_register):
             raise ValueError(
                 f"Argument {name} is not a positional-or-keyword or keyword-only arg"
             )
-        
+
         # Type annotations are manadatory for `tensorrt.plugin.register` args
         if param.annotation == inspect.Parameter.empty:
             raise ValueError(
@@ -105,7 +109,7 @@ def _parse_register_inputs(register_func, lazy_register):
             raise ValueError(
                 f"Argument {name} has a default value. Default values are not supported yet."
             )
-        
+
 
         if issubclass(param.annotation, TensorDesc):
             if saw_first_attr:
@@ -275,6 +279,120 @@ def _validate_impl(impl_func, plugin_def):
 
     return impl_attr_names, found_tactic
 
+def _validate_aot_impl(aot_impl_func, plugin_def):
+    aot_impl_attr_names = []
+
+    sig = inspect.signature(aot_impl_func)
+    registered_attr_names = plugin_def.input_attrs.keys()
+
+    # input arg annotations are optional, but we will validate if provided
+    for name, param in sig.parameters.items():
+        if param.annotation != inspect.Parameter.empty:
+            if name == "outputs":
+                if typing.get_origin(param.annotation) is not tuple:
+                    raise ValueError(
+                        f"'outputs' should be of type Tuple[TensorDesc]. Received {param.annotation}."
+                    )
+                args = typing.get_args(param.annotation)
+                for arg in args:
+                    if not issubclass(arg, TensorDesc):
+                        raise ValueError(
+                            f"Argument for receiving output TensorDesc, '{name}' contains a {param.annotation}. '{name}' should be a Tuple[TensorDesc]."
+                        )
+            elif name == "tactic":
+                if not issubclass(param.annotation, int):
+                    raise ValueError("'tactic' input argument should be an int")
+            elif issubclass(param.annotation, TensorDesc):
+                if name not in plugin_def.input_tensor_names:
+                    raise ValueError(
+                        f"Unexpected tensor '{name}' specified in autotune function. Expected one of {plugin_def.input_tensor_names}."
+                    )
+            else:
+                if name not in plugin_def.input_attrs:
+                    raise ValueError(
+                        f"Unexpected attribute '{name}' specified in aot_impl function. Expected one of {list(registered_attr_names)}."
+                    )
+
+                if param.annotation != plugin_def.input_attrs[name]:
+                    raise ValueError(
+                        f"Attribute '{name}' has a type annotation different from the one specified at registration. Expected '{plugin_def.input_attrs[name]}'."
+                    )
+
+                aot_impl_attr_names.append(name)
+        else:
+            if name in plugin_def.input_attrs:
+                aot_impl_attr_names.append(name)
+
+    # Expected attribute schema should be constructed in the order they appeared in the register function
+    expected_attr_schema_chunks = [
+        n for n in registered_attr_names if n in aot_impl_attr_names
+    ]
+
+    expected_schema = (
+        "("
+        + _join_with(plugin_def.input_tensor_names)
+        + _join_with(expected_attr_schema_chunks, True)
+        + ", outputs, tactic)"
+    )
+
+    if f"({', '.join(sig.parameters.keys())})" != expected_schema:
+        raise ValueError(
+            f"Signature of the aot_impl function '{sig}' does not match the expected input arg schema: {expected_schema}"
+        )
+
+    ret_annotation = sig.return_annotation
+
+    if ret_annotation == inspect.Parameter.empty:
+        raise ValueError(
+            f"No return annotation found for aot_impl function. Received signature {sig}."
+        )
+
+    expected_return_schema = "tuple[str | bytes, str | bytes, tensorrt.plugin.KernelLaunchParams, tensorrt.plugin.SymIntExprs]"
+
+    # Return annotation is optional, but we will validate if one is specified
+    if ret_annotation != inspect.Parameter.empty:
+        if typing.get_origin(ret_annotation) is not tuple:
+            raise ValueError(
+                f"Return annotation is {ret_annotation}. Expected {expected_return_schema}."
+            )
+        else:
+            args = typing.get_args(ret_annotation)
+
+            if len(args) != 4:
+                raise ValueError(
+                    f"Return annotation is {ret_annotation}. Expected {expected_return_schema}."
+                )
+
+            def validate_union_str_or_bytes(index):
+                def validate_str_or_bytes(arg_):
+                    if (arg_ is not str) and (arg_ is not bytes):
+                        raise ValueError(
+                            f"Return annotation for argument at {index} is '{arg_}'. Expected 'str' or 'bytes'."
+                        )
+
+                orig = typing.get_origin(args[index])
+                # orig is `typing.Union` when annotation uses typing module (e.g, Union[str, bytes])
+                # orig is `types.UnionType` when annotation is of the new (3.10+) native syntax (e.g, str | bytes)
+                if orig is typing.Union or orig is types.UnionType:
+                    for a in typing.get_args(args[index]):
+                        validate_str_or_bytes(a)
+                else:
+                # when annoted with `str` or `bytes`
+                    validate_str_or_bytes(args[index])
+
+            # kernel name should be str or bytes encoding
+            validate_union_str_or_bytes(0)
+            # kernel PTX should be str or bytes encoding
+            validate_union_str_or_bytes(1)
+
+            if not issubclass(args[2], KernelLaunchParams):
+                raise ValueError(f"Argument at index 2 of return annotation is '{args[2]}'. Expected 'tensorrt.plugin.KernelLaunchParams'.")
+
+            if not issubclass(args[3], SymExprs):
+                raise ValueError(f"Argument at index 3 of return annotation is '{args[3]}'. Expected a descendent of tensorrt.plugin.SymExprs.")
+
+    return aot_impl_attr_names
+
 
 def _validate_autotune(autotune_func, plugin_def):
 

python/packaging/frontend_sdist/setup.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -32,7 +32,9 @@ tensorrt_submodules = [
     "{}_bindings=={}".format(distribution_package_name, tensorrt_version),
 ]
 nvidia_pip_index_url = os.environ.get("NVIDIA_PIP_INDEX_URL", "https://pypi.nvidia.com")
-disable_internal_pip = os.environ.get("NVIDIA_TENSORRT_DISABLE_INTERNAL_PIP", False)
+
+DISABLE_INTERNAL_PIP_FLAG = "NVIDIA_TENSORRT_DISABLE_INTERNAL_PIP"
+disable_internal_pip = os.environ.get(DISABLE_INTERNAL_PIP_FLAG, "1") == "1"
 
 
 def run_pip_command(args, call_func):
@@ -64,9 +66,7 @@ if sys.platform not in ("linux", "win32"):
 if sys.implementation.name != "cpython":
     raise RuntimeError("TensorRT currently only builds wheels for CPython")
 if platform.machine() not in ("x86_64", "AMD64", "aarch64"):
-    raise RuntimeError(
-        "TensorRT currently only builds wheels for x86_64 and ARM SBSA processors"
-    )
+    raise RuntimeError("TensorRT currently only builds wheels for x86_64 and ARM SBSA processors")
 if "tegra" in platform.release():
     raise RuntimeError("TensorRT does not currently build wheels for Tegra systems")
 
@@ -108,20 +108,14 @@ def parent_command_line():
         pass
     # fall back to shell
     try:
-        return subprocess.check_output(
-            ["ps", "-p", str(pid), "-o", "command", "--no-headers"]
-        ).decode()
+        return subprocess.check_output(["ps", "-p", str(pid), "-o", "command", "--no-headers"]).decode()
     except:
         return ""
 
 
 # use pip-inside-pip hack only if the nvidia index is not set in the environment
 install_requires = []
-if (
-    disable_internal_pip
-    or nvidia_pip_index_url in parent_command_line()
-    or nvidia_pip_index_url in pip_config_list()
-):
+if disable_internal_pip or nvidia_pip_index_url in parent_command_line() or nvidia_pip_index_url in pip_config_list():
     install_requires.extend(tensorrt_submodules)
     cmdclass = {}
 else:
@@ -135,21 +129,13 @@ setup(
     long_description="""
 NVIDIA TensorRT is an SDK that facilitates high-performance machine learning inference. It is designed to work in a complementary fashion with training frameworks such as TensorFlow, PyTorch, and MXNet. It focuses specifically on running an already-trained network quickly and efficiently on NVIDIA hardware.
 
-**IMPORTANT:** This is a special release of TensorRT designed to work only with TensorRT-LLM.
-Please refrain from upgrading to this version if you are not using TensorRT-LLM.
-
-To install, please execute the following:
+If the dependencies of this package cannot be correctly installed from PyPI for any reason, you can try using the NVIDIA package index instead:
 ```
-pip install tensorrt --extra-index-url {}
-```
-Or add the index URL to the (space-separated) PIP_EXTRA_INDEX_URL environment variable:
-```
-export PIP_EXTRA_INDEX_URL='{}'
+export {}=0
 pip install tensorrt
 ```
-When the extra index url does not contain `{}`, a nested `pip install` will run with the proper extra index url hard-coded.
 """.format(
-        nvidia_pip_index_url, nvidia_pip_index_url, nvidia_pip_index_url
+        DISABLE_INTERNAL_PIP_FLAG
     ),
     long_description_content_type="text/markdown",
     author="NVIDIA Corporation",

python/packaging/metapackage/setup.cfg

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2020-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

python/packaging/metapackage/setup.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -21,11 +21,24 @@ from setuptools import setup
 distribution_package_name = "##TENSORRT_MODULE##"
 plugin_import_package_name = f"{distribution_package_name}.plugin"
 
+DISABLE_INTERNAL_PIP_FLAG = "NVIDIA_TENSORRT_DISABLE_INTERNAL_PIP"
+
 setup(
     name=distribution_package_name,
     version="##TENSORRT_PYTHON_VERSION##",
     description="TensorRT Metapackage",
-    long_description="TensorRT Metapackage",
+    long_description="""
+Metapackage for NVIDIA TensorRT, which is an SDK that facilitates high-performance machine learning inference. It is designed to work in a complementary fashion with training frameworks such as TensorFlow, PyTorch, and MXNet. It focuses specifically on running an already-trained network quickly and efficiently on NVIDIA hardware.
+
+If the dependencies of this package cannot be correctly installed from PyPI for any reason, you can try using the NVIDIA package index instead:
+```
+export {}=0
+pip install tensorrt
+```
+""".format(
+        DISABLE_INTERNAL_PIP_FLAG
+    ),
+    long_description_content_type="text/markdown",
     author="NVIDIA Corporation",
     license="Proprietary",
     classifiers=[
@@ -34,9 +47,7 @@ setup(
         "Programming Language :: Python :: 3",
     ],
     packages=[plugin_import_package_name],
-    install_requires=[
-        "##TENSORRT_MODULE##_cu##CUDA_MAJOR##==##TENSORRT_PYTHON_VERSION##"
-    ],
+    install_requires=["##TENSORRT_MODULE##_cu##CUDA_MAJOR##==##TENSORRT_PYTHON_VERSION##"],
     include_package_data=True,
     zip_safe=True,
     keywords="nvidia tensorrt deeplearning inference",