aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider

FlashAttention-4 kernel provider.

Module Contents

Classes

Name	Description
`FlashAttention4KernelProvider`	Run FlashAttention-4 with a sample-derived argument and layout plan.

Functions

Name	Description
`_call_flash_attention4`	Call FlashAttention-4 after copying HND inputs to NHD layout.
`_call_flash_attention4_copy_free`	Call FlashAttention-4 when HND inputs are NHD-contiguous views.
`_extract_qkv`	Extract query, key and value from supported SDPA arguments.
`_hnd_to_nhd`	Map PyTorch SDPA HND layout to contiguous FlashAttention NHD layout.
`_hnd_to_nhd_view`	Map HND to NHD without forcing contiguous storage.
`_import_flash_attention4`	Import and cache the optional FlashAttention-4 module.
`_nhd_to_hnd`	Map FlashAttention NHD output back to PyTorch SDPA HND layout.
`_normalize_sdpa_inputs`	Normalize PyTorch SDPA arguments to FlashAttention-4 arguments.
`_qkv_source`	Return how q/k/v are passed in the representative sample.
`_sample_replacement_plan`	Build an inference plan for one representative sample.
`_validate_qkv`	Validate the subset of PyTorch SDPA inputs handled by this provider.

Data

_CALL_IMPLS

_KEYWORD_QKV

_POSITIONAL_QKV

_RUNTIME_UNAVAILABLE_MESSAGE

logger

API

class aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider.FlashAttention4KernelProvider()

Bases: KernelProvider

Run FlashAttention-4 with a sample-derived argument and layout plan.

_backend

Callable[..., Any]

Load and validate the FlashAttention-4 runtime function lazily.

_version

str

Resolve the FlashAttention-4 version once.

flash_kwargs

dict[str, Any] = {}

name

str

Human-readable provider name.

supported_function

str

Name of the function replaced by this provider.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider.FlashAttention4KernelProvider._from_dict(
    state_dict: dict[str, typing.Any]
) -> aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider.FlashAttention4KernelProvider

classmethod

Restore the selected inference plan.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider.FlashAttention4KernelProvider._infer(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Run the selected FlashAttention-4 inference plan.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider.FlashAttention4KernelProvider._load_runtime_dependencies() -> None

Load FlashAttention-4 dependencies before graph capture.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider.FlashAttention4KernelProvider._prepare(
    samples: list[aitune.torch.module.sample_store.Sample]
) -> bool

Select one inference plan for all representative samples.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider.FlashAttention4KernelProvider._to_dict() -> dict[str, typing.Any]

Serialize the selected inference plan.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._call_flash_attention4(
    flash_attn_func,
    q,
    k,
    v,
    flash_kwargs
)

Call FlashAttention-4 after copying HND inputs to NHD layout.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._call_flash_attention4_copy_free(
    flash_attn_func,
    q,
    k,
    v,
    flash_kwargs
)

Call FlashAttention-4 when HND inputs are NHD-contiguous views.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._extract_qkv(
    args,
    kwargs
)

Extract query, key and value from supported SDPA arguments.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._hnd_to_nhd(
    tensor: torch.Tensor
) -> torch.Tensor

Map PyTorch SDPA HND layout to contiguous FlashAttention NHD layout.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._hnd_to_nhd_view(
    tensor: torch.Tensor
) -> torch.Tensor

Map HND to NHD without forcing contiguous storage.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._import_flash_attention4() -> types.ModuleType | None

Import and cache the optional FlashAttention-4 module.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._nhd_to_hnd(
    tensor: torch.Tensor
) -> torch.Tensor

Map FlashAttention NHD output back to PyTorch SDPA HND layout.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._normalize_sdpa_inputs(
    args,
    kwargs
)

Normalize PyTorch SDPA arguments to FlashAttention-4 arguments.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._qkv_source(
    args,
    kwargs
) -> str

Return how q/k/v are passed in the representative sample.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._sample_replacement_plan(
    sample: aitune.torch.module.sample_store.Sample
) -> tuple[str, tuple[tuple[str, typing.Any], ...], bool]

Build an inference plan for one representative sample.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._validate_qkv(
    q,
    k,
    v,
    enable_gqa: bool
) -> bool

Validate the subset of PyTorch SDPA inputs handled by this provider.

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._CALL_IMPLS = {False: _call_flash_attention4, True: _call_flash_attention4_copy_free}

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._KEYWORD_QKV = 'keyword_qkv'

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._POSITIONAL_QKV = 'positional_qkv'

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider._RUNTIME_UNAVAILABLE_MESSAGE = 'FlashAttention-4 is not available. Install it with `pip install --pre flash-att...

aitune.torch.kernel_forge.kernel_provider.flash_attention4_provider.logger = logging.getLogger(__name__)