Kernels:

phanerozoic
/

bitnet-cpu

Kernel card Files Files and versions

bitnet-cpu / torch-ext /torch_binding.cpp

phanerozoic's picture

bitnet-cpu v1: kernel-builder source tree

ce8679e verified 3 months ago

Raw History Blame

689 Bytes

	#include <torch/library.h>

	#include "registration.h"
	#include "torch_binding.h"

	TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
	ops.def("quantize_act(Tensor! Aq, Tensor! scale, Tensor a) -> ()");
	ops.def("bitnet_gemv_fused(Tensor! out, Tensor a, Tensor w_packed, Tensor scale_wt) -> ()");
	ops.def("bitnet_gemm(Tensor! out, Tensor a_int8, Tensor w_packed, Tensor scale_act, Tensor scale_wt, Tensor? scratch) -> ()");

	#if defined(CPU_KERNEL)
	ops.impl("quantize_act", torch::kCPU, &quantize_act);
	ops.impl("bitnet_gemv_fused", torch::kCPU, &bitnet_gemv_fused);
	ops.impl("bitnet_gemm", torch::kCPU, &bitnet_gemm);
	#endif
	}

	REGISTER_EXTENSION(TORCH_EXTENSION_NAME)