Function quant_matmul_v2

quantization/triton_kernels.py:103–134

(a, qw, qzeros, scales, *, M, N, K, pack_num, group_size, w_bit, offset)

def quant_matmul_v2(a, qw, qzeros, scales, *, M, N, K, pack_num, group_size, w_bit, offset):
    c = torch.empty((M, N), dtype=torch.float16, device=a.device)
    assert qw.shape == (K // pack_num, N)
    # assert qzeros.shape == (K // group_size // pack_num, N)
    # assert scales.shape == (K // group_size, N)
    assert all(x.is_contiguous() for x in [a, qw, c, qzeros, scales])
    # BLOCK_SIZE_K has possible values of 32, 64
    # group_size, K must be divisible by BLOCK_SIZE_K
    assert group_size % 64 == 0, f"group_size {group_size} is not a multiple of 64"
    assert K % 64 == 0, f"K {K} is not a multiple of 64"
    # BLOCK_SIZE_N has possible values of 32, 64, 128, 256
    # N must be divisible by BLOCK_SIZE_N
    assert N % 256 == 0, f"N {N} is not a multiple of 256"

    grid_1d = lambda META: (
        triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
    )
    quant_matmul_kernel[grid_1d](
        a_ptr=a,
        qw_ptr=qw,
        c_ptr=c,
        scales_ptr=scales,
        zeros_ptr=qzeros,
        M=M,
        N=N,
        K=K,
        pack_num=pack_num,
        w_bit=w_bit,
        group_size=group_size,
        offset=offset
    )
    return c