ViperEkura
3067a8e1a6
- Add attention() functional entry delegating to active backend - GQA/MLA forward calls attention() instead of inline cache/SDPA - CUDA kernels support 2D/3D/4D mask via mask_h_stride field - CudaBackend.fwd_decode builds 2D padding mask for mixed seq_lens - KVCache.max_len precomputed in bind_tasks to avoid GPU sync - batch==1 decode short-circuits mask=None - Split tests into conftest, test_backend, test_backend_equivalence, test_kernel_mask - 440 tests pass, L20 decode 1.44-1.60x speedup vs torch native
46 lines
1.2 KiB
Python
46 lines
1.2 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,
|
|
CudaBackend,
|
|
TorchNativeBackend,
|
|
attn_backend,
|
|
get_backend,
|
|
)
|
|
|
|
|
|
def test_default_backend_is_torch_native():
|
|
backend = get_backend()
|
|
assert isinstance(backend, TorchNativeBackend)
|
|
|
|
|
|
def test_attn_backend_context_with_enum():
|
|
with attn_backend(ATTN_BACKEND.CUDA):
|
|
assert isinstance(get_backend(), CudaBackend)
|
|
assert isinstance(get_backend(), TorchNativeBackend)
|
|
|
|
|
|
def test_attn_backend_context_with_class():
|
|
with attn_backend(CudaBackend):
|
|
assert isinstance(get_backend(), CudaBackend)
|
|
assert isinstance(get_backend(), TorchNativeBackend)
|
|
|
|
|
|
def test_attn_backend_context_with_instance():
|
|
custom = CudaBackend()
|
|
with attn_backend(custom):
|
|
assert get_backend() is custom
|
|
assert isinstance(get_backend(), TorchNativeBackend)
|
|
|
|
|
|
def test_cudabackend_is_context_manager():
|
|
with CudaBackend():
|
|
assert isinstance(get_backend(), CudaBackend)
|
|
assert isinstance(get_backend(), TorchNativeBackend)
|