@triton.jit
def softmax_kernel_online(output_ptr, input_ptr, M, N, K, TILE_N: tl.
    constexpr, TILE_K: tl.constexpr):
    pid_k = tl.program_id(0)
    pid_m = tl.program_id(1)
    k_offsets = pid_k * TILE_K + tl.arange(0, TILE_K)
    m = tl.full([TILE_N, TILE_K], value=float('-inf'), dtype=tl.float32)
    z = tl.full([TILE_N, TILE_K], value=0.0, dtype=tl.float32)
    for start_n in range(0, N, TILE_N):
        n_offsets = start_n + tl.arange(0, TILE_N)
        offsets = pid_m * N * K + n_offsets[:, None] * K + k_offsets
        mask = (n_offsets[:, None] < N) & (k_offsets < K)
        inp = tl.load(input_ptr + offsets, mask=mask, other=-float('inf'))
        m_new = tl.maximum(m, inp)
        alpha = tl.exp(m - m_new)
        z = z * alpha + tl.exp(inp - m_new)
        m = m_new
    m_reduced = tl.max(m, 0)
    z = tl.sum(z * tl.exp(m - m_reduced[None, :]), 0)
    m = m_reduced
    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)
        offsets = pid_m * N * K + n_offsets[:, None] * K + k_offsets
        mask = (n_offsets[:, None] < N) & (k_offsets[None, :] < K)
        inp = tl.load(input_ptr + offsets, mask=mask, other=-float('inf'))
        o = tl.exp(inp - m[None, :]) / z[None, :]
        tl.store(output_ptr + offsets, o, mask=mask)
