code: def _softmax(
      Y,
      X,
      M,
      stride_ym,
      stride_yn,
      stride_xm,
      stride_xn,
      stride_mn,
      K,
      # Meta-params
      depth: tl.constexpr,
      causal: tl.constexpr,
      use_mask: tl.constexpr,
      log: tl.constexpr,
  ):
      # fmt: om
  
      """
      Fused softmax kernel over a 3d tensor.
      The softmax is applied over the last dimension, meaning that this is equivalent to torch.softmax(tensor, dim=-1)
  
      Note, if the last dimension is large, say 128K elements, the kernel compile time can shot up to many minutes when
      the kernel is run for the first time.
      """
  
      m = tl.program_id(0)
      n = tl.program_id(1)
  
      # col indices
      k = tl.arange(0, depth)
  
      # the memory address of all the elements that we want to load can be computed as follows
      x_ptrs = X + m * stride_xm + n * stride_xn + k
  
      # load input data; pad out-of-bounds elements with 0
      io_mask = k < K
  
      # Causal - 1: skip on the loads directly
      if causal:
          io_mask = io_mask & (k <= n)
  
      x = tl.load(x_ptrs, mask=io_mask, other=float("-inf")).to(
          tl.float32
      )
  
      # Causal - 2: enforce correctness over a couple of misloaded values
      if causal:
          off = float("-inf")
          off = off.to(x.dtype)  # type: ignore
          x = tl.where(k > n, off, x)
  
      if use_mask:
          mask_ptrs = M + n * stride_mn + k
          add_mask = tl.load(
              mask_ptrs, io_mask, other=float("-inf")
          ).to(tl.float32)
          x += add_mask
  
      # compute numerically-stable softmax
      z = x - tl.max(x, axis=0)
      num = tl.exp(z)
      denom = tl.sum(num, axis=0)
  
      if log:
          y = z - tl.log(denom)
      else:
          y = num / denom
  
      # write back to Y.
      # we only write once, hence the "fused" softmax naming
      y_ptrs = Y + m * stride_ym + n * stride_yn + k
  
      # technically we could write only the lower triangular matrix in the causal case
      # but this is deemed to error prone
      tl.store(y_ptrs, y, mask=k < K)
  
imports: import triton
  import triton.language
label_desc: Use image node_id 2 for calling a global function with example usage: _softmax(Y, X, M, stride_ym, stride_yn, stride_xm, stride_xn, stride_mn, K, depth, causal, use_mask, log) without return types