@triton.jit
def softmax_kernel(output_ptr, input_ptr, M, N, TILE_N: tl.constexpr):
    pid_m = tl.program_id(0)
    n_offsets = tl.arange(0, TILE_N)
    offset = pid_m * N + n_offsets
    input_ptrs = input_ptr + offset
    mask = n_offsets < N
    inp = tl.load(input_ptrs, mask=mask, other=-float('inf')).to(output_ptr
        .dtype.element_ty)
    m = tl.max(inp, 0)
    e = tl.exp(inp - m)
    z = tl.sum(e, 0)
    out = e / z
    output_ptrs = output_ptr + offset
    tl.store(output_ptrs, out, mask=mask)
