code: def _attn_bwd(
      Q,
      K,
      V,
      sm_scale,  #
      DO,  #
      DQ,
      DK,
      DV,  #
      M,
      D,
      # shared by Q/K/V/DO.
      stride_z,
      stride_h,
      stride_tok,
      stride_d,  #
      H,
      N_CTX,  #
      BLOCK_M1: tl.constexpr,  #
      BLOCK_N1: tl.constexpr,  #
      BLOCK_M2: tl.constexpr,  #
      BLOCK_N2: tl.constexpr,  #
      BLK_SLICE_FACTOR: tl.constexpr,  #
      BLOCK_DMODEL: tl.constexpr,
  ):
      LN2: tl.constexpr = 0.6931471824645996  # = ln(2)
  
      bhid = tl.program_id(2)
      off_chz = (bhid * N_CTX).to(tl.int64)
      adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(
          tl.int64
      )
      pid = tl.program_id(0)
  
      # offset pointers for batch/head
      Q += adj
      K += adj
      V += adj
      DO += adj
      DQ += adj
      DK += adj
      DV += adj
      M += off_chz
      D += off_chz
  
      # load scales
      offs_k = tl.arange(0, BLOCK_DMODEL)
  
      start_n = pid * BLOCK_N1
      start_m = start_n
  
      MASK_BLOCK_M1: tl.constexpr = BLOCK_M1 // BLK_SLICE_FACTOR
      offs_n = start_n + tl.arange(0, BLOCK_N1)
  
      dv = tl.zeros([BLOCK_N1, BLOCK_DMODEL], dtype=tl.float32)
      dk = tl.zeros([BLOCK_N1, BLOCK_DMODEL], dtype=tl.float32)
  
      # load K and V: they stay in SRAM throughout the inner loop.
      k = tl.load(
          K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
      )
      v = tl.load(
          V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
      )
  
      num_steps = BLOCK_N1 // MASK_BLOCK_M1
  
      dk, dv = _attn_bwd_dkdv(
          dk,
          dv,  #
          Q,
          k,
          v,
          sm_scale,  #
          DO,  #
          M,
          D,  #
          stride_tok,
          stride_d,  #
          H,
          N_CTX,  #
          MASK_BLOCK_M1,
          BLOCK_N1,
          BLOCK_DMODEL,  #
          start_n,
          start_m,
          num_steps,  #
          MASK=True,  #
      )
  
      start_m += num_steps * MASK_BLOCK_M1
      num_steps = (N_CTX - start_m) // BLOCK_M1
  
      # Compute dK and dV for non-masked blocks.
      dk, dv = _attn_bwd_dkdv(  #
          dk,
          dv,  #
          Q,
          k,
          v,
          sm_scale,  #
          DO,  #
          M,
          D,  #
          stride_tok,
          stride_d,  #
          H,
          N_CTX,  #
          BLOCK_M1,
          BLOCK_N1,
          BLOCK_DMODEL,  #
          start_n,
          start_m,
          num_steps,  #
          MASK=False,  #
      )
  
      dv_ptrs = (
          DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
      )
      tl.store(dv_ptrs, dv)
  
      # Write back dK.
      dk *= sm_scale
      dk_ptrs = (
          DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
      )
      tl.store(dk_ptrs, dk)
  
      # THIS BLOCK DOES DQ:
      start_m = pid * BLOCK_M2
      end_n = start_m + BLOCK_M2
  
      MASK_BLOCK_N2: tl.constexpr = BLOCK_N2 // BLK_SLICE_FACTOR
      offs_m = start_m + tl.arange(0, BLOCK_M2)
  
      q = tl.load(
          Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
      )
      dq = tl.zeros([BLOCK_M2, BLOCK_DMODEL], dtype=tl.float32)
      do = tl.load(
          DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
      )
  
      m = tl.load(M + offs_m)
      m = m[:, None]
  
      # Compute dQ for masked (diagonal) blocks.
      # NOTE: This code scans each row of QK^T backward (from right to left,
      # but inside each call to _attn_bwd_dq, from left to right), but that's
      # not due to anything important.  I just wanted to reuse the loop
      # structure for dK & dV above as much as possible.
      num_steps = BLOCK_M2 // MASK_BLOCK_N2
      dq = _attn_bwd_dq(
          dq,
          q,
          K,
          V,  #
          do,
          m,
          D,  #
          stride_tok,
          stride_d,  #
          H,
          N_CTX,  #
          BLOCK_M2,
          MASK_BLOCK_N2,
          BLOCK_DMODEL,  #
          start_m,
          end_n - num_steps * MASK_BLOCK_N2,
          num_steps,  #
          MASK=True,  #
      )
      end_n -= num_steps * MASK_BLOCK_N2
      # stage 2
      num_steps = end_n // BLOCK_N2
      dq = _attn_bwd_dq(
          dq,
          q,
          K,
          V,  #
          do,
          m,
          D,  #
          stride_tok,
          stride_d,  #
          H,
          N_CTX,  #
          BLOCK_M2,
          BLOCK_N2,
          BLOCK_DMODEL,  #
          start_m,
          end_n - num_steps * BLOCK_N2,
          num_steps,  #
          MASK=False,  #
      )
      # Write back dQ.
      dq_ptrs = (
          DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
      )
      dq *= LN2
      tl.store(dq_ptrs, dq)
  
file_path: triton/python/tutorials/06-fused-attention.py
imports: import pytest
  import torch
  import triton
  import triton.language
label_desc: Use image node_id 6 for calling a global function with example usage: _attn_bwd(Q, K, V, sm_scale, DO, DQ, DK, DV, M, D, stride_z, stride_h, stride_tok, stride_d, H, N_CTX, BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2, BLK_SLICE_FACTOR, BLOCK_DMODEL) without return types