@triton.jit
def softmax_kernel(out_ptr, in_ptr, logz_ptr, M, N, TILE_N: tl.constexpr):
    pid_n = tl.program_id(0)
    pid_m = tl.program_id(1)
    n_offsets = pid_n * TILE_N + tl.arange(0, TILE_N)
    offset = pid_m * N + n_offsets
    mask = n_offsets < N
    inp = tl.load(in_ptr + offset, mask=mask, other=-float('inf')).to(out_ptr
        .dtype.element_ty)
    logz = tl.load(logz_ptr + pid_m).to(out_ptr.dtype.element_ty)
    out = tl.exp(inp - logz)
    tl.store(out_ptr + offset, out, mask=mask)
