code: def _bwd_kernel_one_col_block(
      start_n,
      Q,
      K,
      V,
      Bias,
      DO,
      DQ,
      DK,
      DV,
      LSE,
      D,
      softmax_scale,
      stride_qm,
      stride_kn,
      stride_vn,
      stride_bm,
      stride_dom,
      stride_dqm,
      stride_dkn,
      stride_dvn,
      seqlen_q,
      seqlen_k,
      headdim,
      ATOMIC_ADD: tl.constexpr,
      BIAS_TYPE: tl.constexpr,
      IS_CAUSAL: tl.constexpr,
      BLOCK_HEADDIM: tl.constexpr,
      EVEN_M: tl.constexpr,
      EVEN_N: tl.constexpr,
      EVEN_HEADDIM: tl.constexpr,
      BLOCK_M: tl.constexpr,
      BLOCK_N: tl.constexpr,
  ):
      # We need to make sure begin_m is a multiple of BLOCK_M (not BLOCK_N)
      begin_m = (
          0
          if not IS_CAUSAL
          else ((start_n * BLOCK_N) // BLOCK_M) * BLOCK_M
      )
      # initialize row/col offsets
      offs_qm = begin_m + tl.arange(0, BLOCK_M)
      offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
      offs_m = tl.arange(0, BLOCK_M)
      offs_d = tl.arange(0, BLOCK_HEADDIM)
      # initialize pointers to value-like data
      q_ptrs = Q + (offs_qm[:, None] * stride_qm + offs_d[None, :])
      k_ptrs = K + (offs_n[:, None] * stride_kn + offs_d[None, :])
      v_ptrs = V + (offs_n[:, None] * stride_vn + offs_d[None, :])
      do_ptrs = DO + (offs_qm[:, None] * stride_dom + offs_d[None, :])
      dq_ptrs = DQ + (offs_qm[:, None] * stride_dqm + offs_d[None, :])
      if BIAS_TYPE == "vector":
          b_ptrs = Bias + offs_n
      elif BIAS_TYPE == "matrix":
          b_ptrs = Bias + (
              offs_qm[:, None] * stride_bm + offs_n[None, :]
          )
      # initialize dv and dk
      dv = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
      dk = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
      # There seems to be some problem with Triton pipelining that makes results wrong for
      # headdim=64, seqlen=(113, 255), bias_type='matrix'. In this case the for loop
      # may have zero step, and pipelining with the bias matrix could screw it up.
      # So we just exit early.
      if begin_m >= seqlen_q:
          dv_ptrs = DV + (
              offs_n[:, None] * stride_dvn + offs_d[None, :]
          )
          dk_ptrs = DK + (
              offs_n[:, None] * stride_dkn + offs_d[None, :]
          )
          _bwd_store_dk_dv(
              dk_ptrs,
              dv_ptrs,
              dk,
              dv,
              offs_n,
              offs_d,
              seqlen_k,
              headdim,
              EVEN_M=EVEN_M,
              EVEN_N=EVEN_N,
              EVEN_HEADDIM=EVEN_HEADDIM,
          )
          return
      # k and v stay in SRAM throughout
      # [2022-10-30] TD: Same bug as the fwd. In the case of EVEN_N=True and EVEN_M=False,
      # if we just call tl.load(k_ptrs), we get the wrong output!
      if EVEN_N & EVEN_M:
          if EVEN_HEADDIM:
              k = tl.load(k_ptrs)
              v = tl.load(v_ptrs)
          else:
              k = tl.load(
                  k_ptrs, mask=offs_d[None, :] < headdim, other=0.0
              )
              v = tl.load(
                  v_ptrs, mask=offs_d[None, :] < headdim, other=0.0
              )
      else:
          if EVEN_HEADDIM:
              k = tl.load(
                  k_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0
              )
              v = tl.load(
                  v_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0
              )
          else:
              k = tl.load(
                  k_ptrs,
                  mask=(offs_n[:, None] < seqlen_k)
                  & (offs_d[None, :] < headdim),
                  other=0.0,
              )
              v = tl.load(
                  v_ptrs,
                  mask=(offs_n[:, None] < seqlen_k)
                  & (offs_d[None, :] < headdim),
                  other=0.0,
              )
      # loop over rows
      num_block_m = tl.cdiv(seqlen_q, BLOCK_M)
      for start_m in range(begin_m, num_block_m * BLOCK_M, BLOCK_M):
          start_m = tl.multiple_of(start_m, BLOCK_M)
          offs_m_curr = start_m + offs_m
          # load q, k, v, do on-chip
          # Same bug as below. Otherwise gives wrong result for headdim=40, seqlen=(128, 117)
          if EVEN_M & EVEN_HEADDIM:
              q = tl.load(q_ptrs)
          else:
              if EVEN_HEADDIM:
                  q = tl.load(
                      q_ptrs,
                      mask=offs_m_curr[:, None] < seqlen_q,
                      other=0.0,
                  )
              else:
                  q = tl.load(
                      q_ptrs,
                      mask=(offs_m_curr[:, None] < seqlen_q)
                      & (offs_d[None, :] < headdim),
                      other=0.0,
                  )
          # recompute p = softmax(qk, dim=-1).T
          qk = tl.dot(q, k, trans_b=True)
          # Trying to combine the two masks seem to make the result wrong
          if (
              not EVEN_N
          ):  # Need to mask out otherwise the softmax is wrong
              qk = tl.where(
                  offs_n[None, :] < seqlen_k, qk, float("-inf")
              )
          if IS_CAUSAL:
              qk = tl.where(
                  offs_m_curr[:, None] >= (offs_n[None, :]),
                  qk,
                  float("-inf"),
              )
          if BIAS_TYPE != "none":
              tl.debug_barrier()  # Race condition otherwise
              if BIAS_TYPE == "vector":
                  if EVEN_N:
                      bias = tl.load(b_ptrs).to(tl.float32)
                  else:
                      bias = tl.load(
                          b_ptrs, mask=offs_n < seqlen_k, other=0.0
                      ).to(tl.float32)
                  bias = bias[None, :]
              elif BIAS_TYPE == "matrix":
                  if EVEN_M & EVEN_N:
                      bias = tl.load(b_ptrs).to(tl.float32)
                  else:
                      bias = tl.load(
                          b_ptrs,
                          mask=(offs_m_curr[:, None] < seqlen_q)
                          & (offs_n[None, :] < seqlen_k),
                          other=0.0,
                      ).to(tl.float32)
              qk = qk * softmax_scale + bias
          # There seems to be a race condition when headdim=48/96, and dq, dk, dv are wrong.
          # Also wrong for headdim=64.
          if not (EVEN_M & EVEN_HEADDIM):
              tl.debug_barrier()
          lse_i = tl.load(LSE + offs_m_curr)
          if BIAS_TYPE == "none":
              p = tl.exp(qk * softmax_scale - lse_i[:, None])
          else:
              p = tl.exp(qk - lse_i[:, None])
          # compute dv
          # [2022-10-30] TD: A Triton bug: if EVEN_M=True and EVEN_HEADDIM=False, if we call
          # do = tl.load(do_ptrs, mask=offs_d[None, :] < headdim, other=0.0), we get wrong outputs
          # in the case of headdim=48/96, seqlen_q & seqlen_k >= 512. If headdim=40 or seqlen < 512,
          # the output is correct.
          if EVEN_M & EVEN_HEADDIM:
              do = tl.load(do_ptrs)
          else:
              # [2022-11-01] TD: Triton bug, there's a race condition if we just use m_mask and not d_mask.
              do = tl.load(
                  do_ptrs,
                  mask=(offs_m_curr[:, None] < seqlen_q)
                  & (offs_d[None, :] < headdim),
                  other=0.0,
              )
          # if EVEN_M:
          #     if EVEN_HEADDIM:
          #         do = tl.load(do_ptrs)
          #     else:
          #         do = tl.load(do_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
          # else:
          #     if EVEN_HEADDIM:
          #         do = tl.load(do_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0)
          #     else:
          #         do = tl.load(do_ptrs, mask=(offs_m_curr[:, None] < seqlen_q)
          #                                    & (offs_d[None, :] < headdim), other=0.0)
          dv += tl.dot(p.to(do.dtype), do, trans_a=True)
          # compute dp = dot(v, do)
          # There seems to be a race condition when headdim=48/96, and dq, dk are wrong.
          # Also wrong for headdim=128, seqlen=(108, 256), and ATOMIC_ADD=True
          # Also wrong for headdim=64, seqlen=(1023, 1024), and ATOMIC_ADD=False
          if not (EVEN_M & EVEN_HEADDIM):
              tl.debug_barrier()
          dp = tl.dot(do, v, trans_b=True)
          # There's a race condition for headdim=48
          if not EVEN_HEADDIM:
              tl.debug_barrier()
          # compute ds = p * (dp - delta[:, None])
          # Putting the subtraction after the dp matmul (instead of before) is slightly faster
          Di = tl.load(D + offs_m_curr)
          # Converting ds to q.dtype here reduces register pressure and makes it much faster
          # for BLOCK_HEADDIM=128
          ds = (p * (dp - Di[:, None]) * softmax_scale).to(q.dtype)
          # compute dk = dot(ds.T, q)
          dk += tl.dot(ds, q, trans_a=True)
          # compute dq
          if not (
              EVEN_M & EVEN_HEADDIM
          ):  # Otherewise there's a race condition when BIAS_TYPE='matrix'
              tl.debug_barrier()
          if not ATOMIC_ADD:
              if (
                  EVEN_M & EVEN_HEADDIM
              ):  # Race condition if we just do EVEN_M
                  dq = tl.load(dq_ptrs, eviction_policy="evict_last")
                  dq += tl.dot(ds, k)
                  tl.store(dq_ptrs, dq, eviction_policy="evict_last")
              else:
                  if EVEN_HEADDIM:
                      dq = tl.load(
                          dq_ptrs,
                          mask=offs_m_curr[:, None] < seqlen_q,
                          other=0.0,
                          eviction_policy="evict_last",
                      )
                      dq += tl.dot(ds, k)
                      tl.store(
                          dq_ptrs,
                          dq,
                          mask=offs_m_curr[:, None] < seqlen_q,
                          eviction_policy="evict_last",
                      )
                  else:
                      dq = tl.load(
                          dq_ptrs,
                          mask=(offs_m_curr[:, None] < seqlen_q)
                          & (offs_d[None, :] < headdim),
                          other=0.0,
                          eviction_policy="evict_last",
                      )
                      dq += tl.dot(ds, k)
                      tl.store(
                          dq_ptrs,
                          dq,
                          mask=(offs_m_curr[:, None] < seqlen_q)
                          & (offs_d[None, :] < headdim),
                          eviction_policy="evict_last",
                      )
          else:  # If we're parallelizing across the seqlen_k dimension
              dq = tl.dot(ds, k)
              if (
                  EVEN_M & EVEN_HEADDIM
              ):  # Race condition if we just do EVEN_M
                  tl.atomic_add(dq_ptrs, dq)
              else:
                  if EVEN_HEADDIM:
                      tl.atomic_add(
                          dq_ptrs,
                          dq,
                          mask=offs_m_curr[:, None] < seqlen_q,
                      )
                  else:
                      tl.atomic_add(
                          dq_ptrs,
                          dq,
                          mask=(offs_m_curr[:, None] < seqlen_q)
                          & (offs_d[None, :] < headdim),
                      )
          # increment pointers
          dq_ptrs += BLOCK_M * stride_dqm
          q_ptrs += BLOCK_M * stride_qm
          do_ptrs += BLOCK_M * stride_dom
          if BIAS_TYPE == "matrix":
              b_ptrs += BLOCK_M * stride_bm
      # write-back
      dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])
      dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])
      _bwd_store_dk_dv(
          dk_ptrs,
          dv_ptrs,
          dk,
          dv,
          offs_n,
          offs_d,
          seqlen_k,
          headdim,
          EVEN_M=EVEN_M,
          EVEN_N=EVEN_N,
          EVEN_HEADDIM=EVEN_HEADDIM,
      )
  
imports: import math
  import torch
  import triton
  import triton.language
label_desc: Use image node_id 4 for calling a global function with example usage: _bwd_kernel_one_col_block(start_n, Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qm, stride_kn, stride_vn, stride_bm, stride_dom, stride_dqm, stride_dkn, stride_dvn, seqlen_q, seqlen_k, headdim, ATOMIC_ADD, BIAS_TYPE, IS_CAUSAL, BLOCK_HEADDIM, EVEN_M, EVEN_N, EVEN_HEADDIM, BLOCK_M, BLOCK_N) without return types