ViperEkura
02469887f5
- merge _generate_streaming/_generate_non_streaming into single _generate() with stream flag - delete dead GenerationRequest class and generate_with_request method - inline _next_token helper into generate_async - replace flash-attn double-checked locking with functools.lru_cache - extract _write_and_gather_kv helper shared by TorchNative/FlashAttn backends - inline _kv_cache_is_contiguous into its sole call site in FlashAttnBackend - change default backend priority from flash>cuda>torch to cuda>flash>torch - add ASTR_BACKEND env var to override default backend at resolve time - add supports_graph() static method to AttentionBackend ABC, override in CudaBackend - replace isinstance(get_backend(), CudaBackend) with get_backend().supports_graph() in executor - add torch.cuda.is_available() guard to CudaBackend.supports()
80 lines
2.1 KiB
Python
80 lines
2.1 KiB
Python
"""Backend selection and context-manager switching tests.
|
|
|
|
These tests do not require CUDA — they only check that the active
|
|
backend is correctly set and restored.
|
|
"""
|
|
|
|
import pytest
|
|
|
|
from astrai.extension import (
|
|
ATTN_BACKEND,
|
|
AttentionBackendFactory,
|
|
CudaBackend,
|
|
attn_backend,
|
|
get_backend,
|
|
)
|
|
|
|
|
|
def test_default_backend_is_torch_native():
|
|
"""Default is the highest-priority available backend (flash > cuda > torch)."""
|
|
from astrai.extension.attention_backend import (
|
|
CudaBackend,
|
|
FlashAttnBackend,
|
|
TorchNativeBackend,
|
|
_resolve_default_backend,
|
|
)
|
|
|
|
backend = get_backend()
|
|
assert isinstance(backend, (CudaBackend, FlashAttnBackend, 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 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 get_backend() is default
|
|
|
|
|
|
def test_attention_backend_factory_lists_builtin_backends():
|
|
assert AttentionBackendFactory.list_registered() == [
|
|
"cuda",
|
|
"flash",
|
|
"torch_native",
|
|
]
|
|
|
|
|
|
def test_attn_backend_rejects_unknown_registered_name():
|
|
with pytest.raises(ValueError, match="Unknown component: 'unknown'"):
|
|
with attn_backend("unknown"):
|
|
pass
|
|
|
|
|
|
def test_attn_backend_context_with_class():
|
|
default = get_backend()
|
|
with attn_backend(CudaBackend):
|
|
assert isinstance(get_backend(), CudaBackend)
|
|
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 get_backend() is default
|
|
|
|
|
|
def test_cudabackend_is_context_manager():
|
|
default = get_backend()
|
|
with CudaBackend():
|
|
assert isinstance(get_backend(), CudaBackend)
|
|
assert get_backend() is default
|