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
+1
View File
@@ -32,6 +32,7 @@ urls = { Homepage = "https://github.com/ViperEkura/AstrAI" }
[project.optional-dependencies]
dev = ["pytest==9.0.2", "ruff", "httpx2"]
flash = ["flash-attn>=2.6"]
[tool.setuptools.packages.find]
where = ["."]