@triton.jit
def softmax_kernel_online_v2(output_ptr, input_ptr, M, N, TILE_N: tl.constexpr
    ):
    pid_m = tl.program_id(0)
    m = tl.full((TILE_N,), value=-float('inf'), dtype=output_ptr.dtype.
        element_ty)
    z = tl.full((TILE_N,), value=0, dtype=output_ptr.dtype.element_ty)
    for start_n in range(0, N, TILE_N):
        n_offsets = start_n + 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)
        new_m = tl.maximum(m, inp)
        new_z = tl.exp(m - new_m) * z + tl.exp(inp - new_m)
        m = new_m
        z = new_z
    final_m = tl.max(m, 0)
    z = tl.sum(tl.exp(m - final_m) * z)
    m = final_m
    previous_multiple = prev_multiple_of(N, TILE_N)
    for start_n in range(0, N, TILE_N):
        n_offsets = previous_multiple - start_n + 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)
        e = tl.exp(inp - m)
        out = e / z
        output_ptrs = output_ptr + offset
        tl.store(output_ptrs, out, mask=mask)
