Add sliding window support to the ragged paged attention.

PiperOrigin-RevId: 738457532

Author: Google-ML-Automation

jax/experimental/pallas/ops/tpu/ragged_paged_attention.py

@@ -19,7 +19,6 @@
 specifications. It supports mixed prefill and decoding, enhancing throughput
 during inference.
 """
-
 import functools
 import jax
 from jax import lax
@@ -81,6 +80,7 @@ def ref_ragged_paged_attention(
     num_seqs: jax.Array,  # i32[1],
     *,
     sm_scale: float = 1.0,
+    sliding_window: int | None = None,
     mask_value: float = DEFAULT_MASK_VALUE,
 ):
   _, _, num_kv_heads, head_dim = k_pages.shape
@@ -105,7 +105,10 @@ def ref_ragged_paged_attention(
         jnp.int32, attn.shape, 1
     )
     kv_span = jax.lax.broadcasted_iota(jnp.int32, attn.shape, 2)
-    attn += jnp.where(q_span < kv_span, mask_value, 0.0)
+    mask = q_span < kv_span
+    if sliding_window is not None:
+      mask = jnp.logical_or(mask, q_span - sliding_window >= kv_span)
+    attn += jnp.where(mask, mask_value, 0.0)
     attn = jax.nn.softmax(attn, axis=-1).astype(v.dtype)
     out = jnp.einsum("hqk,khd->qhd", attn, v).astype(queries.dtype)
     outputs.append(out)
@@ -122,6 +125,7 @@ def validate_inputs_on_runtime(
     page_indices: jax.Array,  # i32[max_num_seqs, pages_per_seq]
     cu_q_lens: jax.Array,  # i32[max_num_seqs + 1]
     num_seqs,  # i32[1]
+    sliding_window: int | None = None,
 ):
   check_inputs_shapes(
       q, k_pages, v_pages, kv_lens, page_indices, cu_q_lens, num_seqs
@@ -150,6 +154,8 @@ def validate_inputs_on_runtime(
       raise ValueError(
           f"{q_len=} must be less or equal to {kv_len=} at sequence {i}."
       )
+  if sliding_window is not None and sliding_window <= 0:
+    raise ValueError(f"{sliding_window=} must be positive.")
 
 
 # Expect to run these checks during compile time.
@@ -221,7 +227,8 @@ def ragged_paged_attention_kernel(
     m_ref,  # [num_kv_heads_per_blk, num_q_per_blk * num_q_heads_per_kv_head, 128]
     *,
     sm_scale: float,
-    mask_value: float,
+    sliding_window: int | None = None,
+    mask_value: float = DEFAULT_MASK_VALUE,
 ):
   num_q_per_blk, num_q_heads_per_blk, head_dim = q_ref.shape
   num_seqs = num_seqs_ref[0]
@@ -373,7 +380,7 @@ def flash_attention(
       def masked_store(ref, val, start, end, group=1):
         iota = lax.broadcasted_iota(jnp.int32, ref.shape, 0) // group
         mask = jnp.logical_and(iota >= start, iota < end)
-        pl.store(ref, tuple(slice(None) for _ in ref.shape), val, mask=mask)
+        pl.store(ref, idx=tuple(slice(None) for _ in ref.shape), val=val, mask=mask)
 
       qk = (
           jnp.einsum("nd,md->nm", q, k, preferred_element_type=jnp.float32)
@@ -422,6 +429,9 @@ def init_scratch_ref():
           1,
       )
       causal_mask = row_ids < col_ids
+      if sliding_window is not None:
+        causal_mask = jnp.logical_or(causal_mask,
+                                     row_ids - sliding_window>=col_ids)
       qk += jnp.where(causal_mask, mask_value, 0.0)
       m_curr = jnp.max(qk, axis=1, keepdims=True)
       s_curr = jnp.exp(qk - m_curr)
@@ -601,6 +611,7 @@ def can_be_xla_fully_tiled(x, packing):
         "num_kv_pages_per_block",
         "num_queries_per_block",
         "vmem_limit_bytes",
+        "sliding_window",
     ],
 )
 def ragged_paged_attention(
@@ -614,6 +625,7 @@ def ragged_paged_attention(
     num_seqs: jax.Array,  # i32[1]
     *,
     sm_scale: float = 1.0,
+    sliding_window: int | None = None,
     mask_value: float = DEFAULT_MASK_VALUE,
     num_kv_pages_per_block: int = 16,
     num_queries_per_block: int = 128,
@@ -632,6 +644,7 @@ def ragged_paged_attention(
       kv_lens, only the first num_seqs+1 values are valid.
     num_seqs: the dynamic number of sequences.
     sm_scale: the softmax scale which will be applied to the Q@K^T.
+    sliding_window: the sliding window size for the attention.
     mask_value: mask value for causal mask.
     num_kv_pages_per_block: number of kv pages to be processed in one flash
       attention block in the pallas kernel.
@@ -705,6 +718,7 @@ def q_index_map(heads_blk_idx, q_blk_idx, *_):
       functools.partial(
           ragged_paged_attention_kernel,
           sm_scale=sm_scale,
+          sliding_window=sliding_window,
           mask_value=mask_value,
       ),
       grid_spec=pltpu.PrefetchScalarGridSpec(
@@ -724,6 +738,7 @@ def q_index_map(heads_blk_idx, q_blk_idx, *_):
       out_shape=jax.ShapeDtypeStruct(shape=q.shape, dtype=jnp.float32),
       name="ragged_paged_attention_kernel",
   )
+
   # TODO(jevinjiang): Use f32 acc scratch for output! So we only need
   # to transfer output with desired dtype back to HBM.
   return kernel(*scalar_prefetches, q, k_pages, v_pages).astype(q.dtype)

tests/pallas/tpu_ragged_paged_attention_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import random
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
@@ -50,6 +51,7 @@ def _test_ragged_paged_attention(
       vmem_limit_bytes=32 * 1024 * 1024,
       max_num_batched_tokens=512,
       max_num_seq=8,
+      sliding_window: int | None = None,
   ):
     if not jtu.is_device_tpu_at_least(version=4):
       self.skipTest("Expect TPUv4+")
@@ -101,8 +103,10 @@ def _test_ragged_paged_attention(
         page_indices,
         cu_q_lens,
         num_seqs,
+        sliding_window=sliding_window,
     )
 
+    actual_num_q_tokens = cu_q_lens[num_seqs[0]]
     output = ragged_paged_attention(
         q,
         k_pages,
@@ -114,7 +118,8 @@ def _test_ragged_paged_attention(
         num_kv_pages_per_block=num_kv_pages_per_block,
         num_queries_per_block=num_queries_per_block,
         vmem_limit_bytes=vmem_limit_bytes,
-    )[: cu_q_lens[num_seqs[0]]]
+        sliding_window=sliding_window,
+    )[: actual_num_q_tokens]
 
     expected = ref_ragged_paged_attention(
         q,
@@ -124,6 +129,7 @@ def _test_ragged_paged_attention(
         page_indices,
         cu_q_lens,
         num_seqs=num_seqs,
+        sliding_window=sliding_window,
     )
     tols = {
         "float32": 0.15,
@@ -266,6 +272,7 @@ def test_ragged_paged_attention_mixed(self, dtype):
       dtype=[jnp.float32, jnp.bfloat16],
       num_kv_pages_per_block=[4, 8],
       num_queries_per_block=[32, 64],
+      sliding_window=[None, 5, 128],
   )
   def test_ragged_paged_attention_complex(
       self,
@@ -274,6 +281,7 @@ def test_ragged_paged_attention_complex(
       dtype,
       num_kv_pages_per_block,
       num_queries_per_block,
+      sliding_window: int | None,
   ):
     seq_lens = []
     for _ in range(num_seqs):
@@ -294,8 +302,38 @@ def test_ragged_paged_attention_complex(
         num_pages,
         num_kv_pages_per_block=num_kv_pages_per_block,
         num_queries_per_block=num_queries_per_block,
+        sliding_window=sliding_window,
     )
 
+  def test_ragged_paged_attention_sliding_window_should_be_positive(self):
+    dtype=jnp.float32
+    seq_lens = [(192, 328), (128, 180), (64, 255)]
+    num_heads = (32, 8)
+    head_dim = 128
+    page_size = 16
+    num_pages = 1000
+
+    with self.assertRaisesRegex(ValueError, "must be positive"):
+      self._test_ragged_paged_attention(
+          seq_lens,
+          num_heads,
+          head_dim,
+          page_size,
+          dtype,
+          num_pages,
+          sliding_window=0,
+      )
+
+    with self.assertRaisesRegex(ValueError, "must be positive"):
+      self._test_ragged_paged_attention(
+          seq_lens,
+          num_heads,
+          head_dim,
+          page_size,
+          dtype,
+          num_pages,
+          sliding_window=-1,
+      )
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())