code: def _splash_attention_forward(
      fwd_mask_info: mask_info_lib.MaskInfo,
      q: jax.Array,
      k: jax.Array,
      v: jax.Array,
      segment_ids: SegmentIds | None,
      mask_value: float,
      is_mqa: bool,
      block_sizes: BlockSizes,
      residual_checkpoint_name: str | None,
      save_residuals: bool,
      mask_function: MaskFunctionType | None,
      attn_logits_soft_cap: float | None = None,
  ) -> SplashCustomReturnType:
      num_q_heads, q_seq_len, head_dim = q.shape
      bq, bkv = block_sizes.block_q, block_sizes.block_kv
      bkv_compute = block_sizes.block_kv_compute
  
      if is_mqa:
          expected_kv_rank = 2
          kv_head_dimension = 1
          kv_seq_len_dimension = 0
          num_kv_heads = 1
      else:
          expected_kv_rank = 3
          kv_head_dimension = 2
          kv_seq_len_dimension = 1
          num_kv_heads = k.shape[0]
  
      if len(k.shape) != expected_kv_rank:
          raise ValueError(
              f"Expected {expected_kv_rank}-dim 'key' tensor for MQA. Instead got a"
              f" {len(k.shape)}-dim one."
          )
  
      if k.shape[kv_head_dimension] != head_dim:
          raise ValueError(
              f"Expected 'key' head dimension to be: {head_dim}. Instead got:"
              f" {k.shape[kv_head_dimension]}."
          )
  
      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."
          )
  
      if bkv % bkv_compute:
          raise ValueError(
              f"{bkv=} must be a multiple of {bkv_compute=}."
          )
      if bkv_compute % NUM_LANES:
          raise ValueError(f"{bkv=} must be a multiple of {NUM_LANES}.")
  
      kv_seq_len = k.shape[kv_seq_len_dimension]
  
      q_heads_per_kv_head = num_q_heads // num_kv_heads
  
      if segment_ids is not None:
          if segment_ids.q.shape != (q_seq_len,):
              raise ValueError(
                  "Invalid shape for q segment_ids: "
                  f"{segment_ids.q.shape}. Expected: {(q_seq_len,)}"
              )
          if segment_ids.kv.shape != (kv_seq_len,):
              raise ValueError(
                  "Invalid shape for kv segment_ids: "
                  f"{segment_ids.kv.shape}. Expected: {(kv_seq_len,)}"
              )
  
      q_layout = block_sizes.q_layout
  
      def q_index_map(
          h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None
      ):
          del j, data_next_ref, mask_next_ref, block_mask_ref
          return from_head_minor((h, i, 0), q_layout)
  
      def out_index_map(
          h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None
      ):
          del j, data_next_ref, mask_next_ref, block_mask_ref
          return h, i, 0
  
      k_layout = block_sizes.k_layout
  
      def k_index_map(
          h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None
      ):
          next_j, *_ = _next_nonzero(
              h, i, j, data_next_ref, block_mask_ref, mask_next_ref
          )
          prefix = () if is_mqa else (_div(h, q_heads_per_kv_head),)
          return from_head_minor((*prefix, next_j, 0), k_layout)
  
      v_layout = block_sizes.v_layout
  
      def v_index_map(
          h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None
      ):
          next_j, *_ = _next_nonzero(
              h, i, j, data_next_ref, block_mask_ref, mask_next_ref
          )
          prefix = () if is_mqa else (_div(h, q_heads_per_kv_head),)
          return from_head_minor((*prefix, next_j, 0), v_layout)
  
      def mask_index_map(
          h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None
      ):
          _, next_m, *_ = _next_nonzero(
              h, i, j, data_next_ref, block_mask_ref, mask_next_ref
          )
          return next_m, 0, 0
  
      def q_segment_ids_index_map(h, i, j, *_):
          del h, j  # Unused.
          return i, 0
  
      def kv_segment_ids_index_map(
          h, i, j, data_next_ref, block_mask_ref, mask_next_ref=None
      ):
          next_j, *_ = _next_nonzero(
              h, i, j, data_next_ref, block_mask_ref, mask_next_ref
          )
          return 0, next_j
  
      # Convert the logical shape from head-minor to sequence-minor.
      in_specs = [
          pl.BlockSpec(
              q_index_map,
              from_head_minor((None, bq, head_dim), q_layout),
          ),
          pl.BlockSpec(
              k_index_map,
              from_head_minor(
                  (bkv, head_dim) if is_mqa else (None, bkv, head_dim),
                  k_layout,
              ),
          ),
          pl.BlockSpec(
              v_index_map,
              from_head_minor(
                  (bkv, head_dim) if is_mqa else (None, bkv, head_dim),
                  v_layout,
              ),
          ),
      ]
      if segment_ids is not None:
          in_specs += [
              pl.BlockSpec(q_segment_ids_index_map, (bq, NUM_LANES)),
              pl.BlockSpec(
                  kv_segment_ids_index_map, (NUM_SUBLANES, bkv)
              ),
          ]
          q_segment_ids = jax.lax.broadcast_in_dim(
              segment_ids.q, (q_seq_len, NUM_LANES), (0,)
          )
          kv_segment_ids = jax.lax.broadcast_in_dim(
              segment_ids.kv, (NUM_SUBLANES, kv_seq_len), (1,)
          )
      else:
          in_specs += [None, None]
          q_segment_ids = kv_segment_ids = None
  
      if fwd_mask_info.partial_mask_blocks is not None:
          in_specs.append(pl.BlockSpec(mask_index_map, (None, bq, bkv)))
      else:
          in_specs.append(None)
  
      assert (
          fwd_mask_info.partial_mask_blocks is None
          or fwd_mask_info.q_sequence is None
      )
  
      if fwd_mask_info.q_sequence is not None:
          q_sequence = jax.lax.broadcast_in_dim(
              fwd_mask_info.q_sequence, (q_seq_len, NUM_LANES), (0,)
          )
          in_specs.append(
              pl.BlockSpec(q_segment_ids_index_map, (bq, NUM_LANES))
          )
      else:
          q_sequence = None
          in_specs.append(None)
  
      num_scalar_prefetch = 3
  
      out_shapes = [
          jax.ShapeDtypeStruct(
              (bq, NUM_LANES), jnp.float32
          ),  # m_scratch
          jax.ShapeDtypeStruct(
              (bq, NUM_LANES), jnp.float32
          ),  # l_scratch
          jax.ShapeDtypeStruct(
              (bq, head_dim), jnp.float32
          ),  # o_scratch
          jax.ShapeDtypeStruct(
              (num_q_heads, q_seq_len, head_dim), q.dtype
          ),
      ]
      out_specs = [
          # TODO(sharadmv): convert m/l to be scratch
          pl.BlockSpec(lambda h, i, j, *_: (0, 0), (bq, NUM_LANES)),
          pl.BlockSpec(lambda h, i, j, *_: (0, 0), (bq, NUM_LANES)),
          pl.BlockSpec(lambda h, i, j, *_: (0, 0), (bq, head_dim)),
          pl.BlockSpec(out_index_map, (None, bq, head_dim)),
      ]
      if save_residuals:
          out_shapes += [
              jax.ShapeDtypeStruct(
                  (num_q_heads, q_seq_len, NUM_LANES), jnp.float32
              ),  # logsumexp
          ]
  
          def logsumexp_index_map(h, i, *_):
              return h, i, 0
  
          out_specs += [
              pl.BlockSpec(logsumexp_index_map, (None, bq, NUM_LANES)),
          ]
      else:
          out_shapes += [None]
          out_specs += [None]
  
      # 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_sizes=dataclasses.asdict(block_sizes),
          is_mqa=is_mqa,
          save_residuals=save_residuals,
          mask_value=mask_value,
          is_segmented=segment_ids is not None,
          attn_logits_soft_cap=attn_logits_soft_cap,
          residual_checkpoint_name=residual_checkpoint_name,
      )
  
      mosaic_params = pltpu.encode_kernel_regeneration_metadata(
          metadata_dict
      )
  
      mosaic_params.update(
          dimension_semantics=("parallel", "arbitrary", "arbitrary"),
          flags={"XLA_TPU_FORCE_LP_LLO_SCHEDULER": True},
      )
  
      kernel_name = get_kernel_name(
          is_mqa, save_residuals, segment_ids is not None, "fwd"
      )
  
      if fwd_mask_info.data_next is not None:
          grid_width = fwd_mask_info.data_next.shape[-1]
      else:
          grid_width = kv_seq_len // bkv
  
      grid = (num_q_heads, q_seq_len // bq, grid_width)
      with jax.named_scope(kernel_name):
          all_out = pl.pallas_call(
              partial(
                  flash_attention_kernel,
                  mask_value=mask_value,
                  grid_width=grid_width,
                  bq=bq,
                  bkv=bkv,
                  bkv_compute=bkv_compute,
                  head_dim=head_dim,
                  q_layout=q_layout,
                  k_layout=k_layout,
                  v_layout=v_layout,
                  attn_logits_soft_cap=attn_logits_soft_cap,
                  mask_function=mask_function,
              ),
              grid_spec=pltpu.PrefetchScalarGridSpec(
                  num_scalar_prefetch=num_scalar_prefetch,
                  in_specs=in_specs,
                  out_specs=out_specs,
                  grid=grid,
              ),
              mosaic_params=mosaic_params,
              out_shape=out_shapes,
              name=kernel_name,
          )(
              fwd_mask_info.data_next,
              fwd_mask_info.block_mask,
              fwd_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,
              fwd_mask_info.partial_mask_blocks,
              q_sequence,
          )
  
      (
          _,
          _,
          _,
          out,
          logsumexp,
      ) = all_out
  
      if save_residuals:
          assert logsumexp is not None
          logsumexp = logsumexp[..., 0]
  
      if residual_checkpoint_name is not None:
          out = ad_checkpoint.checkpoint_name(
              out, name=residual_checkpoint_name
          )
          if logsumexp is not None:
              logsumexp = ad_checkpoint.checkpoint_name(
                  logsumexp, name=residual_checkpoint_name
              )
      if save_residuals:
          return out, (logsumexp,)
      return out
  
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 18 for calling a global function with example usage: _splash_attention_forward(fwd_mask_info, q, k, v, segment_ids, mask_value, is_mqa, block_sizes, residual_checkpoint_name, save_residuals, mask_function, attn_logits_soft_cap) and returns: out, from_head_minor, h,  i,  int, from_head_minor, from_head_minor, next_m,  int,  int, i,  int, int,  next_j, out,  logsumexp, h,  i,  int