NVIDIA Dynamo · AITune

AITune

One API that finds the fastest validated backend for every module of your PyTorch model.

Apache 2.0PyTorch 2.8+v0.6.0

The problem · deployment entropy

Every model has dozens of paths to a fast engine, and they keep moving

Pain

7core technologies

~26runtime × precision configurations

4owners: NVIDIA · Meta · Microsoft · Lightning AI

Result: teams fall back to vanilla PyTorch or the simplest path, and leave GPU performance on the table.

Before → after

From a manual loop to one call

✗ Manual · per module, per tool
torch.onnx.export(unet, ...)          # trace or dynamo?
modelopt.quantize(...)                # pin ORT version
trtexec --onnx=unet.onnx --stronglyTyped \
        --minShapes=... --optShapes=...
torch_tensorrt.compile(vae, ...)      # 1 of 3 APIs
torch.compile(text_enc, mode="max-autotune")
# benchmark · compare outputs · repeat

✓ JIT · automatic# no code changes: one env var
AUTOWRAPT_BOOTSTRAP=aitune_enable_jit_tuning \
  python your_script.py

✓ AOT · tune once, save for deploymentimport aitune.torch as ait
modules = ait.inspect(pipe, inputs)
pipe    = ait.wrap(pipe, modules.get_modules())
ait.tune(pipe, inputs)
ait.save(pipe, "tuned_pipe.ait")

What AITune searches

One frontend over every NVIDIA and PyTorch inference path

Click a framework to see the paths and precisions AITune can try.

Architecture

Inspect, wrap, tune, validate, select, deploy

Press → to walk through the pipeline.

↳ falls back to child modules on a graph breakprefill / decode tuned as separate graphsmulti-GPU: TP / CPbring your own Backend

Pipeline from the repo README and the FlexTensor + AITune deck. Inspect finds tunable nn.Modules and records their inputs. Wrap sets per-module strategies and backend configs. Tune builds each candidate (optionally with PTQ, kernel selection or CUDA graphs).
Validate: outputs must be finite and shapes must match on static, dynamic and batch axes; a performance gate profiles a Torch eager baseline and, by default, rejects candidates that are not faster.
Select: pick per module by strategy (max throughput, min latency, or throughput under a latency budget), then run in Python or save a .ait artifact. Modules are tuned one at a time with the others deactivated to free memory.
Robustness: hierarchical conversion falls back recursively to child modules on a conversion error or graph break (dynamic control flow, unsupported inputs). Different input signatures become separate execution graphs, each with its own backend; for an LLM, prefill and decode get separate cache-aware backends. Modules with DTensor / torch.distributed state get multi-GPU backends (TP/CP, 0.6.0).
Extensibility: custom backends subclass Backend and BackendConfig (implement _build, _activate, _infer, _deactivate, _deploy); custom strategies subclass TuneStrategy._tune(...). The experimental Kernel Selector swaps in SageAttention or FlashAttention-4 and builds a delegate backend.

Hybrid deployment

No backend wins everywhere, so AITune picks per module

Diffusion pipeline · speedup vs eager per module

Click a module to see the candidates it raced.

How often each backend is the best, across ~300 models

Internal dashboard, ~300 TRT-priority models on H100. AITune is 100% by construction: it picks the winner.

Selection strategies

Tell it the objective; it picks a different winner

Objective

max_latency_ms18

Strategy classes: MaxThroughputStrategy (default) · MinLatencyStrategy · LatencyBudgetStrategy · FirstWinsStrategy · OneBackendStrategy

Two modes

JIT to try it in seconds, AOT to ship it anywhere

# unchanged application
pipe = FluxPipeline.from_pretrained(...)
pipe("A cat holding a sign")

# one env var turns tuning on
AUTOWRAPT_BOOTSTRAP=aitune_enable_jit_tuning \
  python your_script.py

Records the first calls, tunes top-down while the app runs, and falls back to child modules on a graph break. Deferred mode suits denoising loops.

# tune once, save, reload anywhere
import aitune.torch as ait
modules = ait.inspect(pipe, inputs)
pipe = ait.wrap(pipe, modules.get_modules())
ait.tune(pipe, inputs)
ait.save(pipe, "tuned_pipe.ait")

pipe = ait.load(pipe, "tuned_pipe.ait")

You choose the modules and per-module strategies; AITune benchmarks across batch sizes and writes a portable, checksummed artifact.

	JIT	AOT
Code changes	None	4 lines
Module selection	Automatic	You choose
Tune time	Quick	Slower
Strategies	Global	Global or per module
Benchmarking	No	Yes
Batch extrapolation	No	Yes
Saved artifact	No	.ait
Next start	Re-tunes	Loads, zero warmup

.aitstate + engines · checksummed · portable

→

Pythonait.load(model, "model.ait")

Dynamo workerOpenAI-compatible · multi-GPU · image, video, audio, embeddings

Tritongenerated model repository

Comparison from docs/guides/jit_tuning.md. JIT records samples until min_samples (default 1), tunes hierarchically from the top module, and recurses into children on graph breaks up to max_depth_level. Deferred mode (jit.tune.deferred()) is for denoising loops. JIT can also be enabled with import aitune.torch.jit.enable or @patch_for_jit_tuning.
AOT: the .ait artifact holds tuned module state plus backend artifacts, with checksums verified on load and relative paths for portability. Model code and dependencies must exist at serve time; multi-GPU artifacts are rank-local.
"Compile once, zero warmup on redeploy" (FlexTensor + AITune deck).
Deployment guide (docs/guides/deployment): Python (load the checkpoint), Dynamo (an AITune worker exposes the inference function through Dynamo's OpenAI-compatible frontend, added in 0.4.0; multi-GPU collective worker in 0.6.0), and Triton (generates a model repository from ONNX Runtime, TensorRT or Inductor-AOT artifacts; repository publishing in 0.7.0, unreleased).
Dynamo modalities: image (FLUX, SD), video (Wan), audio (Qwen3-TTS), embeddings (E5 Large, BGE).

Performance · internal benchmarks

Faster than torch.compile, without touching the model

1.46×average speedup vs eager on ~100 Hugging Face models torch.compile: 1.22×

2.23×P75 speedup torch.compile: 1.79×

18backends searched · 128 models · 2 GPU platforms · 20k+ runs

Sort by

eager 1.0×torch.compileAITune* with experimental Kernel Selector

Get started

Shipping every 1–2 months. Try it today.

pip install --extra-index-url \
  https://pypi.nvidia.com aitune

# zero code changes
AUTOWRAPT_BOOTSTRAP=aitune_enable_jit_tuning \
  python your_script.py

GitHub →Docs →Recipes →

13 ready recipes

FLUX.1-devSD3 MediumSDXLWan 2.1 T2VQwen3 0.6BParakeet CTCParakeet RNNTBERT-baseE5 Large V2ESM-2 650MResNet-50YOLOv10n

Each recipe tunes, validates, benchmarks and deploys to Dynamo or Triton.