code: def _splash_attention_bwd_dkv(
      q,
      k,
      v,
      segment_ids,
      logsumexp,
      do,
      di,
      *,
      bq: int,
      bkv: int,
      bkv_compute: int,
      is_mqa: bool,
      mask_info: mask_info_lib.MaskInfo,
      mask_value: float,
      attn_logits_soft_cap: float | None,
      use_fused_bwd_kernel: bool,
      q_layout: QKVLayout,
      k_layout: QKVLayout,
      v_layout: QKVLayout,
      mask_function: MaskFunctionType | None,
  ):
      num_q_heads, q_seq_len, head_dim = q.shape
      if is_mqa:
          num_kv_heads, kv_seq_len = 1, k.shape[0]
      else:
          num_kv_heads, kv_seq_len, _ = k.shape
  
      if bq > q_seq_len:
          raise ValueError(
              f"{bq=} should not be greater than {q_seq_len=}"
          )
      if bkv > kv_seq_len:
          raise ValueError(
              f"{bkv=} should not be greater than {kv_seq_len=}"
          )
      if bkv_compute > bkv:
          raise ValueError(
              f"{bkv_compute=} should not be greater than {bkv=}"
          )
      if bkv % bkv_compute:
          raise ValueError(
              f"{bkv=} should be a multiple of {bkv_compute=}"
          )
  
      if not is_mqa and num_q_heads % num_kv_heads != 0:
          raise ValueError(
              f"In MHA, expected number of 'key' heads ({num_kv_heads}) to be a"
              f" multiple of the number of 'query' heads ({num_q_heads})"
          )
  
      if k.shape != v.shape:
          raise ValueError(
              f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same"
              " shape."
          )
  
      q_heads_per_kv_head = num_q_heads // num_kv_heads
  
      if mask_info.data_next is not None:
          grid_width = mask_info.data_next.shape[-2]
      else:
          grid_width = q_seq_len // bq
  
      grid = (
          kv_seq_len // bkv,
          num_q_heads,
          grid_width,
      )
  
      def o_index_map(
          kv_index,
          head_index,
          q_index,
          data_next_ref,
          block_mask_ref,
          mask_next_ref=None,
      ):
          next_i, *_ = _next_nonzero(
              head_index,
              q_index,
              kv_index,
              data_next_ref,
              block_mask_ref,
              mask_next_ref,
              next_i=True,
          )
          return head_index, next_i, 0
  
      o_spec = pl.BlockSpec(o_index_map, (None, bq, head_dim))
  
      def q_index_map(
          kv_index,
          head_index,
          q_index,
          data_next_ref,
          block_mask_ref,
          mask_next_ref=None,
      ):
          next_i, *_ = _next_nonzero(
              head_index,
              q_index,
              kv_index,
              data_next_ref,
              block_mask_ref,
              mask_next_ref,
              next_i=True,
          )
          return from_head_minor((head_index, next_i, 0), q_layout)
  
      q_spec = pl.BlockSpec(
          q_index_map, from_head_minor((None, bq, head_dim), q_layout)
      )
  
      def k_index_map(kv_index, head_index, *_):
          prefix = (
              () if is_mqa else (_div(head_index, q_heads_per_kv_head),)
          )
          return from_head_minor((*prefix, kv_index, 0), k_layout)
  
      k_spec = pl.BlockSpec(
          k_index_map,
          from_head_minor(
              (bkv, head_dim) if is_mqa else (None, bkv, head_dim),
              k_layout,
          ),
      )
  
      def v_index_map(kv_index, head_index, *_):
          prefix = (
              () if is_mqa else (_div(head_index, q_heads_per_kv_head),)
          )
          return from_head_minor((*prefix, kv_index, 0), v_layout)
  
      v_spec = pl.BlockSpec(
          v_index_map,
          from_head_minor(
              (bkv, head_dim) if is_mqa else (None, bkv, head_dim),
              v_layout,
          ),
      )
  
      if use_fused_bwd_kernel:
  
          def dq_index_map(kv_index, head_index, q_index, *_):
              return (kv_index, head_index, q_index, 0)
  
          dq_spec = pl.BlockSpec(
              dq_index_map, (None, None, bq, head_dim)
          )
          dq_shape = jax.ShapeDtypeStruct(
              (kv_seq_len // bkv, *q.shape), q.dtype
          )
          if bkv == bkv_compute:
              dq_scratch_spec = dq_scratch_shape = None
          else:
              dq_scratch_spec = pl.BlockSpec(
                  lambda *_: (0, 0), (bq, head_dim)
              )
              dq_scratch_shape = jax.ShapeDtypeStruct(
                  (bq, head_dim), jnp.float32
              )
      else:
          dq_spec = dq_shape = dq_scratch_spec = dq_scratch_shape = None
  
      def dkv_index_map(kv_index, head_index, *_):
          prefix = (
              () if is_mqa else (_div(head_index, q_heads_per_kv_head),)
          )
          return (*prefix, kv_index, 0)
  
      dk_spec = dv_spec = pl.BlockSpec(
          dkv_index_map,
          (bkv, head_dim) if is_mqa else (None, bkv, head_dim),
      )
  
      def mask_index_map(
          kv_index,
          head_index,
          q_index,
          data_next_ref,
          block_mask_ref,
          mask_next_ref,
      ):
          _, next_m, *_ = _next_nonzero(
              head_index,
              q_index,
              kv_index,
              data_next_ref,
              block_mask_ref,
              mask_next_ref,
              next_i=True,
          )
          return next_m, 0, 0
  
      mask_spec = pl.BlockSpec(mask_index_map, (None, bkv, bq))
  
      def q_segment_ids_index_map(
          kv_index,
          head_index,
          q_index,
          data_next_ref,
          block_mask_ref,
          mask_next_ref=None,
      ):
          next_i, *_ = _next_nonzero(
              head_index,
              q_index,
              kv_index,
              data_next_ref,
              block_mask_ref,
              mask_next_ref,
              next_i=True,
          )
          return 0, next_i
  
      if segment_ids is not None:
  
          def kv_segment_ids_index_map(kv_index, *_):
              return kv_index, 0
  
          q_segment_spec = pl.BlockSpec(
              q_segment_ids_index_map, (NUM_SUBLANES, bq)
          )
          kv_segment_spec = pl.BlockSpec(
              kv_segment_ids_index_map, (bkv, NUM_LANES)
          )
          q_segment_ids = jax.lax.broadcast_in_dim(
              segment_ids.q, (NUM_SUBLANES, q_seq_len), (1,)
          )
          kv_segment_ids = jax.lax.broadcast_in_dim(
              segment_ids.kv, (kv_seq_len, NUM_LANES), (0,)
          )
      else:
          q_segment_spec = kv_segment_spec = None
          q_segment_ids = kv_segment_ids = None
  
      do_spec = o_spec
  
      def logsumexp_index_map(
          kv_index,
          head_index,
          q_index,
          data_next_ref,
          block_mask_ref,
          mask_next_ref=None,
      ):
          next_i, *_ = _next_nonzero(
              head_index,
              q_index,
              kv_index,
              data_next_ref,
              block_mask_ref,
              mask_next_ref,
              next_i=True,
          )
          return head_index, 0, next_i
  
      assert logsumexp.shape == di.shape == (num_q_heads, q_seq_len)
      # TODO(apaszke): Remove the sublane expansion once Mosaic has all retilings
      logsumexp_shape = (num_q_heads, NUM_SUBLANES, q_seq_len)
      logsumexp = jnp.broadcast_to(
          jnp.expand_dims(logsumexp, -2), logsumexp_shape
      )
      logsumexp_spec = pl.BlockSpec(
          logsumexp_index_map, (None, NUM_SUBLANES, bq)
      )
      assert logsumexp.ndim == len(logsumexp_spec.block_shape)
  
      # TODO(apaszke): Remove the sublane expansion once Mosaic has all retilings
      di = jnp.broadcast_to(jnp.expand_dims(di, -2), logsumexp_shape)
      di_spec = pl.BlockSpec(
          logsumexp_index_map, (None, NUM_SUBLANES, bq)
      )
      assert di.ndim == len(di_spec.block_shape)
  
      in_specs = [
          q_spec,
          k_spec,
          v_spec,
          q_segment_spec,
          kv_segment_spec,
          logsumexp_spec,
          do_spec,
          di_spec,
      ]
      if mask_info.partial_mask_blocks is not None:
          in_specs.append(mask_spec)
      else:
          in_specs.append(None)
  
      if mask_info.q_sequence is not None:
          in_specs.append(
              pl.BlockSpec(q_segment_ids_index_map, (NUM_SUBLANES, bq))
          )
          q_sequence = jax.lax.broadcast_in_dim(
              mask_info.q_sequence, (NUM_SUBLANES, q_seq_len), (1,)
          )
      else:
          q_sequence = None
          in_specs.append(None)
  
      out_shapes = [
          dq_scratch_shape,
          jax.ShapeDtypeStruct((bkv, head_dim), jnp.float32),
          jax.ShapeDtypeStruct((bkv, head_dim), jnp.float32),
          dq_shape,
          jax.ShapeDtypeStruct(k.shape, k.dtype),
          jax.ShapeDtypeStruct(v.shape, v.dtype),
      ]
      out_specs = [
          dq_scratch_spec,
          pl.BlockSpec(lambda *_: (0, 0), (bkv, head_dim)),
          pl.BlockSpec(lambda *_: (0, 0), (bkv, head_dim)),
          dq_spec,
          dk_spec,
          dv_spec,
      ]
  
      kernel = functools.partial(
          _flash_attention_dkv_kernel,
          mask_value=mask_value,
          num_q_heads=num_q_heads,
          num_kv_heads=num_kv_heads,
          is_mqa=is_mqa,
          grid_width=grid_width,
          bq=bq,
          bkv_compute=bkv_compute,
          attn_logits_soft_cap=attn_logits_soft_cap,
          q_layout=q_layout,
          k_layout=k_layout,
          v_layout=v_layout,
          bkv=bkv,
          mask_function=mask_function,
      )
      num_scalar_prefetch = 3
  
      # Attach useful metadata to the custom-call HLO op.
      # Having this information available in an HLO-dump or xprof is valuable for
      # debugging and performance investigation.
      metadata_dict = dict(
          block_q_dkv=bq,
          block_kv_dkv=bkv,
          block_kv_dkv_compute=bkv_compute,
          q_layout=q_layout,
          k_layout=k_layout,
          v_layout=v_layout,
          use_fused_bwd_kernel=use_fused_bwd_kernel,
          is_mqa=is_mqa,
          mask_value=mask_value,
          is_segmented=segment_ids is not None,
          attn_logits_soft_cap=attn_logits_soft_cap,
      )
  
      mosaic_params = pltpu.encode_kernel_regeneration_metadata(
          metadata_dict
      )
      # We set all dimensions to arbitrary because:
      # 1) for kv_seq_len, the splash attention prefetch schedule assumes no
      #    megacore
      # 2) for heads, we are reducing over heads
      # 3) for q_seq_len, we are reducing over it to compute dkv
      mosaic_params.update(
          dimension_semantics=("arbitrary", "arbitrary", "arbitrary"),
          flags={"XLA_TPU_FORCE_LP_LLO_SCHEDULER": True},
      )
  
      kernel_name = get_kernel_name(
          is_mqa, False, segment_ids is not None, "dkv"
      )
      with jax.named_scope(kernel_name):
          _, _, _, dq_unreduced, dk, dv = pl.pallas_call(
              kernel,
              grid_spec=pltpu.PrefetchScalarGridSpec(
                  num_scalar_prefetch=num_scalar_prefetch,
                  in_specs=in_specs,
                  out_specs=out_specs,
                  grid=grid,
              ),
              out_shape=out_shapes,
              mosaic_params=mosaic_params,
              name=kernel_name,
          )(
              mask_info.data_next,
              mask_info.block_mask,
              mask_info.mask_next,
              q
              if q_layout == QKVLayout.HEAD_DIM_MINOR
              else q.swapaxes(-1, -2),
              k
              if k_layout == QKVLayout.HEAD_DIM_MINOR
              else k.swapaxes(-1, -2),
              v
              if v_layout == QKVLayout.HEAD_DIM_MINOR
              else v.swapaxes(-1, -2),
              q_segment_ids,
              kv_segment_ids,
              logsumexp,
              do,
              di,
              mask_info.partial_mask_blocks,
              q_sequence,
          )
      if use_fused_bwd_kernel:
          assert dq_unreduced is not None
          dq = dq_unreduced.sum(axis=0)
      else:
          assert dq_unreduced is None
          dq = None
      return dq, dk, dv
  
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 24 for calling a global function with example usage: _splash_attention_bwd_dkv(q, k, v, segment_ids, logsumexp, do, di) and returns: dq,  dk,  dv, head_index,  next_i,  int, from_head_minor, from_head_minor, from_head_minor, unknown,  kv_index,  int, next_m,  int,  int, int,  next_i, head_index,  int,  next_i, kv_index,  head_index,  q_index,  int, kv_index,  int