fix: default backend race, raise on explicit fallback

- _default_backend lazy init protected with threading.Lock
- attention() raises when explicit backend cannot handle call
- FlashAttnBackend rejects prefill with non-None attn_mask
- training test uses TORCH_NATIVE backend directly

astrai/extension/attention_backend.py

@@ -34,6 +34,7 @@ import enum
 import functools
 import importlib
 import os
+import threading
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Optional, Union
@@ -101,6 +102,7 @@ class ATTN_BACKEND(enum.Enum):
 
 
 _default_backend: Optional["AttentionBackend"] = None
+_default_backend_lock = threading.Lock()
 
 
 def _priority_backends() -> list["AttentionBackend"]:
@@ -139,7 +141,7 @@ def _backend_supports(
             return False
         if q.size(1) == 1 and kv_cache is not None:
             return True
-        return not (attn_mask is not None and not is_causal)
+        return attn_mask is None
     return True
 
 
@@ -174,6 +176,8 @@ def get_backend() -> "AttentionBackend":
         return _current_backend.get()
     except LookupError:
         global _default_backend
+        if _default_backend is None:
+            with _default_backend_lock:
                 if _default_backend is None:
                     _default_backend = _resolve_default_backend()
         return _default_backend
@@ -279,9 +283,18 @@ def attention(
     """
     backend = get_backend()
     if not _backend_supports(backend, q, kv_cache, attn_mask, is_causal):
-        # The active backend cannot run this call (e.g. CUDA on a training /
+        try:
-        # fp32 / unsupported-head_dim input) — fall back to the highest-
+            explicit = _current_backend.get()
-        # priority backend that can, ending at torch SDPA.
+        except LookupError:
+            explicit = None
+        if explicit is not None:
+            raise RuntimeError(
+                f"Explicitly-set backend {type(backend).__name__} cannot "
+                f"handle this attention call (shape={q.shape}, "
+                f"dtype={q.dtype}, kv_cache={'none' if kv_cache is None else 'present'}, "
+                f"attn_mask={'none' if attn_mask is None else 'present'}). "
+                f"Remove the attn_backend() context or switch to a compatible backend."
+            )
         for candidate in _priority_backends():
             if isinstance(candidate, type(backend)):
                 continue

tests/extension/test_backend_equivalence.py

@@ -25,27 +25,26 @@ def _ws(pool: PagePool) -> InferenceWorkspace:
 
 @skip_no_kernel
 def test_training_forward_matches_torch(cuda_model):
-    """Training forward (kv_cache=None) should produce identical logits.
+    """Training forward (kv_cache=None) uses torch-native SDPA.
 
-    CudaBackend is now safe as a default: for training (``kv_cache=None``)
+    CudaBackend does not support training (requires kv_cache).
-    or non-bf16 inputs it falls back to torch SDPA.  Verify the fallback
+    Torch-native backend must match default (which falls back to torch).
-    path matches the torch-native forward exactly.
     """
 
     model, _ = cuda_model
     input_ids = torch.randint(0, 1000, (2, 16), device="cuda")
 
+    with torch.no_grad():
+        out_default = model(input_ids)
+
+    with attn_backend(ATTN_BACKEND.TORCH_NATIVE):
         with torch.no_grad():
             out_torch = model(input_ids)
 
-    with attn_backend(ATTN_BACKEND.CUDA):
-        with torch.no_grad():
-            out_cuda = model(input_ids)
-
     torch.testing.assert_close(
-        out_cuda["logits"], out_torch["logits"], atol=1e-6, rtol=1e-6
+        out_torch["logits"], out_default["logits"], atol=1e-6, rtol=1e-6
     )
-    assert out_torch["logits"].shape[0] == 2
+    assert out_default["logits"].shape[0] == 2
 
 
 @skip_no_kernel