code: def _flash_attention_impl(
      q,
      k,
      v,
      ab,
      segment_ids,
      save_residuals,
      causal,
      sm_scale,
      block_b,
      block_q,
      block_k_major,
      block_k,
      debug,
  ):
      batch_size, num_heads, q_seq_len, head_dim = q.shape
      _, _, kv_seq_len, _ = k.shape
      _verify_block(
          "block_q",
          "q_seq_len",
          block_q,
          q_seq_len,
          should_divide=False,
      )
      _verify_block(
          "block_k_major", "kv_seq_len", block_k_major, kv_seq_len
      )
      _verify_block("block_k", "kv_seq_len", block_k, kv_seq_len)
      _verify_block(
          "block_b", "batch", block_b, batch_size, should_divide=False
      )
  
      # TODO(apaszke): Tile over heads as well.
      grid = (
          pl.cdiv(batch_size, block_b),
          num_heads,
          pl.cdiv(q_seq_len, block_q),
          kv_seq_len // block_k_major,
      )
  
      def q_index_map(batch_index, head_index, q_seq_index, _):
          return (batch_index, head_index, q_seq_index, 0)
  
      def kv_index_map(
          batch_index, head_index, q_seq_index, kv_seq_index
      ):
          if causal:
              # If the kv block is skipped, prefetch the next valid kv block, i.e. the
              # 0th one to be used for the next block_q rows.
              next_kv_index = lax.select(
                  below_or_on_diag(
                      q_seq_index, block_q, kv_seq_index, block_k_major
                  ),
                  kv_seq_index,
                  0,
              )
          else:
              next_kv_index = kv_seq_index
          return (batch_index, head_index, next_kv_index, 0)
  
      def ab_index_map(
          batch_index, head_index, q_seq_index, kv_seq_index
      ):
          if causal:
              should_run = below_or_on_diag(
                  q_seq_index, block_q, kv_seq_index, block_k_major
              )
              # If the ab block is skipped, prefetch the next valid ab block, i.e. the
              # 0th kv to be used for the next block_q rows.
              next_q_index = lax.select(
                  should_run,
                  q_seq_index,
                  lax.select(
                      q_seq_index == (q_seq_len // block_q) - 1,
                      0,
                      q_seq_index + 1,
                  ),
              )
              next_kv_index = lax.select(should_run, kv_seq_index, 0)
          else:
              next_q_index = q_seq_index
              next_kv_index = kv_seq_index
  
          return (batch_index, head_index, next_q_index, next_kv_index)
  
      def o_index_map(batch_index, head_index, q_seq_index, _):
          return (batch_index, head_index, q_seq_index, 0)
  
      def lm_index_map(batch_index, head_index, q_seq_index, _):
          return (batch_index, head_index, q_seq_index, 0)
  
      kernel = functools.partial(
          _flash_attention_kernel,
          causal=causal,
          mask_value=DEFAULT_MASK_VALUE,
          sm_scale=sm_scale,
          block_k=block_k,
          kv_seq_len=kv_seq_len,
      )
      out_shape = jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype)
      out_shape = [out_shape]
      out_specs = [
          pl.BlockSpec(o_index_map, (block_b, 1, block_q, head_dim))
      ]
  
      if block_k != kv_seq_len:
          scratch_shape = functools.partial(
              jax.ShapeDtypeStruct, dtype=jnp.float32
          )
          m_scratch = scratch_shape(
              (block_b, 1, block_q, MIN_BLOCK_SIZE)
          )
          l_scratch = scratch_shape(
              (block_b, 1, block_q, MIN_BLOCK_SIZE)
          )
          acc_scratch = scratch_shape((block_b, 1, block_q, head_dim))
          out_shape += [m_scratch, l_scratch, acc_scratch]
          out_specs += [
              pl.BlockSpec(lambda *_: (0, 0, 0, 0), m_scratch.shape),
              pl.BlockSpec(lambda *_: (0, 0, 0, 0), l_scratch.shape),
              pl.BlockSpec(lambda *_: (0, 0, 0, 0), acc_scratch.shape),
          ]
      else:
          out_shape += [None, None, None]
          out_specs += [None, None, None]
  
      if save_residuals:
          out_specs = [
              *out_specs,
              pl.BlockSpec(
                  lm_index_map, (block_b, 1, block_q, MIN_BLOCK_SIZE)
              ),
              pl.BlockSpec(
                  lm_index_map, (block_b, 1, block_q, MIN_BLOCK_SIZE)
              ),
          ]
          l = jax.ShapeDtypeStruct(
              (batch_size, num_heads, q_seq_len, MIN_BLOCK_SIZE),
              dtype=jnp.float32,
          )
          m = jax.ShapeDtypeStruct(
              (batch_size, num_heads, q_seq_len, MIN_BLOCK_SIZE),
              dtype=jnp.float32,
          )
          out_shape = (*out_shape, l, m)
  
      ab_block_spec = (
          pl.BlockSpec(
              ab_index_map, (block_b, 1, block_q, block_k_major)
          )
          if ab is not None
          else None
      )
  
      q_segment_ids_spec = kv_segment_ids_spec = None
      q_segment_ids = kv_segment_ids = None
      if segment_ids is not None:
  
          def q_segment_ids_index_map(
              batch_index, head_index, q_seq_index, _
          ):
              del head_index
              return (batch_index, q_seq_index, 0)
  
          def kv_segment_ids_index_map(
              batch_index, head_index, q_seq_index, kv_seq_index
          ):
              del head_index
              if causal:
                  next_kv_index = lax.select(
                      below_or_on_diag(
                          q_seq_index,
                          block_q,
                          kv_seq_index,
                          block_k_major,
                      ),
                      kv_seq_index,
                      0,
                  )
              else:
                  next_kv_index = kv_seq_index
              return (batch_index, 0, next_kv_index)
  
          q_segment_ids_spec = pl.BlockSpec(
              q_segment_ids_index_map, (block_b, block_q, NUM_LANES)
          )
          kv_segment_ids_spec = pl.BlockSpec(
              kv_segment_ids_index_map,
              (block_b, NUM_SUBLANES, block_k_major),
          )
  
          q_segment_ids = jax.lax.broadcast_in_dim(
              segment_ids.q,
              (batch_size, q_seq_len, NUM_LANES),
              (
                  0,
                  1,
              ),
          )
          kv_segment_ids = jax.lax.broadcast_in_dim(
              segment_ids.kv,
              (batch_size, NUM_SUBLANES, kv_seq_len),
              (
                  0,
                  2,
              ),
          )
  
      in_specs = [
          pl.BlockSpec(q_index_map, (block_b, 1, block_q, head_dim)),
          pl.BlockSpec(
              kv_index_map, (block_b, 1, block_k_major, head_dim)
          ),
          pl.BlockSpec(
              kv_index_map, (block_b, 1, block_k_major, head_dim)
          ),
          ab_block_spec,
          q_segment_ids_spec,
          kv_segment_ids_spec,
      ]
  
      o, *aux = pl.pallas_call(
          kernel,
          out_shape=out_shape,
          in_specs=in_specs,
          out_specs=out_specs,
          grid=grid,
          debug=debug,
          mosaic_params=dict(
              dimension_semantics=(
                  "parallel",
                  "parallel",
                  "parallel",
                  "arbitrary",
              )
          ),
      )(q, k, v, ab, q_segment_ids, kv_segment_ids)
      if save_residuals:
          l, m = (v[..., 0] for v in aux[-2:])
          return (o, l, m)
      else:
          return o
  
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 9 for calling a global function with example usage: _flash_attention_impl(q, k, v, ab, segment_ids, save_residuals, causal, sm_scale, block_b, block_q, block_k_major, block_k, debug) and returns: batch_index,  head_index,  q_seq_index,  int, batch_index,  head_index,  next_kv_index,  int, batch_index,  head_index,  next_q_index,  next_kv_index, batch_index,  head_index,  q_seq_index,  int, batch_index,  head_index,  q_seq_index,  int, o,  l,  m, o, batch_index,  q_seq_index,  int, batch_index,  int,  next_kv_index