aitune.torch.kernel_forge.kernel_provider_runtime

Runtime activation of a selected kernel optimization plan.

Module Contents

Classes

Name	Description
`KernelProviderRuntime`	Install a kernel optimization plan by swapping `torch.nn.functional` calls.

API

class aitune.torch.kernel_forge.kernel_provider_runtime.KernelProviderRuntime(
    module: torch.nn.Module,
    plan: aitune.torch.kernel_forge.kernel_optimization_plan.KernelOptimizationPlan
)

Install a kernel optimization plan by swapping torch.nn.functional calls.

_function_stacks

dict[str, list[Callable]] = defaultdict(list)

_hooks

list = []

_provider_functions

is_active

bool

Whether runtime hooks are currently installed.

aitune.torch.kernel_forge.kernel_provider_runtime.KernelProviderRuntime._create_post_hook()

Create a hook that restores functions after each module forward.

aitune.torch.kernel_forge.kernel_provider_runtime.KernelProviderRuntime._create_pre_hook()

Create a hook that installs providers before each module forward.

aitune.torch.kernel_forge.kernel_provider_runtime.KernelProviderRuntime._create_provider_function(
    provider: aitune.torch.kernel_forge.kernel_provider.KernelProvider
) -> collections.abc.Callable

staticmethod

Expose a callable provider as a regular Python function.

aitune.torch.kernel_forge.kernel_provider_runtime.KernelProviderRuntime.activate() -> None

Activate the selected providers idempotently.

aitune.torch.kernel_forge.kernel_provider_runtime.KernelProviderRuntime.applied() -> collections.abc.Generator[None, None, None]

Temporarily activate the plan for inference while preserving prior runtime state.

aitune.torch.kernel_forge.kernel_provider_runtime.KernelProviderRuntime.deactivate() -> None

Deactivate providers and restore outstanding function patches idempotently.

Graph capture may invoke a module’s forward pre-hook without invoking its matching post-hook. Unwind any patches left on the function stacks so a provider cannot remain installed after the runtime is deactivated.