@triton.jit
def softmax_kernel(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)
    n_offsets = tl.arange(0, TILE_N)
    offset = pid_m * N * K + n_offsets[:, None] * K + k_offsets
    mask = (n_offsets[:, None] < N) & (k_offsets < K)
    input_ptrs = input_ptr + offset
    inp = tl.load(input_ptrs, mask=mask, other=-float('inf')).to(output_ptr
        .type.element_ty)
    m = tl.max(inp, 0)
    e = tl.exp(inp - m[None, :])
    z = tl.sum(e, 0)
    out = e / z
    output_ptrs = output_ptr + offset
    tl.store(output_ptrs, out, mask=mask)
