code: def _flash_attention_dkv_kernel(
      # Prefetched inputs
      data_next_ref,
      block_mask_ref,
      mask_next_ref,
      # Inputs
      q_ref,
      k_ref,
      v_ref,
      q_segment_ids_ref,
      kv_segment_ids_ref,
      logsumexp_ref,
      do_ref,
      di_ref,
      mask_ref,
      q_sequence_ref,
      # Outputs
      dq_scratch_ref,
      dk_scratch_ref,
      dv_scratch_ref,
      dq_ref,
      dk_ref,
      dv_ref,
      *,
      num_q_heads: int,
      num_kv_heads: int,
      mask_value: float,
      grid_width: int,
      bq: int,
      bkv_compute: int,
      is_mqa: bool,
      attn_logits_soft_cap: float | None,
      q_layout: QKVLayout,
      k_layout: QKVLayout,
      v_layout: QKVLayout,
      bkv: int,
      mask_function: MaskFunctionType | None,
  ):
      HEAD_DIM_MINOR = QKVLayout.HEAD_DIM_MINOR
      kv_index, q_head_index, q_index = (
          pl.program_id(0),
          pl.program_id(1),
          pl.program_id(2),
      )
      should_initialize = q_index == 0
  
      q_heads_per_kv_heads = None
      q_head_index_per_kv_head = None
  
      # Consider this situation:
      # Q_heads:   0, 1, 2, 3, 4, 5, 6, 7
      # KV_heads:  0,    1,    2,    3
      # The gradient scratch buffers should be initialized for Q_heads 0, 2, 4, 6
      # (first Q_heads to 'see' a new KV_head).
      # The gradient output buffers should be written for Q_heads 1, 3, 5, 7 (last
      # Q_heads to 'see' the current KV_head).
  
      # We can use the same logic for both MQA and GA (grouped attention).
      # But for MQA there is no need for the rem instruction, so we skip it.
      if is_mqa:
          should_initialize = jnp.logical_and(
              should_initialize, q_head_index == 0
          )
      elif num_kv_heads < num_q_heads:
          q_heads_per_kv_heads = num_q_heads // num_kv_heads
          q_head_index_per_kv_head = lax.rem(
              q_head_index, q_heads_per_kv_heads
          )
          should_initialize = jnp.logical_and(
              should_initialize, q_head_index_per_kv_head == 0
          )
  
      @pl.when(should_initialize)
      def init():
          dk_scratch_ref[...] = jnp.zeros_like(dk_scratch_ref)
          dv_scratch_ref[...] = jnp.zeros_like(dv_scratch_ref)
  
      _, _, should_run, should_not_mask = _next_nonzero(
          q_head_index,
          q_index,
          kv_index,
          data_next_ref,
          block_mask_ref,
          mask_next_ref,
          next_i=True,
      )
  
      def body(i, _):
          slice_k = pl.ds(i * bkv_compute, bkv_compute)
          q = q_ref[
              ...
          ]  # We keep q potentially transposed, since it's always RHS
  
          def _load_kv(ref, layout):
              if layout == HEAD_DIM_MINOR:
                  return pl.load(ref, (slice_k, slice(None)))
              return pl.load(ref, (slice(None), slice_k)).T
  
          k = _load_kv(k_ref, k_layout)
          v = _load_kv(v_ref, v_layout)
          logsumexp = pl.load(logsumexp_ref, (pl.ds(1), slice(None)))
          do = do_ref[...]
          di = pl.load(di_ref, (pl.ds(1), slice(None)))
  
          qk_dims = (
              NT_DIM_NUMBERS
              if q_layout == HEAD_DIM_MINOR
              else NN_DIM_NUMBERS
          )
          qk_uncapped = lax.dot_general(
              k, q, qk_dims, preferred_element_type=jnp.float32
          )
  
          qk = _apply_mask_and_soft_cap(
              qk_uncapped,
              mask_value,
              should_not_mask,
              mask_ref,
              q_sequence_ref,
              q_segment_ids_ref,
              kv_segment_ids_ref,
              attn_logits_soft_cap=attn_logits_soft_cap,
              k_slice=slice_k,
              k_offset=kv_index * bkv + i * bkv_compute,
              bq=bq,
              k_in_lanes=False,
              mask_function=mask_function,
          )
          p = jnp.exp(qk - logsumexp)
          dv = lax.dot(
              p.astype(do.dtype), do, preferred_element_type=jnp.float32
          )
          dv = dv.astype(dv_scratch_ref.dtype) + pl.load(
              dv_scratch_ref, (slice_k, slice(None))
          )
          pl.store(dv_scratch_ref, (slice_k, slice(None)), dv)
  
          dp = lax.dot_general(
              v,
              do,
              NT_DIM_NUMBERS,
              preferred_element_type=jnp.float32,
          )
          ds = (dp - di) * p
          if attn_logits_soft_cap is not None:
              normalized = qk_uncapped / attn_logits_soft_cap
              d = jnp.tanh(normalized)
              g = ds * (1 - d)
              ds = g + g * d
          dk_dims = (
              NN_DIM_NUMBERS
              if q_layout == HEAD_DIM_MINOR
              else NT_DIM_NUMBERS
          )
          dk = lax.dot_general(
              ds.astype(do.dtype),
              q,
              dk_dims,
              preferred_element_type=jnp.float32,
          )
          dk = dk.astype(dk_scratch_ref.dtype) + pl.load(
              dk_scratch_ref, (slice_k, slice(None))
          )
          pl.store(dk_scratch_ref, (slice_k, slice(None)), dk)
          if dq_scratch_ref is not None or dq_ref is not None:
              dq = lax.dot_general(
                  ds.T.astype(k.dtype),
                  k,
                  NN_DIM_NUMBERS,
                  preferred_element_type=jnp.float32,
              )
              if dq_scratch_ref is not None:
                  # Compute block size != memory block size
                  dq_scratch_ref[...] += dq
              else:
                  # Compute block size == memory block size
                  assert dq_ref is not None
                  dq_ref[...] = dq.astype(dq_ref.dtype)
  
      if dq_scratch_ref is not None:
          dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref)
      elif dq_scratch_ref is None and dq_ref is not None:
          dq_ref[...] = jnp.zeros_like(dq_ref)
  
      @pl.when(should_run)
      def run():
          num_iters = (
              k_ref.shape[0 if k_layout is HEAD_DIM_MINOR else 1]
              // bkv_compute
          )
          lax.fori_loop(0, num_iters, body, None, unroll=True)
  
      if dq_scratch_ref is not None:
          assert dq_ref is not None
          dq_ref[...] = dq_scratch_ref[...].astype(dq_ref.dtype)
  
      should_write = q_index == grid_width - 1
      if is_mqa:
          should_write = jnp.logical_and(
              should_write, q_head_index == num_q_heads - 1
          )
      elif num_kv_heads < num_q_heads:
          should_write = jnp.logical_and(
              should_write,
              q_head_index_per_kv_head == q_heads_per_kv_heads - 1,
          )
  
      @pl.when(should_write)
      def end():
          dk_ref[...] = dk_scratch_ref[...].astype(dk_ref.dtype)
          dv_ref[...] = dv_scratch_ref[...].astype(dv_ref.dtype)
          if dq_scratch_ref is not None:
              dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref)
  
          dk_scratch_ref[...] = jnp.zeros_like(dk_scratch_ref)
          dv_scratch_ref[...] = jnp.zeros_like(dv_scratch_ref)
  
imports: from __future__ import annotations
  import dataclasses
  import enum
  import functools
  from typing import Any, Callable, Literal, NamedTuple, Union, Optional, overload
  import jax
  from jax import ad_checkpoint
  from jax import lax
  from jax import tree_util
  from jax.experimental import pallas
  from jax.experimental.pallas import tpu
  from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask
  from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask_info
  import jax.numpy
  import numpy
label_desc: Use image node_id 23 for calling a global function with example usage: _flash_attention_dkv_kernel(data_next_ref, block_mask_ref, mask_next_ref, q_ref, k_ref, v_ref, q_segment_ids_ref, kv_segment_ids_ref, logsumexp_ref, do_ref, di_ref, mask_ref, q_sequence_ref, dq_scratch_ref, dk_scratch_ref, dv_scratch_ref, dq_ref, dk_ref, dv_ref) and returns: pl, pl