feat: auto-select best available attention backend

- Default backend resolves to highest-priority available: flash -> cuda -> torch
- attention() falls back per-call for training/fp32/unsupported head_dim
- Re-apply index_copy_ for CUDA KV cache writes (index_put_ race mitigation)

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

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