From 8152760b5fae7b60d349b5287406156b835c8407 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 5 Aug 2026 15:37:22 +0800 Subject: [PATCH] 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 | 2 ++ astrai/extension/attention_backend.py | 26 +++++++++++++++----------- scripts/tools/benchmark.py | 15 +++++---------- tests/extension/test_backend.py | 21 +++++++++++++++++++++ 4 files changed, 43 insertions(+), 21 deletions(-) diff --git a/astrai/extension/__init__.py b/astrai/extension/__init__.py index 99bb03b..068e0c6 100644 --- a/astrai/extension/__init__.py +++ b/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", diff --git a/astrai/extension/attention_backend.py b/astrai/extension/attention_backend.py index eec63cc..bdfe668 100644 --- a/astrai/extension/attention_backend.py +++ b/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, -} diff --git a/scripts/tools/benchmark.py b/scripts/tools/benchmark.py index ecc7118..a00be18 100644 --- a/scripts/tools/benchmark.py +++ b/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"] - -_BACKEND_MAP = { - "cuda": ATTN_BACKEND.CUDA, - "torch_native": ATTN_BACKEND.TORCH_NATIVE, -} +_BACKENDS = AttentionBackendFactory.list_registered() 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( diff --git a/tests/extension/test_backend.py b/tests/extension/test_backend.py index 6f7e40a..b010701 100644 --- a/tests/extension/test_backend.py +++ b/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)