feat: TorchTRT Cuda generated kernels plugin support

[bowang007]

Description

This PR introduces torch_tensorrt.annotation, an experimental module for registering hand-written CUDA C++ kernels as both PyTorch custom ops (for eager execution) and TensorRT Quick Deployable Plugins with AOT support (for torch_tensorrt.compile).

Usage

  import torch, torch_tensorrt
  import torch_tensorrt.kernels as ttk                                                                                                                           
   
  CU = """                                                                                                                                                          
  extern "C" __global__ void my_sigmoid(const float* x, int n, float* y) {
      int i = blockIdx.x * blockDim.x + threadIdx.x;
      if (i < n) y[i] = 1.0f / (1.0f + __expf(-x[i]));                                                                                                              
  }
  """                                                                                                                                                               
                  
ttk.cuda_kernel_op(
    "kern_ex::sigmoid",
    ttk.KernelSpec(
        kernel_source=CU_SIGMOID,
        kernel_name="my_sigmoid",
        inputs=[ttk.InputDecl("x")],
        outputs=[ttk.OutputDecl("y", shape=ttk.SameAs("x"))],
        extras=[ttk.Numel("x")],
        geometry=ttk.Elementwise(block=(256,), layout="flat"),
    ),
    supports_dynamic_shapes=True,
)

After this call, torch.ops.ann_ex.sigmoid is available in eager and is embedded as a TensorRT plugin during torch_tensorrt.compile. The meta function, eager
launch, AOT implementation, and PyTorch schema are all derived from the KernelSpec.

API Surface

The module exposes two primary entry points, layered by declarativeness:

auto_cuda_kernel_plugin is the recommended default. The caller supplies a KernelSpec dataclass describing the kernel's inputs, outputs (with a shape relation such
as SameAs or ReduceDims), scalar extras (Numel, DimSize), and launch geometry (Elementwise or Reduction). The framework derives the meta function, eager CUDA
launch, TensorRT AOT implementation, and PyTorch schema. This path covers pointwise kernels (1-D flat or N-D grid launches), reductions (with optional keepdim),
multi-input kernels, and scalar (non-tensor) kernel arguments via ScalarInput.

manual_cuda_kernel_plugin is the lower-level alternative for kernels outside the declarative DSL — shape-changing outputs, multi-output kernels, or non-standard
launch geometries. The caller provides eager_fn and aot_fn directly; the decorator still registers the PyTorch op, TRT plugin, AOT implementation, and converter
in a single call.

A Custom(fn=...) geometry is also available for callers who want the declarative path's schema/meta derivation but need to hand-write the TRT KernelLaunchParams.

Type of change

• New feature (non-breaking change which adds functionality)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[narendasan]

Cool I think this is getting really close. I think we just have a few naming things to make this more user friendly and I think we should let users provide PTX directly in addition to the cuda apis. Also did you add nvrtc as an optional dependency in the pyproject.toml) (maybe under an a extras called kernels)?

[narendasan]

maybe we can call this something like torch_tensorrt.kernels.cuda_kernel_op

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/plugins/_generate_plugin_converter.py	2026-05-11 16:39:32.423025+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/plugins/_generate_plugin_converter.py	2026-05-11 16:39:55.856400+00:00
@@ -36,11 +36,13 @@


def _coerce_plugin_attr_for_qdp(value: Any, attr_annotation: Any) -> Any:
    """Convert Python scalars to the serialized type expected by QDP."""
    if _is_numpy_attr_annotation(attr_annotation):
-        return np.asarray(_unwrap_scalar_attr(value), dtype=_numpy_attr_dtype(attr_annotation))
+        return np.asarray(
+            _unwrap_scalar_attr(value), dtype=_numpy_attr_dtype(attr_annotation)
+        )
    return value


def _is_numpy_attr_annotation(annotation: Any) -> bool:
    return annotation is np.ndarray or typing.get_origin(annotation) is np.ndarray
--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/kernels/_kernel_plugin.py	2026-05-11 16:39:32.436980+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/kernels/_kernel_plugin.py	2026-05-11 16:39:57.457501+00:00
@@ -371,16 +371,14 @@
            annotations[d.name] = d.py_type
        else:
            sig_pieces.append(f"{d.name}: 'torch.Tensor'")
            annotations[d.name] = torch.Tensor
    sig_src = ", ".join(sig_pieces)
-    body = textwrap.dedent(
-        f"""
+    body = textwrap.dedent(f"""
        def _wrapper({sig_src}) -> 'torch.Tensor':
            return _fn({", ".join(param_names)})
-        """
-    )
+        """)
    ns: Dict[str, Any] = {"_fn": fn, "torch": torch}
    exec(compile(body, "<cuda_kernel_op>", "exec"), ns)
    wrapper: Callable[..., Any] = ns["_wrapper"]
    wrapper.__annotations__ = dict(annotations)
    wrapper.__annotations__["return"] = torch.Tensor
@@ -713,8 +711,6 @@
        precompiled_ptx=ptx,
        use_aot_if_available=not any(
            isinstance(input_spec, ScalarInput) for input_spec in spec.inputs
        ),
    )
-    _LOGGER.info(
-        "cuda_kernel_op '%s' registered (schema: %s)", op_name, schema
-    )
+    _LOGGER.info("cuda_kernel_op '%s' registered (schema: %s)", op_name, schema)