ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: auto-select best available attention backend

Browse Source 
- 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)
This commit is contained in:
ViperEkura
2026-08-07 13:48:59 +08:00
parent a59ae8f32e
commit 6f49738991
3 changed files with 105 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/attention_backend.py
+77 -9
View File
@@ -44,6 +44,7 @@ from astrai.extension.attention_ops import (
attn_paged_decode,
attn_paged_prefill,
)
from astrai.extension.loader import is_available
from astrai.factory import BaseFactory
if TYPE_CHECKING:
@@ -136,15 +137,65 @@ class ATTN_BACKEND(enum.Enum):
FLASH = "flash"
_default_backend: Optional["AttentionBackend"] = None
def _priority_backends() -> list["AttentionBackend"]:
"""Available backends in priority order: flash -> cuda -> torch."""
backends: list[AttentionBackend] = []
if flash_attn_available():
backends.append(FlashAttnBackend())
if is_available("attn_paged_decode") and is_available("attn_paged_prefill"):
backends.append(CudaBackend())
backends.append(TorchNativeBackend())
return backends
def _backend_supports(
backend: "AttentionBackend",
q: Tensor,
kv_cache: Optional["KVCache"],
attn_mask: Optional[Tensor],
is_causal: bool,
) -> bool:
"""Whether ``backend`` can run this attention call.
The CUDA kernels are bf16-only, support head_dim in 32/64/128/256, and
need a KV cache (decode/prefill); everything else falls back to torch.
"""
if isinstance(backend, CudaBackend):
return (
kv_cache is not None
and q.dtype == torch.bfloat16
and q.size(-1) in (32, 64, 128, 256)
)
if isinstance(backend, FlashAttnBackend):
return flash_attn_available() and not (attn_mask is not None and not is_causal)
return True
def _resolve_default_backend() -> "AttentionBackend":
"""Pick the highest-priority available backend: flash -> cuda -> torch.
Resolved lazily on first ``get_backend()`` (flash/cuda availability is
checked once and cached). Per-call capability fallback happens in
``attention()``, so this default is safe for training and fp32 models.
"""
return _priority_backends()[0]
def get_backend() -> "AttentionBackend":
"""Return the active backend for the current thread/context.
Falls back to a ``TorchNativeBackend`` singleton when no backend
has been activated via ``with``.
Falls back to the highest-priority available backend (flash -> cuda ->
torch_native) when no backend has been activated via ``with``.
"""
try:
return _current_backend.get()
except LookupError:
global _default_backend
if _default_backend is None:
_default_backend = _resolve_default_backend()
return _default_backend
@@ -225,6 +276,16 @@ def attention(
[batch, q_len, n_heads * head_dim]
"""
backend = get_backend()
if not _backend_supports(backend, q, kv_cache, attn_mask, is_causal):
# The active backend cannot run this call (e.g. CUDA on a training /
# fp32 / unsupported-head_dim input) — fall back to the highest-
# priority backend that can, ending at torch SDPA.
for candidate in _priority_backends():
if isinstance(candidate, type(backend)):
continue
if _backend_supports(candidate, q, kv_cache, attn_mask, is_causal):
backend = candidate
break
return backend.forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
@@ -393,7 +454,7 @@ class TorchNativeBackend(AttentionBackend):
return out
_default_backend = TorchNativeBackend()
_default_backend = None
@AttentionBackendFactory.register(ATTN_BACKEND.CUDA.value)
@@ -407,8 +468,9 @@ class CudaBackend(AttentionBackend):
``attn_paged_prefill`` with ragged-batch support via qo_indptr +
kv_indptr.
``kv_cache is None`` (training) is not handled — use
``TorchNativeBackend`` for training.
``kv_cache is None`` (training) raises — the per-call fallback to
torch SDPA for training / fp32 / unsupported head_dim happens in the
``attention()`` entry point.
Raises ``RuntimeError`` if the required kernel is not available.
"""
@@ -426,8 +488,9 @@ class CudaBackend(AttentionBackend):
if kv_cache is None:
raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
loc = kv_cache.out_cache_loc[:, 0]
kv_cache.k_buffer[layer_id].index_copy_(0, loc, k[:, 0])
kv_cache.v_buffer[layer_id].index_copy_(0, loc, v[:, 0])
b = q.size(0)
q_3d = q.squeeze(1)
@@ -461,8 +524,13 @@ class CudaBackend(AttentionBackend):
if kv_cache is None:
raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
loc = kv_cache.out_cache_loc.reshape(-1)
kv_cache.k_buffer[layer_id].index_copy_(
0, loc, k.reshape(-1, k.size(2), k.size(3))
)
kv_cache.v_buffer[layer_id].index_copy_(
0, loc, v.reshape(-1, v.size(2), v.size(3))
)
b = q.size(0)
q_len = q.size(1)
tests/extension/test_backend.py
+19 -6
View File
@@ -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
+9 -7
View File
@@ -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