Torch Inductor AOT backend.

Module Contents

Classes

Name	Description
`TorchInductorAotBackend`	Backend that compiles models using AOT Inductor (`torch._inductor.aoti_compile_and_package`).
`TorchInductorAotBackendConfig`	Configuration for TorchInductorAotBackend.
`TorchInductorAotBuildStep`	Identifiers for discrete sub-steps of a TorchInductorAot backend build.

Data

logger

API

class aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBackend(
    config: aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBackendConfig | None = None
)

Bases: Backend

Backend that compiles models using AOT Inductor (torch._inductor.aoti_compile_and_package).

Requires PyTorch >= 2.6.

Performs Ahead-of-Time compilation to a portable .pt2 artifact saved to disk. Dynamic batch shapes are inferred automatically from graph_spec when a batch axis is detected. At inference time the artifact is loaded via torch._inductor.aoti_load_package, enabling deployment without Python-interpreter overhead.

Workflow::

backend = TorchInductorAotBackend()

build / tune as usual through ait.tune()

ait.save(model, “model.ait”)

later

ait.load(model, “model.ait”)

STATE_COMPILED_MODEL_PATH

= 'compiled_model_path'

STATE_DEVICE

= 'device'

STATE_TYPE

= 'type'

_build_mode

= BuildMode.AHEAD_OF_TIME

_compiled_model_artifact

ArtifactPath | None = None

_config

= config or TorchInductorAotBackendConfig()

_execution_modes

aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBackend._activate()

Load the compiled model from disk.

aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBackend._build(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    samples: aitune.torch.module.sample_store.SampleStore,
    cache_dir: pathlib.Path
) -> aitune.torch.backend.backend.Backend

Export and compile the model with AOT Inductor, then load the runner.

aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBackend._deactivate()

Deactivate backend.

aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBackend._deploy()

Deploy backend.

aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBackend._infer(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Run inference with the compiled AOT Inductor runner.

Parameters:

*args

AnyDefaults to ()

Inference arguments.

**kwargs

AnyDefaults to {}

Inference keyword arguments.

Returns: Any

The result of the inference.

aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBackend.describe() -> str

Returns the description of the backend.

aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBackend.from_dict(
    module: torch.nn.Module | None,
    state_dict: dict
) -> aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBackend

classmethod

Creates a backend from a state_dict.

aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBackend.key() -> str

Returns the key of the backend.

aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBackend.to_dict() -> dict

Returns the state_dict of the backend.

class aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBackendConfig(
    inductor_configs: dict[str, typing.Any] | None = None
)

Dataclass

Bases: BackendConfig

Configuration for TorchInductorAotBackend.

Parameters:

inductor_configs

dictDefaults to None

Mapping of torch._inductor.config attribute names to values, passed to torch._inductor.aoti_compile_and_package(). Call torch._inductor.list_options() to see all available keys. Example: {"max_autotune": True, "coordinate_descent_tuning": True}

inductor_configs

dict[str, Any] | None = None

class aitune.torch.backend.torch_inductor_aot_backend.TorchInductorAotBuildStep()

Bases: BackendBuildStep

Identifiers for discrete sub-steps of a TorchInductorAot backend build.

AOT_COMPILE

= 'AOT compile'

TORCH_EXPORT

= 'Torch export'

aitune.torch.backend.torch_inductor_aot_backend.logger = getLogger(__name__)