@triton.autotune(configs=[Config({'BLOCK_SIZE': 512}), Config({'BLOCK_SIZE':
    1024}), Config({'BLOCK_SIZE': 2048}), Config({'BLOCK_SIZE': 4096}),
    Config({'BLOCK_SIZE': 8192})], key=['K'])
@triton.jit
def _kernel_quantize_fp8_row(A, A_scale, A_fp8, scale_ub, B, M, N, K,
    stride_ab, stride_am, stride_an, stride_ak, stride_ob, stride_om,
    stride_on, stride_ok, TL_FP8_DTYPE: tl.constexpr, MAX_FP8: tl.constexpr,
    EPS: tl.constexpr, CLAMP_MAX: tl.constexpr, BLOCK_SIZE: tl.constexpr,
    USE_INT64: tl.constexpr) ->None:
    """Quantize and scale each row.

    Scale per row i is computed as MAX_FP8 / max(abs(A[i, :]))

    Kernel naively iterates through  matrix with [1, BLOCK_SIZE] tiles
    in a max pass then scale/quantize pass.

    Todo:
        * Better tiling schemes.

    Args:
        A (Tensor): higher precision input tensor of 4 dimension.
        A_scale (Tensor): [B * M * N] reciprocal scale tensor per row.
        A_fp8 (Tensor): fp8 scaled tensor. A_fp8 = A / a_scale
        scale_ub (Tensor): [1] Maximum value allowed for scale.
        B (int): Size of dimenion 0
        M (int): Size of dimenion 1
        N (int): Size of dimenion 2
        K (int): Size of dimenion 3
        stride_ab (int): Stride of b dimension of A.
        stride_am (int): Stride of m dimension of A.
        stride_an (int): Stride of n dimension of A.
        stride_ak (int): Stride of k dimension of A.
        stride_ob (int): Stride of b dimension of output.
        stride_om (int): Stride of m dimension of output.
        stride_on (int): Stride of n dimension of output.
        stride_ok (int): Stride of k dimension of output.
        TL_FP8_DTYPE (tl.dtype): Target fp8 datatype.
        MAX_FP8 (float): Maxmimum expressible value for FP8.
        EPS (float): Epsilon value for numerical stability.
        CLAMP_MAX (bool): Whethar to apply scale_ub.
        BLOCK_SIZE (int): Block size for reduction.
        USE_INT64 (bool): Whether to use int64 indexing for large inputs.
    """
    pid = tl.program_id(0)
    if USE_INT64:
        pid = pid.to(tl.int64)
    n_offset = tl.arange(0, BLOCK_SIZE)
    a_offset_base = pid // (M * N) * stride_ab + pid % (M * N
        ) // N * stride_am + pid % (M * N) % N * stride_an
    a_fp8_offset_base = pid // (M * N) * stride_ob + pid % (M * N
        ) // N * stride_om + pid % (M * N) % N * stride_on
    cur_max = 0.0
    for _k in range(0, tl.cdiv(K, BLOCK_SIZE)):
        a = tl.load(A + a_offset_base + n_offset * stride_ak, mask=n_offset <
            K, other=0.0)
        tile_max = tl.max(tl.abs(a))
        cur_max = tl.maximum(tile_max, cur_max)
        n_offset += BLOCK_SIZE
    if CLAMP_MAX:
        ub = tl.load(scale_ub)
        cur_max = tl.clamp(cur_max, EPS, ub)
    else:
        cur_max = tl.maximum(cur_max, EPS)
    a_scale = MAX_FP8 / cur_max
    tl.store(A_scale + pid, 1.0 / a_scale)
    n_offset = tl.arange(0, BLOCK_SIZE)
    for _k in range(0, tl.cdiv(K, BLOCK_SIZE)):
        a = tl.load(A + a_offset_base + n_offset * stride_ak, mask=n_offset <
            K, other=0.0)
        a_fp8 = a * a_scale
        a_fp8 = tl.clamp(a_fp8, -MAX_FP8, MAX_FP8)
        a_fp8.to(TL_FP8_DTYPE)
        tl.store(A_fp8 + a_fp8_offset_base + n_offset * stride_ok, a_fp8,
            mask=n_offset < K)
        n_offset += BLOCK_SIZE
