code: def _fwd_kernel(
      Q,
      K,
      V,
      sm_scale,  #
      L,
      M,  #
      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_qh
          + offs_n[None, :] * stride_kn
          + offs_d[:, None] * stride_kk
      )
      off_v = (
          off_hz * stride_qh
          + 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
      m_prev = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
      l_prev = 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, (start_m + 1) * BLOCK_M, BLOCK_N):
          # -- compute qk ----
          k = tl.load(k_ptrs)
          qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
          qk += tl.dot(q, k)
          qk *= sm_scale
          qk = tl.where(
              offs_m[:, None] >= (start_n + offs_n[None, :]),
              qk,
              float("-inf"),
          )
          # compute new m
          m_curr = tl.maximum(tl.max(qk, 1), m_prev)
          # correct old l
          l_prev *= tl.exp(m_prev - m_curr)
          # attention weights
          p = tl.exp(qk - m_curr[:, None])
          l_curr = tl.sum(p, 1) + l_prev
          # rescale operands of matmuls
          l_rcp = 1.0 / l_curr
          p *= l_rcp[:, None]
          acc *= (l_prev * l_rcp)[:, None]
          # update acc
          p = p.to(Q.dtype.element_ty)
          v = tl.load(v_ptrs)
          acc += tl.dot(p, v)
          # update m_i and l_i
          l_prev = l_curr
          m_prev = m_curr
          # update pointers
          k_ptrs += BLOCK_N * stride_kn
          v_ptrs += BLOCK_N * stride_vk
      # rematerialize offsets to save registers
      start_m = tl.program_id(0)
      offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
      # write back l and m
      l_ptrs = L + off_hz * N_CTX + offs_m
      m_ptrs = M + off_hz * N_CTX + offs_m
      tl.store(l_ptrs, l_prev)
      tl.store(m_ptrs, m_prev)
      # 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)
  
file_path: triton/python/test/unit/hopper/test_persistent_warp_specialized_fused-attention.py
imports: import pytest
  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, L, M, 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