aitune.torch.kernel_forge.kernel_provider.diffusers_attention_provider

Diffusers attention dispatcher kernel provider.

Module Contents

Classes

Name	Description
`DiffusersAttentionBackend`	Diffusers attention backends compatible with SDPA replacement.
`DiffusersAttentionKernelProvider`	Run SDPA with a selected Diffusers attention dispatcher backend.

Functions

Name	Description
`_import_diffusers_attention_dispatch`	Import and cache the optional Diffusers attention dispatcher module.

Data

_RUNTIME_UNAVAILABLE_MESSAGE

API

class aitune.torch.kernel_forge.kernel_provider.diffusers_attention_provider.DiffusersAttentionBackend

Bases: enum.Enum

Diffusers attention backends compatible with SDPA replacement.

FLASH

= 'flash'

FLASH_3

= '_flash_3'

FLASH_3_HUB

= '_flash_3_hub'

FLASH_4_HUB

= 'flash_4_hub'

FLASH_HUB

= 'flash_hub'

FLEX

= 'flex'

SAGE

= 'sage'

SAGE_HUB

= 'sage_hub'

SAGE_QK_INT8_PV_FP16_CUDA

= '_sage_qk_int8_pv_fp16_cuda'

SAGE_QK_INT8_PV_FP16_TRITON

= '_sage_qk_int8_pv_fp16_triton'

SAGE_QK_INT8_PV_FP8_CUDA

= '_sage_qk_int8_pv_fp8_cuda'

SAGE_QK_INT8_PV_FP8_CUDA_SM90

= '_sage_qk_int8_pv_fp8_cuda_sm90'

XFORMERS

= 'xformers'

class aitune.torch.kernel_forge.kernel_provider.diffusers_attention_provider.DiffusersAttentionKernelProvider(
    backend: aitune.torch.kernel_forge.kernel_provider.diffusers_attention_provider.DiffusersAttentionBackend
)

Bases: KernelProvider

Run SDPA with a selected Diffusers attention dispatcher backend.

_backend

Callable[..., Tensor]

Load the Diffusers dispatcher bound to the selected backend lazily.

name

str

Human-readable provider name.

supported_function

str

Name of the function replaced by this provider.

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

classmethod

Restore the selected Diffusers attention backend.

aitune.torch.kernel_forge.kernel_provider.diffusers_attention_provider.DiffusersAttentionKernelProvider._infer(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attn_mask: torch.Tensor | None = None,
    dropout_p: float = 0.0,
    is_causal: bool = False,
    scale: float | None = None,
    enable_gqa: bool = False
) -> torch.Tensor

Run SDPA through the selected Diffusers attention backend.

aitune.torch.kernel_forge.kernel_provider.diffusers_attention_provider.DiffusersAttentionKernelProvider._load_runtime_dependencies() -> None

Load Diffusers attention dependencies before graph capture.

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

Prepare the provider without running inference.

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

Serialize the selected Diffusers attention backend.

aitune.torch.kernel_forge.kernel_provider.diffusers_attention_provider._import_diffusers_attention_dispatch() -> types.ModuleType | None

Import and cache the optional Diffusers attention dispatcher module.

aitune.torch.kernel_forge.kernel_provider.diffusers_attention_provider._RUNTIME_UNAVAILABLE_MESSAGE = "Diffusers attention dispatcher is not available. Install it with `pip install '...