code: def _bwd_kernel(
      Q,
      K,
      V,
      sm_scale,  #
      Out,
      DO,  #
      DQ,
      DK,
      DV,  #
      L,  #
      D,  #
      stride_dqa,
      stride_qz,
      stride_qh,
      stride_qm,
      stride_qk,  #
      stride_kz,
      stride_kh,
      stride_kn,
      stride_kk,  #
      stride_vz,
      stride_vh,
      stride_vn,
      stride_vk,  #
      Z,
      H,
      N_CTX,  #
      Z_H_N_CTX,  #
      SQ_Z_H_N_CTX,  #
      BLOCK_M: tl.constexpr,
      BLOCK_DMODEL: tl.constexpr,  #
      BLOCK_N: tl.constexpr,  #
      SEQUENCE_PARALLEL: tl.constexpr,  #
      CAUSAL: tl.constexpr,  #
      MMA_V3: tl.constexpr,  #
  ):
      qk_scale = sm_scale * 1.44269504
      off_hz = tl.program_id(0)
      off_z = off_hz // H
      off_h = off_hz % H
  
      Q_block_ptr = tl.make_block_ptr(
          base=Q,
          shape=(Z_H_N_CTX, BLOCK_DMODEL),
          strides=(stride_qm, stride_qk),
          offsets=(0, 0),
          block_shape=(BLOCK_M, BLOCK_DMODEL),
          order=(1, 0),
      )
      K_block_ptr = tl.make_block_ptr(
          base=K,
          shape=(Z_H_N_CTX, BLOCK_DMODEL),
          strides=(stride_kn, stride_kk),
          offsets=(0, 0),
          block_shape=(BLOCK_M, BLOCK_DMODEL),
          order=(1, 0),
      )
      V_block_ptr = tl.make_block_ptr(
          base=V,
          shape=(Z_H_N_CTX, BLOCK_DMODEL),
          strides=(stride_vn, stride_vk),
          offsets=(0, 0),
          block_shape=(BLOCK_M, BLOCK_DMODEL),
          order=(1, 0),
      )
      DO_block_ptr = tl.make_block_ptr(
          base=DO,
          shape=(Z_H_N_CTX, BLOCK_DMODEL),
          strides=(stride_qm, stride_qk),
          offsets=(0, 0),
          block_shape=(BLOCK_M, BLOCK_DMODEL),
          order=(1, 0),
      )
      if SEQUENCE_PARALLEL:
          DQ_block_ptr = tl.make_block_ptr(
              base=DQ,
              shape=(SQ_Z_H_N_CTX, BLOCK_DMODEL),
              strides=(stride_qm, stride_qk),
              offsets=(0, 0),
              block_shape=(BLOCK_M, BLOCK_DMODEL),
              order=(1, 0),
          )
      else:
          DQ_block_ptr = tl.make_block_ptr(
              base=DQ,
              shape=(Z_H_N_CTX, BLOCK_DMODEL),
              strides=(stride_qm, stride_qk),
              offsets=(0, 0),
              block_shape=(BLOCK_M, BLOCK_DMODEL),
              order=(1, 0),
          )
  
      DK_block_ptr = tl.make_block_ptr(
          base=DK,
          shape=(Z_H_N_CTX, BLOCK_DMODEL),
          strides=(stride_kn, stride_kk),
          offsets=(0, 0),
          block_shape=(BLOCK_M, BLOCK_DMODEL),
          order=(1, 0),
      )
      DV_block_ptr = tl.make_block_ptr(
          base=DV,
          shape=(Z_H_N_CTX, BLOCK_DMODEL),
          strides=(stride_vn, stride_vk),
          offsets=(0, 0),
          block_shape=(BLOCK_M, BLOCK_DMODEL),
          order=(1, 0),
      )
  
      num_block_n = tl.cdiv(N_CTX, BLOCK_N)
      if not SEQUENCE_PARALLEL:
          for start_n in range(0, num_block_n):
              _bwd_kernel_one_col_block(
                  Q,
                  K,
                  V,
                  sm_scale,
                  qk_scale,
                  Out,
                  DO,  #
                  DQ,
                  DK,
                  DV,  #
                  L,  #
                  D,  #
                  Q_block_ptr,
                  K_block_ptr,
                  V_block_ptr,  #
                  DO_block_ptr,
                  DQ_block_ptr,
                  DK_block_ptr,
                  DV_block_ptr,  #
                  stride_dqa,
                  stride_qz,
                  stride_qh,
                  stride_qm,
                  stride_qk,  #
                  stride_kz,
                  stride_kh,
                  stride_kn,
                  stride_kk,  #
                  stride_vz,
                  stride_vh,
                  stride_vn,
                  stride_vk,  #
                  Z,
                  H,
                  N_CTX,  #
                  off_h,
                  off_z,
                  off_hz,
                  start_n,
                  num_block_n,  #
                  BLOCK_M=BLOCK_M,
                  BLOCK_DMODEL=BLOCK_DMODEL,  #
                  BLOCK_N=BLOCK_N,  #
                  SEQUENCE_PARALLEL=SEQUENCE_PARALLEL,  #
                  CAUSAL=CAUSAL,  #
                  MMA_V3=MMA_V3,  #
              )
      else:
          start_n = tl.program_id(1)
          _bwd_kernel_one_col_block(
              Q,
              K,
              V,
              sm_scale,
              qk_scale,
              Out,
              DO,  #
              DQ,
              DK,
              DV,  #
              L,  #
              D,  #
              Q_block_ptr,
              K_block_ptr,
              V_block_ptr,  #
              DO_block_ptr,
              DQ_block_ptr,
              DK_block_ptr,
              DV_block_ptr,  #
              stride_dqa,
              stride_qz,
              stride_qh,
              stride_qm,
              stride_qk,  #
              stride_kz,
              stride_kh,
              stride_kn,
              stride_kk,  #
              stride_vz,
              stride_vh,
              stride_vn,
              stride_vk,  #
              Z,
              H,
              N_CTX,  #
              off_h,
              off_z,
              off_hz,
              start_n,
              num_block_n,  #
              BLOCK_M=BLOCK_M,
              BLOCK_DMODEL=BLOCK_DMODEL,  #
              BLOCK_N=BLOCK_N,  #
              SEQUENCE_PARALLEL=SEQUENCE_PARALLEL,  #
              CAUSAL=CAUSAL,  #
              MMA_V3=MMA_V3,  #
          )
  
file_path: triton/python/triton/ops/flash_attention.py
imports: import torch
  from ..None import cdiv, jit
  from ..None import language
label_desc: Use image node_id 4 for calling a global function with example usage: _bwd_kernel(Q, K, V, sm_scale, Out, DO, DQ, DK, DV, L, D, stride_dqa, stride_qz, stride_qh, stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vn, stride_vk, Z, H, N_CTX, Z_H_N_CTX, SQ_Z_H_N_CTX, BLOCK_M, BLOCK_DMODEL, BLOCK_N, SEQUENCE_PARALLEL, CAUSAL, MMA_V3) without return types