tpu-inference 0.11.1.dev202511150811__py3-none-any.whl

tpu_inference/kernels/flash_attention/kernel.py

@@ -0,0 +1,772 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Flash Attention TPU kernel."""
+from __future__ import annotations
+
+import dataclasses
+import functools
+import math
+from typing import Any, NamedTuple
+
+import jax
+import jax.numpy as jnp
+from jax import lax
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import tpu as pltpu
+
+DEFAULT_MASK_VALUE = -0.7 * float(jnp.finfo(jnp.dtype("float32")).max)
+NUM_LANES = 128
+NUM_SUBLANES = 8
+
+
+class SegmentIds(NamedTuple):
+    """SegmentIds for Q and KV sequences.
+
+  SegmentIds are used to generate segment mask, which prevents attention between
+  different segments in the input sequence. Each array is a list of ids
+  (integers).
+  Only the token with the same id can attend to each other.
+
+  Attributes:
+    q: segment ids along the Q sequence.
+    kv: segment ids along the KV sequence.
+  """
+
+    q: jax.Array  # [batch_size, q_seq_len]
+    kv: jax.Array  # [batch_size, kv_seq_len]
+
+
+@dataclasses.dataclass(frozen=True)
+class BlockSizes:
+    """Tile sizes parameterizing FlashAttention kernels.
+
+  Those parameters have negligible effect on numerics, but affect performance
+  greatly.
+  """
+    block_q: int
+    block_k_major: int
+    block_k: int
+    block_b: int
+
+    def __post_init__(self):
+
+        def verify_major_minor(prefix, suffix, major, minor):
+            if minor > major:
+                raise ValueError(
+                    f"{prefix}{suffix}={minor} should be smaller than"
+                    f" {prefix}_major{suffix}={major}")
+            if major % minor != 0:
+                raise ValueError(f"{prefix}{suffix}={minor} should divide"
+                                 f" {prefix}_major{suffix}={major}")
+
+        verify_major_minor("block_k", "", self.block_k_major, self.block_k)
+
+    @classmethod
+    def get_default(cls, batch_size, num_heads, q_seq_len, kv_len, d_model):
+        # TODO(apaszke,sharadmv): Select better parameters based on a heuristic.
+        del batch_size, num_heads, q_seq_len, kv_len, d_model  # Unused.
+        return BlockSizes(
+            block_q=128,
+            block_k_major=128,
+            block_k=128,
+            block_b=1,
+        )
+
+
+@functools.partial(
+    jax.jit,
+    static_argnames=[
+        "causal",
+        "sm_scale",
+        "block_sizes",
+        "vmem_limit_bytes",
+        "debug",
+    ],
+)
+def flash_attention(
+    q,  # [batch_size, num_heads, q_seq_len, d_model]
+    k,  # [batch_size, num_heads, kv_seq_len, d_model]
+    v,  # [batch_size, num_heads, kv_seq_len, d_model]
+    ab=None,  # [batch_size, num_heads, q_seq_len, kv_seq_len]
+    segment_ids=None,  # q of [batch_size, q_seq_len] and kv of [batch_size, kv_seq_len]
+    *,
+    causal: bool = False,
+    sm_scale: float = 1.0,
+    block_sizes: BlockSizes | None = None,
+    vmem_limit_bytes: int,
+    debug: bool = False,
+):
+    batch_size, num_heads, q_seq_len, d_model = q.shape
+    batch_size_k, num_heads_k, kv_seq_len, d_model_k = k.shape
+    batch_size_v, num_heads_v, kv_seq_len_v, d_model_v = v.shape
+    if batch_size != batch_size_k or batch_size != batch_size_v:
+        raise ValueError(
+            f"Batch size mismatch: got {batch_size}, {batch_size_k} and"
+            f" {batch_size_v} (for q, k, v respectively)")
+    if num_heads != num_heads_k or num_heads != num_heads_v:
+        raise ValueError(
+            f"Head count mismatch: got {num_heads}, {num_heads_k},"
+            f" {num_heads_v} (for q, k, v respectively)")
+    if d_model != d_model_k:
+        raise ValueError(
+            f"Model dimension mismatch: got {d_model} and {d_model_k} (for q and k"
+            " respectively)")
+    if d_model != d_model_v:
+        raise NotImplementedError(
+            "V model dimension unequal to KV model dimension unsupported")
+    if kv_seq_len != kv_seq_len_v:
+        raise ValueError(
+            f"KV sequence length mismatch: got {kv_seq_len} and {kv_seq_len_v}"
+        )
+    if ab is not None:
+        if ab.shape != (batch_size, num_heads, q_seq_len, kv_seq_len):
+            raise ValueError(
+                f"Attention bias shape mismatch: expected ({batch_size=},"
+                f" {num_heads=}, {q_seq_len=}, {kv_seq_len=}), got {ab.shape}")
+    if segment_ids is not None:
+        if segment_ids.q.shape != (batch_size, q_seq_len):
+            raise ValueError(
+                f"Q segment ids shape mismatch: expected ({batch_size=},"
+                f" {q_seq_len=},), got {segment_ids.q.shape}")
+        if segment_ids.kv.shape != (batch_size, kv_seq_len):
+            raise ValueError(
+                f"KV segment ids shape mismatch: expected ({batch_size=},"
+                f" {kv_seq_len=},), got {segment_ids.kv.shape}")
+    if block_sizes is None:
+        block_sizes = BlockSizes.get_default(batch_size, num_heads, q_seq_len,
+                                             kv_seq_len, d_model)
+        # TODO (KWang1998 & hfan): tune the block sizes properly.
+        if kv_seq_len <= 92800:
+            # Override block_k/block_k_major to use `_flash_attention_kernel_single_batch_single_step`.
+            block_sizes = BlockSizes(block_q=block_sizes.block_q,
+                                     block_b=block_sizes.block_b,
+                                     block_k_major=kv_seq_len,
+                                     block_k=kv_seq_len)
+    return _flash_attention(q, k, v, ab, segment_ids, False, causal, sm_scale,
+                            block_sizes, vmem_limit_bytes, debug)
+
+
+def _flash_attention(
+    q,
+    k,
+    v,
+    ab,
+    segment_ids,
+    save_residuals,
+    causal,
+    sm_scale,
+    block_sizes,
+    vmem_limit_bytes,
+    debug,
+):
+    return _flash_attention_impl(
+        q,
+        k,
+        v,
+        ab,
+        segment_ids,
+        save_residuals,
+        causal,
+        sm_scale,
+        block_sizes.block_b,
+        block_sizes.block_q,
+        block_sizes.block_k_major,
+        block_sizes.block_k,
+        vmem_limit_bytes,
+        debug,
+    )
+
+
+MIN_BLOCK_SIZE = 128
+TRANS_B_DIM_NUMBERS = (((1, ), (1, )), ((), ()))
+
+
+def below_or_on_diag(r, r_blk_size, c, c_blk_size):
+    # A block is considered below or on diagonal as long as the bottom left
+    # corner of the block is below or on diagonal.
+    return ((r + 1) * r_blk_size - 1) > (c * c_blk_size)
+
+
+def _flash_attention_kernel(q_tile_ref, *args, **kwargs):
+    block_b = q_tile_ref.shape[0]
+    # If we're not going to tile the softmax, then we can avoid a bunch of VPU ops.
+    if kwargs["block_k"] == kwargs["kv_seq_len"]:
+        kernel = _flash_attention_kernel_single_batch_single_step
+    else:
+        kernel = _flash_attention_kernel_single_batch
+    for batch_idx in range(block_b):
+        kernel((batch_idx, 0), q_tile_ref, *args, **kwargs)
+
+
+def _flash_attention_kernel_single_batch(
+    batch_idx: tuple[int, ...],
+    q_tile_ref,
+    k_tile_ref,
+    v_tile_ref,
+    ab_tile_ref,
+    q_segment_ids_tile_ref,
+    kv_segment_ids_tile_ref,  # Input arrays
+    o_tile_ref,  # Output arrays
+    l_ref,
+    m_ref,
+    m_scratch_ref,
+    l_scratch_ref,
+    acc_scratch_ref,
+    *,
+    causal,
+    sm_scale,
+    block_k,
+    kv_seq_len,
+    mask_value,
+):
+    block_k_major = k_tile_ref.shape[2]
+    block_q = q_tile_ref.shape[2]
+    head_dim = q_tile_ref.shape[-1]
+
+    kv_seq_idx = pl.program_id(3)
+
+    @pl.when(kv_seq_idx == 0)
+    def start_new_sequence():
+        m_scratch_ref[batch_idx] = jnp.full(m_scratch_ref.shape[2:], -jnp.inf,
+                                            jnp.float32)
+        l_scratch_ref[batch_idx] = jnp.zeros(l_scratch_ref.shape[2:],
+                                             jnp.float32)
+        acc_scratch_ref[batch_idx] = jnp.zeros(acc_scratch_ref.shape[2:],
+                                               jnp.float32)
+
+    q_seq_idx = pl.program_id(2)
+    if causal:
+        should_run = below_or_on_diag(q_seq_idx, block_q, kv_seq_idx,
+                                      block_k_major)
+    else:
+        should_run = True
+
+    @pl.when(should_run)
+    def run():
+
+        @pl.loop(0, block_k_major, step=block_k, unroll=True)
+        def _body(start_k):
+            m_prev = m_scratch_ref[batch_idx]
+            l_prev = l_scratch_ref[batch_idx]
+            q = q_tile_ref[batch_idx]  # [block_q, head_dim]
+            k = k_tile_ref[(*batch_idx, pl.dslice(start_k, block_k),
+                            slice(None))]  # [block_k, head_dim]
+
+            s = jax.lax.dot_general(
+                q, k, TRANS_B_DIM_NUMBERS,
+                preferred_element_type=jnp.float32)  # [block_q, block_k]
+
+            # Add attention bias if needed.
+            # TODO(tanburn) Should the attention bias be added before or after
+            # multiplication by sm_scale?
+            if ab_tile_ref is not None:
+                ab = ab_tile_ref[(*batch_idx, pl.dslice(None),
+                                  pl.dslice(start_k,
+                                            block_k))].astype(jnp.float32)
+                s += ab
+
+            if sm_scale != 1.0:
+                s *= sm_scale
+
+            mask = None
+            if q_segment_ids_tile_ref is not None:
+                repeats, rem = divmod(block_k, NUM_LANES)
+                if rem:
+                    raise NotImplementedError(
+                        f"kv block size must be a multiple of {NUM_LANES}")
+                q_segment_ids = pltpu.repeat(
+                    q_segment_ids_tile_ref[batch_idx[0]], repeats,
+                    axis=1)  # [block_q, block_k].
+                kv_segment_ids = kv_segment_ids_tile_ref[
+                    batch_idx[0], :1,
+                    pl.dslice(start_k, block_k)]  # [1, block_k].
+                mask = jnp.equal(q_segment_ids,
+                                 kv_segment_ids).astype(jnp.bool_)
+
+            if causal:
+                mask_shape = (block_q, block_k)
+                row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0)
+                row_ids += q_seq_idx * block_q
+                col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1)
+                col_ids += kv_seq_idx * block_k_major + start_k
+                causal_mask = col_ids <= row_ids
+                mask = (causal_mask if mask is None else jnp.logical_and(
+                    mask, causal_mask))
+
+            s = s if mask is None else s + jnp.where(mask, 0.0, mask_value)
+
+            m_curr = jnp.max(s, axis=1)[:,
+                                        None]  # Row max, shape [block_q, 1].
+            m_next = jnp.maximum(m_prev, m_curr)  # Shape [block_q, 128].
+
+            block_k_repeats, rem = divmod(block_k, MIN_BLOCK_SIZE)
+            if rem:
+                raise NotImplementedError(
+                    f"{block_k=} should be a multiple of {MIN_BLOCK_SIZE}")
+            p = jnp.exp(s - pltpu.repeat(m_next, block_k_repeats, 1))
+
+            alpha = jnp.exp(m_prev - m_next)  # Shape [block_q, 128].
+
+            l_corr = alpha * l_prev
+
+            l_next = jnp.sum(p, axis=1)[:,
+                                        None] + l_corr  # Shape [block_q, 128]
+
+            head_dim_repeats, rem = divmod(head_dim, MIN_BLOCK_SIZE)
+            l_broadcast = lambda l: pltpu.repeat(l, head_dim_repeats, 1)
+            if rem:
+                if head_dim_repeats == 0:
+                    l_broadcast = lambda l: l[:, :head_dim]
+                else:
+                    raise NotImplementedError(
+                        f"{head_dim=} should be a multiple of {MIN_BLOCK_SIZE} if larger"
+                    )
+            l_scratch_ref[batch_idx] = l_next
+            m_scratch_ref[batch_idx] = m_next
+
+            l_next_inv_safe = jnp.where(l_next == 0.0, 1.0, 1.0 / l_next)
+            acc_scratch_ref[batch_idx] *= l_broadcast(l_corr * l_next_inv_safe)
+            v = v_tile_ref[(*batch_idx, pl.dslice(start_k,
+                                                  block_k), slice(None))]
+            o_curr = jax.lax.dot(p.astype(v.dtype),
+                                 v,
+                                 preferred_element_type=jnp.float32)
+            acc_scratch_ref[batch_idx] += o_curr * l_broadcast(l_next_inv_safe)
+
+    @pl.when(kv_seq_idx == (kv_seq_len // block_k_major) - 1)
+    def store_output():
+        o_tile_ref[batch_idx] = acc_scratch_ref[batch_idx].astype(
+            o_tile_ref.dtype)
+        if l_ref is not None:
+            l_ref[batch_idx] = l_scratch_ref[batch_idx].astype(l_ref.dtype)
+        if m_ref is not None:
+            m_ref[batch_idx] = m_scratch_ref[batch_idx].astype(m_ref.dtype)
+
+
+# ruff: noqa #731
+# ruff: noqa #741
+def _flash_attention_kernel_single_batch_single_step(
+    batch_idx: tuple[int, ...],
+    q_tile_ref,
+    k_tile_ref,
+    v_tile_ref,
+    ab_tile_ref,
+    q_segment_ids_tile_ref,
+    kv_segment_ids_tile_ref,  # Input arrays
+    o_tile_ref,  # Output arrays
+    l_ref: Any | None = None,
+    m_ref: Any | None = None,
+    *,
+    causal,
+    sm_scale,
+    block_k,
+    kv_seq_len,
+    mask_value,
+):
+    block_k_major = k_tile_ref.shape[2]
+    block_q = q_tile_ref.shape[2]
+
+    assert kv_seq_len == block_k_major == block_k
+
+    q = q_tile_ref[batch_idx]  # [block_q, head_dim]
+    k = k_tile_ref[batch_idx]  # [block_k, head_dim]
+    s = jax.lax.dot_general(
+        q, k, TRANS_B_DIM_NUMBERS,
+        preferred_element_type=jnp.float32)  # [block_q, block_k]
+
+    if ab_tile_ref is not None:
+        s += ab_tile_ref[batch_idx].astype(jnp.float32)
+    if sm_scale != 1.0:
+        s *= sm_scale
+
+    mask = None
+    if q_segment_ids_tile_ref is not None:
+        repeats, rem = divmod(block_k, NUM_LANES)
+        if rem:
+            raise NotImplementedError(
+                f"kv block size must be a multiple of {NUM_LANES}")
+        q_segment_ids = q_segment_ids_tile_ref[
+            batch_idx[0]]  # [block_q, NUM_LANES].
+        q_segment_ids = pltpu.repeat(q_segment_ids, repeats,
+                                     axis=1)  # [block_q, block_k].
+        kv_segment_ids = kv_segment_ids_tile_ref[batch_idx[0], :
+                                                 1]  # [1, block_k].
+        mask = jnp.equal(q_segment_ids, kv_segment_ids).astype(jnp.bool_)
+
+    if causal:
+        q_seq_idx = pl.program_id(2)
+        mask_shape = (block_q, block_k)
+        row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0)
+        row_ids += q_seq_idx * block_q
+        col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1)
+        causal_mask = col_ids <= row_ids
+        mask = causal_mask if mask is None else jnp.logical_and(
+            mask, causal_mask)
+    s = s if mask is None else s + jnp.where(mask, 0.0, mask_value)
+
+    m = jnp.max(s, axis=1)[:, None]
+    p = jnp.exp(s - m)
+    l = jnp.sum(p, axis=1)[:, None]
+    p /= l
+
+    if m_ref is not None:
+        m_ref[batch_idx] = lax.broadcast_in_dim(m, m_ref.shape[2:], range(2))
+    if l_ref is not None:
+        l_ref[batch_idx] = lax.broadcast_in_dim(l, l_ref.shape[2:], range(2))
+
+    v = v_tile_ref[batch_idx]
+    o_tile_ref[batch_idx] = jax.lax.dot(
+        p.astype(v.dtype), v,
+        preferred_element_type=jnp.float32).astype(o_tile_ref.dtype)
+
+
+def _bytes(x: jax.Array | jax.ShapeDtypeStruct) -> int:
+    return math.prod(x.shape) * x.dtype.itemsize
+
+
+def _fwd_cost_estimate(
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    ab: jax.Array | None,
+    segment_ids: SegmentIds | None,
+    *,
+    causal: bool,
+    sm_scale: jax.Array | None,
+    kernel_inputs_specs,
+    kernel_outputs_specs,
+) -> pl.CostEstimate | None:
+    body_cost = pl.estimate_cost(mha_reference,
+                                 q,
+                                 k,
+                                 v,
+                                 ab,
+                                 segment_ids,
+                                 causal=causal,
+                                 sm_scale=sm_scale)
+    input_bytes = sum(_bytes(x) for x in jax.tree.leaves(kernel_inputs_specs))
+    output_bytes = sum(
+        _bytes(x) for x in jax.tree.leaves(kernel_outputs_specs))
+    return pl.CostEstimate(
+        flops=body_cost.flops,
+        transcendentals=body_cost.transcendentals,
+        bytes_accessed=input_bytes + output_bytes,
+    )
+
+
+def _flash_attention_impl(
+    q,
+    k,
+    v,
+    ab,
+    segment_ids,
+    save_residuals,
+    causal,
+    sm_scale,
+    block_b,
+    block_q,
+    block_k_major,
+    block_k,
+    vmem_limit_bytes,
+    debug,
+):
+    batch_size, num_heads, q_seq_len, head_dim = q.shape
+    _, _, kv_seq_len, _ = k.shape
+    _verify_block("block_q",
+                  "q_seq_len",
+                  block_q,
+                  q_seq_len,
+                  should_divide=False)
+    _verify_block("block_k_major", "kv_seq_len", block_k_major, kv_seq_len)
+    _verify_block("block_k", "kv_seq_len", block_k, kv_seq_len)
+    _verify_block("block_b", "batch", block_b, batch_size, should_divide=False)
+
+    # TODO(apaszke): Tile over heads as well.
+    grid = (
+        pl.cdiv(batch_size, block_b),
+        num_heads,
+        pl.cdiv(q_seq_len, block_q),
+        kv_seq_len // block_k_major,
+    )
+
+    def q_index_map(batch_index, head_index, q_seq_index, _):
+        return (batch_index, head_index, q_seq_index, 0)
+
+    def kv_index_map(batch_index, head_index, q_seq_index, kv_seq_index):
+        if causal:
+            # If the kv block is skipped, prefetch the next valid kv block, i.e. the
+            # 0th one to be used for the next block_q rows.
+            next_kv_index = lax.select(
+                below_or_on_diag(q_seq_index, block_q, kv_seq_index,
+                                 block_k_major),
+                kv_seq_index,
+                0,
+            )
+        else:
+            next_kv_index = kv_seq_index
+        return (batch_index, head_index, next_kv_index, 0)
+
+    def ab_index_map(batch_index, head_index, q_seq_index, kv_seq_index):
+        if causal:
+            should_run = below_or_on_diag(q_seq_index, block_q, kv_seq_index,
+                                          block_k_major)
+            # If the ab block is skipped, prefetch the next valid ab block, i.e. the
+            # 0th kv to be used for the next block_q rows.
+            next_q_index = lax.select(
+                should_run,
+                q_seq_index,
+                lax.select(q_seq_index == (q_seq_len // block_q) - 1, 0,
+                           q_seq_index + 1),
+            )
+            next_kv_index = lax.select(should_run, kv_seq_index, 0)
+        else:
+            next_q_index = q_seq_index
+            next_kv_index = kv_seq_index
+
+        return (batch_index, head_index, next_q_index, next_kv_index)
+
+    def o_index_map(batch_index, head_index, q_seq_index, _):
+        return (batch_index, head_index, q_seq_index, 0)
+
+    def lm_index_map(batch_index, head_index, q_seq_index, _):
+        return (batch_index, head_index, q_seq_index, 0)
+
+    kernel = functools.partial(
+        _flash_attention_kernel,
+        causal=causal,
+        mask_value=DEFAULT_MASK_VALUE,
+        sm_scale=sm_scale,
+        block_k=block_k,
+        kv_seq_len=kv_seq_len,
+    )
+    out_shape = jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype)
+    out_shape = [out_shape]
+    out_specs = [pl.BlockSpec((block_b, 1, block_q, head_dim), o_index_map)]
+
+    if block_k != kv_seq_len:
+        m_scratch = pltpu.VMEM((block_b, 1, block_q, MIN_BLOCK_SIZE),
+                               jnp.float32)
+        l_scratch = pltpu.VMEM((block_b, 1, block_q, MIN_BLOCK_SIZE),
+                               jnp.float32)
+        acc_scratch = pltpu.VMEM((block_b, 1, block_q, head_dim), jnp.float32)
+        scratch_shapes = [m_scratch, l_scratch, acc_scratch]
+    else:
+        scratch_shapes = []
+
+    if save_residuals:
+        out_specs = [
+            *out_specs,
+            pl.BlockSpec((block_b, 1, block_q, MIN_BLOCK_SIZE), lm_index_map),
+            pl.BlockSpec((block_b, 1, block_q, MIN_BLOCK_SIZE), lm_index_map),
+        ]
+        l = jax.ShapeDtypeStruct(
+            (batch_size, num_heads, q_seq_len, MIN_BLOCK_SIZE),
+            dtype=jnp.float32)
+        m = jax.ShapeDtypeStruct(
+            (batch_size, num_heads, q_seq_len, MIN_BLOCK_SIZE),
+            dtype=jnp.float32)
+        out_shape = (*out_shape, l, m)
+    else:
+        out_specs = [*out_specs, None, None]
+        out_shape = (*out_shape, None, None)
+
+    ab_block_spec = (pl.BlockSpec(
+        (block_b, 1, block_q,
+         block_k_major), ab_index_map) if ab is not None else None)
+
+    q_segment_ids_spec = kv_segment_ids_spec = None
+    q_segment_ids = kv_segment_ids = None
+    if segment_ids is not None:
+
+        def q_segment_ids_index_map(batch_index, head_index, q_seq_index, _):
+            del head_index
+            return (batch_index, q_seq_index, 0)
+
+        def kv_segment_ids_index_map(batch_index, head_index, q_seq_index,
+                                     kv_seq_index):
+            del head_index
+            if causal:
+                next_kv_index = lax.select(
+                    below_or_on_diag(q_seq_index, block_q, kv_seq_index,
+                                     block_k_major),
+                    kv_seq_index,
+                    0,
+                )
+            else:
+                next_kv_index = kv_seq_index
+            return (batch_index, 0, next_kv_index)
+
+        q_segment_ids_spec = pl.BlockSpec((block_b, block_q, NUM_LANES),
+                                          q_segment_ids_index_map)
+        kv_segment_ids_spec = pl.BlockSpec(
+            (block_b, NUM_SUBLANES, block_k_major), kv_segment_ids_index_map)
+
+        q_segment_ids = jax.lax.broadcast_in_dim(
+            segment_ids.q,
+            (batch_size, q_seq_len, NUM_LANES),
+            (
+                0,
+                1,
+            ),
+        )
+        kv_segment_ids = jax.lax.broadcast_in_dim(
+            segment_ids.kv,
+            (batch_size, NUM_SUBLANES, kv_seq_len),
+            (
+                0,
+                2,
+            ),
+        )
+
+    in_specs = [
+        pl.BlockSpec((block_b, 1, block_q, head_dim), q_index_map),
+        pl.BlockSpec((block_b, 1, block_k_major, head_dim), kv_index_map),
+        pl.BlockSpec((block_b, 1, block_k_major, head_dim), kv_index_map),
+        ab_block_spec,
+        q_segment_ids_spec,
+        kv_segment_ids_spec,
+    ]
+
+    o, *aux = pl.pallas_call(
+        kernel,
+        grid_spec=pltpu.PrefetchScalarGridSpec(
+            num_scalar_prefetch=0,
+            grid=grid,
+            in_specs=in_specs,
+            out_specs=out_specs,
+            scratch_shapes=scratch_shapes,
+        ),
+        out_shape=out_shape,
+        debug=debug,
+        compiler_params=pltpu.CompilerParams(
+            dimension_semantics=(
+                "parallel",
+                "parallel",
+                "parallel",
+                "arbitrary",
+            ),
+            vmem_limit_bytes=vmem_limit_bytes,
+        ),
+        cost_estimate=_fwd_cost_estimate(
+            q,
+            k,
+            v,
+            ab,
+            segment_ids,
+            causal=causal,
+            sm_scale=sm_scale,
+            kernel_inputs_specs=(q, k, v, ab, q_segment_ids, kv_segment_ids),
+            kernel_outputs_specs=out_shape,
+        ),
+    )(q, k, v, ab, q_segment_ids, kv_segment_ids)
+    if save_residuals:
+        l, m = (v[..., 0] for v in aux[-2:])
+        return (o, l, m)
+    else:
+        return o
+
+
+# For autograd testing.
+def mha_reference_no_custom_vjp(
+    q,
+    k,
+    v,
+    ab: jax.Array | None = None,
+    segment_ids: SegmentIds | None = None,
+    *,
+    causal: bool = False,
+    mask_value: float = DEFAULT_MASK_VALUE,
+    sm_scale: float = 1.0,
+    save_residuals: bool = False,
+):
+    logits = jnp.einsum("bhqc,bhkc->bhqk", q, k)
+    if ab is not None:
+        logits += ab
+    if sm_scale != 1.0:
+        logits *= sm_scale
+
+    mask = None
+    if segment_ids is not None:
+        mask = segment_ids.q[:, :, None] == segment_ids.kv[:, None, :]
+        mask = mask[:, None, :, :]
+
+    if causal:
+        _, _, q_seq_len, _ = q.shape
+        _, _, kv_seq_len, _ = k.shape
+        mask_shape = (q_seq_len, kv_seq_len)
+        row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0)
+        col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1)
+        causal_mask = (col_ids <= row_ids)[None, None, :, :]
+        mask = causal_mask if mask is None else jnp.logical_and(
+            mask, causal_mask)
+
+    logits = logits if mask is None else logits + jnp.where(
+        mask, 0.0, mask_value)
+
+    m = logits.max(axis=-1)
+    unnormalized = jnp.exp(logits - m[..., None])
+    l = unnormalized.sum(axis=-1)
+    weights = unnormalized / l[..., None]
+    out = jnp.einsum("bhqk,bhkc->bhqc", weights, v)
+    if save_residuals:
+        return out, l, m
+    return out
+
+
+@functools.partial(jax.jit,
+                   static_argnames=["causal", "mask_value", "sm_scale"])
+@jax.default_matmul_precision("bfloat16")
+def mha_reference(
+    q,
+    k,
+    v,
+    ab,
+    segment_ids: SegmentIds | None = None,
+    causal: bool = False,
+    mask_value: float = DEFAULT_MASK_VALUE,
+    sm_scale=1.0,
+):
+    return _mha_reference(
+        q,
+        k,
+        v,
+        ab,
+        segment_ids,
+        causal=causal,
+        mask_value=mask_value,
+        sm_scale=sm_scale,
+        save_residuals=False,
+    )
+
+
+def _mha_reference(
+    q,
+    k,
+    v,
+    ab,
+    segment_ids: SegmentIds | None,
+    causal: bool,
+    mask_value: float,
+    sm_scale: float,
+    save_residuals: bool,
+):
+    return mha_reference_no_custom_vjp(
+        q,
+        k,
+        v,
+        ab,
+        segment_ids,
+        causal=causal,
+        mask_value=mask_value,
+        sm_scale=sm_scale,
+        save_residuals=save_residuals,
+    )
+
+
+def _verify_block(block_name, dim_name, block, dim, should_divide=True):
+    if block > dim:
+        raise ValueError(
+            f"{block_name}={block} should be smaller or equal to {dim_name}={dim}"
+        )
+    if should_divide and dim % block != 0:
+        raise ValueError(
+            f"{dim_name}={dim} should be divisible by {block_name}={block}")