code: def _attn_bwd_dkdv(
      dk,
      dv,  #
      Q,
      k,
      v,
      sm_scale,  #
      DO,  #
      M,
      D,  #
      # shared by Q/K/V/DO.
      stride_tok,
      stride_d,  #
      H,
      N_CTX,
      BLOCK_M1: tl.constexpr,  #
      BLOCK_N1: tl.constexpr,  #
      BLOCK_DMODEL: tl.constexpr,  #
      # Filled in by the wrapper.
      start_n,
      start_m,
      num_steps,  #
      MASK: tl.constexpr,
  ):
      offs_m = start_m + tl.arange(0, BLOCK_M1)
      offs_n = start_n + tl.arange(0, BLOCK_N1)
      offs_k = tl.arange(0, BLOCK_DMODEL)
      qT_ptrs = (
          Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
      )
      do_ptrs = (
          DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
      )
      # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
      tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
      curr_m = start_m
      step_m = BLOCK_M1
      for blk_idx in range(num_steps):
          qT = tl.load(qT_ptrs)
          # Load m before computing qk to reduce pipeline stall.
          offs_m = curr_m + tl.arange(0, BLOCK_M1)
          m = tl.load(M + offs_m)
          qkT = tl.dot(k, qT)
          pT = tl.math.exp2(qkT - m[None, :])
          # Autoregressive masking.
          if MASK:
              mask = offs_m[None, :] >= offs_n[:, None]
              pT = tl.where(mask, pT, 0.0)
          do = tl.load(do_ptrs)
          # Compute dV.
          ppT = pT
          ppT = ppT.to(tl.float16)
          dv += tl.dot(ppT, do)
          # D (= delta) is pre-divided by ds_scale.
          Di = tl.load(D + offs_m)
          # Compute dP and dS.
          dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
          dsT = pT * (dpT - Di[None, :])
          dsT = dsT.to(tl.float16)
          dk += tl.dot(dsT, tl.trans(qT))
          # Increment pointers.
          curr_m += step_m
          qT_ptrs += step_m * stride_tok
          do_ptrs += step_m * stride_tok
      return dk, dv
  
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 4 for calling a global function with example usage: _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) and returns: dk,  dv