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
25 lines
608 B
Python
25 lines
608 B
Python
from astrai.model.components.attention import GQA, MLA
|
|
from astrai.model.components.decoder_block import DecoderBlock
|
|
from astrai.model.components.embedding import Embedding
|
|
from astrai.model.components.linear import Linear
|
|
from astrai.model.components.mlp import MLP
|
|
from astrai.model.components.norm import RMSNorm
|
|
from astrai.model.components.rope import (
|
|
RotaryEmbedding,
|
|
apply_rotary_emb,
|
|
get_rotary_emb,
|
|
)
|
|
|
|
__all__ = [
|
|
"Linear",
|
|
"RMSNorm",
|
|
"MLP",
|
|
"Embedding",
|
|
"GQA",
|
|
"MLA",
|
|
"DecoderBlock",
|
|
"RotaryEmbedding",
|
|
"apply_rotary_emb",
|
|
"get_rotary_emb",
|
|
]
|