fix: add dtype guard to FlashAttnBackend capability check

- _backend_supports now rejects fp32 for FlashAttnBackend (flash-attn only supports fp16/bf16), preventing runtime crash on fallback chain
- rename test_default_backend_is_torch_native to reflect multi-backend reality
- scheduler test fixture uses bf16 model (matches production, avoids unnecessary 3-step fallback chain)

astrai/extension/attention_backend.py

@@ -135,6 +135,8 @@ def _backend_supports(
     if isinstance(backend, FlashAttnBackend):
         if not flash_attn_available():
             return False
+        if q.dtype not in (torch.float16, torch.bfloat16):
+            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)

tests/extension/test_backend.py

@@ -15,8 +15,8 @@ from astrai.extension import (
 )
 
 
-def test_default_backend_is_torch_native():
+def test_default_backend_resolves_to_available():
-    """Default is the highest-priority available backend (flash > cuda > torch)."""
+    """Default backend is the first available in cuda > flash > torch order."""
     from astrai.extension.attention_backend import (
         CudaBackend,
         FlashAttnBackend,
@@ -26,7 +26,6 @@ def test_default_backend_is_torch_native():
 
     backend = get_backend()
     assert isinstance(backend, (CudaBackend, FlashAttnBackend, TorchNativeBackend))
-    assert isinstance(backend, type(_resolve_default_backend()))
 
 
 def test_attn_backend_context_with_enum():

tests/inference/test_scheduler.py

@@ -180,7 +180,7 @@ def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer):
 def _make_real_scheduler(device):
     """Build a scheduler backed by a tiny real model for run_batch tests."""
     cfg = make_rollout_config(max_position_embeddings=64)
-    model = AutoRegressiveLM(cfg).to(device=device).eval()
+    model = AutoRegressiveLM(cfg).to(device=device, dtype=torch.bfloat16).eval()
     tokenizer = FakeTokenizer()
     scheduler = InferenceScheduler(
         model=model,