tpu-inference 0.12.0.dev20251222__py3-none-any.whl

tests/kernels/ragged_paged_attention_kernel_v3_test.py

@@ -0,0 +1,520 @@
+# Copyright 2025 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
+#
+#     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.
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from absl.testing import absltest, parameterized
+from jax._src import dtypes
+from jax._src import test_util as jtu
+
+from tpu_inference.kernels.ragged_paged_attention.v3.kernel import (
+    ragged_paged_attention, ref_ragged_paged_attention)
+from tpu_inference.kernels.ragged_paged_attention.v3.util import (
+    align_to, cdiv, get_dtype_packing)
+
+jax.config.parse_flags_with_absl()
+
+
+@jtu.with_config(jax_numpy_dtype_promotion="standard")
+class RaggedPagedAttentionKernelTest(jtu.JaxTestCase):
+
+    def _test_ragged_paged_attention(
+        self,
+        seq_lens,  # List[(q_len, kv_len)]
+        num_heads,  # [num_q_heads, num_kv_heads]
+        head_dim,
+        page_size,
+        q_dtype,
+        kv_dtype,
+        num_pages,
+        *,
+        num_kv_pages_per_block=8,
+        num_queries_per_block=64,
+        vmem_limit_bytes=100 * 1024 * 1024,
+        max_num_batched_tokens=512,
+        max_num_seq=8,
+        sliding_window: int | None = None,
+        soft_cap: float | None = None,
+        q_scale: float | None = None,
+        k_scale: float | None = None,
+        v_scale: float | None = None,
+    ):
+        rng = np.random.default_rng(1234)
+
+        def gen_random(shape, dtype):
+            return jnp.array(rng.random(size=shape,
+                                        dtype=np.float32)).astype(dtype)
+
+        if not jtu.is_device_tpu_at_least(version=4):
+            self.skipTest("Expect TPUv4+")
+        cu_q_lens = [0]
+        kv_lens = []
+        for q_len, kv_len in seq_lens:
+            assert q_len <= kv_len
+            cu_q_lens.append(cu_q_lens[-1] + q_len)
+            kv_lens.append(kv_len)
+
+        max_num_batched_tokens = max(align_to(cu_q_lens[-1], 128),
+                                     max_num_batched_tokens)
+        max_num_seq = max(align_to(len(seq_lens), 8), max_num_seq)
+        max_kv_len = max(kv_lens)
+        pages_per_seq = cdiv(max_kv_len, page_size)
+        num_q_heads, num_kv_heads = num_heads
+
+        q = gen_random((max_num_batched_tokens, num_q_heads, head_dim),
+                       q_dtype)
+        k = gen_random((max_num_batched_tokens, num_kv_heads, head_dim),
+                       kv_dtype)
+        v = gen_random((max_num_batched_tokens, num_kv_heads, head_dim),
+                       kv_dtype)
+        page_cnt = 0
+        page_indices_list = []
+        kv_pages_list = []
+        kv_packing = get_dtype_packing(kv_dtype)
+        padded_head_dim = align_to(head_dim, 128)
+        num_kv_heads_x2 = align_to(num_kv_heads * 2, kv_packing)
+        for kv_len in kv_lens:
+            kv = gen_random((
+                kv_len,
+                num_kv_heads_x2 // kv_packing,
+                kv_packing,
+                padded_head_dim,
+            ), kv_dtype)
+            kv = jnp.pad(
+                kv,
+                (
+                    (
+                        0,
+                        cdiv(kv_len, page_size) * page_size - kv_len,
+                    ),
+                    (0, 0),
+                    (0, 0),
+                    (0, 0),
+                ),
+                constant_values=jnp.nan,
+            ).reshape(
+                -1,
+                page_size,
+                num_kv_heads_x2 // kv_packing,
+                kv_packing,
+                padded_head_dim,
+            )
+            indices = page_cnt + jnp.arange(kv.shape[0], dtype=jnp.int32)
+            indices = jnp.pad(
+                indices,
+                ((0, pages_per_seq - indices.shape[0]), ),
+                constant_values=jnp.nan,
+            )
+            page_indices_list.append(indices)
+            page_cnt += kv.shape[0]
+            kv_pages_list.append(kv)
+
+        kv_cache = jnp.concatenate(kv_pages_list, axis=0)
+        kv_cache = jnp.pad(
+            kv_cache,
+            ((0, num_pages - kv_cache.shape[0]), (0, 0), (0, 0), (0, 0),
+             (0, 0)),
+            constant_values=jnp.nan,
+        )
+        page_indices = jnp.stack(page_indices_list, axis=0)
+        page_indices = jnp.pad(
+            page_indices,
+            ((0, max_num_seq - page_indices.shape[0]), (0, 0)),
+            constant_values=jnp.nan,
+        )
+        page_indices = page_indices.reshape(-1)
+
+        cu_q_lens = jnp.array(cu_q_lens, dtype=jnp.int32)
+        cu_q_lens = jnp.pad(cu_q_lens,
+                            (0, max_num_seq + 1 - cu_q_lens.shape[0]))
+        kv_lens = jnp.array(kv_lens, dtype=jnp.int32)
+        kv_lens = jnp.pad(kv_lens, (0, max_num_seq - kv_lens.shape[0]))
+        distribution = jnp.array([0, 0, len(seq_lens)], dtype=jnp.int32)
+
+        args = (
+            q,
+            k,
+            v,
+            kv_cache,
+            kv_lens,
+            page_indices,
+            cu_q_lens,
+            distribution,
+        )
+
+        kwargs = {
+            "sliding_window": sliding_window,
+            "soft_cap": soft_cap,
+            "q_scale": q_scale,
+            "k_scale": k_scale,
+            "v_scale": v_scale,
+        }
+
+        expected, expected_kv_cache = ref_ragged_paged_attention(
+            *args,
+            **kwargs,
+        )
+
+        output, updated_kv_cache = ragged_paged_attention(
+            *args,
+            **kwargs,
+            num_kv_pages_per_block=num_kv_pages_per_block,
+            num_queries_per_block=num_queries_per_block,
+            vmem_limit_bytes=vmem_limit_bytes,
+        )
+        output = output[:cu_q_lens[distribution[-1]]]
+
+        dtype_bits = (dtypes.bit_width(jnp.dtype(kv_dtype)) if hasattr(
+            dtypes, "bit_width") else dtypes.itemsize_bits(
+                jnp.dtype(kv_dtype)))
+        tols = {
+            32: 0.15,
+            16: 0.2,
+            8: 0.2,
+            4: 0.2,
+        }
+        tol = tols[dtype_bits]
+        self.assertAllClose(output, expected, atol=tol, rtol=tol)
+        mask = ~jnp.isnan(expected_kv_cache)
+        self.assertArraysEqual(updated_kv_cache[mask], expected_kv_cache[mask])
+        self.assertEqual(output.shape[-1], head_dim)
+
+    @parameterized.product(dtype=[jnp.float32, jnp.bfloat16], )
+    def test_ragged_paged_attention_basic(self, dtype):
+        seq_lens = [(192, 328), (128, 180), (64, 255)]
+        num_heads = (32, 8)
+        head_dim = 128
+        page_size = 16
+        num_pages = 1000
+
+        self._test_ragged_paged_attention(
+            seq_lens,
+            num_heads,
+            head_dim,
+            page_size,
+            dtype,
+            dtype,
+            num_pages,
+        )
+
+    # TODO: support integer (int8, int4) and fp4 kv cache
+    @parameterized.product(
+        q_dtype=[jnp.bfloat16],
+        kv_dtype=[jnp.float8_e5m2, jnp.float8_e4m3fn],
+        kv_scales=[(0.5, 0.5), (1.0, 1.0)],
+    )
+    def test_ragged_paged_attention_quantized_kv_cache(self, q_dtype, kv_dtype,
+                                                       kv_scales):
+        if not jtu.is_device_tpu_at_least(version=5):
+            self.skipTest("Expect TPUv5+")
+        seq_lens = [(192, 328), (128, 180), (64, 255)]
+        num_heads = (32, 8)
+        head_dim = 128
+        page_size = 16
+        num_pages = 1000
+        k_scale, v_scale = kv_scales
+
+        self._test_ragged_paged_attention(
+            seq_lens,
+            num_heads,
+            head_dim,
+            page_size,
+            q_dtype,
+            kv_dtype,
+            num_pages,
+            k_scale=k_scale,
+            v_scale=v_scale,
+        )
+
+    @parameterized.product(
+        q_dtype=[jnp.bfloat16],
+        kv_dtype=[jnp.float8_e5m2, jnp.float8_e4m3fn],
+        q_scale=[0.5, 1.0],
+        kv_scales=[(0.5, 0.5), (1.0, 1.0)],
+    )
+    def test_ragged_paged_attention_quantized_attention(
+            self, q_dtype, kv_dtype, q_scale, kv_scales):
+        if not jtu.is_device_tpu_at_least(version=5):
+            self.skipTest("Expect TPUv5+")
+        seq_lens = [(192, 328), (128, 180), (64, 255)]
+        num_heads = (32, 8)
+        head_dim = 128
+        page_size = 16
+        num_pages = 1000
+        k_scale, v_scale = kv_scales
+
+        self._test_ragged_paged_attention(
+            seq_lens,
+            num_heads,
+            head_dim,
+            page_size,
+            q_dtype,
+            kv_dtype,
+            num_pages,
+            q_scale=q_scale,
+            k_scale=k_scale,
+            v_scale=v_scale,
+        )
+
+    @parameterized.product(dtype=[jnp.float32, jnp.bfloat16], )
+    def test_ragged_paged_attention_decode_only(self, dtype):
+        seq_lens = [
+            (1, 18),
+            (1, 129),
+            (1, 597),
+            (1, 122),
+            (1, 64),
+            (1, 322),
+            (1, 463),
+            (1, 181),
+            (1, 1107),
+            (1, 123),
+            (1, 31),
+            (1, 18),
+            (1, 1229),
+            (1, 229),
+            (1, 87),
+            (1, 1328),
+        ]
+        num_heads = (32, 8)
+        head_dim = 128
+        page_size = 16
+        num_pages = 1000
+
+        self._test_ragged_paged_attention(
+            seq_lens,
+            num_heads,
+            head_dim,
+            page_size,
+            dtype,
+            dtype,
+            num_pages,
+        )
+
+    @parameterized.product(dtype=[jnp.float32, jnp.bfloat16], )
+    def test_ragged_paged_attention_prefill_only(self, dtype):
+        seq_lens = [
+            (5, 18),
+            (15, 129),
+            (120, 597),
+            (100, 122),
+            (21, 64),
+            (32, 322),
+            (251, 463),
+            (40, 181),
+            (64, 1107),
+            (99, 123),
+            (10, 31),
+            (5, 18),
+            (3, 1229),
+            (120, 229),
+            (9, 87),
+            (2, 1328),
+        ]
+        num_heads = (32, 8)
+        head_dim = 128
+        page_size = 16
+        num_pages = 1000
+
+        self._test_ragged_paged_attention(
+            seq_lens,
+            num_heads,
+            head_dim,
+            page_size,
+            dtype,
+            dtype,
+            num_pages,
+        )
+
+    @parameterized.product(dtype=[jnp.float32, jnp.bfloat16], )
+    def test_ragged_paged_attention_mixed(self, dtype):
+        seq_lens = [
+            (5, 18),
+            (1, 129),
+            (120, 597),
+            (1, 122),
+            (1, 64),
+            (32, 322),
+            (251, 463),
+            (1, 181),
+            (1, 1107),
+            (99, 123),
+            (1, 31),
+            (5, 18),
+            (3, 1229),
+            (117, 229),
+            (1, 87),
+            (1, 1328),
+        ]
+        num_heads = (32, 8)
+        head_dim = 128
+        page_size = 16
+        num_pages = 1000
+
+        self._test_ragged_paged_attention(
+            seq_lens,
+            num_heads,
+            head_dim,
+            page_size,
+            dtype,
+            dtype,
+            num_pages,
+        )
+
+    @parameterized.product(
+        num_seqs=[1, 17],
+        num_heads=[(32, 8), (12, 2), (5, 1), (3, 3)],
+        head_dim=[80, 240],
+        dtype=[jnp.float32, jnp.bfloat16],
+        # num_kv_pages_per_block=[8, 16],
+        # num_queries_per_block=[16, 32],
+    )
+    def test_ragged_paged_attention_complex(
+        self,
+        num_seqs,
+        num_heads,
+        head_dim,
+        dtype,
+        # num_kv_pages_per_block,
+        # num_queries_per_block,
+    ):
+        rng = np.random.default_rng(1234)
+        q_lens = rng.integers(1, 100, num_seqs)
+        kv_lens = q_lens + rng.integers(0, 50, num_seqs)
+        seq_lens = list(zip(q_lens.tolist(), kv_lens.tolist()))
+        page_size = 16
+        num_pages = 1000
+
+        self._test_ragged_paged_attention(
+            seq_lens,
+            num_heads,
+            head_dim,
+            page_size,
+            dtype,
+            dtype,
+            num_pages,
+            # num_kv_pages_per_block=num_kv_pages_per_block,
+            # num_queries_per_block=num_queries_per_block,
+        )
+
+    @parameterized.product(sliding_window=[None, 5, 128], )
+    def test_ragged_paged_attention_sliding_window(
+        self,
+        sliding_window: int | None,
+    ):
+        num_seqs = 5
+        num_heads = (4, 4)
+        dtype = jnp.float32
+        rng = np.random.default_rng(1234)
+        q_lens = rng.integers(1, 100, num_seqs)
+        kv_lens = q_lens + rng.integers(0, 50, num_seqs)
+        seq_lens = list(zip(q_lens.tolist(), kv_lens.tolist()))
+        head_dim = 128
+        page_size = 16
+        num_pages = 1000
+
+        self._test_ragged_paged_attention(
+            seq_lens,
+            num_heads,
+            head_dim,
+            page_size,
+            dtype,
+            dtype,
+            num_pages,
+            sliding_window=sliding_window,
+        )
+
+    @parameterized.product(soft_cap=[None, 50.0], )
+    def test_ragged_paged_attention_logit_soft_capping(
+        self,
+        soft_cap: float | None,
+    ):
+        num_heads = (16, 2)
+        num_seqs = 2
+        dtype = jnp.float32
+        rng = np.random.default_rng(1234)
+        q_lens = rng.integers(1, 100, num_seqs)
+        kv_lens = q_lens + rng.integers(0, 50, num_seqs)
+        seq_lens = list(zip(q_lens.tolist(), kv_lens.tolist()))
+        head_dim = 128
+        page_size = 16
+        num_pages = 1000
+
+        self._test_ragged_paged_attention(
+            seq_lens,
+            num_heads,
+            head_dim,
+            page_size,
+            dtype,
+            dtype,
+            num_pages,
+            soft_cap=soft_cap,
+        )
+
+    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,
+                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,
+                dtype,
+                num_pages,
+                sliding_window=-1,
+            )
+
+    def test_ragged_paged_attention_soft_cap_cannot_be_zero(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 not be 0.0"):
+            self._test_ragged_paged_attention(
+                seq_lens,
+                num_heads,
+                head_dim,
+                page_size,
+                dtype,
+                dtype,
+                num_pages,
+                soft_cap=0.0,
+            )
+
+
+if __name__ == "__main__":
+    absltest.main(testLoader=jtu.JaxTestLoader())

tests/layers/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 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
+#
+#     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.

tests/layers/common/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 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
+#
+#     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.

tests/layers/common/test_attention_interface.py

@@ -0,0 +1,156 @@
+# Copyright 2025 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
+#
+#     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.
+
+from unittest.mock import MagicMock
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from jax.sharding import Mesh
+
+from tpu_inference.layers.common.attention_interface import attention
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.runner.kv_cache import get_kv_cache_shape_with_mesh
+
+# ---- Test Configuration & Constants ----
+
+# Total number of tokens across all sequences in the batch
+TOTAL_TOKENS = 10
+# Number of sequences in the batch
+NUM_SEQS = 2
+# Padded maximum number of sequences
+MAX_NUM_SEQS = 4
+# Number of attention heads (Query)
+NUM_HEADS = 8
+# Number of attention heads (Key/Value) - for Grouped-Query Attention
+NUM_KV_HEADS = 4
+# Total number of blocks in the KV cache
+NUM_BLOCKS = 32
+# Number of tokens per block
+BLOCK_SIZE = 16
+# Maximum number of blocks a single sequence can occupy
+MAX_BLOCKS_PER_SEQ = 8
+
+
+@pytest.fixture
+def mesh():
+    """Provides a mock 1D JAX mesh for testing."""
+    # Create a mesh with available devices, useful for running on CPU/GPU/TPU
+    # For this test, it will likely be a single CPU device.
+    devices = np.array(jax.local_devices()[:1])
+    if not devices.any():
+        # Add a mock device if no devices are present (e.g., in a CI environment)
+        devices = np.array([jax.devices("cpu")[0]])
+    return Mesh(devices.reshape((-1, 1, 1)), ("data", "attn_dp", "model"))
+
+
+# ---- Test for `attention` ----
+
+
+def _test_attention(monkeypatch, mesh, head_dim, use_sinks=False):
+    """
+    Tests the main `attention` function.
+
+    Verifies that:
+    1. It calls the `sharded_ragged_paged_attention` kernel with correct metadata.
+    2. The final outputs (kv_cache and attention output) have the correct shapes.
+    """
+    # 1. Arrange
+
+    # Create input tensors
+    q_dtype = jnp.float32
+    kv_dtype = jnp.float32
+    q = jnp.ones((TOTAL_TOKENS, NUM_HEADS, head_dim), dtype=q_dtype)
+    k = jnp.ones((TOTAL_TOKENS, NUM_KV_HEADS, head_dim), dtype=kv_dtype)
+    v = jnp.ones((TOTAL_TOKENS, NUM_KV_HEADS, head_dim), dtype=kv_dtype)
+    sinks = jnp.ones((NUM_HEADS, ), dtype=jnp.float32) if use_sinks else None
+
+    kv_cache_shape = get_kv_cache_shape_with_mesh(
+        mesh,
+        NUM_BLOCKS,
+        BLOCK_SIZE,
+        NUM_KV_HEADS,
+        head_dim,
+        kv_dtype,
+    )
+    kv_cache = jnp.zeros(kv_cache_shape, dtype=kv_dtype)
+
+    # Mock ragged_paged_attention to return a tensor of the correct shape
+    mock_paged_attn_kernel = MagicMock(return_value=(jnp.ones(
+        (TOTAL_TOKENS, NUM_HEADS, head_dim)), kv_cache), )
+
+    if head_dim == 64:
+        monkeypatch.setattr(
+            "tpu_inference.layers.common.attention_interface.ragged_paged_attention_hd64",
+            mock_paged_attn_kernel,
+        )
+    else:
+        monkeypatch.setattr(
+            "tpu_inference.layers.common.attention_interface.ragged_paged_attention",
+            mock_paged_attn_kernel,
+        )
+
+    # Create AttentionMetadata
+    attention_metadata = AttentionMetadata(
+        input_positions=jnp.arange(TOTAL_TOKENS, dtype=jnp.int32),
+        block_tables=jnp.zeros((MAX_NUM_SEQS * MAX_BLOCKS_PER_SEQ, ),
+                               dtype=jnp.int32),
+        seq_lens=jnp.array([5, 5, 0, 0], dtype=jnp.int32),
+        query_start_loc=jnp.array([0, 5, 10, 10, 10], dtype=jnp.int32),
+        request_distribution=jnp.array([0, 0, NUM_SEQS], dtype=jnp.int32),
+    )
+
+    # 2. Act
+    final_kv_cache, output = attention(
+        kv_cache=kv_cache,
+        q=q,
+        k=k,
+        v=v,
+        attention_metadata=attention_metadata,
+        mesh=mesh,
+        head_dim_original=head_dim,
+        sinks=sinks,
+    )
+
+    # 3. Assert
+    # Check that both mocked kernels were called
+    mock_paged_attn_kernel.assert_called_once()
+
+    # Check output shapes
+    assert final_kv_cache.shape == kv_cache.shape
+    assert output.shape == q.shape
+
+    # Check that the output is the one from our mock
+    assert jnp.all(output == 1.0)
+
+
+def test_attention(monkeypatch, mesh):
+    _test_attention(monkeypatch, mesh, 128)
+
+
+def test_attention_hd64(monkeypatch, mesh):
+    _test_attention(monkeypatch, mesh, 64)
+
+
+def test_attention_sink(monkeypatch, mesh):
+    _test_attention(monkeypatch, mesh, 64, True)
+
+
+def test_attention_sink_no_64_raises_error(monkeypatch, mesh):
+    with pytest.raises(
+            NotImplementedError,
+            match="Attention sink support is only available when head_dim==64"
+    ):
+        _test_attention(monkeypatch, mesh, 128, True)