Torch-TensorRT JIT Backend Guide

The Torch-TensorRT JIT backend integrates TensorRT acceleration through torch.compile(backend="torch_tensorrt"). This provides a seamless JIT (Just-In-Time) compilation experience without needing intermediate model formats.

Overview

• JIT Compilation: Compiles at runtime using torch.compile
• No Intermediate Formats: No ONNX export or separate engine files
• PyTorch Native: Stays within the PyTorch ecosystem
• Dynamic Recompilation: Automatically recompiles on shape changes
• FP16 Support: Built-in mixed precision support

Quick Start

1
from aitune.torch.backend import TorchTensorRTJitBackend, TorchTensorRTJitBackendConfig, TorchTensorRTConfig
2
import aitune.torch as ait
3
import torch
4
5
# Configure backend
6
config = TorchTensorRTJitBackendConfig(
7
    compile_config=TorchTensorRTConfig(),
8
)
9
backend = TorchTensorRTJitBackend(config)
10
11
# Use in tuning
12
strategy = ait.OneBackendStrategy(backend=backend)
13
14
model = ait.Module(model, "my-model", strategy=strategy)
15
ait.tune(model, input_data)

Configuration Options

TorchTensorRTJitBackendConfig

1
@dataclass
2
class TorchTensorRTJitBackendConfig(BackendConfig):
3
    compile_config: TorchTensorRTConfig
4
    fullgraph: bool = False
5
    dynamic: bool | None = None
6
    autocast_enabled: bool = False
7
    autocast_dtype: torch.dtype | None = None

compile_config

TensorRT compilation settings from torch_tensorrt:

1
from torch_tensorrt.dynamo import CompilationSettings
2
3
config = TorchTensorRTJitBackendConfig(
4
    compile_config=CompilationSettings(
5
        workspace_size=1 << 30,  # 1GB
6
    )
7
)

By default, the engine matches the model’s loaded dtype. To request FP16 kernels via the legacy weak-typing path (deprecated in TensorRT 10.12), pair enabled_precisions with use_explicit_typing=False:

1
config = TorchTensorRTJitBackendConfig(
2
    compile_config=CompilationSettings(
3
        enabled_precisions={torch.float16},
4
        use_explicit_typing=False,
5
    )
6
)

Common options:

• workspace_size: Maximum workspace memory in bytes
• enabled_precisions: Kernel dtype precisions TRT may pick (requires use_explicit_typing=False; legacy weak typing)
• use_explicit_typing: Respect graph dtypes (default True); set False to enable legacy weak typing

fullgraph

Require the entire function to be captured in a single graph.

1
config = TorchTensorRTJitBackendConfig(
2
    fullgraph=True,  # Raise error if graph breaks occur
3
)

Use cases:

• False (default): Allow partial compilation
• True: Ensure complete compilation or fail

dynamic

Enable dynamic shape tracing:

1
config = TorchTensorRTJitBackendConfig(
2
    dynamic=True,  # Enable dynamic shapes
3
)

Options:

• True: Generate dynamic kernels up-front
• False: Always specialize
• None (default): Auto-detect and recompile

autocast_enabled

Enable automatic mixed precision:

1
config = TorchTensorRTJitBackendConfig(
2
    autocast_enabled=True,
3
    autocast_dtype=torch.float16,
4
)

JIT vs AOT Torch-TensorRT

Understanding JIT/AOT Terminology

It’s important to distinguish between two uses of “JIT” and “AOT” in AITune:

AITune Tuning Modes

• Ahead-of-Time Tuning: The declarative approach using inspect(), wrap(), and tune()

  • You explicitly select modules to tune
  • Full control over the tuning process
  • Works with any backend (JIT or AOT)

• Just-in-Time Tuning: The automatic approach using environment variables or imports

  • No code changes required
  • AITune automatically discovers and tunes modules
  • Works with any backend (JIT or AOT)

Torch-TensorRT Backend Types

• TorchTensorRTJitBackend (this page): Uses torch.compile(backend="torch_tensorrt")

  • Compiles at runtime on first inference
  • Does not save compiled artifacts separately
  • Recompiles automatically on shape changes

• TorchTensorRTAotBackend: Uses torch.export.export() and torch_tensorrt.dynamo.compile()

  • Compiles during the tune() call
  • Saves compiled model to disk
  • Fixed compilation (no automatic recompilation)

Combining Them

You can use any combination:

1
# AOT Tuning + JIT Backend
2
# Explicit tuning with runtime compilation
3
wrapped_model = ait.Module(model, "model", strategy=ait.OneBackendStrategy(TorchTensorRTJitBackend()))
4
ait.tune(wrapped_model, data)
5
# the tuned model can be saved but during loading JIT backend will tune again
6
7
# AOT Tuning + AOT Backend
8
# Explicit tuning with ahead-of-time compilation (saved model)
9
wrapped_model = ait.Module(model, "model", strategy=ait.OneBackendStrategy(TorchTensorRTAotBackend()))
10
ait.tune(wrapped_model, data)
11
# the tuned model can be saved, during loading AOT backend will be load from disk
12
13
# JIT Tuning + JIT Backend
14
# Automatic tuning with runtime compilation
15
from aitune.torch import jit_config
16
from aitune.torch.tune_strategy import FirstWinsStrategy
17
jit_config.strategy = FirstWinsStrategy(backends=[TorchTensorRTJitBackend()])

$
export AUTOWRAPT_BOOTSTRAP=aitune_enable_jit_tuning
$
# each time you start the script JIT tuning starts all over again, JIT backend will tune again the module

1
# JIT Tuning + AOT Backend
2
# Automatic tuning with ahead-of-time compilation
3
from aitune.torch import jit_config
4
from aitune.torch.tune_strategy import FirstWinsStrategy
5
jit_config.strategy = FirstWinsStrategy(backends=[TorchTensorRTAotBackend()])

$
export AUTOWRAPT_BOOTSTRAP=aitune_enable_jit_tuning
$
# each time you start the script JIT tuning starts all over again, currently AITune does not reuse past AOT backend artifact, it will start AOT tuning from scratch

Key Takeaway: AITune’s tuning mode (JIT/AOT) is independent from the backend type (JIT/AOT). Choose based on your needs:

• Tuning mode: How you want to control tuning (automatic vs explicit)
• Backend type: How the model gets compiled and stored (runtime vs saved)

Best Practices

1. Use FP16: Enable FP16 for better performance
2. Dynamic Shapes: Enable if input sizes vary frequently
3. Fullgraph for Production: Use fullgraph=True to catch issues early
4. Warmup: Run a few inference calls before benchmarking

Troubleshooting

Issue: Compilation fails

Solution: Try with partial compilation:

1
config = TorchTensorRTJitBackendConfig(
2
    fullgraph=False,  # Allow partial compilation
3
)

Issue: Slow first inference

Cause: JIT compilation happens on the first run.

Solution: This is expected. Subsequent inferences will be fast.

Next Steps

• Compare with Torch-TensorRT AOT Backend
• Learn about TensorRT Backend
• Explore Tune Strategies