From a29bdfae464abd2b7c96d5e15c1711d9abdc32c1 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sun, 23 Aug 2026 14:47:02 +0800 Subject: [PATCH] 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 | 310 ++++++++++++++------ astrai/inference/runtime/executor.py | 4 +- docs/developer/internals.md | 2 +- docs/guides/inference.md | 4 +- tests/extension/test_backend.py | 144 ++++++++- tests/extension/test_backend_equivalence.py | 18 +- 6 files changed, 372 insertions(+), 110 deletions(-) diff --git a/astrai/extension/backend/attention.py b/astrai/extension/backend/attention.py index ee1a8f8..82ef03f 100644 --- a/astrai/extension/backend/attention.py +++ b/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() diff --git a/astrai/inference/runtime/executor.py b/astrai/inference/runtime/executor.py index e8c0584..7d4b611 100644 --- a/astrai/inference/runtime/executor.py +++ b/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, diff --git a/docs/developer/internals.md b/docs/developer/internals.md index 2a8ff74..584f183 100644 --- a/docs/developer/internals.md +++ b/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. diff --git a/docs/guides/inference.md b/docs/guides/inference.md index ea335b3..1add499 100644 --- a/docs/guides/inference.md +++ b/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`): diff --git a/tests/extension/test_backend.py b/tests/extension/test_backend.py index 56c7e92..1f4a55a 100644 --- a/tests/extension/test_backend.py +++ b/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(): diff --git a/tests/extension/test_backend_equivalence.py b/tests/extension/test_backend_equivalence.py index 73b383c..e860880 100644 --- a/tests/extension/test_backend_equivalence.py +++ b/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