| 30 | self.initialize(in_features, out_features, groupsize, double_groupsize, bits, v1, asym) |
| 31 | |
| 32 | def initialize(self, in_features, out_features, groupsize, double_quantize_groupsize, bits, v1, asym): |
| 33 | |
| 34 | if asym: |
| 35 | self.register_buffer('qzeros', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features / 256 * (bits * 8))), dtype=torch.int32)) |
| 36 | if bits == 4: |
| 37 | self.register_buffer('qscales', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features/double_quantize_groupsize), double_quantize_groupsize), dtype=torch.uint8)) |
| 38 | else: |
| 39 | self.register_buffer('qscales', torch.empty((math.ceil(in_features/groupsize), out_features), dtype=torch.uint8)) |
| 40 | |
| 41 | else: |
| 42 | self.register_buffer('qstatistic', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features/double_quantize_groupsize), double_quantize_groupsize), dtype=torch.uint8)) |
| 43 | self.register_buffer('qzeros_zeros', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features/double_quantize_groupsize), 1), dtype=torch.half)) |
| 44 | self.register_buffer('qzeros_scales', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features/double_quantize_groupsize), 1), dtype=torch.half)) |
| 45 | |
| 46 | if not v1: |
| 47 | self.register_buffer('qscales_zeros', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features/double_quantize_groupsize), 1), dtype=torch.half)) |
| 48 | self.register_buffer('qscales_scales', torch.empty((math.ceil(in_features/groupsize), math.ceil(out_features/double_quantize_groupsize), 1), dtype=torch.half)) |
| 49 | else: |
| 50 | self.register_buffer('qscales_zeros', torch.empty((1, out_features, 1), dtype=torch.half)) |
| 51 | self.register_buffer('qscales_scales', torch.empty((1, out_features, 1), dtype=torch.half)) |
| 52 | |
| 53 | self.register_buffer('g_idx', torch.tensor([i // groupsize for i in range(in_features)], dtype=torch.int32)) |
| 54 | self.register_buffer('qweight', torch.empty(math.ceil(in_features / 256 * (bits * 8)), out_features, dtype=torch.int32)) |
| 55 | self.register_buffer("wf", torch.tensor(list(range(0,32,bits)), dtype=torch.int32).unsqueeze(0)) |
| 56 | self.register_buffer('bias', torch.empty(out_features)) |
| 57 | |
| 58 | def forward(self, x): |
| 59 | if self.bits in [2, 4, 8, 16]: |