tpu-inference 0.12.0.dev20251213__py3-none-any.whl → 0.13.2.dev20251230__py3-none-any.whl

tests/layers/jax/sample/test_sampling.py

@@ -0,0 +1,115 @@
+# 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.
+
+# /home/pooyam/tpu_inference/tests/models/jax/layers/test_sampling.py
+import jax.numpy as jnp
+import numpy as np
+from vllm.v1.outputs import LogprobsTensors
+
+from tpu_inference.layers.jax.sample.sampling import (compute_logprobs,
+                                                      gather_logprobs)
+
+
+class TestSampling:
+
+    def test_compute_logprobs(self):
+        logits = jnp.array([[1.0, 2.0, 3.0], [3.0, 2.0, 1.0]],
+                           dtype=jnp.float32)
+        logprobs = compute_logprobs(logits)
+
+        # Expected values computed with scipy.special.log_softmax
+        expected_logprobs = np.array(
+            [
+                [-2.40760596, -1.40760596, -0.40760596],
+                [-0.40760596, -1.40760596, -2.40760596],
+            ],
+            dtype=np.float32,
+        )
+        assert np.allclose(logprobs, expected_logprobs, atol=1e-6)
+
+    def test_gather_logprobs(self):
+        logprobs = jnp.array(
+            [
+                [-2.40760596, -1.40760596, -0.40760596, -3.40760596],
+                [-0.40760596, -1.40760596, -2.40760596, -3.40760596],
+            ],
+            dtype=jnp.float32,
+        )
+        token_ids = jnp.array([2, 0], dtype=jnp.int32)
+        num_logprobs = 2
+
+        result: LogprobsTensors = gather_logprobs(logprobs, token_ids,
+                                                  num_logprobs)
+
+        # check indices
+        expected_indices = np.array(
+            [
+                [2, 2, 1],  # token id 2, top-k are 2, 1
+                [0, 0, 1],  # token id 0, top-k are 0, 1
+            ],
+            dtype=np.int32,
+        )
+        assert np.array_equal(result.logprob_token_ids, expected_indices)
+
+        # check logprobs
+        expected_logprobs_values = np.array(
+            [
+                [-0.40760596, -0.40760596, -1.40760596],
+                [-0.40760596, -0.40760596, -1.40760596],
+            ],
+            dtype=np.float32,
+        )
+        assert np.allclose(result.logprobs,
+                           expected_logprobs_values,
+                           atol=1e-6)
+
+        # check ranks
+        expected_ranks = np.array([1, 1], dtype=np.int32)
+        assert np.array_equal(result.selected_token_ranks, expected_ranks)
+
+    def test_gather_logprobs_with_ties(self):
+        logprobs = jnp.array(
+            [
+                [-1.0, -1.0, -2.0, -2.0],
+            ],
+            dtype=jnp.float32,
+        )
+        token_ids = jnp.array([1], dtype=jnp.int32)
+        num_logprobs = 3
+
+        result: LogprobsTensors = gather_logprobs(logprobs, token_ids,
+                                                  num_logprobs)
+
+        # check logprobs
+        expected_logprobs_values = np.array(
+            [
+                [-1.0, -1.0, -1.0, -2.0],
+            ],
+            dtype=np.float32,
+        )
+        assert np.allclose(result.logprobs,
+                           expected_logprobs_values,
+                           atol=1e-6)
+
+        # check ranks
+        # rank of token 1 is 2 because there are 2 values >= -1.0
+        expected_ranks = np.array([2], dtype=np.int32)
+        assert np.array_equal(result.selected_token_ranks, expected_ranks)
+
+        # check indices
+        # The order of tied elements is not guaranteed.
+        # token id is 1. top-k indices are a permutation of {0, 1, 2} or {0, 1, 3}.
+        assert result.logprob_token_ids[0, 0] == 1
+        top_k_indices = sorted(result.logprob_token_ids[0, 1:].tolist())
+        assert top_k_indices == [0, 1, 2] or top_k_indices == [0, 1, 3]

tests/layers/jax/sample/test_sampling_metadata.py

@@ -0,0 +1,254 @@
+# 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 dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from jax.sharding import Mesh, NamedSharding, PartitionSpec
+
+from tpu_inference.layers.jax.sample.sampling_metadata import (
+    DEFAULT_SAMPLING_PARAMS, TPUSupportedSamplingMetadata)
+
+## Mocks and Fixtures
+
+
+@dataclass
+class MockInputBatch:
+    """A mock of the InputBatch class, using NumPy arrays for CPU tensors."""
+
+    all_greedy: bool
+    num_reqs: int = 0
+    temperature_cpu: np.ndarray = None
+    top_k_cpu: np.ndarray = None
+    top_p_cpu: np.ndarray = None
+    max_num_logprobs: int = None
+
+
+@pytest.fixture(scope="module")
+def mesh() -> Mesh:
+    """Creates a 1D JAX mesh for testing on available devices."""
+    if not jax.devices():
+        pytest.skip("No JAX devices available for testing.")
+    return Mesh(np.array(jax.devices()), axis_names=("data", ))
+
+
+## Test Cases
+
+
+def test_from_input_batch_all_greedy(mesh: Mesh):
+    """
+    Tests TPUSupportedSamplingMetadata.from_input_batch when **all_greedy is True**.
+
+    It should return an object with `do_sampling=False` and `None` for the tensors.
+    """
+    mock_batch = MockInputBatch(all_greedy=True)
+    padded_num_reqs = 4
+
+    metadata = TPUSupportedSamplingMetadata.from_input_batch(
+        mesh=mesh, input_batch=mock_batch, padded_num_reqs=padded_num_reqs)
+
+    assert not metadata.do_sampling, "do_sampling should be False for greedy requests"
+    assert metadata.temperature is None
+    assert metadata.top_k is None
+    assert metadata.top_p is None
+
+
+def test_from_input_batch_with_sampling_and_padding(mesh: Mesh):
+    """
+    Tests TPUSupportedSamplingMetadata.from_input_batch with sampling enabled,
+    requiring the tensors to be **padded** to the correct shape.
+    """
+    num_reqs = 2
+    padded_num_reqs = 4
+
+    # Input tensors must be large enough to hold the padded values.
+    temp_tensor = np.array([0.7, 0.8, 0.0, 0.0], dtype=np.float32)
+    top_k_tensor = np.array([10, 20, 0, 0], dtype=np.int32)
+    top_p_tensor = np.array([0.9, 0.95, 0.0, 0.0], dtype=np.float32)
+
+    mock_batch = MockInputBatch(
+        all_greedy=False,
+        num_reqs=num_reqs,
+        temperature_cpu=temp_tensor,
+        top_k_cpu=top_k_tensor,
+        top_p_cpu=top_p_tensor,
+    )
+
+    metadata = TPUSupportedSamplingMetadata.from_input_batch(
+        mesh=mesh, input_batch=mock_batch, padded_num_reqs=padded_num_reqs)
+
+    # 1. Check metadata flags and types
+    assert metadata.do_sampling, "do_sampling should be True"
+    assert isinstance(metadata.temperature, jnp.ndarray)
+    assert isinstance(metadata.top_k, jnp.ndarray)
+    assert isinstance(metadata.top_p, jnp.ndarray)
+
+    # 2. Check shapes
+    assert metadata.temperature.shape == (padded_num_reqs, )
+    assert metadata.top_k.shape == (padded_num_reqs, )
+    assert metadata.top_p.shape == (padded_num_reqs, )
+
+    # 3. Check sharding (should be fully replicated)
+    expected_sharding = NamedSharding(mesh, PartitionSpec(None))
+    assert metadata.temperature.sharding == expected_sharding
+    assert metadata.top_k.sharding == expected_sharding
+    assert metadata.top_p.sharding == expected_sharding
+
+    # 4. Check that values were correctly padded
+    expected_temp = np.array(
+        [
+            0.7, 0.8, DEFAULT_SAMPLING_PARAMS["temperature"],
+            DEFAULT_SAMPLING_PARAMS["temperature"]
+        ],
+        dtype=np.float32,
+    )
+    expected_top_k = np.array(
+        [
+            10, 20, DEFAULT_SAMPLING_PARAMS["top_k"],
+            DEFAULT_SAMPLING_PARAMS["top_k"]
+        ],
+        dtype=np.int32,
+    )
+    expected_top_p = np.array(
+        [
+            0.9, 0.95, DEFAULT_SAMPLING_PARAMS["top_p"],
+            DEFAULT_SAMPLING_PARAMS["top_p"]
+        ],
+        dtype=np.float32,
+    )
+
+    np.testing.assert_allclose(np.asarray(metadata.temperature), expected_temp)
+    np.testing.assert_array_equal(np.asarray(metadata.top_k), expected_top_k)
+    np.testing.assert_allclose(np.asarray(metadata.top_p), expected_top_p)
+
+
+def test_from_input_batch_no_padding_needed(mesh: Mesh):
+    """
+    Tests the case where `num_reqs` equals `padded_num_reqs`, so **no padding** should occur.
+    """
+    num_reqs = 4
+    padded_num_reqs = 4
+
+    temp_tensor = np.array([0.7, 0.8, 0.6, 0.5], dtype=np.float32)
+    top_k_tensor = np.array([10, 20, 30, 40], dtype=np.int32)
+    top_p_tensor = np.array([0.9, 0.95, 0.85, 0.8], dtype=np.float32)
+
+    mock_batch = MockInputBatch(
+        all_greedy=False,
+        num_reqs=num_reqs,
+        temperature_cpu=temp_tensor,
+        top_k_cpu=top_k_tensor,
+        top_p_cpu=top_p_tensor,
+    )
+
+    metadata = TPUSupportedSamplingMetadata.from_input_batch(
+        mesh=mesh, input_batch=mock_batch, padded_num_reqs=padded_num_reqs)
+
+    assert metadata.do_sampling
+    # Check that values are identical to the input, since no padding was needed
+    np.testing.assert_allclose(np.asarray(metadata.temperature), temp_tensor)
+    np.testing.assert_array_equal(np.asarray(metadata.top_k), top_k_tensor)
+    np.testing.assert_allclose(np.asarray(metadata.top_p), top_p_tensor)
+
+
+def test_jax_tree_util_registration():
+    """
+    Tests that the dataclass is correctly registered as a **JAX PyTree**,
+    meaning `jax.tree_util` functions can operate on it as expected. 🌳
+    """
+    metadata = TPUSupportedSamplingMetadata(
+        temperature=jnp.array([0.7]),
+        top_k=jnp.array([10]),
+        top_p=jnp.array([0.9]),
+        do_sampling=True,
+    )
+
+    # Flatten the PyTree
+    leaves, treedef = jax.tree_util.tree_flatten(metadata)
+
+    # The leaves should be the "data_fields" specified in the decorator
+    assert len(leaves) == 3
+    np.testing.assert_array_equal(leaves[0], jnp.array([0.7]))
+    np.testing.assert_array_equal(leaves[1], jnp.array([10]))
+    np.testing.assert_array_equal(leaves[2], jnp.array([0.9]))
+
+    # Reconstruct the PyTree from leaves
+    new_metadata = jax.tree_util.tree_unflatten(treedef, leaves)
+
+    # The reconstructed object should match the original
+    assert new_metadata.do_sampling == metadata.do_sampling
+    np.testing.assert_array_equal(new_metadata.temperature,
+                                  metadata.temperature)
+    np.testing.assert_array_equal(new_metadata.top_k, metadata.top_k)
+    np.testing.assert_array_equal(new_metadata.top_p, metadata.top_p)
+
+
+def test_from_input_batch_with_logprobs(mesh: Mesh):
+    """
+    Tests that the `logprobs` flag is correctly set based on `max_num_logprobs`.
+    """
+    # Case 1: Logprobs are requested
+    mock_batch_with_logprobs = MockInputBatch(all_greedy=True,
+                                              max_num_logprobs=5)
+    metadata_with = TPUSupportedSamplingMetadata.from_input_batch(
+        mesh=mesh,
+        input_batch=mock_batch_with_logprobs,
+        padded_num_reqs=4,
+    )
+    assert metadata_with.logprobs, "logprobs should be True when max_num_logprobs > 0"
+
+    # Case 2: Logprobs are not requested (max_num_logprobs is 0)
+    mock_batch_no_logprobs_zero = MockInputBatch(all_greedy=True,
+                                                 max_num_logprobs=0)
+    metadata_without_zero = TPUSupportedSamplingMetadata.from_input_batch(
+        mesh=mesh,
+        input_batch=mock_batch_no_logprobs_zero,
+        padded_num_reqs=4,
+    )
+    assert not metadata_without_zero.logprobs, "logprobs should be False when max_num_logprobs is 0"
+
+    # Case 3: Logprobs are not requested (max_num_logprobs is None)
+    mock_batch_no_logprobs_none = MockInputBatch(all_greedy=True,
+                                                 max_num_logprobs=None)
+    metadata_without_none = TPUSupportedSamplingMetadata.from_input_batch(
+        mesh=mesh,
+        input_batch=mock_batch_no_logprobs_none,
+        padded_num_reqs=4,
+    )
+    assert not metadata_without_none.logprobs, "logprobs should be False when max_num_logprobs is None"
+
+
+def test_from_input_batch_sampling_with_logprobs(mesh: Mesh):
+    """
+    Tests enabling both sampling and logprobs simultaneously.
+    """
+    num_reqs = 2
+    padded_num_reqs = 4
+    mock_batch = MockInputBatch(
+        all_greedy=False,
+        num_reqs=num_reqs,
+        temperature_cpu=np.zeros((padded_num_reqs, ), dtype=np.float32),
+        top_k_cpu=np.zeros((padded_num_reqs, ), dtype=np.int32),
+        top_p_cpu=np.zeros((padded_num_reqs, ), dtype=np.float32),
+        max_num_logprobs=10,
+    )
+
+    metadata = TPUSupportedSamplingMetadata.from_input_batch(
+        mesh=mesh, input_batch=mock_batch, padded_num_reqs=padded_num_reqs)
+
+    assert metadata.do_sampling, "do_sampling should be True"
+    assert metadata.logprobs, "logprobs should be True"

tests/layers/jax/test_layers.py

@@ -0,0 +1,155 @@
+# 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 unittest
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from flax import nnx
+from jax.sharding import Mesh
+
+from tpu_inference.layers.jax.layers import DenseFFW, Embedder, RMSNorm
+
+
+class TestLayers(unittest.TestCase):
+    """Unit test suite for common JAX layer blocks."""
+
+    def setUp(self):
+        """Sets up the testing environment before each test."""
+        self.mesh = Mesh(
+            np.array(jax.devices()).reshape(1, -1),
+            axis_names=(
+                "expert",
+                "model",
+            ),
+        )
+
+    def test_rmsnorm_forward_pass(self):
+        """Tests the forward pass of the RMSNorm module."""
+        with jax.set_mesh(self.mesh):
+            dims = 512
+            epsilon = 1e-5
+
+            norm = RMSNorm(
+                dims=dims,
+                random_init=True,
+                epsilon=epsilon,
+                rngs=nnx.Rngs(0),
+                dtype=jnp.float32,
+            )
+
+            seq_len = 128
+            x = jax.random.normal(jax.random.PRNGKey(42), (seq_len, dims))
+
+            output = norm(x)
+
+            self.assertEqual(output.shape, x.shape)
+            self.assertEqual(output.dtype, jnp.float32)
+
+            mean_of_squares = jnp.mean(jnp.square(output), axis=-1)
+            self.assertTrue(
+                jnp.allclose(mean_of_squares, 1.0, atol=1e-5).all())
+
+    def test_denseffw_forward_pass(self):
+        """Tests the forward pass of the DenseFFW module."""
+        with jax.set_mesh(self.mesh):
+            hidden_size = 512
+            intermediate_size = 2048
+
+            ffw_layer = DenseFFW(
+                random_init=True,
+                dtype=jnp.bfloat16,
+                hidden_act="silu",
+                hidden_size=hidden_size,
+                intermediate_size=intermediate_size,
+                rngs=nnx.Rngs(0),
+            )
+
+            seq_len = 128
+            x = jnp.ones((seq_len, hidden_size), dtype=jnp.bfloat16)
+
+            output = ffw_layer(x)
+
+            self.assertEqual(output.shape, x.shape)
+            self.assertEqual(output.dtype, x.dtype)
+
+    def test_embedder_forward_pass(self):
+        """Tests both the encode and decode passes of the Embedder module."""
+        with jax.set_mesh(self.mesh):
+            hidden_size = 512
+            vocab_size = 32000
+            dtype = jnp.bfloat16
+
+            embedder = Embedder(
+                vocab_size=vocab_size,
+                hidden_size=hidden_size,
+                dtype=dtype,
+                random_init=True,
+                rngs=nnx.Rngs(0),
+            )
+
+            seq_len = 128
+            token_ids = jnp.arange(seq_len, dtype=jnp.int32) % vocab_size
+            embeddings = embedder(token_ids, decode=False)
+            self.assertEqual(embeddings.shape, (seq_len, hidden_size))
+            self.assertEqual(embeddings.dtype, dtype)
+
+            hidden_states = jnp.ones((seq_len, hidden_size),
+                                     dtype=jnp.bfloat16)
+            logits = embedder(hidden_states, decode=True)
+            self.assertEqual(logits.shape, (seq_len, vocab_size))
+            self.assertEqual(logits.dtype, dtype)
+
+    def test_embedder_normalization(self):
+        """Tests the embedding normalization feature."""
+        with jax.set_mesh(self.mesh):
+            hidden_size = 512
+            vocab_size = 32000
+
+            rngs_1 = nnx.Rngs(42)
+            rngs_2 = nnx.Rngs(42)
+
+            embedder_norm = Embedder(
+                vocab_size=vocab_size,
+                hidden_size=hidden_size,
+                dtype=jnp.float32,
+                normalize_embeddings=True,
+                random_init=True,
+                rngs=rngs_1,
+            )
+
+            embedder_no_norm = Embedder(
+                vocab_size=vocab_size,
+                hidden_size=hidden_size,
+                dtype=jnp.float32,
+                random_init=True,
+                rngs=rngs_2,
+            )
+
+            token_ids = jnp.arange(10, dtype=jnp.int32)
+
+            embeddings_norm = embedder_norm(token_ids, decode=False)
+            embeddings_no_norm = embedder_no_norm(token_ids, decode=False)
+
+            scaling_factor = jnp.sqrt(hidden_size)
+            expected_embeddings = embeddings_no_norm * scaling_factor
+
+            self.assertTrue(
+                jnp.allclose(embeddings_norm, expected_embeddings,
+                             atol=1e-6).all())
+
+
+if __name__ == "__main__":
+    unittest.main()