20 lines
402 B
Python
20 lines
402 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.rotary import rotary_emb
|
|
|
|
__all__ = [
|
|
"TensorLayout",
|
|
"attn_decode",
|
|
"attn_paged_decode",
|
|
"attn_paged_prefill",
|
|
"attn_prefill",
|
|
"rotary_emb",
|
|
]
|