Kernels:

phanerozoic
/

bitnet-tc

Kernel card Files Files and versions

bitnet-tc / torch-ext /torch_binding.cpp

phanerozoic's picture

bitnet-tc v1: kernel-builder source tree

e4ee85d verified 26 days ago

722 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_bf16, 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(CUDA_KERNEL) \|\| defined(ROCM_KERNEL)
	ops.impl("quantize_act", torch::kCUDA, &quantize_act);
	ops.impl("bitnet_gemv_fused", torch::kCUDA, &bitnet_gemv_fused);
	ops.impl("bitnet_gemm", torch::kCUDA, &bitnet_gemm);
	#endif
	}

	REGISTER_EXTENSION(TORCH_EXTENSION_NAME)