feat: add optional FlashAttention (FA2/FA3) backend
- add FlashAttnBackend (ATTN_BACKEND.FLASH) using flash_attn_func with KV-cache gather + GQA, mirroring TorchNativeBackend - add flash_attn_available() probe gated on compute capability plus a real-kernel smoke test, cached at first use - lazy-import flash-attn via importlib so it stays an optional dependency, raising clear errors when unusable - add 'flash' optional extra (flash-attn>=2.6) and export the new backend
This commit is contained in:
@@ -19,6 +19,7 @@ from astrai.extension.attention_backend import (
|
||||
ATTN_BACKEND,
|
||||
AttentionBackend,
|
||||
CudaBackend,
|
||||
FlashAttnBackend,
|
||||
TorchNativeBackend,
|
||||
attention,
|
||||
attn_backend,
|
||||
@@ -38,6 +39,7 @@ __all__ = [
|
||||
"AttentionBackend",
|
||||
"CudaBackend",
|
||||
"TorchNativeBackend",
|
||||
"FlashAttnBackend",
|
||||
"TensorLayout",
|
||||
"attention",
|
||||
"attn_backend",
|
||||
|
||||
Reference in New Issue
Block a user