code: def _flash_attention_dkv_kernel(
      q_tile_ref,
      k_tile_ref,
      v_tile_ref,
      ab_tile_ref,
      q_segment_ids_tile_ref,
      kv_segment_ids_tile_ref,
      l_tile_ref,
      m_tile_ref,
      do_tile_ref,
      di_tile_ref,
      dk_tile_ref,
      dv_tile_ref,
      dk_scratch_ref,
      dv_scratch_ref,
      *,
      sm_scale: float,
      causal: bool,
      mask_value: float,
      q_seq_len: int,
      block_q: int,
      block_k: int,
  ):
      _, _, block_q_major, _ = q_tile_ref.shape
      _, _, block_k_major, _ = k_tile_ref.shape
  
      q_seq_index = pl.program_id(axis=3)
      kv_seq_index = pl.program_id(axis=2)
  
      @pl.when(q_seq_index == 0)
      def start_new_sequence():
          dk_scratch_ref[:, :] = jnp.zeros(
              dk_scratch_ref.shape, dk_scratch_ref.dtype
          )
          dv_scratch_ref[:, :] = jnp.zeros(
              dv_scratch_ref.shape, dv_scratch_ref.dtype
          )
  
      def q_body(j, _):
          start_q = j * block_q
  
          def k_body(i, _):
              start_k = i * block_k
              k = pl.load(
                  k_tile_ref,
                  (0, 0, pl.ds(start_k, block_k), slice(None)),
              )
              v = pl.load(
                  v_tile_ref,
                  (0, 0, pl.ds(start_k, block_k), slice(None)),
              )
              q = pl.load(
                  q_tile_ref,
                  (0, 0, pl.ds(start_q, block_q), slice(None)),
              )  # [block_q, head_dim]
              l = pl.load(
                  l_tile_ref,
                  (0, 0, pl.ds(start_q, block_q), slice(None)),
              )  # [block_q, 128]
              m = pl.load(
                  m_tile_ref,
                  (0, 0, pl.ds(start_q, block_q), slice(None)),
              )  # [block_q, 128]
              do = pl.load(
                  do_tile_ref,
                  (0, 0, pl.ds(start_q, block_q), slice(None)),
              )  # [block_q, 128]
              di = pl.load(
                  di_tile_ref,
                  (0, 0, pl.ds(start_q, block_q), slice(None)),
              ).astype(
                  jnp.float32
              )  # [block_q, 128]
  
              capped_logits = lax.dot_general(
                  q,
                  k,
                  TRANS_B_DIM_NUMBERS,
                  preferred_element_type=jnp.float32,
              )  # [block_q_major, block_k]
  
              if ab_tile_ref is not None:
                  ab = pl.load(
                      ab_tile_ref,
                      (
                          0,
                          0,
                          pl.dslice(j * block_q, block_q),
                          pl.dslice(i * block_k, block_k),
                      ),
                  ).astype(jnp.float32)
                  capped_logits += ab
  
              if sm_scale != 1.0:
                  capped_logits *= sm_scale
  
              mask = None
              if q_segment_ids_tile_ref is not None:
                  repeats, rem = divmod(block_k, NUM_LANES)
                  if rem:
                      raise NotImplementedError()
                  q_segment_ids = pl.load(
                      q_segment_ids_tile_ref,
                      (0, pl.ds(start_q, block_q), slice(None)),
                  )  # [block_q, NUM_LANES].
                  q_segment_ids = pltpu.repeat(
                      q_segment_ids, repeats, axis=1
                  )  # [block_q, block_k].
                  kv_segment_ids = pl.load(
                      kv_segment_ids_tile_ref,
                      (slice(None), 0, pl.ds(start_k, block_k)),
                  )  # [1, block_k].
                  mask = jnp.equal(
                      q_segment_ids, kv_segment_ids
                  ).astype(jnp.bool_)
  
              if causal:
                  mask_shape = (block_q, block_k)
                  row_ids = jax.lax.broadcasted_iota(
                      jnp.int32, mask_shape, 0
                  )
                  row_ids += q_seq_index * block_q_major + start_q
                  col_ids = jax.lax.broadcasted_iota(
                      jnp.int32, mask_shape, 1
                  )
                  col_ids += kv_seq_index * block_k_major + start_k
                  causal_mask = col_ids <= row_ids
                  mask = (
                      causal_mask
                      if mask is None
                      else jnp.logical_and(mask, causal_mask)
                  )
  
              capped_logits = (
                  capped_logits
                  if mask is None
                  else capped_logits + jnp.where(mask, 0.0, mask_value)
              )
  
              p = jnp.exp(
                  capped_logits
                  - pltpu.repeat(m, block_k // MIN_BLOCK_SIZE, axis=1)
              )
              p = p * pltpu.repeat(
                  1 / l, block_k // MIN_BLOCK_SIZE, axis=1
              )  # [block_q_major, block_k_major]
              dv = lax.dot(
                  p.T.astype(do.dtype),
                  do,
                  preferred_element_type=jnp.float32,
              )
              pl.store(
                  dv_scratch_ref,
                  (pl.ds(start_k, block_k), slice(None)),
                  pl.load(
                      dv_scratch_ref,
                      (pl.ds(start_k, block_k), slice(None)),
                  )
                  + dv.astype(dv_scratch_ref.dtype),
              )
  
              # di: [block_q, 128]
              # do: [block_q, head_dim]
              # v: [block_k_major, head_dim]
              dp = lax.dot_general(
                  do,
                  v,
                  TRANS_B_DIM_NUMBERS,
                  preferred_element_type=jnp.float32,
              )
              ds = (
                  dp
                  - pltpu.repeat(di, block_k // MIN_BLOCK_SIZE, axis=1)
              ) * p
  
              if sm_scale != 1.0:
                  ds = ds * sm_scale
  
              # ds: [block_q_major, block_k_major]
              # q: [block_q_major, head_dim]
              dk = lax.dot(
                  ds.T.astype(do.dtype),
                  q,
                  preferred_element_type=jnp.float32,
              )
              pl.store(
                  dk_scratch_ref,
                  (pl.ds(start_k, block_k), slice(None)),
                  pl.load(
                      dk_scratch_ref,
                      (pl.ds(start_k, block_k), slice(None)),
                  )
                  + dk.astype(dk_scratch_ref.dtype),
              )
  
          lax.fori_loop(
              0, block_k_major // block_k, k_body, None, unroll=True
          )
  
      if causal:
          should_run = below_or_on_diag(
              q_seq_index, block_q_major, kv_seq_index, block_k_major
          )
      else:
          should_run = True
  
      @pl.when(should_run)
      def run():
          lax.fori_loop(
              0, block_q_major // block_q, q_body, None, unroll=True
          )
  
      @pl.when(q_seq_index == q_seq_len // block_q_major - 1)
      def end_of_q_sequence():
          dv_tile_ref[0, 0, :, :] = dv_scratch_ref[...].astype(
              dv_tile_ref
          )
          dk_tile_ref[0, 0, :, :] = dk_scratch_ref[...].astype(
              dk_tile_ref
          )
  
imports: from __future__ import annotations
  import dataclasses
  import functools
  from typing import Any, NamedTuple
  import jax
  from jax import lax
  from jax.experimental import pallas
  from jax.experimental.pallas import tpu
  import jax.numpy
label_desc: Use image node_id 10 for calling a global function with example usage: _flash_attention_dkv_kernel(q_tile_ref, k_tile_ref, v_tile_ref, ab_tile_ref, q_segment_ids_tile_ref, kv_segment_ids_tile_ref, l_tile_ref, m_tile_ref, do_tile_ref, di_tile_ref, dk_tile_ref, dv_tile_ref, dk_scratch_ref, dv_scratch_ref) without return types