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
+2 -2
View File
@@ -205,8 +205,8 @@ class Executor:
max_q_heads = config.num_attention_heads
head_dim = config.hidden_size // config.num_attention_heads
backend = get_backend()
self._graph_supported = backend.supports_graph() and CudaBackend.supports(
head_dim=head_dim
self._graph_supported = backend.supports_graph() and (
CudaBackend.available() and head_dim in CudaBackend.HEAD_DIMS
)
self._workspace = InferenceWorkspace(
max_batch_size=kv_cache.max_batch_size,