ViperEkura
6709534d64
- delete csrc/kernels/gemm.cu and swiglu.cu and drop their CMake and setup.py registration - remove the ops wrappers plus backend/linear.py and backend/swiglu.py so Linear and MLP call F.linear directly - drop the four gemm and swiglu kernel test files and prune the stale cuda_kernels.md sections - add csrc/bench benchmarks for the remaining kernels: attention decode prefill paged decode paged prefill versus single-launch SDPA references, rotary versus the torch fallback, fp8 quantize and mm_fp8 versus torch baselines - attention, rotary_emb, and fp8_ops kernels are unchanged
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",
|
|
]
|