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)
This commit is contained in:
2026-08-07 23:08:22 +08:00
parent 184fbbce5c
commit 1b1f1a0707
3 changed files with 5 additions and 4 deletions
+2 -3
View File
@@ -15,8 +15,8 @@ from astrai.extension import (
)
def test_default_backend_is_torch_native():
"""Default is the highest-priority available backend (flash > cuda > torch)."""
def test_default_backend_resolves_to_available():
"""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():
+1 -1
View File
@@ -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,