0z5a
a144d7f306
- add decode-shape benchmark harness - add bf16 GEMV CUDA primitive with head-dim generic kernel - dispatch decode-time linear layers to gemv for M=1 - extend gemv coverage to small decode batches
30 lines
659 B
Python
30 lines
659 B
Python
"""Backend selection, fallbacks, and execution policies."""
|
|
|
|
from astrai.extension.backend.attention import (
|
|
ATTN_BACKEND,
|
|
AttentionBackend,
|
|
AttentionBackendFactory,
|
|
CudaBackend,
|
|
FlashAttnBackend,
|
|
TorchNativeBackend,
|
|
attention,
|
|
attn_backend,
|
|
get_backend,
|
|
)
|
|
from astrai.extension.backend.linear import linear
|
|
from astrai.extension.backend.rotary import apply_rotary_emb
|
|
|
|
__all__ = [
|
|
"ATTN_BACKEND",
|
|
"AttentionBackend",
|
|
"AttentionBackendFactory",
|
|
"CudaBackend",
|
|
"FlashAttnBackend",
|
|
"TorchNativeBackend",
|
|
"apply_rotary_emb",
|
|
"attention",
|
|
"attn_backend",
|
|
"get_backend",
|
|
"linear",
|
|
]
|