Skill: Write a TensorRT Plugin Integration Plan

Purpose

Guide the agent through planning how a custom CUDA kernel will be wrapped as a TensorRT plugin so it can be invoked from inside a TensorRT engine — covering API choice (IPluginV3 vs IPluginV2DynamicExt), the plugin lifecycle, dynamic shape handling, serialization, mixed precision (FP16/INT8/FP8), workspace management, CUDA graph compatibility, and the C++/Python binding strategy. The output is an integration plan with explicit decisions, not the plugin source code itself.

Use this when

• A custom CUDA kernel must be deployed inside a TensorRT engine (production inference) and no built-in TRT layer covers it.
• An ONNX model contains an unsupported op that must be lowered to a custom plugin during ONNX → TRT conversion.
• A TRT-LLM-style codebase needs a new operator wrapped as a plugin to participate in engine builds.
• Migrating an existing IPluginV2DynamicExt plugin to IPluginV3 for a TensorRT 9.x or 10.x upgrade.

Do not use this when

• The operation is already covered by a built-in TRT layer (most GEMMs, conv shapes, softmax, layernorm fused into MHA, common activations) — the built-in path is almost always faster than a plugin because it participates in TRT's graph optimizer and tactic selection.
• The op can be expressed as a composition of existing TRT layers (e.g. IElementWiseLayer, IMatrixMultiplyLayer, IShuffleLayer) — TRT's optimizer fuses these well; an opaque plugin blocks fusion.
• The deployment target is not TensorRT. For vLLM use a vLLM custom op; for ONNX Runtime use an ORT custom op; for raw PyTorch use torch.library.
• The kernel relies on JIT autotuning at runtime (e.g. Triton with Python-side autotune). TRT engines are built ahead of time and cannot host Python or runtime JIT inside enqueue.
• The kernel allocates host memory or synchronizes inside the launch path — incompatible with CUDA graph capture, which is how production TRT inference servers reach low latency.

Inputs the agent should gather first

• TensorRT version: 8.x, 9.x, or 10.x. The plugin API changes are major across these — IPluginV3 was introduced in 9.x and is the only forward-compatible API in 10.x. Targeting TRT 8.6 means IPluginV2DynamicExt; targeting 9.3+ means IPluginV3 unless matching an existing codebase's convention.
• Target precision(s): FP32, FP16, BF16, INT8, FP8 (E4M3/E5M2). Each precision the plugin must support adds a row to the format combination table and a code path in the kernel.
• Dynamic shape requirements: which input dimensions are dynamic, what are the min/opt/max profiles, and is multi-profile support required? Static-shape-only plugins are simpler but rarely sufficient for transformer inference.
• Engine save/load required: production engines are built once and loaded many times. If yes, full serialization (getSerializationSize, serialize, deserialize) must be designed correctly.
• ONNX import path: if the engine is built from an ONNX model, the plugin must register an ONNX op spec via trtexec --plugins or the ONNX-TRT parser's plugin namespace.
• CUDA graph capture required: most modern TRT inference (Triton Inference Server, TRT-LLM) captures into CUDA graphs. This rules out host allocations and stream synchronization inside enqueue.
• C++ vs Python build path: the engine builder is typically C++ for production, Python for prototyping. Plugin must be loadable from both.
• Existing kernel's launch surface: what does the kernel need at launch time — input/output pointers, shapes, scales, optional masks, workspace? This determines the plugin's attribute set and enqueue signature.

Required reasoning process

1. Choose the plugin API. This is the first decision and constrains everything downstream:

   API	TRT version	Recommended for	Notes
   IPluginV3	9.x+, required in 10.x	New plugins, long-term codebases	Splits responsibilities across IPluginV3OneCore, IPluginV3OneBuild, IPluginV3OneRuntime. Cleaner lifecycle, better forward compatibility.
   IPluginV2DynamicExt	8.x, 9.x (deprecated path)	Matching an existing codebase that uses it (TRT-LLM 0.x, NVIDIA samples)	Single-class API. Still functional in 9.x but flagged for removal.
   IPluginV2Ext / IPluginV2	Legacy	Static-shape-only inference, almost never appropriate	No dynamic shape support; do not use for transformer or vision-with-variable-resolution workloads.

   Default for a greenfield plugin on TRT ≥ 9.3: IPluginV3. Default when extending a codebase already on IPluginV2DynamicExt: match the convention; do not mix APIs in one project.

2. Map the plugin lifecycle to your kernel. For IPluginV3OneBuild + IPluginV3OneRuntime, every method has a specific role and serializing the wrong state into the wrong method is the most common bug class:

   • getNbOutputs, getOutputDataTypes, getOutputShapes: build-time, deterministic. Output shapes must be expressed symbolically in terms of input shapes (TRT supports a limited shape arithmetic via IExprBuilder).
   • supportsFormatCombination: build-time. Returns whether the plugin accepts a given combination of dtype + tensor format for each input/output. Get this wrong and TRT will either reject the plugin or insert reformat ops around it that destroy the performance gain.
   • configurePlugin: called once per engine build with concrete shapes from the optimization profile. Use it to compute and cache attributes derived from shapes, not runtime tensor pointers.
   • getWorkspaceSize: build-time. Return an upper bound on scratch memory needed for any shape in the profile. TRT owns and pools this memory across plugin instances.
   • enqueue: runtime. Receives cudaStream_t, void* const* inputs, void* const* outputs, void* workspace, and PluginTensorDesc. Launch the kernel here. Do not call cudaMalloc, cudaMemcpy from host, or any synchronizing API in this path.
   • clone: must produce a fresh plugin instance with the same attributes. Required because TRT may instantiate one plugin per execution context. Forgetting clone causes data corruption when running multiple contexts concurrently.

3. Design the dynamic-shape contract. TRT engines are built against optimization profiles, each with min/opt/max shapes. The plugin must:

   • Compute output shapes symbolically. If the output is [B, H, S_q, S_kv] from inputs [B, H, S_q, D] and [B, H, S_kv, D], the symbolic expression goes in getOutputShapes using IExprBuilder::operation(...).
   • Validate alignment requirements via supportsFormatCombination. If the kernel requires S_kv % 16 == 0, the plugin cannot silently accept arbitrary S_kv — declare the constraint or assert at runtime.
   • Handle the entire profile range. A plugin that works at the opt shape but segfaults at min (e.g. when S_q == 1) is the single most common production failure mode.

4. Design the serialization contract. Engines are built once and loaded N times in production:

   • Serialize attributes (kernel hyperparameters, precision flags, head dimension, scale factors set at construction time). These are static for the engine's lifetime.
   • Never serialize runtime state — workspace pointers, configured shapes from configurePlugin, CUDA stream handles. These are reconstructed on engine load.
   • The plugin creator (IPluginCreatorV3One for IPluginV3) is registered globally via REGISTER_TENSORRT_PLUGIN and looked up at deserialize time by getPluginName() + getPluginVersion() + getPluginNamespace(). A version mismatch silently fails with a confusing "plugin not found" error.
   • Bump the plugin version when changing the serialization format. Old engines must not deserialize against an incompatible new plugin.

5. Plan the format combination table. This is a table of (input_dtype, input_format, output_dtype, output_format) tuples the plugin advertises as supported. Decisions:

   • For an FP16 kernel: advertise (kHALF, kLINEAR, kHALF, kLINEAR) only. Do not advertise FP32 if the kernel does not support it — TRT will pick FP32 at the boundary and silently insert casts.
   • For INT8: advertise (kINT8, kLINEAR, kINT8, kLINEAR) and provide per-tensor scales via the calibration path. Mixed-precision (INT8 input, FP16 output for an epilogue) is a common pattern but doubles the format combinations.
   • For FP8 (E4M3 input, E5M2 output is unusual; usually both are E4M3): only on TRT 9.x+ targeting Hopper (sm_90a) or Blackwell. Advertise kFP8 with kLINEAR and verify the minimum TRT version supports it.
   • Reformat-free input: declaring fewer format combinations forces TRT to insert reformat ops; declaring too many means the plugin must implement them all. Declare exactly the formats the kernel actually handles.

6. Plan workspace memory. TRT owns workspace; the plugin requests it:

   • getWorkspaceSize returns the maximum bytes needed across the profile range. Underestimating corrupts other ops' workspace (silent data corruption, hard to debug). Overestimating wastes device memory but is safe.
   • Inside enqueue, the workspace pointer is provided. Carve it up with explicit offsets; do not call cudaMalloc. The pointer is valid only for the duration of the enqueue call.
   • If the kernel needs persistent state across calls (rare — autoregressive KV cache is the main case), it must be passed in as an explicit input/output tensor, not held inside the plugin instance, because TRT may schedule plugin calls across multiple streams.

7. Plan for CUDA graph capture. Production TRT inference uses enqueueV3 with CUDA graphs to amortize launch overhead:

   • No host allocation in enqueue. No std::vector resize, no new, no cudaMalloc, no cudaMallocAsync. All allocation must happen at configurePlugin or via getWorkspaceSize.
   • No synchronization in enqueue. No cudaStreamSynchronize, no cudaDeviceSynchronize, no event wait that targets a non-captured stream.
   • All kernel launches inside enqueue must use the provided cudaStream_t.
   • Triton autotuned kernels (Python-side autotune that JIT-compiles on first call) violate every one of these rules and are not TRT-plugin-compatible. The kernel must be ahead-of-time compiled CUDA C++ or pre-tuned PTX.

8. Plan the build and binding artifacts.

   • Build the plugin as a shared library (libmyplugin.so) linked against libnvinfer.so and libcudart.so. Export the creator via REGISTER_TENSORRT_PLUGIN.
   • Loading at runtime: nvinfer1::initLibNvInferPlugins(&logger, "") registers all built-in plugins; for custom plugins set LD_LIBRARY_PATH and call dlopen or use IPluginRegistry::loadLibrary.
   • For production, prefer static linking the plugin into the engine builder binary so version skew between the builder and the runtime cannot occur silently.
   • Python access: after the .so is on LD_LIBRARY_PATH, tensorrt.init_libnvinfer_plugins(logger, namespace) exposes the creator to the Python API. Use this for prototyping; production stays in C++.

9. Validate the plan against existing alternatives. Before committing:

   • Is there a built-in TRT layer that covers this op? Check the TRT operator reference for the target version. Built-in layers participate in tactic selection and graph fusion; plugins do not.
   • Can the op be expressed as a composition of existing layers? TRT's optimizer fuses elementwise chains and adjacent matmul + bias + activation patterns aggressively.
   • Does TRT-LLM already ship a similar plugin (RMSNorm, GPT attention, RoPE, paged KV cache)? Reuse it instead of writing a parallel implementation.
   • Has the kernel been benchmarked against the TRT built-in equivalent? If TRT's built-in is within 10% of the custom kernel, the maintenance cost of a plugin is rarely justified.

Kernel design rules

• Treat the plugin class as a thin wrapper: no algorithmic logic in the plugin C++ file. The CUDA kernel lives in its own translation unit; the plugin's enqueue only marshals pointers and launches.
• Make every attribute that affects shape inference, output dtype, or workspace size a constructor parameter that is serialized. Runtime-derived state (configured shapes, cached descriptors) lives in members but is never serialized.
• Make clone() a deep copy of attributes. Never share mutable state between clones.
• Match plugin name, version, and namespace exactly between the creator and the plugin instance. Mismatches surface as deserialization failures with confusing messages.
• Implement terminate() and destroy() to release any device resources allocated in configurePlugin (e.g. cuBLAS handles, persistent device buffers). TRT calls these when the engine is destroyed.
• Document every assumption the kernel makes about input layout (contiguous, NCHW vs NHWC, alignment of S_kv, etc.) in a comment at the top of the plugin source. Future maintainers will not infer this from the kernel.

Correctness requirements

• The plugin must produce bitwise-identical output to the standalone kernel for the same inputs and stream, given identical RNG state. Any divergence indicates either a marshaling bug in enqueue or a precision mismatch in the format combination table.
• The plugin must handle every shape in the optimization profile range, not just the opt shape. Test explicitly at min, opt, and max per profile.
• Engine save/load must round-trip. After serializing an engine to disk and loading it in a fresh process, the plugin output must match the freshly-built engine's output for the same inputs.
• Multiple concurrent execution contexts must produce independent, correct outputs. This validates clone() and the absence of shared mutable state.
• INT8 and FP8 paths must validate scales: a plugin that quietly drops the input scale produces garbage output that may still be in-range and pass shape checks.

Performance requirements

• The plugin must not block TRT's graph fusion across its boundary unnecessarily. If declaring an extra format combination would let TRT fuse a preceding cast into the plugin, declare it.
• Workspace size must be tight. Returning gigabytes for a small kernel pessimizes co-located plugins and can cause OOM on multi-engine deployments.
• enqueue must not introduce host-device synchronization. Any sync (even an implicit one from a host-blocking API) destroys CUDA graph compatibility and adds tens of microseconds per call.
• Document the expected per-call overhead of the plugin (kernel time + TRT dispatch) and compare against the equivalent built-in or composition-of-layers path. A plugin that is slower than the built-in alternative should not be deployed.
• Do not claim performance parity with a hand-tuned TRT built-in without measurement on the target hardware and TRT version.

Output format

The final response must include:

1. API decision: IPluginV3 vs IPluginV2DynamicExt with explicit justification (TRT version, codebase convention, forward compatibility).
2. Lifecycle table: each lifecycle method mapped to what the plugin does in it, including which attributes are read/written.
3. Format combination table: the exact (input_dtype, format, output_dtype, format) tuples the plugin advertises, with a precision support row for FP32, FP16, BF16, INT8, FP8 as applicable.
4. Dynamic shape contract: which input dims are dynamic, the symbolic output shape expression, and any alignment constraints declared in supportsFormatCombination.
5. Serialization spec: the full list of fields serialized, with sizes; explicit confirmation that no runtime state is serialized.
6. Workspace plan: upper-bound formula for getWorkspaceSize, with the layout of how the workspace is partitioned inside enqueue.
7. CUDA graph compatibility checklist: confirmation that no host alloc, no sync, all launches on the provided stream.
8. Build/binding plan: shared library name, link line, registration mechanism, and Python loading path if applicable.
9. Risk register: list of known failure modes for this specific plugin (precision edge cases, profile boundary shapes, version compatibility).

Common failure modes

• Serializing runtime state instead of attributes: the plugin saves shape descriptors from configurePlugin into the serialized blob. Engine fails to load on a second run because the descriptors point to freed memory.
• Wrong workspace size: getWorkspaceSize returns the size for the opt shape only. At max shape the plugin overruns into the next op's workspace, causing data corruption that manifests two layers downstream and is nearly impossible to debug.
• Plugin name/version mismatch between creator and instance: deserialization fails with "no creator registered for plugin X version Y" and the actual mismatch is a typo or a forgotten version bump.
• Dynamic shape edge case at min profile: kernel works at S_q == 64 (opt) but segfaults at S_q == 1 (min). The min profile is rarely tested in development; production hits it on the first short request.
• Forgetting to override clone(): TRT silently shares the plugin instance across execution contexts; concurrent inference produces interleaved outputs that look like a numerical bug.
• Host allocation inside enqueue: the plugin uses std::vector<int> tmp(N) as scratch. Single-stream inference works; CUDA graph capture fails with cudaErrorStreamCaptureUnsupported at deployment time.
• Triton-autotuned kernel wrapped as a plugin: the kernel JIT-compiles on first call. The first call inside a captured graph fails because Python is not callable from a captured stream. The agent must be told upfront that Triton autotune and TRT plugins are incompatible.
• Reformat ops inserted around the plugin: the plugin advertises only kLINEAR kHALF but its neighbors run in kHWC8 kHALF. TRT inserts reformat ops on both sides; the plugin's measured speedup vanishes.
• INT8 plugin without proper scale handling: the plugin reads the input pointer but ignores PluginTensorDesc::scale. Output is numerically wrong but in-range, passes basic shape tests, and is caught only by an end-to-end accuracy regression.
• Missing ONNX op registration: the plugin works when building the engine from C++ but the ONNX → TRT path cannot find the op. Requires registering the op name in the ONNX parser's plugin namespace.

Review checklist

• Has the TRT version been pinned and the API choice (IPluginV3 vs IPluginV2DynamicExt) justified against it?
• Is every lifecycle method mapped to a specific responsibility, with a clear separation between build-time and runtime work?
• Are output shapes expressed symbolically via IExprBuilder rather than computed only at runtime?
• Does the format combination table list exactly the precision/format pairs the kernel actually supports — no more, no less?
• Is getWorkspaceSize an upper bound across the entire optimization profile range?
• Has it been confirmed that nothing serialized is runtime-derived state?
• Has plugin name, version, and namespace been chosen and held consistent between creator and instance?
• Has clone() been designed and tested for multi-context concurrency?
• Is enqueue free of host allocations and synchronization, with all launches on the provided stream?
• If INT8 or FP8: is the scale path explicit, with the calibration or quantization spec documented?
• Has every shape in min/opt/max for every profile been listed as a test case?
• Has the alternative of using a built-in TRT layer or a composition of existing layers been ruled out with measurement, not assumption?
• If the kernel uses runtime autotuning (Triton or otherwise), has the plan replaced it with an ahead-of-time compiled equivalent before plugin wrapping?