code: def _bwd_kernel(
      Q,
      K,
      V,
      sm_scale,
      Out,
      DO,  #
      DQ,
      DK,
      DV,  #
      L,
      M,  #
      D,  #
      stride_qz,
      stride_qh,
      stride_qm,
      stride_qk,  #
      stride_kz,
      stride_kh,
      stride_kn,
      stride_kk,  #
      stride_vz,
      stride_vh,
      stride_vk,
      stride_vn,  #
      Z,
      H,
      N_CTX,  #
      num_block,  #
      BLOCK_M: tl.constexpr,
      BLOCK_DMODEL: tl.constexpr,  #
      BLOCK_N: tl.constexpr,  #
  ):
      off_hz = tl.program_id(0)
      off_z = off_hz // H
      off_h = off_hz % H
      # offset pointers for batch/head
      Q += off_z * stride_qz + off_h * stride_qh
      K += off_z * stride_qz + off_h * stride_qh
      V += off_z * stride_qz + off_h * stride_qh
      DO += off_z * stride_qz + off_h * stride_qh
      DQ += off_z * stride_qz + off_h * stride_qh
      DK += off_z * stride_qz + off_h * stride_qh
      DV += off_z * stride_qz + off_h * stride_qh
      for start_n in range(0, num_block):
          lo = start_n * BLOCK_M
          # initialize row/col offsets
          offs_qm = lo + tl.arange(0, BLOCK_M)
          offs_n = start_n * BLOCK_M + tl.arange(0, BLOCK_M)
          offs_m = tl.arange(0, BLOCK_N)
          offs_k = tl.arange(0, BLOCK_DMODEL)
          # initialize pointers to value-like data
          q_ptrs = Q + (
              offs_qm[:, None] * stride_qm + offs_k[None, :] * stride_qk
          )
          k_ptrs = K + (
              offs_n[:, None] * stride_kn + offs_k[None, :] * stride_kk
          )
          v_ptrs = V + (
              offs_n[:, None] * stride_qm + offs_k[None, :] * stride_qk
          )
          do_ptrs = DO + (
              offs_qm[:, None] * stride_qm + offs_k[None, :] * stride_qk
          )
          dq_ptrs = DQ + (
              offs_qm[:, None] * stride_qm + offs_k[None, :] * stride_qk
          )
          # pointer to row-wise quantities in value-like data
          D_ptrs = D + off_hz * N_CTX
          m_ptrs = M + off_hz * N_CTX
          # initialize dv amd dk
          dv = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
          dk = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
          # k and v stay in SRAM throughout
          k = tl.load(k_ptrs)
          v = tl.load(v_ptrs)
          # loop over rows
          for start_m in range(lo, num_block * BLOCK_M, BLOCK_M):
              offs_m_curr = start_m + offs_m
              # load q, k, v, do on-chip
              q = tl.load(q_ptrs)
              # recompute p = softmax(qk, dim=-1).T
              # NOTE: `do` is pre-divided by `l`; no normalization here
              qk = tl.dot(q, tl.trans(k))
              qk = tl.where(
                  offs_m_curr[:, None] >= (offs_n[None, :]),
                  qk,
                  float("-inf"),
              )
              m = tl.load(m_ptrs + offs_m_curr)
              p = tl.exp(qk * sm_scale - m[:, None])
              # compute dv
              do = tl.load(do_ptrs)
              dv += tl.dot(tl.trans(p.to(Q.dtype.element_ty)), do)
              # compute dp = dot(v, do)
              Di = tl.load(D_ptrs + offs_m_curr)
              dp = (
                  tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
                  - Di[:, None]
              )
              dp += tl.dot(do, tl.trans(v))
              # compute ds = p * (dp - delta[:, None])
              ds = p * dp * sm_scale
              # compute dk = dot(ds.T, q)
              dk += tl.dot(tl.trans(ds.to(Q.dtype.element_ty)), q)
              # compute dq
              dq = tl.load(dq_ptrs)
              dq += tl.dot(ds.to(Q.dtype.element_ty), k)
              tl.store(dq_ptrs, dq)
              # increment pointers
              dq_ptrs += BLOCK_M * stride_qm
              q_ptrs += BLOCK_M * stride_qm
              do_ptrs += BLOCK_M * stride_qm
          # write-back
          dv_ptrs = DV + (
              offs_n[:, None] * stride_qm + offs_k[None, :] * stride_qk
          )
          dk_ptrs = DK + (
              offs_n[:, None] * stride_kn + offs_k[None, :] * stride_kk
          )
          tl.store(dv_ptrs, dv)
          tl.store(dk_ptrs, dk)
  
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 3 for calling a global function with example usage: _bwd_kernel(Q, K, V, sm_scale, Out, DO, DQ, DK, DV, L, M, D, stride_qz, stride_qh, stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vk, stride_vn, Z, H, N_CTX, num_block, BLOCK_M, BLOCK_DMODEL, BLOCK_N) without return types