Repository: google/jax
Pull request #19200 (2024-01-04 20:43:12)

[Pallas/TPU] Open source "Splash Attention" (Sparse Flash Attention), a general purpose attention kernel where you can specify an attention mask using NumPy.
[Pallas/TPU] Open source "Splash Attention" (Sparse Flash Attention), a general purpose attention kernel where you can specify an attention mask using NumPy.



Patch:
diff --git a/jax/BUILD b/jax/BUILD
index d46f1b3a6..957c272e2 100644
--- a/jax/BUILD
+++ b/jax/BUILD
@@ -565,7 +565,7 @@ pytype_strict_library(
         exclude = [
             "experimental/pallas/gpu.py",
             "experimental/pallas/tpu.py",
-            "experimental/pallas/ops/tpu/*.py",
+            "experimental/pallas/ops/tpu/**/*.py",
         ],
     ),
     visibility = [
@@ -593,7 +593,7 @@ pytype_strict_library(
 
 pytype_strict_library(
     name = "pallas_tpu_ops",
-    srcs = glob(["experimental/pallas/ops/tpu/*.py"]),
+    srcs = glob(["experimental/pallas/ops/tpu/**/*.py"]),
     visibility = [
         ":pallas_tpu_users",
     ],
@@ -601,7 +601,7 @@ pytype_strict_library(
         ":jax",
         ":pallas",
         ":pallas_tpu",
-    ],
+    ] + py_deps("numpy"),
 )
 
 pytype_strict_library(
diff --git a/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py b/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py
new file mode 100644
index 000000000..602c07558
--- /dev/null
+++ b/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py
@@ -0,0 +1,2460 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Implementation of Sparse Flash Attention, a.k.a. "Splash" attention."""
+
+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 as pl
+from jax.experimental.pallas import tpu as pltpu
+from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask as mask_lib
+from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask_info as mask_info_lib
+import jax.numpy as jnp
+import numpy as np
+
+partial = functools.partial
+DEFAULT_MASK_VALUE = -0.7 * float(np.finfo(np.dtype("float32")).max)
+NUM_LANES = 128
+NUM_SUBLANES = 8
+# We predefine some useful dimension numbers for dot_general
+NN_DIM_NUMBERS = (((1,), (0,)), ((), ()))  # standard matmul
+NT_DIM_NUMBERS = (((1,), (1,)), ((), ()))  # RHS transposed
+
+# mypy: ignore-errors
+
+class SegmentIds(NamedTuple):
+  """SegmentIds for Q and KV sequences.
+
+  SegmentIds are a mechanims to ensure that there is no cross-attention between
+  segments (fraction of a sequence) that have been concatenated together into a
+  sequence. Each array is a list of ids (integers). Only tokens with the same
+  id are allowed to attend to each other.
+
+  The static mask (e.g. causal) is "and-ed" with the segment id mask to form
+  the actual attention mask. It is important that the latter does not have any
+  all-zero rows (along dimension kv). Otherwise it would result in a invalid
+  softmax (the denominator would be 0).
+  This condition holds for causal self-attention because in this case segment
+  ids form a block diagonal matrix so at least one element in each row is set.
+  It is easy to break this condition with non-self-attention configurations.
+  Attributes:
+    q: segment ids along the Q sequence
+    kv: segment ids along the KV sequence
+  """
+
+  q: jax.Array  # [q_seq_len]
+  kv: jax.Array  # [kv_seq_len]
+
+
+# Return type of SplashAttention function that implements the custom vjp rule.
+SplashCustomReturnType = Union[
+    # out, no residuals
+    jax.Array,
+    # out, residuals:
+    tuple[jax.Array, tuple[jax.Array,]]
+]
+
+SplashResidualsType = tuple[
+    jax.Array,  # q
+    jax.Array,  # k
+    jax.Array,  # v
+    Optional[SegmentIds],  # segment_ids
+    jax.Array,  # out
+    jax.Array,  # logsumexp
+    Optional[mask_info_lib.MaskInfo],  # dq_mask_info
+    Optional[mask_info_lib.MaskInfo],  # dkv_mask_info
+]
+
+MaskFunctionType = Callable[..., jax.Array]
+
+
+def get_kernel_name(
+    is_mqa: bool, save_residuals: bool, is_segmented: bool, phase: str
+) -> str:
+  """Returns a unique name for all SplashAttention kernel variants."""
+
+  assert phase == "dq" or phase == "dkv" or phase == "fwd"
+  # Saving residuals is supported only for the fwd phase.
+  assert not save_residuals or phase == "fwd"
+  residuals = ""
+  if save_residuals:
+    residuals = "_residuals"
+  elif phase == "fwd":
+    residuals = "_no_residuals"
+  attention_type = "mqa" if is_mqa else "mha"
+  segments = "_segmented" if is_segmented else ""
+  return f"splash_{attention_type}_{phase}{segments}{residuals}"
+
+
+# Reference attention implementations
+
+
+@overload
+def _attention_reference(
+    mask: jax.Array,
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    segment_ids: SegmentIds | None,
+    save_residuals: Literal[False],
+    mask_value: float,
+    custom_type: str,
+    attn_logits_soft_cap: float | None,
+) -> jax.Array:
+  ...
+
+
+@overload
+def _attention_reference(
+    mask: jax.Array,
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    segment_ids: SegmentIds | None,
+    save_residuals: Literal[True],
+    mask_value: float,
+    custom_type: str,
+    attn_logits_soft_cap: float | None,
+) -> tuple[jax.Array, tuple[jax.Array]]:
+  ...
+
+
+def _attention_reference(
+    mask: jax.Array,  # [q_seq_len, kv_seq_len]
+    q: jax.Array,  # [q_seq_len, head_dim]
+    k: jax.Array,  # [kv_seq_len, head_dim]
+    v: jax.Array,  # [kv_seq_len, head_dim]
+    segment_ids: SegmentIds | None,
+    mask_value: float,
+    save_residuals: bool,
+    custom_type: str,
+    attn_logits_soft_cap: float | None,
+):
+  return _attention_reference_default(  # pytype: disable=bad-return-type
+      mask,
+      q,
+      k,
+      v,
+      segment_ids,
+      mask_value,
+      save_residuals,
+      custom_type,
+      attn_logits_soft_cap,
+  )
+
+
+def _attention_reference_default(
+    mask: jax.Array,  # [q_seq_len, kv_seq_len]
+    q: jax.Array,  # [q_seq_len, head_dim]
+    k: jax.Array,  # [kv_seq_len, head_dim]
+    v: jax.Array,  # [kv_seq_len, head_dim]
+    segment_ids: SegmentIds | None,
+    mask_value: float,
+    save_residuals: bool,
+    custom_type: str,
+    attn_logits_soft_cap: float | None,
+):
+  del custom_type
+  logits = jnp.einsum("sd,td->st", q.astype(jnp.float32), k.astype(jnp.float32))
+
+  if segment_ids is not None:
+    mask = jnp.logical_and(
+        mask, segment_ids.q[:, None] == segment_ids.kv[None, :]
+    )
+
+  if attn_logits_soft_cap is not None:
+    logits = jnp.tanh(logits / attn_logits_soft_cap)
+    logits = logits * attn_logits_soft_cap
+
+  logits = jnp.where(mask, logits, mask_value)
+  m = logits.max(axis=-1)
+  s = jnp.exp(logits - m[..., None])
+  l = s.sum(axis=-1)
+  s = s / l[..., None]
+
+  o = jnp.einsum("st,td->sd", s, v.astype(jnp.float32))
+
+  logsumexp = m + jnp.log(l)
+  if save_residuals:
+    return o, (logsumexp,)
+  return o
+
+
+def attention_reference(
+    mask: jax.Array,  # [q_seq_len, kv_seq_len]
+    q: jax.Array,  # [q_seq_len, head_dim]
+    k: jax.Array,  # [kv_seq_len, head_dim]
+    v: jax.Array,  # [kv_seq_len, head_dim]
+    segment_ids: SegmentIds | None,
+    *,
+    mask_value: float = DEFAULT_MASK_VALUE,
+    save_residuals: bool = False,
+    custom_type: str = "flash",
+    attn_logits_soft_cap: float | None = None,
+) -> SplashCustomReturnType:
+  return _attention_reference(  # pytype: disable=wrong-arg-types
+      mask,
+      q,
+      k,
+      v,
+      segment_ids,
+      mask_value=mask_value,
+      save_residuals=save_residuals,
+      custom_type=custom_type,
+      attn_logits_soft_cap=attn_logits_soft_cap,
+  )
+
+
+def _attention_reference_custom_fwd(
+    mask: jax.Array,  # [q_seq_len, kv_seq_len]
+    q: jax.Array,  # [q_seq_len, head_dim]
+    k: jax.Array,  # [kv_seq_len, head_dim]
+    v: jax.Array,  # [kv_seq_len, head_dim]
+    segment_ids: SegmentIds | None,
+    mask_value: float,
+    save_residuals: bool,
+    custom_type: str,
+    attn_logits_soft_cap: float | None,
+):
+  if save_residuals:
+    raise NotImplementedError("Higher-order AD not supported.")
+
+  o, (logsumexp,) = _attention_reference(
+      mask,
+      q,
+      k,
+      v,
+      segment_ids,
+      mask_value=mask_value,
+      save_residuals=True,
+      custom_type=custom_type,
+      attn_logits_soft_cap=attn_logits_soft_cap,
+  )
+  return o, (mask, q, k, v, segment_ids, o, logsumexp)
+
+
+def _attention_reference_custom_bwd(
+    mask_value: float,
+    save_residuals: bool,
+    custom_type: str,
+    attn_logits_soft_cap: float | None,
+    res,
+    do: jax.Array,
+) -> tuple[None, jax.Array, jax.Array, jax.Array, None]:
+  del save_residuals
+  mask, q, k, v, segment_ids, o, logsumexp = res
+
+  uncapped_logits = jnp.einsum(
+      "qc,kc->qk", q, k, preferred_element_type=jnp.float32)
+
+  if attn_logits_soft_cap is not None:
+    logits = jnp.tanh(uncapped_logits / attn_logits_soft_cap)
+    logits = logits * attn_logits_soft_cap
+  else:
+    logits = uncapped_logits
+
+  if segment_ids is not None:
+    mask = jnp.logical_and(
+        mask, segment_ids.q[:, None] == segment_ids.kv[None, :]
+    )
+  logits = jnp.where(mask, logits, mask_value)
+
+  p = jnp.exp(logits - logsumexp[..., None])
+  do = do.astype(jnp.float32)  # pytype: disable=attribute-error
+  dv = jnp.einsum("pt,pd->td", p, do).astype(v.dtype)
+  dp = jnp.einsum("pd,td->pt", do, v.astype(jnp.float32))
+
+  # These two ways of computing ds are mathematically equivalent. The first
+  # involves reducing over the head_dim dimension and the second involves
+  # reducing over a sequence dimension. They tend to produce slightly different
+  # numerics.
+  if custom_type == "flash":
+    di = jnp.sum(o.astype(jnp.float32) * do, axis=-1)[..., None]
+  else:
+    di = jnp.einsum("st,st->s", dp, p)[:, None]
+  ds = (dp - di) * p
+  if attn_logits_soft_cap is not None:
+    normalized = uncapped_logits / attn_logits_soft_cap
+    d = jnp.tanh(normalized)
+    g = ds * (1 - d)
+    ds = g + g * d
+  dk = jnp.einsum("sd,st->td", q.astype(jnp.float32), ds).astype(k.dtype)
+  dq = jnp.einsum("st,td->sd", ds, k.astype(jnp.float32)).astype(q.dtype)
+  return None, dq, dk, dv, None
+
+
+_attention_reference_custom = jax.custom_vjp(
+    _attention_reference, nondiff_argnums=(5, 6, 7, 8)
+)
+_attention_reference_custom.defvjp(_attention_reference_custom_fwd,
+                                   _attention_reference_custom_bwd)
+
+
+def attention_reference_custom(
+    mask: jax.Array,  # [q_seq_len, kv_seq_len]
+    q: jax.Array,  # [q_seq_len, head_dim]
+    k: jax.Array,  # [kv_seq_len, head_dim]
+    v: jax.Array,  # [kv_seq_len, head_dim]
+    segment_ids: SegmentIds | None,
+    *,
+    mask_value: float = DEFAULT_MASK_VALUE,
+    save_residuals: bool = False,
+    custom_type: str = "flash",
+    attn_logits_soft_cap: float | None = None,
+):
+  return _attention_reference_custom(
+      mask,
+      q,
+      k,
+      v,
+      segment_ids,
+      mask_value,
+      save_residuals,
+      custom_type=custom_type,
+      attn_logits_soft_cap=attn_logits_soft_cap,
+  )
+
+
+def make_attention_reference(
+    mask: mask_lib.Mask | np.ndarray,
+    is_mqa: bool,
+    backward_impl: str = "vanilla",
+    **params: Any,
+) -> Callable:
+  @partial(
+      jax.jit,
+      static_argnames=[
+          "mask_value",
+          "save_residuals",
+          "attn_logits_soft_cap",
+      ],
+  )
+  def _wrapped(
+      mask: jax.Array,
+      q: jax.Array,
+      k: jax.Array,
+      v: jax.Array,
+      segment_ids: SegmentIds | None = None,
+      *,
+      mask_value: float = DEFAULT_MASK_VALUE,
+      save_residuals: bool = False,
+      attn_logits_soft_cap: float | None = None,
+  ):
+    if backward_impl == "custom":
+      attn_impl = partial(
+          attention_reference_custom, custom_type="flash",
+      )
+    elif backward_impl == "custom_vanilla":
+      attn_impl = partial(
+          attention_reference_custom, custom_type="vanilla",
+      )
+    else:
+      attn_impl = attention_reference
+    func = partial(
+        attn_impl,
+        mask_value=mask_value,
+        save_residuals=save_residuals,
+        attn_logits_soft_cap=attn_logits_soft_cap,
+        **params,
+    )
+
+    if is_mqa:
+      func = jax.vmap(func, in_axes=(0, 0, None, None, None))
+      is_grouped = False
+    else:
+      # In grouped attention (1 < num_kv_heads && num_kv_heads < num_q_heads).
+      # We interleave the KV heads across the Q heads.
+      # For example: for 8 Q heads and 4 KV heads:
+      # Q head [0, 1] see KV head 0
+      # Q head [2, 3] see KV head 1
+      # Q head [4, 5] see KV head 2
+      # Q head [6, 7] see KV head 3
+      #
+      # The following implementation reshapes Q to expose KV heads and vmaps
+      # Across the Q heads so it is similar to MQA.
+      # Alternatively we can replicate K/V to match Q like so:
+      # k = jnp.repeat(k, q_heads_per_kv_head, axis=0)
+      # v = jnp.repeat(v, q_heads_per_kv_head, axis=0)
+
+      kv_heads = k.shape[0]
+      assert kv_heads == v.shape[0]
+      q_heads, q_seq_len, head_dim = q.shape
+      is_grouped = kv_heads < q_heads
+      if is_grouped:
+        assert q_heads % kv_heads == 0
+        assert mask.shape[0] == q_heads
+        q_heads_per_kv_head = q_heads // kv_heads
+        q = q.reshape((kv_heads, q_heads_per_kv_head, q_seq_len, head_dim))
+        mask = mask.reshape((kv_heads, q_heads_per_kv_head, *mask.shape[1:]))
+
+        # Inner-most vmap: iterate over the q heads.
+        func = jax.vmap(func, in_axes=(0, 0, None, None, None))
+
+      # Outer-most vmap: iterate over the kv heads.
+      func = jax.vmap(func, in_axes=(0, 0, 0, 0, None))
+
+    out = func(mask, q, k, v, segment_ids)
+
+    if is_grouped:
+
+      def reshape_activations(activations):
+        if activations.ndim == 4:
+          kv_heads, q_heads_per_kv_head, q_seq_len, head_dim = activations.shape
+          return activations.reshape(
+              kv_heads * q_heads_per_kv_head, q_seq_len, head_dim
+          )
+        return activations
+
+      def reshape_residuals(residuals):
+        if residuals.ndim == 3:
+          kv_heads, q_heads_per_kv_head, q_seq_len = residuals.shape
+          return residuals.reshape(kv_heads * q_heads_per_kv_head, q_seq_len)
+        return residuals
+
+      if save_residuals:
+        assert isinstance(out, tuple)
+        assert isinstance(out[1], tuple)
+
+        return (reshape_activations(out[0]), (reshape_residuals(out[1][0]),))
+      else:
+        return reshape_activations(out)
+    else:
+      return out
+
+  return functools.partial(_wrapped, jnp.array(mask[:, :, :]))
+
+
+make_masked_mha_reference = partial(make_attention_reference, is_mqa=False)
+make_masked_mqa_reference = partial(make_attention_reference, is_mqa=True)
+
+
+# Splash attention implementation
+
+# We use an IntEnum to make it JSON serializable as regen metadata.
+class QKVLayout(enum.IntEnum):
+  HEAD_DIM_MINOR = enum.auto()  # [..., seq_len, head_dim]
+  SEQ_MINOR = enum.auto()  # [..., head_dim, seq_len]
+
+
+def from_head_minor(vals: tuple[Any, ...], layout: QKVLayout):
+  if layout == QKVLayout.HEAD_DIM_MINOR:
+    return vals
+  return (*vals[:-2], vals[-1], vals[-2])
+
+
+@dataclasses.dataclass(unsafe_hash=True)
+class BlockSizes:
+  """Tile sizes parameterizing SplashAttention kernels.
+
+  Those parameters have negligible effect on numerics, but affect performance
+  greatly.
+
+  Note that changing the layouts only influences the physical layout that the
+  kernel will enforce. The logical interface to splash attention always takes
+  the head dimension as the minormost one.
+  """
+  block_q: int
+  block_kv: int
+  block_kv_compute: int | None = None
+
+  block_q_dkv: int | None = None
+  block_kv_dkv: int | None = None
+  block_kv_dkv_compute: int | None = None
+
+  block_q_dq: int | None = None
+  block_kv_dq: int | None = None
+
+  use_fused_bwd_kernel: bool = False
+
+  q_layout: QKVLayout = QKVLayout.HEAD_DIM_MINOR
+  k_layout: QKVLayout = QKVLayout.HEAD_DIM_MINOR
+  v_layout: QKVLayout = QKVLayout.HEAD_DIM_MINOR
+
+  def __post_init__(self):
+    if self.block_kv_compute is None:
+      self.block_kv_compute = self.block_kv
+    if self.block_kv_dkv_compute is None:
+      self.block_kv_dkv_compute = self.block_kv_dkv
+    if self.use_fused_bwd_kernel:
+      if self.block_q_dq is not None or self.block_kv_dq is not None:
+        raise ValueError(
+            "Block sizes for dq kernel are not needed with a fused kernel."
+        )
+
+  @property
+  def has_backward_blocks(self) -> bool:
+    backward_blocks = (
+        self.block_q_dkv, self.block_kv_dkv, self.block_kv_dkv_compute,
+    )
+    if not self.use_fused_bwd_kernel:
+      backward_blocks += (self.block_q_dq, self.block_kv_dq)
+    return all(b is not None for b in backward_blocks)
+
+  @classmethod
+  def get_default(cls):
+    # TODO(apaszke,sharadmv): Select better parameters based on a heuristic.
+    return BlockSizes(
+        block_q=128,
+        block_kv=128,
+        block_kv_compute=128,
+        block_q_dkv=128,
+        block_kv_dkv=128,
+        block_kv_dkv_compute=128,
+        block_q_dq=128,
+        block_kv_dq=128,
+    )
+
+
+def _next_nonzero(
+    h,
+    i,
+    j,
+    data_next_ref,
+    block_mask_ref,
+    m_next_ref,
+    next_i=False,
+):
+  assert (data_next_ref is None) == (block_mask_ref is None)
+
+  if data_next_ref is None and block_mask_ref is None:
+    # Handle the case in which we have no masking nor next data information.
+    # Simply fetch the next data and apply the mask for every block.
+    assert m_next_ref is None
+    next_data = i if next_i else j
+    return (
+        next_data,
+        None,  # next mask
+        True,  # should run
+        False,  # should not mask
+    )
+
+  assert data_next_ref.shape == block_mask_ref.shape
+  assert m_next_ref is None or data_next_ref.shape[0] == m_next_ref.shape[0]
+
+  # We are working with one head only. Force the head index to 0.
+  if data_next_ref.shape[0] == 1:
+    h = 0
+
+  # When scalar-memory data is of types smaller than int32, then we have to
+  # upcast it back to use it in the kernel.
+
+  to_i32 = lambda x: x.astype(jnp.int32)
+
+  is_nonzero = to_i32(block_mask_ref[h, i, j]) > 0
+  if m_next_ref is None:
+    should_not_mask = True
+    next_m = None
+  else:
+    should_not_mask = to_i32(block_mask_ref[h, i, j]) != 1
+    next_m = to_i32(m_next_ref[h, i, j])
+  next_j = to_i32(data_next_ref[h, i, j])
+  return next_j, next_m, is_nonzero, should_not_mask
+
+
+def _apply_mask_and_soft_cap(
+    qk: jax.Array,
+    mask_value: float,
+    should_not_mask,
+    mask_ref,
+    q_sequence_ref,
+    q_segment_ids_ref,
+    kv_segment_ids_ref,
+    *,
+    attn_logits_soft_cap: float,
+    k_slice: pl.Slice,
+    k_offset: int,
+    bq: int,
+    k_in_lanes=True,
+    mask_function=None,
+) -> jax.Array | tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
+  assert mask_ref is None or q_sequence_ref is None
+  assert (q_sequence_ref is None) == (mask_function is None)
+
+  masks = []
+  if mask_ref is not None:
+    if k_in_lanes:
+      mask = pl.load(mask_ref, (slice(None), k_slice))
+    else:
+      mask = pl.load(mask_ref, (k_slice, slice(None)))
+
+    snm = jnp.where(should_not_mask, 1, 0)
+    masks.append(jnp.bitwise_or(mask, jnp.broadcast_to(snm, mask.shape)))
+
+  if mask_function is not None:
+    # Compute the mask using the given q_sequence indices.
+    # KV indices are computed on the fly. This works because we only support Q
+    # sequence sharding. If we wanted to compute Q indices too, then we would
+    # need to keep into account the current shard along Q sequence.
+
+    if k_in_lanes:
+      assert q_sequence_ref.shape == (bq, NUM_LANES)
+
+      k_sequence = k_offset + jax.lax.broadcasted_iota(
+          jnp.int32, (bq, k_slice.size), 1
+      )
+
+      repeats, rem = divmod(k_slice.size, NUM_LANES)
+      assert rem == 0
+      q_sequence = pltpu.repeat(
+          q_sequence_ref[...], repeats, axis=1
+      )  # [bq, k_slice.size]
+    else:
+      assert q_sequence_ref.shape == (NUM_SUBLANES, bq)
+
+      k_sequence = k_offset + jax.lax.broadcasted_iota(
+          jnp.int32, (k_slice.size, bq), 0
+      )
+      q_sequence = pl.load(q_sequence_ref, (pl.ds(1), slice(None)))  # [1, bq]
+      q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq))
+
+    assert q_sequence.shape == k_sequence.shape
+    masks.append(mask_function(q_sequence, k_sequence))  # pytype: disable=wrong-arg-count
+
+  if q_segment_ids_ref is not None:
+    if k_in_lanes:
+      kv_ids = pl.load(kv_segment_ids_ref, (pl.ds(1), k_slice))  # [1, k_slice]
+      repeats, rem = divmod(kv_ids.shape[1], NUM_LANES)
+      if rem:
+        raise NotImplementedError(f"block_kv must be a multiple of {NUM_LANES}")
+      q_ids = pltpu.repeat(q_segment_ids_ref[:], repeats, axis=1)  # [bq, bkv]
+    else:
+      assert bq == q_segment_ids_ref.shape[-1]
+      repeats, rem = divmod(bq, NUM_LANES)
+      if rem:
+        raise NotImplementedError(f"block_q must be a multiple of {NUM_LANES}")
+      kv_ids = pltpu.repeat(
+          pl.load(kv_segment_ids_ref, (k_slice, slice(None))), repeats, axis=1
+      )  # [k_slice, bq]
+      q_ids = pl.load(q_segment_ids_ref, (pl.ds(1), slice(None)))  # [1, bq]
+    masks.append(q_ids == kv_ids)
+
+  def cap_logits(logits):
+    if attn_logits_soft_cap is not None:
+      logits = jnp.tanh(qk / attn_logits_soft_cap)
+      return logits * attn_logits_soft_cap
+    else:
+      return logits
+
+  if masks:
+    mask = functools.reduce(jnp.logical_and, masks)
+    qk = cap_logits(qk)
+    qk = jnp.where(mask, qk, mask_value)
+  else:
+    qk = cap_logits(qk)
+  return qk
+
+
+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)
+
+
+@overload
+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,
+    mask_function: MaskFunctionType | None,
+    save_residuals: Literal[False] = False,
+    attn_logits_soft_cap: float | None = None,
+) -> jax.Array:
+  ...
+
+
+@overload
+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,
+    mask_function: MaskFunctionType | None,
+    save_residuals: Literal[True],
+    attn_logits_soft_cap: float | None = None,
+) -> SplashCustomReturnType:
+  ...
+
+
+def _div(dividend: int, divisor: int):
+  if divisor == 1:
+    return dividend
+
+  return lax.div(dividend, divisor)
+
+
+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
+
+
+@partial(jax.custom_vjp, nondiff_argnums=(7, 8, 9, 10, 11, 12, 13))
+def _splash_attention_custom(
+    fwd_mask_info: mask_info_lib.MaskInfo,
+    dq_mask_info: mask_info_lib.MaskInfo | None,
+    dkv_mask_info: mask_info_lib.MaskInfo | None,
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    segment_ids: SegmentIds | None,
+    save_residuals: bool,
+    mask_value: float,
+    is_mqa: bool,
+    block_sizes: BlockSizes,
+    residual_checkpoint_name: str | None,
+    mask_function: MaskFunctionType | None,
+    attn_logits_soft_cap: float | None = None,
+) -> SplashCustomReturnType:
+  # The forward function does not use the dq and dkv MaskInfos, it just forwards
+  # them to the backward function as residuals. This is a way to communicate
+  # arbitrary Arrays to the backward function. Since the three MaskInfos are
+  # constants there is no overhead in passing them to the backward function as
+  # residuals. When sharding computation MaskInfos are partitioned so both the
+  # forward and the backward kernels need to work on the relevant slice. If we
+  # recomputed the backward MaskInfos in the backward function from the numpy
+  # mask then we would not work with the MaskInfo slice relevant to the current
+  # device.
+  del dq_mask_info, dkv_mask_info
+
+  return _splash_attention_forward(  # pytype: disable=wrong-arg-types
+      fwd_mask_info,
+      q,
+      k,
+      v,
+      segment_ids,
+      mask_value=mask_value,
+      is_mqa=is_mqa,
+      block_sizes=block_sizes,
+      save_residuals=save_residuals,
+      attn_logits_soft_cap=attn_logits_soft_cap,
+      residual_checkpoint_name=residual_checkpoint_name,
+      mask_function=mask_function,
+  )
+
+
+def _splash_attention_fwd(
+    fwd_mask_info: mask_info_lib.MaskInfo,
+    dq_mask_info: mask_info_lib.MaskInfo | None,
+    dkv_mask_info: mask_info_lib.MaskInfo | None,
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    segment_ids: SegmentIds | None,
+    save_residuals: bool,
+    mask_value: float,
+    is_mqa: bool,
+    block_sizes: BlockSizes,
+    residual_checkpoint_name: str | None,
+    mask_function: MaskFunctionType | None,
+    attn_logits_soft_cap: float | None = None,
+) -> tuple[
+    tuple[jax.Array],
+    SplashResidualsType,
+]:
+  if save_residuals:
+    raise NotImplementedError("Higher-order AD not supported")
+
+  out, (logsumexp,) = _splash_attention_forward(
+      fwd_mask_info,
+      q,
+      k,
+      v,
+      segment_ids,
+      mask_value=mask_value,
+      is_mqa=is_mqa,
+      block_sizes=block_sizes,
+      save_residuals=True,
+      attn_logits_soft_cap=attn_logits_soft_cap,
+      residual_checkpoint_name=residual_checkpoint_name,
+      mask_function=mask_function,
+  )
+  return out, (
+      q,
+      k,
+      v,
+      segment_ids,
+      out,
+      logsumexp,
+      dq_mask_info,
+      dkv_mask_info,
+  )
+
+
+def _flash_attention_dq_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,
+    logsumexp_ref,
+    do_ref,
+    di_ref,
+    mask_ref,
+    q_sequence_ref,
+    # Outputs
+    dq_scratch_ref,
+    dq_ref,
+    *,
+    mask_value: float,
+    grid_width: int,
+    bq: int,
+    bkv: int,
+    attn_logits_soft_cap: float | None = None,
+    q_layout: QKVLayout,
+    k_layout: QKVLayout,
+    v_layout: QKVLayout,
+    mask_function: MaskFunctionType | None,
+):
+  float32 = jnp.float32
+  HEAD_DIM_MINOR = QKVLayout.HEAD_DIM_MINOR
+
+  h, i, j = pl.program_id(0), pl.program_id(1), pl.program_id(2)
+  @pl.when(j == 0)
+  def init():
+    dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref)
+
+  _, _, should_run, should_not_mask = _next_nonzero(
+      h, i, j, data_next_ref, block_mask_ref, mask_next_ref
+  )
+  @pl.when(should_run)
+  def run():
+    q = q_ref[...] if q_layout == HEAD_DIM_MINOR else q_ref[...].T
+    # We keep k and v possibly transposed, since they are RHS of dots.
+    k = k_ref[...]
+    v = v_ref[...]
+    logsumexp = jnp.expand_dims(logsumexp_ref[0], -1)
+    do = do_ref[...]
+    di = jnp.expand_dims(di_ref[0], -1)
+
+    qk_dims = NT_DIM_NUMBERS if k_layout == HEAD_DIM_MINOR else NN_DIM_NUMBERS
+    qk_uncapped = lax.dot_general(q, k, qk_dims, preferred_element_type=float32)
+
+    qk = _apply_mask_and_soft_cap(
+        qk_uncapped,
+        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=pl.ds(0, bkv),
+        k_offset=j * bkv,
+        bq=bq,
+        mask_function=mask_function,
+    )
+    p = jnp.exp(qk - logsumexp)
+    dp_dims = NT_DIM_NUMBERS if v_layout == HEAD_DIM_MINOR else NN_DIM_NUMBERS
+    dp = lax.dot_general(
+        do.astype(v.dtype), v, dp_dims, preferred_element_type=jnp.float32,
+    )
+    ds = (dp - di) * p
+    if attn_logits_soft_cap is not None:
+      normalized = qk_uncapped / attn_logits_soft_cap
+      d = jnp.tanh(normalized)
+      g = ds * (1 - d)
+      ds = g + g * d
+
+    dq_dims = NN_DIM_NUMBERS if k_layout == HEAD_DIM_MINOR else NT_DIM_NUMBERS
+    dq_scratch_ref[...] += lax.dot_general(
+        ds.astype(k.dtype), k, dq_dims,
+        preferred_element_type=jnp.float32,
+    )
+
+  @pl.when(j == grid_width - 1)
+  def end():
+    dq_ref[...] = dq_scratch_ref[...].astype(dq_ref.dtype)
+    dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref)
+
+
+def _splash_attention_bwd_dq(
+    q,
+    k,
+    v,
+    segment_ids,
+    logsumexp,
+    do,
+    di,
+    *,
+    bq: int,
+    bkv: int,
+    is_mqa: bool,
+    mask_info: mask_info_lib.MaskInfo,
+    mask_value: float,
+    attn_logits_soft_cap: float | None,
+    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:
+    kv_seq_len = k.shape[0]
+    num_kv_heads = 1
+  else:
+    kv_seq_len = k.shape[1]
+    num_kv_heads = k.shape[0]
+
+  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 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 % NUM_LANES:
+    raise ValueError(f"{bkv=} must be a multiple of {NUM_LANES}.")
+
+  # TODO(amagni/sharadmv): when adding block_compute, make sure that is a
+  # multiple of NUM_LANES.
+
+  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[-1]
+  else:
+    grid_width = kv_seq_len // bkv
+
+  grid = (num_q_heads, q_seq_len // bq, grid_width)
+
+  def o_index_map(h, i, *_):
+    return h, i, 0
+
+  o_spec = pl.BlockSpec(o_index_map, (None, bq, head_dim))
+
+  def q_index_map(h, i, *_):
+    return from_head_minor((h, i, 0), q_layout)
+
+  q_spec = pl.BlockSpec(
+      q_index_map, from_head_minor((None, bq, head_dim), q_layout)
+  )
+
+  def k_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_ref, *_):
+    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)
+
+  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(h, i, j, data_next_ref, block_mask_ref, mask_next_ref, *_):
+    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)
+
+  v_spec = pl.BlockSpec(
+      v_index_map,
+      from_head_minor(
+          (bkv, head_dim) if is_mqa else (None, bkv, head_dim), v_layout),
+  )
+
+  def mask_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_ref, *_):
+    _, next_m, *_ = _next_nonzero(
+        h, i, j, data_next_ref, block_mask_ref, mask_next_ref
+    )
+    return next_m, 0, 0
+
+  mask_spec = pl.BlockSpec(mask_index_map, (None, bq, bkv))
+
+  def q_segment_ids_index_map(h, i, j, *_):
+    del h, j  # Unused.
+    return i, 0
+
+  if segment_ids is not None:
+
+    def kv_segment_ids_index_map(
+        h, i, j, data_next_ref, block_mask_ref, mask_next_ref, *_
+    ):
+      next_j, *_ = _next_nonzero(
+          h, i, j, data_next_ref, block_mask_ref, mask_next_ref
+      )
+      return 0, next_j
+
+    q_segment_spec = pl.BlockSpec(q_segment_ids_index_map, (bq, NUM_LANES))
+    kv_segment_spec = 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:
+    q_segment_spec = kv_segment_spec = None
+    q_segment_ids = kv_segment_ids = None
+
+  do_spec = dq_spec = o_spec
+
+  def logsumexp_index_map(h, i, *_):
+    return h, 0, i
+
+  logsumexp = jnp.expand_dims(logsumexp, axis=-2)
+  logsumexp_spec = pl.BlockSpec(logsumexp_index_map, (None, 1, bq))
+  assert logsumexp.ndim == len(logsumexp_spec.block_shape)
+
+  di = jnp.expand_dims(di, axis=-2)
+  di_spec = pl.BlockSpec(logsumexp_index_map, (None, 1, 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)
+
+  assert mask_info.partial_mask_blocks is None or mask_info.q_sequence is None
+
+  if mask_info.q_sequence is not None:
+    q_sequence = jax.lax.broadcast_in_dim(
+        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)
+
+  out_shapes = [
+      jax.ShapeDtypeStruct((bq, head_dim), jnp.float32),
+      jax.ShapeDtypeStruct(q.shape, q.dtype),
+  ]
+  out_specs = [
+      pl.BlockSpec(lambda *_: (0, 0), (bq, head_dim)),
+      dq_spec,
+  ]
+
+  kernel = functools.partial(
+      _flash_attention_dq_kernel,
+      grid_width=grid_width,
+      mask_value=mask_value,
+      bq=bq,
+      bkv=bkv,
+      attn_logits_soft_cap=attn_logits_soft_cap,
+      q_layout=q_layout,
+      k_layout=k_layout,
+      v_layout=v_layout,
+      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_dq=bq,
+      block_kv_dq=bkv,
+      q_layout=q_layout,
+      k_layout=k_layout,
+      v_layout=v_layout,
+      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)
+  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, "dq")
+  with jax.named_scope(kernel_name):
+    _, dq = 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,
+    )
+  return dq
+
+
+def _flash_attention_dkv_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,
+    logsumexp_ref,
+    do_ref,
+    di_ref,
+    mask_ref,
+    q_sequence_ref,
+    # Outputs
+    dq_scratch_ref,
+    dk_scratch_ref,
+    dv_scratch_ref,
+    dq_ref,
+    dk_ref,
+    dv_ref,
+    *,
+    num_q_heads: int,
+    num_kv_heads: int,
+    mask_value: float,
+    grid_width: int,
+    bq: int,
+    bkv_compute: int,
+    is_mqa: bool,
+    attn_logits_soft_cap: float | None,
+    q_layout: QKVLayout,
+    k_layout: QKVLayout,
+    v_layout: QKVLayout,
+    bkv: int,
+    mask_function: MaskFunctionType | None,
+):
+  HEAD_DIM_MINOR = QKVLayout.HEAD_DIM_MINOR
+  kv_index, q_head_index, q_index = (
+      pl.program_id(0),
+      pl.program_id(1),
+      pl.program_id(2),
+  )
+  should_initialize = q_index == 0
+
+  q_heads_per_kv_heads = None
+  q_head_index_per_kv_head = None
+
+  # Consider this situation:
+  # Q_heads:   0, 1, 2, 3, 4, 5, 6, 7
+  # KV_heads:  0,    1,    2,    3
+  # The gradient scratch buffers should be initialized for Q_heads 0, 2, 4, 6
+  # (first Q_heads to 'see' a new KV_head).
+  # The gradient output buffers should be written for Q_heads 1, 3, 5, 7 (last
+  # Q_heads to 'see' the current KV_head).
+
+  # We can use the same logic for both MQA and GA (grouped attention).
+  # But for MQA there is no need for the rem instruction, so we skip it.
+  if is_mqa:
+    should_initialize = jnp.logical_and(should_initialize, q_head_index == 0)
+  elif num_kv_heads < num_q_heads:
+    q_heads_per_kv_heads = num_q_heads // num_kv_heads
+    q_head_index_per_kv_head = lax.rem(q_head_index, q_heads_per_kv_heads)
+    should_initialize = jnp.logical_and(
+        should_initialize, q_head_index_per_kv_head == 0
+    )
+  @pl.when(should_initialize)
+  def init():
+    dk_scratch_ref[...] = jnp.zeros_like(dk_scratch_ref)
+    dv_scratch_ref[...] = jnp.zeros_like(dv_scratch_ref)
+
+  _, _, should_run, should_not_mask = _next_nonzero(
+      q_head_index,
+      q_index,
+      kv_index,
+      data_next_ref,
+      block_mask_ref,
+      mask_next_ref,
+      next_i=True,
+  )
+
+  def body(i, _):
+
+    slice_k = pl.ds(i * bkv_compute, bkv_compute)
+    q = q_ref[...]  # We keep q potentially transposed, since it's always RHS
+    def _load_kv(ref, layout):
+      if layout == HEAD_DIM_MINOR:
+        return pl.load(ref, (slice_k, slice(None)))
+      return pl.load(ref, (slice(None), slice_k)).T
+    k = _load_kv(k_ref, k_layout)
+    v = _load_kv(v_ref, v_layout)
+    logsumexp = pl.load(logsumexp_ref, (pl.ds(1), slice(None)))
+    do = do_ref[...]
+    di = pl.load(di_ref, (pl.ds(1), slice(None)))
+
+    qk_dims = NT_DIM_NUMBERS if q_layout == HEAD_DIM_MINOR else NN_DIM_NUMBERS
+    qk_uncapped = lax.dot_general(
+        k, q, qk_dims, preferred_element_type=jnp.float32
+    )
+
+    qk = _apply_mask_and_soft_cap(
+        qk_uncapped,
+        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=kv_index * bkv + i * bkv_compute,
+        bq=bq,
+        k_in_lanes=False,
+        mask_function=mask_function,
+    )
+    p = jnp.exp(qk - logsumexp)
+    dv = lax.dot(p.astype(do.dtype), do, preferred_element_type=jnp.float32)
+    dv = dv.astype(dv_scratch_ref.dtype) + pl.load(
+        dv_scratch_ref, (slice_k, slice(None))
+    )
+    pl.store(dv_scratch_ref, (slice_k, slice(None)), dv)
+
+    dp = lax.dot_general(
+        v, do, NT_DIM_NUMBERS,
+        preferred_element_type=jnp.float32,
+    )
+    ds = (dp - di) * p
+    if attn_logits_soft_cap is not None:
+      normalized = qk_uncapped / attn_logits_soft_cap
+      d = jnp.tanh(normalized)
+      g = ds * (1 - d)
+      ds = g + g * d
+    dk_dims = NN_DIM_NUMBERS if q_layout == HEAD_DIM_MINOR else NT_DIM_NUMBERS
+    dk = lax.dot_general(
+        ds.astype(do.dtype), q, dk_dims, preferred_element_type=jnp.float32
+    )
+    dk = dk.astype(dk_scratch_ref.dtype) + pl.load(
+        dk_scratch_ref, (slice_k, slice(None))
+    )
+    pl.store(dk_scratch_ref, (slice_k, slice(None)), dk)
+    if dq_scratch_ref is not None or dq_ref is not None:
+      dq = lax.dot_general(
+          ds.T.astype(k.dtype), k, NN_DIM_NUMBERS,
+          preferred_element_type=jnp.float32,
+      )
+      if dq_scratch_ref is not None:
+        # Compute block size != memory block size
+        dq_scratch_ref[...] += dq
+      else:
+        # Compute block size == memory block size
+        assert dq_ref is not None
+        dq_ref[...] = dq.astype(dq_ref.dtype)
+
+  if dq_scratch_ref is not None:
+    dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref)
+  elif dq_scratch_ref is None and dq_ref is not None:
+    dq_ref[...] = jnp.zeros_like(dq_ref)
+
+  @pl.when(should_run)
+  def run():
+    num_iters = (
+        k_ref.shape[0 if k_layout is HEAD_DIM_MINOR else 1] // bkv_compute
+    )
+    lax.fori_loop(0, num_iters, body, None, unroll=True)
+  if dq_scratch_ref is not None:
+    assert dq_ref is not None
+    dq_ref[...] = dq_scratch_ref[...].astype(dq_ref.dtype)
+
+  should_write = q_index == grid_width - 1
+  if is_mqa:
+    should_write = jnp.logical_and(
+        should_write, q_head_index == num_q_heads - 1
+    )
+  elif num_kv_heads < num_q_heads:
+    should_write = jnp.logical_and(
+        should_write, q_head_index_per_kv_head == q_heads_per_kv_heads - 1
+    )
+
+  @pl.when(should_write)
+  def end():
+    dk_ref[...] = dk_scratch_ref[...].astype(dk_ref.dtype)
+    dv_ref[...] = dv_scratch_ref[...].astype(dv_ref.dtype)
+    if dq_scratch_ref is not None:
+      dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref)
+
+    dk_scratch_ref[...] = jnp.zeros_like(dk_scratch_ref)
+    dv_scratch_ref[...] = jnp.zeros_like(dv_scratch_ref)
+
+
+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
+
+
+def _splash_attention_bwd(
+    save_residuals: bool,
+    mask_value: float,
+    is_mqa: bool,
+    block_sizes: BlockSizes,
+    residual_checkpoint_name: str | None,
+    mask_function: MaskFunctionType | None,
+    attn_logits_soft_cap: float | None,
+    res: SplashResidualsType,
+    do: jax.Array,
+) -> tuple[
+    mask_info_lib.MaskInfo | None,  # fwd_mask_info
+    mask_info_lib.MaskInfo | None,  # dq_mask_info
+    mask_info_lib.MaskInfo | None,  # dvk_mask_info
+    jax.Array,  # q
+    jax.Array,  # k
+    jax.Array,  # v
+    SegmentIds | None,  # segmend_ids
+]:
+  del save_residuals, residual_checkpoint_name
+  if not block_sizes.has_backward_blocks:
+    raise ValueError("Need to specify backward blocks.")
+  bq_dq, bkv_dq = block_sizes.block_q_dq, block_sizes.block_kv_dq
+  bq_dkv, bkv_dkv_memory, bkv_dkv_compute = (
+      block_sizes.block_q_dkv,
+      block_sizes.block_kv_dkv,
+      block_sizes.block_kv_dkv_compute,
+  )
+  use_fused_bwd_kernel = block_sizes.use_fused_bwd_kernel
+  (
+      q,
+      k,
+      v,
+      segment_ids,
+      o,
+      logsumexp,
+      dq_mask_info,
+      dkv_mask_info,
+  ) = res
+
+  # di: [num_heads, q_seq_len]
+  di = jnp.einsum("hsd,hsd->hs", o.astype(jnp.float32), do.astype(jnp.float32))  # pytype: disable=attribute-error
+  dq, dk, dv = _splash_attention_bwd_dkv(
+      q,
+      k,
+      v,
+      segment_ids,
+      logsumexp,
+      do,
+      di,
+      bq=bq_dkv,
+      bkv=bkv_dkv_memory,
+      bkv_compute=bkv_dkv_compute,
+      is_mqa=is_mqa,
+      mask_info=dkv_mask_info,
+      mask_value=mask_value,
+      attn_logits_soft_cap=attn_logits_soft_cap,
+      use_fused_bwd_kernel=use_fused_bwd_kernel,
+      q_layout=block_sizes.q_layout,
+      k_layout=block_sizes.k_layout,
+      v_layout=block_sizes.v_layout,
+      mask_function=mask_function,
+  )
+  if not use_fused_bwd_kernel:
+    assert dq is None
+    dq = _splash_attention_bwd_dq(
+        q,
+        k,
+        v,
+        segment_ids,
+        logsumexp,
+        do,
+        di,
+        bq=bq_dq,
+        bkv=bkv_dq,
+        is_mqa=is_mqa,
+        mask_info=dq_mask_info,
+        mask_value=mask_value,
+        attn_logits_soft_cap=attn_logits_soft_cap,
+        q_layout=block_sizes.q_layout,
+        k_layout=block_sizes.k_layout,
+        v_layout=block_sizes.v_layout,
+        mask_function=mask_function,
+    )
+  # Match the signature of the fwd function.
+  assert dq is not None
+  return (
+      None,  # fwd_mask_info
+      None,  # dq_mask_info
+      None,  # dvk_mak_info
+      dq,  # q
+      dk,  # k
+      dv,  # v
+      None,  # segment_ids
+  )
+
+
+_splash_attention_custom.defvjp(_splash_attention_fwd, _splash_attention_bwd)
+
+
+@partial(
+    jax.jit,
+    static_argnames=[
+        "is_mqa",
+        "block_sizes",
+        "save_residuals",
+        "mask_value",
+        "attn_logits_soft_cap",
+        "residual_checkpoint_name",
+        "mask_function",
+    ],
+)
+def _splash_attention(
+    fwd_mask_info: mask_info_lib.MaskInfo,
+    dq_mask_info: mask_info_lib.MaskInfo | None,
+    dkv_mask_info: mask_info_lib.MaskInfo | None,
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    segment_ids: SegmentIds | None = None,
+    *,
+    is_mqa: bool,
+    block_sizes: BlockSizes | None,
+    save_residuals: bool,
+    mask_value: float,
+    attn_logits_soft_cap: float | None,
+    residual_checkpoint_name: str | None,
+    mask_function: MaskFunctionType | None,
+) -> SplashCustomReturnType:
+  return _splash_attention_custom(
+      fwd_mask_info,
+      dq_mask_info,
+      dkv_mask_info,
+      q,
+      k,
+      v,
+      segment_ids,
+      mask_value=mask_value,
+      is_mqa=is_mqa,
+      block_sizes=block_sizes,
+      save_residuals=save_residuals,
+      attn_logits_soft_cap=attn_logits_soft_cap,
+      residual_checkpoint_name=residual_checkpoint_name,
+      mask_function=mask_function,
+  )
+
+
+@jax.tree_util.register_pytree_node_class
+class SplashAttentionKernel:
+
+  def __init__(
+      self,
+      fwd_mask_info: mask_info_lib.MaskInfo,
+      dq_mask_info: mask_info_lib.MaskInfo | None,
+      dkv_mask_info: mask_info_lib.MaskInfo | None,
+      **kwargs,
+  ):
+    self.kwargs = kwargs
+    self.fwd_mask_info = fwd_mask_info
+    self.dq_mask_info = dq_mask_info
+    self.dkv_mask_info = dkv_mask_info
+
+  def __call__(self, *args, **kwargs) -> SplashCustomReturnType:
+    return _splash_attention(
+        self.fwd_mask_info,
+        self.dq_mask_info,
+        self.dkv_mask_info,
+        *args,
+        **kwargs,
+        **self.kwargs,
+    )
+
+  def manual_sharding_spec(self, sharding: jax.sharding.NamedSharding):
+    """Returns a value that can be used as a shard_map partition spec for the kernel."""
+    if self.fwd_mask_info.data_next is not None:
+      block_mask_shape = self.fwd_mask_info.data_next.shape
+      try:
+        shard_shape = sharding.shard_shape(block_mask_shape)
+      except ValueError as exc:
+        raise ValueError(
+            "The sharding must divide the mask blocks evenly between devices"
+        ) from exc
+      if block_mask_shape[-1] != shard_shape[-1]:
+        raise ValueError("Sharding the kv sequence dimension is not supported")
+    spec = sharding.spec
+    assert len(spec) == 2
+    replicated = jax.sharding.PartitionSpec()
+    # Shard q_sequence over the sequence dimension only.
+    q_sequence_spec = jax.sharding.PartitionSpec(spec[1])
+    mask_info_specs = mask_info_lib.MaskInfo(  # pytype: disable=wrong-arg-types
+        data_next=spec if self.fwd_mask_info.data_next is not None else None,
+        mask_next=spec if self.fwd_mask_info.mask_next is not None else None,
+        block_mask=spec if self.fwd_mask_info.block_mask is not None else None,
+        partial_mask_blocks=replicated
+        if self.fwd_mask_info.partial_mask_blocks is not None
+        else None,
+        q_sequence=q_sequence_spec
+        if self.fwd_mask_info.q_sequence is not None
+        else None,
+    )
+    return SplashAttentionKernel(
+        mask_info_specs,
+        mask_info_specs if self.dq_mask_info is not None else None,
+        mask_info_specs if self.dkv_mask_info is not None else None,
+        **self.kwargs,
+    )
+
+  def tree_flatten(self):
+    return (
+        (self.fwd_mask_info, self.dq_mask_info, self.dkv_mask_info),
+        self.kwargs,
+    )
+
+  @classmethod
+  def tree_unflatten(cls, kwargs, values):
+    fwd_mask_info, dq_mask_info, dkv_mask_info = values
+    # NamedTuples are not preserved during pytree serialization.
+    dq_mask_info = (
+        mask_info_lib.MaskInfo(*dq_mask_info)
+        if dq_mask_info is not None
+        else None
+    )
+    dkv_mask_info = (
+        mask_info_lib.MaskInfo(*dkv_mask_info)
+        if dkv_mask_info is not None
+        else None
+    )
+    return SplashAttentionKernel(
+        mask_info_lib.MaskInfo(*fwd_mask_info),
+        dq_mask_info,
+        dkv_mask_info,
+        **kwargs,
+    )
+
+
+def _make_splash_attention(
+    mask: np.ndarray | mask_lib.MultiHeadMask,
+    *,
+    block_sizes: BlockSizes | None = None,
+    is_mqa: bool,
+    save_residuals: bool = False,
+    mask_value: float = DEFAULT_MASK_VALUE,
+    attn_logits_soft_cap: float | None = None,
+    downcast_smem_data: bool = True,
+    head_shards: int,
+    q_seq_shards: int,
+    residual_checkpoint_name: str | None = None,
+):
+  if len(mask.shape) != 3:
+    raise ValueError(f'Unexpected mask shape: {mask.shape}')
+
+  if isinstance(mask, np.ndarray):
+    mask = mask_lib.MultiHeadMask(
+        [mask_lib.NumpyMask(head_mask) for head_mask in mask]
+    )
+
+  if block_sizes is None:
+    block_sizes = BlockSizes.get_default()
+  fwd_mask_info, mask_function_fwd = mask_info_lib.process_mask(
+      mask,
+      (block_sizes.block_q, block_sizes.block_kv),
+      downcast_smem_data=downcast_smem_data,
+      head_shards=head_shards,
+      q_seq_shards=q_seq_shards,
+  )
+
+  fwd_mask_info = tree_util.tree_map(jnp.array, fwd_mask_info)
+
+  dq_mask_info = None
+  dkv_mask_info = None
+  if block_sizes.has_backward_blocks:
+    if block_sizes.use_fused_bwd_kernel:
+      dq_mask_info = None
+    else:
+      bq_dq, bkv_dq = block_sizes.block_q_dq, block_sizes.block_kv_dq
+      dq_mask_info, mask_function_dq = mask_info_lib.process_mask(
+          mask,
+          (bq_dq, bkv_dq),
+          downcast_smem_data=downcast_smem_data,
+          head_shards=head_shards,
+          q_seq_shards=q_seq_shards,
+      )
+      assert (mask_function_fwd is None) == (mask_function_dq is None)
+      dq_mask_info = tree_util.tree_map(jnp.array, dq_mask_info)
+    bq_dkv, bkv_dkv = block_sizes.block_q_dkv, block_sizes.block_kv_dkv
+    dkv_mask_info, mask_function_dkv = mask_info_lib.process_mask_dkv(
+        mask,
+        (bq_dkv, bkv_dkv),
+        downcast_smem_data=downcast_smem_data,
+        head_shards=head_shards,
+        q_seq_shards=q_seq_shards,
+        shrink_grid=not block_sizes.use_fused_bwd_kernel,
+    )
+    assert (mask_function_fwd is None) == (mask_function_dkv is None)
+
+    dkv_mask_info = tree_util.tree_map(jnp.array, dkv_mask_info)
+
+  return SplashAttentionKernel(
+      fwd_mask_info,
+      dq_mask_info,
+      dkv_mask_info,
+      block_sizes=block_sizes,
+      is_mqa=is_mqa,
+      save_residuals=save_residuals,
+      mask_value=mask_value,
+      attn_logits_soft_cap=attn_logits_soft_cap,
+      residual_checkpoint_name=residual_checkpoint_name,
+      mask_function=mask_function_fwd,
+  )
+
+
+make_splash_mha = partial(_make_splash_attention, is_mqa=False)
+make_splash_mqa = partial(_make_splash_attention, is_mqa=True)
+
+make_splash_mha_single_device = partial(
+    make_splash_mha, is_mqa=False, head_shards=1, q_seq_shards=1
+)
+
+make_splash_mqa_single_device = partial(
+    make_splash_mha, is_mqa=True, head_shards=1, q_seq_shards=1
+)
diff --git a/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py b/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py
new file mode 100644
index 000000000..e65d9b073
--- /dev/null
+++ b/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py
@@ -0,0 +1,490 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Mini-mask creation library."""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Any, Callable, Sequence, Tuple
+import numpy as np
+
+# mypy: ignore-errors
+
+class Mask:
+  """A base class for splash attention masks."""
+
+  @property
+  def shape(self) -> Tuple[int, ...]:
+    raise NotImplementedError
+
+  def __getitem__(self, idx) -> np.ndarray:
+    raise NotImplementedError
+
+  def __bool__(self) -> bool:
+    raise NotImplementedError(
+        'Conversion to bool is unsupported. Could be caused by using logical'
+        ' instead of bitwise operations on masks.'
+    )
+
+  def __or__(self, other: 'Mask') -> 'Mask':
+    if self.shape != other.shape:
+      raise ValueError(
+          f'Invalid shape for other: {other.shape}, expected: {self.shape}'
+      )
+    return LogicalOr(self, other)
+
+  def __and__(self, other: 'Mask') -> 'Mask':
+    if self.shape != other.shape:
+      raise ValueError(
+          f'Invalid shape for other: {other.shape}, expected: {self.shape}'
+      )
+    return LogicalAnd(self, other)
+
+
+def make_causal_mask(shape: Tuple[int, int], offset: int = 0) -> np.ndarray:
+  """Makes a causal attention mask.
+
+  Args:
+    shape: Shape of the 2-dim mask: (q_seq_len, kv_seq_len).
+    offset: Offset of q start wrt kv. A positive offset shifts the bottom
+      triangle upward, a negative one shifts it downward. A negative offset
+      makes the first 'offset' rows of the attention matrix all 0s which leads
+      to undefined softmax.
+
+  Returns:
+    The causal mask.
+  """
+  q_seq_len, kv_seq_len = shape
+  q_idx = np.arange(q_seq_len, dtype=np.int32)
+  kv_idx = np.arange(kv_seq_len, dtype=np.int32)
+  return (q_idx[:, None] + offset >= kv_idx[None, :]).astype(np.bool_)
+
+
+def make_local_attention_mask(
+    shape: Tuple[int, int],
+    window_size: Tuple[int | None, int | None],
+    *,
+    offset: int = 0,
+) -> np.ndarray:
+  """Makes a local attention mask."""
+  q_seq_len, kv_seq_len = shape
+  q_idx = np.arange(q_seq_len, dtype=np.int32)
+  kv_idx = np.arange(kv_seq_len, dtype=np.int32)
+  mask = np.ones((q_seq_len, kv_seq_len), dtype=np.bool_)
+  left, right = window_size
+  if left is not None:
+    mask = mask & (q_idx[:, None] - left + offset <= kv_idx[None, :])
+  if right is not None:
+    mask = mask & (q_idx[:, None] + right + offset >= kv_idx[None, :])
+  return mask.astype(np.bool_)
+
+
+def make_random_mask(
+    shape: Tuple[int, int], sparsity: float, seed: int
+) -> np.ndarray:
+  """Makes a random attention mask."""
+  np.random.seed(seed)
+  return np.random.binomial(n=1, p=1.0 - sparsity, size=shape).astype(np.bool_)
+
+
+@dataclasses.dataclass
+class LogicalOr(Mask):
+  left: Mask
+  right: Mask
+
+  def __init__(self, left: Mask, right: Mask):
+    if left.shape != right.shape:
+      raise ValueError('Masks must have the same shape')
+    self.left = left
+    self.right = right
+
+  @property
+  def shape(self) -> Tuple[int, ...]:
+    return self.left.shape
+
+  def __getitem__(self, idx) -> np.ndarray:
+    return self.left[idx] | self.right[idx]
+
+  def __hash__(self):
+    return hash((type(self),) + (self.left, self.right))
+
+
+@dataclasses.dataclass
+class LogicalAnd(Mask):
+  left: Mask
+  right: Mask
+
+  def __init__(self, left: Mask, right: Mask):
+    if left.shape != right.shape:
+      raise ValueError('Masks must have the same shape')
+    self.left = left
+    self.right = right
+
+  @property
+  def shape(self) -> Tuple[int, ...]:
+    return self.left.shape
+
+  def __getitem__(self, idx) -> np.ndarray:
+    return self.left[idx] & self.right[idx]
+
+  def __hash__(self):
+    return hash((type(self),) + (self.left, self.right))
+
+
+@dataclasses.dataclass
+class MultiHeadMask(Mask):
+  """Lazy multihead mask, combines multiple lazy masks one per head."""
+
+  masks: Sequence[Mask]
+
+  def __post_init__(self):
+    if not self.masks:
+      raise ValueError('Unsupported empty tuple of masks')
+
+    shape = self.masks[0].shape
+    for mask in self.masks[1:]:
+      if shape != mask.shape:
+        raise ValueError(
+            f'Unexpected mask shape, got: {mask.shape}, expected: {shape}'
+        )
+
+    if not all(isinstance(mask, Mask) for mask in self.masks):
+      raise ValueError('masks should be of type Mask')
+
+    if any(isinstance(mask, MultiHeadMask) for mask in self.masks):
+      raise ValueError('Nesting MultiHeadMasks is not supported')
+
+  @property
+  def shape(self) -> Tuple[int, ...]:
+    return (len(self.masks),) + self.masks[0].shape
+
+  def __getitem__(self, idx) -> np.ndarray:
+    if len(idx) != 3:
+      raise NotImplementedError(f'Unsupported slice: {idx}')
+
+    head_slice = idx[0]
+    if isinstance(head_slice, int):
+      assert head_slice >= 0 and head_slice <= len(self.masks)
+      return self.masks[head_slice][idx[1:]]
+    else:
+      slice_masks = [mask[idx[1:]] for mask in self.masks[head_slice]]
+      return np.stack(slice_masks)
+
+  def __eq__(self, other: object):
+    if not isinstance(other, type(self)):
+      return NotImplemented
+
+    return self.masks == other.masks
+
+  def __hash__(self):
+    return hash((type(self),) + tuple(hash(mask) for mask in self.masks))
+
+
+class _ComputableMask(Mask):
+  """Superclass for all masks that can be computed inside the kernel using a callable object.
+
+  Attributes:
+    _shape: Shape of the 2-dim mask: (q_seq_len, kv_seq_len).
+    offset: Offset of q start wrt kv. A positive offset shifts the bottom
+      triangle upward, a negative one shifts it downward. A negative offset
+      makes the first 'offset' rows of the attention matrix all 0s which leads
+      to undefined softmax.
+    q_sequence: Indices of Q sequence.
+      q_sequence is reused across __getitem__ calls which is important for
+      compile-time performance.
+    mask_function: Function used by the SplashAttention kernel to compute the
+      mask rather than loading it.
+  """
+
+  _shape: Tuple[int, int]
+  q_sequence: np.ndarray
+  mask_function: Callable[..., Any]
+
+  def __init__(
+      self,
+      shape: Tuple[int, int],
+      mask_function: Callable[..., Any],
+      shard_count: int = 1,
+  ):
+    self._shape = shape
+    self.mask_function = mask_function
+    q_seq_len = self.shape[0]
+
+    if q_seq_len % (shard_count * shard_count) != 0:
+      raise ValueError(
+          f'Shard count squared ({shard_count * shard_count}) must'
+          f' divide Q seq_len ({self.shape[0]}) evenly.'
+      )
+
+    self.q_sequence = np.arange(q_seq_len, dtype=np.int32)
+
+  @property
+  def shape(self) -> Tuple[int, ...]:
+    return self._shape
+
+  def __getitem__(self, idx) -> np.ndarray:
+    if len(idx) != 2:
+      raise NotImplementedError(f'Unsupported slice: {idx}')
+
+    q_slice, kv_slice = idx
+    if not isinstance(q_slice, slice) or not isinstance(kv_slice, slice):
+      raise NotImplementedError(f'Unsupported slice: {idx}')
+
+    q_slice = _fill_slice(q_slice, self.shape[0])
+    kv_slice = _fill_slice(kv_slice, self.shape[1])
+
+    rows = self.q_sequence[q_slice]
+    cols = np.arange(kv_slice.start, kv_slice.stop)
+
+    return self.mask_function(rows[:, None], cols[None, :])
+
+  def __eq__(self, other: object):
+    raise NotImplementedError()
+
+  def __hash__(self):
+    raise NotImplementedError()
+
+
+class CausalMask(_ComputableMask):
+  """Lazy causal mask, prevents the model from attending to future tokens.
+
+  Attributes:
+    offset: Offset of q start wrt kv. A positive offset shifts the bottom
+      triangle upward, a negative one shifts it downward. A negative offset
+      makes the first 'offset' rows of the attention matrix all 0s which leads
+      to undefined softmax.
+  """
+
+  offset: int
+
+  def __init__(
+      self,
+      shape: Tuple[int, int],
+      offset: int = 0,
+      shard_count: int = 1,
+  ):
+    self.offset = offset
+
+    def causal_mask_function(q_ids, kv_ids):
+      # When evaluating the mask in _process_mask we typically work with numpy
+      # array views.
+      # Avoid the addition when possible to avoid instantiating an actual array.
+      if self.offset == 0:
+        return q_ids >= kv_ids
+      else:
+        return q_ids + self.offset >= kv_ids
+
+    mask_function = causal_mask_function
+
+    super().__init__(
+        shape=shape,
+        mask_function=mask_function,
+        shard_count=shard_count,
+    )
+
+  def __eq__(self, other: object):
+    if not isinstance(other, type(self)):
+      return NotImplemented
+
+    return (
+        self.shape == other.shape
+        and self.offset == other.offset
+        and np.array_equal(self.q_sequence, other.q_sequence)
+    )
+
+  def __hash__(self):
+    return hash((
+        type(self),
+        self.shape,
+        self.offset,
+        self.q_sequence.tobytes() if self.q_sequence is not None else None,
+    ))
+
+
+class LocalMask(Mask):
+  """Lazy local mask, prevents model from attending to tokens outside window.
+
+  Attributes:
+    _shape: Shape of the 2-dim mask: (q_seq_len, kv_seq_len).
+    window_size: Size of the two sides of the local window (None identifes no
+      limit for the given side).
+    offset: Offset of q start wrt kv. A positive offset shifts the bottom
+      triangle upward, a negative one shifts it downward. A negative offset
+      makes the first 'offset' rows of the attention matrix all 0s which leads
+      to undefined softmax.
+    _q_sequence: Important for performance.
+  """
+
+  # TODO(amagni): Transform LocalMask into a _ComputableMask.
+
+  _shape: Tuple[int, int]
+  window_size: Tuple[int | None, int | None]
+  offset: int
+  _q_sequence: np.ndarray | None = None
+
+  def __init__(
+      self,
+      shape: Tuple[int, int],
+      window_size: Tuple[int | None, int | None],
+      offset: int,
+      shard_count: int = 1,
+  ):
+    self._shape = shape
+    self.window_size = window_size
+    self.offset = offset
+
+    if self.shape[0] % (shard_count * shard_count) != 0:
+      raise ValueError(
+          f'Shard count squared ({shard_count * shard_count}) must'
+          f' divide Q seq_len ({self.shape[0]}) evenly.'
+      )
+
+  @property
+  def shape(self) -> Tuple[int, int]:
+    return self._shape
+
+  def __getitem__(self, idx) -> np.ndarray:
+    if len(idx) != 2:
+      raise NotImplementedError(f'Unsupported slice: {idx}')
+    q_slice, kv_slice = idx
+    if not isinstance(q_slice, slice) or not isinstance(kv_slice, slice):
+      raise NotImplementedError(f'Unsupported slice: {idx}')
+
+    q_slice = _fill_slice(q_slice, self.shape[0])
+    kv_slice = _fill_slice(kv_slice, self.shape[1])
+
+    if self._q_sequence is None:
+      rows = np.arange(q_slice.start, q_slice.stop)
+    else:
+      rows = self._q_sequence[q_slice]
+
+    cols = np.arange(kv_slice.start, kv_slice.stop)
+
+    left_size, right_size = self.window_size
+
+    if left_size is None and right_size is None:
+      return np.ones((rows.shape[0], cols.shape[0]), dtype=np.bool_)
+    else:
+      expanded_cols = cols[None, :]
+      if self.offset != 0:
+        expanded_rows = rows[:, None] + self.offset
+      else:
+        expanded_rows = rows[:, None]
+      if left_size is not None and right_size is not None:
+        return (expanded_rows <= expanded_cols + left_size) & (
+            expanded_cols - right_size <= expanded_rows
+        )
+
+      elif left_size is not None and right_size is None:
+        return expanded_rows <= expanded_cols + left_size
+      else:
+        assert left_size is None and right_size is not None
+        return expanded_cols - right_size <= expanded_rows
+
+  def __eq__(self, other: object):
+    if not isinstance(other, type(self)):
+      return NotImplemented
+
+    return (
+        self.shape == other.shape
+        and self.window_size == other.window_size
+        and self.offset == other.offset
+        and (True if self._q_sequence is None else
+             np.array_equal(self._q_sequence, other._q_sequence))
+    )
+
+  def __hash__(self):
+    return hash((
+        type(self),
+        self.shape,
+        self.window_size,
+        self.offset,
+        self._q_sequence.tobytes() if self._q_sequence is not None else None,
+    ))
+
+
+@dataclasses.dataclass
+class NumpyMask(Mask):
+  """A mask backed by a dense numpy array."""
+
+  array: np.ndarray
+
+  def __post_init__(self):
+    if self.array.ndim != 2:
+      raise ValueError('Expected a 2-dim array')
+
+    if self.array.dtype != np.bool_:
+      raise ValueError('Mask must be a boolean array')
+
+  @property
+  def shape(self) -> Tuple[int, ...]:
+    return self.array.shape
+
+  def __getitem__(self, idx) -> np.ndarray:
+    return self.array[idx]
+
+  def __eq__(self, other: object):
+    if not isinstance(other, type(self)):
+      return NotImplemented
+
+    return np.array_equal(self.array, other.array, equal_nan=True)
+
+  def __hash__(self):
+    return hash((type(self), self.array.tobytes()))
+
+
+def _fill_slice(inp_slice: slice, size: int) -> slice:
+  assert inp_slice.step is None or inp_slice.step == 1
+  start = 0 if inp_slice.start is None else inp_slice.start
+  stop = size if inp_slice.stop is None else inp_slice.stop
+  assert start >= 0
+  assert stop <= size
+  return slice(start, stop, None)
+
+
+@dataclasses.dataclass(frozen=True)
+class FullMask(Mask):
+  """Lazy full mask, allows all tokens to attend to all other tokens."""
+
+  # TODO(amagni): Transform FullMask into a _ComputableMask.
+
+  _shape: tuple[int, int]
+
+  def __post_init__(self):
+    if not isinstance(self.shape, tuple):
+      raise ValueError(f'Unsupported shape type: {type(self.shape)}')
+
+  @property
+  def shape(self) -> Tuple[int, ...]:
+    return self._shape
+
+  def __getitem__(self, idx) -> np.ndarray:
+    if len(idx) != 2:
+      raise NotImplementedError(f'Unsupported slice: {idx}')
+    i, j = idx
+    if not isinstance(i, slice) or not isinstance(j, slice):
+      raise NotImplementedError(f'Unsupported slice: {idx}')
+    i = _fill_slice(i, self.shape[0])
+    j = _fill_slice(j, self.shape[1])
+    return np.ones((i.stop - i.start, j.stop - j.start), dtype=np.bool_)
+
+  def __eq__(self, other: object):
+    if not isinstance(other, type(self)):
+      return NotImplemented
+
+    return self.shape == other.shape
+
+  def __hash__(self):
+    return hash((type(self), self.shape))
diff --git a/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py b/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py
new file mode 100644
index 000000000..bc1b99e5a
--- /dev/null
+++ b/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py
@@ -0,0 +1,932 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Mini-mask creation library."""
+from __future__ import annotations
+
+import collections
+import functools
+from typing import Callable, Dict, List, NamedTuple, Set, Tuple
+from jax import util as jax_util
+from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask as mask_lib
+import numpy as np
+
+# mypy: ignore-errors
+
+# Logic for processing NumPy masks for kernels
+class MaskInfo(NamedTuple):
+  """Contains runtime masking information for the Splash attention kernel.
+
+  The arrays data_next, mask_next and block_mask are placed in TPU
+  scalar-memory. This is a scarse resource so the mask creation logic attempts
+  to shrink the data-type of these arrays to the smallest possible one.
+  This can be: np.int32, np.int16 or np.int8.
+
+  For the arrays data_next, mask_next and block_mask the size of the first
+  dimension can be one of the two following values: num_head or
+  num_head_shards.
+  The first dimension has size:
+  * num_head_shards when there is only one unique mask for each head in a shard.
+  In this case the three arrays are broadcasted to all the heads in the shard.
+  * num_heads when there is more than one unique mask for each head in the
+  shard.
+
+  Attributes:
+    data_next: An integer[num_heads_or_shards, num_q_blocks, num_kv_blocks]
+      NumPy array where each entry contains the next `kv` block index to
+      prefetch.
+    mask_next: An integer[num_heads_or_shards, num_q_blocks, num_kv_blocks]
+      NumPy array where each entry contains the next mask block index in
+      `partial_mask_blocks` to prefetch.
+    block_mask: An integer[num_heads_or_shards, num_q_blocks, num_kv_blocks]
+      NumPy array whose entries can be 0, 1 or 2. An entry of 0 indicates that
+      the corresponding block in the full mask was all zeros. An entry of 1
+      indicates that the corresponding block in the full mask contained both
+      zeros and ones. An entry of 2 indicates the corresponding block was
+      entirely ones.
+    partial_mask_blocks: A i32[num_partial_blocks, block_q, block_kv] NumPy
+      array that contains the blocks of the original mask that contained both
+      zeros and ones. The entries in `mask_next` point to indices in the first
+      axis of this array.
+    q_sequence: A i32[q_sequence_length] NumPy array. When using causal masking,
+      this contains the list of indices that correspond to q tokens. For plain
+      causal this is just np.arange(q_sequence_length).
+  """
+
+  data_next: np.ndarray | None
+  mask_next: np.ndarray | None
+  block_mask: np.ndarray | None
+  partial_mask_blocks: np.ndarray | None
+  q_sequence: np.ndarray | None
+
+
+def _downcast_to_small_type(array: np.ndarray) -> np.ndarray:
+  """Downcast numpy array.
+
+  If possible, downcast the data-type of the input array to the smallest numpy
+  type (among np.int16 and np.int8) that fits the content of the array.
+
+  Args:
+    array: the array to downcast
+
+  Returns:
+    The downcasted array.
+
+  Raises:
+    ValueError: if the input array is not np.int32 or if its elements are not
+    all positive.
+  """
+
+  if array.dtype != np.int32:
+    raise ValueError('Expected int32 input.')
+
+  if not np.all(array >= 0):
+    raise ValueError('Expected non-negative array.')
+
+  if array.size == 0:
+    return array
+
+  max_value = np.max(array)
+
+  if max_value <= np.iinfo(np.int8).max:
+    return array.astype(np.int8)
+  elif max_value <= np.iinfo(np.int16).max:
+    return array.astype(np.int16)
+  else:
+    return array.astype(np.int32)
+
+
+def _check_mask(mask: mask_lib.Mask) -> None:
+  """Check that the given mask is valid.
+
+  A row of all zeros along the kv dimension would result in a division by zero
+  when computing the softmax. This function is meant to protect against that
+  case.
+
+  Args:
+    mask: the mask to check.
+
+  Raises:
+    ValueError: the mask is invalid.
+  """
+
+  assert len(mask.shape) == 2
+
+  exception_message = (
+      'Some rows of the mask (along the kv dimension) are all zeros.\nThis is'
+      ' would result in a division by zero when computing the attention'
+      ' softmax.'
+  )
+
+  is_row_non_zero = np.zeros(mask.shape[0], dtype=np.bool_)
+  for col in range(mask.shape[1]):
+    # Mask only supports slice indices.
+    is_row_non_zero = np.logical_or(
+        is_row_non_zero,
+        mask[(slice(0, mask.shape[0]), slice(col, col + 1))][:, 0],
+    )
+  if not is_row_non_zero.all():
+    raise ValueError(exception_message)
+
+
+class _HashableNDArray:
+  """Helper to make a numpy array hashable: can be added associative containers.
+
+  Attributes:
+    array: The underlying numpy array.
+  """
+  array: np.ndarray
+
+  def __init__(self, array: np.ndarray):
+    self.array = array
+
+  def __hash__(self):
+    return hash(self.array.tobytes())
+
+  def __eq__(self, other: object) -> bool:
+    if not isinstance(other, _HashableNDArray):
+      return NotImplemented
+    return np.array_equal(self.array, other.array, equal_nan=True)
+
+
+def _get_mask_info_for_shard(
+    output_shape: Tuple[int, int, int],
+    has_mask_next: bool,
+    mask: mask_lib.MultiHeadMask,
+    block_shape: Tuple[int, int],
+    coords_to_partial_mask_block_index: Dict[Tuple[int, int, int], int],
+    masks_per_head_shard: int,
+    head_start: int,
+    num_heads: int,
+    q_seq_start: int,
+    q_seq_shard_size: int,
+    blocked_q_seq_start: int,
+    is_dkv: bool,
+) -> Tuple[np.ndarray, np.ndarray | None]:
+  """Process a slice of the mask to compute data_next and mask_next.
+
+  Args:
+    output_shape: The shape of the data_next and mask_next to return
+    has_mask_next: Whether mask_next should be constructed. If False None is
+      returned for mask_next.
+    mask: The full mask to be sliced according to the head and sequence ranges
+    block_shape: Shape of the Pallas grid block.
+    coords_to_partial_mask_block_index: Mapping between the pallas launch grid
+      coordinates and the index of the corresponding block in partial mask block
+      list.
+    masks_per_head_shard: Number of masks per head shards
+    head_start: First head of the current shard.
+    num_heads: Number of heads in the shard.
+    q_seq_start: Start index along the Q sequence for the current shard (in
+      number of tokens).
+    q_seq_shard_size: Number of tokens along the Q sequence for the current
+      shard.
+    blocked_q_seq_start: Start index along the Q sequence for the current shard
+      (in number of grid blocks)
+    is_dkv: True if we are processing the dKV mask
+
+  Returns:
+    Slice of data_next and mask_next (if required) that correspond to the
+    current mask slice.
+  """
+  _, _, kv_seq_len = mask.shape
+
+  q_block_size, kv_block_size = block_shape
+  q_block_count, q_mod = divmod(q_seq_shard_size, q_block_size)
+  kv_block_count, kv_mod = divmod(kv_seq_len, kv_block_size)
+
+  assert q_mod == 0
+  assert kv_mod == 0
+
+  blocked_shape = (
+      (kv_block_count, num_heads, q_block_count)
+      if is_dkv
+      else (num_heads, q_block_count, kv_block_count)
+  )
+
+  data_coords = []
+  mask_coords = []
+  for idx in np.ndindex(blocked_shape):
+    if is_dkv:
+      kv_index, h_index, q_index = idx
+    else:
+      h_index, q_index, kv_index = idx
+
+    h_index = h_index if masks_per_head_shard == 1 else head_start + h_index
+
+    chunk = mask[(
+        h_index,
+        slice(
+            q_seq_start + q_index * q_block_size,
+            q_seq_start + (q_index + 1) * q_block_size,
+        ),
+        slice(kv_index * kv_block_size, (kv_index + 1) * kv_block_size),
+    )]
+    if chunk.any():
+      data_coords.append(idx)
+      if not chunk.all():
+        mask_coords.append(idx)
+
+  # Initialize the output arrays.
+  mask_next = None
+  if has_mask_next:
+    mask_next = np.zeros(output_shape, dtype=np.int32)
+  data_next = np.zeros(output_shape, dtype=np.int32)
+
+  data_coords_iter = iter(data_coords)
+  first_j = coord_j = next(data_coords_iter)
+  if mask_next is not None and mask_coords:
+    mask_coords_iter = iter(mask_coords)
+    first_m = coord_m = next(mask_coords_iter)
+  else:
+    first_m, coord_m, mask_coords_iter = None, None, None
+
+  for idx in np.ndindex(blocked_shape):
+    if is_dkv:
+      kv_index, h_index, q_index = idx
+      chunk_idx: tuple[int, ...] = (h_index, q_index, kv_index)
+      data_dim = 2
+    else:
+      chunk_idx = idx
+      data_dim = 2
+
+    is_next = idx > coord_j
+
+    if is_next:
+      try:
+        coord_j = next(data_coords_iter)
+      except StopIteration:
+        coord_j = first_j
+    data_next[chunk_idx] = coord_j[data_dim]
+
+    if mask_next is not None and mask_coords:
+      assert coord_m is not None
+      is_next_mask = idx > coord_m
+      if is_next_mask:
+        try:
+          coord_m = next(mask_coords_iter)  # type: ignore
+        except StopIteration:
+          coord_m = first_m
+
+      if is_dkv:
+        assert coord_m is not None
+        coord_m_global = (
+            coord_m[1] + head_start,
+            coord_m[2] + blocked_q_seq_start,
+            coord_m[0],
+        )
+      else:
+        assert coord_m is not None
+        coord_m_global = (
+            coord_m[0] + head_start,
+            coord_m[1] + blocked_q_seq_start,
+            coord_m[2],
+        )
+
+      mask_next[chunk_idx] = coords_to_partial_mask_block_index[coord_m_global]
+
+  return data_next, mask_next
+
+
+# When used in a transformer network with multiple layers, the SplashAttention
+# kernel is created serveral times with the same mask. Cache MaskInfo to avoid
+# blowing up compile times. Ideally the size of the cache should be determined
+# by the client.
+@functools.lru_cache(maxsize=12)
+def _process_mask(
+    mask: mask_lib.MultiHeadMask,  # [num_heads, q_seq_len, kv_seq_len]
+    block_shape: Tuple[int, int],
+    is_dkv: bool,
+    *,
+    downcast_smem_data: bool = True,
+    head_shards: int = 1,
+    q_seq_shards: int = 1,
+    shrink_grid: bool = True,
+) -> tuple[MaskInfo, jax_util.HashableFunction | None]:
+  """Transform a dense mask into a sparse representation.
+
+  The number of head and Q sequence shards are needed to create a MaskInfo
+  object that is partitionable (with shmap or PartIR) along these two dimension.
+  In particular for dKV MaskInfo, for each shard the indices of in the data_next
+  array are relative to the current shard.
+  The fwd and dQ MaskInfo objects do not change when sharding along the head or
+  Q dimensions, they would be different if we were to shard along the KV
+  dimension, but the kernel does not support that.
+
+  Args:
+    mask: Dense mask to process.
+    block_shape: Shape of the Pallas grid block.
+    is_dkv: True if we are processing the dKV mask
+    downcast_smem_data: If True, downcast the scalar-memory data of MaskInfo to
+      a data type smaller than np.int32 (if possible).
+    head_shards: Number of head shards of the mesh in which the kernel is
+      launched.
+    q_seq_shards: Number of Q sequence shards of the mesh in which the kernel is
+      launched.
+    shrink_grid: Whether or not we should apply the grid shrinking optimization.
+
+  Returns:
+    `MaskInfo`, a sparse representation of the dense mask.
+    `MaskCallable`: a callable that, given in input Q and KV indices, returns
+      the value of the mask at those coordinates.
+
+  Raises:
+    ValueError: if the input mask is invalid or the block sizes are not
+    compatible with the mask sizes.
+  """
+
+  if len(mask.shape) != 3:
+    raise ValueError(f'Expected a 3-dim mask, instead got: {mask.shape=}')
+
+  head_count, q_seq_len, kv_seq_len = mask.shape
+  q_block_size, kv_block_size = block_shape
+  q_blocks_count, q_mod = divmod(q_seq_len, q_block_size)
+  kv_blocks_count, kv_mod = divmod(kv_seq_len, kv_block_size)
+
+  if q_mod != 0:
+    raise ValueError(f'{q_block_size=} should divide {q_seq_len=}.')
+  if kv_mod != 0:
+    raise ValueError(f'{kv_block_size=} should divide {kv_seq_len=}.')
+
+  q_seq_len_per_shard, mod = divmod(q_seq_len, q_seq_shards)
+  if mod != 0:
+    raise ValueError(f'{q_seq_shards=} should divide {q_seq_len=}.')
+
+  q_blocks_per_shard, mod = divmod(q_seq_len_per_shard, q_block_size)
+  if mod != 0:
+    raise ValueError(f'{q_block_size=} should divide {q_seq_len_per_shard=}.')
+
+  heads_per_shard, mod = divmod(head_count, head_shards)
+  if mod != 0:
+    raise ValueError(f'{head_shards=} should divide {head_count=}.')
+
+  first_mask_size = mask.masks[0].shape
+  for h in range(head_count):
+    if mask.masks[h].shape != first_mask_size:
+      raise ValueError(
+          f'First head mask has shape {first_mask_size}, but head mask {h} has'
+          f' shape {mask.masks[h].shape}. All head masks must have the same'
+          ' shape.'
+      )
+
+  # Uniquify the masks.
+  # Create a collection of the unique head masks in the input multi-head mask.
+  # This avoids processing the same mask multiple times and it enables
+  # deduplicating the partial mask blocks information.
+  # Assign a unique ID to every unique mask.
+
+  def assign_unique_ids(objects):
+    id_map = collections.defaultdict(lambda: len(id_map))
+    return {obj: id_map[obj] for obj in objects}
+
+  unique_masks_dict: Dict[mask_lib.Mask, int] = assign_unique_ids(
+      head_mask for head_mask in mask.masks
+  )
+
+  # Build a mapping of heads to unique masks and masks to unique masks.
+
+  head_to_mask_id: List[int] = [0] * head_count
+  head_shard_to_mask_ids: List[Set[int]] = [set() for _ in range(head_shards)]
+  mask_id_to_heads: List[List[int]] = [
+      [] for _ in range(len(unique_masks_dict))
+  ]
+  mask_id_to_head_shards: List[Set[int]] = [
+      set() for _ in range(len(unique_masks_dict))
+  ]
+
+  for head in range(head_count):
+    mask_id = unique_masks_dict[mask.masks[head]]
+    head_to_mask_id[head] = mask_id
+    head_shard = head // heads_per_shard
+    head_shard_to_mask_ids[head_shard].add(mask_id)
+    mask_id_to_heads[mask_id].append(head)
+    mask_id_to_head_shards[mask_id].add(head_shard)
+
+  # If we have at most one unique mask per each head shard, then we can brodcast
+  # the mask to all the heads in the shard. This is the common case.
+  # If we have more than one mask in each head shard, then the optimization
+  # cannot kick in and we use one mask for each head.
+  # TODO(sharadmv,amagni): In the future we could think of a dynamic mapping of
+  # heads to masks. This will likely introduce an additional field in the
+  # MaskInfo class and runtime overhead to perform an indirect lookup. Since
+  # having multiple masks per head-shard is not a common case we leave this for
+  # future work.
+  max_masks_per_head_shard = max([len(x) for x in head_shard_to_mask_ids])
+  masks_per_head_shard = 1 if max_masks_per_head_shard == 1 else heads_per_shard
+
+  unique_masks = [
+      pair[0] for pair in sorted(unique_masks_dict.items(), key=lambda x: x[1])
+  ]
+
+  # TODO(amagni): checking the validity of the masks is slow for large masks.
+  # Disable it for now, reevalute in the future.
+
+  partial_mask_block_ids: Dict[_HashableNDArray, int] = collections.defaultdict(
+      lambda: len(partial_mask_block_ids)
+  )
+  block_id_to_block_coords: Dict[int, List[Tuple[int, ...]]] = (
+      collections.defaultdict(list)
+  )
+
+  # Shape of the block_mask, mask_next and data_next fields of MaskInfo.
+  block_mask_shape = (
+      head_shards if masks_per_head_shard == 1 else head_count,
+      q_blocks_count,
+      kv_blocks_count,
+  )
+
+  block_mask = np.zeros(block_mask_shape, dtype=np.int32)
+
+  def set_block_mask(mask_id: int, q_index: int, kv_index: int, value: int):
+    if masks_per_head_shard == 1:
+      for shard_index in mask_id_to_head_shards[mask_id]:
+        block_mask[shard_index, q_index, kv_index] = value
+    else:
+      for head_index in mask_id_to_heads[mask_id]:
+        block_mask[head_index, q_index, kv_index] = value
+
+  q_sequence = None
+  mask_function = None
+  if len(unique_masks) == 1:
+    unique_mask = unique_masks[0]
+    # The mask object either define q_sequence and mask_function or none of
+    # them.
+    assert hasattr(unique_mask, 'q_sequence') == hasattr(
+        unique_mask, 'mask_function'
+    )
+
+    # If the mask object defines a q_sequence and a mask_function, then make use
+    # of these in the kernel rather. This is preferable over loading the mask
+    # from memory. When using a mask_function, then mask_next and
+    # partial_mask_blocks are left undefined and not used in the kernel.
+    if hasattr(unique_mask, 'q_sequence') and hasattr(
+        unique_mask, 'mask_function'
+    ):
+      q_sequence = unique_mask.q_sequence
+      mask_function = unique_mask.mask_function
+
+  # Identify the partial mask blocks and the value of the block mask for each
+  # block.
+  # The partial mask blocks are "global", meaning it is a collection of such
+  # blocks across all the heads. Partial mask blocks are uniquified.
+  # When partitioning, all partial mask blocks are replicated across shards.
+  # As a future extension we could assign to each head shard only the relevant
+  # partial blocks. This would reduce the amount of data each shard needs
+  # from: num_unique_partial_block_blocks to
+  # max(unique_partial_blocks_per_shard)
+  for mask_id, unique_mask in enumerate(unique_masks):
+    for coords in np.ndindex((q_blocks_count, kv_blocks_count)):
+      (q_index, kv_index) = coords
+      chunk = unique_mask[(
+          slice(q_index * q_block_size, (q_index + 1) * q_block_size),
+          slice(kv_index * kv_block_size, (kv_index + 1) * kv_block_size),
+      )]
+      has_nonzero = chunk.any()
+      if has_nonzero:
+        all_nonzero = chunk.all()
+        if not all_nonzero:
+          set_block_mask(mask_id, q_index, kv_index, 1)
+          partial_mask_block_id = partial_mask_block_ids[
+              _HashableNDArray(chunk)
+          ]
+          for head_index in mask_id_to_heads[mask_id]:
+            block_id_to_block_coords[partial_mask_block_id].append(
+                (head_index,) + coords
+            )
+        else:
+          set_block_mask(mask_id, q_index, kv_index, 2)
+
+  unique_partial_mask_blocks = [
+      pair[0]
+      for pair in sorted(partial_mask_block_ids.items(), key=lambda x: x[1])
+  ]
+
+  # For each position in the in q, kv grid that contains a partial mask block
+  # record the index of its partial mask block.
+  coords_to_partial_mask_block_index = {}
+  for partial_mask_block_id, coords in block_id_to_block_coords.items():
+    for coo in coords:
+      coords_to_partial_mask_block_index[coo] = partial_mask_block_id
+
+  partial_mask_blocks = None
+  has_mask_next = False
+  if len(unique_partial_mask_blocks) == 1:
+    partial_mask_blocks = [x.array for x in unique_partial_mask_blocks]
+    partial_mask_blocks = partial_mask_blocks[0][None].astype(np.int32)
+    has_mask_next = True
+  elif len(unique_partial_mask_blocks) > 1:
+    partial_mask_blocks = [x.array for x in unique_partial_mask_blocks]
+    partial_mask_blocks = np.stack(partial_mask_blocks, axis=0).astype(np.int32)
+    has_mask_next = True
+  if is_dkv and partial_mask_blocks is not None:
+    partial_mask_blocks = np.swapaxes(partial_mask_blocks, -1, -2)
+
+  all_head_shards_identical = all(
+      head_shard_to_mask_ids[0] == x and len(x) == 1
+      for x in head_shard_to_mask_ids
+  )
+  # When all the head shards are identical, we can process only one and
+  # broadcast the result across all the head shards, this avoids redundant work.
+  shards_to_process = 1 if all_head_shards_identical else head_shards
+
+  # Iterate over the shards.
+  # Work on a fraction of the mask at the time to compute the mask and the data
+  # indices. This is needed to compute the correct data indices, which are
+  # relative to the current slice of the mask.
+
+  q_sequence_axis = 1
+  head_axis = 0
+
+  # Collect mask_info shards along the head dimension, concatentate (or
+  # broadcast) them after the loop.
+  data_next_per_head_list, mask_next_per_head_list = [], []
+  for head_shard in range(shards_to_process):
+    # Collect mask_info shards along the q_sequence dimension, concatenate them
+    # after the loop.
+    data_next_sequence_slices, mask_next_sequence_slices = [], []
+    for q_seq_len_shard in range(q_seq_shards):
+      head_start = head_shard * heads_per_shard
+      q_seq_len_shard_size = q_blocks_per_shard * q_block_size
+      q_seq_len_start = q_seq_len_shard * q_seq_len_shard_size
+
+      blocked_q_seq_len_start = q_seq_len_shard * q_blocks_per_shard
+      blocked_q_seq_len_slice = slice(
+          blocked_q_seq_len_start,
+          (q_seq_len_shard + 1) * q_blocks_per_shard,
+      )
+
+      if masks_per_head_shard == 1:
+        unique_mask = unique_masks[head_to_mask_id[head_start]]
+        unique_mask = mask_lib.MultiHeadMask((unique_mask,))
+        current_mask = unique_mask
+        mask_head_slice = slice(head_shard, head_shard + 1)
+      else:
+        current_mask = mask
+        mask_head_slice = slice(head_start, (head_shard + 1) * heads_per_shard)
+
+      # The current iteration of the loop processes a slice of the mask info
+      # tensors of this shape:
+      mask_info_slice_shape = (
+          mask_head_slice.stop - mask_head_slice.start,
+          blocked_q_seq_len_slice.stop - blocked_q_seq_len_slice.start,
+          kv_blocks_count,
+      )
+      # Generate data_next and mask_next for the current slice.
+      data_next_slice, mask_next_slice = _get_mask_info_for_shard(
+          output_shape=mask_info_slice_shape,
+          has_mask_next=has_mask_next,
+          mask=current_mask,
+          block_shape=block_shape,
+          coords_to_partial_mask_block_index=coords_to_partial_mask_block_index,
+          head_start=head_start,
+          masks_per_head_shard=masks_per_head_shard,
+          num_heads=1 if masks_per_head_shard == 1 else heads_per_shard,
+          q_seq_start=q_seq_len_start,
+          q_seq_shard_size=q_seq_len_shard_size,
+          blocked_q_seq_start=blocked_q_seq_len_start,
+          is_dkv=is_dkv,
+      )
+      data_next_sequence_slices.append(data_next_slice)
+      mask_next_sequence_slices.append(mask_next_slice)
+
+    # Concatenate the sequence shards.
+    data_next_per_head = np.concatenate(
+        data_next_sequence_slices, axis=q_sequence_axis
+    )
+    data_next_per_head_list.append(data_next_per_head)
+    if has_mask_next:
+      mask_next_per_head = np.concatenate(
+          mask_next_sequence_slices, axis=q_sequence_axis
+      )
+      mask_next_per_head_list.append(mask_next_per_head)
+
+  # Concatenate (or broadcast) the head shards.
+  mask_next = None
+  if all_head_shards_identical:
+    assert len(data_next_per_head_list) == 1
+    data_next_shard = data_next_per_head_list[0]
+    assert data_next_shard.shape == (1, q_blocks_count, kv_blocks_count)
+    data_next = np.broadcast_to(
+        data_next_shard,
+        (head_shards, q_blocks_count, kv_blocks_count),
+    )
+    if has_mask_next:
+      assert len(mask_next_per_head_list) == 1
+      mask_next_shard = mask_next_per_head_list[0]
+      assert mask_next_shard.shape == (1, q_blocks_count, kv_blocks_count)
+      mask_next = np.broadcast_to(
+          mask_next_shard,
+          (head_shards, q_blocks_count, kv_blocks_count),
+      )
+  else:
+    data_next = np.concatenate(data_next_per_head_list, axis=head_axis)
+    if has_mask_next:
+      mask_next = np.concatenate(mask_next_per_head_list, axis=head_axis)
+
+  # Shrink the width of the mask info when possible.
+  if (
+      shrink_grid
+      and block_mask.shape[0] == head_shards
+      and len(unique_masks) == 1
+  ):
+    rows_per_q_shard = block_mask.shape[1] // q_seq_shards
+    block_mask_shards = []
+    data_next_shards = []
+    mask_next_shards = []
+
+    # Slice the mask info arrays along the Q dimension. This simulates Q
+    # sharding and enabled shrinking each Q shard of the MaskInfo arrays
+    # independently.
+    for q_seq_len_shard in range(q_seq_shards):
+      rows = slice(
+          q_seq_len_shard * rows_per_q_shard,
+          (q_seq_len_shard + 1) * rows_per_q_shard,
+      )
+      current_block_mask = block_mask[:, rows, :]
+      current_data_next = data_next[:, rows, :]
+      current_mask_next = (
+          mask_next[:, rows, :] if mask_next is not None else None
+      )
+
+      shrink_function = _shrink_mask_info_dkv if is_dkv else _shrink_mask_info
+
+      current_block_mask, current_data_next, current_mask_next = (
+          shrink_function(
+              block_mask=current_block_mask,
+              data_next=current_data_next,
+              mask_next=current_mask_next,
+              head_shards=head_shards,
+          )
+      )
+
+      assert current_block_mask.size > 0
+      assert current_data_next.size > 0
+      assert current_mask_next is None or current_mask_next.size > 0
+      assert current_block_mask.shape == current_data_next.shape
+      assert (
+          current_mask_next is None
+          or current_block_mask.shape == current_mask_next.shape
+      )
+
+      block_mask_shards.append(current_block_mask)
+      data_next_shards.append(current_data_next)
+      mask_next_shards.append(current_mask_next)
+
+    if q_seq_shards == 1:
+      block_mask = block_mask_shards[0]
+      data_next = data_next_shards[0]
+      mask_next = mask_next_shards[0]
+    else:
+      # Since shrinking happens independently for each shard along Q we might
+      # end up with uneven mask info shards. Insert padding where necessary to
+      # maintain the SPMD paradigm.
+      padding_axis = 1 if is_dkv else 2
+
+      max_size = max([x.shape[padding_axis] for x in block_mask_shards])
+      padded_block_mask_shards = []
+      padded_data_next_shards = []
+      padded_mask_next_shards = []
+      assert (
+          len(block_mask_shards)
+          == len(data_next_shards)
+          == len(mask_next_shards)
+      )
+      for (
+          current_block_mask,
+          current_data_next,
+          current_mask_next,
+      ) in zip(
+          block_mask_shards, data_next_shards, mask_next_shards
+      ):
+        # For dKV shrinking happens along axis Q (the rows of MaskInfo), for
+        # fwd and dQ shrinking happens along axis KV (the columns of MaskInfo).
+        if is_dkv:
+          pad_width = (
+              (0, 0),
+              (0, max_size - current_block_mask.shape[padding_axis]),
+              (0, 0),
+          )
+        else:
+          pad_width = (
+              (0, 0),
+              (0, 0),
+              (0, max_size - current_block_mask.shape[padding_axis]),
+          )
+
+        padded_block_mask_shards.append(
+            np.pad(current_block_mask, pad_width=pad_width, constant_values=0)
+        )
+        # By padding data_next and mask_next with 'edge' policy we avoid
+        # fetching new blocks thanks to 'immediate revisiting'.
+        padded_data_next_shards.append(
+            np.pad(current_data_next, pad_width=pad_width, mode='edge')
+        )
+        if current_mask_next is not None:
+          padded_mask_next_shards.append(
+              np.pad(current_mask_next, pad_width=pad_width, mode='edge')
+          )
+
+      block_mask = np.concatenate(padded_block_mask_shards, axis=1)
+      data_next = np.concatenate(padded_data_next_shards, axis=1)
+      mask_next = np.concatenate(padded_mask_next_shards, axis=1)
+
+  if downcast_smem_data:
+    data_next = _downcast_to_small_type(data_next)
+    block_mask = _downcast_to_small_type(block_mask)
+    if mask_next is not None:
+      mask_next = _downcast_to_small_type(mask_next)
+
+  assert (mask_function is not None) == (q_sequence is not None)
+  # When the mask can be computed inside the kernel with a mask_function,
+  # there is no need to load it from memory. So mask_next and
+  # partial_mask_blocks are unused.
+  return (
+      MaskInfo(
+          data_next=data_next,
+          mask_next=mask_next if mask_function is None else None,
+          block_mask=block_mask,
+          partial_mask_blocks=partial_mask_blocks
+          if mask_function is None
+          else None,
+          q_sequence=q_sequence,
+      ),
+      mask_function,
+  )
+
+
+def _shrink_mask_info(
+    *,
+    block_mask: np.ndarray,
+    data_next: np.ndarray,
+    mask_next: np.ndarray,
+    head_shards: int,
+):
+  assert block_mask.ndim == 3
+  assert data_next.ndim == 3
+  assert mask_next is None or mask_next.ndim == 3
+
+  head_block_mask = block_mask[0]
+  non_zero_rows, non_zero_cols = np.nonzero(head_block_mask)
+  # Group non-zero columns based on which row they belong to.
+  unique_row_indices = np.unique(non_zero_rows, return_index=True)[1]
+  grouped_non_zero_cols = np.split(non_zero_cols, unique_row_indices)[1:]
+
+  assert all(len(x) != 0 for x in grouped_non_zero_cols)
+
+  # Pad each row in the non-zero indices to match the width of the longest
+  # row. This avoids having jagged rows.
+  max_non_zero_cols = max([len(x) for x in grouped_non_zero_cols])
+  padded_non_zero_cols = []
+  padding = -1
+  for row in grouped_non_zero_cols:
+    padded_non_zero_cols.append(
+        np.pad(
+            row,
+            pad_width=(0, max_non_zero_cols - row.shape[0]),
+            constant_values=padding,
+        )
+    )
+
+  padded_non_zero_cols = np.stack(padded_non_zero_cols, axis=0)
+
+  assert padded_non_zero_cols.shape[0] == block_mask.shape[1], (
+      padded_non_zero_cols.shape,
+      block_mask.shape,
+  )
+
+  # For each row of array, select the columns indices in padded_non_zero_cols,
+  # ignore padding.
+  def select_cols(array):
+    assert array.ndim == 2
+    assert padded_non_zero_cols.ndim == 2
+    assert array.shape[0] == padded_non_zero_cols.shape[0]
+    assert array.shape[1] >= padded_non_zero_cols.shape[1]
+    selected_rows = []
+    for row in range(array.shape[0]):
+      col = padded_non_zero_cols[row]
+      selected = array[row][col]
+      selected = np.where(col != padding, selected, 0)
+      selected_rows.append(selected)
+
+    return np.stack(selected_rows, axis=0)
+
+  return _slice_mask_info(
+      block_mask=block_mask,
+      data_next=data_next,
+      mask_next=mask_next,
+      head_shards=head_shards,
+      slice_function=select_cols,
+  )
+
+
+def _shrink_mask_info_dkv(
+    *,
+    block_mask: np.ndarray,
+    data_next: np.ndarray,
+    mask_next: np.ndarray,
+    head_shards: int,
+):
+  assert block_mask.ndim == 3
+  assert data_next.ndim == 3
+  assert mask_next is None or mask_next.ndim == 3
+
+  head_block_mask = block_mask[0]
+  grouped_non_zero_rows = []
+  # Group non-zero rows based on which column they belong to.
+  for col_index in range(head_block_mask.shape[1]):
+    col = head_block_mask[:, col_index]
+    non_zero_rows = np.nonzero(col)[0]
+    grouped_non_zero_rows.append(non_zero_rows)
+
+  # Pad each col in the non-zero indices to match the height of the longest
+  # col. This avoids having jagged cols.
+  max_non_zero_rows = max([len(x) for x in grouped_non_zero_rows])
+  padded_non_zero_rows = []
+  padding = -1
+  for col in grouped_non_zero_rows:
+    padded_non_zero_rows.append(
+        np.pad(
+            col,
+            pad_width=(max_non_zero_rows - col.shape[0], 0),
+            constant_values=padding,
+        )
+    )
+
+  padded_non_zero_rows = np.stack(padded_non_zero_rows, axis=1)
+
+  assert padded_non_zero_rows.shape[1] == block_mask.shape[2], (
+      padded_non_zero_rows.shape,
+      block_mask.shape,
+  )
+
+  # For each col of array, select the rows indices in padded_non_zero_rows,
+  # ignore padding.
+  def select_rows(array):
+    assert array.ndim == 2
+    assert padded_non_zero_rows.ndim == 2
+    assert array.shape[1] == padded_non_zero_rows.shape[1]
+    assert array.shape[0] >= padded_non_zero_rows.shape[0]
+    selected_cols = []
+    for col in range(array.shape[1]):
+      row = padded_non_zero_rows[:, col]
+      selected = array[:, col][row]
+      selected = np.where(row != padding, selected, 0)
+      selected_cols.append(selected)
+
+    return np.stack(selected_cols, axis=1)
+
+  return _slice_mask_info(
+      block_mask=block_mask,
+      data_next=data_next,
+      mask_next=mask_next,
+      head_shards=head_shards,
+      slice_function=select_rows,
+  )
+
+
+def _slice_mask_info(
+    *,
+    block_mask: np.ndarray,
+    data_next: np.ndarray,
+    mask_next: np.ndarray,
+    head_shards: int,
+    slice_function: Callable[[np.ndarray], np.ndarray],
+):
+  new_block_mask = []
+  new_data_next = []
+  new_mask_next = []
+  for head_shard in range(head_shards):
+    head_block_mask = block_mask[head_shard]
+    head_block_mask = slice_function(head_block_mask)
+    new_block_mask.append(head_block_mask)
+
+    head_data_next = data_next[head_shard]
+    head_data_next = slice_function(head_data_next)
+    new_data_next.append(head_data_next)
+
+    if mask_next is not None:
+      head_mask_next = mask_next[head_shard]
+      head_mask_next = slice_function(head_mask_next)
+      new_mask_next.append(head_mask_next)
+
+  block_mask = np.stack(new_block_mask, axis=0)
+  data_next = np.stack(new_data_next, axis=0)
+  if mask_next is not None:
+    mask_next = np.stack(new_mask_next, axis=0)
+
+  return block_mask, data_next, mask_next
+
+
+process_mask = functools.partial(_process_mask, is_dkv=False)
+process_mask_dkv = functools.partial(_process_mask, is_dkv=True)


Test patch:
diff --git a/tests/pallas/BUILD b/tests/pallas/BUILD
index 368689f0a..cad4c103f 100644
--- a/tests/pallas/BUILD
+++ b/tests/pallas/BUILD
@@ -78,3 +78,32 @@ jax_test(
         "//third_party/py/jax:pallas_tpu_ops",
     ] + py_deps("absl/testing") + py_deps("numpy") + py_deps("hypothesis"),
 )
+
+jax_test(
+    name = "splash_attention_kernel_test",
+    srcs = [
+        "splash_attention_kernel_test.py",
+    ],
+    disable_backends = [
+        "cpu",
+        "gpu",
+    ],
+    shard_count = 18,
+    deps = [
+        "//third_party/py/jax:pallas_tpu_ops",
+    ] + py_deps("absl/testing") + py_deps("numpy") + py_deps("hypothesis"),
+)
+
+jax_test(
+    name = "splash_attention_mask_test",
+    srcs = [
+        "splash_attention_mask_test.py",
+    ],
+    disable_backends = [
+        "cpu",
+        "gpu",
+    ],
+    deps = [
+        "//third_party/py/jax:pallas_tpu_ops",
+    ] + py_deps("absl/testing") + py_deps("numpy") + py_deps("hypothesis"),
+)
diff --git a/tests/pallas/splash_attention_kernel_test.py b/tests/pallas/splash_attention_kernel_test.py
new file mode 100644
index 000000000..05397f125
--- /dev/null
+++ b/tests/pallas/splash_attention_kernel_test.py
@@ -0,0 +1,641 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for splash_attention."""
+from __future__ import annotations
+
+import dataclasses
+import functools
+from typing import Any, Callable, TypeVar
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import hypothesis as hp
+import hypothesis.strategies as hps
+import jax
+from jax import random
+from jax import tree_util
+from jax._src import test_util as jtu
+from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_kernel as splash
+from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask as mask_lib
+import jax.numpy as jnp
+import numpy as np
+
+
+jax.config.parse_flags_with_absl()
+
+hp.settings.register_profile(
+    "deterministic",
+    database=None,
+    derandomize=True,
+    deadline=None,
+    max_examples=30,
+    print_blob=True,
+)
+hp.settings.load_profile("deterministic")
+
+partial = functools.partial
+Draw = TypeVar("Draw", bound=Callable[[hps.SearchStrategy[Any]], Any])
+
+
+@hps.composite
+def segment_ids_strategy(draw, seq_len: int) -> splash.SegmentIds:
+  boundaries = hps.sets(hps.integers(1, seq_len - 1), min_size=1, max_size=4)
+  bounds = sorted(draw(boundaries))
+  ids_array = np.empty((seq_len,), dtype=np.int32)
+  for i, (start, end) in enumerate(zip((0, *bounds), (*bounds, seq_len))):
+    # Not sure why, but short segments can trip things up
+    if end - start < 2:
+      end = start + 2
+    ids_array[start:end] = i
+  return splash.SegmentIds(ids_array, ids_array)
+
+
+def seed_strategy() -> hps.SearchStrategy[int]:
+  return hps.integers(min_value=0, max_value=4)
+
+
+class Mask:
+
+  def get_mask(self) -> mask_lib.Mask:
+    raise NotImplementedError()
+
+
+def full_mask_strategy(
+    q_seq_len: int, kv_seq_len: int
+) -> hps.SearchStrategy[Mask]:
+  return hps.just(FullMask(q_seq_len, kv_seq_len))
+
+
+@dataclasses.dataclass
+class FullMask(Mask):
+  q_seq_len: int
+  kv_seq_len: int
+
+  def get_mask(self) -> mask_lib.Mask:
+    return mask_lib.FullMask((self.q_seq_len, self.kv_seq_len))
+
+
+def causal_mask_strategy(
+    q_seq_len: int, kv_seq_len: int
+) -> hps.SearchStrategy[Mask]:
+  return hps.just(CausalMask(q_seq_len, kv_seq_len))
+
+
+@dataclasses.dataclass
+class CausalMask(Mask):
+  q_seq_len: int
+  kv_seq_len: int
+
+  def get_mask(self) -> mask_lib.Mask:
+    return mask_lib.CausalMask((self.q_seq_len, self.kv_seq_len))
+
+
+@dataclasses.dataclass
+class LocalAttentionMask(Mask):
+  seq_len: int
+  left: int | None
+  right: int | None
+  offset: int
+
+  def get_mask(self) -> mask_lib.Mask:
+    mask = mask_lib.LocalMask(
+        (self.seq_len, self.seq_len),
+        (self.left, self.right),
+        offset=self.offset,
+    )
+    # Make sure that no row is full of zeros as this is leads to undefined
+    # softmax.
+    diagonal = mask_lib.NumpyMask(np.identity(self.seq_len, dtype=np.bool_))
+    return mask | diagonal
+
+
+@hps.composite
+def local_attention_mask_strategy(draw: Draw, seq_len: int) -> Mask:
+  left_window = draw(
+      hps.one_of(hps.none(), hps.integers(min_value=0, max_value=seq_len))
+  )
+  right_window = draw(
+      hps.one_of(hps.none(), hps.integers(min_value=0, max_value=seq_len))
+  )
+  offset = draw(hps.integers(min_value=-seq_len, max_value=seq_len - 1))
+  return LocalAttentionMask(seq_len, left_window, right_window, offset=offset)
+
+
+@dataclasses.dataclass
+class RandomMask(Mask):
+  q_seq_len: int
+  kv_seq_len: int
+  sparsity: float
+  seed: int
+
+  def get_mask(self) -> mask_lib.Mask:
+    mask = mask_lib.make_random_mask(
+        (self.q_seq_len, self.kv_seq_len), self.sparsity, self.seed
+    )
+    # Make sure that no row is full of zeros as this is leads to undefined
+    # softmax.
+    mask[:, 0] = True
+
+    return mask_lib.NumpyMask(mask)
+
+
+@hps.composite
+def random_mask_strategy(draw: Draw, q_seq_len: int, kv_seq_len: int) -> Mask:
+  seed = draw(hps.integers(min_value=0, max_value=2**32 - 1))
+  sparsity = draw(hps.floats(min_value=0.0, max_value=0.5))
+  return RandomMask(q_seq_len, kv_seq_len, sparsity, seed)
+
+
+@dataclasses.dataclass
+class ComposeMask(Mask):
+  left: Mask
+  right: Mask
+  op: Callable[[mask_lib.Mask, mask_lib.Mask], mask_lib.Mask]
+
+  def get_mask(self) -> mask_lib.Mask:
+    return self.op(self.left.get_mask(), self.right.get_mask())
+
+
+@hps.composite
+def compose_mask_strategy(draw: Draw, q_seq_len: int, kv_seq_len: int) -> Mask:
+  mask1 = draw(mask_strategy(q_seq_len, kv_seq_len))
+  mask2 = draw(mask_strategy(q_seq_len, kv_seq_len))
+  op = draw(
+      hps.one_of(hps.just(mask_lib.LogicalOr), hps.just(mask_lib.LogicalAnd))
+  )
+  return ComposeMask(mask1, mask2, op)
+
+
+@hps.composite
+def mask_strategy(draw: Draw, q_seq_len: int, kv_seq_len: int) -> Mask:
+  oneof = [
+      causal_mask_strategy(q_seq_len, kv_seq_len),
+      full_mask_strategy(q_seq_len, kv_seq_len),
+      random_mask_strategy(q_seq_len, kv_seq_len),
+      # TODO(amagni) Composing masks creates masks that produce minor numerical
+      # differences. We should investigate this in the future.
+      # compose_mask_strategy(q_seq_len, kv_seq_len),
+  ]
+
+  if q_seq_len == kv_seq_len:
+    oneof.append(local_attention_mask_strategy(q_seq_len))
+
+  return draw(hps.one_of(oneof))
+
+
+@hps.composite
+def mha_mask_strategy(
+    draw: Draw, q_seq_len: int, kv_seq_len: int, num_heads: int
+) -> np.ndarray:
+  return np.stack(
+      [draw(mask_strategy(q_seq_len, kv_seq_len)) for _ in range(num_heads)]
+  )
+
+
+@hps.composite
+def sequence_length_strategy(draw: Draw) -> tuple[int, int]:
+  q_seq_len = draw(hps.sampled_from([1024, 2048, 4096]))
+  kv_seq_len = draw(hps.sampled_from([1024, 2048, 4096]))
+  return q_seq_len, kv_seq_len
+
+
+@hps.composite
+def attention_strategy(draw: Draw) -> tuple[int, int, int, np.dtype]:
+  q_seq_len, kv_seq_len = draw(sequence_length_strategy())
+  head_dim = draw(hps.sampled_from([128, 256]))
+  dtype = draw(hps.sampled_from([np.dtype("float32"), np.dtype(jnp.bfloat16)]))
+  return q_seq_len, kv_seq_len, head_dim, dtype
+
+
+@hps.composite
+def mha_strategy(draw: Draw) -> tuple[int, int, int, int, int, np.dtype]:
+  q_seq_len, kv_seq_len, head_dim, dtype = draw(attention_strategy())
+  num_q_heads, num_kv_heads = draw(
+      hps.sampled_from([(1, 1), (2, 2), (4, 1), (8, 4), (6, 2)])
+  )
+  return q_seq_len, kv_seq_len, num_q_heads, num_kv_heads, head_dim, dtype
+
+
+@hps.composite
+def block_sizes_strategy(
+    draw: Draw,
+    q_seq_len: int,
+    kv_seq_len: int,
+    include_bwd_blocks: bool = False,
+    use_fused_bwd_kernel: bool = False,
+) -> splash.BlockSizes:
+  all_block_shapes = [128, 256, 512]
+  q_layout = draw(hps.sampled_from(splash.QKVLayout))
+  k_layout = draw(hps.sampled_from(splash.QKVLayout))
+  v_layout = draw(hps.sampled_from(splash.QKVLayout))
+  layouts = dict(q_layout=q_layout, k_layout=k_layout, v_layout=v_layout)
+  q_valid_block_shapes = [bs for bs in all_block_shapes if bs <= q_seq_len]
+  kv_valid_block_shapes = [bs for bs in all_block_shapes if bs <= kv_seq_len]
+  bq, bkv = (
+      draw(hps.sampled_from(q_valid_block_shapes)),
+      draw(hps.sampled_from(kv_valid_block_shapes)),
+  )
+  bkv_compute = draw(
+      hps.sampled_from([None, *[b for b in kv_valid_block_shapes if b <= bkv]])
+  )
+  if not include_bwd_blocks:
+    return splash.BlockSizes(
+        block_q=bq, block_kv=bkv, block_kv_compute=bkv_compute, **layouts
+    )
+  all_block_shapes = [128, 256]
+  q_valid_block_shapes = [bs for bs in all_block_shapes if bs <= q_seq_len]
+  kv_valid_block_shapes = [bs for bs in all_block_shapes if bs <= kv_seq_len]
+  bq_dkv, bkv_dkv = (
+      draw(hps.sampled_from(q_valid_block_shapes)),
+      draw(hps.sampled_from(kv_valid_block_shapes)),
+  )
+  if use_fused_bwd_kernel:
+    bq_dq, bkv_dq = None, None
+  else:
+    bq_dq = draw(hps.sampled_from(q_valid_block_shapes))
+    bkv_dq = draw(hps.sampled_from(kv_valid_block_shapes))
+  block_kv_dkv_compute = draw(
+      hps.sampled_from(
+          [None, *[b for b in kv_valid_block_shapes if b <= bkv_dkv]]
+      )
+  )
+  return splash.BlockSizes(
+      block_q=bq,
+      block_kv=bkv,
+      block_kv_compute=bkv_compute,
+      block_q_dkv=bq_dkv,
+      block_kv_dkv=bkv_dkv,
+      block_kv_dkv_compute=block_kv_dkv_compute,
+      block_q_dq=bq_dq,
+      block_kv_dq=bkv_dq,
+      use_fused_bwd_kernel=use_fused_bwd_kernel,
+      **layouts,
+  )
+
+
+def attn_logits_soft_cap_strategy() -> hps.SearchStrategy[float | None]:
+  return hps.one_of(hps.just(None), hps.floats(min_value=1.0, max_value=50.0))
+
+
+class AttentionTest(jtu.JaxTestCase):
+
+  def setUp(self):
+    super().setUp()
+    if not jtu.test_device_matches(["tpu"]):
+      self.skipTest("Need TPU devices")
+
+  def _assert_allclose(self, x, y, **kwargs):
+    if x.dtype == np.dtype(jnp.bfloat16):
+      x = x.astype(np.float32)
+    if y.dtype == np.dtype(jnp.bfloat16):
+      y = y.astype(np.float32)
+    self.assertEqual(x.dtype, y.dtype)
+    self.assertTupleEqual(x.shape, y.shape)
+    np.testing.assert_allclose(x, y, **kwargs)
+
+
+class SplashAttentionTest(AttentionTest):
+
+  @parameterized.product(
+      is_mqa=(False, True),
+      is_segmented=(False, True),
+  )
+  @hp.given(hps.data())
+  def test_splash_attention(self, is_mqa, is_segmented, data):
+    seed = data.draw(seed_strategy())
+    key = random.key(seed)
+    k1, k2, k3 = random.split(key, 3)
+
+    q_seq_len, kv_seq_len, num_q_heads, num_kv_heads, head_dim, dtype = (
+        data.draw(mha_strategy())
+    )
+
+    # Avoid segment ids for rectangular matrices, as its hard to enforce
+    # valid masks (non-0 rows).
+    hp.assume(q_seq_len == kv_seq_len or not is_segmented)
+
+    q = random.uniform(k1, (num_q_heads, q_seq_len, head_dim), dtype=dtype)
+    if is_mqa:
+      k = random.uniform(k2, (kv_seq_len, head_dim), dtype=dtype)
+      v = random.uniform(k3, (kv_seq_len, head_dim), dtype=dtype)
+    else:
+      k = random.uniform(k2, (num_kv_heads, kv_seq_len, head_dim), dtype=dtype)
+      v = random.uniform(k3, (num_kv_heads, kv_seq_len, head_dim), dtype=dtype)
+
+    segment_ids = None
+    if is_segmented:
+      assert q_seq_len == kv_seq_len
+      segment_ids = data.draw(segment_ids_strategy(q_seq_len))
+
+    attn_logits_soft_cap = data.draw(attn_logits_soft_cap_strategy())
+    masks = data.draw(mha_mask_strategy(q_seq_len, kv_seq_len, num_q_heads))
+    mask = mask_lib.MultiHeadMask(tuple((m.get_mask() for m in masks)))
+    block_sizes = data.draw(block_sizes_strategy(q_seq_len, kv_seq_len))
+
+    if is_mqa:
+      attn_ref = splash.make_masked_mqa_reference(mask)
+      attn = splash.make_splash_mqa_single_device(
+          mask,
+          block_sizes=block_sizes,
+          attn_logits_soft_cap=attn_logits_soft_cap,
+      )
+    else:
+      attn_ref = splash.make_masked_mha_reference(mask)
+      attn = splash.make_splash_mha_single_device(
+          mask,
+          block_sizes=block_sizes,
+          attn_logits_soft_cap=attn_logits_soft_cap,
+      )
+    o = attn(q, k, v, segment_ids)
+    o_ref = attn_ref(
+        q.astype(np.float32),
+        k.astype(np.float32),
+        v.astype(np.float32),
+        segment_ids,
+        attn_logits_soft_cap=attn_logits_soft_cap,
+    )
+    self._assert_allclose(o, o_ref, atol=3e-3, rtol=3e-3)
+
+  @parameterized.product(
+      is_mqa=(False, True),
+      is_segmented=(False, True),
+  )
+  @hp.given(hps.data())
+  def test_splash_attention_fwd(
+      self, is_mqa, is_segmented, data
+  ):
+    seed = data.draw(seed_strategy())
+    key = random.key(seed)
+    k1, k2, k3 = random.split(key, 3)
+
+    q_seq_len, kv_seq_len, num_q_heads, num_kv_heads, head_dim, dtype = (
+        data.draw(mha_strategy())
+    )
+
+    # Avoid segment ids for rectangular matrices, as its hard to enforce
+    # valid masks (non-0 rows).
+    hp.assume(q_seq_len == kv_seq_len or not is_segmented)
+
+    q = random.uniform(k1, (num_q_heads, q_seq_len, head_dim), dtype=dtype)
+    if is_mqa:
+      k = random.uniform(k2, (kv_seq_len, head_dim), dtype=dtype)
+      v = random.uniform(k3, (kv_seq_len, head_dim), dtype=dtype)
+    else:
+      k = random.uniform(k2, (num_kv_heads, kv_seq_len, head_dim), dtype=dtype)
+      v = random.uniform(k3, (num_kv_heads, kv_seq_len, head_dim), dtype=dtype)
+
+    segment_ids = None
+    if is_segmented:
+      assert q_seq_len == kv_seq_len
+      segment_ids = data.draw(segment_ids_strategy(q_seq_len))
+    attn_logits_soft_cap = data.draw(attn_logits_soft_cap_strategy())
+    masks = data.draw(mha_mask_strategy(q_seq_len, kv_seq_len, num_q_heads))
+    mask = mask_lib.MultiHeadMask(tuple((m.get_mask() for m in masks)))
+    block_sizes = data.draw(block_sizes_strategy(q_seq_len, kv_seq_len))
+    if is_mqa:
+      attn_ref = splash.make_masked_mqa_reference(mask)
+      attn = splash.make_splash_mqa_single_device(
+          mask,
+          block_sizes=block_sizes,
+          save_residuals=True,
+          attn_logits_soft_cap=attn_logits_soft_cap,
+      )
+    else:
+      attn_ref = splash.make_masked_mha_reference(mask)
+      attn = splash.make_splash_mha_single_device(
+          mask,
+          block_sizes=block_sizes,
+          save_residuals=True,
+          attn_logits_soft_cap=attn_logits_soft_cap,
+      )
+    attn_ref = partial(
+        attn_ref,
+        save_residuals=True,
+        attn_logits_soft_cap=attn_logits_soft_cap,
+    )
+    o, (logsumexp,) = attn(q, k, v, segment_ids)
+    o_ref, (logsumexp_ref,) = attn_ref(
+        q.astype(jnp.float32),
+        k.astype(jnp.float32),
+        v.astype(jnp.float32),
+        segment_ids,
+    )
+    self._assert_allclose(o, o_ref, atol=3e-3, rtol=3e-3)
+    self._assert_allclose(logsumexp, logsumexp_ref, atol=1e-3, rtol=1e-3)
+
+  @parameterized.product(
+      is_segmented=(False, True),
+  )
+  @hp.given(hps.data())
+  def test_splash_attention_custom_bwd(self, is_segmented, data):
+    seed = data.draw(seed_strategy(), label="seed")
+    key = random.key(1 + seed)
+    k1, k2, k3, k4 = random.split(key, 4)
+
+    q_seq_len, kv_seq_len, head_dim, dtype = data.draw(attention_strategy())
+
+    # Avoid segment ids for rectangular matrices, as its hard to enforce
+    # valid masks (non-0 rows).
+    hp.assume(q_seq_len == kv_seq_len or not is_segmented)
+
+    q = random.uniform(k1, (q_seq_len, head_dim), dtype=dtype)
+    k = random.uniform(k2, (kv_seq_len, head_dim), dtype=dtype)
+    v = random.uniform(k3, (kv_seq_len, head_dim), dtype=dtype)
+    segment_ids = None
+    if is_segmented:
+      assert q_seq_len == kv_seq_len
+      segment_ids = data.draw(segment_ids_strategy(q_seq_len))
+    masks = data.draw(mha_mask_strategy(q_seq_len, kv_seq_len, 1))
+    mask = jnp.array(masks[0].get_mask()[:, :])
+    attn_logits_soft_cap = data.draw(attn_logits_soft_cap_strategy(),
+                                     label="logit_cap")
+    attn_ref = partial(splash.attention_reference, mask,
+                       attn_logits_soft_cap=attn_logits_soft_cap)
+    attn_custom = partial(splash.attention_reference_custom, mask,
+                          attn_logits_soft_cap=attn_logits_soft_cap)
+    attn_custom_vanilla = partial(splash.attention_reference_custom, mask,
+                                  custom_type="vanilla",
+                                  attn_logits_soft_cap=attn_logits_soft_cap)
+    o_ref, attn_vjp_ref = jax.vjp(attn_ref, q, k, v, segment_ids)
+    q32, k32, v32 = tree_util.tree_map(lambda x: x.astype(jnp.float32),
+                                       (q, k, v))
+    o_custom = attn_custom(q32, k32, v32, segment_ids)
+    _, attn_vjp = jax.vjp(attn_custom, q32, k32, v32, segment_ids)
+    _, attn_vanilla_vjp = jax.vjp(attn_custom_vanilla, q32, k32, v32,
+                                  segment_ids)
+    do = random.uniform(k4, o_custom.shape, dtype=o_custom.dtype) / 10.
+    # These should be identical
+    self._assert_allclose(o_custom, o_ref, atol=1e-5, rtol=1e-5)
+    dq, dk, dv, _ = attn_vjp(do)
+    dq_vanilla, dk_vanilla, dv_vanilla, _ = attn_vanilla_vjp(do)
+    dq_ref, dk_ref, dv_ref, _ = attn_vjp_ref(do)
+    # These will be different because of reassociation
+    if dtype == jnp.bfloat16:
+      atols = {"dv": 4e-3, "dq": 0.05, "dk": 0.05}
+      atols_v = {"dv": 8e-3, "dq": 2e-2, "dk": 2e-2}
+      rtols = {"dv": 4e-3, "dq": 0.05, "dk": 0.05}
+      rtols_v = {"dv": 8e-3, "dq": 2e-2, "dk": 2e-2}
+      if jtu.is_device_tpu(version=5):
+        atols["dk"] = 0.065
+    elif dtype == jnp.float32:
+      atols = {"dv": 3e-3, "dq": 0.05, "dk": 0.05}
+      atols_v = {"dv": 4e-4, "dq": 2e-3, "dk": 3e-3}
+      rtols = {"dv": 3e-3, "dq": 0.05, "dk": 0.05}
+      rtols_v = {"dv": 8e-3, "dq": 5e-4, "dk": 5e-4}
+      if jtu.is_device_tpu(version=4):
+        atols["dk"] = 0.09
+    else:
+      raise NotImplementedError
+    with self.subTest("dv"):
+      self._assert_allclose(
+          dv_vanilla, dv_ref, atol=atols_v["dv"], rtol=rtols_v["dv"]
+      )
+      self._assert_allclose(dv, dv_ref, atol=atols["dv"], rtol=rtols["dv"])
+    with self.subTest("dq"):
+      self._assert_allclose(
+          dq_vanilla, dq_ref, atol=atols_v["dq"], rtol=rtols_v["dq"]
+      )
+      self._assert_allclose(dq, dq_ref, atol=atols["dq"], rtol=rtols["dq"])
+    with self.subTest("dk"):
+      self._assert_allclose(
+          dk_vanilla, dk_ref, atol=atols_v["dk"], rtol=rtols_v["dk"]
+      )
+      self._assert_allclose(dk, dk_ref, atol=atols["dk"], rtol=rtols["dk"])
+
+  @parameterized.product(
+      is_mqa=(False, True),
+      is_segmented=(False, True),
+      downcast_smem_data=(False, True),
+      use_fused_bwd_kernel=(False, True),
+  )
+  @hp.given(hps.data())
+  def test_splash_attention_bwd(
+      self, is_mqa, is_segmented, downcast_smem_data, use_fused_bwd_kernel, data
+  ):
+    seed = data.draw(seed_strategy())
+    key = random.key(seed)
+    k1, k2, k3, k4 = random.split(key, 4)
+
+    q_seq_len, kv_seq_len, num_q_heads, num_kv_heads, head_dim, dtype = (
+        data.draw(mha_strategy())
+    )
+
+    # Avoid segment ids for rectangular matrices, as its hard to enforce
+    # valid masks (non-0 rows).
+    hp.assume(q_seq_len == kv_seq_len or not is_segmented)
+
+    q = random.uniform(k1, (num_q_heads, q_seq_len, head_dim), dtype=dtype)
+    if is_mqa:
+      k = random.uniform(k2, (kv_seq_len, head_dim), dtype=dtype)
+      v = random.uniform(k3, (kv_seq_len, head_dim), dtype=dtype)
+    else:
+      k = random.uniform(k2, (num_kv_heads, kv_seq_len, head_dim), dtype=dtype)
+      v = random.uniform(k3, (num_kv_heads, kv_seq_len, head_dim), dtype=dtype)
+
+    segment_ids = None
+    if is_segmented:
+      assert q_seq_len == kv_seq_len
+      segment_ids = data.draw(segment_ids_strategy(q_seq_len))
+    attn_logits_soft_cap = data.draw(attn_logits_soft_cap_strategy())
+    masks = data.draw(mha_mask_strategy(q_seq_len, kv_seq_len, num_q_heads))
+    mask = mask_lib.MultiHeadMask(tuple(m.get_mask() for m in masks))
+    block_sizes = data.draw(
+        block_sizes_strategy(q_seq_len, kv_seq_len, include_bwd_blocks=True,
+                             use_fused_bwd_kernel=use_fused_bwd_kernel)
+    )
+    if is_mqa:
+      attn_ref = splash.make_masked_mqa_reference(mask, backward_impl="custom")
+      attn = splash.make_splash_mqa_single_device(
+          mask,
+          block_sizes=block_sizes,
+          downcast_smem_data=downcast_smem_data,
+          attn_logits_soft_cap=attn_logits_soft_cap,
+      )
+    else:
+      attn_ref = splash.make_masked_mha_reference(mask, backward_impl="custom")
+      attn = splash.make_splash_mha_single_device(
+          mask,
+          block_sizes=block_sizes,
+          downcast_smem_data=downcast_smem_data,
+          attn_logits_soft_cap=attn_logits_soft_cap,
+      )
+    o, attn_vjp = jax.vjp(attn, q, k, v, segment_ids)
+    q32, k32, v32 = tree_util.tree_map(
+        lambda x: x.astype(jnp.float32), (q, k, v)
+    )
+    o_ref, (logsumexp,) = attn_ref(
+        q32,
+        k32,
+        v32,
+        segment_ids,
+        save_residuals=True,
+        attn_logits_soft_cap=attn_logits_soft_cap,
+    )
+    self._assert_allclose(o, o_ref, atol=3e-3, rtol=3e-3)
+
+    do = random.uniform(k4, o.shape, dtype=o.dtype)
+    dq, dk, dv, _ = attn_vjp(do)
+    def bwd(
+        mask, q, k, v, segment_ids, o, logsumexp, do
+    ) -> tuple[jax.Array, jax.Array, jax.Array]:
+      _, dq, dk, dv, _ = splash._attention_reference_custom_bwd(
+          splash.DEFAULT_MASK_VALUE,
+          False,
+          "flash",
+          attn_logits_soft_cap,
+          (mask, q, k, v, segment_ids, o, logsumexp),
+          do,
+      )
+      return dq, dk, dv
+
+    is_grouped = not is_mqa and num_kv_heads < num_q_heads
+    assert num_q_heads % num_kv_heads == 0
+    head_multiplier = num_q_heads // num_kv_heads
+    if is_mqa:
+      bwd = jax.vmap(bwd, in_axes=(0, 0, None, None, None, 0, 0, 0))
+    else:
+      bwd = jax.vmap(bwd, in_axes=(0, 0, 0, 0, None, 0, 0, 0))
+      # Interleave the KV heads to match the corresponding Q heads.
+      if is_grouped:
+        k32 = jnp.repeat(k32, head_multiplier, axis=0)
+        v32 = jnp.repeat(v32, head_multiplier, axis=0)
+
+    dq_ref, dk_ref, dv_ref = bwd(
+        mask[:, :, :],
+        q32,
+        k32,
+        v32,
+        segment_ids,
+        o.astype(jnp.float32),
+        logsumexp,
+        do.astype(jnp.float32),
+    )
+    if is_mqa:
+      dk_ref, dv_ref = dk_ref.sum(axis=0), dv_ref.sum(axis=0)
+    elif is_grouped:
+      # Perform the sum reduction across the head_multiplier dimension only.
+      # So that the output still has KV heads.
+      dk_ref = dk_ref.reshape(num_kv_heads, head_multiplier, *dk_ref.shape[1:])
+      dv_ref = dv_ref.reshape(num_kv_heads, head_multiplier, *dv_ref.shape[1:])
+
+      dk_ref, dv_ref = dk_ref.sum(axis=1), dv_ref.sum(axis=1)
+
+    self._assert_allclose(dv, dv_ref, atol=2e-2, rtol=3e-2)
+    self._assert_allclose(dq, dq_ref, atol=2e-2, rtol=3e-2)
+    self._assert_allclose(dk, dk_ref, atol=2e-2, rtol=3e-2)
+
+
+if __name__ == "__main__":
+  absltest.main(testLoader=jtu.JaxTestLoader())
diff --git a/tests/pallas/splash_attention_mask_test.py b/tests/pallas/splash_attention_mask_test.py
new file mode 100644
index 000000000..a40887210
--- /dev/null
+++ b/tests/pallas/splash_attention_mask_test.py
@@ -0,0 +1,2097 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for splash_attention_masks."""
+from __future__ import annotations
+
+from typing import List
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask as mask_lib
+from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask_info as mask_info_lib
+from jax._src import test_util as jtu
+import numpy as np
+
+jax.config.parse_flags_with_absl()
+
+
+def _make_lazy_causal_mask(*args, **kwargs):
+  mask = mask_lib.CausalMask(*args, **kwargs)
+  return mask[:, :]
+
+
+def _make_causal_mask(*args, **kwargs):
+  return mask_lib.make_causal_mask(*args, **kwargs)
+
+
+def _make_lazy_local_attention_mask(*args, **kwargs):
+  mask = mask_lib.LocalMask(*args, **kwargs)
+  return mask[:, :]
+
+
+def _make_local_attention_mask(*args, **kwargs):
+  return mask_lib.make_local_attention_mask(*args, **kwargs)
+
+
+class SplashAttentionMaskTest(jtu.JaxTestCase):
+
+  @parameterized.parameters([_make_lazy_causal_mask, _make_causal_mask])
+  def test_causal_mask(self, make_causal_mask):
+    expected = np.array([[1]], dtype=np.bool_)
+    actual = make_causal_mask((1, 1))
+
+    with self.subTest("unit"):
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 0, 0, 0],
+            [1, 1, 0, 0],
+            [1, 1, 1, 0],
+            [1, 1, 1, 1],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_causal_mask((4, 4))
+
+    with self.subTest("square"):
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 0, 0, 0, 0, 0],
+            [1, 1, 0, 0, 0, 0],
+            [1, 1, 1, 0, 0, 0],
+            [1, 1, 1, 1, 0, 0],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_causal_mask((4, 6))
+
+    with self.subTest("wide_rectangle"):
+      self.assertArraysEqual(actual, expected)
+
+    actual = make_causal_mask((6, 4))
+    expected = np.array(
+        [
+            [1, 0, 0, 0],
+            [1, 1, 0, 0],
+            [1, 1, 1, 0],
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+        ],
+        dtype=np.bool_,
+    )
+
+    with self.subTest("tall_rectangle"):
+      self.assertArraysEqual(actual, expected)
+
+    actual = make_causal_mask((4, 4), -1)
+    expected = np.array(
+        [
+            [0, 0, 0, 0],
+            [1, 0, 0, 0],
+            [1, 1, 0, 0],
+            [1, 1, 1, 0],
+        ],
+        dtype=np.bool_,
+    )
+
+    with self.subTest("negative_offset"):
+      self.assertArraysEqual(actual, expected)
+
+    actual = make_causal_mask((4, 4), 1)
+    expected = np.array(
+        [
+            [1, 1, 0, 0],
+            [1, 1, 1, 0],
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+        ],
+        dtype=np.bool_,
+    )
+
+    with self.subTest("positive_offset"):
+      self.assertArraysEqual(actual, expected)
+
+  @parameterized.parameters(
+      [_make_lazy_local_attention_mask, _make_local_attention_mask]
+  )
+  def test_local_attention_mask(self, make_local_attention_mask):
+    expected = np.array([[1]], dtype=np.bool_)
+    actual = make_local_attention_mask((1, 1), (0, None), offset=0)
+    self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+            [0, 1, 1, 1],
+            [0, 0, 1, 1],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((4, 4), (1, None), offset=0)
+    with self.subTest("left_1"):
+      self.assertArraysEqual(actual, expected)
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 1, 1, 0],
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((4, 4), (None, 2), offset=0)
+    with self.subTest("right_2"):
+      self.assertArraysEqual(actual, expected)
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 1, 0, 0],
+            [1, 1, 1, 0],
+            [0, 1, 1, 1],
+            [0, 0, 1, 1],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((4, 4), (1, 1), offset=0)
+    with self.subTest("left_1_right_1"):
+      self.assertArraysEqual(actual, expected)
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 0, 0, 0],
+            [1, 1, 0, 0],
+            [0, 1, 1, 0],
+            [0, 0, 1, 1],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((4, 4), (1, 0), offset=0)
+    with self.subTest("left_1_right_0"):
+      self.assertArraysEqual(actual, expected)
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 1, 1, 0],
+            [0, 1, 1, 1],
+            [0, 0, 1, 1],
+            [0, 0, 0, 1],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((4, 4), (0, 2), offset=0)
+    with self.subTest("left_0_right_2"):
+      self.assertArraysEqual(actual, expected)
+
+  @parameterized.parameters(
+      [_make_lazy_local_attention_mask, _make_local_attention_mask]
+  )
+  def test_local_attention_mask_wide_rectangle(self, make_local_attention_mask):
+    expected = np.array(
+        [
+            [1, 1, 1, 1, 1, 1],
+            [1, 1, 1, 1, 1, 1],
+            [0, 1, 1, 1, 1, 1],
+            [0, 0, 1, 1, 1, 1],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((4, 6), (1, None), offset=0)
+    with self.subTest("left_1"):
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 1, 1, 0, 0, 0],
+            [1, 1, 1, 1, 0, 0],
+            [1, 1, 1, 1, 1, 0],
+            [1, 1, 1, 1, 1, 1],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((4, 6), (None, 2), offset=0)
+    with self.subTest("right_2"):
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 1, 0, 0, 0, 0],
+            [1, 1, 1, 0, 0, 0],
+            [0, 1, 1, 1, 0, 0],
+            [0, 0, 1, 1, 1, 0],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((4, 6), (1, 1), offset=0)
+    with self.subTest("left_1_right_1"):
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 0, 0, 0, 0, 0],
+            [1, 1, 0, 0, 0, 0],
+            [0, 1, 1, 0, 0, 0],
+            [0, 0, 1, 1, 0, 0],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((4, 6), (1, 0), offset=0)
+    with self.subTest("left_1_right_0"):
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 1, 1, 0, 0, 0],
+            [0, 1, 1, 1, 0, 0],
+            [0, 0, 1, 1, 1, 0],
+            [0, 0, 0, 1, 1, 1],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((4, 6), (0, 2), offset=0)
+    with self.subTest("left_0_right_2"):
+      self.assertArraysEqual(actual, expected)
+
+  @parameterized.parameters(
+      [_make_lazy_local_attention_mask, _make_local_attention_mask]
+  )
+  def test_local_attention_mask_tall_rectangle(self, make_local_attention_mask):
+    expected = np.array(
+        [
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+            [0, 1, 1, 1],
+            [0, 0, 1, 1],
+            [0, 0, 0, 1],
+            [0, 0, 0, 0],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((6, 4), (1, None), offset=0)
+    with self.subTest("left_1"):
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 1, 1, 0],
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((6, 4), (None, 2), offset=0)
+    with self.subTest("right_2"):
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 1, 0, 0],
+            [1, 1, 1, 0],
+            [0, 1, 1, 1],
+            [0, 0, 1, 1],
+            [0, 0, 0, 1],
+            [0, 0, 0, 0],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((6, 4), (1, 1), offset=0)
+    with self.subTest("left_1_right_1"):
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 0, 0, 0],
+            [1, 1, 0, 0],
+            [0, 1, 1, 0],
+            [0, 0, 1, 1],
+            [0, 0, 0, 1],
+            [0, 0, 0, 0],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((6, 4), (1, 0), offset=0)
+    with self.subTest("left_1_right_0"):
+      self.assertArraysEqual(actual, expected)
+
+    expected = np.array(
+        [
+            [1, 1, 1, 0],
+            [0, 1, 1, 1],
+            [0, 0, 1, 1],
+            [0, 0, 0, 1],
+            [0, 0, 0, 0],
+            [0, 0, 0, 0],
+        ],
+        dtype=np.bool_,
+    )
+    actual = make_local_attention_mask((6, 4), (0, 2), offset=0)
+    with self.subTest("left_0_right_2"):
+      self.assertArraysEqual(actual, expected)
+
+  @parameterized.product(
+      block_size=[(256, 256), (256, 128), (128, 256)],
+      shape=[(1024, 1024), (1024, 2048), (2048, 1024)],
+  )
+  def test_lazy_causal_mask_chunking(
+      self, block_size: tuple[int, int], shape: tuple[int, int]
+  ):
+    dense_mask = mask_lib.make_causal_mask(shape=shape)
+    self._compare_masks(
+        dense_mask,
+        mask_lib.CausalMask(shape),
+        block_size,
+    )
+
+  @parameterized.parameters([
+      ((256, 256), (1024, 1024), (128, None), 0),
+      ((256, 128), (1024, 1024), (128, None), 16),
+      ((128, 256), (1024, 1024), (128, None), 16),
+      ((256, 256), (1024, 1024), (128, 256), 0),
+      ((256, 128), (1024, 1024), (128, 256), 0),
+      ((128, 256), (1024, 1024), (128, 256), 16),
+      ((256, 256), (1024, 1024), (None, 256), 0),
+      ((256, 128), (1024, 1024), (None, 256), 32),
+      ((128, 256), (1024, 1024), (None, 256), 32),
+      #
+      ((256, 256), (1024, 2048), (128, None), 0),
+      ((256, 128), (1024, 2048), (128, None), 16),
+      ((128, 256), (1024, 2048), (128, None), 16),
+      ((256, 256), (1024, 2048), (128, 256), 0),
+      ((256, 128), (1024, 2048), (128, 256), 0),
+      ((128, 256), (1024, 2048), (128, 256), 16),
+      ((256, 256), (1024, 2048), (None, 256), 0),
+      ((256, 128), (1024, 2048), (None, 256), 32),
+      ((128, 256), (1024, 2048), (None, 256), 32),
+      #
+      ((256, 256), (2048, 1024), (128, None), 0),
+      ((256, 128), (2048, 1024), (128, None), 16),
+      ((128, 256), (2048, 1024), (128, None), 16),
+      ((256, 256), (2048, 1024), (128, 256), 0),
+      ((256, 128), (2048, 1024), (128, 256), 0),
+      ((128, 256), (2048, 1024), (128, 256), 16),
+      ((256, 256), (2048, 1024), (None, 256), 0),
+      ((256, 128), (2048, 1024), (None, 256), 32),
+      ((128, 256), (2048, 1024), (None, 256), 32),
+  ])
+  def test_lazy_local_mask_chunking(
+      self,
+      block_size: tuple[int, int],
+      shape: tuple[int, int],
+      window_size: tuple[int | None, int | None],
+      offset: int,
+  ):
+    dense_mask = mask_lib.make_local_attention_mask(
+        shape, window_size, offset=offset
+    )
+    self._compare_masks(
+        dense_mask,
+        mask_lib.LocalMask(shape, window_size, offset),
+        block_size,
+    )
+
+  def test_using_logical_operators_raises_exception(self):
+    mask_1 = mask_lib.NumpyMask(
+        mask_lib.make_random_mask((256, 256), 0.5, seed=1)
+    )
+    mask_2 = mask_lib.NumpyMask(
+        mask_lib.make_random_mask((256, 256), 0.5, seed=2)
+    )
+
+    with self.subTest("logical_or"):
+      with self.assertRaises(NotImplementedError):
+        res = mask_1 or mask_2
+        del res
+
+    with self.subTest("logical_and"):
+      with self.assertRaises(NotImplementedError):
+        res = mask_1 and mask_2
+        del res
+
+  @parameterized.parameters([((256, 256),), ((512, 256),), ((512, 256),)])
+  def test_lazy_mask_or(self, shape: tuple[int, int]):
+    mask_1 = mask_lib.make_random_mask(shape, 0.5, seed=1)
+    mask_2 = mask_lib.make_random_mask(shape, 0.5, seed=2)
+
+    lazy_or = mask_lib.NumpyMask(mask_1) | mask_lib.NumpyMask(mask_2)
+    dense = np.logical_or(mask_1, mask_2)
+
+    self._compare_masks(dense, lazy_or, (256, 256))
+
+  @parameterized.parameters([((256, 256),), ((512, 256),), ((512, 256),)])
+  def test_lazy_mask_and(self, shape: tuple[int, int]):
+    mask_1 = mask_lib.make_random_mask(shape, 0.5, seed=1)
+    mask_2 = mask_lib.make_random_mask(shape, 0.5, seed=2)
+
+    lazy_and = mask_lib.NumpyMask(mask_1) & mask_lib.NumpyMask(mask_2)
+    dense = np.logical_and(mask_1, mask_2)
+
+    self._compare_masks(dense, lazy_and, (256, 256))
+
+  @parameterized.parameters([((256, 256),), ((512, 256),), ((512, 256),)])
+  def test_lazy_multi_head_mask(self, shape: tuple[int, int]):
+    mask_1 = mask_lib.make_random_mask(shape, 0.5, seed=1)
+    mask_2 = mask_lib.make_random_mask(shape, 0.5, seed=2)
+
+    lazy_multi_head = mask_lib.MultiHeadMask(
+        (mask_lib.NumpyMask(mask_1), mask_lib.NumpyMask(mask_2))
+    )
+    dense = np.stack((mask_1, mask_2), axis=0)
+
+    self._compare_masks(dense, lazy_multi_head, (256, 256))
+
+  @parameterized.parameters([((256, 256),), ((512, 256),), ((512, 256),)])
+  def test_lazy_full_mask(self, shape: tuple[int, int]):
+    lazy_full = mask_lib.FullMask(shape)
+    dense = np.ones(shape, dtype=np.bool_)
+
+    self._compare_masks(dense, lazy_full, (256, 256))
+
+  def _compare_masks(
+      self,
+      dense_mask: np.ndarray,
+      lazy_mask: mask_lib.Mask,
+      block_size: tuple[int, int],
+  ):
+    self.assertEqual(dense_mask.shape, lazy_mask.shape)
+
+    *prefix, width, height = dense_mask.shape
+
+    assert width % block_size[0] == 0
+    assert height % block_size[1] == 0
+
+    full_lazy_mask = lazy_mask[
+        (*[slice(p) for p in prefix], slice(None), slice(None))
+    ]
+    self.assertArraysEqual(dense_mask, full_lazy_mask)
+    for i, j in np.ndindex(width // block_size[0], height // block_size[1]):
+      indexer = (
+          *[slice(p) for p in prefix],
+          slice(i * block_size[0], (i + 1) * block_size[0]),
+          slice(j * block_size[1], (j + 1) * block_size[1]),
+      )
+      dense_chunk = dense_mask[indexer]
+      lazy_chunk = lazy_mask[indexer]
+      self.assertArraysEqual(dense_chunk, lazy_chunk)
+
+
+class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
+  """Check the construction of MaskInfo from Mask."""
+
+  def _assert_mask_info_match(
+      self,
+      actual: mask_info_lib.MaskInfo,
+      expected: mask_info_lib.MaskInfo,
+  ):
+
+    def assert_array_is_positive(array: np.ndarray | None):
+      if array is None:
+        return
+
+      is_positive = np.all(array >= 0)
+      self.assertTrue(is_positive)
+
+    assert_array_is_positive(actual.mask_next)
+    assert_array_is_positive(actual.partial_mask_blocks)
+    assert_array_is_positive(actual.block_mask)
+    assert_array_is_positive(actual.data_next)
+
+    self.assertEqual(
+        actual.data_next is not None, expected.data_next is not None
+    )
+    self.assertEqual(
+        actual.block_mask is not None, expected.block_mask is not None
+    )
+    self.assertEqual(
+        actual.mask_next is not None, expected.mask_next is not None
+    )
+    self.assertEqual(
+        actual.partial_mask_blocks is not None,
+        expected.partial_mask_blocks is not None,
+    )
+
+    self.assertEqual(
+        actual.q_sequence is not None, expected.q_sequence is not None
+    )
+
+    if actual.partial_mask_blocks is not None:
+      self.assertArraysEqual(
+          actual.partial_mask_blocks,
+          expected.partial_mask_blocks,
+          err_msg="partial_mask_blocks",
+          verbose=True,
+      )
+
+    if actual.q_sequence is not None:
+      self.assertArraysEqual(
+          actual.q_sequence,
+          expected.q_sequence,
+          err_msg="q_sequence",
+          verbose=True,
+      )
+
+    self.assertArraysEqual(
+        actual.block_mask,
+        expected.block_mask,
+        err_msg="block_mask",
+        verbose=True,
+    )
+
+    if actual.data_next is not None and actual.block_mask is not None:
+      self.assertEqual(actual.data_next.shape, actual.block_mask.shape)
+
+    if actual.block_mask is not None:
+      is_non_zero_block = np.where(expected.block_mask > 0, True, False)
+
+      self.assertArraysEqual(
+          np.where(is_non_zero_block, actual.data_next, -1),
+          expected.data_next,
+          err_msg="data_next",
+          verbose=True,
+      )
+
+    if actual.mask_next is not None:
+      is_partial_block = np.where(expected.block_mask == 1, True, False)
+      self.assertArraysEqual(
+          np.where(is_partial_block, actual.mask_next, -1),
+          expected.mask_next,
+          err_msg="mask_next",
+          verbose=True,
+      )
+
+  def _process_mask(self, *args, **kwargs):
+    mask_info, mask_function = mask_info_lib.process_mask(*args, **kwargs)
+    mask_info_dkv, dkv_mask_function = mask_info_lib.process_mask_dkv(
+        *args, **kwargs
+    )
+    self.assertEqual(mask_function, dkv_mask_function)
+    return mask_info, mask_info_dkv, mask_function
+
+  _expected_full_block_mask = np.array(
+      [
+          [2, 2, 2, 2],
+          [2, 2, 2, 2],
+          [2, 2, 2, 2],
+          [2, 2, 2, 2],
+      ],
+      dtype=np.int8,
+  )
+
+  _expected_full_block_mask_dkv = _expected_full_block_mask
+
+  _expected_full_data_next = np.array(
+      [
+          [0, 1, 2, 3],
+          [0, 1, 2, 3],
+          [0, 1, 2, 3],
+          [0, 1, 2, 3],
+      ],
+      dtype=np.int8,
+  )
+
+  _expected_full_data_next_dkv = np.array(
+      [
+          [0, 0, 0, 0],
+          [1, 1, 1, 1],
+          [2, 2, 2, 2],
+          [3, 3, 3, 3],
+      ],
+      dtype=np.int8,
+  )
+
+  # The mask_next array for a full mask is typically empty. The exception to
+  # this is when one head has a full mask and other heads have non-full masks.
+  # In that case the mask_next array is full, but none of its elements are
+  # actually relevant (they are never read).
+  def _expected_full_mask_next(self):
+    return np.array(
+        [
+            [-1, -1, -1, -1],
+            [-1, -1, -1, -1],
+            [-1, -1, -1, -1],
+            [-1, -1, -1, -1],
+        ],
+        dtype=np.int8,
+    )
+
+  _expected_full_mask_next_dkv = _expected_full_mask_next
+
+  _expected_causal_block_mask = np.array(
+      [
+          [1, 0, 0, 0],
+          [2, 1, 0, 0],
+          [2, 2, 1, 0],
+          [2, 2, 2, 1],
+      ],
+      dtype=np.int8,
+  )
+
+  _expected_causal_block_mask_dkv = _expected_causal_block_mask
+
+  _expected_causal_data_next = np.array(
+      [
+          [0, -1, -1, -1],
+          [0, 1, -1, -1],
+          [0, 1, 2, -1],
+          [0, 1, 2, 3],
+      ],
+      dtype=np.int8,
+  )
+
+  _expected_causal_data_next_dkv = np.array(
+      [
+          [0, -1, -1, -1],
+          [1, 1, -1, -1],
+          [2, 2, 2, -1],
+          [3, 3, 3, 3],
+      ],
+      dtype=np.int8,
+  )
+
+  def _expected_causal_mask_next(self, mask_base_index: int):
+    zero = mask_base_index
+    return np.array(
+        [
+            [zero, -1, -1, -1],
+            [-1, zero, -1, -1],
+            [-1, -1, zero, -1],
+            [-1, -1, -1, zero],
+        ],
+        dtype=np.int8,
+    )
+
+  _expected_causal_mask_next_dkv = _expected_causal_mask_next
+
+  _expected_local_block_mask = np.array(
+      [
+          [1, 1, 0, 0],
+          [1, 1, 1, 0],
+          [0, 1, 1, 1],
+          [0, 0, 1, 1],
+      ],
+      dtype=np.int8,
+  )
+
+  _expected_local_block_mask_dkv = _expected_local_block_mask
+
+  _expected_local_data_next = np.array(
+      [
+          [0, 1, -1, -1],
+          [0, 1, 2, -1],
+          [-1, 1, 2, 3],
+          [-1, -1, 2, 3],
+      ],
+      dtype=np.int8,
+  )
+
+  _expected_local_data_next_dkv = np.array(
+      [
+          [0, 0, -1, -1],
+          [1, 1, 1, -1],
+          [-1, 2, 2, 2],
+          [-1, -1, 3, 3],
+      ],
+      dtype=np.int8,
+  )
+
+  def _expected_local_mask_next(self, mask_base_index: int):
+    zero = mask_base_index
+    one = mask_base_index + 1
+    two = mask_base_index + 2
+    return np.array(
+        [
+            [zero, one, -1, -1],
+            [two, zero, one, -1],
+            [-1, two, zero, one],
+            [-1, -1, two, zero],
+        ],
+        dtype=np.int8,
+    )
+
+  _expected_local_mask_next_dkv = _expected_local_mask_next
+
+  def _stack(self, arrays: List[np.ndarray]) -> np.ndarray:
+    return np.stack(arrays, axis=0)
+
+  # For each test, check both the lazy and the dense versions of the mask.
+  @parameterized.parameters((True,), (False,))
+  def test_full_mask(self, is_lazy_mask: bool):
+    sequence_lengths = (64, 64)
+    block_shape = (16, 16)
+
+    if is_lazy_mask:
+      full_mask = mask_lib.MultiHeadMask((mask_lib.FullMask(sequence_lengths),))
+    else:
+      full_mask = mask_lib.MultiHeadMask((
+          mask_lib.NumpyMask(np.ones(sequence_lengths, dtype=np.bool_)),
+      ))
+
+    mask_info, mask_info_dkv, mask_function = self._process_mask(
+        full_mask, block_shape
+    )
+    self.assertIsNone(mask_function)
+
+    expected_mask_info = mask_info_lib.MaskInfo(
+        self._expected_full_data_next[None],
+        None,
+        self._expected_full_block_mask[None],
+        None,
+        None,
+    )
+
+    expected_mask_info_dkv = mask_info_lib.MaskInfo(
+        self._expected_full_data_next_dkv[None],
+        None,
+        self._expected_full_block_mask[None],
+        None,
+        None,
+    )
+
+    self._assert_mask_info_match(mask_info, expected_mask_info)
+    self._assert_mask_info_match(mask_info_dkv, expected_mask_info_dkv)
+
+  @parameterized.parameters((True,), (False,))
+  def test_two_causal_masks(self, is_lazy_mask: bool):
+    sequence_lengths = (64, 64)
+    block_shape = (16, 16)
+
+    if is_lazy_mask:
+      causal_mask = mask_lib.CausalMask(sequence_lengths)
+    else:
+      causal_mask = mask_lib.NumpyMask(
+          mask_lib.make_causal_mask(sequence_lengths)
+      )
+
+    multi_head = mask_lib.MultiHeadMask((causal_mask, causal_mask))
+
+    mask_info, mask_info_dkv, mask_function = self._process_mask(
+        multi_head, block_shape
+    )
+    if is_lazy_mask:
+      self.assertIsNotNone(mask_function)
+    else:
+      self.assertIsNone(mask_function)
+
+    expected_mask_info = mask_info_lib.MaskInfo(
+        self._expected_causal_data_next[None],
+        self._expected_causal_mask_next(0)[None] if not is_lazy_mask else None,
+        self._expected_causal_block_mask[None],
+        np.expand_dims(np.tril(np.ones(block_shape, dtype=np.int32)), 0)
+        if not is_lazy_mask
+        else None,
+        np.arange(sequence_lengths[0], dtype=np.int32)
+        if is_lazy_mask
+        else None,
+    )
+
+    expected_mask_info_dkv = mask_info_lib.MaskInfo(
+        self._expected_causal_data_next_dkv[None],
+        self._expected_causal_mask_next_dkv(0)[None]
+        if not is_lazy_mask
+        else None,
+        self._expected_causal_block_mask_dkv[None],
+        np.expand_dims(
+            np.tril(np.ones(block_shape, dtype=np.int32)), 0
+        ).swapaxes(-1, -2)
+        if not is_lazy_mask
+        else None,
+        np.arange(sequence_lengths[0], dtype=np.int32)
+        if is_lazy_mask
+        else None,
+    )
+
+    self._assert_mask_info_match(mask_info, expected_mask_info)
+    self._assert_mask_info_match(mask_info_dkv, expected_mask_info_dkv)
+
+  @parameterized.parameters((True,), (False,))
+  def test_rectangular_wide_causal_mask(self, is_lazy_mask: bool):
+    sequence_lengths = (64, 128)
+    block_shape = (16, 16)
+
+    if is_lazy_mask:
+      causal_mask = mask_lib.CausalMask(sequence_lengths)
+    else:
+      causal_mask = mask_lib.NumpyMask(
+          mask_lib.make_causal_mask(sequence_lengths)
+      )
+
+    multi_head = mask_lib.MultiHeadMask((causal_mask,))
+
+    mask_info, mask_info_dkv, mask_function = self._process_mask(
+        multi_head, block_shape
+    )
+    if is_lazy_mask:
+      self.assertIsNotNone(mask_function)
+    else:
+      self.assertIsNone(mask_function)
+
+    expected_mask_info = mask_info_lib.MaskInfo(
+        self._expected_causal_data_next[None],
+        self._expected_causal_mask_next(0)[None] if not is_lazy_mask else None,
+        self._expected_causal_block_mask[None],
+        np.expand_dims(np.tril(np.ones(block_shape, dtype=np.int32)), 0)
+        if not is_lazy_mask
+        else None,
+        np.arange(sequence_lengths[0], dtype=np.int32)
+        if is_lazy_mask
+        else None,
+    )
+
+    expected_causal_data_next_dkv = np.array(
+        [[
+            [0, -1, -1, -1, -1, -1, -1, -1],
+            [1, 1, -1, -1, -1, -1, -1, -1],
+            [2, 2, 2, -1, -1, -1, -1, -1],
+            [3, 3, 3, 3, -1, -1, -1, -1],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_causal_mask_next_dkv = np.array(
+        [[
+            [0, -1, -1, -1, -1, -1, -1, -1],
+            [-1, 0, -1, -1, -1, -1, -1, -1],
+            [-1, -1, 0, -1, -1, -1, -1, -1],
+            [-1, -1, -1, 0, -1, -1, -1, -1],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_causal_block_mask_dkv = np.array(
+        [[
+            [1, 0, 0, 0, 0, 0, 0, 0],
+            [2, 1, 0, 0, 0, 0, 0, 0],
+            [2, 2, 1, 0, 0, 0, 0, 0],
+            [2, 2, 2, 1, 0, 0, 0, 0],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_mask_info_dkv = mask_info_lib.MaskInfo(
+        expected_causal_data_next_dkv,
+        expected_causal_mask_next_dkv if not is_lazy_mask else None,
+        expected_causal_block_mask_dkv,
+        np.expand_dims(
+            np.tril(np.ones(block_shape, dtype=np.int32)), 0
+        ).swapaxes(-1, -2)
+        if not is_lazy_mask
+        else None,
+        np.arange(sequence_lengths[0], dtype=np.int32)
+        if is_lazy_mask
+        else None,
+    )
+
+    self._assert_mask_info_match(mask_info, expected_mask_info)
+    self._assert_mask_info_match(mask_info_dkv, expected_mask_info_dkv)
+
+  @parameterized.parameters((True,), (False,))
+  def test_rectangular_tall_causal_mask(self, is_lazy_mask: bool):
+    sequence_lengths = (128, 64)
+    block_shape = (16, 16)
+
+    if is_lazy_mask:
+      causal_mask = mask_lib.CausalMask(sequence_lengths)
+    else:
+      causal_mask = mask_lib.NumpyMask(
+          mask_lib.make_causal_mask(sequence_lengths)
+      )
+
+    multi_head = mask_lib.MultiHeadMask((causal_mask,))
+
+    mask_info, mask_info_dkv, mask_function = self._process_mask(
+        multi_head, block_shape
+    )
+    if is_lazy_mask:
+      self.assertIsNotNone(mask_function)
+    else:
+      self.assertIsNone(mask_function)
+
+    expected_causal_data_next = np.array(
+        [[
+            [0, -1, -1, -1],
+            [0, 1, -1, -1],
+            [0, 1, 2, -1],
+            [0, 1, 2, 3],
+            [0, 1, 2, 3],
+            [0, 1, 2, 3],
+            [0, 1, 2, 3],
+            [0, 1, 2, 3],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_causal_mask_next = np.array(
+        [[
+            [0, -1, -1, -1],
+            [-1, 0, -1, -1],
+            [-1, -1, 0, -1],
+            [-1, -1, -1, 0],
+            [-1, -1, -1, -1],
+            [-1, -1, -1, -1],
+            [-1, -1, -1, -1],
+            [-1, -1, -1, -1],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_causal_block_mask = np.array(
+        [[
+            [1, 0, 0, 0],
+            [2, 1, 0, 0],
+            [2, 2, 1, 0],
+            [2, 2, 2, 1],
+            [2, 2, 2, 2],
+            [2, 2, 2, 2],
+            [2, 2, 2, 2],
+            [2, 2, 2, 2],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_mask_info = mask_info_lib.MaskInfo(
+        expected_causal_data_next,
+        expected_causal_mask_next if not is_lazy_mask else None,
+        expected_causal_block_mask,
+        np.expand_dims(np.tril(np.ones(block_shape, dtype=np.int32)), 0)
+        if not is_lazy_mask
+        else None,
+        np.arange(sequence_lengths[0], dtype=np.int32)
+        if is_lazy_mask
+        else None,
+    )
+
+    expected_causal_data_next_dkv = np.array(
+        [[
+            [0, -1, -1, -1],
+            [1, 1, -1, -1],
+            [2, 2, 2, -1],
+            [3, 3, 3, 3],
+            [4, 4, 4, 4],
+            [5, 5, 5, 5],
+            [6, 6, 6, 6],
+            [7, 7, 7, 7],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_causal_mask_next_dkv = np.array(
+        [[
+            [0, -1, -1, -1],
+            [-1, 0, -1, -1],
+            [-1, -1, 0, -1],
+            [-1, -1, -1, 0],
+            [-1, -1, -1, -1],
+            [-1, -1, -1, -1],
+            [-1, -1, -1, -1],
+            [-1, -1, -1, -1],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_causal_block_mask_dkv = np.array(
+        [[
+            [1, 0, 0, 0],
+            [2, 1, 0, 0],
+            [2, 2, 1, 0],
+            [2, 2, 2, 1],
+            [2, 2, 2, 2],
+            [2, 2, 2, 2],
+            [2, 2, 2, 2],
+            [2, 2, 2, 2],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_mask_info_dkv = mask_info_lib.MaskInfo(
+        expected_causal_data_next_dkv,
+        expected_causal_mask_next_dkv if not is_lazy_mask else None,
+        expected_causal_block_mask_dkv,
+        np.expand_dims(
+            np.tril(np.ones(block_shape, dtype=np.int32)), 0
+        ).swapaxes(-1, -2)
+        if not is_lazy_mask
+        else None,
+        np.arange(sequence_lengths[0], dtype=np.int32)
+        if is_lazy_mask
+        else None,
+    )
+
+    self._assert_mask_info_match(mask_info, expected_mask_info)
+    self._assert_mask_info_match(mask_info_dkv, expected_mask_info_dkv)
+
+  @parameterized.parameters((True,), (False,))
+  def test_local_mask(self, is_lazy_mask: bool):
+    sequence_lengths = (64, 64)
+    block_shape = (16, 16)
+    window_size = 8
+    if is_lazy_mask:
+      local_mask = mask_lib.LocalMask(
+          sequence_lengths,
+          window_size=(window_size, window_size),
+          offset=0,
+      )
+    else:
+      local_mask = mask_lib.NumpyMask(
+          mask_lib.make_local_attention_mask(
+              sequence_lengths, window_size=(window_size, window_size), offset=0
+          )
+      )
+
+    multi_head = mask_lib.MultiHeadMask((local_mask,))
+
+    mask_info, mask_info_dkv, mask_function = self._process_mask(
+        multi_head, block_shape
+    )
+    self.assertIsNone(mask_function)
+
+    expected_partial_mask_blocks = self._stack(
+        [
+            np.triu(
+                np.tri(*block_shape, window_size, dtype=np.int32), -window_size
+            ),
+            np.tri(*block_shape, -window_size, dtype=np.int32),
+            np.triu(np.ones(block_shape, dtype=np.int32), window_size),
+        ],
+    )
+
+    expected_local_data_next = np.array(
+        [[
+            [0, 1, -1],
+            [0, 1, 2],
+            [1, 2, 3],
+            [2, 3, -1],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_local_mask_next = np.array(
+        [[
+            [0, 1, -1],
+            [2, 0, 1],
+            [2, 0, 1],
+            [2, 0, -1],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_local_block_mask = np.array(
+        [[
+            [1, 1, 0],
+            [1, 1, 1],
+            [1, 1, 1],
+            [1, 1, 0],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_mask_info = mask_info_lib.MaskInfo(
+        expected_local_data_next,
+        expected_local_mask_next,
+        expected_local_block_mask,
+        expected_partial_mask_blocks,
+        None,
+    )
+
+    expected_local_data_next_dkv = np.array(
+        [[
+            [-1, 0, 1, -1],
+            [0, 1, 2, 2],
+            [1, 2, 3, 3],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_local_mask_next_dkv = np.array(
+        [[
+            [-1, 1, 1, -1],
+            [0, 0, 0, 1],
+            [2, 2, 2, 0],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_local_block_mask_dkv = np.array(
+        [[
+            [0, 1, 1, 0],
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_mask_info_dkv = mask_info_lib.MaskInfo(
+        expected_local_data_next_dkv,
+        expected_local_mask_next_dkv,
+        expected_local_block_mask_dkv,
+        expected_partial_mask_blocks.swapaxes(-1, -2),
+        None,
+    )
+
+    self._assert_mask_info_match(mask_info, expected_mask_info)
+    self._assert_mask_info_match(mask_info_dkv, expected_mask_info_dkv)
+
+  @parameterized.parameters((True,), (False,))
+  def test_local_mask_narrow(self, is_lazy_mask: bool):
+    sequence_lengths = (64, 64)
+    block_shape = (16, 16)
+    window_size = 8
+    if is_lazy_mask:
+      local_mask = mask_lib.LocalMask(
+          sequence_lengths,
+          window_size=(window_size, 0),
+          offset=0,
+      )
+    else:
+      local_mask = mask_lib.NumpyMask(
+          mask_lib.make_local_attention_mask(
+              sequence_lengths, window_size=(window_size, 0), offset=0
+          )
+      )
+
+    multi_head = mask_lib.MultiHeadMask((local_mask,))
+
+    mask_info, mask_info_dkv, mask_function = self._process_mask(
+        multi_head, block_shape
+    )
+    self.assertIsNone(mask_function)
+
+    expected_partial_mask_blocks = self._stack(
+        [
+            np.triu(np.tri(*block_shape, 0, dtype=np.int32), -window_size),
+            np.triu(np.ones(block_shape, dtype=np.int32), window_size),
+        ],
+    )
+
+    expected_local_data_next = np.array(
+        [[
+            [0, -1],
+            [0, 1],
+            [1, 2],
+            [2, 3],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_local_mask_next = np.array(
+        [[
+            [0, -1],
+            [1, 0],
+            [1, 0],
+            [1, 0],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_local_block_mask = np.array(
+        [[
+            [1, 0],
+            [1, 1],
+            [1, 1],
+            [1, 1],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_mask_info = mask_info_lib.MaskInfo(
+        expected_local_data_next,
+        expected_local_mask_next,
+        expected_local_block_mask,
+        expected_partial_mask_blocks,
+        None,
+    )
+
+    expected_local_data_next_dkv = np.array(
+        [[
+            [0, 1, 2, -1],
+            [1, 2, 3, 3],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_local_mask_next_dkv = np.array(
+        [[
+            [0, 0, 0, -1],
+            [1, 1, 1, 0],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_local_block_mask_dkv = np.array(
+        [[
+            [1, 1, 1, 0],
+            [1, 1, 1, 1],
+        ]],
+        dtype=np.int8,
+    )
+
+    expected_mask_info_dkv = mask_info_lib.MaskInfo(
+        expected_local_data_next_dkv,
+        expected_local_mask_next_dkv,
+        expected_local_block_mask_dkv,
+        expected_partial_mask_blocks.swapaxes(-1, -2),
+        None,
+    )
+
+    self._assert_mask_info_match(mask_info, expected_mask_info)
+    self._assert_mask_info_match(mask_info_dkv, expected_mask_info_dkv)
+
+  @parameterized.parameters((True,), (False,))
+  def test_two_head_shards_one_causal_one_local(self, is_lazy_mask: bool):
+    sequence_lengths = (64, 64)
+    block_shape = (16, 16)
+    window_size = 8
+
+    if is_lazy_mask:
+      causal_mask = mask_lib.CausalMask(sequence_lengths)
+      local_mask = mask_lib.LocalMask(
+          sequence_lengths,
+          window_size=(window_size, window_size),
+          offset=0,
+      )
+    else:
+      causal_mask = mask_lib.NumpyMask(
+          mask_lib.make_causal_mask(sequence_lengths)
+      )
+      local_mask = mask_lib.NumpyMask(
+          mask_lib.make_local_attention_mask(
+              sequence_lengths, window_size=(window_size, window_size), offset=0
+          )
+      )
+
+    mask = mask_lib.MultiHeadMask((causal_mask, local_mask))
+
+    mask_info, mask_info_dkv, mask_function = self._process_mask(
+        mask, block_shape, head_shards=2
+    )
+    self.assertIsNone(mask_function)
+
+    expected_block_mask = self._stack(
+        [self._expected_causal_block_mask, self._expected_local_block_mask]
+    )
+    expected_data_next = self._stack(
+        [self._expected_causal_data_next, self._expected_local_data_next]
+    )
+    expected_mask_next = self._stack(
+        [self._expected_causal_mask_next(0), self._expected_local_mask_next(1)],
+    )
+
+    expected_partial_mask_blocks = self._stack([
+        np.tril(np.ones(block_shape, dtype=np.int32)),
+        np.triu(
+            np.tri(*block_shape, window_size, dtype=np.int32),
+            -window_size,
+        ),
+        np.tri(*block_shape, -window_size, dtype=np.int32),
+        np.triu(np.ones(block_shape, dtype=np.int32), window_size),
+    ])
+
+    expected_block_mask_dkv = self._stack(
+        [
+            self._expected_causal_block_mask_dkv,
+            self._expected_local_block_mask_dkv,
+        ],
+    )
+    expected_data_next_dkv = self._stack(
+        [
+            self._expected_causal_data_next_dkv,
+            self._expected_local_data_next_dkv,
+        ],
+    )
+    expected_mask_next_dkv = self._stack(
+        [
+            self._expected_causal_mask_next_dkv(0),
+            self._expected_local_mask_next_dkv(1),
+        ],
+    )
+
+    expected_mask_info = mask_info_lib.MaskInfo(
+        expected_data_next,
+        expected_mask_next,
+        expected_block_mask,
+        expected_partial_mask_blocks,
+        None,
+    )
+
+    expected_mask_info_dkv = mask_info_lib.MaskInfo(
+        expected_data_next_dkv,
+        expected_mask_next_dkv,
+        expected_block_mask_dkv,
+        expected_partial_mask_blocks.swapaxes(-1, -2),
+        None,
+    )
+
+    self._assert_mask_info_match(mask_info, expected_mask_info)
+    self._assert_mask_info_match(mask_info_dkv, expected_mask_info_dkv)
+
+  @parameterized.parameters((True,), (False,))
+  def test_two_head_shards_causal_full(self, is_lazy_mask: bool):
+    sequence_lengths = (64, 64)
+    block_shape = (16, 16)
+
+    if is_lazy_mask:
+      causal_mask = mask_lib.CausalMask(sequence_lengths)
+      full_mask = mask_lib.FullMask(sequence_lengths)
+    else:
+      causal_mask = mask_lib.NumpyMask(
+          mask_lib.make_causal_mask(sequence_lengths)
+      )
+      full_mask = mask_lib.NumpyMask(np.ones(sequence_lengths, dtype=np.bool_))
+
+    mask = mask_lib.MultiHeadMask((causal_mask, full_mask))
+
+    mask_info, mask_info_dkv, mask_function = self._process_mask(
+        mask, block_shape, head_shards=2
+    )
+    self.assertIsNone(mask_function)
+
+    expected_block_mask = self._stack(
+        [
+            self._expected_causal_block_mask,
+            self._expected_full_block_mask,
+        ],
+    )
+
+    expected_data_next = self._stack([
+        self._expected_causal_data_next,
+        self._expected_full_data_next,
+    ])
+
+    expected_mask_next = self._stack([
+        self._expected_causal_mask_next(0),
+        self._expected_full_mask_next(),
+    ])
+
+    expected_partial_mask_blocks = np.expand_dims(
+        np.tril(np.ones(block_shape, dtype=np.int32)), 0
+    )
+
+    expected_mask_info = mask_info_lib.MaskInfo(
+        expected_data_next,
+        expected_mask_next,
+        expected_block_mask,
+        expected_partial_mask_blocks,
+        None,
+    )
+
+    expected_block_mask_dkv = self._stack([
+        self._expected_causal_block_mask_dkv,
+        self._expected_full_block_mask_dkv,
+    ])
+    expected_data_next_dkv = self._stack(
+        [self._expected_causal_data_next_dkv, self._expected_full_data_next_dkv]
+    )
+
+    expected_mask_next_dkv = self._stack([
+        self._expected_causal_mask_next_dkv(0),
+        self._expected_full_mask_next_dkv(),
+    ])
+
+    expected_mask_info_dkv = mask_info_lib.MaskInfo(
+        expected_data_next_dkv,
+        expected_mask_next_dkv,
+        expected_block_mask_dkv,
+        expected_partial_mask_blocks.swapaxes(-1, -2),
+        None,
+    )
+
+    self._assert_mask_info_match(mask_info, expected_mask_info)
+    self._assert_mask_info_match(mask_info_dkv, expected_mask_info_dkv)
+
+  @parameterized.parameters((True,), (False,))
+  def test_two_qseq_shards_causal_local(self, is_lazy_mask: bool):
+    sequence_lengths = (64, 64)
+    block_shape = (16, 16)
+    window_size = 8
+
+    if is_lazy_mask:
+      causal_mask = mask_lib.CausalMask(sequence_lengths)
+      local_mask = mask_lib.LocalMask(
+          sequence_lengths,
+          window_size=(window_size, window_size),
+          offset=0,
+      )
+    else:
+      causal_mask = mask_lib.NumpyMask(
+          mask_lib.make_causal_mask(sequence_lengths)
+      )
+      local_mask = mask_lib.NumpyMask(
+          mask_lib.make_local_attention_mask(
+              sequence_lengths, window_size=(window_size, window_size), offset=0
+          )
+      )
+
+    mask = mask_lib.MultiHeadMask((causal_mask, local_mask))
+
+    mask_info, mask_info_dkv, mask_function = self._process_mask(
+        mask, block_shape, q_seq_shards=2
+    )
+    self.assertIsNone(mask_function)
+
+    expected_block_mask = self._stack(
+        [self._expected_causal_block_mask, self._expected_local_block_mask]
+    )
+    expected_data_next = self._stack(
+        [self._expected_causal_data_next, self._expected_local_data_next]
+    )
+    expected_mask_next = self._stack(
+        [self._expected_causal_mask_next(0), self._expected_local_mask_next(1)]
+    )
+
+    expected_partial_mask_blocks = self._stack([
+        np.tril(np.ones(block_shape, dtype=np.int32)),
+        np.triu(
+            np.tri(*block_shape, window_size, dtype=np.int32),
+            -window_size,
+        ),
+        np.tri(*block_shape, -window_size, dtype=np.int32),
+        np.triu(np.ones(block_shape, dtype=np.int32), window_size),
+    ])
+
+    expected_mask_info = mask_info_lib.MaskInfo(
+        expected_data_next,
+        expected_mask_next,
+        expected_block_mask,
+        expected_partial_mask_blocks,
+        None,
+    )
+
+    expected_block_mask_dkv = self._stack([
+        self._expected_causal_block_mask_dkv,
+        self._expected_local_block_mask_dkv,
+    ])
+    expected_data_next_dkv = np.array(
+        [
+            [
+                [0, -1, -1, -1],
+                [1, 1, -1, -1],
+                [0, 0, 0, -1],
+                [1, 1, 1, 1],
+            ],
+            [
+                [0, 0, -1, -1],
+                [1, 1, 1, -1],
+                [-1, 0, 0, 0],
+                [-1, -1, 1, 1],
+            ],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_mask_next_dkv = self._stack([
+        self._expected_causal_mask_next_dkv(0),
+        self._expected_local_mask_next_dkv(1),
+    ])
+
+    expected_mask_info_dkv = mask_info_lib.MaskInfo(
+        expected_data_next_dkv,
+        expected_mask_next_dkv,
+        expected_block_mask_dkv,
+        expected_partial_mask_blocks.swapaxes(-1, -2),
+        None,
+    )
+
+    self._assert_mask_info_match(mask_info, expected_mask_info)
+    self._assert_mask_info_match(mask_info_dkv, expected_mask_info_dkv)
+
+  def test_two_qseq_shards_causal_local_stacked(self):
+    sequence_lengths = (64, 64)
+    block_shape = (16, 16)
+    window_size = 8
+
+    causal_mask = mask_lib.make_causal_mask(sequence_lengths)
+    local_mask = mask_lib.make_local_attention_mask(
+        sequence_lengths, window_size=(window_size, window_size), offset=0
+    )
+    mask = np.concatenate((causal_mask, local_mask), axis=0)
+    mask = mask_lib.NumpyMask(mask)
+    mask = mask_lib.MultiHeadMask((mask,))
+
+    mask_info, mask_info_dkv, mask_function = self._process_mask(
+        mask, block_shape, q_seq_shards=2
+    )
+    self.assertIsNone(mask_function)
+
+    expected_local_block_mask = np.array(
+        [
+            [1, 1, 0, 0],
+            [1, 1, 1, 0],
+            [1, 1, 1, 0],
+            [1, 1, 0, 0],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_local_data_next = np.array(
+        [
+            [0, 1, -1, -1],
+            [0, 1, 2, -1],
+            [1, 2, 3, -1],
+            [2, 3, -1, -1],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_local_mask_next = np.array(
+        [
+            [1, 2, -1, -1],
+            [3, 1, 2, -1],
+            [3, 1, 2, -1],
+            [3, 1, -1, -1],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_block_mask = np.concatenate(
+        [self._expected_causal_block_mask, expected_local_block_mask],
+        axis=0,
+    )
+    expected_data_next = np.concatenate(
+        [self._expected_causal_data_next, expected_local_data_next],
+        axis=0,
+    )
+    expected_mask_next = np.concatenate(
+        [self._expected_causal_mask_next(0), expected_local_mask_next],
+        axis=0,
+    )
+
+    expected_partial_mask_blocks = self._stack([
+        np.tril(np.ones(block_shape, dtype=np.int32)),
+        np.triu(
+            np.tri(*block_shape, window_size, dtype=np.int32),
+            -window_size,
+        ),
+        np.tri(*block_shape, -window_size, dtype=np.int32),
+        np.triu(np.ones(block_shape, dtype=np.int32), window_size),
+    ])
+
+    expected_mask_info = mask_info_lib.MaskInfo(
+        expected_data_next[None],
+        expected_mask_next[None],
+        expected_block_mask[None],
+        expected_partial_mask_blocks,
+        None,
+    )
+
+    # TODO(amagni): this mask can be improved by bringing all the padding on one
+    # side.
+    expected_local_block_mask_dkv = np.array(
+        [
+            [0, 1, 1, 0],
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+            [0, 0, 0, 0],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_block_mask_dkv = np.concatenate([
+        self._expected_causal_block_mask_dkv,
+        expected_local_block_mask_dkv,
+    ])
+
+    expected_local_data_next_dkv = np.array(
+        [
+            [-1, 0, 1, -1],
+            [0, 1, 2, 2],
+            [1, 2, 3, 3],
+            [-1, -1, -1, -1],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_data_next_dkv = np.concatenate([
+        self._expected_causal_data_next_dkv,
+        expected_local_data_next_dkv,
+    ])
+
+    expected_local_mask_next_dkv = np.array(
+        [
+            [-1, 2, 2, -1],
+            [1, 1, 1, 2],
+            [3, 3, 3, 1],
+            [-1, -1, -1, -1],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_mask_next_dkv = np.concatenate([
+        self._expected_causal_mask_next_dkv(0),
+        expected_local_mask_next_dkv,
+    ])
+
+    expected_mask_info_dkv = mask_info_lib.MaskInfo(
+        expected_data_next_dkv[None],
+        expected_mask_next_dkv[None],
+        expected_block_mask_dkv[None],
+        expected_partial_mask_blocks.swapaxes(-1, -2),
+        None,
+    )
+
+    self._assert_mask_info_match(mask_info, expected_mask_info)
+    self._assert_mask_info_match(mask_info_dkv, expected_mask_info_dkv)
+
+  def test_two_qseq_shards_local_wide_local_narrow_stacked(self):
+    sequence_lengths = (64, 64)
+    block_shape = (16, 16)
+    window_size = 8
+
+    local_mask_wide = mask_lib.make_local_attention_mask(
+        sequence_lengths, window_size=(window_size, window_size), offset=0
+    )
+    local_mask_narrow = mask_lib.make_local_attention_mask(
+        sequence_lengths, window_size=(window_size, 0), offset=0
+    )
+
+    mask = np.concatenate((local_mask_wide, local_mask_narrow), axis=0)
+    mask = mask_lib.NumpyMask(mask)
+    mask = mask_lib.MultiHeadMask((mask,))
+
+    mask_info, mask_info_dkv, mask_function = self._process_mask(
+        mask, block_shape, q_seq_shards=2
+    )
+    self.assertIsNone(mask_function)
+
+    expected_local_wide_block_mask = np.array(
+        [
+            [1, 1, 0],
+            [1, 1, 1],
+            [1, 1, 1],
+            [1, 1, 0],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_local_wide_data_next = np.array(
+        [
+            [0, 1, -1],
+            [0, 1, 2],
+            [1, 2, 3],
+            [2, 3, -1],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_local_wide_mask_next = np.array(
+        [
+            [0, 1, -1],
+            [2, 0, 1],
+            [2, 0, 1],
+            [2, 0, -1],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_local_narrow_data_next = np.array(
+        [
+            [0, -1, -1],
+            [0, 1, -1],
+            [1, 2, -1],
+            [2, 3, -1],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_local_narrow_mask_next = np.array(
+        [
+            [3, -1, -1],
+            [2, 3, -1],
+            [2, 3, -1],
+            [2, 3, -1],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_local_narrow_block_mask = np.array(
+        [
+            [1, 0, 0],
+            [1, 1, 0],
+            [1, 1, 0],
+            [1, 1, 0],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_block_mask = np.concatenate(
+        [expected_local_wide_block_mask, expected_local_narrow_block_mask],
+        axis=0,
+    )
+    expected_data_next = np.concatenate(
+        [expected_local_wide_data_next, expected_local_narrow_data_next],
+        axis=0,
+    )
+    expected_mask_next = np.concatenate(
+        [expected_local_wide_mask_next, expected_local_narrow_mask_next],
+        axis=0,
+    )
+
+    expected_partial_mask_blocks = self._stack([
+        # Wide
+        np.triu(
+            np.tri(*block_shape, window_size, dtype=np.int32),
+            -window_size,
+        ),
+        np.tri(*block_shape, -window_size, dtype=np.int32),
+        np.triu(np.ones(block_shape, dtype=np.int32), window_size),
+        # Narrow
+        np.triu(np.tri(*block_shape, 0, dtype=np.int32), -window_size),
+    ])
+
+    expected_mask_info = mask_info_lib.MaskInfo(
+        expected_data_next[None],
+        expected_mask_next[None],
+        expected_block_mask[None],
+        expected_partial_mask_blocks,
+        None,
+    )
+
+    expected_local_wide_block_mask_dkv = np.array(
+        [
+            [0, 1, 1, 0],
+            [1, 1, 1, 1],
+            [1, 1, 1, 1],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_local_wide_data_next_dkv = np.array(
+        [
+            [-1, 0, 1, -1],
+            [0, 1, 2, 2],
+            [1, 2, 3, 3],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_local_wide_mask_next_dkv = np.array(
+        [
+            [-1, 1, 1, -1],
+            [0, 0, 0, 1],
+            [2, 2, 2, 0],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_local_narrow_data_next_dkv = np.array(
+        [
+            [0, 1, 2, -1],
+            [1, 2, 3, 3],
+            [-1, -1, -1, -1],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_local_narrow_mask_next_dkv = np.array(
+        [
+            [3, 3, 3, -1],
+            [2, 2, 2, 3],
+            [-1, -1, -1, -1],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_local_narrow_block_mask_dkv = np.array(
+        [
+            [1, 1, 1, 0],
+            [1, 1, 1, 1],
+            [0, 0, 0, 0],
+        ],
+        dtype=np.int8,
+    )
+
+    expected_block_mask_dkv = np.concatenate(
+        [
+            expected_local_wide_block_mask_dkv,
+            expected_local_narrow_block_mask_dkv,
+        ],
+        axis=0,
+    )
+
+    expected_data_next_dkv = np.concatenate(
+        [
+            expected_local_wide_data_next_dkv,
+            expected_local_narrow_data_next_dkv,
+        ],
+        axis=0,
+    )
+
+    expected_mask_next_dkv = np.concatenate(
+        [
+            expected_local_wide_mask_next_dkv,
+            expected_local_narrow_mask_next_dkv,
+        ],
+        axis=0,
+    )
+
+    expected_mask_info_dkv = mask_info_lib.MaskInfo(
+        expected_data_next_dkv[None],
+        expected_mask_next_dkv[None],
+        expected_block_mask_dkv[None],
+        expected_partial_mask_blocks.swapaxes(-1, -2),
+        None,
+    )
+
+    self._assert_mask_info_match(mask_info, expected_mask_info)
+    self._assert_mask_info_match(mask_info_dkv, expected_mask_info_dkv)
+
+  @parameterized.parameters((True,), (False,))
+  def test_two_head_shards_causal_mask(self, is_lazy_mask: bool):
+    sequence_lengths = (64, 64)
+    block_shape = (16, 16)
+
+    if is_lazy_mask:
+      causal_mask = mask_lib.CausalMask(sequence_lengths)
+    else:
+      causal_mask = mask_lib.NumpyMask(
+          mask_lib.make_causal_mask(sequence_lengths)
+      )
+
+    mask = mask_lib.MultiHeadMask((causal_mask, causal_mask))
+
+    mask_info, mask_info_dkv, mask_function = self._process_mask(
+        mask, block_shape, head_shards=2
+    )
+    if is_lazy_mask:
+      self.assertIsNotNone(mask_function)
+    else:
+      self.assertIsNone(mask_function)
+
+    expected_block_mask = self._stack(
+        [self._expected_causal_block_mask, self._expected_causal_block_mask]
+    )
+
+    expected_data_next = self._stack(
+        [self._expected_causal_data_next, self._expected_causal_data_next]
+    )
+
+    expected_mask_next = self._stack(
+        [self._expected_causal_mask_next(0), self._expected_causal_mask_next(0)]
+    )
+
+    expected_partial_mask_blocks = np.expand_dims(
+        np.tril(np.ones(block_shape, dtype=np.int32)), 0
+    )
+
+    expected_mask_info = mask_info_lib.MaskInfo(
+        expected_data_next,
+        expected_mask_next if not is_lazy_mask else None,
+        expected_block_mask,
+        expected_partial_mask_blocks if not is_lazy_mask else None,
+        np.arange(sequence_lengths[0], dtype=np.int32)
+        if is_lazy_mask
+        else None,
+    )
+
+    expected_block_mask_dkv = self._stack([
+        self._expected_causal_block_mask_dkv,
+        self._expected_causal_block_mask_dkv,
+    ])
+
+    expected_data_next_dkv = self._stack([
+        self._expected_causal_data_next_dkv,
+        self._expected_causal_data_next_dkv,
+    ])
+
+    expected_mask_next_dkv = self._stack([
+        self._expected_causal_mask_next_dkv(0),
+        self._expected_causal_mask_next_dkv(0),
+    ])
+
+    expected_mask_info_dkv = mask_info_lib.MaskInfo(
+        expected_data_next_dkv,
+        expected_mask_next_dkv if not is_lazy_mask else None,
+        expected_block_mask_dkv,
+        expected_partial_mask_blocks.swapaxes(-1, -2)
+        if not is_lazy_mask
+        else None,
+        np.arange(sequence_lengths[0], dtype=np.int32)
+        if is_lazy_mask
+        else None,
+    )
+
+    self._assert_mask_info_match(mask_info, expected_mask_info)
+    self._assert_mask_info_match(mask_info_dkv, expected_mask_info_dkv)
+
+  @parameterized.parameters((True,), (False,))
+  def test_two_head_shards_two_causal_two_local(self, is_lazy_mask: bool):
+    sequence_lengths = (64, 64)
+    block_shape = (16, 16)
+    window_size = 8
+
+    if is_lazy_mask:
+      causal_mask = mask_lib.CausalMask(sequence_lengths)
+      local_mask = mask_lib.LocalMask(
+          sequence_lengths,
+          window_size=(window_size, window_size),
+          offset=0,
+      )
+    else:
+      causal_mask = mask_lib.NumpyMask(
+          mask_lib.make_causal_mask(sequence_lengths)
+      )
+      local_mask = mask_lib.NumpyMask(
+          mask_lib.make_local_attention_mask(
+              sequence_lengths, window_size=(window_size, window_size), offset=0
+          )
+      )
+
+    mask = mask_lib.MultiHeadMask(
+        (causal_mask, causal_mask, local_mask, local_mask)
+    )
+
+    mask_info, mask_info_dkv, mask_function = self._process_mask(
+        mask, block_shape, head_shards=2
+    )
+    self.assertIsNone(mask_function)
+
+    expected_block_mask = self._stack(
+        [self._expected_causal_block_mask, self._expected_local_block_mask]
+    )
+
+    expected_data_next = self._stack(
+        [self._expected_causal_data_next, self._expected_local_data_next]
+    )
+
+    expected_mask_next = self._stack(
+        [self._expected_causal_mask_next(0), self._expected_local_mask_next(1)]
+    )
+
+    expected_partial_mask_blocks = self._stack(
+        [
+            np.tril(np.ones(block_shape, dtype=np.int32)),
+            np.triu(
+                np.tri(*block_shape, window_size, dtype=np.int32),
+                -window_size,
+            ),
+            np.tri(*block_shape, -window_size, dtype=np.int32),
+            np.triu(np.ones(block_shape, dtype=np.int32), window_size),
+        ],
+    )
+
+    expected_mask_info = mask_info_lib.MaskInfo(
+        expected_data_next,
+        expected_mask_next,
+        expected_block_mask,
+        expected_partial_mask_blocks,
+        None,
+    )
+
+    expected_block_mask_dkv = self._stack([
+        self._expected_causal_block_mask_dkv,
+        self._expected_local_block_mask_dkv,
+    ])
+
+    expected_data_next_dkv = self._stack([
+        self._expected_causal_data_next_dkv,
+        self._expected_local_data_next_dkv,
+    ])
+
+    expected_mask_next_dkv = self._stack([
+        self._expected_causal_mask_next_dkv(0),
+        self._expected_local_mask_next_dkv(1),
+    ])
+
+    expected_mask_info_dkv = mask_info_lib.MaskInfo(
+        expected_data_next_dkv,
+        expected_mask_next_dkv,
+        expected_block_mask_dkv,
+        expected_partial_mask_blocks.swapaxes(-1, -2),
+        None,
+    )
+
+    self._assert_mask_info_match(mask_info, expected_mask_info)
+    self._assert_mask_info_match(mask_info_dkv, expected_mask_info_dkv)
+
+  def test_huge_mask(self):
+    # Don't go too high with the mask size to avoid timeouts. Prefer covering
+    # multiple cases rather one very large one. This configuration replicates
+    # a realistic training shape. In particular, a large number of head shards
+    # and interleaving contribute to increasing processing time.
+    sequence_length = (32 * 1024, 32 * 1024)
+    block_shape = (512, 1024)
+
+    num_shards = 16
+    causal_mask = mask_lib.CausalMask(
+        sequence_length, 0, shard_count=num_shards
+    )
+
+    multi_head = mask_lib.MultiHeadMask((causal_mask,) * 64)
+
+    mask_info, mask_function = mask_info_lib.process_mask(
+        multi_head, block_shape, head_shards=8, q_seq_shards=16
+    )
+
+    self.assertIsNotNone(mask_function)
+    self.assertIsNotNone(mask_info.block_mask)
+    self.assertIsNotNone(mask_info.data_next)
+    self.assertIsNone(mask_info.mask_next)
+    self.assertIsNone(mask_info.partial_mask_blocks)
+    self.assertIsNotNone(mask_info.q_sequence)
+
+  def test_huge_mask2(self):
+    sequence_lengths = (32 * 1024, 32 * 1024)
+    block_shape = (1024, 1024)
+    window_size = 8
+
+    local_mask = mask_lib.LocalMask(
+        sequence_lengths,
+        window_size=(window_size, window_size),
+        offset=0,
+    )
+
+    multi_head = mask_lib.MultiHeadMask((local_mask,) * 32)
+
+    mask_info, mask_function = mask_info_lib.process_mask(
+        multi_head, block_shape
+    )
+
+    self.assertIsNone(mask_function)
+    self.assertIsNotNone(mask_info.block_mask)
+    self.assertIsNotNone(mask_info.data_next)
+    self.assertIsNotNone(mask_info.mask_next)
+    self.assertIsNotNone(mask_info.partial_mask_blocks)
+
+  def test_process_invalid_mask(self):
+    """Masks with of an all-0 row causes undefined softmax, reject them."""
+    sequence_length = 32
+
+    invalid_mask = np.ones(
+        (4, sequence_length, sequence_length), dtype=np.bool_
+    )
+    invalid_mask[2, 14, :] = False
+
+    invalid_mask = mask_lib.MultiHeadMask(
+        [mask_lib.NumpyMask(head_mask) for head_mask in invalid_mask]
+    )
+
+    with self.assertRaises(ValueError) as ctx:
+      for mask in invalid_mask.masks:
+        mask_info_lib._check_mask(mask)
+
+    self.assertIn("softmax", str(ctx.exception))
+
+
+if __name__ == "__main__":
+  absltest.main(testLoader=jtu.JaxTestLoader())