@triton.autotune(configs=[triton.Config({'threads_per_warp': 32}, num_warps
    =32), triton.Config({'threads_per_warp': 32}, num_warps=16), triton.
    Config({'threads_per_warp': 32}, num_warps=8), triton.Config({
    'threads_per_warp': 32}, num_warps=4), triton.Config({
    'threads_per_warp': 16}, num_warps=64), triton.Config({
    'threads_per_warp': 16}, num_warps=32), triton.Config({
    'threads_per_warp': 16}, num_warps=16), triton.Config({
    'threads_per_warp': 16}, num_warps=8), triton.Config({
    'threads_per_warp': 16}, num_warps=4)], key=['BLOCK_SIZE_X',
    'BLOCK_SIZE_Y'])
@triton.jit
def softmax_kernel(output_ptr, input_ptr, input_row_stride,
    output_row_stride, n_cols, BLOCK_SIZE_X: tl.constexpr, BLOCK_SIZE_Y: tl
    .constexpr):
    row_idx = tl.program_id(0) * BLOCK_SIZE_Y
    row_start_ptr = input_ptr + row_idx * input_row_stride
    col_offsets = tl.arange(0, BLOCK_SIZE_X)
    row_offsets = tl.arange(0, BLOCK_SIZE_Y)
    offsets = col_offsets[None, :] + row_offsets[:, None] * input_row_stride
    input_ptrs = row_start_ptr + offsets
    mask = col_offsets[None, :] < n_cols
    row = tl.load(input_ptrs, mask=mask, other=-float('inf'))
    row_minus_max = row - tl.max(row, axis=1)[:, None]
    numerator = tl.exp(row_minus_max)
    denominator = tl.sum(numerator, axis=1)[:, None]
    softmax_output = numerator / denominator
    output_row_start_ptr = output_ptr + row_idx * output_row_stride
    output_ptrs = output_row_start_ptr + offsets
    tl.store(output_ptrs, softmax_output, mask=mask)
