File size: 683 Bytes
ce8679e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
#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);