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

astrai/extension/backend/attention.py

@@ -21,9 +21,20 @@ Usage — mirroring ``torch.nn.attention.sdpa_kernel``:
         ...
 
 Thread-safe via ``contextvars`` — each scheduler thread gets its own
-active backend. ``get_backend()`` returns the active one, falling back
-to a process-wide default (cuda > flash > torch, overridable via
-``ASTR_BACKEND``).
+active backend. Backend resolution follows a strict precedence:
+
+1. explicit ``attn_backend(...)`` context (wins over everything),
+2. the process-wide ``ASTR_BACKEND`` environment override,
+3. an implicit default picked from the available backends
+   (cuda > flash > torch).
+
+Capability is polymorphic: every backend declares ``available()``
+(machine-level) and ``supports_call(...)`` (per-call), so adding a new
+backend requires no changes to the resolution logic. Training calls
+(``fwd=None``, no KV cache) resolve through the same priority list: the
+CUDA cache kernels cannot run without a cache, so they fall back to
+flash (when it can handle the call — mask-free/causal only) and finally
+to the reference ``TorchNativeBackend``.
 
 Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
 (blhd). The backend returns ``[batch, seq_len, n_heads * head_dim]``.
@@ -32,11 +43,12 @@ Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
 import contextvars
 import enum
 import functools
+import logging
 import os
 import threading
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -57,8 +69,9 @@ except Exception:
 if TYPE_CHECKING:
     from astrai.inference.cache import KVCache
 
+logger = logging.getLogger(__name__)
+
 
-_default_backend: Optional["AttentionBackend"] = None
 _default_backend_lock = threading.Lock()
 _env_backend_name: Optional[str] = None
 _env_backend: Optional["AttentionBackend"] = None
@@ -66,6 +79,10 @@ _current_backend: contextvars.ContextVar[Optional["AttentionBackend"]] = (
     contextvars.ContextVar("attn_backend", default=None)
 )
 
+# Backends are stateless — one canonical instance per class, created lazily
+# and reused everywhere (resolution, fallback, context managers).
+_singletons: Dict[type, "AttentionBackend"] = {}
+
 
 @functools.lru_cache(maxsize=1)
 def flash_attn_available() -> bool:
@@ -102,58 +119,40 @@ class ATTN_BACKEND(enum.Enum):
     FLASH = "flash"
 
 
-def _priority_backends() -> list["AttentionBackend"]:
-    """Available backends in priority order: cuda -> flash -> torch."""
-    backends: list[AttentionBackend] = []
-    if is_available("attn_paged_decode") and is_available("attn_paged_prefill"):
-        backends.append(CudaBackend())
-    if flash_attn_available():
-        backends.append(FlashAttnBackend())
-    backends.append(TorchNativeBackend())
-    return backends
+def _instance(backend_cls: type) -> "AttentionBackend":
+    """Return the canonical singleton instance for a backend class.
 
-
-def _backend_supports(
-    backend: "AttentionBackend",
-    q: Tensor,
-    kv_cache: Optional["KVCache"],
-    attn_mask: Optional[Tensor],
-    is_causal: bool,
-    fwd: Optional[str],
-) -> bool:
-    """Whether ``backend`` can run this attention call.
-
-    The CUDA kernels are bf16-only, support head_dim in 32/64/128/256, and
-    need a KV cache (decode/prefill); everything else falls back to torch.
+    Backends hold no per-instance state, so a single cached instance is
+    safe and avoids per-call allocation on the attention hot path.
     """
-    if isinstance(backend, CudaBackend):
-        return (
-            fwd in ("prefill", "decode")
-            and kv_cache is not None
-            and q.ndim == 3
-            and q.dtype == torch.bfloat16
-            and q.size(-1) in (32, 64, 128, 256)
-            and is_available(f"attn_paged_{fwd}")
-        )
-    if isinstance(backend, FlashAttnBackend):
-        if not flash_attn_available():
-            return False
-        if q.dtype not in (torch.float16, torch.bfloat16):
-            return False
-        if fwd is not None:
-            return q.ndim == 3 and hasattr(_flash_attn, "flash_attn_varlen_func")
-        if attn_mask is None or is_causal:
-            return True
-        return attn_mask.dim() == 4
-    return True
+    backend = _singletons.get(backend_cls)
+    if backend is None:
+        backend = backend_cls()
+        _singletons[backend_cls] = backend
+    return backend
+
+
+@functools.lru_cache(maxsize=1)
+def _priority_backends() -> Tuple["AttentionBackend", ...]:
+    """Available backends in priority order: cuda -> flash -> torch.
+
+    Computed once (machine availability cannot change at runtime) and
+    cached forever; the tuple always ends with ``TorchNativeBackend``,
+    which is unconditionally available.
+    """
+    return tuple(
+        _instance(cls)
+        for cls in (CudaBackend, FlashAttnBackend, TorchNativeBackend)
+        if cls.available()
+    )
 
 
 def _resolve_default_backend() -> "AttentionBackend":
     """Pick the highest-priority available backend (cuda -> flash -> torch).
 
-    Resolved lazily on first ``get_backend()`` and cached.  Per-call
-    capability fallback happens in ``attention()``, so the default is
-    safe for training and fp32 models.
+    Resolved lazily on first use and cached via ``_priority_backends``.
+    Per-call capability fallback happens in ``attention()``, so the
+    default is safe for training and fp32 models.
     """
     return _priority_backends()[0]
 
@@ -168,9 +167,14 @@ def _environment_backend() -> Optional["AttentionBackend"]:
         with _default_backend_lock:
             if name != _env_backend_name:
                 try:
-                    _env_backend = AttentionBackendFactory.create(name)
+                    _env_backend = _resolve_backend(name)
                 except (ValueError, RuntimeError):
                     _env_backend = None
+                    logger.warning(
+                        "ASTR_BACKEND=%r is not a registered attention backend; "
+                        "falling back to default resolution",
+                        name,
+                    )
                 _env_backend_name = name
     return _env_backend
 
@@ -178,43 +182,46 @@ def _environment_backend() -> Optional["AttentionBackend"]:
 def _resolve_backend(
     backend: Optional[Union[str, ATTN_BACKEND, "AttentionBackend", type]] = None,
 ) -> "AttentionBackend":
-    """Resolve a backend configuration, defaulting to the process policy."""
+    """Resolve a backend configuration to its canonical instance.
+
+    Accepts a registered name, ``ATTN_BACKEND`` enum value, backend class,
+    or instance.  Names/classes resolve to the shared singleton; a caller
+    may still pass its own instance to opt out of sharing.
+    """
     if backend is not None:
         if isinstance(backend, ATTN_BACKEND):
-            return AttentionBackendFactory.create(backend.value)
+            return _instance(AttentionBackendFactory.get_component_class(backend.value))
         if isinstance(backend, str):
-            return AttentionBackendFactory.create(backend)
+            return _instance(AttentionBackendFactory.get_component_class(backend))
         if isinstance(backend, type) and issubclass(backend, AttentionBackend):
-            return backend()
+            return _instance(backend)
         if isinstance(backend, AttentionBackend):
             return backend
         raise TypeError(
             f"expected a registered name, ATTN_BACKEND, AttentionBackend type, "
             f"or instance, got {type(backend).__name__}"
         )
-
-    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
+    return _resolve_default_backend()
 
 
 def get_backend(
     use_default: bool = True,
 ) -> Optional["AttentionBackend"]:
-    """Return the context override, optionally falling back to the process default.
+    """Resolve the active backend: explicit context > env > default.
 
-    ``ASTR_BACKEND`` is a process-wide override and takes precedence over the
-    context value. Pass ``use_default=False`` at request submission to retain
-    only an environment override or the caller's :func:`attn_backend` value.
+    An ``attn_backend(...)`` context is the caller's explicit choice and
+    always wins.  ``ASTR_BACKEND`` is a process-wide override consulted
+    only when no context is set.  Pass ``use_default=False`` at request
+    submission to retain only an environment override or the caller's
+    :func:`attn_backend` value.
     """
-    return (
-        _environment_backend()
-        or _current_backend.get()
-        or (_resolve_backend() if use_default else None)
-    )
+    context_backend = _current_backend.get()
+    if context_backend is not None:
+        return context_backend
+    env_backend = _environment_backend()
+    if env_backend is not None:
+        return env_backend
+    return _resolve_default_backend() if use_default else None
 
 
 @contextmanager
@@ -262,13 +269,23 @@ def attention(
     attn_mask: Optional[Tensor] = None,
     is_causal: bool = False,
     fwd: Optional[str] = None,
+    backend: Optional[Union[str, ATTN_BACKEND, "AttentionBackend", type]] = None,
 ) -> Tensor:
     """Functional attention entry point — mirrors ``F.scaled_dot_product_attention``.
 
-    Delegates to the active backend (set via ``with attn_backend(...)``).
+    Delegates to the active backend.  ``backend`` (optional) is an explicit
+    escape hatch; when omitted the backend is resolved as
+    explicit context > ``ASTR_BACKEND`` env > default (cuda > flash > torch).
     Handles KV cache I/O, GQA head expansion, and causal masking so the
     caller only needs to provide projected q/k/v.
 
+    Training calls (``fwd=None``, ``kv_cache=None``) resolve through the
+    same capability chain — the CUDA cache kernels cannot run without a
+    cache, so they fall back to flash (mask-free/causal calls only) and
+    finally to torch SDPA.  An explicitly-selected backend that cannot
+    handle the call raises — an implicit one falls back down the priority
+    list to the first capable backend.
+
     Args:
         q: [batch, q_len, n_heads, head_dim] (blhd)
         k: [batch, q_len, n_kv_heads, head_dim] (blhd)
@@ -277,30 +294,43 @@ def attention(
         layer_id: transformer layer index for buffer access.
         attn_mask: pre-built attention mask (SDPA-compatible).
         is_causal: whether to apply causal masking.
+        fwd: "prefill" / "decode" for inference, None for training.
+        backend: optional explicit backend (name, enum, class, or instance).
 
     Returns:
         [batch, q_len, n_heads * head_dim]
     """
-    explicit = get_backend(use_default=False)
-    backend = get_backend()
-    if fwd is None and explicit is None:
-        backend = TorchNativeBackend()
-    if not _backend_supports(backend, q, kv_cache, attn_mask, is_causal, fwd):
-        if explicit is not None:
+    if backend is not None:
+        selected = _resolve_backend(backend)
+        explicit = True
+    else:
+        context_backend = _current_backend.get()
+        explicit = context_backend is not None
+        # Resolve through the same chain as inference: explicit context >
+        # ASTR_BACKEND env > default.  Training calls (fwd=None, no cache)
+        # land on the CUDA backend and fall back by capability below —
+        # flash when it can handle the call, else torch SDPA.
+        selected = get_backend()
+        assert selected is not None
+
+    if not selected.supports_call(q, kv_cache, attn_mask, is_causal, fwd):
+        if explicit:
             raise RuntimeError(
-                f"Explicitly-set backend {type(backend).__name__} cannot "
+                f"Explicitly-set backend {type(selected).__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
-            if _backend_supports(candidate, q, kv_cache, attn_mask, is_causal, fwd):
-                backend = candidate
-                break
-    return backend.forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal, fwd)
+        selected = next(
+            (
+                candidate
+                for candidate in _priority_backends()
+                if candidate.supports_call(q, kv_cache, attn_mask, is_causal, fwd)
+            ),
+            _instance(TorchNativeBackend),
+        )
+    return selected.forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal, fwd)
 
 
 class AttentionBackend(ABC):
@@ -310,6 +340,17 @@ class AttentionBackend(ABC):
     ``fwd_prefill`` (q_len > 1, with or without cache). The public
     ``forward`` method dispatches based on q_len.
 
+    Capability contract — every backend declares:
+
+    * ``available()`` — machine-level: can this backend exist here
+      (kernel ``.so`` loaded, flash-attn present, GPU available)?
+      Used once to build the default priority list.
+    * ``supports_call(q, kv_cache, attn_mask, is_causal, fwd)`` — can this
+      backend run this *specific* call (shape/dtype/cache/mask)?  Used by
+      ``attention()`` for the per-call fallback.  Resolution logic never
+      checks concrete backend types, so adding a backend requires no
+      changes outside its own class.
+
     Three equivalent ways to activate a backend::
 
         with attn_backend(ATTN_BACKEND.TORCH_NATIVE):  # enum
@@ -327,6 +368,30 @@ class AttentionBackend(ABC):
     def __exit__(self, *exc) -> None:
         _current_backend.reset(self._token)
 
+    @classmethod
+    @abstractmethod
+    def available(cls) -> bool:
+        """Return True if this backend can run on the current machine.
+
+        Checks static availability only (compiled kernels, optional
+        packages, GPU presence) — not call-specific constraints.
+        """
+
+    @abstractmethod
+    def supports_call(
+        self,
+        q: Tensor,
+        kv_cache: Optional["KVCache"],
+        attn_mask: Optional[Tensor],
+        is_causal: bool,
+        fwd: Optional[str],
+    ) -> bool:
+        """Return True if this backend can run this specific attention call.
+
+        Called on the canonical singleton instance (or a caller-provided
+        one); must be side-effect free.
+        """
+
     def forward(
         self,
         q: Tensor,
@@ -412,8 +477,18 @@ class TorchNativeBackend(AttentionBackend):
     runs SDPA directly on the projected q/k/v.
     """
 
-    @staticmethod
-    def supports(**kwargs) -> bool:
+    @classmethod
+    def available(cls) -> bool:
+        return True
+
+    def supports_call(
+        self,
+        q: Tensor,
+        kv_cache: Optional["KVCache"],
+        attn_mask: Optional[Tensor],
+        is_causal: bool,
+        fwd: Optional[str],
+    ) -> bool:
         return True
 
     def fwd_decode(
@@ -516,16 +591,37 @@ class CudaBackend(AttentionBackend):
     Raises ``RuntimeError`` if the required kernel is not available.
     """
 
-    @staticmethod
-    def supports(**kwargs) -> bool:
-        head_dim = kwargs.get("head_dim", -1)
+    # Head dims supported by the CUDA kernels (single source of truth).
+    HEAD_DIMS = (32, 64, 128, 256)
+
+    @classmethod
+    def available(cls) -> bool:
         return (
             torch.cuda.is_available()
-            and head_dim in (32, 64, 128, 256)
             and is_available("attn_paged_decode")
             and is_available("attn_paged_prefill")
         )
 
+    def supports_call(
+        self,
+        q: Tensor,
+        kv_cache: Optional["KVCache"],
+        attn_mask: Optional[Tensor],
+        is_causal: bool,
+        fwd: Optional[str],
+    ) -> bool:
+        # The CUDA kernels are bf16-only, support head_dim in
+        # HEAD_DIMS, and need a KV cache (decode/prefill); everything
+        # else falls back down the priority list to torch.
+        return (
+            fwd in ("prefill", "decode")
+            and kv_cache is not None
+            and q.ndim == 3
+            and q.dtype == torch.bfloat16
+            and q.size(-1) in self.HEAD_DIMS
+            and is_available(f"attn_paged_{fwd}")
+        )
+
     @staticmethod
     def supports_graph() -> bool:
         return True
@@ -606,10 +702,30 @@ class FlashAttnBackend(AttentionBackend):
     ``flash_attn_func``.
     """
 
-    @staticmethod
-    def supports(**kwargs) -> bool:
+    @classmethod
+    def available(cls) -> bool:
         return flash_attn_available()
 
+    def supports_call(
+        self,
+        q: Tensor,
+        kv_cache: Optional["KVCache"],
+        attn_mask: Optional[Tensor],
+        is_causal: bool,
+        fwd: Optional[str],
+    ) -> bool:
+        if not self.available():
+            return False
+        if q.dtype not in (torch.float16, torch.bfloat16):
+            return False
+        if fwd is not None:
+            return q.ndim == 3 and hasattr(_flash_attn, "flash_attn_varlen_func")
+        # Dense (training) path: flash_attn_func cannot apply a custom
+        # mask, so only mask-free calls are supported — ``is_causal`` is
+        # a flag, not a mask.  Masked training (SFT/DPO/GRPO) must fall
+        # back to TorchNativeBackend instead of silently ignoring the mask.
+        return attn_mask is None
+
     def fwd_decode(
         self,
         q: Tensor,
@@ -649,9 +765,9 @@ class FlashAttnBackend(AttentionBackend):
             k = repeat_kv(k, n_rep)
             v = repeat_kv(v, n_rep)
 
-        if attn_mask is not None and not is_causal and attn_mask.dim() != 4:
+        if attn_mask is not None:
             raise ValueError(
-                "FlashAttnBackend does not support a custom attention mask; "
+                "FlashAttnBackend cannot handle a custom attention mask; "
                 "use a causal mask or select TorchNativeBackend."
             )
         fa = _flash_attn
@@ -664,7 +780,7 @@ class FlashAttnBackend(AttentionBackend):
             q.contiguous(),
             k.contiguous(),
             v.contiguous(),
-            causal=is_causal or (attn_mask is not None and attn_mask.dim() == 4),
+            causal=is_causal,
         )
         return out.contiguous()
 

astrai/inference/runtime/executor.py

@@ -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,

docs/developer/internals.md

@@ -188,7 +188,7 @@ Attention computation is decoupled from the model via `AttentionBackend` ABC (`a
 - **`FlashAttnBackend`**: optional flash-attn dispatch with `flash_attn_with_kvcache` fast path for contiguous cache; falls back to KV gather + `flash_attn_func`.
 - **`TorchNativeBackend`** (always-available fallback): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
 - The `attention(...)` entry point uses cuda > flash > torch priority and chooses another compatible backend when an automatically selected backend cannot handle a call.
-- `ASTR_BACKEND=cuda|torch_native|flash` and `attn_backend(...)` are explicit selections; incompatible calls raise instead of silently changing backend.
+- Resolution precedence is: explicit `attn_backend(...)` context > `ASTR_BACKEND` env > default. An explicit `attn_backend(...)` selection is strict (incompatible calls raise); `ASTR_BACKEND` is a default-level override that falls back to a compatible backend when incapable. Training calls (`fwd=None`, no KV cache) resolve by capability: the CUDA cache kernels cannot run without a cache, so they fall back to flash (mask-free/causal calls only) and finally to torch SDPA.
 
 Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/backend/rotary.py`, which auto-dispatches to the fused CUDA kernel (`rotary_emb.cu`) during inference or torch complex multiply during training (for autograd compatibility). Both attention backends share the same rotary dispatch.
 

docs/guides/inference.md

@@ -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`):
 

tests/extension/test_backend.py

@@ -2,20 +2,33 @@
 
 These tests do not require CUDA — they only check that the active
 backend is correctly set and restored.
+
+Resolution precedence under test: explicit ``attn_backend(...)``
+context > ``ASTR_BACKEND`` env override > implicit default.  Training
+calls (``fwd=None``, no KV cache) resolve by capability: the CUDA cache
+kernels cannot run without a cache, so they fall back to flash (mask-free
+calls only) and finally to torch SDPA.
 """
 
+import importlib
+
 import pytest
+import torch
 
 from astrai.extension import (
     ATTN_BACKEND,
+    AttentionBackend,
     AttentionBackendFactory,
     CudaBackend,
     FlashAttnBackend,
     TorchNativeBackend,
+    attention,
     attn_backend,
     get_backend,
 )
 
+_attn_module = importlib.import_module("astrai.extension.backend.attention")
+
 
 def test_default_backend_resolves_to_available():
     """Default backend is the first available in cuda > flash > torch order."""
@@ -23,6 +36,10 @@ def test_default_backend_resolves_to_available():
     assert isinstance(backend, (CudaBackend, FlashAttnBackend, TorchNativeBackend))
 
 
+def test_default_backend_is_cached_singleton():
+    assert get_backend() is get_backend()
+
+
 def test_attn_backend_context_with_enum():
     default = get_backend()
     with attn_backend(ATTN_BACKEND.CUDA):
@@ -44,11 +61,132 @@ def test_backend_can_read_only_context_selection():
     assert get_backend(use_default=False) is None
 
 
-def test_environment_backend_overrides_context(monkeypatch):
+def test_context_beats_environment_backend(monkeypatch):
+    """An explicit attn_backend() context wins over ASTR_BACKEND."""
     monkeypatch.setenv("ASTR_BACKEND", "torch_native")
     with attn_backend("cuda"):
-        assert type(get_backend()).__name__ == "TorchNativeBackend"
-        assert type(get_backend(use_default=False)).__name__ == "TorchNativeBackend"
+        assert isinstance(get_backend(), CudaBackend)
+        assert isinstance(get_backend(use_default=False), CudaBackend)
+
+
+def test_environment_backend_used_without_context(monkeypatch):
+    monkeypatch.setenv("ASTR_BACKEND", "torch_native")
+    assert isinstance(get_backend(), TorchNativeBackend)
+    assert isinstance(get_backend(use_default=False), TorchNativeBackend)
+
+
+def test_environment_backend_does_not_break_training(monkeypatch):
+    """Training (fwd=None, no cache) must not steer onto cache-only kernels.
+
+    Regression: with ASTR_BACKEND=cuda, a training forward used to raise
+    because the env override was treated as an explicit selection.
+    """
+    monkeypatch.setenv("ASTR_BACKEND", "cuda")
+    q = torch.zeros(1, 2, 4, 8, dtype=torch.bfloat16)
+    out = attention(q, q, q)
+    assert out.shape == q.shape
+
+
+def test_explicit_backend_mismatch_raises(monkeypatch):
+    monkeypatch.delenv("ASTR_BACKEND", raising=False)
+    q = torch.zeros(1, 2, 4, 8, dtype=torch.bfloat16)
+    with pytest.raises(RuntimeError, match="Explicitly-set backend"):
+        with attn_backend("cuda"):
+            attention(q, q, q)  # cuda + no KV cache -> cannot handle
+
+
+def test_implicit_backend_falls_back_when_incapable(monkeypatch):
+    """An implicit (env) backend that cannot run the call falls back."""
+    monkeypatch.setenv("ASTR_BACKEND", "cuda")
+    q = torch.zeros(1, 2, 4, 8, dtype=torch.float32)  # fp32: cuda kernels can't
+    out = attention(q, q, q, fwd="prefill", is_causal=True)
+    assert out.shape == q.shape
+
+
+def _flash_available(monkeypatch) -> None:
+    """Pretend flash-attn is usable and rebuild the priority list."""
+    monkeypatch.setattr(_attn_module, "flash_attn_available", lambda: True)
+    _attn_module._priority_backends.cache_clear()
+
+
+def test_training_falls_back_to_flash_before_torch_when_capable(monkeypatch):
+    """Training (no cache) prefers flash over torch when flash can run the call."""
+    _flash_available(monkeypatch)
+    try:
+        prio = _attn_module._priority_backends()
+        names = [type(b).__name__ for b in prio]
+        assert "FlashAttnBackend" in names
+        assert names.index("FlashAttnBackend") < names.index("TorchNativeBackend")
+
+        q = torch.zeros(1, 2, 4, 8, dtype=torch.bfloat16)
+        # Mask-free training call resolves to flash, not torch.
+        resolved = next(b for b in prio if b.supports_call(q, None, None, False, None))
+        assert isinstance(resolved, FlashAttnBackend)
+    finally:
+        _attn_module._priority_backends.cache_clear()
+
+
+def test_flash_dense_supports_only_mask_free_calls(monkeypatch):
+    """FlashAttnBackend cannot apply custom masks in the dense path."""
+    _flash_available(monkeypatch)
+    flash = _attn_module._instance(_attn_module.FlashAttnBackend)
+    q = torch.zeros(1, 2, 4, 8, dtype=torch.bfloat16)
+    mask_4d = torch.zeros(1, 1, 2, 2, dtype=torch.bool)
+
+    assert flash.supports_call(q, None, None, False, None) is True
+    assert flash.supports_call(q, None, None, True, None) is True
+    assert flash.supports_call(q, None, mask_4d, False, None) is False
+
+
+def test_flash_dense_rejects_custom_mask(monkeypatch):
+    """A masked dense call must fail loudly, never silently ignore the mask."""
+    _flash_available(monkeypatch)
+    flash = _attn_module._instance(_attn_module.FlashAttnBackend)
+    q = torch.zeros(1, 2, 4, 8, dtype=torch.bfloat16)
+    mask_4d = torch.zeros(1, 1, 2, 2, dtype=torch.bool)
+
+    with pytest.raises(ValueError, match="custom attention mask"):
+        flash._forward_dense(q, q, q, attn_mask=mask_4d, is_causal=False)
+
+
+def test_backend_resolution_returns_shared_singletons():
+    with attn_backend("cuda") as first:
+        pass
+    with attn_backend("cuda") as second:
+        assert first is second
+
+
+class _DummyBackend(AttentionBackend):
+    """Minimal backend used only to prove capability is polymorphic."""
+
+    @classmethod
+    def available(cls) -> bool:
+        return True
+
+    def supports_call(self, q, kv_cache, attn_mask, is_causal, fwd) -> bool:
+        return True
+
+    def fwd_decode(
+        self, q, k, v, kv_cache=None, layer_id=0, attn_mask=None, is_causal=False
+    ):
+        return q
+
+    def fwd_prefill(
+        self, q, k, v, kv_cache=None, layer_id=0, attn_mask=None, is_causal=False
+    ):
+        return q
+
+
+def test_custom_backend_usable_without_touching_resolution():
+    """A third-party backend plugs in via context or explicit param."""
+    custom = _DummyBackend()
+    q = torch.zeros(1, 2, 4, 8)
+
+    with attn_backend(custom):
+        assert get_backend() is custom
+
+    out = attention(q, q, q, backend=custom)
+    assert out is q
 
 
 def test_attention_backend_factory_lists_builtin_backends():

tests/extension/test_backend_equivalence.py

@@ -34,10 +34,13 @@ def _ws(pool: PagePool) -> InferenceWorkspace:
 
 @skip_no_kernel
 def test_training_forward_matches_torch(cuda_model):
-    """Training forward (kv_cache=None) uses torch-native SDPA.
+    """Training forward (kv_cache=None) resolves to a capable dense backend.
 
-    CudaBackend does not support training (requires kv_cache).
-    Torch-native backend must match default (which falls back to torch).
+    CudaBackend cannot run training (requires a KV cache), so the default
+    falls back by capability — flash when it can handle the call
+    (mask-free/causal), otherwise torch SDPA.  The default path must not
+    raise and must produce finite logits; explicitly selected torch SDPA
+    must be deterministic across runs.
     """
 
     model, _ = cuda_model
@@ -48,12 +51,15 @@ def test_training_forward_matches_torch(cuda_model):
 
     with attn_backend(ATTN_BACKEND.TORCH_NATIVE):
         with torch.no_grad():
-            out_torch = model(input_ids)
+            out_torch_a = model(input_ids)
+        with torch.no_grad():
+            out_torch_b = model(input_ids)
 
+    assert out_default["logits"].shape == out_torch_a["logits"].shape
+    assert torch.isfinite(out_default["logits"]).all()
     torch.testing.assert_close(
-        out_torch["logits"], out_default["logits"], atol=1e-6, rtol=1e-6
+        out_torch_a["logits"], out_torch_b["logits"], atol=0, rtol=0
     )
-    assert out_default["logits"].shape[0] == 2
 
 
 @skip_no_kernel