diff --git a/astrai/extension/attention_backend.py b/astrai/extension/attention_backend.py index 3dfe99a..c31417a 100644 --- a/astrai/extension/attention_backend.py +++ b/astrai/extension/attention_backend.py @@ -44,6 +44,7 @@ from astrai.extension.attention_ops import ( attn_paged_decode, attn_paged_prefill, ) +from astrai.extension.loader import is_available from astrai.factory import BaseFactory if TYPE_CHECKING: @@ -136,15 +137,65 @@ class ATTN_BACKEND(enum.Enum): FLASH = "flash" +_default_backend: Optional["AttentionBackend"] = None + + +def _priority_backends() -> list["AttentionBackend"]: + """Available backends in priority order: flash -> cuda -> torch.""" + backends: list[AttentionBackend] = [] + if flash_attn_available(): + backends.append(FlashAttnBackend()) + if is_available("attn_paged_decode") and is_available("attn_paged_prefill"): + backends.append(CudaBackend()) + backends.append(TorchNativeBackend()) + return backends + + +def _backend_supports( + backend: "AttentionBackend", + q: Tensor, + kv_cache: Optional["KVCache"], + attn_mask: Optional[Tensor], + is_causal: bool, +) -> 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. + """ + if isinstance(backend, CudaBackend): + return ( + kv_cache is not None + and q.dtype == torch.bfloat16 + and q.size(-1) in (32, 64, 128, 256) + ) + if isinstance(backend, FlashAttnBackend): + return flash_attn_available() and not (attn_mask is not None and not is_causal) + return True + + +def _resolve_default_backend() -> "AttentionBackend": + """Pick the highest-priority available backend: flash -> cuda -> torch. + + Resolved lazily on first ``get_backend()`` (flash/cuda availability is + checked once and cached). Per-call capability fallback happens in + ``attention()``, so this default is safe for training and fp32 models. + """ + return _priority_backends()[0] + + def get_backend() -> "AttentionBackend": """Return the active backend for the current thread/context. - Falls back to a ``TorchNativeBackend`` singleton when no backend - has been activated via ``with``. + Falls back to the highest-priority available backend (flash -> cuda -> + torch_native) when no backend has been activated via ``with``. """ try: return _current_backend.get() except LookupError: + global _default_backend + if _default_backend is None: + _default_backend = _resolve_default_backend() return _default_backend @@ -225,6 +276,16 @@ def attention( [batch, q_len, n_heads * head_dim] """ 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 / + # fp32 / unsupported-head_dim input) — fall back to the highest- + # priority backend that can, ending at torch SDPA. + for candidate in _priority_backends(): + if isinstance(candidate, type(backend)): + continue + if _backend_supports(candidate, q, kv_cache, attn_mask, is_causal): + backend = candidate + break return backend.forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal) @@ -393,7 +454,7 @@ class TorchNativeBackend(AttentionBackend): return out -_default_backend = TorchNativeBackend() +_default_backend = None @AttentionBackendFactory.register(ATTN_BACKEND.CUDA.value) @@ -407,8 +468,9 @@ class CudaBackend(AttentionBackend): ``attn_paged_prefill`` with ragged-batch support via qo_indptr + kv_indptr. - ``kv_cache is None`` (training) is not handled — use - ``TorchNativeBackend`` for training. + ``kv_cache is None`` (training) raises — the per-call fallback to + torch SDPA for training / fp32 / unsupported head_dim happens in the + ``attention()`` entry point. Raises ``RuntimeError`` if the required kernel is not available. """ @@ -426,8 +488,9 @@ class CudaBackend(AttentionBackend): if kv_cache is None: raise RuntimeError("CudaBackend does not support training (kv_cache=None)") - kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k - kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v + loc = kv_cache.out_cache_loc[:, 0] + kv_cache.k_buffer[layer_id].index_copy_(0, loc, k[:, 0]) + kv_cache.v_buffer[layer_id].index_copy_(0, loc, v[:, 0]) b = q.size(0) q_3d = q.squeeze(1) @@ -461,8 +524,13 @@ class CudaBackend(AttentionBackend): if kv_cache is None: raise RuntimeError("CudaBackend does not support training (kv_cache=None)") - kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k - kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v + loc = kv_cache.out_cache_loc.reshape(-1) + kv_cache.k_buffer[layer_id].index_copy_( + 0, loc, k.reshape(-1, k.size(2), k.size(3)) + ) + kv_cache.v_buffer[layer_id].index_copy_( + 0, loc, v.reshape(-1, v.size(2), v.size(3)) + ) b = q.size(0) q_len = q.size(1) diff --git a/tests/extension/test_backend.py b/tests/extension/test_backend.py index b010701..1dabe7b 100644 --- a/tests/extension/test_backend.py +++ b/tests/extension/test_backend.py @@ -17,20 +17,30 @@ from astrai.extension import ( def test_default_backend_is_torch_native(): + """Default is the highest-priority available backend (flash > cuda > torch).""" + from astrai.extension.attention_backend import ( + CudaBackend, + TorchNativeBackend, + _resolve_default_backend, + ) + backend = get_backend() - assert isinstance(backend, TorchNativeBackend) + assert isinstance(backend, (CudaBackend, TorchNativeBackend)) + assert isinstance(backend, type(_resolve_default_backend())) def test_attn_backend_context_with_enum(): + default = get_backend() with attn_backend(ATTN_BACKEND.CUDA): assert isinstance(get_backend(), CudaBackend) - assert isinstance(get_backend(), TorchNativeBackend) + assert get_backend() is default def test_attn_backend_context_with_registered_name(): + default = get_backend() with attn_backend("cuda"): assert isinstance(get_backend(), CudaBackend) - assert isinstance(get_backend(), TorchNativeBackend) + assert get_backend() is default def test_attention_backend_factory_lists_builtin_backends(): @@ -48,19 +58,22 @@ def test_attn_backend_rejects_unknown_registered_name(): def test_attn_backend_context_with_class(): + default = get_backend() with attn_backend(CudaBackend): assert isinstance(get_backend(), CudaBackend) - assert isinstance(get_backend(), TorchNativeBackend) + assert get_backend() is default def test_attn_backend_context_with_instance(): custom = CudaBackend() + default = get_backend() with attn_backend(custom): assert get_backend() is custom - assert isinstance(get_backend(), TorchNativeBackend) + assert get_backend() is default def test_cudabackend_is_context_manager(): + default = get_backend() with CudaBackend(): assert isinstance(get_backend(), CudaBackend) - assert isinstance(get_backend(), TorchNativeBackend) + assert get_backend() is default diff --git a/tests/extension/test_backend_equivalence.py b/tests/extension/test_backend_equivalence.py index 981889f..9f9e59e 100644 --- a/tests/extension/test_backend_equivalence.py +++ b/tests/extension/test_backend_equivalence.py @@ -27,9 +27,9 @@ def _ws(pool: PagePool) -> InferenceWorkspace: def test_training_forward_matches_torch(cuda_model): """Training forward (kv_cache=None) should produce identical logits. - CudaBackend is inference-only: it raises when kv_cache is None. Training - must use TorchNativeBackend (the default). Verify the torch path is - stable and that CudaBackend rejects the training path explicitly. + CudaBackend is now safe as a default: for training (``kv_cache=None``) + or non-bf16 inputs it falls back to torch SDPA. Verify the fallback + path matches the torch-native forward exactly. """ import pytest @@ -39,11 +39,13 @@ def test_training_forward_matches_torch(cuda_model): with torch.no_grad(): out_torch = model(input_ids) - with pytest.raises(RuntimeError, match="does not support training"): - with attn_backend(ATTN_BACKEND.CUDA): - with torch.no_grad(): - 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 + ) assert out_torch["logits"].shape[0] == 2