code: def _fwd_kernel(
      Q,
      K,
      V,
      sm_scale,
      TMP,
      Out,
      stride_qz,
      stride_qh,
      stride_qm,
      stride_qk,
      stride_kz,
      stride_kh,
      stride_kn,
      stride_kk,
      stride_vz,
      stride_vh,
      stride_vk,
      stride_vn,
      stride_oz,
      stride_oh,
      stride_om,
      stride_on,
      Z,
      H,
      N_CTX,
      BLOCK_M: tl.constexpr,
      BLOCK_DMODEL: tl.constexpr,
      BLOCK_N: tl.constexpr,
  ):
      start_m = tl.program_id(0)
      off_hz = tl.program_id(1)
      # initialize offsets
      offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
      offs_n = tl.arange(0, BLOCK_N)
      offs_d = tl.arange(0, BLOCK_DMODEL)
      off_q = (
          off_hz * stride_qh
          + offs_m[:, None] * stride_qm
          + offs_d[None, :] * stride_qk
      )
      off_k = (
          off_hz * stride_kh
          + offs_n[:, None] * stride_kn
          + offs_d[None, :] * stride_kk
      )
      off_v = (
          off_hz * stride_vh
          + offs_n[:, None] * stride_qm
          + offs_d[None, :] * stride_qk
      )
      # Initialize pointers to Q, K, V
      q_ptrs = Q + off_q
      k_ptrs = K + off_k
      v_ptrs = V + off_v
      # initialize pointer to m and l
      t_ptrs = TMP + off_hz * N_CTX + offs_m
      m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
      l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
      acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
      # load q: it will stay in SRAM throughout
      q = tl.load(q_ptrs)
      # loop over k, v and update accumulator
      for start_n in range(0, N_CTX, BLOCK_N):
          start_n = tl.multiple_of(start_n, BLOCK_N)
          # -- compute qk ----
          k = tl.load(k_ptrs + start_n * stride_kn)
  
          qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
          qk += tl.dot(q, k, trans_b=True)
          qk *= sm_scale
          # -- compute m_ij, p, l_ij
          m_ij = tl.max(qk, 1)
          p = tl.exp(qk - m_ij[:, None])
          l_ij = tl.sum(p, 1)
          # -- update m_i and l_i
          m_i_new = tl.maximum(m_i, m_ij)
          alpha = tl.exp(m_i - m_i_new)
          beta = tl.exp(m_ij - m_i_new)
          l_i_new = alpha * l_i + beta * l_ij
          # -- update output accumulator --
          # scale p
          p_scale = beta / l_i_new
          p = p * p_scale[:, None]
          # scale acc
          acc_scale = l_i / l_i_new * alpha
          tl.store(t_ptrs, acc_scale)
          acc_scale = tl.load(
              t_ptrs
          )  # BUG: have to store and immediately load
          acc = acc * acc_scale[:, None]
          # update acc
          v = tl.load(v_ptrs + start_n * stride_vk)
          p = p.to(tl.float16)
          acc += tl.dot(p, v)
          # update m_i and l_i
          l_i = l_i_new
          m_i = m_i_new
      # initialize pointers to output
      offs_n = tl.arange(0, BLOCK_DMODEL)
      off_o = (
          off_hz * stride_oh
          + offs_m[:, None] * stride_om
          + offs_n[None, :] * stride_on
      )
      out_ptrs = Out + off_o
      tl.store(out_ptrs, acc)
  
imports: import torch
  import triton
  import triton.language
label_desc: Use image node_id 1 for calling a global function with example usage: _fwd_kernel(Q, K, V, sm_scale, TMP, Out, stride_qz, stride_qh, stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vk, stride_vn, stride_oz, stride_oh, stride_om, stride_on, Z, H, N_CTX, BLOCK_M, BLOCK_DMODEL, BLOCK_N) without return types