code: def flash_attention_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,
      mask_ref,
      q_sequence_ref,
      # Outputs
      m_scratch_ref,
      l_scratch_ref,
      o_scratch_ref,
      o_ref,
      logsumexp_ref=None,
      *,
      mask_value: float,
      grid_width: int,
      bq: int,
      bkv: int,
      bkv_compute: int,
      head_dim: int,
      q_layout: QKVLayout,
      k_layout: QKVLayout,
      v_layout: QKVLayout,
      attn_logits_soft_cap: float | None,
      mask_function: MaskFunctionType | None,
  ):
      float32 = jnp.float32
      HEAD_DIM_MINOR = QKVLayout.HEAD_DIM_MINOR
  
      head_dim_repeats, rem = divmod(head_dim, NUM_LANES)
      if rem != 0:
          raise NotImplementedError(
              f"{head_dim=} should be a multiple of {NUM_LANES}"
          )
  
      h, i, j = pl.program_id(0), pl.program_id(1), pl.program_id(2)
  
      @pl.when(j == 0)
      def init():
          o_scratch_ref[...] = jnp.zeros_like(o_scratch_ref)
          m_scratch_ref[...] = jnp.full_like(m_scratch_ref, mask_value)
          l_scratch_ref[...] = jnp.zeros_like(l_scratch_ref)
  
      _, _, should_run, should_not_mask = _next_nonzero(
          h,
          i,
          j,
          data_next_ref,
          block_mask_ref,
          mask_next_ref,
      )
  
      def body(kv_compute_index, _):
          slice_k = pl.ds(kv_compute_index * bkv_compute, bkv_compute)
          m_prev, l_prev = m_scratch_ref[...], l_scratch_ref[...]
          assert m_prev.shape == (bq, NUM_LANES)
          assert l_prev.shape == (bq, NUM_LANES)
  
          q = q_ref[...] if q_layout == HEAD_DIM_MINOR else q_ref[...].T
          qk_dims = (
              NT_DIM_NUMBERS
              if k_layout == HEAD_DIM_MINOR
              else NN_DIM_NUMBERS
          )
          if k_layout == HEAD_DIM_MINOR:
              k = pl.load(k_ref, (slice_k, slice(None)))
          else:
              k = pl.load(k_ref, (slice(None), slice_k))
          qk = lax.dot_general(
              q, k, qk_dims, preferred_element_type=float32
          )
  
          assert qk.shape == (bq, bkv_compute)
          apply_mask_and_soft_cap = functools.partial(
              _apply_mask_and_soft_cap,
              qk,
              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=j * bkv + kv_compute_index * bkv_compute,
              bq=bq,
              mask_function=mask_function,
          )
  
          qk = apply_mask_and_soft_cap()
  
          m_curr = qk.max(axis=-1)[:, None]
          assert m_curr.shape == (bq, 1)
          m_next = jnp.maximum(m_prev, m_curr)
          assert m_next.shape == (bq, NUM_LANES)
  
          bkv_repeats, rem = divmod(bkv_compute, NUM_LANES)
          if rem != 0:
              raise NotImplementedError(
                  f"{bkv_compute=} should be a multiple of {NUM_LANES}"
              )
  
          s_curr = jnp.exp(
              qk - pltpu.repeat(m_next, bkv_repeats, axis=1)
          )
          assert s_curr.shape == (bq, bkv_compute)
  
          l_curr = jax.lax.broadcast_in_dim(
              s_curr.sum(axis=-1), l_prev.shape, (0,)
          )
          assert l_curr.shape == (bq, NUM_LANES)
  
          alpha = jnp.exp(m_prev - m_next)
          l_next = l_curr + alpha * l_prev
          m_scratch_ref[...], l_scratch_ref[...] = m_next, l_next
  
          sv_dims = (
              NN_DIM_NUMBERS
              if v_layout == HEAD_DIM_MINOR
              else NT_DIM_NUMBERS
          )
          if v_layout == HEAD_DIM_MINOR:
              v = pl.load(v_ref, (slice_k, slice(None)))
          else:
              v = pl.load(v_ref, (slice(None), slice_k))
          v = v.astype(float32)
          o_curr = lax.dot_general(s_curr, v, sv_dims)
  
          alpha_o = pltpu.repeat(alpha, head_dim_repeats, axis=1)
          o_scratch_ref[:] = alpha_o * o_scratch_ref[:] + o_curr
  
      @pl.when(should_run)
      def run():
          assert bkv % bkv_compute == 0
          num_iters = (
              k_ref.shape[0 if k_layout == HEAD_DIM_MINOR else 1]
              // bkv_compute
          )
          lax.fori_loop(0, num_iters, body, None, unroll=True)
  
      @pl.when(j == grid_width - 1)
      def end():
          l = l_scratch_ref[...]
          l_inv = pltpu.repeat(1.0 / l, head_dim_repeats, axis=1)
          o_ref[...] = (o_scratch_ref[...] * l_inv).astype(o_ref.dtype)
          if logsumexp_ref is not None:
              assert logsumexp_ref.shape == (bq, NUM_LANES)
              logsumexp_ref[...] = (
                  jnp.log(l) + m_scratch_ref[...]
              ).astype(logsumexp_ref.dtype)
  
          m_scratch_ref[...] = jnp.zeros_like(m_scratch_ref)
          l_scratch_ref[...] = jnp.zeros_like(l_scratch_ref)
          o_scratch_ref[...] = jnp.zeros_like(o_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 14 for calling a global function with example usage: flash_attention_kernel(data_next_ref, block_mask_ref, mask_next_ref, q_ref, k_ref, v_ref, q_segment_ids_ref, kv_segment_ids_ref, mask_ref, q_sequence_ref, m_scratch_ref, l_scratch_ref, o_scratch_ref, o_ref, logsumexp_ref) without return types