Backends

NVIDIA AITune supports multiple tuning backends, each with different characteristics and use cases. The backends align with a common interface for the build and inference process.

TensorRT Backend

The TensorRT backend provides highly optimized inference using NVIDIA’s TensorRT engine. It offers the best performance for production deployments. The backend integrates TensorRT Model Optimizer in a seamless flow.

1
from aitune.torch.backend import TensorRTBackend, TensorRTBackendConfig, ONNXAutoCastConfig
2
3
config = TensorRTBackendConfig(quantization_config=ONNXAutoCastConfig()) # FP16 autocast through ModelOpt
4
backend = TensorRTBackend(config)

CUDA Graphs Support

The TensorRT backend supports CUDA Graphs for reduced CPU overhead and improved inference performance. CUDA Graphs automatically capture and replay GPU operations, eliminating kernel launch overhead for repeated inference calls. This feature is disabled by default.

Keep in mind that graphs are automatically recaptured when input shapes change.

1
from aitune.torch.backend import TensorRTBackend, TensorRTBackendConfig
2
3
# Enable CUDA Graphs for optimized inference
4
config = TensorRTBackendConfig(use_cuda_graphs=True)
5
backend = TensorRTBackend(config)

Torch-TensorRT Backend (JIT)

The Torch-TensorRT JIT backend integrates TensorRT tuning directly into PyTorch, providing seamless tuning without model conversion through torch.compile.

1
import torch
2
from aitune.torch.backend import TorchTensorRTJitBackend, TorchTensorRTJitBackendConfig, TorchTensorRTConfig
3
4
config = TorchTensorRTJitBackendConfig(compile_config=TorchTensorRTConfig())
5
backend = TorchTensorRTJitBackend(config)

Torch-TensorRT Backend (AOT)

The Torch-TensorRT backend integrates TensorRT tuning directly into PyTorch, providing seamless tuning without model conversion through torch_tensorrt.compile.

1
import torch
2
from aitune.torch.backend import TorchTensorRTAotBackend, TorchTensorRTAotBackendConfig, TorchTensorRTConfig
3
4
config = TorchTensorRTAotBackendConfig(compile_config=TorchTensorRTConfig())
5
backend = TorchTensorRTAotBackend(config)

TorchAO Backend

The TorchAO backend leverages PyTorch’s AO (Accelerated Optimization) framework for model tuning.

1
from aitune.torch.backend import TorchAOBackend
2
3
backend = TorchAOBackend()

Torch Inductor Backend (JIT)

The Torch Inductor JIT backend uses PyTorch’s Inductor compiler through torch.compile for model tuning.

1
from aitune.torch.backend import TorchInductorJitBackend
2
3
backend = TorchInductorJitBackend()

Torch Inductor Backend (AOT)

The Torch Inductor AOT backend uses PyTorch’s AOT Inductor compiler to produce a compiled artifact that can be saved and loaded with AITune checkpoints.

1
from aitune.torch.backend import TorchInductorAotBackend
2
3
backend = TorchInductorAotBackend()

ONNXRuntime Backend

The ONNXRuntime backend exports the selected module to ONNX and runs inference through ONNX Runtime with CUDA or TensorRT execution providers.

1
from aitune.torch.backend import ONNXRuntimeBackend, ONNXRuntimeBackendConfig, ONNXExecutionProvider
2
3
config = ONNXRuntimeBackendConfig(execution_provider=ONNXExecutionProvider.CUDA)
4
backend = ONNXRuntimeBackend(config)