code: def mha_backward_kernel(
      # Inputs
      q_ref,
      k_ref,
      v_ref,
      segment_ids_ref: jax.Array | None,
      out_ref,
      do_scaled_ref,
      l_ref,
      m_ref,
      delta_ref,
      _,
      # Outputs
      dq_ref,
      dk_ref,
      dv_ref,
      *,
      sm_scale: float,
      causal: bool,
      block_q: int,
      block_d: int,
      block_k: int,
  ):
      del out_ref, l_ref  # Not needed
      seq_len = q_ref.shape[0]
  
      def outer_loop(start_k, _):
          dv = jnp.zeros([block_k, block_d], dtype=jnp.float32)
          dk = jnp.zeros([block_k, block_d], dtype=jnp.float32)
          k = pl.load(
              k_ref, (pl.ds(start_k * block_k, block_k), slice(None))
          )
          v = pl.load(
              v_ref, (pl.ds(start_k * block_k, block_k), slice(None))
          )
          span_k = start_k * block_k + jnp.arange(block_k)
          kv_segment_ids = (
              None
              if segment_ids_ref is None
              else pl.load(
                  segment_ids_ref, (pl.ds(start_k * block_k, block_k),)
              )
          )
  
          def inner_loop(start_q, carry):
              dv, dk = carry
              q = pl.load(
                  q_ref,
                  (pl.ds(start_q * block_q, block_q), slice(None)),
              )
              qk = pl.dot(q, k.T)
              qk = qk.astype(q_ref.dtype)
              qk = qk.astype(jnp.float32)
              if sm_scale != 1.0:
                  qk *= sm_scale
  
              q_segment_ids = (
                  None
                  if segment_ids_ref is None
                  else pl.load(
                      segment_ids_ref,
                      (pl.ds(start_q * block_q, block_q),),
                  )
              )
  
              if causal or segment_ids_ref is not None:
                  mask = None
                  if segment_ids_ref is not None:
                      mask = segment_mask(q_segment_ids, kv_segment_ids)
  
                  if causal:
                      span_q = start_q * block_q + jnp.arange(block_q)
                      causal_mask = span_q[:, None] >= span_k[None, :]
                      mask = (
                          causal_mask
                          if mask is None
                          else jnp.logical_and(mask, causal_mask)
                      )
                  qk = jnp.where(mask, qk, DEFAULT_MASK_VALUE)
  
              m = pl.load(m_ref, (pl.ds(start_q * block_q, block_q),))
              p = jnp.exp(qk - m[:, None])
              do = pl.load(
                  do_scaled_ref,
                  (pl.ds(start_q * block_q, block_q), slice(None)),
              )
              dv = dv + pl.dot(p.astype(do.dtype).T, do)
              di = pl.load(
                  delta_ref, (pl.ds(start_q * block_q, block_q),)
              )
              dp = (
                  jnp.zeros((block_q, block_k), dtype=jnp.float32)
                  - di[:, None]
              )
              dp = dp + pl.dot(do, v.T)
              ds = p * dp
              if sm_scale != 1.0:
                  ds = ds * sm_scale
              dk = dk + pl.dot(ds.astype(q_ref.dtype).T, q)
              dq = pl.load(
                  dq_ref,
                  (pl.ds(start_q * block_q, block_q), slice(None)),
                  eviction_policy="evict_last",
              )
              dq = dq + pl.dot(ds.astype(k.dtype), k).astype(dq.dtype)
              pl.store(
                  dq_ref,
                  (pl.ds(start_q * block_q, block_q), slice(None)),
                  dq,
                  eviction_policy="evict_last",
              )
              return dv, dk
  
          if causal:
              lower_bound = lax.div(start_k * block_k, block_q)
          else:
              lower_bound = 0
          dv, dk = lax.fori_loop(
              lower_bound,
              pl.cdiv(seq_len, block_q),
              inner_loop,
              (dv, dk),
          )
          pl.store(
              dv_ref,
              (pl.ds(start_k * block_k, block_k), slice(None)),
              dv.astype(dv_ref.dtype),
          )
          pl.store(
              dk_ref,
              (pl.ds(start_k * block_k, block_k), slice(None)),
              dk.astype(dk_ref.dtype),
          )
  
      lax.fori_loop(0, pl.cdiv(seq_len, block_k), outer_loop, None)
  
file_path: jax/jax/experimental/pallas/ops/attention.py
imports: from __future__ import annotations
  import functools
  from typing import Any, Optional
  import jax
  from jax import lax
  from jax.experimental import pallas
  import jax.numpy
  import numpy
label_desc: Use image node_id 7 for calling a global function with example usage: mha_backward_kernel(q_ref, k_ref, v_ref, segment_ids_ref, out_ref, do_scaled_ref, l_ref, m_ref, delta_ref, _, dq_ref, dk_ref, dv_ref) and returns: dv,  dk