ViperEkura
3e67b4f88d
- Single-kernel rotary embedding (cos/sin lookup + rotation) replaces PyTorch complex-multiply path (3 kernel launches + f32 upcast per call) - RotaryEmbedding now stores cos_table/sin_table and returns (cos, sin) f32 tuple instead of a complex tensor - apply_rotary_emb in rotary_backend.py auto-dispatches: CUDA kernel if available, else torch complex-multiply fallback; backend-agnostic (both attention backends benefit) - Kernel: 256-thread blocks, grid-stride loop, vectorized __nv_bfloat162 load/store, f32 compute, bf16 out - Standalone kernel 6-9x faster than torch across decode/prefill shapes, max diff 0 (decode) to 3e-2 (large prefill, bf16) - Benchmark (L20, bf16, CUDA backend): B=1 9.48->7.25ms (+31%), B=4 10.73->7.67ms (+40%), B=8 10.77->7.81ms (+38%), B=16 10.79->7.83ms (+38%)
25 lines
647 B
Python
25 lines
647 B
Python
from astrai.extension.rotary_backend import apply_rotary_emb
|
|
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,
|
|
get_rotary_emb,
|
|
)
|
|
|
|
__all__ = [
|
|
"Linear",
|
|
"RMSNorm",
|
|
"MLP",
|
|
"Embedding",
|
|
"GQA",
|
|
"MLA",
|
|
"DecoderBlock",
|
|
"RotaryEmbedding",
|
|
"apply_rotary_emb",
|
|
"get_rotary_emb",
|
|
]
|