| 19 | bitnet_lib = ctypes.CDLL('bitnet_kernels/libbitnet.so') |
| 20 | |
| 21 | def bitnet_int8xint2_linear(input0, input1, s, ws): |
| 22 | out_shape = list(input0.shape) |
| 23 | out_shape[-1] = input1.shape[0] |
| 24 | |
| 25 | stream = torch.cuda.current_stream() |
| 26 | |
| 27 | M = input0.shape[0] |
| 28 | if len(out_shape) == 3: |
| 29 | M *= input0.shape[1] |
| 30 | N = input1.shape[0] |
| 31 | K = input1.shape[1] * 4 |
| 32 | |
| 33 | ret = torch.zeros(*out_shape, dtype=torch.bfloat16, device=input0.device) |
| 34 | |
| 35 | bitnet_lib.bitlinear_int8xint2(*[ctypes.c_void_p(input0.data_ptr()), ctypes.c_void_p(input1.data_ptr()), ctypes.c_void_p(ret.data_ptr()), ctypes.c_void_p(s.data_ptr()), ctypes.c_void_p(ws.data_ptr()), ctypes.c_int(M), ctypes.c_int(N), ctypes.c_int(K), ctypes.c_void_p(stream.cuda_stream)]) |
| 36 | |
| 37 | return ret |
| 38 | |
| 39 | @dataclass |
| 40 | class ModelArgs: |