diff --git a/JAXBench/benchmark/level2/11p_Megablox_GMM/baseline.py b/JAXBench/benchmark/level2/11p_Megablox_GMM/baseline.py new file mode 100644 index 0000000..d7c411c --- /dev/null +++ b/JAXBench/benchmark/level2/11p_Megablox_GMM/baseline.py @@ -0,0 +1,881 @@ +# Copyright 2024 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. + +"""Pallas grouped matrix multiplication kernel — Qwen3-235B-A22B MoE dimensions. + +Upstream kernel from jax.experimental.pallas.ops.tpu.megablox.gmm, wrapped as a +JAXBench workload with CONFIG / create_inputs / workload. + +Metadata/utilities imported from the installed JAX package (not optimizable). +""" + +import numpy as np +import time +from collections.abc import Callable +import functools +from typing import Any, Optional + +import jax +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +from jax.experimental.pallas.ops.tpu.megablox import common +import jax.numpy as jnp + + +partial = functools.partial + + +def _validate_args( + *, + lhs: jnp.ndarray, + rhs: jnp.ndarray, + group_sizes: jnp.ndarray, + expected_rhs_dims: int = 3, +) -> tuple[jnp.ndarray, jnp.ndarray, jnp.dtype]: + """Validates the arguments for the gmm function.""" + # Validate 'lhs'. + if lhs.ndim != 2: + raise ValueError(f"Expected 2-tensor for 'lhs' but got {lhs.ndim}-tensor.") + common.assert_is_supported_dtype(lhs.dtype) + + # Validate 'rhs'. + if rhs.ndim != expected_rhs_dims: + raise ValueError( + f"Expected {expected_rhs_dims}-tensor for 'rhs' but got" + f" {rhs.ndim}-tensor." + ) + common.assert_is_supported_dtype(rhs.dtype) + + # Validate 'group_sizes'. + if group_sizes.dtype != jnp.int32: + raise ValueError( + f"Expected 32-bit integer 'group_sizes' but got {group_sizes.dtype}." + ) + + return lhs, group_sizes, common.select_input_dtype(lhs, rhs) + + +def _calculate_num_tiles(x: int, tx: int) -> int: + tiles, rem = divmod(x, tx) + if rem: + raise ValueError(f"{x} must be divisible by x-dimension tile size ({tx}).") + return tiles + + +def _calculate_irregular_num_tiles(x: int, tx: int) -> tuple[int, int]: + tiles, rem = divmod(x, tx) + if rem: + tiles += 1 + return tiles, rem + + +GroupMetadata = Any # TODO(enriqueps): Clean this up and use a namedtuple + + +def make_group_metadata( + *, + group_sizes: jnp.ndarray, + m: int, + tm: int, + start_group: jnp.ndarray, + num_nonzero_groups: int, + visit_empty_groups: bool = True, +) -> GroupMetadata: + """Create the metadata needed for grouped matmul computation. + + Args: + group_sizes: A 1d, jnp.ndarray with shape [num_groups] and jnp.int32 dtype. + m: The number of rows in lhs. + tm: The m-dimension tile size being used. + start_group: The group in group sizes to start computing from. This is + particularly useful for when rhs num_groups is sharded. + num_nonzero_groups: Number of groups in group sizes to compute on. Useful in + combination with group_offset. + visit_empty_groups: If True, do not squeeze tiles for empty groups out of + the metadata. This is necessary for tgmm, where we at least need to zero + the output for each group. + + Returns: + tuple of: + group_offsets: A 1d, jnp.ndarray with shape [num_groups+1] and jnp.int32 + dtype. group_offsets[i] indicates the row at which group [i] starts in + the lhs matrix and group_offsets[i-1] = m. + group_ids: A 1d, jnp.ndarray with shape [m_tiles + num_groups] and + jnp.int32 dtype. group_ids[i] indicates which group grid index 'i' will + work on. + m_tile_ids: A 1d, jnp.ndarray with shape [m_tiles + num_groups] and + jnp.int32. m_tile_ids[i] indicates which m-dimension tile grid index 'i' + will work on. + num_tiles: The number of m-dimension tiles to execute. + """ + num_groups = group_sizes.shape[0] + end_group = start_group + num_nonzero_groups - 1 + + # Calculate the offset of each group, starting at zero. This metadata is + # similar to row offsets in a CSR matrix. The following properties hold: + # + # group_offsets.shape = [num_groups + 1] + # group_offsets[0] = 0 + # group_offsets[num_groups] = m + # + # The row at which group 'i' starts is group_offsets[i]. + group_ends = jnp.cumsum(group_sizes) + group_offsets = jnp.concatenate([jnp.zeros(1, dtype=jnp.int32), group_ends]) + + # Assign a group id to each grid index. + # + # If a group starts somewhere other than the start of a tile or ends somewhere + # other than the end of a tile we need to compute that full tile. Calculate + # the number of tiles for each group by rounding their end up to the nearest + # 'tm' and their start down to the nearest 'tm'. + + # (1) Round the group_ends up to the nearest multiple of 'tm'. + # + # NOTE: This does not change group_offsets[num_groups], which is m + # (because we enforce m is divisible by tm). + rounded_group_ends = ((group_ends + tm - 1) // tm * tm).astype(jnp.int32) + + # (2) Round the group_starts down to the nearest multiple of 'tm'. + group_starts = jnp.concatenate( + [jnp.zeros(1, dtype=jnp.int32), group_ends[:-1]] + ) + rounded_group_starts = group_starts // tm * tm + + # (3) Calculate the number of rows in each group. + # + # NOTE: Handle zero-sized groups as a special case. If the start for a + # zero-sized group is not divisible by 'tm' its start will be rounded down and + # its end will be rounded up such that its size will become 1 tile here. + rounded_group_sizes = rounded_group_ends - rounded_group_starts + rounded_group_sizes = jnp.where(group_sizes == 0, 0, rounded_group_sizes) + + # (4) Convert the group sizes from units of rows to unit of 'tm' sized tiles. + # + # An m-dimension tile is 'owned' by group 'i' if the first row of the tile + # belongs to group 'i'. In addition to owned tiles, each group can have 0 or 1 + # initial partial tiles if it's first row does not occur in the first row of a + # tile. The '0-th' group never has a partial tile because it always starts at + # the 0-th row. + # + # If no group has a partial tile, the total number of tiles is equal to + # 'm // tm'. If every group has a partial except the 0-th group, the total + # number of tiles is equal to 'm // tm + num_groups - 1'. Thus we know that + # + # tiles_m <= group_tiles.sum() <= tiles_m + num_groups - 1 + # + # Where tiles_m = m // tm. + # + # NOTE: All group sizes are divisible by 'tm' because of the rounding in steps + # (1) and (2) so this division is exact. + group_tiles = rounded_group_sizes // tm + + if visit_empty_groups: + # Insert one tile for empty groups. + group_tiles = jnp.where(group_sizes == 0, 1, group_tiles) + + # Create the group ids for each grid index based on the tile counts for each + # group. + # + # NOTE: This repeat(...) will pad group_ids with the final group id if + # group_tiles.sum() < tiles_m + num_groups - 1. The kernel grid will be sized + # such that we only execute the necessary number of tiles. + tiles_m = _calculate_num_tiles(m, tm) + group_ids = jnp.repeat( + jnp.arange(num_groups, dtype=jnp.int32), + group_tiles, + total_repeat_length=tiles_m + num_groups - 1, + ) + + # Assign an m-dimension tile id to each grid index. + # + # NOTE: Output tiles can only be re-visited consecutively. The following + # procedure guarantees that m-dimension tile indices respect this. + + # (1) Calculate how many times each m-dimension tile will be visited. + # + # Each tile is guaranteed to be visited once by the group that owns the tile. + # The remaining possible visits occur when a group starts inside of a tile at + # a position other than the first row. We can calculate which m-dimension tile + # each group starts in by floor-dividing its offset with `tm` and then count + # tile visits with a histogram. + # + # To avoid double counting tile visits from the group that owns the tile, + # filter these out by assigning their tile id to `tile_m` (one beyond the max) + # such that they're ignored by the subsequent histogram. Also filter out any + # group which is empty. + # + # TODO(tgale): Invert the 'partial_tile_mask' predicates to be more clear. + partial_tile_mask = jnp.logical_or( + (group_offsets[:-1] % tm) == 0, group_sizes == 0 + ) + + # Explicitly enable tiles for zero sized groups, if specified. This covers + # zero sized groups that start on a tile-aligned row and those that do not. + if visit_empty_groups: + partial_tile_mask = jnp.where(group_sizes == 0, 0, partial_tile_mask) + + partial_tile_ids = jnp.where( + partial_tile_mask, tiles_m, group_offsets[:-1] // tm + ) + + tile_visits = ( + jnp.histogram(partial_tile_ids, bins=tiles_m, range=(0, tiles_m - 1))[0] + + 1 + ) + + # Create the m-dimension tile ids for each grid index based on the visit + # counts for each tile. + m_tile_ids = jnp.repeat( + jnp.arange(tiles_m, dtype=jnp.int32), + tile_visits.astype(jnp.int32), + total_repeat_length=tiles_m + num_groups - 1, + ) + + # Account for sharding. + # + # Find the start of the groups owned by our shard and shift the group_ids and + # m_tile_ids s.t. the metadata for our tiles are at the front of the arrays. + # + # TODO(tgale): Move this offset into the kernel to avoid these rolls. + first_tile_in_shard = (group_ids < start_group).sum() + group_ids = jnp.roll(group_ids, shift=-first_tile_in_shard, axis=0) + m_tile_ids = jnp.roll(m_tile_ids, shift=-first_tile_in_shard, axis=0) + + # Calculate the number of tiles we need to compute for our shard. + # + # Remove tile visits that belong to a group not in our shard. + iota = jnp.arange(num_groups, dtype=jnp.int32) + active_group_mask = jnp.logical_and(iota <= end_group, iota >= start_group) + group_tiles = jnp.where(active_group_mask, group_tiles, 0) + num_tiles = group_tiles.sum() + return (group_offsets, group_ids, m_tile_ids), num_tiles + + +def _get_group_size( + *, grid_id: jnp.ndarray, group_metadata: GroupMetadata +) -> jnp.ndarray: + """Calculate the number of rows in the current group.""" + group_offsets, group_ids = group_metadata[:2] + group_id = group_ids[grid_id] + group_start = group_offsets[group_id] + group_end = group_offsets[group_id + 1] + return group_end - group_start + + +def _get_store_mask( + *, + grid_id: jnp.ndarray, + group_metadata: GroupMetadata, + tm: int, + tn: int, +) -> jnp.ndarray: + """Mask for rows that belong to the current group in the current tile.""" + group_offsets, group_ids, m_tile_ids = group_metadata[:3] + group_id = group_ids[grid_id] + group_start = group_offsets[group_id] + group_end = group_offsets[group_id + 1] + m_id = m_tile_ids[grid_id] * tm + iota = jax.lax.broadcasted_iota(jnp.int32, (tm, tn), 0) + m_id + return jnp.logical_and(iota >= group_start, iota < group_end) + + +def _zero_uninitialized_memory( + out: jnp.ndarray, + *, + start_group: jnp.ndarray, + num_nonzero_groups: int, + group_metadata: GroupMetadata, +) -> jnp.ndarray: + """Zero out uninitialized memory from output.""" + group_offsets = group_metadata[0] + group_start = group_offsets[start_group] + group_end = group_offsets[start_group + num_nonzero_groups] + valid_mask = jax.lax.broadcasted_iota(jnp.int32, (out.shape[0],), 0) + valid_mask = (valid_mask >= group_start) & (valid_mask < group_end) + return jnp.where(valid_mask[:, None], out, 0) + + +LutFn = Callable[[int, int, int], Optional[tuple[int, int, int]]] + + +@functools.partial( + jax.jit, + static_argnames=[ + "preferred_element_type", + "tiling", + "transpose_rhs", + "interpret", + ], +) +def gmm( + lhs: jnp.ndarray, + rhs: jnp.ndarray, + group_sizes: jnp.ndarray, + preferred_element_type: jnp.dtype = jnp.float32, + tiling: tuple[int, int, int] | LutFn | None = (128, 128, 128), + group_offset: jnp.ndarray | None = None, + existing_out: jnp.ndarray | None = None, + transpose_rhs: bool = False, + interpret: bool = False, +) -> jnp.ndarray: + """Compute lhs[sizes[i-1]:sizes[i], :] @ rhs for each group 'i'. + + Args: + lhs: A 2d, jnp.ndarray with shape [m, k]. + rhs: A 3d, jnp.ndarray with shape [num_groups, k, n]. + group_sizes: A 1d, jnp.ndarray with shape [num_groups] and jnp.int32 dtype. + preferred_element_type: jnp.dtype, the element type for the output matrix. + tiling: 3-tuple of ints. The m, k and n-dimension tile sizes. + group_offset: The group in group sizes to start computing from. This is + particularly useful for when rhs num_groups is sharded. + existing_out: Existing output to write to. + transpose_rhs: True if the rhs needs to be transposed. + interpret: Whether or not to run the kernel in interpret mode, helpful for + testing and debugging. + + Returns: + A 2d, jnp.ndarray with shape [m, n]. + """ + + if existing_out is not None: + assert isinstance(existing_out, jax.Array) + expected_dtype = existing_out.dtype + if expected_dtype != preferred_element_type: + raise ValueError( + "Existing output dtype must match preferred_element_type." + ) + if group_offset is None: + group_offset = jnp.array([0], dtype=jnp.int32) + else: + if group_offset.shape: + raise ValueError( + f"group_offset must be a ()-shaped array. Got: {group_offset.shape}." + ) + group_offset = group_offset[None] + num_current_groups = rhs.shape[0] + num_total_groups = group_sizes.shape[0] + lhs, group_sizes, input_dtype = _validate_args( + lhs=lhs, rhs=rhs, group_sizes=group_sizes + ) + + # Gather shape information. + m, k, n = (lhs.shape[0], lhs.shape[1], rhs.shape[2]) + if transpose_rhs: + n = rhs.shape[1] + + # If tiling is callable, look up the problem dimensions in the LUT. If no tuned + # tile dimensions are available throw an error. + if callable(tiling): + tiling = tiling(m, k, n) + + if tiling is None: + raise ValueError(f"No tuned tiling found for (m, k, n) = ({m}, {k}, {n})") + + tm, tk, tn = tiling + tiles_k, k_rem = _calculate_irregular_num_tiles(k, tk) + tiles_n, n_rem = _calculate_irregular_num_tiles(n, tn) + del n_rem + + # Create the metadata we need for computation. + group_metadata, num_active_tiles = make_group_metadata( # pylint: disable=unbalanced-tuple-unpacking + group_sizes=group_sizes, + m=m, + tm=tm, + start_group=group_offset[0], + num_nonzero_groups=rhs.shape[0], + visit_empty_groups=False, + ) + + def kernel( + group_metadata, + group_offset, + lhs, + rhs, + existing_out, + out, + acc_scratch, + ): + group_offsets, group_ids, m_tile_ids = group_metadata + del group_offsets, group_ids, group_offset + + grid_id = pl.program_id(1) + k_i = pl.program_id(2) + + @pl.when(k_i == 0) + def _zero_acc(): + acc_scratch[...] = jnp.zeros_like(acc_scratch) + + if existing_out is not None: + prev_grid_id = jnp.where(grid_id > 0, grid_id - 1, 0) + is_first_processed_group = grid_id == 0 + m_tile_changed = m_tile_ids[grid_id] != m_tile_ids[prev_grid_id] + first_time_seeing_out = jnp.logical_or( + is_first_processed_group, m_tile_changed + ) + + @pl.when(first_time_seeing_out) + def _init_out(): + out[...] = existing_out[...] + + def mask_k_rem(x, *, dim): + if k_rem == 0: + return x + + orig_dtype = x.dtype + iota = lax.broadcasted_iota(jnp.int32, x.shape, dim) + x = x.astype(jnp.float32) + return jnp.where(iota < k_rem, x, 0).astype(orig_dtype) + + def _store_accum(): + mask = _get_store_mask( + grid_id=grid_id, + group_metadata=group_metadata, + tm=tm, + tn=tn, + ) + to_store = acc_scratch[...] + out[...] = jax.lax.select( + mask[...], to_store, out[...].astype(jnp.float32) + ).astype(preferred_element_type) + + def _accum(is_last_k_tile): + if is_last_k_tile: + mask_k_rem_lhs = partial(mask_k_rem, dim=1) + mask_k_rem_rhs = partial(mask_k_rem, dim=int(transpose_rhs)) + else: + mask_k_rem_lhs = lambda x: x + mask_k_rem_rhs = lambda x: x + + if transpose_rhs: + dot_general_dims = (((1,), (1,)), ((), ())) + else: + dot_general_dims = (((1,), (0,)), ((), ())) + + loaded_lhs = lhs[...] + loaded_rhs = rhs[...] + acc_scratch[...] += lax.dot_general( + mask_k_rem_lhs(loaded_lhs).astype(input_dtype), + mask_k_rem_rhs(loaded_rhs).astype(input_dtype), + preferred_element_type=jnp.float32, + dimension_numbers=dot_general_dims, + ) + + if is_last_k_tile: + _store_accum() + + lax.cond( + k_i == tiles_k - 1, + partial(_accum, True), + partial(_accum, False), + ) + + def lhs_transform_indices(n_i, grid_id, k_i, group_metadata, group_offset): + # lhs is (m, k). Load the [tm, tk] matrix for this m-tile. + group_offsets, group_ids, m_tile_ids = group_metadata + del n_i, group_offsets, group_ids, group_offset + return m_tile_ids[grid_id], k_i + + def rhs_transform_indices(n_i, grid_id, k_i, group_metadata, group_offset): + # rhs is (num_groups, k, n). Load the [tk, tn] matrix based on the group id + # for this m-tile. + group_offsets, group_ids, m_tile_ids = group_metadata + del group_offsets, m_tile_ids + if transpose_rhs: + k_i, n_i = n_i, k_i + + # NOTE: If we're working on only a shard of the rhs we need to adjust the + # group index we load from to account for this. The group_ids are in the + # "unsharded" domain. + return group_ids[grid_id] - group_offset[0], k_i, n_i + + def out_transform_indices(n_i, grid_id, k_i, group_metadata, group_offset): + # out is (m, n). Load the [tm, tn] matrix for this m-tile. + group_offsets, group_ids, m_tile_ids = group_metadata + del k_i, group_offsets, group_ids, group_offset + return m_tile_ids[grid_id], n_i + + out_block_spec = pl.BlockSpec((tm, tn), out_transform_indices) + if existing_out is None: + in_out_block_spec: Any = None + input_output_aliases = {} + else: + in_out_block_spec = out_block_spec + input_output_aliases = {6: 0} + + lhs_block_spec = pl.BlockSpec((tm, tk), lhs_transform_indices) + if transpose_rhs: + rhs_block_spec = pl.BlockSpec((None, tn, tk), rhs_transform_indices) + else: + rhs_block_spec = pl.BlockSpec((None, tk, tn), rhs_transform_indices) + + lhs_bytes = lhs.size * lhs.itemsize + rhs_bytes = (k * n) * rhs.itemsize # We don't read all of rhs + out_bytes = (m * n) * jnp.dtype(preferred_element_type).itemsize + max_active_tiles = group_metadata[1].size + bytes_accessed = ( + (lhs_bytes * tiles_n) + (rhs_bytes * max_active_tiles) + out_bytes + ) + flops = 2 * m * k * n + cost_estimate = pl.CostEstimate( + flops=flops, bytes_accessed=bytes_accessed, transcendentals=0 + ) + call_gmm = pl.pallas_call( + kernel, + out_shape=jax.ShapeDtypeStruct((m, n), preferred_element_type), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=2, + in_specs=[ + lhs_block_spec, + rhs_block_spec, + in_out_block_spec, + ], + out_specs=out_block_spec, + grid=(tiles_n, num_active_tiles, tiles_k), + scratch_shapes=[pltpu.VMEM((tm, tn), jnp.float32)], + ), + input_output_aliases=input_output_aliases, + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary")), + interpret=interpret, + cost_estimate=cost_estimate, + ) + + out = call_gmm( + group_metadata, + group_offset, + lhs, + rhs, + existing_out, + ) + if existing_out is None and num_current_groups < num_total_groups: + out = _zero_uninitialized_memory( + out, + start_group=group_offset[0], + num_nonzero_groups=rhs.shape[0], + group_metadata=group_metadata, + ) + return out + + +@functools.partial( + jax.jit, + static_argnames=[ + "preferred_element_type", + "tiling", + "num_actual_groups", + "interpret", + ], +) +def tgmm( + lhs: jnp.ndarray, + rhs: jnp.ndarray, + group_sizes: jnp.ndarray, + preferred_element_type: jnp.dtype = jnp.float32, + tiling: tuple[int, int, int] | LutFn | None = (128, 128, 128), + group_offset: jnp.ndarray | None = None, + num_actual_groups: int | None = None, + existing_out: jnp.ndarray | None = None, + interpret: bool = False, +) -> jnp.ndarray: + """Compute lhs[:, sizes[i-1]:sizes[i]] @ rhs[sizes[i-1]:sizes[i], :]. + + Args: + lhs: A 2d, jnp.ndarray with shape [k, m]. + rhs: A 2d, jnp.ndarray with shape [m, n]. + group_sizes: A 1d, jnp.ndarray with shape [num_groups] and jnp.int32 dtype. + preferred_element_type: jnp.dtype, the element type for the output matrix. + tiling: 3-tuple of ints. The m, k and n-dimension tile sizes. + group_offset: The group in group sizes to start computing from. This is + particularly useful for when rhs num_groups is sharded. + num_actual_groups: For when num_groups is sharded and we should only compute + the groups that are local, starting from group_offset. + existing_out: Existing output to write to. + interpret: Whether or not to run the kernel in interpret mode, helpful for + testing and debugging. + + Returns: + A 3d, jnp.ndarray with shape [num_groups, k, n]. + """ + if group_offset is None: + group_offset = jnp.array([0], dtype=jnp.int32) + else: + group_offset = group_offset[None] + lhs, group_sizes, input_dtype = _validate_args( + lhs=lhs, rhs=rhs, group_sizes=group_sizes, expected_rhs_dims=2 + ) + + # Gather shape information. + k, m, n = (lhs.shape[0], lhs.shape[1], rhs.shape[1]) + num_groups = group_sizes.shape[0] + num_actual_groups = ( + num_actual_groups if num_actual_groups is not None else num_groups + ) + + # If tiling is callable, look up the problem dimensions in the LUT. If no tuned + # tile dimensions are available throw an error. + if callable(tiling): + tiling = tiling(m, k, n) + + if tiling is None: + raise ValueError(f"No tuned tiling found for (m, k, n) = ({m}, {k}, {n})") + + tm, tk, tn = tiling + tiles_k, k_rem = _calculate_irregular_num_tiles(k, tk) + del k_rem + tiles_n, n_rem = _calculate_irregular_num_tiles(n, tn) + del n_rem + + # Create the metadata we need for computation. + group_metadata, num_active_tiles = make_group_metadata( + group_sizes=group_sizes, + m=m, + tm=tm, + start_group=group_offset[0], + num_nonzero_groups=num_actual_groups, + visit_empty_groups=True, + ) + + def kernel( + group_metadata, + group_offset, + lhs, + rhs, + existing_out, + out, + acc_scratch, + ): + grid_id = pl.program_id(2) + group_offsets, group_ids, m_tile_ids = group_metadata + del group_offsets, group_offset, m_tile_ids + + group = group_ids[grid_id] + prev_grid_id = jnp.where(grid_id > 0, grid_id - 1, 0) + prev_group = group_ids[prev_grid_id] + + group_has_changed = jnp.logical_or(grid_id == 0, prev_group != group) + + @pl.when(group_has_changed) + def _zero_acc(): + acc_scratch[...] = jnp.zeros_like(acc_scratch) + + # We'll only do computation if our group has a nonzero number of rows in it. + dont_skip = ( + _get_group_size(grid_id=grid_id, group_metadata=group_metadata) > 0 + ) + + @pl.when(dont_skip) + def _do(): + rhs_mask = _get_store_mask( + grid_id=grid_id, + group_metadata=group_metadata, + tm=tm, + tn=tn, + ) + lhs_mask = _get_store_mask( + grid_id=grid_id, + group_metadata=group_metadata, + tm=tm, + tn=tk, + ) + + loaded_lhs = lhs[...] + loaded_rhs = rhs[...] + loaded_lhs = lax.select( + lhs_mask[...], + loaded_lhs.astype(jnp.float32), + jnp.zeros_like(lhs, jnp.float32), + ).swapaxes(0, 1) + loaded_rhs = lax.select( + rhs_mask[...], + loaded_rhs.astype(jnp.float32), + jnp.zeros_like(rhs, jnp.float32), + ) + + acc_scratch[...] += lax.dot( + loaded_lhs.astype(input_dtype), + loaded_rhs.astype(input_dtype), + preferred_element_type=jnp.float32, + ) + + is_end_of_grid = grid_id == (pl.num_programs(2) - 1) + next_grid_id = jnp.where(is_end_of_grid, grid_id, grid_id + 1) + next_group = group_ids[next_grid_id] + + group_is_changing = jnp.logical_or(is_end_of_grid, group != next_group) + + @pl.when(group_is_changing) + def _store_accum(): + to_store = acc_scratch[...] + if existing_out is not None: + to_store += existing_out[...].astype(jnp.float32) + out[...] = to_store.astype(preferred_element_type) + + def lhs_transform_indices(n_i, k_i, grid_id, group_metadata, group_offset): + # lhs is (m, k). Load the [tm, tk] matrix for this m-tile. + group_offsets, group_ids, m_tile_ids = group_metadata + del n_i, group_offsets, group_ids, group_offset + return m_tile_ids[grid_id], k_i + + def rhs_transform_indices(n_i, k_i, grid_id, group_metadata, group_offset): + # rhs is (m, n). Load the [tm, tn] matrix for this m-tile. + group_offsets, group_ids, m_tile_ids = group_metadata + del k_i, group_offsets, group_ids, group_offset + return m_tile_ids[grid_id], n_i + + def out_transform_indices(n_i, k_i, grid_id, group_metadata, group_offset): + # out is (num_groups, k, n). Load the [tk, tn] matrix based on the group id + # for this m-tile. + group_offsets, group_ids, m_tile_ids = group_metadata + del group_offsets, m_tile_ids + + # NOTE: If we're working on only a shard of the output we need to adjust the + # group index we load from to account for this. The group_ids are in the + # "unsharded" domain. + return group_ids[grid_id] - group_offset[0], k_i, n_i + + out_block_spec = pl.BlockSpec((None, tk, tn), out_transform_indices) + if existing_out is None: + in_out_block_spec: Any = None + input_output_aliases = {} + else: + in_out_block_spec = out_block_spec + input_output_aliases = {6: 0} + + lhs_block_spec = pl.BlockSpec((tm, tk), lhs_transform_indices) + rhs_block_spec = pl.BlockSpec((tm, tn), rhs_transform_indices) + + lhs_bytes = lhs.size * lhs.itemsize + rhs_bytes = rhs.size * rhs.itemsize + out_bytewidth = jnp.dtype(preferred_element_type).itemsize + out_bytes = (num_actual_groups * k * n) * out_bytewidth + bytes_accessed = ( + (lhs_bytes * tiles_n) + (rhs_bytes * tiles_k) + out_bytes + ) + flops = 2 * m * k * n + cost_estimate = pl.CostEstimate( + flops=flops, bytes_accessed=bytes_accessed, transcendentals=0 + ) + lhs = lhs.swapaxes(0, 1) + call_gmm = pl.pallas_call( + kernel, + out_shape=jax.ShapeDtypeStruct( + (num_actual_groups, k, n), preferred_element_type + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=2, + in_specs=[ + lhs_block_spec, + rhs_block_spec, + in_out_block_spec, + ], + out_specs=out_block_spec, + grid=(tiles_n, tiles_k, num_active_tiles), + scratch_shapes=[pltpu.VMEM((tk, tn), jnp.float32)], + ), + input_output_aliases=input_output_aliases, + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary")), + interpret=interpret, + cost_estimate=cost_estimate, + ) + + out = call_gmm( + group_metadata, + group_offset, + lhs, + rhs, + existing_out, + ) + return out + + +CONFIG = { + 'name': 'pallas_megablox_gmm_qwen3_235b', + 'model': 'Qwen3-235B-A22B', + 'operator': 'pallas_gmm', + 'num_experts': 128, + 'num_experts_per_tok': 8, + 'emb_dim': 4096, + 'moe_mlp_dim': 1536, + 'seq_len': 4096, + 'atol': 1e-3, + 'rtol': 1e-2, +} + +# Tuned by autotune_block_sizes.py. Re-run to update. +TUNED_PARAMS = {'tiling': [256, 1024, 1024]} + + +def get_flops(): + top_k = CONFIG['num_experts_per_tok'] + K = CONFIG['emb_dim'] + N = CONFIG['moe_mlp_dim'] + S = CONFIG['seq_len'] + M = S * top_k + return 2 * M * K * N + + +def create_inputs(dtype=jnp.bfloat16): + key = jax.random.key(42) + k1, k2 = jax.random.split(key, 2) + G = CONFIG['num_experts'] + top_k = CONFIG['num_experts_per_tok'] + K = CONFIG['emb_dim'] + N = CONFIG['moe_mlp_dim'] + S = CONFIG['seq_len'] + M = S * top_k + lhs = jax.random.normal(k1, (M, K), dtype=dtype) + lhs = lhs.astype(jnp.bfloat16).astype(dtype) + rhs = jax.random.normal(k2, (G, K, N), dtype=dtype) * 0.02 + rhs = rhs.astype(jnp.bfloat16).astype(dtype) + tokens_per_expert = M // G + group_sizes = jnp.full((G,), tokens_per_expert, dtype=jnp.int32) + return lhs, rhs, group_sizes + + +def workload(lhs, rhs, group_sizes): + return gmm(lhs, rhs, group_sizes, tiling=tuple(TUNED_PARAMS['tiling'])) + + +def benchmark(num_warmup=5, num_iters=100): + """Benchmark and return results dict.""" + inputs = create_inputs() + fn = jax.jit(workload) + for _ in range(num_warmup): + out = fn(*inputs) + out.block_until_ready() + times = [] + for _ in range(num_iters): + t0 = time.perf_counter() + out = fn(*inputs) + out.block_until_ready() + times.append(time.perf_counter() - t0) + times = np.array(times) * 1000 + avg = float(np.mean(times)) + return { + 'name': CONFIG['name'], + 'model': CONFIG['model'], + 'operator': CONFIG['operator'], + 'config': {k: v for k, v in CONFIG.items() if k not in ('name', 'model', 'operator', 'atol', 'rtol')}, + 'time_ms': round(avg, 4), + 'std_ms': round(float(np.std(times)), 4), + 'output_shape': list(out.shape) if hasattr(out, 'shape') else [], + 'status': 'success', + } + + +if __name__ == '__main__': + import json + print(json.dumps(benchmark())) diff --git a/JAXBench/benchmark/level2/1p_Flash_Attention/baseline.py b/JAXBench/benchmark/level2/1p_Flash_Attention/baseline.py new file mode 100644 index 0000000..1d7c572 --- /dev/null +++ b/JAXBench/benchmark/level2/1p_Flash_Attention/baseline.py @@ -0,0 +1,1817 @@ +# 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. + +"""Flash Attention TPU kernel.""" + +import numpy as np +import time +import dataclasses +import functools +import math +from typing import Any, NamedTuple + +import jax +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +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 + + block_q_major_dkv: int | None = None + block_k_major_dkv: int | None = None + block_k_dkv: int | None = None + block_q_dkv: int | None = None + + block_k_major_dq: int | None = None + block_k_dq: int | None = None + block_q_dq: int | None = None + + 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) + if self.block_q_major_dkv is not None and self.block_q_dkv is not None: + verify_major_minor( + "block_q", "_dkv", self.block_q_major_dkv, self.block_q_dkv + ) + if self.block_k_major_dkv is not None and self.block_k_dkv is not None: + verify_major_minor( + "block_k", "_dkv", self.block_k_major_dkv, self.block_k_dkv + ) + if self.block_k_major_dq is not None and self.block_k_dq is not None: + verify_major_minor( + "block_k", "_dq", self.block_k_major_dq, self.block_k_dq + ) + + @property + def has_backward_blocks(self) -> bool: + backward_blocks = ( + self.block_q_major_dkv, + self.block_k_major_dkv, + self.block_q_dkv, + self.block_k_dkv, + self.block_k_major_dq, + self.block_k_dq, + self.block_q_dq, + ) + return all(b is not None for b in backward_blocks) + + @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, + block_q_major_dkv=128, + block_k_major_dkv=128, + block_k_dkv=128, + block_q_dkv=128, + block_k_major_dq=128, + block_k_dq=128, + block_q_dq=128, + ) + + +@functools.partial( + jax.jit, + static_argnames=[ + "causal", + "sm_scale", + "block_sizes", + "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, + 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 + ) + return _flash_attention( + q, k, v, ab, segment_ids, False, causal, sm_scale, block_sizes, debug + ) + + +@functools.partial(jax.custom_vjp, nondiff_argnames=("save_residuals", "causal", "sm_scale", "block_sizes", "debug")) +def _flash_attention( + q, + k, + v, + ab, + segment_ids, + save_residuals, + causal, + sm_scale, + block_sizes, + 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, + debug, + ) + + +def _flash_attention_fwd( + q, + k, + v, + ab, + segment_ids, + save_residuals, + causal, + sm_scale, + block_sizes, + debug, +): + if save_residuals: + raise NotImplementedError("Higher-order AD not supported") + o, l, m = _flash_attention( + q, k, v, ab, segment_ids, True, causal, sm_scale, block_sizes, debug + ) + return o, (q, k, v, ab, segment_ids, o, l, m) + + +def _flash_attention_bwd( + save_residuals: bool, + causal: bool, + sm_scale: float, + block_sizes: BlockSizes, + debug: bool, + residuals, + do, +): + """VJP rule for FlashAttention.""" + if save_residuals: + raise NotImplementedError("Higher-order AD not supported") + (q, k, v, ab, segment_ids, o, l, m) = residuals + if not block_sizes.has_backward_blocks: + raise ValueError( + "Program is being differentiated, but not all backward blocks are" + " specified" + ) + + di = jnp.sum( + o.astype(jnp.float32) * do.astype(jnp.float32), axis=-1 + ) # [batch_size, num_heads, q_seq_len] + + dk, dv = _flash_attention_bwd_dkv( + q, + k, + v, + ab, + segment_ids, + l, + m, + do, + di, + block_q_major=block_sizes.block_q_major_dkv, + block_k_major=block_sizes.block_k_major_dkv, + block_k=block_sizes.block_k_dkv, + block_q=block_sizes.block_q_dkv, + sm_scale=sm_scale, + causal=causal, + mask_value=DEFAULT_MASK_VALUE, + debug=debug, + ) + + dq, ds = _flash_attention_bwd_dq( + q, + k, + v, + ab, + segment_ids, + l, + m, + do, + di, + block_q_major=block_sizes.block_q_dq, + block_k_major=block_sizes.block_k_major_dq, + block_k=block_sizes.block_k_dq, + sm_scale=sm_scale, + causal=causal, + mask_value=DEFAULT_MASK_VALUE, + debug=debug, + ) + return dq, dk, dv, ds, None + + +_flash_attention.defvjp(fwd=_flash_attention_fwd, bwd=_flash_attention_bwd) + + +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 = jnp.tile( + q_segment_ids_tile_ref[batch_idx[0]], (1, repeats) + ) # [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 - jnp.tile(m_next, (1, block_k_repeats))) + + 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: jnp.tile(l, (1, head_dim_repeats)) + 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) + + +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 = jnp.tile( + q_segment_ids, (1, repeats) + ) # [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, + 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", + ) + ), + 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 + + +def _flash_attention_dkv_kernel( + q_tile_ref, + k_tile_ref, + v_tile_ref, + ab_tile_ref, + q_segment_ids_tile_ref, + kv_segment_ids_tile_ref, + l_tile_ref, + m_tile_ref, + do_tile_ref, + di_tile_ref, + dk_tile_ref, + dv_tile_ref, + dk_scratch_ref, + dv_scratch_ref, + *, + sm_scale: float, + causal: bool, + mask_value: float, + q_seq_len: int, + block_q: int, + block_k: int, +): + _, _, block_q_major, _ = q_tile_ref.shape + _, _, block_k_major, _ = k_tile_ref.shape + + q_seq_index = pl.program_id(axis=3) + kv_seq_index = pl.program_id(axis=2) + + @pl.when(q_seq_index == 0) + def start_new_sequence(): + dk_scratch_ref[:, :] = jnp.zeros(dk_scratch_ref.shape, dk_scratch_ref.dtype) + dv_scratch_ref[:, :] = jnp.zeros(dv_scratch_ref.shape, dv_scratch_ref.dtype) + + def q_body(j, _): + start_q = j * block_q + def k_body(i, _): + start_k = i * block_k + k = k_tile_ref[0, 0, pl.ds(start_k, block_k), :] + v = v_tile_ref[0, 0, pl.ds(start_k, block_k), :] + q = q_tile_ref[0, 0, pl.ds(start_q, block_q), :] # [block_q, head_dim] + l = l_tile_ref[0, 0, pl.ds(start_q, block_q), :] # [block_q, 128] + m = m_tile_ref[0, 0, pl.ds(start_q, block_q), :] # [block_q, 128] + do = do_tile_ref[0, 0, pl.ds(start_q, block_q), :] # [block_q, 128] + di = di_tile_ref[0, 0, pl.ds(start_q, block_q), :].astype( + jnp.float32 + ) # [block_q, 128] + + capped_logits = lax.dot_general( + q, k, TRANS_B_DIM_NUMBERS, preferred_element_type=jnp.float32 + ) # [block_q_major, block_k] + + if ab_tile_ref is not None: + ab = ab_tile_ref[ + 0, + 0, + pl.dslice(j * block_q, block_q), + pl.dslice(i * block_k, block_k), + ].astype(jnp.float32) + capped_logits += ab + + if sm_scale != 1.0: + capped_logits *= sm_scale + + mask = None + if q_segment_ids_tile_ref is not None: + repeats, rem = divmod(block_k, NUM_LANES) + if rem: + raise NotImplementedError( + ) + q_segment_ids = q_segment_ids_tile_ref[ + 0, pl.ds(start_q, block_q), : + ] # [block_q, NUM_LANES]. + q_segment_ids = jnp.tile( + q_segment_ids, (1, repeats) + ) # [block_q, block_k]. + kv_segment_ids = kv_segment_ids_tile_ref[ + :, 0, pl.ds(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_index * block_q_major + start_q + col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1) + col_ids += kv_seq_index * 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) + ) + + capped_logits = ( + capped_logits + if mask is None + else capped_logits + jnp.where(mask, 0.0, mask_value) + ) + + p = jnp.exp( + capped_logits - jnp.tile(m, (1, block_k // MIN_BLOCK_SIZE)) + ) + p = p * jnp.tile( + 1 / l, (1, block_k // MIN_BLOCK_SIZE) + ) # [block_q_major, block_k_major] + dv = lax.dot(p.T.astype(do.dtype), do, preferred_element_type=jnp.float32) + dv_scratch_ref[pl.ds(start_k, block_k), :] += dv.astype( + dv_scratch_ref.dtype + ) + + # di: [block_q, 128] + # do: [block_q, head_dim] + # v: [block_k_major, head_dim] + dp = lax.dot_general( + do, v, TRANS_B_DIM_NUMBERS, preferred_element_type=jnp.float32 + ) + ds = (dp - jnp.tile(di, (1, block_k // MIN_BLOCK_SIZE))) * p + + if sm_scale != 1.0: + ds = ds * sm_scale + + # ds: [block_q_major, block_k_major] + # q: [block_q_major, head_dim] + dk = lax.dot(ds.T.astype(do.dtype), q, preferred_element_type=jnp.float32) + dk_scratch_ref[pl.ds(start_k, block_k), :] += dk.astype( + dk_scratch_ref.dtype + ) + lax.fori_loop(0, block_k_major // block_k, k_body, None, unroll=True) + + if causal: + should_run = below_or_on_diag( + q_seq_index, block_q_major, kv_seq_index, block_k_major + ) + else: + should_run = True + + @pl.when(should_run) + def run(): + lax.fori_loop(0, block_q_major // block_q, q_body, None, unroll=True) + + @pl.when(q_seq_index == q_seq_len // block_q_major - 1) + def end_of_q_sequence(): + dv_tile_ref[0, 0, :, :] = dv_scratch_ref[...].astype(dv_tile_ref.dtype) + dk_tile_ref[0, 0, :, :] = dk_scratch_ref[...].astype(dk_tile_ref.dtype) + + +def _flash_attention_bwd_dkv( + q, + k, + v, + ab, + segment_ids, + l, + m, + do, + di, + *, + block_q_major: int | None, + block_q: int | None, + block_k_major: int | None, + block_k: int | None, + sm_scale: float, + causal: bool = False, + mask_value: float = DEFAULT_MASK_VALUE, + debug: bool = False, +): + batch_size, num_heads, q_seq_len, head_dim = q.shape + _, _, kv_seq_len, _ = k.shape + _verify_block("block_q_major_dkv", "q_seq_len", block_q_major, q_seq_len) + _verify_block("block_q_dkv", "q_seq_len", block_q, q_seq_len) + _verify_block("block_k_major_dkv", "kv_seq_len", block_k_major, kv_seq_len) + _verify_block("block_k_dkv", "kv_seq_len", block_k, kv_seq_len) + + # Broadcast out scalar values + m = jnp.broadcast_to(m[..., None], (*m.shape, MIN_BLOCK_SIZE)) + l = jnp.broadcast_to(l[..., None], (*l.shape, MIN_BLOCK_SIZE)) + # Preprocess contraction for bwd pass + di = jnp.broadcast_to(di[..., None], (*di.shape, MIN_BLOCK_SIZE)) + + # kv index needs to be before q index since q index is the contractng + # dimension. + grid = ( + batch_size, + num_heads, + kv_seq_len // block_k_major, + q_seq_len // block_q_major, + ) + + def qo_index_map(batch_index, head_index, kv_seq_index, q_seq_index): + if causal: + # If the q block is skipped, stay at the 0th q block. + next_q_index = lax.select( + below_or_on_diag( + q_seq_index, block_q_major, kv_seq_index, block_k_major + ), + q_seq_index, + 0, + ) + else: + next_q_index = q_seq_index + + return (batch_index, head_index, next_q_index, 0) + + qo_spec = pl.BlockSpec((1, 1, block_q_major, head_dim), qo_index_map) + assert qo_spec.block_shape is not None + assert q.ndim == len(qo_spec.block_shape) + do_spec = qo_spec + assert do.ndim == len(qo_spec.block_shape) + + def kv_index_map(batch_index, head_index, kv_seq_index, _): + return (batch_index, head_index, kv_seq_index, 0) + + kv_spec = pl.BlockSpec((1, 1, block_k_major, head_dim), kv_index_map) + assert kv_spec.block_shape is not None + assert k.ndim == len(kv_spec.block_shape) + assert v.ndim == len(kv_spec.block_shape) + + def lm_index_map(batch_index, head_index, _, q_seq_index): + return (batch_index, head_index, q_seq_index, 0) + + lm_spec = pl.BlockSpec((1, 1, block_q_major, MIN_BLOCK_SIZE), lm_index_map) + assert lm_spec.block_shape is not None + assert l.ndim == len(lm_spec.block_shape) + assert m.ndim == len(lm_spec.block_shape) + + di_spec = pl.BlockSpec((1, 1, block_q_major, MIN_BLOCK_SIZE), qo_index_map) + assert di_spec.block_shape is not None + assert di.ndim == len(di_spec.block_shape) + + def ab_index_map(batch_index, head_index, kv_seq_index, q_seq_index): + return (batch_index, head_index, q_seq_index, kv_seq_index) + + dab_spec = ( + pl.BlockSpec((1, 1, block_q_major, 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, kv_seq_index, q_seq_index + ): + del head_index + if causal: + next_q_index = lax.select( + below_or_on_diag( + q_seq_index, block_q_major, kv_seq_index, block_k_major + ), + q_seq_index, + 0, + ) + else: + next_q_index = q_seq_index + return (batch_index, next_q_index, 0) + + def kv_segment_ids_index_map(batch_index, head_index, kv_seq_index, _): + del head_index + return (batch_index, 0, kv_seq_index) + + q_segment_ids_spec = pl.BlockSpec( + (1, block_q_major, NUM_LANES), q_segment_ids_index_map + ) + kv_segment_ids_spec = pl.BlockSpec( + (1, 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 = [ + qo_spec, + kv_spec, + kv_spec, + dab_spec, + q_segment_ids_spec, + kv_segment_ids_spec, + lm_spec, + lm_spec, + do_spec, + di_spec, + ] + + out_shapes = [ + jax.ShapeDtypeStruct((batch_size, num_heads, kv_seq_len, head_dim), + k.dtype), + jax.ShapeDtypeStruct((batch_size, num_heads, kv_seq_len, head_dim), + v.dtype), + ] + def dkv_index_map(batch_index, head_index, kv_seq_index, _): + return (batch_index, head_index, kv_seq_index, 0) + + dkv_spec = pl.BlockSpec((1, 1, block_k_major, head_dim), dkv_index_map) + out_specs = [dkv_spec, dkv_spec] + scratch_shapes = [ + pltpu.VMEM((block_k_major, head_dim), jnp.float32), # type: ignore + pltpu.VMEM((block_k_major, head_dim), jnp.float32), # type: ignore + ] + + kernel = functools.partial( + _flash_attention_dkv_kernel, + block_q=block_q, # type: ignore + block_k=block_k, # type: ignore + sm_scale=sm_scale, + causal=causal, + mask_value=mask_value, + q_seq_len=q_seq_len, + ) + name_scope = f"flash_mha_bwd_dkv_{block_q_major=}_{block_q=}_{block_k_major=}_{block_k=}" + with jax.named_scope(name_scope): + dk, dv = 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_shapes, + debug=debug, + compiler_params=pltpu.CompilerParams( + dimension_semantics=( + "parallel", + "parallel", + "parallel", + "arbitrary", + ) + ), + )(q, k, v, ab, q_segment_ids, kv_segment_ids, l, m, do, di) + assert dk.shape == k.shape + assert dv.shape == v.shape + return dk, dv + + +def _flash_attention_dq_kernel( + q_tile_ref, + k_tile_ref, + v_tile_ref, + ab_tile_ref, + q_segment_ids_tile_ref, + kv_segment_ids_tile_ref, + l_tile_ref, + m_tile_ref, + do_tile_ref, + di_tile_ref, + dq_tile_ref, + ds_tile_ref, + dq_scratch_ref, + *, + sm_scale: float, + causal: bool, + mask_value: float, + kv_seq_len: int, + block_k: int, +): + _, _, block_k_major, _ = k_tile_ref.shape + _, _, block_q_major, _ = q_tile_ref.shape + + kv_seq_index = pl.program_id(axis=3) + q_seq_index = pl.program_id(axis=2) + + @pl.when(kv_seq_index == 0) + def start_new_sequence(): + dq_scratch_ref[:, :] = jnp.zeros(dq_scratch_ref.shape, dq_scratch_ref.dtype) + + def body(i, _): + k_slice = pl.ds(i * block_k, block_k) + q = q_tile_ref[0, 0, :, :] + k = k_tile_ref[0, 0, k_slice, :] # [block_k, head_dim] + v = v_tile_ref[0, 0, k_slice, :] # [block_k, head_dim] + l = l_tile_ref[0, 0, :, :] # [block_q_major, 128] + m = m_tile_ref[0, 0, :, :] # [block_q_major, 128] + do = do_tile_ref[0, 0, :, :] # [block_q_major, head_dim] + di = di_tile_ref[0, 0, :].astype(jnp.float32) # [block_q_major, 128] + + capped_logits = jax.lax.dot_general( + q, k, TRANS_B_DIM_NUMBERS, preferred_element_type=jnp.float32 + ) + + if ab_tile_ref is not None: + ab = ab_tile_ref[0, 0, :, pl.dslice(i * block_k, block_k)].astype( + jnp.float32 + ) + capped_logits += ab + + if sm_scale != 1.0: + capped_logits *= 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 = jnp.tile( + q_segment_ids_tile_ref[0], (1, repeats) + ) # [block_q, block_k]. + kv_segment_ids = kv_segment_ids_tile_ref[:, 0, k_slice] # [1, block_k]. + mask = jnp.equal(q_segment_ids, kv_segment_ids).astype(jnp.bool_) + + if causal: + mask_shape = (block_q_major, block_k) + row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0) + row_ids += q_seq_index * block_q_major + col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1) + col_ids += kv_seq_index * block_k_major + i * block_k + causal_mask = col_ids <= row_ids + mask = causal_mask if mask is None else jnp.logical_and(mask, causal_mask) + capped_logits = ( + capped_logits + if mask is None + else capped_logits + jnp.where(mask, 0.0, mask_value) + ) + + p = jnp.exp( + capped_logits - jnp.tile(m, (1, block_k // MIN_BLOCK_SIZE)) + ) + p = p * jnp.tile( + 1 / l, (1, block_k // MIN_BLOCK_SIZE) + ) # [block_q_major, block_k] + + # di: [block_q_major, 128] + # do: [block_q_major, head_dim] + # v: [block_k_major, head_dim] + dp = jax.lax.dot_general( + do, + v, + TRANS_B_DIM_NUMBERS, + preferred_element_type=jnp.float32, + ) + ds = (dp - jnp.tile(di, (1, block_k // MIN_BLOCK_SIZE))) * p + # dp = jnp.dot(do, v.T) + # ds = (dp - (dp * p).sum(axis=1)[:, None]) * p + + if sm_scale != 1.0: + ds = ds * sm_scale + + if ds_tile_ref is not None: + ds_tile_ref[0, 0, :, pl.dslice(i * block_k, block_k)] = ds.astype( + ds_tile_ref.dtype + ) + + # dp: [block_q_major, block_k] + # k: [block_k, head_dim] + dq_scratch_ref[:, :] += lax.dot( + ds.astype(k.dtype), + k, + preferred_element_type=jnp.float32, + ).astype(dq_scratch_ref.dtype) + + if causal: + should_run = below_or_on_diag( + q_seq_index, block_q_major, kv_seq_index, block_k_major + ) + should_not_run = lax.select(should_run, False, True) + else: + should_run = True + should_not_run = False # type: ignore + + @pl.when(should_run) + def run(): + lax.fori_loop(0, block_k_major // block_k, body, None, unroll=True) + + @pl.when(should_not_run) + def zero_out_ds(): + if ds_tile_ref is not None: + ds_tile_ref[...] = jnp.zeros_like(ds_tile_ref) + + @pl.when(kv_seq_index == kv_seq_len // block_k_major - 1) + def end_of_kv_sequence(): + dq_tile_ref[0, 0, :, :] = dq_scratch_ref[...].astype(dq_tile_ref.dtype) + dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref) + + +def _flash_attention_bwd_dq( + q, + k, + v, + ab, + segment_ids, + l, + m, + do, + di, + *, + block_q_major: int | None, + block_k_major: int | None, + block_k: int | None, + sm_scale: float, + causal: bool, + mask_value: float, + debug: bool, +): + batch_size, num_heads, q_seq_len, head_dim = q.shape + _, _, kv_seq_len, _ = k.shape + _verify_block("block_q_dq", "q_seq_len", block_q_major, q_seq_len) + _verify_block("block_k_major_dq", "kv_seq_len", block_k_major, kv_seq_len) + _verify_block("block_k_dq", "block_k", block_k, kv_seq_len) + + # Broadcast out scalar values + m = jnp.broadcast_to(m[..., None], (*m.shape, MIN_BLOCK_SIZE)) + l = jnp.broadcast_to(l[..., None], (*l.shape, MIN_BLOCK_SIZE)) + # Preprocess contraction for bwd pass + di = jnp.broadcast_to(di[..., None], (*di.shape, block_k_major)) + + grid = ( + batch_size, + num_heads, + q_seq_len // block_q_major, + kv_seq_len // block_k_major, + ) + + def qo_index_map(batch_index, head_index, q_seq_index, _): + return (batch_index, head_index, q_seq_index, 0) + + qo_spec = pl.BlockSpec((1, 1, block_q_major, head_dim), qo_index_map) + do_spec = qo_spec + + 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_major, 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) + + kv_spec = pl.BlockSpec((1, 1, block_k_major, head_dim), kv_index_map) + assert kv_spec.block_shape is not None + assert k.ndim == len(kv_spec.block_shape) + assert v.ndim == len(kv_spec.block_shape) + + def lm_index_map(batch_index, head_index, q_seq_index, _): + return (batch_index, head_index, q_seq_index, 0) + + lm_spec = pl.BlockSpec((1, 1, block_q_major, MIN_BLOCK_SIZE), lm_index_map) + assert lm_spec.block_shape is not None + assert l.ndim == len(lm_spec.block_shape) + assert m.ndim == len(lm_spec.block_shape) + + di_spec = pl.BlockSpec((1, 1, block_q_major, MIN_BLOCK_SIZE), qo_index_map) + assert di_spec.block_shape is not None + assert di.ndim == len(di_spec.block_shape) + + def ab_index_map(batch_index, head_index, q_seq_index, kv_seq_index): + return (batch_index, head_index, q_seq_index, kv_seq_index) + + dab_spec = ( + pl.BlockSpec((1, 1, block_q_major, 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: + # 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_major, 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( + (1, block_q_major, NUM_LANES), q_segment_ids_index_map + ) + kv_segment_ids_spec = pl.BlockSpec( + (1, 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 = [ + qo_spec, + kv_spec, + kv_spec, + dab_spec, + q_segment_ids_spec, + kv_segment_ids_spec, + lm_spec, + lm_spec, + do_spec, + di_spec, + ] + + out_shapes = [ + jax.ShapeDtypeStruct(q.shape, q.dtype), + jax.ShapeDtypeStruct(ab.shape, ab.dtype) if ab is not None else None, + ] + dq_spec = pl.BlockSpec((1, 1, block_q_major, head_dim), qo_index_map) + out_specs = [ + dq_spec, + dab_spec, + ] + scratch_shapes = [pltpu.VMEM((block_q_major, head_dim), jnp.float32)] # type: ignore + + kernel = functools.partial( + _flash_attention_dq_kernel, + sm_scale=sm_scale, + causal=causal, + mask_value=mask_value, + block_k=block_k, # type: ignore + kv_seq_len=kv_seq_len, + ) + name_scope = f"flash_mha_bwd_dq_{block_q_major=}_{block_k_major=}_{block_k=}" + with jax.named_scope(name_scope): + dq, ds = 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_shapes, + debug=debug, + compiler_params=pltpu.CompilerParams( + dimension_semantics=( + "parallel", + "parallel", + "parallel", + "arbitrary", + ) + ), + )(q, k, v, ab, q_segment_ids, kv_segment_ids, l, m, do, di) + + # dab is just ds + return dq, ds + + +# 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, + ) + + +@functools.partial(jax.custom_vjp, nondiff_argnames=("causal", "mask_value", "sm_scale", "save_residuals")) +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 _mha_reference_fwd( + q, + k, + v, + ab, + segment_ids: SegmentIds | None, + causal: bool, + mask_value: float, + sm_scale: float, + save_residuals: bool, +): + if save_residuals: + raise NotImplementedError + res = _mha_reference( + q, + k, + v, + ab, + segment_ids, + causal=causal, + mask_value=mask_value, + sm_scale=sm_scale, + save_residuals=True, + ) + assert isinstance(res, tuple) + out, l, m = res + return out, (q, k, v, ab, segment_ids, out, l, m) + + +@functools.partial( + jax.jit, + static_argnames=[ + "causal", + "mask_value", + "sm_scale", + ], +) +def mha_reference_bwd( + q, + k, + v, + ab, + segment_ids: SegmentIds | None, + o, + l, + m, + do, + causal: bool = False, + mask_value: float = DEFAULT_MASK_VALUE, + sm_scale: float = 1.0, +): + if sm_scale != 1.0: + raise NotImplementedError + + logits = jnp.einsum( + "bhqc,bhkc->bhqk", + q.astype(jnp.float32), + k.astype(jnp.float32), + ) + if ab is not None: + logits += ab + + 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) + + unnormalized = jnp.exp(logits - m[..., None]) + p = unnormalized / l[..., None] + dv = jnp.einsum("bhpt,bhpd->bhtd", p, do.astype(jnp.float32)).astype(v.dtype) + + dp = jnp.einsum( + "bhpd,bhtd->bhpt", do.astype(jnp.float32), v.astype(jnp.float32) + ) + + di = jnp.sum(o.astype(jnp.float32) * do.astype(jnp.float32), axis=-1)[ + ..., None + ] # [batch_size, num_heads, q_seq_len] + + ds = (dp - di) * p + dk = jnp.einsum("bhsd,bhst->bhtd", q.astype(jnp.float32), ds).astype(k.dtype) + dq = jnp.einsum("bhst,bhtd->bhsd", ds, k.astype(jnp.float32)).astype(q.dtype) + + # dab is just ds + dab = ds if ab is not None else None + return dq, dk, dv, dab + + +def _mha_reference_bwd( + causal: bool, + mask_value: float, + sm_scale: float, + save_residuals: bool, + residuals, + do, +): + del save_residuals + q, k, v, ab, segment_ids, o, l, m = residuals + dq, dk, dv, dab = mha_reference_bwd( + q, + k, + v, + ab, + segment_ids, + o, + l, + m, + do, + causal=causal, + mask_value=mask_value, + sm_scale=sm_scale, + ) + return dq, dk, dv, dab, None + + +_mha_reference.defvjp(fwd=_mha_reference_fwd, bwd=_mha_reference_bwd) + + +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}" + ) + + +CONFIG = { + 'name': 'pallas_flash_attention_llama70b', + 'model': 'Llama-3.1-70B', + 'operator': 'pallas_flash_attention', + 'batch': 4, + 'seq_len': 4096, + 'num_heads': 64, + 'head_dim': 128, + 'atol': 2e-3, + 'rtol': 2e-3, +} + +# Tuned by autotune_block_sizes.py. Re-run to update. +TUNED_PARAMS = { + # Autotuned (forward pass). + 'block_q': 2048, + 'block_k_major': 2048, + 'block_k': 1024, + # Not autotuned (batch=1, backward-only). + 'block_b': 1, + 'block_q_major_dkv': 128, + 'block_k_major_dkv': 128, + 'block_k_dkv': 128, + 'block_q_dkv': 128, + 'block_k_major_dq': 128, + 'block_k_dq': 128, + 'block_q_dq': 128, +} + + +def get_flops(): + B, S = CONFIG['batch'], CONFIG['seq_len'] + H, D = CONFIG['num_heads'], CONFIG['head_dim'] + return 4 * B * H * S * S * D + + +def create_inputs(dtype=jnp.bfloat16): + key = jax.random.key(42) + k1, k2, k3 = jax.random.split(key, 3) + B = CONFIG['batch'] + H = CONFIG['num_heads'] + S = CONFIG['seq_len'] + D = CONFIG['head_dim'] + q = jax.random.normal(k1, (B, H, S, D), dtype=dtype) + k = jax.random.normal(k2, (B, H, S, D), dtype=dtype) + v = jax.random.normal(k3, (B, H, S, D), dtype=dtype) + return q, k, v + + +def workload(q, k, v): + sm_scale = 1.0 / math.sqrt(CONFIG['head_dim']) + block_sizes = BlockSizes( + block_q=TUNED_PARAMS['block_q'], + block_k_major=TUNED_PARAMS['block_k_major'], + block_k=TUNED_PARAMS['block_k'], + block_b=TUNED_PARAMS['block_b'], + block_q_major_dkv=TUNED_PARAMS['block_q_major_dkv'], + block_k_major_dkv=TUNED_PARAMS['block_k_major_dkv'], + block_k_dkv=TUNED_PARAMS['block_k_dkv'], + block_q_dkv=TUNED_PARAMS['block_q_dkv'], + block_k_major_dq=TUNED_PARAMS['block_k_major_dq'], + block_k_dq=TUNED_PARAMS['block_k_dq'], + block_q_dq=TUNED_PARAMS['block_q_dq'], + ) + return flash_attention( + q, k, v, causal=True, sm_scale=sm_scale, block_sizes=block_sizes, + ) + + +def benchmark(num_warmup=5, num_iters=100): + """Benchmark and return results dict.""" + inputs = create_inputs() + fn = jax.jit(workload) + for _ in range(num_warmup): + out = fn(*inputs) + out.block_until_ready() + times = [] + for _ in range(num_iters): + t0 = time.perf_counter() + out = fn(*inputs) + out.block_until_ready() + times.append(time.perf_counter() - t0) + times = np.array(times) * 1000 + avg = float(np.mean(times)) + return { + 'name': CONFIG['name'], + 'model': CONFIG['model'], + 'operator': CONFIG['operator'], + 'config': {k: v for k, v in CONFIG.items() if k not in ('name', 'model', 'operator', 'atol', 'rtol')}, + 'time_ms': round(avg, 4), + 'std_ms': round(float(np.std(times)), 4), + 'output_shape': list(out.shape) if hasattr(out, 'shape') else [], + 'status': 'success', + } + + +if __name__ == '__main__': + import json + print(json.dumps(benchmark())) diff --git a/JAXBench/benchmark/level2/2p_GQA_Attention/baseline.py b/JAXBench/benchmark/level2/2p_GQA_Attention/baseline.py new file mode 100644 index 0000000..14359ea --- /dev/null +++ b/JAXBench/benchmark/level2/2p_GQA_Attention/baseline.py @@ -0,0 +1,2648 @@ +# 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.""" + +import time +from collections.abc import Callable, Mapping +import dataclasses +import enum +import functools +from typing import Any, Literal, NamedTuple, Optional, Union, 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 mechanism 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( + block_metadata: Mapping[str, Any], + 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}_" + "_".join( + f"{k}={v}" for k, v in sorted(block_metadata.items()) + ) + + +# 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_argnames=( + "mask_value", "save_residuals", "custom_type", "attn_logits_soft_cap") +) +_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: # pytype: disable=attribute-error + kv_heads, q_heads_per_kv_head, q_seq_len, head_dim = activations.shape # pytype: disable=attribute-error + return activations.reshape( # pytype: disable=attribute-error + 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(frozen=True, slots=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: + object.__setattr__(self, "block_kv_compute", self.block_kv) + if self.block_kv_dkv_compute is None: + object.__setattr__(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 | jax.Array, + 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 = mask_ref[:, k_slice] + else: + mask = mask_ref[k_slice, :] + + masks.append( + jnp.bitwise_or(mask, jnp.broadcast_to(should_not_mask, 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 = jnp.tile( + q_sequence_ref[...], (1, repeats) + ) # [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 = q_sequence_ref[:1, :] # [1, bq] + q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq)) + + assert q_sequence.shape == k_sequence.shape + computed_mask = mask_function(q_sequence, k_sequence) # pytype: disable=wrong-arg-count + if computed_mask.dtype != jnp.dtype(jnp.bool_): + raise ValueError( + "Mask function must return a boolean-valued array, but got:" + f" {computed_mask.dtype}" + ) + masks.append(computed_mask) + + if q_segment_ids_ref is not None: + if k_in_lanes: + kv_ids = kv_segment_ids_ref[: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 = jnp.tile(q_segment_ids_ref[:], (1, repeats)) # [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 = jnp.tile( + kv_segment_ids_ref[k_slice, :], (1, repeats) + ) # [k_slice, bq] + q_ids = q_segment_ids_ref[:1, :] # [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_v: 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_v_repeats, rem = divmod(head_dim_v, NUM_LANES) + if rem != 0: + raise NotImplementedError( + f"{head_dim_v=} 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) + + global_kv_index, _, 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 = k_ref[slice_k, :] + else: + k = k_ref[:, 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, + # When the iteration space is shrunk (for local attention for example), + # the kv_index program_id does not correspond to the actual coordinates + # of the KV data. Make sure to use the 'unshrunk' index (coming from the + # data_next array) when computing the mask. + k_offset=global_kv_index * 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] # pytype: disable=attribute-error + 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 - jnp.tile(m_next, (1, bkv_repeats))) + 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 = v_ref[slice_k, :] + else: + v = v_ref[:, slice_k] + v = v.astype(float32) + o_curr = lax.dot_general(s_curr, v, sv_dims) + + alpha_o = jnp.tile(alpha, (1, head_dim_v_repeats)) + 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 = jnp.tile(1.0 / l, (1, head_dim_v_repeats)) + 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, + interpret: bool = False +) -> SplashCustomReturnType: + num_q_heads, q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + 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] + + partial_mask_blocks = fwd_mask_info.partial_mask_blocks + if ( + partial_mask_blocks is not None + and jnp.dtype(partial_mask_blocks.dtype) != np.bool_ + ): + raise ValueError( + "partial_mask_blocks must be of type np.bool_ but got" + f" {partial_mask_blocks.dtype}" + ) + + 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_qk: + raise ValueError( + f"Expected 'key' head dimension to be: {head_dim_qk}. 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[:-1] != v.shape[:-1]: + raise ValueError( + f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " + "leading dimensions." + ) + + if bkv % bkv_compute: + raise ValueError(f"{bkv=} must be a multiple of {bkv_compute=}.") + if bkv_compute % NUM_LANES: + raise ValueError(f"{bkv_compute=} 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( + from_head_minor((None, bq, head_dim_qk), q_layout), q_index_map + ), + pl.BlockSpec( + from_head_minor( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), k_layout + ), + k_index_map, + ), + pl.BlockSpec( + from_head_minor( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), v_layout + ), + v_index_map, + ), + ] + if segment_ids is not None: + in_specs += [ + pl.BlockSpec((bq, NUM_LANES), q_segment_ids_index_map), + pl.BlockSpec((NUM_SUBLANES, bkv), kv_segment_ids_index_map), + ] + 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((None, bq, bkv), mask_index_map)) + 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((bq, NUM_LANES), q_segment_ids_index_map)) + 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_v), jnp.float32), # o_scratch + jax.ShapeDtypeStruct((num_q_heads, q_seq_len, head_dim_v), q.dtype), + ] + out_specs = [ + # TODO(sharadmv): convert m/l to be scratch + pl.BlockSpec((bq, NUM_LANES), lambda h, i, j, *_: (0, 0)), + pl.BlockSpec((bq, NUM_LANES), lambda h, i, j, *_: (0, 0)), + pl.BlockSpec((bq, head_dim_v), lambda h, i, j, *_: (0, 0)), + pl.BlockSpec((None, bq, head_dim_v), out_index_map), + ] + 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((None, bq, NUM_LANES), logsumexp_index_map), + ] + else: + out_shapes += [None] + out_specs += [None] + + kernel_name = get_kernel_name( + dataclasses.asdict(block_sizes), + is_mqa=is_mqa, + save_residuals=save_residuals, + is_segmented=segment_ids is not None, + phase="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_v=head_dim_v, + 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, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary"), + ), + out_shape=out_shapes, + name=kernel_name, + interpret=interpret, + )( + 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_argnames=( + "save_residuals", "mask_value", "is_mqa", "block_sizes", + "residual_checkpoint_name", "mask_function", "attn_logits_soft_cap", + "interpret") +) +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, + interpret: bool = False, +) -> 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, + residual_checkpoint_name=residual_checkpoint_name, + save_residuals=save_residuals, + mask_function=mask_function, + attn_logits_soft_cap=attn_logits_soft_cap, + interpret=interpret, + ) + + +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, + interpret: bool = False, +) -> tuple[ + tuple[jax.Array], + SplashResidualsType, +]: + if save_residuals: + raise NotImplementedError("Higher-order AD not supported") + + out, (logsumexp,) = _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, + residual_checkpoint_name=residual_checkpoint_name, + save_residuals=True, + mask_function=mask_function, + attn_logits_soft_cap=attn_logits_soft_cap, + interpret=interpret, + ) + 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) + + global_kv_index, _, 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), + # When the iteration space is shrunk (for local attention for example), + # the kv_index program_id does not correspond to the actual coordinates + # of the KV data. Make sure to use the 'unshrunk' index (coming from the + # data_next array) when computing the mask. + k_offset=global_kv_index * 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, + interpret: bool, +): + num_q_heads, q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + 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[:-1] != v.shape[:-1]: + raise ValueError( + f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " + "leading dimensions." + ) + + 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((None, bq, head_dim_v), o_index_map) + + def q_index_map(h, i, *_): + return from_head_minor((h, i, 0), q_layout) + + q_spec = pl.BlockSpec( + from_head_minor((None, bq, head_dim_qk), q_layout), q_index_map + ) + + 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( + from_head_minor( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), k_layout + ), + k_index_map, + ) + + 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( + from_head_minor( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), v_layout + ), + v_index_map, + ) + + 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((None, bq, bkv), mask_index_map) + + 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((bq, NUM_LANES), q_segment_ids_index_map) + kv_segment_spec = pl.BlockSpec( + (NUM_SUBLANES, bkv), kv_segment_ids_index_map + ) + 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 = o_spec + + def logsumexp_index_map(h, i, *_): + return h, 0, i + + logsumexp = jnp.expand_dims(logsumexp, axis=-2) + logsumexp_spec = pl.BlockSpec((None, 1, bq), logsumexp_index_map) + assert logsumexp.ndim == len(logsumexp_spec.block_shape) + + di = jnp.expand_dims(di, axis=-2) + di_spec = pl.BlockSpec((None, 1, bq), logsumexp_index_map) + 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((bq, NUM_LANES), q_segment_ids_index_map)) + else: + q_sequence = None + in_specs.append(None) + + out_shapes = [ + jax.ShapeDtypeStruct((bq, head_dim_qk), jnp.float32), + jax.ShapeDtypeStruct(q.shape, q.dtype), + ] + out_specs = [ + pl.BlockSpec((bq, head_dim_qk), lambda *_: (0, 0)), + pl.BlockSpec((None, bq, head_dim_qk), lambda h, i, *_: (h, i, 0)), + ] + + 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 + + kernel_name = get_kernel_name( + 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, + save_residuals=False, + is_segmented=segment_ids is not None, + phase="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, + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary", "arbitrary", "arbitrary"), + ), + name=kernel_name, + interpret=interpret, + )( + 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 ref[slice_k, :] + return ref[:, slice_k].T + k = _load_kv(k_ref, k_layout) + v = _load_kv(v_ref, v_layout) + logsumexp = logsumexp_ref[:1, :] + do = do_ref[...] + di = di_ref[:1, :] + + 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) + dv_scratch_ref[slice_k, :] + dv_scratch_ref[slice_k, :] = 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) + dk_scratch_ref[slice_k, :] + dk_scratch_ref[slice_k, :] = 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, + interpret: bool, +): + num_q_heads, q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + 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[:-1] != v.shape[:-1]: + raise ValueError( + f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " + "leading dimensions." + ) + + 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((None, bq, head_dim_v), o_index_map) + + 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( + from_head_minor((None, bq, head_dim_qk), q_layout), q_index_map + ) + + 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( + from_head_minor( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), + k_layout, + ), + k_index_map, + ) + + 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( + from_head_minor( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), + v_layout, + ), + v_index_map, + ) + + 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((None, None, bq, head_dim_qk), dq_index_map) + 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((bq, head_dim_qk), lambda *_: (0, 0)) + dq_scratch_shape = jax.ShapeDtypeStruct((bq, head_dim_qk), 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 = pl.BlockSpec( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), + dkv_index_map, + ) + + dv_spec = pl.BlockSpec( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), + dkv_index_map, + ) + + 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((None, bkv, bq), mask_index_map) + + 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((NUM_SUBLANES, bq), q_segment_ids_index_map) + kv_segment_spec = pl.BlockSpec((bkv, NUM_LANES), kv_segment_ids_index_map) + 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((None, NUM_SUBLANES, bq), logsumexp_index_map) + 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((None, NUM_SUBLANES, bq), logsumexp_index_map) + 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((NUM_SUBLANES, bq), q_segment_ids_index_map)) + 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_qk), jnp.float32), + jax.ShapeDtypeStruct((bkv, head_dim_v), jnp.float32), + dq_shape, + jax.ShapeDtypeStruct(k.shape, k.dtype), + jax.ShapeDtypeStruct(v.shape, v.dtype), + ] + out_specs = [ + dq_scratch_spec, + pl.BlockSpec((bkv, head_dim_qk), lambda *_: (0, 0)), + pl.BlockSpec((bkv, head_dim_v), lambda *_: (0, 0)), + 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 + + kernel_name = get_kernel_name( + 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, + ), + is_mqa=is_mqa, + save_residuals=False, + is_segmented=segment_ids is not None, + phase="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, + # 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 + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary", "arbitrary", "arbitrary"), + ), + name=kernel_name, + interpret=interpret, + )( + 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, + interpret: bool, + 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, + interpret=interpret, + ) + 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, + interpret=interpret, + ) + # 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", + "interpret", + ], +) +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, + interpret: bool, +) -> SplashCustomReturnType: + """ + For dynamic masks, `partial_mask_blocks` has shape (head_count, q_blocks, kv_blocks, block_q, block_kv). + This shape allows sharding across both head count and query sequence dimensions. + + Note: The leading dimensions (head_count, q_blocks, kv_blocks) must be + collapsed into a single dimension before being passed to the kernel. + """ + def _collapse_partial_mask_blocks(mask_info: mask_info_lib.MaskInfo | None): + if mask_info is None or mask_info.partial_mask_blocks is None: + return mask_info + + return mask_info._replace( + partial_mask_blocks=mask_info.partial_mask_blocks.reshape( + -1, *mask_info.partial_mask_blocks.shape[-2:] + ) + ) + + fwd_mask_info = _collapse_partial_mask_blocks(fwd_mask_info) + dq_mask_info = _collapse_partial_mask_blocks(dq_mask_info) + dkv_mask_info = _collapse_partial_mask_blocks(dkv_mask_info) + 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, + interpret=interpret, + ) + + +@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() + partial_mask_blocks_spec = ( + spec if self.fwd_mask_info.is_dynamic_mask else replicated + ) + # 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=partial_mask_blocks_spec + 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 | jax.Array | 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, + interpret: bool = False, +): + 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() + + process_mask_fn = ( + mask_info_lib.process_dynamic_mask + if isinstance(mask, jax.Array) + else mask_info_lib.process_mask + ) + + process_mask_dvk_fn = ( + mask_info_lib.process_dynamic_mask_dkv + if isinstance(mask, jax.Array) + else mask_info_lib.process_mask_dkv + ) + + fwd_mask_info, mask_function_fwd = process_mask_fn( + 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 = process_mask_fn( + 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 = process_mask_dvk_fn( + 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, + interpret=interpret, + ) + + +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 +) + + +CONFIG = { + 'name': 'llama3_405b_gqa_optimized', + 'model': 'Llama-3.1-405B', + 'operator': 'gqa_attention', + 'batch': 4, + 'seq_len': 4096, + 'num_query_heads': 128, + 'num_kv_heads': 8, + 'head_dim': 128, + 'emb_dim': 16384, +} + +# Tuned by autotune_block_sizes.py. Re-run to update. +TUNED_PARAMS = { + # Autotuned for 128 query heads / 8 KV heads. + 'block_q': 2048, + 'block_kv': 2048, + 'block_kv_compute': 1024, + 'q_layout': 1, # QKVLayout.HEAD_DIM_MINOR=1, SEQ_MINOR=2 + 'k_layout': 1, + 'v_layout': 1, + 'head_shards': 1, + 'q_seq_shards': 1, + # Not autotuned (backward-only). + 'block_q_dkv': None, + 'block_kv_dkv': None, + 'block_kv_dkv_compute': None, + 'block_q_dq': None, + 'block_kv_dq': None, +} + + +def get_flops(): + B, S = CONFIG['batch'], CONFIG['seq_len'] + Hq, D = CONFIG['num_query_heads'], CONFIG['head_dim'] + return 4 * B * Hq * S * S * D + + +def create_inputs(dtype=jnp.bfloat16): + """Returns (query, key, value) matching GQA baseline: (B, S, H, D) layout.""" + key = jax.random.key(42) + k1, k2, k3 = jax.random.split(key, 3) + B, S = CONFIG['batch'], CONFIG['seq_len'] + Hq, Hkv, D = CONFIG['num_query_heads'], CONFIG['num_kv_heads'], CONFIG['head_dim'] + query = jax.random.normal(k1, (B, S, Hq, D), dtype=dtype) + key_t = jax.random.normal(k2, (B, S, Hkv, D), dtype=dtype) + value = jax.random.normal(k3, (B, S, Hkv, D), dtype=dtype) + return query, key_t, value + + +def workload(query, key, value): + """GQA with Pallas splash attention (autotuned block sizes).""" + # Transpose from BSHD to BHSD for splash attention + q = query.transpose(0, 2, 1, 3) # (B, H_q, S, D) + k = key.transpose(0, 2, 1, 3) # (B, H_kv, S, D) + v = value.transpose(0, 2, 1, 3) # (B, H_kv, S, D) + + B, H_q, S, D = q.shape + q = q * (D ** -0.5) + H_kv = v.shape[1] + heads_per_group = H_q // H_kv + mask = mask_lib.CausalMask(shape=(S, S)) + multi_head_mask = mask_lib.MultiHeadMask([mask] * H_q) + block_sizes = BlockSizes( + block_q=TUNED_PARAMS['block_q'], + block_kv=TUNED_PARAMS['block_kv'], + block_kv_compute=TUNED_PARAMS['block_kv_compute'], + q_layout=QKVLayout(TUNED_PARAMS['q_layout']), + k_layout=QKVLayout(TUNED_PARAMS['k_layout']), + v_layout=QKVLayout(TUNED_PARAMS['v_layout']), + block_q_dkv=TUNED_PARAMS['block_q_dkv'], + block_kv_dkv=TUNED_PARAMS['block_kv_dkv'], + block_kv_dkv_compute=TUNED_PARAMS['block_kv_dkv_compute'], + block_q_dq=TUNED_PARAMS['block_q_dq'], + block_kv_dq=TUNED_PARAMS['block_kv_dq'], + ) + splash_kernel = _make_splash_attention( + multi_head_mask, block_sizes=block_sizes, + is_mqa=False, + head_shards=TUNED_PARAMS['head_shards'], + q_seq_shards=TUNED_PARAMS['q_seq_shards'], + ) + @jax.vmap + def _attend(q_batch, k_batch, v_batch): + k_repeated = jnp.repeat(k_batch, heads_per_group, axis=0) + v_repeated = jnp.repeat(v_batch, heads_per_group, axis=0) + return splash_kernel(q_batch, k_repeated, v_repeated) + out = _attend(q, k, v) # (B, H_q, S, D) + return out.transpose(0, 2, 1, 3) # (B, S, H_q, D) to match baseline + + +def benchmark(num_warmup=5, num_iters=100): + """Benchmark and return results dict.""" + inputs = create_inputs() + fn = jax.jit(workload) + for _ in range(num_warmup): + out = fn(*inputs) + out.block_until_ready() + times = [] + for _ in range(num_iters): + t0 = time.perf_counter() + out = fn(*inputs) + out.block_until_ready() + times.append(time.perf_counter() - t0) + times = np.array(times) * 1000 + avg = float(np.mean(times)) + return { + 'name': CONFIG['name'], + 'model': CONFIG['model'], + 'operator': CONFIG['operator'], + 'config': {k: v for k, v in CONFIG.items() if k not in ('name', 'model', 'operator', 'atol', 'rtol')}, + 'time_ms': round(avg, 4), + 'std_ms': round(float(np.std(times)), 4), + 'output_shape': list(out.shape) if hasattr(out, 'shape') else [], + 'status': 'success', + } + + +if __name__ == '__main__': + import json + print(json.dumps(benchmark())) diff --git a/JAXBench/benchmark/level2/3p_MLA_Attention/baseline.py b/JAXBench/benchmark/level2/3p_MLA_Attention/baseline.py new file mode 100644 index 0000000..157632f --- /dev/null +++ b/JAXBench/benchmark/level2/3p_MLA_Attention/baseline.py @@ -0,0 +1,1615 @@ +# Copyright 2025 Google LLC +# +# 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 +# +# http://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. +"""TPU-Friendly MLA Ragged Paged Attention kernel.""" + +import functools + +import time +import numpy as np +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) + +DEFAULT_VMEM_LIMIT_BYTES = 100 * 1024 * 1024 + + +def cdiv(a, b): + assert b != 0 + return (a + b - 1) // b + + +def align_to(x, a): + return cdiv(x, a) * a + + +def get_dtype_bitwidth(dtype): + return jax.dtypes.itemsize_bits(dtype) + + +def get_dtype_packing(dtype): + bits = get_dtype_bitwidth(dtype) + return 32 // bits + + +def get_kv_cache_shape( + total_num_pages, + page_size, + kv_dim, + kv_dtype, +): + kv_packing = get_dtype_packing(kv_dtype) + return ( + total_num_pages, + align_to(page_size, kv_packing) // kv_packing, + kv_packing, + align_to(kv_dim, 128), + ) + + +# Expect to run this validation during compile time. +def static_validate_inputs( + ql_nope: jax.Array, # [max_num_tokens, actual_num_q_heads, actual_lkv_dim] + q_pe: jax.Array, # [max_num_tokens, actual_num_q_heads, actual_r_dim] + new_kv_c: jax.Array, # [max_num_tokens, actual_lkv_dim] + new_k_pe: jax.Array, # [max_num_tokens, actual_r_dim] + cache_kv: jax.Array, # [total_num_pages, page_size_per_kv_packing, kv_packing, lkv_dim] + kv_lens: jax.Array, # i32[max_num_seqs] + page_indices: jax.Array, # i32[max_num_seqs * pages_per_seq] + cu_q_lens: jax.Array, # i32[max_num_seqs + 1] + distribution: jax.Array, # i32[3] + *, + sm_scale: float = 1.0, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = DEFAULT_MASK_VALUE, + q_scale: float | None = None, + k_scale: float | None = None, + v_scale: float | None = None, + # Kernel optimization params. + chunk_prefill_size: int | None = None, + # Kernel tuning params. + num_kv_pages_per_block: int | None = None, + num_queries_per_block: int | None = None, + vmem_limit_bytes: int | None = None, + # Debug params. + debug_mode: bool = False, +): + """Validate inputs to the MLA RPA kernel statically.""" + if len(ql_nope.shape) != 3: + raise ValueError(f"Expected 3D array for {ql_nope.shape=}") + if len(q_pe.shape) != 3: + raise ValueError(f"Expected 3D array for {q_pe.shape=}") + if len(new_kv_c.shape) != 2: + raise ValueError(f"Expected 2D array for {new_kv_c.shape=}") + if len(new_k_pe.shape) != 2: + raise ValueError(f"Expected 2D array for {new_k_pe.shape=}") + + if ql_nope.shape[:2] != q_pe.shape[:2]: + raise ValueError( + f"Expected {ql_nope.shape[:2]=} to be equal to {q_pe.shape[:2]=}") + if ql_nope.shape[0] != new_kv_c.shape[0]: + raise ValueError( + f"Expected {ql_nope.shape[0]=} to be equal to {new_kv_c.shape[0]=}" + ) + if new_kv_c.shape[0] != new_k_pe.shape[0]: + raise ValueError( + f"Expected {new_kv_c.shape[0]=} to be equal to {new_k_pe.shape[0]=}" + ) + if ql_nope.shape[2] != new_kv_c.shape[1]: + raise ValueError( + f"Expected {ql_nope.shape[2]=} to be equal to {new_kv_c.shape[1]=}" + ) + if q_pe.shape[2] != new_k_pe.shape[1]: + raise ValueError( + f"Expected {q_pe.shape[2]=} to be equal to {new_k_pe.shape[1]=}") + + actual_lkv_dim = ql_nope.shape[2] + actual_r_dim = q_pe.shape[2] + lkv_dim = align_to(actual_lkv_dim, 128) + r_dim = align_to(actual_r_dim, 128) + + ( + _, + page_size_per_kv_packing, + kv_packing, + kv_dim, + ) = cache_kv.shape + + if lkv_dim + r_dim != kv_dim: + raise ValueError( + f"Expected {lkv_dim=} + {r_dim=} to be equal to {kv_dim=}") + + if not (cache_kv.dtype == new_kv_c.dtype): + raise ValueError( + f"Expected {cache_kv.dtype=} to be equal to {new_kv_c.dtype=}.") + if not (cache_kv.dtype == new_k_pe.dtype): + raise ValueError( + f"Expected {cache_kv.dtype=} to be equal to {new_k_pe.dtype=}.") + + # Integer kv quantization is currently not supported. + if not jnp.issubdtype(cache_kv.dtype, jnp.floating): + raise ValueError(f"Expected {cache_kv.dtype=} to be a floating point.") + + if kv_packing != get_dtype_packing(cache_kv.dtype): + raise ValueError( + f"{kv_packing=} does not match with {cache_kv.dtype=}") + + if not (jnp.int32 == kv_lens.dtype == page_indices.dtype == cu_q_lens.dtype + == distribution.dtype): + raise ValueError( + f"Expected int32 dtype for {kv_lens.dtype=}, {page_indices.dtype=}," + f" {cu_q_lens.dtype=}, {distribution.dtype=}") + + if not (len(kv_lens.shape) == len(page_indices.shape) == len( + cu_q_lens.shape) == 1): + raise ValueError( + f"Expected 1D array for {kv_lens.shape=}, {page_indices.shape=}," + f" {cu_q_lens.shape=}") + + max_num_seqs = kv_lens.shape[0] + num_page_indices = page_indices.shape[0] + if num_page_indices % max_num_seqs != 0: + raise ValueError( + f"Expected {num_page_indices=} to be divisible by {max_num_seqs=}." + ) + if cu_q_lens.shape != (max_num_seqs + 1, ): + raise ValueError( + f"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},).") + if distribution.shape != (3, ): + raise ValueError(f"Expected {distribution.shape=} to be (3,).") + + page_size = page_size_per_kv_packing * kv_packing + if page_size % kv_packing != 0: + raise ValueError(f"{page_size=} must be divisible by {kv_packing=}.") + if sliding_window is not None and sliding_window <= 0: + raise ValueError(f"{sliding_window=} must be positive.") + if soft_cap is not None and soft_cap == 0.0: + raise ValueError(f"{soft_cap=} must not be 0.0.") + if chunk_prefill_size is not None and chunk_prefill_size <= 0: + raise ValueError(f"{chunk_prefill_size=} must be positive.") + if num_kv_pages_per_block is not None: + if num_kv_pages_per_block <= 0: + raise ValueError(f"{num_kv_pages_per_block=} must be positive.") + if num_queries_per_block is not None: + if num_queries_per_block <= 0: + raise ValueError(f"{num_queries_per_block=} must be positive.") + if vmem_limit_bytes is not None and vmem_limit_bytes <= 0: + raise ValueError(f"{vmem_limit_bytes=} must be positive.") + + # No constraints for the following inputs. + del sm_scale + del mask_value + del q_scale + del k_scale + del v_scale + del debug_mode + + +def _mla_ragged_paged_attention_kernel( + # Prefetch + kv_lens_ref, # [max_num_seqs] + page_indices_ref, # [max_num_seqs * pages_per_seq] + cu_q_lens_ref, # [max_num_seqs + 1] + distribution_ref, # [3] (decode_end, prefill_end, mixed_end) + sem_ids_ref, # [3] (bq_sem_idx, bkv_sem_idx, bo_sem_idx) + bo_ids_ref, # [4] (bo_sem_0_seq_idx, bo_sem_1_seq_idx, bo_sem_0_bo_idx, bo_sem_1_bo_idx) + bkv_update_ids_ref, # [6] (bkv_sem_0_seq_idx, bkv_sem_1_seq_idx, bkv_sem_0_offset, bkv_sem_1_offset, bkv_sem_0_sz, bkv_sem_1_sz) + # Input + ql_nope_hbm_ref, # [max_num_tokens, num_q_heads_per_q_packing, q_packing, lkv_dim] + q_pe_hbm_ref, # [max_num_tokens, num_q_heads_per_q_packing, q_packing, r_dim] + new_kv_c_hbm_ref, # [max_num_tokens_per_kv_packing, kv_packing, lkv_dim] + new_k_pe_hbm_ref, # [max_num_tokens_per_kv_packing, kv_packing, r_dim] + cache_kv_hbm_ref, # [total_num_pages, page_size_per_kv_packing, kv_packing, align_to(lkv_dim + r_dim, 128)] + # Output + o_hbm_ref, # [max_num_tokens, num_q_heads_per_q_packing, q_packing, lkv_dim] + updated_cache_kv_hbm_ref, # [total_num_pages, page_size_per_kv_packing, kv_packing, align_to(lkv_dim + r_dim, 128)] + # Scratch + bkvc_x2_ref, # [2, bkv_buf_sz_per_kv_packing, kv_packing, lkv_dim] + bkpe_x2_ref, # [2, bkv_buf_sz_per_kv_packing, kv_packing, r_dim] + bq_nope_x2_ref, # [2, bq_sz, num_q_heads_per_q_packing, q_packing, lkv_dim] + bq_rope_x2_ref, # [2, bq_sz, num_q_heads_per_q_packing, q_packing, r_dim] + bo_x2_ref, # [2, bq_sz, num_q_heads_per_q_packing, q_packing, lkv_dim] + sems, # [4, 2] + l_ref, # [bq_sz * num_q_heads, 128], + m_ref, # [bq_sz * num_q_heads, 128], + acc_ref, # [bq_sz * num_q_heads, lkv_dim], + *, + sm_scale: float, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float = DEFAULT_MASK_VALUE, + q_scale: float | None = None, + k_scale: float | None = None, + v_scale: float | None = None, + chunk_prefill_size: int | None = None, + bkv_p, + bq_sz, + debug_mode: bool = False, +): + assert ql_nope_hbm_ref.shape == o_hbm_ref.shape + # Validation checks on the dimensions + nope_dim = ql_nope_hbm_ref.shape[-1] + pe_dim = q_pe_hbm_ref.shape[-1] + assert nope_dim + pe_dim == cache_kv_hbm_ref.shape[-1] + + _, num_q_heads_per_q_packing, q_packing, lkv_dim = ql_nope_hbm_ref.shape + r_dim = q_pe_hbm_ref.shape[-1] + num_q_heads = num_q_heads_per_q_packing * q_packing + total_num_pages, page_size_per_kv_packing, kv_packing, _ = ( + cache_kv_hbm_ref.shape) + max_num_seqs = kv_lens_ref.shape[0] + num_page_indices = page_indices_ref.shape[0] + + assert num_page_indices % max_num_seqs == 0 + pages_per_seq = num_page_indices // max_num_seqs + q_dtype = ql_nope_hbm_ref.dtype + # Validate against the KV dtype. + kv_dtype = cache_kv_hbm_ref.dtype + assert q_pe_hbm_ref.dtype == q_dtype + assert o_hbm_ref.dtype == q_dtype + assert get_dtype_packing(q_dtype) == q_packing + assert get_dtype_packing(kv_dtype) == kv_packing + assert lkv_dim % 128 == 0 + assert r_dim % 128 == 0 + bkv_sz_per_kv_packing = bkv_p * page_size_per_kv_packing + bkv_sz = bkv_sz_per_kv_packing * kv_packing + page_size = page_size_per_kv_packing * kv_packing + seq_idx = pl.program_id(0) + num_seqs = pl.num_programs(0) + decode_end = distribution_ref[0] + prefill_end = distribution_ref[1] + mixed_end = distribution_ref[2] + + q_start = cu_q_lens_ref[seq_idx] + q_end = cu_q_lens_ref[seq_idx + 1] + q_len = q_end - q_start + kv_len = kv_lens_ref[seq_idx] + + def debug_print(msg, *args): + if debug_mode: + pl.debug_print(msg, *args) + + debug_print("[RPA debug] ======= In loop seq_idx={}", seq_idx) + debug_print("[RPA debug] num_seqs={}", num_seqs) + debug_print("[RPA debug] decode_end={}", decode_end) + debug_print("[RPA debug] prefill_end={}", prefill_end) + debug_print("[RPA debug] mixed_end={}", mixed_end) + debug_print("[RPA debug] bkv_p={}", bkv_p) + debug_print("[RPA debug] page_size={}", page_size) + debug_print("[RPA debug] pages_per_seq={}", pages_per_seq) + debug_print("[RPA debug] bkv_sz_per_kv_packing={}", bkv_sz_per_kv_packing) + debug_print("[RPA debug] bq_sz={}", bq_sz) + debug_print("[RPA debug] q_start={}", q_start) + debug_print("[RPA debug] q_end={}", q_end) + debug_print("[RPA debug] q_len={}", q_len) + debug_print("[RPA debug] kv_len={}", kv_len) + + def flash_attention( + ql_nope, # [actual_bq_sz * num_q_heads, lkv_dim] + q_pe, # [actual_bq_sz * num_q_heads, r_dim] + kv_c, # [bkv_sz, lkv_dim] <- Correspond to data from bkvc_x2_ref + k_pe, # [bkv_sz, r_dim] <- Correspond to data from bpe_x2_ref + *, + bq_idx, + bkv_idx, + ): + assert len(ql_nope.shape) == 2 + assert len(q_pe.shape) == 2 + assert len(kv_c.shape) == 2 + assert len(k_pe.shape) == 2 + assert ql_nope.shape[0] % num_q_heads == 0 + assert ql_nope.shape[0] == q_pe.shape[0] + assert q_pe.shape[0] % bq_sz == 0 + assert ql_nope.shape[1] == lkv_dim + assert q_pe.shape[1] == r_dim + assert kv_c.shape == (bkv_sz, lkv_dim) + assert k_pe.shape == (bkv_sz, r_dim) + head_l_ref = l_ref.at[:ql_nope.shape[0]] + head_m_ref = m_ref.at[:ql_nope.shape[0]] + head_acc_ref = acc_ref.at[:ql_nope.shape[0]] + + def load_with_init(ref, init_val): + return jnp.where(bkv_idx == 0, jnp.full_like(ref, init_val), + ref[...]) + + # Follow FlashAttention-2 forward pass. + s_nope = jnp.einsum("nd,md->nm", + ql_nope, + kv_c, + preferred_element_type=jnp.float32) + s_pe = jnp.einsum("nd,md->nm", + q_pe, + k_pe, + preferred_element_type=jnp.float32) + s = s_nope + s_pe + s *= sm_scale + if k_scale is not None: + s *= k_scale + if q_scale is not None: + s *= q_scale + + q_span = (kv_len - q_len + bq_idx * bq_sz + + lax.broadcasted_iota(jnp.int32, s.shape, 0) // num_q_heads) + k_span = bkv_idx * bkv_sz + lax.broadcasted_iota(jnp.int32, s.shape, 1) + mask = q_span < k_span + if sliding_window is not None: + mask = jnp.logical_or(mask, q_span - sliding_window >= k_span) + + if soft_cap is not None: + s = soft_cap * jnp.tanh(s / soft_cap) + s = jnp.where(mask, mask_value, s) + s_rowmax = jnp.max(s, axis=1, keepdims=True) + m_prev = load_with_init(head_m_ref, -jnp.inf) + m_curr = jnp.maximum(m_prev, s_rowmax) + head_m_ref[...] = m_curr + p = jnp.exp(s - broadcast_minor(m_curr, s.shape)) + + pv = jnp.einsum("nm,md->nd", + p, + kv_c, + preferred_element_type=jnp.float32) + if v_scale is not None: + pv *= v_scale + + p_rowsum = jnp.sum(p, axis=1, keepdims=True) + exp_m_diff = jnp.exp(m_prev - m_curr) + l_prev = load_with_init(head_l_ref, 0.0) + l_curr = exp_m_diff * l_prev + p_rowsum + head_l_ref[...] = l_curr + o_prev = load_with_init(head_acc_ref, 0.0) + o_curr = broadcast_minor(exp_m_diff, o_prev.shape) * o_prev + pv + head_acc_ref[...] = o_curr + + def _async_copy(src, dst, sem, wait): + if debug_mode: + # Skip DMA if debug mode is enabled. + return + cp = pltpu.make_async_copy(src, dst, sem) + if wait: + cp.wait() + else: + cp.start() + + def _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, *, wait=False): + sem = sems.at[0, bkv_sem_idx] + # bkvc_x2_ref shape: [2, bkv_sz_per_kv_packing + 2, kv_packing, lkv_dim] + bkvc_vmem_ref = bkvc_x2_ref.at[bkv_sem_idx] + bkvpe_vmem_ref = bkpe_x2_ref.at[bkv_sem_idx] + + # [total_num_pages, page_size_per_kv_packing, kv_packing, align_to(lkv_dim + r_dim, 128)] + # [total_num_pages * page_size_per_kv_packing, kv_packing, align_to(lkv_dim + r_dim, 128)] + reshaped_cache_hbm_ref = cache_kv_hbm_ref.reshape( + total_num_pages * page_size_per_kv_packing, + *cache_kv_hbm_ref.shape[2:], + ) + + kv_len = kv_lens_ref[seq_idx] + kv_len_start = bkv_idx * bkv_sz + kv_p_start = bkv_idx * bkv_p + + q_start = cu_q_lens_ref[seq_idx] + q_end = cu_q_lens_ref[seq_idx + 1] + q_len = q_end - q_start + + kv_left = kv_len - kv_len_start + kv_left_frm_cache = jnp.maximum(kv_left - q_len, 0) + kv_left_frm_cache_per_kv_packing = cdiv(kv_left_frm_cache, kv_packing) + kv_left_frm_new = kv_left - kv_left_frm_cache + + bkv_sz_frm_cache = jnp.minimum(kv_left_frm_cache, bkv_sz) + bkv_sz_frm_new = jnp.minimum(bkv_sz - bkv_sz_frm_cache, + kv_left_frm_new) + bkv_sz_frm_cache_per_kv_packing = cdiv(bkv_sz_frm_cache, kv_packing) + bkv_sz_frm_new_per_kv_packing = cdiv(bkv_sz_frm_new, kv_packing) + page_indices_offset = seq_idx * pages_per_seq + kv_p_start + + new_kv_len_start = q_end - kv_left_frm_new + new_kv_len_start_per_kv_packing = new_kv_len_start // kv_packing + bkv_sz_frm_new_kv_packing_to_fetch = jnp.where( + bkv_sz_frm_new > 0, + cdiv(new_kv_len_start + bkv_sz_frm_new, kv_packing) - + new_kv_len_start_per_kv_packing, + 0, + ) + dma_bkv_sz = (bkv_sz_frm_cache_per_kv_packing + + bkv_sz_frm_new_kv_packing_to_fetch) + + debug_print( + "[RPA debug]" + f" -----------{'wait' if wait else 'start'}_fetch_bkv-----------") + debug_print("[RPA debug] seq_idx={}", seq_idx) + debug_print("[RPA debug] bkv_idx={}", bkv_idx) + debug_print("[RPA debug] bkv_sem_idx={}", bkv_sem_idx) + debug_print("[RPA debug] kv_len_start={}", kv_len_start) + debug_print("[RPA debug] kv_p_start={}", kv_p_start) + debug_print("[RPA debug] kv_left={}", kv_left) + debug_print("[RPA debug] kv_left_frm_cache={}", kv_left_frm_cache) + debug_print("[RPA debug] kv_left_frm_new={}", kv_left_frm_new) + debug_print("[RPA debug] bkv_sz_frm_cache={}", bkv_sz_frm_cache) + debug_print( + "[RPA debug] bkv_sz_frm_cache_per_kv_packing={}", + bkv_sz_frm_cache_per_kv_packing, + ) + debug_print( + "[RPA debug] bkv_sz_frm_new_per_kv_packing={}", + bkv_sz_frm_new_per_kv_packing, + ) + debug_print("[RPA debug] page_indices_offset={}", page_indices_offset) + debug_print(f"[RPA debug] bkvc_vmem_ref.shape: {bkvc_vmem_ref.shape}") + debug_print( + f"[RPA debug] bkvpe_vmem_ref.shape: {bkvpe_vmem_ref.shape}") + + if not wait: + # Make sure the current bkv buffer is safe to overwrite. + wait_update_kv_cache(bkv_sem_idx) + + # Fetch effective kv from kv cache. To pipeline multiple DMA calls, we + # utilize static for loop instead of dynamic for loop. + # Loop through all pages in a block + for i in range(bkv_p): + # Ensure only effective kvs are copied and we don't go negative. + sz_per_kv_packing = jnp.clip( + kv_left_frm_cache_per_kv_packing - + i * page_size_per_kv_packing, + 0, + page_size_per_kv_packing, + ) + # If the page index is out of bound, we set page_idx to the last page. + # And there will be no copy since sz will be 0. + page_idx = jnp.minimum(page_indices_offset + i, + num_page_indices - 1) + _async_copy( + reshaped_cache_hbm_ref.at[ + pl.ds( + page_indices_ref[page_idx] * + page_size_per_kv_packing, + sz_per_kv_packing, + ), + ..., + :nope_dim, + ], + # [bkv_sz_per_kv_packing + 2, kv_packing, lkv_dim]. + bkvc_vmem_ref.at[pl.ds(i * page_size_per_kv_packing, + sz_per_kv_packing)], + sem, + wait, + ) + _async_copy( + reshaped_cache_hbm_ref.at[ + pl.ds( + page_indices_ref[page_idx] * + page_size_per_kv_packing, + sz_per_kv_packing, + ), + ..., + nope_dim:, + ], + # [bkv_sz_per_kv_packing + 2, kv_packing, r_dim]. + bkvpe_vmem_ref.at[pl.ds(i * page_size_per_kv_packing, + sz_per_kv_packing)], + sem, + wait, + ) + debug_print( + "[RPA debug] loop_body bkv_p={}, i={}, page_size_per_kv_packing={}," + " sz_per_kv_packing={}, page_idx={}, page_indices_ref[page_idx]={}", + bkv_p, + i, + page_size_per_kv_packing, + sz_per_kv_packing, + page_idx, + page_indices_ref[page_idx], + ) + + # Fetch new KVs by appending to the existing vmem buffers. + # Fetch either up to the end of the buffer or kv_left_frm_new, whichever + # is smaller. Since DMAs are word-aligned based on kv_packing, and the + # boundary between the old cache and the new KV tokens might not be + # word-aligned, we append the new KV words right after the last word + # containing old cache data. This can create "holes" (misalignments + # within the words), which we will shift and pack correctly later. + debug_print("[RPA debug] new_kv_len_start={}", new_kv_len_start) + debug_print( + "[RPA debug] new_kv_len_start_per_kv_packing={}", + new_kv_len_start_per_kv_packing, + ) + debug_print(f"new_kv_c_hbm_ref.shape: {new_kv_c_hbm_ref.shape}") + debug_print(f"new_k_pe_hbm_ref.shape: {new_k_pe_hbm_ref.shape}") + _async_copy( + new_kv_c_hbm_ref.at[pl.ds( + new_kv_len_start_per_kv_packing, + bkv_sz_frm_new_kv_packing_to_fetch, + )], + bkvc_vmem_ref.at[pl.ds( + bkv_sz_frm_cache_per_kv_packing, + bkv_sz_frm_new_kv_packing_to_fetch, + )], + sem, + wait, + ) + _async_copy( + new_k_pe_hbm_ref.at[pl.ds( + new_kv_len_start_per_kv_packing, + bkv_sz_frm_new_kv_packing_to_fetch, + )], + bkvpe_vmem_ref.at[pl.ds( + bkv_sz_frm_cache_per_kv_packing, + bkv_sz_frm_new_kv_packing_to_fetch, + )], + sem, + wait, + ) + + else: + # When we wait, we can use a dummy copy to wait for DMAs to complete where + # src == dst. However, the dma size must be correct. + dst_kvc = bkvc_vmem_ref.at[pl.ds(0, dma_bkv_sz)] + _async_copy( + src=dst_kvc, + dst=dst_kvc, + sem=sem, + wait=True, + ) + dst_kvpe = bkvpe_vmem_ref.at[pl.ds(0, dma_bkv_sz)] + _async_copy( + src=dst_kvpe, + dst=dst_kvpe, + sem=sem, + wait=True, + ) + + # This returns the (offset, size) in units of tokens: + # offset: starting token index where the new KV should be stored + # size: number of tokens of the new KV, which is 1 in decode. + return kv_len_start + bkv_sz_frm_cache, bkv_sz_frm_new + + def _pack_new_kv(bkv_sem_idx, offset, update_sz): + """Packs newly computed KVs into the correct sub-word alignment in VMEM. + + When new KV tokens are DMA'd from HBM into VMEM, they are copied at the + granularity of packed words (e.g., 4 tokens per word for fp8) by head + dimension (mapped to lanes). The starting token `offset` in the KV cache, + however, might not fall exactly on a word boundary. This means the elements + within the packed words might be misaligned relative to their final + destination in the cache. + + This function corrects this alignment by: + 1. Computing the bit-shift amount needed based on the difference between the + destination token offset (`kv_packing_offset`) and the source token + offset (`new_kv_packing_offset`). + 2. Looping over the affected words and using bitwise shifts and logical ORs + to realign the elements across word boundaries. + 3. Merging the correctly aligned new KV elements into the VMEM buffer using + a mask, leaving existing (older) KV elements intact. + + Args: + bkv_sem_idx: The semaphore index for the current KV block. + offset: The starting token offset in the KV cache where the new KVs begin. + update_sz: The number of new tokens to be packed. + """ + # shape: [bkv_sz_per_kv_packing + 2, kv_packing, lkv_dim] + bkvc_vmem_ref = bkvc_x2_ref.at[bkv_sem_idx] + # shape: [bkv_sz_per_kv_packing + 2, kv_packing, r_dim] + bkvpe_vmem_ref = bkpe_x2_ref.at[bkv_sem_idx] + + update_kv_packing_iters = cdiv((offset % kv_packing) + update_sz, + kv_packing) + kv_packing_offset = offset % kv_packing + new_kv_len_start = q_end - kv_len + offset + new_kv_packing_offset = new_kv_len_start % kv_packing + + token_offset_in_bkv = offset % bkv_sz + kv_packing_idx = token_offset_in_bkv // kv_packing + + # Compute the shift amount for each word in bits + shift_amount = kv_packing_offset - new_kv_packing_offset + bits_per_element = get_dtype_bitwidth(bkvc_vmem_ref.dtype) + shift_bits = bits_per_element * (shift_amount % kv_packing) + shift_bits = shift_bits.astype(jnp.uint32) + + # Calculate the starting index in the KV buffer corresponding to the new KV + # to fetch the data from. This index accounts for the potential offset + # caused by the shift_amount. + # (-shift_amount) // kv_packing will be: + # 0 if new_kv_packing_offset <= kv_packing_offset + # -1 if new_kv_packing_offset > kv_packing_offset. + kv_packing_idx_new = (cdiv(token_offset_in_bkv, kv_packing) + + (-shift_amount) // kv_packing) + curr_kvc_reg = bkvc_vmem_ref[kv_packing_idx_new, :, :] + curr_kpe_reg = bkvpe_vmem_ref[kv_packing_idx_new, :, :] + next_kvc_reg = bkvc_vmem_ref[kv_packing_idx_new + 1, :, :] + next_kpe_reg = bkvpe_vmem_ref[kv_packing_idx_new + 1, :, :] + + def merge_loop_body(i, vals): + ( + kv_packing_idx, + kv_packing_idx_new, + curr_kvc_reg, + curr_kpe_reg, + next_kvc_reg, + next_kpe_reg, + ) = vals + curr_kvc_reg_u32 = pltpu.bitcast(curr_kvc_reg, jnp.uint32) + curr_kpe_reg_u32 = pltpu.bitcast(curr_kpe_reg, jnp.uint32) + next_kvc_reg_u32 = pltpu.bitcast(next_kvc_reg, jnp.uint32) + next_kpe_reg_u32 = pltpu.bitcast(next_kpe_reg, jnp.uint32) + + shifted_kvc_u32 = lax.bitwise_or( + lax.shift_right_logical(curr_kvc_reg_u32, 32 - shift_bits), + lax.shift_left(next_kvc_reg_u32, shift_bits), + ) + shifted_kpe_u32 = lax.bitwise_or( + lax.shift_right_logical(curr_kpe_reg_u32, 32 - shift_bits), + lax.shift_left(next_kpe_reg_u32, shift_bits), + ) + + # If shift_bits is 0, we should use the current word. Otherwise, + # shifting by 32 bits would result in shifted_*_u32 becoming + # next_*_reg_u32, which is incorrect. + rotated_kvc_u32 = lax.select(shift_bits == 0, curr_kvc_reg_u32, + shifted_kvc_u32) + rotated_kpe_u32 = lax.select(shift_bits == 0, curr_kpe_reg_u32, + shifted_kpe_u32) + + next_kvc_reg_shifted = pltpu.bitcast(rotated_kvc_u32, + next_kvc_reg.dtype) + next_kpe_reg_shifted = pltpu.bitcast(rotated_kpe_u32, + next_kpe_reg.dtype) + + offset_in_word = i * kv_packing + lax.broadcasted_iota( + dtype=jnp.int32, shape=[kv_packing, lkv_dim], dimension=0) + kvc_mask = jnp.logical_and( + offset_in_word >= kv_packing_offset, + offset_in_word < kv_packing_offset + update_sz, + ) + updated_kvc_reg = lax.select( + kvc_mask, + next_kvc_reg_shifted, + bkvc_vmem_ref[kv_packing_idx, :, :], + ) + offset_in_word_pe = i * kv_packing + lax.broadcasted_iota( + dtype=jnp.int32, shape=[kv_packing, r_dim], dimension=0) + kpe_mask = jnp.logical_and( + offset_in_word_pe >= kv_packing_offset, + offset_in_word_pe < kv_packing_offset + update_sz, + ) + updated_kpe_reg = lax.select( + kpe_mask, + next_kpe_reg_shifted, + bkvpe_vmem_ref[kv_packing_idx, :, :], + ) + + # Store back the merged word + bkvc_vmem_ref[kv_packing_idx, :, :] = updated_kvc_reg + bkvpe_vmem_ref[kv_packing_idx, :, :] = updated_kpe_reg + + # Move to the next word. + kv_packing_idx += 1 + kv_packing_idx_new += 1 + curr_kvc_reg = next_kvc_reg + curr_kpe_reg = next_kpe_reg + next_kvc_reg = bkvc_vmem_ref[kv_packing_idx_new + 1, :, :] + next_kpe_reg = bkvpe_vmem_ref[kv_packing_idx_new + 1, :, :] + return ( + kv_packing_idx, + kv_packing_idx_new, + curr_kvc_reg, + curr_kpe_reg, + next_kvc_reg, + next_kpe_reg, + ) + + lax.fori_loop( + 0, + update_kv_packing_iters, + merge_loop_body, + ( + kv_packing_idx, + kv_packing_idx_new, + curr_kvc_reg, + curr_kpe_reg, + next_kvc_reg, + next_kpe_reg, + ), + ) + + def _update_kv_cache( + seq_idx, + bkv_sem_idx, + offset, # In units of tokens. + update_sz, # In units of tokens. + *, + wait=False, + ): + sem = sems.at[3, bkv_sem_idx] + # shape: [bkv_sz_per_kv_packing + 2, kv_packing, lkv_dim] + bkvc_vmem_ref = bkvc_x2_ref.at[bkv_sem_idx] + # shape: [bkv_sz_per_kv_packing + 2, kv_packing, r_dim] + bkvpe_vmem_ref = bkpe_x2_ref.at[bkv_sem_idx] + + update_kv_packing_iters = cdiv((offset % kv_packing) + update_sz, + kv_packing) + + # Expected shape: + # [total_num_pages, page_size_per_kv_packing, kv_packing, + # align_to(lkv_dim + r_dim, 128)] + cache_kv_hbm_shape = updated_cache_kv_hbm_ref.shape + reshaped_cache_kv_hbm_ref = updated_cache_kv_hbm_ref.reshape( + cache_kv_hbm_shape[0] * cache_kv_hbm_shape[1], + *cache_kv_hbm_shape[2:], + ) + + if not wait: + # Issue DMA copy for the updated parts, page by page. + kv_p_start = offset // page_size + kv_p_end = cdiv(offset + update_sz, page_size) + start_word_in_page = (offset % page_size) // kv_packing + start_word_in_vmem = (offset % bkv_sz) // kv_packing + words_to_transfer = update_kv_packing_iters + page_indices_offset = seq_idx * pages_per_seq + kv_p_start + + def loop_body(i, states): + curr_word_in_page, words_to_transfer, curr_word_in_vmem = states + sz_words = jnp.minimum( + page_size_per_kv_packing - curr_word_in_page, + words_to_transfer) + page_idx = page_indices_ref[page_indices_offset + i] + + _async_copy( + # bkvc_vmem_ref shape: + # [bkv_sz_per_kv_packing+2, kv_packing, lkv_dim] + bkvc_vmem_ref.at[pl.ds(curr_word_in_vmem, sz_words)], + reshaped_cache_kv_hbm_ref.at[ + pl.ds( + page_idx * page_size_per_kv_packing + + curr_word_in_page, + sz_words, + ), + ..., + :nope_dim, + ], + sem, + wait=False, + ) + _async_copy( + # bkvpe_vmem_ref shape: [bkv_sz_per_kv_packing+2, kv_packing, r_dim] + bkvpe_vmem_ref.at[pl.ds(curr_word_in_vmem, sz_words)], + reshaped_cache_kv_hbm_ref.at[ + pl.ds( + page_idx * page_size_per_kv_packing + + curr_word_in_page, + sz_words, + ), + ..., + nope_dim:, + ], + sem, + wait=False, + ) + return 0, words_to_transfer - sz_words, curr_word_in_vmem + sz_words + + lax.fori_loop( + 0, + kv_p_end - kv_p_start, + loop_body, + ( + start_word_in_page, + words_to_transfer, + start_word_in_vmem, + ), # initial states + unroll=False, + ) + else: # Wait + dma_sz_words = update_kv_packing_iters + # bkvc_vmem_ref shape: [bkv_sz_per_kv_packing + 2, kv_packing, lkv_dim] + dst_kv = bkvc_vmem_ref.at[pl.ds(0, dma_sz_words)] + _async_copy( + src=dst_kv, + dst=dst_kv, + sem=sem, + wait=True, + ) + dst_kv = bkvpe_vmem_ref.at[pl.ds(0, dma_sz_words)] + _async_copy( + src=dst_kv, + dst=dst_kv, + sem=sem, + wait=True, + ) + + def _fetch_bq(seq_idx, bq_idx, bq_sem_idx, *, wait=False): + sem = sems.at[1, bq_sem_idx] + bq_nope_vmem_ref = bq_nope_x2_ref.at[bq_sem_idx] + bq_rope_vmem_ref = bq_rope_x2_ref.at[bq_sem_idx] + + q_len_start = cu_q_lens_ref[seq_idx] + bq_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + debug_print( + "[RPA debug]" + f" -----------{'wait' if wait else 'start'}_fetch_bq-----------") + debug_print("[RPA debug] seq_idx={}", seq_idx) + debug_print("[RPA debug] bq_idx={}", bq_idx) + debug_print("[RPA debug] bq_sem_idx={}", bq_sem_idx) + debug_print("[RPA debug] q_len_start={}", q_len_start) + debug_print("[RPA debug] q_end={}", q_end) + debug_print("[RPA debug] sz={}", sz) + + _async_copy( + ql_nope_hbm_ref.at[pl.ds(q_len_start, sz)], + bq_nope_vmem_ref.at[pl.ds(0, sz)], + sem, + wait, + ) + + _async_copy( + q_pe_hbm_ref.at[pl.ds(q_len_start, sz)], + bq_rope_vmem_ref.at[pl.ds(0, sz)], + sem, + wait, + ) + + def _send_bo(seq_idx, bo_idx, bo_sem_idx, *, wait=False): + sem = sems.at[2, bo_sem_idx] + vmem_ref = bo_x2_ref.at[bo_sem_idx] + q_len_start = cu_q_lens_ref[seq_idx] + bo_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + debug_print( + "[RPA debug]" + f" -----------{'wait' if wait else 'start'}_send_bo-----------") + debug_print("[RPA debug] seq_idx={}", seq_idx) + debug_print("[RPA debug] bo_idx={}", bo_idx) + debug_print("[RPA debug] bo_sem_idx={}", bo_sem_idx) + debug_print("[RPA debug] q_len_start={}", q_len_start) + debug_print("[RPA debug] q_end={}", q_end) + debug_print("[RPA debug] sz={}", sz) + + _async_copy( + vmem_ref.at[pl.ds(0, sz)], + o_hbm_ref.at[pl.ds(q_len_start, sz)], + sem, + wait, + ) + + def start_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx): + return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx) + + def wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx): + return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, wait=True) + + def start_fetch_bq(seq_idx, bq_idx, bq_sem_idx): + return _fetch_bq(seq_idx, bq_idx, bq_sem_idx) + + def wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx): + return _fetch_bq(seq_idx, bq_idx, bq_sem_idx, wait=True) + + def start_send_bo(seq_idx, bo_idx, bo_sem_idx): + bo_ids_ref[bo_sem_idx] = seq_idx + bo_ids_ref[bo_sem_idx + 2] = bo_idx + _send_bo(seq_idx, bo_idx, bo_sem_idx) + + def wait_send_bo(bo_sem_idx): + old_seq_idx = bo_ids_ref[bo_sem_idx] + old_bo_idx = bo_ids_ref[bo_sem_idx + 2] + + @pl.when(jnp.logical_and(0 <= old_seq_idx, old_seq_idx <= seq_idx)) + def _(): + _send_bo(old_seq_idx, old_bo_idx, bo_sem_idx, wait=True) + + def start_update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz): + bkv_update_ids_ref[bkv_sem_idx] = seq_idx + bkv_update_ids_ref[bkv_sem_idx + 2] = offset + bkv_update_ids_ref[bkv_sem_idx + 4] = update_sz + _update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz) + + def wait_update_kv_cache(bkv_sem_idx): + update_sz = bkv_update_ids_ref[bkv_sem_idx + 4] + + @pl.when(update_sz > 0) + def _(): + seq_idx = bkv_update_ids_ref[bkv_sem_idx] + offset = bkv_update_ids_ref[bkv_sem_idx + 2] + bkv_update_ids_ref[bkv_sem_idx + 4] = 0 + _update_kv_cache(seq_idx, + bkv_sem_idx, + offset, + update_sz, + wait=True) + + def load_bq(bq_sem_idx, *, actual_bq_sz=bq_sz): + q_nope_ref = (bq_nope_x2_ref.bitcast( + jnp.uint32).at[bq_sem_idx].reshape( + bq_sz * num_q_heads_per_q_packing, lkv_dim)) + q_nope_vec = pltpu.bitcast( + q_nope_ref[:actual_bq_sz * num_q_heads_per_q_packing], + q_dtype, + ).reshape(actual_bq_sz * num_q_heads, lkv_dim) + q_rope_ref = (bq_rope_x2_ref.bitcast( + jnp.uint32).at[bq_sem_idx].reshape( + bq_sz * num_q_heads_per_q_packing, r_dim)) + q_rope_vec = pltpu.bitcast( + q_rope_ref[:actual_bq_sz * num_q_heads_per_q_packing], + q_dtype, + ).reshape(actual_bq_sz * num_q_heads, r_dim) + return q_nope_vec, q_rope_vec + + def load_bkv(bkv_sem_idx, *, bkvc_mask, bkpe_mask): + bkvc_ref = (bkvc_x2_ref.bitcast( + jnp.uint32).at[bkv_sem_idx, :bkv_sz_per_kv_packing].reshape( + bkv_sz_per_kv_packing, lkv_dim)) + bkvc_vec = pltpu.bitcast(bkvc_ref[...], + kv_dtype).reshape(bkv_sz, lkv_dim) + bkvc_vec = lax.select(bkvc_mask, bkvc_vec, jnp.zeros_like(bkvc_vec)) + + bkpe_ref = (bkpe_x2_ref.bitcast( + jnp.uint32).at[bkv_sem_idx, :bkv_sz_per_kv_packing].reshape( + bkv_sz_per_kv_packing, r_dim)) + bkpe_vec = pltpu.bitcast(bkpe_ref[...], + kv_dtype).reshape(bkv_sz, r_dim) + bkpe_vec = lax.select(bkpe_mask, bkpe_vec, jnp.zeros_like(bkpe_vec)) + + return bkvc_vec, bkpe_vec + + def broadcast_minor(src, shape): + if src.shape == shape: + return src + assert src.shape[:-1] == shape[:-1] + assert src.shape[-1] % 128 == 0 + target_minor = align_to(shape[-1], src.shape[-1]) + # no-op concatenation. + return jnp.concatenate( + [src for _ in range(target_minor // src.shape[-1])], + axis=-1)[..., :shape[-1]] + + def process(static_q_len=None): + num_bkv = cdiv(kv_len, bkv_sz) + if static_q_len is None: + actual_bq_sz = bq_sz + num_bq = cdiv(q_len, actual_bq_sz) + else: + actual_bq_sz = min(bq_sz, static_q_len) + num_bq = cdiv(static_q_len, actual_bq_sz) + + debug_print("[RPA debug] process") + debug_print("[RPA debug] num_bkv={}", num_bkv) # num_bkv=3, bkv_sz=512 + debug_print("[RPA debug] bkv_sz={}", bkv_sz) + debug_print("[RPA debug] num_bq={}", num_bq) + debug_print("[RPA debug] kv_len={}", kv_len) + + def get_next_bq_ids(seq_idx, bq_idx, bq_sem_idx): + next_bq_idx = bq_idx + 1 + is_last_bq = next_bq_idx == num_bq + next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx) + next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx) + next_bq_sem_idx = lax.select(bq_sem_idx == 0, 1, 0) + return next_seq_idx, next_bq_idx, next_bq_sem_idx + + def get_next_bkv_ids(seq_idx, bq_idx, bkv_idx, bkv_sem_idx): + next_bkv_idx = bkv_idx + 1 + is_last_bkv = next_bkv_idx == num_bkv + next_bkv_idx = lax.select(is_last_bkv, 0, next_bkv_idx) + next_bq_idx = lax.select(is_last_bkv, bq_idx + 1, bq_idx) + is_last_bq = next_bq_idx == num_bq + next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx) + next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx) + next_bkv_sem_idx = lax.select(bkv_sem_idx == 0, 1, 0) + return next_seq_idx, next_bq_idx, next_bkv_idx, next_bkv_sem_idx + + def compute_with_bq(bq_idx, _): + bq_sem_idx = sem_ids_ref[0] + next_seq_idx, next_bq_idx, next_bq_sem_idx = get_next_bq_ids( + seq_idx, bq_idx, bq_sem_idx) + + # Prefetch next bq + @pl.when(next_seq_idx < num_seqs) + def prefetch_next_bq(): + sem_ids_ref[0] = next_bq_sem_idx + start_fetch_bq(next_seq_idx, next_bq_idx, next_bq_sem_idx) + + def compute_with_bkv(bkv_idx, _): + # Create bitmask for KV. + assert bkv_sz % kv_packing == 0 + actual_bkv_sz = jnp.minimum(bkv_sz, kv_len - bkv_idx * bkv_sz) + bkvc_shape = (bkv_sz, lkv_dim) + bkvc_mask = (lax.broadcasted_iota(jnp.int32, bkvc_shape, 0) + < actual_bkv_sz) + bkpe_shape = (bkv_sz, r_dim) + bkpe_mask = (lax.broadcasted_iota(jnp.int32, bkpe_shape, 0) + < actual_bkv_sz) + + # Get next bkv ids. + bkv_sem_idx = sem_ids_ref[1] + next_seq_idx, _, next_bkv_idx, next_bkv_sem_idx = get_next_bkv_ids( + seq_idx, bq_idx, bkv_idx, bkv_sem_idx) + + # Prefetch next bkv + @pl.when(next_seq_idx < num_seqs) + def prefetch_next_bkv(): + sem_ids_ref[1] = next_bkv_sem_idx + start_fetch_bkv(next_seq_idx, next_bkv_idx, + next_bkv_sem_idx) + + # Wait for cur bq if not ready yet + @pl.when(bkv_idx == 0) + def wait_cur_bq(): + wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx) + + # Wait for cur bkv + offset, update_sz = wait_fetch_bkv(seq_idx, bkv_idx, + bkv_sem_idx) + + # Pack and align new KVs in VMEM if the block has new KVs. + # We may have to do this for each block of KV in VMEM. + @pl.when(update_sz > 0) + def pack_new_kv(): + _pack_new_kv(bkv_sem_idx, offset, update_sz) + + # Start updating bkv to kv cache if applicable. + # Only needed in first bq loop. + @pl.when(jnp.logical_and(update_sz > 0, bq_idx == 0)) + def update_cur_bkv_to_cache(): + start_update_kv_cache(seq_idx, bkv_sem_idx, offset, + update_sz) + + # Load bkv into vreg + bkvc, bkpe = load_bkv(bkv_sem_idx, + bkvc_mask=bkvc_mask, + bkpe_mask=bkpe_mask) + + bq_nope_vec, bq_pe_vec = load_bq(bq_sem_idx, + actual_bq_sz=actual_bq_sz) + + debug_print("[RPA debug] flash attention") + debug_print( + "[RPA debug] bq_nope_vec.shape={}, {}", + bq_nope_vec.shape[0], + bq_nope_vec.shape[1], + ) # num_bkv=3, bkv_sz=512 + debug_print( + "[RPA debug] bq_pe_vec.shape={}, {}", + bq_pe_vec.shape[0], + bq_pe_vec.shape[1], + ) + debug_print("[RPA debug] bkvc.shape={}, {}", bkvc.shape[0], + bkvc.shape[1]) + debug_print("[RPA debug] bkpe.shape={}, {}", bkpe.shape[0], + bkpe.shape[1]) + + if debug_mode: + return + + flash_attention( + bq_nope_vec, + bq_pe_vec, + bkvc, + bkpe, + bq_idx=bq_idx, + bkv_idx=bkv_idx, + ) + + lax.fori_loop(0, num_bkv, compute_with_bkv, None, unroll=False) + + # Load acc and calculate final output. + acc = acc_ref[...] + l = broadcast_minor(l_ref[...], acc.shape) # noqa + out = (lax.div(acc, l) if q_dtype == jnp.float32 else + (acc * pl.reciprocal(l, approx=True)).astype(q_dtype)) + + # Wait for previous bo to be fully sent before storing new bo. + bo_sem_idx = sem_ids_ref[2] + sem_ids_ref[2] = lax.select(bo_sem_idx == 0, 1, 0) + wait_send_bo(bo_sem_idx) + + # Store output from acc to bo. + bo_x2_ref.at[bo_sem_idx].bitcast(jnp.int32).reshape( + bq_sz * num_q_heads_per_q_packing, + lkv_dim, + )[...] = pltpu.bitcast(out, jnp.int32) + + # Send cur bo + start_send_bo(seq_idx, bq_idx, bo_sem_idx) + + lax.fori_loop(0, num_bq, compute_with_bq, None, unroll=False) + + ### ------- Kernel start ------- ### + + @pl.when(seq_idx == 0) + def prologue(): + start_fetch_bq(0, 0, 0) + start_fetch_bkv(0, 0, 0) + + @pl.when(seq_idx < decode_end) + def process_decode(): + process(static_q_len=1) + + @pl.when(jnp.logical_and(decode_end <= seq_idx, seq_idx < prefill_end)) + def process_prefill(): + process(static_q_len=chunk_prefill_size) + + @pl.when(jnp.logical_and(prefill_end <= seq_idx, seq_idx < mixed_end)) + def process_mixed(): + process() + + @pl.when(seq_idx == num_seqs - 1) + def epilogue(): + for i in range(2): + wait_send_bo(i) + wait_update_kv_cache(i) + + ### ------- Kernel end ------- ### + + +def prepare_q_inputs( + q: jax.Array, # [max_num_tokens, actual_num_q_heads, actual_head_dim], +): + max_num_tokens, actual_num_q_heads, actual_head_dim = q.shape + q_packing = get_dtype_packing(q.dtype) + num_q_heads = align_to(actual_num_q_heads, q_packing) + head_dim = align_to(actual_head_dim, 128) + q = jnp.pad( + q.reshape( + max_num_tokens, + actual_num_q_heads, + actual_head_dim, + ), + ( + (0, 0), + (0, num_q_heads - actual_num_q_heads), + (0, head_dim - actual_head_dim), + ), + constant_values=0, + ).reshape( + max_num_tokens, + num_q_heads // q_packing, + q_packing, + head_dim, + ) + return q + + +def prepare_kv_inputs(kv: jax.Array): + max_num_tokens, actual_head_dim = kv.shape + kv_packing = get_dtype_packing(kv.dtype) + # Pad to packing + if max_num_tokens % kv_packing != 0: + pad = kv_packing - (max_num_tokens % kv_packing) + kv = jnp.pad(kv, ((0, pad), (0, 0)), constant_values=0) + + head_dim = align_to(actual_head_dim, 128) + kv = kv.reshape(-1, kv_packing, actual_head_dim) + kv = jnp.pad(kv, ((0, 0), (0, 0), (0, head_dim - actual_head_dim)), + constant_values=0) + return kv + + +def prepare_outputs( + out, # [max_num_tokens, num_q_heads // q_packing, q_packing, head_dim] + actual_num_q_heads: int, + actual_head_dim: int, +): + ( + max_num_tokens, + num_q_heads_per_q_packing, + q_packing, + head_dim, + ) = out.shape + return out.reshape( + max_num_tokens, + num_q_heads_per_q_packing * q_packing, + head_dim, + )[:, :actual_num_q_heads, :actual_head_dim] + +CONFIG = { + 'name': 'MLA', + 'batch_size': 128, + 'q_len': 1, + 'kv_len_val': 9216, + 'page_size': 256, + 'symbol': 'd', +} + +def create_inputs(): + key = jax.random.PRNGKey(0) + + num_heads = 128 + lkv_dim = 512 + r_dim = 64 + q_dtype = jnp.bfloat16 + kv_dtype = jnp.bfloat16 + + padded_r_dim = align_to(r_dim, 128) + padded_lkv_dim = align_to(lkv_dim, 128) + padded_kv_dim = padded_lkv_dim + padded_r_dim + packing = get_dtype_packing(kv_dtype) + + def gen_random(k, shape, dtype): + return jax.random.uniform(k, shape, dtype=jnp.float32).astype(dtype) + + total_kv_tokens = CONFIG['batch_size'] * CONFIG['kv_len_val'] + num_pages = cdiv(total_kv_tokens, CONFIG['page_size']) + CONFIG['batch_size'] + + total_q_len = CONFIG['batch_size'] * CONFIG['q_len'] + cu_q_lens_list = [i * CONFIG['q_len'] for i in range(CONFIG['batch_size'] + 1)] + + pages_per_seq = cdiv(CONFIG['kv_len_val'], CONFIG['page_size']) + page_indices_list = [] + page_count = 0 + for _ in range(CONFIG['batch_size']): + num_seq_pages = cdiv(CONFIG['kv_len_val'], CONFIG['page_size']) + indices = list(range(page_count, page_count + num_seq_pages)) + page_indices_list.extend(indices + [-1] * (pages_per_seq - num_seq_pages)) + page_count += num_seq_pages + + total_num_pages = max(num_pages, page_count) + + key, k1, k2, k3, k4, k5 = jax.random.split(key, 6) + ql_nope = gen_random(k1, (total_q_len, num_heads, lkv_dim), q_dtype) + q_pe = gen_random(k2, (total_q_len, num_heads, r_dim), q_dtype) + new_kv_c = gen_random(k3, (total_q_len, lkv_dim), kv_dtype) + new_k_pe = gen_random(k4, (total_q_len, r_dim), kv_dtype) + + cache_kv = gen_random( + k5, + (total_num_pages, CONFIG['page_size'] // packing, packing, padded_kv_dim), + kv_dtype, + ) + + kv_lens = jnp.array([CONFIG['kv_len_val']] * CONFIG['batch_size'], dtype=jnp.int32) + page_indices = jnp.array(page_indices_list, dtype=jnp.int32) + cu_q_lens = jnp.array(cu_q_lens_list, dtype=jnp.int32) + + num_decode_seqs = CONFIG['batch_size'] if CONFIG['q_len'] == 1 else 0 + distribution = jnp.array([num_decode_seqs, num_decode_seqs, CONFIG['batch_size']], dtype=jnp.int32) + + return ( + ql_nope, q_pe, new_kv_c, new_k_pe, cache_kv, kv_lens, + page_indices, cu_q_lens, distribution + ) + +@functools.partial( + jax.jit, + static_argnames=( + "sm_scale", + "sliding_window", + "soft_cap", + "mask_value", + "q_scale", + "k_scale", + "v_scale", + "chunk_prefill_size", + "num_kv_pages_per_block", + "num_queries_per_block", + "vmem_limit_bytes", + "debug_mode", + ), + donate_argnames=("cache_kv", ), +) +def mla_ragged_paged_attention( + ql_nope: jax.Array, # [max_num_tokens, actual_num_q_heads, actual_lkv_dim] + q_pe: jax.Array, # [max_num_tokens, actual_num_q_heads, actual_r_dim] + new_kv_c: jax.Array, # [max_num_tokens, actual_lkv_dim] + new_k_pe: jax.Array, # [max_num_tokens, actual_r_dim] + cache_kv: jax. + Array, # [total_num_pages, page_size_per_kv_packing, kv_packing, align_to(lkv_dim, 128)] + kv_lens: jax.Array, # i32[max_num_seqs] + page_indices: jax.Array, # i32[max_num_seqs * pages_per_seq] + cu_q_lens: jax.Array, # i32[max_num_seqs + 1] + distribution: jax.Array, # i32[3] + *, + sm_scale: float = 1.0, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = DEFAULT_MASK_VALUE, + q_scale: float | None = None, + k_scale: float | None = None, + v_scale: float | None = None, + # Kernel optimization params. + chunk_prefill_size: int | None = None, + # Kernel tuning params. + num_kv_pages_per_block: int | None = None, + num_queries_per_block: int | None = None, + vmem_limit_bytes: int | None = None, + # Debug params. + debug_mode: bool = False, +) -> tuple[ + jax.Array, # [max_num_tokens, actual_num_q_heads, actual_lkv_dim] + jax. + Array, # [total_num_pages, page_size_per_kv_packing, kv_packing, align_to(lkv_dim, 128) + align_to(r_dim, 128)] +]: + """MLA Ragged paged attention that supports mixed prefill and decode. + + Args: + ql_nope: concatenated all sequences' queries. + q_pe: concatenated all sequences' rope. + new_kv_c: concatenated all sequences' kv_c values + new_k_pe: concatenated all sequences' k_pe values + cache_kv: the current kv cache. + kv_lens: the length of each sequence in the kv cache. + page_indices: flattened page indices look-up table by (seq_id, page_id). + cu_q_lens: the cumulative sum of the effective query lengths. Similar to + kv_lens, only the first num_seqs+1 values are valid. + distribution: (i, j, k) represents that sequences[0:i] are decode-only, + sequences[i:j] are chunked-prefill-only, and sequences[j:k] are mixed. The + k is also the total number of sequences. + sm_scale: the softmax scale which will be applied to the Q@K^T. + sliding_window: the sliding window size for the attention. + soft_cap: the logit soft cap for the attention. + mask_value: mask value for causal mask. + q_scale: the scale for the query. + k_scale: the scale for the key cache. + v_scale: the scale for the value cache. + num_kv_pages_per_block: number of kv pages to be processed in one flash + attention block in the pallas kernel. + num_queries_per_block: number of kv pages to be processed in one flash + attention block in the pallas kernel. + vmem_limit_bytes: the vmem limit for the pallas kernel. + debug_mode: if true, RPA does not issue any DMAs or run flash attention but + print debug info. Need to compile with `--xla_tpu_enable_log_recorder`. + + Returns: + The output of attention and the updated kv cache. + """ + if num_kv_pages_per_block is None or num_queries_per_block is None: + raise ValueError( + "num_kv_pages_per_block and num_queries_per_block must be specified." + ) + static_validate_inputs( + ql_nope, + q_pe, + new_kv_c, + new_k_pe, + cache_kv, + kv_lens, + page_indices, + cu_q_lens, + distribution, + sm_scale=sm_scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + mask_value=mask_value, + q_scale=q_scale, + k_scale=k_scale, + v_scale=v_scale, + chunk_prefill_size=chunk_prefill_size, + num_kv_pages_per_block=num_kv_pages_per_block, + num_queries_per_block=num_queries_per_block, + vmem_limit_bytes=vmem_limit_bytes, + debug_mode=debug_mode, + ) + + _, actual_num_q_heads, actual_lkv_dim = ql_nope.shape + + ql_nope = prepare_q_inputs( + ql_nope + ) # [max_num_tokens, num_q_heads_per_q_packing, q_packing, lkv_dim] + q_pe = prepare_q_inputs( + q_pe) # [max_num_tokens, num_q_heads_per_q_packing, q_packing, r_dim] + new_kv_c = prepare_kv_inputs( + new_kv_c) # [max_num_tokens_per_kv_packing, kv_packing, lkv_dim] + new_k_pe = prepare_kv_inputs( + new_k_pe) # [max_num_tokens_per_kv_packing, kv_packing, r_dim] + lkv_dim = new_kv_c.shape[-1] + r_dim = new_k_pe.shape[-1] + + _, page_size_per_kv_packing, kv_packing, _ = cache_kv.shape + page_size = page_size_per_kv_packing * kv_packing + _, num_q_heads_per_q_packing, q_packing, _ = ql_nope.shape + max_num_seqs = kv_lens.shape[0] + num_page_indices = page_indices.shape[0] + assert num_page_indices % max_num_seqs == 0 + num_q_heads = num_q_heads_per_q_packing * q_packing + + bkv_p = num_kv_pages_per_block + bq_sz = num_queries_per_block + bkv_sz_per_kv_packing = bkv_p * page_size_per_kv_packing + # Add 2 additional words of buffering to accommodate misaligned new KV. + # We need two additional words because the beginning and end of the new KV may + # both not be aligned to kv_packing boundaries. + # Example: + # + # T0 T4 K2 K6 + # T1 K3 + # T2 K0 K4 + # T3 K1 K5 + # + # - Ti is existing KV tokens and Ki is the new KV. + # - Each column is a 32-bit word. + # - KV packing is 4 + # + # We have 12 total tokens, so normally we would only allocate 12/4=3 words + # But due to misalignment, we need to allocate 5 words. + bkv_buf_sz_per_kv_packing = bkv_sz_per_kv_packing + 2 + grid = (distribution[2], ) + + in_specs = [ + pl.BlockSpec(memory_space=pltpu.HBM), # ql_nope + pl.BlockSpec(memory_space=pltpu.HBM), # q_pe + pl.BlockSpec(memory_space=pltpu.HBM), # new_kv_c + pl.BlockSpec(memory_space=pltpu.HBM), # new_k_pe + pl.BlockSpec(memory_space=pltpu.HBM), # cache_kv + ] + + out_specs = [ + pl.BlockSpec(memory_space=pltpu.HBM), # o + pl.BlockSpec(memory_space=pltpu.HBM), # updated_cache_kv + ] + + bkvc_double_buf = pltpu.VMEM( + (2, bkv_buf_sz_per_kv_packing, kv_packing, lkv_dim), + cache_kv.dtype, + ) + + bkpe_double_buf = pltpu.VMEM( + (2, bkv_buf_sz_per_kv_packing, kv_packing, r_dim), + cache_kv.dtype, + ) + bq_nope_double_buf = pltpu.VMEM( + (2, bq_sz, num_q_heads_per_q_packing, q_packing, lkv_dim), + ql_nope.dtype, + ) + + bq_rope_double_buf = pltpu.VMEM( + (2, bq_sz, num_q_heads_per_q_packing, q_packing, r_dim), + q_pe.dtype, + ) + + bo_double_buf = bq_nope_double_buf + + l_scratch = pltpu.VMEM( + (bq_sz * num_q_heads, 128), + jnp.float32, + ) + m_scratch = l_scratch + + acc_scratch = pltpu.VMEM( + (bq_sz * num_q_heads, lkv_dim), + jnp.float32, + ) + + scratch_shapes = [ + bkvc_double_buf, + bkpe_double_buf, + bq_nope_double_buf, + bq_rope_double_buf, + bo_double_buf, # Double buffering for output block. + # Semaphores for double buffering of bkv, bq, bo and bkv_update. + pltpu.SemaphoreType.DMA((4, 2)), + # Intermediate buffers per kv head for flash attention. + l_scratch, + m_scratch, + acc_scratch, + ] + + scalar_prefetches = ( + kv_lens, + page_indices, + cu_q_lens, + distribution, + # (bq_sem_idx, bkv_sem_idx, bo_sem_idx) + jnp.zeros((3, ), jnp.int32), + # (bo_sem_0_seq_idx, bo_sem_1_seq_idx, bo_sem_0_bo_idx, bo_sem_1_bo_idx) + jnp.full((4, ), -1, jnp.int32), + # (bkv_sem_0_seq_idx, bkv_sem_1_seq_idx, bkv_sem_0_offset, bkv_sem_1_offset, bkv_sem_0_sz, bkv_sem_1_sz) + jnp.full((6, ), -1, jnp.int32), + ) + + scope_name = f"MLA-RPA-bq_{bq_sz}-bkvp_{bkv_p}-p_{page_size}" + kernel = jax.named_scope(scope_name)( + pl.pallas_call( + functools.partial( + _mla_ragged_paged_attention_kernel, + sm_scale=sm_scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + mask_value=mask_value, + q_scale=q_scale, + k_scale=k_scale, + v_scale=v_scale, + chunk_prefill_size=chunk_prefill_size, + bq_sz=bq_sz, + bkv_p=bkv_p, + debug_mode=debug_mode, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=len(scalar_prefetches), + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + scratch_shapes=scratch_shapes, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary", ), + vmem_limit_bytes=vmem_limit_bytes, + ), + out_shape=[ + jax.ShapeDtypeStruct(shape=ql_nope.shape, dtype=ql_nope.dtype), + jax.ShapeDtypeStruct(shape=cache_kv.shape, + dtype=cache_kv.dtype), + ], + input_output_aliases={ + 7: 0, # Alias output activation with ql_nope + 11: 1, # Aliasing cache_kv with updated_cache_kv + }, + name=scope_name, + )) + + output, updated_kv = kernel( + *scalar_prefetches, + ql_nope, + q_pe, + new_kv_c, + new_k_pe, + cache_kv, + ) + output = prepare_outputs( + output, actual_num_q_heads, + actual_lkv_dim) # [max_num_tokens, actual_num_q_heads, actual_lkv_dim] + + return output, updated_kv + + +def workload( + ql_nope: jax.Array, + q_pe: jax.Array, + new_kv_c: jax.Array, + new_k_pe: jax.Array, + cache_kv: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + distribution: jax.Array, +): + return mla_ragged_paged_attention( + ql_nope, + q_pe, + new_kv_c, + new_k_pe, + cache_kv, + kv_lens, + page_indices, + cu_q_lens, + distribution, + num_kv_pages_per_block=16, + num_queries_per_block=4, + vmem_limit_bytes=DEFAULT_VMEM_LIMIT_BYTES + ) + + +def benchmark(num_warmup=5, num_iters=100): + """Benchmark and return results dict.""" + inputs = create_inputs() + fn = jax.jit(workload) + for _ in range(num_warmup): + out = fn(*inputs) + jax.block_until_ready(out) + times = [] + for _ in range(num_iters): + t0 = time.perf_counter() + out = fn(*inputs) + jax.block_until_ready(out) + times.append(time.perf_counter() - t0) + return { + 'times': [round(float(t * 1000), 4) for t in times], + 'time_ms': round(float(np.mean(times) * 1000), 4), + 'std_ms': round(float(np.std(times) * 1000), 4), + 'output_shape': [list(out[0].shape), list(out[1].shape)], + 'status': 'success', + } diff --git a/JAXBench/benchmark/level2/4p_Sparse_Attention/baseline.py b/JAXBench/benchmark/level2/4p_Sparse_Attention/baseline.py new file mode 100644 index 0000000..4f1995a --- /dev/null +++ b/JAXBench/benchmark/level2/4p_Sparse_Attention/baseline.py @@ -0,0 +1,2645 @@ +# 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.""" + +import time +from collections.abc import Callable, Mapping +import dataclasses +import enum +import functools +from typing import Any, Literal, NamedTuple, Optional, Union, 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 mechanism 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( + block_metadata: Mapping[str, Any], + 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}_" + "_".join( + f"{k}={v}" for k, v in sorted(block_metadata.items()) + ) + + +# 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_argnames=( + "mask_value", "save_residuals", "custom_type", "attn_logits_soft_cap") +) +_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: # pytype: disable=attribute-error + kv_heads, q_heads_per_kv_head, q_seq_len, head_dim = activations.shape # pytype: disable=attribute-error + return activations.reshape( # pytype: disable=attribute-error + 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(frozen=True, slots=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: + object.__setattr__(self, "block_kv_compute", self.block_kv) + if self.block_kv_dkv_compute is None: + object.__setattr__(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 | jax.Array, + 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 = mask_ref[:, k_slice] + else: + mask = mask_ref[k_slice, :] + + masks.append( + jnp.bitwise_or(mask, jnp.broadcast_to(should_not_mask, 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 = jnp.tile( + q_sequence_ref[...], (1, repeats) + ) # [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 = q_sequence_ref[:1, :] # [1, bq] + q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq)) + + assert q_sequence.shape == k_sequence.shape + computed_mask = mask_function(q_sequence, k_sequence) # pytype: disable=wrong-arg-count + if computed_mask.dtype != jnp.dtype(jnp.bool_): + raise ValueError( + "Mask function must return a boolean-valued array, but got:" + f" {computed_mask.dtype}" + ) + masks.append(computed_mask) + + if q_segment_ids_ref is not None: + if k_in_lanes: + kv_ids = kv_segment_ids_ref[: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 = jnp.tile(q_segment_ids_ref[:], (1, repeats)) # [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 = jnp.tile( + kv_segment_ids_ref[k_slice, :], (1, repeats) + ) # [k_slice, bq] + q_ids = q_segment_ids_ref[:1, :] # [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_v: 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_v_repeats, rem = divmod(head_dim_v, NUM_LANES) + if rem != 0: + raise NotImplementedError( + f"{head_dim_v=} 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) + + global_kv_index, _, 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 = k_ref[slice_k, :] + else: + k = k_ref[:, 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, + # When the iteration space is shrunk (for local attention for example), + # the kv_index program_id does not correspond to the actual coordinates + # of the KV data. Make sure to use the 'unshrunk' index (coming from the + # data_next array) when computing the mask. + k_offset=global_kv_index * 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] # pytype: disable=attribute-error + 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 - jnp.tile(m_next, (1, bkv_repeats))) + 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 = v_ref[slice_k, :] + else: + v = v_ref[:, slice_k] + v = v.astype(float32) + o_curr = lax.dot_general(s_curr, v, sv_dims) + + alpha_o = jnp.tile(alpha, (1, head_dim_v_repeats)) + 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 = jnp.tile(1.0 / l, (1, head_dim_v_repeats)) + 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, + interpret: bool = False +) -> SplashCustomReturnType: + num_q_heads, q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + 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] + + partial_mask_blocks = fwd_mask_info.partial_mask_blocks + if ( + partial_mask_blocks is not None + and jnp.dtype(partial_mask_blocks.dtype) != np.bool_ + ): + raise ValueError( + "partial_mask_blocks must be of type np.bool_ but got" + f" {partial_mask_blocks.dtype}" + ) + + 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_qk: + raise ValueError( + f"Expected 'key' head dimension to be: {head_dim_qk}. 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[:-1] != v.shape[:-1]: + raise ValueError( + f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " + "leading dimensions." + ) + + if bkv % bkv_compute: + raise ValueError(f"{bkv=} must be a multiple of {bkv_compute=}.") + if bkv_compute % NUM_LANES: + raise ValueError(f"{bkv_compute=} 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( + from_head_minor((None, bq, head_dim_qk), q_layout), q_index_map + ), + pl.BlockSpec( + from_head_minor( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), k_layout + ), + k_index_map, + ), + pl.BlockSpec( + from_head_minor( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), v_layout + ), + v_index_map, + ), + ] + if segment_ids is not None: + in_specs += [ + pl.BlockSpec((bq, NUM_LANES), q_segment_ids_index_map), + pl.BlockSpec((NUM_SUBLANES, bkv), kv_segment_ids_index_map), + ] + 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((None, bq, bkv), mask_index_map)) + 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((bq, NUM_LANES), q_segment_ids_index_map)) + 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_v), jnp.float32), # o_scratch + jax.ShapeDtypeStruct((num_q_heads, q_seq_len, head_dim_v), q.dtype), + ] + out_specs = [ + # TODO(sharadmv): convert m/l to be scratch + pl.BlockSpec((bq, NUM_LANES), lambda h, i, j, *_: (0, 0)), + pl.BlockSpec((bq, NUM_LANES), lambda h, i, j, *_: (0, 0)), + pl.BlockSpec((bq, head_dim_v), lambda h, i, j, *_: (0, 0)), + pl.BlockSpec((None, bq, head_dim_v), out_index_map), + ] + 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((None, bq, NUM_LANES), logsumexp_index_map), + ] + else: + out_shapes += [None] + out_specs += [None] + + kernel_name = get_kernel_name( + dataclasses.asdict(block_sizes), + is_mqa=is_mqa, + save_residuals=save_residuals, + is_segmented=segment_ids is not None, + phase="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_v=head_dim_v, + 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, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary"), + ), + out_shape=out_shapes, + name=kernel_name, + interpret=interpret, + )( + 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_argnames=( + "save_residuals", "mask_value", "is_mqa", "block_sizes", + "residual_checkpoint_name", "mask_function", "attn_logits_soft_cap", + "interpret") +) +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, + interpret: bool = False, +) -> 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, + residual_checkpoint_name=residual_checkpoint_name, + save_residuals=save_residuals, + mask_function=mask_function, + attn_logits_soft_cap=attn_logits_soft_cap, + interpret=interpret, + ) + + +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, + interpret: bool = False, +) -> tuple[ + tuple[jax.Array], + SplashResidualsType, +]: + if save_residuals: + raise NotImplementedError("Higher-order AD not supported") + + out, (logsumexp,) = _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, + residual_checkpoint_name=residual_checkpoint_name, + save_residuals=True, + mask_function=mask_function, + attn_logits_soft_cap=attn_logits_soft_cap, + interpret=interpret, + ) + 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) + + global_kv_index, _, 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), + # When the iteration space is shrunk (for local attention for example), + # the kv_index program_id does not correspond to the actual coordinates + # of the KV data. Make sure to use the 'unshrunk' index (coming from the + # data_next array) when computing the mask. + k_offset=global_kv_index * 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, + interpret: bool, +): + num_q_heads, q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + 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[:-1] != v.shape[:-1]: + raise ValueError( + f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " + "leading dimensions." + ) + + 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((None, bq, head_dim_v), o_index_map) + + def q_index_map(h, i, *_): + return from_head_minor((h, i, 0), q_layout) + + q_spec = pl.BlockSpec( + from_head_minor((None, bq, head_dim_qk), q_layout), q_index_map + ) + + 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( + from_head_minor( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), k_layout + ), + k_index_map, + ) + + 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( + from_head_minor( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), v_layout + ), + v_index_map, + ) + + 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((None, bq, bkv), mask_index_map) + + 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((bq, NUM_LANES), q_segment_ids_index_map) + kv_segment_spec = pl.BlockSpec( + (NUM_SUBLANES, bkv), kv_segment_ids_index_map + ) + 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 = o_spec + + def logsumexp_index_map(h, i, *_): + return h, 0, i + + logsumexp = jnp.expand_dims(logsumexp, axis=-2) + logsumexp_spec = pl.BlockSpec((None, 1, bq), logsumexp_index_map) + assert logsumexp.ndim == len(logsumexp_spec.block_shape) + + di = jnp.expand_dims(di, axis=-2) + di_spec = pl.BlockSpec((None, 1, bq), logsumexp_index_map) + 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((bq, NUM_LANES), q_segment_ids_index_map)) + else: + q_sequence = None + in_specs.append(None) + + out_shapes = [ + jax.ShapeDtypeStruct((bq, head_dim_qk), jnp.float32), + jax.ShapeDtypeStruct(q.shape, q.dtype), + ] + out_specs = [ + pl.BlockSpec((bq, head_dim_qk), lambda *_: (0, 0)), + pl.BlockSpec((None, bq, head_dim_qk), lambda h, i, *_: (h, i, 0)), + ] + + 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 + + kernel_name = get_kernel_name( + 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, + save_residuals=False, + is_segmented=segment_ids is not None, + phase="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, + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary", "arbitrary", "arbitrary"), + ), + name=kernel_name, + interpret=interpret, + )( + 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 ref[slice_k, :] + return ref[:, slice_k].T + k = _load_kv(k_ref, k_layout) + v = _load_kv(v_ref, v_layout) + logsumexp = logsumexp_ref[:1, :] + do = do_ref[...] + di = di_ref[:1, :] + + 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) + dv_scratch_ref[slice_k, :] + dv_scratch_ref[slice_k, :] = 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) + dk_scratch_ref[slice_k, :] + dk_scratch_ref[slice_k, :] = 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, + interpret: bool, +): + num_q_heads, q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + 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[:-1] != v.shape[:-1]: + raise ValueError( + f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " + "leading dimensions." + ) + + 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((None, bq, head_dim_v), o_index_map) + + 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( + from_head_minor((None, bq, head_dim_qk), q_layout), q_index_map + ) + + 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( + from_head_minor( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), + k_layout, + ), + k_index_map, + ) + + 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( + from_head_minor( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), + v_layout, + ), + v_index_map, + ) + + 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((None, None, bq, head_dim_qk), dq_index_map) + 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((bq, head_dim_qk), lambda *_: (0, 0)) + dq_scratch_shape = jax.ShapeDtypeStruct((bq, head_dim_qk), 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 = pl.BlockSpec( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), + dkv_index_map, + ) + + dv_spec = pl.BlockSpec( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), + dkv_index_map, + ) + + 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((None, bkv, bq), mask_index_map) + + 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((NUM_SUBLANES, bq), q_segment_ids_index_map) + kv_segment_spec = pl.BlockSpec((bkv, NUM_LANES), kv_segment_ids_index_map) + 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((None, NUM_SUBLANES, bq), logsumexp_index_map) + 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((None, NUM_SUBLANES, bq), logsumexp_index_map) + 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((NUM_SUBLANES, bq), q_segment_ids_index_map)) + 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_qk), jnp.float32), + jax.ShapeDtypeStruct((bkv, head_dim_v), jnp.float32), + dq_shape, + jax.ShapeDtypeStruct(k.shape, k.dtype), + jax.ShapeDtypeStruct(v.shape, v.dtype), + ] + out_specs = [ + dq_scratch_spec, + pl.BlockSpec((bkv, head_dim_qk), lambda *_: (0, 0)), + pl.BlockSpec((bkv, head_dim_v), lambda *_: (0, 0)), + 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 + + kernel_name = get_kernel_name( + 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, + ), + is_mqa=is_mqa, + save_residuals=False, + is_segmented=segment_ids is not None, + phase="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, + # 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 + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary", "arbitrary", "arbitrary"), + ), + name=kernel_name, + interpret=interpret, + )( + 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, + interpret: bool, + 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, + interpret=interpret, + ) + 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, + interpret=interpret, + ) + # 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", + "interpret", + ], +) +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, + interpret: bool, +) -> SplashCustomReturnType: + """ + For dynamic masks, `partial_mask_blocks` has shape (head_count, q_blocks, kv_blocks, block_q, block_kv). + This shape allows sharding across both head count and query sequence dimensions. + + Note: The leading dimensions (head_count, q_blocks, kv_blocks) must be + collapsed into a single dimension before being passed to the kernel. + """ + def _collapse_partial_mask_blocks(mask_info: mask_info_lib.MaskInfo | None): + if mask_info is None or mask_info.partial_mask_blocks is None: + return mask_info + + return mask_info._replace( + partial_mask_blocks=mask_info.partial_mask_blocks.reshape( + -1, *mask_info.partial_mask_blocks.shape[-2:] + ) + ) + + fwd_mask_info = _collapse_partial_mask_blocks(fwd_mask_info) + dq_mask_info = _collapse_partial_mask_blocks(dq_mask_info) + dkv_mask_info = _collapse_partial_mask_blocks(dkv_mask_info) + 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, + interpret=interpret, + ) + + +@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() + partial_mask_blocks_spec = ( + spec if self.fwd_mask_info.is_dynamic_mask else replicated + ) + # 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=partial_mask_blocks_spec + 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 | jax.Array | 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, + interpret: bool = False, +): + 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() + + process_mask_fn = ( + mask_info_lib.process_dynamic_mask + if isinstance(mask, jax.Array) + else mask_info_lib.process_mask + ) + + process_mask_dvk_fn = ( + mask_info_lib.process_dynamic_mask_dkv + if isinstance(mask, jax.Array) + else mask_info_lib.process_mask_dkv + ) + + fwd_mask_info, mask_function_fwd = process_mask_fn( + 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 = process_mask_fn( + 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 = process_mask_dvk_fn( + 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, + interpret=interpret, + ) + + +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 +) + + +CONFIG = { + 'name': 'pallas_splash_attention_llama70b', + 'model': 'Llama-3.1-70B', + 'operator': 'pallas_splash_attention', + 'batch': 4, + 'seq_len': 4096, + 'num_query_heads': 64, + 'num_kv_heads': 8, + 'head_dim': 128, + 'atol': 3e-3, + 'rtol': 3e-3, +} + +# Tuned by autotune_block_sizes.py. Re-run to update. +TUNED_PARAMS = { + # Autotuned (forward pass). + 'block_q': 2048, + 'block_kv': 2048, + 'block_kv_compute': 1024, + 'q_layout': 1, # QKVLayout.HEAD_DIM_MINOR=1, SEQ_MINOR=2 + 'k_layout': 1, + 'v_layout': 1, + 'head_shards': 1, + 'q_seq_shards': 1, + # Not autotuned (backward-only). + 'block_q_dkv': None, + 'block_kv_dkv': None, + 'block_kv_dkv_compute': None, + 'block_q_dq': None, + 'block_kv_dq': None, +} + + +def get_flops(): + B = CONFIG['batch'] + H_q = CONFIG['num_query_heads'] + S = CONFIG['seq_len'] + D = CONFIG['head_dim'] + return 4 * B * H_q * S * S * D + + +def create_inputs(dtype=jnp.bfloat16): + key = jax.random.key(42) + k1, k2, k3 = jax.random.split(key, 3) + B = CONFIG['batch'] + S = CONFIG['seq_len'] + H_q = CONFIG['num_query_heads'] + H_kv = CONFIG['num_kv_heads'] + D = CONFIG['head_dim'] + q = jax.random.normal(k1, (B, H_q, S, D), dtype=dtype) * (D ** -0.5) + k = jax.random.normal(k2, (B, H_kv, S, D), dtype=dtype) * 0.02 + v = jax.random.normal(k3, (B, H_kv, S, D), dtype=dtype) * 0.02 + return q, k, v + + +def workload(q, k, v): + B, H_q, S, D = q.shape + H_kv = v.shape[1] + heads_per_group = H_q // H_kv + mask = mask_lib.CausalMask(shape=(S, S)) + multi_head_mask = mask_lib.MultiHeadMask([mask] * H_q) + block_sizes = BlockSizes( + block_q=TUNED_PARAMS['block_q'], + block_kv=TUNED_PARAMS['block_kv'], + block_kv_compute=TUNED_PARAMS['block_kv_compute'], + q_layout=QKVLayout(TUNED_PARAMS['q_layout']), + k_layout=QKVLayout(TUNED_PARAMS['k_layout']), + v_layout=QKVLayout(TUNED_PARAMS['v_layout']), + block_q_dkv=TUNED_PARAMS['block_q_dkv'], + block_kv_dkv=TUNED_PARAMS['block_kv_dkv'], + block_kv_dkv_compute=TUNED_PARAMS['block_kv_dkv_compute'], + block_q_dq=TUNED_PARAMS['block_q_dq'], + block_kv_dq=TUNED_PARAMS['block_kv_dq'], + ) + splash_kernel = _make_splash_attention( + multi_head_mask, block_sizes=block_sizes, + is_mqa=False, + head_shards=TUNED_PARAMS['head_shards'], + q_seq_shards=TUNED_PARAMS['q_seq_shards'], + ) + @jax.vmap + def _attend(q_batch, k_batch, v_batch): + k_repeated = jnp.repeat(k_batch, heads_per_group, axis=0) + v_repeated = jnp.repeat(v_batch, heads_per_group, axis=0) + return splash_kernel(q_batch, k_repeated, v_repeated) + return _attend(q, k, v) + + +def benchmark(num_warmup=5, num_iters=100): + """Benchmark and return results dict.""" + inputs = create_inputs() + fn = jax.jit(workload) + for _ in range(num_warmup): + out = fn(*inputs) + out.block_until_ready() + times = [] + for _ in range(num_iters): + t0 = time.perf_counter() + out = fn(*inputs) + out.block_until_ready() + times.append(time.perf_counter() - t0) + times = np.array(times) * 1000 + avg = float(np.mean(times)) + return { + 'name': CONFIG['name'], + 'model': CONFIG['model'], + 'operator': CONFIG['operator'], + 'config': {k: v for k, v in CONFIG.items() if k not in ('name', 'model', 'operator', 'atol', 'rtol')}, + 'time_ms': round(avg, 4), + 'std_ms': round(float(np.std(times)), 4), + 'output_shape': list(out.shape) if hasattr(out, 'shape') else [], + 'status': 'success', + } + + +if __name__ == '__main__': + import json + print(json.dumps(benchmark())) diff --git a/JAXBench/benchmark/level2/51p_DeepSeek_V4_CSA/baseline.py b/JAXBench/benchmark/level2/51p_DeepSeek_V4_CSA/baseline.py new file mode 100644 index 0000000..586b7e8 --- /dev/null +++ b/JAXBench/benchmark/level2/51p_DeepSeek_V4_CSA/baseline.py @@ -0,0 +1,1406 @@ +# Imports +import numpy as np +import time +import functools +from enum import Enum +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 +from jax.experimental.pallas import tpu_sc as plsc + +# Initialization +CONFIG = { + 'name': 'dsv4_csa', + 'model': 'deepseek', + 'operator': 'compressed_sparse_attention', + 'configs': { + "decode_large_batch": (256, 1, 9216, 1024, 1024), + "decode_medium_batch": (128, 1, 9216, 1024, 1024), + "decode_small_pages": (128, 1, 9216, 256, 1024), + "decode_random_access": (256, 1, 9216, 64, 512), + "prefill_full": (1, 1024, 1024, 1024, 1024), + "prefill_short": (1, 256, 1024, 1024, 1024), + "prefill_short_small_pages": (1, 256, 1024, 256, 1024), + "prefill_mid_chunk_512": (1, 512, 4096, 1024, 1024), + "prefill_mid_chunk_1024": (1, 1024, 8192, 1024, 1024), + "prefill_mid_chunk_small_pages": (1, 1024, 8192, 64, 512), + }, + 'atol': 0.02, + 'rtol': 0.02, +} + + +def create_inputs(): + configs = CONFIG['configs'] + if isinstance(configs, dict): + configs = [(name, *params) for name, params in configs.items()] + + # --- DSV4 CSA geometry ------------------------------------------------- + NOPE_DIM = 448 + NUM_SCALES = NOPE_DIM // 64 # 7 + ROPE_DIM = 64 + TOKEN_BYTES = 512 # (4, 128) uint8 per token in the nope cache + HEAD_DIM = NOPE_DIM + ROPE_DIM # 512 + NUM_Q_HEADS = 8 + Q_DTYPE = jnp.bfloat16 + SM_SCALE = float(HEAD_DIM ** -0.5) + ATTN_BATCH = 16 + VMEM_LIMIT_BYTES = 100 * 1024 * 1024 + + # --- Cache contents ---------------------------------------------------- + CACHE_TILE = 8192 + k_nope, k_scale, k_rope, k_cfg = jax.random.split(jax.random.key(0), 4) + + nope_f8 = jax.random.normal( + k_nope, (CACHE_TILE, NOPE_DIM), jnp.float32 + ).astype(jnp.float8_e4m3fn) + nope_bytes = jax.lax.bitcast_convert_type(nope_f8, jnp.uint8) + scale_bytes = jax.random.randint( + k_scale, (CACHE_TILE, NUM_SCALES), 125, 130, dtype=jnp.int32 + ).astype(jnp.uint8) + pad_bytes = jnp.zeros( + (CACHE_TILE, TOKEN_BYTES - NOPE_DIM - NUM_SCALES), jnp.uint8 + ) + + nope_tile = jnp.concatenate( + [nope_bytes, scale_bytes, pad_bytes], axis=1 + ) + + rope_bits = jax.lax.bitcast_convert_type( + jax.random.normal( + k_rope, (CACHE_TILE, ROPE_DIM), jnp.float32 + ).astype(jnp.bfloat16), + jnp.uint16, + ).astype(jnp.uint32) + rope_tile = jnp.concatenate( + [ + (rope_bits >> 8).astype(jnp.uint8), + (rope_bits & 0xFF).astype(jnp.uint8), + ], + axis=1, + ) + + def fill(tile, num_slots): + reps = -(-num_slots // tile.shape[0]) + return jnp.tile(tile, (reps, 1))[:num_slots] + + def build_case(name, batch_size, q_len, kv_len, page_size, csa_topk, key): + k_page, k_topk, k_q, k_sink, k_acc, k_l, k_m = jax.random.split(key, 7) + + num_tokens = batch_size * q_len + pages_per_seq = -(-kv_len // page_size) + total_num_pages = batch_size * pages_per_seq + assert num_tokens % ATTN_BATCH == 0 + + q = jax.random.normal( + k_q, (num_tokens, NUM_Q_HEADS, HEAD_DIM), jnp.float32 + ).astype(Q_DTYPE) + + num_slots = total_num_pages * page_size + cache_kv_nope = fill(nope_tile, num_slots).reshape( + total_num_pages, page_size, 4, 128 + ) + cache_kv_rope = fill(rope_tile, num_slots).reshape( + total_num_pages, page_size // 4, 4, 128 + ) + + # Distinct physical pages per sequence: no KV sharing, so the gather + # sees the full footprint a serving batch would. + page_indices = jax.random.permutation( + k_page, total_num_pages + ).astype(jnp.int32) + + # Ragged query offsets. Every sequence contributes q_len tokens. + cu_q_lens = (jnp.arange(batch_size + 1, dtype=jnp.int32) * q_len) + # (decode_end, prefill_end, num_seqs); only [2] is read, but keep the + # decode/prefill split honest. + num_decode = batch_size if q_len == 1 else 0 + distribution = jnp.array( + [num_decode, batch_size, batch_size], jnp.int32 + ) + + pos_in_seq = jnp.arange(num_tokens, dtype=jnp.int32) % q_len + causal_len = (kv_len - q_len + pos_in_seq + 1)[:, None] + slot = jnp.arange(csa_topk, dtype=jnp.int32)[None, :] + scattered = jax.random.randint( + k_topk, (num_tokens, csa_topk), 0, 1 << 30, dtype=jnp.int32 + ) % causal_len + topk_indices = jnp.where( + slot < causal_len, + jnp.where(causal_len < csa_topk, slot, scattered), + -1, + ).astype(jnp.int32) + + # Sliding-window partial from the preceding kernel: running max, + # running denominator (a sum of exponentials, so >= 1) and the + # unnormalised accumulator. + attention_sinks = jax.random.normal( + k_sink, (NUM_Q_HEADS,), jnp.float32 + ) + swa_accumution = jax.random.normal( + k_acc, (num_tokens, NUM_Q_HEADS, HEAD_DIM), jnp.float32 + ).astype(Q_DTYPE) + swa_l = jax.random.uniform( + k_l, (num_tokens, NUM_Q_HEADS), jnp.float32, 1.0, 64.0 + ) + swa_m = jax.random.normal(k_m, (num_tokens, NUM_Q_HEADS), jnp.float32) + + args = [ + q, + cache_kv_nope, + cache_kv_rope, + topk_indices, + page_indices, + cu_q_lens, + distribution, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + ] + return args + + keys = jax.random.split(k_cfg, len(configs)) + return [build_case(*cfg, key) for cfg, key in zip(configs, keys)] + + +# Computation +def main_kernel( + nope_in_hbm_ref: jax.Ref, + rope_in_hbm_ref: jax.Ref, + indices_hbm_ref: jax.Ref, + nope_out_hbm_ref: jax.Ref, + rope_out_hbm_ref: jax.Ref, + *, + core_axis_name: str, + subcore_axis_name: str, + num_row_subchunks: int, + num_streams: int, +): + tpu_info = pltpu.get_tpu_info() + sc_info = tpu_info.sparse_core + assert sc_info is not None + num_simd_lanes = sc_info.num_lanes + num_cores = jax.lax.axis_size((core_axis_name, subcore_axis_name)) + row_subchunk_size = num_simd_lanes + row_chunk_size = row_subchunk_size * num_row_subchunks + block_size = row_chunk_size * num_cores + num_blocks = pl.cdiv(indices_hbm_ref.shape[0], block_size) + + # Inputs are 8-bit; + # nope output stays uint8 (4/int32), rope is unpacked to bf16 (2/int32). + in_bits = jax.dtypes.itemsize_bits(nope_in_hbm_ref.dtype) + in_packing = 32 // in_bits + in_mask = (1 << in_bits) - 1 # 0xFF for 8-bit. + nope_out_packing = 32 // jax.dtypes.itemsize_bits(nope_out_hbm_ref.dtype) + rope_out_bits = jax.dtypes.itemsize_bits(rope_out_hbm_ref.dtype) + rope_out_packing = 32 // rope_out_bits + core_index = lax.axis_index((core_axis_name, subcore_axis_name)) + + # SparseCore gather 32-bit words + nope_in_i32 = nope_in_hbm_ref.bitcast(jnp.int32) + rope_in_i32 = rope_in_hbm_ref.bitcast(jnp.int32) + nope_out_i32 = nope_out_hbm_ref.bitcast(jnp.int32) + rope_out_i32 = rope_out_hbm_ref.bitcast(jnp.int32) + + nope_in_cols = nope_in_i32.shape[1] + rope_in_cols = rope_in_i32.shape[1] + nope_out_cols = nope_out_hbm_ref.shape[1] + rope_out_cols = rope_out_hbm_ref.shape[1] + + def _delta_swap(x, y, shift, swap_mask): + """Exchanges selected bits between two words. + + Swaps x's bits at positions ``swap_mask << shift`` with y's bits at + positions ``swap_mask``; all other bits are left untouched. Worked + example with ``shift=8, swap_mask=0x00FF00FF`` on byte-quads + ``x = [x0 x1 x2 x3]`` and ``y = [y0 y1 y2 y3]`` (byte 0 = least + significant):: + + swap_mask << 8 = 0xFF00FF00 -> x's bytes 1 and 3 + swap_mask = 0x00FF00FF -> y's bytes 0 and 2 + result: x = [x0 y0 x2 y2], y = [x1 y1 x3 y3] + + i.e. x's odd bytes trade places with y's even bytes. The identity + ``t = ((x >> s) ^ y) & mask; x ^= t << s; y ^= t`` does this in 6 ops + with only `t` as scratch. The arithmetic right shift is safe: its + sign-extended high bits are dropped by the `& swap_mask`. + """ + t = jnp.bitwise_and( + jnp.bitwise_xor(jnp.bitwise_right_shift(x, shift), y), swap_mask + ) + x = jnp.bitwise_xor(x, jnp.left_shift(t, shift)) + y = jnp.bitwise_xor(y, t) + return x, y + + def process_nope(gather_ref, out_ref, out_row_base=0): + # A more efficient implementation of `process_nope_reference`, with fewer + # ALU ops. We've seen that SparseCore's Integer ALU FLOPs being the + # bottleneck of the kernel, this is the preferred implementation. + col_slice = pl.ds(0, 128) + low_16 = jnp.int32(0x0000FFFF) # one 16-bit half of each word + low_byte_of_each_half = jnp.int32(0x00FF00FF) # bytes 0 and 2 + for i in range(num_simd_lanes // nope_out_packing): + d = [ + gather_ref[pl.ds(i * nope_out_packing + j, 1), col_slice] + for j in range(nope_out_packing) + ] + # Round 1: transpose the four 2x2 byte blocks -- exchange the high + # 16 bits of each word with the low 16 bits of its distance-2 peer. + d[0], d[2] = _delta_swap(d[0], d[2], 16, low_16) + d[1], d[3] = _delta_swap(d[1], d[3], 16, low_16) + # Round 2: transpose within each 2x2 block -- exchange the odd bytes + # of each word with the even bytes of its adjacent peer. + d[0], d[1] = _delta_swap(d[0], d[1], 8, low_byte_of_each_half) + d[2], d[3] = _delta_swap(d[2], d[3], 8, low_byte_of_each_half) + # d[m] now holds byte lane m of all 4 inputs -> output sub-row m.s + for m in range(nope_out_packing): + out_ref[pl.ds(out_row_base + i, 1), pl.ds(m * 128, 128)] = d[m] + + def process_rope(gather_ref, out_ref, idx_sub, out_row_base=0): + # one (1, 128) uint8 is one token's rope data, which encodes 64 bf16 + # values. (0, 64) are the high bits for bf16 data, (64, 128) are the low. + half = rope_out_cols + col_hi = pl.ds(0, half) + col_lo = pl.ds(half, half) + + def bf16_bits(k): + sub = lax.rem(idx_sub[k], in_packing) + hi = jnp.bitwise_and( + jnp.bitwise_right_shift( + gather_ref[pl.ds(k, 1), col_hi], in_bits * sub + ), + in_mask, + ) + lo = jnp.bitwise_and( + jnp.bitwise_right_shift( + gather_ref[pl.ds(k, 1), col_lo], in_bits * sub + ), + in_mask, + ) + return jnp.bitwise_or(jnp.left_shift(hi, in_bits), lo) + + for t in range(num_simd_lanes // rope_out_packing): + packed = jnp.zeros((1, half), dtype=jnp.int32) + for pk in range(rope_out_packing): + k = t * rope_out_packing + pk + packed = jnp.bitwise_or( + packed, jnp.left_shift(bf16_bits(k), pk * rope_out_bits) + ) + out_ref[pl.ds(out_row_base + t, 1), pl.ds(0, half)] = packed + + def outer_pipeline(idx_ref): + b = pl.program_id(0) + out_row_base = (b * num_cores + core_index) * num_row_subchunks + + # Subchunk handled by stream `s` at inner step `r`. `num_streams` + # independent `pl.Indirect` gathers run concurrently per step, keeping + # several gather DMAs in flight to raise effective read bandwidth. + def subchunk(r, s): + return r * num_streams + s + + def idx_window(r, s): + return idx_ref[ + pl.ds(subchunk(r, s) * row_subchunk_size, row_subchunk_size) + ] + + # Rows produced per stream in each output (nope packs 4/int32, rope 2). + nope_rows_per_stream = row_subchunk_size // nope_out_packing + rope_rows_per_stream = row_subchunk_size // rope_out_packing + + def _body(*refs): + r = pl.program_id(0) + nope_g = refs[0 * num_streams : 1 * num_streams] + rope_g = refs[1 * num_streams : 2 * num_streams] + nope_o = refs[2 * num_streams] + rope_o = refs[2 * num_streams + 1] + for s in range(num_streams): + process_nope( + gather_ref=nope_g[s], + out_ref=nope_o, + out_row_base=s * nope_rows_per_stream, + ) + process_rope( + gather_ref=rope_g[s], + out_ref=rope_o, + idx_sub=idx_window(r, s), + out_row_base=s * rope_rows_per_stream, + ) + + # Have multiple parallel `pl.Indirect` to hide random access read latency. + # Output contiguous memory access, multiple output parallel DMAs not help + # with performance. + + # nope: gather int32 row == index (1 int32 row per entry). + nope_in_specs = tuple( + pl.BlockSpec( + (pl.Indirect(row_subchunk_size), nope_in_cols), + lambda r, s=s: (idx_window(r, s), 0), + ) + for s in range(num_streams) + ) + # rope: gather int32 row == index // in_packing (in_packing entries/row). + rope_in_specs = tuple( + pl.BlockSpec( + (pl.Indirect(row_subchunk_size), rope_in_cols), + lambda r, s=s: (lax.div(idx_window(r, s), in_packing), 0), + ) + for s in range(num_streams) + ) + # One merged output block per cache, covering all `num_streams` + # subchunks. + nope_out_spec = pl.BlockSpec( + (num_streams * nope_rows_per_stream, nope_out_cols), + lambda r: (out_row_base // num_streams + r, 0), + ) + rope_out_spec = pl.BlockSpec( + (num_streams * rope_rows_per_stream, rope_out_cols), + lambda r: (out_row_base // num_streams + r, 0), + ) + pltpu.emit_pipeline( + _body, + grid=(num_row_subchunks // num_streams,), + in_specs=nope_in_specs + rope_in_specs, + out_specs=(nope_out_spec, rope_out_spec), + )( + *([nope_in_i32] * num_streams), + *([rope_in_i32] * num_streams), + nope_out_i32, + rope_out_i32, + ) + + pltpu.emit_pipeline( + outer_pipeline, + grid=(num_blocks,), + in_specs=pl.BlockSpec( + (row_chunk_size,), + lambda b: (b * num_cores + core_index,), + ), + )(indices_hbm_ref) + + +@functools.partial(jax.jit) +def csa_gather( + nope_cache: jax.Array, + rope_cache: jax.Array, + indices: jax.Array, +) -> tuple[jax.Array, jax.Array]: + """Fused SparseCore gather of the nope and rope caches. + + Args: + nope_cache: (total_pages, page_size, 4, 128) uint8. Each (4, 128) uint8 is + token's nope + nope scales. It encodes 448 fp8 + 7 e8m0 scales + padding. + rope_cache: (total_pages, page_size // 4, 4, 128) uint8. Each (1, 128) uint8 + is token's rope. It encodes 64 bf16. + indices: (N,) int32. Token indices into the caches. + + Returns: + nope_out: (N, 512) uint8. + Each (512) uint8 is token's nope. + rope_out: (N, 64) bf16. + Each (64) bf16 is token's rope. + """ + assert indices.ndim == 1, "Indices must be 1D." + assert nope_cache.dtype == rope_cache.dtype, "Caches must share a dtype." + assert nope_cache.dtype == jnp.uint8, "Caches must be uint8." + + # Flatten both caches to 128-wide rows and view as raw bytes. + nope_cache = nope_cache.reshape(-1, nope_cache.shape[3]) + rope_cache = rope_cache.reshape(-1, rope_cache.shape[3]) + sc_info = pltpu.get_tpu_info().sparse_core + assert sc_info is not None, "SparseCore info is missing." + out_size = indices.size + nope_out_cols = 512 + # rope: each 128-byte entry encodes 64 bf16 (high bytes [0:64], low [64:128]). + rope_out_cols = 64 + num_simd_lanes = sc_info.num_lanes + num_cores = sc_info.num_cores * sc_info.num_subcores + + # `num_streams` independent `pl.Indirect` gathers are issued per + # pipeline step to keep multiple gather DMAs in flight. + # See `outer_pipeline` for details. + num_streams = 2 + num_row_subchunks = 32 + assert ( + num_row_subchunks % num_streams == 0 + ), f"{num_streams=} must divide {num_row_subchunks=}." + row_subchunk_size = num_simd_lanes + row_chunk_size = row_subchunk_size * num_row_subchunks + block_size = row_chunk_size * num_cores + out_pad_size = ( + (out_size + block_size - 1) // block_size + ) * block_size - out_size + indices = jnp.pad(indices, ((0, out_pad_size))) + vector_mesh = plsc.VectorSubcoreMesh( + num_cores=sc_info.num_cores, + num_subcores=sc_info.num_subcores, + core_axis_name="core", + subcore_axis_name="subcore", + ) + nope_out, rope_out = pl.kernel( + functools.partial( + main_kernel, + core_axis_name=vector_mesh.core_axis_name, + subcore_axis_name=vector_mesh.subcore_axis_name, + num_row_subchunks=num_row_subchunks, + num_streams=num_streams, + ), + out_type=( + jax.ShapeDtypeStruct( + (out_size + out_pad_size, nope_out_cols), jnp.uint8 + ), + jax.ShapeDtypeStruct( + (out_size + out_pad_size, rope_out_cols), jnp.bfloat16 + ), + ), + compiler_params=pltpu.CompilerParams( + use_tc_tiling_on_sc=True, + needs_layout_passes=True, + disable_bounds_checks=True, + ), + mesh=vector_mesh, + name="sc_csa_gather", + )(nope_cache, rope_cache, indices) + return ( + nope_out[:out_size], + rope_out[:out_size], + ) + + + +DEFAULT_VMEM_LIMIT_BYTES = 100 * 1024 * 1024 + + +def cdiv(a, b): + assert b != 0 + return (a + b - 1) // b + + +def align_to(x, a): + return cdiv(x, a) * a + + +def get_dtype_bitwidth(dtype): + return jax.dtypes.itemsize_bits(dtype) + + +def get_dtype_packing(dtype): + bits = get_dtype_bitwidth(dtype) + return 32 // bits + + +def get_kv_cache_shape( + total_num_pages, + page_size, + kv_dim, + kv_dtype, +): + kv_packing = get_dtype_packing(kv_dtype) + return ( + total_num_pages, + align_to(page_size, kv_packing) // kv_packing, + kv_packing, + align_to(kv_dim, 128), + ) + + +_GATHER_PAGE_CHUNK = 128 + + +def _gather_page_ids_kernel(windows_ref, logical_ref, out_ref, *, num_chunks): + logical = logical_ref[...] # i32[block_tokens, topk] + out = jnp.zeros_like(logical) + for c in range(num_chunks): + window_chunk = windows_ref[ + :, c * _GATHER_PAGE_CHUNK : (c + 1) * _GATHER_PAGE_CHUNK + ] # i32[block_tokens, 128] + local = logical - c * _GATHER_PAGE_CHUNK + gathered = jnp.take_along_axis( + window_chunk, jnp.clip(local, 0, _GATHER_PAGE_CHUNK - 1), axis=1 + ) + out = jnp.where((local >= 0) & (local < _GATHER_PAGE_CHUNK), gathered, out) + out_ref[...] = out + + +def gather_page_ids( + page_indices: jax.Array, # i32[max_num_seqs * pages_per_seq] + seq_page_ids: jax.Array, # i32[num_tokens, topk] (logical page within seq) + seq_ids_segment: jax.Array, # i32[num_tokens] (token -> seq id) + max_num_seqs: int, + *, + block_tokens: int = 8, +) -> jax.Array: + """Gathers physical page ids for the CSA top-k tokens.""" + num_tokens, topk = seq_page_ids.shape + pages_per_seq = page_indices.shape[0] // max_num_seqs + num_chunks = cdiv(pages_per_seq, _GATHER_PAGE_CHUNK) + padded_pps = num_chunks * _GATHER_PAGE_CHUNK + + page_table = page_indices.reshape(max_num_seqs, pages_per_seq) + if padded_pps != pages_per_seq: + page_table = jnp.pad(page_table, ((0, 0), (0, padded_pps - pages_per_seq))) + # Per-token page-table window. This is a whole-row gather. + windows = page_table[seq_ids_segment] # i32[num_tokens, padded_pps] + logical = jnp.clip(seq_page_ids, 0, pages_per_seq - 1) + + padded_tokens = align_to(num_tokens, block_tokens) + if padded_tokens != num_tokens: + pad = padded_tokens - num_tokens + windows = jnp.pad(windows, ((0, pad), (0, 0))) + logical = jnp.pad(logical, ((0, pad), (0, 0))) + + out = pl.pallas_call( + functools.partial(_gather_page_ids_kernel, num_chunks=num_chunks), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=0, + in_specs=[ + pl.BlockSpec((block_tokens, padded_pps), lambda t: (t, 0)), + pl.BlockSpec((block_tokens, topk), lambda t: (t, 0)), + ], + out_specs=pl.BlockSpec((block_tokens, topk), lambda t: (t, 0)), + grid=(padded_tokens // block_tokens,), + ), + out_shape=jax.ShapeDtypeStruct((padded_tokens, topk), jnp.int32), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary",), + disable_bounds_checks=True, + ), + name="gather_page_ids", + )(windows, logical) + return out[:num_tokens] + + +def _dequant_dsv4_fp8(bkv_nope: jax.Array): + """Dequantize FP8 values to BF16.""" + nope_fp8 = pltpu.bitcast(bkv_nope[:, :448], jnp.float8_e4m3fn).astype( + jnp.bfloat16 + ) + nope_scales = pltpu.bitcast( + bkv_nope[:, 448 : 448 + 7], jnp.float8_e8m0fnu + ).astype(jnp.bfloat16) + nope_scales = jnp.repeat(nope_scales.T, 64, axis=0).T + nope = (nope_fp8 * nope_scales).astype(jnp.bfloat16) + return nope + + +def _attention_kernel( + # Prefetch + kv_lens_ref, # [max_num_seqs] + start_end_seq_idx_ref, # [2] (start_seq_idx, end_seq_idx) + sem_ids_ref, # [2] (bi_sem_idx, bo_sem_idx) + bo_ids_ref, # [2, batch_size] (bo_sem_0_seq_idx, bo_sem_1_seq_idx) + # Input + attention_sinks_ref, # float32[num_q_heads] + q_hbm_ref, # [max_num_tokens, num_q_heads, head_dim] + cache_kv_nope_hbm_ref, # [total_num_pages, page_size, nope_dim] + cache_kv_rope_hbm_ref, # [total_num_pages, page_size, rope_dim] + swa_accumution_hbm_ref, # [max_num_tokens, num_q_heads, head_dim] + swa_l_hbm_ref, # [max_num_tokens, num_l_heads] + swa_m_hbm_ref, # [max_num_tokens, num_l_heads] + # Output + o_hbm_ref, # [max_num_tokens, num_q_heads, head_dim] + # Scratch + bkv_nope_x2_ref, # [2, batch_size, page_size, nope_dim] + bkv_rope_x2_ref, # [2, batch_size, page_size, rope_dim] + bq_x2_ref, # [2, batch_size, num_q_heads, head_dim] + bo_x2_ref, # [2, batch_size, num_q_heads, head_dim] + bl_x2_ref, # [2, batch_size, num_l_heads] + bm_x2_ref, # [2, batch_size, num_l_heads] + swa_acc_x2_ref, # [2, batch_size, num_q_heads, head_dim] + sems, # [7, 2, batch_size] + *, + sm_scale: float, + batch_size: int = 1, +): + assert q_hbm_ref.shape == o_hbm_ref.shape + + num_tokens, num_q_heads, head_dim = q_hbm_ref.shape + _, page_size, _ = cache_kv_nope_hbm_ref.shape + assert kv_lens_ref.shape[0] == num_tokens + bkv_sz = page_size + + q_dtype = q_hbm_ref.dtype + q_packing = get_dtype_packing(q_dtype) + # Validate against the KV dtype. + assert o_hbm_ref.dtype == q_dtype + + assert head_dim % 128 == 0 + assert num_q_heads % q_packing == 0 + + start_seq_idx = start_end_seq_idx_ref[0] + end_seq_idx = start_end_seq_idx_ref[1] + + batch_start_seq_idx = start_seq_idx + pl.program_id(0) * batch_size + batch_end_seq_idx = batch_start_seq_idx + batch_size - 1 + + def flash_attention_step1_qk_softmax( + q, # [bq_sz * num_q_heads, head_dim] + kv, # [bkv_sz, head_dim] <- Correspond to data from bkv_*_x2_ref + swa_m, # [bq_sz * num_q_heads], + swa_l, # [bq_sz * num_q_heads], + attention_sinks, # [num_q_heads] + ): + assert len(q.shape) == 2 + assert len(kv.shape) == 2 + assert q.shape[0] % num_q_heads == 0 + assert q.shape[1] == head_dim + assert kv.shape == (bkv_sz, head_dim) + + # Follow FlashAttention-2 forward pass. + s = jnp.einsum("nd,md->nm", q, kv, preferred_element_type=jnp.float32) + s *= sm_scale + + s_rowmax = jnp.max(s, axis=1, keepdims=True) + m_prev = swa_m + m_curr = jnp.maximum(m_prev, s_rowmax) + p = jnp.exp(s - m_curr) + exp_m_diff = jnp.exp(m_prev - m_curr) + p_rowsum = jnp.sum(p, axis=1, keepdims=True) + l_prev = swa_l + l_curr = exp_m_diff * l_prev + p_rowsum + exp_attention_sinks = jnp.exp(attention_sinks - m_curr) + l = l_curr + exp_attention_sinks + + return p, exp_m_diff, l + + def flash_attention_step2_pv( + p, + kv, + exp_m_diff, + swa_acc, + l, + ): + pv = jnp.einsum("nm,md->nd", p, kv, preferred_element_type=jnp.float32) + + o_prev = swa_acc + acc = exp_m_diff * o_prev + pv + out = ( + lax.div(acc, l) + if q_dtype == jnp.float32 + else (acc * pl.reciprocal(l, approx=True)).astype(q_dtype) + ) + return out + + def _async_copy(src, dst, sem, wait): + cp = pltpu.make_async_copy(src, dst, sem) + if wait: + cp.wait() + else: + cp.start() + + def _fetch_bkv(seq_idx, bkv_sem_idx, batch_idx, *, wait=False): + sem_nope = sems.at[0, bkv_sem_idx, batch_idx] + sem_rope = sems.at[6, bkv_sem_idx, batch_idx] + + bkv_nope_vmem_ref = bkv_nope_x2_ref.at[bkv_sem_idx, batch_idx] + bkv_rope_vmem_ref = bkv_rope_x2_ref.at[bkv_sem_idx, batch_idx] + + # The index into cache_kv_hbm_ref should be relative to the current + # chunk. + page_idx = seq_idx - start_seq_idx + if not wait: + _async_copy( + cache_kv_nope_hbm_ref.at[page_idx], + bkv_nope_vmem_ref, + sem_nope, + wait, + ) + _async_copy( + cache_kv_rope_hbm_ref.at[page_idx], + bkv_rope_vmem_ref, + sem_rope, + wait, + ) + else: + # When we wait, we can use a dummy copy to wait for DMAs to complete where + # src == dst. However, the dma size must be correct. + dst_nope = bkv_nope_vmem_ref + _async_copy( + src=dst_nope, + dst=dst_nope, + sem=sem_nope, + wait=True, + ) + dst_rope = bkv_rope_vmem_ref + _async_copy( + src=dst_rope, + dst=dst_rope, + sem=sem_rope, + wait=True, + ) + + def _fetch_bq(seq_idx, bq_sem_idx, batch_idx, *, wait=False): + sem = sems.at[1, bq_sem_idx, batch_idx] + bq_vmem_ref = bq_x2_ref.at[bq_sem_idx, batch_idx] + + _async_copy( + q_hbm_ref.at[seq_idx], + bq_vmem_ref, + sem, + wait, + ) + + def _send_bo(seq_idx, bo_sem_idx, batch_idx, *, wait=False): + sem = sems.at[2, bo_sem_idx, batch_idx] + vmem_ref = bo_x2_ref.at[bo_sem_idx, batch_idx] + + _async_copy( + vmem_ref, + o_hbm_ref.at[seq_idx], + sem, + wait, + ) + + def _fetch_swa(seq_idx, bq_sem_idx, batch_idx, *, wait=False): + sem_acc = sems.at[3, bq_sem_idx, batch_idx] + sem_l = sems.at[4, bq_sem_idx, batch_idx] + sem_m = sems.at[5, bq_sem_idx, batch_idx] + + if not wait: + _async_copy( + swa_accumution_hbm_ref.at[seq_idx], + swa_acc_x2_ref.at[bq_sem_idx, batch_idx], + sem_acc, + wait=False, + ) + _async_copy( + swa_l_hbm_ref.at[seq_idx], + bl_x2_ref.at[bq_sem_idx, batch_idx], + sem_l, + wait=False, + ) + _async_copy( + swa_m_hbm_ref.at[seq_idx], + bm_x2_ref.at[bq_sem_idx, batch_idx], + sem_m, + wait=False, + ) + + else: + dst_acc = swa_acc_x2_ref.at[bq_sem_idx, batch_idx] + _async_copy(src=dst_acc, dst=dst_acc, sem=sem_acc, wait=True) + + dst_l = bl_x2_ref.at[bq_sem_idx, batch_idx] + _async_copy(src=dst_l, dst=dst_l, sem=sem_l, wait=True) + + dst_m = bm_x2_ref.at[bq_sem_idx, batch_idx] + _async_copy(src=dst_m, dst=dst_m, sem=sem_m, wait=True) + + def start_fetch_bkv(seq_idx, bkv_sem_idx, batch_idx): + return _fetch_bkv(seq_idx, bkv_sem_idx, batch_idx) + + def wait_fetch_bkv(seq_idx, bkv_sem_idx, batch_idx): + return _fetch_bkv(seq_idx, bkv_sem_idx, batch_idx, wait=True) + + def start_fetch_bq(seq_idx, bq_sem_idx, batch_idx): + return _fetch_bq(seq_idx, bq_sem_idx, batch_idx) + + def wait_fetch_bq(seq_idx, bq_sem_idx, batch_idx): + return _fetch_bq(seq_idx, bq_sem_idx, batch_idx, wait=True) + + def start_fetch_swa(seq_idx, bq_sem_idx, batch_idx): + return _fetch_swa(seq_idx, bq_sem_idx, batch_idx) + + def wait_fetch_swa(seq_idx, bq_sem_idx, batch_idx): + return _fetch_swa(seq_idx, bq_sem_idx, batch_idx, wait=True) + + def start_send_bo(seq_idx, bo_sem_idx, batch_idx): + bo_ids_ref[bo_sem_idx, batch_idx] = seq_idx + _send_bo(seq_idx, bo_sem_idx, batch_idx) + + def wait_send_bo(bo_sem_idx, batch_idx): + old_seq_idx = bo_ids_ref[bo_sem_idx, batch_idx] + + @pl.when(0 <= old_seq_idx) + def _(): + _send_bo(old_seq_idx, bo_sem_idx, batch_idx, wait=True) + + def load_bq(bq_sem_idx, batch_idx): + q = bq_x2_ref.at[bq_sem_idx, batch_idx][...] + return q + + def load_bkv(bkv_sem_idx, batch_idx): + bkv_nope = bkv_nope_x2_ref.at[bkv_sem_idx, batch_idx][...] + bkv_nope = _dequant_dsv4_fp8(bkv_nope) + + bkv_rope = bkv_rope_x2_ref.at[bkv_sem_idx, batch_idx][...] + bkv = jnp.concatenate([bkv_nope, bkv_rope], axis=-1) + + # In vLLM, multiple caches may overlay on the same KV Tensor. For example, + # compressor state cache write data in bfloat16 / float32 format, certain + # byte pattern are interpreted as NaN in FP8, e.g. float8_e8m0fnu byte 0xFF + # decodes to NaN. + # We need to mask out the data by the actual kv_len to avoid NaN propagting + # to the downstream computation. + kv_len = kv_lens_ref[batch_start_seq_idx + batch_idx] + k_span = lax.broadcasted_iota(jnp.int32, bkv.shape, 0) + bkv = jnp.where(k_span < kv_len, bkv, 0) + return bkv + + def load_swa_output(bq_sem_idx, batch_idx): + swa_acc = swa_acc_x2_ref[bq_sem_idx, batch_idx, ...] + swa_l = bl_x2_ref[bq_sem_idx, batch_idx, :num_q_heads][..., None] + swa_m = bm_x2_ref[bq_sem_idx, batch_idx, :num_q_heads][..., None] + return swa_acc, swa_l, swa_m + + def process(): + + def get_next_seq_ids(seq_idx, bi_sem_idx): + next_seq_idx = seq_idx + batch_size + next_bi_sem_idx = lax.select(bi_sem_idx == 0, 1, 0) + return next_seq_idx, next_bi_sem_idx + + bi_sem_idx = sem_ids_ref[0] + next_seq_idx, next_bi_sem_idx = get_next_seq_ids( + batch_start_seq_idx, bi_sem_idx + ) + + # Prefetch next seq + @pl.when(next_seq_idx < end_seq_idx) + def prefetch_next_seq(): + sem_ids_ref[0] = next_bi_sem_idx + for batch_idx in range(batch_size): + start_fetch_bq(next_seq_idx + batch_idx, next_bi_sem_idx, batch_idx) + start_fetch_swa(next_seq_idx + batch_idx, next_bi_sem_idx, batch_idx) + start_fetch_bkv(next_seq_idx + batch_idx, next_bi_sem_idx, batch_idx) + + bo_sem_idx = sem_ids_ref[1] + sem_ids_ref[1] = lax.select(bo_sem_idx == 0, 1, 0) + attention_sinks = attention_sinks_ref[...][..., None] + + prev_p = None + prev_bkv = None + prev_exp_m_diff = None + prev_l = None + prev_swa_acc = None + + for batch_idx in range(batch_size): + + # Wait for cur blocks if not ready yet + wait_fetch_bq(batch_start_seq_idx + batch_idx, bi_sem_idx, batch_idx) + wait_fetch_swa(batch_start_seq_idx + batch_idx, bi_sem_idx, batch_idx) + wait_fetch_bkv(batch_start_seq_idx + batch_idx, bi_sem_idx, batch_idx) + + bkv = load_bkv(bi_sem_idx, batch_idx) + bq = load_bq(bi_sem_idx, batch_idx) + swa_acc, swa_l, swa_m = load_swa_output(bi_sem_idx, batch_idx) + + p, exp_m_diff, l = flash_attention_step1_qk_softmax( + bq, + bkv, + swa_m, + swa_l, + attention_sinks, + ) + + if prev_p is not None: + assert prev_bkv is not None + assert prev_exp_m_diff is not None + assert prev_l is not None + out = flash_attention_step2_pv( + prev_p, + prev_bkv, + prev_exp_m_diff, + prev_swa_acc, + prev_l, + ) + + # Wait for previous bo to be fully sent before storing new bo. + wait_send_bo(bo_sem_idx, batch_idx - 1) + # Store output from acc to bo. + bo_x2_ref.at[bo_sem_idx, batch_idx - 1][...] = out + # Send cur bo + start_send_bo( + batch_start_seq_idx + batch_idx - 1, bo_sem_idx, batch_idx - 1 + ) + + prev_p = p + prev_bkv = bkv + prev_exp_m_diff = exp_m_diff + prev_l = l + prev_swa_acc = swa_acc + + # end of pipelining loop + assert prev_p is not None + assert prev_bkv is not None + assert prev_exp_m_diff is not None + assert prev_l is not None + out = flash_attention_step2_pv( + prev_p, + prev_bkv, + prev_exp_m_diff, + prev_swa_acc, + prev_l, + ) + + # Wait for previous bo to be fully sent before storing new bo. + wait_send_bo(bo_sem_idx, batch_size - 1) + # Store output from acc to bo. + bo_x2_ref.at[bo_sem_idx, batch_size - 1][...] = out + # Send cur bo + start_send_bo( + batch_start_seq_idx + batch_size - 1, bo_sem_idx, batch_size - 1 + ) + + ### ------- Kernel start ------- ### + + @pl.when(batch_start_seq_idx == start_seq_idx) + def prologue(): + for batch_idx in range(batch_size): + start_fetch_bq(batch_start_seq_idx + batch_idx, 0, batch_idx) + start_fetch_swa(batch_start_seq_idx + batch_idx, 0, batch_idx) + start_fetch_bkv(batch_start_seq_idx + batch_idx, 0, batch_idx) + + process() + + @pl.when(batch_end_seq_idx == end_seq_idx - 1) + def epilogue(): + for i in range(2): + for batch_idx in range(batch_size): + wait_send_bo(i, batch_idx) + + ### ------- Kernel end ------- ### + + +def prepare_q_inputs( + q: jax.Array, # [max_num_tokens, actual_num_q_heads, actual_head_dim], +): + _, actual_num_q_heads, actual_head_dim = q.shape + q_packing = get_dtype_packing(q.dtype) + num_q_heads = align_to(actual_num_q_heads, q_packing) + head_dim = align_to(actual_head_dim, 128) + q = jnp.pad( + q, + ( + (0, 0), + (0, num_q_heads - actual_num_q_heads), + (0, head_dim - actual_head_dim), + ), + constant_values=0, + ) + return q + + +def prepare_swa_inputs( + swa_accumution: jax.Array, # [max_num_tokens, num_q_heads, head_dim] + swa_l: jax.Array, # [max_num_tokens, num_q_heads] + swa_m: jax.Array, # [max_num_tokens, num_q_heads] +): + _, actual_num_q_heads, actual_head_dim = swa_accumution.shape + swa_packing = get_dtype_packing(swa_accumution.dtype) + num_q_heads = align_to(actual_num_q_heads, swa_packing) + head_dim = align_to(actual_head_dim, 128) + swa_accumution = jnp.pad( + swa_accumution, + ( + (0, 0), + (0, num_q_heads - actual_num_q_heads), + (0, head_dim - actual_head_dim), + ), + constant_values=0, + ) + num_l_heads = align_to(num_q_heads, 128) + swa_l = jnp.pad( + swa_l, + ( + (0, 0), + (0, num_l_heads - actual_num_q_heads), + ), + constant_values=0, + ) + swa_m = jnp.pad( + swa_m, + ( + (0, 0), + (0, num_l_heads - actual_num_q_heads), + ), + constant_values=0, + ) + return swa_accumution, swa_l, swa_m + + +def prepare_outputs( + out, # [max_num_tokens, num_q_heads, head_dim] + actual_num_q_heads: int, + actual_head_dim: int, +): + return out[:, :actual_num_q_heads, :actual_head_dim] + + +# Main Attention kernel for DeepSeek V4 CSA (gather and attention) +# Note that the compressed kv tokens of current batch (current forward pass) +# have been written to the `cache_kv` by the compressor module before calling +# this function, `kv_lens` reflects the length after compressed kv cache write. +@functools.partial( + jax.jit, + static_argnames=( + "sm_scale", + "attention_kernel_batch_size", + "gather_and_attention_chunk_size", + "vmem_limit_bytes", + ), +) +def sparse_ragged_paged_attention( + q: jax.Array, # [max_num_tokens, actual_num_q_heads, head_dim] + cache_kv_nope: jax.Array, # [total_num_pages, page_size, 4, 128] + cache_kv_rope: jax.Array, # [total_num_pages, page_size // 4, 4, 128] + topk_indices: jax.Array, # i32[max_num_tokens, csa_topk] + page_indices: jax.Array, # i32[max_num_seqs * pages_per_seq] + cu_q_lens: jax.Array, # i32[max_num_seqs + 1] + distribution: jax.Array, # i32[3] + attention_sinks: jax.Array, # float32[actual_num_q_heads] + swa_accumution: jax.Array, # bf16[max_num_tokens, num_q_heads, head_dim] + swa_l: jax.Array, # float32[max_num_tokens, num_q_heads] + swa_m: jax.Array, # float32[max_num_tokens, num_q_heads] + *, + sm_scale: float = 1.0, + # Kernel optimization params. + gather_and_attention_chunk_size: int | None = None, + attention_kernel_batch_size: int = 16, + vmem_limit_bytes: int = DEFAULT_VMEM_LIMIT_BYTES, +) -> jax.Array: + """MLA Ragged paged attention that supports mixed prefill and decode. + + Args: + q: concatenated all sequences' queries. + cache_kv_nope: the current kv cache for nope. + cache_kv_rope: the current kv cache for rope. + topk_indices: for each query token, the indices of the top k key tokens to + attend to. + page_indices: flattened page indices look-up table by (seq_id, page_id). + cu_q_lens: the cumulative sum of the effective query lengths. Similar to + kv_lens, only the first num_seqs+1 values are valid. + distribution: (i, j, k) represents that sequences[0:i] are decode-only, + sequences[i:j] are chunked-prefill-only, and sequences[j:k] are mixed. The + k is also the total number of sequences. + sm_scale: the softmax scale which will be applied to the Q@K^T. + vmem_limit_bytes: the vmem limit for the pallas kernel. + + Returns: + The output of attention. + """ + # The cache is DSV4 FP8 format. + # nope_cache contains 448 fp8 + 7 fp8 scales, + # rope_cache contains 64 bf16 + assert cache_kv_nope.dtype == jnp.uint8 + assert cache_kv_rope.dtype == jnp.uint8 + if gather_and_attention_chunk_size is None: + gather_and_attention_chunk_size = q.shape[0] + + _, actual_num_q_heads, actual_head_dim = q.shape + + q = prepare_q_inputs(q) # [max_num_tokens, num_q_heads, head_dim] + head_dim = q.shape[-1] + attention_sinks = jnp.pad( + attention_sinks, + (0, q.shape[1] - actual_num_q_heads), + constant_values=jnp.finfo(attention_sinks.dtype).min, + ) + assert swa_accumution.dtype == q.dtype + swa_accumution, swa_l, swa_m = prepare_swa_inputs( + swa_accumution, swa_l, swa_m + ) + + _, page_size, _, _ = cache_kv_nope.shape + + _, num_q_heads, _ = q.shape + max_num_seqs = cu_q_lens.shape[0] - 1 + num_page_indices = page_indices.shape[0] + assert num_page_indices % max_num_seqs == 0 + + def run_mla_kernel( + q: jax.Array, # [max_num_tokens, num_q_heads, head_dim] + cache_kv_nope: jax.Array, # [total_num_pages, page_size, nope_dim] + cache_kv_rope: jax.Array, # [total_num_pages, page_size, rope_dim] + kv_lens: jax.Array, # i32[max_num_seqs] + attention_sinks: jax.Array, # float32[num_q_heads] + swa_accumution: jax.Array, # bf16[max_num_tokens, num_q_heads, head_dim] + swa_l: jax.Array, # float32[max_num_tokens, num_l_heads] + swa_m: jax.Array, # float32[max_num_tokens, num_l_heads] + start_seq_idx: jax.Array, # i32 + end_seq_idx: jax.Array, # i32 + kernel_batch_size: int, + ): + batch_size = kernel_batch_size + end_seq_idx = jnp.maximum(start_seq_idx, end_seq_idx) + grid = (cdiv(end_seq_idx - start_seq_idx, batch_size),) + in_specs = [ + pl.BlockSpec(memory_space=pltpu.VMEM), # attention_sinks + pl.BlockSpec(memory_space=pltpu.HBM), # q + pl.BlockSpec(memory_space=pltpu.HBM), # cache_kv_nope + pl.BlockSpec(memory_space=pltpu.HBM), # cache_kv_rope + pl.BlockSpec(memory_space=pltpu.HBM), # swa_accumution + pl.BlockSpec(memory_space=pltpu.HBM), # swa_l + pl.BlockSpec(memory_space=pltpu.HBM), # swa_m + ] + + out_specs = pl.BlockSpec(memory_space=pltpu.HBM) # o + + page_size = cache_kv_nope.shape[1] + bkv_nope_double_buf = pltpu.VMEM( + (2, batch_size, page_size, *cache_kv_nope.shape[2:]), + cache_kv_nope.dtype, + ) + bkv_rope_double_buf = pltpu.VMEM( + (2, batch_size, page_size, *cache_kv_rope.shape[2:]), + cache_kv_rope.dtype, + ) + + bq_double_bufq = pltpu.VMEM( + (2, batch_size, num_q_heads, head_dim), + q.dtype, + ) + + bo_double_buf = bq_double_bufq + + num_l_heads = align_to(num_q_heads, 128) + bl_double_buf = pltpu.VMEM( + (2, batch_size, num_l_heads), + jnp.float32, + ) + bm_double_buf = bl_double_buf + + swa_acc_double_buf = pltpu.VMEM( + (2, batch_size, num_q_heads, head_dim), + q.dtype, + ) + + scratch_shapes = [ + bkv_nope_double_buf, + bkv_rope_double_buf, + bq_double_bufq, + bo_double_buf, # Double buffering for output block. + bl_double_buf, # Double buffering for l output. + bm_double_buf, # Double buffering for m output. + swa_acc_double_buf, # Buffer for swa_accumution. + # Semaphores for double buffering of bkv_nope, bq, bo, swa_acc, swa_l, swa_m, bkv_rope + pltpu.SemaphoreType.DMA((7, 2, batch_size)), + ] + + scalar_prefetches = ( + kv_lens, + jnp.array([start_seq_idx, end_seq_idx], jnp.int32), + # (bi_sem_idx, bo_sem_idx) + jnp.zeros((2,), jnp.int32), + # (bo_sem_0_seq_idx, bo_sem_1_seq_idx) + jnp.full((2, batch_size), -1, jnp.int32), + ) + + scope_name = f"MLA-p_{page_size}" + kernel = jax.named_scope(scope_name)( + pl.pallas_call( + functools.partial( + _attention_kernel, + sm_scale=sm_scale, + batch_size=batch_size, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=len(scalar_prefetches), + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + scratch_shapes=scratch_shapes, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary",), + vmem_limit_bytes=vmem_limit_bytes, + disable_bounds_checks=True, + ), + out_shape=jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype), + input_output_aliases={ + 5: 0, # Alias output activation with q + }, + name=scope_name, + ) + ) + return kernel( + *scalar_prefetches, + attention_sinks, + q, + cache_kv_nope, + cache_kv_rope, + swa_accumution, + swa_l, + swa_m, + ) + + tokens_per_seq = cu_q_lens[1:] - cu_q_lens[:-1] + seq_ids_segment = jnp.repeat( + jnp.arange(max_num_seqs), tokens_per_seq, total_repeat_length=q.shape[0] + ) + assert topk_indices is not None + # TODO: skip gather for padding tokens in topk_indices. + kv_lens = jnp.sum(topk_indices != -1, axis=-1) + + seq_page_ids = topk_indices // page_size + token_offset = topk_indices % page_size + topk = topk_indices.shape[-1] + page_ids = gather_page_ids( + page_indices, seq_page_ids, seq_ids_segment, max_num_seqs + ) + + # For the "-1" padding elements in topk_indices, we scatter the corresponding + # page_ids and token_offset to avoid gather memory access hotspotting. + is_padding = topk_indices == -1 + total_num_pages = cache_kv_nope.shape[0] + flat_element_index = jnp.arange(q.shape[0] * topk, dtype=jnp.int32).reshape( + q.shape[0], topk + ) + # 104729 and 15485863 are randomly chosen large prime numbers. + scattered_page_ids = (flat_element_index * 104729) % total_num_pages + scattered_token_offset = (flat_element_index * 15485863) % page_size + page_ids = jnp.where(is_padding, scattered_page_ids, page_ids) + token_offset = jnp.where( + is_padding, + scattered_token_offset, + token_offset, + ) + + assert page_ids.shape == (q.shape[0], topk) + + # TODO: handle the case where q.shape[0] is not divisible by + # gather_and_attention_chunk_size. + assert q.shape[0] % gather_and_attention_chunk_size == 0 + num_chunks = q.shape[0] // gather_and_attention_chunk_size + + for i in range(num_chunks): + start_pos = i * gather_and_attention_chunk_size + end_pos = start_pos + gather_and_attention_chunk_size + indices = ( + page_ids[start_pos:end_pos, ...] * page_size + + token_offset[start_pos:end_pos, ...] + ).reshape(-1) + + # For prefilling of short sequences (or early in the sequence), there are + # very few number of KVs in the sequence, so different qs' selected topk + # would have large overlap. This causes gather read hotspotting. We've seen + # 30%+ performance degradation compared to the no-duplicate-indices case. + # + # TODO: we could consider let the caller (tpu-runner) to sort the sequences + # based on their lengths. For the sequences-segment below certain length, + # we use a different kernel (dense attention and mask), for the rest of + # sequences, we use this gather-and-attention kernel. + gathered_nope_buffer, gathered_rope_buffer = csa_gather( + cache_kv_nope, + cache_kv_rope, + indices, + ) + gathered_nope_buffer = gathered_nope_buffer.reshape( + gather_and_attention_chunk_size, topk, -1 + ) + gathered_rope_buffer = gathered_rope_buffer.reshape( + gather_and_attention_chunk_size, topk, -1 + ) + # We treat each query token as a one independent sequence, attend to their + # respective gathered kv tokens in the `gathered_kv_buffer`. + # -1 in topk_indices is padded elements at the end of each row. + # Batching + assert gather_and_attention_chunk_size % attention_kernel_batch_size == 0 + batch_end = ( + cdiv( + jnp.minimum( + cu_q_lens[distribution[2]], + start_pos + gather_and_attention_chunk_size, + ), + attention_kernel_batch_size, + ) + * attention_kernel_batch_size + ) + q = run_mla_kernel( + q, + gathered_nope_buffer, + gathered_rope_buffer, + kv_lens, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + start_seq_idx=start_pos, + end_seq_idx=batch_end, + kernel_batch_size=attention_kernel_batch_size, + ) + return prepare_outputs( + q, actual_num_q_heads, actual_head_dim + ) # [max_num_tokens, actual_num_q_heads, actual_head_dim] + +def workload( + q: jax.Array, + cache_kv_nope: jax.Array, + cache_kv_rope: jax.Array, + topk_indices: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + distribution: jax.Array, + attention_sinks: jax.Array, + swa_accumution: jax.Array, + swa_l: jax.Array, + swa_m: jax.Array, +): + sm_scale = float(512 ** -0.5) + gather_and_attention_chunk_size = None + attention_kernel_batch_size = 16 + vmem_limit_bytes = 100 * 1024 * 1024 + + return sparse_ragged_paged_attention( + q, + cache_kv_nope, + cache_kv_rope, + topk_indices, + page_indices, + cu_q_lens, + distribution, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + sm_scale=sm_scale, + gather_and_attention_chunk_size=gather_and_attention_chunk_size, + attention_kernel_batch_size=attention_kernel_batch_size, + vmem_limit_bytes=vmem_limit_bytes, + ) + +def benchmark(num_warmup=5, num_iters=100): + """Benchmark and return results dict.""" + inputs_list = create_inputs() + fn = jax.jit(workload) + times_list = [] + std_ms_list = [] + output_shape_list = [] + for inputs in inputs_list: + for _ in range(num_warmup): + out = fn(*inputs) + jax.block_until_ready(out) + times = [] + for _ in range(num_iters): + t0 = time.perf_counter() + out = fn(*inputs) + jax.block_until_ready(out) + times.append(time.perf_counter() - t0) + times = np.array(times) * 1000 + avg = float(np.mean(times)) + times_list.append(round(avg, 4)) + std_ms_list.append(round(float(np.std(times)), 4)) + if hasattr(out, 'shape'): + out_shape = list(out.shape) + elif isinstance(out, (tuple, list)): + out_shape = [list(x.shape) if hasattr(x, 'shape') else [] for x in out] + else: + out_shape = [] + output_shape_list.append(out_shape) + return { + 'name': CONFIG['name'], + 'model': CONFIG['model'], + 'operator': CONFIG['operator'], + 'config': {k: v for k, v in CONFIG.items() if k not in ('name', 'model', 'operator', 'atol', 'rtol')}, + 'time_ms': times_list, + 'std_ms': std_ms_list, + 'output_shape': output_shape_list, + 'status': 'success', + } + + +if __name__ == '__main__': + import json + print(json.dumps(benchmark())) diff --git a/JAXBench/benchmark/level2/52p_DeepSeek_V4_HCA/baseline.py b/JAXBench/benchmark/level2/52p_DeepSeek_V4_HCA/baseline.py new file mode 100644 index 0000000..0237a0a --- /dev/null +++ b/JAXBench/benchmark/level2/52p_DeepSeek_V4_HCA/baseline.py @@ -0,0 +1,1090 @@ +"""TPU-Friendly MLA Ragged Paged Attention kernel.""" + +from enum import Enum +import functools + +import jax +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +import numpy as np +import time + + +DEFAULT_VMEM_LIMIT_BYTES = 100 * 1024 * 1024 + + +CONFIG = { + 'name': 'dsv4_hca', + 'model': 'deepseek', + 'operator': 'hybrid_context_attention', + 'configs': { + "decode_large_batch": (256, 1, 9216, 1024, 16, 1), + "prefill_full": (1, 1024, 1024, 1024, 16, 32), + "prefill_full_8x": (1, 1024, 8192, 1024, 16, 32), + "decode_medium_batch": (128, 1, 9216, 1024, 16, 1), + "prefill_short": (1, 256, 1024, 1024, 16, 32), + "prefill_mid_chunk_512": (1, 512, 4096, 1024, 16, 32), + "decode_small_pages": (128, 1, 9216, 256, 16, 1), + "prefill_short_small_pages": (1, 256, 1024, 256, 16, 32), + "decode_large_pages": (256, 1, 9216, 2048, 16, 1), + "prefill_large_pages": (1, 1024, 8192, 2048, 4, 32), + }, + 'atol': 0.005, + 'rtol': 0.02, +} + +def create_inputs(): + configs = CONFIG['configs'] + if isinstance(configs, dict): + configs = [(name, *params) for name, params in configs.items()] + + HEAD_DIM = 512 + NUM_Q_HEADS = 8 + Q_DTYPE = jnp.bfloat16 + + CACHE_TILE = 8192 + k_kv, k_cfg = jax.random.split(jax.random.key(0), 2) + + kv_tile = jax.random.normal( + k_kv, (CACHE_TILE, 4, 128), jnp.float32 + ).astype(Q_DTYPE) + bits = jax.lax.bitcast_convert_type(kv_tile, jnp.uint16) + kv_tile_bytes = jnp.stack( + [(bits & 0xFF).astype(jnp.uint8), (bits >> 8).astype(jnp.uint8)], + axis=2, + ).reshape(CACHE_TILE, 8, 128) + + def build_case(name, batch_size, q_len, kv_len, page_size, bkv_pages, bq, key): + k_page, k_q, k_sink, k_acc, k_l, k_m = jax.random.split(key, 6) + + num_tokens = batch_size * q_len + pages_per_seq = -(-kv_len // page_size) + total_num_pages = batch_size * pages_per_seq + num_slots = total_num_pages * page_size + + q = jax.random.normal( + k_q, (num_tokens, NUM_Q_HEADS, HEAD_DIM), jnp.float32 + ).astype(Q_DTYPE) + + reps = -(-num_slots // CACHE_TILE) + cache_kv = jnp.tile(kv_tile_bytes, (reps, 1, 1))[:num_slots].reshape( + total_num_pages, page_size * 2, 4, 128 + ) + + page_indices = jax.random.permutation( + k_page, total_num_pages + ).astype(jnp.int32) + + kv_lens = jnp.full((batch_size,), kv_len, jnp.int32) + pos_in_seq = jnp.arange(num_tokens, dtype=jnp.int32) % q_len + kv_lens_to_attend = kv_len - q_len + pos_in_seq + 1 + cu_q_lens = jnp.arange(batch_size + 1, dtype=jnp.int32) * q_len + + if q_len == 1: + distribution = jnp.array( + [batch_size, batch_size, batch_size], jnp.int32 + ) + chunk_prefill_size = None + else: + distribution = jnp.array([0, batch_size, batch_size], jnp.int32) + chunk_prefill_size = q_len + + attention_sinks = jax.random.normal( + k_sink, (NUM_Q_HEADS,), jnp.float32 + ) + swa_accumution = jax.random.normal( + k_acc, (num_tokens, NUM_Q_HEADS, HEAD_DIM), jnp.float32 + ).astype(Q_DTYPE) + swa_l = jax.random.uniform( + k_l, (num_tokens, NUM_Q_HEADS), jnp.float32, 1.0, 64.0 + ) + swa_m = jax.random.normal(k_m, (num_tokens, NUM_Q_HEADS), jnp.float32) + + args = [ + q, + cache_kv, + kv_lens, + kv_lens_to_attend, + page_indices, + cu_q_lens, + distribution, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + ] + return args + + keys = jax.random.split(k_cfg, len(configs)) + return [build_case(*cfg, key) for cfg, key in zip(configs, keys)] + + + + +def cdiv(a, b): + assert b != 0 + return (a + b - 1) // b + + +def align_to(x, a): + return cdiv(x, a) * a + + +def get_dtype_bitwidth(dtype): + return jax.dtypes.itemsize_bits(dtype) + + +def get_dtype_packing(dtype): + bits = get_dtype_bitwidth(dtype) + return 32 // bits + + +class MlaCase(Enum): + """Represents the different cases for MLA. + + - DECODE: Sequences are in decode-only mode (q_len = 1). + - PREFILL: Sequences are in prefill-only mode (q_len > 1, static). + - MIXED: Sequences can be a mix of prefill and decode (q_len > 1, dynamic). + """ + + DECODE = 0 + PREFILL = 1 + MIXED = 2 + + @property + def symbol(self): + return { + MlaCase.DECODE: "d", + MlaCase.PREFILL: "p", + MlaCase.MIXED: "m", + }[self] + + +def _mla_ragged_paged_attention_kernel( + # Prefetch + kv_lens_ref, # [max_num_seqs] + kv_lens_to_attend_ref, # [max_num_tokens] + page_indices_ref, # [max_num_seqs * pages_per_seq] + cu_q_lens_ref, # [max_num_seqs + 1] + start_end_seq_idx_ref, # [2] (start_seq_idx, end_seq_idx) + sem_ids_ref, # [3] (bq_sem_idx, bkv_sem_idx, bo_sem_idx) + bo_ids_ref, # [4] (bo_sem_0_seq_idx, bo_sem_1_seq_idx, bo_sem_0_bo_idx, bo_sem_1_bo_idx) + # Input + attention_sinks_ref, # float32[num_q_heads] + q_hbm_ref, # [max_num_tokens, num_q_heads, head_dim] + cache_kv_hbm_ref, # [total_num_pages, page_size_per_kv_packing, kv_packing, lkv_dim] + swa_accumution_hbm_ref, # [max_num_tokens, num_q_heads, head_dim] + swa_l_hbm_ref, # [max_num_tokens, num_l_heads] + swa_m_hbm_ref, # [max_num_tokens, num_l_heads] + # Output + o_hbm_ref, # [max_num_tokens, num_q_heads, head_dim] + # Scratch + bkv_x2_ref, # [2, bkv_buf_sz_per_kv_packing, kv_packing, lkv_dim] + bq_x2_ref, # [2, bq_sz, num_q_heads, head_dim] + bo_x2_ref, # [2, bq_sz, num_q_heads, head_dim] + bl_x2_ref, # [2, bq_sz, num_l_heads] + bm_x2_ref, # [2, bq_sz, num_l_heads] + swa_acc_x2_ref, # [2, bq_sz, num_q_heads, head_dim] + sems, # [7, 2] + l_ref, # [bq_sz * num_q_heads, 128], + m_ref, # [bq_sz * num_q_heads, 128], + acc_ref, # [bq_sz * num_q_heads, head_dim], + *, + static_q_len: int, + sm_scale: float, + bkv_p, + bq_sz, +): + assert q_hbm_ref.shape == o_hbm_ref.shape + + _, num_q_heads, head_dim = q_hbm_ref.shape + total_num_pages, num_slots, kv_packing, lkv_dim = cache_kv_hbm_ref.shape + num_slots_per_token = 2 + assert num_slots % num_slots_per_token == 0 + page_size = num_slots // num_slots_per_token + max_num_seqs = kv_lens_ref.shape[0] + num_page_indices = page_indices_ref.shape[0] + + assert num_page_indices % max_num_seqs == 0 + pages_per_seq = num_page_indices // max_num_seqs + q_dtype = q_hbm_ref.dtype + q_packing = get_dtype_packing(q_dtype) + # Validate against the KV dtype. + kv_dtype = cache_kv_hbm_ref.dtype + assert o_hbm_ref.dtype == q_dtype + assert get_dtype_packing(kv_dtype) == kv_packing + assert lkv_dim % 128 == 0 + assert head_dim % 128 == 0 + bkv_sz = bkv_p * page_size + assert num_q_heads % q_packing == 0 + num_q_heads_per_q_packing = num_q_heads // q_packing + + start_seq_idx = start_end_seq_idx_ref[0] + end_seq_idx = start_end_seq_idx_ref[1] + seq_idx = pl.program_id(0) + start_seq_idx + q_start = cu_q_lens_ref[seq_idx] + q_end = cu_q_lens_ref[seq_idx + 1] + q_len = q_end - q_start + kv_len = kv_lens_ref[seq_idx] + + def flash_attention( + q, # [bq_sz * num_q_heads, head_dim] + kv, # [bkv_sz, head_dim] <- Correspond to data from bkv_x2_ref + *, + bq_idx, + bkv_idx, + kv_lens_to_attend_segment, + ): + assert len(q.shape) == 2 + assert len(kv.shape) == 2 + assert q.shape[0] % num_q_heads == 0 + assert q.shape[1] == head_dim + assert kv.shape == (bkv_sz, head_dim) + head_l_ref = l_ref.at[: q.shape[0]] + head_m_ref = m_ref.at[: q.shape[0]] + head_acc_ref = acc_ref.at[: q.shape[0]] + + # Follow FlashAttention-2 forward pass. + s = jnp.einsum("nd,md->nm", q, kv, preferred_element_type=jnp.float32) + s *= sm_scale + + k_span = bkv_idx * bkv_sz + lax.broadcasted_iota(jnp.int32, s.shape, 1) + mask = kv_lens_to_attend_segment.reshape(s.shape) <= k_span + + s = jnp.where(mask, jnp.finfo(s.dtype).min, s) + s_rowmax = jnp.max(s, axis=1, keepdims=True) + m_prev = head_m_ref[...] + m_curr = jnp.maximum(m_prev, s_rowmax) + head_m_ref[...] = m_curr + p = jnp.exp(s - broadcast_minor(m_curr, s.shape)) + + pv = jnp.einsum("nm,md->nd", p, kv, preferred_element_type=jnp.float32) + + p_rowsum = jnp.sum(p, axis=1, keepdims=True) + exp_m_diff = jnp.exp(m_prev - m_curr) + l_prev = head_l_ref[...] + l_curr = exp_m_diff * l_prev + p_rowsum + head_l_ref[...] = l_curr + o_prev = head_acc_ref[...] + o_curr = broadcast_minor(exp_m_diff, o_prev.shape) * o_prev + pv + head_acc_ref[...] = o_curr + + def _async_copy(src, dst, sem, wait): + cp = pltpu.make_async_copy(src, dst, sem) + if wait: + cp.wait() + else: + cp.start() + + def _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, *, wait=False): + sem = sems.at[0, bkv_sem_idx] + # bkv_x2_ref shape: [2, bkv_sz, num_slots_per_token * kv_packing, lkv_dim] + bkv_vmem_ref = bkv_x2_ref.at[bkv_sem_idx] + + reshaped_cache_hbm_ref = cache_kv_hbm_ref.reshape( + total_num_pages * page_size, + num_slots_per_token * kv_packing, + lkv_dim, + ) + + kv_len = kv_lens_ref[seq_idx] + kv_len_start = bkv_idx * bkv_sz + kv_p_start = bkv_idx * bkv_p + + kv_left = kv_len - kv_len_start + dma_bkv_sz = jnp.minimum(kv_left, bkv_sz) + page_indices_offset = seq_idx * pages_per_seq + kv_p_start + + if not wait: + # Fetch effective kv from kv cache. To pipeline multiple DMA calls, we + # utilize static for loop instead of dynamic for loop. + # Loop through all pages in a block + for i in range(bkv_p): + # Ensure only effective kvs are copied and we don't go negative. + sz = jnp.clip( + kv_left - i * page_size, + 0, + page_size, + ) + # If the page index is out of bound, we set page_idx to the last page. + # And there will be no copy since sz will be 0. + page_idx = jnp.minimum(page_indices_offset + i, num_page_indices - 1) + _async_copy( + reshaped_cache_hbm_ref.at[ + pl.ds( + page_indices_ref[page_idx] * page_size, + sz, + ), + ], + bkv_vmem_ref.at[pl.ds(i * page_size, sz)], + sem, + wait, + ) + + else: + # When we wait, we can use a dummy copy to wait for DMAs to complete where + # src == dst. However, the dma size must be correct. + dst_kv = bkv_vmem_ref.at[pl.ds(0, dma_bkv_sz)] + _async_copy( + src=dst_kv, + dst=dst_kv, + sem=sem, + wait=True, + ) + + def _fetch_bq(seq_idx, bq_idx, bq_sem_idx, *, wait=False): + sem = sems.at[1, bq_sem_idx] + bq_vmem_ref = bq_x2_ref.at[bq_sem_idx] + + q_len_start = cu_q_lens_ref[seq_idx] + bq_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy( + q_hbm_ref.at[pl.ds(q_len_start, sz)], + bq_vmem_ref.at[pl.ds(0, sz)], + sem, + wait, + ) + + def _send_bo(seq_idx, bo_idx, bo_sem_idx, *, wait=False): + sem = sems.at[2, bo_sem_idx] + vmem_ref = bo_x2_ref.at[bo_sem_idx] + q_len_start = cu_q_lens_ref[seq_idx] + bo_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy( + vmem_ref.at[pl.ds(0, sz)], + o_hbm_ref.at[pl.ds(q_len_start, sz)], + sem, + wait, + ) + + def _fetch_swa(seq_idx, bq_idx, bq_sem_idx, *, wait=False): + sem_acc = sems.at[3, bq_sem_idx] + sem_l = sems.at[4, bq_sem_idx] + sem_m = sems.at[5, bq_sem_idx] + + q_len_start = cu_q_lens_ref[seq_idx] + bq_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + if not wait: + _async_copy( + swa_accumution_hbm_ref.at[pl.ds(q_len_start, sz)], + swa_acc_x2_ref.at[bq_sem_idx, pl.ds(0, sz)], + sem_acc, + wait=False, + ) + _async_copy( + swa_l_hbm_ref.at[pl.ds(q_len_start, sz)], + bl_x2_ref.at[bq_sem_idx, pl.ds(0, sz)], + sem_l, + wait=False, + ) + _async_copy( + swa_m_hbm_ref.at[pl.ds(q_len_start, sz)], + bm_x2_ref.at[bq_sem_idx, pl.ds(0, sz)], + sem_m, + wait=False, + ) + + else: + dst_acc = swa_acc_x2_ref.at[bq_sem_idx, pl.ds(0, sz)] + _async_copy(src=dst_acc, dst=dst_acc, sem=sem_acc, wait=True) + + dst_l = bl_x2_ref.at[bq_sem_idx, pl.ds(0, sz)] + _async_copy(src=dst_l, dst=dst_l, sem=sem_l, wait=True) + + dst_m = bm_x2_ref.at[bq_sem_idx, pl.ds(0, sz)] + _async_copy(src=dst_m, dst=dst_m, sem=sem_m, wait=True) + + acc_ref[...] = ( + swa_acc_x2_ref[bq_sem_idx, ...] + .astype(jnp.float32) + .reshape(bq_sz * num_q_heads, head_dim) + ) + bl = jnp.concat( + [bl_x2_ref[bq_sem_idx, i, :num_q_heads] for i in range(bq_sz)] + )[..., None] + l_ref[...] = jnp.concat([bl for _ in range(128)], axis=-1) + + bm = jnp.concat( + [bm_x2_ref[bq_sem_idx, i, :num_q_heads] for i in range(bq_sz)] + )[..., None] + m_ref[...] = jnp.concat([bm for _ in range(128)], axis=-1) + + def start_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx): + return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx) + + def wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx): + return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, wait=True) + + def start_fetch_bq(seq_idx, bq_idx, bq_sem_idx): + return _fetch_bq(seq_idx, bq_idx, bq_sem_idx) + + def wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx): + return _fetch_bq(seq_idx, bq_idx, bq_sem_idx, wait=True) + + def start_fetch_swa(seq_idx, bq_idx, bq_sem_idx): + return _fetch_swa(seq_idx, bq_idx, bq_sem_idx) + + def wait_fetch_swa(seq_idx, bq_idx, bq_sem_idx): + return _fetch_swa(seq_idx, bq_idx, bq_sem_idx, wait=True) + + def start_send_bo(seq_idx, bo_idx, bo_sem_idx): + bo_ids_ref[bo_sem_idx] = seq_idx + bo_ids_ref[bo_sem_idx + 2] = bo_idx + _send_bo(seq_idx, bo_idx, bo_sem_idx) + + def wait_send_bo(bo_sem_idx): + old_seq_idx = bo_ids_ref[bo_sem_idx] + old_bo_idx = bo_ids_ref[bo_sem_idx + 2] + + @pl.when(jnp.logical_and(0 <= old_seq_idx, old_seq_idx <= seq_idx)) + def _(): + _send_bo(old_seq_idx, old_bo_idx, bo_sem_idx, wait=True) + + def load_bq(bq_sem_idx): + q_ref = ( + bq_x2_ref.bitcast(jnp.uint32) + .at[bq_sem_idx] + .reshape(bq_sz * num_q_heads_per_q_packing, head_dim) + ) + q = pltpu.bitcast( + q_ref[: bq_sz * num_q_heads_per_q_packing], + q_dtype, + ).reshape(bq_sz * num_q_heads, head_dim) + return q + + def load_bkv(bkv_sem_idx, bkv_idx): + bkv_u8 = pltpu.bitcast(bkv_x2_ref.at[bkv_sem_idx][...], jnp.uint8) + assert bkv_u8.shape[-1] == cache_kv_hbm_ref.shape[-1] + bkv_bf16 = pltpu.bitcast(bkv_u8, jnp.bfloat16) + bkv = bkv_bf16.reshape(-1, head_dim) + assert bkv.shape == (bkv_sz, head_dim) + return bkv + + def broadcast_minor(src, shape): + if src.shape == shape: + return src + assert src.shape[:-1] == shape[:-1] + assert src.shape[-1] % 128 == 0 + target_minor = align_to(shape[-1], src.shape[-1]) + # no-op concatenation. + return jnp.concatenate( + [src for _ in range(target_minor // src.shape[-1])], axis=-1 + )[..., : shape[-1]] + + def process(): + # Force at least one bkv block and one bq block per sequence: the + # double-buffered DMA pipeline hands the bkv and bq semaphore across + # sequence boundaries and assumes every sequence runs >=1 bkv and bq + # iteration. + num_bkv = jnp.maximum(1, cdiv(kv_len, bkv_sz)) + if static_q_len is None: + num_bq = jnp.maximum(1, cdiv(q_len, bq_sz)) + else: + num_bq = jnp.maximum(1, cdiv(static_q_len, bq_sz)) + + def get_next_bq_ids(seq_idx, bq_idx, bq_sem_idx): + next_bq_idx = bq_idx + 1 + is_last_bq = next_bq_idx == num_bq + next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx) + next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx) + next_bq_sem_idx = lax.select(bq_sem_idx == 0, 1, 0) + return next_seq_idx, next_bq_idx, next_bq_sem_idx + + def get_next_bkv_ids(seq_idx, bq_idx, bkv_idx, bkv_sem_idx): + next_bkv_idx = bkv_idx + 1 + is_last_bkv = next_bkv_idx == num_bkv + next_bkv_idx = lax.select(is_last_bkv, 0, next_bkv_idx) + next_bq_idx = lax.select(is_last_bkv, bq_idx + 1, bq_idx) + is_last_bq = next_bq_idx == num_bq + next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx) + next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx) + next_bkv_sem_idx = lax.select(bkv_sem_idx == 0, 1, 0) + return next_seq_idx, next_bq_idx, next_bkv_idx, next_bkv_sem_idx + + def compute_with_bq(bq_idx, _): + bq_sem_idx = sem_ids_ref[0] + next_seq_idx, next_bq_idx, next_bq_sem_idx = get_next_bq_ids( + seq_idx, bq_idx, bq_sem_idx + ) + + kv_lens_to_attend_segment = jnp.broadcast_to( + jnp.stack([ + kv_lens_to_attend_ref[q_start + bq_idx * bq_sz + i] + for i in range(bq_sz) + ])[:, None, None], + (bq_sz, num_q_heads, bkv_sz), + ) + + # Prefetch next bq + @pl.when(next_seq_idx < end_seq_idx) + def prefetch_next_bq(): + sem_ids_ref[0] = next_bq_sem_idx + start_fetch_bq(next_seq_idx, next_bq_idx, next_bq_sem_idx) + start_fetch_swa(next_seq_idx, next_bq_idx, next_bq_sem_idx) + + def compute_with_bkv(bkv_idx, carry): + kv_lens_to_attend_segment = carry[0] + + # Get next bkv ids. + bkv_sem_idx = sem_ids_ref[1] + next_seq_idx, _, next_bkv_idx, next_bkv_sem_idx = get_next_bkv_ids( + seq_idx, bq_idx, bkv_idx, bkv_sem_idx + ) + + # Prefetch next bkv + @pl.when(next_seq_idx < end_seq_idx) + def prefetch_next_bkv(): + sem_ids_ref[1] = next_bkv_sem_idx + start_fetch_bkv(next_seq_idx, next_bkv_idx, next_bkv_sem_idx) + + # Wait for cur bkv + wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx) + + # Load bkv into vreg. There is no need to mask out invalid k/v entries, + # because the score of invalid Q.K^T pairs are masked (to be zero) in + # flash attention, so that the invalid kv entries + # (as long as they are not NaN or inf) won't affect to the output. + bkv = load_bkv(bkv_sem_idx, bkv_idx) + + bq = load_bq(bq_sem_idx) + + flash_attention( + bq, + bkv, + bq_idx=bq_idx, + bkv_idx=bkv_idx, + kv_lens_to_attend_segment=kv_lens_to_attend_segment, + ) + return (kv_lens_to_attend_segment,) + + # Wait for cur bq if not ready yet + wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx) + wait_fetch_swa(seq_idx, bq_idx, bq_sem_idx) + jax.lax.fori_loop( + 0, + num_bkv, + compute_with_bkv, + (kv_lens_to_attend_segment,), + unroll=False, + ) + + # Load acc and calculate final output. + acc = acc_ref[...] + attention_sinks = jnp.concat( + [attention_sinks_ref[...] for _ in range(bq_sz)] + )[..., None] + exp_attention_sinks = jnp.exp(attention_sinks - m_ref[...]) + l = l_ref[...] + exp_attention_sinks + l = broadcast_minor(l, acc.shape) + out = ( + lax.div(acc, l) + if q_dtype == jnp.float32 + else (acc * pl.reciprocal(l, approx=True)).astype(q_dtype) + ) + + # Wait for previous bo to be fully sent before storing new bo. + bo_sem_idx = sem_ids_ref[2] + sem_ids_ref[2] = lax.select(bo_sem_idx == 0, 1, 0) + wait_send_bo(bo_sem_idx) + + # Store output from acc to bo. + bo_x2_ref.at[bo_sem_idx].bitcast(jnp.int32).reshape( + bq_sz * num_q_heads_per_q_packing, + head_dim, + )[...] = pltpu.bitcast(out, jnp.int32) + + # Send cur bo + start_send_bo(seq_idx, bq_idx, bo_sem_idx) + + lax.fori_loop(0, num_bq, compute_with_bq, None, unroll=False) + + ### ------- Kernel start ------- ### + + @pl.when(seq_idx == start_seq_idx) + def prologue(): + start_fetch_bq(start_seq_idx, 0, 0) + start_fetch_swa(start_seq_idx, 0, 0) + + # Initialize bkv_x2_ref to avoid NaN issues from accessing uninitialized + # memory + bkv_zeros = jnp.zeros(bkv_x2_ref.shape[1:], bkv_x2_ref.dtype) + bkv_x2_ref[0] = bkv_zeros + start_fetch_bkv(start_seq_idx, 0, 0) + bkv_x2_ref[1] = bkv_zeros + + process() + + @pl.when(seq_idx == end_seq_idx - 1) + def epilogue(): + for i in range(2): + wait_send_bo(i) + + ### ------- Kernel end ------- ### + + +def prepare_q_inputs( + q: jax.Array, # [max_num_tokens, actual_num_q_heads, actual_head_dim], +): + _, actual_num_q_heads, actual_head_dim = q.shape + q_packing = get_dtype_packing(q.dtype) + num_q_heads = align_to(actual_num_q_heads, q_packing) + head_dim = align_to(actual_head_dim, 128) + q = jnp.pad( + q, + ( + (0, 0), + (0, num_q_heads - actual_num_q_heads), + (0, head_dim - actual_head_dim), + ), + constant_values=0, + ) + return q + + +def prepare_swa_inputs( + swa_accumution: jax.Array, # [max_num_tokens, num_q_heads, head_dim] + swa_l: jax.Array, # [max_num_tokens, num_q_heads] + swa_m: jax.Array, # [max_num_tokens, num_q_heads] +): + _, actual_num_q_heads, actual_head_dim = swa_accumution.shape + swa_packing = get_dtype_packing(swa_accumution.dtype) + num_q_heads = align_to(actual_num_q_heads, swa_packing) + head_dim = align_to(actual_head_dim, 128) + swa_accumution = jnp.pad( + swa_accumution, + ( + (0, 0), + (0, num_q_heads - actual_num_q_heads), + (0, head_dim - actual_head_dim), + ), + constant_values=0, + ) + num_l_heads = align_to(num_q_heads, 128) + swa_l = jnp.pad( + swa_l, + ( + (0, 0), + (0, num_l_heads - actual_num_q_heads), + ), + constant_values=0, + ) + swa_m = jnp.pad( + swa_m, + ( + (0, 0), + (0, num_l_heads - actual_num_q_heads), + ), + constant_values=0, + ) + return swa_accumution, swa_l, swa_m + + +def prepare_outputs( + out, # [max_num_tokens, num_q_heads, head_dim] + actual_num_q_heads: int, + actual_head_dim: int, +): + return out[:, :actual_num_q_heads, :actual_head_dim] + + +# TODO: support batching decode q tokens as performance optimization. + + +# Main Attention kernel for DeepSeek V4 HCA. +# Note that the compressed kv tokens of current batch (current forward pass) +# have been written to the `cache_kv` by the compressor module before calling +# this function, `kv_lens` reflects the length after compressed kv cache write. + + +# Quantize and dequantize into / from +# DSv4 fp8 format (448 fp8, 64 bf16, 7 fp8 scales, 7 e8m0 scale for 448 fp8) is +# quite expensive for TPU. +# For HCA, we just skip the quantization and dequantization, the kv cache stores +# bf16 data. +# HCA's compression ratio is 128, the overall size of HCA's compressed kv cache +# is very small compared to other caches such as CSA's compressed cache. Extra +# storage for storing KV cache in bf16 is trivial. +@functools.partial( + jax.jit, + static_argnames=( + "sm_scale", + "chunk_prefill_size", + "num_kv_pages_per_block", + "num_queries_per_block", + "vmem_limit_bytes", + ), +) +def mla_ragged_paged_attention( + q: jax.Array, # [max_num_tokens, actual_num_q_heads, head_dim] + cache_kv: jax.Array, # [total_num_pages, page_size * 2, 2, 128] uint8 + kv_lens: jax.Array, # i32[max_num_seqs] + kv_lens_to_attend: jax.Array, # i32[max_num_tokens] + page_indices: jax.Array, # i32[max_num_seqs * pages_per_seq] + cu_q_lens: jax.Array, # i32[max_num_seqs + 1] + distribution: jax.Array, # i32[3] + attention_sinks: jax.Array, # float32[actual_num_q_heads] + swa_accumution: jax.Array, # bf16[max_num_tokens, num_q_heads, head_dim] + swa_l: jax.Array, # float32[max_num_tokens, num_q_heads] + swa_m: jax.Array, # float32[max_num_tokens, num_q_heads] + *, + sm_scale: float = 1.0, + # Kernel optimization params. + chunk_prefill_size: int | None = None, + # Kernel tuning params for decode, prefill, and mixed cases. + # If passsed in as int, all cases are the same. + num_kv_pages_per_block: tuple[int, int, int] | int | None = None, + num_queries_per_block: tuple[int, int, int] | int | None = None, + vmem_limit_bytes: int = DEFAULT_VMEM_LIMIT_BYTES, +) -> jax.Array: + """MLA Ragged paged attention that supports mixed prefill and decode. + + Args: + q: concatenated all sequences' queries. + cache_kv: the current kv cache. + kv_lens: the length of each sequence in the kv cache. + kv_lens_to_attend: for each query token, the length of kv sequence to attend + to. The attend to length is <= kv_lens[seq_id] for that query token. + page_indices: flattened page indices look-up table by (seq_id, page_id). + cu_q_lens: the cumulative sum of the effective query lengths. Similar to + kv_lens, only the first num_seqs+1 values are valid. + distribution: (i, j, k) represents that sequences[0:i] are decode-only, + sequences[i:j] are chunked-prefill-only, and sequences[j:k] are mixed. The + k is also the total number of sequences. + sm_scale: the softmax scale which will be applied to the Q@K^T. + num_kv_pages_per_block: number of kv pages to be processed in one flash + attention block in the pallas kernel. This is a tuple of (decode, prefill, + mixed) cases. + num_queries_per_block: number of queries to be processed in one flash + attention block in the pallas kernel. This is a tuple of (decode, prefill, + mixed) cases. + vmem_limit_bytes: the vmem limit for the pallas kernel. + + Returns: + The output of attention. + """ + assert cache_kv.dtype == jnp.uint8 + + if num_kv_pages_per_block is None or num_queries_per_block is None: + raise ValueError( + "num_kv_pages_per_block and num_queries_per_block must be specified." + ) + if isinstance(num_kv_pages_per_block, int): + num_kv_pages_per_blocks = [num_kv_pages_per_block for _ in range(3)] + else: + num_kv_pages_per_blocks = num_kv_pages_per_block + + if isinstance(num_queries_per_block, int): + num_queries_per_blocks = [num_queries_per_block for _ in range(3)] + else: + num_queries_per_blocks = num_queries_per_block + + _, actual_num_q_heads, actual_head_dim = q.shape + + q = prepare_q_inputs(q) # [max_num_tokens, num_q_heads, head_dim] + head_dim = q.shape[-1] + attention_sinks = jnp.pad( + attention_sinks, + (0, q.shape[1] - actual_num_q_heads), + constant_values=jnp.finfo(attention_sinks.dtype).min, + ) + assert swa_accumution.dtype == q.dtype + swa_accumution, swa_l, swa_m = prepare_swa_inputs( + swa_accumution, swa_l, swa_m + ) + + _, num_slots_per_page, kv_packing, lkv_dim = cache_kv.shape + num_slots_per_tokens = 2 + assert num_slots_per_page % num_slots_per_tokens == 0 + page_size = num_slots_per_page // num_slots_per_tokens + _, num_q_heads, _ = q.shape + max_num_seqs = cu_q_lens.shape[0] - 1 + num_page_indices = page_indices.shape[0] + assert num_page_indices % max_num_seqs == 0 + + def run_mla_kernel( + q: jax.Array, # [max_num_tokens, num_q_heads, head_dim] + cache_kv: jax.Array, # [total_num_pages, page_size * 2, kv_packing, lkv_dim] + kv_lens: jax.Array, # i32[max_num_seqs] + kv_lens_to_attend: jax.Array | None, # i32[max_num_tokens] + page_indices: jax.Array, # i32[max_num_seqs * pages_per_seq] + cu_q_lens: jax.Array, # i32[max_num_seqs + 1] + attention_sinks: jax.Array, # float32[num_q_heads] + swa_accumution: jax.Array, # bf16[max_num_tokens, num_q_heads, head_dim] + swa_l: jax.Array, # float32[max_num_tokens, num_l_heads] + swa_m: jax.Array, # float32[max_num_tokens, num_l_heads] + start_seq_idx: jax.Array, # i32 + end_seq_idx: jax.Array, # i32 + static_q_len: int | None, + num_kv_pages_per_block: int, + num_queries_per_block: int, + case: MlaCase = MlaCase.MIXED, + ): + + bkv_p = num_kv_pages_per_block + if static_q_len is not None: + bq_sz = min(num_queries_per_block, static_q_len) + else: + bq_sz = num_queries_per_block + + grid = (end_seq_idx - start_seq_idx,) + in_specs = [ + pl.BlockSpec(memory_space=pltpu.VMEM), # attention_sinks + pl.BlockSpec(memory_space=pltpu.HBM), # q + pl.BlockSpec(memory_space=pltpu.HBM), # cache_kv + pl.BlockSpec(memory_space=pltpu.HBM), # swa_accumution + pl.BlockSpec(memory_space=pltpu.HBM), # swa_l + pl.BlockSpec(memory_space=pltpu.HBM), # swa_m + ] + + out_specs = pl.BlockSpec(memory_space=pltpu.HBM) # o + + # last 2 dimension mapped to one tokens's kv + assert ( + num_slots_per_tokens * kv_packing * lkv_dim + == head_dim * get_dtype_bitwidth(q.dtype) // 8 + ) + bkv_double_buf = pltpu.VMEM( + (2, bkv_p * page_size, num_slots_per_tokens * kv_packing, lkv_dim), + cache_kv.dtype, + ) + + bq_double_bufq = pltpu.VMEM( + (2, bq_sz, num_q_heads, head_dim), + q.dtype, + ) + + bo_double_buf = bq_double_bufq + + num_l_heads = align_to(num_q_heads, 128) + bl_double_buf = pltpu.VMEM( + (2, bq_sz, num_l_heads), + jnp.float32, + ) + bm_double_buf = bl_double_buf + + swa_acc_double_buf = pltpu.VMEM( + (2, bq_sz, num_q_heads, head_dim), + q.dtype, + ) + + l_scratch = pltpu.VMEM( + (bq_sz * num_q_heads, 128), + jnp.float32, + ) + m_scratch = l_scratch + + acc_scratch = pltpu.VMEM( + (bq_sz * num_q_heads, head_dim), + jnp.float32, + ) + + scratch_shapes = [ + bkv_double_buf, + bq_double_bufq, + bo_double_buf, # Double buffering for output block. + bl_double_buf, # Double buffering for l output. + bm_double_buf, # Double buffering for m output. + swa_acc_double_buf, # Buffer for swa_accumution. + # Semaphores for double buffering of bkv, bq, bo, swa_acc, swa_l, swa_m, topk. + pltpu.SemaphoreType.DMA((6, 2)), + # Intermediate buffers per kv head for flash attention. + l_scratch, + m_scratch, + acc_scratch, + ] + + scalar_prefetches = ( + kv_lens, + kv_lens_to_attend, + page_indices, + cu_q_lens, + jnp.array([start_seq_idx, end_seq_idx], jnp.int32), + # (bq_sem_idx, bkv_sem_idx, bo_sem_idx) + jnp.zeros((3,), jnp.int32), + # (bo_sem_0_seq_idx, bo_sem_1_seq_idx, bo_sem_0_bo_idx, bo_sem_1_bo_idx) + jnp.full((4,), -1, jnp.int32), + ) + + scope_name = f"MLA-{case.symbol}-bq_{bq_sz}-bkvp_{bkv_p}-p_{page_size}" + kernel = jax.named_scope(scope_name)( + pl.pallas_call( + functools.partial( + _mla_ragged_paged_attention_kernel, + sm_scale=sm_scale, + static_q_len=static_q_len, + bq_sz=bq_sz, + bkv_p=bkv_p, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=len(scalar_prefetches), + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + scratch_shapes=scratch_shapes, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary",), + vmem_limit_bytes=vmem_limit_bytes, + disable_bounds_checks=True, + ), + out_shape=jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype), + input_output_aliases={ + 8: 0, # Alias output activation with q + }, + name=scope_name, + ) + ) + return kernel( + *scalar_prefetches, + attention_sinks, + q, + cache_kv, + swa_accumution, + swa_l, + swa_m, + ) + + # Decode-only + q = run_mla_kernel( + q, + cache_kv, + kv_lens, + kv_lens_to_attend, + page_indices, + cu_q_lens, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + num_kv_pages_per_block=num_kv_pages_per_blocks[0], + num_queries_per_block=num_queries_per_blocks[0], + start_seq_idx=jnp.array(0), + end_seq_idx=distribution[0], + static_q_len=1, + case=MlaCase.DECODE, + ) + + if chunk_prefill_size is not None: + # Handle prefill where the query length is fixed per sequence. + q = run_mla_kernel( + q, + cache_kv, + kv_lens, + kv_lens_to_attend, + page_indices, + cu_q_lens, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + num_kv_pages_per_block=num_kv_pages_per_blocks[1], + num_queries_per_block=num_queries_per_blocks[1], + start_seq_idx=distribution[0], + end_seq_idx=distribution[1], + static_q_len=chunk_prefill_size, + case=MlaCase.PREFILL, + ) + + # Handle mixed case where the query length per sequence is variable. + q = run_mla_kernel( + q, + cache_kv, + kv_lens, + kv_lens_to_attend, + page_indices, + cu_q_lens, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + num_kv_pages_per_block=num_kv_pages_per_blocks[2], + num_queries_per_block=num_queries_per_blocks[2], + start_seq_idx=distribution[1], + end_seq_idx=distribution[2], + static_q_len=None, + case=MlaCase.MIXED, + ) + output = prepare_outputs( + q, actual_num_q_heads, actual_head_dim + ) # [max_num_tokens, actual_num_q_heads, actual_head_dim] + + return output + + +def workload( + q, cache_kv, kv_lens, kv_lens_to_attend, page_indices, cu_q_lens, + distribution, attention_sinks, swa_accumution, swa_l, swa_m, +): + num_tokens = q.shape[0] + batch_size = cu_q_lens.shape[0] - 1 + q_len = num_tokens // batch_size + page_size = cache_kv.shape[1] // 2 + + chunk_prefill_size = None if q_len == 1 else q_len + bq = 1 if q_len == 1 else 32 + bkv_pages = 4 if (page_size == 2048 and q_len > 1) else 16 + + sm_scale = float(512 ** -0.5) + vmem_limit_bytes = 100 * 1024 * 1024 + + return mla_ragged_paged_attention( + q, cache_kv, kv_lens, kv_lens_to_attend, page_indices, cu_q_lens, + distribution, attention_sinks, swa_accumution, swa_l, swa_m, + sm_scale=sm_scale, + chunk_prefill_size=chunk_prefill_size, + num_kv_pages_per_block=bkv_pages, + num_queries_per_block=bq, + vmem_limit_bytes=vmem_limit_bytes, + ) + +def benchmark(num_warmup=5, num_iters=100): + """Benchmark and return results dict.""" + inputs_list = create_inputs() + fn = jax.jit(workload) + times_list = [] + std_ms_list = [] + output_shape_list = [] + for inputs in inputs_list: + for _ in range(num_warmup): + out = fn(*inputs) + jax.block_until_ready(out) + times = [] + for _ in range(num_iters): + t0 = time.perf_counter() + out = fn(*inputs) + jax.block_until_ready(out) + times.append(time.perf_counter() - t0) + times = np.array(times) * 1000 + avg = float(np.mean(times)) + times_list.append(round(avg, 4)) + std_ms_list.append(round(float(np.std(times)), 4)) + if hasattr(out, 'shape'): + out_shape = list(out.shape) + elif isinstance(out, (tuple, list)): + out_shape = [list(x.shape) if hasattr(x, 'shape') else [] for x in out] + else: + out_shape = [] + output_shape_list.append(out_shape) + return { + 'name': CONFIG['name'], + 'model': CONFIG['model'], + 'operator': CONFIG['operator'], + 'config': {k: v for k, v in CONFIG.items() if k not in ('name', 'model', 'operator', 'atol', 'rtol')}, + 'time_ms': times_list, + 'std_ms': std_ms_list, + 'output_shape': output_shape_list, + 'status': 'success', + } + +if __name__ == '__main__': + import json + print(json.dumps(benchmark())) diff --git a/JAXBench/benchmark/level2/53p_DeepSeek_V4_SWA/baseline.py b/JAXBench/benchmark/level2/53p_DeepSeek_V4_SWA/baseline.py new file mode 100644 index 0000000..66af65f --- /dev/null +++ b/JAXBench/benchmark/level2/53p_DeepSeek_V4_SWA/baseline.py @@ -0,0 +1,973 @@ +# Imports +import numpy as np +import time +import functools +from enum import Enum +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 + +# Initialization +CONFIG = { + 'name': 'deepseek_swa', + 'model': 'deepseek', + 'operator': 'sliding_window_attention', + 'configs': { + "decode_8k_1k": (256, 1, 9216, 1024, 8, 512, 128), + "prefill_first_chunk": (1, 1024, 1024, 1024, 8, 512, 128), + "prefill_last_chunk": (1, 1024, 8192, 1024, 8, 512, 128), + "medium_batch_decode": (128, 1, 9216, 1024, 8, 512, 128), + "prefill_256": (1, 256, 1024, 1024, 8, 512, 128), + "prefill_512": (1, 512, 4096, 1024, 8, 512, 128), + "alt_page_decode": (128, 1, 9216, 256, 8, 512, 128), + "alt_page_prefill": (1, 256, 1024, 256, 8, 512, 128), + }, + 'atol': 0.01, + 'rtol': 0.01, +} + +def cdiv_val(a, b): + assert b != 0 + return (a + b - 1) // b + + +def create_inputs(): + configs = CONFIG['configs'] + if isinstance(configs, dict): + configs = [(name, *params) for name, params in configs.items()] + inputs = [] + key = jax.random.PRNGKey(42) + + for ( + name, + batch_size, + q_len, + kv_len_val, + page_size, + num_q_heads, + head_dim, + sliding_window, + ) in configs: + key, k1, k2, k3 = jax.random.split(key, 4) + + num_tokens = batch_size * q_len + kv_lens = jnp.full((batch_size,), kv_len_val, dtype=jnp.int32) + cu_q_lens = jnp.arange(0, num_tokens + 1, q_len, dtype=jnp.int32) + + pages_per_seq = cdiv_val(kv_len_val, page_size) + 2 + total_pages = batch_size * pages_per_seq + page_indices = jnp.arange(total_pages, dtype=jnp.int32) + + q = jax.random.normal( + k1, (num_tokens, num_q_heads, head_dim), dtype=jnp.bfloat16 + ) + new_kv = jax.random.normal(k2, (num_tokens, head_dim), dtype=jnp.bfloat16) + attention_sinks = jax.random.normal(k3, (num_q_heads,), dtype=jnp.float32) + + num_decode_seqs = batch_size if q_len == 1 else 0 + distribution = jnp.array( + [num_decode_seqs, num_decode_seqs, batch_size], dtype=jnp.int32 + ) + kernel_cache = jnp.zeros((total_pages, page_size * 2, 4, 128), dtype=jnp.uint8) + + args = [ + q, + new_kv, + kernel_cache, + kv_lens, + page_indices, + cu_q_lens, + distribution, + attention_sinks, + ] + inputs.append(args) + + return inputs + +# Computation +def cdiv(a, b): + assert b != 0 + return (a + b - 1) // b + +def align_to(x, a): + return cdiv(x, a) * a + +def get_dtype_bitwidth(dtype): + return jax.dtypes.itemsize_bits(dtype) + +def get_dtype_packing(dtype): + bits = get_dtype_bitwidth(dtype) + return 32 // bits + +class MlaCase(Enum): + DECODE = 0 + PREFILL = 1 + MIXED = 2 + + @property + def symbol(self): + return { + MlaCase.DECODE: "d", + MlaCase.PREFILL: "p", + MlaCase.MIXED: "m", + }[self] + +def _mla_sliding_window_ragged_paged_attention_kernel( + kv_lens_ref, + page_indices_ref, + cu_q_lens_ref, + start_end_seq_idx_ref, + sem_ids_ref, + bo_ids_ref, + bkv_update_ids_ref, + attention_sinks_ref, + q_hbm_ref, + new_kv_hbm_ref, + cache_kv_hbm_ref, + in_output_hbm_ref, + in_l_hbm_ref, + in_m_hbm_ref, + o_hbm_ref, + updated_cache_kv_hbm_ref, + l_hbm_ref, + m_hbm_ref, + bkv_x2_ref, + bq_x2_ref, + bo_x2_ref, + bl_x2_ref, + bm_x2_ref, + sems, + l_ref, + m_ref, + acc_ref, + *, + static_q_len: int, + sm_scale: float, + sliding_window: int, + logical_page_size: int, + unnormalized_output: bool, + q_compute_block_size: int | None, + bkv_p, + bq_sz, +): + assert q_hbm_ref.shape == o_hbm_ref.shape + assert sliding_window > 0 + + _, num_q_heads, head_dim = q_hbm_ref.shape + q_packing = get_dtype_packing(q_hbm_ref.dtype) + assert num_q_heads % q_packing == 0 + num_q_heads_per_q_packing = num_q_heads // q_packing + + total_num_pages, physical_page_size_per_kv_packing, kv_packing, lkv_dim = cache_kv_hbm_ref.shape + q_dtype = q_hbm_ref.dtype + assert o_hbm_ref.dtype == q_dtype + assert head_dim % 128 == 0 + token_bytes = head_dim * get_dtype_bitwidth(q_dtype) // 8 + slot_bytes = kv_packing * lkv_dim + assert token_bytes % slot_bytes == 0 + slots_per_token = token_bytes // slot_bytes + phys_tokens_per_page = physical_page_size_per_kv_packing // slots_per_token + + max_num_seqs = kv_lens_ref.shape[0] + num_page_indices = page_indices_ref.shape[0] + assert num_page_indices % max_num_seqs == 0 + pages_per_seq = num_page_indices // max_num_seqs + + bkv_sz = bkv_p * logical_page_size + page_size = logical_page_size + + max_bkv_blocks = cdiv(bq_sz + sliding_window - 1, bkv_sz) + single_bkv_block = max_bkv_blocks == 1 + + start_seq_idx = start_end_seq_idx_ref[0] + end_seq_idx = start_end_seq_idx_ref[1] + seq_idx = pl.program_id(0) + start_seq_idx + q_start = cu_q_lens_ref[seq_idx] + q_end = cu_q_lens_ref[seq_idx + 1] + q_len = q_end - q_start + kv_len = kv_lens_ref[seq_idx] + + def flash_attention( + q, + kv, + *, + bq_idx, + bkv_idx, + start_offset, + ): + assert len(q.shape) == 2 + assert len(kv.shape) == 2 + assert q.shape[0] % num_q_heads == 0 + assert q.shape[1] == head_dim + assert kv.shape == (bkv_sz, head_dim) + n = q.shape[0] // num_q_heads + + if q_compute_block_size is None: + chunk_sz = n + else: + chunk_sz = q_compute_block_size if n % q_compute_block_size == 0 else n + num_chunks = n // chunk_sz + + def load_with_init(ref, init_val): + if single_bkv_block: + return jnp.full_like(ref, init_val) + else: + return jnp.where(bkv_idx == 0, jnp.full_like(ref, init_val), ref[...]) + + k_span = start_offset + bkv_idx * bkv_sz + lax.broadcasted_iota(jnp.int32, (1, bkv_sz), 1) + k_pos = start_offset + bkv_idx * bkv_sz + lax.broadcasted_iota(jnp.int32, (bkv_sz, 1), 0) + kv = jnp.where(k_pos < kv_len, kv, 0.0) + + chunk_size = chunk_sz * num_q_heads + for c in range(num_chunks): + start_row = c * chunk_size + qc = q[start_row : start_row + chunk_size] + cl_ref = l_ref.at[start_row : start_row + chunk_size] + cm_ref = m_ref.at[start_row : start_row + chunk_size] + cacc_ref = acc_ref.at[start_row : start_row + chunk_size] + + s = jnp.einsum("nd,md->nm", qc, kv, preferred_element_type=jnp.float32) + s *= sm_scale + + q_span = kv_len - q_len + bq_idx * bq_sz + (start_row + lax.broadcasted_iota(jnp.int32, (chunk_size, 1), 0)) // num_q_heads + keep = (q_span - k_span).astype(jnp.uint32) < jnp.uint32(sliding_window) + + s = jnp.where(keep, s, jnp.finfo(s.dtype).min) + s_rowmax = jnp.max(s, axis=1, keepdims=True) + m_prev = load_with_init(cm_ref, jnp.finfo(jnp.float32).min) + m_curr = jnp.maximum(m_prev, s_rowmax) + cm_ref[...] = m_curr + p = jnp.exp(s - broadcast_minor(m_curr, s.shape)) + p = jnp.where(keep, p, 0.0) + + pv = jnp.einsum("nm,md->nd", p, kv, preferred_element_type=jnp.float32) + + p_rowsum = jnp.sum(p, axis=1, keepdims=True) + exp_m_diff = jnp.exp(m_prev - m_curr) + l_prev = load_with_init(cl_ref, 0.0) + l_curr = exp_m_diff * l_prev + p_rowsum + cl_ref[...] = l_curr + o_prev = load_with_init(cacc_ref, 0.0) + o_curr = broadcast_minor(exp_m_diff, o_prev.shape) * o_prev + pv + cacc_ref[...] = o_curr + + def _async_copy(src, dst, sem, wait): + cp = pltpu.make_async_copy(src, dst, sem) + if wait: + cp.wait() + else: + cp.start() + + def _get_kv_len(seq_idx): + return jnp.where(seq_idx < end_seq_idx, kv_lens_ref[seq_idx], 0) + + def _get_q_len(seq_idx): + return jnp.where(seq_idx < end_seq_idx, cu_q_lens_ref[seq_idx + 1] - cu_q_lens_ref[seq_idx], 0) + + def _start_offset(seq_idx, bq_idx): + return jnp.maximum(_get_kv_len(seq_idx) - _get_q_len(seq_idx) + bq_idx * bq_sz - sliding_window + 1, 0) + + def _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, start_offset, *, wait=False): + sem = sems.at[0, bkv_sem_idx] + bkv_vmem_ref = bkv_x2_ref.at[bkv_sem_idx] + + reshaped_cache_hbm_ref = cache_kv_hbm_ref.reshape(total_num_pages * phys_tokens_per_page, slots_per_token * kv_packing, lkv_dim) + + kv_len = kv_lens_ref[seq_idx] + kv_len_start = start_offset + bkv_idx * bkv_sz + kv_p_start = kv_len_start // page_size + page_off = kv_len_start - kv_p_start * page_size + + q_start = cu_q_lens_ref[seq_idx] + q_end = cu_q_lens_ref[seq_idx + 1] + q_len = q_end - q_start + + kv_left = jnp.maximum(kv_len - kv_len_start, 0) + kv_left_frm_cache = jnp.maximum(kv_left - q_len, 0) + kv_left_frm_new = kv_left - kv_left_frm_cache + + bkv_sz_frm_cache = jnp.minimum(kv_left_frm_cache, bkv_sz) + bkv_sz_frm_new = jnp.minimum(bkv_sz - bkv_sz_frm_cache, kv_left_frm_new) + page_indices_offset = seq_idx * pages_per_seq + kv_p_start + + new_kv_len_start = q_end - kv_left_frm_new + dma_bkv_sz = bkv_sz_frm_cache + bkv_sz_frm_new + + if not wait: + wait_update_kv_cache(bkv_sem_idx) + + for i in range(bkv_p + 1): + if i == 0: + in_page_off = page_off + vmem_off = jnp.int32(0) + avail = page_size - page_off + else: + in_page_off = jnp.int32(0) + vmem_off = i * page_size - page_off + avail = jnp.int32(page_size) + sz = jnp.clip(bkv_sz_frm_cache - vmem_off, 0, avail) + page_idx = jnp.minimum(page_indices_offset + i, num_page_indices - 1) + _async_copy( + reshaped_cache_hbm_ref.at[pl.ds(page_indices_ref[page_idx] * phys_tokens_per_page + in_page_off, sz)], + bkv_vmem_ref.at[pl.ds(vmem_off, sz)], + sem, + wait, + ) + + _async_copy( + new_kv_hbm_ref.at[pl.ds(new_kv_len_start, bkv_sz_frm_new)], + bkv_vmem_ref.at[pl.ds(bkv_sz_frm_cache, bkv_sz_frm_new)], + sem, + wait, + ) + + else: + dst_kv = bkv_vmem_ref.at[pl.ds(0, dma_bkv_sz)] + _async_copy(src=dst_kv, dst=dst_kv, sem=sem, wait=True) + + return kv_len_start + bkv_sz_frm_cache, bkv_sz_frm_new, bkv_sz_frm_cache + + def _update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz, *, in_vmem_start=0, wait=False): + sem = sems.at[3, bkv_sem_idx] + bkv_vmem_ref = bkv_x2_ref.at[bkv_sem_idx] + + update_kv_packing_iters = update_sz + + reshaped_cache_kv_hbm_ref = updated_cache_kv_hbm_ref.reshape(total_num_pages * phys_tokens_per_page, slots_per_token * kv_packing, lkv_dim) + + if not wait: + kv_p_start = offset // page_size + kv_p_end = cdiv(offset + update_sz, page_size) + start_word_in_page = offset % page_size + start_word_in_vmem = in_vmem_start + words_to_transfer = update_kv_packing_iters + page_indices_offset = seq_idx * pages_per_seq + kv_p_start + + def loop_body(i, states): + curr_word_in_page, words_to_transfer, curr_word_in_vmem = states + sz = jnp.minimum(page_size - curr_word_in_page, words_to_transfer) + page_idx = page_indices_ref[page_indices_offset + i] + + _async_copy( + bkv_vmem_ref.at[pl.ds(curr_word_in_vmem, sz)], + reshaped_cache_kv_hbm_ref.at[pl.ds(page_idx * phys_tokens_per_page + curr_word_in_page, sz)], + sem, + wait=False, + ) + return 0, words_to_transfer - sz, curr_word_in_vmem + sz + + lax.fori_loop( + 0, + kv_p_end - kv_p_start, + loop_body, + (start_word_in_page, words_to_transfer, start_word_in_vmem), + unroll=False, + ) + else: + dma_sz_words = update_kv_packing_iters + dst_kv = bkv_vmem_ref.at[pl.ds(0, dma_sz_words)] + _async_copy(src=dst_kv, dst=dst_kv, sem=sem, wait=True) + + def _fetch_bq(seq_idx, bq_idx, bq_sem_idx, *, wait=False): + sem = sems.at[1, bq_sem_idx] + bq_vmem_ref = bq_x2_ref.at[bq_sem_idx] + + q_len_start = cu_q_lens_ref[seq_idx] + bq_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy(q_hbm_ref.at[pl.ds(q_len_start, sz)], bq_vmem_ref.at[pl.ds(0, sz)], sem, wait) + + def _send_bo(seq_idx, bo_idx, bo_sem_idx, *, wait=False): + sem = sems.at[2, bo_sem_idx] + vmem_ref = bo_x2_ref.at[bo_sem_idx] + q_len_start = cu_q_lens_ref[seq_idx] + bo_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy(vmem_ref.at[pl.ds(0, sz)], o_hbm_ref.at[pl.ds(q_len_start, sz)], sem, wait) + + def _send_l(seq_idx, bo_idx, bo_sem_idx, *, wait=False): + sem = sems.at[4, bo_sem_idx] + vmem_ref = bl_x2_ref.at[bo_sem_idx] + q_len_start = cu_q_lens_ref[seq_idx] + bo_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy(vmem_ref.at[pl.ds(0, sz)], l_hbm_ref.at[pl.ds(q_len_start, sz)], sem, wait) + + def _send_m(seq_idx, bo_idx, bo_sem_idx, *, wait=False): + sem = sems.at[5, bo_sem_idx] + vmem_ref = bm_x2_ref.at[bo_sem_idx] + q_len_start = cu_q_lens_ref[seq_idx] + bo_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy(vmem_ref.at[pl.ds(0, sz)], m_hbm_ref.at[pl.ds(q_len_start, sz)], sem, wait) + + def start_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, start_offset): + return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, start_offset) + + def wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, start_offset): + return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, start_offset, wait=True) + + def start_fetch_bq(seq_idx, bq_idx, bq_sem_idx): + return _fetch_bq(seq_idx, bq_idx, bq_sem_idx) + + def wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx): + return _fetch_bq(seq_idx, bq_idx, bq_sem_idx, wait=True) + + def start_send_bo(seq_idx, bo_idx, bo_sem_idx): + bo_ids_ref[bo_sem_idx] = seq_idx + bo_ids_ref[bo_sem_idx + 2] = bo_idx + _send_bo(seq_idx, bo_idx, bo_sem_idx) + _send_l(seq_idx, bo_idx, bo_sem_idx) + _send_m(seq_idx, bo_idx, bo_sem_idx) + + def wait_send_bo(bo_sem_idx): + old_seq_idx = bo_ids_ref[bo_sem_idx] + old_bo_idx = bo_ids_ref[bo_sem_idx + 2] + + @pl.when(jnp.logical_and(0 <= old_seq_idx, old_seq_idx <= seq_idx)) + def _(): + _send_bo(old_seq_idx, old_bo_idx, bo_sem_idx, wait=True) + _send_l(old_seq_idx, old_bo_idx, bo_sem_idx, wait=True) + _send_m(old_seq_idx, old_bo_idx, bo_sem_idx, wait=True) + + def start_update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz, vmem_start): + bkv_update_ids_ref[bkv_sem_idx] = seq_idx + bkv_update_ids_ref[bkv_sem_idx + 2] = offset + bkv_update_ids_ref[bkv_sem_idx + 4] = update_sz + _update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz, in_vmem_start=vmem_start) + + def wait_update_kv_cache(bkv_sem_idx): + update_sz = bkv_update_ids_ref[bkv_sem_idx + 4] + + @pl.when(update_sz > 0) + def _(): + seq_idx = bkv_update_ids_ref[bkv_sem_idx] + offset = bkv_update_ids_ref[bkv_sem_idx + 2] + bkv_update_ids_ref[bkv_sem_idx + 4] = 0 + _update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz, wait=True) + + def load_bq(bq_sem_idx): + q_ref = bq_x2_ref.bitcast(jnp.uint32).at[bq_sem_idx].reshape(bq_sz * num_q_heads_per_q_packing, head_dim) + q = pltpu.bitcast(q_ref[: bq_sz * num_q_heads_per_q_packing], q_dtype).reshape(bq_sz * num_q_heads, head_dim) + return q + + def load_bkv(bkv_sem_idx, bkv_idx, start_offset): + bkv_u8 = bkv_x2_ref.at[bkv_sem_idx][...] + bkv = pltpu.bitcast(bkv_u8, jnp.bfloat16).reshape(bkv_sz, head_dim) + return bkv + + def broadcast_minor(src, shape): + if src.shape == shape: + return src + assert src.shape[:-1] == shape[:-1] + assert src.shape[-1] % 128 == 0 + target_minor = align_to(shape[-1], src.shape[-1]) + return jnp.concatenate([src for _ in range(target_minor // src.shape[-1])], axis=-1)[..., : shape[-1]] + + def process(): + if static_q_len is None: + num_bq = jnp.maximum(1, cdiv(q_len, bq_sz)) + else: + num_bq = jnp.maximum(1, cdiv(static_q_len, bq_sz)) + + def get_next_bq_ids(seq_idx, bq_idx, bq_sem_idx): + next_bq_idx = bq_idx + 1 + is_last_bq = next_bq_idx == num_bq + next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx) + next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx) + next_bq_sem_idx = lax.select(bq_sem_idx == 0, 1, 0) + return next_seq_idx, next_bq_idx, next_bq_sem_idx + + def compute_with_bq(bq_idx, _): + cur_start_offset = _start_offset(seq_idx, bq_idx) + start_bkv_idx = 0 + if single_bkv_block: + end_bkv_idx = 1 + else: + end_bkv_idx = jnp.maximum(cdiv(jnp.minimum(kv_len - q_len + (bq_idx + 1) * bq_sz, kv_len) - cur_start_offset, bkv_sz), 1) + + def get_next_bkv_ids(seq_idx, bq_idx, bkv_idx, bkv_sem_idx): + next_bkv_idx = bkv_idx + 1 + is_last_bkv = next_bkv_idx == end_bkv_idx + next_bq_idx = lax.select(is_last_bkv, bq_idx + 1, bq_idx) + is_last_bq = next_bq_idx == num_bq + next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx) + next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx) + next_bkv_idx = lax.select(is_last_bkv, 0, next_bkv_idx) + next_bkv_sem_idx = lax.select(bkv_sem_idx == 0, 1, 0) + return next_seq_idx, next_bq_idx, next_bkv_idx, next_bkv_sem_idx + + bq_sem_idx = sem_ids_ref[0] + next_seq_idx, next_bq_idx, next_bq_sem_idx = get_next_bq_ids(seq_idx, bq_idx, bq_sem_idx) + + @pl.when(next_seq_idx < end_seq_idx) + def prefetch_next_bq(): + sem_ids_ref[0] = next_bq_sem_idx + start_fetch_bq(next_seq_idx, next_bq_idx, next_bq_sem_idx) + + def compute_with_bkv(bkv_idx, _): + bkv_sem_idx = sem_ids_ref[1] + next_seq_idx, next_bq_idx, next_bkv_idx, next_bkv_sem_idx = get_next_bkv_ids(seq_idx, bq_idx, bkv_idx, bkv_sem_idx) + + @pl.when(next_seq_idx < end_seq_idx) + def prefetch_next_bkv(): + sem_ids_ref[1] = next_bkv_sem_idx + next_start_offset = _start_offset(next_seq_idx, next_bq_idx) + start_fetch_bkv(next_seq_idx, next_bkv_idx, next_bkv_sem_idx, next_start_offset) + + offset, update_sz, vmem_start = wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, cur_start_offset) + + @pl.when(update_sz > 0) + def update_cur_bkv_to_cache(): + start_update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz, vmem_start) + + bkv = load_bkv(bkv_sem_idx, bkv_idx, cur_start_offset) + bq = load_bq(bq_sem_idx) + + flash_attention(bq, bkv, bq_idx=bq_idx, bkv_idx=bkv_idx, start_offset=cur_start_offset) + + wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx) + if single_bkv_block: + compute_with_bkv(0, None) + else: + lax.fori_loop(start_bkv_idx, end_bkv_idx, compute_with_bkv, None, unroll=False) + + acc = acc_ref[...] + + if unnormalized_output: + l = broadcast_minor(l_ref[...], acc.shape) + out = acc.astype(q_dtype) + else: + attention_sinks = jnp.concat([attention_sinks_ref[...] for _ in range(bq_sz)])[..., None] + exp_attention_sinks = jnp.exp(attention_sinks - m_ref[...]) + l = l_ref[...] + exp_attention_sinks + l = broadcast_minor(l, acc.shape) + out = lax.div(acc, l) if q_dtype == jnp.float32 else (acc * pl.reciprocal(l, approx=True)).astype(q_dtype) + + bo_sem_idx = sem_ids_ref[2] + sem_ids_ref[2] = lax.select(bo_sem_idx == 0, 1, 0) + wait_send_bo(bo_sem_idx) + + bo_x2_ref.at[bo_sem_idx].bitcast(jnp.int32).reshape(bq_sz * num_q_heads_per_q_packing, head_dim)[...] = pltpu.bitcast(out, jnp.int32) + bl_x2_ref.at[bo_sem_idx][:bq_sz, :num_q_heads] = l_ref[..., 0].reshape(bq_sz, num_q_heads) + bm_x2_ref.at[bo_sem_idx][:bq_sz, :num_q_heads] = m_ref[..., 0].reshape(bq_sz, num_q_heads) + + start_send_bo(seq_idx, bq_idx, bo_sem_idx) + + lax.fori_loop(0, num_bq, compute_with_bq, None, unroll=False) + + @pl.when(seq_idx == start_seq_idx) + def prologue(): + start_fetch_bq(start_seq_idx, 0, 0) + start_fetch_bkv(start_seq_idx, 0, 0, _start_offset(start_seq_idx, 0)) + + process() + + @pl.when(seq_idx == end_seq_idx - 1) + def epilogue(): + for i in range(2): + wait_send_bo(i) + wait_update_kv_cache(i) + +def prepare_q_inputs(q: jax.Array): + max_num_tokens, actual_num_q_heads, actual_head_dim = q.shape + q_packing = get_dtype_packing(q.dtype) + num_q_heads = align_to(actual_num_q_heads, q_packing) + head_dim = align_to(actual_head_dim, 128) + q = jnp.pad( + q, + ((0, 0), (0, num_q_heads - actual_num_q_heads), (0, head_dim - actual_head_dim)), + constant_values=0, + ) + return q + +def prepare_kv_inputs(kv: jax.Array): + assert kv.dtype == jnp.bfloat16 + tokens, head_dim = kv.shape + assert head_dim % 128 == 0 + kv_u16 = jax.lax.bitcast_convert_type(kv, jnp.uint16) + kv_hi = ((kv_u16 >> 8) & 0xFF).astype(jnp.uint8) + kv_lo = (kv_u16 & 0xFF).astype(jnp.uint8) + nb = head_dim // 128 + kv_hi = kv_hi.reshape(tokens, nb, 128) + kv_lo = kv_lo.reshape(tokens, nb, 128) + interleaved = jnp.stack([kv_lo, kv_hi], axis=2) + return interleaved.reshape(tokens, head_dim * 2) + +def prepare_outputs(out, actual_num_q_heads: int, actual_head_dim: int): + return out[:, :actual_num_q_heads, :actual_head_dim] + +@functools.partial( + jax.jit, + static_argnames=( + "sm_scale", + "sliding_window", + "chunk_prefill_size", + "num_kv_pages_per_block", + "num_queries_per_block", + "vmem_limit_bytes", + "logical_page_size", + "unnormalized_output", + "q_compute_block_size", + ), + donate_argnames=("cache_kv",), +) +def mla_sliding_window_ragged_paged_attention( + q: jax.Array, + new_kv: jax.Array, + cache_kv: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + distribution: jax.Array, + attention_sinks: jax.Array, + *, + sm_scale: float = 1.0, + sliding_window: int, + logical_page_size: int, + chunk_prefill_size: int | None = None, + num_kv_pages_per_block: tuple[int, int, int] | int | None = None, + num_queries_per_block: tuple[int, int, int] | int | None = None, + q_compute_block_size: int | None = None, + vmem_limit_bytes: int = 100 * 1024 * 1024, + unnormalized_output: bool = False, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]: + if num_kv_pages_per_block is None or num_queries_per_block is None: + raise ValueError("num_kv_pages_per_block and num_queries_per_block must be specified.") + + if isinstance(num_kv_pages_per_block, int): + num_kv_pages_per_blocks = [num_kv_pages_per_block for _ in range(3)] + else: + num_kv_pages_per_blocks = num_kv_pages_per_block + + if isinstance(num_queries_per_block, int): + num_queries_per_blocks = [num_queries_per_block for _ in range(3)] + else: + num_queries_per_blocks = num_queries_per_block + + _, actual_num_q_heads, actual_head_dim = q.shape + + q = prepare_q_inputs(q) + attention_sinks = jnp.pad( + attention_sinks, + (0, q.shape[1] - actual_num_q_heads), + constant_values=jnp.finfo(attention_sinks.dtype).min, + ) + assert new_kv.dtype == jnp.bfloat16 + assert cache_kv.dtype == jnp.uint8 + head_dim = q.shape[-1] + _, physical_page_size_per_kv_packing, kv_packing, lkv_dim = cache_kv.shape + + slot_bytes = kv_packing * lkv_dim + token_bytes = head_dim * get_dtype_bitwidth(new_kv.dtype) // 8 + assert token_bytes % slot_bytes == 0 + slots_per_token = token_bytes // slot_bytes + phys_tokens_per_page = physical_page_size_per_kv_packing // slots_per_token + + new_kv = prepare_kv_inputs(new_kv) + new_kv = new_kv.reshape(new_kv.shape[0], slots_per_token * kv_packing, lkv_dim) + assert logical_page_size <= phys_tokens_per_page + + _, num_q_heads, _ = q.shape + max_num_seqs = kv_lens.shape[0] + num_page_indices = page_indices.shape[0] + assert num_page_indices % max_num_seqs == 0 + + def run_mla_kernel( + q: jax.Array, + new_kv: jax.Array, + cache_kv: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + start_seq_idx: jax.Array, + end_seq_idx: jax.Array, + in_output: jax.Array, + in_l: jax.Array, + in_m: jax.Array, + attention_sinks: jax.Array, + static_q_len: int | None, + unnormalized_output: bool, + num_kv_pages_per_block: int, + num_queries_per_block: int, + case: MlaCase = MlaCase.MIXED, + ): + bkv_p = num_kv_pages_per_block + if static_q_len is not None: + bq_sz = min(num_queries_per_block, static_q_len) + else: + bq_sz = num_queries_per_block + bkv_sz = bkv_p * logical_page_size + grid = (end_seq_idx - start_seq_idx,) + + in_specs = [ + pl.BlockSpec(memory_space=pltpu.VMEM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + ] + + out_specs = [ + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + ] + + bkv_double_buf = pltpu.VMEM((2, bkv_sz, slots_per_token * kv_packing, lkv_dim), cache_kv.dtype) + bq_double_bufq = pltpu.VMEM((2, bq_sz, num_q_heads, head_dim), q.dtype) + bo_double_buf = bq_double_bufq + + num_l_heads = align_to(num_q_heads, 128) + bl_double_buf = pltpu.VMEM((2, bq_sz, num_l_heads), jnp.float32) + bm_double_buf = bl_double_buf + + l_scratch = pltpu.VMEM((bq_sz * num_q_heads, 128), jnp.float32) + m_scratch = l_scratch + acc_scratch = pltpu.VMEM((bq_sz * num_q_heads, head_dim), jnp.float32) + + scratch_shapes = [ + bkv_double_buf, + bq_double_bufq, + bo_double_buf, + bl_double_buf, + bm_double_buf, + pltpu.SemaphoreType.DMA((6, 2)), + l_scratch, + m_scratch, + acc_scratch, + ] + + scalar_prefetches = ( + kv_lens, + page_indices, + cu_q_lens, + jnp.array([start_seq_idx, end_seq_idx], jnp.int32), + jnp.zeros((3,), jnp.int32), + jnp.full((4,), -1, jnp.int32), + jnp.full((6,), -1, jnp.int32), + ) + + scope_name = f"SWA-{case.symbol}-bq_{bq_sz}-bkvp_{bkv_p}" + kernel = jax.named_scope(scope_name)( + pl.pallas_call( + functools.partial( + _mla_sliding_window_ragged_paged_attention_kernel, + sm_scale=sm_scale, + sliding_window=sliding_window, + static_q_len=static_q_len, + bq_sz=bq_sz, + bkv_p=bkv_p, + logical_page_size=logical_page_size, + unnormalized_output=unnormalized_output, + q_compute_block_size=q_compute_block_size, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=len(scalar_prefetches), + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + scratch_shapes=scratch_shapes, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary",), + vmem_limit_bytes=vmem_limit_bytes, + disable_bounds_checks=True, + ), + out_shape=[ + jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype), + jax.ShapeDtypeStruct(shape=cache_kv.shape, dtype=cache_kv.dtype), + jax.ShapeDtypeStruct(shape=(q.shape[0], num_l_heads), dtype=jnp.float32), + jax.ShapeDtypeStruct(shape=(q.shape[0], num_l_heads), dtype=jnp.float32), + ], + input_output_aliases={ + 11: 0, + 10: 1, + 12: 2, + 13: 3, + }, + name=scope_name, + ) + ) + return kernel( + *scalar_prefetches, + attention_sinks, + q, + new_kv, + cache_kv, + in_output, + in_l, + in_m, + ) + + num_l_heads = align_to(num_q_heads, 128) + if unnormalized_output: + l = jnp.zeros((q.shape[0], num_l_heads), dtype=jnp.float32) + m = jnp.full((q.shape[0], num_l_heads), jnp.finfo(jnp.float32).min, dtype=jnp.float32) + in_output = jnp.zeros_like(q) + else: + l = jnp.zeros((q.shape[0], num_l_heads), dtype=jnp.float32) + m = jnp.zeros((q.shape[0], num_l_heads), dtype=jnp.float32) + in_output = jnp.zeros_like(q) + + output, updated_kv, out_l, out_m = run_mla_kernel( + q, + new_kv, + cache_kv, + kv_lens, + page_indices, + cu_q_lens, + num_kv_pages_per_block=num_kv_pages_per_blocks[0], + num_queries_per_block=num_queries_per_blocks[0], + start_seq_idx=jnp.array(0), + end_seq_idx=distribution[0], + in_output=in_output, + in_l=l, + in_m=m, + attention_sinks=attention_sinks, + static_q_len=1, + unnormalized_output=unnormalized_output, + case=MlaCase.DECODE, + ) + + if chunk_prefill_size is not None: + output, updated_kv, out_l, out_m = run_mla_kernel( + q, + new_kv, + updated_kv, + kv_lens, + page_indices, + cu_q_lens, + num_kv_pages_per_block=num_kv_pages_per_blocks[1], + num_queries_per_block=num_queries_per_blocks[1], + start_seq_idx=distribution[0], + end_seq_idx=distribution[1], + in_output=output, + in_l=out_l, + in_m=out_m, + attention_sinks=attention_sinks, + static_q_len=chunk_prefill_size, + unnormalized_output=unnormalized_output, + case=MlaCase.PREFILL, + ) + + output, updated_kv, out_l, out_m = run_mla_kernel( + q, + new_kv, + updated_kv, + kv_lens, + page_indices, + cu_q_lens, + num_kv_pages_per_block=num_kv_pages_per_blocks[2], + num_queries_per_block=num_queries_per_blocks[2], + start_seq_idx=distribution[1], + end_seq_idx=distribution[2], + in_output=output, + in_l=out_l, + in_m=out_m, + attention_sinks=attention_sinks, + static_q_len=None, + unnormalized_output=unnormalized_output, + case=MlaCase.MIXED, + ) + + output = prepare_outputs(output, actual_num_q_heads, actual_head_dim) + out_l = out_l[:, :actual_num_q_heads] + return output, updated_kv, out_l, out_m + +def workload( + q: jax.Array, + new_kv: jax.Array, + cache_kv: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + distribution: jax.Array, + attention_sinks: jax.Array, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]: + sm_scale = 1.0 + sliding_window = 128 + logical_page_size = 128 + chunk_prefill_size = None + num_kv_pages_per_block = 2 + num_queries_per_block = 32 + q_compute_block_size = 2 + vmem_limit_bytes = 100 * 1024 * 1024 + unnormalized_output = True + + if cache_kv.shape[1] != logical_page_size * 2: + total_pages = cache_kv.shape[0] + cache_kv = jnp.zeros((total_pages, logical_page_size * 2, 4, 128), dtype=cache_kv.dtype) + + return mla_sliding_window_ragged_paged_attention( + q, + new_kv, + cache_kv, + kv_lens, + page_indices, + cu_q_lens, + distribution, + attention_sinks, + sm_scale=sm_scale, + sliding_window=sliding_window, + logical_page_size=logical_page_size, + chunk_prefill_size=chunk_prefill_size, + num_kv_pages_per_block=num_kv_pages_per_block, + num_queries_per_block=num_queries_per_block, + q_compute_block_size=q_compute_block_size, + vmem_limit_bytes=vmem_limit_bytes, + unnormalized_output=unnormalized_output, + ) + +def benchmark(num_warmup=5, num_iters=100): + """Benchmark and return results dict.""" + inputs_list = create_inputs() + fn = jax.jit(workload) + times_list = [] + std_ms_list = [] + output_shape_list = [] + for inputs in inputs_list: + for _ in range(num_warmup): + out = fn(*inputs) + jax.block_until_ready(out) + times = [] + for _ in range(num_iters): + t0 = time.perf_counter() + out = fn(*inputs) + jax.block_until_ready(out) + times.append(time.perf_counter() - t0) + times = np.array(times) * 1000 + avg = float(np.mean(times)) + times_list.append(round(avg, 4)) + std_ms_list.append(round(float(np.std(times)), 4)) + if hasattr(out, 'shape'): + out_shape = list(out.shape) + elif isinstance(out, (tuple, list)): + out_shape = [list(x.shape) if hasattr(x, 'shape') else [] for x in out] + else: + out_shape = [] + output_shape_list.append(out_shape) + return { + 'name': CONFIG['name'], + 'model': CONFIG['model'], + 'operator': CONFIG['operator'], + 'config': {k: v for k, v in CONFIG.items() if k not in ('name', 'model', 'operator', 'atol', 'rtol')}, + 'time_ms': times_list, + 'std_ms': std_ms_list, + 'output_shape': output_shape_list, + 'status': 'success', + } + + +if __name__ == '__main__': + import json + print(json.dumps(benchmark())) diff --git a/JAXBench/benchmark/level2/6p_Paged_Attention/baseline.py b/JAXBench/benchmark/level2/6p_Paged_Attention/baseline.py new file mode 100644 index 0000000..7cddbf7 --- /dev/null +++ b/JAXBench/benchmark/level2/6p_Paged_Attention/baseline.py @@ -0,0 +1,766 @@ +# Copyright 2024 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. + +"""Pallas PagedAttention TPU kernel — Llama-3.1-70B decode dimensions. + +Upstream kernel from jax.experimental.pallas.ops.tpu.paged_attention, wrapped +as a JAXBench workload with CONFIG / create_inputs / workload. + +quantization_utils imported from the installed JAX package (not optimizable). +""" + +import time +from collections.abc import Sequence +import functools +from typing import Literal + +import jax +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +from jax.experimental.pallas.ops.tpu.paged_attention import quantization_utils +import jax.numpy as jnp +import numpy as np + + +DEFAULT_MASK_VALUE = -0.7 * float(np.finfo(np.dtype("float32")).max) + + +class MultiPageAsyncCopyDescriptor: + """Descriptor for async copy of multiple K/V pages from HBM.""" + + def __init__( + self, + pages_hbm_ref, + scales_pages_hbm_ref, + vmem_buffer, + scales_vmem_buffer, + sem, + page_indices, + page_indices_start_offset, + num_pages_to_load, + head_index, + ): + self._vmem_buffer = vmem_buffer + self._scales_vmem_buffer = scales_vmem_buffer + self._num_pages_to_load = num_pages_to_load + if head_index is not None: + self._pages_hbm_ref = pages_hbm_ref.at[head_index] + if scales_pages_hbm_ref is not None: + self._scales_pages_hbm_ref = scales_pages_hbm_ref.at[head_index] + else: + self._scales_pages_hbm_ref = None + else: + self._pages_hbm_ref = pages_hbm_ref + self._scales_pages_hbm_ref = scales_pages_hbm_ref + self._sem = sem + self._page_indices = page_indices + self._page_indices_start_offset = page_indices_start_offset + self._async_copies = [ + self._make_async_copy(i) for i in range(self._num_pages_to_load) + ] + if ( + self._scales_pages_hbm_ref is not None + and self._scales_vmem_buffer is not None + ): + self._async_copies += [ + self._make_scales_async_copy(i) + for i in range(self._num_pages_to_load) + ] + + def _make_async_copy(self, i): + page_index = self._page_indices[self._page_indices_start_offset + i] + return pltpu.make_async_copy( + self._pages_hbm_ref.at[page_index], self._vmem_buffer.at[i], self._sem + ) + + def _make_scales_async_copy(self, i): + page_index = self._page_indices[self._page_indices_start_offset + i] + return pltpu.make_async_copy( + self._scales_pages_hbm_ref.at[page_index], # pytype: disable=attribute-error + self._scales_vmem_buffer.at[i], # pytype: disable=attribute-error + self._sem, + ) + + def start(self): + """Starts the async copies.""" + for async_copy in self._async_copies: + async_copy.start() + + def _maybe_dequantize(self, x, x_scale, dtype=jnp.bfloat16): + if x_scale is None: + return x.astype(dtype) + return quantization_utils.from_int8(x, x_scale, dtype=dtype) + + def wait_and_get_loaded(self) -> jax.Array: + """Wait async copies and gets the loaded buffer as a jax.Array.""" + for async_copy in self._async_copies: + async_copy.wait() + head_dim = self._vmem_buffer.shape[-1] + jax_array = self._vmem_buffer[...].astype(jnp.float32) + if self._scales_vmem_buffer is not None: + scales_jax_array = self._scales_vmem_buffer[...].astype(jnp.float32) + else: + scales_jax_array = None + jax_array = self._maybe_dequantize(jax_array, scales_jax_array) + return jax_array.reshape(-1, head_dim) + + +def paged_flash_attention_kernel( + lengths_ref, + page_indices_ref, + buffer_index_ref, + init_flag_ref, + q_ref, + k_pages_hbm_ref, + k_scales_pages_hbm_ref, + v_pages_hbm_ref, + v_scales_pages_hbm_ref, + o_ref, + m_ref, + l_ref, + k_vmem_buffer, + k_scales_vmem_buffer, + v_vmem_buffer, + v_scales_vmem_buffer, + k_sems, + v_sems, + *, + batch_size: int, + pages_per_compute_block: int, + pages_per_sequence: int, + mask_value: float, + attn_logits_soft_cap: float | None, + megacore_mode: str | None, + program_ids=(), +): + """Pallas kernel for paged attention.""" + if program_ids: + core_index, b, h, i = program_ids + else: + core_index, b, h, i = ( + pl.program_id(0), + pl.program_id(1), + pl.program_id(2), + pl.program_id(3), + ) + num_kv_heads, _, page_size, _ = k_pages_hbm_ref.shape + bk = page_size * pages_per_compute_block + num_cores = pl.num_programs(0) + + b_step = num_cores if megacore_mode == "batch" else 1 + b_start = core_index if megacore_mode == "batch" else 0 + h_step = num_cores if megacore_mode == "kv_head" else 1 + h_start = core_index if megacore_mode == "kv_head" else 0 + + h = h * h_step + h_start + b = b * b_step + b_start + length = lengths_ref[b] + + def compute_block_indices(b, h, i): + + def advance_b(): + next_b = b + b_step + + def advance_to_next_non_zero_length(): + next_next_b = next_b + b_step + return lax.fori_loop( + lax.div(next_next_b, b_step), + lax.div(batch_size, b_step), + lambda _, b: jnp.where(lengths_ref[b] == 0, b + b_step, b), + next_next_b, + ) + + return ( + lax.cond( + jnp.logical_and( + next_b < batch_size, + lengths_ref[lax.clamp(0, next_b, batch_size - 1)] == 0), + advance_to_next_non_zero_length, + lambda: next_b, + ), + h_start, + 0, + ) + + def advance_h(): + next_h = h + h_step + return lax.cond(next_h < num_kv_heads, lambda: (b, next_h, 0), advance_b) + + return lax.cond(i * bk < lengths_ref[b], lambda: (b, h, i), advance_h) + + def create_kv_async_copy_descriptors(b, h, i, buffer_index): + page_offset = b * pages_per_sequence + i * pages_per_compute_block + pages_to_load = pages_per_compute_block + async_copy_k = MultiPageAsyncCopyDescriptor( + k_pages_hbm_ref, + k_scales_pages_hbm_ref, + k_vmem_buffer.at[buffer_index], + k_scales_vmem_buffer.at[buffer_index] + if k_scales_vmem_buffer is not None + else None, + k_sems.at[buffer_index], + page_indices_ref, + page_offset, + pages_to_load, + h, + ) + async_copy_v = MultiPageAsyncCopyDescriptor( + v_pages_hbm_ref, + v_scales_pages_hbm_ref, + v_vmem_buffer.at[buffer_index], + v_scales_vmem_buffer.at[buffer_index] + if v_scales_vmem_buffer is not None + else None, + v_sems.at[buffer_index], + page_indices_ref, + page_offset, + pages_to_load, + h, + ) + return async_copy_k, async_copy_v + + @pl.when(i * bk < length) + def flash_attention(): # pylint: disable=unused-variable + init_flag = init_flag_ref[0] + init_flag_ref[0] = 0 + buffer_index = buffer_index_ref[0] + next_b, next_h, next_i = compute_block_indices(b, h, i + 1) + + @pl.when(init_flag) + def prefetch_first_block(): # pylint: disable=unused-variable + async_copy_k, async_copy_v = create_kv_async_copy_descriptors( + b, h, i, buffer_index + ) + async_copy_k.start() + async_copy_v.start() + + @pl.when(i == 0) + def init(): # pylint: disable=unused-variable + m_ref[...] = jnp.full_like(m_ref, -jnp.inf) + l_ref[...] = jnp.zeros_like(l_ref) + o_ref[...] = jnp.zeros_like(o_ref) + + @pl.when(next_b < batch_size) + def prefetch_next_block(): # pylint: disable=unused-variable + next_buffer_index = jnp.where(buffer_index == 0, 1, 0) + async_copy_next_k, async_copy_next_v = create_kv_async_copy_descriptors( + next_b, next_h, next_i, next_buffer_index + ) + async_copy_next_k.start() + async_copy_next_v.start() + buffer_index_ref[0] = next_buffer_index + + async_copy_k, async_copy_v = create_kv_async_copy_descriptors( + b, h, i, buffer_index + ) + q = q_ref[...].astype(jnp.float32) + k = async_copy_k.wait_and_get_loaded() + qk = jnp.einsum("gd,td->gt", q, k, preferred_element_type=jnp.float32) + if attn_logits_soft_cap is not None: + capped_qk = jnp.tanh(qk / attn_logits_soft_cap) + qk = capped_qk * attn_logits_soft_cap + + mask = i * bk + jax.lax.broadcasted_iota(jnp.int32, qk.shape, 1) < length + qk = qk + jnp.where(mask, 0.0, mask_value) + m_curr = qk.max(axis=-1) + + s_curr = jnp.exp(qk - m_curr[..., None]) + m_prev, l_prev = m_ref[...], l_ref[...] + l_curr = jax.lax.broadcast_in_dim(s_curr.sum(axis=-1), l_prev.shape, (0,)) + m_curr = jax.lax.broadcast_in_dim(m_curr, m_prev.shape, (0,)) + m_next = jnp.maximum(m_prev, m_curr) + alpha = jnp.exp(m_prev - m_next) + beta = jnp.exp(m_curr - m_next) + l_next = alpha * l_prev + beta * l_curr + m_ref[...], l_ref[...] = m_next, l_next + + v = async_copy_v.wait_and_get_loaded() + o_curr = jnp.einsum("gt,td->gd", s_curr, v) + + o_ref[...] = ( + (l_prev * alpha * o_ref[...] + beta * o_curr) / l_next + ).astype(o_ref.dtype) + + +def paged_flash_attention_kernel_inline_seq_dim( + lengths_ref, + page_indices_ref, + buffer_index_ref, + init_flag_ref, + q_ref, + k_pages_hbm_ref, + k_scales_pages_hbm_ref, + v_pages_hbm_ref, + v_scales_pages_hbm_ref, + o_ref, + m_ref, + l_ref, + k_vmem_buffer, + k_scales_vmem_buffer, + v_vmem_buffer, + v_scales_vmem_buffer, + k_sems, + v_sems, + *, + batch_size: int, + pages_per_compute_block: int, + pages_per_sequence: int, + mask_value: float, + attn_logits_soft_cap: float | None, + megacore_mode: str | None, +): + core_index, b, h = pl.program_id(0), pl.program_id(1), pl.program_id(2) + + # Initialize the output HBM buffers to avoid accessing garbage memory inside + # the kernel body below. + m_ref[...] = jnp.full_like(m_ref, -jnp.inf) + l_ref[...] = jnp.zeros_like(l_ref) + o_ref[...] = jnp.zeros_like(o_ref) + + def body(i, _): + paged_flash_attention_kernel( + lengths_ref, + page_indices_ref, + buffer_index_ref, + init_flag_ref, + q_ref, + k_pages_hbm_ref, + k_scales_pages_hbm_ref, + v_pages_hbm_ref, + v_scales_pages_hbm_ref, + o_ref, + m_ref, + l_ref, + k_vmem_buffer, + k_scales_vmem_buffer, + v_vmem_buffer, + v_scales_vmem_buffer, + k_sems, + v_sems, + batch_size=batch_size, + pages_per_compute_block=pages_per_compute_block, + pages_per_sequence=pages_per_sequence, + mask_value=mask_value, + attn_logits_soft_cap=attn_logits_soft_cap, + megacore_mode=megacore_mode, + program_ids=(core_index, b, h, i), + ) + return () + + bk = pages_per_compute_block * k_pages_hbm_ref.shape[-2] + + if megacore_mode == "batch": + num_cores = pl.num_programs(0) + length = lengths_ref[b * num_cores + core_index] + else: + length = lengths_ref[b] + + lax.fori_loop(0, lax.div(length + bk - 1, bk), body, ()) + + +@functools.partial( + jax.jit, + static_argnames=[ + "pages_per_compute_block", + "attn_logits_soft_cap", + "mask_value", + "megacore_mode", + "inline_seq_dim", + ], +) +def paged_attention( + q: jax.Array, + k_pages: jax.Array | quantization_utils.QuantizedTensor, + v_pages: jax.Array | quantization_utils.QuantizedTensor, + lengths: jax.Array, + page_indices: jax.Array, + *, + mask_value: float = DEFAULT_MASK_VALUE, + attn_logits_soft_cap: float | None = None, + pages_per_compute_block: int, + megacore_mode: str | None = None, + inline_seq_dim: bool = True, +) -> jax.Array: + """Paged grouped query attention. + + Args: + q: A [batch_size, num_q_heads, head_dim] jax.Array. + k_pages: A [num_kv_heads, total_num_pages, page_size, head_dim] jax.Array. + v_pages: A [num_kv_heads, total_num_pages, page_size, head_dim] jax.Array. + lengths: A i32[batch_size] jax.Array the length of each example. + page_indices: A i32[batch_size, pages_per_sequence] jax.Array. Each entry + should be in the range of [0, total_num_pages), indicating where to locate + the page in `k_pages` or `v_pages`. + mask_value: The value used for padding in attention. By default it is a very + negative floating point number. + attn_logits_soft_cap: The value used for soft capping the attention logits. + pages_per_compute_block: how many pages to be processed in one flash + attention block in the pallas kernel. + megacore_mode: if set, enable megacore to parallelize the computation. Must + be one of ['kv_head', 'batch', None]. Caveat: set this only if megacore is + enabled, otherwise the kernel may hang. If you are not sure, leave it to + None. + * None: disable megacore parallelism. + * kv_head: megacore parallelism on KV heads; requires number of KV heads + divisible by 2. + * batch: megacore parallelism on batch dimension; requires batch divisible + by 2. + inline_seq_dim: whether to fuse kernel instances along the sequence dim into + one kernel. + + Returns: + The output of attention([batch_size, num_q_heads, head_dim]). + """ + if isinstance(k_pages, quantization_utils.QuantizedTensor): + k_pages, k_scales_pages = k_pages.weight, k_pages.scales + assert isinstance(k_scales_pages, jax.Array) # For typing. + k_scales_pages = jnp.broadcast_to( + k_scales_pages, (*k_scales_pages.shape[:-1], k_pages.shape[-1]) + ) + else: + k_scales_pages = None + if isinstance(v_pages, quantization_utils.QuantizedTensor): + v_pages, v_scales_pages = v_pages.weight, v_pages.scales + assert isinstance(v_scales_pages, jax.Array) # For typing. + v_scales_pages = jnp.broadcast_to( + v_scales_pages, (*v_scales_pages.shape[:-1], v_pages.shape[-1]) + ) + else: + v_scales_pages = None + + batch_size, num_q_heads, head_dim = q.shape + num_kv_heads, _, page_size, head_dim_k = k_pages.shape + batch_size_paged_indices, pages_per_sequence = page_indices.shape + + if k_pages.shape != v_pages.shape: + raise ValueError( + f"k_pages and v_pages must have the same shape. Got {k_pages.shape} and" + f" {v_pages.shape}" # pytype: disable=attribute-error + ) + if num_q_heads % num_kv_heads != 0: + raise ValueError( + "Number of Q heads must be divisible by number of KV heads. Got" + f" {num_q_heads} and {num_kv_heads}." + ) + if head_dim_k != head_dim: + raise ValueError( + "head_dim of Q must be the same as that of K/V. Got" + f" {head_dim} and {head_dim_k}." + ) + if pages_per_sequence % pages_per_compute_block != 0: + raise ValueError( + "pages_per_compute_block must be divisible by pages per sequence. Got" + f" {pages_per_compute_block} and {pages_per_sequence}." + ) + if lengths.shape != (batch_size,): + raise ValueError("`lengths` and `q` must have the same batch size") + if batch_size_paged_indices != batch_size: + raise ValueError("`page_indices` and `q` must have the same batch size") + if lengths.dtype != jnp.int32: + raise ValueError( + f"The dtype of `lengths` must be int32. Got {lengths.dtype}" + ) + + # TODO(dinghua): get the actual cores per chip once there's an official API. + if megacore_mode == "kv_head": + if num_kv_heads % 2 != 0: + raise ValueError( + "number of KV heads must be even when megacore_mode is 'kv_head'" + ) + num_cores = 2 + elif megacore_mode == "batch": + if batch_size % 2 != 0: + raise ValueError("batch size must be even when megacore_mode is 'batch'") + num_cores = 2 + elif megacore_mode is None: + num_cores = 1 + else: + raise ValueError("megacore_mode must be one of ['kv_head', 'batch', None]") + + num_groups = num_q_heads // num_kv_heads + if (num_groups) % 8 != 0: + # Reshape q to hint XLA to pick a <1x128> layout otherwise it will pick a + # <8x128> layout for a <1x128> memref inside the kernel and error out. + q = q.reshape(batch_size, num_q_heads, 1, head_dim) + if megacore_mode == "kv_head": + q_block_spec = pl.BlockSpec( + (None, num_groups, None, head_dim), + lambda core_index, b, h, *_: (b, h * num_cores + core_index, 0, 0), + ) + elif megacore_mode == "batch": + q_block_spec = pl.BlockSpec( + (None, num_groups, None, head_dim), + lambda core_index, b, h, *_: (b * num_cores + core_index, h, 0, 0), + ) + else: + q_block_spec = pl.BlockSpec( + (None, num_groups, None, head_dim), + lambda core_index, b, h, *_: (b, h, 0, 0), + ) + q_dtype_for_kernel_launch = jnp.float32 + else: + if megacore_mode == "kv_head": + q_block_spec = pl.BlockSpec( + (None, num_groups, head_dim), + lambda core_index, b, h, *_: (b, h * num_cores + core_index, 0), + ) + elif megacore_mode == "batch": + q_block_spec = pl.BlockSpec( + (None, num_groups, head_dim), + lambda core_index, b, h, *_: (b * num_cores + core_index, h, 0), + ) + else: + q_block_spec = pl.BlockSpec( + (None, num_groups, head_dim), + lambda core_index, b, h, *_: (b, h, 0), + ) + q_dtype_for_kernel_launch = q.dtype + + dimension_semantics: Sequence[Literal["parallel", "arbitrary"]] + if inline_seq_dim: + kernel = paged_flash_attention_kernel_inline_seq_dim + grid = ( + num_cores, + batch_size // num_cores if megacore_mode == "batch" else batch_size, + num_kv_heads // num_cores + if megacore_mode == "kv_head" + else num_kv_heads, + ) + dimension_semantics = ("parallel", "arbitrary", "arbitrary") + else: + kernel = paged_flash_attention_kernel + grid = ( + num_cores, + batch_size // num_cores if megacore_mode == "batch" else batch_size, + num_kv_heads // num_cores + if megacore_mode == "kv_head" + else num_kv_heads, + pages_per_sequence // pages_per_compute_block, + ) # type: ignore + dimension_semantics = ("parallel", "arbitrary", "arbitrary", "arbitrary") + + if k_scales_pages is not None and v_scales_pages is not None: + in_specs = [ + q_block_spec, + pl.BlockSpec(memory_space=pl.ANY), + pl.BlockSpec(memory_space=pl.ANY), + pl.BlockSpec(memory_space=pl.ANY), + pl.BlockSpec(memory_space=pl.ANY), + ] + scratch_shapes = ( + pltpu.VMEM( + ( + 2, # For double buffering during DMA copies. + pages_per_compute_block, + page_size, + head_dim, + ), + k_pages.dtype, + ), # k_pages buffer + pltpu.VMEM( + ( + 2, # For double buffering during DMA copies. + pages_per_compute_block, + page_size, + head_dim, + ), + k_scales_pages.dtype, # pytype: disable=attribute-error + ), # k_scales_pages buffer + pltpu.VMEM( + ( + 2, # For double buffering during DMA copies. + pages_per_compute_block, + page_size, + head_dim, + ), + v_pages.dtype, + ), # v_pages buffer + pltpu.VMEM( + ( + 2, # For double buffering during DMA copies. + pages_per_compute_block, + page_size, + head_dim, + ), + v_scales_pages.dtype, # pytype: disable=attribute-error + ), # v_scales_pages buffer + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + ) + else: + in_specs = [ + q_block_spec, + pl.BlockSpec(memory_space=pl.ANY), + None, # type: ignore[list-item] + pl.BlockSpec(memory_space=pl.ANY), + None, # type: ignore[list-item] + ] + scratch_shapes = ( + pltpu.VMEM( + ( + 2, # For double buffering during DMA copies. + pages_per_compute_block, + page_size, + head_dim, + ), + k_pages.dtype, + ), # k_pages buffer + None, + pltpu.VMEM( + ( + 2, # For double buffering during DMA copies. + pages_per_compute_block, + page_size, + head_dim, + ), + v_pages.dtype, + ), # v_pages buffer + None, + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + ) + + out, _, _ = pl.pallas_call( + functools.partial( + kernel, + pages_per_sequence=pages_per_sequence, + batch_size=batch_size, + pages_per_compute_block=pages_per_compute_block, + mask_value=mask_value, + attn_logits_soft_cap=attn_logits_soft_cap, + megacore_mode=megacore_mode, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + # There are 4 scalars prefetched per kernel call: `lengths_ref`, + # `page_indices_ref`, `buffer_index_ref`, `init_flag_ref` + num_scalar_prefetch=4, + in_specs=in_specs, + out_specs=[ + q_block_spec, + q_block_spec, + q_block_spec, + ], + grid=grid, + scratch_shapes=scratch_shapes, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=dimension_semantics + ), + out_shape=[ + jax.ShapeDtypeStruct(q.shape, q_dtype_for_kernel_launch), + jax.ShapeDtypeStruct((*q.shape[:-1], 1), jnp.float32), + jax.ShapeDtypeStruct((*q.shape[:-1], 1), jnp.float32), + ], + )( + lengths, + page_indices.reshape(-1), + jnp.zeros((1,), jnp.int32), # buffer index + jnp.ones((1,), jnp.int32), # init flag + q.astype(q_dtype_for_kernel_launch), + k_pages, + k_scales_pages, + v_pages, + v_scales_pages, + ) + return out.reshape(batch_size, num_q_heads, head_dim).astype(q.dtype) + + +CONFIG = { + 'name': 'pallas_paged_attention_llama70b', + 'model': 'Llama-3.1-70B', + 'operator': 'pallas_paged_attention', + 'batch': 64, + 'num_q_heads': 64, + 'num_kv_heads': 8, + 'head_dim': 128, + 'page_size': 16, + 'pages_per_seq': 256, + 'atol': 1e-2, + 'rtol': 2e-2, +} + +# Tuned by autotune_block_sizes.py. Re-run to update. +TUNED_PARAMS = {'pages_per_compute_block': 128} + + +def get_flops(): + B = CONFIG['batch'] + H_q = CONFIG['num_q_heads'] + D = CONFIG['head_dim'] + seq_len = CONFIG['pages_per_seq'] * CONFIG['page_size'] + return B * H_q * (4 * seq_len * D) + + +def create_inputs(dtype=jnp.bfloat16): + key = jax.random.key(42) + keys = jax.random.split(key, 5) + B = CONFIG['batch'] + H_q = CONFIG['num_q_heads'] + H_kv = CONFIG['num_kv_heads'] + D = CONFIG['head_dim'] + page_size = CONFIG['page_size'] + pages_per_seq = CONFIG['pages_per_seq'] + total_num_pages = B * pages_per_seq + + q = jax.random.normal(keys[0], (B, H_q, D), dtype=dtype) + k_pages = jax.random.normal(keys[1], (total_num_pages, page_size, H_kv, D), dtype=dtype) * 0.02 + v_pages = jax.random.normal(keys[2], (total_num_pages, page_size, H_kv, D), dtype=dtype) * 0.02 + + k_pages = k_pages.transpose(2, 0, 1, 3) + v_pages = v_pages.transpose(2, 0, 1, 3) + + seq_len = pages_per_seq * page_size + lengths = jnp.full((B,), seq_len, dtype=jnp.int32) + page_indices = jnp.arange(total_num_pages, dtype=jnp.int32).reshape(B, pages_per_seq) + return q, k_pages, v_pages, lengths, page_indices + + +def workload(q, k_pages, v_pages, lengths, page_indices): + return paged_attention( + q, k_pages, v_pages, lengths, page_indices, + pages_per_compute_block=TUNED_PARAMS['pages_per_compute_block'], + ) + + +def benchmark(num_warmup=5, num_iters=100): + """Benchmark and return results dict.""" + inputs = create_inputs() + fn = jax.jit(workload) + for _ in range(num_warmup): + out = fn(*inputs) + out.block_until_ready() + times = [] + for _ in range(num_iters): + t0 = time.perf_counter() + out = fn(*inputs) + out.block_until_ready() + times.append(time.perf_counter() - t0) + times = np.array(times) * 1000 + avg = float(np.mean(times)) + return { + 'name': CONFIG['name'], + 'model': CONFIG['model'], + 'operator': CONFIG['operator'], + 'config': {k: v for k, v in CONFIG.items() if k not in ('name', 'model', 'operator', 'atol', 'rtol')}, + 'time_ms': round(avg, 4), + 'std_ms': round(float(np.std(times)), 4), + 'output_shape': list(out.shape) if hasattr(out, 'shape') else [], + 'status': 'success', + } + + +if __name__ == '__main__': + import json + print(json.dumps(benchmark())) diff --git a/JAXBench/benchmark/level2/7p_Ragged_Paged_Attention/baseline.py b/JAXBench/benchmark/level2/7p_Ragged_Paged_Attention/baseline.py new file mode 100644 index 0000000..82e9546 --- /dev/null +++ b/JAXBench/benchmark/level2/7p_Ragged_Paged_Attention/baseline.py @@ -0,0 +1,1035 @@ +# Copyright 2025 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. + +"""TPU-Friendly Ragged Paged Attention kernel. + +This kernel offers a highly optimized implementation of ragged paged attention, +specifically designed for TPU and compatible with a wide range of model +specifications. It supports mixed prefill and decoding, enhancing throughput +during inference. +""" + +import numpy as np +import time +import functools +import jax +from jax import lax +from jax._src import dtypes +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +from jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes import get_tuned_block_sizes +import jax.numpy as jnp + + +DEFAULT_MASK_VALUE = -0.7 * float(jnp.finfo(jnp.dtype("float32")).max) + + +class MultiPageAsyncCopyDescriptor: + """Descriptor for async copy of multiple K/V pages from HBM.""" + + def __init__( + self, + pages_hbm_ref, # [total_num_pages, page_size, num_combined_kv_heads_per_blk, head_dim] + vmem_buf, # [num_kv_pages_per_blk, page_size, num_combined_kv_heads_per_blk, head_dim] + sem, + page_indices_ref, # i32[max_num_seqs, pages_per_seq] + metadata, # [seq_idx, start_page_idx, end_page_idx] + ): + self._vmem_buf = vmem_buf + seq_id, start_page_idx, end_page_idx = metadata + self._async_copies = [] + # TODO(jevinjiang): Only fetch dynamic shape in need! This will insert + # a bunch of if-ops. Check the performance when we have benchmarking setup. + for i in range(vmem_buf.shape[0]): + page_idx = start_page_idx + i + page_idx = jax.lax.select(page_idx < end_page_idx, page_idx, 0) + self._async_copies.append( + pltpu.make_async_copy( + pages_hbm_ref.at[page_indices_ref[seq_id, page_idx]], + vmem_buf.at[i], + sem, + ) + ) + + def start(self): + """Starts the async copies.""" + for async_copy in self._async_copies: + async_copy.start() + + def wait(self): + for async_copy in self._async_copies: + async_copy.wait() + return self._vmem_buf + + +def ref_ragged_paged_attention( + queries: jax.Array, # [max_num_batched_tokens, num_q_heads, head_dim] + kv_pages: jax.Array, # [total_num_pages, page_size, num_combined_kv_heads, head_dim] + kv_lens: jax.Array, # i32[max_num_seqs] + page_indices: jax.Array, # i32[max_num_seqs, pages_per_seq] + cu_q_lens: jax.Array, # i32[max_num_seqs + 1] + num_seqs: jax.Array, # i32[1], + *, + sm_scale: float = 1.0, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = DEFAULT_MASK_VALUE, + k_scale: float | None = None, + v_scale: float | None = None, +): + static_validate_inputs( + queries, + kv_pages, + kv_lens, + page_indices, + cu_q_lens, + num_seqs, + sm_scale=sm_scale, + k_scale=k_scale, + v_scale=v_scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + mask_value=mask_value, + ) + if mask_value is None: + mask_value = DEFAULT_MASK_VALUE + _, _, num_combined_kv_heads, head_dim = kv_pages.shape + assert num_combined_kv_heads % 2 == 0 + num_kv_heads = num_combined_kv_heads // 2 + num_q_heads = queries.shape[1] + assert num_q_heads % num_kv_heads == 0 + num_query_per_kv = num_q_heads // num_kv_heads + outputs = [] + for i in range(num_seqs[0]): + q_start = cu_q_lens[i] + q_end = cu_q_lens[i + 1] + q_len = q_end - q_start + kv_len = kv_lens[i] + indices = page_indices[i] + q = queries[q_start:q_end] + k = kv_pages[indices, :, 0::2, :].reshape(-1, num_kv_heads, head_dim)[ + :kv_len + ] + v = kv_pages[indices, :, 1::2, :].reshape(-1, num_kv_heads, head_dim)[ + :kv_len + ] + if k_scale is not None: + k = k.astype(jnp.float32) * k_scale + k = k.astype(q.dtype) + if v_scale is not None: + v = v.astype(jnp.float32) * v_scale + v = v.astype(q.dtype) + k = jnp.repeat(k, num_query_per_kv, axis=1) + v = jnp.repeat(v, num_query_per_kv, axis=1) + attn = jnp.einsum("qhd,khd->hqk", q, k, preferred_element_type=jnp.float32) + attn *= sm_scale + q_span = (kv_len - q_len) + jax.lax.broadcasted_iota( + jnp.int32, attn.shape, 1 + ) + kv_span = jax.lax.broadcasted_iota(jnp.int32, attn.shape, 2) + mask = q_span < kv_span + if sliding_window is not None: + mask = jnp.logical_or(mask, q_span - sliding_window >= kv_span) + if soft_cap is not None: + attn = soft_cap * jnp.tanh(attn / soft_cap) + attn += jnp.where(mask, mask_value, 0.0) + attn = jax.nn.softmax(attn, axis=-1).astype(v.dtype) + out = jnp.einsum("hqk,khd->qhd", attn, v).astype(queries.dtype) + outputs.append(out) + + return jnp.concatenate(outputs, axis=0) + + +# Expect to run these checks during runtime. +def dynamic_validate_inputs( + q: jax.Array, # [max_num_batched_tokens, num_q_heads, head_dim] + kv_pages: jax.Array, # [total_num_pages, page_size, num_combined_kv_heads, head_dim] + kv_lens: jax.Array, # i32[max_num_seqs] + page_indices: jax.Array, # i32[max_num_seqs, pages_per_seq] + cu_q_lens: jax.Array, # i32[max_num_seqs + 1] + num_seqs: jax.Array, # i32[1] + *, + # These inputs are optional. If not specified, we will not validate them. + sm_scale: float | None = None, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = None, + k_scale: float | None = None, + v_scale: float | None = None, + # Kernel tuning params. + num_kv_pages_per_block: int | None = None, + num_queries_per_block: int | None = None, + vmem_limit_bytes: int | None = None, +): + static_validate_inputs( + q, + kv_pages, + kv_lens, + page_indices, + cu_q_lens, + num_seqs, + sm_scale=sm_scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + mask_value=mask_value, + k_scale=k_scale, + v_scale=v_scale, + num_kv_pages_per_block=num_kv_pages_per_block, + num_queries_per_block=num_queries_per_block, + vmem_limit_bytes=vmem_limit_bytes, + ) + max_num_batched_tokens = q.shape[0] + page_size = kv_pages.shape[1] + max_num_seqs, pages_per_seq = page_indices.shape + if num_seqs[0] > max_num_seqs: + raise ValueError(f"{num_seqs[0]=} must be less or equal to {max_num_seqs=}") + max_kv_len = jnp.max(kv_lens) + min_pages_per_seq = pl.cdiv(max_kv_len, page_size) + if pages_per_seq < min_pages_per_seq: + raise ValueError( + f"{pages_per_seq=} must be greater or equal to" + f" {min_pages_per_seq=} given {max_kv_len=} and {page_size=}." + ) + if cu_q_lens[num_seqs[0]] > max_num_batched_tokens: + raise ValueError( + f"Total q tokens {cu_q_lens[num_seqs[0]]} must be less or equal to" + f" {max_num_batched_tokens=}." + ) + for i in range(num_seqs[0]): + q_len = cu_q_lens[i + 1] - cu_q_lens[i] + kv_len = kv_lens[i] + if q_len > kv_len: + raise ValueError( + f"{q_len=} must be less or equal to {kv_len=} at sequence {i}." + ) + + +# Expect to run these checks during compile time. +def static_validate_inputs( + q: jax.Array, # [max_num_batched_tokens, num_q_heads, head_dim] + kv_pages: jax.Array, # [total_num_pages, page_size, num_combined_kv_heads, head_dim] + kv_lens: jax.Array, # i32[max_num_seqs] + page_indices: jax.Array, # i32[max_num_seqs, pages_per_seq] + cu_q_lens: jax.Array, # i32[max_num_seqs + 1] + num_seqs: jax.Array, # i32[1] + *, + # These inputs are optional. If not specified, we will not validate them. + sm_scale: float | None = None, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = None, + k_scale: float | None = None, + v_scale: float | None = None, + # Kernel tuning params. + num_kv_pages_per_block: int | None = None, + num_queries_per_block: int | None = None, + vmem_limit_bytes: int | None = None, +): + _, num_q_heads, head_dim = q.shape + _, _, num_combined_kv_heads, head_dim_k = kv_pages.shape + assert num_combined_kv_heads % 2 == 0 + assert isinstance(k_scale, float) or k_scale is None + assert isinstance(v_scale, float) or v_scale is None + num_kv_heads = num_combined_kv_heads // 2 + max_num_seqs, pages_per_seq = page_indices.shape + if num_seqs.shape != (1,): + raise ValueError(f"{num_seqs.shape=} must be (1,)") + if head_dim_k != head_dim: + raise ValueError( + f"Q head_dim {head_dim} must be the same as that of K/V {head_dim_k}." + ) + if kv_lens.shape != (max_num_seqs,): + raise ValueError( + f"Expected {kv_lens.shape=} to be ({max_num_seqs},) where" + " `max_num_seqs` is `page_indices.shape[0]`." + ) + if cu_q_lens.shape != (max_num_seqs + 1,): + raise ValueError( + f"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},) where" + " `max_num_seqs` is `page_indices.shape[0]`." + ) + if ( + kv_lens.dtype != jnp.int32 + or page_indices.dtype != jnp.int32 + or cu_q_lens.dtype != jnp.int32 + ): + raise ValueError( + "The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must be" + f" int32. Got {kv_lens.dtype=}, {page_indices.dtype=}," + f" {cu_q_lens.dtype=}." + ) + if num_q_heads % num_kv_heads != 0: + raise ValueError(f"{num_q_heads=} must be divisible by {num_kv_heads=}") + if sliding_window is not None and sliding_window <= 0: + raise ValueError(f"{sliding_window=} must be positive.") + if soft_cap is not None and soft_cap == 0.0: + raise ValueError(f"{soft_cap=} must not be 0.0.") + if ( + num_kv_pages_per_block is not None + and not 0 < num_kv_pages_per_block <= pages_per_seq + ): + raise ValueError( + f"{num_kv_pages_per_block=} must be in range (0, {pages_per_seq}]." + ) + if num_queries_per_block is not None and num_queries_per_block <= 0: + raise ValueError(f"{num_queries_per_block=} must be positive.") + if vmem_limit_bytes is not None and vmem_limit_bytes <= 0: + raise ValueError(f"{vmem_limit_bytes=} must be positive.") + del sm_scale # No constraints on sm_scale. + del mask_value # No consstraints on mask_value. + + +def ragged_paged_attention_kernel( + # Prefetch + kv_lens_ref, # [max_num_seqs] + page_indices_ref, # [max_num_seqs, pages_per_seq] + cu_q_lens_ref, # [max_num_seqs + 1] + seq_buf_idx_ref, + # TODO(jevinjiang): if OOM in SMEM, consider pack to other scalar refs. + num_seqs_ref, + # Input + q_ref, # [num_q_per_blk, num_q_heads_per_blk, head_dim] + kv_pages_hbm_ref, # [total_num_pages, page_size, num_combined_kv_heads, head_dim] + # Output + o_ref, # [num_q_per_blk, num_q_heads_per_blk, head_dim] + # Scratch + kv_bufs, # [2, num_kv_pages_per_blk, page_size, num_combined_kv_heads_per_blk, head_dim] + sems, # [2, 2] + l_ref, # [num_kv_heads_per_blk, num_q_per_blk * num_q_heads_per_kv_head, 128] + m_ref, # [num_kv_heads_per_blk, num_q_per_blk * num_q_heads_per_kv_head, 128] + acc_ref, # [num_q_per_blk, num_q_heads_per_blk, head_dim] + *, + sm_scale: float, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = DEFAULT_MASK_VALUE, + k_scale: float | None = None, + v_scale: float | None = None, +): + if mask_value is None: + mask_value = DEFAULT_MASK_VALUE + num_q_per_blk, num_q_heads_per_blk, head_dim = q_ref.shape + pages_per_seq = page_indices_ref.shape[-1] + num_seqs = num_seqs_ref[0] + _, num_kv_pages_per_blk, page_size, num_combined_kv_heads_per_blk, _ = ( + kv_bufs.shape + ) + num_kv_heads_per_blk = num_combined_kv_heads_per_blk // 2 + num_kv_per_blk = num_kv_pages_per_blk * page_size + num_q_heads_per_kv_head = num_q_heads_per_blk // num_kv_heads_per_blk + heads_blk_idx, q_blk_idx = ( + pl.program_id(0), + pl.program_id(1), + ) + num_heads_blks = pl.num_programs(0) + init_seq_idx = seq_buf_idx_ref[0] + init_buf_idx = seq_buf_idx_ref[1] + q_len_start = q_blk_idx * num_q_per_blk + q_len_end = q_len_start + num_q_per_blk + + def create_kv_async_copy_descriptors( + heads_blk_idx, seq_idx, kv_blk_idx, buf_idx + ): + start_kv_page_idx = kv_blk_idx * num_kv_pages_per_blk + end_kv_page_idx = jnp.minimum( + pages_per_seq, pl.cdiv(kv_lens_ref[seq_idx], page_size) + ) + metadata = (seq_idx, start_kv_page_idx, end_kv_page_idx) + heads_start = heads_blk_idx * num_combined_kv_heads_per_blk + async_copy_kv = MultiPageAsyncCopyDescriptor( + kv_pages_hbm_ref.at[ + :, :, pl.ds(heads_start, num_combined_kv_heads_per_blk), : + ], + kv_bufs.at[buf_idx], + sems.at[buf_idx], + page_indices_ref, + metadata, + ) + return async_copy_kv + + # TODO(jevinjiang): Add these to Mosaic: + # 1. Support arbitrary strided load/store for int4 and int8 dtype. + # 2. Support arbitrary strided load/store for any last dimension. + def strided_load_kv(ref, start, step): + packing = get_dtype_packing(ref.dtype) + if packing == 1: + return [ref[start::step, :]], [ref[start + 1 :: step, :]] + assert packing in (2, 4, 8) + assert step % packing == 0 + k_list, v_list = [], [] + b_start = start // packing + b_step = step // packing + b_ref = ref.bitcast(jnp.uint32) + b = b_ref[b_start::b_step, :] + + # TODO(chengjiyao): use the general strided loading logic for bf16 after + # fixing the issue in mosaic's infer vector layout pass + if ref.dtype == jnp.bfloat16: + bk = b << 16 + bv = b & jnp.uint32(0xFFFF0000) + k = pltpu.bitcast(bk, jnp.float32).astype(jnp.bfloat16) + v = pltpu.bitcast(bv, jnp.float32).astype(jnp.bfloat16) + k_list.append(k) + v_list.append(v) + else: + bitwidth = 32 // packing + bitcast_dst_dtype = jnp.dtype(f"uint{bitwidth}") + for i in range(0, packing, 2): + bk = b >> (i * bitwidth) + k = pltpu.bitcast(bk.astype(bitcast_dst_dtype), ref.dtype) + k_list.append(k) + bv = b >> ((i + 1) * bitwidth) + v = pltpu.bitcast(bv.astype(bitcast_dst_dtype), ref.dtype) + v_list.append(v) + + return k_list, v_list + + def fold_on_2nd_minor(vec): + assert vec.dtype == jnp.bfloat16 or vec.dtype == jnp.float32 + assert len(vec.shape) >= 2 + last_dim = vec.shape[-1] + packing = get_dtype_packing(vec.dtype) + if vec.shape[-2] % packing != 0: + vec = vec.astype(jnp.float32) + return vec.reshape(-1, last_dim) + + @pl.when(heads_blk_idx + q_blk_idx == 0) + def prefetch_first_kv_blk(): + async_copy_kv = create_kv_async_copy_descriptors( + heads_blk_idx, init_seq_idx, 0, init_buf_idx + ) + async_copy_kv.start() + + def is_cur_q_blk_needed(q_states): + done, cur_seq_idx, _ = q_states + should_run = jnp.logical_and(q_len_start < cu_q_lens_ref[num_seqs], + cur_seq_idx < num_seqs) + return jnp.logical_and(done == 0, should_run) + + def compute_with_cur_q_blk(q_states): + done, cur_seq_idx, cur_buf_idx = q_states + q_start = cu_q_lens_ref[cur_seq_idx] + q_end = cu_q_lens_ref[cur_seq_idx + 1] + q_len = q_end - q_start + kv_len = kv_lens_ref[cur_seq_idx] + + def get_next_prefetch_ids( + heads_blk_idx, cur_seq_idx, kv_blk_idx, cur_buf_idx + ): + next_kv_blk_idx = kv_blk_idx + 1 + is_last_kv_blk = next_kv_blk_idx * num_kv_per_blk >= kv_len + next_kv_blk_idx = lax.select( + is_last_kv_blk, + 0, + next_kv_blk_idx, + ) + is_cur_seq_end_in_cur_q_blk = q_end <= q_len_end + next_seq_idx = lax.select( + is_last_kv_blk, + lax.select(is_cur_seq_end_in_cur_q_blk, cur_seq_idx + 1, cur_seq_idx), + cur_seq_idx, + ) + is_last_seq = next_seq_idx == num_seqs + next_seq_idx = lax.select( + is_last_seq, + 0, + next_seq_idx, + ) + next_heads_blk_idx = lax.select( + is_last_seq, + heads_blk_idx + 1, + heads_blk_idx, + ) + next_buf_idx = lax.select(cur_buf_idx == 0, 1, 0) + return next_heads_blk_idx, next_seq_idx, next_kv_blk_idx, next_buf_idx + + def flash_attention( + q, # [num_q_per_blk * num_q_heads_per_kv_head, head_dim] + k, # [num_kv_per_blk, head_dim] + v, # [num_kv_per_blk, head_dim] + head_l_ref, # [num_q_per_blk * num_q_heads_per_kv_head, 128] + head_m_ref, # [num_q_per_blk * num_q_heads_per_kv_head, 128] + head_acc_ref, # [num_q_per_blk, num_q_heads_per_kv_head, head_dim] + *, + kv_blk_idx, + ): + assert q.shape == ( + num_q_per_blk * num_q_heads_per_kv_head, + head_dim, + ) + assert ( + k.shape + == v.shape + == ( + num_kv_per_blk, + head_dim, + ) + ) + assert k.dtype == v.dtype + assert ( + head_m_ref.shape + == head_l_ref.shape + == ( + num_q_per_blk * num_q_heads_per_kv_head, + 128, + ) + ) + assert head_acc_ref.shape == ( + num_q_per_blk, + num_q_heads_per_kv_head, + head_dim, + ) + kv_len_start = kv_blk_idx * num_kv_per_blk + + def masked_store(ref, val, start, end, group=1): + iota = lax.broadcasted_iota(jnp.int32, ref.shape, 0) // group + pltpu.store(ref, val, mask=jnp.logical_and(iota >= start, iota < end)) + + def load_with_init(ref, init_val): + return jnp.where( + kv_blk_idx == 0, jnp.full_like(ref, init_val), ref[...] + ) + + # kv lens will be contracting dim, we should mask out the NaNs. + kv_mask = ( + lax.broadcasted_iota(jnp.int32, k.shape, 0) < kv_len - kv_len_start + ) + k = jnp.where(kv_mask, k.astype(jnp.float32), 0).astype(k.dtype) + v = jnp.where(kv_mask, v.astype(jnp.float32), 0).astype(v.dtype) + + qk = ( + jnp.einsum("nd,md->nm", q, k, preferred_element_type=jnp.float32) + * sm_scale + ) + store_start = jnp.maximum(q_start - q_len_start, 0) + store_end = jnp.minimum(q_end - q_len_start, num_q_per_blk) + + row_ids = ( + (kv_len - q_len) + + q_len_start + - q_start + + jax.lax.broadcasted_iota( + jnp.int32, + (num_q_per_blk * num_q_heads_per_kv_head, num_kv_per_blk), + 0, + ) + // num_q_heads_per_kv_head + ) + col_ids = kv_len_start + jax.lax.broadcasted_iota( + jnp.int32, + (num_q_per_blk * num_q_heads_per_kv_head, num_kv_per_blk), + 1, + ) + causal_mask = row_ids < col_ids + if sliding_window is not None: + causal_mask = jnp.logical_or(causal_mask, + row_ids - sliding_window >= col_ids) + if soft_cap is not None: + qk = soft_cap * jnp.tanh(qk / soft_cap) + qk += jnp.where(causal_mask, mask_value, 0.0) + m_curr = jnp.max(qk, axis=1, keepdims=True) + s_curr = jnp.exp(qk - m_curr) + qkv = jnp.dot(s_curr, v, preferred_element_type=jnp.float32) + lm_store_shape = head_m_ref.shape + m_curr = jnp.broadcast_to(m_curr, lm_store_shape) + l_curr = jnp.broadcast_to( + s_curr.sum(axis=1, keepdims=True), lm_store_shape + ) + m_prev = load_with_init(head_m_ref, -jnp.inf) + l_prev = load_with_init(head_l_ref, 0.0) + m_next = jnp.maximum(m_prev, m_curr) + masked_store( + head_m_ref, m_next, store_start, store_end, num_q_heads_per_kv_head + ) + alpha = jnp.exp(m_prev - m_next) + beta = jnp.exp(m_curr - m_next) + l_alpha = alpha * l_prev + l_next = l_alpha + beta * l_curr + l_next_safe = jnp.where(l_next == 0.0, 1.0, l_next) + masked_store( + head_l_ref, + l_next_safe, + store_start, + store_end, + num_q_heads_per_kv_head, + ) + + def broadcast_to_shape(arr, shape): + if arr.shape == shape: + return arr + assert len(arr.shape) == len(shape) + assert arr.shape[0] == shape[0] + assert shape[1] % arr.shape[1] == 0 + # no-op concatenation. + return jnp.concatenate( + [arr for _ in range(shape[1] // arr.shape[1])], axis=1 + ) + + o_curr = load_with_init(head_acc_ref, 0.0).reshape(-1, head_dim) + l_alpha = broadcast_to_shape(l_alpha, qkv.shape) + beta = broadcast_to_shape(beta, qkv.shape) + l_next_safe = broadcast_to_shape(l_next_safe, qkv.shape) + out = lax.div( + l_alpha * o_curr + beta * qkv, + l_next_safe, + ) + masked_store( + head_acc_ref, + out.reshape(head_acc_ref.shape), + store_start, + store_end, + ) + + def is_valid_kv_blk_in_cur_seq(kv_states): + kv_blk_idx, _ = kv_states + return kv_blk_idx * num_kv_per_blk < kv_len + + def compute_with_kv_blk_in_cur_seq(kv_states): + kv_blk_idx, cur_buf_idx = kv_states + next_heads_blk_idx, next_seq_idx, next_kv_blk_idx, next_buf_idx = ( + get_next_prefetch_ids( + heads_blk_idx, cur_seq_idx, kv_blk_idx, cur_buf_idx + ) + ) + + @pl.when(next_heads_blk_idx < num_heads_blks) + def prefetch_next_kv_blk(): + # TODO(jevinjiang): reuse the same buffer if it is already prefetched! + # TODO(jevinjiang): only fetch effective dynamic size to hold kv_len and + # DMA to fixed size buffer! + next_async_copy_kv = create_kv_async_copy_descriptors( + next_heads_blk_idx, next_seq_idx, next_kv_blk_idx, next_buf_idx + ) + next_async_copy_kv.start() + + cur_async_copy_kv = create_kv_async_copy_descriptors( + heads_blk_idx, cur_seq_idx, kv_blk_idx, cur_buf_idx + ) + kv_ref = cur_async_copy_kv.wait().reshape( + num_kv_pages_per_blk * page_size * num_combined_kv_heads_per_blk, + head_dim, + ) + kv_packing = get_dtype_packing(kv_ref.dtype) + # NOTE: kv_packing is divided by 2 because k and v are packed together. + kv_load_step = max(1, kv_packing // 2) + for kv_head_chunk_idx in range(0, num_kv_heads_per_blk, kv_load_step): + k_list, v_list = strided_load_kv( + kv_ref, kv_head_chunk_idx * 2, num_combined_kv_heads_per_blk + ) + for step_idx in range(kv_load_step): + k = k_list[step_idx] + v = v_list[step_idx] + if k_scale is not None: + # NOTE: Conversion between arbitrary data types is not supported. + # That's why it is converted to float32 first. + k = k.astype(jnp.float32) * k_scale + k = k.astype(q_ref.dtype) + if v_scale is not None: + v = v.astype(jnp.float32) * v_scale + v = v.astype(q_ref.dtype) + kv_head_idx = kv_head_chunk_idx + step_idx + q_head_idx = kv_head_idx * num_q_heads_per_kv_head + # TODO(jevinjiang): extra handling for packed type that can start at + # unaligned position! + q = fold_on_2nd_minor( + q_ref[:, q_head_idx : q_head_idx + num_q_heads_per_kv_head, :] + ) + flash_attention( + q, + k, + v, + l_ref.at[kv_head_idx], + m_ref.at[kv_head_idx], + acc_ref.at[ + :, q_head_idx : q_head_idx + num_q_heads_per_kv_head, : + ], + kv_blk_idx=kv_blk_idx, + ) + return kv_blk_idx + 1, next_buf_idx + + _, next_buf_idx = lax.while_loop( + is_valid_kv_blk_in_cur_seq, + compute_with_kv_blk_in_cur_seq, + (0, cur_buf_idx), # (kv_blk_idx, buf_idx) + ) + next_seq_idx = lax.select(q_end <= q_len_end, cur_seq_idx + 1, cur_seq_idx) + done = lax.select(q_end < q_len_end, done, 1) + return done, next_seq_idx, next_buf_idx + + _, seq_idx, buf_idx = lax.while_loop( + is_cur_q_blk_needed, + compute_with_cur_q_blk, + (0, init_seq_idx, init_buf_idx), # (done, seq_idx, buf_idx) + ) + # Reset seq_idx for next kv_heads_blk if run out of seqs! + seq_buf_idx_ref[0] = lax.select(seq_idx < num_seqs, seq_idx, 0) + seq_buf_idx_ref[1] = buf_idx + # acc_ref rows outside the packed query prefix are never written. Mask them + # explicitly so inactive capacity has deterministic zeros rather than + # uninitialized VMEM contents (including the tail of the final partial block). + q_ids = q_len_start + lax.broadcasted_iota(jnp.int32, o_ref.shape, 0) + q_is_active = q_ids < cu_q_lens_ref[num_seqs] + o_ref[...] = jnp.where( + q_is_active, + acc_ref[...].astype(q_ref.dtype), + jnp.zeros(o_ref.shape, dtype=q_ref.dtype), + ) + + + +def get_dtype_packing(dtype): + bits = dtypes.itemsize_bits(dtype) + return 32 // bits + + +def get_min_heads_per_blk( + num_q_heads, num_combined_kv_heads, q_dtype, kv_dtype +): + q_packing = get_dtype_packing(q_dtype) + kv_packing = get_dtype_packing(kv_dtype) + + def can_be_xla_fully_tiled(x, packing): + if x % packing != 0: + return False + x //= packing + return x in (1, 2, 4, 8) or x % 8 == 0 + + # TODO(jevinjiang): support unaligned number of heads! + if not can_be_xla_fully_tiled(num_combined_kv_heads, kv_packing): + raise ValueError( + f"Not implemented: {num_combined_kv_heads=} can not be XLA fully tiled." + ) + assert num_combined_kv_heads % 2 == 0 + num_kv_heads = num_combined_kv_heads // 2 + assert num_q_heads % num_kv_heads == 0 + ratio = num_q_heads // num_kv_heads + # TODO(jevinjiang): we can choose smaller tiling for packed type if large + # second minor tiling is not on. + max_combined_kv_tiling = 8 * kv_packing + min_combined_kv_heads = ( + max_combined_kv_tiling + if num_combined_kv_heads % max_combined_kv_tiling == 0 + else num_combined_kv_heads + ) + min_q_heads = min_combined_kv_heads // 2 * ratio + if can_be_xla_fully_tiled(min_q_heads, q_packing): + return min_q_heads, min_combined_kv_heads + return num_q_heads, num_combined_kv_heads + + +@functools.partial( + jax.jit, + static_argnames=[ + "sm_scale", + "mask_value", + "num_kv_pages_per_block", + "num_queries_per_block", + "vmem_limit_bytes", + "sliding_window", + "soft_cap", + "k_scale", + "v_scale", + ], +) +def ragged_paged_attention( + q: jax.Array, # [max_num_batched_tokens, num_q_heads, head_dim] + # TODO(jevinjiang): create a write_to_kv_cache kernel! + kv_pages: jax.Array, # [total_num_pages, page_size, num_combined_kv_heads, head_dim] + kv_lens: jax.Array, # i32[max_num_seqs] + page_indices: jax.Array, # i32[max_num_seqs, pages_per_seq] + cu_q_lens: jax.Array, # i32[max_num_seqs + 1] + num_seqs: jax.Array, # i32[1] + *, + sm_scale: float = 1.0, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = DEFAULT_MASK_VALUE, + k_scale: float | None = None, + v_scale: float | None = None, + num_kv_pages_per_block: int | None = None, + num_queries_per_block: int | None = None, + vmem_limit_bytes: int | None = None, +): + """Ragged paged attention that supports mixed prefill and decode. + + Args: + q: concatenated all sequences' queries. + kv_pages: paged KV cache. Normally in HBM. + kv_lens: padded kv lengths. Only the first num_seqs values are valid. + page_indices: the first index indicates which page to use in the kv cache + for each sequence. Only the first num_seqs values are valid. + cu_q_lens: the cumulative sum of the effective query lengths. Similar to + kv_lens, only the first num_seqs+1 values are valid. + num_seqs: the dynamic number of sequences. + sm_scale: the softmax scale which will be applied to the Q@K^T. + sliding_window: the sliding window size for the attention. + soft_cap: the logit soft cap for the attention. + mask_value: mask value for causal mask. + k_scale: the scale for the key cache. + v_scale: the scale for the value cache. + num_kv_pages_per_block: number of kv pages to be processed in one flash + attention block in the pallas kernel. + num_queries_per_block: number of kv pages to be processed in one flash + attention block in the pallas kernel. + vmem_limit_bytes: the vmem limit for the pallas kernel. + + Returns: + The output of the attention. + """ + static_validate_inputs( + q, + kv_pages, + kv_lens, + page_indices, + cu_q_lens, + num_seqs, + sm_scale=sm_scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + mask_value=mask_value, + k_scale=k_scale, + v_scale=v_scale, + num_kv_pages_per_block=num_kv_pages_per_block, + num_queries_per_block=num_queries_per_block, + vmem_limit_bytes=vmem_limit_bytes, + ) + if mask_value is None: + mask_value = DEFAULT_MASK_VALUE + num_q_tokens, num_q_heads, head_dim = q.shape + _, page_size, num_combined_kv_heads, _ = kv_pages.shape + assert num_combined_kv_heads % 2 == 0 + num_kv_heads = num_combined_kv_heads // 2 + _, pages_per_seq = page_indices.shape + num_q_heads_per_blk, num_combined_kv_heads_per_blk = get_min_heads_per_blk( + num_q_heads, num_combined_kv_heads, q.dtype, kv_pages.dtype + ) + num_q_per_blk = num_queries_per_block + num_kv_pages_per_blk = num_kv_pages_per_block + if num_q_per_blk is None or num_kv_pages_per_blk is None: + num_kv_pages_per_blk, num_q_per_blk = get_tuned_block_sizes( + q.dtype, + kv_pages.dtype, + num_q_heads_per_blk, + num_combined_kv_heads_per_blk // 2, + head_dim, + page_size, + num_q_tokens, + pages_per_seq, + ) + num_q_heads_per_kv_head = num_q_heads // num_kv_heads + num_q_blks = pl.cdiv(num_q_tokens, num_q_per_blk) + assert num_combined_kv_heads_per_blk % 2 == 0 + num_kv_heads_per_blk = num_combined_kv_heads_per_blk // 2 + assert num_q_heads_per_blk % num_q_heads_per_kv_head == 0 + num_heads_blks = num_q_heads // num_q_heads_per_blk + grid = (num_heads_blks, num_q_blks) + + def q_index_map(heads_blk_idx, q_blk_idx, *_): + return (q_blk_idx, heads_blk_idx, 0) + + q_block_spec = pl.BlockSpec( + (num_q_per_blk, num_q_heads_per_blk, head_dim), + q_index_map, + ) + in_specs = [ + q_block_spec, + pl.BlockSpec(memory_space=pl.ANY), + ] + out_specs = q_block_spec + lm_scratch = pltpu.VMEM( + # TODO(jevinjiang): use 128 instead of 1 is due to Mosaic does not support + # unaligned slicing! + (num_kv_heads_per_blk, num_q_per_blk * num_q_heads_per_kv_head, 128), + jnp.float32, + ) + acc_scratch = pltpu.VMEM( + (num_q_per_blk, num_q_heads_per_blk, head_dim), + jnp.float32, + ) + double_buf_scratch = pltpu.VMEM( + ( + 2, # For double buffering during DMA copies. + num_kv_pages_per_blk, + page_size, + num_combined_kv_heads_per_blk, + head_dim, + ), + kv_pages.dtype, + ) + scratch_shapes = [ + double_buf_scratch, # kv_bufs + pltpu.SemaphoreType.DMA((2,)), # Semaphores for double buffers. + lm_scratch, # l_ref + lm_scratch, # m_ref + acc_scratch, + ] + scalar_prefetches = ( + kv_lens, + page_indices, + cu_q_lens, + jnp.array((0, 0), jnp.int32), # seq_idx, buf_idx + num_seqs, + ) + kernel = pl.pallas_call( + functools.partial( + ragged_paged_attention_kernel, + sm_scale=sm_scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + mask_value=mask_value, + k_scale=k_scale, + v_scale=v_scale, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=len(scalar_prefetches), + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + scratch_shapes=scratch_shapes, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=( + "arbitrary", + "arbitrary", + ), + vmem_limit_bytes=vmem_limit_bytes, + ), + out_shape=jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype), + name="ragged_paged_attention_kernel", + ) + + return kernel(*scalar_prefetches, q, kv_pages) + + +import math + +CONFIG = { + 'name': 'pallas_ragged_paged_attention_llama70b', + 'model': 'Llama-3.1-70B', + 'operator': 'pallas_ragged_paged_attention', + 'max_num_batched_tokens': 4096, + 'max_num_seqs': 64, + 'num_q_heads': 64, + 'num_kv_heads': 8, + 'head_dim': 128, + 'page_size': 16, + 'pages_per_seq': 256, + 'atol': 0.2, + 'rtol': 0.2, +} + +# IMPORTANT: This benchmark tests ONE representative serving scenario, not a +# distribution of all possible RPA inputs. See baseline.py for the scenario +# definition: 48 decodes plus eight chunked prefills fill a 4096-token step. +ACTIVE_Q_LENS = (1,) * 48 + (512,) * 7 + (464,) +ACTIVE_KV_LENS = tuple( + 257 + ((i * 73) % 240) * 16 for i in range(48) +) + (1023, 1535, 2047, 2559, 3071, 3583, 4095, 4095) +NUM_ACTIVE_SEQS = len(ACTIVE_Q_LENS) + +# Inherited from the former uniform workload. Re-run tune_pallas.py for this +# ragged trace before treating the optimized latency as a tuned human baseline. +TUNED_PARAMS = { + 'num_kv_pages_per_block': 64, # autotuned (was 32) + 'num_queries_per_block': 64, # autotuned + 'vmem_limit_bytes': 33554432, # not autotuned (hardware constraint) +} + + +def get_flops(): + H_q = CONFIG['num_q_heads'] + D = CONFIG['head_dim'] + return H_q * 4 * D * sum( + q_len * kv_len + for q_len, kv_len in zip(ACTIVE_Q_LENS, ACTIVE_KV_LENS) + ) + + +def create_inputs(dtype=jnp.bfloat16): + key = jax.random.key(42) + k1, k2, k3 = jax.random.split(key, 3) + max_tokens = CONFIG['max_num_batched_tokens'] + max_seqs = CONFIG['max_num_seqs'] + H_q = CONFIG['num_q_heads'] + H_kv = CONFIG['num_kv_heads'] + D = CONFIG['head_dim'] + page_size = CONFIG['page_size'] + pages_per_seq = CONFIG['pages_per_seq'] + num_combined_kv_heads = 2 * H_kv + total_num_pages = max_seqs * pages_per_seq + q = jax.random.normal(k1, (max_tokens, H_q, D), dtype=dtype) + kv_pages = jax.random.normal( + k2, (total_num_pages, page_size, num_combined_kv_heads, D), dtype=dtype + ) + q_lens = jnp.array(ACTIVE_Q_LENS, dtype=jnp.int32) + kv_lens = jnp.pad( + jnp.array(ACTIVE_KV_LENS, dtype=jnp.int32), + (0, max_seqs - NUM_ACTIVE_SEQS), + ) + active_cu_q_lens = jnp.concatenate( + (jnp.zeros((1,), dtype=jnp.int32), jnp.cumsum(q_lens)) + ) + cu_q_lens = jnp.pad( + active_cu_q_lens, + (0, max_seqs + 1 - active_cu_q_lens.shape[0]), + constant_values=active_cu_q_lens[-1], + ) + page_indices = jax.random.permutation( + k3, total_num_pages, independent=True + ).astype(jnp.int32).reshape(max_seqs, pages_per_seq) + num_seqs = jnp.array([NUM_ACTIVE_SEQS], dtype=jnp.int32) + return q, kv_pages, kv_lens, page_indices, cu_q_lens, num_seqs + + +def workload(q, kv_pages, kv_lens, page_indices, cu_q_lens, num_seqs): + sm_scale = 1.0 / math.sqrt(CONFIG['head_dim']) + return ragged_paged_attention( + q, kv_pages, kv_lens, page_indices, cu_q_lens, num_seqs, + sm_scale=sm_scale, + num_kv_pages_per_block=TUNED_PARAMS['num_kv_pages_per_block'], + num_queries_per_block=TUNED_PARAMS['num_queries_per_block'], + vmem_limit_bytes=TUNED_PARAMS['vmem_limit_bytes'], + ) + + +def benchmark(num_warmup=5, num_iters=100): + """Benchmark and return results dict.""" + inputs = create_inputs() + fn = jax.jit(workload) + for _ in range(num_warmup): + out = fn(*inputs) + out.block_until_ready() + times = [] + for _ in range(num_iters): + t0 = time.perf_counter() + out = fn(*inputs) + out.block_until_ready() + times.append(time.perf_counter() - t0) + times = np.array(times) * 1000 + avg = float(np.mean(times)) + return { + 'name': CONFIG['name'], + 'model': CONFIG['model'], + 'operator': CONFIG['operator'], + 'config': {k: v for k, v in CONFIG.items() if k not in ('name', 'model', 'operator', 'atol', 'rtol')}, + 'time_ms': round(avg, 4), + 'std_ms': round(float(np.std(times)), 4), + 'output_shape': list(out.shape) if hasattr(out, 'shape') else [], + 'status': 'success', + } + + +if __name__ == '__main__': + import json + print(json.dumps(benchmark())) diff --git a/JAXBench/benchmark/level2/8p_GEMM/baseline.py b/JAXBench/benchmark/level2/8p_GEMM/baseline.py new file mode 100644 index 0000000..e4c4752 --- /dev/null +++ b/JAXBench/benchmark/level2/8p_GEMM/baseline.py @@ -0,0 +1,153 @@ +# 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. + +"""Pallas matmul TPU kernel — Llama-3.1-70B FFN dimensions. + +Upstream kernel from jax.experimental.pallas.ops.tpu.matmul, wrapped as a +JAXBench workload with CONFIG / create_inputs / workload. + +See discussion in https://docs.jax.dev/en/latest/pallas/tpu/matmul.html. +""" + +import numpy as np +import time +import functools + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +CONFIG = { + 'name': 'pallas_matmul_llama70b', + 'model': 'Llama-3.1-70B', + 'operator': 'pallas_matmul', + 'M': 8192, + 'K': 8192, + 'N': 28672, + 'atol': 1e-3, + 'rtol': 1e-2, +} + +# Tuned by autotune_block_sizes.py. Re-run to update. +TUNED_PARAMS = {'block_shape': [1024, 2048], 'block_k': 1024} + + +def matmul_kernel(x_tile_ref, y_tile_ref, o_tile_ref, acc_ref): + @pl.when(pl.program_id(2) == 0) + def init(): + acc_ref[...] = jnp.zeros_like(acc_ref) + + acc_ref[...] = acc_ref[...] + jnp.dot( + x_tile_ref[...], + y_tile_ref[...], + preferred_element_type=acc_ref.dtype, + ) + # It is possible to make this conditional but in general this bundle packs + # quite well for a simple matmul kernel + o_tile_ref[...] = acc_ref[...].astype(o_tile_ref.dtype) + + +@functools.partial( + jax.jit, static_argnames=["block_shape", "block_k", "debug", "out_dtype"] +) +def matmul( + x: jax.Array, + y: jax.Array, + *, + block_shape, + block_k: int = 256, + out_dtype: jnp.dtype | None = None, + debug: bool = False, +) -> jax.Array: + if out_dtype is None: + if x.dtype != y.dtype: + # TODO(tlongeri): Maybe we could use a deduction similar to jnp.dot + raise TypeError( + f"Cannot deduce output dtype for different input dtypes: {x.dtype}," + f" {y.dtype}" + ) + out_dtype = x.dtype + acc_dtype = jnp.float32 + if x.dtype in [jnp.int8, jnp.int4, jnp.uint8, jnp.uint4]: + acc_dtype = jnp.int32 + + l, r = block_shape + return pl.pallas_call( + matmul_kernel, + out_shape=jax.ShapeDtypeStruct((x.shape[0], y.shape[1]), out_dtype), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=0, + in_specs=[ + pl.BlockSpec((l, block_k), lambda i, _, k: (i, k)), + pl.BlockSpec((block_k, r), lambda _, j, k: (k, j)), + ], + out_specs=pl.BlockSpec((l, r), lambda i, j, k: (i, j)), + grid=(x.shape[0] // l, y.shape[1] // r, x.shape[1] // block_k), + scratch_shapes=[pltpu.VMEM((l, r), acc_dtype)], + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "parallel", "arbitrary")), + debug=debug, + )(x, y) + + +def get_flops(): + M, K, N = CONFIG['M'], CONFIG['K'], CONFIG['N'] + return 2 * M * K * N + + +def create_inputs(dtype=jnp.bfloat16): + key = jax.random.key(42) + k1, k2 = jax.random.split(key, 2) + M, K, N = CONFIG['M'], CONFIG['K'], CONFIG['N'] + x = jax.random.normal(k1, (M, K), dtype=dtype) + y = jax.random.normal(k2, (K, N), dtype=dtype) * 0.02 + return x, y + + +def workload(x, y): + return matmul(x, y, block_shape=tuple(TUNED_PARAMS['block_shape']), block_k=TUNED_PARAMS['block_k']) + + +def benchmark(num_warmup=5, num_iters=100): + """Benchmark and return results dict.""" + inputs = create_inputs() + fn = jax.jit(workload) + for _ in range(num_warmup): + out = fn(*inputs) + out.block_until_ready() + times = [] + for _ in range(num_iters): + t0 = time.perf_counter() + out = fn(*inputs) + out.block_until_ready() + times.append(time.perf_counter() - t0) + times = np.array(times) * 1000 + avg = float(np.mean(times)) + return { + 'name': CONFIG['name'], + 'model': CONFIG['model'], + 'operator': CONFIG['operator'], + 'config': {k: v for k, v in CONFIG.items() if k not in ('name', 'model', 'operator', 'atol', 'rtol')}, + 'time_ms': round(avg, 4), + 'std_ms': round(float(np.std(times)), 4), + 'output_shape': list(out.shape) if hasattr(out, 'shape') else [], + 'status': 'success', + } + + +if __name__ == '__main__': + import json + print(json.dumps(benchmark())) diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/11p_Megablox_GMM/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/11p_Megablox_GMM/kernel_task.yaml new file mode 100644 index 0000000..a6be5e8 --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/11p_Megablox_GMM/kernel_task.yaml @@ -0,0 +1,37 @@ +task_id: 11p_Megablox_GMM +description: Kernel task for 11p_Megablox_GMM +input_gen_code: |- + def get_inputs(dtype=jnp.bfloat16): + import jax + import jax.numpy as jnp + + CONFIG = { + 'name': 'megablox_gmm_qwen3_235b', + 'model': 'Qwen3-235B-A22B', + 'operator': 'grouped_matmul', + 'num_experts': 128, + 'num_experts_per_tok': 8, + 'emb_dim': 4096, + 'moe_mlp_dim': 1536, + 'seq_len': 4096, + } + key = jax.random.key(42) + k1, k2 = jax.random.split(key, 2) + G = CONFIG['num_experts'] + top_k = CONFIG['num_experts_per_tok'] + K = CONFIG['emb_dim'] + N = CONFIG['moe_mlp_dim'] + S = CONFIG['seq_len'] + M = S * top_k + lhs = jax.random.normal(k1, (M, K), dtype=dtype) + rhs = jax.random.normal(k2, (G, K, N), dtype=dtype) * 0.02 + max_expert_size = M // G + group_sizes = jnp.full((G,), max_expert_size, dtype=jnp.int32) + + dynamic_args = [lhs, rhs, group_sizes] + static_args = [max_expert_size] + + return dynamic_args, static_args + +rtol: 0.01 +atol: 0.01 diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/11p_Megablox_GMM/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/11p_Megablox_GMM/reference.py new file mode 100644 index 0000000..db99fcf --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/11p_Megablox_GMM/reference.py @@ -0,0 +1,627 @@ +# Imports +from collections.abc import Callable +import functools +from typing import Any, Optional + +import jax +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +from jax.experimental.pallas.ops.tpu.megablox import common +import jax.numpy as jnp + +partial = functools.partial + +# Initialization +def get_inputs(dtype=jnp.bfloat16): + CONFIG = { + 'name': 'megablox_gmm_qwen3_235b', + 'model': 'Qwen3-235B-A22B', + 'operator': 'grouped_matmul', + 'num_experts': 128, + 'num_experts_per_tok': 8, + 'emb_dim': 4096, + 'moe_mlp_dim': 1536, + 'seq_len': 4096, + } + key = jax.random.key(42) + k1, k2 = jax.random.split(key, 2) + G = CONFIG['num_experts'] + top_k = CONFIG['num_experts_per_tok'] + K = CONFIG['emb_dim'] + N = CONFIG['moe_mlp_dim'] + S = CONFIG['seq_len'] + M = S * top_k + lhs = jax.random.normal(k1, (M, K), dtype=dtype) + lhs = lhs.astype(jnp.bfloat16).astype(dtype) + rhs = jax.random.normal(k2, (G, K, N), dtype=dtype) * 0.02 + rhs = rhs.astype(jnp.bfloat16).astype(dtype) + max_expert_size = M // G + group_sizes = jnp.full((G,), max_expert_size, dtype=jnp.int32) + + dynamic_args = [lhs, rhs, group_sizes] + static_args = [max_expert_size] + + return dynamic_args, static_args + +# Computation +GroupMetadata = Any + +def _validate_args( + *, + lhs: jnp.ndarray, + rhs: jnp.ndarray, + group_sizes: jnp.ndarray, + expected_rhs_dims: int = 3, +) -> tuple[jnp.ndarray, jnp.ndarray, jnp.dtype]: + if lhs.ndim != 2: + raise ValueError(f"Expected 2-tensor for 'lhs' but got {lhs.ndim}-tensor.") + common.assert_is_supported_dtype(lhs.dtype) + + if rhs.ndim != expected_rhs_dims: + raise ValueError( + f"Expected {expected_rhs_dims}-tensor for 'rhs' but got" + f" {rhs.ndim}-tensor." + ) + common.assert_is_supported_dtype(rhs.dtype) + + if group_sizes.dtype != jnp.int32: + raise ValueError( + f"Expected 32-bit integer 'group_sizes' but got {group_sizes.dtype}." + ) + + return lhs, group_sizes, common.select_input_dtype(lhs, rhs) + +def _calculate_num_tiles(x: int, tx: int) -> int: + tiles, rem = divmod(x, tx) + if rem: + raise ValueError(f"{x} must be divisible by x-dimension tile size ({tx}).") + return tiles + +def _calculate_irregular_num_tiles(x: int, tx: int) -> tuple[int, int]: + tiles, rem = divmod(x, tx) + if rem: + tiles += 1 + return tiles, rem + +def make_group_metadata( + *, + group_sizes: jnp.ndarray, + m: int, + tm: int, + start_group: jnp.ndarray, + num_nonzero_groups: int, + visit_empty_groups: bool = True, +) -> GroupMetadata: + num_groups = group_sizes.shape[0] + end_group = start_group + num_nonzero_groups - 1 + + group_ends = jnp.cumsum(group_sizes) + group_offsets = jnp.concatenate([jnp.zeros(1, dtype=jnp.int32), group_ends]) + + rounded_group_ends = ((group_ends + tm - 1) // tm * tm).astype(jnp.int32) + + group_starts = jnp.concatenate( + [jnp.zeros(1, dtype=jnp.int32), group_ends[:-1]] + ) + rounded_group_starts = group_starts // tm * tm + + rounded_group_sizes = rounded_group_ends - rounded_group_starts + rounded_group_sizes = jnp.where(group_sizes == 0, 0, rounded_group_sizes) + + group_tiles = rounded_group_sizes // tm + + if visit_empty_groups: + group_tiles = jnp.where(group_sizes == 0, 1, group_tiles) + + tiles_m = _calculate_num_tiles(m, tm) + group_ids = jnp.repeat( + jnp.arange(num_groups, dtype=jnp.int32), + group_tiles, + total_repeat_length=tiles_m + num_groups - 1, + ) + + partial_tile_mask = jnp.logical_or( + (group_offsets[:-1] % tm) == 0, group_sizes == 0 + ) + + if visit_empty_groups: + partial_tile_mask = jnp.where(group_sizes == 0, 0, partial_tile_mask) + + partial_tile_ids = jnp.where( + partial_tile_mask, tiles_m, group_offsets[:-1] // tm + ) + + tile_visits = ( + jnp.histogram(partial_tile_ids, bins=tiles_m, range=(0, tiles_m - 1))[0] + + 1 + ) + + m_tile_ids = jnp.repeat( + jnp.arange(tiles_m, dtype=jnp.int32), + tile_visits.astype(jnp.int32), + total_repeat_length=tiles_m + num_groups - 1, + ) + + first_tile_in_shard = (group_ids < start_group).sum() + group_ids = jnp.roll(group_ids, shift=-first_tile_in_shard, axis=0) + m_tile_ids = jnp.roll(m_tile_ids, shift=-first_tile_in_shard, axis=0) + + iota = jnp.arange(num_groups, dtype=jnp.int32) + active_group_mask = jnp.logical_and(iota <= end_group, iota >= start_group) + group_tiles = jnp.where(active_group_mask, group_tiles, 0) + num_tiles = group_tiles.sum() + return (group_offsets, group_ids, m_tile_ids), num_tiles + +def _get_group_size( + *, grid_id: jnp.ndarray, group_metadata: GroupMetadata +) -> jnp.ndarray: + group_offsets, group_ids = group_metadata[:2] + group_id = group_ids[grid_id] + group_start = group_offsets[group_id] + group_end = group_offsets[group_id + 1] + return group_end - group_start + +def _get_store_mask( + *, + grid_id: jnp.ndarray, + group_metadata: GroupMetadata, + tm: int, + tn: int, +) -> jnp.ndarray: + group_offsets, group_ids, m_tile_ids = group_metadata[:3] + group_id = group_ids[grid_id] + group_start = group_offsets[group_id] + group_end = group_offsets[group_id + 1] + m_id = m_tile_ids[grid_id] * tm + iota = jax.lax.broadcasted_iota(jnp.int32, (tm, tn), 0) + m_id + return jnp.logical_and(iota >= group_start, iota < group_end) + +def _zero_uninitialized_memory( + out: jnp.ndarray, + *, + start_group: jnp.ndarray, + num_nonzero_groups: int, + group_metadata: GroupMetadata, +) -> jnp.ndarray: + group_offsets = group_metadata[0] + group_start = group_offsets[start_group] + group_end = group_offsets[start_group + num_nonzero_groups] + valid_mask = jax.lax.broadcasted_iota(jnp.int32, (out.shape[0],), 0) + valid_mask = (valid_mask >= group_start) & (valid_mask < group_end) + return jnp.where(valid_mask[:, None], out, 0) + +LutFn = Callable[[int, int, int], Optional[tuple[int, int, int]]] + +@functools.partial( + jax.jit, + static_argnames=[ + "preferred_element_type", + "tiling", + "transpose_rhs", + "interpret", + ], +) +def gmm( + lhs: jnp.ndarray, + rhs: jnp.ndarray, + group_sizes: jnp.ndarray, + preferred_element_type: jnp.dtype = jnp.float32, + tiling: tuple[int, int, int] | LutFn | None = (128, 128, 128), + group_offset: jnp.ndarray | None = None, + existing_out: jnp.ndarray | None = None, + transpose_rhs: bool = False, + interpret: bool = False, +) -> jnp.ndarray: + + if existing_out is not None: + assert isinstance(existing_out, jax.Array) + expected_dtype = existing_out.dtype + if expected_dtype != preferred_element_type: + raise ValueError( + "Existing output dtype must match preferred_element_type." + ) + if group_offset is None: + group_offset = jnp.array([0], dtype=jnp.int32) + else: + if group_offset.shape: + raise ValueError( + f"group_offset must be a ()-shaped array. Got: {group_offset.shape}." + ) + group_offset = group_offset[None] + num_current_groups = rhs.shape[0] + num_total_groups = group_sizes.shape[0] + lhs, group_sizes, input_dtype = _validate_args( + lhs=lhs, rhs=rhs, group_sizes=group_sizes + ) + + m, k, n = (lhs.shape[0], lhs.shape[1], rhs.shape[2]) + if transpose_rhs: + n = rhs.shape[1] + + if callable(tiling): + tiling = tiling(m, k, n) + + if tiling is None: + raise ValueError(f"No tuned tiling found for (m, k, n) = ({m}, {k}, {n})") + + tm, tk, tn = tiling + tiles_k, k_rem = _calculate_irregular_num_tiles(k, tk) + tiles_n, n_rem = _calculate_irregular_num_tiles(n, tn) + del n_rem + + group_metadata, num_active_tiles = make_group_metadata( + group_sizes=group_sizes, + m=m, + tm=tm, + start_group=group_offset[0], + num_nonzero_groups=rhs.shape[0], + visit_empty_groups=False, + ) + + def kernel( + group_metadata, + group_offset, + lhs, + rhs, + existing_out, + out, + acc_scratch, + ): + group_offsets, group_ids, m_tile_ids = group_metadata + del group_offsets, group_ids, group_offset + + grid_id = pl.program_id(1) + k_i = pl.program_id(2) + + @pl.when(k_i == 0) + def _zero_acc(): + acc_scratch[...] = jnp.zeros_like(acc_scratch) + + if existing_out is not None: + prev_grid_id = jnp.where(grid_id > 0, grid_id - 1, 0) + is_first_processed_group = grid_id == 0 + m_tile_changed = m_tile_ids[grid_id] != m_tile_ids[prev_grid_id] + first_time_seeing_out = jnp.logical_or( + is_first_processed_group, m_tile_changed + ) + + @pl.when(first_time_seeing_out) + def _init_out(): + out[...] = existing_out[...] + + def mask_k_rem(x, *, dim): + if k_rem == 0: + return x + + orig_dtype = x.dtype + iota = lax.broadcasted_iota(jnp.int32, x.shape, dim) + x = x.astype(jnp.float32) + return jnp.where(iota < k_rem, x, 0).astype(orig_dtype) + + def _store_accum(): + mask = _get_store_mask( + grid_id=grid_id, + group_metadata=group_metadata, + tm=tm, + tn=tn, + ) + to_store = acc_scratch[...] + out[...] = jax.lax.select( + mask[...], to_store, out[...].astype(jnp.float32) + ).astype(preferred_element_type) + + def _accum(is_last_k_tile): + if is_last_k_tile: + mask_k_rem_lhs = partial(mask_k_rem, dim=1) + mask_k_rem_rhs = partial(mask_k_rem, dim=int(transpose_rhs)) + else: + mask_k_rem_lhs = lambda x: x + mask_k_rem_rhs = lambda x: x + + if transpose_rhs: + dot_general_dims = (((1,), (1,)), ((), ())) + else: + dot_general_dims = (((1,), (0,)), ((), ())) + + loaded_lhs = lhs[...] + loaded_rhs = rhs[...] + acc_scratch[...] += lax.dot_general( + mask_k_rem_lhs(loaded_lhs).astype(input_dtype), + mask_k_rem_rhs(loaded_rhs).astype(input_dtype), + preferred_element_type=jnp.float32, + dimension_numbers=dot_general_dims, + ) + + if is_last_k_tile: + _store_accum() + + lax.cond( + k_i == tiles_k - 1, + partial(_accum, True), + partial(_accum, False), + ) + + def lhs_transform_indices(n_i, grid_id, k_i, group_metadata, group_offset): + group_offsets, group_ids, m_tile_ids = group_metadata + del n_i, group_offsets, group_ids, group_offset + return m_tile_ids[grid_id], k_i + + def rhs_transform_indices(n_i, grid_id, k_i, group_metadata, group_offset): + group_offsets, group_ids, m_tile_ids = group_metadata + del group_offsets, m_tile_ids + if transpose_rhs: + k_i, n_i = n_i, k_i + + return group_ids[grid_id] - group_offset[0], k_i, n_i + + def out_transform_indices(n_i, grid_id, k_i, group_metadata, group_offset): + group_offsets, group_ids, m_tile_ids = group_metadata + del k_i, group_offsets, group_ids, group_offset + return m_tile_ids[grid_id], n_i + + out_block_spec = pl.BlockSpec((tm, tn), out_transform_indices) + if existing_out is None: + in_out_block_spec: Any = None + input_output_aliases = {} + else: + in_out_block_spec = out_block_spec + input_output_aliases = {6: 0} + + lhs_block_spec = pl.BlockSpec((tm, tk), lhs_transform_indices) + if transpose_rhs: + rhs_block_spec = pl.BlockSpec((None, tn, tk), rhs_transform_indices) + else: + rhs_block_spec = pl.BlockSpec((None, tk, tn), rhs_transform_indices) + + lhs_bytes = lhs.size * lhs.itemsize + rhs_bytes = (k * n) * rhs.itemsize + out_bytes = (m * n) * jnp.dtype(preferred_element_type).itemsize + max_active_tiles = group_metadata[1].size + bytes_accessed = ( + (lhs_bytes * tiles_n) + (rhs_bytes * max_active_tiles) + out_bytes + ) + flops = 2 * m * k * n + cost_estimate = pl.CostEstimate( + flops=flops, bytes_accessed=bytes_accessed, transcendentals=0 + ) + call_gmm = pl.pallas_call( + kernel, + out_shape=jax.ShapeDtypeStruct((m, n), preferred_element_type), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=2, + in_specs=[ + lhs_block_spec, + rhs_block_spec, + in_out_block_spec, + ], + out_specs=out_block_spec, + grid=(tiles_n, num_active_tiles, tiles_k), + scratch_shapes=[pltpu.VMEM((tm, tn), jnp.float32)], + ), + input_output_aliases=input_output_aliases, + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary")), + interpret=interpret, + cost_estimate=cost_estimate, + ) + + out = call_gmm( + group_metadata, + group_offset, + lhs, + rhs, + existing_out, + ) + if existing_out is None and num_current_groups < num_total_groups: + out = _zero_uninitialized_memory( + out, + start_group=group_offset[0], + num_nonzero_groups=rhs.shape[0], + group_metadata=group_metadata, + ) + return out + +@functools.partial( + jax.jit, + static_argnames=[ + "preferred_element_type", + "tiling", + "num_actual_groups", + "interpret", + ], +) +def tgmm( + lhs: jnp.ndarray, + rhs: jnp.ndarray, + group_sizes: jnp.ndarray, + preferred_element_type: jnp.dtype = jnp.float32, + tiling: tuple[int, int, int] | LutFn | None = (128, 128, 128), + group_offset: jnp.ndarray | None = None, + num_actual_groups: int | None = None, + existing_out: jnp.ndarray | None = None, + interpret: bool = False, +) -> jnp.ndarray: + if group_offset is None: + group_offset = jnp.array([0], dtype=jnp.int32) + else: + group_offset = group_offset[None] + lhs, group_sizes, input_dtype = _validate_args( + lhs=lhs, rhs=rhs, group_sizes=group_sizes, expected_rhs_dims=2 + ) + + k, m, n = (lhs.shape[0], lhs.shape[1], rhs.shape[1]) + num_groups = group_sizes.shape[0] + num_actual_groups = ( + num_actual_groups if num_actual_groups is not None else num_groups + ) + + if callable(tiling): + tiling = tiling(m, k, n) + + if tiling is None: + raise ValueError(f"No tuned tiling found for (m, k, n) = ({m}, {k}, {n})") + + tm, tk, tn = tiling + tiles_k, k_rem = _calculate_irregular_num_tiles(k, tk) + del k_rem + tiles_n, n_rem = _calculate_irregular_num_tiles(n, tn) + del n_rem + + group_metadata, num_active_tiles = make_group_metadata( + group_sizes=group_sizes, + m=m, + tm=tm, + start_group=group_offset[0], + num_nonzero_groups=num_actual_groups, + visit_empty_groups=True, + ) + + def kernel( + group_metadata, + group_offset, + lhs, + rhs, + existing_out, + out, + acc_scratch, + ): + grid_id = pl.program_id(2) + group_offsets, group_ids, m_tile_ids = group_metadata + del group_offsets, group_offset, m_tile_ids + + group = group_ids[grid_id] + prev_grid_id = jnp.where(grid_id > 0, grid_id - 1, 0) + prev_group = group_ids[prev_grid_id] + + group_has_changed = jnp.logical_or(grid_id == 0, prev_group != group) + + @pl.when(group_has_changed) + def _zero_acc(): + acc_scratch[...] = jnp.zeros_like(acc_scratch) + + dont_skip = ( + _get_group_size(grid_id=grid_id, group_metadata=group_metadata) > 0 + ) + + @pl.when(dont_skip) + def _do(): + rhs_mask = _get_store_mask( + grid_id=grid_id, + group_metadata=group_metadata, + tm=tm, + tn=tn, + ) + lhs_mask = _get_store_mask( + grid_id=grid_id, + group_metadata=group_metadata, + tm=tm, + tn=tk, + ) + + loaded_lhs = lhs[...] + loaded_rhs = rhs[...] + loaded_lhs = lax.select( + lhs_mask[...], + loaded_lhs.astype(jnp.float32), + jnp.zeros_like(lhs, jnp.float32), + ).swapaxes(0, 1) + loaded_rhs = lax.select( + rhs_mask[...], + loaded_rhs.astype(jnp.float32), + jnp.zeros_like(rhs, jnp.float32), + ) + + acc_scratch[...] += lax.dot( + loaded_lhs.astype(input_dtype), + loaded_rhs.astype(input_dtype), + preferred_element_type=jnp.float32, + ) + + is_end_of_grid = grid_id == (pl.num_programs(2) - 1) + next_grid_id = jnp.where(is_end_of_grid, grid_id, grid_id + 1) + next_group = group_ids[next_grid_id] + + group_is_changing = jnp.logical_or(is_end_of_grid, group != next_group) + + @pl.when(group_is_changing) + def _store_accum(): + to_store = acc_scratch[...] + if existing_out is not None: + to_store += existing_out[...].astype(jnp.float32) + out[...] = to_store.astype(preferred_element_type) + + def lhs_transform_indices(n_i, k_i, grid_id, group_metadata, group_offset): + group_offsets, group_ids, m_tile_ids = group_metadata + del n_i, group_offsets, group_ids, group_offset + return m_tile_ids[grid_id], k_i + + def rhs_transform_indices(n_i, k_i, grid_id, group_metadata, group_offset): + group_offsets, group_ids, m_tile_ids = group_metadata + del k_i, group_offsets, group_ids, group_offset + return m_tile_ids[grid_id], n_i + + def out_transform_indices(n_i, k_i, grid_id, group_metadata, group_offset): + group_offsets, group_ids, m_tile_ids = group_metadata + del group_offsets, m_tile_ids + + return group_ids[grid_id] - group_offset[0], k_i, n_i + + out_block_spec = pl.BlockSpec((None, tk, tn), out_transform_indices) + if existing_out is None: + in_out_block_spec: Any = None + input_output_aliases = {} + else: + in_out_block_spec = out_block_spec + input_output_aliases = {6: 0} + + lhs_block_spec = pl.BlockSpec((tm, tk), lhs_transform_indices) + rhs_block_spec = pl.BlockSpec((tm, tn), rhs_transform_indices) + + lhs_bytes = lhs.size * lhs.itemsize + rhs_bytes = rhs.size * rhs.itemsize + out_bytewidth = jnp.dtype(preferred_element_type).itemsize + out_bytes = (num_actual_groups * k * n) * out_bytewidth + bytes_accessed = ( + (lhs_bytes * tiles_n) + (rhs_bytes * tiles_k) + out_bytes + ) + flops = 2 * m * k * n + cost_estimate = pl.CostEstimate( + flops=flops, bytes_accessed=bytes_accessed, transcendentals=0 + ) + lhs = lhs.swapaxes(0, 1) + call_gmm = pl.pallas_call( + kernel, + out_shape=jax.ShapeDtypeStruct( + (num_actual_groups, k, n), preferred_element_type + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=2, + in_specs=[ + lhs_block_spec, + rhs_block_spec, + in_out_block_spec, + ], + out_specs=out_block_spec, + grid=(tiles_n, tiles_k, num_active_tiles), + scratch_shapes=[pltpu.VMEM((tk, tn), jnp.float32)], + ), + input_output_aliases=input_output_aliases, + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary")), + interpret=interpret, + cost_estimate=cost_estimate, + ) + + out = call_gmm( + group_metadata, + group_offset, + lhs, + rhs, + existing_out, + ) + return out + +def computation(lhs, rhs, group_sizes, max_expert_size): + TUNED_PARAMS = {'tiling': [256, 1024, 1024]} + return gmm(lhs, rhs, group_sizes, tiling=tuple(TUNED_PARAMS['tiling'])) \ No newline at end of file diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/1p_Flash_Attention/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/1p_Flash_Attention/kernel_task.yaml new file mode 100644 index 0000000..88826be --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/1p_Flash_Attention/kernel_task.yaml @@ -0,0 +1,23 @@ +task_id: 1p_Flash_Attention +description: Kernel task for 1p_Flash_Attention +input_gen_code: |- + def get_inputs(): + import jax + import jax.numpy as jnp + + dtype = jnp.bfloat16 + key = jax.random.key(42) + k1, k2, k3 = jax.random.split(key, 3) + B = 4 + S = 4096 + H = 64 + D = 128 + query = jax.random.normal(k1, (B, H, S, D), dtype=dtype) + key_t = jax.random.normal(k2, (B, H, S, D), dtype=dtype) + value = jax.random.normal(k3, (B, H, S, D), dtype=dtype) + dynamic_args = [query, key_t, value] + static_args = [] + return dynamic_args, static_args + +rtol: 0.01 +atol: 0.05 diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/1p_Flash_Attention/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/1p_Flash_Attention/reference.py new file mode 100644 index 0000000..07593be --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/1p_Flash_Attention/reference.py @@ -0,0 +1,1674 @@ +# Imports +import dataclasses +import functools +import math +from typing import Any, NamedTuple +import jax +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +# Initialization +def get_inputs(): + dtype = jnp.bfloat16 + key = jax.random.key(42) + k1, k2, k3 = jax.random.split(key, 3) + B = 4 + S = 4096 + H = 64 + D = 128 + query = jax.random.normal(k1, (B, H, S, D), dtype=dtype) + key_t = jax.random.normal(k2, (B, H, S, D), dtype=dtype) + value = jax.random.normal(k3, (B, H, S, D), dtype=dtype) + dynamic_args = [query, key_t, value] + static_args = [] + return dynamic_args, static_args + +# Computation +def computation(q, k, v): + DEFAULT_MASK_VALUE = -0.7 * float(jnp.finfo(jnp.dtype("float32")).max) + NUM_LANES = 128 + NUM_SUBLANES = 8 + MIN_BLOCK_SIZE = 128 + TRANS_B_DIM_NUMBERS = (((1,), (1,)), ((), ())) + + CONFIG = { + 'name': 'pallas_flash_attention_llama70b', + 'model': 'Llama-3.1-70B', + 'operator': 'pallas_flash_attention', + 'batch': 4, + 'seq_len': 4096, + 'num_heads': 64, + 'head_dim': 128, + 'atol': 2e-3, + 'rtol': 2e-3, + } + + TUNED_PARAMS = { + 'block_q': 2048, + 'block_k_major': 2048, + 'block_k': 1024, + 'block_b': 1, + 'block_q_major_dkv': 128, + 'block_k_major_dkv': 128, + 'block_k_dkv': 128, + 'block_q_dkv': 128, + 'block_k_major_dq': 128, + 'block_k_dq': 128, + 'block_q_dq': 128, + } + + class SegmentIds(NamedTuple): + q: jax.Array + kv: jax.Array + + @dataclasses.dataclass(frozen=True) + class BlockSizes: + block_q: int + block_k_major: int + block_k: int + block_b: int + block_q_major_dkv: int | None = None + block_k_major_dkv: int | None = None + block_k_dkv: int | None = None + block_q_dkv: int | None = None + block_k_major_dq: int | None = None + block_k_dq: int | None = None + block_q_dq: int | None = None + + 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) + if self.block_q_major_dkv is not None and self.block_q_dkv is not None: + verify_major_minor( + "block_q", "_dkv", self.block_q_major_dkv, self.block_q_dkv + ) + if self.block_k_major_dkv is not None and self.block_k_dkv is not None: + verify_major_minor( + "block_k", "_dkv", self.block_k_major_dkv, self.block_k_dkv + ) + if self.block_k_major_dq is not None and self.block_k_dq is not None: + verify_major_minor( + "block_k", "_dq", self.block_k_major_dq, self.block_k_dq + ) + + @property + def has_backward_blocks(self) -> bool: + backward_blocks = ( + self.block_q_major_dkv, + self.block_k_major_dkv, + self.block_q_dkv, + self.block_k_dkv, + self.block_k_major_dq, + self.block_k_dq, + self.block_q_dq, + ) + return all(b is not None for b in backward_blocks) + + @classmethod + def get_default(cls, batch_size, num_heads, q_seq_len, kv_len, d_model): + del batch_size, num_heads, q_seq_len, kv_len, d_model + return BlockSizes( + block_q=128, + block_k_major=128, + block_k=128, + block_b=1, + block_q_major_dkv=128, + block_k_major_dkv=128, + block_k_dkv=128, + block_q_dkv=128, + block_k_major_dq=128, + block_k_dq=128, + block_q_dq=128, + ) + + @functools.partial( + jax.jit, + static_argnames=[ + "causal", + "sm_scale", + "block_sizes", + "debug", + ], + ) + def flash_attention( + q, + k, + v, + ab=None, + segment_ids=None, + *, + causal: bool = False, + sm_scale: float = 1.0, + block_sizes: BlockSizes | None = None, + 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 + ) + return _flash_attention( + q, k, v, ab, segment_ids, False, causal, sm_scale, block_sizes, debug + ) + + @functools.partial(jax.custom_vjp, nondiff_argnames=("save_residuals", "causal", "sm_scale", "block_sizes", "debug")) + def _flash_attention( + q, + k, + v, + ab, + segment_ids, + save_residuals, + causal, + sm_scale, + block_sizes, + 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, + debug, + ) + + def _flash_attention_fwd( + q, + k, + v, + ab, + segment_ids, + save_residuals, + causal, + sm_scale, + block_sizes, + debug, + ): + if save_residuals: + raise NotImplementedError("Higher-order AD not supported") + o, l, m = _flash_attention( + q, k, v, ab, segment_ids, True, causal, sm_scale, block_sizes, debug + ) + return o, (q, k, v, ab, segment_ids, o, l, m) + + def _flash_attention_bwd( + save_residuals: bool, + causal: bool, + sm_scale: float, + block_sizes: BlockSizes, + debug: bool, + residuals, + do, + ): + if save_residuals: + raise NotImplementedError("Higher-order AD not supported") + (q, k, v, ab, segment_ids, o, l, m) = residuals + if not block_sizes.has_backward_blocks: + raise ValueError( + "Program is being differentiated, but not all backward blocks are" + " specified" + ) + + di = jnp.sum( + o.astype(jnp.float32) * do.astype(jnp.float32), axis=-1 + ) + + dk, dv = _flash_attention_bwd_dkv( + q, + k, + v, + ab, + segment_ids, + l, + m, + do, + di, + block_q_major=block_sizes.block_q_major_dkv, + block_k_major=block_sizes.block_k_major_dkv, + block_k=block_sizes.block_k_dkv, + block_q=block_sizes.block_q_dkv, + sm_scale=sm_scale, + causal=causal, + mask_value=DEFAULT_MASK_VALUE, + debug=debug, + ) + + dq, ds = _flash_attention_bwd_dq( + q, + k, + v, + ab, + segment_ids, + l, + m, + do, + di, + block_q_major=block_sizes.block_q_dq, + block_k_major=block_sizes.block_k_major_dq, + block_k=block_sizes.block_k_dq, + sm_scale=sm_scale, + causal=causal, + mask_value=DEFAULT_MASK_VALUE, + debug=debug, + ) + return dq, dk, dv, ds, None + + _flash_attention.defvjp(fwd=_flash_attention_fwd, bwd=_flash_attention_bwd) + + def below_or_on_diag(r, r_blk_size, c, c_blk_size): + 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 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, + o_tile_ref, + 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] + k = k_tile_ref[ + (*batch_idx, pl.dslice(start_k, block_k), slice(None)) + ] + + s = jax.lax.dot_general( + q, k, TRANS_B_DIM_NUMBERS, preferred_element_type=jnp.float32 + ) + + 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 = jnp.tile( + q_segment_ids_tile_ref[batch_idx[0]], (1, repeats) + ) + kv_segment_ids = kv_segment_ids_tile_ref[ + batch_idx[0], :1, pl.dslice(start_k, 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] + m_next = jnp.maximum(m_prev, m_curr) + + 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 - jnp.tile(m_next, (1, block_k_repeats))) + + alpha = jnp.exp(m_prev - m_next) + + l_corr = alpha * l_prev + + l_next = jnp.sum(p, axis=1)[:, None] + l_corr + + head_dim_repeats, rem = divmod(head_dim, MIN_BLOCK_SIZE) + l_broadcast = lambda l: jnp.tile(l, (1, head_dim_repeats)) + 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) + + 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, + o_tile_ref, + 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] + k = k_tile_ref[batch_idx] + s = jax.lax.dot_general( + q, k, TRANS_B_DIM_NUMBERS, preferred_element_type=jnp.float32 + ) + + 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] + ] + q_segment_ids = jnp.tile( + q_segment_ids, (1, repeats) + ) + kv_segment_ids = kv_segment_ids_tile_ref[batch_idx[0], :1] + 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, + 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) + + 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: + 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 + ) + 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", + ) + ), + 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 + + def _flash_attention_dkv_kernel( + q_tile_ref, + k_tile_ref, + v_tile_ref, + ab_tile_ref, + q_segment_ids_tile_ref, + kv_segment_ids_tile_ref, + l_tile_ref, + m_tile_ref, + do_tile_ref, + di_tile_ref, + dk_tile_ref, + dv_tile_ref, + dk_scratch_ref, + dv_scratch_ref, + *, + sm_scale: float, + causal: bool, + mask_value: float, + q_seq_len: int, + block_q: int, + block_k: int, + ): + _, _, block_q_major, _ = q_tile_ref.shape + _, _, block_k_major, _ = k_tile_ref.shape + + q_seq_index = pl.program_id(axis=3) + kv_seq_index = pl.program_id(axis=2) + + @pl.when(q_seq_index == 0) + def start_new_sequence(): + dk_scratch_ref[:, :] = jnp.zeros(dk_scratch_ref.shape, dk_scratch_ref.dtype) + dv_scratch_ref[:, :] = jnp.zeros(dv_scratch_ref.shape, dv_scratch_ref.dtype) + + def q_body(j, _): + start_q = j * block_q + def k_body(i, _): + start_k = i * block_k + k = k_tile_ref[0, 0, pl.ds(start_k, block_k), :] + v = v_tile_ref[0, 0, pl.ds(start_k, block_k), :] + q = q_tile_ref[0, 0, pl.ds(start_q, block_q), :] + l = l_tile_ref[0, 0, pl.ds(start_q, block_q), :] + m = m_tile_ref[0, 0, pl.ds(start_q, block_q), :] + do = do_tile_ref[0, 0, pl.ds(start_q, block_q), :] + di = di_tile_ref[0, 0, pl.ds(start_q, block_q), :].astype( + jnp.float32 + ) + + capped_logits = lax.dot_general( + q, k, TRANS_B_DIM_NUMBERS, preferred_element_type=jnp.float32 + ) + + if ab_tile_ref is not None: + ab = ab_tile_ref[ + 0, + 0, + pl.dslice(j * block_q, block_q), + pl.dslice(i * block_k, block_k), + ].astype(jnp.float32) + capped_logits += ab + + if sm_scale != 1.0: + capped_logits *= sm_scale + + mask = None + if q_segment_ids_tile_ref is not None: + repeats, rem = divmod(block_k, NUM_LANES) + if rem: + raise NotImplementedError( + ) + q_segment_ids = q_segment_ids_tile_ref[ + 0, pl.ds(start_q, block_q), : + ] + q_segment_ids = jnp.tile( + q_segment_ids, (1, repeats) + ) + kv_segment_ids = kv_segment_ids_tile_ref[ + :, 0, pl.ds(start_k, 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_index * block_q_major + start_q + col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1) + col_ids += kv_seq_index * 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) + ) + + capped_logits = ( + capped_logits + if mask is None + else capped_logits + jnp.where(mask, 0.0, mask_value) + ) + + p = jnp.exp( + capped_logits - jnp.tile(m, (1, block_k // MIN_BLOCK_SIZE)) + ) + p = p * jnp.tile( + 1 / l, (1, block_k // MIN_BLOCK_SIZE) + ) + dv = lax.dot(p.T.astype(do.dtype), do, preferred_element_type=jnp.float32) + dv_scratch_ref[pl.ds(start_k, block_k), :] += dv.astype( + dv_scratch_ref.dtype + ) + + dp = lax.dot_general( + do, v, TRANS_B_DIM_NUMBERS, preferred_element_type=jnp.float32 + ) + ds = (dp - jnp.tile(di, (1, block_k // MIN_BLOCK_SIZE))) * p + + if sm_scale != 1.0: + ds = ds * sm_scale + + dk = lax.dot(ds.T.astype(do.dtype), q, preferred_element_type=jnp.float32) + dk_scratch_ref[pl.ds(start_k, block_k), :] += dk.astype( + dk_scratch_ref.dtype + ) + lax.fori_loop(0, block_k_major // block_k, k_body, None, unroll=True) + + if causal: + should_run = below_or_on_diag( + q_seq_index, block_q_major, kv_seq_index, block_k_major + ) + else: + should_run = True + + @pl.when(should_run) + def run(): + lax.fori_loop(0, block_q_major // block_q, q_body, None, unroll=True) + + @pl.when(q_seq_index == q_seq_len // block_q_major - 1) + def end_of_q_sequence(): + dv_tile_ref[0, 0, :, :] = dv_scratch_ref[...].astype(dv_tile_ref.dtype) + dk_tile_ref[0, 0, :, :] = dk_scratch_ref[...].astype(dk_tile_ref.dtype) + + def _flash_attention_bwd_dkv( + q, + k, + v, + ab, + segment_ids, + l, + m, + do, + di, + *, + block_q_major: int | None, + block_q: int | None, + block_k_major: int | None, + block_k: int | None, + sm_scale: float, + causal: bool = False, + mask_value: float = DEFAULT_MASK_VALUE, + debug: bool = False, + ): + batch_size, num_heads, q_seq_len, head_dim = q.shape + _, _, kv_seq_len, _ = k.shape + _verify_block("block_q_major_dkv", "q_seq_len", block_q_major, q_seq_len) + _verify_block("block_q_dkv", "q_seq_len", block_q, q_seq_len) + _verify_block("block_k_major_dkv", "kv_seq_len", block_k_major, kv_seq_len) + _verify_block("block_k_dkv", "kv_seq_len", block_k, kv_seq_len) + + m = jnp.broadcast_to(m[..., None], (*m.shape, MIN_BLOCK_SIZE)) + l = jnp.broadcast_to(l[..., None], (*l.shape, MIN_BLOCK_SIZE)) + di = jnp.broadcast_to(di[..., None], (*di.shape, MIN_BLOCK_SIZE)) + + grid = ( + batch_size, + num_heads, + kv_seq_len // block_k_major, + q_seq_len // block_q_major, + ) + + def qo_index_map(batch_index, head_index, kv_seq_index, q_seq_index): + if causal: + next_q_index = lax.select( + below_or_on_diag( + q_seq_index, block_q_major, kv_seq_index, block_k_major + ), + q_seq_index, + 0, + ) + else: + next_q_index = q_seq_index + + return (batch_index, head_index, next_q_index, 0) + + qo_spec = pl.BlockSpec((1, 1, block_q_major, head_dim), qo_index_map) + assert qo_spec.block_shape is not None + assert q.ndim == len(qo_spec.block_shape) + do_spec = qo_spec + assert do.ndim == len(qo_spec.block_shape) + + def kv_index_map(batch_index, head_index, kv_seq_index, _): + return (batch_index, head_index, kv_seq_index, 0) + + kv_spec = pl.BlockSpec((1, 1, block_k_major, head_dim), kv_index_map) + assert kv_spec.block_shape is not None + assert k.ndim == len(kv_spec.block_shape) + assert v.ndim == len(kv_spec.block_shape) + + def lm_index_map(batch_index, head_index, _, q_seq_index): + return (batch_index, head_index, q_seq_index, 0) + + lm_spec = pl.BlockSpec((1, 1, block_q_major, MIN_BLOCK_SIZE), lm_index_map) + assert lm_spec.block_shape is not None + assert l.ndim == len(lm_spec.block_shape) + assert m.ndim == len(lm_spec.block_shape) + + di_spec = pl.BlockSpec((1, 1, block_q_major, MIN_BLOCK_SIZE), qo_index_map) + assert di_spec.block_shape is not None + assert di.ndim == len(di_spec.block_shape) + + def ab_index_map(batch_index, head_index, kv_seq_index, q_seq_index): + return (batch_index, head_index, q_seq_index, kv_seq_index) + + dab_spec = ( + pl.BlockSpec((1, 1, block_q_major, 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, kv_seq_index, q_seq_index + ): + del head_index + if causal: + next_q_index = lax.select( + below_or_on_diag( + q_seq_index, block_q_major, kv_seq_index, block_k_major + ), + q_seq_index, + 0, + ) + else: + next_q_index = q_seq_index + return (batch_index, next_q_index, 0) + + def kv_segment_ids_index_map(batch_index, head_index, kv_seq_index, _): + del head_index + return (batch_index, 0, kv_seq_index) + + q_segment_ids_spec = pl.BlockSpec( + (1, block_q_major, NUM_LANES), q_segment_ids_index_map + ) + kv_segment_ids_spec = pl.BlockSpec( + (1, 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 = [ + qo_spec, + kv_spec, + kv_spec, + dab_spec, + q_segment_ids_spec, + kv_segment_ids_spec, + lm_spec, + lm_spec, + do_spec, + di_spec, + ] + + out_shapes = [ + jax.ShapeDtypeStruct((batch_size, num_heads, kv_seq_len, head_dim), + k.dtype), + jax.ShapeDtypeStruct((batch_size, num_heads, kv_seq_len, head_dim), + v.dtype), + ] + def dkv_index_map(batch_index, head_index, kv_seq_index, _): + return (batch_index, head_index, kv_seq_index, 0) + + dkv_spec = pl.BlockSpec((1, 1, block_k_major, head_dim), dkv_index_map) + out_specs = [dkv_spec, dkv_spec] + scratch_shapes = [ + pltpu.VMEM((block_k_major, head_dim), jnp.float32), + pltpu.VMEM((block_k_major, head_dim), jnp.float32), + ] + + kernel = functools.partial( + _flash_attention_dkv_kernel, + block_q=block_q, + block_k=block_k, + sm_scale=sm_scale, + causal=causal, + mask_value=mask_value, + q_seq_len=q_seq_len, + ) + name_scope = f"flash_mha_bwd_dkv_{block_q_major=}_{block_q=}_{block_k_major=}_{block_k=}" + with jax.named_scope(name_scope): + dk, dv = 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_shapes, + debug=debug, + compiler_params=pltpu.CompilerParams( + dimension_semantics=( + "parallel", + "parallel", + "parallel", + "arbitrary", + ) + ), + )(q, k, v, ab, q_segment_ids, kv_segment_ids, l, m, do, di) + assert dk.shape == k.shape + assert dv.shape == v.shape + return dk, dv + + def _flash_attention_dq_kernel( + q_tile_ref, + k_tile_ref, + v_tile_ref, + ab_tile_ref, + q_segment_ids_tile_ref, + kv_segment_ids_tile_ref, + l_tile_ref, + m_tile_ref, + do_tile_ref, + di_tile_ref, + dq_tile_ref, + ds_tile_ref, + dq_scratch_ref, + *, + sm_scale: float, + causal: bool, + mask_value: float, + kv_seq_len: int, + block_k: int, + ): + _, _, block_k_major, _ = k_tile_ref.shape + _, _, block_q_major, _ = q_tile_ref.shape + + kv_seq_index = pl.program_id(axis=3) + q_seq_index = pl.program_id(axis=2) + + @pl.when(kv_seq_index == 0) + def start_new_sequence(): + dq_scratch_ref[:, :] = jnp.zeros(dq_scratch_ref.shape, dq_scratch_ref.dtype) + + def body(i, _): + k_slice = pl.ds(i * block_k, block_k) + q = q_tile_ref[0, 0, :, :] + k = k_tile_ref[0, 0, k_slice, :] + v = v_tile_ref[0, 0, k_slice, :] + l = l_tile_ref[0, 0, :, :] + m = m_tile_ref[0, 0, :, :] + do = do_tile_ref[0, 0, :, :] + di = di_tile_ref[0, 0, :].astype(jnp.float32) + + capped_logits = jax.lax.dot_general( + q, k, TRANS_B_DIM_NUMBERS, preferred_element_type=jnp.float32 + ) + + if ab_tile_ref is not None: + ab = ab_tile_ref[0, 0, :, pl.dslice(i * block_k, block_k)].astype( + jnp.float32 + ) + capped_logits += ab + + if sm_scale != 1.0: + capped_logits *= 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 = jnp.tile( + q_segment_ids_tile_ref[0], (1, repeats) + ) + kv_segment_ids = kv_segment_ids_tile_ref[:, 0, k_slice] + mask = jnp.equal(q_segment_ids, kv_segment_ids).astype(jnp.bool_) + + if causal: + mask_shape = (block_q_major, block_k) + row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0) + row_ids += q_seq_index * block_q_major + col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1) + col_ids += kv_seq_index * block_k_major + i * block_k + causal_mask = col_ids <= row_ids + mask = causal_mask if mask is None else jnp.logical_and(mask, causal_mask) + capped_logits = ( + capped_logits + if mask is None + else capped_logits + jnp.where(mask, 0.0, mask_value) + ) + + p = jnp.exp( + capped_logits - jnp.tile(m, (1, block_k // MIN_BLOCK_SIZE)) + ) + p = p * jnp.tile( + 1 / l, (1, block_k // MIN_BLOCK_SIZE) + ) + + dp = jax.lax.dot_general( + do, + v, + TRANS_B_DIM_NUMBERS, + preferred_element_type=jnp.float32, + ) + ds = (dp - jnp.tile(di, (1, block_k // MIN_BLOCK_SIZE))) * p + + if sm_scale != 1.0: + ds = ds * sm_scale + + if ds_tile_ref is not None: + ds_tile_ref[0, 0, :, pl.dslice(i * block_k, block_k)] = ds.astype( + ds_tile_ref.dtype + ) + + dq_scratch_ref[:, :] += lax.dot( + ds.astype(k.dtype), + k, + preferred_element_type=jnp.float32, + ).astype(dq_scratch_ref.dtype) + + if causal: + should_run = below_or_on_diag( + q_seq_index, block_q_major, kv_seq_index, block_k_major + ) + should_not_run = lax.select(should_run, False, True) + else: + should_run = True + should_not_run = False + + @pl.when(should_run) + def run(): + lax.fori_loop(0, block_k_major // block_k, body, None, unroll=True) + + @pl.when(should_not_run) + def zero_out_ds(): + if ds_tile_ref is not None: + ds_tile_ref[...] = jnp.zeros_like(ds_tile_ref) + + @pl.when(kv_seq_index == kv_seq_len // block_k_major - 1) + def end_of_kv_sequence(): + dq_tile_ref[0, 0, :, :] = dq_scratch_ref[...].astype(dq_tile_ref.dtype) + dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref) + + def _flash_attention_bwd_dq( + q, + k, + v, + ab, + segment_ids, + l, + m, + do, + di, + *, + block_q_major: int | None, + block_k_major: int | None, + block_k: int | None, + sm_scale: float, + causal: bool, + mask_value: float, + debug: bool, + ): + batch_size, num_heads, q_seq_len, head_dim = q.shape + _, _, kv_seq_len, _ = k.shape + _verify_block("block_q_dq", "q_seq_len", block_q_major, q_seq_len) + _verify_block("block_k_major_dq", "kv_seq_len", block_k_major, kv_seq_len) + _verify_block("block_k_dq", "block_k", block_k, kv_seq_len) + + m = jnp.broadcast_to(m[..., None], (*m.shape, MIN_BLOCK_SIZE)) + l = jnp.broadcast_to(l[..., None], (*l.shape, MIN_BLOCK_SIZE)) + di = jnp.broadcast_to(di[..., None], (*di.shape, block_k_major)) + + grid = ( + batch_size, + num_heads, + q_seq_len // block_q_major, + kv_seq_len // block_k_major, + ) + + def qo_index_map(batch_index, head_index, q_seq_index, _): + return (batch_index, head_index, q_seq_index, 0) + + qo_spec = pl.BlockSpec((1, 1, block_q_major, head_dim), qo_index_map) + do_spec = qo_spec + + def kv_index_map(batch_index, head_index, q_seq_index, kv_seq_index): + if causal: + next_kv_index = lax.select( + below_or_on_diag( + q_seq_index, block_q_major, 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) + + kv_spec = pl.BlockSpec((1, 1, block_k_major, head_dim), kv_index_map) + assert kv_spec.block_shape is not None + assert k.ndim == len(kv_spec.block_shape) + assert v.ndim == len(kv_spec.block_shape) + + def lm_index_map(batch_index, head_index, q_seq_index, _): + return (batch_index, head_index, q_seq_index, 0) + + lm_spec = pl.BlockSpec((1, 1, block_q_major, MIN_BLOCK_SIZE), lm_index_map) + assert lm_spec.block_shape is not None + assert l.ndim == len(lm_spec.block_shape) + assert m.ndim == len(lm_spec.block_shape) + + di_spec = pl.BlockSpec((1, 1, block_q_major, MIN_BLOCK_SIZE), qo_index_map) + assert di_spec.block_shape is not None + assert di.ndim == len(di_spec.block_shape) + + def ab_index_map(batch_index, head_index, q_seq_index, kv_seq_index): + return (batch_index, head_index, q_seq_index, kv_seq_index) + + dab_spec = ( + pl.BlockSpec((1, 1, block_q_major, 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_major, 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( + (1, block_q_major, NUM_LANES), q_segment_ids_index_map + ) + kv_segment_ids_spec = pl.BlockSpec( + (1, 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 = [ + qo_spec, + kv_spec, + kv_spec, + dab_spec, + q_segment_ids_spec, + kv_segment_ids_spec, + lm_spec, + lm_spec, + do_spec, + di_spec, + ] + + out_shapes = [ + jax.ShapeDtypeStruct(q.shape, q.dtype), + jax.ShapeDtypeStruct(ab.shape, ab.dtype) if ab is not None else None, + ] + dq_spec = pl.BlockSpec((1, 1, block_q_major, head_dim), qo_index_map) + out_specs = [ + dq_spec, + dab_spec, + ] + scratch_shapes = [pltpu.VMEM((block_q_major, head_dim), jnp.float32)] + + kernel = functools.partial( + _flash_attention_dq_kernel, + sm_scale=sm_scale, + causal=causal, + mask_value=mask_value, + block_k=block_k, + kv_seq_len=kv_seq_len, + ) + name_scope = f"flash_mha_bwd_dq_{block_q_major=}_{block_k_major=}_{block_k=}" + with jax.named_scope(name_scope): + dq, ds = 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_shapes, + debug=debug, + compiler_params=pltpu.CompilerParams( + dimension_semantics=( + "parallel", + "parallel", + "parallel", + "arbitrary", + ) + ), + )(q, k, v, ab, q_segment_ids, kv_segment_ids, l, m, do, di) + + return dq, ds + + 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, + ) + + @functools.partial(jax.custom_vjp, nondiff_argnames=("causal", "mask_value", "sm_scale", "save_residuals")) + 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 _mha_reference_fwd( + q, + k, + v, + ab, + segment_ids: SegmentIds | None, + causal: bool, + mask_value: float, + sm_scale: float, + save_residuals: bool, + ): + if save_residuals: + raise NotImplementedError + res = _mha_reference( + q, + k, + v, + ab, + segment_ids, + causal=causal, + mask_value=mask_value, + sm_scale=sm_scale, + save_residuals=True, + ) + assert isinstance(res, tuple) + out, l, m = res + return out, (q, k, v, ab, segment_ids, out, l, m) + + @functools.partial( + jax.jit, + static_argnames=[ + "causal", + "mask_value", + "sm_scale", + ], + ) + def mha_reference_bwd( + q, + k, + v, + ab, + segment_ids: SegmentIds | None, + o, + l, + m, + do, + causal: bool = False, + mask_value: float = DEFAULT_MASK_VALUE, + sm_scale: float = 1.0, + ): + if sm_scale != 1.0: + raise NotImplementedError + + logits = jnp.einsum( + "bhqc,bhkc->bhqk", + q.astype(jnp.float32), + k.astype(jnp.float32), + ) + if ab is not None: + logits += ab + + 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) + + unnormalized = jnp.exp(logits - m[..., None]) + p = unnormalized / l[..., None] + dv = jnp.einsum("bhpt,bhpd->bhtd", p, do.astype(jnp.float32)).astype(v.dtype) + + dp = jnp.einsum( + "bhpd,bhtd->bhpt", do.astype(jnp.float32), v.astype(jnp.float32) + ) + + di = jnp.sum(o.astype(jnp.float32) * do.astype(jnp.float32), axis=-1)[ + ..., None + ] + + ds = (dp - di) * p + dk = jnp.einsum("bhsd,bhst->bhtd", q.astype(jnp.float32), ds).astype(k.dtype) + dq = jnp.einsum("bhst,bhtd->bhsd", ds, k.astype(jnp.float32)).astype(q.dtype) + + dab = ds if ab is not None else None + return dq, dk, dv, dab + + def _mha_reference_bwd( + causal: bool, + mask_value: float, + sm_scale: float, + save_residuals: bool, + residuals, + do, + ): + del save_residuals + q, k, v, ab, segment_ids, o, l, m = residuals + dq, dk, dv, dab = mha_reference_bwd( + q, + k, + v, + ab, + segment_ids, + o, + l, + m, + do, + causal=causal, + mask_value=mask_value, + sm_scale=sm_scale, + ) + return dq, dk, dv, dab, None + + _mha_reference.defvjp(fwd=_mha_reference_fwd, bwd=_mha_reference_bwd) + + 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}" + ) + + sm_scale = 1.0 / math.sqrt(CONFIG['head_dim']) + block_sizes = BlockSizes( + block_q=TUNED_PARAMS['block_q'], + block_k_major=TUNED_PARAMS['block_k_major'], + block_k=TUNED_PARAMS['block_k'], + block_b=TUNED_PARAMS['block_b'], + block_q_major_dkv=TUNED_PARAMS['block_q_major_dkv'], + block_k_major_dkv=TUNED_PARAMS['block_k_major_dkv'], + block_k_dkv=TUNED_PARAMS['block_k_dkv'], + block_q_dkv=TUNED_PARAMS['block_q_dkv'], + block_k_major_dq=TUNED_PARAMS['block_k_major_dq'], + block_k_dq=TUNED_PARAMS['block_k_dq'], + block_q_dq=TUNED_PARAMS['block_q_dq'], + ) + return flash_attention( + q, k, v, causal=True, sm_scale=sm_scale, block_sizes=block_sizes, + ) \ No newline at end of file diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/2p_GQA_Attention/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/2p_GQA_Attention/kernel_task.yaml new file mode 100644 index 0000000..5999897 --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/2p_GQA_Attention/kernel_task.yaml @@ -0,0 +1,32 @@ +task_id: 2p_GQA_Attention +description: Kernel task for 2p_GQA_Attention +input_gen_code: |- + def get_inputs(dtype=jnp.bfloat16): + import jax + import jax.numpy as jnp + from functools import partial + + CONFIG = { + 'name': 'llama3_405b_gqa', + 'model': 'Llama-3.1-405B', + 'operator': 'gqa_attention', + 'batch': 4, + 'seq_len': 4096, + 'num_query_heads': 128, + 'num_kv_heads': 8, + 'head_dim': 128, + 'emb_dim': 16384, + } + key = jax.random.key(42) + k1, k2, k3 = jax.random.split(key, 3) + B, S = CONFIG['batch'], CONFIG['seq_len'] + Hq, Hkv, D = CONFIG['num_query_heads'], CONFIG['num_kv_heads'], CONFIG['head_dim'] + query = jax.random.normal(k1, (B, S, Hq, D), dtype=dtype) + key_t = jax.random.normal(k2, (B, S, Hkv, D), dtype=dtype) + value = jax.random.normal(k3, (B, S, Hkv, D), dtype=dtype) + dynamic_args = [query, key_t, value] + static_args = [] + return dynamic_args, static_args + +rtol: 0.05 +atol: 0.05 \ No newline at end of file diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/2p_GQA_Attention/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/2p_GQA_Attention/reference.py new file mode 100644 index 0000000..52f2ec1 --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/2p_GQA_Attention/reference.py @@ -0,0 +1,2414 @@ +# Imports +from collections.abc import Callable, Mapping +import dataclasses +import enum +import functools +from typing import Any, Literal, NamedTuple, Optional, Union, 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 + +# Initialization +def get_inputs(dtype=jnp.bfloat16): + CONFIG = { + 'name': 'llama3_405b_gqa', + 'model': 'Llama-3.1-405B', + 'operator': 'gqa_attention', + 'batch': 4, + 'seq_len': 4096, + 'num_query_heads': 128, + 'num_kv_heads': 8, + 'head_dim': 128, + 'emb_dim': 16384, + } + key = jax.random.key(42) + k1, k2, k3 = jax.random.split(key, 3) + B, S = CONFIG['batch'], CONFIG['seq_len'] + Hq, Hkv, D = CONFIG['num_query_heads'], CONFIG['num_kv_heads'], CONFIG['head_dim'] + query = jax.random.normal(k1, (B, S, Hq, D), dtype=dtype) + key_t = jax.random.normal(k2, (B, S, Hkv, D), dtype=dtype) + value = jax.random.normal(k3, (B, S, Hkv, D), dtype=dtype) + dynamic_args = [query, key_t, value] + static_args = [] + return dynamic_args, static_args + +# Computation +def computation(query, key, value): + TUNED_PARAMS = { + 'block_q': 2048, + 'block_kv': 2048, + 'block_kv_compute': 1024, + 'q_layout': 1, + 'k_layout': 1, + 'v_layout': 1, + 'head_shards': 1, + 'q_seq_shards': 1, + 'block_q_dkv': None, + 'block_kv_dkv': None, + 'block_kv_dkv_compute': None, + 'block_q_dq': None, + 'block_kv_dq': None, + } + partial = functools.partial + DEFAULT_MASK_VALUE = -0.7 * float(np.finfo(np.dtype("float32")).max) + NUM_LANES = 128 + NUM_SUBLANES = 8 + NN_DIM_NUMBERS = (((1,), (0,)), ((), ())) + NT_DIM_NUMBERS = (((1,), (1,)), ((), ())) + + class SegmentIds(NamedTuple): + q: jax.Array + kv: jax.Array + + SplashCustomReturnType = Union[ + jax.Array, + tuple[jax.Array, tuple[jax.Array,]] + ] + + SplashResidualsType = tuple[ + jax.Array, + jax.Array, + jax.Array, + Optional[SegmentIds], + jax.Array, + jax.Array, + Optional[mask_info_lib.MaskInfo], + Optional[mask_info_lib.MaskInfo], + ] + + MaskFunctionType = Callable[..., jax.Array] + + def get_kernel_name( + block_metadata: Mapping[str, Any], + is_mqa: bool, + save_residuals: bool, + is_segmented: bool, + phase: str, + ) -> str: + assert phase == "dq" or phase == "dkv" or phase == "fwd" + 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}_" + "_".join( + f"{k}={v}" for k, v in sorted(block_metadata.items()) + ) + + @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: jax.Array, + k: jax.Array, + v: jax.Array, + segment_ids: SegmentIds | None, + mask_value: float, + save_residuals: bool, + custom_type: str, + attn_logits_soft_cap: float | None, + ): + return _attention_reference_default( + mask, + q, + k, + v, + segment_ids, + mask_value, + save_residuals, + custom_type, + attn_logits_soft_cap, + ) + + def _attention_reference_default( + mask: jax.Array, + q: jax.Array, + k: jax.Array, + v: jax.Array, + 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: jax.Array, + k: jax.Array, + v: jax.Array, + 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( + 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: jax.Array, + k: jax.Array, + v: jax.Array, + 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) + dv = jnp.einsum("pt,pd->td", p, do).astype(v.dtype) + dp = jnp.einsum("pd,td->pt", do, v.astype(jnp.float32)) + + 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_argnames=( + "mask_value", "save_residuals", "custom_type", "attn_logits_soft_cap") + ) + _attention_reference_custom.defvjp(_attention_reference_custom_fwd, + _attention_reference_custom_bwd) + + def attention_reference_custom( + mask: jax.Array, + q: jax.Array, + k: jax.Array, + v: jax.Array, + 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: + 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:])) + + func = jax.vmap(func, in_axes=(0, 0, None, None, None)) + + 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) + + class QKVLayout(enum.IntEnum): + HEAD_DIM_MINOR = enum.auto() + SEQ_MINOR = enum.auto() + + 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(frozen=True, slots=True) + class BlockSizes: + 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: + object.__setattr__(self, "block_kv_compute", self.block_kv) + if self.block_kv_dkv_compute is None: + object.__setattr__(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): + 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: + assert m_next_ref is None + next_data = i if next_i else j + return ( + next_data, + None, + True, + False, + ) + + 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] + + if data_next_ref.shape[0] == 1: + h = 0 + + 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 | jax.Array, + 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 = mask_ref[:, k_slice] + else: + mask = mask_ref[k_slice, :] + + masks.append( + jnp.bitwise_or(mask, jnp.broadcast_to(should_not_mask, mask.shape)) + ) + if mask_function is not None: + 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 = jnp.tile( + q_sequence_ref[...], (1, repeats) + ) + 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 = q_sequence_ref[:1, :] + q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq)) + + assert q_sequence.shape == k_sequence.shape + computed_mask = mask_function(q_sequence, k_sequence) + if computed_mask.dtype != jnp.dtype(jnp.bool_): + raise ValueError( + "Mask function must return a boolean-valued array, but got:" + f" {computed_mask.dtype}" + ) + masks.append(computed_mask) + + if q_segment_ids_ref is not None: + if k_in_lanes: + kv_ids = kv_segment_ids_ref[: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 = jnp.tile(q_segment_ids_ref[:], (1, repeats)) + 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 = jnp.tile( + kv_segment_ids_ref[k_slice, :], (1, repeats) + ) + q_ids = q_segment_ids_ref[:1, :] + 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( + data_next_ref, + block_mask_ref, + mask_next_ref, + q_ref, + k_ref, + v_ref, + q_segment_ids_ref, + kv_segment_ids_ref, + mask_ref, + q_sequence_ref, + m_scratch_ref, + l_scratch_ref, + o_scratch_ref, + o_ref, + logsumexp_ref=None, + *, + mask_value: float, + grid_width: int, + bq: int, + bkv: int, + bkv_compute: int, + head_dim_v: 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_v_repeats, rem = divmod(head_dim_v, NUM_LANES) + if rem != 0: + raise NotImplementedError( + f"{head_dim_v=} 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) + + global_kv_index, _, 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 = k_ref[slice_k, :] + else: + k = k_ref[:, 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=global_kv_index * 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 - jnp.tile(m_next, (1, bkv_repeats))) + 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 = v_ref[slice_k, :] + else: + v = v_ref[:, slice_k] + v = v.astype(float32) + o_curr = lax.dot_general(s_curr, v, sv_dims) + + alpha_o = jnp.tile(alpha, (1, head_dim_v_repeats)) + 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 = jnp.tile(1.0 / l, (1, head_dim_v_repeats)) + 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, + interpret: bool = False + ) -> SplashCustomReturnType: + num_q_heads, q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + 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] + + partial_mask_blocks = fwd_mask_info.partial_mask_blocks + if ( + partial_mask_blocks is not None + and jnp.dtype(partial_mask_blocks.dtype) != np.bool_ + ): + raise ValueError( + "partial_mask_blocks must be of type np.bool_ but got" + f" {partial_mask_blocks.dtype}" + ) + + 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_qk: + raise ValueError( + f"Expected 'key' head dimension to be: {head_dim_qk}. 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[:-1] != v.shape[:-1]: + raise ValueError( + f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " + "leading dimensions." + ) + + if bkv % bkv_compute: + raise ValueError(f"{bkv=} must be a multiple of {bkv_compute=}.") + if bkv_compute % NUM_LANES: + raise ValueError(f"{bkv_compute=} 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 + 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 + + in_specs = [ + pl.BlockSpec( + from_head_minor((None, bq, head_dim_qk), q_layout), q_index_map + ), + pl.BlockSpec( + from_head_minor( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), k_layout + ), + k_index_map, + ), + pl.BlockSpec( + from_head_minor( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), v_layout + ), + v_index_map, + ), + ] + if segment_ids is not None: + in_specs += [ + pl.BlockSpec((bq, NUM_LANES), q_segment_ids_index_map), + pl.BlockSpec((NUM_SUBLANES, bkv), kv_segment_ids_index_map), + ] + 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((None, bq, bkv), mask_index_map)) + 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((bq, NUM_LANES), q_segment_ids_index_map)) + else: + q_sequence = None + in_specs.append(None) + + num_scalar_prefetch = 3 + + out_shapes = [ + jax.ShapeDtypeStruct((bq, NUM_LANES), jnp.float32), + jax.ShapeDtypeStruct((bq, NUM_LANES), jnp.float32), + jax.ShapeDtypeStruct((bq, head_dim_v), jnp.float32), + jax.ShapeDtypeStruct((num_q_heads, q_seq_len, head_dim_v), q.dtype), + ] + out_specs = [ + pl.BlockSpec((bq, NUM_LANES), lambda h, i, j, *_: (0, 0)), + pl.BlockSpec((bq, NUM_LANES), lambda h, i, j, *_: (0, 0)), + pl.BlockSpec((bq, head_dim_v), lambda h, i, j, *_: (0, 0)), + pl.BlockSpec((None, bq, head_dim_v), out_index_map), + ] + if save_residuals: + out_shapes += [ + jax.ShapeDtypeStruct( + (num_q_heads, q_seq_len, NUM_LANES), jnp.float32 + ), + ] + + def logsumexp_index_map(h, i, *_): + return h, i, 0 + + out_specs += [ + pl.BlockSpec((None, bq, NUM_LANES), logsumexp_index_map), + ] + else: + out_shapes += [None] + out_specs += [None] + + kernel_name = get_kernel_name( + dataclasses.asdict(block_sizes), + is_mqa=is_mqa, + save_residuals=save_residuals, + is_segmented=segment_ids is not None, + phase="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_v=head_dim_v, + 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, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary"), + ), + out_shape=out_shapes, + name=kernel_name, + interpret=interpret, + )( + 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_argnames=( + "save_residuals", "mask_value", "is_mqa", "block_sizes", + "residual_checkpoint_name", "mask_function", "attn_logits_soft_cap", + "interpret") + ) + 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, + interpret: bool = False, + ) -> SplashCustomReturnType: + del dq_mask_info, dkv_mask_info + + return _splash_attention_forward( + fwd_mask_info, + q, + k, + v, + segment_ids, + mask_value=mask_value, + is_mqa=is_mqa, + block_sizes=block_sizes, + residual_checkpoint_name=residual_checkpoint_name, + save_residuals=save_residuals, + mask_function=mask_function, + attn_logits_soft_cap=attn_logits_soft_cap, + interpret=interpret, + ) + + 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, + interpret: bool = False, + ) -> 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, + residual_checkpoint_name=residual_checkpoint_name, + save_residuals=True, + mask_function=mask_function, + attn_logits_soft_cap=attn_logits_soft_cap, + interpret=interpret, + ) + return out, ( + q, + k, + v, + segment_ids, + out, + logsumexp, + dq_mask_info, + dkv_mask_info, + ) + + def _flash_attention_dq_kernel( + data_next_ref, + block_mask_ref, + mask_next_ref, + 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, + 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) + + global_kv_index, _, 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 + 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=global_kv_index * 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, + interpret: bool, + ): + num_q_heads, q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + 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[:-1] != v.shape[:-1]: + raise ValueError( + f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " + "leading dimensions." + ) + + if bkv % NUM_LANES: + raise ValueError(f"{bkv=} must be 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((None, bq, head_dim_v), o_index_map) + + def q_index_map(h, i, *_): + return from_head_minor((h, i, 0), q_layout) + + q_spec = pl.BlockSpec( + from_head_minor((None, bq, head_dim_qk), q_layout), q_index_map + ) + + 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( + from_head_minor( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), k_layout + ), + k_index_map, + ) + + 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( + from_head_minor( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), v_layout + ), + v_index_map, + ) + + 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((None, bq, bkv), mask_index_map) + + def q_segment_ids_index_map(h, i, j, *_): + del h, j + 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((bq, NUM_LANES), q_segment_ids_index_map) + kv_segment_spec = pl.BlockSpec( + (NUM_SUBLANES, bkv), kv_segment_ids_index_map + ) + 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 = o_spec + + def logsumexp_index_map(h, i, *_): + return h, 0, i + + logsumexp = jnp.expand_dims(logsumexp, axis=-2) + logsumexp_spec = pl.BlockSpec((None, 1, bq), logsumexp_index_map) + assert logsumexp.ndim == len(logsumexp_spec.block_shape) + + di = jnp.expand_dims(di, axis=-2) + di_spec = pl.BlockSpec((None, 1, bq), logsumexp_index_map) + 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((bq, NUM_LANES), q_segment_ids_index_map)) + else: + q_sequence = None + in_specs.append(None) + + out_shapes = [ + jax.ShapeDtypeStruct((bq, head_dim_qk), jnp.float32), + jax.ShapeDtypeStruct(q.shape, q.dtype), + ] + out_specs = [ + pl.BlockSpec((bq, head_dim_qk), lambda *_: (0, 0)), + pl.BlockSpec((None, bq, head_dim_qk), lambda h, i, *_: (h, i, 0)), + ] + + 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 + + kernel_name = get_kernel_name( + 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, + save_residuals=False, + is_segmented=segment_ids is not None, + phase="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, + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary", "arbitrary", "arbitrary"), + ), + name=kernel_name, + interpret=interpret, + )( + 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( + data_next_ref, + block_mask_ref, + mask_next_ref, + 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, + 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 + + 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[...] + def _load_kv(ref, layout): + if layout == HEAD_DIM_MINOR: + return ref[slice_k, :] + return ref[:, slice_k].T + k = _load_kv(k_ref, k_layout) + v = _load_kv(v_ref, v_layout) + logsumexp = logsumexp_ref[:1, :] + do = do_ref[...] + di = di_ref[:1, :] + + 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) + dv_scratch_ref[slice_k, :] + dv_scratch_ref[slice_k, :] = 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) + dk_scratch_ref[slice_k, :] + dk_scratch_ref[slice_k, :] = 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: + dq_scratch_ref[...] += dq + else: + 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, + interpret: bool, + ): + num_q_heads, q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + 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[:-1] != v.shape[:-1]: + raise ValueError( + f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " + "leading dimensions." + ) + + 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((None, bq, head_dim_v), o_index_map) + + 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( + from_head_minor((None, bq, head_dim_qk), q_layout), q_index_map + ) + + 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( + from_head_minor( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), + k_layout, + ), + k_index_map, + ) + + 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( + from_head_minor( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), + v_layout, + ), + v_index_map, + ) + + 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((None, None, bq, head_dim_qk), dq_index_map) + 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((bq, head_dim_qk), lambda *_: (0, 0)) + dq_scratch_shape = jax.ShapeDtypeStruct((bq, head_dim_qk), 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 = pl.BlockSpec( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), + dkv_index_map, + ) + + dv_spec = pl.BlockSpec( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), + dkv_index_map, + ) + + 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((None, bkv, bq), mask_index_map) + + 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((NUM_SUBLANES, bq), q_segment_ids_index_map) + kv_segment_spec = pl.BlockSpec((bkv, NUM_LANES), kv_segment_ids_index_map) + 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) + 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((None, NUM_SUBLANES, bq), logsumexp_index_map) + assert logsumexp.ndim == len(logsumexp_spec.block_shape) + + di = jnp.broadcast_to(jnp.expand_dims(di, -2), logsumexp_shape) + di_spec = pl.BlockSpec((None, NUM_SUBLANES, bq), logsumexp_index_map) + 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((NUM_SUBLANES, bq), q_segment_ids_index_map)) + 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_qk), jnp.float32), + jax.ShapeDtypeStruct((bkv, head_dim_v), jnp.float32), + dq_shape, + jax.ShapeDtypeStruct(k.shape, k.dtype), + jax.ShapeDtypeStruct(v.shape, v.dtype), + ] + out_specs = [ + dq_scratch_spec, + pl.BlockSpec((bkv, head_dim_qk), lambda *_: (0, 0)), + pl.BlockSpec((bkv, head_dim_v), lambda *_: (0, 0)), + 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 + + kernel_name = get_kernel_name( + 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, + ), + is_mqa=is_mqa, + save_residuals=False, + is_segmented=segment_ids is not None, + phase="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, + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary", "arbitrary", "arbitrary"), + ), + name=kernel_name, + interpret=interpret, + )( + 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, + interpret: bool, + res: SplashResidualsType, + do: jax.Array, + ) -> tuple[ + mask_info_lib.MaskInfo | None, + mask_info_lib.MaskInfo | None, + mask_info_lib.MaskInfo | None, + jax.Array, + jax.Array, + jax.Array, + SegmentIds | None, + ]: + 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 = jnp.einsum("hsd,hsd->hs", o.astype(jnp.float32), do.astype(jnp.float32)) + 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, + interpret=interpret, + ) + 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, + interpret=interpret, + ) + assert dq is not None + return ( + None, + None, + None, + dq, + dk, + dv, + None, + ) + + _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", + "interpret", + ], + ) + 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, + interpret: bool, + ) -> SplashCustomReturnType: + def _collapse_partial_mask_blocks(mask_info: mask_info_lib.MaskInfo | None): + if mask_info is None or mask_info.partial_mask_blocks is None: + return mask_info + + return mask_info._replace( + partial_mask_blocks=mask_info.partial_mask_blocks.reshape( + -1, *mask_info.partial_mask_blocks.shape[-2:] + ) + ) + + fwd_mask_info = _collapse_partial_mask_blocks(fwd_mask_info) + dq_mask_info = _collapse_partial_mask_blocks(dq_mask_info) + dkv_mask_info = _collapse_partial_mask_blocks(dkv_mask_info) + 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, + interpret=interpret, + ) + + @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): + 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() + partial_mask_blocks_spec = ( + spec if self.fwd_mask_info.is_dynamic_mask else replicated + ) + q_sequence_spec = jax.sharding.PartitionSpec(spec[1]) + mask_info_specs = mask_info_lib.MaskInfo( + 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=partial_mask_blocks_spec + 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 + 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 | jax.Array | 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, + interpret: bool = False, + ): + 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() + + process_mask_fn = ( + mask_info_lib.process_dynamic_mask + if isinstance(mask, jax.Array) + else mask_info_lib.process_mask + ) + + process_mask_dvk_fn = ( + mask_info_lib.process_dynamic_mask_dkv + if isinstance(mask, jax.Array) + else mask_info_lib.process_mask_dkv + ) + + fwd_mask_info, mask_function_fwd = process_mask_fn( + 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 = process_mask_fn( + 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 = process_mask_dvk_fn( + 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, + interpret=interpret, + ) + + 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 + ) + + q = query.transpose(0, 2, 1, 3) + k = key.transpose(0, 2, 1, 3) + v = value.transpose(0, 2, 1, 3) + + B, H_q, S, D = q.shape + q = q * (D ** -0.5) + H_kv = v.shape[1] + heads_per_group = H_q // H_kv + mask = mask_lib.CausalMask(shape=(S, S)) + multi_head_mask = mask_lib.MultiHeadMask([mask] * H_q) + block_sizes = BlockSizes( + block_q=TUNED_PARAMS['block_q'], + block_kv=TUNED_PARAMS['block_kv'], + block_kv_compute=TUNED_PARAMS['block_kv_compute'], + q_layout=QKVLayout(TUNED_PARAMS['q_layout']), + k_layout=QKVLayout(TUNED_PARAMS['k_layout']), + v_layout=QKVLayout(TUNED_PARAMS['v_layout']), + block_q_dkv=TUNED_PARAMS['block_q_dkv'], + block_kv_dkv=TUNED_PARAMS['block_kv_dkv'], + block_kv_dkv_compute=TUNED_PARAMS['block_kv_dkv_compute'], + block_q_dq=TUNED_PARAMS['block_q_dq'], + block_kv_dq=TUNED_PARAMS['block_kv_dq'], + ) + splash_kernel = _make_splash_attention( + multi_head_mask, block_sizes=block_sizes, + is_mqa=False, + head_shards=TUNED_PARAMS['head_shards'], + q_seq_shards=TUNED_PARAMS['q_seq_shards'], + ) + @jax.vmap + def _attend(q_batch, k_batch, v_batch): + k_repeated = jnp.repeat(k_batch, heads_per_group, axis=0) + v_repeated = jnp.repeat(v_batch, heads_per_group, axis=0) + return splash_kernel(q_batch, k_repeated, v_repeated) + out = _attend(q, k, v) + return out.transpose(0, 2, 1, 3) \ No newline at end of file diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/3p_MLA_Attention/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/3p_MLA_Attention/kernel_task.yaml new file mode 100644 index 0000000..41a1391 --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/3p_MLA_Attention/kernel_task.yaml @@ -0,0 +1,90 @@ +task_id: 3p_MLA_Attention +description: Kernel task for 3p_MLA_Attention +input_gen_code: |- + def get_inputs(): + import jax + import jax.numpy as jnp + from functools import partial + + def cdiv(a, b): + assert b != 0 + return (a + b - 1) // b + + def align_to(x, a): + return cdiv(x, a) * a + + def get_dtype_packing(dtype): + bits = jax.dtypes.itemsize_bits(dtype) + return 32 // bits + + CONFIG = { + 'name': 'MLA', + 'batch_size': 128, + 'q_len': 1, + 'kv_len_val': 9216, + 'page_size': 256, + 'symbol': 'd', + } + + key = jax.random.PRNGKey(0) + + num_heads = 128 + lkv_dim = 512 + r_dim = 64 + q_dtype = jnp.bfloat16 + kv_dtype = jnp.bfloat16 + + padded_r_dim = align_to(r_dim, 128) + padded_lkv_dim = align_to(lkv_dim, 128) + padded_kv_dim = padded_lkv_dim + padded_r_dim + packing = get_dtype_packing(kv_dtype) + + def gen_random(k, shape, dtype): + return jax.random.uniform(k, shape, dtype=jnp.float32).astype(dtype) + + total_kv_tokens = CONFIG['batch_size'] * CONFIG['kv_len_val'] + num_pages = cdiv(total_kv_tokens, CONFIG['page_size']) + CONFIG['batch_size'] + + total_q_len = CONFIG['batch_size'] * CONFIG['q_len'] + cu_q_lens_list = [i * CONFIG['q_len'] for i in range(CONFIG['batch_size'] + 1)] + + pages_per_seq = cdiv(CONFIG['kv_len_val'], CONFIG['page_size']) + page_indices_list = [] + page_count = 0 + for _ in range(CONFIG['batch_size']): + num_seq_pages = cdiv(CONFIG['kv_len_val'], CONFIG['page_size']) + indices = list(range(page_count, page_count + num_seq_pages)) + page_indices_list.extend(indices + [-1] * (pages_per_seq - num_seq_pages)) + page_count += num_seq_pages + + total_num_pages = max(num_pages, page_count) + + key, k1, k2, k3, k4, k5 = jax.random.split(key, 6) + ql_nope = gen_random(k1, (total_q_len, num_heads, lkv_dim), q_dtype) + q_pe = gen_random(k2, (total_q_len, num_heads, r_dim), q_dtype) + new_kv_c = gen_random(k3, (total_q_len, lkv_dim), kv_dtype) + new_k_pe = gen_random(k4, (total_q_len, r_dim), kv_dtype) + + cache_kv = gen_random( + k5, + (total_num_pages, CONFIG['page_size'] // packing, packing, padded_kv_dim), + kv_dtype, + ) + + kv_lens = jnp.array([CONFIG['kv_len_val']] * CONFIG['batch_size'], dtype=jnp.int32) + page_indices = jnp.array(page_indices_list, dtype=jnp.int32) + cu_q_lens = jnp.array(cu_q_lens_list, dtype=jnp.int32) + + num_decode_seqs = CONFIG['batch_size'] if CONFIG['q_len'] == 1 else 0 + distribution = jnp.array([num_decode_seqs, num_decode_seqs, CONFIG['batch_size']], dtype=jnp.int32) + + dynamic_args = [ + ql_nope, q_pe, new_kv_c, new_k_pe, cache_kv, kv_lens, + page_indices, cu_q_lens, distribution + ] + static_args = [] + + return dynamic_args, static_args + +rtol: 0.01 +atol: 0.05 diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/3p_MLA_Attention/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/3p_MLA_Attention/reference.py new file mode 100644 index 0000000..1616ce1 --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/3p_MLA_Attention/reference.py @@ -0,0 +1,1161 @@ +# Imports +import functools +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 +from functools import partial + +# Initialization +def get_inputs(): + def cdiv(a, b): + assert b != 0 + return (a + b - 1) // b + + def align_to(x, a): + return cdiv(x, a) * a + + def get_dtype_packing(dtype): + bits = jax.dtypes.itemsize_bits(dtype) + return 32 // bits + + CONFIG = { + 'name': 'config_8', + 'batch_size': 128, + 'q_len': 1, + 'kv_len_val': 9216, + 'page_size': 256, + 'symbol': 'd', + } + + key = jax.random.PRNGKey(0) + + num_heads = 128 + lkv_dim = 512 + r_dim = 64 + q_dtype = jnp.bfloat16 + kv_dtype = jnp.bfloat16 + + padded_r_dim = align_to(r_dim, 128) + padded_lkv_dim = align_to(lkv_dim, 128) + padded_kv_dim = padded_lkv_dim + padded_r_dim + packing = get_dtype_packing(kv_dtype) + + def gen_random(k, shape, dtype): + return jax.random.uniform(k, shape, dtype=jnp.float32).astype(dtype) + + total_kv_tokens = CONFIG['batch_size'] * CONFIG['kv_len_val'] + num_pages = cdiv(total_kv_tokens, CONFIG['page_size']) + CONFIG['batch_size'] + + total_q_len = CONFIG['batch_size'] * CONFIG['q_len'] + cu_q_lens_list = [i * CONFIG['q_len'] for i in range(CONFIG['batch_size'] + 1)] + + pages_per_seq = cdiv(CONFIG['kv_len_val'], CONFIG['page_size']) + page_indices_list = [] + page_count = 0 + for _ in range(CONFIG['batch_size']): + num_seq_pages = cdiv(CONFIG['kv_len_val'], CONFIG['page_size']) + indices = list(range(page_count, page_count + num_seq_pages)) + page_indices_list.extend(indices + [-1] * (pages_per_seq - num_seq_pages)) + page_count += num_seq_pages + + total_num_pages = max(num_pages, page_count) + + key, k1, k2, k3, k4, k5 = jax.random.split(key, 6) + ql_nope = gen_random(k1, (total_q_len, num_heads, lkv_dim), q_dtype) + q_pe = gen_random(k2, (total_q_len, num_heads, r_dim), q_dtype) + new_kv_c = gen_random(k3, (total_q_len, lkv_dim), kv_dtype) + new_k_pe = gen_random(k4, (total_q_len, r_dim), kv_dtype) + + cache_kv = gen_random( + k5, + (total_num_pages, CONFIG['page_size'] // packing, packing, padded_kv_dim), + kv_dtype, + ) + + kv_lens = jnp.array([CONFIG['kv_len_val']] * CONFIG['batch_size'], dtype=jnp.int32) + page_indices = jnp.array(page_indices_list, dtype=jnp.int32) + cu_q_lens = jnp.array(cu_q_lens_list, dtype=jnp.int32) + + num_decode_seqs = CONFIG['batch_size'] if CONFIG['q_len'] == 1 else 0 + distribution = jnp.array([num_decode_seqs, num_decode_seqs, CONFIG['batch_size']], dtype=jnp.int32) + + dynamic_args = [ + ql_nope, q_pe, new_kv_c, new_k_pe, cache_kv, kv_lens, + page_indices, cu_q_lens, distribution + ] + static_args = [] + + return dynamic_args, static_args + +# Computation +def computation(ql_nope, q_pe, new_kv_c, new_k_pe, cache_kv, kv_lens, page_indices, cu_q_lens, distribution): + DEFAULT_MASK_VALUE = -0.7 * float(jnp.finfo(jnp.dtype("float32")).max) + DEFAULT_VMEM_LIMIT_BYTES = 100 * 1024 * 1024 + + def cdiv(a, b): + assert b != 0 + return (a + b - 1) // b + + def align_to(x, a): + return cdiv(x, a) * a + + def get_dtype_bitwidth(dtype): + return jax.dtypes.itemsize_bits(dtype) + + def get_dtype_packing(dtype): + bits = get_dtype_bitwidth(dtype) + return 32 // bits + + def get_kv_cache_shape( + total_num_pages, + page_size, + kv_dim, + kv_dtype, + ): + kv_packing = get_dtype_packing(kv_dtype) + return ( + total_num_pages, + align_to(page_size, kv_packing) // kv_packing, + kv_packing, + align_to(kv_dim, 128), + ) + + def static_validate_inputs( + ql_nope: jax.Array, + q_pe: jax.Array, + new_kv_c: jax.Array, + new_k_pe: jax.Array, + cache_kv: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + distribution: jax.Array, + *, + sm_scale: float = 1.0, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = DEFAULT_MASK_VALUE, + q_scale: float | None = None, + k_scale: float | None = None, + v_scale: float | None = None, + chunk_prefill_size: int | None = None, + num_kv_pages_per_block: int | None = None, + num_queries_per_block: int | None = None, + vmem_limit_bytes: int | None = None, + debug_mode: bool = False, + ): + if len(ql_nope.shape) != 3: + raise ValueError(f"Expected 3D array for {ql_nope.shape=}") + if len(q_pe.shape) != 3: + raise ValueError(f"Expected 3D array for {q_pe.shape=}") + if len(new_kv_c.shape) != 2: + raise ValueError(f"Expected 2D array for {new_kv_c.shape=}") + if len(new_k_pe.shape) != 2: + raise ValueError(f"Expected 2D array for {new_k_pe.shape=}") + + if ql_nope.shape[:2] != q_pe.shape[:2]: + raise ValueError(f"Expected {ql_nope.shape[:2]=} to be equal to {q_pe.shape[:2]=}") + if ql_nope.shape[0] != new_kv_c.shape[0]: + raise ValueError(f"Expected {ql_nope.shape[0]=} to be equal to {new_kv_c.shape[0]=}") + if new_kv_c.shape[0] != new_k_pe.shape[0]: + raise ValueError(f"Expected {new_kv_c.shape[0]=} to be equal to {new_k_pe.shape[0]=}") + if ql_nope.shape[2] != new_kv_c.shape[1]: + raise ValueError(f"Expected {ql_nope.shape[2]=} to be equal to {new_kv_c.shape[1]=}") + if q_pe.shape[2] != new_k_pe.shape[1]: + raise ValueError(f"Expected {q_pe.shape[2]=} to be equal to {new_k_pe.shape[1]=}") + + actual_lkv_dim = ql_nope.shape[2] + actual_r_dim = q_pe.shape[2] + lkv_dim = align_to(actual_lkv_dim, 128) + r_dim = align_to(actual_r_dim, 128) + + _, page_size_per_kv_packing, kv_packing, kv_dim = cache_kv.shape + + if lkv_dim + r_dim != kv_dim: + raise ValueError(f"Expected {lkv_dim=} + {r_dim=} to be equal to {kv_dim=}") + + if not (cache_kv.dtype == new_kv_c.dtype): + raise ValueError(f"Expected {cache_kv.dtype=} to be equal to {new_kv_c.dtype=}.") + if not (cache_kv.dtype == new_k_pe.dtype): + raise ValueError(f"Expected {cache_kv.dtype=} to be equal to {new_k_pe.dtype=}.") + + if not jnp.issubdtype(cache_kv.dtype, jnp.floating): + raise ValueError(f"Expected {cache_kv.dtype=} to be a floating point.") + + if kv_packing != get_dtype_packing(cache_kv.dtype): + raise ValueError(f"{kv_packing=} does not match with {cache_kv.dtype=}") + + if not (jnp.int32 == kv_lens.dtype == page_indices.dtype == cu_q_lens.dtype == distribution.dtype): + raise ValueError(f"Expected int32 dtype for {kv_lens.dtype=}, {page_indices.dtype=}, {cu_q_lens.dtype=}, {distribution.dtype=}") + + if not (len(kv_lens.shape) == len(page_indices.shape) == len(cu_q_lens.shape) == 1): + raise ValueError(f"Expected 1D array for {kv_lens.shape=}, {page_indices.shape=}, {cu_q_lens.shape=}") + + max_num_seqs = kv_lens.shape[0] + num_page_indices = page_indices.shape[0] + if num_page_indices % max_num_seqs != 0: + raise ValueError(f"Expected {num_page_indices=} to be divisible by {max_num_seqs=}.") + if cu_q_lens.shape != (max_num_seqs + 1, ): + raise ValueError(f"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},).") + if distribution.shape != (3, ): + raise ValueError(f"Expected {distribution.shape=} to be (3,).") + + page_size = page_size_per_kv_packing * kv_packing + if page_size % kv_packing != 0: + raise ValueError(f"{page_size=} must be divisible by {kv_packing=}.") + if sliding_window is not None and sliding_window <= 0: + raise ValueError(f"{sliding_window=} must be positive.") + if soft_cap is not None and soft_cap == 0.0: + raise ValueError(f"{soft_cap=} must not be 0.0.") + if chunk_prefill_size is not None and chunk_prefill_size <= 0: + raise ValueError(f"{chunk_prefill_size=} must be positive.") + if num_kv_pages_per_block is not None: + if num_kv_pages_per_block <= 0: + raise ValueError(f"{num_kv_pages_per_block=} must be positive.") + if num_queries_per_block is not None: + if num_queries_per_block <= 0: + raise ValueError(f"{num_queries_per_block=} must be positive.") + if vmem_limit_bytes is not None and vmem_limit_bytes <= 0: + raise ValueError(f"{vmem_limit_bytes=} must be positive.") + + del sm_scale + del mask_value + del q_scale + del k_scale + del v_scale + del debug_mode + + def _mla_ragged_paged_attention_kernel( + kv_lens_ref, + page_indices_ref, + cu_q_lens_ref, + distribution_ref, + sem_ids_ref, + bo_ids_ref, + bkv_update_ids_ref, + ql_nope_hbm_ref, + q_pe_hbm_ref, + new_kv_c_hbm_ref, + new_k_pe_hbm_ref, + cache_kv_hbm_ref, + o_hbm_ref, + updated_cache_kv_hbm_ref, + bkvc_x2_ref, + bkpe_x2_ref, + bq_nope_x2_ref, + bq_rope_x2_ref, + bo_x2_ref, + sems, + l_ref, + m_ref, + acc_ref, + *, + sm_scale: float, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float = DEFAULT_MASK_VALUE, + q_scale: float | None = None, + k_scale: float | None = None, + v_scale: float | None = None, + chunk_prefill_size: int | None = None, + bkv_p, + bq_sz, + debug_mode: bool = False, + ): + assert ql_nope_hbm_ref.shape == o_hbm_ref.shape + nope_dim = ql_nope_hbm_ref.shape[-1] + pe_dim = q_pe_hbm_ref.shape[-1] + assert nope_dim + pe_dim == cache_kv_hbm_ref.shape[-1] + + _, num_q_heads_per_q_packing, q_packing, lkv_dim = ql_nope_hbm_ref.shape + r_dim = q_pe_hbm_ref.shape[-1] + num_q_heads = num_q_heads_per_q_packing * q_packing + total_num_pages, page_size_per_kv_packing, kv_packing, _ = cache_kv_hbm_ref.shape + max_num_seqs = kv_lens_ref.shape[0] + num_page_indices = page_indices_ref.shape[0] + + assert num_page_indices % max_num_seqs == 0 + pages_per_seq = num_page_indices // max_num_seqs + q_dtype = ql_nope_hbm_ref.dtype + kv_dtype = cache_kv_hbm_ref.dtype + assert q_pe_hbm_ref.dtype == q_dtype + assert o_hbm_ref.dtype == q_dtype + assert get_dtype_packing(q_dtype) == q_packing + assert get_dtype_packing(kv_dtype) == kv_packing + assert lkv_dim % 128 == 0 + assert r_dim % 128 == 0 + bkv_sz_per_kv_packing = bkv_p * page_size_per_kv_packing + bkv_sz = bkv_sz_per_kv_packing * kv_packing + page_size = page_size_per_kv_packing * kv_packing + seq_idx = pl.program_id(0) + num_seqs = pl.num_programs(0) + decode_end = distribution_ref[0] + prefill_end = distribution_ref[1] + mixed_end = distribution_ref[2] + + q_start = cu_q_lens_ref[seq_idx] + q_end = cu_q_lens_ref[seq_idx + 1] + q_len = q_end - q_start + kv_len = kv_lens_ref[seq_idx] + + def flash_attention( + ql_nope, + q_pe, + kv_c, + k_pe, + *, + bq_idx, + bkv_idx, + ): + assert len(ql_nope.shape) == 2 + assert len(q_pe.shape) == 2 + assert len(kv_c.shape) == 2 + assert len(k_pe.shape) == 2 + assert ql_nope.shape[0] % num_q_heads == 0 + assert ql_nope.shape[0] == q_pe.shape[0] + assert q_pe.shape[0] % bq_sz == 0 + assert ql_nope.shape[1] == lkv_dim + assert q_pe.shape[1] == r_dim + assert kv_c.shape == (bkv_sz, lkv_dim) + assert k_pe.shape == (bkv_sz, r_dim) + head_l_ref = l_ref.at[:ql_nope.shape[0]] + head_m_ref = m_ref.at[:ql_nope.shape[0]] + head_acc_ref = acc_ref.at[:ql_nope.shape[0]] + + def load_with_init(ref, init_val): + return jnp.where(bkv_idx == 0, jnp.full_like(ref, init_val), ref[...]) + + s_nope = jnp.einsum("nd,md->nm", ql_nope, kv_c, preferred_element_type=jnp.float32) + s_pe = jnp.einsum("nd,md->nm", q_pe, k_pe, preferred_element_type=jnp.float32) + s = s_nope + s_pe + s *= sm_scale + if k_scale is not None: + s *= k_scale + if q_scale is not None: + s *= q_scale + + q_span = (kv_len - q_len + bq_idx * bq_sz + lax.broadcasted_iota(jnp.int32, s.shape, 0) // num_q_heads) + k_span = bkv_idx * bkv_sz + lax.broadcasted_iota(jnp.int32, s.shape, 1) + mask = q_span < k_span + if sliding_window is not None: + mask = jnp.logical_or(mask, q_span - sliding_window >= k_span) + + if soft_cap is not None: + s = soft_cap * jnp.tanh(s / soft_cap) + s = jnp.where(mask, mask_value, s) + s_rowmax = jnp.max(s, axis=1, keepdims=True) + m_prev = load_with_init(head_m_ref, -jnp.inf) + m_curr = jnp.maximum(m_prev, s_rowmax) + head_m_ref[...] = m_curr + p = jnp.exp(s - broadcast_minor(m_curr, s.shape)) + + pv = jnp.einsum("nm,md->nd", p, kv_c, preferred_element_type=jnp.float32) + if v_scale is not None: + pv *= v_scale + + p_rowsum = jnp.sum(p, axis=1, keepdims=True) + exp_m_diff = jnp.exp(m_prev - m_curr) + l_prev = load_with_init(head_l_ref, 0.0) + l_curr = exp_m_diff * l_prev + p_rowsum + head_l_ref[...] = l_curr + o_prev = load_with_init(head_acc_ref, 0.0) + o_curr = broadcast_minor(exp_m_diff, o_prev.shape) * o_prev + pv + head_acc_ref[...] = o_curr + + def _async_copy(src, dst, sem, wait): + if debug_mode: + return + cp = pltpu.make_async_copy(src, dst, sem) + if wait: + cp.wait() + else: + cp.start() + + def _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, *, wait=False): + sem = sems.at[0, bkv_sem_idx] + bkvc_vmem_ref = bkvc_x2_ref.at[bkv_sem_idx] + bkvpe_vmem_ref = bkpe_x2_ref.at[bkv_sem_idx] + + reshaped_cache_hbm_ref = cache_kv_hbm_ref.reshape( + total_num_pages * page_size_per_kv_packing, + *cache_kv_hbm_ref.shape[2:], + ) + + kv_len = kv_lens_ref[seq_idx] + kv_len_start = bkv_idx * bkv_sz + kv_p_start = bkv_idx * bkv_p + + q_start = cu_q_lens_ref[seq_idx] + q_end = cu_q_lens_ref[seq_idx + 1] + q_len = q_end - q_start + + kv_left = kv_len - kv_len_start + kv_left_frm_cache = jnp.maximum(kv_left - q_len, 0) + kv_left_frm_cache_per_kv_packing = cdiv(kv_left_frm_cache, kv_packing) + kv_left_frm_new = kv_left - kv_left_frm_cache + + bkv_sz_frm_cache = jnp.minimum(kv_left_frm_cache, bkv_sz) + bkv_sz_frm_new = jnp.minimum(bkv_sz - bkv_sz_frm_cache, kv_left_frm_new) + bkv_sz_frm_cache_per_kv_packing = cdiv(bkv_sz_frm_cache, kv_packing) + bkv_sz_frm_new_per_kv_packing = cdiv(bkv_sz_frm_new, kv_packing) + page_indices_offset = seq_idx * pages_per_seq + kv_p_start + + new_kv_len_start = q_end - kv_left_frm_new + new_kv_len_start_per_kv_packing = new_kv_len_start // kv_packing + bkv_sz_frm_new_kv_packing_to_fetch = jnp.where( + bkv_sz_frm_new > 0, + cdiv(new_kv_len_start + bkv_sz_frm_new, kv_packing) - new_kv_len_start_per_kv_packing, + 0, + ) + dma_bkv_sz = bkv_sz_frm_cache_per_kv_packing + bkv_sz_frm_new_kv_packing_to_fetch + + if not wait: + wait_update_kv_cache(bkv_sem_idx) + + for i in range(bkv_p): + sz_per_kv_packing = jnp.clip( + kv_left_frm_cache_per_kv_packing - i * page_size_per_kv_packing, + 0, + page_size_per_kv_packing, + ) + page_idx = jnp.minimum(page_indices_offset + i, num_page_indices - 1) + _async_copy( + reshaped_cache_hbm_ref.at[ + pl.ds(page_indices_ref[page_idx] * page_size_per_kv_packing, sz_per_kv_packing), + ..., + :nope_dim, + ], + bkvc_vmem_ref.at[pl.ds(i * page_size_per_kv_packing, sz_per_kv_packing)], + sem, + wait, + ) + _async_copy( + reshaped_cache_hbm_ref.at[ + pl.ds(page_indices_ref[page_idx] * page_size_per_kv_packing, sz_per_kv_packing), + ..., + nope_dim:, + ], + bkvpe_vmem_ref.at[pl.ds(i * page_size_per_kv_packing, sz_per_kv_packing)], + sem, + wait, + ) + + _async_copy( + new_kv_c_hbm_ref.at[pl.ds(new_kv_len_start_per_kv_packing, bkv_sz_frm_new_kv_packing_to_fetch)], + bkvc_vmem_ref.at[pl.ds(bkv_sz_frm_cache_per_kv_packing, bkv_sz_frm_new_kv_packing_to_fetch)], + sem, + wait, + ) + _async_copy( + new_k_pe_hbm_ref.at[pl.ds(new_kv_len_start_per_kv_packing, bkv_sz_frm_new_kv_packing_to_fetch)], + bkvpe_vmem_ref.at[pl.ds(bkv_sz_frm_cache_per_kv_packing, bkv_sz_frm_new_kv_packing_to_fetch)], + sem, + wait, + ) + + else: + dst_kvc = bkvc_vmem_ref.at[pl.ds(0, dma_bkv_sz)] + _async_copy(src=dst_kvc, dst=dst_kvc, sem=sem, wait=True) + dst_kvpe = bkvpe_vmem_ref.at[pl.ds(0, dma_bkv_sz)] + _async_copy(src=dst_kvpe, dst=dst_kvpe, sem=sem, wait=True) + + return kv_len_start + bkv_sz_frm_cache, bkv_sz_frm_new + + def _pack_new_kv(bkv_sem_idx, offset, update_sz): + bkvc_vmem_ref = bkvc_x2_ref.at[bkv_sem_idx] + bkvpe_vmem_ref = bkpe_x2_ref.at[bkv_sem_idx] + + update_kv_packing_iters = cdiv((offset % kv_packing) + update_sz, kv_packing) + kv_packing_offset = offset % kv_packing + new_kv_len_start = q_end - kv_len + offset + new_kv_packing_offset = new_kv_len_start % kv_packing + + token_offset_in_bkv = offset % bkv_sz + kv_packing_idx = token_offset_in_bkv // kv_packing + + shift_amount = kv_packing_offset - new_kv_packing_offset + bits_per_element = get_dtype_bitwidth(bkvc_vmem_ref.dtype) + shift_bits = bits_per_element * (shift_amount % kv_packing) + shift_bits = shift_bits.astype(jnp.uint32) + + kv_packing_idx_new = cdiv(token_offset_in_bkv, kv_packing) + (-shift_amount) // kv_packing + curr_kvc_reg = bkvc_vmem_ref[kv_packing_idx_new, :, :] + curr_kpe_reg = bkvpe_vmem_ref[kv_packing_idx_new, :, :] + next_kvc_reg = bkvc_vmem_ref[kv_packing_idx_new + 1, :, :] + next_kpe_reg = bkvpe_vmem_ref[kv_packing_idx_new + 1, :, :] + + def merge_loop_body(i, vals): + ( + kv_packing_idx, + kv_packing_idx_new, + curr_kvc_reg, + curr_kpe_reg, + next_kvc_reg, + next_kpe_reg, + ) = vals + curr_kvc_reg_u32 = pltpu.bitcast(curr_kvc_reg, jnp.uint32) + curr_kpe_reg_u32 = pltpu.bitcast(curr_kpe_reg, jnp.uint32) + next_kvc_reg_u32 = pltpu.bitcast(next_kvc_reg, jnp.uint32) + next_kpe_reg_u32 = pltpu.bitcast(next_kpe_reg, jnp.uint32) + + shifted_kvc_u32 = lax.bitwise_or( + lax.shift_right_logical(curr_kvc_reg_u32, 32 - shift_bits), + lax.shift_left(next_kvc_reg_u32, shift_bits), + ) + shifted_kpe_u32 = lax.bitwise_or( + lax.shift_right_logical(curr_kpe_reg_u32, 32 - shift_bits), + lax.shift_left(next_kpe_reg_u32, shift_bits), + ) + + rotated_kvc_u32 = lax.select(shift_bits == 0, curr_kvc_reg_u32, shifted_kvc_u32) + rotated_kpe_u32 = lax.select(shift_bits == 0, curr_kpe_reg_u32, shifted_kpe_u32) + + next_kvc_reg_shifted = pltpu.bitcast(rotated_kvc_u32, next_kvc_reg.dtype) + next_kpe_reg_shifted = pltpu.bitcast(rotated_kpe_u32, next_kpe_reg.dtype) + + offset_in_word = i * kv_packing + lax.broadcasted_iota(dtype=jnp.int32, shape=[kv_packing, lkv_dim], dimension=0) + kvc_mask = jnp.logical_and( + offset_in_word >= kv_packing_offset, + offset_in_word < kv_packing_offset + update_sz, + ) + updated_kvc_reg = lax.select( + kvc_mask, + next_kvc_reg_shifted, + bkvc_vmem_ref[kv_packing_idx, :, :], + ) + offset_in_word_pe = i * kv_packing + lax.broadcasted_iota(dtype=jnp.int32, shape=[kv_packing, r_dim], dimension=0) + kpe_mask = jnp.logical_and( + offset_in_word_pe >= kv_packing_offset, + offset_in_word_pe < kv_packing_offset + update_sz, + ) + updated_kpe_reg = lax.select( + kpe_mask, + next_kpe_reg_shifted, + bkvpe_vmem_ref[kv_packing_idx, :, :], + ) + + bkvc_vmem_ref[kv_packing_idx, :, :] = updated_kvc_reg + bkvpe_vmem_ref[kv_packing_idx, :, :] = updated_kpe_reg + + kv_packing_idx += 1 + kv_packing_idx_new += 1 + curr_kvc_reg = next_kvc_reg + curr_kpe_reg = next_kpe_reg + next_kvc_reg = bkvc_vmem_ref[kv_packing_idx_new + 1, :, :] + next_kpe_reg = bkvpe_vmem_ref[kv_packing_idx_new + 1, :, :] + return ( + kv_packing_idx, + kv_packing_idx_new, + curr_kvc_reg, + curr_kpe_reg, + next_kvc_reg, + next_kpe_reg, + ) + + lax.fori_loop( + 0, + update_kv_packing_iters, + merge_loop_body, + ( + kv_packing_idx, + kv_packing_idx_new, + curr_kvc_reg, + curr_kpe_reg, + next_kvc_reg, + next_kpe_reg, + ), + ) + + def _update_kv_cache( + seq_idx, + bkv_sem_idx, + offset, + update_sz, + *, + wait=False, + ): + sem = sems.at[3, bkv_sem_idx] + bkvc_vmem_ref = bkvc_x2_ref.at[bkv_sem_idx] + bkvpe_vmem_ref = bkpe_x2_ref.at[bkv_sem_idx] + + update_kv_packing_iters = cdiv((offset % kv_packing) + update_sz, kv_packing) + + cache_kv_hbm_shape = updated_cache_kv_hbm_ref.shape + reshaped_cache_kv_hbm_ref = updated_cache_kv_hbm_ref.reshape( + cache_kv_hbm_shape[0] * cache_kv_hbm_shape[1], + *cache_kv_hbm_shape[2:], + ) + + if not wait: + kv_p_start = offset // page_size + kv_p_end = cdiv(offset + update_sz, page_size) + start_word_in_page = (offset % page_size) // kv_packing + start_word_in_vmem = (offset % bkv_sz) // kv_packing + words_to_transfer = update_kv_packing_iters + page_indices_offset = seq_idx * pages_per_seq + kv_p_start + + def loop_body(i, states): + curr_word_in_page, words_to_transfer, curr_word_in_vmem = states + sz_words = jnp.minimum(page_size_per_kv_packing - curr_word_in_page, words_to_transfer) + page_idx = page_indices_ref[page_indices_offset + i] + + _async_copy( + bkvc_vmem_ref.at[pl.ds(curr_word_in_vmem, sz_words)], + reshaped_cache_kv_hbm_ref.at[ + pl.ds(page_idx * page_size_per_kv_packing + curr_word_in_page, sz_words), + ..., + :nope_dim, + ], + sem, + wait=False, + ) + _async_copy( + bkvpe_vmem_ref.at[pl.ds(curr_word_in_vmem, sz_words)], + reshaped_cache_kv_hbm_ref.at[ + pl.ds(page_idx * page_size_per_kv_packing + curr_word_in_page, sz_words), + ..., + nope_dim:, + ], + sem, + wait=False, + ) + return 0, words_to_transfer - sz_words, curr_word_in_vmem + sz_words + + lax.fori_loop( + 0, + kv_p_end - kv_p_start, + loop_body, + ( + start_word_in_page, + words_to_transfer, + start_word_in_vmem, + ), + unroll=False, + ) + else: + dma_sz_words = update_kv_packing_iters + dst_kv = bkvc_vmem_ref.at[pl.ds(0, dma_sz_words)] + _async_copy(src=dst_kv, dst=dst_kv, sem=sem, wait=True) + dst_kv = bkvpe_vmem_ref.at[pl.ds(0, dma_sz_words)] + _async_copy(src=dst_kv, dst=dst_kv, sem=sem, wait=True) + + def _fetch_bq(seq_idx, bq_idx, bq_sem_idx, *, wait=False): + sem = sems.at[1, bq_sem_idx] + bq_nope_vmem_ref = bq_nope_x2_ref.at[bq_sem_idx] + bq_rope_vmem_ref = bq_rope_x2_ref.at[bq_sem_idx] + + q_len_start = cu_q_lens_ref[seq_idx] + bq_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy( + ql_nope_hbm_ref.at[pl.ds(q_len_start, sz)], + bq_nope_vmem_ref.at[pl.ds(0, sz)], + sem, + wait, + ) + + _async_copy( + q_pe_hbm_ref.at[pl.ds(q_len_start, sz)], + bq_rope_vmem_ref.at[pl.ds(0, sz)], + sem, + wait, + ) + + def _send_bo(seq_idx, bo_idx, bo_sem_idx, *, wait=False): + sem = sems.at[2, bo_sem_idx] + vmem_ref = bo_x2_ref.at[bo_sem_idx] + q_len_start = cu_q_lens_ref[seq_idx] + bo_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy( + vmem_ref.at[pl.ds(0, sz)], + o_hbm_ref.at[pl.ds(q_len_start, sz)], + sem, + wait, + ) + + def start_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx): + return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx) + + def wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx): + return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, wait=True) + + def start_fetch_bq(seq_idx, bq_idx, bq_sem_idx): + return _fetch_bq(seq_idx, bq_idx, bq_sem_idx) + + def wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx): + return _fetch_bq(seq_idx, bq_idx, bq_sem_idx, wait=True) + + def start_send_bo(seq_idx, bo_idx, bo_sem_idx): + bo_ids_ref[bo_sem_idx] = seq_idx + bo_ids_ref[bo_sem_idx + 2] = bo_idx + _send_bo(seq_idx, bo_idx, bo_sem_idx) + + def wait_send_bo(bo_sem_idx): + old_seq_idx = bo_ids_ref[bo_sem_idx] + old_bo_idx = bo_ids_ref[bo_sem_idx + 2] + + @pl.when(jnp.logical_and(0 <= old_seq_idx, old_seq_idx <= seq_idx)) + def _(): + _send_bo(old_seq_idx, old_bo_idx, bo_sem_idx, wait=True) + + def start_update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz): + bkv_update_ids_ref[bkv_sem_idx] = seq_idx + bkv_update_ids_ref[bkv_sem_idx + 2] = offset + bkv_update_ids_ref[bkv_sem_idx + 4] = update_sz + _update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz) + + def wait_update_kv_cache(bkv_sem_idx): + update_sz = bkv_update_ids_ref[bkv_sem_idx + 4] + + @pl.when(update_sz > 0) + def _(): + seq_idx = bkv_update_ids_ref[bkv_sem_idx] + offset = bkv_update_ids_ref[bkv_sem_idx + 2] + bkv_update_ids_ref[bkv_sem_idx + 4] = 0 + _update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz, wait=True) + + def load_bq(bq_sem_idx, *, actual_bq_sz=bq_sz): + q_nope_ref = (bq_nope_x2_ref.bitcast(jnp.uint32).at[bq_sem_idx].reshape(bq_sz * num_q_heads_per_q_packing, lkv_dim)) + q_nope_vec = pltpu.bitcast(q_nope_ref[:actual_bq_sz * num_q_heads_per_q_packing], q_dtype).reshape(actual_bq_sz * num_q_heads, lkv_dim) + q_rope_ref = (bq_rope_x2_ref.bitcast(jnp.uint32).at[bq_sem_idx].reshape(bq_sz * num_q_heads_per_q_packing, r_dim)) + q_rope_vec = pltpu.bitcast(q_rope_ref[:actual_bq_sz * num_q_heads_per_q_packing], q_dtype).reshape(actual_bq_sz * num_q_heads, r_dim) + return q_nope_vec, q_rope_vec + + def load_bkv(bkv_sem_idx, *, bkvc_mask, bkpe_mask): + bkvc_ref = (bkvc_x2_ref.bitcast(jnp.uint32).at[bkv_sem_idx, :bkv_sz_per_kv_packing].reshape(bkv_sz_per_kv_packing, lkv_dim)) + bkvc_vec = pltpu.bitcast(bkvc_ref[...], kv_dtype).reshape(bkv_sz, lkv_dim) + bkvc_vec = lax.select(bkvc_mask, bkvc_vec, jnp.zeros_like(bkvc_vec)) + + bkpe_ref = (bkpe_x2_ref.bitcast(jnp.uint32).at[bkv_sem_idx, :bkv_sz_per_kv_packing].reshape(bkv_sz_per_kv_packing, r_dim)) + bkpe_vec = pltpu.bitcast(bkpe_ref[...], kv_dtype).reshape(bkv_sz, r_dim) + bkpe_vec = lax.select(bkpe_mask, bkpe_vec, jnp.zeros_like(bkpe_vec)) + + return bkvc_vec, bkpe_vec + + def broadcast_minor(src, shape): + if src.shape == shape: + return src + assert src.shape[:-1] == shape[:-1] + assert src.shape[-1] % 128 == 0 + target_minor = align_to(shape[-1], src.shape[-1]) + return jnp.concatenate([src for _ in range(target_minor // src.shape[-1])], axis=-1)[..., :shape[-1]] + + def process(static_q_len=None): + num_bkv = cdiv(kv_len, bkv_sz) + if static_q_len is None: + actual_bq_sz = bq_sz + num_bq = cdiv(q_len, actual_bq_sz) + else: + actual_bq_sz = min(bq_sz, static_q_len) + num_bq = cdiv(static_q_len, actual_bq_sz) + + def get_next_bq_ids(seq_idx, bq_idx, bq_sem_idx): + next_bq_idx = bq_idx + 1 + is_last_bq = next_bq_idx == num_bq + next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx) + next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx) + next_bq_sem_idx = lax.select(bq_sem_idx == 0, 1, 0) + return next_seq_idx, next_bq_idx, next_bq_sem_idx + + def get_next_bkv_ids(seq_idx, bq_idx, bkv_idx, bkv_sem_idx): + next_bkv_idx = bkv_idx + 1 + is_last_bkv = next_bkv_idx == num_bkv + next_bkv_idx = lax.select(is_last_bkv, 0, next_bkv_idx) + next_bq_idx = lax.select(is_last_bkv, bq_idx + 1, bq_idx) + is_last_bq = next_bq_idx == num_bq + next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx) + next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx) + next_bkv_sem_idx = lax.select(bkv_sem_idx == 0, 1, 0) + return next_seq_idx, next_bq_idx, next_bkv_idx, next_bkv_sem_idx + + def compute_with_bq(bq_idx, _): + bq_sem_idx = sem_ids_ref[0] + next_seq_idx, next_bq_idx, next_bq_sem_idx = get_next_bq_ids(seq_idx, bq_idx, bq_sem_idx) + + @pl.when(next_seq_idx < num_seqs) + def prefetch_next_bq(): + sem_ids_ref[0] = next_bq_sem_idx + start_fetch_bq(next_seq_idx, next_bq_idx, next_bq_sem_idx) + + def compute_with_bkv(bkv_idx, _): + assert bkv_sz % kv_packing == 0 + actual_bkv_sz = jnp.minimum(bkv_sz, kv_len - bkv_idx * bkv_sz) + bkvc_shape = (bkv_sz, lkv_dim) + bkvc_mask = (lax.broadcasted_iota(jnp.int32, bkvc_shape, 0) < actual_bkv_sz) + bkpe_shape = (bkv_sz, r_dim) + bkpe_mask = (lax.broadcasted_iota(jnp.int32, bkpe_shape, 0) < actual_bkv_sz) + + bkv_sem_idx = sem_ids_ref[1] + next_seq_idx, _, next_bkv_idx, next_bkv_sem_idx = get_next_bkv_ids(seq_idx, bq_idx, bkv_idx, bkv_sem_idx) + + @pl.when(next_seq_idx < num_seqs) + def prefetch_next_bkv(): + sem_ids_ref[1] = next_bkv_sem_idx + start_fetch_bkv(next_seq_idx, next_bkv_idx, next_bkv_sem_idx) + + @pl.when(bkv_idx == 0) + def wait_cur_bq(): + wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx) + + offset, update_sz = wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx) + + @pl.when(update_sz > 0) + def pack_new_kv(): + _pack_new_kv(bkv_sem_idx, offset, update_sz) + + @pl.when(jnp.logical_and(update_sz > 0, bq_idx == 0)) + def update_cur_bkv_to_cache(): + start_update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz) + + bkvc, bkpe = load_bkv(bkv_sem_idx, bkvc_mask=bkvc_mask, bkpe_mask=bkpe_mask) + bq_nope_vec, bq_pe_vec = load_bq(bq_sem_idx, actual_bq_sz=actual_bq_sz) + + if debug_mode: + return + + flash_attention( + bq_nope_vec, + bq_pe_vec, + bkvc, + bkpe, + bq_idx=bq_idx, + bkv_idx=bkv_idx, + ) + + lax.fori_loop(0, num_bkv, compute_with_bkv, None, unroll=False) + + acc = acc_ref[...] + l = broadcast_minor(l_ref[...], acc.shape) + out = (lax.div(acc, l) if q_dtype == jnp.float32 else (acc * pl.reciprocal(l, approx=True)).astype(q_dtype)) + + bo_sem_idx = sem_ids_ref[2] + sem_ids_ref[2] = lax.select(bo_sem_idx == 0, 1, 0) + wait_send_bo(bo_sem_idx) + + bo_x2_ref.at[bo_sem_idx].bitcast(jnp.int32).reshape( + bq_sz * num_q_heads_per_q_packing, + lkv_dim, + )[...] = pltpu.bitcast(out, jnp.int32) + + start_send_bo(seq_idx, bq_idx, bo_sem_idx) + + lax.fori_loop(0, num_bq, compute_with_bq, None, unroll=False) + + @pl.when(seq_idx == 0) + def prologue(): + start_fetch_bq(0, 0, 0) + start_fetch_bkv(0, 0, 0) + + @pl.when(seq_idx < decode_end) + def process_decode(): + process(static_q_len=1) + + @pl.when(jnp.logical_and(decode_end <= seq_idx, seq_idx < prefill_end)) + def process_prefill(): + process(static_q_len=chunk_prefill_size) + + @pl.when(jnp.logical_and(prefill_end <= seq_idx, seq_idx < mixed_end)) + def process_mixed(): + process() + + @pl.when(seq_idx == num_seqs - 1) + def epilogue(): + for i in range(2): + wait_send_bo(i) + wait_update_kv_cache(i) + + def prepare_q_inputs(q: jax.Array): + max_num_tokens, actual_num_q_heads, actual_head_dim = q.shape + q_packing = get_dtype_packing(q.dtype) + num_q_heads = align_to(actual_num_q_heads, q_packing) + head_dim = align_to(actual_head_dim, 128) + q = jnp.pad( + q.reshape( + max_num_tokens, + actual_num_q_heads, + actual_head_dim, + ), + ( + (0, 0), + (0, num_q_heads - actual_num_q_heads), + (0, head_dim - actual_head_dim), + ), + constant_values=0, + ).reshape( + max_num_tokens, + num_q_heads // q_packing, + q_packing, + head_dim, + ) + return q + + def prepare_kv_inputs(kv: jax.Array): + max_num_tokens, actual_head_dim = kv.shape + kv_packing = get_dtype_packing(kv.dtype) + if max_num_tokens % kv_packing != 0: + pad = kv_packing - (max_num_tokens % kv_packing) + kv = jnp.pad(kv, ((0, pad), (0, 0)), constant_values=0) + + head_dim = align_to(actual_head_dim, 128) + kv = kv.reshape(-1, kv_packing, actual_head_dim) + kv = jnp.pad(kv, ((0, 0), (0, 0), (0, head_dim - actual_head_dim)), constant_values=0) + return kv + + def prepare_outputs( + out, + actual_num_q_heads: int, + actual_head_dim: int, + ): + ( + max_num_tokens, + num_q_heads_per_q_packing, + q_packing, + head_dim, + ) = out.shape + return out.reshape( + max_num_tokens, + num_q_heads_per_q_packing * q_packing, + head_dim, + )[:, :actual_num_q_heads, :actual_head_dim] + + @functools.partial( + jax.jit, + static_argnames=( + "sm_scale", + "sliding_window", + "soft_cap", + "mask_value", + "q_scale", + "k_scale", + "v_scale", + "chunk_prefill_size", + "num_kv_pages_per_block", + "num_queries_per_block", + "vmem_limit_bytes", + "debug_mode", + ), + donate_argnames=("cache_kv", ), + ) + def mla_ragged_paged_attention( + ql_nope: jax.Array, + q_pe: jax.Array, + new_kv_c: jax.Array, + new_k_pe: jax.Array, + cache_kv: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + distribution: jax.Array, + *, + sm_scale: float = 1.0, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = DEFAULT_MASK_VALUE, + q_scale: float | None = None, + k_scale: float | None = None, + v_scale: float | None = None, + chunk_prefill_size: int | None = None, + num_kv_pages_per_block: int | None = None, + num_queries_per_block: int | None = None, + vmem_limit_bytes: int | None = None, + debug_mode: bool = False, + ) -> tuple[jax.Array, jax.Array]: + if num_kv_pages_per_block is None or num_queries_per_block is None: + raise ValueError("num_kv_pages_per_block and num_queries_per_block must be specified.") + static_validate_inputs( + ql_nope, + q_pe, + new_kv_c, + new_k_pe, + cache_kv, + kv_lens, + page_indices, + cu_q_lens, + distribution, + sm_scale=sm_scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + mask_value=mask_value, + q_scale=q_scale, + k_scale=k_scale, + v_scale=v_scale, + chunk_prefill_size=chunk_prefill_size, + num_kv_pages_per_block=num_kv_pages_per_block, + num_queries_per_block=num_queries_per_block, + vmem_limit_bytes=vmem_limit_bytes, + debug_mode=debug_mode, + ) + + _, actual_num_q_heads, actual_lkv_dim = ql_nope.shape + + ql_nope = prepare_q_inputs(ql_nope) + q_pe = prepare_q_inputs(q_pe) + new_kv_c = prepare_kv_inputs(new_kv_c) + new_k_pe = prepare_kv_inputs(new_k_pe) + lkv_dim = new_kv_c.shape[-1] + r_dim = new_k_pe.shape[-1] + + _, page_size_per_kv_packing, kv_packing, _ = cache_kv.shape + page_size = page_size_per_kv_packing * kv_packing + _, num_q_heads_per_q_packing, q_packing, _ = ql_nope.shape + max_num_seqs = kv_lens.shape[0] + num_page_indices = page_indices.shape[0] + assert num_page_indices % max_num_seqs == 0 + num_q_heads = num_q_heads_per_q_packing * q_packing + + bkv_p = num_kv_pages_per_block + bq_sz = num_queries_per_block + bkv_sz_per_kv_packing = bkv_p * page_size_per_kv_packing + bkv_buf_sz_per_kv_packing = bkv_sz_per_kv_packing + 2 + grid = (distribution[2], ) + + in_specs = [ + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + ] + + out_specs = [ + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + ] + + bkvc_double_buf = pltpu.VMEM( + (2, bkv_buf_sz_per_kv_packing, kv_packing, lkv_dim), + cache_kv.dtype, + ) + + bkpe_double_buf = pltpu.VMEM( + (2, bkv_buf_sz_per_kv_packing, kv_packing, r_dim), + cache_kv.dtype, + ) + bq_nope_double_buf = pltpu.VMEM( + (2, bq_sz, num_q_heads_per_q_packing, q_packing, lkv_dim), + ql_nope.dtype, + ) + + bq_rope_double_buf = pltpu.VMEM( + (2, bq_sz, num_q_heads_per_q_packing, q_packing, r_dim), + q_pe.dtype, + ) + + bo_double_buf = bq_nope_double_buf + + l_scratch = pltpu.VMEM( + (bq_sz * num_q_heads, 128), + jnp.float32, + ) + m_scratch = l_scratch + + acc_scratch = pltpu.VMEM( + (bq_sz * num_q_heads, lkv_dim), + jnp.float32, + ) + + scratch_shapes = [ + bkvc_double_buf, + bkpe_double_buf, + bq_nope_double_buf, + bq_rope_double_buf, + bo_double_buf, + pltpu.SemaphoreType.DMA((4, 2)), + l_scratch, + m_scratch, + acc_scratch, + ] + + scalar_prefetches = ( + kv_lens, + page_indices, + cu_q_lens, + distribution, + jnp.zeros((3, ), jnp.int32), + jnp.full((4, ), -1, jnp.int32), + jnp.full((6, ), -1, jnp.int32), + ) + + scope_name = f"MLA-RPA-bq_{bq_sz}-bkvp_{bkv_p}-p_{page_size}" + kernel = jax.named_scope(scope_name)( + pl.pallas_call( + functools.partial( + _mla_ragged_paged_attention_kernel, + sm_scale=sm_scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + mask_value=mask_value, + q_scale=q_scale, + k_scale=k_scale, + v_scale=v_scale, + chunk_prefill_size=chunk_prefill_size, + bq_sz=bq_sz, + bkv_p=bkv_p, + debug_mode=debug_mode, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=len(scalar_prefetches), + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + scratch_shapes=scratch_shapes, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary", ), + vmem_limit_bytes=vmem_limit_bytes, + ), + out_shape=[ + jax.ShapeDtypeStruct(shape=ql_nope.shape, dtype=ql_nope.dtype), + jax.ShapeDtypeStruct(shape=cache_kv.shape, dtype=cache_kv.dtype), + ], + input_output_aliases={ + 7: 0, + 11: 1, + }, + name=scope_name, + )) + + output, updated_kv = kernel( + *scalar_prefetches, + ql_nope, + q_pe, + new_kv_c, + new_k_pe, + cache_kv, + ) + output = prepare_outputs(output, actual_num_q_heads, actual_lkv_dim) + + return output, updated_kv + + def workload( + ql_nope: jax.Array, + q_pe: jax.Array, + new_kv_c: jax.Array, + new_k_pe: jax.Array, + cache_kv: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + distribution: jax.Array, + ): + return mla_ragged_paged_attention( + ql_nope, + q_pe, + new_kv_c, + new_k_pe, + cache_kv, + kv_lens, + page_indices, + cu_q_lens, + distribution, + num_kv_pages_per_block=16, + num_queries_per_block=4, + vmem_limit_bytes=DEFAULT_VMEM_LIMIT_BYTES + ) + + return workload(ql_nope, q_pe, new_kv_c, new_k_pe, cache_kv, kv_lens, page_indices, cu_q_lens, distribution) \ No newline at end of file diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/4p_Sparse_Attention/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/4p_Sparse_Attention/kernel_task.yaml new file mode 100644 index 0000000..0b8e676 --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/4p_Sparse_Attention/kernel_task.yaml @@ -0,0 +1,34 @@ +task_id: 4p_Sparse_Attention +description: Kernel task for 4p_Sparse_Attention +input_gen_code: |- + def get_inputs(dtype=jnp.bfloat16): + import jax + import jax.numpy as jnp + + CONFIG = { + 'name': 'llama3_70b_sparse_attention', + 'model': 'Llama-3.1-70B', + 'operator': 'sparse_attention', + 'batch': 4, + 'seq_len': 4096, + 'num_query_heads': 64, + 'num_kv_heads': 8, + 'head_dim': 128, + } + key = jax.random.key(42) + k1, k2, k3 = jax.random.split(key, 3) + B = CONFIG["batch"] + S = CONFIG["seq_len"] + H_q = CONFIG["num_query_heads"] + H_kv = CONFIG["num_kv_heads"] + D = CONFIG["head_dim"] + q = jax.random.normal(k1, (B, H_q, S, D), dtype=dtype) * (D ** -0.5) + k = jax.random.normal(k2, (B, H_kv, S, D), dtype=dtype) * 0.02 + v = jax.random.normal(k3, (B, H_kv, S, D), dtype=dtype) * 0.02 + + dynamic_args = [q, k, v] + static_args = [S, H_q, H_kv, D] + return dynamic_args, static_args + +rtol: 0.01 +atol: 0.01 diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/4p_Sparse_Attention/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/4p_Sparse_Attention/reference.py new file mode 100644 index 0000000..d66b622 --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/4p_Sparse_Attention/reference.py @@ -0,0 +1,2397 @@ +# Imports +from collections.abc import Callable, Mapping +import dataclasses +import enum +import functools +from typing import Any, Literal, NamedTuple, Optional, Union, 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 + +# Initialization +def get_inputs(dtype=jnp.bfloat16): + CONFIG = { + 'name': 'llama3_70b_sparse_attention', + 'model': 'Llama-3.1-70B', + 'operator': 'sparse_attention', + 'batch': 4, + 'seq_len': 4096, + 'num_query_heads': 64, + 'num_kv_heads': 8, + 'head_dim': 128, + } + key = jax.random.key(42) + k1, k2, k3 = jax.random.split(key, 3) + B = CONFIG['batch'] + S = CONFIG['seq_len'] + H_q = CONFIG['num_query_heads'] + H_kv = CONFIG['num_kv_heads'] + D = CONFIG['head_dim'] + q = jax.random.normal(k1, (B, H_q, S, D), dtype=dtype) * (D ** -0.5) + k = jax.random.normal(k2, (B, H_kv, S, D), dtype=dtype) * 0.02 + v = jax.random.normal(k3, (B, H_kv, S, D), dtype=dtype) * 0.02 + + dynamic_args = [q, k, v] + static_args = [S, H_q, H_kv, D] + return dynamic_args, static_args + +# Computation +class SegmentIds(NamedTuple): + q: jax.Array + kv: jax.Array + +SplashCustomReturnType = Union[ + jax.Array, + tuple[jax.Array, tuple[jax.Array,]] +] + +SplashResidualsType = tuple[ + jax.Array, + jax.Array, + jax.Array, + Optional[SegmentIds], + jax.Array, + jax.Array, + Optional[mask_info_lib.MaskInfo], + Optional[mask_info_lib.MaskInfo], +] + +MaskFunctionType = Callable[..., jax.Array] + +def get_kernel_name( + block_metadata: Mapping[str, Any], + is_mqa: bool, + save_residuals: bool, + is_segmented: bool, + phase: str, +) -> str: + assert phase == "dq" or phase == "dkv" or phase == "fwd" + 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}_" + "_".join( + f"{k}={v}" for k, v in sorted(block_metadata.items()) + ) + +@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: jax.Array, + k: jax.Array, + v: jax.Array, + segment_ids: SegmentIds | None, + mask_value: float, + save_residuals: bool, + custom_type: str, + attn_logits_soft_cap: float | None, +): + return _attention_reference_default( + mask, + q, + k, + v, + segment_ids, + mask_value, + save_residuals, + custom_type, + attn_logits_soft_cap, + ) + +def _attention_reference_default( + mask: jax.Array, + q: jax.Array, + k: jax.Array, + v: jax.Array, + 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: jax.Array, + k: jax.Array, + v: jax.Array, + segment_ids: SegmentIds | None, + *, + mask_value: float = -0.7 * float(np.finfo(np.dtype("float32")).max), + save_residuals: bool = False, + custom_type: str = "flash", + attn_logits_soft_cap: float | None = None, +) -> SplashCustomReturnType: + return _attention_reference( + 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: jax.Array, + k: jax.Array, + v: jax.Array, + 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) + dv = jnp.einsum("pt,pd->td", p, do).astype(v.dtype) + dp = jnp.einsum("pd,td->pt", do, v.astype(jnp.float32)) + + 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_argnames=( + "mask_value", "save_residuals", "custom_type", "attn_logits_soft_cap") +) +_attention_reference_custom.defvjp(_attention_reference_custom_fwd, + _attention_reference_custom_bwd) + +def attention_reference_custom( + mask: jax.Array, + q: jax.Array, + k: jax.Array, + v: jax.Array, + segment_ids: SegmentIds | None, + *, + mask_value: float = -0.7 * float(np.finfo(np.dtype("float32")).max), + 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: + @functools.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 = -0.7 * float(np.finfo(np.dtype("float32")).max), + save_residuals: bool = False, + attn_logits_soft_cap: float | None = None, + ): + if backward_impl == "custom": + attn_impl = functools.partial( + attention_reference_custom, custom_type="flash", + ) + elif backward_impl == "custom_vanilla": + attn_impl = functools.partial( + attention_reference_custom, custom_type="vanilla", + ) + else: + attn_impl = attention_reference + func = functools.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: + 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:])) + + func = jax.vmap(func, in_axes=(0, 0, None, None, None)) + + 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 = functools.partial(make_attention_reference, is_mqa=False) +make_masked_mqa_reference = functools.partial(make_attention_reference, is_mqa=True) + +class QKVLayout(enum.IntEnum): + HEAD_DIM_MINOR = enum.auto() + SEQ_MINOR = enum.auto() + +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(frozen=True, slots=True) +class BlockSizes: + 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: + object.__setattr__(self, "block_kv_compute", self.block_kv) + if self.block_kv_dkv_compute is None: + object.__setattr__(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): + 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: + assert m_next_ref is None + next_data = i if next_i else j + return ( + next_data, + None, + True, + False, + ) + + 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] + + if data_next_ref.shape[0] == 1: + h = 0 + + 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 | jax.Array, + 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 = mask_ref[:, k_slice] + else: + mask = mask_ref[k_slice, :] + + masks.append( + jnp.bitwise_or(mask, jnp.broadcast_to(should_not_mask, mask.shape)) + ) + if mask_function is not None: + if k_in_lanes: + assert q_sequence_ref.shape == (bq, 128) + + k_sequence = k_offset + jax.lax.broadcasted_iota( + jnp.int32, (bq, k_slice.size), 1 + ) + + repeats, rem = divmod(k_slice.size, 128) + assert rem == 0 + q_sequence = jnp.tile( + q_sequence_ref[...], (1, repeats) + ) + else: + assert q_sequence_ref.shape == (8, bq) + + k_sequence = k_offset + jax.lax.broadcasted_iota( + jnp.int32, (k_slice.size, bq), 0 + ) + q_sequence = q_sequence_ref[:1, :] + q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq)) + + assert q_sequence.shape == k_sequence.shape + computed_mask = mask_function(q_sequence, k_sequence) + if computed_mask.dtype != jnp.dtype(jnp.bool_): + raise ValueError( + "Mask function must return a boolean-valued array, but got:" + f" {computed_mask.dtype}" + ) + masks.append(computed_mask) + + if q_segment_ids_ref is not None: + if k_in_lanes: + kv_ids = kv_segment_ids_ref[:1, k_slice] + repeats, rem = divmod(kv_ids.shape[1], 128) + if rem: + raise NotImplementedError(f"block_kv must be a multiple of {128}") + q_ids = jnp.tile(q_segment_ids_ref[:], (1, repeats)) + else: + assert bq == q_segment_ids_ref.shape[-1] + repeats, rem = divmod(bq, 128) + if rem: + raise NotImplementedError(f"block_q must be a multiple of {128}") + kv_ids = jnp.tile( + kv_segment_ids_ref[k_slice, :], (1, repeats) + ) + q_ids = q_segment_ids_ref[:1, :] + 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( + data_next_ref, + block_mask_ref, + mask_next_ref, + q_ref, + k_ref, + v_ref, + q_segment_ids_ref, + kv_segment_ids_ref, + mask_ref, + q_sequence_ref, + m_scratch_ref, + l_scratch_ref, + o_scratch_ref, + o_ref, + logsumexp_ref=None, + *, + mask_value: float, + grid_width: int, + bq: int, + bkv: int, + bkv_compute: int, + head_dim_v: 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_v_repeats, rem = divmod(head_dim_v, 128) + if rem != 0: + raise NotImplementedError( + f"{head_dim_v=} should be a multiple of {128}" + ) + + 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) + + global_kv_index, _, 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, 128) + assert l_prev.shape == (bq, 128) + + q = q_ref[...] if q_layout == HEAD_DIM_MINOR else q_ref[...].T + qk_dims = (((1,), (1,)), ((), ())) if k_layout == HEAD_DIM_MINOR else (((1,), (0,)), ((), ())) + if k_layout == HEAD_DIM_MINOR: + k = k_ref[slice_k, :] + else: + k = k_ref[:, 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=global_kv_index * 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, 128) + + bkv_repeats, rem = divmod(bkv_compute, 128) + if rem != 0: + raise NotImplementedError( + f"{bkv_compute=} should be a multiple of {128}" + ) + + s_curr = jnp.exp(qk - jnp.tile(m_next, (1, bkv_repeats))) + 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, 128) + + 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 = (((1,), (0,)), ((), ())) if v_layout == HEAD_DIM_MINOR else (((1,), (1,)), ((), ())) + if v_layout == HEAD_DIM_MINOR: + v = v_ref[slice_k, :] + else: + v = v_ref[:, slice_k] + v = v.astype(float32) + o_curr = lax.dot_general(s_curr, v, sv_dims) + + alpha_o = jnp.tile(alpha, (1, head_dim_v_repeats)) + 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 = jnp.tile(1.0 / l, (1, head_dim_v_repeats)) + o_ref[...] = (o_scratch_ref[...] * l_inv).astype(o_ref.dtype) + if logsumexp_ref is not None: + assert logsumexp_ref.shape == (bq, 128) + 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, + interpret: bool = False +) -> SplashCustomReturnType: + num_q_heads, q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + 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] + + partial_mask_blocks = fwd_mask_info.partial_mask_blocks + if ( + partial_mask_blocks is not None + and jnp.dtype(partial_mask_blocks.dtype) != np.bool_ + ): + raise ValueError( + "partial_mask_blocks must be of type np.bool_ but got" + f" {partial_mask_blocks.dtype}" + ) + + 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_qk: + raise ValueError( + f"Expected 'key' head dimension to be: {head_dim_qk}. 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[:-1] != v.shape[:-1]: + raise ValueError( + f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " + "leading dimensions." + ) + + if bkv % bkv_compute: + raise ValueError(f"{bkv=} must be a multiple of {bkv_compute=}.") + if bkv_compute % 128: + raise ValueError(f"{bkv_compute=} must be a multiple of {128}.") + + 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 + 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 + + in_specs = [ + pl.BlockSpec( + from_head_minor((None, bq, head_dim_qk), q_layout), q_index_map + ), + pl.BlockSpec( + from_head_minor( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), k_layout + ), + k_index_map, + ), + pl.BlockSpec( + from_head_minor( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), v_layout + ), + v_index_map, + ), + ] + if segment_ids is not None: + in_specs += [ + pl.BlockSpec((bq, 128), q_segment_ids_index_map), + pl.BlockSpec((8, bkv), kv_segment_ids_index_map), + ] + q_segment_ids = jax.lax.broadcast_in_dim( + segment_ids.q, (q_seq_len, 128), (0,) + ) + kv_segment_ids = jax.lax.broadcast_in_dim( + segment_ids.kv, (8, 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((None, bq, bkv), mask_index_map)) + 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, 128), (0,) + ) + in_specs.append(pl.BlockSpec((bq, 128), q_segment_ids_index_map)) + else: + q_sequence = None + in_specs.append(None) + + num_scalar_prefetch = 3 + + out_shapes = [ + jax.ShapeDtypeStruct((bq, 128), jnp.float32), + jax.ShapeDtypeStruct((bq, 128), jnp.float32), + jax.ShapeDtypeStruct((bq, head_dim_v), jnp.float32), + jax.ShapeDtypeStruct((num_q_heads, q_seq_len, head_dim_v), q.dtype), + ] + out_specs = [ + pl.BlockSpec((bq, 128), lambda h, i, j, *_: (0, 0)), + pl.BlockSpec((bq, 128), lambda h, i, j, *_: (0, 0)), + pl.BlockSpec((bq, head_dim_v), lambda h, i, j, *_: (0, 0)), + pl.BlockSpec((None, bq, head_dim_v), out_index_map), + ] + if save_residuals: + out_shapes += [ + jax.ShapeDtypeStruct( + (num_q_heads, q_seq_len, 128), jnp.float32 + ), + ] + + def logsumexp_index_map(h, i, *_): + return h, i, 0 + + out_specs += [ + pl.BlockSpec((None, bq, 128), logsumexp_index_map), + ] + else: + out_shapes += [None] + out_specs += [None] + + kernel_name = get_kernel_name( + dataclasses.asdict(block_sizes), + is_mqa=is_mqa, + save_residuals=save_residuals, + is_segmented=segment_ids is not None, + phase="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( + functools.partial( + flash_attention_kernel, + mask_value=mask_value, + grid_width=grid_width, + bq=bq, + bkv=bkv, + bkv_compute=bkv_compute, + head_dim_v=head_dim_v, + 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, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary", "arbitrary"), + ), + out_shape=out_shapes, + name=kernel_name, + interpret=interpret, + )( + 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 + +@functools.partial(jax.custom_vjp, nondiff_argnames=( + "save_residuals", "mask_value", "is_mqa", "block_sizes", + "residual_checkpoint_name", "mask_function", "attn_logits_soft_cap", + "interpret") +) +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, + interpret: bool = False, +) -> SplashCustomReturnType: + del dq_mask_info, dkv_mask_info + + return _splash_attention_forward( + fwd_mask_info, + q, + k, + v, + segment_ids, + mask_value=mask_value, + is_mqa=is_mqa, + block_sizes=block_sizes, + residual_checkpoint_name=residual_checkpoint_name, + save_residuals=save_residuals, + mask_function=mask_function, + attn_logits_soft_cap=attn_logits_soft_cap, + interpret=interpret, + ) + +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, + interpret: bool = False, +) -> 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, + residual_checkpoint_name=residual_checkpoint_name, + save_residuals=True, + mask_function=mask_function, + attn_logits_soft_cap=attn_logits_soft_cap, + interpret=interpret, + ) + return out, ( + q, + k, + v, + segment_ids, + out, + logsumexp, + dq_mask_info, + dkv_mask_info, + ) + +def _flash_attention_dq_kernel( + data_next_ref, + block_mask_ref, + mask_next_ref, + 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, + 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) + + global_kv_index, _, 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 + 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 = (((1,), (1,)), ((), ())) if k_layout == HEAD_DIM_MINOR else (((1,), (0,)), ((), ())) + 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=global_kv_index * bkv, + bq=bq, + mask_function=mask_function, + ) + p = jnp.exp(qk - logsumexp) + dp_dims = (((1,), (1,)), ((), ())) if v_layout == HEAD_DIM_MINOR else (((1,), (0,)), ((), ())) + 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 = (((1,), (0,)), ((), ())) if k_layout == HEAD_DIM_MINOR else (((1,), (1,)), ((), ())) + 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, + interpret: bool, +): + num_q_heads, q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + 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[:-1] != v.shape[:-1]: + raise ValueError( + f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " + "leading dimensions." + ) + + if bkv % 128: + raise ValueError(f"{bkv=} must be a multiple of {128}.") + + 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((None, bq, head_dim_v), o_index_map) + + def q_index_map(h, i, *_): + return from_head_minor((h, i, 0), q_layout) + + q_spec = pl.BlockSpec( + from_head_minor((None, bq, head_dim_qk), q_layout), q_index_map + ) + + 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( + from_head_minor( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), k_layout + ), + k_index_map, + ) + + 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( + from_head_minor( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), v_layout + ), + v_index_map, + ) + + 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((None, bq, bkv), mask_index_map) + + def q_segment_ids_index_map(h, i, j, *_): + del h, j + 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((bq, 128), q_segment_ids_index_map) + kv_segment_spec = pl.BlockSpec( + (8, bkv), kv_segment_ids_index_map + ) + q_segment_ids = jax.lax.broadcast_in_dim( + segment_ids.q, (q_seq_len, 128), (0,) + ) + kv_segment_ids = jax.lax.broadcast_in_dim( + segment_ids.kv, (8, kv_seq_len), (1,) + ) + else: + q_segment_spec = kv_segment_spec = None + q_segment_ids = kv_segment_ids = None + + do_spec = o_spec + + def logsumexp_index_map(h, i, *_): + return h, 0, i + + logsumexp = jnp.expand_dims(logsumexp, axis=-2) + logsumexp_spec = pl.BlockSpec((None, 1, bq), logsumexp_index_map) + assert logsumexp.ndim == len(logsumexp_spec.block_shape) + + di = jnp.expand_dims(di, axis=-2) + di_spec = pl.BlockSpec((None, 1, bq), logsumexp_index_map) + 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, 128), (0,) + ) + in_specs.append(pl.BlockSpec((bq, 128), q_segment_ids_index_map)) + else: + q_sequence = None + in_specs.append(None) + + out_shapes = [ + jax.ShapeDtypeStruct((bq, head_dim_qk), jnp.float32), + jax.ShapeDtypeStruct(q.shape, q.dtype), + ] + out_specs = [ + pl.BlockSpec((bq, head_dim_qk), lambda *_: (0, 0)), + pl.BlockSpec((None, bq, head_dim_qk), lambda h, i, *_: (h, i, 0)), + ] + + 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 + + kernel_name = get_kernel_name( + 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, + save_residuals=False, + is_segmented=segment_ids is not None, + phase="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, + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary", "arbitrary", "arbitrary"), + ), + name=kernel_name, + interpret=interpret, + )( + 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( + data_next_ref, + block_mask_ref, + mask_next_ref, + 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, + 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 + + 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[...] + def _load_kv(ref, layout): + if layout == HEAD_DIM_MINOR: + return ref[slice_k, :] + return ref[:, slice_k].T + k = _load_kv(k_ref, k_layout) + v = _load_kv(v_ref, v_layout) + logsumexp = logsumexp_ref[:1, :] + do = do_ref[...] + di = di_ref[:1, :] + + qk_dims = (((1,), (1,)), ((), ())) if q_layout == HEAD_DIM_MINOR else (((1,), (0,)), ((), ())) + 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) + dv_scratch_ref[slice_k, :] + dv_scratch_ref[slice_k, :] = dv + + dp = lax.dot_general( + v, do, (((1,), (1,)), ((), ())), + 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 = (((1,), (0,)), ((), ())) if q_layout == HEAD_DIM_MINOR else (((1,), (1,)), ((), ())) + dk = lax.dot_general( + ds.astype(do.dtype), q, dk_dims, preferred_element_type=jnp.float32 + ) + dk = dk.astype(dk_scratch_ref.dtype) + dk_scratch_ref[slice_k, :] + dk_scratch_ref[slice_k, :] = dk + if dq_scratch_ref is not None or dq_ref is not None: + dq = lax.dot_general( + ds.T.astype(k.dtype), k, (((1,), (0,)), ((), ())), + preferred_element_type=jnp.float32, + ) + if dq_scratch_ref is not None: + dq_scratch_ref[...] += dq + else: + 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, + interpret: bool, +): + num_q_heads, q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + 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[:-1] != v.shape[:-1]: + raise ValueError( + f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " + "leading dimensions." + ) + + 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((None, bq, head_dim_v), o_index_map) + + 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( + from_head_minor((None, bq, head_dim_qk), q_layout), q_index_map + ) + + 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( + from_head_minor( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), + k_layout, + ), + k_index_map, + ) + + 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( + from_head_minor( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), + v_layout, + ), + v_index_map, + ) + + 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((None, None, bq, head_dim_qk), dq_index_map) + 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((bq, head_dim_qk), lambda *_: (0, 0)) + dq_scratch_shape = jax.ShapeDtypeStruct((bq, head_dim_qk), 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 = pl.BlockSpec( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), + dkv_index_map, + ) + + dv_spec = pl.BlockSpec( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), + dkv_index_map, + ) + + 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((None, bkv, bq), mask_index_map) + + 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((8, bq), q_segment_ids_index_map) + kv_segment_spec = pl.BlockSpec((bkv, 128), kv_segment_ids_index_map) + q_segment_ids = jax.lax.broadcast_in_dim( + segment_ids.q, (8, q_seq_len), (1,) + ) + kv_segment_ids = jax.lax.broadcast_in_dim( + segment_ids.kv, (kv_seq_len, 128), (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) + logsumexp_shape = (num_q_heads, 8, q_seq_len) + logsumexp = jnp.broadcast_to(jnp.expand_dims(logsumexp, -2), logsumexp_shape) + logsumexp_spec = pl.BlockSpec((None, 8, bq), logsumexp_index_map) + assert logsumexp.ndim == len(logsumexp_spec.block_shape) + + di = jnp.broadcast_to(jnp.expand_dims(di, -2), logsumexp_shape) + di_spec = pl.BlockSpec((None, 8, bq), logsumexp_index_map) + 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((8, bq), q_segment_ids_index_map)) + q_sequence = jax.lax.broadcast_in_dim( + mask_info.q_sequence, (8, q_seq_len), (1,) + ) + else: + q_sequence = None + in_specs.append(None) + + out_shapes = [ + dq_scratch_shape, + jax.ShapeDtypeStruct((bkv, head_dim_qk), jnp.float32), + jax.ShapeDtypeStruct((bkv, head_dim_v), jnp.float32), + dq_shape, + jax.ShapeDtypeStruct(k.shape, k.dtype), + jax.ShapeDtypeStruct(v.shape, v.dtype), + ] + out_specs = [ + dq_scratch_spec, + pl.BlockSpec((bkv, head_dim_qk), lambda *_: (0, 0)), + pl.BlockSpec((bkv, head_dim_v), lambda *_: (0, 0)), + 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 + + kernel_name = get_kernel_name( + 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, + ), + is_mqa=is_mqa, + save_residuals=False, + is_segmented=segment_ids is not None, + phase="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, + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary", "arbitrary", "arbitrary"), + ), + name=kernel_name, + interpret=interpret, + )( + 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, + interpret: bool, + res: SplashResidualsType, + do: jax.Array, +) -> tuple[ + mask_info_lib.MaskInfo | None, + mask_info_lib.MaskInfo | None, + mask_info_lib.MaskInfo | None, + jax.Array, + jax.Array, + jax.Array, + SegmentIds | None, +]: + 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 = jnp.einsum("hsd,hsd->hs", o.astype(jnp.float32), do.astype(jnp.float32)) + 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, + interpret=interpret, + ) + 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, + interpret=interpret, + ) + assert dq is not None + return ( + None, + None, + None, + dq, + dk, + dv, + None, + ) + +_splash_attention_custom.defvjp(_splash_attention_fwd, _splash_attention_bwd) + +@functools.partial( + jax.jit, + static_argnames=[ + "is_mqa", + "block_sizes", + "save_residuals", + "mask_value", + "attn_logits_soft_cap", + "residual_checkpoint_name", + "mask_function", + "interpret", + ], +) +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, + interpret: bool, +) -> SplashCustomReturnType: + def _collapse_partial_mask_blocks(mask_info: mask_info_lib.MaskInfo | None): + if mask_info is None or mask_info.partial_mask_blocks is None: + return mask_info + + return mask_info._replace( + partial_mask_blocks=mask_info.partial_mask_blocks.reshape( + -1, *mask_info.partial_mask_blocks.shape[-2:] + ) + ) + + fwd_mask_info = _collapse_partial_mask_blocks(fwd_mask_info) + dq_mask_info = _collapse_partial_mask_blocks(dq_mask_info) + dkv_mask_info = _collapse_partial_mask_blocks(dkv_mask_info) + 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, + interpret=interpret, + ) + +@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): + 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() + partial_mask_blocks_spec = ( + spec if self.fwd_mask_info.is_dynamic_mask else replicated + ) + q_sequence_spec = jax.sharding.PartitionSpec(spec[1]) + mask_info_specs = mask_info_lib.MaskInfo( + 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=partial_mask_blocks_spec + 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 + 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 | jax.Array | mask_lib.MultiHeadMask, + *, + block_sizes: BlockSizes | None = None, + is_mqa: bool, + save_residuals: bool = False, + mask_value: float = -0.7 * float(np.finfo(np.dtype("float32")).max), + 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, + interpret: bool = False, +): + 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() + + process_mask_fn = ( + mask_info_lib.process_dynamic_mask + if isinstance(mask, jax.Array) + else mask_info_lib.process_mask + ) + + process_mask_dvk_fn = ( + mask_info_lib.process_dynamic_mask_dkv + if isinstance(mask, jax.Array) + else mask_info_lib.process_mask_dkv + ) + + fwd_mask_info, mask_function_fwd = process_mask_fn( + 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 = process_mask_fn( + 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 = process_mask_dvk_fn( + 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, + interpret=interpret, + ) + +make_splash_mha = functools.partial(_make_splash_attention, is_mqa=False) +make_splash_mqa = functools.partial(_make_splash_attention, is_mqa=True) + +make_splash_mha_single_device = functools.partial( + make_splash_mha, is_mqa=False, head_shards=1, q_seq_shards=1 +) + +make_splash_mqa_single_device = functools.partial( + make_splash_mha, is_mqa=True, head_shards=1, q_seq_shards=1 +) + +def computation(q, k, v, S, H_q, H_kv, D): + TUNED_PARAMS = { + 'block_q': 2048, + 'block_kv': 2048, + 'block_kv_compute': 1024, + 'q_layout': 1, + 'k_layout': 1, + 'v_layout': 1, + 'head_shards': 1, + 'q_seq_shards': 1, + 'block_q_dkv': None, + 'block_kv_dkv': None, + 'block_kv_dkv_compute': None, + 'block_q_dq': None, + 'block_kv_dq': None, + } + heads_per_group = H_q // H_kv + mask = mask_lib.CausalMask(shape=(S, S)) + multi_head_mask = mask_lib.MultiHeadMask([mask] * H_q) + block_sizes = BlockSizes( + block_q=TUNED_PARAMS['block_q'], + block_kv=TUNED_PARAMS['block_kv'], + block_kv_compute=TUNED_PARAMS['block_kv_compute'], + q_layout=QKVLayout(TUNED_PARAMS['q_layout']), + k_layout=QKVLayout(TUNED_PARAMS['k_layout']), + v_layout=QKVLayout(TUNED_PARAMS['v_layout']), + block_q_dkv=TUNED_PARAMS['block_q_dkv'], + block_kv_dkv=TUNED_PARAMS['block_kv_dkv'], + block_kv_dkv_compute=TUNED_PARAMS['block_kv_dkv_compute'], + block_q_dq=TUNED_PARAMS['block_q_dq'], + block_kv_dq=TUNED_PARAMS['block_kv_dq'], + ) + splash_kernel = _make_splash_attention( + multi_head_mask, block_sizes=block_sizes, + is_mqa=False, + head_shards=TUNED_PARAMS['head_shards'], + q_seq_shards=TUNED_PARAMS['q_seq_shards'], + ) + @jax.vmap + def _attend(q_batch, k_batch, v_batch): + k_repeated = jnp.repeat(k_batch, heads_per_group, axis=0) + v_repeated = jnp.repeat(v_batch, heads_per_group, axis=0) + return splash_kernel(q_batch, k_repeated, v_repeated) + return _attend(q, k, v) \ No newline at end of file diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/51p_DeepSeek_V4_CSA/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/51p_DeepSeek_V4_CSA/kernel_task.yaml new file mode 100644 index 0000000..014066b --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/51p_DeepSeek_V4_CSA/kernel_task.yaml @@ -0,0 +1,149 @@ +task_id: 51p_DeepSeek_V4_CSA +description: |- + DeepSeek-V4 compressed sparse attention (CSA) +input_gen_code: |- + def get_inputs(): + import jax + import jax.numpy as jnp + + configs = [ + ("decode_large_batch", 256, 1, 9216, 1024, 1024), + ("decode_medium_batch", 128, 1, 9216, 1024, 1024), + ("decode_small_pages", 128, 1, 9216, 256, 1024), + ("decode_random_access", 256, 1, 9216, 64, 512), + ("prefill_full", 1, 1024, 1024, 1024, 1024), + ("prefill_short", 1, 256, 1024, 1024, 1024), + ("prefill_short_small_pages", 1, 256, 1024, 256, 1024), + ("prefill_mid_chunk_512", 1, 512, 4096, 1024, 1024), + ("prefill_mid_chunk_1024", 1, 1024, 8192, 1024, 1024), + ("prefill_mid_chunk_small_pages", 1, 1024, 8192, 64, 512), + ] + + outputs = [] + k_cfg = jax.random.PRNGKey(42) + + # --- DSV4 CSA geometry ------------------------------------------------- + NOPE_DIM = 448 + NUM_SCALES = NOPE_DIM // 64 # 7 + ROPE_DIM = 64 + TOKEN_BYTES = 512 # (4, 128) uint8 per token in the nope cache + HEAD_DIM = NOPE_DIM + ROPE_DIM # 512 + NUM_Q_HEADS = 8 + Q_DTYPE = jnp.bfloat16 + SM_SCALE = float(HEAD_DIM ** -0.5) + ATTN_BATCH = 16 + VMEM_LIMIT_BYTES = 100 * 1024 * 1024 + + # --- Cache contents ---------------------------------------------------- + CACHE_TILE = 8192 + k_nope, k_scale, k_rope, k_cfg = jax.random.split(k_cfg, 4) + + nope_f8 = jax.random.normal( + k_nope, (CACHE_TILE, NOPE_DIM), jnp.float32 + ).astype(jnp.float8_e4m3fn) + nope_bytes = jax.lax.bitcast_convert_type(nope_f8, jnp.uint8) + + scale_bytes = jax.random.randint( + k_scale, (CACHE_TILE, NUM_SCALES), 125, 130, dtype=jnp.int32 + ).astype(jnp.uint8) + + pad_bytes = jnp.zeros( + (CACHE_TILE, TOKEN_BYTES - NOPE_DIM - NUM_SCALES), jnp.uint8 + ) + + nope_tile = jnp.concatenate( + [nope_bytes, scale_bytes, pad_bytes], axis=1 + ) + + rope_bits = jax.lax.bitcast_convert_type( + jax.random.normal( + k_rope, (CACHE_TILE, ROPE_DIM), jnp.float32 + ).astype(jnp.bfloat16), + jnp.uint16, + ).astype(jnp.uint32) + rope_tile = jnp.concatenate( + [ + (rope_bits >> 8).astype(jnp.uint8), + (rope_bits & 0xFF).astype(jnp.uint8), + ], + axis=1, + ) + + def fill(tile, num_slots): + reps = -(-num_slots // tile.shape[0]) + return jnp.tile(tile, (reps, 1))[:num_slots] + + keys = jax.random.split(k_cfg, len(configs)) + for cfg, key in zip(configs, keys): + name, batch_size, q_len, kv_len, page_size, csa_topk = cfg + k_page, k_topk, k_q, k_sink, k_acc, k_l, k_m = jax.random.split(key, 7) + + num_tokens = batch_size * q_len + pages_per_seq = -(-kv_len // page_size) + total_num_pages = batch_size * pages_per_seq + assert num_tokens % ATTN_BATCH == 0 + + q = jax.random.normal( + k_q, (num_tokens, NUM_Q_HEADS, HEAD_DIM), jnp.float32 + ).astype(Q_DTYPE) + + num_slots = total_num_pages * page_size + cache_kv_nope = fill(nope_tile, num_slots).reshape( + total_num_pages, page_size, 4, 128 + ) + cache_kv_rope = fill(rope_tile, num_slots).reshape( + total_num_pages, page_size // 4, 4, 128 + ) + + page_indices = jax.random.permutation( + k_page, total_num_pages + ).astype(jnp.int32) + + cu_q_lens = (jnp.arange(batch_size + 1, dtype=jnp.int32) * q_len) + num_decode = batch_size if q_len == 1 else 0 + distribution = jnp.array( + [num_decode, batch_size, batch_size], jnp.int32 + ) + + pos_in_seq = jnp.arange(num_tokens, dtype=jnp.int32) % q_len + causal_len = (kv_len - q_len + pos_in_seq + 1)[:, None] + slot = jnp.arange(csa_topk, dtype=jnp.int32)[None, :] + scattered = jax.random.randint( + k_topk, (num_tokens, csa_topk), 0, 1 << 30, dtype=jnp.int32 + ) % causal_len + topk_indices = jnp.where( + slot < causal_len, + jnp.where(causal_len < csa_topk, slot, scattered), + -1, + ).astype(jnp.int32) + + attention_sinks = jax.random.normal( + k_sink, (NUM_Q_HEADS,), jnp.float32 + ) + swa_accumution = jax.random.normal( + k_acc, (num_tokens, NUM_Q_HEADS, HEAD_DIM), jnp.float32 + ).astype(Q_DTYPE) + swa_l = jax.random.uniform( + k_l, (num_tokens, NUM_Q_HEADS), jnp.float32, 1.0, 64.0 + ) + swa_m = jax.random.normal(k_m, (num_tokens, NUM_Q_HEADS), jnp.float32) + + dynamic_args = [ + q, + cache_kv_nope, + cache_kv_rope, + topk_indices, + page_indices, + cu_q_lens, + distribution, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + ] + outputs.append((dynamic_args, [])) + + return outputs + +rtol: 0.02 +atol: 0.02 diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/51p_DeepSeek_V4_CSA/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/51p_DeepSeek_V4_CSA/reference.py new file mode 100644 index 0000000..8065af9 --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/51p_DeepSeek_V4_CSA/reference.py @@ -0,0 +1,1351 @@ +# Imports +import numpy as np +import time +import functools +from enum import Enum +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 +from jax.experimental.pallas import tpu_sc as plsc + +# Initialization +def get_inputs(): + import jax + import jax.numpy as jnp + + configs = [ + ("decode_large_batch", 256, 1, 9216, 1024, 1024), + ("decode_medium_batch", 128, 1, 9216, 1024, 1024), + ("decode_small_pages", 128, 1, 9216, 256, 1024), + ("decode_random_access", 256, 1, 9216, 64, 512), + ("prefill_full", 1, 1024, 1024, 1024, 1024), + ("prefill_short", 1, 256, 1024, 1024, 1024), + ("prefill_short_small_pages", 1, 256, 1024, 256, 1024), + ("prefill_mid_chunk_512", 1, 512, 4096, 1024, 1024), + ("prefill_mid_chunk_1024", 1, 1024, 8192, 1024, 1024), + ("prefill_mid_chunk_small_pages", 1, 1024, 8192, 64, 512), + ] + + outputs = [] + k_cfg = jax.random.PRNGKey(0) + + # --- DSV4 CSA geometry ------------------------------------------------- + NOPE_DIM = 448 + NUM_SCALES = NOPE_DIM // 64 # 7 + ROPE_DIM = 64 + TOKEN_BYTES = 512 # (4, 128) uint8 per token in the nope cache + HEAD_DIM = NOPE_DIM + ROPE_DIM # 512 + NUM_Q_HEADS = 8 + Q_DTYPE = jnp.bfloat16 + SM_SCALE = float(HEAD_DIM ** -0.5) + ATTN_BATCH = 16 + VMEM_LIMIT_BYTES = 100 * 1024 * 1024 + + # --- Cache contents ---------------------------------------------------- + CACHE_TILE = 8192 + k_nope, k_scale, k_rope, k_cfg = jax.random.split(k_cfg, 4) + + nope_f8 = jax.random.normal( + k_nope, (CACHE_TILE, NOPE_DIM), jnp.float32 + ).astype(jnp.float8_e4m3fn) + nope_bytes = jax.lax.bitcast_convert_type(nope_f8, jnp.uint8) + + scale_bytes = jax.random.randint( + k_scale, (CACHE_TILE, NUM_SCALES), 125, 130, dtype=jnp.int32 + ).astype(jnp.uint8) + + pad_bytes = jnp.zeros( + (CACHE_TILE, TOKEN_BYTES - NOPE_DIM - NUM_SCALES), jnp.uint8 + ) + + nope_tile = jnp.concatenate( + [nope_bytes, scale_bytes, pad_bytes], axis=1 + ) + + rope_bits = jax.lax.bitcast_convert_type( + jax.random.normal( + k_rope, (CACHE_TILE, ROPE_DIM), jnp.float32 + ).astype(jnp.bfloat16), + jnp.uint16, + ).astype(jnp.uint32) + rope_tile = jnp.concatenate( + [ + (rope_bits >> 8).astype(jnp.uint8), + (rope_bits & 0xFF).astype(jnp.uint8), + ], + axis=1, + ) + + def fill(tile, num_slots): + reps = -(-num_slots // tile.shape[0]) + return jnp.tile(tile, (reps, 1))[:num_slots] + + keys = jax.random.split(k_cfg, len(configs)) + for cfg, key in zip(configs, keys): + name, batch_size, q_len, kv_len, page_size, csa_topk = cfg + k_page, k_topk, k_q, k_sink, k_acc, k_l, k_m = jax.random.split(key, 7) + + num_tokens = batch_size * q_len + pages_per_seq = -(-kv_len // page_size) + total_num_pages = batch_size * pages_per_seq + assert num_tokens % ATTN_BATCH == 0 + + q = jax.random.normal( + k_q, (num_tokens, NUM_Q_HEADS, HEAD_DIM), jnp.float32 + ).astype(Q_DTYPE) + + num_slots = total_num_pages * page_size + cache_kv_nope = fill(nope_tile, num_slots).reshape( + total_num_pages, page_size, 4, 128 + ) + cache_kv_rope = fill(rope_tile, num_slots).reshape( + total_num_pages, page_size // 4, 4, 128 + ) + + page_indices = jax.random.permutation( + k_page, total_num_pages + ).astype(jnp.int32) + + cu_q_lens = (jnp.arange(batch_size + 1, dtype=jnp.int32) * q_len) + num_decode = batch_size if q_len == 1 else 0 + distribution = jnp.array( + [num_decode, batch_size, batch_size], jnp.int32 + ) + + pos_in_seq = jnp.arange(num_tokens, dtype=jnp.int32) % q_len + causal_len = (kv_len - q_len + pos_in_seq + 1)[:, None] + slot = jnp.arange(csa_topk, dtype=jnp.int32)[None, :] + scattered = jax.random.randint( + k_topk, (num_tokens, csa_topk), 0, 1 << 30, dtype=jnp.int32 + ) % causal_len + topk_indices = jnp.where( + slot < causal_len, + jnp.where(causal_len < csa_topk, slot, scattered), + -1, + ).astype(jnp.int32) + + attention_sinks = jax.random.normal( + k_sink, (NUM_Q_HEADS,), jnp.float32 + ) + swa_accumution = jax.random.normal( + k_acc, (num_tokens, NUM_Q_HEADS, HEAD_DIM), jnp.float32 + ).astype(Q_DTYPE) + swa_l = jax.random.uniform( + k_l, (num_tokens, NUM_Q_HEADS), jnp.float32, 1.0, 64.0 + ) + swa_m = jax.random.normal(k_m, (num_tokens, NUM_Q_HEADS), jnp.float32) + + dynamic_args = [ + q, + cache_kv_nope, + cache_kv_rope, + topk_indices, + page_indices, + cu_q_lens, + distribution, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + ] + outputs.append((dynamic_args, [])) + + return outputs + +# Computation + +def main_kernel( + nope_in_hbm_ref: jax.Ref, + rope_in_hbm_ref: jax.Ref, + indices_hbm_ref: jax.Ref, + nope_out_hbm_ref: jax.Ref, + rope_out_hbm_ref: jax.Ref, + *, + core_axis_name: str, + subcore_axis_name: str, + num_row_subchunks: int, + num_streams: int, +): + tpu_info = pltpu.get_tpu_info() + sc_info = tpu_info.sparse_core + assert sc_info is not None + num_simd_lanes = sc_info.num_lanes + num_cores = jax.lax.axis_size((core_axis_name, subcore_axis_name)) + row_subchunk_size = num_simd_lanes + row_chunk_size = row_subchunk_size * num_row_subchunks + block_size = row_chunk_size * num_cores + num_blocks = pl.cdiv(indices_hbm_ref.shape[0], block_size) + + # Inputs are 8-bit; + # nope output stays uint8 (4/int32), rope is unpacked to bf16 (2/int32). + in_bits = jax.dtypes.itemsize_bits(nope_in_hbm_ref.dtype) + in_packing = 32 // in_bits + in_mask = (1 << in_bits) - 1 # 0xFF for 8-bit. + nope_out_packing = 32 // jax.dtypes.itemsize_bits(nope_out_hbm_ref.dtype) + rope_out_bits = jax.dtypes.itemsize_bits(rope_out_hbm_ref.dtype) + rope_out_packing = 32 // rope_out_bits + core_index = lax.axis_index((core_axis_name, subcore_axis_name)) + + # SparseCore gather 32-bit words + nope_in_i32 = nope_in_hbm_ref.bitcast(jnp.int32) + rope_in_i32 = rope_in_hbm_ref.bitcast(jnp.int32) + nope_out_i32 = nope_out_hbm_ref.bitcast(jnp.int32) + rope_out_i32 = rope_out_hbm_ref.bitcast(jnp.int32) + + nope_in_cols = nope_in_i32.shape[1] + rope_in_cols = rope_in_i32.shape[1] + nope_out_cols = nope_out_hbm_ref.shape[1] + rope_out_cols = rope_out_hbm_ref.shape[1] + + def _delta_swap(x, y, shift, swap_mask): + """Exchanges selected bits between two words. + + Swaps x's bits at positions ``swap_mask << shift`` with y's bits at + positions ``swap_mask``; all other bits are left untouched. Worked + example with ``shift=8, swap_mask=0x00FF00FF`` on byte-quads + ``x = [x0 x1 x2 x3]`` and ``y = [y0 y1 y2 y3]`` (byte 0 = least + significant):: + + swap_mask << 8 = 0xFF00FF00 -> x's bytes 1 and 3 + swap_mask = 0x00FF00FF -> y's bytes 0 and 2 + result: x = [x0 y0 x2 y2], y = [x1 y1 x3 y3] + + i.e. x's odd bytes trade places with y's even bytes. The identity + ``t = ((x >> s) ^ y) & mask; x ^= t << s; y ^= t`` does this in 6 ops + with only `t` as scratch. The arithmetic right shift is safe: its + sign-extended high bits are dropped by the `& swap_mask`. + """ + t = jnp.bitwise_and( + jnp.bitwise_xor(jnp.bitwise_right_shift(x, shift), y), swap_mask + ) + x = jnp.bitwise_xor(x, jnp.left_shift(t, shift)) + y = jnp.bitwise_xor(y, t) + return x, y + + def process_nope(gather_ref, out_ref, out_row_base=0): + # A more efficient implementation of `process_nope_reference`, with fewer + # ALU ops. We've seen that SparseCore's Integer ALU FLOPs being the + # bottleneck of the kernel, this is the preferred implementation. + col_slice = pl.ds(0, 128) + low_16 = jnp.int32(0x0000FFFF) # one 16-bit half of each word + low_byte_of_each_half = jnp.int32(0x00FF00FF) # bytes 0 and 2 + for i in range(num_simd_lanes // nope_out_packing): + d = [ + gather_ref[pl.ds(i * nope_out_packing + j, 1), col_slice] + for j in range(nope_out_packing) + ] + # Round 1: transpose the four 2x2 byte blocks -- exchange the high + # 16 bits of each word with the low 16 bits of its distance-2 peer. + d[0], d[2] = _delta_swap(d[0], d[2], 16, low_16) + d[1], d[3] = _delta_swap(d[1], d[3], 16, low_16) + # Round 2: transpose within each 2x2 block -- exchange the odd bytes + # of each word with the even bytes of its adjacent peer. + d[0], d[1] = _delta_swap(d[0], d[1], 8, low_byte_of_each_half) + d[2], d[3] = _delta_swap(d[2], d[3], 8, low_byte_of_each_half) + # d[m] now holds byte lane m of all 4 inputs -> output sub-row m.s + for m in range(nope_out_packing): + out_ref[pl.ds(out_row_base + i, 1), pl.ds(m * 128, 128)] = d[m] + + def process_rope(gather_ref, out_ref, idx_sub, out_row_base=0): + # one (1, 128) uint8 is one token's rope data, which encodes 64 bf16 + # values. (0, 64) are the high bits for bf16 data, (64, 128) are the low. + half = rope_out_cols + col_hi = pl.ds(0, half) + col_lo = pl.ds(half, half) + + def bf16_bits(k): + sub = lax.rem(idx_sub[k], in_packing) + hi = jnp.bitwise_and( + jnp.bitwise_right_shift( + gather_ref[pl.ds(k, 1), col_hi], in_bits * sub + ), + in_mask, + ) + lo = jnp.bitwise_and( + jnp.bitwise_right_shift( + gather_ref[pl.ds(k, 1), col_lo], in_bits * sub + ), + in_mask, + ) + return jnp.bitwise_or(jnp.left_shift(hi, in_bits), lo) + + for t in range(num_simd_lanes // rope_out_packing): + packed = jnp.zeros((1, half), dtype=jnp.int32) + for pk in range(rope_out_packing): + k = t * rope_out_packing + pk + packed = jnp.bitwise_or( + packed, jnp.left_shift(bf16_bits(k), pk * rope_out_bits) + ) + out_ref[pl.ds(out_row_base + t, 1), pl.ds(0, half)] = packed + + def outer_pipeline(idx_ref): + b = pl.program_id(0) + out_row_base = (b * num_cores + core_index) * num_row_subchunks + + # Subchunk handled by stream `s` at inner step `r`. `num_streams` + # independent `pl.Indirect` gathers run concurrently per step, keeping + # several gather DMAs in flight to raise effective read bandwidth. + def subchunk(r, s): + return r * num_streams + s + + def idx_window(r, s): + return idx_ref[ + pl.ds(subchunk(r, s) * row_subchunk_size, row_subchunk_size) + ] + + # Rows produced per stream in each output (nope packs 4/int32, rope 2). + nope_rows_per_stream = row_subchunk_size // nope_out_packing + rope_rows_per_stream = row_subchunk_size // rope_out_packing + + def _body(*refs): + r = pl.program_id(0) + nope_g = refs[0 * num_streams : 1 * num_streams] + rope_g = refs[1 * num_streams : 2 * num_streams] + nope_o = refs[2 * num_streams] + rope_o = refs[2 * num_streams + 1] + for s in range(num_streams): + process_nope( + gather_ref=nope_g[s], + out_ref=nope_o, + out_row_base=s * nope_rows_per_stream, + ) + process_rope( + gather_ref=rope_g[s], + out_ref=rope_o, + idx_sub=idx_window(r, s), + out_row_base=s * rope_rows_per_stream, + ) + + # Have multiple parallel `pl.Indirect` to hide random access read latency. + # Output contiguous memory access, multiple output parallel DMAs not help + # with performance. + + # nope: gather int32 row == index (1 int32 row per entry). + nope_in_specs = tuple( + pl.BlockSpec( + (pl.Indirect(row_subchunk_size), nope_in_cols), + lambda r, s=s: (idx_window(r, s), 0), + ) + for s in range(num_streams) + ) + # rope: gather int32 row == index // in_packing (in_packing entries/row). + rope_in_specs = tuple( + pl.BlockSpec( + (pl.Indirect(row_subchunk_size), rope_in_cols), + lambda r, s=s: (lax.div(idx_window(r, s), in_packing), 0), + ) + for s in range(num_streams) + ) + # One merged output block per cache, covering all `num_streams` + # subchunks. + nope_out_spec = pl.BlockSpec( + (num_streams * nope_rows_per_stream, nope_out_cols), + lambda r: (out_row_base // num_streams + r, 0), + ) + rope_out_spec = pl.BlockSpec( + (num_streams * rope_rows_per_stream, rope_out_cols), + lambda r: (out_row_base // num_streams + r, 0), + ) + pltpu.emit_pipeline( + _body, + grid=(num_row_subchunks // num_streams,), + in_specs=nope_in_specs + rope_in_specs, + out_specs=(nope_out_spec, rope_out_spec), + )( + *([nope_in_i32] * num_streams), + *([rope_in_i32] * num_streams), + nope_out_i32, + rope_out_i32, + ) + + pltpu.emit_pipeline( + outer_pipeline, + grid=(num_blocks,), + in_specs=pl.BlockSpec( + (row_chunk_size,), + lambda b: (b * num_cores + core_index,), + ), + )(indices_hbm_ref) + + +@functools.partial(jax.jit) +def csa_gather( + nope_cache: jax.Array, + rope_cache: jax.Array, + indices: jax.Array, +) -> tuple[jax.Array, jax.Array]: + """Fused SparseCore gather of the nope and rope caches. + + Args: + nope_cache: (total_pages, page_size, 4, 128) uint8. Each (4, 128) uint8 is + token's nope + nope scales. It encodes 448 fp8 + 7 e8m0 scales + padding. + rope_cache: (total_pages, page_size // 4, 4, 128) uint8. Each (1, 128) uint8 + is token's rope. It encodes 64 bf16. + indices: (N,) int32. Token indices into the caches. + + Returns: + nope_out: (N, 512) uint8. + Each (512) uint8 is token's nope. + rope_out: (N, 64) bf16. + Each (64) bf16 is token's rope. + """ + assert indices.ndim == 1, "Indices must be 1D." + assert nope_cache.dtype == rope_cache.dtype, "Caches must share a dtype." + assert nope_cache.dtype == jnp.uint8, "Caches must be uint8." + + # Flatten both caches to 128-wide rows and view as raw bytes. + nope_cache = nope_cache.reshape(-1, nope_cache.shape[3]) + rope_cache = rope_cache.reshape(-1, rope_cache.shape[3]) + sc_info = pltpu.get_tpu_info().sparse_core + assert sc_info is not None, "SparseCore info is missing." + out_size = indices.size + nope_out_cols = 512 + # rope: each 128-byte entry encodes 64 bf16 (high bytes [0:64], low [64:128]). + rope_out_cols = 64 + num_simd_lanes = sc_info.num_lanes + num_cores = sc_info.num_cores * sc_info.num_subcores + + # `num_streams` independent `pl.Indirect` gathers are issued per + # pipeline step to keep multiple gather DMAs in flight. + # See `outer_pipeline` for details. + num_streams = 2 + num_row_subchunks = 32 + assert ( + num_row_subchunks % num_streams == 0 + ), f"{num_streams=} must divide {num_row_subchunks=}." + row_subchunk_size = num_simd_lanes + row_chunk_size = row_subchunk_size * num_row_subchunks + block_size = row_chunk_size * num_cores + out_pad_size = ( + (out_size + block_size - 1) // block_size + ) * block_size - out_size + indices = jnp.pad(indices, ((0, out_pad_size))) + vector_mesh = plsc.VectorSubcoreMesh( + num_cores=sc_info.num_cores, + num_subcores=sc_info.num_subcores, + core_axis_name="core", + subcore_axis_name="subcore", + ) + nope_out, rope_out = pl.kernel( + functools.partial( + main_kernel, + core_axis_name=vector_mesh.core_axis_name, + subcore_axis_name=vector_mesh.subcore_axis_name, + num_row_subchunks=num_row_subchunks, + num_streams=num_streams, + ), + out_type=( + jax.ShapeDtypeStruct( + (out_size + out_pad_size, nope_out_cols), jnp.uint8 + ), + jax.ShapeDtypeStruct( + (out_size + out_pad_size, rope_out_cols), jnp.bfloat16 + ), + ), + compiler_params=pltpu.CompilerParams( + use_tc_tiling_on_sc=True, + needs_layout_passes=True, + disable_bounds_checks=True, + ), + mesh=vector_mesh, + name="sc_csa_gather", + )(nope_cache, rope_cache, indices) + return ( + nope_out[:out_size], + rope_out[:out_size], + ) + + + +DEFAULT_VMEM_LIMIT_BYTES = 100 * 1024 * 1024 + + +def cdiv(a, b): + assert b != 0 + return (a + b - 1) // b + + +def align_to(x, a): + return cdiv(x, a) * a + + +def get_dtype_bitwidth(dtype): + return jax.dtypes.itemsize_bits(dtype) + + +def get_dtype_packing(dtype): + bits = get_dtype_bitwidth(dtype) + return 32 // bits + + +def get_kv_cache_shape( + total_num_pages, + page_size, + kv_dim, + kv_dtype, +): + kv_packing = get_dtype_packing(kv_dtype) + return ( + total_num_pages, + align_to(page_size, kv_packing) // kv_packing, + kv_packing, + align_to(kv_dim, 128), + ) + + +_GATHER_PAGE_CHUNK = 128 + + +def _gather_page_ids_kernel(windows_ref, logical_ref, out_ref, *, num_chunks): + logical = logical_ref[...] # i32[block_tokens, topk] + out = jnp.zeros_like(logical) + for c in range(num_chunks): + window_chunk = windows_ref[ + :, c * _GATHER_PAGE_CHUNK : (c + 1) * _GATHER_PAGE_CHUNK + ] # i32[block_tokens, 128] + local = logical - c * _GATHER_PAGE_CHUNK + gathered = jnp.take_along_axis( + window_chunk, jnp.clip(local, 0, _GATHER_PAGE_CHUNK - 1), axis=1 + ) + out = jnp.where((local >= 0) & (local < _GATHER_PAGE_CHUNK), gathered, out) + out_ref[...] = out + + +def gather_page_ids( + page_indices: jax.Array, # i32[max_num_seqs * pages_per_seq] + seq_page_ids: jax.Array, # i32[num_tokens, topk] (logical page within seq) + seq_ids_segment: jax.Array, # i32[num_tokens] (token -> seq id) + max_num_seqs: int, + *, + block_tokens: int = 8, +) -> jax.Array: + """Gathers physical page ids for the CSA top-k tokens.""" + num_tokens, topk = seq_page_ids.shape + pages_per_seq = page_indices.shape[0] // max_num_seqs + num_chunks = cdiv(pages_per_seq, _GATHER_PAGE_CHUNK) + padded_pps = num_chunks * _GATHER_PAGE_CHUNK + + page_table = page_indices.reshape(max_num_seqs, pages_per_seq) + if padded_pps != pages_per_seq: + page_table = jnp.pad(page_table, ((0, 0), (0, padded_pps - pages_per_seq))) + # Per-token page-table window. This is a whole-row gather. + windows = page_table[seq_ids_segment] # i32[num_tokens, padded_pps] + logical = jnp.clip(seq_page_ids, 0, pages_per_seq - 1) + + padded_tokens = align_to(num_tokens, block_tokens) + if padded_tokens != num_tokens: + pad = padded_tokens - num_tokens + windows = jnp.pad(windows, ((0, pad), (0, 0))) + logical = jnp.pad(logical, ((0, pad), (0, 0))) + + out = pl.pallas_call( + functools.partial(_gather_page_ids_kernel, num_chunks=num_chunks), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=0, + in_specs=[ + pl.BlockSpec((block_tokens, padded_pps), lambda t: (t, 0)), + pl.BlockSpec((block_tokens, topk), lambda t: (t, 0)), + ], + out_specs=pl.BlockSpec((block_tokens, topk), lambda t: (t, 0)), + grid=(padded_tokens // block_tokens,), + ), + out_shape=jax.ShapeDtypeStruct((padded_tokens, topk), jnp.int32), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary",), + disable_bounds_checks=True, + ), + name="gather_page_ids", + )(windows, logical) + return out[:num_tokens] + + +def _dequant_dsv4_fp8(bkv_nope: jax.Array): + """Dequantize FP8 values to BF16.""" + nope_fp8 = pltpu.bitcast(bkv_nope[:, :448], jnp.float8_e4m3fn).astype( + jnp.bfloat16 + ) + nope_scales = pltpu.bitcast( + bkv_nope[:, 448 : 448 + 7], jnp.float8_e8m0fnu + ).astype(jnp.bfloat16) + nope_scales = jnp.repeat(nope_scales.T, 64, axis=0).T + nope = (nope_fp8 * nope_scales).astype(jnp.bfloat16) + return nope + + +def _attention_kernel( + # Prefetch + kv_lens_ref, # [max_num_seqs] + start_end_seq_idx_ref, # [2] (start_seq_idx, end_seq_idx) + sem_ids_ref, # [2] (bi_sem_idx, bo_sem_idx) + bo_ids_ref, # [2, batch_size] (bo_sem_0_seq_idx, bo_sem_1_seq_idx) + # Input + attention_sinks_ref, # float32[num_q_heads] + q_hbm_ref, # [max_num_tokens, num_q_heads, head_dim] + cache_kv_nope_hbm_ref, # [total_num_pages, page_size, nope_dim] + cache_kv_rope_hbm_ref, # [total_num_pages, page_size, rope_dim] + swa_accumution_hbm_ref, # [max_num_tokens, num_q_heads, head_dim] + swa_l_hbm_ref, # [max_num_tokens, num_l_heads] + swa_m_hbm_ref, # [max_num_tokens, num_l_heads] + # Output + o_hbm_ref, # [max_num_tokens, num_q_heads, head_dim] + # Scratch + bkv_nope_x2_ref, # [2, batch_size, page_size, nope_dim] + bkv_rope_x2_ref, # [2, batch_size, page_size, rope_dim] + bq_x2_ref, # [2, batch_size, num_q_heads, head_dim] + bo_x2_ref, # [2, batch_size, num_q_heads, head_dim] + bl_x2_ref, # [2, batch_size, num_l_heads] + bm_x2_ref, # [2, batch_size, num_l_heads] + swa_acc_x2_ref, # [2, batch_size, num_q_heads, head_dim] + sems, # [7, 2, batch_size] + *, + sm_scale: float, + batch_size: int = 1, +): + assert q_hbm_ref.shape == o_hbm_ref.shape + + num_tokens, num_q_heads, head_dim = q_hbm_ref.shape + _, page_size, _ = cache_kv_nope_hbm_ref.shape + assert kv_lens_ref.shape[0] == num_tokens + bkv_sz = page_size + + q_dtype = q_hbm_ref.dtype + q_packing = get_dtype_packing(q_dtype) + # Validate against the KV dtype. + assert o_hbm_ref.dtype == q_dtype + + assert head_dim % 128 == 0 + assert num_q_heads % q_packing == 0 + + start_seq_idx = start_end_seq_idx_ref[0] + end_seq_idx = start_end_seq_idx_ref[1] + + batch_start_seq_idx = start_seq_idx + pl.program_id(0) * batch_size + batch_end_seq_idx = batch_start_seq_idx + batch_size - 1 + + def flash_attention_step1_qk_softmax( + q, # [bq_sz * num_q_heads, head_dim] + kv, # [bkv_sz, head_dim] <- Correspond to data from bkv_*_x2_ref + swa_m, # [bq_sz * num_q_heads], + swa_l, # [bq_sz * num_q_heads], + attention_sinks, # [num_q_heads] + ): + assert len(q.shape) == 2 + assert len(kv.shape) == 2 + assert q.shape[0] % num_q_heads == 0 + assert q.shape[1] == head_dim + assert kv.shape == (bkv_sz, head_dim) + + # Follow FlashAttention-2 forward pass. + s = jnp.einsum("nd,md->nm", q, kv, preferred_element_type=jnp.float32) + s *= sm_scale + + s_rowmax = jnp.max(s, axis=1, keepdims=True) + m_prev = swa_m + m_curr = jnp.maximum(m_prev, s_rowmax) + p = jnp.exp(s - m_curr) + exp_m_diff = jnp.exp(m_prev - m_curr) + p_rowsum = jnp.sum(p, axis=1, keepdims=True) + l_prev = swa_l + l_curr = exp_m_diff * l_prev + p_rowsum + exp_attention_sinks = jnp.exp(attention_sinks - m_curr) + l = l_curr + exp_attention_sinks + + return p, exp_m_diff, l + + def flash_attention_step2_pv( + p, + kv, + exp_m_diff, + swa_acc, + l, + ): + pv = jnp.einsum("nm,md->nd", p, kv, preferred_element_type=jnp.float32) + + o_prev = swa_acc + acc = exp_m_diff * o_prev + pv + out = ( + lax.div(acc, l) + if q_dtype == jnp.float32 + else (acc * pl.reciprocal(l, approx=True)).astype(q_dtype) + ) + return out + + def _async_copy(src, dst, sem, wait): + cp = pltpu.make_async_copy(src, dst, sem) + if wait: + cp.wait() + else: + cp.start() + + def _fetch_bkv(seq_idx, bkv_sem_idx, batch_idx, *, wait=False): + sem_nope = sems.at[0, bkv_sem_idx, batch_idx] + sem_rope = sems.at[6, bkv_sem_idx, batch_idx] + + bkv_nope_vmem_ref = bkv_nope_x2_ref.at[bkv_sem_idx, batch_idx] + bkv_rope_vmem_ref = bkv_rope_x2_ref.at[bkv_sem_idx, batch_idx] + + # The index into cache_kv_hbm_ref should be relative to the current + # chunk. + page_idx = seq_idx - start_seq_idx + if not wait: + _async_copy( + cache_kv_nope_hbm_ref.at[page_idx], + bkv_nope_vmem_ref, + sem_nope, + wait, + ) + _async_copy( + cache_kv_rope_hbm_ref.at[page_idx], + bkv_rope_vmem_ref, + sem_rope, + wait, + ) + else: + # When we wait, we can use a dummy copy to wait for DMAs to complete where + # src == dst. However, the dma size must be correct. + dst_nope = bkv_nope_vmem_ref + _async_copy( + src=dst_nope, + dst=dst_nope, + sem=sem_nope, + wait=True, + ) + dst_rope = bkv_rope_vmem_ref + _async_copy( + src=dst_rope, + dst=dst_rope, + sem=sem_rope, + wait=True, + ) + + def _fetch_bq(seq_idx, bq_sem_idx, batch_idx, *, wait=False): + sem = sems.at[1, bq_sem_idx, batch_idx] + bq_vmem_ref = bq_x2_ref.at[bq_sem_idx, batch_idx] + + _async_copy( + q_hbm_ref.at[seq_idx], + bq_vmem_ref, + sem, + wait, + ) + + def _send_bo(seq_idx, bo_sem_idx, batch_idx, *, wait=False): + sem = sems.at[2, bo_sem_idx, batch_idx] + vmem_ref = bo_x2_ref.at[bo_sem_idx, batch_idx] + + _async_copy( + vmem_ref, + o_hbm_ref.at[seq_idx], + sem, + wait, + ) + + def _fetch_swa(seq_idx, bq_sem_idx, batch_idx, *, wait=False): + sem_acc = sems.at[3, bq_sem_idx, batch_idx] + sem_l = sems.at[4, bq_sem_idx, batch_idx] + sem_m = sems.at[5, bq_sem_idx, batch_idx] + + if not wait: + _async_copy( + swa_accumution_hbm_ref.at[seq_idx], + swa_acc_x2_ref.at[bq_sem_idx, batch_idx], + sem_acc, + wait=False, + ) + _async_copy( + swa_l_hbm_ref.at[seq_idx], + bl_x2_ref.at[bq_sem_idx, batch_idx], + sem_l, + wait=False, + ) + _async_copy( + swa_m_hbm_ref.at[seq_idx], + bm_x2_ref.at[bq_sem_idx, batch_idx], + sem_m, + wait=False, + ) + + else: + dst_acc = swa_acc_x2_ref.at[bq_sem_idx, batch_idx] + _async_copy(src=dst_acc, dst=dst_acc, sem=sem_acc, wait=True) + + dst_l = bl_x2_ref.at[bq_sem_idx, batch_idx] + _async_copy(src=dst_l, dst=dst_l, sem=sem_l, wait=True) + + dst_m = bm_x2_ref.at[bq_sem_idx, batch_idx] + _async_copy(src=dst_m, dst=dst_m, sem=sem_m, wait=True) + + def start_fetch_bkv(seq_idx, bkv_sem_idx, batch_idx): + return _fetch_bkv(seq_idx, bkv_sem_idx, batch_idx) + + def wait_fetch_bkv(seq_idx, bkv_sem_idx, batch_idx): + return _fetch_bkv(seq_idx, bkv_sem_idx, batch_idx, wait=True) + + def start_fetch_bq(seq_idx, bq_sem_idx, batch_idx): + return _fetch_bq(seq_idx, bq_sem_idx, batch_idx) + + def wait_fetch_bq(seq_idx, bq_sem_idx, batch_idx): + return _fetch_bq(seq_idx, bq_sem_idx, batch_idx, wait=True) + + def start_fetch_swa(seq_idx, bq_sem_idx, batch_idx): + return _fetch_swa(seq_idx, bq_sem_idx, batch_idx) + + def wait_fetch_swa(seq_idx, bq_sem_idx, batch_idx): + return _fetch_swa(seq_idx, bq_sem_idx, batch_idx, wait=True) + + def start_send_bo(seq_idx, bo_sem_idx, batch_idx): + bo_ids_ref[bo_sem_idx, batch_idx] = seq_idx + _send_bo(seq_idx, bo_sem_idx, batch_idx) + + def wait_send_bo(bo_sem_idx, batch_idx): + old_seq_idx = bo_ids_ref[bo_sem_idx, batch_idx] + + @pl.when(0 <= old_seq_idx) + def _(): + _send_bo(old_seq_idx, bo_sem_idx, batch_idx, wait=True) + + def load_bq(bq_sem_idx, batch_idx): + q = bq_x2_ref.at[bq_sem_idx, batch_idx][...] + return q + + def load_bkv(bkv_sem_idx, batch_idx): + bkv_nope = bkv_nope_x2_ref.at[bkv_sem_idx, batch_idx][...] + bkv_nope = _dequant_dsv4_fp8(bkv_nope) + + bkv_rope = bkv_rope_x2_ref.at[bkv_sem_idx, batch_idx][...] + bkv = jnp.concatenate([bkv_nope, bkv_rope], axis=-1) + + # In vLLM, multiple caches may overlay on the same KV Tensor. For example, + # compressor state cache write data in bfloat16 / float32 format, certain + # byte pattern are interpreted as NaN in FP8, e.g. float8_e8m0fnu byte 0xFF + # decodes to NaN. + # We need to mask out the data by the actual kv_len to avoid NaN propagting + # to the downstream computation. + kv_len = kv_lens_ref[batch_start_seq_idx + batch_idx] + k_span = lax.broadcasted_iota(jnp.int32, bkv.shape, 0) + bkv = jnp.where(k_span < kv_len, bkv, 0) + return bkv + + def load_swa_output(bq_sem_idx, batch_idx): + swa_acc = swa_acc_x2_ref[bq_sem_idx, batch_idx, ...] + swa_l = bl_x2_ref[bq_sem_idx, batch_idx, :num_q_heads][..., None] + swa_m = bm_x2_ref[bq_sem_idx, batch_idx, :num_q_heads][..., None] + return swa_acc, swa_l, swa_m + + def process(): + + def get_next_seq_ids(seq_idx, bi_sem_idx): + next_seq_idx = seq_idx + batch_size + next_bi_sem_idx = lax.select(bi_sem_idx == 0, 1, 0) + return next_seq_idx, next_bi_sem_idx + + bi_sem_idx = sem_ids_ref[0] + next_seq_idx, next_bi_sem_idx = get_next_seq_ids( + batch_start_seq_idx, bi_sem_idx + ) + + # Prefetch next seq + @pl.when(next_seq_idx < end_seq_idx) + def prefetch_next_seq(): + sem_ids_ref[0] = next_bi_sem_idx + for batch_idx in range(batch_size): + start_fetch_bq(next_seq_idx + batch_idx, next_bi_sem_idx, batch_idx) + start_fetch_swa(next_seq_idx + batch_idx, next_bi_sem_idx, batch_idx) + start_fetch_bkv(next_seq_idx + batch_idx, next_bi_sem_idx, batch_idx) + + bo_sem_idx = sem_ids_ref[1] + sem_ids_ref[1] = lax.select(bo_sem_idx == 0, 1, 0) + attention_sinks = attention_sinks_ref[...][..., None] + + prev_p = None + prev_bkv = None + prev_exp_m_diff = None + prev_l = None + prev_swa_acc = None + + for batch_idx in range(batch_size): + + # Wait for cur blocks if not ready yet + wait_fetch_bq(batch_start_seq_idx + batch_idx, bi_sem_idx, batch_idx) + wait_fetch_swa(batch_start_seq_idx + batch_idx, bi_sem_idx, batch_idx) + wait_fetch_bkv(batch_start_seq_idx + batch_idx, bi_sem_idx, batch_idx) + + bkv = load_bkv(bi_sem_idx, batch_idx) + bq = load_bq(bi_sem_idx, batch_idx) + swa_acc, swa_l, swa_m = load_swa_output(bi_sem_idx, batch_idx) + + p, exp_m_diff, l = flash_attention_step1_qk_softmax( + bq, + bkv, + swa_m, + swa_l, + attention_sinks, + ) + + if prev_p is not None: + assert prev_bkv is not None + assert prev_exp_m_diff is not None + assert prev_l is not None + out = flash_attention_step2_pv( + prev_p, + prev_bkv, + prev_exp_m_diff, + prev_swa_acc, + prev_l, + ) + + # Wait for previous bo to be fully sent before storing new bo. + wait_send_bo(bo_sem_idx, batch_idx - 1) + # Store output from acc to bo. + bo_x2_ref.at[bo_sem_idx, batch_idx - 1][...] = out + # Send cur bo + start_send_bo( + batch_start_seq_idx + batch_idx - 1, bo_sem_idx, batch_idx - 1 + ) + + prev_p = p + prev_bkv = bkv + prev_exp_m_diff = exp_m_diff + prev_l = l + prev_swa_acc = swa_acc + + # end of pipelining loop + assert prev_p is not None + assert prev_bkv is not None + assert prev_exp_m_diff is not None + assert prev_l is not None + out = flash_attention_step2_pv( + prev_p, + prev_bkv, + prev_exp_m_diff, + prev_swa_acc, + prev_l, + ) + + # Wait for previous bo to be fully sent before storing new bo. + wait_send_bo(bo_sem_idx, batch_size - 1) + # Store output from acc to bo. + bo_x2_ref.at[bo_sem_idx, batch_size - 1][...] = out + # Send cur bo + start_send_bo( + batch_start_seq_idx + batch_size - 1, bo_sem_idx, batch_size - 1 + ) + + ### ------- Kernel start ------- ### + + @pl.when(batch_start_seq_idx == start_seq_idx) + def prologue(): + for batch_idx in range(batch_size): + start_fetch_bq(batch_start_seq_idx + batch_idx, 0, batch_idx) + start_fetch_swa(batch_start_seq_idx + batch_idx, 0, batch_idx) + start_fetch_bkv(batch_start_seq_idx + batch_idx, 0, batch_idx) + + process() + + @pl.when(batch_end_seq_idx == end_seq_idx - 1) + def epilogue(): + for i in range(2): + for batch_idx in range(batch_size): + wait_send_bo(i, batch_idx) + + ### ------- Kernel end ------- ### + + +def prepare_q_inputs( + q: jax.Array, # [max_num_tokens, actual_num_q_heads, actual_head_dim], +): + _, actual_num_q_heads, actual_head_dim = q.shape + q_packing = get_dtype_packing(q.dtype) + num_q_heads = align_to(actual_num_q_heads, q_packing) + head_dim = align_to(actual_head_dim, 128) + q = jnp.pad( + q, + ( + (0, 0), + (0, num_q_heads - actual_num_q_heads), + (0, head_dim - actual_head_dim), + ), + constant_values=0, + ) + return q + + +def prepare_swa_inputs( + swa_accumution: jax.Array, # [max_num_tokens, num_q_heads, head_dim] + swa_l: jax.Array, # [max_num_tokens, num_q_heads] + swa_m: jax.Array, # [max_num_tokens, num_q_heads] +): + _, actual_num_q_heads, actual_head_dim = swa_accumution.shape + swa_packing = get_dtype_packing(swa_accumution.dtype) + num_q_heads = align_to(actual_num_q_heads, swa_packing) + head_dim = align_to(actual_head_dim, 128) + swa_accumution = jnp.pad( + swa_accumution, + ( + (0, 0), + (0, num_q_heads - actual_num_q_heads), + (0, head_dim - actual_head_dim), + ), + constant_values=0, + ) + num_l_heads = align_to(num_q_heads, 128) + swa_l = jnp.pad( + swa_l, + ( + (0, 0), + (0, num_l_heads - actual_num_q_heads), + ), + constant_values=0, + ) + swa_m = jnp.pad( + swa_m, + ( + (0, 0), + (0, num_l_heads - actual_num_q_heads), + ), + constant_values=0, + ) + return swa_accumution, swa_l, swa_m + + +def prepare_outputs( + out, # [max_num_tokens, num_q_heads, head_dim] + actual_num_q_heads: int, + actual_head_dim: int, +): + return out[:, :actual_num_q_heads, :actual_head_dim] + + +# Main Attention kernel for DeepSeek V4 CSA (gather and attention) +# Note that the compressed kv tokens of current batch (current forward pass) +# have been written to the `cache_kv` by the compressor module before calling +# this function, `kv_lens` reflects the length after compressed kv cache write. +@functools.partial( + jax.jit, + static_argnames=( + "sm_scale", + "attention_kernel_batch_size", + "gather_and_attention_chunk_size", + "vmem_limit_bytes", + ), +) +def sparse_ragged_paged_attention( + q: jax.Array, # [max_num_tokens, actual_num_q_heads, head_dim] + cache_kv_nope: jax.Array, # [total_num_pages, page_size, 4, 128] + cache_kv_rope: jax.Array, # [total_num_pages, page_size // 4, 4, 128] + topk_indices: jax.Array, # i32[max_num_tokens, csa_topk] + page_indices: jax.Array, # i32[max_num_seqs * pages_per_seq] + cu_q_lens: jax.Array, # i32[max_num_seqs + 1] + distribution: jax.Array, # i32[3] + attention_sinks: jax.Array, # float32[actual_num_q_heads] + swa_accumution: jax.Array, # bf16[max_num_tokens, num_q_heads, head_dim] + swa_l: jax.Array, # float32[max_num_tokens, num_q_heads] + swa_m: jax.Array, # float32[max_num_tokens, num_q_heads] + *, + sm_scale: float = 1.0, + # Kernel optimization params. + gather_and_attention_chunk_size: int | None = None, + attention_kernel_batch_size: int = 16, + vmem_limit_bytes: int = DEFAULT_VMEM_LIMIT_BYTES, +) -> jax.Array: + """MLA Ragged paged attention that supports mixed prefill and decode. + + Args: + q: concatenated all sequences' queries. + cache_kv_nope: the current kv cache for nope. + cache_kv_rope: the current kv cache for rope. + topk_indices: for each query token, the indices of the top k key tokens to + attend to. + page_indices: flattened page indices look-up table by (seq_id, page_id). + cu_q_lens: the cumulative sum of the effective query lengths. Similar to + kv_lens, only the first num_seqs+1 values are valid. + distribution: (i, j, k) represents that sequences[0:i] are decode-only, + sequences[i:j] are chunked-prefill-only, and sequences[j:k] are mixed. The + k is also the total number of sequences. + sm_scale: the softmax scale which will be applied to the Q@K^T. + vmem_limit_bytes: the vmem limit for the pallas kernel. + + Returns: + The output of attention. + """ + # The cache is DSV4 FP8 format. + # nope_cache contains 448 fp8 + 7 fp8 scales, + # rope_cache contains 64 bf16 + assert cache_kv_nope.dtype == jnp.uint8 + assert cache_kv_rope.dtype == jnp.uint8 + if gather_and_attention_chunk_size is None: + gather_and_attention_chunk_size = q.shape[0] + + _, actual_num_q_heads, actual_head_dim = q.shape + + q = prepare_q_inputs(q) # [max_num_tokens, num_q_heads, head_dim] + head_dim = q.shape[-1] + attention_sinks = jnp.pad( + attention_sinks, + (0, q.shape[1] - actual_num_q_heads), + constant_values=jnp.finfo(attention_sinks.dtype).min, + ) + assert swa_accumution.dtype == q.dtype + swa_accumution, swa_l, swa_m = prepare_swa_inputs( + swa_accumution, swa_l, swa_m + ) + + _, page_size, _, _ = cache_kv_nope.shape + + _, num_q_heads, _ = q.shape + max_num_seqs = cu_q_lens.shape[0] - 1 + num_page_indices = page_indices.shape[0] + assert num_page_indices % max_num_seqs == 0 + + def run_mla_kernel( + q: jax.Array, # [max_num_tokens, num_q_heads, head_dim] + cache_kv_nope: jax.Array, # [total_num_pages, page_size, nope_dim] + cache_kv_rope: jax.Array, # [total_num_pages, page_size, rope_dim] + kv_lens: jax.Array, # i32[max_num_seqs] + attention_sinks: jax.Array, # float32[num_q_heads] + swa_accumution: jax.Array, # bf16[max_num_tokens, num_q_heads, head_dim] + swa_l: jax.Array, # float32[max_num_tokens, num_l_heads] + swa_m: jax.Array, # float32[max_num_tokens, num_l_heads] + start_seq_idx: jax.Array, # i32 + end_seq_idx: jax.Array, # i32 + kernel_batch_size: int, + ): + batch_size = kernel_batch_size + end_seq_idx = jnp.maximum(start_seq_idx, end_seq_idx) + grid = (cdiv(end_seq_idx - start_seq_idx, batch_size),) + in_specs = [ + pl.BlockSpec(memory_space=pltpu.VMEM), # attention_sinks + pl.BlockSpec(memory_space=pltpu.HBM), # q + pl.BlockSpec(memory_space=pltpu.HBM), # cache_kv_nope + pl.BlockSpec(memory_space=pltpu.HBM), # cache_kv_rope + pl.BlockSpec(memory_space=pltpu.HBM), # swa_accumution + pl.BlockSpec(memory_space=pltpu.HBM), # swa_l + pl.BlockSpec(memory_space=pltpu.HBM), # swa_m + ] + + out_specs = pl.BlockSpec(memory_space=pltpu.HBM) # o + + page_size = cache_kv_nope.shape[1] + bkv_nope_double_buf = pltpu.VMEM( + (2, batch_size, page_size, *cache_kv_nope.shape[2:]), + cache_kv_nope.dtype, + ) + bkv_rope_double_buf = pltpu.VMEM( + (2, batch_size, page_size, *cache_kv_rope.shape[2:]), + cache_kv_rope.dtype, + ) + + bq_double_bufq = pltpu.VMEM( + (2, batch_size, num_q_heads, head_dim), + q.dtype, + ) + + bo_double_buf = bq_double_bufq + + num_l_heads = align_to(num_q_heads, 128) + bl_double_buf = pltpu.VMEM( + (2, batch_size, num_l_heads), + jnp.float32, + ) + bm_double_buf = bl_double_buf + + swa_acc_double_buf = pltpu.VMEM( + (2, batch_size, num_q_heads, head_dim), + q.dtype, + ) + + scratch_shapes = [ + bkv_nope_double_buf, + bkv_rope_double_buf, + bq_double_bufq, + bo_double_buf, # Double buffering for output block. + bl_double_buf, # Double buffering for l output. + bm_double_buf, # Double buffering for m output. + swa_acc_double_buf, # Buffer for swa_accumution. + # Semaphores for double buffering of bkv_nope, bq, bo, swa_acc, swa_l, swa_m, bkv_rope + pltpu.SemaphoreType.DMA((7, 2, batch_size)), + ] + + scalar_prefetches = ( + kv_lens, + jnp.array([start_seq_idx, end_seq_idx], jnp.int32), + # (bi_sem_idx, bo_sem_idx) + jnp.zeros((2,), jnp.int32), + # (bo_sem_0_seq_idx, bo_sem_1_seq_idx) + jnp.full((2, batch_size), -1, jnp.int32), + ) + + scope_name = f"MLA-p_{page_size}" + kernel = jax.named_scope(scope_name)( + pl.pallas_call( + functools.partial( + _attention_kernel, + sm_scale=sm_scale, + batch_size=batch_size, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=len(scalar_prefetches), + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + scratch_shapes=scratch_shapes, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary",), + vmem_limit_bytes=vmem_limit_bytes, + disable_bounds_checks=True, + ), + out_shape=jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype), + input_output_aliases={ + 5: 0, # Alias output activation with q + }, + name=scope_name, + ) + ) + return kernel( + *scalar_prefetches, + attention_sinks, + q, + cache_kv_nope, + cache_kv_rope, + swa_accumution, + swa_l, + swa_m, + ) + + tokens_per_seq = cu_q_lens[1:] - cu_q_lens[:-1] + seq_ids_segment = jnp.repeat( + jnp.arange(max_num_seqs), tokens_per_seq, total_repeat_length=q.shape[0] + ) + assert topk_indices is not None + # TODO: skip gather for padding tokens in topk_indices. + kv_lens = jnp.sum(topk_indices != -1, axis=-1) + + seq_page_ids = topk_indices // page_size + token_offset = topk_indices % page_size + topk = topk_indices.shape[-1] + page_ids = gather_page_ids( + page_indices, seq_page_ids, seq_ids_segment, max_num_seqs + ) + + # For the "-1" padding elements in topk_indices, we scatter the corresponding + # page_ids and token_offset to avoid gather memory access hotspotting. + is_padding = topk_indices == -1 + total_num_pages = cache_kv_nope.shape[0] + flat_element_index = jnp.arange(q.shape[0] * topk, dtype=jnp.int32).reshape( + q.shape[0], topk + ) + # 104729 and 15485863 are randomly chosen large prime numbers. + scattered_page_ids = (flat_element_index * 104729) % total_num_pages + scattered_token_offset = (flat_element_index * 15485863) % page_size + page_ids = jnp.where(is_padding, scattered_page_ids, page_ids) + token_offset = jnp.where( + is_padding, + scattered_token_offset, + token_offset, + ) + + assert page_ids.shape == (q.shape[0], topk) + + # TODO: handle the case where q.shape[0] is not divisible by + # gather_and_attention_chunk_size. + assert q.shape[0] % gather_and_attention_chunk_size == 0 + num_chunks = q.shape[0] // gather_and_attention_chunk_size + + for i in range(num_chunks): + start_pos = i * gather_and_attention_chunk_size + end_pos = start_pos + gather_and_attention_chunk_size + indices = ( + page_ids[start_pos:end_pos, ...] * page_size + + token_offset[start_pos:end_pos, ...] + ).reshape(-1) + + # For prefilling of short sequences (or early in the sequence), there are + # very few number of KVs in the sequence, so different qs' selected topk + # would have large overlap. This causes gather read hotspotting. We've seen + # 30%+ performance degradation compared to the no-duplicate-indices case. + # + # TODO: we could consider let the caller (tpu-runner) to sort the sequences + # based on their lengths. For the sequences-segment below certain length, + # we use a different kernel (dense attention and mask), for the rest of + # sequences, we use this gather-and-attention kernel. + gathered_nope_buffer, gathered_rope_buffer = csa_gather( + cache_kv_nope, + cache_kv_rope, + indices, + ) + gathered_nope_buffer = gathered_nope_buffer.reshape( + gather_and_attention_chunk_size, topk, -1 + ) + gathered_rope_buffer = gathered_rope_buffer.reshape( + gather_and_attention_chunk_size, topk, -1 + ) + # We treat each query token as a one independent sequence, attend to their + # respective gathered kv tokens in the `gathered_kv_buffer`. + # -1 in topk_indices is padded elements at the end of each row. + # Batching + assert gather_and_attention_chunk_size % attention_kernel_batch_size == 0 + batch_end = ( + cdiv( + jnp.minimum( + cu_q_lens[distribution[2]], + start_pos + gather_and_attention_chunk_size, + ), + attention_kernel_batch_size, + ) + * attention_kernel_batch_size + ) + q = run_mla_kernel( + q, + gathered_nope_buffer, + gathered_rope_buffer, + kv_lens, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + start_seq_idx=start_pos, + end_seq_idx=batch_end, + kernel_batch_size=attention_kernel_batch_size, + ) + return prepare_outputs( + q, actual_num_q_heads, actual_head_dim + ) # [max_num_tokens, actual_num_q_heads, actual_head_dim] + +def computation( + q: jax.Array, + cache_kv_nope: jax.Array, + cache_kv_rope: jax.Array, + topk_indices: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + distribution: jax.Array, + attention_sinks: jax.Array, + swa_accumution: jax.Array, + swa_l: jax.Array, + swa_m: jax.Array, +): + sm_scale = float(512 ** -0.5) + gather_and_attention_chunk_size = None + attention_kernel_batch_size = 16 + vmem_limit_bytes = 100 * 1024 * 1024 + + return sparse_ragged_paged_attention( + q, + cache_kv_nope, + cache_kv_rope, + topk_indices, + page_indices, + cu_q_lens, + distribution, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + sm_scale=sm_scale, + gather_and_attention_chunk_size=gather_and_attention_chunk_size, + attention_kernel_batch_size=attention_kernel_batch_size, + vmem_limit_bytes=vmem_limit_bytes, + ) diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/52p_DeepSeek_V4_HCA/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/52p_DeepSeek_V4_HCA/kernel_task.yaml new file mode 100644 index 0000000..af7f247 --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/52p_DeepSeek_V4_HCA/kernel_task.yaml @@ -0,0 +1,107 @@ +task_id: 52p_DeepSeek_V4_HCA +description: DeepSeek-V4 hybrid context attention (HCA) +input_gen_code: |- + def get_inputs(): + import jax + import jax.numpy as jnp + + HEAD_DIM = 512 + NUM_Q_HEADS = 8 + Q_DTYPE = jnp.bfloat16 + SM_SCALE = float(HEAD_DIM ** -0.5) + VMEM_LIMIT_BYTES = 100 * 1024 * 1024 + + CONFIGS = [ + (256, 1, 9216, 1024, 16, 1), + (1, 1024, 1024, 1024, 16, 32), + (1, 1024, 8192, 1024, 16, 32), + (128, 1, 9216, 1024, 16, 1), + (1, 256, 1024, 1024, 16, 32), + (1, 512, 4096, 1024, 16, 32), + (128, 1, 9216, 256, 16, 1), + (1, 256, 1024, 256, 16, 32), + (256, 1, 9216, 2048, 16, 1), + (1, 1024, 8192, 2048, 4, 32), + ] + + CACHE_TILE = 8192 + k_kv, k_cfg = jax.random.split(jax.random.key(0), 2) + + kv_tile = jax.random.normal( + k_kv, (CACHE_TILE, 4, 128), jnp.float32 + ).astype(Q_DTYPE) + bits = jax.lax.bitcast_convert_type(kv_tile, jnp.uint16) + kv_tile_bytes = jnp.stack( + [(bits & 0xFF).astype(jnp.uint8), (bits >> 8).astype(jnp.uint8)], + axis=2, + ).reshape(CACHE_TILE, 8, 128) + + outputs = [] + + keys = jax.random.split(k_cfg, len(CONFIGS)) + for cfg, key in zip(CONFIGS, keys): + batch_size, q_len, kv_len, page_size, bkv_pages, bq = cfg + k_page, k_q, k_sink, k_acc, k_l, k_m = jax.random.split(key, 6) + + num_tokens = batch_size * q_len + pages_per_seq = -(-kv_len // page_size) + total_num_pages = batch_size * pages_per_seq + num_slots = total_num_pages * page_size + + q = jax.random.normal( + k_q, (num_tokens, NUM_Q_HEADS, HEAD_DIM), jnp.float32 + ).astype(Q_DTYPE) + + reps = -(-num_slots // CACHE_TILE) + cache_kv = jnp.tile(kv_tile_bytes, (reps, 1, 1))[:num_slots].reshape( + total_num_pages, page_size * 2, 4, 128 + ) + + page_indices = jax.random.permutation( + k_page, total_num_pages + ).astype(jnp.int32) + + kv_lens = jnp.full((batch_size,), kv_len, jnp.int32) + pos_in_seq = jnp.arange(num_tokens, dtype=jnp.int32) % q_len + kv_lens_to_attend = kv_len - q_len + pos_in_seq + 1 + cu_q_lens = jnp.arange(batch_size + 1, dtype=jnp.int32) * q_len + + if q_len == 1: + distribution = jnp.array( + [batch_size, batch_size, batch_size], jnp.int32 + ) + chunk_prefill_size = None + else: + distribution = jnp.array([0, batch_size, batch_size], jnp.int32) + chunk_prefill_size = q_len + + attention_sinks = jax.random.normal( + k_sink, (NUM_Q_HEADS,), jnp.float32 + ) + swa_accumution = jax.random.normal( + k_acc, (num_tokens, NUM_Q_HEADS, HEAD_DIM), jnp.float32 + ).astype(Q_DTYPE) + swa_l = jax.random.uniform( + k_l, (num_tokens, NUM_Q_HEADS), jnp.float32, 1.0, 64.0 + ) + swa_m = jax.random.normal(k_m, (num_tokens, NUM_Q_HEADS), jnp.float32) + + dynamic_args = [ + q, + cache_kv, + kv_lens, + kv_lens_to_attend, + page_indices, + cu_q_lens, + distribution, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + ] + outputs.append((dynamic_args, [])) + + return outputs + +rtol: 0.01 +atol: 0.01 diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/52p_DeepSeek_V4_HCA/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/52p_DeepSeek_V4_HCA/reference.py new file mode 100644 index 0000000..6fb4b67 --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/52p_DeepSeek_V4_HCA/reference.py @@ -0,0 +1,1028 @@ + +import functools +from enum import Enum +import jax +from jax import lax +import jax.numpy as jnp +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu + +DEFAULT_VMEM_LIMIT_BYTES = 100 * 1024 * 1024 + +def get_inputs(): + import jax + import jax.numpy as jnp + + HEAD_DIM = 512 + NUM_Q_HEADS = 8 + Q_DTYPE = jnp.bfloat16 + + CONFIGS = [ + (256, 1, 9216, 1024, 16, 1), + (1, 1024, 1024, 1024, 16, 32), + (1, 1024, 8192, 1024, 16, 32), + (128, 1, 9216, 1024, 16, 1), + (1, 256, 1024, 1024, 16, 32), + (1, 512, 4096, 1024, 16, 32), + (128, 1, 9216, 256, 16, 1), + (1, 256, 1024, 256, 16, 32), + (256, 1, 9216, 2048, 16, 1), + (1, 1024, 8192, 2048, 4, 32), + ] + + CACHE_TILE = 8192 + k_kv, k_cfg = jax.random.split(jax.random.key(0), 2) + + kv_tile = jax.random.normal( + k_kv, (CACHE_TILE, 4, 128), jnp.float32 + ).astype(Q_DTYPE) + bits = jax.lax.bitcast_convert_type(kv_tile, jnp.uint16) + kv_tile_bytes = jnp.stack( + [(bits & 0xFF).astype(jnp.uint8), (bits >> 8).astype(jnp.uint8)], + axis=2, + ).reshape(CACHE_TILE, 8, 128) + + outputs = [] + + keys = jax.random.split(k_cfg, len(CONFIGS)) + for cfg, key in zip(CONFIGS, keys): + batch_size, q_len, kv_len, page_size, bkv_pages, bq = cfg + k_page, k_q, k_sink, k_acc, k_l, k_m = jax.random.split(key, 6) + + num_tokens = batch_size * q_len + pages_per_seq = -(-kv_len // page_size) + total_num_pages = batch_size * pages_per_seq + num_slots = total_num_pages * page_size + + q = jax.random.normal( + k_q, (num_tokens, NUM_Q_HEADS, HEAD_DIM), jnp.float32 + ).astype(Q_DTYPE) + + reps = -(-num_slots // CACHE_TILE) + cache_kv = jnp.tile(kv_tile_bytes, (reps, 1, 1))[:num_slots].reshape( + total_num_pages, page_size * 2, 4, 128 + ) + + page_indices = jax.random.permutation( + k_page, total_num_pages + ).astype(jnp.int32) + + kv_lens = jnp.full((batch_size,), kv_len, jnp.int32) + pos_in_seq = jnp.arange(num_tokens, dtype=jnp.int32) % q_len + kv_lens_to_attend = kv_len - q_len + pos_in_seq + 1 + cu_q_lens = jnp.arange(batch_size + 1, dtype=jnp.int32) * q_len + + if q_len == 1: + distribution = jnp.array( + [batch_size, batch_size, batch_size], jnp.int32 + ) + chunk_prefill_size = None + else: + distribution = jnp.array([0, batch_size, batch_size], jnp.int32) + chunk_prefill_size = q_len + + attention_sinks = jax.random.normal( + k_sink, (NUM_Q_HEADS,), jnp.float32 + ) + swa_accumution = jax.random.normal( + k_acc, (num_tokens, NUM_Q_HEADS, HEAD_DIM), jnp.float32 + ).astype(Q_DTYPE) + swa_l = jax.random.uniform( + k_l, (num_tokens, NUM_Q_HEADS), jnp.float32, 1.0, 64.0 + ) + swa_m = jax.random.normal(k_m, (num_tokens, NUM_Q_HEADS), jnp.float32) + + dynamic_args = [ + q, + cache_kv, + kv_lens, + kv_lens_to_attend, + page_indices, + cu_q_lens, + distribution, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + ] + outputs.append((dynamic_args, [])) + + return outputs + +def cdiv(a, b): + assert b != 0 + return (a + b - 1) // b + +def align_to(x, a): + return cdiv(x, a) * a + +def get_dtype_bitwidth(dtype): + return jax.dtypes.itemsize_bits(dtype) + +def get_dtype_packing(dtype): + bits = get_dtype_bitwidth(dtype) + return 32 // bits + +class MlaCase(Enum): + """Represents the different cases for MLA. + + - DECODE: Sequences are in decode-only mode (q_len = 1). + - PREFILL: Sequences are in prefill-only mode (q_len > 1, static). + - MIXED: Sequences can be a mix of prefill and decode (q_len > 1, dynamic). + """ + + DECODE = 0 + PREFILL = 1 + MIXED = 2 + + @property + def symbol(self): + return { + MlaCase.DECODE: "d", + MlaCase.PREFILL: "p", + MlaCase.MIXED: "m", + }[self] + + +def _mla_ragged_paged_attention_kernel( + # Prefetch + kv_lens_ref, # [max_num_seqs] + kv_lens_to_attend_ref, # [max_num_tokens] + page_indices_ref, # [max_num_seqs * pages_per_seq] + cu_q_lens_ref, # [max_num_seqs + 1] + start_end_seq_idx_ref, # [2] (start_seq_idx, end_seq_idx) + sem_ids_ref, # [3] (bq_sem_idx, bkv_sem_idx, bo_sem_idx) + bo_ids_ref, # [4] (bo_sem_0_seq_idx, bo_sem_1_seq_idx, bo_sem_0_bo_idx, bo_sem_1_bo_idx) + # Input + attention_sinks_ref, # float32[num_q_heads] + q_hbm_ref, # [max_num_tokens, num_q_heads, head_dim] + cache_kv_hbm_ref, # [total_num_pages, page_size_per_kv_packing, kv_packing, lkv_dim] + swa_accumution_hbm_ref, # [max_num_tokens, num_q_heads, head_dim] + swa_l_hbm_ref, # [max_num_tokens, num_l_heads] + swa_m_hbm_ref, # [max_num_tokens, num_l_heads] + # Output + o_hbm_ref, # [max_num_tokens, num_q_heads, head_dim] + # Scratch + bkv_x2_ref, # [2, bkv_buf_sz_per_kv_packing, kv_packing, lkv_dim] + bq_x2_ref, # [2, bq_sz, num_q_heads, head_dim] + bo_x2_ref, # [2, bq_sz, num_q_heads, head_dim] + bl_x2_ref, # [2, bq_sz, num_l_heads] + bm_x2_ref, # [2, bq_sz, num_l_heads] + swa_acc_x2_ref, # [2, bq_sz, num_q_heads, head_dim] + sems, # [7, 2] + l_ref, # [bq_sz * num_q_heads, 128], + m_ref, # [bq_sz * num_q_heads, 128], + acc_ref, # [bq_sz * num_q_heads, head_dim], + *, + static_q_len: int, + sm_scale: float, + bkv_p, + bq_sz, +): + assert q_hbm_ref.shape == o_hbm_ref.shape + + _, num_q_heads, head_dim = q_hbm_ref.shape + total_num_pages, num_slots, kv_packing, lkv_dim = cache_kv_hbm_ref.shape + num_slots_per_token = 2 + assert num_slots % num_slots_per_token == 0 + page_size = num_slots // num_slots_per_token + max_num_seqs = kv_lens_ref.shape[0] + num_page_indices = page_indices_ref.shape[0] + + assert num_page_indices % max_num_seqs == 0 + pages_per_seq = num_page_indices // max_num_seqs + q_dtype = q_hbm_ref.dtype + q_packing = get_dtype_packing(q_dtype) + # Validate against the KV dtype. + kv_dtype = cache_kv_hbm_ref.dtype + assert o_hbm_ref.dtype == q_dtype + assert get_dtype_packing(kv_dtype) == kv_packing + assert lkv_dim % 128 == 0 + assert head_dim % 128 == 0 + bkv_sz = bkv_p * page_size + assert num_q_heads % q_packing == 0 + num_q_heads_per_q_packing = num_q_heads // q_packing + + start_seq_idx = start_end_seq_idx_ref[0] + end_seq_idx = start_end_seq_idx_ref[1] + seq_idx = pl.program_id(0) + start_seq_idx + q_start = cu_q_lens_ref[seq_idx] + q_end = cu_q_lens_ref[seq_idx + 1] + q_len = q_end - q_start + kv_len = kv_lens_ref[seq_idx] + + def flash_attention( + q, # [bq_sz * num_q_heads, head_dim] + kv, # [bkv_sz, head_dim] <- Correspond to data from bkv_x2_ref + *, + bq_idx, + bkv_idx, + kv_lens_to_attend_segment, + ): + assert len(q.shape) == 2 + assert len(kv.shape) == 2 + assert q.shape[0] % num_q_heads == 0 + assert q.shape[1] == head_dim + assert kv.shape == (bkv_sz, head_dim) + head_l_ref = l_ref.at[: q.shape[0]] + head_m_ref = m_ref.at[: q.shape[0]] + head_acc_ref = acc_ref.at[: q.shape[0]] + + # Follow FlashAttention-2 forward pass. + s = jnp.einsum("nd,md->nm", q, kv, preferred_element_type=jnp.float32) + s *= sm_scale + + k_span = bkv_idx * bkv_sz + lax.broadcasted_iota(jnp.int32, s.shape, 1) + mask = kv_lens_to_attend_segment.reshape(s.shape) <= k_span + + s = jnp.where(mask, jnp.finfo(s.dtype).min, s) + s_rowmax = jnp.max(s, axis=1, keepdims=True) + m_prev = head_m_ref[...] + m_curr = jnp.maximum(m_prev, s_rowmax) + head_m_ref[...] = m_curr + p = jnp.exp(s - broadcast_minor(m_curr, s.shape)) + + pv = jnp.einsum("nm,md->nd", p, kv, preferred_element_type=jnp.float32) + + p_rowsum = jnp.sum(p, axis=1, keepdims=True) + exp_m_diff = jnp.exp(m_prev - m_curr) + l_prev = head_l_ref[...] + l_curr = exp_m_diff * l_prev + p_rowsum + head_l_ref[...] = l_curr + o_prev = head_acc_ref[...] + o_curr = broadcast_minor(exp_m_diff, o_prev.shape) * o_prev + pv + head_acc_ref[...] = o_curr + + def _async_copy(src, dst, sem, wait): + cp = pltpu.make_async_copy(src, dst, sem) + if wait: + cp.wait() + else: + cp.start() + + def _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, *, wait=False): + sem = sems.at[0, bkv_sem_idx] + # bkv_x2_ref shape: [2, bkv_sz, num_slots_per_token * kv_packing, lkv_dim] + bkv_vmem_ref = bkv_x2_ref.at[bkv_sem_idx] + + reshaped_cache_hbm_ref = cache_kv_hbm_ref.reshape( + total_num_pages * page_size, + num_slots_per_token * kv_packing, + lkv_dim, + ) + + kv_len = kv_lens_ref[seq_idx] + kv_len_start = bkv_idx * bkv_sz + kv_p_start = bkv_idx * bkv_p + + kv_left = kv_len - kv_len_start + dma_bkv_sz = jnp.minimum(kv_left, bkv_sz) + page_indices_offset = seq_idx * pages_per_seq + kv_p_start + + if not wait: + # Fetch effective kv from kv cache. To pipeline multiple DMA calls, we + # utilize static for loop instead of dynamic for loop. + # Loop through all pages in a block + for i in range(bkv_p): + # Ensure only effective kvs are copied and we don't go negative. + sz = jnp.clip( + kv_left - i * page_size, + 0, + page_size, + ) + # If the page index is out of bound, we set page_idx to the last page. + # And there will be no copy since sz will be 0. + page_idx = jnp.minimum(page_indices_offset + i, num_page_indices - 1) + _async_copy( + reshaped_cache_hbm_ref.at[ + pl.ds( + page_indices_ref[page_idx] * page_size, + sz, + ), + ], + bkv_vmem_ref.at[pl.ds(i * page_size, sz)], + sem, + wait, + ) + + else: + # When we wait, we can use a dummy copy to wait for DMAs to complete where + # src == dst. However, the dma size must be correct. + dst_kv = bkv_vmem_ref.at[pl.ds(0, dma_bkv_sz)] + _async_copy( + src=dst_kv, + dst=dst_kv, + sem=sem, + wait=True, + ) + + def _fetch_bq(seq_idx, bq_idx, bq_sem_idx, *, wait=False): + sem = sems.at[1, bq_sem_idx] + bq_vmem_ref = bq_x2_ref.at[bq_sem_idx] + + q_len_start = cu_q_lens_ref[seq_idx] + bq_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy( + q_hbm_ref.at[pl.ds(q_len_start, sz)], + bq_vmem_ref.at[pl.ds(0, sz)], + sem, + wait, + ) + + def _send_bo(seq_idx, bo_idx, bo_sem_idx, *, wait=False): + sem = sems.at[2, bo_sem_idx] + vmem_ref = bo_x2_ref.at[bo_sem_idx] + q_len_start = cu_q_lens_ref[seq_idx] + bo_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy( + vmem_ref.at[pl.ds(0, sz)], + o_hbm_ref.at[pl.ds(q_len_start, sz)], + sem, + wait, + ) + + def _fetch_swa(seq_idx, bq_idx, bq_sem_idx, *, wait=False): + sem_acc = sems.at[3, bq_sem_idx] + sem_l = sems.at[4, bq_sem_idx] + sem_m = sems.at[5, bq_sem_idx] + + q_len_start = cu_q_lens_ref[seq_idx] + bq_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + if not wait: + _async_copy( + swa_accumution_hbm_ref.at[pl.ds(q_len_start, sz)], + swa_acc_x2_ref.at[bq_sem_idx, pl.ds(0, sz)], + sem_acc, + wait=False, + ) + _async_copy( + swa_l_hbm_ref.at[pl.ds(q_len_start, sz)], + bl_x2_ref.at[bq_sem_idx, pl.ds(0, sz)], + sem_l, + wait=False, + ) + _async_copy( + swa_m_hbm_ref.at[pl.ds(q_len_start, sz)], + bm_x2_ref.at[bq_sem_idx, pl.ds(0, sz)], + sem_m, + wait=False, + ) + + else: + dst_acc = swa_acc_x2_ref.at[bq_sem_idx, pl.ds(0, sz)] + _async_copy(src=dst_acc, dst=dst_acc, sem=sem_acc, wait=True) + + dst_l = bl_x2_ref.at[bq_sem_idx, pl.ds(0, sz)] + _async_copy(src=dst_l, dst=dst_l, sem=sem_l, wait=True) + + dst_m = bm_x2_ref.at[bq_sem_idx, pl.ds(0, sz)] + _async_copy(src=dst_m, dst=dst_m, sem=sem_m, wait=True) + + acc_ref[...] = ( + swa_acc_x2_ref[bq_sem_idx, ...] + .astype(jnp.float32) + .reshape(bq_sz * num_q_heads, head_dim) + ) + bl = jnp.concat( + [bl_x2_ref[bq_sem_idx, i, :num_q_heads] for i in range(bq_sz)] + )[..., None] + l_ref[...] = jnp.concat([bl for _ in range(128)], axis=-1) + + bm = jnp.concat( + [bm_x2_ref[bq_sem_idx, i, :num_q_heads] for i in range(bq_sz)] + )[..., None] + m_ref[...] = jnp.concat([bm for _ in range(128)], axis=-1) + + def start_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx): + return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx) + + def wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx): + return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, wait=True) + + def start_fetch_bq(seq_idx, bq_idx, bq_sem_idx): + return _fetch_bq(seq_idx, bq_idx, bq_sem_idx) + + def wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx): + return _fetch_bq(seq_idx, bq_idx, bq_sem_idx, wait=True) + + def start_fetch_swa(seq_idx, bq_idx, bq_sem_idx): + return _fetch_swa(seq_idx, bq_idx, bq_sem_idx) + + def wait_fetch_swa(seq_idx, bq_idx, bq_sem_idx): + return _fetch_swa(seq_idx, bq_idx, bq_sem_idx, wait=True) + + def start_send_bo(seq_idx, bo_idx, bo_sem_idx): + bo_ids_ref[bo_sem_idx] = seq_idx + bo_ids_ref[bo_sem_idx + 2] = bo_idx + _send_bo(seq_idx, bo_idx, bo_sem_idx) + + def wait_send_bo(bo_sem_idx): + old_seq_idx = bo_ids_ref[bo_sem_idx] + old_bo_idx = bo_ids_ref[bo_sem_idx + 2] + + @pl.when(jnp.logical_and(0 <= old_seq_idx, old_seq_idx <= seq_idx)) + def _(): + _send_bo(old_seq_idx, old_bo_idx, bo_sem_idx, wait=True) + + def load_bq(bq_sem_idx): + q_ref = ( + bq_x2_ref.bitcast(jnp.uint32) + .at[bq_sem_idx] + .reshape(bq_sz * num_q_heads_per_q_packing, head_dim) + ) + q = pltpu.bitcast( + q_ref[: bq_sz * num_q_heads_per_q_packing], + q_dtype, + ).reshape(bq_sz * num_q_heads, head_dim) + return q + + def load_bkv(bkv_sem_idx, bkv_idx): + bkv_u8 = pltpu.bitcast(bkv_x2_ref.at[bkv_sem_idx][...], jnp.uint8) + assert bkv_u8.shape[-1] == cache_kv_hbm_ref.shape[-1] + bkv_bf16 = pltpu.bitcast(bkv_u8, jnp.bfloat16) + bkv = bkv_bf16.reshape(-1, head_dim) + assert bkv.shape == (bkv_sz, head_dim) + return bkv + + def broadcast_minor(src, shape): + if src.shape == shape: + return src + assert src.shape[:-1] == shape[:-1] + assert src.shape[-1] % 128 == 0 + target_minor = align_to(shape[-1], src.shape[-1]) + # no-op concatenation. + return jnp.concatenate( + [src for _ in range(target_minor // src.shape[-1])], axis=-1 + )[..., : shape[-1]] + + def process(): + # Force at least one bkv block and one bq block per sequence: the + # double-buffered DMA pipeline hands the bkv and bq semaphore across + # sequence boundaries and assumes every sequence runs >=1 bkv and bq + # iteration. + num_bkv = jnp.maximum(1, cdiv(kv_len, bkv_sz)) + if static_q_len is None: + num_bq = jnp.maximum(1, cdiv(q_len, bq_sz)) + else: + num_bq = jnp.maximum(1, cdiv(static_q_len, bq_sz)) + + def get_next_bq_ids(seq_idx, bq_idx, bq_sem_idx): + next_bq_idx = bq_idx + 1 + is_last_bq = next_bq_idx == num_bq + next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx) + next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx) + next_bq_sem_idx = lax.select(bq_sem_idx == 0, 1, 0) + return next_seq_idx, next_bq_idx, next_bq_sem_idx + + def get_next_bkv_ids(seq_idx, bq_idx, bkv_idx, bkv_sem_idx): + next_bkv_idx = bkv_idx + 1 + is_last_bkv = next_bkv_idx == num_bkv + next_bkv_idx = lax.select(is_last_bkv, 0, next_bkv_idx) + next_bq_idx = lax.select(is_last_bkv, bq_idx + 1, bq_idx) + is_last_bq = next_bq_idx == num_bq + next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx) + next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx) + next_bkv_sem_idx = lax.select(bkv_sem_idx == 0, 1, 0) + return next_seq_idx, next_bq_idx, next_bkv_idx, next_bkv_sem_idx + + def compute_with_bq(bq_idx, _): + bq_sem_idx = sem_ids_ref[0] + next_seq_idx, next_bq_idx, next_bq_sem_idx = get_next_bq_ids( + seq_idx, bq_idx, bq_sem_idx + ) + + kv_lens_to_attend_segment = jnp.broadcast_to( + jnp.stack([ + kv_lens_to_attend_ref[q_start + bq_idx * bq_sz + i] + for i in range(bq_sz) + ])[:, None, None], + (bq_sz, num_q_heads, bkv_sz), + ) + + # Prefetch next bq + @pl.when(next_seq_idx < end_seq_idx) + def prefetch_next_bq(): + sem_ids_ref[0] = next_bq_sem_idx + start_fetch_bq(next_seq_idx, next_bq_idx, next_bq_sem_idx) + start_fetch_swa(next_seq_idx, next_bq_idx, next_bq_sem_idx) + + def compute_with_bkv(bkv_idx, carry): + kv_lens_to_attend_segment = carry[0] + + # Get next bkv ids. + bkv_sem_idx = sem_ids_ref[1] + next_seq_idx, _, next_bkv_idx, next_bkv_sem_idx = get_next_bkv_ids( + seq_idx, bq_idx, bkv_idx, bkv_sem_idx + ) + + # Prefetch next bkv + @pl.when(next_seq_idx < end_seq_idx) + def prefetch_next_bkv(): + sem_ids_ref[1] = next_bkv_sem_idx + start_fetch_bkv(next_seq_idx, next_bkv_idx, next_bkv_sem_idx) + + # Wait for cur bkv + wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx) + + # Load bkv into vreg. There is no need to mask out invalid k/v entries, + # because the score of invalid Q.K^T pairs are masked (to be zero) in + # flash attention, so that the invalid kv entries + # (as long as they are not NaN or inf) won't affect to the output. + bkv = load_bkv(bkv_sem_idx, bkv_idx) + + bq = load_bq(bq_sem_idx) + + flash_attention( + bq, + bkv, + bq_idx=bq_idx, + bkv_idx=bkv_idx, + kv_lens_to_attend_segment=kv_lens_to_attend_segment, + ) + return (kv_lens_to_attend_segment,) + + # Wait for cur bq if not ready yet + wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx) + wait_fetch_swa(seq_idx, bq_idx, bq_sem_idx) + jax.lax.fori_loop( + 0, + num_bkv, + compute_with_bkv, + (kv_lens_to_attend_segment,), + unroll=False, + ) + + # Load acc and calculate final output. + acc = acc_ref[...] + attention_sinks = jnp.concat( + [attention_sinks_ref[...] for _ in range(bq_sz)] + )[..., None] + exp_attention_sinks = jnp.exp(attention_sinks - m_ref[...]) + l = l_ref[...] + exp_attention_sinks + l = broadcast_minor(l, acc.shape) + out = ( + lax.div(acc, l) + if q_dtype == jnp.float32 + else (acc * pl.reciprocal(l, approx=True)).astype(q_dtype) + ) + + # Wait for previous bo to be fully sent before storing new bo. + bo_sem_idx = sem_ids_ref[2] + sem_ids_ref[2] = lax.select(bo_sem_idx == 0, 1, 0) + wait_send_bo(bo_sem_idx) + + # Store output from acc to bo. + bo_x2_ref.at[bo_sem_idx].bitcast(jnp.int32).reshape( + bq_sz * num_q_heads_per_q_packing, + head_dim, + )[...] = pltpu.bitcast(out, jnp.int32) + + # Send cur bo + start_send_bo(seq_idx, bq_idx, bo_sem_idx) + + lax.fori_loop(0, num_bq, compute_with_bq, None, unroll=False) + + ### ------- Kernel start ------- ### + + @pl.when(seq_idx == start_seq_idx) + def prologue(): + start_fetch_bq(start_seq_idx, 0, 0) + start_fetch_swa(start_seq_idx, 0, 0) + + # Initialize bkv_x2_ref to avoid NaN issues from accessing uninitialized + # memory + bkv_zeros = jnp.zeros(bkv_x2_ref.shape[1:], bkv_x2_ref.dtype) + bkv_x2_ref[0] = bkv_zeros + start_fetch_bkv(start_seq_idx, 0, 0) + bkv_x2_ref[1] = bkv_zeros + + process() + + @pl.when(seq_idx == end_seq_idx - 1) + def epilogue(): + for i in range(2): + wait_send_bo(i) + + ### ------- Kernel end ------- ### + + +def prepare_q_inputs( + q: jax.Array, # [max_num_tokens, actual_num_q_heads, actual_head_dim], +): + _, actual_num_q_heads, actual_head_dim = q.shape + q_packing = get_dtype_packing(q.dtype) + num_q_heads = align_to(actual_num_q_heads, q_packing) + head_dim = align_to(actual_head_dim, 128) + q = jnp.pad( + q, + ( + (0, 0), + (0, num_q_heads - actual_num_q_heads), + (0, head_dim - actual_head_dim), + ), + constant_values=0, + ) + return q + + +def prepare_swa_inputs( + swa_accumution: jax.Array, # [max_num_tokens, num_q_heads, head_dim] + swa_l: jax.Array, # [max_num_tokens, num_q_heads] + swa_m: jax.Array, # [max_num_tokens, num_q_heads] +): + _, actual_num_q_heads, actual_head_dim = swa_accumution.shape + swa_packing = get_dtype_packing(swa_accumution.dtype) + num_q_heads = align_to(actual_num_q_heads, swa_packing) + head_dim = align_to(actual_head_dim, 128) + swa_accumution = jnp.pad( + swa_accumution, + ( + (0, 0), + (0, num_q_heads - actual_num_q_heads), + (0, head_dim - actual_head_dim), + ), + constant_values=0, + ) + num_l_heads = align_to(num_q_heads, 128) + swa_l = jnp.pad( + swa_l, + ( + (0, 0), + (0, num_l_heads - actual_num_q_heads), + ), + constant_values=0, + ) + swa_m = jnp.pad( + swa_m, + ( + (0, 0), + (0, num_l_heads - actual_num_q_heads), + ), + constant_values=0, + ) + return swa_accumution, swa_l, swa_m + + +def prepare_outputs( + out, # [max_num_tokens, num_q_heads, head_dim] + actual_num_q_heads: int, + actual_head_dim: int, +): + return out[:, :actual_num_q_heads, :actual_head_dim] + + +# TODO: support batching decode q tokens as performance optimization. + + +# Main Attention kernel for DeepSeek V4 HCA. +# Note that the compressed kv tokens of current batch (current forward pass) +# have been written to the `cache_kv` by the compressor module before calling +# this function, `kv_lens` reflects the length after compressed kv cache write. + + +# Quantize and dequantize into / from +# DSv4 fp8 format (448 fp8, 64 bf16, 7 fp8 scales, 7 e8m0 scale for 448 fp8) is +# quite expensive for TPU. +# For HCA, we just skip the quantization and dequantization, the kv cache stores +# bf16 data. +# HCA's compression ratio is 128, the overall size of HCA's compressed kv cache +# is very small compared to other caches such as CSA's compressed cache. Extra +# storage for storing KV cache in bf16 is trivial. +@functools.partial( + jax.jit, + static_argnames=( + "sm_scale", + "chunk_prefill_size", + "num_kv_pages_per_block", + "num_queries_per_block", + "vmem_limit_bytes", + ), +) +def mla_ragged_paged_attention( + q: jax.Array, # [max_num_tokens, actual_num_q_heads, head_dim] + cache_kv: jax.Array, # [total_num_pages, page_size * 2, 2, 128] uint8 + kv_lens: jax.Array, # i32[max_num_seqs] + kv_lens_to_attend: jax.Array, # i32[max_num_tokens] + page_indices: jax.Array, # i32[max_num_seqs * pages_per_seq] + cu_q_lens: jax.Array, # i32[max_num_seqs + 1] + distribution: jax.Array, # i32[3] + attention_sinks: jax.Array, # float32[actual_num_q_heads] + swa_accumution: jax.Array, # bf16[max_num_tokens, num_q_heads, head_dim] + swa_l: jax.Array, # float32[max_num_tokens, num_q_heads] + swa_m: jax.Array, # float32[max_num_tokens, num_q_heads] + *, + sm_scale: float = 1.0, + # Kernel optimization params. + chunk_prefill_size: int | None = None, + # Kernel tuning params for decode, prefill, and mixed cases. + # If passsed in as int, all cases are the same. + num_kv_pages_per_block: tuple[int, int, int] | int | None = None, + num_queries_per_block: tuple[int, int, int] | int | None = None, + vmem_limit_bytes: int = DEFAULT_VMEM_LIMIT_BYTES, +) -> jax.Array: + """MLA Ragged paged attention that supports mixed prefill and decode. + + Args: + q: concatenated all sequences' queries. + cache_kv: the current kv cache. + kv_lens: the length of each sequence in the kv cache. + kv_lens_to_attend: for each query token, the length of kv sequence to attend + to. The attend to length is <= kv_lens[seq_id] for that query token. + page_indices: flattened page indices look-up table by (seq_id, page_id). + cu_q_lens: the cumulative sum of the effective query lengths. Similar to + kv_lens, only the first num_seqs+1 values are valid. + distribution: (i, j, k) represents that sequences[0:i] are decode-only, + sequences[i:j] are chunked-prefill-only, and sequences[j:k] are mixed. The + k is also the total number of sequences. + sm_scale: the softmax scale which will be applied to the Q@K^T. + num_kv_pages_per_block: number of kv pages to be processed in one flash + attention block in the pallas kernel. This is a tuple of (decode, prefill, + mixed) cases. + num_queries_per_block: number of queries to be processed in one flash + attention block in the pallas kernel. This is a tuple of (decode, prefill, + mixed) cases. + vmem_limit_bytes: the vmem limit for the pallas kernel. + + Returns: + The output of attention. + """ + assert cache_kv.dtype == jnp.uint8 + + if num_kv_pages_per_block is None or num_queries_per_block is None: + raise ValueError( + "num_kv_pages_per_block and num_queries_per_block must be specified." + ) + if isinstance(num_kv_pages_per_block, int): + num_kv_pages_per_blocks = [num_kv_pages_per_block for _ in range(3)] + else: + num_kv_pages_per_blocks = num_kv_pages_per_block + + if isinstance(num_queries_per_block, int): + num_queries_per_blocks = [num_queries_per_block for _ in range(3)] + else: + num_queries_per_blocks = num_queries_per_block + + _, actual_num_q_heads, actual_head_dim = q.shape + + q = prepare_q_inputs(q) # [max_num_tokens, num_q_heads, head_dim] + head_dim = q.shape[-1] + attention_sinks = jnp.pad( + attention_sinks, + (0, q.shape[1] - actual_num_q_heads), + constant_values=jnp.finfo(attention_sinks.dtype).min, + ) + assert swa_accumution.dtype == q.dtype + swa_accumution, swa_l, swa_m = prepare_swa_inputs( + swa_accumution, swa_l, swa_m + ) + + _, num_slots_per_page, kv_packing, lkv_dim = cache_kv.shape + num_slots_per_tokens = 2 + assert num_slots_per_page % num_slots_per_tokens == 0 + page_size = num_slots_per_page // num_slots_per_tokens + _, num_q_heads, _ = q.shape + max_num_seqs = cu_q_lens.shape[0] - 1 + num_page_indices = page_indices.shape[0] + assert num_page_indices % max_num_seqs == 0 + + def run_mla_kernel( + q: jax.Array, # [max_num_tokens, num_q_heads, head_dim] + cache_kv: jax.Array, # [total_num_pages, page_size * 2, kv_packing, lkv_dim] + kv_lens: jax.Array, # i32[max_num_seqs] + kv_lens_to_attend: jax.Array | None, # i32[max_num_tokens] + page_indices: jax.Array, # i32[max_num_seqs * pages_per_seq] + cu_q_lens: jax.Array, # i32[max_num_seqs + 1] + attention_sinks: jax.Array, # float32[num_q_heads] + swa_accumution: jax.Array, # bf16[max_num_tokens, num_q_heads, head_dim] + swa_l: jax.Array, # float32[max_num_tokens, num_l_heads] + swa_m: jax.Array, # float32[max_num_tokens, num_l_heads] + start_seq_idx: jax.Array, # i32 + end_seq_idx: jax.Array, # i32 + static_q_len: int | None, + num_kv_pages_per_block: int, + num_queries_per_block: int, + case: MlaCase = MlaCase.MIXED, + ): + + bkv_p = num_kv_pages_per_block + if static_q_len is not None: + bq_sz = min(num_queries_per_block, static_q_len) + else: + bq_sz = num_queries_per_block + + grid = (end_seq_idx - start_seq_idx,) + in_specs = [ + pl.BlockSpec(memory_space=pltpu.VMEM), # attention_sinks + pl.BlockSpec(memory_space=pltpu.HBM), # q + pl.BlockSpec(memory_space=pltpu.HBM), # cache_kv + pl.BlockSpec(memory_space=pltpu.HBM), # swa_accumution + pl.BlockSpec(memory_space=pltpu.HBM), # swa_l + pl.BlockSpec(memory_space=pltpu.HBM), # swa_m + ] + + out_specs = pl.BlockSpec(memory_space=pltpu.HBM) # o + + # last 2 dimension mapped to one tokens's kv + assert ( + num_slots_per_tokens * kv_packing * lkv_dim + == head_dim * get_dtype_bitwidth(q.dtype) // 8 + ) + bkv_double_buf = pltpu.VMEM( + (2, bkv_p * page_size, num_slots_per_tokens * kv_packing, lkv_dim), + cache_kv.dtype, + ) + + bq_double_bufq = pltpu.VMEM( + (2, bq_sz, num_q_heads, head_dim), + q.dtype, + ) + + bo_double_buf = bq_double_bufq + + num_l_heads = align_to(num_q_heads, 128) + bl_double_buf = pltpu.VMEM( + (2, bq_sz, num_l_heads), + jnp.float32, + ) + bm_double_buf = bl_double_buf + + swa_acc_double_buf = pltpu.VMEM( + (2, bq_sz, num_q_heads, head_dim), + q.dtype, + ) + + l_scratch = pltpu.VMEM( + (bq_sz * num_q_heads, 128), + jnp.float32, + ) + m_scratch = l_scratch + + acc_scratch = pltpu.VMEM( + (bq_sz * num_q_heads, head_dim), + jnp.float32, + ) + + scratch_shapes = [ + bkv_double_buf, + bq_double_bufq, + bo_double_buf, # Double buffering for output block. + bl_double_buf, # Double buffering for l output. + bm_double_buf, # Double buffering for m output. + swa_acc_double_buf, # Buffer for swa_accumution. + # Semaphores for double buffering of bkv, bq, bo, swa_acc, swa_l, swa_m, topk. + pltpu.SemaphoreType.DMA((6, 2)), + # Intermediate buffers per kv head for flash attention. + l_scratch, + m_scratch, + acc_scratch, + ] + + scalar_prefetches = ( + kv_lens, + kv_lens_to_attend, + page_indices, + cu_q_lens, + jnp.array([start_seq_idx, end_seq_idx], jnp.int32), + # (bq_sem_idx, bkv_sem_idx, bo_sem_idx) + jnp.zeros((3,), jnp.int32), + # (bo_sem_0_seq_idx, bo_sem_1_seq_idx, bo_sem_0_bo_idx, bo_sem_1_bo_idx) + jnp.full((4,), -1, jnp.int32), + ) + + scope_name = f"MLA-{case.symbol}-bq_{bq_sz}-bkvp_{bkv_p}-p_{page_size}" + kernel = jax.named_scope(scope_name)( + pl.pallas_call( + functools.partial( + _mla_ragged_paged_attention_kernel, + sm_scale=sm_scale, + static_q_len=static_q_len, + bq_sz=bq_sz, + bkv_p=bkv_p, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=len(scalar_prefetches), + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + scratch_shapes=scratch_shapes, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary",), + vmem_limit_bytes=vmem_limit_bytes, + disable_bounds_checks=True, + ), + out_shape=jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype), + input_output_aliases={ + 8: 0, # Alias output activation with q + }, + name=scope_name, + ) + ) + return kernel( + *scalar_prefetches, + attention_sinks, + q, + cache_kv, + swa_accumution, + swa_l, + swa_m, + ) + + # Decode-only + q = run_mla_kernel( + q, + cache_kv, + kv_lens, + kv_lens_to_attend, + page_indices, + cu_q_lens, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + num_kv_pages_per_block=num_kv_pages_per_blocks[0], + num_queries_per_block=num_queries_per_blocks[0], + start_seq_idx=jnp.array(0), + end_seq_idx=distribution[0], + static_q_len=1, + case=MlaCase.DECODE, + ) + + if chunk_prefill_size is not None: + # Handle prefill where the query length is fixed per sequence. + q = run_mla_kernel( + q, + cache_kv, + kv_lens, + kv_lens_to_attend, + page_indices, + cu_q_lens, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + num_kv_pages_per_block=num_kv_pages_per_blocks[1], + num_queries_per_block=num_queries_per_blocks[1], + start_seq_idx=distribution[0], + end_seq_idx=distribution[1], + static_q_len=chunk_prefill_size, + case=MlaCase.PREFILL, + ) + + # Handle mixed case where the query length per sequence is variable. + q = run_mla_kernel( + q, + cache_kv, + kv_lens, + kv_lens_to_attend, + page_indices, + cu_q_lens, + attention_sinks, + swa_accumution, + swa_l, + swa_m, + num_kv_pages_per_block=num_kv_pages_per_blocks[2], + num_queries_per_block=num_queries_per_blocks[2], + start_seq_idx=distribution[1], + end_seq_idx=distribution[2], + static_q_len=None, + case=MlaCase.MIXED, + ) + output = prepare_outputs( + q, actual_num_q_heads, actual_head_dim + ) # [max_num_tokens, actual_num_q_heads, actual_head_dim] + + return output + + +def computation( + q, cache_kv, kv_lens, kv_lens_to_attend, page_indices, cu_q_lens, + distribution, attention_sinks, swa_accumution, swa_l, swa_m, +): + num_tokens = q.shape[0] + batch_size = cu_q_lens.shape[0] - 1 + q_len = num_tokens // batch_size + page_size = cache_kv.shape[1] // 2 + + chunk_prefill_size = None if q_len == 1 else q_len + bq = 1 if q_len == 1 else 32 + bkv_pages = 4 if (page_size == 2048 and q_len > 1) else 16 + + sm_scale = float(512 ** -0.5) + vmem_limit_bytes = 100 * 1024 * 1024 + + return mla_ragged_paged_attention( + q, cache_kv, kv_lens, kv_lens_to_attend, page_indices, cu_q_lens, + distribution, attention_sinks, swa_accumution, swa_l, swa_m, + sm_scale=sm_scale, + chunk_prefill_size=chunk_prefill_size, + num_kv_pages_per_block=bkv_pages, + num_queries_per_block=bq, + vmem_limit_bytes=vmem_limit_bytes, + ) diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/53p_DeepSeek_V4_SWA/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/53p_DeepSeek_V4_SWA/kernel_task.yaml new file mode 100644 index 0000000..682f45c --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/53p_DeepSeek_V4_SWA/kernel_task.yaml @@ -0,0 +1,74 @@ +task_id: 53p_DeepSeek_V4_SWA +description: DeepSeek Sliding Window Attention (SWA) +input_gen_code: |- + def get_inputs(): + import jax + import jax.numpy as jnp + + configs = [ + ("decode_8k_1k", 256, 1, 9216, 1024, 8, 512, 128), + ("prefill_first_chunk", 1, 1024, 1024, 1024, 8, 512, 128), + ("prefill_last_chunk", 1, 1024, 8192, 1024, 8, 512, 128), + ("medium_batch_decode", 128, 1, 9216, 1024, 8, 512, 128), + ("prefill_256", 1, 256, 1024, 1024, 8, 512, 128), + ("prefill_512", 1, 512, 4096, 1024, 8, 512, 128), + ("alt_page_decode", 128, 1, 9216, 256, 8, 512, 128), + ("alt_page_prefill", 1, 256, 1024, 256, 8, 512, 128), + ] + + outputs = [] + key = jax.random.PRNGKey(42) + + def cdiv_val(a, b): + assert b != 0 + return (a + b - 1) // b + + for ( + name, + batch_size, + q_len, + kv_len_val, + page_size, + num_q_heads, + head_dim, + sliding_window, + ) in configs: + key, k1, k2, k3 = jax.random.split(key, 4) + + num_tokens = batch_size * q_len + kv_lens = jnp.full((batch_size,), kv_len_val, dtype=jnp.int32) + cu_q_lens = jnp.arange(0, num_tokens + 1, q_len, dtype=jnp.int32) + + pages_per_seq = cdiv_val(kv_len_val, page_size) + 2 + total_pages = batch_size * pages_per_seq + page_indices = jnp.arange(total_pages, dtype=jnp.int32) + + q = jax.random.normal( + k1, (num_tokens, num_q_heads, head_dim), dtype=jnp.bfloat16 + ) + new_kv = jax.random.normal(k2, (num_tokens, head_dim), dtype=jnp.bfloat16) + attention_sinks = jax.random.normal(k3, (num_q_heads,), dtype=jnp.float32) + + num_decode_seqs = batch_size if q_len == 1 else 0 + distribution = jnp.array( + [num_decode_seqs, num_decode_seqs, batch_size], dtype=jnp.int32 + ) + kernel_cache = jnp.zeros((total_pages, page_size * 2, 4, 128), dtype=jnp.uint8) + + dynamic_args = [ + q, + new_kv, + kernel_cache, + kv_lens, + page_indices, + cu_q_lens, + distribution, + attention_sinks, + ] + static_args = [] + outputs.append((dynamic_args, static_args)) + + return outputs + +rtol: 0.01 +atol: 0.01 diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/53p_DeepSeek_V4_SWA/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/53p_DeepSeek_V4_SWA/reference.py new file mode 100644 index 0000000..f28eaf5 --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/53p_DeepSeek_V4_SWA/reference.py @@ -0,0 +1,921 @@ +# Imports +import functools +from enum import Enum +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 + +# Initialization +def get_inputs(): + import jax + import jax.numpy as jnp + + configs = [ + ("decode_8k_1k", 256, 1, 9216, 1024, 8, 512, 128), + ("prefill_first_chunk", 1, 1024, 1024, 1024, 8, 512, 128), + ("prefill_last_chunk", 1, 1024, 8192, 1024, 8, 512, 128), + ("medium_batch_decode", 128, 1, 9216, 1024, 8, 512, 128), + ("prefill_256", 1, 256, 1024, 1024, 8, 512, 128), + ("prefill_512", 1, 512, 4096, 1024, 8, 512, 128), + ("alt_page_decode", 128, 1, 9216, 256, 8, 512, 128), + ("alt_page_prefill", 1, 256, 1024, 256, 8, 512, 128), + ] + + outputs = [] + key = jax.random.PRNGKey(42) + + def cdiv_val(a, b): + assert b != 0 + return (a + b - 1) // b + + for ( + name, + batch_size, + q_len, + kv_len_val, + page_size, + num_q_heads, + head_dim, + sliding_window, + ) in configs: + key, k1, k2, k3 = jax.random.split(key, 4) + + num_tokens = batch_size * q_len + kv_lens = jnp.full((batch_size,), kv_len_val, dtype=jnp.int32) + cu_q_lens = jnp.arange(0, num_tokens + 1, q_len, dtype=jnp.int32) + + pages_per_seq = cdiv_val(kv_len_val, page_size) + 2 + total_pages = batch_size * pages_per_seq + page_indices = jnp.arange(total_pages, dtype=jnp.int32) + + q = jax.random.normal( + k1, (num_tokens, num_q_heads, head_dim), dtype=jnp.bfloat16 + ) + new_kv = jax.random.normal(k2, (num_tokens, head_dim), dtype=jnp.bfloat16) + attention_sinks = jax.random.normal(k3, (num_q_heads,), dtype=jnp.float32) + + num_decode_seqs = batch_size if q_len == 1 else 0 + distribution = jnp.array( + [num_decode_seqs, num_decode_seqs, batch_size], dtype=jnp.int32 + ) + kernel_cache = jnp.zeros((total_pages, page_size * 2, 4, 128), dtype=jnp.uint8) + + dynamic_args = [ + q, + new_kv, + kernel_cache, + kv_lens, + page_indices, + cu_q_lens, + distribution, + attention_sinks, + ] + static_args = [] + outputs.append((dynamic_args, static_args)) + + return outputs + +# Computation +def cdiv(a, b): + assert b != 0 + return (a + b - 1) // b + +def align_to(x, a): + return cdiv(x, a) * a + +def get_dtype_bitwidth(dtype): + return jax.dtypes.itemsize_bits(dtype) + +def get_dtype_packing(dtype): + bits = get_dtype_bitwidth(dtype) + return 32 // bits + +class MlaCase(Enum): + DECODE = 0 + PREFILL = 1 + MIXED = 2 + + @property + def symbol(self): + return { + MlaCase.DECODE: "d", + MlaCase.PREFILL: "p", + MlaCase.MIXED: "m", + }[self] + +def _mla_sliding_window_ragged_paged_attention_kernel( + kv_lens_ref, + page_indices_ref, + cu_q_lens_ref, + start_end_seq_idx_ref, + sem_ids_ref, + bo_ids_ref, + bkv_update_ids_ref, + attention_sinks_ref, + q_hbm_ref, + new_kv_hbm_ref, + cache_kv_hbm_ref, + in_output_hbm_ref, + in_l_hbm_ref, + in_m_hbm_ref, + o_hbm_ref, + updated_cache_kv_hbm_ref, + l_hbm_ref, + m_hbm_ref, + bkv_x2_ref, + bq_x2_ref, + bo_x2_ref, + bl_x2_ref, + bm_x2_ref, + sems, + l_ref, + m_ref, + acc_ref, + *, + static_q_len: int, + sm_scale: float, + sliding_window: int, + logical_page_size: int, + unnormalized_output: bool, + q_compute_block_size: int | None, + bkv_p, + bq_sz, +): + assert q_hbm_ref.shape == o_hbm_ref.shape + assert sliding_window > 0 + + _, num_q_heads, head_dim = q_hbm_ref.shape + q_packing = get_dtype_packing(q_hbm_ref.dtype) + assert num_q_heads % q_packing == 0 + num_q_heads_per_q_packing = num_q_heads // q_packing + + total_num_pages, physical_page_size_per_kv_packing, kv_packing, lkv_dim = cache_kv_hbm_ref.shape + q_dtype = q_hbm_ref.dtype + assert o_hbm_ref.dtype == q_dtype + assert head_dim % 128 == 0 + token_bytes = head_dim * get_dtype_bitwidth(q_dtype) // 8 + slot_bytes = kv_packing * lkv_dim + assert token_bytes % slot_bytes == 0 + slots_per_token = token_bytes // slot_bytes + phys_tokens_per_page = physical_page_size_per_kv_packing // slots_per_token + + max_num_seqs = kv_lens_ref.shape[0] + num_page_indices = page_indices_ref.shape[0] + assert num_page_indices % max_num_seqs == 0 + pages_per_seq = num_page_indices // max_num_seqs + + bkv_sz = bkv_p * logical_page_size + page_size = logical_page_size + + max_bkv_blocks = cdiv(bq_sz + sliding_window - 1, bkv_sz) + single_bkv_block = max_bkv_blocks == 1 + + start_seq_idx = start_end_seq_idx_ref[0] + end_seq_idx = start_end_seq_idx_ref[1] + seq_idx = pl.program_id(0) + start_seq_idx + q_start = cu_q_lens_ref[seq_idx] + q_end = cu_q_lens_ref[seq_idx + 1] + q_len = q_end - q_start + kv_len = kv_lens_ref[seq_idx] + + def flash_attention( + q, + kv, + *, + bq_idx, + bkv_idx, + start_offset, + ): + assert len(q.shape) == 2 + assert len(kv.shape) == 2 + assert q.shape[0] % num_q_heads == 0 + assert q.shape[1] == head_dim + assert kv.shape == (bkv_sz, head_dim) + n = q.shape[0] // num_q_heads + + if q_compute_block_size is None: + chunk_sz = n + else: + chunk_sz = q_compute_block_size if n % q_compute_block_size == 0 else n + num_chunks = n // chunk_sz + + def load_with_init(ref, init_val): + if single_bkv_block: + return jnp.full_like(ref, init_val) + else: + return jnp.where(bkv_idx == 0, jnp.full_like(ref, init_val), ref[...]) + + k_span = start_offset + bkv_idx * bkv_sz + lax.broadcasted_iota(jnp.int32, (1, bkv_sz), 1) + k_pos = start_offset + bkv_idx * bkv_sz + lax.broadcasted_iota(jnp.int32, (bkv_sz, 1), 0) + kv = jnp.where(k_pos < kv_len, kv, 0.0) + + chunk_size = chunk_sz * num_q_heads + for c in range(num_chunks): + start_row = c * chunk_size + qc = q[start_row : start_row + chunk_size] + cl_ref = l_ref.at[start_row : start_row + chunk_size] + cm_ref = m_ref.at[start_row : start_row + chunk_size] + cacc_ref = acc_ref.at[start_row : start_row + chunk_size] + + s = jnp.einsum("nd,md->nm", qc, kv, preferred_element_type=jnp.float32) + s *= sm_scale + + q_span = kv_len - q_len + bq_idx * bq_sz + (start_row + lax.broadcasted_iota(jnp.int32, (chunk_size, 1), 0)) // num_q_heads + keep = (q_span - k_span).astype(jnp.uint32) < jnp.uint32(sliding_window) + + s = jnp.where(keep, s, jnp.finfo(s.dtype).min) + s_rowmax = jnp.max(s, axis=1, keepdims=True) + m_prev = load_with_init(cm_ref, jnp.finfo(jnp.float32).min) + m_curr = jnp.maximum(m_prev, s_rowmax) + cm_ref[...] = m_curr + p = jnp.exp(s - broadcast_minor(m_curr, s.shape)) + p = jnp.where(keep, p, 0.0) + + pv = jnp.einsum("nm,md->nd", p, kv, preferred_element_type=jnp.float32) + + p_rowsum = jnp.sum(p, axis=1, keepdims=True) + exp_m_diff = jnp.exp(m_prev - m_curr) + l_prev = load_with_init(cl_ref, 0.0) + l_curr = exp_m_diff * l_prev + p_rowsum + cl_ref[...] = l_curr + o_prev = load_with_init(cacc_ref, 0.0) + o_curr = broadcast_minor(exp_m_diff, o_prev.shape) * o_prev + pv + cacc_ref[...] = o_curr + + def _async_copy(src, dst, sem, wait): + cp = pltpu.make_async_copy(src, dst, sem) + if wait: + cp.wait() + else: + cp.start() + + def _get_kv_len(seq_idx): + return jnp.where(seq_idx < end_seq_idx, kv_lens_ref[seq_idx], 0) + + def _get_q_len(seq_idx): + return jnp.where(seq_idx < end_seq_idx, cu_q_lens_ref[seq_idx + 1] - cu_q_lens_ref[seq_idx], 0) + + def _start_offset(seq_idx, bq_idx): + return jnp.maximum(_get_kv_len(seq_idx) - _get_q_len(seq_idx) + bq_idx * bq_sz - sliding_window + 1, 0) + + def _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, start_offset, *, wait=False): + sem = sems.at[0, bkv_sem_idx] + bkv_vmem_ref = bkv_x2_ref.at[bkv_sem_idx] + + reshaped_cache_hbm_ref = cache_kv_hbm_ref.reshape(total_num_pages * phys_tokens_per_page, slots_per_token * kv_packing, lkv_dim) + + kv_len = kv_lens_ref[seq_idx] + kv_len_start = start_offset + bkv_idx * bkv_sz + kv_p_start = kv_len_start // page_size + page_off = kv_len_start - kv_p_start * page_size + + q_start = cu_q_lens_ref[seq_idx] + q_end = cu_q_lens_ref[seq_idx + 1] + q_len = q_end - q_start + + kv_left = jnp.maximum(kv_len - kv_len_start, 0) + kv_left_frm_cache = jnp.maximum(kv_left - q_len, 0) + kv_left_frm_new = kv_left - kv_left_frm_cache + + bkv_sz_frm_cache = jnp.minimum(kv_left_frm_cache, bkv_sz) + bkv_sz_frm_new = jnp.minimum(bkv_sz - bkv_sz_frm_cache, kv_left_frm_new) + page_indices_offset = seq_idx * pages_per_seq + kv_p_start + + new_kv_len_start = q_end - kv_left_frm_new + dma_bkv_sz = bkv_sz_frm_cache + bkv_sz_frm_new + + if not wait: + wait_update_kv_cache(bkv_sem_idx) + + for i in range(bkv_p + 1): + if i == 0: + in_page_off = page_off + vmem_off = jnp.int32(0) + avail = page_size - page_off + else: + in_page_off = jnp.int32(0) + vmem_off = i * page_size - page_off + avail = jnp.int32(page_size) + sz = jnp.clip(bkv_sz_frm_cache - vmem_off, 0, avail) + page_idx = jnp.minimum(page_indices_offset + i, num_page_indices - 1) + _async_copy( + reshaped_cache_hbm_ref.at[pl.ds(page_indices_ref[page_idx] * phys_tokens_per_page + in_page_off, sz)], + bkv_vmem_ref.at[pl.ds(vmem_off, sz)], + sem, + wait, + ) + + _async_copy( + new_kv_hbm_ref.at[pl.ds(new_kv_len_start, bkv_sz_frm_new)], + bkv_vmem_ref.at[pl.ds(bkv_sz_frm_cache, bkv_sz_frm_new)], + sem, + wait, + ) + + else: + dst_kv = bkv_vmem_ref.at[pl.ds(0, dma_bkv_sz)] + _async_copy(src=dst_kv, dst=dst_kv, sem=sem, wait=True) + + return kv_len_start + bkv_sz_frm_cache, bkv_sz_frm_new, bkv_sz_frm_cache + + def _update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz, *, in_vmem_start=0, wait=False): + sem = sems.at[3, bkv_sem_idx] + bkv_vmem_ref = bkv_x2_ref.at[bkv_sem_idx] + + update_kv_packing_iters = update_sz + + reshaped_cache_kv_hbm_ref = updated_cache_kv_hbm_ref.reshape(total_num_pages * phys_tokens_per_page, slots_per_token * kv_packing, lkv_dim) + + if not wait: + kv_p_start = offset // page_size + kv_p_end = cdiv(offset + update_sz, page_size) + start_word_in_page = offset % page_size + start_word_in_vmem = in_vmem_start + words_to_transfer = update_kv_packing_iters + page_indices_offset = seq_idx * pages_per_seq + kv_p_start + + def loop_body(i, states): + curr_word_in_page, words_to_transfer, curr_word_in_vmem = states + sz = jnp.minimum(page_size - curr_word_in_page, words_to_transfer) + page_idx = page_indices_ref[page_indices_offset + i] + + _async_copy( + bkv_vmem_ref.at[pl.ds(curr_word_in_vmem, sz)], + reshaped_cache_kv_hbm_ref.at[pl.ds(page_idx * phys_tokens_per_page + curr_word_in_page, sz)], + sem, + wait=False, + ) + return 0, words_to_transfer - sz, curr_word_in_vmem + sz + + lax.fori_loop( + 0, + kv_p_end - kv_p_start, + loop_body, + (start_word_in_page, words_to_transfer, start_word_in_vmem), + unroll=False, + ) + else: + dma_sz_words = update_kv_packing_iters + dst_kv = bkv_vmem_ref.at[pl.ds(0, dma_sz_words)] + _async_copy(src=dst_kv, dst=dst_kv, sem=sem, wait=True) + + def _fetch_bq(seq_idx, bq_idx, bq_sem_idx, *, wait=False): + sem = sems.at[1, bq_sem_idx] + bq_vmem_ref = bq_x2_ref.at[bq_sem_idx] + + q_len_start = cu_q_lens_ref[seq_idx] + bq_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy(q_hbm_ref.at[pl.ds(q_len_start, sz)], bq_vmem_ref.at[pl.ds(0, sz)], sem, wait) + + def _send_bo(seq_idx, bo_idx, bo_sem_idx, *, wait=False): + sem = sems.at[2, bo_sem_idx] + vmem_ref = bo_x2_ref.at[bo_sem_idx] + q_len_start = cu_q_lens_ref[seq_idx] + bo_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy(vmem_ref.at[pl.ds(0, sz)], o_hbm_ref.at[pl.ds(q_len_start, sz)], sem, wait) + + def _send_l(seq_idx, bo_idx, bo_sem_idx, *, wait=False): + sem = sems.at[4, bo_sem_idx] + vmem_ref = bl_x2_ref.at[bo_sem_idx] + q_len_start = cu_q_lens_ref[seq_idx] + bo_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy(vmem_ref.at[pl.ds(0, sz)], l_hbm_ref.at[pl.ds(q_len_start, sz)], sem, wait) + + def _send_m(seq_idx, bo_idx, bo_sem_idx, *, wait=False): + sem = sems.at[5, bo_sem_idx] + vmem_ref = bm_x2_ref.at[bo_sem_idx] + q_len_start = cu_q_lens_ref[seq_idx] + bo_idx * bq_sz + q_end = cu_q_lens_ref[seq_idx + 1] + sz = jnp.minimum(bq_sz, q_end - q_len_start) + + _async_copy(vmem_ref.at[pl.ds(0, sz)], m_hbm_ref.at[pl.ds(q_len_start, sz)], sem, wait) + + def start_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, start_offset): + return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, start_offset) + + def wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, start_offset): + return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, start_offset, wait=True) + + def start_fetch_bq(seq_idx, bq_idx, bq_sem_idx): + return _fetch_bq(seq_idx, bq_idx, bq_sem_idx) + + def wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx): + return _fetch_bq(seq_idx, bq_idx, bq_sem_idx, wait=True) + + def start_send_bo(seq_idx, bo_idx, bo_sem_idx): + bo_ids_ref[bo_sem_idx] = seq_idx + bo_ids_ref[bo_sem_idx + 2] = bo_idx + _send_bo(seq_idx, bo_idx, bo_sem_idx) + _send_l(seq_idx, bo_idx, bo_sem_idx) + _send_m(seq_idx, bo_idx, bo_sem_idx) + + def wait_send_bo(bo_sem_idx): + old_seq_idx = bo_ids_ref[bo_sem_idx] + old_bo_idx = bo_ids_ref[bo_sem_idx + 2] + + @pl.when(jnp.logical_and(0 <= old_seq_idx, old_seq_idx <= seq_idx)) + def _(): + _send_bo(old_seq_idx, old_bo_idx, bo_sem_idx, wait=True) + _send_l(old_seq_idx, old_bo_idx, bo_sem_idx, wait=True) + _send_m(old_seq_idx, old_bo_idx, bo_sem_idx, wait=True) + + def start_update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz, vmem_start): + bkv_update_ids_ref[bkv_sem_idx] = seq_idx + bkv_update_ids_ref[bkv_sem_idx + 2] = offset + bkv_update_ids_ref[bkv_sem_idx + 4] = update_sz + _update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz, in_vmem_start=vmem_start) + + def wait_update_kv_cache(bkv_sem_idx): + update_sz = bkv_update_ids_ref[bkv_sem_idx + 4] + + @pl.when(update_sz > 0) + def _(): + seq_idx = bkv_update_ids_ref[bkv_sem_idx] + offset = bkv_update_ids_ref[bkv_sem_idx + 2] + bkv_update_ids_ref[bkv_sem_idx + 4] = 0 + _update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz, wait=True) + + def load_bq(bq_sem_idx): + q_ref = bq_x2_ref.bitcast(jnp.uint32).at[bq_sem_idx].reshape(bq_sz * num_q_heads_per_q_packing, head_dim) + q = pltpu.bitcast(q_ref[: bq_sz * num_q_heads_per_q_packing], q_dtype).reshape(bq_sz * num_q_heads, head_dim) + return q + + def load_bkv(bkv_sem_idx, bkv_idx, start_offset): + bkv_u8 = bkv_x2_ref.at[bkv_sem_idx][...] + bkv = pltpu.bitcast(bkv_u8, jnp.bfloat16).reshape(bkv_sz, head_dim) + return bkv + + def broadcast_minor(src, shape): + if src.shape == shape: + return src + assert src.shape[:-1] == shape[:-1] + assert src.shape[-1] % 128 == 0 + target_minor = align_to(shape[-1], src.shape[-1]) + return jnp.concatenate([src for _ in range(target_minor // src.shape[-1])], axis=-1)[..., : shape[-1]] + + def process(): + if static_q_len is None: + num_bq = jnp.maximum(1, cdiv(q_len, bq_sz)) + else: + num_bq = jnp.maximum(1, cdiv(static_q_len, bq_sz)) + + def get_next_bq_ids(seq_idx, bq_idx, bq_sem_idx): + next_bq_idx = bq_idx + 1 + is_last_bq = next_bq_idx == num_bq + next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx) + next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx) + next_bq_sem_idx = lax.select(bq_sem_idx == 0, 1, 0) + return next_seq_idx, next_bq_idx, next_bq_sem_idx + + def compute_with_bq(bq_idx, _): + cur_start_offset = _start_offset(seq_idx, bq_idx) + start_bkv_idx = 0 + if single_bkv_block: + end_bkv_idx = 1 + else: + end_bkv_idx = jnp.maximum(cdiv(jnp.minimum(kv_len - q_len + (bq_idx + 1) * bq_sz, kv_len) - cur_start_offset, bkv_sz), 1) + + def get_next_bkv_ids(seq_idx, bq_idx, bkv_idx, bkv_sem_idx): + next_bkv_idx = bkv_idx + 1 + is_last_bkv = next_bkv_idx == end_bkv_idx + next_bq_idx = lax.select(is_last_bkv, bq_idx + 1, bq_idx) + is_last_bq = next_bq_idx == num_bq + next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx) + next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx) + next_bkv_idx = lax.select(is_last_bkv, 0, next_bkv_idx) + next_bkv_sem_idx = lax.select(bkv_sem_idx == 0, 1, 0) + return next_seq_idx, next_bq_idx, next_bkv_idx, next_bkv_sem_idx + + bq_sem_idx = sem_ids_ref[0] + next_seq_idx, next_bq_idx, next_bq_sem_idx = get_next_bq_ids(seq_idx, bq_idx, bq_sem_idx) + + @pl.when(next_seq_idx < end_seq_idx) + def prefetch_next_bq(): + sem_ids_ref[0] = next_bq_sem_idx + start_fetch_bq(next_seq_idx, next_bq_idx, next_bq_sem_idx) + + def compute_with_bkv(bkv_idx, _): + bkv_sem_idx = sem_ids_ref[1] + next_seq_idx, next_bq_idx, next_bkv_idx, next_bkv_sem_idx = get_next_bkv_ids(seq_idx, bq_idx, bkv_idx, bkv_sem_idx) + + @pl.when(next_seq_idx < end_seq_idx) + def prefetch_next_bkv(): + sem_ids_ref[1] = next_bkv_sem_idx + next_start_offset = _start_offset(next_seq_idx, next_bq_idx) + start_fetch_bkv(next_seq_idx, next_bkv_idx, next_bkv_sem_idx, next_start_offset) + + offset, update_sz, vmem_start = wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, cur_start_offset) + + @pl.when(update_sz > 0) + def update_cur_bkv_to_cache(): + start_update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz, vmem_start) + + bkv = load_bkv(bkv_sem_idx, bkv_idx, cur_start_offset) + bq = load_bq(bq_sem_idx) + + flash_attention(bq, bkv, bq_idx=bq_idx, bkv_idx=bkv_idx, start_offset=cur_start_offset) + + wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx) + if single_bkv_block: + compute_with_bkv(0, None) + else: + lax.fori_loop(start_bkv_idx, end_bkv_idx, compute_with_bkv, None, unroll=False) + + acc = acc_ref[...] + + if unnormalized_output: + l = broadcast_minor(l_ref[...], acc.shape) + out = acc.astype(q_dtype) + else: + attention_sinks = jnp.concat([attention_sinks_ref[...] for _ in range(bq_sz)])[..., None] + exp_attention_sinks = jnp.exp(attention_sinks - m_ref[...]) + l = l_ref[...] + exp_attention_sinks + l = broadcast_minor(l, acc.shape) + out = lax.div(acc, l) if q_dtype == jnp.float32 else (acc * pl.reciprocal(l, approx=True)).astype(q_dtype) + + bo_sem_idx = sem_ids_ref[2] + sem_ids_ref[2] = lax.select(bo_sem_idx == 0, 1, 0) + wait_send_bo(bo_sem_idx) + + bo_x2_ref.at[bo_sem_idx].bitcast(jnp.int32).reshape(bq_sz * num_q_heads_per_q_packing, head_dim)[...] = pltpu.bitcast(out, jnp.int32) + bl_x2_ref.at[bo_sem_idx][:bq_sz, :num_q_heads] = l_ref[..., 0].reshape(bq_sz, num_q_heads) + bm_x2_ref.at[bo_sem_idx][:bq_sz, :num_q_heads] = m_ref[..., 0].reshape(bq_sz, num_q_heads) + + start_send_bo(seq_idx, bq_idx, bo_sem_idx) + + lax.fori_loop(0, num_bq, compute_with_bq, None, unroll=False) + + @pl.when(seq_idx == start_seq_idx) + def prologue(): + start_fetch_bq(start_seq_idx, 0, 0) + start_fetch_bkv(start_seq_idx, 0, 0, _start_offset(start_seq_idx, 0)) + + process() + + @pl.when(seq_idx == end_seq_idx - 1) + def epilogue(): + for i in range(2): + wait_send_bo(i) + wait_update_kv_cache(i) + +def prepare_q_inputs(q: jax.Array): + max_num_tokens, actual_num_q_heads, actual_head_dim = q.shape + q_packing = get_dtype_packing(q.dtype) + num_q_heads = align_to(actual_num_q_heads, q_packing) + head_dim = align_to(actual_head_dim, 128) + q = jnp.pad( + q, + ((0, 0), (0, num_q_heads - actual_num_q_heads), (0, head_dim - actual_head_dim)), + constant_values=0, + ) + return q + +def prepare_kv_inputs(kv: jax.Array): + assert kv.dtype == jnp.bfloat16 + tokens, head_dim = kv.shape + assert head_dim % 128 == 0 + kv_u16 = jax.lax.bitcast_convert_type(kv, jnp.uint16) + kv_hi = ((kv_u16 >> 8) & 0xFF).astype(jnp.uint8) + kv_lo = (kv_u16 & 0xFF).astype(jnp.uint8) + nb = head_dim // 128 + kv_hi = kv_hi.reshape(tokens, nb, 128) + kv_lo = kv_lo.reshape(tokens, nb, 128) + interleaved = jnp.stack([kv_lo, kv_hi], axis=2) + return interleaved.reshape(tokens, head_dim * 2) + +def prepare_outputs(out, actual_num_q_heads: int, actual_head_dim: int): + return out[:, :actual_num_q_heads, :actual_head_dim] + +@functools.partial( + jax.jit, + static_argnames=( + "sm_scale", + "sliding_window", + "chunk_prefill_size", + "num_kv_pages_per_block", + "num_queries_per_block", + "vmem_limit_bytes", + "logical_page_size", + "unnormalized_output", + "q_compute_block_size", + ), + donate_argnames=("cache_kv",), +) +def mla_sliding_window_ragged_paged_attention( + q: jax.Array, + new_kv: jax.Array, + cache_kv: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + distribution: jax.Array, + attention_sinks: jax.Array, + *, + sm_scale: float = 1.0, + sliding_window: int, + logical_page_size: int, + chunk_prefill_size: int | None = None, + num_kv_pages_per_block: tuple[int, int, int] | int | None = None, + num_queries_per_block: tuple[int, int, int] | int | None = None, + q_compute_block_size: int | None = None, + vmem_limit_bytes: int = 100 * 1024 * 1024, + unnormalized_output: bool = False, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]: + if num_kv_pages_per_block is None or num_queries_per_block is None: + raise ValueError("num_kv_pages_per_block and num_queries_per_block must be specified.") + + if isinstance(num_kv_pages_per_block, int): + num_kv_pages_per_blocks = [num_kv_pages_per_block for _ in range(3)] + else: + num_kv_pages_per_blocks = num_kv_pages_per_block + + if isinstance(num_queries_per_block, int): + num_queries_per_blocks = [num_queries_per_block for _ in range(3)] + else: + num_queries_per_blocks = num_queries_per_block + + _, actual_num_q_heads, actual_head_dim = q.shape + + q = prepare_q_inputs(q) + attention_sinks = jnp.pad( + attention_sinks, + (0, q.shape[1] - actual_num_q_heads), + constant_values=jnp.finfo(attention_sinks.dtype).min, + ) + assert new_kv.dtype == jnp.bfloat16 + assert cache_kv.dtype == jnp.uint8 + head_dim = q.shape[-1] + _, physical_page_size_per_kv_packing, kv_packing, lkv_dim = cache_kv.shape + + slot_bytes = kv_packing * lkv_dim + token_bytes = head_dim * get_dtype_bitwidth(new_kv.dtype) // 8 + assert token_bytes % slot_bytes == 0 + slots_per_token = token_bytes // slot_bytes + phys_tokens_per_page = physical_page_size_per_kv_packing // slots_per_token + + new_kv = prepare_kv_inputs(new_kv) + new_kv = new_kv.reshape(new_kv.shape[0], slots_per_token * kv_packing, lkv_dim) + assert logical_page_size <= phys_tokens_per_page + + _, num_q_heads, _ = q.shape + max_num_seqs = kv_lens.shape[0] + num_page_indices = page_indices.shape[0] + assert num_page_indices % max_num_seqs == 0 + + def run_mla_kernel( + q: jax.Array, + new_kv: jax.Array, + cache_kv: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + start_seq_idx: jax.Array, + end_seq_idx: jax.Array, + in_output: jax.Array, + in_l: jax.Array, + in_m: jax.Array, + attention_sinks: jax.Array, + static_q_len: int | None, + unnormalized_output: bool, + num_kv_pages_per_block: int, + num_queries_per_block: int, + case: MlaCase = MlaCase.MIXED, + ): + bkv_p = num_kv_pages_per_block + if static_q_len is not None: + bq_sz = min(num_queries_per_block, static_q_len) + else: + bq_sz = num_queries_per_block + bkv_sz = bkv_p * logical_page_size + grid = (end_seq_idx - start_seq_idx,) + + in_specs = [ + pl.BlockSpec(memory_space=pltpu.VMEM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + ] + + out_specs = [ + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + pl.BlockSpec(memory_space=pltpu.HBM), + ] + + bkv_double_buf = pltpu.VMEM((2, bkv_sz, slots_per_token * kv_packing, lkv_dim), cache_kv.dtype) + bq_double_bufq = pltpu.VMEM((2, bq_sz, num_q_heads, head_dim), q.dtype) + bo_double_buf = bq_double_bufq + + num_l_heads = align_to(num_q_heads, 128) + bl_double_buf = pltpu.VMEM((2, bq_sz, num_l_heads), jnp.float32) + bm_double_buf = bl_double_buf + + l_scratch = pltpu.VMEM((bq_sz * num_q_heads, 128), jnp.float32) + m_scratch = l_scratch + acc_scratch = pltpu.VMEM((bq_sz * num_q_heads, head_dim), jnp.float32) + + scratch_shapes = [ + bkv_double_buf, + bq_double_bufq, + bo_double_buf, + bl_double_buf, + bm_double_buf, + pltpu.SemaphoreType.DMA((6, 2)), + l_scratch, + m_scratch, + acc_scratch, + ] + + scalar_prefetches = ( + kv_lens, + page_indices, + cu_q_lens, + jnp.array([start_seq_idx, end_seq_idx], jnp.int32), + jnp.zeros((3,), jnp.int32), + jnp.full((4,), -1, jnp.int32), + jnp.full((6,), -1, jnp.int32), + ) + + scope_name = f"SWA-{case.symbol}-bq_{bq_sz}-bkvp_{bkv_p}" + kernel = jax.named_scope(scope_name)( + pl.pallas_call( + functools.partial( + _mla_sliding_window_ragged_paged_attention_kernel, + sm_scale=sm_scale, + sliding_window=sliding_window, + static_q_len=static_q_len, + bq_sz=bq_sz, + bkv_p=bkv_p, + logical_page_size=logical_page_size, + unnormalized_output=unnormalized_output, + q_compute_block_size=q_compute_block_size, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=len(scalar_prefetches), + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + scratch_shapes=scratch_shapes, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("arbitrary",), + vmem_limit_bytes=vmem_limit_bytes, + disable_bounds_checks=True, + ), + out_shape=[ + jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype), + jax.ShapeDtypeStruct(shape=cache_kv.shape, dtype=cache_kv.dtype), + jax.ShapeDtypeStruct(shape=(q.shape[0], num_l_heads), dtype=jnp.float32), + jax.ShapeDtypeStruct(shape=(q.shape[0], num_l_heads), dtype=jnp.float32), + ], + input_output_aliases={ + 11: 0, + 10: 1, + 12: 2, + 13: 3, + }, + name=scope_name, + ) + ) + return kernel( + *scalar_prefetches, + attention_sinks, + q, + new_kv, + cache_kv, + in_output, + in_l, + in_m, + ) + + num_l_heads = align_to(num_q_heads, 128) + if unnormalized_output: + l = jnp.zeros((q.shape[0], num_l_heads), dtype=jnp.float32) + m = jnp.full((q.shape[0], num_l_heads), jnp.finfo(jnp.float32).min, dtype=jnp.float32) + in_output = jnp.zeros_like(q) + else: + l = jnp.zeros((q.shape[0], num_l_heads), dtype=jnp.float32) + m = jnp.zeros((q.shape[0], num_l_heads), dtype=jnp.float32) + in_output = jnp.zeros_like(q) + + output, updated_kv, out_l, out_m = run_mla_kernel( + q, + new_kv, + cache_kv, + kv_lens, + page_indices, + cu_q_lens, + num_kv_pages_per_block=num_kv_pages_per_blocks[0], + num_queries_per_block=num_queries_per_blocks[0], + start_seq_idx=jnp.array(0), + end_seq_idx=distribution[0], + in_output=in_output, + in_l=l, + in_m=m, + attention_sinks=attention_sinks, + static_q_len=1, + unnormalized_output=unnormalized_output, + case=MlaCase.DECODE, + ) + + if chunk_prefill_size is not None: + output, updated_kv, out_l, out_m = run_mla_kernel( + q, + new_kv, + updated_kv, + kv_lens, + page_indices, + cu_q_lens, + num_kv_pages_per_block=num_kv_pages_per_blocks[1], + num_queries_per_block=num_queries_per_blocks[1], + start_seq_idx=distribution[0], + end_seq_idx=distribution[1], + in_output=output, + in_l=out_l, + in_m=out_m, + attention_sinks=attention_sinks, + static_q_len=chunk_prefill_size, + unnormalized_output=unnormalized_output, + case=MlaCase.PREFILL, + ) + + output, updated_kv, out_l, out_m = run_mla_kernel( + q, + new_kv, + updated_kv, + kv_lens, + page_indices, + cu_q_lens, + num_kv_pages_per_block=num_kv_pages_per_blocks[2], + num_queries_per_block=num_queries_per_blocks[2], + start_seq_idx=distribution[1], + end_seq_idx=distribution[2], + in_output=output, + in_l=out_l, + in_m=out_m, + attention_sinks=attention_sinks, + static_q_len=None, + unnormalized_output=unnormalized_output, + case=MlaCase.MIXED, + ) + + output = prepare_outputs(output, actual_num_q_heads, actual_head_dim) + out_l = out_l[:, :actual_num_q_heads] + return output, updated_kv, out_l, out_m + +def computation( + q: jax.Array, + new_kv: jax.Array, + cache_kv: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + distribution: jax.Array, + attention_sinks: jax.Array, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]: + + sm_scale = 1.0 + sliding_window = 128 + logical_page_size = 128 + chunk_prefill_size = None + num_kv_pages_per_block = 2 + num_queries_per_block = 32 + q_compute_block_size = 2 + vmem_limit_bytes = 100 * 1024 * 1024 + unnormalized_output = True + + if cache_kv.shape[1] != logical_page_size * 2: + total_pages = cache_kv.shape[0] + cache_kv = jnp.zeros((total_pages, logical_page_size * 2, 4, 128), dtype=cache_kv.dtype) + + return mla_sliding_window_ragged_paged_attention( + q, + new_kv, + cache_kv, + kv_lens, + page_indices, + cu_q_lens, + distribution, + attention_sinks, + sm_scale=sm_scale, + sliding_window=sliding_window, + logical_page_size=logical_page_size, + chunk_prefill_size=chunk_prefill_size, + num_kv_pages_per_block=num_kv_pages_per_block, + num_queries_per_block=num_queries_per_block, + q_compute_block_size=q_compute_block_size, + vmem_limit_bytes=vmem_limit_bytes, + unnormalized_output=unnormalized_output, + ) diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/6p_Paged_Attention/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/6p_Paged_Attention/kernel_task.yaml new file mode 100644 index 0000000..c910fde --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/6p_Paged_Attention/kernel_task.yaml @@ -0,0 +1,47 @@ +task_id: 6p_Paged_Attention +description: Kernel task for 6p_Paged_Attention +input_gen_code: |- + def get_inputs(dtype=jnp.bfloat16): + import jax + import jax.numpy as jnp + + CONFIG = { + 'name': 'llama3_70b_paged_attention', + 'model': 'Llama-3.1-70B', + 'operator': 'paged_attention', + 'num_seqs': 64, + 'max_seq_len': 4096, + 'num_query_heads': 64, + 'num_kv_heads': 8, + 'head_dim': 128, + 'page_size': 16, + 'pages_per_seq': 256, + } + + key = jax.random.key(42) + keys = jax.random.split(key, 5) + num_seqs = CONFIG['num_seqs'] + num_q_heads = CONFIG['num_query_heads'] + num_kv_heads = CONFIG['num_kv_heads'] + head_dim = CONFIG['head_dim'] + page_size = CONFIG['page_size'] + pages_per_seq = CONFIG['pages_per_seq'] + total_pages = num_seqs * pages_per_seq + max_seq_len_derived = pages_per_seq * page_size + + max_num_tokens = num_seqs + queries = jax.random.normal(keys[0], (max_num_tokens, num_q_heads, head_dim), dtype=dtype) + k_pages = jax.random.normal(keys[1], (total_pages, page_size, num_kv_heads, head_dim), dtype=dtype) * 0.02 + v_pages = jax.random.normal(keys[2], (total_pages, page_size, num_kv_heads, head_dim), dtype=dtype) * 0.02 + + kv_lens = jnp.full((num_seqs,), max_seq_len_derived, dtype=jnp.int32) + page_indices = jnp.arange(total_pages, dtype=jnp.int32).reshape(num_seqs, pages_per_seq) + cu_q_lens = jnp.arange(num_seqs + 1, dtype=jnp.int32) + + dynamic_args = [queries, k_pages, v_pages, kv_lens, page_indices, cu_q_lens] + static_args = [num_seqs, num_q_heads, num_kv_heads, head_dim, max_seq_len_derived] + + return dynamic_args, static_args + +rtol: 0.01 +atol: 0.01 diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/6p_Paged_Attention/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/6p_Paged_Attention/reference.py new file mode 100644 index 0000000..5c1c50f --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/6p_Paged_Attention/reference.py @@ -0,0 +1,694 @@ +# Imports +from collections.abc import Sequence +import functools +from typing import Literal + +import jax +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +from jax.experimental.pallas.ops.tpu.paged_attention import quantization_utils +import jax.numpy as jnp +import numpy as np + +# Initialization +def get_inputs(dtype=jnp.bfloat16): + CONFIG = { + "name": "llama3_70b_paged_attention", + "model": "Llama-3.1-70B", + "operator": "paged_attention", + "num_seqs": 64, + "max_seq_len": 4096, + "num_query_heads": 64, + "num_kv_heads": 8, + "head_dim": 128, + "page_size": 16, + "pages_per_seq": 256, + } + + key = jax.random.key(42) + keys = jax.random.split(key, 5) + num_seqs = CONFIG["num_seqs"] + num_q_heads = CONFIG["num_query_heads"] + num_kv_heads = CONFIG["num_kv_heads"] + head_dim = CONFIG["head_dim"] + page_size = CONFIG["page_size"] + pages_per_seq = CONFIG["pages_per_seq"] + total_pages = num_seqs * pages_per_seq + max_seq_len_derived = pages_per_seq * page_size + + max_num_tokens = num_seqs + queries = jax.random.normal( + keys[0], (max_num_tokens, num_q_heads, head_dim), dtype=dtype + ) + k_pages = ( + jax.random.normal( + keys[1], (total_pages, page_size, num_kv_heads, head_dim), dtype=dtype + ) + * 0.02 + ) + v_pages = ( + jax.random.normal( + keys[2], (total_pages, page_size, num_kv_heads, head_dim), dtype=dtype + ) + * 0.02 + ) + + kv_lens = jnp.full((num_seqs,), max_seq_len_derived, dtype=jnp.int32) + page_indices = jnp.arange(total_pages, dtype=jnp.int32).reshape( + num_seqs, pages_per_seq + ) + cu_q_lens = jnp.arange(num_seqs + 1, dtype=jnp.int32) + + dynamic_args = [queries, k_pages, v_pages, kv_lens, page_indices, cu_q_lens] + static_args = [ + num_seqs, + num_q_heads, + num_kv_heads, + head_dim, + max_seq_len_derived, + ] + + return dynamic_args, static_args + +# Computation +class MultiPageAsyncCopyDescriptor: + def __init__( + self, + pages_hbm_ref, + scales_pages_hbm_ref, + vmem_buffer, + scales_vmem_buffer, + sem, + page_indices, + page_indices_start_offset, + num_pages_to_load, + head_index, + ): + self._vmem_buffer = vmem_buffer + self._scales_vmem_buffer = scales_vmem_buffer + self._num_pages_to_load = num_pages_to_load + if head_index is not None: + self._pages_hbm_ref = pages_hbm_ref.at[head_index] + if scales_pages_hbm_ref is not None: + self._scales_pages_hbm_ref = scales_pages_hbm_ref.at[head_index] + else: + self._scales_pages_hbm_ref = None + else: + self._pages_hbm_ref = pages_hbm_ref + self._scales_pages_hbm_ref = scales_pages_hbm_ref + self._sem = sem + self._page_indices = page_indices + self._page_indices_start_offset = page_indices_start_offset + self._async_copies = [ + self._make_async_copy(i) for i in range(self._num_pages_to_load) + ] + if ( + self._scales_pages_hbm_ref is not None + and self._scales_vmem_buffer is not None + ): + self._async_copies += [ + self._make_scales_async_copy(i) + for i in range(self._num_pages_to_load) + ] + + def _make_async_copy(self, i): + page_index = self._page_indices[self._page_indices_start_offset + i] + return pltpu.make_async_copy( + self._pages_hbm_ref.at[page_index], self._vmem_buffer.at[i], self._sem + ) + + def _make_scales_async_copy(self, i): + page_index = self._page_indices[self._page_indices_start_offset + i] + return pltpu.make_async_copy( + self._scales_pages_hbm_ref.at[page_index], + self._scales_vmem_buffer.at[i], + self._sem, + ) + + def start(self): + for async_copy in self._async_copies: + async_copy.start() + + def _maybe_dequantize(self, x, x_scale, dtype=jnp.bfloat16): + if x_scale is None: + return x.astype(dtype) + return quantization_utils.from_int8(x, x_scale, dtype=dtype) + + def wait_and_get_loaded(self) -> jax.Array: + for async_copy in self._async_copies: + async_copy.wait() + head_dim = self._vmem_buffer.shape[-1] + jax_array = self._vmem_buffer[...].astype(jnp.float32) + if self._scales_vmem_buffer is not None: + scales_jax_array = self._scales_vmem_buffer[...].astype(jnp.float32) + else: + scales_jax_array = None + jax_array = self._maybe_dequantize(jax_array, scales_jax_array) + return jax_array.reshape(-1, head_dim) + + +def paged_flash_attention_kernel( + lengths_ref, + page_indices_ref, + buffer_index_ref, + init_flag_ref, + q_ref, + k_pages_hbm_ref, + k_scales_pages_hbm_ref, + v_pages_hbm_ref, + v_scales_pages_hbm_ref, + o_ref, + m_ref, + l_ref, + k_vmem_buffer, + k_scales_vmem_buffer, + v_vmem_buffer, + v_scales_vmem_buffer, + k_sems, + v_sems, + *, + batch_size: int, + pages_per_compute_block: int, + pages_per_sequence: int, + mask_value: float, + attn_logits_soft_cap: float | None, + megacore_mode: str | None, + program_ids=(), +): + if program_ids: + core_index, b, h, i = program_ids + else: + core_index, b, h, i = ( + pl.program_id(0), + pl.program_id(1), + pl.program_id(2), + pl.program_id(3), + ) + num_kv_heads, _, page_size, _ = k_pages_hbm_ref.shape + bk = page_size * pages_per_compute_block + num_cores = pl.num_programs(0) + + b_step = num_cores if megacore_mode == "batch" else 1 + b_start = core_index if megacore_mode == "batch" else 0 + h_step = num_cores if megacore_mode == "kv_head" else 1 + h_start = core_index if megacore_mode == "kv_head" else 0 + + h = h * h_step + h_start + b = b * b_step + b_start + length = lengths_ref[b] + + def compute_block_indices(b, h, i): + + def advance_b(): + next_b = b + b_step + + def advance_to_next_non_zero_length(): + next_next_b = next_b + b_step + return lax.fori_loop( + lax.div(next_next_b, b_step), + lax.div(batch_size, b_step), + lambda _, b: jnp.where(lengths_ref[b] == 0, b + b_step, b), + next_next_b, + ) + + return ( + lax.cond( + jnp.logical_and( + next_b < batch_size, + lengths_ref[lax.clamp(0, next_b, batch_size - 1)] == 0), + advance_to_next_non_zero_length, + lambda: next_b, + ), + h_start, + 0, + ) + + def advance_h(): + next_h = h + h_step + return lax.cond(next_h < num_kv_heads, lambda: (b, next_h, 0), advance_b) + + return lax.cond(i * bk < lengths_ref[b], lambda: (b, h, i), advance_h) + + def create_kv_async_copy_descriptors(b, h, i, buffer_index): + page_offset = b * pages_per_sequence + i * pages_per_compute_block + pages_to_load = pages_per_compute_block + async_copy_k = MultiPageAsyncCopyDescriptor( + k_pages_hbm_ref, + k_scales_pages_hbm_ref, + k_vmem_buffer.at[buffer_index], + k_scales_vmem_buffer.at[buffer_index] + if k_scales_vmem_buffer is not None + else None, + k_sems.at[buffer_index], + page_indices_ref, + page_offset, + pages_to_load, + h, + ) + async_copy_v = MultiPageAsyncCopyDescriptor( + v_pages_hbm_ref, + v_scales_pages_hbm_ref, + v_vmem_buffer.at[buffer_index], + v_scales_vmem_buffer.at[buffer_index] + if v_scales_vmem_buffer is not None + else None, + v_sems.at[buffer_index], + page_indices_ref, + page_offset, + pages_to_load, + h, + ) + return async_copy_k, async_copy_v + + @pl.when(i * bk < length) + def flash_attention(): + init_flag = init_flag_ref[0] + init_flag_ref[0] = 0 + buffer_index = buffer_index_ref[0] + next_b, next_h, next_i = compute_block_indices(b, h, i + 1) + + @pl.when(init_flag) + def prefetch_first_block(): + async_copy_k, async_copy_v = create_kv_async_copy_descriptors( + b, h, i, buffer_index + ) + async_copy_k.start() + async_copy_v.start() + + @pl.when(i == 0) + def init(): + m_ref[...] = jnp.full_like(m_ref, -jnp.inf) + l_ref[...] = jnp.zeros_like(l_ref) + o_ref[...] = jnp.zeros_like(o_ref) + + @pl.when(next_b < batch_size) + def prefetch_next_block(): + next_buffer_index = jnp.where(buffer_index == 0, 1, 0) + async_copy_next_k, async_copy_next_v = create_kv_async_copy_descriptors( + next_b, next_h, next_i, next_buffer_index + ) + async_copy_next_k.start() + async_copy_next_v.start() + buffer_index_ref[0] = next_buffer_index + + async_copy_k, async_copy_v = create_kv_async_copy_descriptors( + b, h, i, buffer_index + ) + q = q_ref[...].astype(jnp.float32) + k = async_copy_k.wait_and_get_loaded() + qk = jnp.einsum("gd,td->gt", q, k, preferred_element_type=jnp.float32) + if attn_logits_soft_cap is not None: + capped_qk = jnp.tanh(qk / attn_logits_soft_cap) + qk = capped_qk * attn_logits_soft_cap + + mask = i * bk + jax.lax.broadcasted_iota(jnp.int32, qk.shape, 1) < length + qk = qk + jnp.where(mask, 0.0, mask_value) + m_curr = qk.max(axis=-1) + + s_curr = jnp.exp(qk - m_curr[..., None]) + m_prev, l_prev = m_ref[...], l_ref[...] + l_curr = jax.lax.broadcast_in_dim(s_curr.sum(axis=-1), l_prev.shape, (0,)) + m_curr = jax.lax.broadcast_in_dim(m_curr, m_prev.shape, (0,)) + m_next = jnp.maximum(m_prev, m_curr) + alpha = jnp.exp(m_prev - m_next) + beta = jnp.exp(m_curr - m_next) + l_next = alpha * l_prev + beta * l_curr + m_ref[...], l_ref[...] = m_next, l_next + + v = async_copy_v.wait_and_get_loaded() + o_curr = jnp.einsum("gt,td->gd", s_curr, v) + + o_ref[...] = ( + (l_prev * alpha * o_ref[...] + beta * o_curr) / l_next + ).astype(o_ref.dtype) + + +def paged_flash_attention_kernel_inline_seq_dim( + lengths_ref, + page_indices_ref, + buffer_index_ref, + init_flag_ref, + q_ref, + k_pages_hbm_ref, + k_scales_pages_hbm_ref, + v_pages_hbm_ref, + v_scales_pages_hbm_ref, + o_ref, + m_ref, + l_ref, + k_vmem_buffer, + k_scales_vmem_buffer, + v_vmem_buffer, + v_scales_vmem_buffer, + k_sems, + v_sems, + *, + batch_size: int, + pages_per_compute_block: int, + pages_per_sequence: int, + mask_value: float, + attn_logits_soft_cap: float | None, + megacore_mode: str | None, +): + core_index, b, h = pl.program_id(0), pl.program_id(1), pl.program_id(2) + + m_ref[...] = jnp.full_like(m_ref, -jnp.inf) + l_ref[...] = jnp.zeros_like(l_ref) + o_ref[...] = jnp.zeros_like(o_ref) + + def body(i, _): + paged_flash_attention_kernel( + lengths_ref, + page_indices_ref, + buffer_index_ref, + init_flag_ref, + q_ref, + k_pages_hbm_ref, + k_scales_pages_hbm_ref, + v_pages_hbm_ref, + v_scales_pages_hbm_ref, + o_ref, + m_ref, + l_ref, + k_vmem_buffer, + k_scales_vmem_buffer, + v_vmem_buffer, + v_scales_vmem_buffer, + k_sems, + v_sems, + batch_size=batch_size, + pages_per_compute_block=pages_per_compute_block, + pages_per_sequence=pages_per_sequence, + mask_value=mask_value, + attn_logits_soft_cap=attn_logits_soft_cap, + megacore_mode=megacore_mode, + program_ids=(core_index, b, h, i), + ) + return () + + bk = pages_per_compute_block * k_pages_hbm_ref.shape[-2] + + if megacore_mode == "batch": + num_cores = pl.num_programs(0) + length = lengths_ref[b * num_cores + core_index] + else: + length = lengths_ref[b] + + lax.fori_loop(0, lax.div(length + bk - 1, bk), body, ()) + + +@functools.partial( + jax.jit, + static_argnames=[ + "pages_per_compute_block", + "attn_logits_soft_cap", + "mask_value", + "megacore_mode", + "inline_seq_dim", + ], +) +def paged_attention( + q: jax.Array, + k_pages: jax.Array | quantization_utils.QuantizedTensor, + v_pages: jax.Array | quantization_utils.QuantizedTensor, + lengths: jax.Array, + page_indices: jax.Array, + *, + mask_value: float, + attn_logits_soft_cap: float | None = None, + pages_per_compute_block: int, + megacore_mode: str | None = None, + inline_seq_dim: bool = True, +) -> jax.Array: + if isinstance(k_pages, quantization_utils.QuantizedTensor): + k_pages, k_scales_pages = k_pages.weight, k_pages.scales + assert isinstance(k_scales_pages, jax.Array) + k_scales_pages = jnp.broadcast_to( + k_scales_pages, (*k_scales_pages.shape[:-1], k_pages.shape[-1]) + ) + else: + k_scales_pages = None + if isinstance(v_pages, quantization_utils.QuantizedTensor): + v_pages, v_scales_pages = v_pages.weight, v_pages.scales + assert isinstance(v_scales_pages, jax.Array) + v_scales_pages = jnp.broadcast_to( + v_scales_pages, (*v_scales_pages.shape[:-1], v_pages.shape[-1]) + ) + else: + v_scales_pages = None + + batch_size, num_q_heads, head_dim = q.shape + num_kv_heads, _, page_size, head_dim_k = k_pages.shape + batch_size_paged_indices, pages_per_sequence = page_indices.shape + + if k_pages.shape != v_pages.shape: + raise ValueError( + f"k_pages and v_pages must have the same shape. Got {k_pages.shape} and" + f" {v_pages.shape}" + ) + if num_q_heads % num_kv_heads != 0: + raise ValueError( + "Number of Q heads must be divisible by number of KV heads. Got" + f" {num_q_heads} and {num_kv_heads}." + ) + if head_dim_k != head_dim: + raise ValueError( + "head_dim of Q must be the same as that of K/V. Got" + f" {head_dim} and {head_dim_k}." + ) + if pages_per_sequence % pages_per_compute_block != 0: + raise ValueError( + "pages_per_compute_block must be divisible by pages per sequence. Got" + f" {pages_per_compute_block} and {pages_per_sequence}." + ) + if lengths.shape != (batch_size,): + raise ValueError("`lengths` and `q` must have the same batch size") + if batch_size_paged_indices != batch_size: + raise ValueError("`page_indices` and `q` must have the same batch size") + if lengths.dtype != jnp.int32: + raise ValueError( + f"The dtype of `lengths` must be int32. Got {lengths.dtype}" + ) + + if megacore_mode == "kv_head": + if num_kv_heads % 2 != 0: + raise ValueError( + "number of KV heads must be even when megacore_mode is 'kv_head'" + ) + num_cores = 2 + elif megacore_mode == "batch": + if batch_size % 2 != 0: + raise ValueError("batch size must be even when megacore_mode is 'batch'") + num_cores = 2 + elif megacore_mode is None: + num_cores = 1 + else: + raise ValueError("megacore_mode must be one of ['kv_head', 'batch', None]") + + num_groups = num_q_heads // num_kv_heads + if (num_groups) % 8 != 0: + q = q.reshape(batch_size, num_q_heads, 1, head_dim) + if megacore_mode == "kv_head": + q_block_spec = pl.BlockSpec( + (None, num_groups, None, head_dim), + lambda core_index, b, h, *_: (b, h * num_cores + core_index, 0, 0), + ) + elif megacore_mode == "batch": + q_block_spec = pl.BlockSpec( + (None, num_groups, None, head_dim), + lambda core_index, b, h, *_: (b * num_cores + core_index, h, 0, 0), + ) + else: + q_block_spec = pl.BlockSpec( + (None, num_groups, None, head_dim), + lambda core_index, b, h, *_: (b, h, 0, 0), + ) + q_dtype_for_kernel_launch = jnp.float32 + else: + if megacore_mode == "kv_head": + q_block_spec = pl.BlockSpec( + (None, num_groups, head_dim), + lambda core_index, b, h, *_: (b, h * num_cores + core_index, 0), + ) + elif megacore_mode == "batch": + q_block_spec = pl.BlockSpec( + (None, num_groups, head_dim), + lambda core_index, b, h, *_: (b * num_cores + core_index, h, 0), + ) + else: + q_block_spec = pl.BlockSpec( + (None, num_groups, head_dim), + lambda core_index, b, h, *_: (b, h, 0), + ) + q_dtype_for_kernel_launch = q.dtype + + dimension_semantics: Sequence[Literal["parallel", "arbitrary"]] + if inline_seq_dim: + kernel = paged_flash_attention_kernel_inline_seq_dim + grid = ( + num_cores, + batch_size // num_cores if megacore_mode == "batch" else batch_size, + num_kv_heads // num_cores + if megacore_mode == "kv_head" + else num_kv_heads, + ) + dimension_semantics = ("parallel", "arbitrary", "arbitrary") + else: + kernel = paged_flash_attention_kernel + grid = ( + num_cores, + batch_size // num_cores if megacore_mode == "batch" else batch_size, + num_kv_heads // num_cores + if megacore_mode == "kv_head" + else num_kv_heads, + pages_per_sequence // pages_per_compute_block, + ) + dimension_semantics = ("parallel", "arbitrary", "arbitrary", "arbitrary") + + if k_scales_pages is not None and v_scales_pages is not None: + in_specs = [ + q_block_spec, + pl.BlockSpec(memory_space=pl.ANY), + pl.BlockSpec(memory_space=pl.ANY), + pl.BlockSpec(memory_space=pl.ANY), + pl.BlockSpec(memory_space=pl.ANY), + ] + scratch_shapes = ( + pltpu.VMEM( + ( + 2, + pages_per_compute_block, + page_size, + head_dim, + ), + k_pages.dtype, + ), + pltpu.VMEM( + ( + 2, + pages_per_compute_block, + page_size, + head_dim, + ), + k_scales_pages.dtype, + ), + pltpu.VMEM( + ( + 2, + pages_per_compute_block, + page_size, + head_dim, + ), + v_pages.dtype, + ), + pltpu.VMEM( + ( + 2, + pages_per_compute_block, + page_size, + head_dim, + ), + v_scales_pages.dtype, + ), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + ) + else: + in_specs = [ + q_block_spec, + pl.BlockSpec(memory_space=pl.ANY), + None, + pl.BlockSpec(memory_space=pl.ANY), + None, + ] + scratch_shapes = ( + pltpu.VMEM( + ( + 2, + pages_per_compute_block, + page_size, + head_dim, + ), + k_pages.dtype, + ), + None, + pltpu.VMEM( + ( + 2, + pages_per_compute_block, + page_size, + head_dim, + ), + v_pages.dtype, + ), + None, + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + ) + + out, _, _ = pl.pallas_call( + functools.partial( + kernel, + pages_per_sequence=pages_per_sequence, + batch_size=batch_size, + pages_per_compute_block=pages_per_compute_block, + mask_value=mask_value, + attn_logits_soft_cap=attn_logits_soft_cap, + megacore_mode=megacore_mode, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=4, + in_specs=in_specs, + out_specs=[ + q_block_spec, + q_block_spec, + q_block_spec, + ], + grid=grid, + scratch_shapes=scratch_shapes, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=dimension_semantics + ), + out_shape=[ + jax.ShapeDtypeStruct(q.shape, q_dtype_for_kernel_launch), + jax.ShapeDtypeStruct((*q.shape[:-1], 1), jnp.float32), + jax.ShapeDtypeStruct((*q.shape[:-1], 1), jnp.float32), + ], + )( + lengths, + page_indices.reshape(-1), + jnp.zeros((1,), jnp.int32), + jnp.ones((1,), jnp.int32), + q.astype(q_dtype_for_kernel_launch), + k_pages, + k_scales_pages, + v_pages, + v_scales_pages, + ) + return out.reshape(batch_size, num_q_heads, head_dim).astype(q.dtype) + + +def computation( + queries, + k_pages, + v_pages, + kv_lens, + page_indices, + cu_q_lens, + num_seqs, + num_q_heads, + num_kv_heads, + head_dim, + max_seq_len, +): + DEFAULT_MASK_VALUE = -0.7 * float(np.finfo(np.dtype("float32")).max) + pages_per_compute_block = 128 + k_pages = k_pages.transpose(2, 0, 1, 3) + v_pages = v_pages.transpose(2, 0, 1, 3) + return paged_attention( + queries, k_pages, v_pages, kv_lens, page_indices, + mask_value=DEFAULT_MASK_VALUE, + pages_per_compute_block=pages_per_compute_block, + ) \ No newline at end of file diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/7p_Ragged_Paged_Attention/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/7p_Ragged_Paged_Attention/kernel_task.yaml new file mode 100644 index 0000000..aea37dd --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/7p_Ragged_Paged_Attention/kernel_task.yaml @@ -0,0 +1,66 @@ +task_id: 7p_Ragged_Paged_Attention +description: Kernel task for 7p_Ragged_Paged_Attention +input_gen_code: |- + def get_inputs(): + CONFIG = { + 'name': 'ragged_paged_attention_llama70b', + 'model': 'Llama-3.1-70B', + 'operator': 'ragged_paged_attention', + 'max_num_batched_tokens': 4096, + 'max_num_seqs': 64, + 'num_q_heads': 64, + 'num_kv_heads': 8, + 'head_dim': 128, + 'page_size': 16, + 'pages_per_seq': 256, + } + + ACTIVE_Q_LENS = (1,) * 48 + (512,) * 7 + (464,) + ACTIVE_KV_LENS = tuple( + 257 + ((i * 73) % 240) * 16 for i in range(48) + ) + (1023, 1535, 2047, 2559, 3071, 3583, 4095, 4095) + NUM_ACTIVE_SEQS = len(ACTIVE_Q_LENS) + + dtype = jnp.bfloat16 + key = jax.random.key(42) + k1, k2, k3 = jax.random.split(key, 3) + max_tokens = CONFIG['max_num_batched_tokens'] + max_seqs = CONFIG['max_num_seqs'] + H_q = CONFIG['num_q_heads'] + H_kv = CONFIG['num_kv_heads'] + D = CONFIG['head_dim'] + page_size = CONFIG['page_size'] + pages_per_seq = CONFIG['pages_per_seq'] + num_combined_kv_heads = 2 * H_kv + total_num_pages = max_seqs * pages_per_seq + + q = jax.random.normal(k1, (max_tokens, H_q, D), dtype=dtype) + kv_pages = jax.random.normal( + k2, (total_num_pages, page_size, num_combined_kv_heads, D), dtype=dtype + ) + q_lens = jnp.array(ACTIVE_Q_LENS, dtype=jnp.int32) + kv_lens = jnp.pad( + jnp.array(ACTIVE_KV_LENS, dtype=jnp.int32), + (0, max_seqs - NUM_ACTIVE_SEQS), + ) + active_cu_q_lens = jnp.concatenate( + (jnp.zeros((1,), dtype=jnp.int32), jnp.cumsum(q_lens)) + ) + cu_q_lens = jnp.pad( + active_cu_q_lens, + (0, max_seqs + 1 - active_cu_q_lens.shape[0]), + constant_values=active_cu_q_lens[-1], + ) + + page_indices = jax.random.permutation( + k3, total_num_pages, independent=True + ).astype(jnp.int32).reshape(max_seqs, pages_per_seq) + + num_seqs = jnp.array([NUM_ACTIVE_SEQS], dtype=jnp.int32) + + dynamic_args = [q, kv_pages, kv_lens, page_indices, cu_q_lens, num_seqs] + static_args = [D, max_tokens] + return dynamic_args, static_args + +rtol: 0.01 +atol: 0.01 diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/7p_Ragged_Paged_Attention/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/7p_Ragged_Paged_Attention/reference.py new file mode 100644 index 0000000..13b1461 --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/7p_Ragged_Paged_Attention/reference.py @@ -0,0 +1,898 @@ +# Imports +import functools +import jax +from jax import lax +from jax._src import dtypes +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +from jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes import get_tuned_block_sizes +import jax.numpy as jnp +import math +import json + +# Initialization +def get_inputs(): + CONFIG = { + 'name': 'ragged_paged_attention_llama70b', + 'model': 'Llama-3.1-70B', + 'operator': 'ragged_paged_attention', + 'max_num_batched_tokens': 4096, + 'max_num_seqs': 64, + 'num_q_heads': 64, + 'num_kv_heads': 8, + 'head_dim': 128, + 'page_size': 16, + 'pages_per_seq': 256, + } + + ACTIVE_Q_LENS = (1,) * 48 + (512,) * 7 + (464,) + ACTIVE_KV_LENS = tuple( + 257 + ((i * 73) % 240) * 16 for i in range(48) + ) + (1023, 1535, 2047, 2559, 3071, 3583, 4095, 4095) + NUM_ACTIVE_SEQS = len(ACTIVE_Q_LENS) + + dtype = jnp.bfloat16 + key = jax.random.key(42) + k1, k2, k3 = jax.random.split(key, 3) + max_tokens = CONFIG['max_num_batched_tokens'] + max_seqs = CONFIG['max_num_seqs'] + H_q = CONFIG['num_q_heads'] + H_kv = CONFIG['num_kv_heads'] + D = CONFIG['head_dim'] + page_size = CONFIG['page_size'] + pages_per_seq = CONFIG['pages_per_seq'] + num_combined_kv_heads = 2 * H_kv + total_num_pages = max_seqs * pages_per_seq + + q = jax.random.normal(k1, (max_tokens, H_q, D), dtype=dtype) + kv_pages = jax.random.normal( + k2, (total_num_pages, page_size, num_combined_kv_heads, D), dtype=dtype + ) + q_lens = jnp.array(ACTIVE_Q_LENS, dtype=jnp.int32) + kv_lens = jnp.pad( + jnp.array(ACTIVE_KV_LENS, dtype=jnp.int32), + (0, max_seqs - NUM_ACTIVE_SEQS), + ) + active_cu_q_lens = jnp.concatenate( + (jnp.zeros((1,), dtype=jnp.int32), jnp.cumsum(q_lens)) + ) + cu_q_lens = jnp.pad( + active_cu_q_lens, + (0, max_seqs + 1 - active_cu_q_lens.shape[0]), + constant_values=active_cu_q_lens[-1], + ) + + page_indices = jax.random.permutation( + k3, total_num_pages, independent=True + ).astype(jnp.int32).reshape(max_seqs, pages_per_seq) + + num_seqs = jnp.array([NUM_ACTIVE_SEQS], dtype=jnp.int32) + + dynamic_args = [q, kv_pages, kv_lens, page_indices, cu_q_lens, num_seqs] + static_args = [D, max_tokens] + return dynamic_args, static_args + +# Computation +class MultiPageAsyncCopyDescriptor: + def __init__( + self, + pages_hbm_ref, + vmem_buf, + sem, + page_indices_ref, + metadata, + ): + self._vmem_buf = vmem_buf + seq_id, start_page_idx, end_page_idx = metadata + self._async_copies = [] + for i in range(vmem_buf.shape[0]): + page_idx = start_page_idx + i + page_idx = jax.lax.select(page_idx < end_page_idx, page_idx, 0) + self._async_copies.append( + pltpu.make_async_copy( + pages_hbm_ref.at[page_indices_ref[seq_id, page_idx]], + vmem_buf.at[i], + sem, + ) + ) + + def start(self): + for async_copy in self._async_copies: + async_copy.start() + + def wait(self): + for async_copy in self._async_copies: + async_copy.wait() + return self._vmem_buf + + +def ref_ragged_paged_attention( + queries: jax.Array, + kv_pages: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + num_seqs: jax.Array, + *, + sm_scale: float = 1.0, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = -0.7 * float(jnp.finfo(jnp.dtype("float32")).max), + k_scale: float | None = None, + v_scale: float | None = None, +): + static_validate_inputs( + queries, + kv_pages, + kv_lens, + page_indices, + cu_q_lens, + num_seqs, + sm_scale=sm_scale, + k_scale=k_scale, + v_scale=v_scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + mask_value=mask_value, + ) + if mask_value is None: + mask_value = -0.7 * float(jnp.finfo(jnp.dtype("float32")).max) + _, _, num_combined_kv_heads, head_dim = kv_pages.shape + assert num_combined_kv_heads % 2 == 0 + num_kv_heads = num_combined_kv_heads // 2 + num_q_heads = queries.shape[1] + assert num_q_heads % num_kv_heads == 0 + num_query_per_kv = num_q_heads // num_kv_heads + outputs = [] + for i in range(num_seqs[0]): + q_start = cu_q_lens[i] + q_end = cu_q_lens[i + 1] + q_len = q_end - q_start + kv_len = kv_lens[i] + indices = page_indices[i] + q = queries[q_start:q_end] + k = kv_pages[indices, :, 0::2, :].reshape(-1, num_kv_heads, head_dim)[ + :kv_len + ] + v = kv_pages[indices, :, 1::2, :].reshape(-1, num_kv_heads, head_dim)[ + :kv_len + ] + if k_scale is not None: + k = k.astype(jnp.float32) * k_scale + k = k.astype(q.dtype) + if v_scale is not None: + v = v.astype(jnp.float32) * v_scale + v = v.astype(q.dtype) + k = jnp.repeat(k, num_query_per_kv, axis=1) + v = jnp.repeat(v, num_query_per_kv, axis=1) + attn = jnp.einsum("qhd,khd->hqk", q, k, preferred_element_type=jnp.float32) + attn *= sm_scale + q_span = (kv_len - q_len) + jax.lax.broadcasted_iota( + jnp.int32, attn.shape, 1 + ) + kv_span = jax.lax.broadcasted_iota(jnp.int32, attn.shape, 2) + mask = q_span < kv_span + if sliding_window is not None: + mask = jnp.logical_or(mask, q_span - sliding_window >= kv_span) + if soft_cap is not None: + attn = soft_cap * jnp.tanh(attn / soft_cap) + attn += jnp.where(mask, mask_value, 0.0) + attn = jax.nn.softmax(attn, axis=-1).astype(v.dtype) + out = jnp.einsum("hqk,khd->qhd", attn, v).astype(queries.dtype) + outputs.append(out) + + return jnp.concatenate(outputs, axis=0) + + +def dynamic_validate_inputs( + q: jax.Array, + kv_pages: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + num_seqs: jax.Array, + *, + sm_scale: float | None = None, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = None, + k_scale: float | None = None, + v_scale: float | None = None, + num_kv_pages_per_block: int | None = None, + num_queries_per_block: int | None = None, + vmem_limit_bytes: int | None = None, +): + static_validate_inputs( + q, + kv_pages, + kv_lens, + page_indices, + cu_q_lens, + num_seqs, + sm_scale=sm_scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + mask_value=mask_value, + k_scale=k_scale, + v_scale=v_scale, + num_kv_pages_per_block=num_kv_pages_per_block, + num_queries_per_block=num_queries_per_block, + vmem_limit_bytes=vmem_limit_bytes, + ) + max_num_batched_tokens = q.shape[0] + page_size = kv_pages.shape[1] + max_num_seqs, pages_per_seq = page_indices.shape + if num_seqs[0] > max_num_seqs: + raise ValueError(f"{num_seqs[0]=} must be less or equal to {max_num_seqs=}") + max_kv_len = jnp.max(kv_lens) + min_pages_per_seq = pl.cdiv(max_kv_len, page_size) + if pages_per_seq < min_pages_per_seq: + raise ValueError( + f"{pages_per_seq=} must be greater or equal to" + f" {min_pages_per_seq=} given {max_kv_len=} and {page_size=}." + ) + if cu_q_lens[num_seqs[0]] > max_num_batched_tokens: + raise ValueError( + f"Total q tokens {cu_q_lens[num_seqs[0]]} must be less or equal to" + f" {max_num_batched_tokens=}." + ) + for i in range(num_seqs[0]): + q_len = cu_q_lens[i + 1] - cu_q_lens[i] + kv_len = kv_lens[i] + if q_len > kv_len: + raise ValueError( + f"{q_len=} must be less or equal to {kv_len=} at sequence {i}." + ) + + +def static_validate_inputs( + q: jax.Array, + kv_pages: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + num_seqs: jax.Array, + *, + sm_scale: float | None = None, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = None, + k_scale: float | None = None, + v_scale: float | None = None, + num_kv_pages_per_block: int | None = None, + num_queries_per_block: int | None = None, + vmem_limit_bytes: int | None = None, +): + _, num_q_heads, head_dim = q.shape + _, _, num_combined_kv_heads, head_dim_k = kv_pages.shape + assert num_combined_kv_heads % 2 == 0 + assert isinstance(k_scale, float) or k_scale is None + assert isinstance(v_scale, float) or v_scale is None + num_kv_heads = num_combined_kv_heads // 2 + max_num_seqs, pages_per_seq = page_indices.shape + if num_seqs.shape != (1,): + raise ValueError(f"{num_seqs.shape=} must be (1,)") + if head_dim_k != head_dim: + raise ValueError( + f"Q head_dim {head_dim} must be the same as that of K/V {head_dim_k}." + ) + if kv_lens.shape != (max_num_seqs,): + raise ValueError( + f"Expected {kv_lens.shape=} to be ({max_num_seqs},) where" + " `max_num_seqs` is `page_indices.shape[0]`." + ) + if cu_q_lens.shape != (max_num_seqs + 1,): + raise ValueError( + f"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},) where" + " `max_num_seqs` is `page_indices.shape[0]`." + ) + if ( + kv_lens.dtype != jnp.int32 + or page_indices.dtype != jnp.int32 + or cu_q_lens.dtype != jnp.int32 + ): + raise ValueError( + "The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must be" + f" int32. Got {kv_lens.dtype=}, {page_indices.dtype=}," + f" {cu_q_lens.dtype=}." + ) + if num_q_heads % num_kv_heads != 0: + raise ValueError(f"{num_q_heads=} must be divisible by {num_kv_heads=}") + if sliding_window is not None and sliding_window <= 0: + raise ValueError(f"{sliding_window=} must be positive.") + if soft_cap is not None and soft_cap == 0.0: + raise ValueError(f"{soft_cap=} must not be 0.0.") + if ( + num_kv_pages_per_block is not None + and not 0 < num_kv_pages_per_block <= pages_per_seq + ): + raise ValueError( + f"{num_kv_pages_per_block=} must be in range (0, {pages_per_seq}]." + ) + if num_queries_per_block is not None and num_queries_per_block <= 0: + raise ValueError(f"{num_queries_per_block=} must be positive.") + if vmem_limit_bytes is not None and vmem_limit_bytes <= 0: + raise ValueError(f"{vmem_limit_bytes=} must be positive.") + del sm_scale + del mask_value + + +def ragged_paged_attention_kernel( + kv_lens_ref, + page_indices_ref, + cu_q_lens_ref, + seq_buf_idx_ref, + num_seqs_ref, + q_ref, + kv_pages_hbm_ref, + o_ref, + kv_bufs, + sems, + l_ref, + m_ref, + acc_ref, + *, + sm_scale: float, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = -0.7 * float(jnp.finfo(jnp.dtype("float32")).max), + k_scale: float | None = None, + v_scale: float | None = None, +): + if mask_value is None: + mask_value = -0.7 * float(jnp.finfo(jnp.dtype("float32")).max) + num_q_per_blk, num_q_heads_per_blk, head_dim = q_ref.shape + pages_per_seq = page_indices_ref.shape[-1] + num_seqs = num_seqs_ref[0] + _, num_kv_pages_per_blk, page_size, num_combined_kv_heads_per_blk, _ = ( + kv_bufs.shape + ) + num_kv_heads_per_blk = num_combined_kv_heads_per_blk // 2 + num_kv_per_blk = num_kv_pages_per_blk * page_size + num_q_heads_per_kv_head = num_q_heads_per_blk // num_kv_heads_per_blk + heads_blk_idx, q_blk_idx = ( + pl.program_id(0), + pl.program_id(1), + ) + num_heads_blks = pl.num_programs(0) + init_seq_idx = seq_buf_idx_ref[0] + init_buf_idx = seq_buf_idx_ref[1] + q_len_start = q_blk_idx * num_q_per_blk + q_len_end = q_len_start + num_q_per_blk + + def create_kv_async_copy_descriptors( + heads_blk_idx, seq_idx, kv_blk_idx, buf_idx + ): + start_kv_page_idx = kv_blk_idx * num_kv_pages_per_blk + end_kv_page_idx = jnp.minimum( + pages_per_seq, pl.cdiv(kv_lens_ref[seq_idx], page_size) + ) + metadata = (seq_idx, start_kv_page_idx, end_kv_page_idx) + heads_start = heads_blk_idx * num_combined_kv_heads_per_blk + async_copy_kv = MultiPageAsyncCopyDescriptor( + kv_pages_hbm_ref.at[ + :, :, pl.ds(heads_start, num_combined_kv_heads_per_blk), : + ], + kv_bufs.at[buf_idx], + sems.at[buf_idx], + page_indices_ref, + metadata, + ) + return async_copy_kv + + def strided_load_kv(ref, start, step): + packing = get_dtype_packing(ref.dtype) + if packing == 1: + return [ref[start::step, :]], [ref[start + 1 :: step, :]] + assert packing in (2, 4, 8) + assert step % packing == 0 + k_list, v_list = [], [] + b_start = start // packing + b_step = step // packing + b_ref = ref.bitcast(jnp.uint32) + b = b_ref[b_start::b_step, :] + + if ref.dtype == jnp.bfloat16: + bk = b << 16 + bv = b & jnp.uint32(0xFFFF0000) + k = pltpu.bitcast(bk, jnp.float32).astype(jnp.bfloat16) + v = pltpu.bitcast(bv, jnp.float32).astype(jnp.bfloat16) + k_list.append(k) + v_list.append(v) + else: + bitwidth = 32 // packing + bitcast_dst_dtype = jnp.dtype(f"uint{bitwidth}") + for i in range(0, packing, 2): + bk = b >> (i * bitwidth) + k = pltpu.bitcast(bk.astype(bitcast_dst_dtype), ref.dtype) + k_list.append(k) + bv = b >> ((i + 1) * bitwidth) + v = pltpu.bitcast(bv.astype(bitcast_dst_dtype), ref.dtype) + v_list.append(v) + + return k_list, v_list + + def fold_on_2nd_minor(vec): + assert vec.dtype == jnp.bfloat16 or vec.dtype == jnp.float32 + assert len(vec.shape) >= 2 + last_dim = vec.shape[-1] + packing = get_dtype_packing(vec.dtype) + if vec.shape[-2] % packing != 0: + vec = vec.astype(jnp.float32) + return vec.reshape(-1, last_dim) + + @pl.when(heads_blk_idx + q_blk_idx == 0) + def prefetch_first_kv_blk(): + async_copy_kv = create_kv_async_copy_descriptors( + heads_blk_idx, init_seq_idx, 0, init_buf_idx + ) + async_copy_kv.start() + + def is_cur_q_blk_needed(q_states): + done, cur_seq_idx, _ = q_states + should_run = jnp.logical_and(q_len_start < cu_q_lens_ref[num_seqs], + cur_seq_idx < num_seqs) + return jnp.logical_and(done == 0, should_run) + + def compute_with_cur_q_blk(q_states): + done, cur_seq_idx, cur_buf_idx = q_states + q_start = cu_q_lens_ref[cur_seq_idx] + q_end = cu_q_lens_ref[cur_seq_idx + 1] + q_len = q_end - q_start + kv_len = kv_lens_ref[cur_seq_idx] + + def get_next_prefetch_ids( + heads_blk_idx, cur_seq_idx, kv_blk_idx, cur_buf_idx + ): + next_kv_blk_idx = kv_blk_idx + 1 + is_last_kv_blk = next_kv_blk_idx * num_kv_per_blk >= kv_len + next_kv_blk_idx = lax.select( + is_last_kv_blk, + 0, + next_kv_blk_idx, + ) + is_cur_seq_end_in_cur_q_blk = q_end <= q_len_end + next_seq_idx = lax.select( + is_last_kv_blk, + lax.select(is_cur_seq_end_in_cur_q_blk, cur_seq_idx + 1, cur_seq_idx), + cur_seq_idx, + ) + is_last_seq = next_seq_idx == num_seqs + next_seq_idx = lax.select( + is_last_seq, + 0, + next_seq_idx, + ) + next_heads_blk_idx = lax.select( + is_last_seq, + heads_blk_idx + 1, + heads_blk_idx, + ) + next_buf_idx = lax.select(cur_buf_idx == 0, 1, 0) + return next_heads_blk_idx, next_seq_idx, next_kv_blk_idx, next_buf_idx + + def flash_attention( + q, + k, + v, + head_l_ref, + head_m_ref, + head_acc_ref, + *, + kv_blk_idx, + ): + assert q.shape == ( + num_q_per_blk * num_q_heads_per_kv_head, + head_dim, + ) + assert ( + k.shape + == v.shape + == ( + num_kv_per_blk, + head_dim, + ) + ) + assert k.dtype == v.dtype + assert ( + head_m_ref.shape + == head_l_ref.shape + == ( + num_q_per_blk * num_q_heads_per_kv_head, + 128, + ) + ) + assert head_acc_ref.shape == ( + num_q_per_blk, + num_q_heads_per_kv_head, + head_dim, + ) + kv_len_start = kv_blk_idx * num_kv_per_blk + + def masked_store(ref, val, start, end, group=1): + iota = lax.broadcasted_iota(jnp.int32, ref.shape, 0) // group + pltpu.store(ref, val, mask=jnp.logical_and(iota >= start, iota < end)) + + def load_with_init(ref, init_val): + return jnp.where( + kv_blk_idx == 0, jnp.full_like(ref, init_val), ref[...] + ) + + kv_mask = ( + lax.broadcasted_iota(jnp.int32, k.shape, 0) < kv_len - kv_len_start + ) + k = jnp.where(kv_mask, k.astype(jnp.float32), 0).astype(k.dtype) + v = jnp.where(kv_mask, v.astype(jnp.float32), 0).astype(v.dtype) + + qk = ( + jnp.einsum("nd,md->nm", q, k, preferred_element_type=jnp.float32) + * sm_scale + ) + store_start = jnp.maximum(q_start - q_len_start, 0) + store_end = jnp.minimum(q_end - q_len_start, num_q_per_blk) + + row_ids = ( + (kv_len - q_len) + + q_len_start + - q_start + + jax.lax.broadcasted_iota( + jnp.int32, + (num_q_per_blk * num_q_heads_per_kv_head, num_kv_per_blk), + 0, + ) + // num_q_heads_per_kv_head + ) + col_ids = kv_len_start + jax.lax.broadcasted_iota( + jnp.int32, + (num_q_per_blk * num_q_heads_per_kv_head, num_kv_per_blk), + 1, + ) + causal_mask = row_ids < col_ids + if sliding_window is not None: + causal_mask = jnp.logical_or(causal_mask, + row_ids - sliding_window >= col_ids) + if soft_cap is not None: + qk = soft_cap * jnp.tanh(qk / soft_cap) + qk += jnp.where(causal_mask, mask_value, 0.0) + m_curr = jnp.max(qk, axis=1, keepdims=True) + s_curr = jnp.exp(qk - m_curr) + qkv = jnp.dot(s_curr, v, preferred_element_type=jnp.float32) + lm_store_shape = head_m_ref.shape + m_curr = jnp.broadcast_to(m_curr, lm_store_shape) + l_curr = jnp.broadcast_to( + s_curr.sum(axis=1, keepdims=True), lm_store_shape + ) + m_prev = load_with_init(head_m_ref, -jnp.inf) + l_prev = load_with_init(head_l_ref, 0.0) + m_next = jnp.maximum(m_prev, m_curr) + masked_store( + head_m_ref, m_next, store_start, store_end, num_q_heads_per_kv_head + ) + alpha = jnp.exp(m_prev - m_next) + beta = jnp.exp(m_curr - m_next) + l_alpha = alpha * l_prev + l_next = l_alpha + beta * l_curr + l_next_safe = jnp.where(l_next == 0.0, 1.0, l_next) + masked_store( + head_l_ref, + l_next_safe, + store_start, + store_end, + num_q_heads_per_kv_head, + ) + + def broadcast_to_shape(arr, shape): + if arr.shape == shape: + return arr + assert len(arr.shape) == len(shape) + assert arr.shape[0] == shape[0] + assert shape[1] % arr.shape[1] == 0 + return jnp.concatenate( + [arr for _ in range(shape[1] // arr.shape[1])], axis=1 + ) + + o_curr = load_with_init(head_acc_ref, 0.0).reshape(-1, head_dim) + l_alpha = broadcast_to_shape(l_alpha, qkv.shape) + beta = broadcast_to_shape(beta, qkv.shape) + l_next_safe = broadcast_to_shape(l_next_safe, qkv.shape) + out = lax.div( + l_alpha * o_curr + beta * qkv, + l_next_safe, + ) + masked_store( + head_acc_ref, + out.reshape(head_acc_ref.shape), + store_start, + store_end, + ) + + def is_valid_kv_blk_in_cur_seq(kv_states): + kv_blk_idx, _ = kv_states + return kv_blk_idx * num_kv_per_blk < kv_len + + def compute_with_kv_blk_in_cur_seq(kv_states): + kv_blk_idx, cur_buf_idx = kv_states + next_heads_blk_idx, next_seq_idx, next_kv_blk_idx, next_buf_idx = ( + get_next_prefetch_ids( + heads_blk_idx, cur_seq_idx, kv_blk_idx, cur_buf_idx + ) + ) + + @pl.when(next_heads_blk_idx < num_heads_blks) + def prefetch_next_kv_blk(): + next_async_copy_kv = create_kv_async_copy_descriptors( + next_heads_blk_idx, next_seq_idx, next_kv_blk_idx, next_buf_idx + ) + next_async_copy_kv.start() + + cur_async_copy_kv = create_kv_async_copy_descriptors( + heads_blk_idx, cur_seq_idx, kv_blk_idx, cur_buf_idx + ) + kv_ref = cur_async_copy_kv.wait().reshape( + num_kv_pages_per_blk * page_size * num_combined_kv_heads_per_blk, + head_dim, + ) + kv_packing = get_dtype_packing(kv_ref.dtype) + kv_load_step = max(1, kv_packing // 2) + for kv_head_chunk_idx in range(0, num_kv_heads_per_blk, kv_load_step): + k_list, v_list = strided_load_kv( + kv_ref, kv_head_chunk_idx * 2, num_combined_kv_heads_per_blk + ) + for step_idx in range(kv_load_step): + k = k_list[step_idx] + v = v_list[step_idx] + if k_scale is not None: + k = k.astype(jnp.float32) * k_scale + k = k.astype(q_ref.dtype) + if v_scale is not None: + v = v.astype(jnp.float32) * v_scale + v = v.astype(q_ref.dtype) + kv_head_idx = kv_head_chunk_idx + step_idx + q_head_idx = kv_head_idx * num_q_heads_per_kv_head + q = fold_on_2nd_minor( + q_ref[:, q_head_idx : q_head_idx + num_q_heads_per_kv_head, :] + ) + flash_attention( + q, + k, + v, + l_ref.at[kv_head_idx], + m_ref.at[kv_head_idx], + acc_ref.at[ + :, q_head_idx : q_head_idx + num_q_heads_per_kv_head, : + ], + kv_blk_idx=kv_blk_idx, + ) + return kv_blk_idx + 1, next_buf_idx + + _, next_buf_idx = lax.while_loop( + is_valid_kv_blk_in_cur_seq, + compute_with_kv_blk_in_cur_seq, + (0, cur_buf_idx), + ) + next_seq_idx = lax.select(q_end <= q_len_end, cur_seq_idx + 1, cur_seq_idx) + done = lax.select(q_end < q_len_end, done, 1) + return done, next_seq_idx, next_buf_idx + + _, seq_idx, buf_idx = lax.while_loop( + is_cur_q_blk_needed, + compute_with_cur_q_blk, + (0, init_seq_idx, init_buf_idx), + ) + seq_buf_idx_ref[0] = lax.select(seq_idx < num_seqs, seq_idx, 0) + seq_buf_idx_ref[1] = buf_idx + q_ids = q_len_start + lax.broadcasted_iota(jnp.int32, o_ref.shape, 0) + q_is_active = q_ids < cu_q_lens_ref[num_seqs] + o_ref[...] = jnp.where( + q_is_active, + acc_ref[...].astype(q_ref.dtype), + jnp.zeros(o_ref.shape, dtype=q_ref.dtype), + ) + + +def get_dtype_packing(dtype): + bits = dtypes.itemsize_bits(dtype) + return 32 // bits + + +def get_min_heads_per_blk( + num_q_heads, num_combined_kv_heads, q_dtype, kv_dtype +): + q_packing = get_dtype_packing(q_dtype) + kv_packing = get_dtype_packing(kv_dtype) + + def can_be_xla_fully_tiled(x, packing): + if x % packing != 0: + return False + x //= packing + return x in (1, 2, 4, 8) or x % 8 == 0 + + if not can_be_xla_fully_tiled(num_combined_kv_heads, kv_packing): + raise ValueError( + f"Not implemented: {num_combined_kv_heads=} can not be XLA fully tiled." + ) + assert num_combined_kv_heads % 2 == 0 + num_kv_heads = num_combined_kv_heads // 2 + assert num_q_heads % num_kv_heads == 0 + ratio = num_q_heads // num_kv_heads + max_combined_kv_tiling = 8 * kv_packing + min_combined_kv_heads = ( + max_combined_kv_tiling + if num_combined_kv_heads % max_combined_kv_tiling == 0 + else num_combined_kv_heads + ) + min_q_heads = min_combined_kv_heads // 2 * ratio + if can_be_xla_fully_tiled(min_q_heads, q_packing): + return min_q_heads, min_combined_kv_heads + return num_q_heads, num_combined_kv_heads + + +@functools.partial( + jax.jit, + static_argnames=[ + "sm_scale", + "mask_value", + "num_kv_pages_per_block", + "num_queries_per_block", + "vmem_limit_bytes", + "sliding_window", + "soft_cap", + "k_scale", + "v_scale", + ], +) +def ragged_paged_attention( + q: jax.Array, + kv_pages: jax.Array, + kv_lens: jax.Array, + page_indices: jax.Array, + cu_q_lens: jax.Array, + num_seqs: jax.Array, + *, + sm_scale: float = 1.0, + sliding_window: int | None = None, + soft_cap: float | None = None, + mask_value: float | None = -0.7 * float(jnp.finfo(jnp.dtype("float32")).max), + k_scale: float | None = None, + v_scale: float | None = None, + num_kv_pages_per_block: int | None = None, + num_queries_per_block: int | None = None, + vmem_limit_bytes: int | None = None, +): + static_validate_inputs( + q, + kv_pages, + kv_lens, + page_indices, + cu_q_lens, + num_seqs, + sm_scale=sm_scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + mask_value=mask_value, + k_scale=k_scale, + v_scale=v_scale, + num_kv_pages_per_block=num_kv_pages_per_block, + num_queries_per_block=num_queries_per_block, + vmem_limit_bytes=vmem_limit_bytes, + ) + if mask_value is None: + mask_value = -0.7 * float(jnp.finfo(jnp.dtype("float32")).max) + num_q_tokens, num_q_heads, head_dim = q.shape + _, page_size, num_combined_kv_heads, _ = kv_pages.shape + assert num_combined_kv_heads % 2 == 0 + num_kv_heads = num_combined_kv_heads // 2 + _, pages_per_seq = page_indices.shape + num_q_heads_per_blk, num_combined_kv_heads_per_blk = get_min_heads_per_blk( + num_q_heads, num_combined_kv_heads, q.dtype, kv_pages.dtype + ) + num_q_per_blk = num_queries_per_block + num_kv_pages_per_blk = num_kv_pages_per_block + if num_q_per_blk is None or num_kv_pages_per_blk is None: + num_kv_pages_per_blk, num_q_per_blk = get_tuned_block_sizes( + q.dtype, + kv_pages.dtype, + num_q_heads_per_blk, + num_combined_kv_heads_per_blk // 2, + head_dim, + page_size, + num_q_tokens, + pages_per_seq, + ) + num_q_heads_per_kv_head = num_q_heads // num_kv_heads + num_q_blks = pl.cdiv(num_q_tokens, num_q_per_blk) + assert num_combined_kv_heads_per_blk % 2 == 0 + num_kv_heads_per_blk = num_combined_kv_heads_per_blk // 2 + assert num_q_heads_per_blk % num_q_heads_per_kv_head == 0 + num_heads_blks = num_q_heads // num_q_heads_per_blk + grid = (num_heads_blks, num_q_blks) + + def q_index_map(heads_blk_idx, q_blk_idx, *_): + return (q_blk_idx, heads_blk_idx, 0) + + q_block_spec = pl.BlockSpec( + (num_q_per_blk, num_q_heads_per_blk, head_dim), + q_index_map, + ) + in_specs = [ + q_block_spec, + pl.BlockSpec(memory_space=pl.ANY), + ] + out_specs = q_block_spec + lm_scratch = pltpu.VMEM( + (num_kv_heads_per_blk, num_q_per_blk * num_q_heads_per_kv_head, 128), + jnp.float32, + ) + acc_scratch = pltpu.VMEM( + (num_q_per_blk, num_q_heads_per_blk, head_dim), + jnp.float32, + ) + double_buf_scratch = pltpu.VMEM( + ( + 2, + num_kv_pages_per_blk, + page_size, + num_combined_kv_heads_per_blk, + head_dim, + ), + kv_pages.dtype, + ) + scratch_shapes = [ + double_buf_scratch, + pltpu.SemaphoreType.DMA((2,)), + lm_scratch, + lm_scratch, + acc_scratch, + ] + scalar_prefetches = ( + kv_lens, + page_indices, + cu_q_lens, + jnp.array((0, 0), jnp.int32), + num_seqs, + ) + kernel = pl.pallas_call( + functools.partial( + ragged_paged_attention_kernel, + sm_scale=sm_scale, + sliding_window=sliding_window, + soft_cap=soft_cap, + mask_value=mask_value, + k_scale=k_scale, + v_scale=v_scale, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=len(scalar_prefetches), + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + scratch_shapes=scratch_shapes, + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=( + "arbitrary", + "arbitrary", + ), + vmem_limit_bytes=vmem_limit_bytes, + ), + out_shape=jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype), + name="ragged_paged_attention_kernel", + ) + + return kernel(*scalar_prefetches, q, kv_pages) + + +def computation(q, kv_pages, kv_lens, page_indices, cu_q_lens, num_seqs, D, max_tokens): + TUNED_PARAMS = { + 'num_kv_pages_per_block': 64, + 'num_queries_per_block': 64, + 'vmem_limit_bytes': 33554432, + } + sm_scale = 1.0 / math.sqrt(D) + return ragged_paged_attention( + q, kv_pages, kv_lens, page_indices, cu_q_lens, num_seqs, + sm_scale=sm_scale, + num_kv_pages_per_block=TUNED_PARAMS['num_kv_pages_per_block'], + num_queries_per_block=TUNED_PARAMS['num_queries_per_block'], + vmem_limit_bytes=TUNED_PARAMS['vmem_limit_bytes'], + ) \ No newline at end of file diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/8p_GEMM/kernel_task.yaml b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/8p_GEMM/kernel_task.yaml new file mode 100644 index 0000000..ec833ac --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/8p_GEMM/kernel_task.yaml @@ -0,0 +1,23 @@ +task_id: 8p_GEMM +description: Kernel task for 8p_GEMM +input_gen_code: |- + def get_inputs(dtype=jnp.bfloat16): + import jax + import jax.numpy as jnp + import time + import numpy as np + import json + + M = 8192 + K = 8192 + N = 28672 + key = jax.random.key(42) + k1, k2 = jax.random.split(key, 2) + A = jax.random.normal(k1, (M, K), dtype=dtype) + B = jax.random.normal(k2, (K, N), dtype=dtype) * 0.02 + dynamic_args = [A, B] + static_args = [] + return dynamic_args, static_args + +rtol: 0.01 +atol: 0.01 diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/8p_GEMM/reference.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/8p_GEMM/reference.py new file mode 100644 index 0000000..1f2516b --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/level2/8p_GEMM/reference.py @@ -0,0 +1,77 @@ +# Imports +import functools +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +# Initialization +def get_inputs(): + M = 8192 + K = 8192 + N = 28672 + dtype = jnp.bfloat16 + key = jax.random.key(42) + k1, k2 = jax.random.split(key, 2) + x = jax.random.normal(k1, (M, K), dtype=dtype) + y = jax.random.normal(k2, (K, N), dtype=dtype) * 0.02 + return [x, y], [] + +# Computation +def matmul_kernel(x_tile_ref, y_tile_ref, o_tile_ref, acc_ref): + @pl.when(pl.program_id(2) == 0) + def init(): + acc_ref[...] = jnp.zeros_like(acc_ref) + + acc_ref[...] = acc_ref[...] + jnp.dot( + x_tile_ref[...], + y_tile_ref[...], + preferred_element_type=acc_ref.dtype, + ) + o_tile_ref[...] = acc_ref[...].astype(o_tile_ref.dtype) + +@functools.partial( + jax.jit, static_argnames=["block_shape", "block_k", "debug", "out_dtype"] +) +def matmul( + x: jax.Array, + y: jax.Array, + *, + block_shape, + block_k: int = 256, + out_dtype: jnp.dtype | None = None, + debug: bool = False, +) -> jax.Array: + if out_dtype is None: + if x.dtype != y.dtype: + raise TypeError( + f"Cannot deduce output dtype for different input dtypes: {x.dtype}," + f" {y.dtype}" + ) + out_dtype = x.dtype + acc_dtype = jnp.float32 + if x.dtype in [jnp.int8, jnp.int4, jnp.uint8, jnp.uint4]: + acc_dtype = jnp.int32 + + l, r = block_shape + return pl.pallas_call( + matmul_kernel, + out_shape=jax.ShapeDtypeStruct((x.shape[0], y.shape[1]), out_dtype), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=0, + in_specs=[ + pl.BlockSpec((l, block_k), lambda i, _, k: (i, k)), + pl.BlockSpec((block_k, r), lambda _, j, k: (k, j)), + ], + out_specs=pl.BlockSpec((l, r), lambda i, j, k: (i, j)), + grid=(x.shape[0] // l, y.shape[1] // r, x.shape[1] // block_k), + scratch_shapes=[pltpu.VMEM((l, r), acc_dtype)], + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "parallel", "arbitrary")), + debug=debug, + )(x, y) + +def computation(x, y): + TUNED_PARAMS = {'block_shape': [1024, 2048], 'block_k': 1024} + return matmul(x, y, block_shape=tuple(TUNED_PARAMS['block_shape']), block_k=TUNED_PARAMS['block_k']) \ No newline at end of file diff --git a/MaxKernel/evaluation/jaxbench_adapted_dataset/verify_with_jaxbench.py b/MaxKernel/evaluation/jaxbench_adapted_dataset/verify_with_jaxbench.py new file mode 100755 index 0000000..4015ae3 --- /dev/null +++ b/MaxKernel/evaluation/jaxbench_adapted_dataset/verify_with_jaxbench.py @@ -0,0 +1,192 @@ +import os +import sys +import importlib.util +import argparse +import jax +import jax.numpy as jnp +import numpy as np + +sys.dont_write_bytecode = True + +def load_module(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + +def baseline_is_single(inputs): + # Single config: a flat sequence of arrays. Multi: a list of arg lists. + return not all(isinstance(a, (list, tuple)) for a in inputs) + +def reference_is_single(ret): + # Single config: (dynamic_args, static_args). Multi: a list of those pairs. + first = ret[0] + return not ( + isinstance(first, (list, tuple)) + and len(first) == 2 + and isinstance(first[0], (list, tuple)) + ) + +def check_match(baseline_out, reference_out, problem_name): + if isinstance(baseline_out, (list, tuple)): + if (not isinstance(reference_out, (list, tuple)) + or len(baseline_out) != len(reference_out)): + print(f" [FAIL] Output structures differ for {problem_name}") + return False + + all_match = True + for i, (b, r) in enumerate(zip(baseline_out, reference_out)): + arr_b = np.array(b) + arr_r = np.array(r) + if not np.array_equal(arr_b, arr_r, equal_nan=True): + try: + diff = np.max(np.abs(arr_b - arr_r)) + except Exception: + diff = "N/A" + print( + f" [FAIL] Element {i} does not match exactly for " + f"{problem_name} (Max diff: {diff})" + ) + all_match = False + + if all_match: + print(f" [PASS] {problem_name}") + return all_match + else: + arr_b = np.array(baseline_out) + arr_r = np.array(reference_out) + if not np.array_equal(arr_b, arr_r, equal_nan=True): + try: + diff = np.max(np.abs(arr_b - arr_r)) + except Exception: + diff = "N/A" + print( + f" [FAIL] Outputs do not match exactly for " + f"{problem_name} (Max diff: {diff})" + ) + return False + + print(f" [PASS] {problem_name}") + return True + +def main(): + parser = argparse.ArgumentParser(description="Verify outputs") + parser.add_argument( + "--level", "-l", type=str, default="level2", + help="Level to test (e.g. level1 or level2)" + ) + parser.add_argument( + "--problem", "-p", type=str, default=None, + help="Name of a specific problem to verify (e.g. 1p_Flash_Attention)" + ) + args_parsed = parser.parse_args() + + script_dir = os.path.dirname(os.path.abspath(__file__)) + jaxbench_dir = os.path.join( + script_dir, f"../../../JAXBench/benchmark/{args_parsed.level}" + ) + adapted_dir = os.path.join(script_dir, f"./{args_parsed.level}") + + if not os.path.exists(adapted_dir): + print(f"Directory not found: {adapted_dir}") + return + + problems = [ + d for d in os.listdir(adapted_dir) + if os.path.isdir(os.path.join(adapted_dir, d)) + ] + problems.sort() + + if args_parsed.problem: + if args_parsed.problem not in problems: + print(f"Problem '{args_parsed.problem}' not found in {adapted_dir}") + return + problems = [args_parsed.problem] + + for problem in problems: + print(f"Verifying {problem}...") + try: + # Load JaxBench baseline + baseline_path = os.path.join(jaxbench_dir, problem, "baseline.py") + if not os.path.exists(baseline_path): + print( + f" [SKIP] No baseline.py found for {problem} in JaxBench" + ) + continue + baseline_mod = load_module(f"{problem}_baseline", baseline_path) + + # Load Adapted reference + reference_path = os.path.join(adapted_dir, problem, "reference.py") + if not os.path.exists(reference_path): + print( + f" [SKIP] No reference.py found for {problem} in " + f"adapted dataset" + ) + continue + reference_mod = load_module(f"{problem}_reference", reference_path) + + # Collect JaxBench inputs (one entry per configuration) + if hasattr(baseline_mod, "create_inputs"): + baseline_inputs = baseline_mod.create_inputs() + else: + print( + f" [ERROR] No create_inputs found in " + f"baseline for {problem}" + ) + continue + + if not isinstance(baseline_inputs, (list, tuple)): + baseline_inputs = (baseline_inputs,) + if baseline_is_single(baseline_inputs): + baseline_inputs = [baseline_inputs] + + # 2. Collect adapted inputs (one entry per configuration) + get_inputs_ret = reference_mod.get_inputs() + if reference_is_single(get_inputs_ret): + get_inputs_ret = [get_inputs_ret] + + if len(baseline_inputs) != len(get_inputs_ret): + print( + f" [ERROR] {len(baseline_inputs)} baseline " + f"configuration(s) but {len(get_inputs_ret)} in the " + f"adapted dataset for {problem}" + ) + continue + + multi_config = len(get_inputs_ret) > 1 + for cfg_idx, (cfg_inputs, cfg_ret) in enumerate( + zip(baseline_inputs, get_inputs_ret) + ): + label = f"{problem}[{cfg_idx}]" if multi_config else problem + try: + if len(cfg_ret) != 2: + print( + f" [ERROR] get_inputs() returned {len(cfg_ret)} " + f"elements instead of 2 for {label}" + ) + continue + dynamic_args, static_args = cfg_ret + + # Run JaxBench version + baseline_out = baseline_mod.workload(*cfg_inputs) + jax.block_until_ready(baseline_out) + + # Run adapted version + args = list(dynamic_args) + if static_args: + args += list(static_args) + reference_out = reference_mod.computation(*args) + jax.block_until_ready(reference_out) + + # Compare outputs + check_match(baseline_out, reference_out, label) + del baseline_out, reference_out + except Exception as e: + print(f" [ERROR] {label}: {e}") + + except Exception as e: + print(f" [ERROR] {problem}: {e}") + +if __name__ == '__main__': + main()