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:
2026-08-05 15:27:26 +08:00
parent 2667b8116d
commit 8c052c99ee
3 changed files with 173 additions and 5 deletions
+2
View File
@@ -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",