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:
@@ -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 = ["."]
|
||||
|
||||
Reference in New Issue
Block a user