Kernels:

phanerozoic
/

bitnet-cpu

Kernel card Files Files and versions

bitnet-cpu / torch-ext /torch_binding.h

phanerozoic's picture

bitnet-cpu v1: kernel-builder source tree

ce8679e verified 26 days ago

683 Bytes

	#pragma once

	#include <torch/torch.h>
	#include <c10/util/Optional.h>

	void quantize_act(torch::Tensor &Aq,
	torch::Tensor &scale,
	torch::Tensor const &A);

	void bitnet_gemv_fused(torch::Tensor &out,
	torch::Tensor const &a,
	torch::Tensor const &w_packed,
	torch::Tensor const &scale_wt);

	void bitnet_gemm(torch::Tensor &out,
	torch::Tensor const &a_int8,
	torch::Tensor const &w_packed,
	torch::Tensor const &scale_act,
	torch::Tensor const &scale_wt,
	c10::optional<torch::Tensor> const &scratch);