Repository: jax-ml/jax
Pull request #27904 (2025-04-10 04:23:24)

Add `ChunkedCausalMask` for Splash Attention to support attention masking similar to Llama4.
Add `ChunkedCausalMask` for Splash Attention to support attention masking similar to Llama4.
Llama4 uses (interleaved) chunk attention to support long context.



Patch:
diff --git a/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py b/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py
index eab2a695d..e43f30e77 100644
--- a/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py
+++ b/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask.py
@@ -92,6 +92,35 @@ def make_local_attention_mask(
   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:
@@ -196,15 +225,20 @@ class MultiHeadMask(Mask):
 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.
+    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.
   """
@@ -314,6 +348,66 @@ class CausalMask(_ComputableMask):
     ))
 
 
+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 accross 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(Mask):
   """Lazy local mask, prevents model from attending to tokens outside window.
 


Test patch:
diff --git a/tests/pallas/tpu_splash_attention_mask_test.py b/tests/pallas/tpu_splash_attention_mask_test.py
index 5379eb109..7c4b53529 100644
--- a/tests/pallas/tpu_splash_attention_mask_test.py
+++ b/tests/pallas/tpu_splash_attention_mask_test.py
@@ -44,6 +44,15 @@ def _make_local_attention_mask(*args, **kwargs):
   return mask_lib.make_local_attention_mask(*args, **kwargs)
 
 
+def _make_lazy_chunked_causal_mask(shape, chunk_size):
+  mask = mask_lib.ChunkedCausalMask(shape=shape, chunk_size=chunk_size)
+  return mask[:, :]
+
+
+def _make_chunked_causal_mask(shape, chunk_size):
+  return mask_lib.make_chunk_attention_mask(shape=shape, chunk_size=chunk_size)
+
+
 class SplashAttentionMaskTest(jtu.JaxTestCase):
 
   @parameterized.parameters([_make_lazy_causal_mask, _make_causal_mask])
@@ -412,6 +421,181 @@ class SplashAttentionMaskTest(jtu.JaxTestCase):
         block_size,
     )
 
+  @parameterized.parameters(
+      [_make_lazy_chunked_causal_mask, _make_chunked_causal_mask]
+  )
+  def test_chunked_causal_mask(self, make_chunked_mask):
+    """Tests the chunked causal mask logic for various shapes and chunk sizes."""
+    with self.subTest("unit"):
+      expected = np.array([[1]], dtype=np.bool_)
+      actual = make_chunked_mask(shape=(1, 1), chunk_size=1)
+      self.assertArraysEqual(actual, expected)
+      actual = make_chunked_mask(shape=(1, 1), chunk_size=2)
+      self.assertArraysEqual(actual, expected)
+
+    with self.subTest("square_exact_chunks"):
+      # Chunk 0: [0, 1], Chunk 1: [2, 3]
+      expected = np.array(
+          [
+              [1, 0, 0, 0],
+              [1, 1, 0, 0],
+              [0, 0, 1, 0],
+              [0, 0, 1, 1],
+          ],
+          dtype=np.bool_,
+      )
+      actual = make_chunked_mask(shape=(4, 4), chunk_size=2)
+      self.assertArraysEqual(actual, expected)
+
+    with self.subTest("square_uneven_chunks"):
+      expected = np.array(
+          [
+              [1, 0, 0, 0, 0],
+              [1, 1, 0, 0, 0],
+              [1, 1, 1, 0, 0],
+              [0, 0, 0, 1, 0],
+              [0, 0, 0, 1, 1],
+          ],
+          dtype=np.bool_,
+      )
+      actual = make_chunked_mask(shape=(5, 5), chunk_size=3)
+      self.assertArraysEqual(actual, expected)
+
+    with self.subTest("wide_rectangle"):
+      expected = np.array(
+          [
+              [1, 0, 0, 0, 0, 0],
+              [1, 1, 0, 0, 0, 0],
+              [1, 1, 1, 0, 0, 0],
+              [0, 0, 0, 1, 0, 0],
+          ],
+          dtype=np.bool_,
+      )
+      actual = make_chunked_mask(shape=(4, 6), chunk_size=3)
+      self.assertArraysEqual(actual, expected)
+
+    with self.subTest("tall_rectangle"):
+      expected = np.array(
+          [
+              [1, 0, 0, 0],
+              [1, 1, 0, 0],
+              [1, 1, 1, 0],
+              [0, 0, 0, 1],
+              [0, 0, 0, 1],
+              [0, 0, 0, 1],
+          ],
+          dtype=np.bool_,
+      )
+      actual = make_chunked_mask(shape=(6, 4), chunk_size=3)
+      self.assertArraysEqual(actual, expected)
+
+    with self.subTest("chunk_size_1"):
+      # Should only allow self-attention q==k and chunk_size == 1
+      expected = np.array(
+          [
+              [1, 0, 0, 0],
+              [0, 1, 0, 0],
+              [0, 0, 1, 0],
+              [0, 0, 0, 1],
+          ],
+          dtype=np.bool_,
+      )
+      actual = make_chunked_mask(shape=(4, 4), chunk_size=1)
+      self.assertArraysEqual(actual, expected)
+
+    with self.subTest("chunk_size_greater_equal_seqlen"):
+      # Should behave like a normal causal mask
+      expected = np.array(
+          [
+              [1, 0, 0, 0],
+              [1, 1, 0, 0],
+              [1, 1, 1, 0],
+              [1, 1, 1, 1],
+          ],
+          dtype=np.bool_,
+      )
+      # Test chunk_size == seqlen
+      actual_eq = make_chunked_mask(shape=(4, 4), chunk_size=4)
+      self.assertArraysEqual(actual_eq, expected)
+      # Test chunk_size > seqlen
+      actual_gt = make_chunked_mask(shape=(4, 4), chunk_size=5)
+      self.assertArraysEqual(actual_gt, expected)
+
+  @parameterized.product(
+      block_size=[(128, 128), (256, 128), (128, 256)],
+      shape=[(512, 512), (512, 1024), (1024, 512)],
+      chunk_size=[64, 128, 256, 512, 1024],
+  )
+  def test_lazy_chunked_causal_mask_chunking(
+      self,
+      block_size: tuple[int, int],
+      shape: tuple[int, int],
+      chunk_size: int,
+  ):
+    """Compares lazy chunked mask evaluation against the dense version block-by-block."""
+    q_len, kv_len = shape
+    # Adjust block size if it exceeds shape dimensions
+    adjusted_block_size = (
+        min(block_size[0], q_len),
+        min(block_size[1], kv_len),
+    )
+
+    if (
+        q_len % adjusted_block_size[0] != 0
+        or kv_len % adjusted_block_size[1] != 0
+    ):
+      self.skipTest(
+          f"Shape {shape} not divisible by block_size {adjusted_block_size}"
+      )
+
+    dense_mask = _make_chunked_causal_mask(shape=shape, chunk_size=chunk_size)
+    lazy_mask = mask_lib.ChunkedCausalMask(shape=shape, chunk_size=chunk_size)
+    self._compare_masks(
+        dense_mask,
+        lazy_mask,
+        adjusted_block_size,
+    )
+
+  def test_chunked_causal_mask_invalid_chunk_size(self):
+    """Tests that invalid chunk_size raises ValueError."""
+    with self.assertRaises(ValueError):
+      mask_lib.ChunkedCausalMask(shape=(10, 10), chunk_size=0)
+    with self.assertRaises(ValueError):
+      mask_lib.ChunkedCausalMask(shape=(10, 10), chunk_size=-1)
+    with self.assertRaises(ValueError):
+      mask_lib.make_chunk_attention_mask(shape=(10, 10), chunk_size=0)
+
+  def test_chunked_causal_mask_minimal_equality_hash(self):
+    """Tests for __eq__ and __hash__ of ChunkedCausalMask."""
+    shape1, chunk_size1 = (128, 256), 16
+    shape2, chunk_size2 = (128, 128), 32  # Different shape/chunk_size
+
+    # Create three masks: two identical, one with different shape/chunk_size.
+    mask1 = mask_lib.ChunkedCausalMask(shape=shape1, chunk_size=chunk_size1)
+    mask2 = mask_lib.ChunkedCausalMask(shape=shape1, chunk_size=chunk_size1)
+    mask_diff_shape = mask_lib.ChunkedCausalMask(
+        shape=shape2, chunk_size=chunk_size1
+    )
+    mask_diff_chunk = mask_lib.ChunkedCausalMask(
+        shape=shape1, chunk_size=chunk_size2
+    )
+    other_obj = object()
+
+    # Test __eq__
+    self.assertEqual(mask1, mask2)
+    self.assertNotEqual(mask1, mask_diff_shape)
+    self.assertNotEqual(mask1, mask_diff_chunk)
+    self.assertNotEqual(mask1, other_obj)
+
+    # Test __hash__ of identical masks
+    self.assertEqual(hash(mask1), hash(mask2))
+
+    mask_set = {mask1, mask2, mask_diff_chunk}
+    self.assertLen(mask_set, 2)  # mask1 and mask2 are duplicates
+    self.assertIn(mask1, mask_set)
+    self.assertIn(mask_diff_chunk, mask_set)
+    self.assertNotIn(mask_diff_shape, mask_set)
+
   def test_using_logical_operators_raises_exception(self):
     mask_1 = mask_lib.NumpyMask(
         mask_lib.make_random_mask((256, 256), 0.5, seed=1)