code: def mha_forward_kernel(
      q_ref,
      k_ref,
      v_ref,  # Input arrays
      segment_ids_ref: jax.Array | None,  # segment_id arrays
      o_ref: Any,  # Output
      *residual_refs: Any,  # Residual outputs
      num_heads: int,
      sm_scale: float,
      causal: bool,
      block_q: int,
      block_d: int,
      block_k: int,
  ):
      seq_len = q_ref.shape[0]
      start_q = pl.program_id(0)
  
      # o is the buffer where we accumulate the output on sram.
      # m_i and l_i (see FlashAttention paper) are updated during the k,v loop.
      m_i = jnp.zeros(block_q, dtype=jnp.float32) - float("inf")
      l_i = jnp.zeros(block_q, dtype=jnp.float32)
      # acc is the buffer where we accumulate the output on sram.
      o = jnp.zeros((block_q, block_d), dtype=jnp.float32)
  
      # Load q: it will stay in L1 throughout. Indices form a matrix because we
      # read, compute, and write all in 2d chunks. 1 element ~= 1 CUDA thread index.
      # q tile has shape [block_q, block_d], block_d == head_dim.
      curr_q_slice = pl.dslice(start_q * block_q, block_q)
      q = pl.load(q_ref, (curr_q_slice, pl.dslice(None)))
      q_segment_ids = (
          None
          if segment_ids_ref is None
          else pl.load(segment_ids_ref, (curr_q_slice,))
      )
  
      # In FlashAttention algorithm 1 there are 2 loops: slow over tiles of kv (size
      # (Bc == block_k here), and fast over blocks of q (size Br == block_q here).
      # Here we only loop over blocks of kv to process entire seq_len, the loop over
      # blocks of q is carried out by the grid.
      def body(start_k, carry):
          o_prev, m_prev, l_prev = carry
          curr_k_slice = pl.dslice(start_k * block_k, block_k)
  
          k = pl.load(k_ref, (curr_k_slice, slice(None)))
          kv_segment_ids = (
              None
              if segment_ids_ref is None
              else pl.load(segment_ids_ref, (curr_k_slice,))
          )
          qk = pl.dot(q, k.T)  # [block_q, block_k]
          if sm_scale != 1.0:
              qk *= sm_scale  # [block_q, block_k]
  
          # Avoids Triton crash.
          # if num_heads > 2:
          #   qk = qk.astype(q_ref.dtype)
          #   qk = qk.astype(jnp.float32)
  
          if causal or segment_ids_ref is not None:
              mask = None
              if segment_ids_ref is not None:
                  mask = segment_mask(q_segment_ids, kv_segment_ids)
              if causal:
                  span_q = start_q * block_q + jnp.arange(block_q)
                  span_k = start_k * block_k + jnp.arange(block_k)
                  causal_mask = span_q[:, None] >= span_k[None, :]
                  mask = (
                      causal_mask
                      if mask is None
                      else jnp.logical_and(mask, causal_mask)
                  )
              # Apply mask to qk.
              qk = jnp.where(mask, qk, DEFAULT_MASK_VALUE)
  
          m_curr = qk.max(axis=-1)
          m_next = jnp.maximum(m_prev, m_curr)
          correction = jnp.exp(m_prev - m_next)
          l_prev_corr = correction * l_prev
          s_curr = jnp.exp(
              qk - m_next[:, None]
          )  # Use m_next instead of m_curr to avoid a correction on l_curr
          l_curr = s_curr.sum(axis=-1)
          l_next = l_prev_corr + l_curr
          l_next_rcp = 1.0 / l_next
          s_curr = s_curr * l_next_rcp[:, None]
          o_prev_corr = (l_prev_corr * l_next_rcp)[:, None] * o_prev
          v = pl.load(v_ref, (curr_k_slice, pl.dslice(block_d)))
          o_curr = pl.dot(s_curr.astype(v.dtype), v)
  
          o_next = o_prev_corr + o_curr
          return o_next, m_next, l_next
  
      if causal:
          # Ceildiv (`pl.cdiv` and `//` do not work due to type of start_q)
          upper_bound = lax.div(
              block_q * (start_q + 1) + block_k - 1, block_k
          )
      else:
          upper_bound = pl.cdiv(seq_len, block_k)  # type: ignore
      o, m_i, l_i = lax.fori_loop(0, upper_bound, body, (o, m_i, l_i))
  
      if residual_refs:
          l_ref, m_ref = residual_refs
          pl.store(l_ref, (curr_q_slice,), l_i)
          pl.store(m_ref, (curr_q_slice,), m_i)
      # Write output to dram.
      o = o.astype(o_ref.dtype)
      pl.store(o_ref, (curr_q_slice, pl.dslice(None)), o)
  
file_path: jax/jax/experimental/pallas/ops/attention.py
imports: from __future__ import annotations
  import functools
  from typing import Any, Optional
  import jax
  from jax import lax
  from jax.experimental import pallas
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: mha_forward_kernel(q_ref, k_ref, v_ref, segment_ids_ref, o_ref) and returns: o_next,  m_next,  l_next