diff --git a/docs/guides/optimization/sharding.md b/docs/guides/optimization/sharding.md index 1dfb685bdf..97abf70070 100644 --- a/docs/guides/optimization/sharding.md +++ b/docs/guides/optimization/sharding.md @@ -256,7 +256,9 @@ Context parallelism is similar to FSDP except we shard the sequence dimension of Care needs to be taken to shard the sequence dimension for attention - only the queries are sharded by sequence, the keys and values need to be all-gathered to perform the full computation. Additionally if we naively shard the sequence dimension then the attention computation is not evenly distributed due to the lower triangular causal mask - shards corresponding to later queries have more non-zero mask and thus become the bottleneck. Instead we “stripe” the inputs, so that the first shard has the first and last chunk of the sequence, the second shard has the second and second to last, etc. This striping is done on the initial data inputs (instead of every layer), so it is a small cost. -Note in general there are many flavors of CP such as ring attention, which in theory can hide all of the comms (as opposed to this implementation where the KV all gathers are probably exposed). This all gather is relatively cheap so we have implemented this flavor for now, a good trade-off of complexity and performance. Currently TPUs only support this all gather strategy `context_parallel_strategy=all_gather`, but GPUs support both an `all_gather` strategy or a `ring` strategy which will perform the computation and communication in chunks and ideally overlap in a collective matmul fashion. This strategy requires extending the online softmax trick from only within chip to additionally apply it across chips. +Note in general there are many flavors of CP such as ring attention, which in theory can hide all of the comms (as opposed to all gather CP where the KV all gathers are probably exposed). This all gather is relatively cheap so it remains a good trade-off of complexity and performance. + +MaxText supports `context_parallel_strategy=all_gather`, and supports `context_parallel_strategy=ring` through GPU Transformer Engine and TPU Tokamax Splash paths; ring performs the computation and communication in chunks and ideally overlaps them in a collective matmul fashion. This strategy requires extending the online softmax trick from only within chip to additionally apply it across chips. ### CP Arithmetic Intensity diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 3dffe5e7ac..2091bf0d8b 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -3261,15 +3261,60 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de context_parallel_size = getattr(self, f"ici_{self.context_sharding}_parallelism", 1) * getattr( self, f"dcn_{self.context_sharding}_parallelism", 1 ) - if context_parallel_size > 1 and self.context_parallel_strategy.lower() == "ring": - if "gpu" not in self.hardware: + context_parallel_strategy = self.context_parallel_strategy.lower() + if ( + context_parallel_strategy == "ring" + and "gpu" not in self.hardware + and "tpu" not in self.hardware + and context_parallel_size > 1 + ): + raise ValueError( + "Ring context parallelism strategy (context_parallel_strategy='ring') is only supported on GPUs " + "or TPU with attention=flash and use_tokamax_splash=True." + ) + if context_parallel_strategy == "ring" and "gpu" not in self.hardware and "tpu" in self.hardware: + if context_parallel_size <= 1: + raise ValueError("TPU Tokamax ring attention requires context_parallel_size > 1.") + if self.context_sharding != "context": + raise ValueError("TPU Tokamax ring attention requires context_sharding='context'.") + if self.dq_reduction_steps not in (0, 3): + raise ValueError("TPU Tokamax ring attention requires dq_reduction_steps to be 0 or 3.") + if self.max_target_length % (context_parallel_size * context_parallel_size) != 0: raise ValueError( - "Ring context parallelism strategy (context_parallel_strategy='ring') is only supported on GPUs." + "TPU Tokamax ring attention requires max_target_length to be divisible by context_parallel_size squared." ) + if self.attention != "flash": + raise ValueError("TPU ring context parallelism requires attention=flash.") + if not self.use_tokamax_splash: + raise ValueError("TPU ring context parallelism requires use_tokamax_splash=True.") + if self.use_jax_splash: + raise ValueError("TPU ring context parallelism requires use_jax_splash=False.") + if self.attention_type != "global": + raise ValueError("TPU Tokamax ring attention is initially supported only for global causal attention.") + if self.packing: + raise ValueError("TPU Tokamax ring attention does not support packing yet.") + if self.context_parallel_load_balance: + raise ValueError("TPU Tokamax ring attention does not support context_parallel_load_balance yet.") + if self.use_ragged_attention: + raise ValueError("TPU Tokamax ring attention does not support ragged attention.") + if self.attention_sink: + raise ValueError("TPU Tokamax ring attention does not support attention sinks.") + if self.use_indexer: + raise ValueError("TPU Tokamax ring attention does not support sparse indexer masks.") + if self.use_chunked_prefill: + raise ValueError("TPU Tokamax ring attention does not support chunked prefill yet.") + if self.moba: + raise ValueError("TPU Tokamax ring attention does not support MoBA.") + if self.use_multimodal: + raise ValueError("TPU Tokamax ring attention does not support multimodal attention.") + if self.use_qk_clip: + raise ValueError("TPU Tokamax ring attention does not support QK-Clip statistics yet.") + if self.enable_dropout and self.dropout_rate > 0.0: + raise ValueError("TPU Tokamax ring attention does not support dropout yet.") # STRIPED reorder strategy is a Transformer Engine feature and is GPU-only. - # The AUTO + packing case (which training resolves to STRIPED) is not validated here - # because test code paths may load the same config but use a different reorder path. - # Training's runtime path in max_utils.reorder_causal_load_balanced enforces this. + # The AUTO + packing case, which training resolves to STRIPED, is not + # validated here because test code paths may load the same config but use a + # different reorder path. Training's runtime path enforces this. if ( context_parallel_size > 1 and "gpu" not in self.hardware diff --git a/src/maxtext/kernels/attention/tokamax_ring_attention.py b/src/maxtext/kernels/attention/tokamax_ring_attention.py new file mode 100644 index 0000000000..e042e77633 --- /dev/null +++ b/src/maxtext/kernels/attention/tokamax_ring_attention.py @@ -0,0 +1,347 @@ +# Copyright 2026 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 +# +# 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. +"""MaxText integration for Tokamax ring attention. + +Tokamax 0.0.12 source: +https://github.com/openxla/tokamax/tree/4936e75/tokamax/_src/ops/experimental/tpu/splash_attention +""" + +from __future__ import annotations + +import dataclasses +import functools +from typing import Any + +import jax +from jax.experimental import pallas as pl + +from maxtext.common.common_types import MODEL_MODE_TRAIN +from maxtext.kernels.tokamax_splash_attention import ring_attention_kernel +from maxtext.kernels.tokamax_splash_attention import splash_attention_kernel as tokamax_splash_kernel +from maxtext.kernels.tokamax_splash_attention import splash_attention_mask as tokamax_splash_mask + + +def is_context_parallel_ring_requested(config: Any) -> bool: + """Returns True when the config requests ring context parallelism.""" + return config.context_parallel_strategy.lower() == "ring" + + +def _mesh_axes_for_dim(axis_names: Any) -> tuple[Any, ...]: + if axis_names is None: + return () + if isinstance(axis_names, str): + return (axis_names,) + return tuple(axis for axis in axis_names if axis is not None) + + +def _mesh_axes_size(mesh: Any, axes: tuple[Any, ...]) -> int: + size = 1 + for axis in axes: + if axis not in mesh.shape: + raise ValueError(f"TPU Tokamax ring attention requires mesh axis {axis!r} to exist.") + size *= mesh.shape[axis] + return size + + +def with_sequence_axis(axis_names: Any, ring_axis: str, sequence_dim: int) -> Any: + """Returns axis names with the sequence dimension set to the ring axis.""" + if axis_names is None: + return None + if len(axis_names) <= sequence_dim: + raise ValueError("TPU Tokamax ring attention expects a sequence sharding dimension.") + axes = list(axis_names) + existing_sequence_axes = _mesh_axes_for_dim(axes[sequence_dim]) + if existing_sequence_axes and existing_sequence_axes != (ring_axis,): + raise ValueError( + "TPU Tokamax ring attention expects the existing sequence sharding to be " + f"unsharded or exactly {(ring_axis,)}, got {existing_sequence_axes}." + ) + axes[sequence_dim] = ring_axis + if isinstance(axis_names, jax.sharding.PartitionSpec): + return jax.sharding.PartitionSpec(*axes, unreduced=axis_names.unreduced, reduced=axis_names.reduced) + if isinstance(axis_names, tuple): + return tuple(axes) + return axes + + +def _validate_ring_axis_only_on_sequence( + axis_names: Any, + *, + tensor_name: str, + sequence_dim: int, + ring_axis: str, +) -> None: + """Raises if the ring mesh axis appears outside the sequence dimension.""" + for dim, axis_name in enumerate(axis_names): + if dim == sequence_dim: + continue + dim_axes = _mesh_axes_for_dim(axis_name) + if ring_axis in dim_axes: + raise ValueError( + "TPU Tokamax ring attention requires the context axis to appear only " + f"on the sequence dimension; got {ring_axis!r} on {tensor_name} dim {dim}." + ) + + +def validate_dkv_sharding( + *, + axis_names_q: Any, + axis_names_kv: Any, + dkv_dim_q: int, + dkv_dim_kv: int, +) -> None: + """Validates that the head-dim/D_KV dimension stays local for ring attention.""" + q_dkv_axes = _mesh_axes_for_dim(axis_names_q[dkv_dim_q]) + kv_dkv_axes = _mesh_axes_for_dim(axis_names_kv[dkv_dim_kv]) + if q_dkv_axes or kv_dkv_axes: + raise ValueError( + "TPU Tokamax ring attention does not support sharding the D_KV/head-dim " + f"dimension; got Q axes {q_dkv_axes} and K/V axes {kv_dkv_axes}." + ) + + +def validate_tokamax_ring_runtime( + *, + model_mode: str, + use_ragged_attention: bool = False, + previous_chunk: Any = None, + sinks: Any = None, + indexer_mask: Any = None, + bidirectional_mask: Any = None, + record_max_logits: bool = False, +) -> None: + """Validates runtime-only constraints for the MaxText ring path.""" + if model_mode != MODEL_MODE_TRAIN: + raise ValueError("TPU Tokamax ring attention is supported only for train mode.") + if use_ragged_attention: + raise ValueError("TPU Tokamax ring attention does not support ragged attention.") + if previous_chunk is not None: + raise ValueError("TPU Tokamax ring attention does not support chunked prefill yet.") + if sinks is not None: + raise ValueError("TPU Tokamax ring attention does not support attention sinks.") + if indexer_mask is not None: + raise ValueError("TPU Tokamax ring attention does not support indexer masks.") + if bidirectional_mask is not None: + raise ValueError("TPU Tokamax ring attention does not support bidirectional masks.") + if record_max_logits: + raise NotImplementedError("TPU Tokamax ring attention does not support record_max_logits yet.") + + +def validate_ring_mesh_axis( + *, + axis_names_q: Any, + axis_names_kv: Any, + sequence_dim_q: int, + sequence_dim_kv: int, + mesh: Any, + ring_axis: str, +) -> None: + """Validates sequence sharding before ring attention.""" + if not ring_axis: + raise ValueError("TPU Tokamax ring attention requires a non-empty context_sharding axis.") + if ring_axis not in mesh.shape: + raise ValueError(f"TPU Tokamax ring attention requires mesh axis {ring_axis!r} to exist.") + + _validate_ring_axis_only_on_sequence( + axis_names_q, + tensor_name="Q", + sequence_dim=sequence_dim_q, + ring_axis=ring_axis, + ) + _validate_ring_axis_only_on_sequence( + axis_names_kv, + tensor_name="K/V", + sequence_dim=sequence_dim_kv, + ring_axis=ring_axis, + ) + expected_axes = (ring_axis,) + q_sequence_axes = _mesh_axes_for_dim(axis_names_q[sequence_dim_q]) + key_value_sequence_axes = _mesh_axes_for_dim(axis_names_kv[sequence_dim_kv]) + if q_sequence_axes != expected_axes: + raise ValueError( + "TPU Tokamax ring attention requires Q sequence sharding to be exactly " + f"{expected_axes}, got {q_sequence_axes}." + ) + if key_value_sequence_axes != expected_axes: + raise ValueError( + "TPU Tokamax ring attention requires K/V sequence sharding to be exactly " + f"{expected_axes}, got {key_value_sequence_axes}." + ) + + +def validate_head_sharding( + *, + axis_names_q: Any, + axis_names_kv: Any, + mesh: Any, + num_query_heads: int, + num_kv_heads: int, + head_dim_q: int, + head_dim_kv: int, +) -> None: + """Validates that local head layout preserves GQA/MQA head mapping.""" + q_head_axes = _mesh_axes_for_dim(axis_names_q[head_dim_q]) + kv_head_axes = _mesh_axes_for_dim(axis_names_kv[head_dim_kv]) + q_head_shards = _mesh_axes_size(mesh, q_head_axes) + kv_head_shards = _mesh_axes_size(mesh, kv_head_axes) + if num_query_heads % q_head_shards != 0: + raise ValueError( + "TPU Tokamax ring attention requires num_query_heads " + f"({num_query_heads}) to be divisible by Q head shards ({q_head_shards})." + ) + if num_kv_heads % kv_head_shards != 0: + raise ValueError( + "TPU Tokamax ring attention requires num_kv_heads " + f"({num_kv_heads}) to be divisible by KV head shards ({kv_head_shards})." + ) + + if num_kv_heads == 1: + if kv_head_axes: + raise ValueError("TPU Tokamax ring attention does not support sharding the single MQA KV head.") + return + + if q_head_axes != kv_head_axes: + raise ValueError( + "TPU Tokamax ring attention requires Q and KV head sharding to match for MHA/GQA, " + f"got Q head axes {q_head_axes} and KV head axes {kv_head_axes}." + ) + local_query_heads = num_query_heads // q_head_shards + local_kv_heads = num_kv_heads // kv_head_shards + if local_query_heads % local_kv_heads != 0: + raise ValueError( + "TPU Tokamax ring attention requires local query heads " + f"({local_query_heads}) to be divisible by local KV heads ({local_kv_heads})." + ) + + +def build_splash_config( + config: Any, + *, + q_seq_len: int, + kv_seq_len: int, + context_parallel_size: int, + attn_logits_soft_cap: float | None = None, +) -> Any: + """Converts MaxText Splash config fields into Tokamax `SplashConfig`.""" + if context_parallel_size <= 1: + raise ValueError("context_parallel_size must be > 1 for ring attention.") + dq_reduction_steps = config.dq_reduction_steps + q_seq_len_per_shard = q_seq_len // context_parallel_size + kv_seq_len_per_shard = kv_seq_len // context_parallel_size + block_q = min(config.sa_block_q, q_seq_len_per_shard) + block_kv = min(config.sa_block_kv, kv_seq_len_per_shard) + block_kv_compute = min(config.sa_block_kv_compute, kv_seq_len_per_shard) + block_q_dkv = min(config.sa_block_q_dkv, q_seq_len_per_shard) + block_kv_dkv = min(config.sa_block_kv_dkv, kv_seq_len_per_shard) + block_kv_dkv_compute = min(config.sa_block_kv_dkv_compute, kv_seq_len_per_shard) + return tokamax_splash_kernel.SplashConfig( + block_q=block_q, + block_kv=block_kv, + block_kv_compute=block_kv_compute, + block_q_dkv=block_q_dkv, + block_kv_dkv=block_kv_dkv, + block_kv_dkv_compute=block_kv_dkv_compute, + use_fused_bwd_kernel=True, + q_layout=tokamax_splash_kernel.QKVLayout[config.sa_q_layout], + k_layout=tokamax_splash_kernel.QKVLayout[config.sa_k_layout], + v_layout=tokamax_splash_kernel.QKVLayout[config.sa_v_layout], + attn_logits_soft_cap=attn_logits_soft_cap, + residual_checkpoint_name="context", + use_base2_exp=False, + fwd_cost_estimate=pl.CostEstimate(flops=config.cost_estimate_flops_fwd, transcendentals=0, bytes_accessed=0) + if config.cost_estimate_flops_fwd >= 0 + else None, + bwd_cost_estimate=pl.CostEstimate(flops=config.cost_estimate_flops_bwd, transcendentals=0, bytes_accessed=0) + if config.cost_estimate_flops_bwd >= 0 + else None, + dq_reduction_steps=dq_reduction_steps if dq_reduction_steps > 0 else None, + use_experimental_scheduler=config.use_splash_scheduler, + ) + + +def _make_causal_mask(shape: tuple[int, int], context_parallel_size: int): + """Builds a lazy causal mask for ring attention.""" + if context_parallel_size <= 1: + raise ValueError("context_parallel_size must be > 1 for ring attention.") + return tokamax_splash_mask.CausalMask(shape=shape, shard_count=context_parallel_size) + + +def make_sharded_ring_attention_kernel( + config: Any, + *, + query: Any, + key: Any, + context_parallel_size: int, + ring_axis: str, + attn_logits_soft_cap: float | None, + maybe_shard_with_pspec: Any, +): + """Builds and shards the Tokamax ring attention kernel for MaxText.""" + splash_config = build_splash_config( + config, + q_seq_len=query.shape[2], + kv_seq_len=key.shape[2], + context_parallel_size=context_parallel_size, + attn_logits_soft_cap=attn_logits_soft_cap, + ) + if config.use_max_logit_estimate > 0: + splash_config = dataclasses.replace(splash_config, max_logit_const=config.use_max_logit_estimate) + + mask = _make_causal_mask( + (query.shape[2], key.shape[2]), + context_parallel_size, + ) + + @functools.partial(jax.jit, static_argnames=["single_head_mask"]) + def wrap_ring_kernel(single_head_mask): + return ring_attention_kernel.make_ring_attention( + single_head_mask, + config=splash_config, + is_mqa=False, + save_residuals=False, + ring_axis=ring_axis, + q_seq_shards=context_parallel_size, + kv_seq_shards=context_parallel_size, + ) + + ring_kernel = wrap_ring_kernel(mask) + ring_kernel_spec = ring_kernel.manual_sharding_spec() + ring_kernel = jax.tree.map( + lambda arr, spec: None if arr is None else maybe_shard_with_pspec(arr, spec), + ring_kernel, + ring_kernel_spec, + is_leaf=lambda x: x is None, + ) + return splash_config, ring_kernel, ring_kernel_spec + + +def call_ring_attention( + query: Any, + key: Any, + value: Any, + decoder_segment_ids_q: Any, + decoder_segment_ids_kv: Any, + ring_kernel: Any, +): + """Calls a Tokamax ring attention kernel over the MaxText batch dimension.""" + if (decoder_segment_ids_q is None) != (decoder_segment_ids_kv is None): + raise ValueError("decoder_segment_ids_q and decoder_segment_ids_kv must both be set or both be None.") + if decoder_segment_ids_q is None: + return jax.vmap(lambda q, k, v: ring_kernel(q, k, v, None), in_axes=(0, 0, 0))(query, key, value) + + def call_one(q, k, v, q_segment_ids, kv_segment_ids): + segment_ids = ring_attention_kernel.SegmentIds(q_segment_ids, kv_segment_ids) + return ring_kernel(q, k, v, segment_ids) + + return jax.vmap(call_one, in_axes=(0, 0, 0, 0, 0))(query, key, value, decoder_segment_ids_q, decoder_segment_ids_kv) diff --git a/src/maxtext/kernels/tokamax_splash_attention/__init__.py b/src/maxtext/kernels/tokamax_splash_attention/__init__.py new file mode 100644 index 0000000000..537d4f943b --- /dev/null +++ b/src/maxtext/kernels/tokamax_splash_attention/__init__.py @@ -0,0 +1,18 @@ +# Copyright 2026 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 +# +# 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. +"""Tokamax SplashAttention runtime. + +Adapted from OpenXLA Tokamax 0.0.12: +https://github.com/openxla/tokamax/tree/4936e75/tokamax/_src/ops/experimental/tpu/splash_attention +""" diff --git a/src/maxtext/kernels/tokamax_splash_attention/base.py b/src/maxtext/kernels/tokamax_splash_attention/base.py new file mode 100644 index 0000000000..df793fcec9 --- /dev/null +++ b/src/maxtext/kernels/tokamax_splash_attention/base.py @@ -0,0 +1,270 @@ +# pylint: skip-file +# Copyright 2025 DeepMind Technologies Limited. All Rights Reserved. +# +# 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. +# ============================================================================== +"""Base functionality for Sparse Flash Attention.""" + +import functools +from typing import Final, NamedTuple, TypeAlias +import jax +import jax.numpy as jnp +import numpy as np +from maxtext.kernels.tokamax_splash_attention import splash_attention_mask_info as mask_info_lib + + +MaskInfo = mask_info_lib.MaskInfo + + +DEFAULT_MASK_VALUE: Final[float] = -0.7 * float(np.finfo(np.dtype("float32")).max) + + +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. For causal self-attention, segment ids form a + block diagonal matrix so at least one element in each row is set. + Non-self-attention configurations can have all-masked rows; Splash treats + those rows as zero-output rows with mask-value logsumexp. + Attributes: + q: segment ids along the Q sequence + kv: segment ids along the KV sequence + """ + + q: jax.Array | jax.sharding.PartitionSpec # [q_seq_len] + kv: jax.Array | jax.sharding.PartitionSpec # [kv_seq_len] + + +# Return type of SplashAttention function that implements the custom vjp rule. +SplashCustomReturnType: TypeAlias = jax.Array | tuple[jax.Array, dict[str, jax.Array]] + +SplashResidualsType = tuple[ + jax.Array, # q + jax.Array, # k + jax.Array, # v + SegmentIds | None, # segment_ids + jax.Array | None, # sinks + jax.Array, # out + jax.Array, # logsumexp + MaskInfo | None, # dkv_mask_info +] + + +def _attention_reference_impl( + q: jax.Array, + k: jax.Array, + v: jax.Array, + mask: jax.Array, + segment_ids: SegmentIds | None, + sinks: jax.Array | None, + mask_value: float, + save_residuals: bool, + attn_logits_soft_cap: float | None, +) -> SplashCustomReturnType: + 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 + + if sinks is not None: + assert sinks.shape == () # should already be vmapped + + logits = jnp.where(mask, logits, mask_value) + m = logits.max(axis=-1) + sinks = None if sinks is None else sinks.astype(logits.dtype) + m = m if sinks is None else jnp.maximum(m, sinks) + s = jnp.exp(logits - m[..., None]) + s = jnp.where(mask, s, 0.0) + l = s.sum(axis=-1) + (0 if sinks is None else jnp.exp(sinks - m)) + safe_l = jnp.where(l == 0.0, 1.0, l) + p = s / safe_l[..., None] + + o = jnp.einsum("st,td->sd", p, v.astype(jnp.float32)) + + if save_residuals: + logsumexp = jnp.where(l == 0.0, mask_value, m + jnp.log(safe_l)) + return o, {"logsumexp": logsumexp, "max_logits": m} + return o + + +def _attention_reference_custom_bwd( + do, + q, + k, + v, + mask, + segment_ids, + sinks, + o, + logsumexp, + mask_value: float = DEFAULT_MASK_VALUE, + backward_impl: str = "vanilla", + attn_logits_soft_cap: float | None = None, +) -> tuple[jax.Array, jax.Array, jax.Array, None, None, jax.Array | None]: + 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]) + p = jnp.where(mask, p, 0.0) + 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 backward_impl == "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) + dsinks = None + if sinks is not None: + sinks_exp = -jnp.exp(sinks[..., None, None].astype(jnp.float32) - logsumexp[..., None].astype(jnp.float32)) + dsinks = jnp.sum(sinks_exp.astype(o.dtype) * o * do, axis=(-1, -2)) + return dq, dk, dv, None, None, dsinks + + +@functools.partial( + jax.jit, + static_argnames=[ + "mask_value", + "save_residuals", + "attn_logits_soft_cap", + "is_mqa", + ], +) +def attention_reference( + q: jax.Array, + k: jax.Array, + v: jax.Array, + mask: jax.Array, + segment_ids: SegmentIds | None = None, + sinks: jax.Array | None = None, + *, + is_mqa: bool, + mask_value: float = DEFAULT_MASK_VALUE, + save_residuals: bool = False, + attn_logits_soft_cap: float | None = None, +): + """A JIT-compiled reference implementation of attention, handles MQA and MHA.""" + attn_impl = functools.partial( + _attention_reference_impl, + mask_value=mask_value, + save_residuals=save_residuals, + attn_logits_soft_cap=attn_logits_soft_cap, + ) + + if is_mqa: + func = jax.vmap(attn_impl, in_axes=(0, None, None, None, None, 0)) + 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 + + kv_heads, q_heads = k.shape[0], q.shape[0] + assert q_heads % kv_heads == 0 + + if kv_heads < q_heads: + # Repeat K and V heads to match the number of Q heads. + q_heads_per_kv = q_heads // kv_heads + k = jnp.repeat(k, repeats=q_heads_per_kv, axis=0) + v = jnp.repeat(v, repeats=q_heads_per_kv, axis=0) + + func = jax.vmap(attn_impl, in_axes=(0, 0, 0, None, None, 0)) + + out = func(q, k, v, mask, segment_ids, sinks) + return out + + +@functools.partial(jax.jit, static_argnames=["is_mqa", "backward_impl", "attn_logits_soft_cap"]) +def attention_reference_vjp( + do, + q, + k, + v, + mask, + segment_ids, + sinks, + o, + logsumexp, + *, + is_mqa: bool, + backward_impl: str = "vanilla", + attn_logits_soft_cap: float | None = None, +): + """Wrapper for backward reference that handles GQA/MQA broadcasting and reduction.""" + bwd = functools.partial( + _attention_reference_custom_bwd, + backward_impl=backward_impl, + attn_logits_soft_cap=attn_logits_soft_cap, + ) + + num_q_heads = q.shape[0] + num_kv_heads = 1 if is_mqa else k.shape[0] + + is_grouped = not is_mqa and num_kv_heads < num_q_heads + assert num_q_heads % num_kv_heads == 0 + head_multiplier = num_q_heads // num_kv_heads + if is_mqa: + bwd = jax.vmap(bwd, in_axes=(0, 0, None, None, None, None, 0, 0, 0)) + else: + bwd = jax.vmap(bwd, in_axes=(0, 0, 0, 0, None, None, 0, 0, 0)) + # Interleave the KV heads to match the corresponding Q heads. + if is_grouped: + k = jnp.repeat(k, head_multiplier, axis=0) + v = jnp.repeat(v, head_multiplier, axis=0) + + dq, dk, dv, _, _, dsinks = bwd(do, q, k, v, mask, segment_ids, sinks, o, logsumexp) + + if is_mqa: + dk, dv = dk.sum(axis=0), dv.sum(axis=0) + elif is_grouped: + # Perform the sum reduction across the head_multiplier dimension only. + # So that the output still has KV heads. + dk = dk.reshape(num_kv_heads, head_multiplier, *dk.shape[1:]) + dv = dv.reshape(num_kv_heads, head_multiplier, *dv.shape[1:]) + dk, dv = dk.sum(axis=1), dv.sum(axis=1) + + return dq, dk, dv, dsinks diff --git a/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py b/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py new file mode 100644 index 0000000000..519629650c --- /dev/null +++ b/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py @@ -0,0 +1,755 @@ +# pylint: skip-file +# Copyright 2025 DeepMind Technologies Limited. All Rights Reserved. +# +# 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. +# ============================================================================== + +"""Implementation of Ring Attention.""" + +import dataclasses +import functools +from typing import Any + +import jax +from jax import lax +from jax import tree_util +import jax.numpy as jnp +import numpy as np +from maxtext.kernels.tokamax_splash_attention import base +from maxtext.kernels.tokamax_splash_attention import ring_attention_utils +from maxtext.kernels.tokamax_splash_attention import splash_attention_kernel as splash_kernel +from maxtext.kernels.tokamax_splash_attention import splash_attention_mask as mask_lib +from maxtext.kernels.tokamax_splash_attention import splash_attention_mask_info as mask_info_lib + +P = jax.P +MaskInfo = mask_info_lib.MaskInfo +partial = functools.partial + +SegmentIds = base.SegmentIds +SplashConfig = splash_kernel.SplashConfig +SplashResidualsType = base.SplashResidualsType +SplashCustomReturnType = base.SplashCustomReturnType +MaskFunctionType = splash_kernel.MaskFunctionType +_splash_attention_forward = splash_kernel._splash_attention_forward # pylint: disable=protected-access +_splash_attention_bwd = splash_kernel._splash_attention_bwd # pylint: disable=protected-access + +_dynamic_slice_mask_info = ring_attention_utils.dynamic_slice_mask_info +_offset_q_sequence_for_kv_shard = ring_attention_utils.offset_q_sequence_for_kv_shard +_has_no_active_blocks = ring_attention_utils.has_no_active_blocks +_has_empty_attention_rows = ring_attention_utils.has_empty_attention_rows +_mask_sparsity = ring_attention_utils.mask_sparsity +_has_axis = ring_attention_utils.has_axis + + +def _validate_ring_axis_size(ring_axis: str, ring_axis_size: int, expected_ring_size: int) -> None: + if ring_axis_size != expected_ring_size: + raise ValueError( + f"Ring axis {ring_axis} has size {ring_axis_size}, but ring attention " + f"was built for {expected_ring_size} sequence shards." + ) + + +def _ring_attention_forward( + fwd_mask_info: MaskInfo, + q: jax.Array, + k: jax.Array, + v: jax.Array, + segment_ids: SegmentIds | None, + mask_value: float, + is_mqa: bool, + config: SplashConfig | None, + mask_function: MaskFunctionType | None, + fwd_mask_sparsity: float, + *, + sinks: jax.Array | None = None, + ring_axis: str, + expected_ring_size: int, +) -> tuple[jax.Array, tuple[jax.Array, jax.Array]]: + + if q.shape[-1] != k.shape[-1]: + raise NotImplementedError("Queries and keys must have the same head dimension.") + + if sinks is not None: + raise NotImplementedError("Sinks aren't supported yet.") + + ring_axis_size = lax.axis_size(ring_axis) + _validate_ring_axis_size(ring_axis, ring_axis_size, expected_ring_size) + ring_axis_idx = lax.axis_index(ring_axis) + + shift = partial( + lax.ppermute, + axis_name=ring_axis, + perm=[(i, (i + 1) % ring_axis_size) for i in range(ring_axis_size)], + ) + # for example, if ring size is 4 + # Device 3 => permute_idx 0, offset (3-0) % 4 = 3, + # permute_idx 1, offset (3-1) % 4 = 2, etc. + # Device 2 => permute_idx 0, offset (2-0) % 4 = 2, + # permute_idx 1, offset (2-1) % 4 = 1, etc. + # Device 1 => permute_idx 0, offset (1-0) % 4 = 1, + # permute_idx 1, offset (1-1) % 4 = 0, etc. + # Device 0 => permute_idx 0, offset (0-0) % 4 = 0, + # permute_idx 1, offset (0-1) % 4 = 3, etc. + + splash_fwd_partial = partial( + _splash_attention_forward, + save_residuals=True, + mask_value=mask_value, + is_mqa=is_mqa, + config=config, + mask_function=mask_function, + fwd_mask_sparsity=fwd_mask_sparsity, + max_logit_value=None, + ) + # Initial accumulator values + o_shape = q.shape[:-1] + (v.shape[-1],) + o_init = jnp.zeros(o_shape, dtype=jnp.float32) + l_init = jnp.zeros((o_shape[0], o_shape[1]), jnp.float32) + m_init = jnp.full_like(l_init, mask_value, dtype=jnp.float32) + + def body(carry, i: int): + m_prev, l_prev, o_prev, k_current, v_current, segment_ids_current = carry + + current_kv_shard_idx = (ring_axis_idx - i) % ring_axis_size + local_fwd_mask_info = _dynamic_slice_mask_info(fwd_mask_info, current_kv_shard_idx, ring_axis_size) + local_fwd_mask_info = _offset_q_sequence_for_kv_shard(local_fwd_mask_info, current_kv_shard_idx, k_current.shape[-2]) + k_next = shift(k_current) + v_next = shift(v_current) + + if segment_ids is not None: + kv_segment_ids_next = shift(segment_ids_current.kv) + segment_ids_next = SegmentIds(segment_ids.q, kv_segment_ids_next) + else: + segment_ids_next = None + + out_curr, stats = splash_fwd_partial( + local_fwd_mask_info, + q, + k_current, + v_current, + segment_ids=segment_ids_current, + sinks=sinks, + ) + no_active_blocks = _has_no_active_blocks(local_fwd_mask_info) + out_curr = jnp.where(no_active_blocks, jnp.zeros_like(out_curr), out_curr) + lse_curr = jnp.where( + no_active_blocks, + jnp.full_like(stats["logsumexp"], mask_value), + stats["logsumexp"], + ) + m_curr = jnp.where( + no_active_blocks, + jnp.full_like(stats["max_logits"], mask_value), + stats["max_logits"], + ) + empty_rows = _has_empty_attention_rows(lse_curr, m_curr, mask_value) + l_curr = jnp.exp(lse_curr - m_curr) + l_curr = jnp.where(empty_rows, jnp.zeros_like(l_curr), l_curr) + o_curr = out_curr.astype(jnp.float32) * l_curr[..., None] + 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 + o_next = alpha[..., None] * o_prev + beta[..., None] * o_curr + return (m_next, l_next, o_next, k_next, v_next, segment_ids_next), None + + initial_carry = (m_init, l_init, o_init, k, v, segment_ids) + (m_final, l_final, o_final, _, _, _), _ = lax.scan( + body, + initial_carry, + xs=jnp.arange(0, ring_axis_size), + length=ring_axis_size, + unroll=False, + ) # type: ignore[arg-type] + # Final normalization + assert l_final.dtype == jnp.float32 + l_inv = jnp.where(l_final == 0.0, 0.0, 1.0 / l_final) + out = (o_final * l_inv[..., None]).astype(q.dtype) + # Final logsumexp for residuals + lse = jnp.log(l_final) + m_final + lse = jnp.where(l_final == 0.0, mask_value, lse) + + return out, (lse, m_final) + + +def _ring_attention_bwd( + mask_value: float, + is_mqa: bool, + config: SplashConfig | None, + mask_function: MaskFunctionType | None, + fwd_mask_sparsity: float, + dkv_mask_sparsity: float, + save_residuals: bool, + ring_axis: str, + expected_ring_size: int, + # Residuals and gradients + res: Any, + do: jax.Array, +): + del save_residuals + (q, k, v, segment_ids, sinks, out, logsumexp, dkv_mask_info) = res + do = do.astype(jnp.float32) + if dkv_mask_info is None: + raise ValueError("Need to specify backward blocks.") + + ring_axis_size = lax.axis_size(ring_axis) + _validate_ring_axis_size(ring_axis, ring_axis_size, expected_ring_size) + ring_axis_idx = lax.axis_index(ring_axis) + + shift = partial( + lax.ppermute, + axis_name=ring_axis, + perm=[(i, (i + 1) % ring_axis_size) for i in range(ring_axis_size)], + ) + dq_accum = jnp.zeros(q.shape, dtype=jnp.float32) + dk_accum = jnp.zeros(k.shape, dtype=jnp.float32) + dv_accum = jnp.zeros(v.shape, dtype=jnp.float32) + dsinks = sinks + + def body(carry, i: int): + ( + dq_accum, + dk_accum, + dv_accum, + k_current, + v_current, + segment_ids_current, + _, + ) = carry + k_next = shift(k_current) + v_next = shift(v_current) + + current_kv_shard_idx = (ring_axis_idx - i) % ring_axis_size + local_dkv_mask_info = _dynamic_slice_mask_info(dkv_mask_info, current_kv_shard_idx, ring_axis_size) + local_dkv_mask_info = _offset_q_sequence_for_kv_shard(local_dkv_mask_info, current_kv_shard_idx, k_current.shape[-2]) + if segment_ids is not None: + kv_segment_ids_next = shift(segment_ids_current.kv) + segment_ids_next = SegmentIds(segment_ids.q, kv_segment_ids_next) + else: + segment_ids_next = None + + residuals_for_chunk = ( + q, + k_current, + v_current, + segment_ids_current, + sinks, + out, + logsumexp, + local_dkv_mask_info, + ) + + attn_bwd = functools.partial( + _splash_attention_bwd, + save_residuals=False, + mask_value=mask_value, + is_mqa=is_mqa, + config=config, + mask_function=mask_function, + fwd_mask_sparsity=fwd_mask_sparsity, + dkv_mask_sparsity=dkv_mask_sparsity, + return_fp32_grads=True, + ) + _, _, dq_i, dk_i, dv_i, _, dsinks, _ = attn_bwd(res=residuals_for_chunk, grads=do) + dv_next = shift(dv_accum + dv_i.astype(jnp.float32)) + dk_next = shift(dk_accum + dk_i.astype(jnp.float32)) + dq_accum = dq_accum + dq_i.astype(jnp.float32) + + return ( + dq_accum, + dk_next, + dv_next, + k_next, + v_next, + segment_ids_next, + dsinks, + ), None + + initial_carry = (dq_accum, dk_accum, dv_accum, k, v, segment_ids, dsinks) + (dq, dk, dv, _, _, _, dsinks), _ = lax.scan( + body, + initial_carry, + xs=jnp.arange(ring_axis_size), + length=ring_axis_size, + unroll=False, + ) + + if sinks is not None: + dsinks = jax.lax.psum(dsinks, axis_name=ring_axis) + + return ( + None, # fwd_mask_info + None, # dkv_mask_info + dq.astype(q.dtype), + dk.astype(k.dtype), + dv.astype(v.dtype), + None, + dsinks, + ) + + +def _ring_attention_fwd( + fwd_mask_info: MaskInfo, + dkv_mask_info: MaskInfo | None, + q: jax.Array, + k: jax.Array, + v: jax.Array, + segment_ids: SegmentIds | None, + sinks: jax.Array | None, + # nondiff_args + mask_value: float, # 1 + is_mqa: bool, # 2 + config: SplashConfig | None, # 3 + mask_function: MaskFunctionType | None, # 4 + fwd_mask_sparsity: float, # 5 + dkv_mask_sparsity: float, # 6 + save_residuals: bool, # 7 + ring_axis: str, # 8 + expected_ring_size: int, # 9 +) -> tuple[jax.Array, SplashResidualsType]: + """Forward pass for the custom VJP of ring attention. + + This function is used by `jax.custom_vjp` to define the forward pass + of the ring attention computation, also returning residuals needed for + the backward pass. + + Args: + fwd_mask_info: Mask information for the forward pass. + dkv_mask_info: Mask information for the backward pass for dK/dV. + q: Query array. + k: Key array. + v: Value array. + segment_ids: Optional segment IDs for packed sequences. + sinks: Optional sink tokens. + mask_value: The value used for masked-out attention scores. + is_mqa: Whether Multi-Query Attention is used. + config: SplashAttention configuration. + mask_function: Optional function to apply additional masking. + fwd_mask_sparsity: Sparsity level of the forward mask. + save_residuals: Whether to save residuals for the backward pass. + ring_axis: The name of the jax axis used for the ring. + + Returns: + A tuple containing: + - The output of the ring attention computation. + - Residuals needed for the backward pass (`SplashResidualsType`). + """ + del dkv_mask_sparsity + if save_residuals: + raise NotImplementedError("Higher-order AD not supported.") + + out, (logsumexp, max_logits) = _ring_attention_forward( + fwd_mask_info, + q, + k, + v, + segment_ids, + sinks=sinks, + mask_value=mask_value, + is_mqa=is_mqa, + config=config, + mask_function=mask_function, + fwd_mask_sparsity=fwd_mask_sparsity, + ring_axis=ring_axis, + expected_ring_size=expected_ring_size, + ) + residuals = (q, k, v, segment_ids, sinks, out, logsumexp, dkv_mask_info) + return out, residuals + + +@partial( + jax.custom_vjp, + nondiff_argnames=( + "mask_value", + "is_mqa", + "config", + "mask_function", + "fwd_mask_sparsity", + "dkv_mask_sparsity", + "save_residuals", + "ring_axis", + "expected_ring_size", + ), +) +def _ring_attention_custom( + fwd_mask_info: MaskInfo, + dkv_mask_info: MaskInfo | None, + q: jax.Array, + k: jax.Array, + v: jax.Array, + segment_ids: SegmentIds | None, + sinks: jax.Array | None, + mask_value: float, + is_mqa: bool, + config: SplashConfig | None, + mask_function: MaskFunctionType | None, + fwd_mask_sparsity: float, + dkv_mask_sparsity: float, + save_residuals: bool, + ring_axis: str, + expected_ring_size: int, +) -> SplashCustomReturnType: + """Performs ring attention with a custom VJP. + + This function is a wrapper around `_ring_attention_forward` and is used + to define the custom gradient for ring attention. + + Args: + fwd_mask_info: Mask information for the forward pass. + dkv_mask_info: Mask information for the backward pass for dK/dV. + q: Query array. + k: Key array. + v: Value array. + segment_ids: Optional segment IDs for packed sequences. + sinks: Optional sink tokens. + mask_value: The value used for masked-out attention scores. + is_mqa: Whether Multi-Query Attention is used. + config: SplashAttention configuration. + mask_function: Optional function to apply additional masking. + fwd_mask_sparsity: Sparsity level of the forward mask. + save_residuals: Whether to save residuals for the backward pass. + ring_axis: The name of the jax axis used for the ring. + + Returns: + The output of the ring attention computation. + """ + del dkv_mask_info, dkv_mask_sparsity + out, _ = _ring_attention_forward( + fwd_mask_info, + q, + k, + v, + segment_ids, + sinks=sinks, + mask_value=mask_value, + is_mqa=is_mqa, + config=config, + mask_function=mask_function, + fwd_mask_sparsity=fwd_mask_sparsity, + ring_axis=ring_axis, + expected_ring_size=expected_ring_size, + ) + return out + + +_ring_attention_custom.defvjp(_ring_attention_fwd, _ring_attention_bwd) + + +@partial( + jax.jit, + static_argnames=[ + "is_mqa", + "config", + "mask_value", + "mask_function", + "fwd_mask_sparsity", + "dkv_mask_sparsity", + "save_residuals", + "ring_axis", + "expected_ring_size", + ], +) +def _ring_attention( + fwd_mask_info: MaskInfo, + dkv_mask_info: MaskInfo | None, + q: jax.Array, + k: jax.Array, + v: jax.Array, + segment_ids: SegmentIds | None = None, + sinks: jax.Array | None = None, + *, + is_mqa: bool, + config: SplashConfig | None, + mask_value: float, + mask_function: MaskFunctionType | None, + fwd_mask_sparsity: float, + dkv_mask_sparsity: float, + save_residuals: bool = False, + ring_axis: str, + expected_ring_size: int, +) -> SplashCustomReturnType: + """Performs ring attention using SplashAttention kernels. + + This function orchestrates the ring attention mechanism by iterating through + shards of keys and values across devices along the specified `ring_axis`, + using `_splash_attention_forward` for each chunk. + + Args: + fwd_mask_info: Mask information for the forward pass. + dkv_mask_info: Mask information for the backward pass for dK/dV. + q: Query array. + k: Key array. + v: Value array. + segment_ids: Optional segment IDs for packed sequences. + sinks: Optional sink tokens. + is_mqa: Whether Multi-Query Attention is used. + config: SplashAttention configuration. + mask_value: The value used for masked-out attention scores. + mask_function: Optional function to apply additional masking. + fwd_mask_sparsity: Sparsity level of the forward mask. + save_residuals: Whether to save residuals for the backward pass. + ring_axis: The name of the jax axis used for the ring. + + Returns: + The output of the ring attention computation. + + Raises: + ValueError: If the specified `ring_axis` does not exist. + """ + if not _has_axis(ring_axis): + raise ValueError(f"Ring axis {ring_axis} does not exist") + + return _ring_attention_custom( + fwd_mask_info, + dkv_mask_info, + q, + k, + v, + segment_ids, + sinks, + is_mqa=is_mqa, + config=config, + mask_value=mask_value, + mask_function=mask_function, + fwd_mask_sparsity=fwd_mask_sparsity, + dkv_mask_sparsity=dkv_mask_sparsity, + save_residuals=save_residuals, + ring_axis=ring_axis, + expected_ring_size=expected_ring_size, + ) + + +@jax.tree_util.register_pytree_node_class +class RingSplashAttentionKernel: + """Implements Ring Attention using SplashAttention for sequence parallelism. + + This kernel computes global attention by keeping Keys and Values distributed + across the `ring_axis`. Instead of gathering full sequences, it rotates K/V + shards between devices and accumulates results incrementally. This allows + processing sequence lengths that exceed single-device memory limits. + + Attributes: + fwd_mask_info: Mask information for the forward pass. + dkv_mask_info: Mask information for the backward pass for dK/dV. + ring_axis: The name of the jax axis used for the ring. + expected_ring_size: Number of sequence shards used when building the + kernel. + kwargs: Additional keyword arguments passed to the SplashAttentionKernel + constructor. + """ + + def __init__( + self, + fwd_mask_info: MaskInfo, + dkv_mask_info: MaskInfo | None, + ring_axis: str, + expected_ring_size: int, + **kwargs, + ): + self.fwd_mask_info = fwd_mask_info + self.dkv_mask_info = dkv_mask_info + self.ring_axis = ring_axis + self.expected_ring_size = expected_ring_size + self.kwargs = kwargs + + def __call__(self, *args, **kwargs): + reserved_kwargs = set(self.kwargs) | {"ring_axis", "expected_ring_size"} + overrides = reserved_kwargs.intersection(kwargs) + if overrides: + raise ValueError( + "Ring attention call-time kwargs cannot override validated kernel " f"configuration: {sorted(overrides)}." + ) + return _ring_attention( + self.fwd_mask_info, + self.dkv_mask_info, + *args, + **dict( + self.kwargs, + **kwargs, + ring_axis=self.ring_axis, + expected_ring_size=self.expected_ring_size, + ), + ) + + def manual_sharding_spec(self): + """Ring attention expects MaskInfo metadata to be sharded by Q shard. + + Each local Q shard slices the current KV shard metadata at each ring step. + Partial mask blocks stay replicated because they are shared by all shards. + """ + + spec = jax.sharding.PartitionSpec(self.ring_axis) + _resolve_spec = lambda x: spec if x is not None else None + + def mask_info_spec(mask_info): + if mask_info is None: + return None + return MaskInfo( # pytype: disable=wrong-arg-types + mask_next=_resolve_spec(mask_info.mask_next), + active_rows=_resolve_spec(mask_info.active_rows), + active_cols=_resolve_spec(mask_info.active_cols), + num_active_blocks=_resolve_spec(mask_info.num_active_blocks), + block_mask=_resolve_spec(mask_info.block_mask), + partial_mask_blocks=jax.sharding.PartitionSpec() # replicated + if mask_info.partial_mask_blocks is not None + else None, + q_sequence=_resolve_spec(mask_info.q_sequence), + ) + + return RingSplashAttentionKernel( + mask_info_spec(self.fwd_mask_info), + mask_info_spec(self.dkv_mask_info), + ring_axis=self.ring_axis, + expected_ring_size=self.expected_ring_size, + **self.kwargs, + ) + + def tree_flatten(self): + children = (self.fwd_mask_info, self.dkv_mask_info) + aux_data = self.kwargs.copy() + aux_data["ring_axis"] = self.ring_axis + aux_data["expected_ring_size"] = self.expected_ring_size + return children, aux_data + + @classmethod + def tree_unflatten(cls, aux_data, children): + fwd_mask_info, dkv_mask_info = children + dkv_mask_info = mask_info_lib.MaskInfo(*dkv_mask_info) if dkv_mask_info is not None else None + return cls( + mask_info_lib.MaskInfo(*fwd_mask_info), + dkv_mask_info, + **aux_data, + ) + + +def make_ring_attention( + mask: np.ndarray | mask_lib.Mask, + *, + config: SplashConfig | None = None, + is_mqa: bool, + save_residuals: bool = False, + mask_value: float = base.DEFAULT_MASK_VALUE, + downcast_smem_data: bool = True, + partial_mask_blocks_dtype: jax.typing.DTypeLike = np.int8, + ring_axis: str, + q_seq_shards: int = 1, + kv_seq_shards: int = 1, +): + """Creates a RingSplashAttentionKernel. + + Args: + mask: The attention mask. Ring attention supports square FullMask and + CausalMask with offset=0. + config: SplashAttention configuration. If None, uses the default config. + is_mqa: Whether the model uses Multi-Query Attention. + save_residuals: Whether to save residuals for the backward pass. + mask_value: The value to use for masked-out attention scores. + downcast_smem_data: Whether to downcast data in shared memory. + partial_mask_blocks_dtype: The dtype for partial mask blocks. + ring_axis: The name of the jax scan axis used for the ring. + q_seq_shards: The number of shards for the query sequence dimension. + kv_seq_shards: The number of shards for the key/value sequence dimension. + + Returns: + A RingSplashAttentionKernel instance. + + Raises: + ValueError: If the mask shape is unexpected or ring_axis is not specified. + NotImplementedError: If a requested mask, config, or residual mode is not + supported by the ring attention implementation. + """ + + if len(mask.shape) != 2: + raise ValueError(f"Unexpected mask shape: {mask.shape}") + if mask.shape[0] != mask.shape[1]: + raise NotImplementedError( + "Ring attention currently supports only square self-attention masks; " f"got shape {mask.shape}." + ) + if save_residuals: + raise NotImplementedError("Ring attention does not support save_residuals=True.") + if q_seq_shards <= 0 or kv_seq_shards <= 0: + raise ValueError("q_seq_shards and kv_seq_shards must be positive.") + if q_seq_shards != kv_seq_shards: + raise NotImplementedError("Ring attention currently requires q_seq_shards == kv_seq_shards.") + + if isinstance(mask, (np.ndarray, mask_lib.NumpyMask)): + raise NotImplementedError( + "Ring attention does not support dense NumpyMask yet; use FullMask or " "CausalMask with offset=0." + ) + + if not isinstance(mask, (mask_lib.FullMask, mask_lib.CausalMask)): + raise NotImplementedError( + "Only FullMask and CausalMask are supported for ring attention; " + "other lazy masks are intentionally unsupported until rotated-shard " + f"semantics are covered. Got {type(mask)}." + ) + if isinstance(mask, mask_lib.CausalMask) and mask.offset != 0: + raise NotImplementedError("Ring attention supports CausalMask only with offset=0; " f"got offset={mask.offset}.") + + default_config = SplashConfig.get_default() + if config is None: + config = dataclasses.replace(default_config, use_base2_exp=False) + elif config.use_base2_exp: + raise NotImplementedError( + "Ring attention does not support use_base2_exp=True because ring " + "normalization currently combines natural-log softmax statistics." + ) + if not config.fuse_reciprocal: + raise NotImplementedError( + "Ring attention requires fuse_reciprocal=True because ring " + "normalization combines already-normalized per-shard outputs." + ) + + process_fn = partial( + mask_info_lib.process_mask, + downcast_smem_data=downcast_smem_data, + partial_mask_blocks_dtype=partial_mask_blocks_dtype, + q_seq_shards=q_seq_shards, + kv_seq_shards=kv_seq_shards, + ) + + fwd_mask_info, mask_function_fwd = process_fn( + mask, + (config.block_q, config.block_kv), + return_dynamic_grid=True, + ) + fwd_mask_sparsity = _mask_sparsity(fwd_mask_info) + fwd_mask_info = tree_util.tree_map(jnp.array, fwd_mask_info) + + dkv_mask_info = None + dkv_mask_sparsity = 0.0 + if config.has_backward_blocks: + bq_dkv, bkv_dkv = config.block_q_dkv, config.block_kv_dkv + dkv_mask_info, mask_function_dkv = process_fn( + mask, + (bq_dkv, bkv_dkv), + is_dkv=True, + return_dynamic_grid=config.dq_reduction_steps == 3, + ) + assert (mask_function_fwd is None) == (mask_function_dkv is None) + dkv_mask_sparsity = _mask_sparsity(dkv_mask_info) + dkv_mask_info = tree_util.tree_map(jnp.array, dkv_mask_info) + + return RingSplashAttentionKernel( + fwd_mask_info, + dkv_mask_info, + ring_axis=ring_axis, + expected_ring_size=q_seq_shards, + config=config, + is_mqa=is_mqa, + save_residuals=save_residuals, + mask_value=mask_value, + mask_function=mask_function_fwd, + fwd_mask_sparsity=fwd_mask_sparsity, + dkv_mask_sparsity=dkv_mask_sparsity, + ) diff --git a/src/maxtext/kernels/tokamax_splash_attention/ring_attention_utils.py b/src/maxtext/kernels/tokamax_splash_attention/ring_attention_utils.py new file mode 100644 index 0000000000..6384335508 --- /dev/null +++ b/src/maxtext/kernels/tokamax_splash_attention/ring_attention_utils.py @@ -0,0 +1,87 @@ +# Copyright 2026 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 +# +# 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. +"""Shared helpers for Tokamax ring-family Splash attention kernels.""" + +import jax +from jax import lax +import jax.numpy as jnp +import numpy as np +from maxtext.kernels.tokamax_splash_attention import splash_attention_mask_info as mask_info_lib + +MaskInfo = mask_info_lib.MaskInfo + + +def dynamic_slice_mask_info(mask_info: MaskInfo, kv_shard_idx: jax.Array, ring_size: int) -> MaskInfo: + """Slices MaskInfo for the current ring step.""" + + def slice_if_exists(arr: jax.Array | None): + if arr is None: + return None + + shard_len = arr.shape[-1] // ring_size + start_idx = kv_shard_idx * shard_len + return lax.dynamic_slice_in_dim(arr, start_idx, shard_len, axis=-1) + + return MaskInfo( + mask_next=slice_if_exists(mask_info.mask_next), + active_rows=slice_if_exists(mask_info.active_rows), + active_cols=slice_if_exists(mask_info.active_cols), + num_active_blocks=slice_if_exists(mask_info.num_active_blocks), + block_mask=slice_if_exists(mask_info.block_mask), + partial_mask_blocks=mask_info.partial_mask_blocks, # partial mask blocks are global + q_sequence=mask_info.q_sequence, # Q sequence stays stationary + ) + + +def offset_q_sequence_for_kv_shard(mask_info: MaskInfo, kv_shard_idx: jax.Array, kv_seq_len: int) -> MaskInfo: + """Converts lazy mask Q ids to the current local KV coordinate frame.""" + if mask_info.q_sequence is None: + return mask_info + + kv_shard_offset = jnp.asarray(kv_shard_idx, dtype=mask_info.q_sequence.dtype) * kv_seq_len + return mask_info._replace(q_sequence=mask_info.q_sequence - kv_shard_offset) + + +def has_no_active_blocks(mask_info: MaskInfo) -> jax.Array | bool: + """Returns whether the mask metadata has no active compute blocks.""" + if mask_info.num_active_blocks is None: + return False + has_no_scheduled_blocks = jnp.all(mask_info.num_active_blocks == 0) + if mask_info.block_mask is None: + return has_no_scheduled_blocks + # Dynamic mask metadata is padded with -1 after num_active_blocks. + block_ids = jnp.arange(mask_info.block_mask.shape[-1]) + active_blocks = block_ids < jnp.max(mask_info.num_active_blocks) + block_mask = jnp.where(active_blocks, mask_info.block_mask, 0) + has_no_compute_blocks = jnp.all(block_mask == 0) + return jnp.logical_or(has_no_scheduled_blocks, has_no_compute_blocks) + + +def has_empty_attention_rows(logsumexp: jax.Array, max_logits: jax.Array, mask_value: float) -> jax.Array: + mask_value = jnp.asarray(mask_value, dtype=logsumexp.dtype) + return jnp.logical_and(logsumexp == mask_value, max_logits == mask_value) + + +def mask_sparsity(mask_info: MaskInfo) -> float: + if mask_info.block_mask is None: + return 1.0 + return float(np.mean(mask_info.block_mask > 0)) + + +def has_axis(axis_name: str) -> bool: + try: + lax.axis_size(axis_name) + return True + except (NameError, ValueError): + return False diff --git a/src/maxtext/kernels/tokamax_splash_attention/splash_attention_kernel.py b/src/maxtext/kernels/tokamax_splash_attention/splash_attention_kernel.py new file mode 100644 index 0000000000..4c29a4bb27 --- /dev/null +++ b/src/maxtext/kernels/tokamax_splash_attention/splash_attention_kernel.py @@ -0,0 +1,2063 @@ +# pylint: skip-file +# Copyright 2025 DeepMind Technologies Limited. All Rights Reserved. +# +# 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. +# ============================================================================== + +"""Implementation of Sparse Flash Attention, a.k.a. "Splash" attention.""" + +from collections.abc import Callable +import dataclasses +import enum +import functools +import json +import math +from typing import Any + +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 +import jax.numpy as jnp +import numpy as np +from maxtext.kernels.tokamax_splash_attention import base +from maxtext.kernels.tokamax_splash_attention import splash_attention_mask as mask_lib +from maxtext.kernels.tokamax_splash_attention import splash_attention_mask_info as mask_info_lib + + +P = jax.P +MaskInfo = mask_info_lib.MaskInfo +partial = functools.partial +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 + +LOG2E = math.log2(math.e) +LOG2E_INV = 1 / LOG2E + +# mypy: ignore-errors + + +def _not(x: jax.Array | bool) -> jax.Array | bool: + if isinstance(x, jax.Array): + return jnp.logical_not(x) + return not x + + +def _base2_stats_to_natural_base(stats: dict[str, jax.Array], mask_value: float) -> dict[str, jax.Array]: + """Converts base-2 residual stats to natural-base user stats.""" + stats["logsumexp"] = jnp.where( + stats["logsumexp"] == mask_value, + mask_value, + stats["logsumexp"] / LOG2E, + ) + stats["max_logits"] = jnp.where( + stats["max_logits"] == mask_value, + mask_value, + stats["max_logits"] / LOG2E, + ) + return stats + + +SegmentIds = base.SegmentIds + +MaskFunctionType = Callable[..., jax.Array] + + +def get_kernel_name(is_mqa: bool, save_residuals: bool, is_segmented: bool, phase: str) -> str: + """Returns a unique name for all SplashAttention kernel variants.""" + assert phase in ["dq", "dkv", "fwd"] + # Saving residuals is supported only for the fwd phase. + assert not save_residuals or phase == "fwd" + residuals = "_residuals" if save_residuals else "_no_residuals" + attention_type = "mqa" if is_mqa else "mha" + segments = "_segmented" if is_segmented else "" + return f"splash_{attention_type}_{phase}{segments}{residuals}" + + +# 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 SplashConfig: + """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 + + # TODO: Remove these 3 params, they're only kept for backwards compatibility. + block_q_dq: int | None = None + block_kv_dq: int | None = None + use_fused_bwd_kernel: bool = True + + q_layout: QKVLayout = QKVLayout.HEAD_DIM_MINOR + k_layout: QKVLayout = QKVLayout.HEAD_DIM_MINOR + v_layout: QKVLayout = QKVLayout.HEAD_DIM_MINOR + + fwd_cost_estimate: pl.CostEstimate | None = None + bwd_cost_estimate: pl.CostEstimate | None = None + + residual_checkpoint_name: str | None = None # whether to checkpoint outputs + attn_logits_soft_cap: float | None = None + fuse_reciprocal: bool = True # whether to compute o / lse inside the kernel + use_base2_exp: bool = True + max_logit_const: float | None = None + interpret: bool = False + # The fused bwd kernel accumulates dq at every grid step. To safely avoid + # read/write conflicts we conservatively avoid *any* in-kernel reductions. + # This parameter allows to override this behavior and specifies the number of + # reduction steps. For now, only 3 or all the kv steps are supported. + dq_reduction_steps: int | None = None + # An experimental scheduler that sometimes produces better softmax overlap. + use_experimental_scheduler: bool = False + + 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.dq_reduction_steps is not None and self.dq_reduction_steps != 3: + raise ValueError(f"Invalid dq_reduction_steps: {self.dq_reduction_steps}, only 3 or" " None are supported.") + if not self.use_fused_bwd_kernel: + raise ValueError("Only the fused bwd kernel is supported.") + + @property + def has_backward_blocks(self) -> bool: + backward_blocks = ( + self.block_q_dkv, + self.block_kv_dkv, + self.block_kv_dkv_compute, + ) + return all(b is not None for b in backward_blocks) + + @classmethod + def get_default(cls): + # TODO: Select better parameters based on a heuristic. + return SplashConfig( + 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, + fuse_reciprocal=True, + ) + + +to_i32 = lambda x: x.astype(jnp.int32) + + +def _apply_mask_and_soft_cap( + qk: jax.Array, + mask_value: float, + mask_ref, + q_sequence_ref, + q_segment_ids_ref, + kv_segment_ids_ref, + *, + attn_logits_soft_cap: float | None, + k_slice: pl.Slice, + k_offset: int | jax.Array, + bq: int, + k_in_lanes=True, + mask_function=None, + has_partial_mask: bool = False, +) -> tuple[jax.Array, jax.Array | None]: + assert mask_ref is None or q_sequence_ref is None + assert (q_sequence_ref is None) == (mask_function is None) + + masks = [] + if has_partial_mask: + if mask_ref is not None: + mask = mask_ref[:, k_slice] if k_in_lanes else mask_ref[k_slice, :] + masks.append(mask) + elif 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(logits / 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) + return qk, mask + else: + qk = cap_logits(qk) + return qk, None + + +def flash_attention_kernel( + # Prefetched inputs + active_rows_ref, + active_cols_ref, + mask_next_ref, + bounds_start_ref, + bounds_end_ref, + block_mask_ref, + # Inputs + q_ref, + k_ref, + v_ref, + q_segment_ids_ref, + kv_segment_ids_ref, + sinks_ref, + mask_ref, + q_sequence_ref, + max_logit_value_ref, + # Outputs + o_ref, + logsumexp_ref, + l_linear_ref, + max_logits_ref, + # Scratch + m_scratch_ref, + l_scratch_ref, + o_scratch_ref, + *, + mask_value: float, + kv_steps: int, + bq: int, + bkv: int, + bkv_compute: int, + head_dim_v: int, + mask_function: MaskFunctionType | None, + fuse_reciprocal: bool, # config.fuse_reciprocal or not save_residuals + config: SplashConfig, +): + del mask_next_ref, active_rows_ref + float32 = jnp.float32 + HEAD_DIM_MINOR = QKVLayout.HEAD_DIM_MINOR + attn_logits_soft_cap = config.attn_logits_soft_cap + if attn_logits_soft_cap is not None and config.use_base2_exp: + attn_logits_soft_cap *= LOG2E + + # If the head_dim_v is not a multiple of the number of lanes, it will be + # padded to that multiple with zeros. + head_dim_v_repeats = pl.cdiv(head_dim_v, NUM_LANES) + + grid_idx = pl.program_id(1) + h = pl.program_id(0) + + if block_mask_ref is not None: + block_mask = block_mask_ref[grid_idx].astype(jnp.int32) + should_not_mask = block_mask != 1 + should_run = block_mask != 0 + should_initialize = bounds_start_ref[grid_idx].astype(jnp.bool_) + should_write = bounds_end_ref[grid_idx].astype(jnp.bool_) + j = active_cols_ref[grid_idx].astype(jnp.int32) + else: + should_not_mask = False + should_run = True + j = grid_idx % kv_steps + should_initialize = j == 0 + should_write = j == kv_steps - 1 + + max_logit_estimate = config.max_logit_const # potentially None + if max_logit_value_ref is not None: # already ensures max_logit_const is None + max_logit_estimate = max_logit_value_ref[0, h] + + if config.use_base2_exp and max_logit_estimate is not None: + max_logit_estimate *= LOG2E + + @pl.when(should_initialize) + def init(): + o_scratch_ref[...] = jnp.zeros_like(o_scratch_ref) + + sink = None + if sinks_ref is not None: + sink = sinks_ref[0, h].astype(m_scratch_ref.dtype) + if config.use_base2_exp: + sink *= LOG2E + + if sinks_ref is None and max_logit_estimate is None: + m_scratch_ref[...] = jnp.full_like(m_scratch_ref, mask_value) + l_scratch_ref[...] = jnp.zeros_like(l_scratch_ref) + elif sinks_ref is None and max_logit_estimate is not None: + m_scratch_ref[...] = jnp.full_like(m_scratch_ref, max_logit_estimate) + l_scratch_ref[...] = jnp.zeros_like(l_scratch_ref) + elif sinks_ref is not None and max_logit_estimate is None: + m_scratch_ref[...] = jnp.full_like(m_scratch_ref, sink) + l_scratch_ref[...] = jnp.ones_like(l_scratch_ref) + else: # sinks_ref is not None and max_logit_estimate is not None + exp = jnp.exp2 if config.use_base2_exp else jnp.exp + m_scratch_ref[...] = jnp.full_like(m_scratch_ref, max_logit_estimate) + l_scratch_ref[...] = exp(sink - jnp.full_like(l_scratch_ref, max_logit_estimate)) + + def body(kv_compute_index, _, has_partial_mask=False): + 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 config.q_layout == HEAD_DIM_MINOR else q_ref[...].T + if config.use_base2_exp: + q *= LOG2E + + qk_dims = NT_DIM_NUMBERS if config.k_layout == HEAD_DIM_MINOR else NN_DIM_NUMBERS + if config.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, + mask_ref, + q_sequence_ref, + q_segment_ids_ref, + kv_segment_ids_ref, + attn_logits_soft_cap=attn_logits_soft_cap, + k_slice=slice_k, + k_offset=j * bkv + kv_compute_index * bkv_compute, + bq=bq, + mask_function=mask_function, + has_partial_mask=has_partial_mask, + ) + + qk, mask = apply_mask_and_soft_cap() + + if max_logit_estimate is None: + 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) + else: + m_next = None + + bkv_repeats, rem = divmod(bkv_compute, NUM_LANES) + if rem != 0: + raise NotImplementedError(f"{bkv_compute=} should be a multiple of {NUM_LANES}") + + exp = jnp.exp2 if config.use_base2_exp else jnp.exp + if max_logit_estimate is None: + s_curr = exp(qk - jnp.tile(m_next, (1, bkv_repeats))) + else: + s_curr = exp(qk - max_logit_estimate) + if mask is not None: + s_curr = jnp.where(mask, s_curr, 0.0) + 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) + + if max_logit_estimate is None: + alpha = exp(m_prev - m_next) + l_next = l_curr + alpha * l_prev + m_scratch_ref[...], l_scratch_ref[...] = m_next, l_next + else: + alpha = None + l_scratch_ref[...] = l_curr + l_prev + + sv_dims = NN_DIM_NUMBERS if config.v_layout == HEAD_DIM_MINOR else NT_DIM_NUMBERS + if config.v_layout == HEAD_DIM_MINOR: + v = v_ref[slice_k, :] + else: + v = v_ref[:, slice_k] + o_curr = lax.dot_general(s_curr, v, sv_dims) + + if max_logit_estimate is None: + alpha_o = jnp.tile(alpha, (1, head_dim_v_repeats)) + alpha_o = alpha_o[..., : o_scratch_ref.shape[-1]] + o_scratch_ref[...] = alpha_o * o_scratch_ref[...] + o_curr + else: + o_scratch_ref[...] = o_scratch_ref[...] + o_curr + + assert bkv % bkv_compute == 0 + num_iters = k_ref.shape[0 if config.k_layout == HEAD_DIM_MINOR else 1] // bkv_compute + + @pl.when(jnp.logical_and(should_not_mask, should_run)) + def _(): + lax.fori_loop(0, num_iters, body, None, unroll=True) + + @pl.when(jnp.logical_and(_not(should_not_mask), should_run)) + def _(): + lax.fori_loop(0, num_iters, partial(body, has_partial_mask=True), None, unroll=True) + + @pl.when(should_write) + def end(): + l = l_scratch_ref[...] + m = m_scratch_ref[...] + safe_l = jnp.where(l == 0.0, 1.0, l) + if fuse_reciprocal: # allows fusing reciprocal out of the kernel + l_inv = jnp.tile(jnp.where(l == 0.0, 0.0, 1.0 / safe_l), (1, head_dim_v_repeats)) + l_inv = l_inv[..., : o_scratch_ref.shape[-1]] + o_ref[...] = (o_scratch_ref[...] * l_inv).astype(o_ref.dtype) + else: + o_ref[...] = o_scratch_ref[...].astype(o_ref.dtype) + if logsumexp_ref is not None: + assert logsumexp_ref.shape == (bq, NUM_LANES) + log = jnp.log2 if config.use_base2_exp else jnp.log + logsumexp = jnp.where(l == 0.0, mask_value, m + log(safe_l)) + logsumexp_ref[...] = logsumexp.astype(logsumexp_ref.dtype) + if l_linear_ref is not None: + assert l_linear_ref.shape == (bq, NUM_LANES) + l_linear_ref[...] = l.astype(l_linear_ref.dtype) + if max_logits_ref is not None: + assert max_logits_ref.shape == (bq, NUM_LANES) + max_logits_ref[...] = m.astype(max_logits_ref.dtype) + + +def _div(dividend: int, divisor: int): + if divisor == 1: + return dividend + + return lax.div(dividend, divisor) + + +def _bytes(x: jax.Array | jax.ShapeDtypeStruct | None) -> int: + if x is None: + return 0 + + if jnp.issubdtype(x.dtype, jnp.floating): + info = jnp.finfo + elif jnp.issubdtype(x.dtype, jnp.integer): + info = jnp.iinfo + else: + raise ValueError(f"Unsupported dtype: {x.dtype}") + return math.ceil(math.prod(x.shape) * info(x.dtype).bits / 8) + + +def _splash_attention_forward( + mask_info: MaskInfo, + q: jax.Array, + k: jax.Array, + v: jax.Array, + segment_ids: base.SegmentIds | None, + sinks: jax.Array | None, + mask_value: float, + is_mqa: bool, + config: SplashConfig, + save_residuals: bool, + mask_function: MaskFunctionType | None, + fwd_mask_sparsity: float, + max_logit_value: jax.Array | None = None, +) -> base.SplashCustomReturnType: + num_q_heads, q_seq_len, head_dim_qk = q.shape + head_dim_v = v.shape[-1] + bq, bkv = config.block_q, config.block_kv + bkv_compute = config.block_kv_compute + fuse_reciprocal = config.fuse_reciprocal or not save_residuals + bounds_start, bounds_end = mask_info_lib.find_bounds(mask_info.active_rows) + + if is_mqa: + expected_kv_rank = 2 + num_kv_heads = 1 + else: + expected_kv_rank = 3 + num_kv_heads = k.shape[0] + + if len(k.shape) != expected_kv_rank: + raise ValueError(f"Expected {expected_kv_rank}-dim 'key' tensor for MQA. Instead got a" f" {len(k.shape)}-dim one.") + + if k.shape[-1] != head_dim_qk: + raise ValueError(f"Expected 'key' head dimension to be: {head_dim_qk}. Instead got:" f" {k.shape[-1]}.") + + if not is_mqa and num_q_heads % num_kv_heads != 0: + raise ValueError( + f"In MHA/GQA, expected number of 'query' heads ({num_q_heads}) to" + f" be a multiple of the number of 'key' heads ({num_kv_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[-2] + kv_steps = kv_seq_len // bkv + q_heads_per_kv_head = num_q_heads // num_kv_heads + dynamic_grid = mask_info.active_rows is not None + + if segment_ids is not None: + assert isinstance(segment_ids.q, jax.Array) # for pytype + assert isinstance(segment_ids.kv, jax.Array) # for pytype + 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,)}") + if config.max_logit_const is not None and max_logit_value is not None: + raise ValueError(f"Only one of {config.max_logit_const=} and" f" {max_logit_value=} can be set.") + if max_logit_value is not None: + if max_logit_value.shape not in ((), (1,), (num_q_heads,)): + raise ValueError( + "max_logit_value should be a 0,1-dim jax.Array of shape (), (1,) or" + f" ({num_q_heads=},) but got {jax.typeof(max_logit_value)}" + ) + max_logit_value = jnp.broadcast_to(jnp.atleast_1d(max_logit_value), (num_q_heads,)) + + q_layout = config.q_layout + k_layout = config.k_layout + v_layout = config.v_layout + + def unravel(f): + def index_map(h, grid_idx, rows_ref, cols_ref, *_): + if dynamic_grid: + i = to_i32(rows_ref[grid_idx]) + j = to_i32(cols_ref[grid_idx]) + else: + i = grid_idx // kv_steps + j = grid_idx % kv_steps + return f(h, i, j) + + return index_map + + def create_kv_index_map(layout): + def index_map(h, i, j): + del i # Unused. + prefix = () if is_mqa else (_div(h, q_heads_per_kv_head),) + return from_head_minor((*prefix, j, 0), layout) + + return index_map + + q_index_map = unravel(lambda h, i, j: from_head_minor((h, i, 0), q_layout)) + out_index_map = unravel(lambda h, i, j: (h, i, 0)) + k_index_map = unravel(create_kv_index_map(k_layout)) + v_index_map = unravel(create_kv_index_map(v_layout)) + + def mask_index_map(h, grid_idx, rows_ref, cols_ref, mask_next_ref=None, *_): + del h, rows_ref, cols_ref # Unused. + next_m = to_i32(mask_next_ref[grid_idx]) + return next_m, 0, 0 + + q_segment_ids_index_map = unravel(lambda h, i, j: (i, 0)) + kv_segment_ids_index_map = unravel(lambda h, i, j: (0, 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 sinks is not None: + assert sinks.shape == (num_q_heads,), f"{sinks.shape=} != {num_q_heads=}" + # align sinks to sublanes to allow vmap and shard_map over the kernel + in_specs += [ + pl.BlockSpec( + (NUM_SUBLANES, num_q_heads), + lambda h, i, j, *_: (0, 0), + memory_space=pltpu.SMEM, + ) + ] + sinks = jnp.broadcast_to(sinks.astype(jnp.float32)[None, :], (NUM_SUBLANES, num_q_heads)) + else: + in_specs += [None] + + if 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 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) + + if max_logit_value is not None: + # reshape to allow sublane selection for vmap-ping and shard_map-ping + max_logit_value = jnp.broadcast_to( + max_logit_value.astype(jnp.float32)[None, :], + (NUM_SUBLANES, num_q_heads), + ) + in_specs += [ + pl.BlockSpec( + (NUM_SUBLANES, num_q_heads), + lambda *_: (0, 0), + memory_space=pltpu.SMEM, + ) + ] + else: + in_specs.append(None) + + out_shapes = [ + jax.ShapeDtypeStruct((num_q_heads, q_seq_len, head_dim_v), q.dtype), + ] + out_specs = [ + pl.BlockSpec((None, bq, head_dim_v), out_index_map), + ] + if save_residuals: + logsumexp_index_map = unravel(lambda h, i, j, *_: (h, i, 0)) + + out_shapes += [ + # logsumexp + jax.ShapeDtypeStruct((num_q_heads, q_seq_len, NUM_LANES), jnp.float32) if fuse_reciprocal else None, + # l_linear + jax.ShapeDtypeStruct((num_q_heads, q_seq_len, NUM_LANES), jnp.float32) if not fuse_reciprocal else None, + # max_logits + jax.ShapeDtypeStruct((num_q_heads, q_seq_len, NUM_LANES), jnp.float32), + ] + out_specs += [ + pl.BlockSpec((None, bq, NUM_LANES), logsumexp_index_map) if fuse_reciprocal else None, + pl.BlockSpec((None, bq, NUM_LANES), logsumexp_index_map) if not fuse_reciprocal else None, + pl.BlockSpec((None, bq, NUM_LANES), logsumexp_index_map), + ] + else: + out_shapes += [None, None, None] + out_specs += [None, None, None] + + kernel_name = get_kernel_name( + is_mqa=is_mqa, + save_residuals=save_residuals, + is_segmented=segment_ids is not None, + phase="fwd", + ) + metadata = {"xprof_metadata": json.dumps(dataclasses.asdict(config))} + + def _fwd_cost_estimate( + q: jax.Array, + k: jax.Array, + v: jax.Array, + q_segment_ids: jax.Array | None, + kv_segment_ids: jax.Array | None, + partial_mask_blocks: jax.Array | None, + out_shapes: list[jax.ShapeDtypeStruct], + mask_sparsity: float, + ) -> pl.CostEstimate: + num_q_heads, q_seq_len, head_dim_qk = q.shape + kv_seq_len, head_dim_v = v.shape[-2:] + + matmul_flops = 2 * q_seq_len * kv_seq_len * head_dim_qk + 2 * q_seq_len * kv_seq_len * head_dim_v + + # This is an upper bound because `mask_sparsity` is actually the mean + # sparsity of the non-fully masked **blocks**. + total_flops = num_q_heads * matmul_flops * mask_sparsity + + # Count expensive exp() calls + transcendentals = num_q_heads * q_seq_len * kv_seq_len * mask_sparsity + + inputs_ = [q, k, v, q_segment_ids, kv_segment_ids, partial_mask_blocks] + input_bytes = sum(map(_bytes, inputs_)) + output_bytes = sum(map(_bytes, out_shapes)) + return pl.CostEstimate( + flops=int(total_flops), + transcendentals=int(transcendentals), + bytes_accessed=int(input_bytes + output_bytes), + ) + + vmem_inputs = [ + q, + k, + v, + q_segment_ids, + kv_segment_ids, + mask_info.partial_mask_blocks, + ] + cost_estimate = config.fwd_cost_estimate or _fwd_cost_estimate(*vmem_inputs, out_shapes, fwd_mask_sparsity) + + if dynamic_grid: + num_active_blocks = mask_info.num_active_blocks[0] + grid = (num_q_heads, num_active_blocks) + is_empty_attention_block = num_active_blocks == 0 + else: + grid = (num_q_heads, kv_steps * (q_seq_len // bq)) + is_empty_attention_block = False + + with jax.named_scope(kernel_name): + all_out = pl.pallas_call( + partial( + flash_attention_kernel, + mask_value=mask_value, + kv_steps=kv_steps, + bq=bq, + bkv=bkv, + bkv_compute=bkv_compute, + head_dim_v=head_dim_v, + # note: fuse_reciprocal can only be False if save_residuals is True + # fuse_reciprocal = (config.fuse_reciprocal or not save_residuals) + fuse_reciprocal=fuse_reciprocal, + config=config, + mask_function=mask_function, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=6, + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + scratch_shapes=[ + pltpu.VMEM((bq, NUM_LANES), jnp.float32), # m_scratch + pltpu.VMEM((bq, NUM_LANES), jnp.float32), # l_scratch + pltpu.VMEM((bq, head_dim_v), jnp.float32), # o_scratch + ], + ), + compiler_params=pltpu.CompilerParams( + dimension_semantics=("parallel", "arbitrary"), + flags={"XLA_TPU_FORCE_LP_LLO_SCHEDULER": (config.use_experimental_scheduler)}, + ), + out_shape=out_shapes, + name=kernel_name, + cost_estimate=cost_estimate, + interpret=config.interpret, + metadata=metadata, + )( + mask_info.active_rows, + mask_info.active_cols, + mask_info.mask_next, + bounds_start, + bounds_end, + mask_info.block_mask, + q if q_layout == QKVLayout.HEAD_DIM_MINOR else q.mT, + k if k_layout == QKVLayout.HEAD_DIM_MINOR else k.mT, + v if v_layout == QKVLayout.HEAD_DIM_MINOR else v.mT, + q_segment_ids, + kv_segment_ids, + sinks, + mask_info.partial_mask_blocks, + q_sequence, + max_logit_value, + ) + out, logsumexp, l_linear, max_logits = all_out + + # If there is no compute to do within an attention block, then we want to + # initialize the output and residuals to default values. Otherwise, we will + # read uninitialized memory. This is a common case in ring attention. + def init_if_empty(x: jax.Array, value: float) -> jax.Array: + if not dynamic_grid: + return x + + return jnp.where(is_empty_attention_block, value, x) + + out = init_if_empty(out, 0.0) + + if save_residuals: + assert max_logits is not None + max_logits = init_if_empty(max_logits[..., 0], mask_value) + + if fuse_reciprocal: + assert logsumexp is not None + logsumexp = init_if_empty(logsumexp[..., 0], mask_value) + else: + assert l_linear is not None + log = jnp.log2 if config.use_base2_exp else jnp.log + + l = init_if_empty(l_linear[..., 0], 0.0) + safe_l = jnp.where(l == 0.0, 1.0, l) + logsumexp = jnp.where(l == 0.0, mask_value, max_logits + log(safe_l)) + l_inv = jnp.where(l == 0.0, 0.0, 1.0 / safe_l) + out = (out * l_inv[..., None]).astype(out.dtype) + else: + # If we're not saving residuals, then we can't fuse the reciprocal + # out of the kernel. + assert fuse_reciprocal + + if config.residual_checkpoint_name is not None: + out = ad_checkpoint.checkpoint_name(out, name=config.residual_checkpoint_name) + if logsumexp is not None: + logsumexp = ad_checkpoint.checkpoint_name(logsumexp, name=config.residual_checkpoint_name) + if save_residuals: + stats = {"logsumexp": logsumexp, "max_logits": max_logits} + stats = jax.tree.map(jax.lax.stop_gradient, stats) + return out, stats + return out + + +@partial( + jax.custom_vjp, + nondiff_argnames=( + "save_residuals", + "mask_value", + "is_mqa", + "config", + "mask_function", + "fwd_mask_sparsity", + "dkv_mask_sparsity", + ), +) +def _splash_attention_custom( + fwd_mask_info: MaskInfo, + dkv_mask_info: MaskInfo | None, + q: jax.Array, + k: jax.Array, + v: jax.Array, + segment_ids: base.SegmentIds | None, + sinks: jax.Array | None, + save_residuals: bool, + mask_value: float, + is_mqa: bool, + config: SplashConfig, + mask_function: MaskFunctionType | None, + fwd_mask_sparsity: float, + dkv_mask_sparsity: float, + max_logit_value: jax.Array | None = None, +) -> base.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 dkv_mask_info + + ret = _splash_attention_forward( # pytype: disable=wrong-arg-types + fwd_mask_info, + q, + k, + v, + segment_ids, + sinks, + mask_value=mask_value, + is_mqa=is_mqa, + config=config, + save_residuals=save_residuals, + mask_function=mask_function, + fwd_mask_sparsity=fwd_mask_sparsity, + max_logit_value=max_logit_value, + ) + if save_residuals: + out, stats = ret + if config.use_base2_exp: # for user, output values in natural base + stats = _base2_stats_to_natural_base(stats, mask_value) + return out, stats + else: + return ret + + +def _splash_attention_fwd( + fwd_mask_info: MaskInfo, + dkv_mask_info: MaskInfo | None, + q: jax.Array, + k: jax.Array, + v: jax.Array, + segment_ids: base.SegmentIds | None, + sinks: jax.Array | None, + save_residuals: bool, + mask_value: float, + is_mqa: bool, + config: SplashConfig, + mask_function: MaskFunctionType | None, + fwd_mask_sparsity: float, + dkv_mask_sparsity: float, + max_logit_value: jax.Array | None = None, +) -> tuple[tuple[jax.Array], base.SplashResidualsType]: + + # TODO: add some higher order AD check that isn't save_residuals based. + # if save_residuals: + # raise NotImplementedError("Higher-order AD not supported.") + + out, stats = _splash_attention_forward( # pytype: disable=wrong-arg-types + fwd_mask_info, + q, + k, + v, + segment_ids, + sinks, + mask_value=mask_value, + is_mqa=is_mqa, + config=config, + save_residuals=True, + mask_function=mask_function, + fwd_mask_sparsity=fwd_mask_sparsity, + max_logit_value=max_logit_value, + ) + logsumexp = stats["logsumexp"] # save in the config base for the bwd pass + if config.use_base2_exp: # for user, output values in natural base + stats = _base2_stats_to_natural_base(stats, mask_value) + residuals = q, k, v, segment_ids, sinks, out, logsumexp, dkv_mask_info + if save_residuals: + return (out, stats), residuals + else: + return out, residuals + + +def _flash_attention_dq_kernel( + # Prefetched inputs + active_rows_ref, + active_cols_ref, + mask_next_ref, + bounds_start_ref, + bounds_end_ref, + block_mask_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, + kv_steps: int, + bq: int, + bkv: int, + mask_function: MaskFunctionType | None, + config: SplashConfig, +): + del mask_next_ref, active_rows_ref + float32 = jnp.float32 + HEAD_DIM_MINOR = QKVLayout.HEAD_DIM_MINOR + attn_logits_soft_cap = config.attn_logits_soft_cap + if attn_logits_soft_cap is not None and config.use_base2_exp: + attn_logits_soft_cap *= LOG2E + + grid_idx = pl.program_id(1) + if block_mask_ref is not None: + kv_index = active_cols_ref[grid_idx].astype(jnp.int32) + should_not_mask = block_mask_ref[grid_idx].astype(jnp.int32) != 1 + should_initialize = bounds_start_ref[grid_idx].astype(jnp.bool_) + should_write = bounds_end_ref[grid_idx].astype(jnp.bool_) + else: + kv_index = grid_idx % kv_steps + should_not_mask = False + should_initialize = kv_index == 0 + should_write = kv_index == kv_steps - 1 + + @pl.when(should_initialize) + def init(): + dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref) + + def body(has_partial_mask: bool = False): + q = q_ref[...] if config.q_layout == HEAD_DIM_MINOR else q_ref[...].T + if config.use_base2_exp: + q *= LOG2E + # 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 config.k_layout == HEAD_DIM_MINOR else NN_DIM_NUMBERS + qk_uncapped = lax.dot_general(q, k, qk_dims, preferred_element_type=float32) + + qk, mask = _apply_mask_and_soft_cap( + qk_uncapped, + mask_value, + 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=kv_index * bkv, + bq=bq, + mask_function=mask_function, + has_partial_mask=has_partial_mask, + ) + exp = jnp.exp2 if config.use_base2_exp else jnp.exp + p = exp(qk - logsumexp) + if mask is not None: + p = jnp.where(mask, p, 0.0) + dp_dims = NT_DIM_NUMBERS if config.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) + ds = ds * (1 - d * d) + + dq_dims = NN_DIM_NUMBERS if config.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(should_not_mask) + def _(): + body() + + @pl.when(jnp.logical_not(should_not_mask)) + def _(): + body(has_partial_mask=True) + + @pl.when(should_write) + def end(): + dq_ref[...] = dq_scratch_ref[...].astype(dq_ref.dtype) + + +def _flash_attention_dkv_kernel( + # Prefetched inputs + active_rows_ref, + active_cols_ref, + mask_next_ref, + bounds_start_ref, + bounds_end_ref, + block_mask_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, + # aliases + dq_alias, + dk_alias, + dv_alias, + # Outputs + dq_ref, + dk_ref, + dv_ref, + # Scratch + dq_scratch_ref, + dk_scratch_ref, + dv_scratch_ref, + *, + mask_value: float, + q_steps: int, + bq: int, + bkv_compute: int, + bkv: int, + mask_function: MaskFunctionType | None, + q_heads_per_kv_head: int, + config: SplashConfig, +): + del mask_next_ref, active_cols_ref + HEAD_DIM_MINOR = QKVLayout.HEAD_DIM_MINOR + attn_logits_soft_cap = config.attn_logits_soft_cap + if attn_logits_soft_cap is not None and config.use_base2_exp: + attn_logits_soft_cap *= LOG2E + + if active_rows_ref is not None: + assert bounds_start_ref is not None + assert bounds_end_ref is not None + grid_idx = pl.program_id(1) + kv_index = active_rows_ref[grid_idx].astype(jnp.int32) + should_initialize = bounds_start_ref[grid_idx].astype(jnp.bool_) + should_write = bounds_end_ref[grid_idx].astype(jnp.bool_) + else: + kv_index, q_head, q_index = ( + pl.program_id(0), + pl.program_id(1), + pl.program_id(2), + ) + grid_idx = (kv_index * q_steps) + q_index + should_initialize = q_index == 0 + should_write = True if q_steps <= 2 else q_index == q_steps - 1 + if q_heads_per_kv_head > 1: + q_head_index_per_kv_head = lax.rem(q_head, q_heads_per_kv_head) + should_initialize = jnp.logical_and(should_initialize, q_head_index_per_kv_head == 0) + should_write = jnp.logical_and(should_write, q_head_index_per_kv_head == q_heads_per_kv_head - 1) + + if block_mask_ref is not None: + should_not_mask = block_mask_ref[grid_idx].astype(jnp.int32) != 1 + should_run = block_mask_ref[grid_idx].astype(jnp.int32) != 0 + else: + should_not_mask = False + should_run = True + + # 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). + + @pl.when(should_initialize) + def init(): + dk_scratch_ref[...] = jnp.zeros_like(dk_scratch_ref) + dv_scratch_ref[...] = jnp.zeros_like(dv_scratch_ref) + + def body(i, _, has_partial_mask=False): + + slice_k = pl.ds(i * bkv_compute, bkv_compute) + q = q_ref[...] # We keep q potentially transposed, since it's always RHS + if config.use_base2_exp: + scaled_q = q * LOG2E + else: + scaled_q = q + + def _load_kv(ref, layout): + if layout == HEAD_DIM_MINOR: + return ref[slice_k, :] + return ref[:, slice_k].T + + k = _load_kv(k_ref, config.k_layout) + v = _load_kv(v_ref, config.v_layout) + logsumexp = logsumexp_ref[:1, :] + do = do_ref[...] + di = di_ref[:1, :] + + qk_dims = NT_DIM_NUMBERS if config.q_layout == HEAD_DIM_MINOR else NN_DIM_NUMBERS + qk_uncapped = lax.dot_general(k, scaled_q, qk_dims, preferred_element_type=jnp.float32) + + qk, mask = _apply_mask_and_soft_cap( + qk_uncapped, + mask_value, + 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, + has_partial_mask=has_partial_mask, + ) + exp = jnp.exp2 if config.use_base2_exp else jnp.exp + p = exp(qk - logsumexp) + if mask is not None: + p = jnp.where(mask, p, 0.0) + 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) + ds = ds * (1 - d * d) + dk_dims = NN_DIM_NUMBERS if config.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 + if dq_alias is not None: + dq_ref[...] = dq_alias[...] + dq.astype(dq_ref.dtype) + else: + 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_alias is not None: + dq_ref[...] = dq_alias[...] + else: + dq_ref[...] = jnp.zeros_like(dq_ref) + + num_iters = k_ref.shape[0 if config.k_layout is HEAD_DIM_MINOR else 1] // bkv_compute + + @pl.when(jnp.logical_and(should_not_mask, should_run)) + def _(): + lax.fori_loop(0, num_iters, body, None, unroll=True) + + @pl.when(jnp.logical_and(_not(should_not_mask), should_run)) + def _(): + lax.fori_loop(0, num_iters, partial(body, has_partial_mask=True), None, unroll=True) + + if dq_scratch_ref is not None: + if dq_alias is not None: + dq_ref[...] = dq_alias[...] + dq_scratch_ref[...].astype(dq_ref.dtype) + else: + dq_ref[...] = dq_scratch_ref[...].astype(dq_ref.dtype) + + if dk_alias is None: + assert dv_alias is None + + @pl.when(should_write) + def _(): + dk_ref[...] = dk_scratch_ref[...].astype(dk_ref.dtype) + dv_ref[...] = dv_scratch_ref[...].astype(dv_ref.dtype) + + else: + q_head = pl.program_id(0) + first_q_head_in_kv_group = lax.rem(q_head, q_heads_per_kv_head) == 0 + + @pl.when(jnp.logical_and(should_write, first_q_head_in_kv_group)) + def _(): + dk_ref[...] = dk_scratch_ref[...].astype(dk_ref.dtype) + dv_ref[...] = dv_scratch_ref[...].astype(dv_ref.dtype) + + @pl.when(jnp.logical_and(should_write, _not(first_q_head_in_kv_group))) + def _(): + dk_ref[...] = dk_alias[...] + dk_scratch_ref[...].astype(dk_ref.dtype) + dv_ref[...] = dv_alias[...] + dv_scratch_ref[...].astype(dv_ref.dtype) + + +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: MaskInfo, + mask_value: float, + mask_function: MaskFunctionType | None, + config: SplashConfig, + dkv_mask_sparsity: float, + return_fp32_grads: bool = False, +): + num_q_heads, q_seq_len, head_dim_qk = q.shape + kv_seq_len, head_dim_v = v.shape[-2:] + num_kv_heads = 1 if is_mqa else k.shape[0] + dynamic_grid = mask_info.active_rows is not None + + bounds_start, bounds_end = mask_info_lib.find_bounds(mask_info.active_rows) + 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/GQA, expected number of 'query' heads ({num_q_heads}) to" + f" be a multiple of the number of 'key' heads ({num_kv_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.") + + kv_steps = kv_seq_len // bkv + q_steps = q_seq_len // bq + q_heads_per_kv_head = num_q_heads // num_kv_heads + + if dynamic_grid: + + def unravel(f): + def index_map(h, grid_idx, rows_ref, cols_ref, *_): + j = to_i32(rows_ref[grid_idx]) + i = to_i32(cols_ref[grid_idx]) + return f(h, i, j) + + return index_map + + grid_size = mask_info.num_active_blocks[0] + grid = (num_q_heads, grid_size) + + def mask_index_map(h, grid_idx, rows_ref, cols_ref, mask_next_ref=None, *_): + del h, rows_ref, cols_ref # Unused. + next_m = to_i32(mask_next_ref[grid_idx]) + return next_m, 0, 0 + + else: + unravel = lambda f: lambda j, h, i, *_: f(h, i, j) + grid = (kv_steps, num_q_heads, q_steps) + + def mask_index_map(j, h, i, rows_ref, cols_ref, mask_next_ref=None, *_): + del h, rows_ref, cols_ref # Unused. + grid_idx = j * q_steps + i + next_m = to_i32(mask_next_ref[grid_idx]) + return next_m, 0, 0 + + q_index_map = unravel(lambda h, i, j: from_head_minor((h, i, 0), config.q_layout)) + o_index_map = unravel(lambda h, i, j: (h, i, 0)) + + def create_kv_index_map(layout): + def index_map(h, i, j, *_): + del i # Unused. + prefix = () if is_mqa else (_div(h, q_heads_per_kv_head),) + return from_head_minor((*prefix, j, 0), layout) + + return index_map + + k_index_map = unravel(create_kv_index_map(config.k_layout)) + v_index_map = unravel(create_kv_index_map(config.v_layout)) + + q_spec = pl.BlockSpec(from_head_minor((None, bq, head_dim_qk), config.q_layout), q_index_map) + + o_spec = pl.BlockSpec((None, bq, head_dim_v), o_index_map) + k_spec = pl.BlockSpec( + from_head_minor( + (bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), + config.k_layout, + ), + k_index_map, + ) + + v_spec = pl.BlockSpec( + from_head_minor( + (bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), + config.v_layout, + ), + v_index_map, + ) + + def create_dkv_index_map(h, i, j, *_): + del i # Unused. + prefix = () if is_mqa else (_div(h, q_heads_per_kv_head),) + return (*prefix, j, 0) + + dkv_index_map = unravel(create_dkv_index_map) + + 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, + ) + mask_spec = pl.BlockSpec((None, bkv, bq), mask_index_map) + + q_segment_ids_index_map = unravel(lambda h, i, j: (0, i)) + if segment_ids is not None: + kv_segment_ids_index_map = unravel(lambda h, i, j: (j, 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 + + logsumexp_index_map = unravel(lambda h, i, j: (h, 0, i)) + + assert logsumexp.shape == di.shape == (num_q_heads, q_seq_len) + # TODO: 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: 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) + + dq_reduction_steps = config.dq_reduction_steps + if not dynamic_grid and kv_steps <= 3 and dq_reduction_steps == 3: + dq_reduction_steps = None + + dq = dq_alias_spec = None + if dq_reduction_steps == 3: + dq_index_map = unravel(lambda h, i, j: (j % 3, h, i, 0)) + dq_spec = pl.BlockSpec((None, None, bq, head_dim_qk), dq_index_map) + dq_alias_spec = dq_spec + dq_dtype = jnp.float32 if return_fp32_grads else q.dtype + dq_shape = jax.ShapeDtypeStruct((3, *q.shape), dq_dtype) + dq = jnp.zeros_like(dq_shape) + else: + dq_index_map = unravel(lambda h, i, j: (j, h, i, 0)) + dq_spec = pl.BlockSpec((None, None, bq, head_dim_qk), dq_index_map) + # Only accumulate in fp32 if there's a small number of reduction steps. + q_dtype = jnp.float32 if return_fp32_grads or kv_steps > 4 else q.dtype + dq_shape = jax.ShapeDtypeStruct((kv_steps, *q.shape), q_dtype) + + in_specs += [dq_alias_spec] + + if bkv == bkv_compute: + dq_scratch = None + else: + dq_scratch = pltpu.VMEM((bq, head_dim_qk), jnp.float32) + + if dynamic_grid and q_heads_per_kv_head != 1: + # in/out aliasing to accumulate within kv groups. + in_specs += [dk_spec, dv_spec] + dk = lax.empty(k.shape, dtype=jnp.float32) + dv = lax.empty(v.shape, dtype=jnp.float32) + # Keep gradients in fp32 when accumulating over head groups. + dk_type = dv_type = jnp.float32 + else: + in_specs += [None, None] + dk, dv = None, None + dk_type = jnp.float32 if return_fp32_grads else k.dtype + dv_type = jnp.float32 if return_fp32_grads else v.dtype + + out_shapes = [ + dq_shape, + jax.ShapeDtypeStruct(k.shape, dk_type), + jax.ShapeDtypeStruct(v.shape, dv_type), + ] + out_specs = [dq_spec, dk_spec, dv_spec] + + kernel = functools.partial( + _flash_attention_dkv_kernel, + mask_value=mask_value, + q_steps=q_steps, + bq=bq, + bkv_compute=bkv_compute, + config=config, + bkv=bkv, + mask_function=mask_function, + q_heads_per_kv_head=q_heads_per_kv_head, + ) + + kernel_name = get_kernel_name( + is_mqa=is_mqa, + save_residuals=False, + is_segmented=segment_ids is not None, + phase="dkv", + ) + metadata = { + "xprof_metadata": json.dumps( + dict( + block_q_dkv=bq, + block_kv_dkv=bkv, + block_kv_dkv_compute=bkv_compute, + q_layout=config.q_layout, + k_layout=config.k_layout, + v_layout=config.v_layout, + use_experimental_scheduler=config.use_experimental_scheduler, + ), + ) + } + args = [ + # scalar prefetch + mask_info.active_rows, + mask_info.active_cols, + mask_info.mask_next, + bounds_start, + bounds_end, + mask_info.block_mask, + # inputs + q if config.q_layout == QKVLayout.HEAD_DIM_MINOR else q.mT, + k if config.k_layout == QKVLayout.HEAD_DIM_MINOR else k.mT, + v if config.v_layout == QKVLayout.HEAD_DIM_MINOR else v.mT, + q_segment_ids, + kv_segment_ids, + logsumexp, + do, + di, + mask_info.partial_mask_blocks, + q_sequence, + ] + num_args = sum(1 for x in args if x is not None) + input_output_aliases = {} + if dq_reduction_steps == 3: + if dynamic_grid and q_heads_per_kv_head != 1: + input_output_aliases = {num_args: 0, num_args + 1: 1, num_args + 2: 2} + else: + input_output_aliases = {num_args: 0} + elif dynamic_grid and q_heads_per_kv_head != 1: + input_output_aliases = {num_args: 1, num_args + 1: 2} + + scratch_shapes = [ + dq_scratch, + pltpu.VMEM((bkv, head_dim_qk), jnp.float32), + pltpu.VMEM((bkv, head_dim_v), jnp.float32), + ] + + def _bwd_cost_estimate( + q: jax.Array, + k: jax.Array, + v: jax.Array, + q_segment_ids: jax.Array | None, + kv_segment_ids: jax.Array | None, + logsumexp: jax.Array, + do: jax.Array, + di: jax.Array, + partial_mask_blocks: jax.Array | None, + q_sequence: jax.Array | None, + out_shapes: list[jax.ShapeDtypeStruct], + mask_sparsity_factor: float, + ) -> pl.CostEstimate: + num_q_heads, q_seq_len, head_dim_qk = q.shape + kv_seq_len, head_dim_v = v.shape[-2:] + + total_matmul_flops_per_head = ( + 2 * q_seq_len * kv_seq_len * head_dim_qk # qk + + 2 * q_seq_len * kv_seq_len * head_dim_v # dv + + 2 * q_seq_len * kv_seq_len * head_dim_v # dp + + 2 * q_seq_len * kv_seq_len * head_dim_qk # dq + + 2 * q_seq_len * kv_seq_len * head_dim_qk # dk + ) + + estimated_flops = int(total_matmul_flops_per_head * num_q_heads * mask_sparsity_factor) + + exp_flops = num_q_heads * q_seq_len * kv_seq_len * mask_sparsity_factor + if config.attn_logits_soft_cap is None: + tanh_flops = 0 + else: + tanh_flops = 2 * num_q_heads * q_seq_len * kv_seq_len * mask_sparsity_factor + estimated_transcendentals = int(exp_flops + tanh_flops) + + inputs_ = [ + q, + k, + v, + q_segment_ids, + kv_segment_ids, + logsumexp, + do, + di, + partial_mask_blocks, + q_sequence, + ] + input_bytes = sum(map(_bytes, inputs_)) + output_bytes = sum(map(_bytes, out_shapes)) + + estimated_bytes = input_bytes + output_bytes + + return pl.CostEstimate( + flops=estimated_flops, + transcendentals=estimated_transcendentals, + bytes_accessed=estimated_bytes, + ) + + cost_estimate = config.bwd_cost_estimate or _bwd_cost_estimate( + q, + k, + v, + q_segment_ids, + kv_segment_ids, + logsumexp, + do, + di, + mask_info.partial_mask_blocks, + q_sequence, + out_shapes, + dkv_mask_sparsity, + ) + + with jax.named_scope(kernel_name): + dq_unreduced, dk, dv = pl.pallas_call( + kernel, + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=6, + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + scratch_shapes=scratch_shapes, + ), + out_shape=out_shapes, + input_output_aliases=input_output_aliases, + # We set all dimensions to arbitrary because: + # 1) for heads, we are reducing over heads + # 2) for kv_seq_len, the splash attention prefetch schedule assumes no + # megacore + # 3) for q_seq_len, we are reducing over it to compute dkv + compiler_params=pltpu.CompilerParams(dimension_semantics=("arbitrary",) * len(grid)), + name=kernel_name, + cost_estimate=cost_estimate, + interpret=config.interpret, + metadata=metadata, + )(*args, dq, dk, dv) + dq = dq_unreduced.sum(axis=0) + if return_fp32_grads: + return dq.astype(jnp.float32), dk.astype(jnp.float32), dv.astype(jnp.float32) + dq = dq.astype(q.dtype) + dk = dk.astype(k.dtype) + dv = dv.astype(v.dtype) + return dq, dk, dv + + +def _splash_attention_bwd( + save_residuals: bool, + mask_value: float, + is_mqa: bool, + config: SplashConfig, + mask_function: MaskFunctionType | None, + fwd_mask_sparsity: float, + dkv_mask_sparsity: float, + res: base.SplashResidualsType, + grads: jax.Array | tuple[jax.Array, dict[str, jax.Array]], + return_fp32_grads: bool = False, +) -> tuple[ + MaskInfo | None, # fwd_mask_info + MaskInfo | None, # dvk_mask_info + jax.Array, # q + jax.Array, # k + jax.Array, # v + base.SegmentIds | None, # segment_ids + jax.Array | None, # segment_ids + jax.Array | None, # max_logit_estimate +]: + # If `save_residuals` is True, `_splash_attention_fwd` returns `(out, stats)`, + # so we unpack the gradients, otherwise it returns `out` and `grads` is just + # `do`. + if save_residuals: + do, _ = grads + else: + do = grads + del save_residuals, fwd_mask_sparsity + if not config.has_backward_blocks: + raise ValueError("Need to specify backward blocks.") + bq_dkv, bkv_dkv_memory, bkv_dkv_compute = ( + config.block_q_dkv, + config.block_kv_dkv, + config.block_kv_dkv_compute, + ) + q, k, v, segment_ids, sinks, o, logsumexp, 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, + mask_function=mask_function, + config=config, + dkv_mask_sparsity=dkv_mask_sparsity, + return_fp32_grads=return_fp32_grads, + ) + dsinks = None + if sinks is not None: + logsumexp_ = (logsumexp / LOG2E) if config.use_base2_exp else logsumexp + sinks_exp = -jnp.exp(sinks[..., None, None].astype(jnp.float32) - logsumexp_[..., None].astype(jnp.float32)) + dsinks = jnp.sum(sinks_exp.astype(o.dtype) * o * do, axis=(-1, -2)) + # Match the signature of the fwd function. + assert dq is not None + return ( + None, # fwd_mask_info + None, # dvk_mak_info + dq, # q + dk, # k + dv, # v + None, # segment_ids + dsinks, # sinks + None, # max_logit_estimate + ) + + +_splash_attention_custom.defvjp(_splash_attention_fwd, _splash_attention_bwd) + + +@partial( + jax.jit, + static_argnames=[ + "is_mqa", + "config", + "save_residuals", + "mask_value", + "mask_function", + "fwd_mask_sparsity", + "dkv_mask_sparsity", + ], +) +def _splash_attention( + fwd_mask_info: MaskInfo, + dkv_mask_info: MaskInfo | None, + q: jax.Array, + k: jax.Array, + v: jax.Array, + segment_ids: base.SegmentIds | None = None, + sinks: jax.Array | None = None, + *, + is_mqa: bool, + config: SplashConfig | None, + save_residuals: bool, + mask_value: float, + max_logit_value: jax.Array | None = None, + mask_function: MaskFunctionType | None, + fwd_mask_sparsity: float, + dkv_mask_sparsity: float, +) -> base.SplashCustomReturnType: + return _splash_attention_custom( + fwd_mask_info, + dkv_mask_info, + q, + k, + v, + segment_ids, + sinks, + mask_value=mask_value, + is_mqa=is_mqa, + save_residuals=save_residuals, + config=config, + max_logit_value=max_logit_value, + mask_function=mask_function, + fwd_mask_sparsity=fwd_mask_sparsity, + dkv_mask_sparsity=dkv_mask_sparsity, + ) + + +@jax.tree_util.register_pytree_node_class +class SplashAttentionKernel: + + def __init__( + self, + fwd_mask_info: MaskInfo, + dkv_mask_info: MaskInfo | None, + **kwargs, + ): + self.kwargs = kwargs + self.fwd_mask_info = fwd_mask_info + self.dkv_mask_info = dkv_mask_info + + def __call__(self, *args, **kwargs) -> base.SplashCustomReturnType: + return _splash_attention( + self.fwd_mask_info, + self.dkv_mask_info, + *args, + **dict(self.kwargs, **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.block_mask is not None: + block_mask_shape = self.fwd_mask_info.block_mask.shape + try: + 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 len(sharding.spec) != 1: + raise ValueError("Only q sequence sharding is supported.") + + _resolve_spec = lambda x: sharding.spec if x is not None else None + + def mask_info_spec(mask_info): + if mask_info is None: + return None + return MaskInfo( # pytype: disable=wrong-arg-types + mask_next=_resolve_spec(mask_info.mask_next), + active_rows=_resolve_spec(mask_info.active_rows), + active_cols=_resolve_spec(mask_info.active_cols), + num_active_blocks=_resolve_spec(mask_info.num_active_blocks), + block_mask=_resolve_spec(mask_info.block_mask), + partial_mask_blocks=jax.sharding.PartitionSpec() # replicated + if mask_info.partial_mask_blocks is not None + else None, + q_sequence=_resolve_spec(mask_info.q_sequence), + ) + + return SplashAttentionKernel( + mask_info_spec(self.fwd_mask_info), + mask_info_spec(self.dkv_mask_info), + **self.kwargs, + ) + + def tree_flatten(self): + return ((self.fwd_mask_info, self.dkv_mask_info), self.kwargs) + + @classmethod + def tree_unflatten(cls, kwargs, values): + fwd_mask_info, dkv_mask_info = values + # NamedTuples are not preserved during pytree serialization. + dkv_mask_info = MaskInfo(*dkv_mask_info) if dkv_mask_info is not None else None + return SplashAttentionKernel(MaskInfo(*fwd_mask_info), dkv_mask_info, **kwargs) + + +def _make_splash_attention( + mask: np.ndarray | mask_lib.Mask, + *, + config: SplashConfig | None = None, + is_mqa: bool, + save_residuals: bool = False, + mask_value: float = base.DEFAULT_MASK_VALUE, + downcast_smem_data: bool = True, + partial_mask_blocks_dtype: jax.typing.DTypeLike = np.int8, + q_seq_shards: int, +): + if len(mask.shape) != 2: + raise ValueError(f"Unexpected mask shape: {mask.shape}") + + if isinstance(mask, np.ndarray): + mask = mask_lib.NumpyMask(mask) + + if config is None: + config = SplashConfig.get_default() + + process_fn = partial( + mask_info_lib.process_mask, + downcast_smem_data=downcast_smem_data, + partial_mask_blocks_dtype=partial_mask_blocks_dtype, + q_seq_shards=q_seq_shards, + ) + + fwd_mask_info, mask_function_fwd = process_fn( + mask, + (config.block_q, config.block_kv), + ) + fwd_mask_sparsity = float(np.mean(fwd_mask_info.block_mask != 0)) + fwd_mask_info = tree_util.tree_map(jnp.array, fwd_mask_info) + + dkv_mask_info = None + if config.has_backward_blocks: + bq_dkv, bkv_dkv = config.block_q_dkv, config.block_kv_dkv + dkv_mask_info, mask_function_dkv = process_fn( + mask, + (bq_dkv, bkv_dkv), + is_dkv=True, + return_dynamic_grid=config.dq_reduction_steps == 3, + ) + + assert (mask_function_fwd is None) == (mask_function_dkv is None) + + dkv_mask_sparsity = float(np.mean(dkv_mask_info.block_mask != 0)) + dkv_mask_info = tree_util.tree_map(jnp.array, dkv_mask_info) + else: + dkv_mask_sparsity = 1.0 + + return SplashAttentionKernel( + fwd_mask_info, + dkv_mask_info, + config=config, + is_mqa=is_mqa, + save_residuals=save_residuals, + mask_value=mask_value, + mask_function=mask_function_fwd, + fwd_mask_sparsity=fwd_mask_sparsity, + dkv_mask_sparsity=dkv_mask_sparsity, + ) + + +def _make_dynamic_splash_attention( + mask: jax.Array, + *, + mesh: jax.sharding.Mesh | None = None, + mask_spec: jax.sharding.PartitionSpec | None = None, + config: SplashConfig | None = None, + is_mqa: bool, + save_residuals: bool = False, + mask_value: float = base.DEFAULT_MASK_VALUE, + downcast_smem_data: bool = True, + partial_mask_blocks_dtype: jax.typing.DTypeLike = np.int8, +): + if (mesh is not None) != (mask_spec is not None): + raise ValueError("Either both or neither of mesh and mask_spec must be specified.") + + if mask_spec is not None and len(mask_spec) != 1: + raise ValueError("Only shard over the query sequence dimension.") + + if len(mask.shape) != 2: + raise ValueError(f"Unexpected mask shape: {mask.shape}") + + if config is None: + config = SplashConfig.get_default() + + # This is the only mode that supports the dynamic grid. + config = dataclasses.replace(config, dq_reduction_steps=3) + + def process_mask_shard(mask): + process_mask_fn = functools.partial( + mask_info_lib._process_dynamic_mask, + downcast_smem_data=downcast_smem_data, + partial_mask_blocks_dtype=partial_mask_blocks_dtype, + ) + + fwd_mask_info = process_mask_fn(mask, (config.block_q, config.block_kv), is_dkv=False) + + dkv_mask_info = None + if config.has_backward_blocks: + dkv_mask_info = process_mask_fn(mask, (config.block_q_dkv, config.block_kv_dkv), is_dkv=True) + + return fwd_mask_info, dkv_mask_info + + kwargs = dict( + config=config, + is_mqa=is_mqa, + save_residuals=save_residuals, + mask_value=mask_value, + mask_function=None, + fwd_mask_sparsity=1.0, + dkv_mask_sparsity=1.0, + ) + + # If the input mask is replicated we don't need to call shard_map. + if mask_spec is None: + fwd_mask_info, dkv_mask_info = process_mask_shard(mask) + kernel = SplashAttentionKernel(fwd_mask_info, dkv_mask_info, **kwargs) + return kernel + + mask_info_specs = MaskInfo( # pytype: disable=wrong-arg-types + mask_next=mask_spec, + active_rows=None, + active_cols=None, + num_active_blocks=None, + block_mask=mask_spec, + partial_mask_blocks=mask_spec, + q_sequence=None, + ) + out_specs = ( + mask_info_specs, + mask_info_specs if config.has_backward_blocks else None, + ) + + @partial( + jax.shard_map, + mesh=mesh, + in_specs=mask_spec, + out_specs=out_specs, + check_vma=False, + ) + def process_all_shards(mask): + return process_mask_shard(mask) + + fwd_mask_info, dkv_mask_info = process_all_shards(mask) + kernel = SplashAttentionKernel(fwd_mask_info, dkv_mask_info, **kwargs) + kernel_spec = SplashAttentionKernel(*out_specs, **kwargs) + + return (kernel, kernel_spec) + + +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, q_seq_shards=1) + +make_splash_mqa_single_device = partial(make_splash_mqa, q_seq_shards=1) + +make_dynamic_splash_mqa = partial(_make_dynamic_splash_attention, is_mqa=True) +make_dynamic_splash_mha = partial(_make_dynamic_splash_attention, is_mqa=False) diff --git a/src/maxtext/kernels/tokamax_splash_attention/splash_attention_mask.py b/src/maxtext/kernels/tokamax_splash_attention/splash_attention_mask.py new file mode 100644 index 0000000000..536d875e62 --- /dev/null +++ b/src/maxtext/kernels/tokamax_splash_attention/splash_attention_mask.py @@ -0,0 +1,506 @@ +# pylint: skip-file +# Copyright 2025 DeepMind Technologies Limited. All Rights Reserved. +# +# 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. +# ============================================================================== + +"""Mini-mask creation library.""" + +from collections.abc import Callable +import dataclasses +from typing import Any, Self + +import numpy as np + +# mypy: ignore-errors + + +class Mask: + """A base class for splash attention masks.""" + + @property + def shape(self) -> tuple[int, ...]: + raise NotImplementedError + + def __getitem__(self, idx) -> np.ndarray: + raise NotImplementedError + + def __bool__(self) -> bool: + raise NotImplementedError( + "Conversion to bool is unsupported. Could be caused by using logical" " instead of bitwise operations on masks." + ) + + def __or__(self, other: Self) -> Self: + if self.shape != other.shape: + raise ValueError(f"Invalid shape for other: {other.shape}, expected: {self.shape}") + return LogicalOr(self, other) + + def __and__(self, other: Self) -> Self: + if self.shape != other.shape: + raise ValueError(f"Invalid shape for other: {other.shape}, expected: {self.shape}") + return LogicalAnd(self, other) + + +def make_causal_mask(shape: tuple[int, int], offset: int = 0) -> np.ndarray: + """Makes a causal attention mask. + + Args: + shape: Shape of the 2-dim mask: (q_seq_len, kv_seq_len). + offset: Offset of q start wrt kv. A positive offset shifts the bottom + triangle upward, a negative one shifts it downward. A negative offset + makes the first 'offset' rows of the attention matrix all 0s which leads + to undefined softmax. + + Returns: + The causal mask. + """ + q_seq_len, kv_seq_len = shape + q_idx = np.arange(q_seq_len, dtype=np.int32) + kv_idx = np.arange(kv_seq_len, dtype=np.int32) + return (q_idx[:, None] + offset >= kv_idx[None, :]).astype(np.bool_) + + +def make_local_attention_mask( + shape: tuple[int, int], + window_size: tuple[int | None, int | None], + *, + offset: int = 0, +) -> np.ndarray: + """Makes a local attention mask.""" + q_seq_len, kv_seq_len = shape + q_idx = np.arange(q_seq_len, dtype=np.int32) + kv_idx = np.arange(kv_seq_len, dtype=np.int32) + mask = np.ones((q_seq_len, kv_seq_len), dtype=np.bool_) + left, right = window_size + if left is not None: + mask = mask & (q_idx[:, None] - left + offset <= kv_idx[None, :]) + if right is not None: + mask = mask & (q_idx[:, None] + right + offset >= kv_idx[None, :]) + return mask.astype(np.bool_) + + +def make_chunk_attention_mask(shape: tuple[int, int], chunk_size: int) -> np.ndarray: + """Makes a chunked causal attention mask. + + Args: + shape: The desired shape of the mask (q_seq_len, kv_seq_len). + chunk_size: The size of the attention chunks. + + Returns: + A boolean mask of shape `mask_shape` where True indicates attention is + allowed according to chunked causal rules, and False otherwise. + + Raises: + ValueError: If chunk_window_size is None or not positive. + """ + if chunk_size <= 0: + raise ValueError("chunk_size must be positive") + + q_seq_len, kv_seq_len = shape + q_idx = np.arange(q_seq_len, dtype=np.int32) + kv_idx = np.arange(kv_seq_len, dtype=np.int32) + + # chunk mask calculation + same_chunk = (q_idx[:, None] // chunk_size) == (kv_idx[None, :] // chunk_size) + mask = same_chunk & (q_idx[:, None] >= kv_idx[None, :]) + return mask + + +def make_random_mask(shape: tuple[int, int], sparsity: float, seed: int) -> np.ndarray: + """Makes a random attention mask.""" + np.random.seed(seed) + return np.random.binomial(n=1, p=1.0 - sparsity, size=shape).astype(np.bool_) + + +@dataclasses.dataclass(slots=True) +class LogicalOr(Mask): + left: Mask + right: Mask + + def __init__(self, left: Mask, right: Mask): + if left.shape != right.shape: + raise ValueError("Masks must have the same shape") + self.left = left + self.right = right + + @property + def shape(self) -> tuple[int, ...]: + return self.left.shape + + def __getitem__(self, idx) -> np.ndarray: + return self.left[idx] | self.right[idx] + + def __hash__(self): + return hash((type(self),) + (self.left, self.right)) + + +@dataclasses.dataclass(slots=True) +class LogicalAnd(Mask): + left: Mask + right: Mask + + def __init__(self, left: Mask, right: Mask): + if left.shape != right.shape: + raise ValueError("Masks must have the same shape") + self.left = left + self.right = right + + @property + def shape(self) -> tuple[int, ...]: + return self.left.shape + + def __getitem__(self, idx) -> np.ndarray: + return self.left[idx] & self.right[idx] + + def __hash__(self): + return hash((type(self),) + (self.left, self.right)) + + +class _ComputableMask(Mask): + """Superclass for all masks that can be computed inside the kernel using a callable object. + + This subclass is designed to be used with Splash Attention. + It allows the mask logic to be computed on-the-fly or fused into the attention + kernel, avoiding the memory cost of materializing the full + (sequence_length, sequence_length) boolean mask array, which can be excessive + for long sequences. + + Attributes: + _shape: Shape of the 2-dim mask: (q_seq_len, kv_seq_len). + offset: Offset of q start wrt kv. A positive offset shifts the bottom + triangle upward, a negative one shifts it downward. A negative offset + makes the first 'offset' rows of the attention matrix all 0s which leads + to undefined softmax. + q_sequence: Indices of Q sequence. q_sequence is reused across __getitem__ + calls which is important for compile-time performance. + mask_function: Function used by the SplashAttention kernel to compute the + mask rather than loading it. + """ + + _shape: tuple[int, int] + q_sequence: np.ndarray + mask_function: Callable[..., Any] + + def __init__( + self, + shape: tuple[int, int], + mask_function: Callable[..., Any], + shard_count: int = 1, + ): + self._shape = shape + self.mask_function = mask_function + q_seq_len = self.shape[0] + + if q_seq_len % (shard_count * shard_count) != 0: + raise ValueError( + f"Shard count squared ({shard_count * shard_count}) must" f" divide Q seq_len ({self.shape[0]}) evenly." + ) + + self.q_sequence = np.arange(q_seq_len, dtype=np.int32) + + @property + def shape(self) -> tuple[int, ...]: + return self._shape + + def __getitem__(self, idx) -> np.ndarray: + if len(idx) != 2: + raise NotImplementedError(f"Unsupported slice: {idx}") + + q_slice, kv_slice = idx + if not isinstance(q_slice, slice) or not isinstance(kv_slice, slice): + raise NotImplementedError(f"Unsupported slice: {idx}") + + q_slice = _fill_slice(q_slice, self.shape[0]) + kv_slice = _fill_slice(kv_slice, self.shape[1]) + + rows = self.q_sequence[q_slice] + cols = np.arange(kv_slice.start, kv_slice.stop) + + return self.mask_function(rows[:, None], cols[None, :]) + + def __eq__(self, other: object): + raise NotImplementedError() + + def __hash__(self): + raise NotImplementedError() + + +class CausalMask(_ComputableMask): + """Lazy causal mask, prevents the model from attending to future tokens. + + Attributes: + offset: Offset of q start wrt kv. A positive offset shifts the bottom + triangle upward, a negative one shifts it downward. A negative offset + makes the first 'offset' rows of the attention matrix all 0s which leads + to undefined softmax. + """ + + offset: int + + def __init__( + self, + shape: tuple[int, int], + offset: int = 0, + shard_count: int = 1, + ): + self.offset = offset + + def causal_mask_function(q_ids, kv_ids): + # When evaluating the mask in _process_mask we typically work with numpy + # array views. + # Avoid the addition when possible to avoid instantiating an actual array. + if self.offset == 0: + return q_ids >= kv_ids + else: + return q_ids + self.offset >= kv_ids + + mask_function = causal_mask_function + + super().__init__( + shape=shape, + mask_function=mask_function, + shard_count=shard_count, + ) + + def __eq__(self, other: object): + if not isinstance(other, type(self)): + return NotImplemented + + return self.shape == other.shape and self.offset == other.offset and np.array_equal(self.q_sequence, other.q_sequence) + + def __hash__(self): + return hash( + ( + type(self), + self.shape, + self.offset, + self.q_sequence.tobytes() if self.q_sequence is not None else None, + ) + ) + + +class ChunkedCausalMask(_ComputableMask): + """Lazy chunked causal mask. + + Attention is causal within each chunk (0, K), (K, 2K), (2K, 3K), ... tokens + attend to each other but not across chunks. + Llama4 models use interleaved chunk attention along with global attention. + + + Attributes: + chunk_size: The size of each attention chunk. + """ + + chunk_size: int + + def __init__( + self, + shape: tuple[int, int], + chunk_size: int, + shard_count: int = 1, + ): + if chunk_size <= 0: + raise ValueError("chunk_size must be positive") + self.chunk_size = chunk_size + + # Define the mask function for chunk attention + def chunked_causal_mask_function(q_ids, kv_ids): + """Computes the mask logic for the given slice indices.""" + # Condition 1: Same chunk + same_chunk = (q_ids // self.chunk_size) == (kv_ids // self.chunk_size) + + # Condition 2: Causal + causal = q_ids >= kv_ids + + return same_chunk & causal + + super().__init__( + shape=shape, + mask_function=chunked_causal_mask_function, + shard_count=shard_count, + ) + + def __eq__(self, other: object): + if not isinstance(other, type(self)): + return NotImplemented + + return ( + self.shape == other.shape + and self.chunk_size == other.chunk_size + and np.array_equal(self.q_sequence, other.q_sequence) + ) + + def __hash__(self): + return hash( + ( + type(self), + self.shape, + self.chunk_size, + self.q_sequence.tobytes() if self.q_sequence is not None else None, + ) + ) + + +class LocalMask(_ComputableMask): + """Lazy local mask, prevents model from attending to tokens outside window. + + Attributes: + window_size: Size of the two sides of the local window (None identifies no + limit for the given side). + offset: Offset of q start wrt kv. A positive offset shifts the bottom + triangle upward, a negative one shifts it downward. A negative offset + makes the first 'offset' rows of the attention matrix all 0s which leads + to undefined softmax. + """ + + window_size: tuple[int | None, int | None] + offset: int + + def __init__( + self, + shape: tuple[int, int], + window_size: tuple[int | None, int | None], + offset: int, + shard_count: int = 1, + ): + self.window_size = window_size + self.offset = offset + + def local_mask_function(q_ids, kv_ids): + """Computes the local attention mask for the given slice indices.""" + left_size, right_size = self.window_size + + assert q_ids.ndim == 2 + assert kv_ids.ndim == 2 + + if left_size is None and right_size is None: + return np.ones((q_ids.shape[0], kv_ids.shape[1]), dtype=np.bool_) + + # Avoid the addition when possible to avoid instantiating an actual array. + if offset != 0: + shifted_q_ids = q_ids + self.offset + else: + shifted_q_ids = q_ids + + mask = None + if left_size is not None: + mask = shifted_q_ids - left_size <= kv_ids + if right_size is not None: + if mask is None: + mask = shifted_q_ids + right_size >= kv_ids + else: + mask &= shifted_q_ids + right_size >= kv_ids + return mask + + super().__init__( + shape=shape, + mask_function=local_mask_function, + shard_count=shard_count, + ) + + def __eq__(self, other: object): + if not isinstance(other, type(self)): + return False + + return ( + self.shape == other.shape + and self.window_size == other.window_size + and self.offset == other.offset + and np.array_equal(self.q_sequence, other.q_sequence) + ) + + def __hash__(self): + return hash( + ( + type(self), + self.shape, + self.window_size, + self.offset, + self.q_sequence.tobytes() if self.q_sequence is not None else None, + ) + ) + + +@dataclasses.dataclass(slots=True) +class NumpyMask(Mask): + """A mask backed by a dense numpy array.""" + + array: np.ndarray + + def __post_init__(self): + if self.array.ndim != 2: + raise ValueError("Expected a 2-dim array") + + if self.array.dtype != np.bool_: + raise ValueError("Mask must be a boolean array") + + @property + def shape(self) -> tuple[int, ...]: + return self.array.shape + + def __getitem__(self, idx) -> np.ndarray: + return self.array[idx] + + def __eq__(self, other: object): + if not isinstance(other, type(self)): + return NotImplemented + + return np.array_equal(self.array, other.array, equal_nan=True) + + def __hash__(self): + return hash((type(self), self.array.tobytes())) + + +def _fill_slice(inp_slice: slice, size: int) -> slice: + assert inp_slice.step is None or inp_slice.step == 1 + start = 0 if inp_slice.start is None else inp_slice.start + stop = size if inp_slice.stop is None else inp_slice.stop + assert start >= 0 + assert stop <= size + return slice(start, stop, None) + + +@dataclasses.dataclass(frozen=True, slots=True) +class FullMask(Mask): + """Lazy full mask, allows all tokens to attend to all other tokens.""" + + # TODO: Transform FullMask into a _ComputableMask. + + _shape: tuple[int, int] + + def __post_init__(self): + if not isinstance(self.shape, tuple): + raise ValueError(f"Unsupported shape type: {type(self.shape)}") + + @property + def shape(self) -> tuple[int, ...]: + return self._shape + + def __getitem__(self, idx) -> np.ndarray: + if len(idx) != 2: + raise NotImplementedError(f"Unsupported slice: {idx}") + i, j = idx + if not isinstance(i, slice) or not isinstance(j, slice): + raise NotImplementedError(f"Unsupported slice: {idx}") + i = _fill_slice(i, self.shape[0]) + j = _fill_slice(j, self.shape[1]) + return np.ones((i.stop - i.start, j.stop - j.start), dtype=np.bool_) + + def __eq__(self, other: object): + if not isinstance(other, type(self)): + return NotImplemented + + return self.shape == other.shape + + def __hash__(self): + return hash((type(self), self.shape)) diff --git a/src/maxtext/kernels/tokamax_splash_attention/splash_attention_mask_info.py b/src/maxtext/kernels/tokamax_splash_attention/splash_attention_mask_info.py new file mode 100644 index 0000000000..986b380c09 --- /dev/null +++ b/src/maxtext/kernels/tokamax_splash_attention/splash_attention_mask_info.py @@ -0,0 +1,607 @@ +# pylint: skip-file +# Copyright 2025 DeepMind Technologies Limited. All Rights Reserved. +# +# 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. +# ============================================================================== + +"""Mini-mask creation library.""" + +import collections +import functools +from typing import Any, NamedTuple + +import jax +import jax.numpy as jnp +import numpy as np +from maxtext.kernels.tokamax_splash_attention import splash_attention_mask as mask_lib + +# mypy: ignore-errors + +lax = jax.lax +MaskCallable = Any + + +def find_bounds( + arr: jax.Array | np.ndarray, +) -> tuple[jax.Array | np.ndarray | None, jax.Array | np.ndarray | None]: + # Find the first and last block of a row to determine when to initialize/store + # the output. + + if arr is None: + return None, None + + bounds_start = (arr != jnp.roll(arr, shift=1, axis=-1)).astype(jnp.int32) + bounds_end = (arr != jnp.roll(arr, shift=-1, axis=-1)).astype(jnp.int32) + bounds_start = bounds_start.at[0].set(1) + bounds_end = bounds_end.at[-1].set(1) + + return bounds_start, bounds_end + + +# Logic for processing NumPy masks for kernels +class MaskInfo(NamedTuple): + """Contains runtime masking information for the Splash attention kernel. + + The arrays, mask_next and block_mask are placed in TPU + scalar-memory. This is a scarse resource so the mask creation logic attempts + to shrink the data-type of these arrays to the smallest possible one. + This can be: np.int32, np.int16 or np.int8. + + Attributes: + mask_next: An integer[num_active_blocks] NumPy array where each entry + contains the next mask block index in `partial_mask_blocks` to prefetch. + active_rows: An integer[num_active_blocks] NumPy array where each entry + contains the row index of the corresponding active block in the original + mask. + active_cols: An integer[num_active_blocks] NumPy array where each entry + contains the column index of the corresponding active block in the + original mask. + block_mask: An integer[num_active_blocks] NumPy array where each entry is + either 1 or 2. 1 means the corresponding block is full and 2 means the + corresponding block is partially masked. + num_active_blocks: An integer[] NumPy array whose entries are the sizes of + the corresponding blocks in the original mask. + partial_mask_blocks: An int8[num_partial_blocks, block_q, block_kv] NumPy + array that contains the blocks of the original mask that contained both + zeros and ones. The entries in `mask_next` point to indices in the first + axis of this array. + q_sequence: A i32[q_sequence_length] NumPy array. When using causal masking, + this contains the list of indices that correspond to q tokens. For plain + causal this is just np.arange(q_sequence_length). + """ + + mask_next: np.ndarray | jax.Array | None + active_rows: np.ndarray | jax.Array | None + active_cols: np.ndarray | jax.Array | None + block_mask: np.ndarray | jax.Array | None + num_active_blocks: np.ndarray | jax.Array | None + partial_mask_blocks: np.ndarray | jax.Array | None + q_sequence: np.ndarray | None + + +def _downcast_to_small_type(array: np.ndarray) -> np.ndarray: + """Downcast numpy array. + + If possible, downcast the data-type of the input array to the smallest numpy + type (among np.int16 and np.int8) that fits the content of the array. + + Args: + array: the array to downcast + + Returns: + The downcasted array. + + Raises: + ValueError: if the input array is not np.int32 or if its elements are not + all positive. + """ + if array.dtype != np.int32: + raise ValueError(f"Expected int32 input, but got {array.dtype}.") + + if not np.all(array >= -1): + # Allow -1 for padding. + raise ValueError("Expected non-negative array.") + + if array.size == 0: + return array + + max_value = np.max(array) + + if max_value <= np.iinfo(np.int8).max: + return array.astype(np.int8) + elif max_value <= np.iinfo(np.int16).max: + return array.astype(np.int16) + else: + return array.astype(np.int32) + + +def _check_mask(mask: mask_lib.Mask) -> None: + """Check that the given mask is valid. + + A row of all zeros along the kv dimension would result in a division by zero + when computing the softmax. This function is meant to protect against that + case. + + Args: + mask: the mask to check. + + Raises: + ValueError: the mask is invalid. + """ + + assert len(mask.shape) == 2 + + exception_message = ( + "Some rows of the mask (along the kv dimension) are all zeros.\nThis is" + " would result in a division by zero when computing the attention" + " softmax." + ) + + is_row_non_zero = np.zeros(mask.shape[0], dtype=np.bool_) + for col in range(mask.shape[1]): + # Mask only supports slice indices. + is_row_non_zero = np.logical_or( + is_row_non_zero, + mask[(slice(0, mask.shape[0]), slice(col, col + 1))][:, 0], + ) + if not is_row_non_zero.all(): + raise ValueError(exception_message) + + +class _HashableNDArray: + """Helper to make a numpy array hashable: can be added associative containers. + + Attributes: + array: The underlying numpy array. + """ + + __slots__ = ("array", "_hash") + array: np.ndarray + + def __init__(self, array: np.ndarray): + self.array = array + self._hash = hash(array.tobytes()) + + def __hash__(self): + return self._hash + + def __eq__(self, other: object) -> bool: + if not isinstance(other, _HashableNDArray): + return NotImplemented + return np.array_equal(self.array, other.array, equal_nan=True) + + +def _generate_shard_metadata( + block_mask: np.ndarray, + partial_blocks: np.ndarray, + is_dkv: bool, + return_dynamic_grid: bool, +): + if is_dkv: + block_mask = block_mask.mT + partial_blocks = partial_blocks.mT + + if return_dynamic_grid: + active_mask = block_mask > 0 + # If an entire row is masked then that output tile won't be visited. + # We extend the grid to visit these tiles to initialize them. + active_mask[:, 0] |= ~active_mask.any(axis=1) + active_indices = np.argwhere(active_mask) + active_rows = active_indices[:, 0].astype(np.int32) + active_cols = active_indices[:, 1].astype(np.int32) + block_mask = block_mask[active_mask > 0] + grid_size = active_rows.size + else: + active_indices = np.ndindex(block_mask.shape) + active_rows = active_cols = grid_size = None + + partial_coords = np.argwhere(partial_blocks != -1) + if partial_coords.size > 0: + mask_next = [] + mask_coords_iter = iter([tuple(c) for c in partial_coords]) + first_m = coord_m = next(mask_coords_iter) + + for idx in active_indices: + is_next_mask = tuple(idx) > tuple(coord_m) + if is_next_mask: + try: + coord_m = next(mask_coords_iter) # type: ignore + except StopIteration: + coord_m = first_m + mask_next.append(partial_blocks[coord_m]) + else: + mask_next = np.full(block_mask.size, -1, dtype=np.int32) + + mask_next = np.array(mask_next, dtype=np.int32) + flat_block_mask = block_mask.flatten() + + return active_rows, active_cols, mask_next, flat_block_mask, grid_size + + +def _causal_state_grid( + mask: mask_lib.CausalMask, + q_block_size: int, + kv_block_size: int, +) -> np.ndarray: + """Returns block states for a causal mask without materializing chunks.""" + q_seq_len, kv_seq_len = mask.shape + q_blocks_count = q_seq_len // q_block_size + kv_blocks_count = kv_seq_len // kv_block_size + q_block_min = np.arange(q_blocks_count, dtype=np.int32) * q_block_size + q_block_max = q_block_min + q_block_size - 1 + kv_block_min = np.arange(kv_blocks_count, dtype=np.int32) * kv_block_size + kv_block_max = kv_block_min + kv_block_size - 1 + + empty = q_block_max[:, None] + mask.offset < kv_block_min[None, :] + full = q_block_min[:, None] + mask.offset >= kv_block_max[None, :] + return np.where(empty, 0, np.where(full, 2, 1)).astype(np.int32) + + +def _process_dynamic_mask( + mask: jax.Array, + block_shape: tuple[int, int], + is_dkv: bool, + *, + downcast_smem_data: bool = True, + partial_mask_blocks_dtype: jax.typing.DTypeLike = np.int8, +) -> MaskInfo: + """Process a dynamic mask to compute it's local sparsity data. + + Note that this operates on a single shard of the mask. + + Args: + mask: [q_seq_len, kv_seq_len] jax.Array representing a dense mask to + process. + block_shape: A Tuple[int, int] representing the shape of the Pallas grid + block. + is_dkv: True if we are processing the dKV mask + downcast_smem_data: If True, downcast the scalar-memory data of MaskInfo to + a data type smaller than np.int32 (if possible). + + Returns: + `MaskInfo`, a sparse representation of the dense mask. + + Raises: + ValueError: if the input mask is invalid or the block sizes are not + compatible with the mask sizes. + """ + if len(mask.shape) != 2: + raise ValueError(f"Expected a 2-dim mask, instead got: {mask.shape}.") + + q_seq_len, kv_seq_len = mask.shape + q_block_size, kv_block_size = block_shape + q_blocks_count, q_mod = divmod(q_seq_len, q_block_size) + kv_blocks_count, kv_mod = divmod(kv_seq_len, kv_block_size) + + if q_mod != 0: + raise ValueError(f"{q_block_size=} should divide {q_seq_len=}.") + if kv_mod != 0: + raise ValueError(f"{kv_block_size=} should divide {kv_seq_len=}.") + + # Tile the last 2 dimensions of the mask into 2D tiles of size `block_shape`. + mask_blocks = ( + mask.reshape( + q_blocks_count, + q_block_size, + kv_blocks_count, + kv_block_size, + ) + .swapaxes(-2, -3) + .astype(partial_mask_blocks_dtype) + ) + + any_mask = jnp.any(mask_blocks, axis=(-1, -2)).astype(np.int32) + all_mask = jnp.all(mask_blocks, axis=(-1, -2)).astype(np.int32) + block_mask = any_mask + all_mask + + block_ids = jnp.arange(block_mask.size, dtype=np.int32).reshape(block_mask.shape) + if is_dkv: + block_mask = block_mask.swapaxes(-1, -2) + block_ids = block_ids.swapaxes(-1, -2) + mask_blocks = mask_blocks.swapaxes(-1, -2) + + active_mask = block_mask > 0 + # If an entire row is masked then that output tile won't be visited. + # We extend the grid to visit these tiles to initialize them. + empty_rows = jnp.all(block_mask == 0, axis=-1) + first_col = jnp.arange(block_mask.shape[1]) == 0 + active_mask |= empty_rows[:, None] & first_col + + num_active_blocks = active_mask.flatten().sum(keepdims=True) + active_indices = jnp.argwhere(active_mask, size=active_mask.size, fill_value=-1) + active_rows = active_indices[:, 0].astype(np.int32) + active_cols = active_indices[:, 1].astype(np.int32) + + block_mask = block_mask[active_rows, active_cols] + mask_next = block_ids.at[active_rows, active_cols].get(wrap_negative_indices=False) + mask_next = jnp.where(block_mask == 1, mask_next, 0) + + # Mask out the blocks that aren't active. + mask = (jnp.arange(block_mask.size) < num_active_blocks).astype(np.int32) + block_mask = block_mask * mask + + # Collapsing because the block ids are linearized. + mask_blocks = lax.collapse(mask_blocks, 0, 2) + + def _downcast(array: jax.Array, max_value: int) -> jax.Array: + if array.size == 0: + return array + + if array.dtype != np.int32: + raise ValueError(f"Expected int32 input, but got {array.dtype}.") + + if max_value <= np.iinfo(np.int8).max: + return array.astype(np.int8) + elif max_value <= np.iinfo(np.int16).max: + return array.astype(np.int16) + else: + return array.astype(np.int32) + + if downcast_smem_data: + block_mask = block_mask.astype(np.int8) # values are in the range [0, 1, 2] + mask_next = _downcast(mask_next, q_blocks_count * kv_blocks_count) + + return MaskInfo( + mask_next=mask_next, + active_rows=active_rows, + active_cols=active_cols, + block_mask=block_mask, + num_active_blocks=num_active_blocks, + partial_mask_blocks=mask_blocks, + q_sequence=None, + ) + + +# When used in a transformer network with multiple layers, the SplashAttention +# kernel is created several times with the same mask. Cache MaskInfo to avoid +# blowing up compile times. Ideally the size of the cache should be determined +# by the client. +@functools.lru_cache(maxsize=12) +def _process_mask( + mask: mask_lib.Mask, # [q_seq_len, kv_seq_len] + block_shape: tuple[int, int], + is_dkv: bool, + *, + downcast_smem_data: bool = True, + partial_mask_blocks_dtype: jax.typing.DTypeLike = np.int8, + q_seq_shards: int = 1, + kv_seq_shards: int = 1, + return_dynamic_grid: bool = True, +) -> tuple[MaskInfo, MaskCallable | None]: + """Transform a dense mask into a sparse representation. + + The number Q sequence shards are needed to create a MaskInfo + object that is partitionable (with shard_map) along that dimension. + Args: + mask: Dense mask to process. + block_shape: Shape of the Pallas grid block. + is_dkv: True if we are processing the dKV mask + downcast_smem_data: If True, downcast the SMEM data of MaskInfo to a data + type smaller if possible. + q_seq_shards: Number of Q sequence shards of the mesh in which the kernel is + launched. + + Returns: + `MaskInfo`, a sparse representation of the dense mask. + `MaskCallable`: a callable that, given Q and KV indices, returns + the value of the mask at those coordinates. + + Raises: + ValueError: if the input mask is invalid or the block sizes are not + compatible with the mask sizes. + """ + + if len(mask.shape) != 2: + raise ValueError(f"Expected a 2-dim mask, instead got: {mask.shape=}") + + q_seq_len, kv_seq_len = mask.shape + q_block_size, kv_block_size = block_shape + q_blocks_count, q_mod = divmod(q_seq_len, q_block_size) + kv_blocks_count, kv_mod = divmod(kv_seq_len, kv_block_size) + + if q_mod != 0: + raise ValueError(f"{q_block_size=} should divide {q_seq_len=}.") + if kv_mod != 0: + raise ValueError(f"{kv_block_size=} should divide {kv_seq_len=}.") + + q_seq_len_per_shard, mod = divmod(q_seq_len, q_seq_shards) + if mod != 0: + raise ValueError(f"{q_seq_shards=} should divide {q_seq_len=}.") + + q_blocks_per_shard, mod = divmod(q_seq_len_per_shard, q_block_size) + if mod != 0: + raise ValueError(f"{q_block_size=} should divide {q_seq_len_per_shard=}.") + + kv_seq_len_per_shard, mod = divmod(kv_seq_len, kv_seq_shards) + if mod != 0: + raise ValueError(f"{kv_seq_shards=} should divide {kv_seq_len=}.") + + kv_blocks_per_shard, mod = divmod(kv_seq_len_per_shard, kv_block_size) + if mod != 0: + raise ValueError(f"{kv_block_size=} should divide {kv_seq_len_per_shard=}.") + + # TODO: checking the validity of the masks is slow for large masks. + # Disable it for now, reevaluate in the future. + + # The mask object either define q_sequence and mask_function or none of + # them. + assert hasattr(mask, "q_sequence") == hasattr(mask, "mask_function") + + # If the mask object defines a q_sequence and a mask_function, then make use + # of these in the kernel rather. This is preferable over loading the mask + # from memory. When using a mask_function, then mask_next and + # partial_mask_blocks are left undefined and not used in the kernel. + if hasattr(mask, "q_sequence") and hasattr(mask, "mask_function"): + q_sequence = mask.q_sequence + mask_function = mask.mask_function + else: + q_sequence = mask_function = None + + if isinstance(mask, mask_lib.CausalMask): + state_grid = _causal_state_grid(mask, q_block_size, kv_block_size) + partial_id_grid = np.full((q_blocks_count, kv_blocks_count), -1, dtype=np.int32) + partial_mask_blocks = None + full_mask = (state_grid == 2).all() + if full_mask: + return ( + MaskInfo( + mask_next=None, + active_rows=None, + active_cols=None, + block_mask=None, + num_active_blocks=None, + partial_mask_blocks=None, + q_sequence=None, + ), + None, + ) + else: + # Identify the partial mask blocks and the value of the block mask for each + # block. + # Partial mask blocks are uniquified. When partitioning, all partial mask + # blocks are replicated across shards. + + blocked_shape = (q_blocks_count, kv_blocks_count) + state_grid = np.zeros(blocked_shape, dtype=np.int32) + partial_id_grid = np.full(blocked_shape, -1, dtype=np.int32) + + partial_blocks_map = collections.defaultdict(lambda: len(partial_blocks_map)) + unique_chunks = [] + + # Partition the dense mask into blocks and categorize them: + # 0 = Empty, 1 = Partial (mixed 0s and 1s), 2 = Full (all 1s). + # Partial blocks are deduplicated and stored in unique_chunks to save memory. + for coords in np.ndindex((q_blocks_count, kv_blocks_count)): + (q_idx, kv_idx) = coords + chunk = mask[ + ( + slice(q_idx * q_block_size, (q_idx + 1) * q_block_size), + slice(kv_idx * kv_block_size, (kv_idx + 1) * kv_block_size), + ) + ] + if chunk.any(): + if chunk.all(): + state_grid[q_idx, kv_idx] = 2 + else: + state_grid[q_idx, kv_idx] = 1 + chunk_id = partial_blocks_map[_HashableNDArray(chunk)] + partial_id_grid[q_idx, kv_idx] = chunk_id + + if chunk_id == len(unique_chunks): + unique_chunks.append(chunk) + + full_mask = (state_grid == 2).all() + if full_mask: + return ( + MaskInfo( + mask_next=None, + active_rows=None, + active_cols=None, + block_mask=None, + num_active_blocks=None, + partial_mask_blocks=None, + q_sequence=None, + ), + None, + ) + + if unique_chunks: + partial_mask_blocks = np.stack(unique_chunks).astype(partial_mask_blocks_dtype) + if is_dkv: + partial_mask_blocks = partial_mask_blocks.mT + else: + partial_mask_blocks = None + + # Work on a fraction of the mask at the time to compute the mask. This is + # needed to compute the correct data indices, which are relative to the + # current slice of the mask. + all_shards_metadata = [] + for q_shard_idx in range(q_seq_shards): + for kv_shard_idx in range(kv_seq_shards): + q_slice = slice( + q_shard_idx * q_blocks_per_shard, + (q_shard_idx + 1) * q_blocks_per_shard, + ) + kv_slice = slice( + kv_shard_idx * kv_blocks_per_shard, + (kv_shard_idx + 1) * kv_blocks_per_shard, + ) + metadata = _generate_shard_metadata( + state_grid[q_slice, kv_slice], + partial_id_grid[q_slice, kv_slice], + is_dkv, + return_dynamic_grid, + ) + all_shards_metadata.append(metadata) + + ( + active_rows_slices, + active_cols_slices, + mask_next_slices, + block_mask_slices, + num_active_blocks, + ) = zip(*all_shards_metadata) + + if return_dynamic_grid: + # Pad each slice to the largest number of active blocks in any shard. + max_size = max(num_active_blocks) + pad_slice = lambda arr: np.pad(arr, (0, max_size - arr.shape[0]), mode="constant", constant_values=-1) + active_rows_slices = list(map(pad_slice, active_rows_slices)) + active_cols_slices = list(map(pad_slice, active_cols_slices)) + mask_next_slices = list(map(pad_slice, mask_next_slices)) + block_mask_slices = list(map(pad_slice, block_mask_slices)) + + # Concatenate the sequence shards. + active_rows = np.concatenate(active_rows_slices, axis=0) + active_cols = np.concatenate(active_cols_slices, axis=0) + num_active_blocks = np.array(num_active_blocks, dtype=np.int32) + + if downcast_smem_data: + active_rows = _downcast_to_small_type(active_rows) + active_cols = _downcast_to_small_type(active_cols) + else: + active_rows = active_cols = num_active_blocks = None + + mask_next = np.concatenate(mask_next_slices, axis=0) + block_mask = np.concatenate(block_mask_slices, axis=0) + + if downcast_smem_data: + mask_next = _downcast_to_small_type(mask_next) + block_mask = _downcast_to_small_type(block_mask) + + if partial_mask_blocks is None: + mask_next = None + + assert (mask_function is not None) == (q_sequence is not None) + # When the mask can be computed inside the kernel with a mask_function, + # there is no need to load it from memory. So mask_next and + # partial_mask_blocks are unused. + return ( + MaskInfo( + mask_next=mask_next if mask_function is None else None, + active_rows=active_rows, + active_cols=active_cols, + block_mask=block_mask, + num_active_blocks=num_active_blocks, + partial_mask_blocks=partial_mask_blocks if mask_function is None else None, + q_sequence=q_sequence, + ), + mask_function, + ) + + +process_mask = functools.partial(_process_mask, is_dkv=False) +process_mask_dkv = functools.partial(_process_mask, is_dkv=True) + +process_dynamic_mask = functools.partial(_process_dynamic_mask, is_dkv=False) +process_dynamic_mask_dkv = functools.partial(_process_dynamic_mask, is_dkv=True) diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 1c015b464b..b3f22fc336 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -66,16 +66,17 @@ ) from maxtext.inference.kvcache import KVQuant, KVTensor from maxtext.kernels.attention import jax_flash_attention +from maxtext.kernels.attention import tokamax_ring_attention from maxtext.kernels.attention.ragged_attention import ragged_gqa from maxtext.kernels.attention.ragged_attention import ragged_mha +from maxtext.kernels.tokamax_splash_attention import splash_attention_kernel as tokamax_splash_kernel +from maxtext.kernels.tokamax_splash_attention import splash_attention_mask as tokamax_splash_mask from maxtext.layers import nnx_wrappers from maxtext.layers.initializers import variable_to_logically_partitioned from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_utils from maxtext.utils.sharding import logical_to_mesh_axes, maybe_shard_with_pspec import numpy as np -from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_kernel as tokamax_splash_kernel -from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_mask as tokamax_splash_mask # pylint: disable=line-too-long, g-doc-args, g-doc-return-or-yield, bad-continuation, g-inconsistent-quotes # pytype: disable=attribute-error @@ -519,6 +520,51 @@ def __init__( self.use_ragged_attention = use_ragged_attention self.ragged_block_size = ragged_block_size self.rngs = rngs + if self.attention_kernel == "flash" and tokamax_ring_attention.is_context_parallel_ring_requested(self.config): + target_hardware = self.mesh.devices[(0,) * self.mesh.devices.ndim].platform + if target_hardware == "tpu": + if not self.config.use_tokamax_splash: + raise ValueError("TPU Tokamax ring attention requires use_tokamax_splash=True.") + if self.config.use_jax_splash: + raise ValueError("TPU Tokamax ring attention requires use_jax_splash=False.") + if self.config.context_parallel_load_balance: + raise ValueError("TPU Tokamax ring attention does not support context_parallel_load_balance yet.") + if self.config.packing: + raise ValueError("TPU Tokamax ring attention does not support packing yet.") + if self.attention_type != AttentionType.GLOBAL: + raise ValueError("TPU Tokamax ring attention is initially supported only for global causal attention.") + + context_axis = self.config.context_sharding + axis_names_q = self._logical_to_mesh_axes(self.flash_axis_names_q) + axis_names_kv = self._logical_to_mesh_axes(self.flash_axis_names_kv) + axis_names_kv = tokamax_ring_attention.with_sequence_axis( + axis_names_kv, + context_axis, + sequence_dim=2, + ) + tokamax_ring_attention.validate_ring_mesh_axis( + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + sequence_dim_q=2, + sequence_dim_kv=2, + mesh=self.mesh, + ring_axis=context_axis, + ) + tokamax_ring_attention.validate_head_sharding( + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + mesh=self.mesh, + num_query_heads=self.num_query_heads, + num_kv_heads=self.num_kv_heads, + head_dim_q=1, + head_dim_kv=1, + ) + tokamax_ring_attention.validate_dkv_sharding( + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + dkv_dim_q=3, + dkv_dim_kv=3, + ) def maybe_create_nnx(einsum, *args): if isinstance(einsum, nn.Module): @@ -947,6 +993,28 @@ def _generate_moba_mask(self, query: Array, key: Array, q_positions: Array) -> A moba_mask = jax.vmap(self.generate_moba_mask_single_item, in_axes=(0, 0, None))(query, key, q_positions) return moba_mask + def _validate_tpu_tokamax_ring_runtime( + self, + *, + model_mode: str, + previous_chunk: Any = None, + bidirectional_mask: Any = None, + sinks: Array | None = None, + indexer_mask: Array | None = None, + use_ragged_attention: bool = False, + record_max_logits: bool = False, + ) -> None: + """Validates runtime constraints for the TPU Tokamax ring path.""" + tokamax_ring_attention.validate_tokamax_ring_runtime( + model_mode=model_mode, + previous_chunk=previous_chunk, + sinks=sinks, + indexer_mask=indexer_mask, + use_ragged_attention=use_ragged_attention, + bidirectional_mask=bidirectional_mask, + record_max_logits=record_max_logits, + ) + def apply_attention( self, query: Array, @@ -971,6 +1039,12 @@ def apply_attention( self.check_attention_inputs(query, key, value) length = query.shape[-3] target_hardware = self.mesh.devices[(0,) * self.mesh.devices.ndim].platform + if ( + target_hardware == "tpu" + and tokamax_ring_attention.is_context_parallel_ring_requested(self.config) + and self.attention_kernel != "flash" + ): + raise ValueError("TPU Tokamax ring attention requires attention_kernel='flash'.") if use_ragged_attention and model_mode == MODEL_MODE_AUTOREGRESSIVE: if lengths is None: @@ -1033,6 +1107,11 @@ def apply_attention( decoder_segment_ids, self.attn_logits_soft_cap, sinks, + indexer_mask=indexer_mask, + model_mode=model_mode, + previous_chunk=previous_chunk, + bidirectional_mask=bidirectional_mask, + use_ragged_attention=use_ragged_attention, record_max_logits=record_max_logits, ) if max_logits is not None: @@ -1214,12 +1293,27 @@ def tpu_flash_attention( attn_logits_soft_cap: float | None = None, sinks: Array | None = None, indexer_mask: Array | None = None, + model_mode: str = MODEL_MODE_TRAIN, + previous_chunk: Any = None, + bidirectional_mask: Any = None, + use_ragged_attention: bool = False, record_max_logits: bool = False, ) -> tuple[Array, Array]: """TPU Flash Attention.""" + use_tokamax_ring = tokamax_ring_attention.is_context_parallel_ring_requested(self.config) cp_size = self.mesh.shape.get(self.config.context_sharding, 1) load_balanced_context_parallel = self.config.context_parallel_load_balance + if use_tokamax_ring: + self._validate_tpu_tokamax_ring_runtime( + model_mode=model_mode, + previous_chunk=previous_chunk, + bidirectional_mask=bidirectional_mask, + sinks=sinks, + indexer_mask=indexer_mask, + use_ragged_attention=use_ragged_attention, + record_max_logits=record_max_logits, + ) # Transpose to ('batch', 'heads', 'length', 'kv') query = jnp.transpose(query, axes=(0, 2, 1, 3)) @@ -1236,6 +1330,23 @@ def tpu_flash_attention( axis_names_q = self._logical_to_mesh_axes(self.flash_axis_names_q) axis_names_kv = self._logical_to_mesh_axes(self.flash_axis_names_kv) indexer_mask_axis_names = self._logical_to_mesh_axes((BATCH_ATTN, Q_LENGTH, KV_LENGTH)) + if use_tokamax_ring: + context_axis = self.config.context_sharding + segment_axis_names_q = tokamax_ring_attention.with_sequence_axis( + segment_axis_names_q, + context_axis, + sequence_dim=1, + ) + axis_names_kv = tokamax_ring_attention.with_sequence_axis( + axis_names_kv, + context_axis, + sequence_dim=2, + ) + segment_axis_names_kv = tokamax_ring_attention.with_sequence_axis( + segment_axis_names_kv, + context_axis, + sequence_dim=1, + ) devices_in_data_fsdp = self.mesh.shape.get("data", 1) * self.mesh.shape.get("fsdp", 1) assert (query.shape[0] / devices_in_data_fsdp).is_integer(), ( @@ -1296,54 +1407,67 @@ def create_sa_config(config, query, key, attn_logits_soft_cap): ) return sa_config - sa_config = create_sa_config(self.config, query, key, attn_logits_soft_cap) - mask_shape = (query.shape[2], key.shape[2]) # (q_seq_len, kv_seq_len) - mask_module = tokamax_splash_mask if self.config.use_tokamax_splash else splash_attention_mask - if self.attention_type == AttentionType.FULL: - mask = mask_module.FullMask(mask_shape) + if use_tokamax_ring: + sa_config, splash_kernel, segment_axis_names_splash_kernel = ( + tokamax_ring_attention.make_sharded_ring_attention_kernel( + self.config, + query=query, + key=key, + context_parallel_size=cp_size, + ring_axis=self.config.context_sharding, + attn_logits_soft_cap=attn_logits_soft_cap, + maybe_shard_with_pspec=self._maybe_shard_with_pspec, + ) + ) else: - mask = mask_module.CausalMask(shape=mask_shape) - - use_load_balanced_cp = cp_size > 1 and load_balanced_context_parallel - if use_load_balanced_cp and self.attention_type != AttentionType.FULL: - mask = LoadBalancedCausalMask(shape=mask_shape, cp_size=cp_size) - - # Apply local masking if local sliding attention is enabled. - if self.attention_type == AttentionType.LOCAL_SLIDING: - if self.sliding_window_size is None: - raise ValueError("Sliding_window_size must be set if Local Sliding attention type") - local_window_size = (self.sliding_window_size - 1, self.sliding_window_size) - if use_load_balanced_cp: - mask &= LoadBalancedLocalMask( - shape=(query.shape[2], key.shape[2]), - window_size=local_window_size, - offset=0, - cp_size=cp_size, - ) + sa_config = create_sa_config(self.config, query, key, attn_logits_soft_cap) + mask_shape = (query.shape[2], key.shape[2]) # (q_seq_len, kv_seq_len) + mask_module = tokamax_splash_mask if self.config.use_tokamax_splash else splash_attention_mask + if self.attention_type == AttentionType.FULL: + mask = mask_module.FullMask(mask_shape) else: - mask &= mask_module.LocalMask( - shape=(query.shape[2], key.shape[2]), - window_size=local_window_size, - offset=0, - ) - elif self.attention_type == AttentionType.CHUNK: - if self.chunk_attn_window_size is None: - raise ValueError("chunk_attn_window_size must be set for chunk attention type") + mask = mask_module.CausalMask(shape=mask_shape) - if use_load_balanced_cp: - mask &= LoadBalancedChunkedCausalMask( - shape=(query.shape[2], key.shape[2]), - chunk_size=self.chunk_attn_window_size, - cp_size=cp_size, - ) - else: - mask &= ChunkedCausalMask( - shape=(query.shape[2], key.shape[2]), - chunk_size=self.chunk_attn_window_size, - ) + use_load_balanced_cp = cp_size > 1 and load_balanced_context_parallel + if use_load_balanced_cp and self.attention_type != AttentionType.FULL: + mask = LoadBalancedCausalMask(shape=mask_shape, cp_size=cp_size) + + # Apply local masking if local sliding attention is enabled. + if self.attention_type == AttentionType.LOCAL_SLIDING: + if self.sliding_window_size is None: + raise ValueError("Sliding_window_size must be set if Local Sliding attention type") + local_window_size = (self.sliding_window_size - 1, self.sliding_window_size) + if use_load_balanced_cp: + mask &= LoadBalancedLocalMask( + shape=(query.shape[2], key.shape[2]), + window_size=local_window_size, + offset=0, + cp_size=cp_size, + ) + else: + mask &= mask_module.LocalMask( + shape=(query.shape[2], key.shape[2]), + window_size=local_window_size, + offset=0, + ) + elif self.attention_type == AttentionType.CHUNK: + if self.chunk_attn_window_size is None: + raise ValueError("chunk_attn_window_size must be set for chunk attention type") + + if use_load_balanced_cp: + mask &= LoadBalancedChunkedCausalMask( + shape=(query.shape[2], key.shape[2]), + chunk_size=self.chunk_attn_window_size, + cp_size=cp_size, + ) + else: + mask &= ChunkedCausalMask( + shape=(query.shape[2], key.shape[2]), + chunk_size=self.chunk_attn_window_size, + ) max_logit_value = None - if self.config.use_tokamax_splash: + if not use_tokamax_ring and self.config.use_tokamax_splash: # Create mask single_head_mask = mask # tokamax now just uses a single mask and assumes broadcast to all heads if self.config.use_max_logit_estimate > 0: @@ -1356,7 +1480,7 @@ def create_sa_config(config, query, key, attn_logits_soft_cap): "single_head_mask", ], ) - def wrap_splash_kernel(single_head_mask): + def wrap_tokamax_splash_kernel(single_head_mask): splash_kernel = tokamax_splash_kernel.make_splash_mha( mask=single_head_mask, config=sa_config, @@ -1364,14 +1488,14 @@ def wrap_splash_kernel(single_head_mask): ) return splash_kernel - splash_kernel = wrap_splash_kernel(single_head_mask) + splash_kernel = wrap_tokamax_splash_kernel(single_head_mask) segment_axis_names_splash_kernel = self._logical_to_mesh_axes((Q_LENGTH,)) splash_kernel = self._maybe_shard_with_pspec(splash_kernel, segment_axis_names_splash_kernel) - elif self.config.use_jax_splash: + elif not use_tokamax_ring and self.config.use_jax_splash: if self.config.use_max_logit_estimate > 0: sa_config = dataclasses.replace(sa_config, max_logit_const=self.config.use_max_logit_estimate) segment_axis_names_splash_kernel = nn.logical_to_mesh_axes((Q_LENGTH,)) - else: + elif not use_tokamax_ring: # Create multi-head mask multi_head_mask = splash_attention_mask.MultiHeadMask(masks=(mask,) * query.shape[1]) @@ -1383,7 +1507,7 @@ def wrap_splash_kernel(single_head_mask): "shard_head_size", ], ) - def wrap_splash_kernel(multi_head_mask, shard_head_size=1): + def wrap_jax_splash_kernel(multi_head_mask, shard_head_size=1): splash_kernel = splash_attention_kernel.make_splash_mha( mask=multi_head_mask, head_shards=shard_head_size, # the size of the axis if sharding over heads @@ -1397,7 +1521,7 @@ def wrap_splash_kernel(multi_head_mask, shard_head_size=1): head_physical_axes = logical_to_mesh_axes((HEAD,), self.mesh)[0] head_physical_axes = (head_physical_axes,) if isinstance(head_physical_axes, str) else (head_physical_axes or ()) shard_head_size = math.prod(self.mesh.shape.get(ax, 1) for ax in head_physical_axes) - splash_kernel = wrap_splash_kernel(multi_head_mask, shard_head_size) + splash_kernel = wrap_jax_splash_kernel(multi_head_mask, shard_head_size) named_sharding = jax.sharding.NamedSharding(self.mesh, axis_names_splash_kernel) segment_axis_names_splash_kernel = splash_kernel.manual_sharding_spec(named_sharding) splash_kernel = jax.tree.map( @@ -1407,14 +1531,10 @@ def wrap_splash_kernel(multi_head_mask, shard_head_size=1): is_leaf=lambda x: x is None, ) - # Now call the function wrap_flash_attention which does the actual computation. - # The splash kernel is passed as a parameter to the function. Since we have the shard map - # decorating the wrap_flash_attention function, the data will be correctly sharded - # meaning q will be sharded over sequence aka context length but K and V will be duplicated - # The shardings are specified in the in_specs and out_specs of the shard_map decorator: - # 'segment_axis_names_q' maps to ['activation_q_length', ['context']] meaning that q is sharded over the context axis - # 'segment_axis_names_kv' maps to ['activation_kv_length', []] meaning that K and V are not sharded - # splash_kernel is sharded over (HEAD, LENGTH) + # wrap_flash_attention receives Q/K/V and mask metadata with the shardings + # specified in the shard_map in_specs below. For the all-gather path Q is + # sequence-sharded and K/V are replicated. For the Tokamax ring path Q, K, + # V, and segment IDs are all sequence-sharded over the context axis. if record_max_logits: # max_logits will share similar sharding as query but last dim is unrelated to model @@ -1462,11 +1582,19 @@ def wrap_flash_attention( sinks, indexer_mask, ): - # If load_balanced_context_parallel is enabled, reorder the key and value tensors - # to ensure that they are contiguous in memory. - # This is necessary for the splash attention kernel to work correctly because it expects - # the K and V to be contiguous. Note that K and V are not sharded over the sequence aka context axis - # This was we get the unsharded unpermuted key and value tensors + if use_tokamax_ring: + attention_output = tokamax_ring_attention.call_ring_attention( + query, + key, + value, + decoder_segment_ids_q, + decoder_segment_ids_kv, + splash_kernel, + ) + return attention_output, None + + # The load-balanced all-gather path restores K/V to contiguous order + # before calling Splash attention. if cp_size > 1 and load_balanced_context_parallel: key = max_utils.reorder_sequence(tensor=key, cp_size=cp_size, seq_dim=2, to_contiguous=True) value = max_utils.reorder_sequence(tensor=value, cp_size=cp_size, seq_dim=2, to_contiguous=True) diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 2841933aa9..95a40345f0 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -23,6 +23,7 @@ from absl.testing import parameterized from flax import nnx +from flax.linen import partitioning as nn_partitioning import jax import jax.numpy as jnp from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask @@ -1139,6 +1140,172 @@ def test_tpu_flash_attention_packed_all_gather_context_parallel(self, context_pa f"are not close. context_parallel_load_balance={context_parallel_load_balance}.", ) + @pytest.mark.tpu_only + def test_tpu_flash_attention_ring_context_parallel(self): + """Test equivalence between dot_product and flash attention + ring context parallelism""" + + cfg_cp = pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + **self.config_arguments, + attention="flash", + context_parallel_strategy="ring", + context_parallel_load_balance=False, + ici_context_parallelism=2, + use_tokamax_splash=True, + use_jax_splash=False, + packing=False, + dtype="float32", + ) + devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp) + mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes) + lnx, decoder_segment_ids, decoder_positions = self.get_data(cfg_cp.dtype) + attention_as_mha_generic = Attention( + config=self.cfg, + num_query_heads=cfg_cp.num_query_heads, + num_kv_heads=cfg_cp.num_kv_heads, + head_dim=cfg_cp.head_dim, + max_target_length=cfg_cp.max_target_length, + max_prefill_predict_length=cfg_cp.max_prefill_predict_length, + inputs_q_shape=lnx.shape, + inputs_kv_shape=lnx.shape, + mesh=self.mesh, + attention_kernel="dot_product", + dtype=cfg_cp.dtype, + dropout_rate=cfg_cp.dropout_rate, + rngs=self.nnx_rng, + ) + mha_generic_output, _ = attention_as_mha_generic( + lnx, + lnx, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + generic_state = nnx.state(attention_as_mha_generic) + + attention_as_mha_flash_cp = Attention( + config=cfg_cp, + num_query_heads=cfg_cp.num_query_heads, + num_kv_heads=cfg_cp.num_kv_heads, + head_dim=cfg_cp.head_dim, + max_target_length=cfg_cp.max_target_length, + max_prefill_predict_length=cfg_cp.max_prefill_predict_length, + inputs_q_shape=lnx.shape, + inputs_kv_shape=lnx.shape, + mesh=mesh_cp, + attention_kernel="flash", + dtype=cfg_cp.dtype, + dropout_rate=cfg_cp.dropout_rate, + model_mode=MODEL_MODE_PREFILL, + rngs=self.nnx_rng, + ) + nnx.update(attention_as_mha_flash_cp, generic_state) + + mha_generic_flash_cp_output = attention_test_util.forward_with_context_expert_parallelism( + cfg_cp, + mesh_cp, + attention_as_mha_flash_cp, + lnx, + decoder_segment_ids, + decoder_positions, + ) + + mha_generic_output = jax.device_get(mha_generic_output) + mha_generic_flash_cp_output = jax.device_get(mha_generic_flash_cp_output) + + self.assertTrue( + jax.numpy.allclose(mha_generic_output, mha_generic_flash_cp_output, rtol=1e-02, atol=1e-02, equal_nan=False), + msg="Logits from generic dot product and flash attention + ring context parallelism are not close.", + ) + + @pytest.mark.tpu_only + def test_tpu_flash_attention_ring_context_parallel_grad(self): + """Test gradient equivalence between dot_product and flash attention + ring context parallelism""" + + cfg_cp = pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + **self.config_arguments, + attention="flash", + context_parallel_strategy="ring", + context_parallel_load_balance=False, + ici_context_parallelism=2, + use_tokamax_splash=True, + use_jax_splash=False, + packing=False, + dtype="float32", + ) + devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp) + mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes) + lnx, decoder_segment_ids, decoder_positions = self.get_data(cfg_cp.dtype) + attention_as_mha_generic = Attention( + config=self.cfg, + num_query_heads=cfg_cp.num_query_heads, + num_kv_heads=cfg_cp.num_kv_heads, + head_dim=cfg_cp.head_dim, + max_target_length=cfg_cp.max_target_length, + max_prefill_predict_length=cfg_cp.max_prefill_predict_length, + inputs_q_shape=lnx.shape, + inputs_kv_shape=lnx.shape, + mesh=self.mesh, + attention_kernel="dot_product", + dtype=cfg_cp.dtype, + dropout_rate=cfg_cp.dropout_rate, + rngs=self.nnx_rng, + ) + generic_state = nnx.state(attention_as_mha_generic) + + attention_as_mha_flash_cp = Attention( + config=cfg_cp, + num_query_heads=cfg_cp.num_query_heads, + num_kv_heads=cfg_cp.num_kv_heads, + head_dim=cfg_cp.head_dim, + max_target_length=cfg_cp.max_target_length, + max_prefill_predict_length=cfg_cp.max_prefill_predict_length, + inputs_q_shape=lnx.shape, + inputs_kv_shape=lnx.shape, + mesh=mesh_cp, + attention_kernel="flash", + dtype=cfg_cp.dtype, + dropout_rate=cfg_cp.dropout_rate, + model_mode=MODEL_MODE_PREFILL, + rngs=self.nnx_rng, + ) + nnx.update(attention_as_mha_flash_cp, generic_state) + + def generic_loss(lnx): + output, _ = attention_as_mha_generic( + lnx, + lnx, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return jnp.mean(output.astype(jnp.float32) ** 2) + + def ring_loss(lnx): + output, _ = attention_as_mha_flash_cp( + lnx, + lnx, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return jnp.mean(output.astype(jnp.float32) ** 2) + + generic_grad = jax.grad(generic_loss)(lnx) + with jax.set_mesh(mesh_cp), nn_partitioning.axis_rules(cfg_cp.logical_axis_rules): + ring_grad = jax.grad(ring_loss)(lnx) + generic_grad = jax.device_get(generic_grad) + ring_grad = jax.device_get(ring_grad) + + self.assertTrue( + jax.numpy.allclose(generic_grad, ring_grad, rtol=1e-02, atol=1e-07, equal_nan=False), + msg="Input gradients from generic dot product and flash attention + ring context parallelism are not close.", + ) + @pytest.mark.tpu_only def test_dot_product_cache_axis_order(self): all_axis_orders = tuple(itertools.permutations(range(4))) diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index 4d61b4caaf..f0504e4dd8 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -96,6 +96,80 @@ def test_validation_error(self): with self.assertRaises(pydantic.ValidationError): pyconfig.initialize(argv) + def test_tpu_tokamax_ring_config_validation_accepts_initial_config(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "attention=flash", + "use_tokamax_splash=True", + "use_jax_splash=False", + "context_parallel_strategy=ring", + "context_parallel_load_balance=False", + "ici_context_parallelism=2", + "hardware=tpu", + "packing=False", + "dataset_type=synthetic", + "skip_jax_distributed_system=True", + ] + mock_devices = [unittest.mock.MagicMock(slice_index=0) for _ in range(8)] + with unittest.mock.patch("jax.devices", return_value=mock_devices): + config = pyconfig.initialize(argv) + + self.assertEqual(config.context_parallel_strategy, "ring") + self.assertEqual(config.ici_context_parallelism, 2) + self.assertEqual(config.attention, "flash") + self.assertTrue(config.use_tokamax_splash) + + def test_tpu_tokamax_ring_config_validation_rejects_unsupported_configs(self): + base_args = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "attention=flash", + "use_tokamax_splash=True", + "use_jax_splash=False", + "context_parallel_strategy=ring", + "context_parallel_load_balance=False", + "ici_context_parallelism=2", + "hardware=tpu", + "packing=False", + "dataset_type=synthetic", + "skip_jax_distributed_system=True", + ] + cases = [ + (["ici_context_parallelism=1"], ["ici_context_parallelism=2"], "context_parallel_size > 1"), + (["context_sharding=expert", "ici_expert_parallelism=2"], [], "context_sharding"), + (["dq_reduction_steps=2"], [], "dq_reduction_steps"), + (["max_target_length=2050"], [], "context_parallel_size squared"), + (["attention=dot_product"], ["attention=flash"], "attention=flash"), + (["use_tokamax_splash=False"], ["use_tokamax_splash=True"], "use_tokamax_splash"), + (["use_jax_splash=True"], ["use_jax_splash=False"], "use_jax_splash"), + (["attention_type=full"], [], "global causal"), + (["packing=True"], ["packing=False"], "packing"), + ( + ["context_parallel_load_balance=True"], + ["context_parallel_load_balance=False"], + "context_parallel_load_balance", + ), + (["use_ragged_attention=True"], [], "ragged attention"), + (["attention_sink=True"], [], "attention sinks"), + (["use_indexer=True", "q_lora_rank=1"], [], "sparse indexer"), + (["use_chunked_prefill=True"], [], "chunked prefill"), + (["moba=True"], [], "MoBA"), + (["use_multimodal=True"], [], "multimodal"), + (["use_qk_clip=True"], [], "QK-Clip"), + (["dropout_rate=0.1"], [], "dropout"), + ] + mock_devices = [unittest.mock.MagicMock(slice_index=0) for _ in range(8)] + for bad_args, args_to_remove, expected_regex in cases: + with self.subTest(bad_args=bad_args): + argv = [arg for arg in base_args if arg not in args_to_remove] + argv.extend(bad_args) + with unittest.mock.patch("jax.devices", return_value=mock_devices): + with self.assertRaisesRegex((ValueError, pydantic.ValidationError), expected_regex): + pyconfig.initialize(argv) + def test_load_balanced_chunk_context_parallel_config(self): argv = [ "", diff --git a/tests/unit/tokamax_ring_attention_test.py b/tests/unit/tokamax_ring_attention_test.py new file mode 100644 index 0000000000..47ddefd9ce --- /dev/null +++ b/tests/unit/tokamax_ring_attention_test.py @@ -0,0 +1,115 @@ +# Copyright 2026 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 +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for MaxText's Tokamax ring attention adapter.""" + +# pylint: disable=protected-access + +from __future__ import annotations + +import types +from unittest import mock + +from absl.testing import absltest +import jax +import jax.numpy as jnp + +from maxtext.kernels.attention import tokamax_ring_attention + + +class TokamaxRingAttentionTest(absltest.TestCase): + + def test_is_context_parallel_ring_requested_accepts_case_insensitive_strategy(self): + config = types.SimpleNamespace(context_parallel_strategy="Ring") + + self.assertTrue(tokamax_ring_attention.is_context_parallel_ring_requested(config)) + + def test_make_causal_mask_uses_local_tokamax_causal_mask(self): + mask = tokamax_ring_attention._make_causal_mask((16, 16), 4) + + self.assertEqual(mask.shape, (16, 16)) + self.assertEqual(mask.q_sequence.tolist(), list(range(16))) + + def test_validate_ring_mesh_axis_requires_key_value_sequence_sharding(self): + mesh = types.SimpleNamespace(shape={"context": 4}) + + with self.assertRaisesRegex(ValueError, "K/V sequence sharding"): + tokamax_ring_attention.validate_ring_mesh_axis( + axis_names_q=(None, None, "context", None), + axis_names_kv=(None, None, None, None), + sequence_dim_q=2, + sequence_dim_kv=2, + mesh=mesh, + ring_axis="context", + ) + + def test_validate_head_sharding_rejects_mismatched_gqa_axes(self): + mesh = types.SimpleNamespace(shape={"tensor": 2}) + + with self.assertRaisesRegex(ValueError, "Q and KV head sharding"): + tokamax_ring_attention.validate_head_sharding( + axis_names_q=(None, "tensor", "context", None), + axis_names_kv=(None, None, "context", None), + mesh=mesh, + num_query_heads=8, + num_kv_heads=4, + head_dim_q=1, + head_dim_kv=1, + ) + + def test_call_ring_attention_uses_ring_kernel_segment_ids(self): + captured = {} + + class RingSegmentIds: + + def __init__(self, q, kv): + self.q = q + self.kv = kv + + def kernel(q, k, v, segment_ids): + captured["segment_ids_type"] = type(segment_ids) + captured["q_segment_shape"] = segment_ids.q.shape + captured["kv_segment_shape"] = segment_ids.kv.shape + return q + k + v + + query = jnp.ones((1, 2, 4, 2)) + key = jnp.ones((1, 2, 4, 2)) + value = jnp.ones((1, 2, 4, 2)) + segment_ids = jnp.ones((1, 4), dtype=jnp.int32) + + with mock.patch.object(tokamax_ring_attention.ring_attention_kernel, "SegmentIds", RingSegmentIds): + out = tokamax_ring_attention.call_ring_attention( + query, + key, + value, + segment_ids, + segment_ids, + kernel, + ) + + self.assertEqual(out.shape, query.shape) + self.assertIs(captured["segment_ids_type"], RingSegmentIds) + self.assertEqual(captured["q_segment_shape"], (4,)) + self.assertEqual(captured["kv_segment_shape"], (4,)) + + def test_with_sequence_axis_preserves_partition_spec_type(self): + spec = jax.sharding.PartitionSpec("data", None, None, "model") + + out = tokamax_ring_attention.with_sequence_axis(spec, "context", sequence_dim=2) + + self.assertIsInstance(out, jax.sharding.PartitionSpec) + self.assertEqual(tuple(out), ("data", None, "context", "model")) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/unit/tokamax_splash_attention_mask_test.py b/tests/unit/tokamax_splash_attention_mask_test.py new file mode 100644 index 0000000000..590e8046b8 --- /dev/null +++ b/tests/unit/tokamax_splash_attention_mask_test.py @@ -0,0 +1,61 @@ +# Copyright 2026 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 +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for local Tokamax SplashAttention masks.""" + +# pylint: disable=protected-access + +from absl.testing import absltest +import numpy as np + +from maxtext.kernels.tokamax_splash_attention import splash_attention_mask +from maxtext.kernels.tokamax_splash_attention import splash_attention_mask_info + + +class TokamaxSplashAttentionMaskTest(absltest.TestCase): + + def test_causal_state_grid_uses_block_extents(self): + mask = splash_attention_mask.CausalMask((8, 8), shard_count=2) + + got = splash_attention_mask_info._causal_state_grid(mask, q_block_size=2, kv_block_size=2) + + expected = np.array( + [ + [1, 0, 0, 0], + [2, 1, 0, 0], + [2, 2, 1, 0], + [2, 2, 2, 1], + ], + dtype=np.int32, + ) + np.testing.assert_array_equal(got, expected) + + def test_process_causal_mask_keeps_lazy_mask_function(self): + mask = splash_attention_mask.CausalMask((8, 8), shard_count=2) + + mask_info, mask_function = splash_attention_mask_info.process_mask( + mask, + block_shape=(2, 2), + q_seq_shards=2, + kv_seq_shards=2, + return_dynamic_grid=True, + ) + + self.assertIs(mask_function, mask.mask_function) + self.assertIsNotNone(mask_info.q_sequence) + self.assertIsNone(mask_info.partial_mask_blocks) + np.testing.assert_array_equal(mask_info.q_sequence, np.arange(8, dtype=np.int32)) + + +if __name__ == "__main__": + absltest.main()