File size: 2,335 Bytes
8752091 6b7efd8 8752091 6577ccb 8752091 6b7efd8 8752091 6577ccb 8752091 2ca20a1 8752091 6b7efd8 8752091 843a654 8752091 6b7efd8 8752091 6577ccb 8752091 6b7efd8 8752091 6b7efd8 8752091 6577ccb 8752091 2ca20a1 8752091 6b7efd8 8752091 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """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",
]
|