refactor: rework attention backend resolution

- explicit attn_backend() context wins over ASTR_BACKEND env
- polymorphic available()/supports_call() replace isinstance dispatch
- cache singleton backend instances to avoid hot-path allocation
- training (fwd=None) resolves cuda > flash > torch by capability
- flash dense supports mask-free calls only; masked training falls back to torch
This commit is contained in:
2026-08-23 14:47:02 +08:00
parent 10fec8dca1
commit a29bdfae46
6 changed files with 372 additions and 110 deletions
+3 -1
View File
@@ -85,7 +85,9 @@ AttentionBackend (ABC)
Default priority is cuda > flash > torch. Automatic selection may choose a
compatible fallback for a particular call. Set
`ASTR_BACKEND=cuda|torch_native|flash` to require one backend process-wide.
`ASTR_BACKEND=cuda|torch_native|flash` to override the default process-wide;
an explicit `attn_backend(...)` context still takes precedence over the env
override.
Select via context manager (mirrors `torch.nn.attention.sdpa_kernel`):