@triton.jit
def _kernel_quantize_mx4(A, out, rand_bits, M, K, GROUPS_PER_ROW,
    GROUPS_PER_THREAD, ROW_PADDING, GROUP_SIZE: tl.constexpr, EBITS: tl.
    constexpr, MBITS: tl.constexpr, ROUNDING_MODE: tl.constexpr,
    STOCHASTIC_CASTING: tl.constexpr, FP4_EXP_BIAS: tl.constexpr,
    GROUP_LOAD: tl.constexpr, USE_INT64: tl.constexpr) ->None:
    """Quantize a 1D float tensor into a packed MX4 tensor.

    Args:
        A (Tensor): [M] float tensor to be quantized.
        out (Tensor): [M / 2 + M / GROUP_SIZE] output containing packed mx4 values.
        rand_bits (Optional Tensor): [M, K / 2] random integers used for stochastic rounding.
        M (int): Number of input rows.
        K (int): Number of input columns.
        GROUPS_PER_ROW (int): Number of groups in each row of the input.
        GROUPS_PER_THREAD (int): Number of groups to process per thread.
        ROW_PADDING (int): Number of elements of padding to insert into each row.
        GROUP_SIZE (int): Size of chunks that use the same shared exponent.
        EBITS (int): Number of exponent bits in target mx4 format.
        MBITS (int): Number of mantissa bits in target mx4 format.
        ROUNDING_MODE (int): Which rounding method to use when calculating shared exponent.
        STOCHASTIC_CASTING (bool): Whether to use stochastic rounding when downcasting.
        FP4_EXP_BIAS (int): Exponent bias of target mx4 format.
        GROUP_LOAD (int): Number of groups to process simultaneously.
        USE_INT64 (bool): Whether to use int64 for indexing. This is needed for large tensors.
    """
    FP32_EXP_MASK: tl.constexpr = 2139095040
    FP32_EXP_OFFSET: tl.constexpr = 23
    FP32_EXP_BIAS: tl.constexpr = 127
    FP32_SIGN_OFFSET: tl.constexpr = 31
    SIGN_MASK: tl.constexpr = 1
    FP32_MANTISSA_MASK: tl.constexpr = 8388607
    MBITS_IMPLICIT: tl.constexpr = MBITS + 1
    MAX_FP32_MANTISSA_BITS: tl.constexpr = 24
    IMPLIED_1_BIT: tl.constexpr = 1 << 23
    FP32_MIN_NORMAL: tl.constexpr = 2 ** -126
    MANTISSA_OVERFLOW_THRESHOLD: tl.constexpr = (1 << MBITS_IMPLICIT) - 1
    EXPONENT_OVERFLOW_THRESHOLD: tl.constexpr = (1 << EBITS) - 1
    IMPLICIT_1_MASK = (1 << MBITS_IMPLICIT - 1) - 1
    RAND_MASK: tl.constexpr = (1 << FP32_EXP_OFFSET - MBITS) - 1
    pid = tl.program_id(0)
    if USE_INT64:
        pid = pid.to(tl.int64)
        M = tl.cast(M, tl.int64)
        K = tl.cast(K, tl.int64)
        GROUPS_PER_THREAD = tl.cast(GROUPS_PER_THREAD, tl.int64)
    PACKED_GROUP_SIZE: tl.constexpr = GROUP_SIZE // 2 + 1
    NUM_GROUPS = M * GROUPS_PER_ROW
    OUTPUT_CHUNK_SIZE = GROUPS_PER_THREAD * GROUP_SIZE // 2 + GROUPS_PER_THREAD
    OUTPUT_SIZE = GROUP_SIZE * NUM_GROUPS // 2 + NUM_GROUPS
    input_start = pid * (GROUPS_PER_THREAD * GROUP_SIZE)
    output_start = pid * OUTPUT_CHUNK_SIZE
    exp_start = output_start + GROUP_SIZE // 2
    input_offset = tl.arange(0, GROUP_LOAD * GROUP_SIZE) + input_start
    output_offset = tl.arange(0, GROUP_LOAD * (GROUP_SIZE // 2))
    if ROUNDING_MODE == 3:
        rand_bits_offset = tl.arange(0, GROUP_LOAD) + pid * GROUPS_PER_THREAD
    else:
        rand_bits_offset = pid * GROUPS_PER_THREAD
    output_offset += output_offset // (GROUP_SIZE // 2) + output_start
    exp_offset = tl.arange(0, GROUP_LOAD) * PACKED_GROUP_SIZE + exp_start
    for _k in range(0, tl.cdiv(GROUPS_PER_THREAD, GROUP_LOAD)):
        pad_mask = input_offset % (GROUPS_PER_ROW * GROUP_SIZE) < K
        if ROW_PADDING != 0:
            padded_input_offset = input_offset - input_offset // (
                GROUPS_PER_ROW * GROUP_SIZE) * ROW_PADDING
        else:
            padded_input_offset = input_offset
        a = tl.load(A + padded_input_offset, mask=(padded_input_offset < M *
            K) & (padded_input_offset < (pid + 1) * GROUPS_PER_THREAD *
            GROUP_SIZE) & pad_mask, other=0)
        a_groups = tl.reshape(a, [GROUP_LOAD, GROUP_SIZE])
        group_max = tl.max(tl.abs(a_groups), axis=1)
        group_max = tl.where(group_max == 0, FP32_MIN_NORMAL, group_max)
        group_rand_bits = None
        if ROUNDING_MODE == 3 or STOCHASTIC_CASTING:
            group_rand_bits = tl.load(rand_bits + rand_bits_offset, mask=
                rand_bits_offset < K // GROUP_SIZE, other=0)
            rand_bits_offset += GROUP_LOAD
        group_exp = _compute_exp(group_max, ROUNDING_MODE, group_rand_bits,
            MBITS)
        group_exp = group_exp - EBITS
        group_exp = tl.clamp(group_exp, -127, 125)
        scale = tl.exp2(group_exp.to(tl.float64)).to(tl.float32)
        scaled_a = tl.reshape(a, [GROUP_LOAD, GROUP_SIZE]) / tl.reshape(scale,
            [GROUP_LOAD, 1])
        scaled_a = tl.reshape(scaled_a, [GROUP_LOAD * GROUP_SIZE])
        tl.store(out + exp_offset, (group_exp + FP32_EXP_BIAS).to(tl.int8),
            mask=(exp_offset < OUTPUT_SIZE) & (exp_offset < 
            OUTPUT_CHUNK_SIZE * (pid + 1)))
        scaled_a = scaled_a.to(tl.int32, bitcast=True)
        if STOCHASTIC_CASTING:
            philox_4x_offset = tl.split(tl.reshape(input_offset, [
                GROUP_LOAD * GROUP_SIZE // 2, 2], can_reorder=True))
            philox_4x_offset = tl.split(tl.reshape(philox_4x_offset, [
                GROUP_LOAD * GROUP_SIZE // 4, 2], can_reorder=True))
            a_4x, b_4x, c_4x, d_4x = tl.randint4x(group_rand_bits,
                philox_4x_offset, n_rounds=7)
            stochastic_round_bits = tl.join(tl.join(a_4x, b_4x), tl.join(
                c_4x, d_4x))
            stochastic_round_bits = tl.reshape(stochastic_round_bits, [
                GROUP_LOAD * GROUP_SIZE]).to(tl.int32, bitcast=True)
            scaled_a = scaled_a + (stochastic_round_bits & RAND_MASK)
        sign_bit = scaled_a >> FP32_SIGN_OFFSET & SIGN_MASK
        biased_exp = (scaled_a & FP32_EXP_MASK) >> FP32_EXP_OFFSET
        trailing_mantissa = scaled_a & FP32_MANTISSA_MASK
        new_biased_exp = biased_exp - FP32_EXP_BIAS + FP4_EXP_BIAS
        exp_diff = tl.where(new_biased_exp <= 0, 1 - new_biased_exp, 0)
        exp_diff = tl.minimum(exp_diff, MAX_FP32_MANTISSA_BITS)
        is_subnorm = biased_exp == 0
        mantissa = tl.where(is_subnorm, trailing_mantissa, 
            trailing_mantissa + IMPLIED_1_BIT)
        fp32_sig_bits = tl.where(is_subnorm, 23, 24).to(tl.int32)
        mantissa = mantissa >> fp32_sig_bits + exp_diff - MBITS_IMPLICIT - 1
        mantissa = mantissa + 1 >> 1
        overflow = mantissa > MANTISSA_OVERFLOW_THRESHOLD
        mantissa = tl.where(overflow and not is_subnorm, mantissa >> 1,
            mantissa)
        new_biased_exp = tl.where(new_biased_exp <= 0 and mantissa == 2, 1,
            new_biased_exp)
        mantissa = mantissa & IMPLICIT_1_MASK
        new_biased_exp = tl.where(overflow, new_biased_exp + 1, new_biased_exp)
        mantissa = tl.where(new_biased_exp > EXPONENT_OVERFLOW_THRESHOLD, 1,
            mantissa)
        new_biased_exp = tl.maximum(tl.minimum(new_biased_exp,
            EXPONENT_OVERFLOW_THRESHOLD), 0)
        mx4_value = new_biased_exp << MBITS_IMPLICIT - 1 | mantissa
        mx4_value = sign_bit << EBITS + MBITS | mx4_value
        low_mx4, high_mx4 = tl.split(tl.reshape(mx4_value, [GROUP_LOAD *
            GROUP_SIZE // 2, 2]))
        packed_mx4 = (high_mx4 << 4 | low_mx4).to(tl.int8)
        tl.store(out + output_offset, packed_mx4, mask=(output_offset <
            OUTPUT_SIZE) & (output_offset < OUTPUT_CHUNK_SIZE * (pid + 1)))
        input_offset += GROUP_LOAD * GROUP_SIZE
        exp_offset += GROUP_LOAD * PACKED_GROUP_SIZE
        output_offset += GROUP_LOAD * PACKED_GROUP_SIZE
