File size: 1,740 Bytes
a3afa1d d2b84ea a3afa1d | 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 | """Shared attention backends, masks, and optional optimized kernels."""
from ._core import (
VALID_ATTENTION_BACKENDS,
AttentionBackend,
BlockMask,
_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,
clear_flex_attention_caches,
create_block_mask,
flex_attention,
get_attention_mask,
get_attn_implementation,
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__ = [
"FASTPLMS_ATTENTION_FUNCTIONS",
"FASTPLMS_ATTENTION_MASKS",
"VALID_ATTENTION_BACKENDS",
"AttentionBackend",
"BlockMask",
"FastPLMsAttentionMixin",
"_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",
"clear_flex_attention_caches",
"create_block_mask",
"flex_attention",
"get_attention_mask",
"get_attn_implementation",
"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",
]
|