(self, x)
| 147 | |
| 148 | @torch.no_grad() |
| 149 | def forward(self, x): |
| 150 | out_shape = x.shape[:-1] + (self.out_features, ) |
| 151 | inputs = x.reshape(-1, x.shape[-1]) |
| 152 | # trans = lambda x: x.T.contiguous() |
| 153 | if inputs.shape[0] > 1: |
| 154 | if USE_TRITON: |
| 155 | out = quant_matmul_v2(inputs, |
| 156 | self.qweight.T.contiguous(), |
| 157 | self.qzeros.T.contiguous(), |
| 158 | self.scales.T.contiguous(), |
| 159 | M=out_shape[-2], |
| 160 | N=out_shape[-1], |
| 161 | K=self.in_features, |
| 162 | pack_num=self.split_k_iters, |
| 163 | group_size=self.group_size, |
| 164 | w_bit=self.w_bit, |
| 165 | offset=self.offset |
| 166 | ) |
| 167 | else: |
| 168 | out = awq_inference_engine.gemm_forward_cuda(inputs, self.qweight, self.scales, self.qzeros, self.group_size, self.split_k_iters) |
| 169 | else: |
| 170 | out = awq_inference_engine.gemv_forward_cuda(inputs, self.qweight, self.scales, self.qzeros, self.w_bit, self.group_size) |
| 171 | out = out + self.bias if self.bias is not None else out |
| 172 | |
| 173 | return out.reshape(out_shape) |
| 174 | |
| 175 | def extra_repr(self) -> str: |
| 176 | return 'in_features={}, out_features={}, bias={}, w_bit={}, group_size={}'.format( |
nothing calls this directly
no test coverage detected