TensorRT Backend Guide

The TensorRT backend provides highly optimized inference using NVIDIA’s TensorRT engine. It offers the best performance for production deployments on NVIDIA GPUs and seamlessly integrates TensorRT Model Optimizer for advanced quantization workflows.

Overview

The TensorRT backend:

• High Performance: Maximum inference speed on NVIDIA GPUs
• Dynamic Shapes: Supports optimization profiles for variable input sizes
• Quantization: INT8, FP8, INT4, FP16/BF16 autocast, and mixed precision support
• CUDA Graphs: Optional CUDA graph capture for reduced CPU overhead
• Model Optimizer Integration: Advanced quantization via TensorRT Model Optimizer
• Flexible Export: Supports both Dynamo and script-based ONNX export

Quick Start

Basic Usage

1
from aitune.torch.backend import TensorRTBackend, TensorRTBackendConfig, ONNXAutoCastConfig, ONNXQuantizationConfig, TorchQuantizationConfig
2
import aitune.torch as ait
3
4
# Configure TensorRT backend
5
config = TensorRTBackendConfig(use_dynamo=True)
6
backend = TensorRTBackend(config)
7
8
# Use with tuning
9
from aitune.torch.tune_strategy import OneBackendStrategy
10
strategy = OneBackendStrategy(backend=backend)
11
12
model = ait.Module(model, "my-model", strategy=strategy)
13
ait.tune(model, input_data)

With FP16 Precision

1
config = TensorRTBackendConfig(
2
    quantization_config=ONNXAutoCastConfig(precision="fp16"),
3
    workspace_size=1 << 30,  # 1GB workspace
4
)
5
backend = TensorRTBackend(config)

With CUDA Graphs

1
config = TensorRTBackendConfig(
2
    use_cuda_graphs=True,  # Enable CUDA graphs
3
)
4
backend = TensorRTBackend(config)

Configuration Options

TensorRTBackendConfig

1
@dataclass
2
class TensorRTBackendConfig(BackendConfig):
3
    use_dynamo: bool = True
4
    workspace_size: int | None = None
5
    opset_version: int | None = None
6
    optimization_level: int | None = None
7
    compatibility_level: int | None = None
8
    timing_cache: Path | None = None
9
    profiles: ProfileMode | list[TensorRTProfile] = ProfileMode.SINGLE
10
    device: str = "cuda"
11
    quantization_config: ONNXAutoCastConfig | ONNXQuantizationConfig | TorchQuantizationConfig | None = None
12
    enable_tf32: bool = True
13
    use_cuda_graphs: bool = False

use_dynamo

Use torch.dynamo for ONNX export (recommended).

1
# Use Dynamo export (recommended)
2
config = TensorRTBackendConfig(use_dynamo=True)
3
4
# Use script-based export (fallback)
5
config = TensorRTBackendConfig(use_dynamo=False)

When to use:

• True (default): Better compatibility with modern PyTorch models
• False: Legacy models or when Dynamo export fails

workspace_size

Maximum memory workspace for TensorRT engine building.

1
config = TensorRTBackendConfig(
2
    workspace_size=1 << 30,  # 1GB
3
)
4
5
# Or larger for complex models
6
config = TensorRTBackendConfig(
7
    workspace_size=4 << 30,  # 4GB
8
)

Guidelines:

• Default: TensorRT chooses automatically
• Larger workspace → More optimization opportunities → Longer build time
• Recommended: 1-4GB for most models

opset_version

ONNX opset version for export.

1
config = TensorRTBackendConfig(
2
    opset_version=17,  # Use ONNX opset 17
3
)

Guidelines:

• Default: Latest stable opset
• Specify only if you need a particular opset for compatibility

optimization_level

TensorRT builder optimization level (0-5).

1
config = TensorRTBackendConfig(
2
    optimization_level=5,  # Maximum optimization
3
)

Levels:

• 0: No optimization
• 3: Default (balanced)
• 5: Maximum optimization (longer build time)

compatibility_level

Hardware compatibility level for the engine.

1
import tensorrt as trt
2
3
config = TensorRTBackendConfig(
4
    compatibility_level=trt.HardwareCompatibilityLevel.AMPERE_PLUS,
5
)

Options:

• None: Optimized for current GPU
• Specific level: Portable across compatible GPUs

timing_cache

Path to timing cache for faster subsequent builds.

1
from pathlib import Path
2
3
config = TensorRTBackendConfig(
4
    timing_cache=Path("/path/to/timing_cache.bin"),
5
)

Benefits:

• Faster engine rebuilds
• Reuse timing information across builds
• Especially useful during development

profiles

Optimization profiles for dynamic shapes.

1
from aitune.torch.backend.tensorrt import ProfileMode, TensorRTProfile
2
3
# Single profile (default)
4
config = TensorRTBackendConfig(
5
    profiles=ProfileMode.SINGLE,
6
)
7
8
# Multiple profiles from samples
9
config = TensorRTBackendConfig(
10
    profiles=ProfileMode.SAMPLES_USED,
11
)
12
13
# Custom profiles
14
config = TensorRTBackendConfig(
15
    profiles=[
16
        TensorRTProfile()
17
            .add_input_shape("input", (1, 3, 224, 224), (4, 3, 224, 224), (8, 3, 224, 224)),
18
    ]
19
)

See Optimization Profiles section for details.

device

Device for TensorRT engine.

1
config = TensorRTBackendConfig(
2
    device="cuda",  # Default
3
)

quantization_config

TensorRT backend supports multiple quantization methods through TensorRT Model Optimizer integration. Use ONNXAutoCastConfig for FP16/BF16 mixed precision, ONNXQuantizationConfig for ONNX INT8/FP8/INT4 quantization, and TorchQuantizationConfig for ModelOpt PyTorch quantization presets.

1
# FP16/BF16 mixed precision autocast
2
config = TensorRTBackendConfig(
3
    quantization_config=ONNXAutoCastConfig(precision="fp16"),
4
)
5
6
# or
7
8
# ONNX quantization
9
config = TensorRTBackendConfig(
10
    quantization_config=ONNXQuantizationConfig(precision="int8", calibration_method="max"),
11
)
12
13
# or
14
15
# ModelOpt PyTorch quantization presets
16
config = TensorRTBackendConfig(
17
    quantization_config=TorchQuantizationConfig(quantization_config="FP8_DEFAULT_CFG"),
18
)

For current ModelOpt preset names and version-specific support, use the Model Optimizer documentation for the installed version instead of copying preset lists into AITune docs.

enable_tf32

Enable TF32 tensor cores on Ampere+ GPUs.

1
config = TensorRTBackendConfig(
2
    enable_tf32=True,  # Default
3
)

Benefits:

• Faster FP32 operations on Ampere and newer GPUs
• No accuracy loss for most models
• Recommended to keep enabled

use_cuda_graphs

Enable CUDA graph capture for inference.

1
config = TensorRTBackendConfig(
2
    use_cuda_graphs=True,
3
)

Benefits:

• Reduced CPU overhead
• Better performance for small models
• Automatic re-capture on shape changes

Limitations:

• First inference is slower (graph capture)
• Shape changes trigger re-capture
• Not beneficial for very large models

Optimization Profiles

Optimization profiles define the range of input shapes TensorRT will optimize for. They are essential for models with dynamic input sizes.

Profile Modes

SINGLE (Default)

Automatically generates a single profile from recorded samples:

1
config = TensorRTBackendConfig(
2
    profiles=ProfileMode.SINGLE,
3
)

• Min shape: Minimum observed across all samples
• Opt shape: Maximum observed across all samples
• Max shape: Maximum observed across all samples

SAMPLES_USED

Generates one profile per unique input shape:

1
config = TensorRTBackendConfig(
2
    profiles=ProfileMode.SAMPLES_USED,
3
)

Important: Increase max_num_samples_stored:

1
from aitune.torch.config import config as global_config
2
3
global_config.max_num_samples_stored = 100  # Or float("inf")

Use case: When you have distinct input shape categories that need separate optimization.

Custom Profiles

Define exact optimization profiles:

These profiles assume the module defines forward(input), so input is the path of its top-level tensor parameter:

1
from aitune.torch.backend.tensorrt import TensorRTProfile
2
3
profiles = [
4
    # Profile for small inputs
5
    TensorRTProfile()
6
        .add_input_shape(
7
            "input",
8
            min_shape=(1, 3, 224, 224),
9
            opt_shape=(4, 3, 224, 224),
10
            max_shape=(8, 3, 224, 224),
11
        ),
12
    # Profile for large inputs
13
    TensorRTProfile()
14
        .add_input_shape(
15
            "input",
16
            min_shape=(1, 3, 512, 512),
17
            opt_shape=(4, 3, 512, 512),
18
            max_shape=(8, 3, 512, 512),
19
        ),
20
]
21
22
config = TensorRTBackendConfig(profiles=profiles)

Finding Input Paths

Input tensor paths are shown in tuning logs. Top-level paths are forward parameter names; nested paths include all dictionary keys, sequence indices, or attributes:

INFO - 🚀 Tuning graph `0` for module `my-model`:
INFO -   graph_spec:
INFO -     input_spec:
 Tensors:
╒═══════════════╤═════════════════╤═══════════════════════════════╤══════════════════╤══════════════════╤═══════════════╕
│ Access Path   │ Semantic Path   │ Shape                         │ Min Shape        │ Max Shape        │ Dtype         │
╞═══════════════╪═════════════════╪═══════════════════════════════╪══════════════════╪══════════════════╪═══════════════╡
│ input         │ input           │ ['batch0', 3, 'dim2', 'dim3'] │ [2, 3, 224, 224] │ [8, 3, 448, 448] │ torch.float32 │
╘═══════════════╧═════════════════╧═══════════════════════════════╧══════════════════╧══════════════════╧═══════════════╛

Use the string shown in the Path column as the profile key. For example, the top-level path input uses "input", while a nested dictionary path uses inputs["tokens"].

Best Practices for Profiles

1. Min < Opt < Max: Ensure min ≤ opt ≤ max for all dimensions
2. Opt = Typical: Set opt to your most common input size
3. Range Coverage: Ensure your runtime inputs fall within [min, max]
4. Multiple Profiles: Use for distinct size categories, not slight variations
5. Test Runtime Shapes: Verify your production shapes are covered

Troubleshooting

Issue: ONNX export fails

Solution: Try disabling Dynamo export:

1
config = TensorRTBackendConfig(use_dynamo=False)

Issue: Engine build fails due to memory

Solution: Reduce workspace size:

1
config = TensorRTBackendConfig(workspace_size=512 << 20)  # 512MB

Issue: Runtime shape not supported

Error: Input shape X exceeds max profile shape Y

Solution: Update profiles to cover your runtime shapes:

1
profiles = [
2
    TensorRTProfile()
3
        .add_input_shape("input", min_shape=(1, 3, 224, 224), opt_shape=(4, 3, 224, 224), max_shape=(16, 3, 224, 224))
4
]
5
config = TensorRTBackendConfig(profiles=profiles)

Issue: Slow first inference

Cause: This is expected when using CUDA graphs (graph capture overhead).

Solution: Warmup with a few inference calls before measuring performance.

Issue: INT8 accuracy drop

Solution: Try different quantization algorithms:

1
# Try 'entropy' instead of 'max'
2
quantization_config = ONNXQuantizationConfig(
3
    precision="int8",
4
    calibration_method="entropy",
5
)

Best Practices

1. Use FP16: Use ONNXAutoCastConfig(precision="fp16") for FP16 mixed precision without a full quantization pass
2. Enable TF32: Keep enable_tf32=True on Ampere+ GPUs
3. Profile Carefully: Ensure optimization profiles cover all runtime shapes
4. Timing Cache: Use timing cache during development for faster iteration
5. CUDA Graphs: Enable for latency-sensitive small models
6. Workspace Size: Start with 1-2 GB and, increase if the build fails
7. Quantization: Validate accuracy with a representative test set

Next Steps

• Learn about Torch-TensorRT JIT Backend
• Learn about Torch-TensorRT AOT Backend
• Explore Tune Strategies
• Review Deployment Guide