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:
@@ -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`):
|
||||
|
||||
|
||||
Reference in New Issue
Block a user