| """Shared attention backends, masks, and optional optimized kernels.""" |
|
|
| from ._auto import ( |
| AUTO_ATTENTION, |
| AttentionCandidate, |
| AttentionExecutionContext, |
| AttentionResolution, |
| ) |
| from ._core import ( |
| LEGACY_CHECKPOINT_ATTENTION_BACKENDS, |
| VALID_ATTENTION_BACKENDS, |
| AttentionBackend, |
| BlockMask, |
| FlashPaddingLayout, |
| _ensure_flash_kernels_loaded, |
| _get_flex_attention_fn, |
| _get_flex_block_mask, |
| _kernels_flash_forward, |
| _kernels_flash_varlen_forward, |
| _unpad_input, |
| bool_to_additive_mask, |
| canonical_checkpoint_attention_backend, |
| clear_flex_attention_caches, |
| create_block_mask, |
| flash_kernel_unsupported_reason, |
| flex_attention, |
| get_attention_mask, |
| get_attn_implementation, |
| get_flash_padding_layout, |
| index_first_axis, |
| index_put_first_axis, |
| kernels_flash_attention_func, |
| pad_input, |
| resolve_attention_backend, |
| resolve_attention_backend_for_call, |
| set_config_attn_implementation, |
| warn_attention_backend_fallback, |
| ) |
| from .interfaces import ( |
| FASTPLMS_ATTENTION_FUNCTIONS, |
| FASTPLMS_ATTENTION_MASKS, |
| FastPLMsAttentionMixin, |
| validate_transformers_attention_interfaces, |
| ) |
|
|
|
|
| __all__ = [ |
| "AUTO_ATTENTION", |
| "FASTPLMS_ATTENTION_FUNCTIONS", |
| "FASTPLMS_ATTENTION_MASKS", |
| "LEGACY_CHECKPOINT_ATTENTION_BACKENDS", |
| "VALID_ATTENTION_BACKENDS", |
| "AttentionBackend", |
| "AttentionCandidate", |
| "AttentionExecutionContext", |
| "AttentionResolution", |
| "BlockMask", |
| "FastPLMsAttentionMixin", |
| "FlashPaddingLayout", |
| "_ensure_flash_kernels_loaded", |
| "_get_flex_attention_fn", |
| "_get_flex_block_mask", |
| "_kernels_flash_forward", |
| "_kernels_flash_varlen_forward", |
| "_unpad_input", |
| "bool_to_additive_mask", |
| "canonical_checkpoint_attention_backend", |
| "clear_flex_attention_caches", |
| "create_block_mask", |
| "flash_kernel_unsupported_reason", |
| "flex_attention", |
| "get_attention_mask", |
| "get_attn_implementation", |
| "get_flash_padding_layout", |
| "index_first_axis", |
| "index_put_first_axis", |
| "kernels_flash_attention_func", |
| "pad_input", |
| "resolve_attention_backend", |
| "resolve_attention_backend_for_call", |
| "set_config_attn_implementation", |
| "validate_transformers_attention_interfaces", |
| "warn_attention_backend_fallback", |
| ] |
|
|
