pub mod activation;
pub mod attention;
pub mod embedding;
pub mod gemm;
pub mod layernorm;
pub mod rmsnorm;
pub mod rope;
pub mod softmax;
pub mod transpose;

pub use activation::{add, gelu, mul, scale, silu, silu_mul};
pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu};
pub use attention::{attention, decode_attention, flash_attention};
pub use embedding::embedding;
pub use gemm::{batched_matmul, matmul, GemmBackend};
pub use layernorm::layernorm;
pub use rmsnorm::{add_rmsnorm, rmsnorm};
pub use rope::{rope_inplace, RopeCache};
pub use softmax::softmax;

/// Register GPU kernels with the tensor crate. Call once at startup.
pub fn init() {
    xserv_tensor::register_gpu_contiguous(strided_to_contiguous_gpu);
}