Support Splash Attention when seq_len is not a multiple of the block size.

GitOrigin-RevId: eceb692

Author: ds-hwang

axlearn/common/attention_bias.py

@@ -732,7 +732,7 @@ def mask(query_position: Tensor, key_position: Tensor):
 def and_masks(*mask_fns: ConfigOr[MaskFn]) -> MaskFn:
     """Returns a MaskFn that's the intersection of provided MaskFn's."""
 
-    def mask(query_position: Tensor, key_position: Tensor):
+    def mask(query_position: Tensor, key_position: Tensor) -> Tensor:
         fns = [maybe_instantiate(arg) for arg in mask_fns]
         result = fns[0](query_position, key_position)
         for mask in fns[1:]:
@@ -754,14 +754,31 @@ def sliding_window_causal_mask(sliding_window_size: int) -> MaskFn:
         sliding_window_size: Left context of sliding window mask.
     """
 
-    def mask(query_position: Tensor, key_position: Tensor):
+    def mask(query_position: Tensor, key_position: Tensor) -> Tensor:
         pos_mask = query_position - key_position <= sliding_window_size
         return pos_mask
 
     fun = and_masks(causal_mask, mask)
     return fun
 
 
+def truncated_key_mask(valid_k_len: int) -> MaskFn:
+    """Creates a mask function that cuts off attention beyond valid_k_len.
+
+    Args:
+        valid_k_len: The valid key length. Positions >= valid_k_len will be masked out.
+
+    Returns:
+        A mask function that returns True for key_position < valid_k_len, False otherwise.
+    """
+
+    def mask_fn(query_position: Tensor, key_position: Tensor) -> Tensor:
+        del query_position
+        return key_position < valid_k_len
+
+    return mask_fn
+
+
 def make_causal_biases(seq_len: int) -> Tensor:
     """Generates attention logit biases for causal masking.
 

axlearn/common/attention_bias_test.py

@@ -16,6 +16,7 @@
     SegmentIdAttentionBias,
     TensorAttentionBias,
     sliding_window_causal_mask,
+    truncated_key_mask,
 )
 from axlearn.common.utils import Tensor
 
@@ -49,6 +50,20 @@ def test_sliding_window_mask_with_time_step(self, time_step, expected):
         out_mask = bool_mask.astype(jnp.int32)
         self.assertEqual(out_mask.tolist(), expected)
 
+    @parameterized.parameters(
+        [3, [[1, 1, 1, 0, 0]]],
+        [5, [[1, 1, 1, 1, 1]]],
+    )
+    def test_key_limit_mask(self, valid_k_len, expected):
+        """Test make_key_limit_mask masks keys beyond valid_k_len."""
+        mask_fn = truncated_key_mask(valid_k_len)
+        seq_len = 5
+        query_positions = jnp.arange(seq_len)[:, None]
+        key_positions = jnp.arange(seq_len)[None, :]
+        bool_mask = mask_fn(query_positions, key_positions)
+        out_mask = bool_mask.astype(jnp.int32)
+        self.assertEqual(out_mask.tolist(), expected)
+
 
 class AttentionBiasTest(test_utils.TestCase):
     @parameterized.parameters(

axlearn/common/flash_attention/common.py

@@ -10,6 +10,7 @@
 from functools import partial
 from typing import Any, Literal, NamedTuple, Optional
 
+import chex
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -163,7 +164,7 @@ class Config(Configurable.Config):
 
     def __init__(self, cfg: Config):
         super().__init__(cfg)
-        self.cfg: BaseFlashAttention.Config = self.config
+        self.cfg = cfg
 
     def get_backend_overrides(self, name: str, default: Any) -> Any:
         return (self.cfg.backend_overrides or {}).get(name, default)
@@ -575,3 +576,64 @@ def get_tpu_dot_precision(dtype) -> jax.lax.Precision:
     if dtype == jnp.bfloat16:
         return jax.lax.Precision.DEFAULT
     raise ValueError(f"Unsupported dtype {dtype}")
+
+
+def maybe_pad_inputs(
+    block_size: int, query: Tensor, key: Tensor, value: Tensor, segment_id: Optional[Tensor]
+) -> tuple[Tensor, Tensor, Tensor, Optional[Tensor]]:
+    """Pads query, key, value, and segment_id tensors to align with block_size requirements.
+
+    This function ensures that the sequence length dimension of input tensors is divisible by
+    block_size, which is required for efficient block-based attention computation. Padding is
+    applied to the sequence length dimension (axis=1) as needed.
+
+    Args:
+        block_size: The block size that sequence lengths must be divisible by. Must be positive.
+        query: Query tensor of shape [batch_size, query_len, num_heads, per_head_dim].
+        key: Key tensor of shape [batch_size, key_len, num_kv_heads, per_head_dim].
+        value: Value tensor of shape [batch_size, value_len, num_kv_heads, per_head_dim].
+            Must have the same shape as key.
+        segment_id: Optional segment ID tensor of shape [batch_size, seq_len]. If provided,
+            it will be padded with -1 values to match query padding.
+
+    Returns:
+        A tuple containing:
+            - Padded query tensor
+            - Padded key tensor
+            - Padded value tensor
+            - Padded segment_id tensor (or None if input was None)
+
+    Raises:
+        AssertionError: If block_size is not positive, or if key and value shapes don't match.
+    """
+    # Input validation with chex assertions
+    chex.assert_scalar_positive(block_size)
+    chex.assert_equal_rank((query, key, value))
+    chex.assert_rank(query, 4)
+    chex.assert_equal_shape([key, value])
+    chex.assert_equal(query.shape[0], key.shape[0])
+
+    if segment_id is not None:
+        chex.assert_rank(segment_id, 2)
+        chex.assert_equal(segment_id.shape, query.shape[:2])
+
+    def pad_fn(x, pad):
+        return jnp.pad(x, ((0, 0), (0, pad), (0, 0), (0, 0)))
+
+    query_pad = -query.shape[1] % block_size
+    if query_pad > 0:
+        # This part will be cut after attention, so any value is fine.
+        query = pad_fn(query, query_pad)
+        if segment_id is not None:
+            # segment_ids == 0 means padding.
+            segment_id = jnp.pad(segment_id, ((0, 0), (0, query_pad)))
+
+    key_pad = -key.shape[1] % block_size
+    if key_pad > 0:
+        key = pad_fn(key, key_pad)
+        value = pad_fn(value, key_pad)
+
+    chex.assert_equal(query.shape[1] % block_size, 0)
+    chex.assert_equal(key.shape[1] % block_size, 0)
+    chex.assert_equal(value.shape[1] % block_size, 0)
+    return query, key, value, segment_id

axlearn/common/flash_attention/common_test.py

@@ -2,10 +2,15 @@
 
 """Tests for common utilities"""
 
-from absl.testing import parameterized
+import jax.numpy as jnp
+from absl.testing import absltest, parameterized
 
 from axlearn.common.attention_bias import sliding_window_causal_mask
-from axlearn.common.flash_attention.common import build_mask, build_sliding_window_mask
+from axlearn.common.flash_attention.common import (
+    build_mask,
+    build_sliding_window_mask,
+    maybe_pad_inputs,
+)
 from axlearn.common.test_utils import TestCase
 
 
@@ -20,3 +25,51 @@ def test_sliding_window_fast_path(self, sliding_window_sz, seq_len, block_size):
         mask = build_mask(sliding_window_causal_mask(sliding_window_sz), **args)
         sliding_mask = build_sliding_window_mask(**args, sliding_window_size=sliding_window_sz)
         self.assertNestedEqual(sliding_mask, mask)
+
+
+class UtilsTest(TestCase):
+    @parameterized.parameters(
+        dict(block_size=8, input_len=7, output_len=8),
+        dict(block_size=8, input_len=16, output_len=16),
+        dict(block_size=16, input_len=15, output_len=16),
+        dict(block_size=64, input_len=63, output_len=64),
+    )
+    def test_maybe_pad_inputs(self, block_size, input_len, output_len):
+        query = jnp.ones((2, input_len, 4, 64))
+        key = jnp.ones((2, input_len, 2, 64))
+        value = jnp.ones((2, input_len, 2, 64))
+        segment_id = jnp.ones((2, input_len), dtype=jnp.int32) * 5
+
+        orig_q_len, orig_k_len = query.shape[1], key.shape[1]
+        padded_query, padded_key, padded_value, padded_segment_id = maybe_pad_inputs(
+            block_size, query, key, value, segment_id
+        )
+
+        self.assertEqual(padded_query.shape[1], output_len)
+        self.assertEqual(padded_key.shape[1], output_len)
+        self.assertEqual(padded_value.shape[1], output_len)
+        self.assertEqual(padded_segment_id.shape[1], output_len)
+        self.assertEqual(orig_q_len, input_len)
+        self.assertEqual(orig_k_len, input_len)
+        self.assertEqual(padded_query.shape[1] % block_size, 0)
+
+        self.assertNestedAllClose(padded_query[:, :input_len], query)
+        self.assertNestedAllClose(padded_segment_id[:, :input_len], segment_id)
+
+        if output_len > input_len:
+            self.assertNestedAllClose(
+                padded_query[:, input_len:], jnp.zeros_like(padded_query[:, input_len:])
+            )
+            self.assertNestedAllClose(
+                padded_segment_id[:, input_len:], jnp.zeros_like(padded_segment_id[:, input_len:])
+            )
+
+        padded_query2, _, _, padded_segment_id2 = maybe_pad_inputs(
+            block_size, query, key, value, None
+        )
+        self.assertEqual(padded_query2.shape[1], output_len)
+        self.assertIsNone(padded_segment_id2)
+
+
+if __name__ == "__main__":
+    absltest.main()