tpu-inference 0.11.1.dev202511150811__py3-none-any.whl

tests/kernels/quantized_matmul_kernel_test.py

@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+
+import jax
+import jax.numpy as jnp
+from absl.testing import absltest, parameterized
+from jax._src import test_util as jtu
+
+from tpu_inference.kernels.quantized_matmul import (kernel, tuned_block_sizes,
+                                                    util)
+
+quantized_matmul_kernel = kernel.quantized_matmul_kernel
+quantize_tensor = util.quantize_tensor
+get_tuned_block_sizes = tuned_block_sizes.get_tuned_block_sizes
+
+jax.config.parse_flags_with_absl()
+
+
+@functools.partial(jax.jit, static_argnames=["quantize_activation"])
+def reference_quantized_matmul(
+    x: jax.Array,
+    w_q: jax.Array,
+    w_scale: jax.Array,
+    quantize_activation=True,
+):
+    if quantize_activation:
+        acc_dtype = jnp.float32
+        if quantize_activation and jnp.issubdtype(w_q.dtype, jnp.integer):
+            acc_dtype = jnp.int32
+
+        x_q, x_scale = quantize_tensor(x, w_q.dtype)
+        out = jax.lax.dot_general(
+            x_q,
+            w_q,
+            dimension_numbers=(((1, ), (1, )), ((), ())),
+            preferred_element_type=acc_dtype,
+        ).astype(jnp.float32)
+        out *= x_scale
+    else:
+        out = jax.lax.dot_general(
+            x,
+            w_q,
+            dimension_numbers=(((1, ), (1, )), ((), ())),
+            preferred_element_type=jnp.float32,
+        )
+    out *= jnp.expand_dims(w_scale, 0)
+    return out.astype(x.dtype)
+
+
+@jtu.with_config(jax_numpy_dtype_promotion="standard")
+class QuantizedMatmulKernelTest(jtu.JaxTestCase):
+
+    def setUp(self):
+        super().setUp()
+        if not jtu.is_device_tpu_at_least(6):
+            self.skipTest("Expect TPUv6+")
+
+    def _test_quantized_matmul(
+        self,
+        dtype: jnp.dtype,
+        q_dtype: jnp.dtype,
+        bs: int,
+        n_input_features: int,
+        n_output_features: int,
+        quantize_activation: bool,
+        tuned_value=None,
+        atol=0.5,
+        rtol=0.5,
+    ):
+
+        prng_key = jax.random.key(1234)
+        k0, k1 = jax.random.split(prng_key, 2)
+        x = jax.random.uniform(k0, (bs, n_input_features),
+                               dtype=dtype,
+                               minval=0,
+                               maxval=1)
+        w = jax.random.uniform(
+            k1,
+            (n_output_features, n_input_features),
+            dtype=dtype,
+            minval=-1,
+            maxval=1,
+        )
+        w_q, w_scale = quantize_tensor(w, q_dtype)
+        w_scale = jnp.squeeze(w_scale)
+        assert w_scale.shape == (n_output_features, )
+
+        x_q_dtype = w_q.dtype if quantize_activation else dtype
+        output = quantized_matmul_kernel(
+            x,
+            w_q,
+            w_scale,
+            x_q_dtype=x_q_dtype,
+            tuned_value=tuned_value,
+        )
+        expected = reference_quantized_matmul(
+            x, w_q, w_scale, quantize_activation=quantize_activation)
+
+        self.assertAllClose(output,
+                            expected,
+                            rtol=rtol,
+                            atol=atol,
+                            check_dtypes=True)
+
+    @parameterized.product(
+        dtype=[jnp.bfloat16, jnp.float32],
+        q_dtype=[jnp.int8, jnp.float8_e4m3fn],
+        bs=[128, 256, 512],
+        n_input_features=[128, 256, 512],
+        n_output_features=[128, 256, 512],
+        quantize_activation=[True],
+    )
+    def test_quantized_matmul_various_input_shapes(
+        self,
+        dtype: jnp.dtype,
+        q_dtype: jnp.dtype,
+        bs: int,
+        n_input_features: int,
+        n_output_features: int,
+        quantize_activation: bool,
+    ):
+        self._test_quantized_matmul(
+            dtype,
+            q_dtype,
+            bs,
+            n_input_features,
+            n_output_features,
+            quantize_activation=quantize_activation,
+            tuned_value=None,
+        )
+
+    @parameterized.product(
+        dtype=[jnp.bfloat16, jnp.float32],
+        q_dtype=[jnp.int8, jnp.float8_e4m3fn],
+        bs=[64, 192],
+        n_input_features=[64, 192],
+        n_output_features=[64, 192],
+        quantize_activation=[True],
+    )
+    def test_quantized_matmul_unaligned_input_shapes(
+        self,
+        dtype: jnp.dtype,
+        q_dtype: jnp.dtype,
+        bs: int,
+        n_input_features: int,
+        n_output_features: int,
+        quantize_activation: bool,
+    ):
+        self._test_quantized_matmul(
+            dtype,
+            q_dtype,
+            bs,
+            n_input_features,
+            n_output_features,
+            quantize_activation=quantize_activation,
+            tuned_value=None,
+        )
+
+    @parameterized.parameters(
+        (jnp.bfloat16, jnp.int8, 128, 1280, 8192, True),
+        (jnp.bfloat16, jnp.int8, 128, 28672, 4096, True),
+        (jnp.bfloat16, jnp.int8, 128, 4096, 14336, True),
+        (jnp.bfloat16, jnp.int8, 128, 4096, 4096, True),
+        (jnp.bfloat16, jnp.int8, 128, 6144, 4096, True),
+        (jnp.bfloat16, jnp.int8, 128, 7168, 8192, True),
+        (jnp.bfloat16, jnp.int8, 128, 8192, 1024, True),
+        (jnp.bfloat16, jnp.int8, 128, 8192, 3584, True),
+    )
+    def test_quantized_matmul_use_tuned_block_sizes(
+        self,
+        dtype: jnp.dtype,
+        q_dtype: jnp.dtype,
+        bs: int,
+        n_input_features: int,
+        n_output_features: int,
+        quantize_activation: bool,
+    ):
+        self._test_quantized_matmul(
+            dtype,
+            q_dtype,
+            bs,
+            n_input_features,
+            n_output_features,
+            quantize_activation=quantize_activation,
+            tuned_value=None,
+        )
+
+
+if __name__ == "__main__":
+    absltest.main(testLoader=jtu.JaxTestLoader())

tests/kernels/ragged_kv_cache_update_v2_test.py

@@ -0,0 +1,234 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+from absl.testing import parameterized
+from jax._src import test_util as jtu
+from jax.sharding import Mesh, NamedSharding
+from jax.sharding import PartitionSpec as P
+
+from tpu_inference.kernels.ragged_paged_attention.v2.ragged_kv_cache_update import \
+    kv_cache_update
+
+
+def kv_cache_update_ref(new_kv, slot_mapping, kv_cache):
+    """Reference implementation of KV cache update."""
+    for i in range(slot_mapping.shape[1]):
+        start_idx, new_kv_idx, slice_len = slot_mapping[:, i]
+        kv_cache = kv_cache.at[start_idx:start_idx + slice_len].set(
+            new_kv[new_kv_idx:new_kv_idx + slice_len])
+    return kv_cache
+
+
+@jtu.with_config(jax_numpy_dtype_promotion="standard")
+class KVCacheUpdateTest(jtu.JaxTestCase):
+
+    def _generate_data(self, page_size, combined_kv_head_num, head_dim):
+        page_num = 20
+        padded_num_tokens = 128
+        prng_key = jax.random.key(1234)
+        kv_cache = jnp.zeros(
+            (page_num * page_size, combined_kv_head_num, head_dim),
+            dtype=jnp.bfloat16)
+        new_kv = jax.random.normal(
+            prng_key, (padded_num_tokens, combined_kv_head_num, head_dim),
+            dtype=jnp.bfloat16)
+        slice_lens = np.array([7, page_size, page_size, 1, 1, 1, 9],
+                              dtype=np.int32)
+        num_slices = jnp.array([len(slice_lens)], dtype=np.int32)
+        kv_cache_start_indices = np.array([
+            page_size * 2 - 7, page_size * 2, page_size * 3, page_size * 4 + 6,
+            page_size * 5 + 7, page_size * 6 + 8, page_size * 15 + 3
+        ],
+                                          dtype=np.int32)
+        new_kv_cache_indices = np.concatenate(
+            [np.array([0], dtype=np.int32),
+             np.cumsum(slice_lens[:-1])])
+        slot_mapping_np = np.stack(
+            [kv_cache_start_indices, new_kv_cache_indices, slice_lens], axis=1)
+        slot_mapping_np = np.transpose(slot_mapping_np)
+        slot_mapping = jnp.array(slot_mapping_np, dtype=jnp.int32)
+        return new_kv, slot_mapping, kv_cache, num_slices
+
+    @parameterized.product(
+        page_size=[32, 33],
+        combined_kv_head_num=[2, 16],
+        head_dim=[128, 256],
+        num_slices_per_block=[None, 8],
+        dynamic_validate_inputs=[False, True],
+    )
+    def test_basic(self, page_size: int, combined_kv_head_num: int,
+                   head_dim: int, num_slices_per_block: int,
+                   dynamic_validate_inputs: bool):
+        new_kv, slot_mapping, kv_cache, num_slices = self._generate_data(
+            page_size, combined_kv_head_num, head_dim)
+        old_kv_cache_copy = kv_cache.copy()
+
+        with jax.disable_jit(disable=dynamic_validate_inputs):
+            updated_kv_cache = kv_cache_update(
+                new_kv,
+                slot_mapping,
+                kv_cache,
+                num_slices,
+                page_size=page_size,
+                num_slices_per_block=num_slices_per_block,
+                dynamic_validate_inputs=dynamic_validate_inputs)
+        updated_kv_cache_ref = kv_cache_update_ref(new_kv,
+                                                   np.asarray(slot_mapping),
+                                                   old_kv_cache_copy)
+        self.assertAllClose(updated_kv_cache,
+                            updated_kv_cache_ref,
+                            atol=1e-4,
+                            rtol=1e-4)
+
+    @parameterized.product(
+        page_size=[32, 33],
+        combined_kv_head_num=[16, 32],
+        head_dim=[128, 256],
+        num_slices_per_block=[None, 8],
+    )
+    def test_torchax_shard_map(self, page_size: int, combined_kv_head_num: int,
+                               head_dim: int, num_slices_per_block: int):
+        new_kv, slot_mapping, kv_cache, num_slices = self._generate_data(
+            page_size, combined_kv_head_num, head_dim)
+        old_kv_cache_copy = kv_cache.copy()
+
+        mesh = Mesh(jax.devices(), 'x')
+        kv_cache_pspec = P(None, 'x', None)
+
+        new_kv = jax.device_put(new_kv, NamedSharding(mesh, kv_cache_pspec))
+        slot_mapping = jax.device_put(slot_mapping, NamedSharding(mesh, P()))
+        kv_cache = jax.device_put(kv_cache,
+                                  NamedSharding(mesh, kv_cache_pspec))
+        num_slices = jax.device_put(num_slices, NamedSharding(mesh, P()))
+
+        updated_kv_cache = kv_cache_update(new_kv, slot_mapping, kv_cache,
+                                           num_slices,
+                                           page_size=page_size,
+                                           num_slices_per_block=\
+                                               num_slices_per_block,
+                                           mesh=mesh,
+                                           kv_cache_pspec=kv_cache_pspec,)
+        updated_kv_cache_ref = kv_cache_update_ref(new_kv,
+                                                   np.asarray(slot_mapping),
+                                                   old_kv_cache_copy)
+        self.assertAllClose(updated_kv_cache,
+                            updated_kv_cache_ref,
+                            atol=1e-4,
+                            rtol=1e-4)
+
+    def test_invalid_inputs(self):
+        # Test all the cases when the inputs are invalid in the `_dynamic_validate_inputs` method
+        page_size = 32
+        combined_kv_head_num = 2
+        head_dim = 128
+
+        new_kv, slot_mapping, kv_cache, num_slices = self._generate_data(
+            page_size, combined_kv_head_num, head_dim)
+
+        with jax.disable_jit():
+            # Case 1: new_kv_start < 0
+            invalid_slot_mapping = slot_mapping.at[1, 0].set(-1)
+            with self.assertRaisesRegex(
+                    ValueError, "new_kv_start=-1 must be greater than"):
+                kv_cache_update(new_kv,
+                                invalid_slot_mapping,
+                                kv_cache,
+                                num_slices,
+                                page_size=page_size,
+                                dynamic_validate_inputs=True)
+
+            # Case 2: kv_cache_start < 0
+            invalid_slot_mapping = slot_mapping.at[0, 0].set(-1)
+            with self.assertRaisesRegex(
+                    ValueError, "kv_cache_start=-1 must be greater than"):
+                kv_cache_update(new_kv,
+                                invalid_slot_mapping,
+                                kv_cache,
+                                num_slices,
+                                page_size=page_size,
+                                dynamic_validate_inputs=True)
+
+            # Case 3: slice_len <= 0
+            invalid_slot_mapping = slot_mapping.at[2, 0].set(0)
+            with self.assertRaisesRegex(
+                    ValueError, "slice_len=0 must be less or equal to"):
+                kv_cache_update(new_kv,
+                                invalid_slot_mapping,
+                                kv_cache,
+                                num_slices,
+                                page_size=page_size,
+                                dynamic_validate_inputs=True)
+
+            # Case 4: slice_len > page_size
+            invalid_slot_mapping = slot_mapping.at[2, 0].set(page_size + 1)
+            with self.assertRaisesRegex(
+                    ValueError,
+                    f"slice_len={page_size + 1} must be less or equal to"):
+                kv_cache_update(new_kv,
+                                invalid_slot_mapping,
+                                kv_cache,
+                                num_slices,
+                                page_size=page_size,
+                                dynamic_validate_inputs=True)
+
+            # Case 5: new_kv_start + slice_len > new_token_num
+            invalid_slot_mapping = slot_mapping.at[1, 0].set(new_kv.shape[0])
+            with self.assertRaisesRegex(
+                    ValueError,
+                    "new_kv_start=128 \+ slice_len=7 must be less or equal to new_token_num=128"
+            ):
+                kv_cache_update(new_kv,
+                                invalid_slot_mapping,
+                                kv_cache,
+                                num_slices,
+                                page_size=page_size,
+                                dynamic_validate_inputs=True)
+
+            # Case 6: kv_cache_start + slice_len > kv_cache_token_num
+            invalid_slot_mapping = slot_mapping.at[0, 0].set(kv_cache.shape[0])
+            with self.assertRaisesRegex(
+                    ValueError,
+                    "kv_cache_start=640 \+ slice_len=7 must be less or equal to kv_cache_token_num=640"
+            ):
+                kv_cache_update(new_kv,
+                                invalid_slot_mapping,
+                                kv_cache,
+                                num_slices,
+                                page_size=page_size,
+                                dynamic_validate_inputs=True)
+
+            # Case 7: Each slice must reside in the same page
+            invalid_slot_mapping = slot_mapping.at[0, 0].set(page_size - 1)
+            invalid_slot_mapping = invalid_slot_mapping.at[2, 0].set(page_size)
+            with self.assertRaisesRegex(
+                    ValueError, "Each slice must reside in the same page"):
+                kv_cache_update(new_kv,
+                                invalid_slot_mapping,
+                                kv_cache,
+                                num_slices,
+                                page_size=page_size,
+                                dynamic_validate_inputs=True)
+
+            # Case 8: new_kv slices are not continuous
+            invalid_slot_mapping = slot_mapping.at[1,
+                                                   1].set(slot_mapping[1, 1] +
+                                                          1)
+            with self.assertRaisesRegex(ValueError, "is expeced to equal to"):
+                kv_cache_update(new_kv,
+                                invalid_slot_mapping,
+                                kv_cache,
+                                num_slices,
+                                page_size=page_size,
+                                dynamic_validate_inputs=True)
+
+            # Case 9: Overlap among the kv cache slices
+            invalid_slot_mapping = slot_mapping.at[0, 4].set(slot_mapping[0,
+                                                                          3])
+            with self.assertRaisesRegex(
+                    ValueError, "Overlap detected in kv_cache intervals"):
+                kv_cache_update(new_kv,
+                                invalid_slot_mapping,
+                                kv_cache,
+                                num_slices,
+                                page_size=page_size,
+                                dynamic_validate_inputs=True)