Function transpose_matmul248

models/quantization.py:260–272

(input, qweight, scales, qzeros, g_idx, bits, maxq)

def transpose_matmul248(input, qweight, scales, qzeros, g_idx, bits, maxq):
    output_dim = (qweight.shape[0] * 32) // bits
    output = torch.empty((input.shape[0], output_dim), device='cuda', dtype=torch.float16)
    grid = lambda META: (
    triton.cdiv(input.shape[0], META['BLOCK_SIZE_M']) * triton.cdiv(output_dim, META['BLOCK_SIZE_K']),)
    transpose_matmul_248_kernel[grid](input, qweight, output,
                                      scales, qzeros, g_idx,
                                      input.shape[0], qweight.shape[1], output_dim, bits, maxq,
                                      input.stride(0), input.stride(1),
                                      qweight.stride(0), qweight.stride(1),
                                      output.stride(0), output.stride(1),
                                      scales.stride(0), qzeros.stride(0))
    return output


class QuantLinearFunction(torch.autograd.Function):