Torch Inductor JIT Backend Guide

1
from aitune.torch.backend import TorchInductorJitBackend, TorchInductorJitBackendConfig
2
import aitune.torch as ait
3
import torch
4
5
# Configure backend
6
config = TorchInductorJitBackendConfig(mode="max-autotune")
7
backend = TorchInductorJitBackend(config)
8
9
# Use in tuning
10
from aitune.torch.tune_strategy import OneBackendStrategy
11
strategy = ait.OneBackendStrategy(backend=backend)
12
13
model = ait.Module(model, "my-model", strategy=strategy)
14
ait.tune(model, input_data)

1
@dataclass
2
class TorchInductorJitBackendConfig(BackendConfig):
3
    fullgraph: bool = False
4
    dynamic: bool | None = None
5
    mode: str | None = None
6
    options: dict | None = None
7
    autocast_enabled: bool = False
8
    autocast_dtype: torch.dtype | None = None

1
# Default mode (balanced)
2
config = TorchInductorJitBackendConfig(mode="default")
3
4
# Reduce Python overhead with CUDA graphs
5
config = TorchInductorJitBackendConfig(mode="reduce-overhead")
6
7
# Maximum auto-tuning
8
config = TorchInductorJitBackendConfig(mode="max-autotune")
9
10
# Max autotune without CUDA graphs
11
config = TorchInductorJitBackendConfig(mode="max-autotune-no-cudagraphs")

1
config = TorchInductorJitBackendConfig(
2
    fullgraph=True,  # Error if graph breaks occur
3
    mode="max-autotune",
4
)

1
# Always generate dynamic kernels
2
config = TorchInductorJitBackendConfig(dynamic=True)
3
4
# Never generate dynamic kernels (always specialize)
5
config = TorchInductorJitBackendConfig(dynamic=False)
6
7
# Auto-detect (default)
8
config = TorchInductorJitBackendConfig(dynamic=None)

1
# See all options: torch._inductor.list_options()
2
config = TorchInductorJitBackendConfig(
3
    options={
4
        "triton.cudagraphs": True,
5
        "max_autotune": True,
6
        "coordinate_descent_tuning": True,
7
    }
8
)

1
config = TorchInductorJitBackendConfig(
2
    mode="max-autotune",
3
    autocast_enabled=True,
4
    autocast_dtype=torch.float16,
5
)

1
# Set environment variables before running
2
import os
3
os.environ['TORCH_LOGS'] = 'dynamic,perf_hints,graph_breaks'
4
5
# Then run tuning
6
ait.tune(wrapped_model, input_data)

1
# See what mode does
2
import torch
3
print(torch._inductor.list_mode_options())
4
5
# See all available options
6
print(torch._inductor.list_options())

1
config = TorchInductorJitBackendConfig(mode="default")

1
config = TorchInductorJitBackendConfig(
2
    mode="default",
3
    dynamic=True,
4
)