| 273 | |
| 274 | |
| 275 | class QuantLinearFunction(torch.autograd.Function): |
| 276 | @staticmethod |
| 277 | @custom_fwd(cast_inputs=torch.float16) |
| 278 | def forward(ctx, input, qweight, scales, qzeros, g_idx, bits, maxq): |
| 279 | output = matmul248(input, qweight, scales, qzeros, g_idx, bits, maxq) |
| 280 | ctx.save_for_backward(qweight, scales, qzeros, g_idx) |
| 281 | ctx.bits, ctx.maxq = bits, maxq |
| 282 | return output |
| 283 | |
| 284 | @staticmethod |
| 285 | @custom_bwd |
| 286 | def backward(ctx, grad_output): |
| 287 | qweight, scales, qzeros, g_idx = ctx.saved_tensors |
| 288 | bits, maxq = ctx.bits, ctx.maxq |
| 289 | grad_input = None |
| 290 | |
| 291 | if ctx.needs_input_grad[0]: |
| 292 | grad_input = transpose_matmul248(grad_output, qweight, scales, qzeros, g_idx, bits, maxq) |
| 293 | return grad_input, None, None, None, None, None, None |
| 294 | |
| 295 | class QuantLinear(nn.Module): |
| 296 | def __init__(self, bits, groupsize, infeatures, outfeatures, bias): |
nothing calls this directly
no outgoing calls
no test coverage detected