@triton.jit
def matmul_quant_kernel(A, B, C, N, M, K, K1, stride_an, stride_ak,
    stride_bk, stride_bm, stride_cn, stride_cm, scale, zero_point, BITS: tl
    .constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr, BLOCK_SIZE_K1: tl.constexpr):
    dtype = C.dtype.element_ty
    zero_point_ty = tl.full((), zero_point, dtype=dtype)
    scale_ty = tl.full((), scale, dtype=dtype)
    QFACT = 8 // BITS
    QMASK = (1 << BITS) - 1
    QMASK_INT = tl.full((), QMASK, dtype=tl.uint8)
    pid_n = tl.program_id(0)
    pid_m = tl.program_id(1)
    rn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    rm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    rk = tl.arange(0, BLOCK_SIZE_K)
    rk1, rq1 = rk // QFACT, rk % QFACT
    a_shift = BITS * rq1
    mask_a = rn[:, None] < N
    mask_b = rm[None, :] < M
    mask_c = mask_a & mask_b
    A1 = A + (rn[:, None] * stride_an + rk1[None, :] * stride_ak)
    B1 = B + (rk[:, None] * stride_bk + rm[None, :] * stride_bm)
    C1 = C + (rn[:, None] * stride_cn + rm[None, :] * stride_cm)
    acc = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=dtype)
    for k in range(0, K, BLOCK_SIZE_K):
        aq = tl.load(A1, mask=mask_a)
        b = tl.load(B1, mask=mask_b)
        ai = aq >> a_shift & QMASK_INT
        a = ai.to(dtype) - zero_point_ty
        b1 = b.to(dtype)
        acc += tl.dot(a, b1, out_dtype=dtype)
        A1 += BLOCK_SIZE_K1 * stride_ak
        B1 += BLOCK_SIZE_K * stride_bk
    acc *= scale_ty
    tl.store(C1, acc, mask=mask_c)
