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
22 lines
467 B
Python
22 lines
467 B
Python
"""Stateless wrappers around compiled extension kernels."""
|
|
|
|
from astrai.extension.ops.attention import (
|
|
TensorLayout,
|
|
attn_decode,
|
|
attn_paged_decode,
|
|
attn_paged_prefill,
|
|
attn_prefill,
|
|
)
|
|
from astrai.extension.ops.gemv import bf16_gemv
|
|
from astrai.extension.ops.rotary import rotary_emb
|
|
|
|
__all__ = [
|
|
"TensorLayout",
|
|
"attn_decode",
|
|
"attn_paged_decode",
|
|
"attn_paged_prefill",
|
|
"attn_prefill",
|
|
"bf16_gemv",
|
|
"rotary_emb",
|
|
]
|