refactor : use factory for attention backends

- register built-in backends through BaseFactory
- derive benchmark choices from registered backends
- cover string selection and invalid backend names

astrai/extension/__init__.py

@@ -18,6 +18,7 @@ SDPA is handled by the attention backend, not the wrapper functions.
 from astrai.extension.attention_backend import (
     ATTN_BACKEND,
     AttentionBackend,
+    AttentionBackendFactory,
     CudaBackend,
     FlashAttnBackend,
     TorchNativeBackend,
@@ -37,6 +38,7 @@ from astrai.extension.rotary_backend import apply_rotary_emb
 __all__ = [
     "ATTN_BACKEND",
     "AttentionBackend",
+    "AttentionBackendFactory",
     "CudaBackend",
     "TorchNativeBackend",
     "FlashAttnBackend",

astrai/extension/attention_backend.py

@@ -44,6 +44,7 @@ from astrai.extension.attention_ops import (
     attn_paged_decode,
     attn_paged_prefill,
 )
+from astrai.factory import BaseFactory
 from astrai.inference.core.cache import KVCache
 
 _current_backend: contextvars.ContextVar["AttentionBackend"] = contextvars.ContextVar(
@@ -146,11 +147,11 @@ def get_backend() -> "AttentionBackend":
 
 
 @contextmanager
-def attn_backend(backend: Union[ATTN_BACKEND, "AttentionBackend", type]):
+def attn_backend(backend: Union[str, ATTN_BACKEND, "AttentionBackend", type]):
     """Context manager to select an attention backend.
 
     Mirrors ``torch.nn.attention.sdpa_kernel``. Accepts an
-    ``ATTN_BACKEND`` enum value, a backend class, or a backend instance.
+    registered name, ``ATTN_BACKEND`` enum value, backend class, or instance.
 
     Examples::
 
@@ -162,14 +163,17 @@ def attn_backend(backend: Union[ATTN_BACKEND, "AttentionBackend", type]):
             ...
     """
     if isinstance(backend, ATTN_BACKEND):
-        instance = _BACKEND_REGISTRY[backend]()
+        instance = AttentionBackendFactory.create(backend.value)
+    elif isinstance(backend, str):
+        instance = AttentionBackendFactory.create(backend)
     elif isinstance(backend, type) and issubclass(backend, AttentionBackend):
         instance = backend()
     elif isinstance(backend, AttentionBackend):
         instance = backend
     else:
         raise TypeError(
-            f"expected ATTN_BACKEND, AttentionBackend type, or instance, "
+            f"expected a registered name, ATTN_BACKEND, AttentionBackend type, "
+            f"or instance, "
             f"got {type(backend).__name__}"
         )
     token = _current_backend.set(instance)
@@ -301,6 +305,11 @@ class AttentionBackend(ABC):
         """Multi-token prefill or training forward."""
 
 
+class AttentionBackendFactory(BaseFactory[AttentionBackend]):
+    """Factory for registered attention backends."""
+
+
+@AttentionBackendFactory.register(ATTN_BACKEND.TORCH_NATIVE.value)
 class TorchNativeBackend(AttentionBackend):
     """Reference backend using torch SDPA with indirect KV cache indexing.
 
@@ -385,6 +394,7 @@ class TorchNativeBackend(AttentionBackend):
 _default_backend = TorchNativeBackend()
 
 
+@AttentionBackendFactory.register(ATTN_BACKEND.CUDA.value)
 class CudaBackend(AttentionBackend):
     """CUDA kernel backend with direct KV cache access.
 
@@ -474,6 +484,7 @@ class CudaBackend(AttentionBackend):
         return out.reshape(b, q_len, q.size(2), q.size(3)).flatten(2)
 
 
+@AttentionBackendFactory.register(ATTN_BACKEND.FLASH.value)
 class FlashAttnBackend(AttentionBackend):
     """FlashAttention (FA2/FA3) backend via the optional ``flash-attn`` package.
 
@@ -557,10 +568,3 @@ class FlashAttnBackend(AttentionBackend):
             q.contiguous(), k.contiguous(), v.contiguous(), causal=is_causal
         )
         return out.contiguous().flatten(2)
-
-
-_BACKEND_REGISTRY: dict[ATTN_BACKEND, type[AttentionBackend]] = {
-    ATTN_BACKEND.TORCH_NATIVE: TorchNativeBackend,
-    ATTN_BACKEND.CUDA: CudaBackend,
-    ATTN_BACKEND.FLASH: FlashAttnBackend,
-}

scripts/tools/benchmark.py

@@ -1,23 +1,18 @@
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Union
 
 import click
 import torch
 
 from astrai import setup_logging
 from astrai.config import AutoRegressiveLMConfig
-from astrai.extension import ATTN_BACKEND, attn_backend
+from astrai.extension import ATTN_BACKEND, AttentionBackendFactory, attn_backend
 from astrai.inference.core.cache import PagePool
 from astrai.model import AutoModel
 
 _DTYPES = ["bfloat16", "float16", "float32"]
 _CACHES = ["contiguous", "paged"]
-_BACKENDS = ["cuda", "torch_native"]
+_BACKENDS = AttentionBackendFactory.list_registered()
-
-_BACKEND_MAP = {
-    "cuda": ATTN_BACKEND.CUDA,
-    "torch_native": ATTN_BACKEND.TORCH_NATIVE,
-}
 
 
 class BenchmarkResult:
@@ -46,7 +41,7 @@ class GenerationBenchmark:
         device: str = "cuda",
         dtype: torch.dtype = torch.bfloat16,
         cache_type: str = "contiguous",
-        backend: ATTN_BACKEND = ATTN_BACKEND.CUDA,
+        backend: Union[str, ATTN_BACKEND] = ATTN_BACKEND.CUDA,
     ):
         self.device = device
         self.dtype = dtype
@@ -297,7 +292,7 @@ def benchmark_command(
             device=device,
             dtype=dtype_map[dtype],
             cache_type=cache,
-            backend=_BACKEND_MAP[name],
+            backend=name,
         )
 
         click.secho(

tests/extension/test_backend.py

@@ -8,6 +8,7 @@ import pytest
 
 from astrai.extension import (
     ATTN_BACKEND,
+    AttentionBackendFactory,
     CudaBackend,
     TorchNativeBackend,
     attn_backend,
@@ -26,6 +27,26 @@ def test_attn_backend_context_with_enum():
     assert isinstance(get_backend(), TorchNativeBackend)
 
 
+def test_attn_backend_context_with_registered_name():
+    with attn_backend("cuda"):
+        assert isinstance(get_backend(), CudaBackend)
+    assert isinstance(get_backend(), TorchNativeBackend)
+
+
+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():
     with attn_backend(CudaBackend):
         assert isinstance(get_backend(), CudaBackend)