Support general mask_fn via ComputableMask

GitOrigin-RevId: b907e3a

Author: ds-hwang

axlearn/common/attention_bias.py

@@ -20,7 +20,6 @@
 import functools
 import typing
 from typing import (
-    Callable,
     Generic,
     Iterable,
     Optional,
@@ -717,28 +716,30 @@ def __eq__(self, other):
         return type(other) is type(self)
 
 
-def _composite_masks(op: Callable[[Tensor, Tensor], Tensor], *mask_fns: ConfigOr[MaskFn]):
-    if len(mask_fns) == 0:
-        raise RuntimeError(f"Input must not be empty: {mask_fns}")
+def or_masks(*mask_fns: ConfigOr[MaskFn]) -> MaskFn:
+    """Returns a MaskFn that's the union of provided MaskFn's."""
 
     def mask(query_position: Tensor, key_position: Tensor):
         fns = [maybe_instantiate(arg) for arg in mask_fns]
         result = fns[0](query_position, key_position)
         for mask in fns[1:]:
-            result = op(result, mask(query_position, key_position))
+            result = result | mask(query_position, key_position)
         return result
 
     return mask
 
 
-def or_masks(*mask_fns: ConfigOr[MaskFn]) -> MaskFn:
-    """Returns a MaskFn that's the union of provided MaskFn's."""
-    return _composite_masks(jnp.logical_or, *mask_fns)
-
-
 def and_masks(*mask_fns: ConfigOr[MaskFn]) -> MaskFn:
     """Returns a MaskFn that's the intersection of provided MaskFn's."""
-    return _composite_masks(jnp.logical_and, *mask_fns)
+
+    def mask(query_position: Tensor, key_position: Tensor):
+        fns = [maybe_instantiate(arg) for arg in mask_fns]
+        result = fns[0](query_position, key_position)
+        for mask in fns[1:]:
+            result = result & mask(query_position, key_position)
+        return result
+
+    return mask
 
 
 def sliding_window_causal_mask(sliding_window_size: int) -> MaskFn:

axlearn/common/flash_attention/splash_attention_mask.py

@@ -0,0 +1,75 @@
+"""Splash attention mask.
+
+This code is adapted from the jax_ml/jax library, specifically from the
+https://github.com/jax-ml/jax/blob/9bcfac6542a330b77f29d5cc5dcf4a57f55b2947/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py
+TODO(dhwang2): Delete this file once JAX is upgraded and LocalMask becomes a computable mask.
+"""
+
+from typing import Callable
+
+import numpy as np
+from jax.experimental.pallas.ops.tpu.splash_attention.splash_attention_mask import _ComputableMask
+
+from axlearn.common.utils import Tensor
+
+
+class ComputableMask(_ComputableMask):
+    """Computable mask for splash attention that supports custom mask functions.
+
+    This mask accepts any Jax/Numpy exchangeable mask function following the MaskFn protocol
+    from attention_bias.py, such as causal_mask or sliding_window_causal_mask.
+
+    Attributes:
+      mask_fn: A callable mask function that takes query_position and key_position
+        tensors and returns a boolean mask tensor.
+    """
+
+    mask_fn: Callable[[Tensor, Tensor], Tensor]
+
+    def __init__(
+        self,
+        shape: tuple[int, int],
+        mask_fn: Callable[[Tensor, Tensor], Tensor],
+        shard_count: int = 1,
+    ):
+        """Initialize ComputableMask.
+
+        Args:
+            shape: The shape of the attention mask (q_len, kv_len).
+            mask_fn: A callable that implements the MaskFn protocol from attention_bias.py.
+                Takes (query_position, key_position) and returns a boolean mask.
+            shard_count: Number of shards.
+        """
+        self.mask_fn = mask_fn
+
+        def mask_function(q_ids, kv_ids):
+            """Computes the attention mask using the provided mask_fn."""
+            assert q_ids.ndim == 2
+            assert kv_ids.ndim == 2
+            return self.mask_fn(q_ids, kv_ids)
+
+        super().__init__(
+            shape=shape,
+            mask_function=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.mask_fn == other.mask_fn
+            and np.array_equal(self.q_sequence, other.q_sequence)
+        )
+
+    def __hash__(self):
+        return hash(
+            (
+                type(self),
+                self.shape,
+                id(self.mask_fn),
+                self.q_sequence.tobytes() if self.q_sequence is not None else None,
+            )
+        )

axlearn/common/flash_attention/splash_attention_mask_test.py

@@ -0,0 +1,98 @@
+"""Tests for splash attention mask.
+
+This code is adapted from the jax_ml/jax library, specifically from the
+https://github.com/jax-ml/jax/blob/9bcfac6542a330b77f29d5cc5dcf4a57f55b2947/tests/pallas/tpu_splash_attention_mask_test.py
+TODO(dhwang2): Delete this file once JAX is upgraded and LocalMask becomes a computable mask.
+"""
+
+import numpy as np
+from absl.testing import absltest, parameterized
+from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask as mask_lib
+
+from axlearn.common.attention_bias import causal_mask, sliding_window_causal_mask
+from axlearn.common.flash_attention.splash_attention_mask import ComputableMask
+from axlearn.common.test_utils import TestCase
+
+
+class SplashAttentionMaskTest(TestCase):
+    @parameterized.parameters(
+        [
+            ((256, 256), (1024, 1024)),
+            ((256, 128), (1024, 1024)),
+            ((128, 256), (1024, 1024)),
+            ((256, 256), (1024, 2048)),
+            ((256, 128), (1024, 2048)),
+            ((128, 256), (1024, 2048)),
+            ((256, 256), (2048, 1024)),
+            ((256, 128), (2048, 1024)),
+            ((128, 256), (2048, 1024)),
+        ]
+    )
+    def test_causal_mask_fn(self, block_size, shape):
+        """Test ComputableMask with causal_mask function from attention_bias.py."""
+        # Create expected dense causal mask
+        q_len, kv_len = shape
+        q_ids = np.arange(q_len)[:, None]
+        kv_ids = np.arange(kv_len)[None, :]
+        dense_mask = causal_mask(q_ids, kv_ids)
+
+        # Create ComputableMask with mask_fn
+        lazy_mask = ComputableMask(shape=shape, mask_fn=causal_mask)
+
+        self._compare_masks(dense_mask, lazy_mask, block_size)
+
+    @parameterized.parameters(
+        [
+            ((256, 256), (1024, 1024), 128),
+            ((256, 128), (1024, 1024), 128),
+            ((128, 256), (1024, 1024), 128),
+            ((256, 256), (1024, 2048), 256),
+            ((256, 128), (1024, 2048), 256),
+            ((128, 256), (1024, 2048), 256),
+            ((256, 256), (2048, 1024), 512),
+            ((256, 128), (2048, 1024), 512),
+            ((128, 256), (2048, 1024), 512),
+        ]
+    )
+    def test_sliding_window_causal_mask_fn(self, block_size, shape, window_size):
+        """Test ComputableMask with sliding_window_causal_mask from attention_bias.py."""
+        # Create expected dense sliding window causal mask
+        q_len, kv_len = shape
+        q_ids = np.arange(q_len)[:, None]
+        kv_ids = np.arange(kv_len)[None, :]
+        mask_fn = sliding_window_causal_mask(sliding_window_size=window_size)
+        dense_mask = mask_fn(q_ids, kv_ids)
+
+        # Create ComputableMask with mask_fn
+        lazy_mask = ComputableMask(shape=shape, mask_fn=mask_fn)
+
+        self._compare_masks(dense_mask, lazy_mask, block_size)
+
+    def _compare_masks(
+        self,
+        dense_mask: np.ndarray,
+        lazy_mask: mask_lib.Mask,
+        block_size: tuple[int, int],
+    ):
+        self.assertEqual(dense_mask.shape, lazy_mask.shape)
+
+        *prefix, width, height = dense_mask.shape
+
+        assert width % block_size[0] == 0
+        assert height % block_size[1] == 0
+
+        full_lazy_mask = lazy_mask[(*[slice(p) for p in prefix], slice(None), slice(None))]
+        self.assertNestedEqual(dense_mask, full_lazy_mask)
+        for i, j in np.ndindex(width // block_size[0], height // block_size[1]):
+            indexer = (
+                *[slice(p) for p in prefix],
+                slice(i * block_size[0], (i + 1) * block_size[0]),
+                slice(j * block_size[1], (j + 1) * block_size[1]),
+            )
+            dense_chunk = dense_mask[indexer]
+            lazy_chunk = lazy_mask[indexer]
+            self.assertNestedEqual(dense_chunk, lazy_chunk)
+
+
+if __name__ == "__main__":
+    absltest.main()

axlearn/common/flash_attention/tpu_attention.py

@@ -8,7 +8,6 @@
 import jax
 import jax.ad_checkpoint
 import jax.numpy as jnp
-import numpy as np
 from jax import lax
 from jax._src.mesh import thread_resources
 from jax.experimental import pallas as pl
@@ -48,6 +47,7 @@
     repeat_kv_heads,
 )
 from axlearn.common.flash_attention.remat import FLASH_ATTN_RESIDUAL_NAME
+from axlearn.common.flash_attention.splash_attention_mask import ComputableMask
 from axlearn.common.flash_attention.types import FlashAttentionWithShardMapSpecs
 from axlearn.common.kv_cache.base_kv_cache import BaseKVCache
 from axlearn.common.utils import Nested, Tensor
@@ -72,21 +72,7 @@ def _to_splash_mask(
         return splash_attention_mask.LocalMask(
             shape=mask_shape, window_size=(left_size, 0), offset=0, shard_count=q_seq_shards
         )
-
-    # Because mask.mask() may use jnp ops. e.g. jnp.logical_and.
-    with jax.ensure_compile_time_eval():
-        # This code is reached only when `kv_cache_type=None` (i.e., forward and prefill) and
-        # `target_len == source_len` (i.e., self-attention) (see `check_tpu_splash_attention`).
-        # `target_positions` and `source_positions` are always in the range [0, seq_len].
-        target_positions = np.arange(mask_shape[0])[None, :, None]
-        source_positions = np.arange(mask_shape[1])[None, None, :]
-        # `mask.mask` expects rank 3 tensors.
-        mask_array = np.asarray(mask.mask(target_positions, source_positions))
-        mask_array = np.squeeze(mask_array, axis=0)
-
-    # NumpyMask is backed by a dense [target_len, source_len] numpy array.
-    # May consume a large amount of host memory for long sequences at compile time.
-    return splash_attention_mask.NumpyMask(array=mask_array)
+    return ComputableMask(shape=mask_shape, shard_count=q_seq_shards, mask_fn=mask.mask)
 
 
 # The following code is adapted from jax-ml/jax:

axlearn/common/flash_attention/tpu_attention_test.py

@@ -29,6 +29,7 @@
 from axlearn.common.flash_attention import tpu_attention
 from axlearn.common.flash_attention.common import ReferenceMHA
 from axlearn.common.flash_attention.layer import default_mha_dim_to_partition_spec
+from axlearn.common.flash_attention.splash_attention_mask import ComputableMask
 from axlearn.common.flash_attention.test_utils import generate_attention_data
 from axlearn.common.test_utils import TestCase, Tolerance, is_supported_mesh_shape
 from axlearn.common.utils import Tensor
@@ -56,6 +57,9 @@ def mask(query_position: Tensor, key_position: Tensor):
     return fun
 
 
+_singleton_mask_fn = jax_fn_mask(5)
+
+
 @skipIfGPU
 class TestFlashAttention(TestCase):
     """Tests FlashAttention layer."""
@@ -107,13 +111,11 @@ def test_sliding_window_mask_equivalence(self, seq_len, sliding_window_size):
         ],
         [
             MaskFnAttentionBias(
-                jax_fn_mask(5),
+                _singleton_mask_fn,
                 target_positions=jnp.arange(8)[None],
                 source_positions=jnp.arange(8)[None],
             ),
-            splash_attention_mask.NumpyMask(
-                array=np.array(jax_fn_mask(5)(jnp.arange(8)[:, None], jnp.arange(8)[None, :]))
-            ),
+            ComputableMask(shape=(8, 8), mask_fn=_singleton_mask_fn),
         ],
     )
     def test_to_splash_mask(self, mask, expected):
@@ -126,6 +128,81 @@ def inside_tracing(mask):
 
         inside_tracing(mask)
 
+    @parameterized.product(sliding_window_size=(None, 16), seq_len=(128, 2048))
+    def test_computable_mask(self, sliding_window_size, seq_len):
+        """Test that ComputableMask with mask_fn produces same results as equivalent splash mask."""
+        batch_size = 2
+        num_heads = 4
+        per_head_dim = 64
+        num_kv_heads = num_heads // 2
+
+        # Generate test data
+        q, k, v, _ = generate_attention_data(
+            batch_size,
+            seq_len,
+            seq_len,
+            num_heads,
+            per_head_dim,
+            num_kv_heads,
+            mask_fn=None,
+            dtype=jnp.bfloat16,
+        )
+
+        cfg = dict(
+            interpret=jax.default_backend() == "cpu",
+            softmax_scale=per_head_dim**-0.5,
+            tpu_block_size=128,
+        )
+        fn = tpu_attention.TPUSplashAttention.default_config().set(**cfg).instantiate()
+
+        # Test with ComputableMask.
+        if sliding_window_size is not None:
+            mask_fn = sliding_window_causal_mask(sliding_window_size=sliding_window_size)
+            bias_with_ref_mask = SlidingWindowAttentionBias(
+                mask=mask_fn,
+                sliding_window_size=sliding_window_size,
+                target_positions=jnp.arange(seq_len)[None],
+                source_positions=jnp.arange(seq_len)[None],
+            )
+        else:
+            mask_fn = causal_mask
+            bias_with_ref_mask = CausalAttentionBias(
+                target_positions=jnp.arange(seq_len)[None],
+                source_positions=jnp.arange(seq_len)[None],
+            )
+
+        prng_key = jax.random.PRNGKey(42)
+        # Outputs by splash mask
+        input_batch_ref = dict(
+            query=q,
+            key=k,
+            value=v,
+            bias=bias_with_ref_mask,
+            prng_key=prng_key,
+            logit_sink=None,
+        )
+        out_ref = fn(input_batch_ref)
+
+        # Outputs by computable mask
+        bias_with_computable_mask = MaskFnAttentionBias(
+            mask=mask_fn,
+            target_positions=jnp.arange(seq_len)[None],
+            source_positions=jnp.arange(seq_len)[None],
+        )
+        input_batch_computable = dict(
+            query=q,
+            key=k,
+            value=v,
+            bias=bias_with_computable_mask,
+            prng_key=prng_key,
+            logit_sink=None,
+        )
+        self.assertTrue(fn.is_supported(input_batch=input_batch_computable, kv_cache_type=None))
+        out_computable = fn(input_batch_computable)
+
+        # Both should produce the same output
+        self.assertNestedAllClose(out_computable, out_ref)
+
     @parameterized.product(
         batch_size=[2],
         kv_len=[128],