Kernel Providers

KernelOptimizer finds expensive torch.nn.functional calls inside a module hierarchy, evaluates compatible kernel providers, and returns a plan containing only candidates that are faster than the original PyTorch functions. The plan can be activated directly, without wrapping the module in an AITune backend.

Use KernelSelectorBackend instead when AITune should own provider selection, apply the selected plan while building another backend, and save the resulting plan or compiled artifact in an AITune checkpoint. The direct API documented on this page is useful when the application should manage the plan and its runtime explicitly.

Use the direct optimizer when you want to:

• optimize individual functional calls without compiling the complete module;
• control which kernel implementations participate in selection;
• optimize a module invoked from a larger model or pipeline;
• inspect, serialize, or activate the selected provider plan yourself.

CUDA is required for profiling and benchmarking.

Provider preparation, correctness validation, and benchmarking run under torch.no_grad(). Applying a plan only manages provider activation, so inference code should explicitly use torch.no_grad().

How kernel providers work

The optimization flow has two separate phases:

1. KernelOptimizer.make_plan() profiles the workload, prepares, validates, and benchmarks candidates, and returns an immutable KernelOptimizationPlan. It does not modify the module.
2. plan.apply(module) installs forward hooks on the selected module for the duration of a context. During its forward pass, matching torch.nn.functional calls are temporarily redirected to the selected providers.

The optimizer:

1. profiles the supplied inference callable and collects representative inputs for functions supported by configured providers or generators;
2. ranks observed functional calls by CUDA kernel time, summarizing up to 100 functions;
3. submits eligible asynchronous generators before evaluating static providers;
4. calls prepare() on compatible static providers to derive the state required for inference;
5. validates every prepared or generated provider against the original PyTorch function under torch.no_grad();
6. benchmarks valid candidates under torch.no_grad() using the observed input distribution;
7. selects the fastest provider for each function only when it beats the original function.

An unavailable optional runtime or a failing candidate is isolated and skipped while the remaining candidates continue. SageAttention and FlashAttention-4 load their runtime functions lazily during provider inference; prepare() only derives an inference plan from the representative samples.

Static providers and generators are considered only when the function meets their configured minimum profiled time share. The 100-function summary limit is an internal safety bound rather than a user-facing tuning option.

Direct optimizer example

The following example evaluates the PyTorch SDPA implementations available on the current GPU. It uses KernelOptimizer directly and does not create an AITune backend:

import logging
import torch
import torch.nn.functional as F
from torch import nn
from torch.nn.attention import SDPBackend
from aitune.torch.kernel_forge import KernelOptimizer
from aitune.torch.kernel_forge.kernel_provider import TorchSDPAKernelProvider
logging.basicConfig(level=logging.INFO, format="%(message)s", force=True)
class Attention(nn.Module):
    def forward(self, query, key, value):
        return F.scaled_dot_product_attention(query, key, value, is_causal=True)
model = Attention().eval().cuda()
sample = tuple(
    torch.randn(2, 8, 256, 64, device="cuda", dtype=torch.float16)
    for _ in range(3)
)
data = [(sample, {})]
optimizer = KernelOptimizer(
    kernel_providers=[
        TorchSDPAKernelProvider(SDPBackend.MATH),
        TorchSDPAKernelProvider(SDPBackend.EFFICIENT_ATTENTION),
        TorchSDPAKernelProvider(SDPBackend.CUDNN_ATTENTION),
        TorchSDPAKernelProvider(SDPBackend.FLASH_ATTENTION),
    ],
)
plan = optimizer.make_plan(model, data)
with plan.apply(model):
    output = model(*sample)

A plan can be empty when no candidate supports every representative input, passes correctness validation, or improves on the baseline. An empty plan is a valid result: plan.apply(model) leaves the module running its original PyTorch functions.

Built-in providers

AITune includes the following static providers:

FlashAttention-4 targets Hopper and Blackwell GPUs, such as H100 and B200. The PyPI sageattention package provides SageAttention V1; for newer implementations, install SageAttention from its source repository.

Optional providers can be configured together. An unavailable or incompatible implementation is skipped while the remaining candidates continue through validation and benchmarking:

from torch.nn.attention import SDPBackend
from aitune.torch.kernel_forge import KernelOptimizer
from aitune.torch.kernel_forge.kernel_provider import (
    FlashAttention4KernelProvider,
    SageAttentionKernelProvider,
    TorchSDPAKernelProvider,
)
optimizer = KernelOptimizer(
    kernel_providers=[
        FlashAttention4KernelProvider(),
        SageAttentionKernelProvider(),
        TorchSDPAKernelProvider(SDPBackend.FLASH_ATTENTION),
    ],
)

Providers specialize their inference plans from representative samples. Inconsistent or unsupported sample plans can cause prepare() to return False; runtime and correctness failures reject the candidate before it can enter the selected plan.

Diffusers attention dispatcher

DiffusersAttentionKernelProvider exposes Diffusers attention dispatcher implementations as kernel candidates. Create one provider for each backend that should be evaluated:

from aitune.torch.kernel_forge.kernel_provider import (
    DiffusersAttentionBackend,
    DiffusersAttentionKernelProvider,
)
providers = [
    DiffusersAttentionKernelProvider(DiffusersAttentionBackend.FLASH),
    DiffusersAttentionKernelProvider(DiffusersAttentionBackend.SAGE),
    DiffusersAttentionKernelProvider(DiffusersAttentionBackend.XFORMERS),
]

DiffusersAttentionBackend lists the dispatcher implementations compatible with SDPA replacement. Their availability and optional dependencies depend on the installed Diffusers version. See the Diffusers attention backend documentation for the current list. The provider converts PyTorch SDPA’s 4D HND layout to the NHD layout expected by the Diffusers dispatcher and converts its output back to HND.

Hub-backed implementations, whose enum names end in _HUB, download their kernels from the Hugging Face Hub and require kernels>=0.12 in addition to Diffusers. Non-Hub implementations do not depend on the kernels package.

The Diffusers native, _native_cudnn, _native_efficient, _native_flash, and _native_math implementations call torch.nn.functional.scaled_dot_product_attention themselves and therefore cannot replace that function without recursion. They are intentionally omitted from DiffusersAttentionBackend; use TorchSDPAKernelProvider for those PyTorch SDPA backends instead.

Provider interface and lifecycle

Import the base provider APIs from aitune.torch.kernel_forge.kernel_provider:

from aitune.torch.kernel_forge.kernel_provider import (
    KernelProvider,
    KernelProviderState,
    kernel_provider_from_dict,
)

A provider implements one torch.nn.functional function and has two states:

Subclasses implement:

The public prepare(samples) method is idempotent. A successful call changes the state from INIT to READY; repeated calls on a ready provider return True without rebuilding its state. Calling the provider or to_dict() before it is ready raises RuntimeError.

The default name and repr(provider) use the provider class name. Providers may override name with a more useful description. Each subclass is registered under its class name when its module is imported, allowing kernel_provider_from_dict() and KernelOptimizationPlan.from_dict() to restore the concrete class. Import custom provider classes before restoring plans that contain them.

Asynchronous kernel generators

Kernel generators produce providers asynchronously and are exported from the same package:

from aitune.torch.kernel_forge.kernel_provider import (
    KernelGenerationResult,
    KernelGenerator,
)

A KernelGenerator implements:

• __repr__() for a human-readable description;
• supports_functions() to list supported torch.nn.functional function names;
• prepare(function, samples) to determine whether generation can handle all samples;
• submit(function, samples) to return a Future[KernelGenerationResult] without waiting for generation to finish.

A KernelGenerationResult contains the function name, description, and exactly one of a generated provider or an error message. A successful generator must return a ready, serializable provider because the optimizer immediately uses it for correctness validation and benchmarking. Generator exceptions are isolated, and unfinished futures are cancelled or ignored after generation_timeout.

Optimizing a module inside a pipeline

The inference callable and optimized module can be different objects. This is useful when a pipeline prepares inputs or invokes the target module internally:

prompt = "A futuristic cityscape"
data = [((prompt,), {})]
plan = optimizer.make_plan(
    pipe,
    data,
    module=pipe.transformer,
)
with plan.apply(pipe.transformer):
    result = pipe(prompt)

function is the callable executed for profiling. module defines the module hierarchy in which functional calls are attributed and later redirected to providers. When function is an nn.Module, it is also used as module by default.

See Module Function Kernel Profiler for details about function attribution and representative input collection.

Runtime lifecycle

For temporary activation, apply the plan directly. Hooks are removed when the context exits, including when inference raises an exception:

with plan.apply(model):
    output = model(*sample)

plan.apply() enters torch.no_grad() for inference and restores the previous gradient state when the context exits.

Use KernelProviderRuntime directly when activation must span multiple contexts or requires explicit lifecycle control:

from aitune.torch.kernel_forge import KernelProviderRuntime
runtime = KernelProviderRuntime(model, plan)
with torch.no_grad():
    runtime.activate()
    try:
        output = model(*sample)
    finally:
        runtime.deactivate()
with runtime.applied():
    output = model(*sample)

Activation and deactivation are idempotent. If a runtime is already active, applied() preserves that state when the context exits. Direct activate() and deactivate() calls only manage provider hooks, so callers using that lifecycle must continue to manage the inference context explicitly.

The runtime temporarily changes process-global torch.nn.functional attributes during the selected module’s forward pass. Do not run concurrent forwards that overlap an active provider scope in different threads.

Saving and restoring a plan

KernelOptimizationPlan.providers is an immutable tuple of prepared providers. Plans serialize it under the "providers" key and can be restored without profiling again:

import json
from pathlib import Path
import torch
from aitune.torch.kernel_forge import KernelOptimizationPlan
plan_path = Path("kernel-plan.json")
plan_path.write_text(json.dumps(plan.to_dict()))
restored_plan = KernelOptimizationPlan.from_dict(
    json.loads(plan_path.read_text())
)
with restored_plan.apply(model):
    output = model(*sample)

Each serialized provider includes a "type" field derived from its class name and its provider-specific inference state. Deserialization restores providers directly in the READY state. Optional runtime packages required by selected providers must be installed when optimized inference runs. A restored plan remains specialized to the functional call patterns represented by the samples used during optimization.

Configuration

The most relevant KernelOptimizer options are:

Representative data uses the same [(args, kwargs), ...] structure as ModuleFunctionKernelProfiler.profile(). Include every input shape, dtype, layout, and argument combination that the applied plan is expected to handle.