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

tests/layers/jax/test_rope.py

@@ -0,0 +1,93 @@
+# 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
+from jax import numpy as jnp
+from jax._src import test_util as jtu
+from jax.sharding import Mesh
+
+from tpu_inference.layers.jax.rope import (DeepseekScalingRotaryEmbedding,
+                                           RotaryEmbedding)
+
+
+class RotaryEmbeddingTest(jtu.JaxTestCase):
+
+    def test_apply_rope(self):
+        head_dim = 2
+        rope_theta = 10000
+        original_max_position_embeddings = 2
+        rope = RotaryEmbedding(
+            rotary_dim=head_dim,
+            rope_theta=rope_theta,
+            original_max_position_embeddings=original_max_position_embeddings,
+            dtype=jnp.float32)
+        rope.initialize_cache()
+        self.assertTrue(
+            rope.sin_cos_cache.shape == (original_max_position_embeddings,
+                                         head_dim))
+        expected_sin_cos = jnp.array([[1, 0], [0.5403023, 0.841471]],
+                                     dtype=jnp.float32)
+        self.assertArraysAllClose(rope.sin_cos_cache, expected_sin_cos)
+
+        num_tokens = 2
+        num_heads = 1
+        positions = jnp.arange(num_tokens)
+        x = jnp.ones((num_tokens, num_heads, head_dim))
+        x_rope = rope.apply_rope(positions, x)
+        expected_x_rope = jnp.array([[[1, 1]], [[-0.30116874, 1.3817732]]],
+                                    dtype=jnp.float32)
+        self.assertTrue(x_rope.shape == x.shape)
+        self.assertArraysAllClose(x_rope, expected_x_rope)
+
+
+class DeepseekScalingRotaryEmbeddingTest(jtu.JaxTestCase):
+
+    def test_apply_rope(self):
+        head_dim = 2
+        rope_theta = 10000
+        original_max_position_embeddings = 1
+        scaling_factor = 2
+        devices = jax.devices()
+        mesh = Mesh(devices, ('data', ))
+
+        rope = DeepseekScalingRotaryEmbedding(
+            rotary_dim=head_dim,
+            rope_theta=rope_theta,
+            original_max_position_embeddings=original_max_position_embeddings,
+            scaling_factor=scaling_factor,
+            dtype=jnp.float32)
+        rope.initialize_cache(mesh)
+        expected_padded_dim = 128
+        self.assertTrue(
+            rope.sin_cos_cache.shape == (scaling_factor *
+                                         original_max_position_embeddings,
+                                         expected_padded_dim))
+
+        valid_cache_slice = rope.sin_cos_cache[:, :head_dim]
+
+        expected_sin_cos = jnp.array([[1.0693147, 0], [0.5777532, 0.8997973]],
+                                     dtype=jnp.float32)
+
+        self.assertArraysAllClose(valid_cache_slice, expected_sin_cos)
+
+        num_tokens = 2
+        num_heads = 1
+        positions = jnp.arange(num_tokens)
+        x = jnp.ones((num_tokens, num_heads, head_dim))
+        x_rope = rope.apply_rope(positions, x)
+        expected_x_rope = jnp.array(
+            [[[1.0693147, 1.0693147]], [[-0.32204413, 1.4775505]]],
+            dtype=jnp.float32)
+        self.assertTrue(x_rope.shape == x.shape)
+        self.assertArraysAllClose(x_rope, expected_x_rope)

tests/layers/jax/test_sharding.py

@@ -0,0 +1,159 @@
+# 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
+from unittest.mock import MagicMock
+
+import jax
+
+from tpu_inference.layers.common.sharding import (Sharding, ShardingConfig,
+                                                  ShardingRulesConfig,
+                                                  ShardingStrategy)
+
+
+class TestSharding(unittest.TestCase):
+    """Unit test suite for the sharding configuration logic."""
+
+    def setUp(self):
+        """Sets up the testing environment before each test."""
+
+        self.mock_devices = [MagicMock(coords=i) for i in range(8)]
+        self.original_jax_devices = jax.devices
+        jax.devices = lambda: self.mock_devices
+
+    def tearDown(self):
+        """Restores the original jax.devices function after tests."""
+        jax.devices = self.original_jax_devices
+
+    def test_sharding_strategy_init(self):
+        """Tests the initialization of the ShardingStrategy."""
+        strategy = ShardingStrategy(
+            tensor_parallelism=2,
+            expert_parallelism=4,
+            data_parallelism=1,
+            sequence_parallelism=1,
+        )
+        self.assertEqual(strategy.tensor_parallelism, 2)
+        self.assertEqual(strategy.expert_parallelism, 4)
+
+    def test_sharding_config_init(self):
+        """Tests the initialization of ShardingConfig."""
+        config = ShardingConfig()
+        self.assertIsInstance(config.prefill_rules, ShardingRulesConfig)
+        self.assertIsInstance(config.generate_rules, ShardingRulesConfig)
+
+        custom_rules = ShardingRulesConfig(activation_ffw_td=("model", None))
+        config_with_rules = ShardingConfig(prefill_rules=custom_rules)
+        self.assertEqual(config_with_rules.prefill_rules.activation_ffw_td,
+                         ("model", None))
+
+    def test_apply_overrides(self):
+        """Tests the _apply_overrides method for valid and invalid keys."""
+        sharding = Sharding(
+            prefill_rules={},
+            generate_rules={},
+        )
+        config_obj = ShardingRulesConfig()
+
+        valid_overrides = {"activation_ffw_td": ("model", None)}
+        sharding._apply_overrides(config_obj, valid_overrides)
+        self.assertEqual(config_obj.activation_ffw_td, ("model", None))
+
+        invalid_overrides = {"non_existent_attribute": (None, "model")}
+        with self.assertRaises(AttributeError):
+            sharding._apply_overrides(config_obj, invalid_overrides)
+
+    def test_default_sharding_config(self):
+        """Tests that default sharding rules are created correctly."""
+        sharding = Sharding(
+            prefill_rules={},
+            generate_rules={},
+        )
+
+        sharding_cfg = sharding.get_sharding_cfg()
+        generate_rules = sharding_cfg.generate_rules
+
+        self.assertEqual(generate_rules.ffw_weight_df, (None, "model"))
+        self.assertEqual(generate_rules.moe_router_de, (None, "model"))
+        self.assertEqual(generate_rules.attn_q_weight_dnh,
+                         (None, "model", None))
+
+    def test_sharding_init_with_overrides(self):
+        """Tests Sharding initialization with programmatic overrides."""
+        generate_overrides = {"logits_tv": ("data", "model")}
+
+        sharding = Sharding(
+            generate_rules=generate_overrides,
+            prefill_rules={},
+        )
+
+        sharding_cfg = sharding.get_sharding_cfg()
+        self.assertNotEqual(sharding_cfg.generate_rules.logits_tv,
+                            (None, "model"))
+        self.assertEqual(sharding_cfg.generate_rules.logits_tv,
+                         ("data", "model"))
+
+    def test_get_overrides_from_vllm_config(self):
+        """Tests fetching sharding overrides from a mock VllmConfig."""
+
+        mock_vllm_config_prefill = MagicMock()
+        mock_vllm_config_prefill.additional_config = {
+            "sharding": {
+                "logical_rules": {
+                    "all": {
+                        "norm_scale": ("model", )
+                    },
+                    "prefill": {
+                        "activation_ffw_td": ("data", "model")
+                    },
+                }
+            }
+        }
+        sharding_prefill = Sharding(
+            vllm_config=mock_vllm_config_prefill,
+            prefill_rules={},
+            generate_rules={},
+        )
+        prefill_overrides = sharding_prefill._get_overrides("prefill")
+
+        self.assertEqual(prefill_overrides["norm_scale"], ("model", ))
+        self.assertEqual(prefill_overrides["activation_ffw_td"],
+                         ("data", "model"))
+
+        mock_vllm_config_generate = MagicMock()
+        mock_vllm_config_generate.additional_config = {
+            "sharding": {
+                "logical_rules": {
+                    "all": {
+                        "norm_scale": ("model", )
+                    },
+                    "prefill": {
+                        "activation_ffw_td": ("data", "model")
+                    },
+                }
+            }
+        }
+        sharding_generate = Sharding(
+            vllm_config=mock_vllm_config_generate,
+            prefill_rules={},
+            generate_rules={},
+        )
+        generate_overrides = sharding_generate._get_overrides("generate")
+
+        self.assertEqual(generate_overrides["norm_scale"], ("model", ))
+        self.assertNotIn("activation_ffw_td", generate_overrides)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/layers/jax/test_transformer_block.py

@@ -0,0 +1,152 @@
+# 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
+from unittest.mock import MagicMock
+
+import jax.numpy as jnp
+from flax import nnx
+
+from tpu_inference.layers.jax.layers import DenseFFW
+from tpu_inference.layers.jax.moe.moe import MoE
+from tpu_inference.layers.jax.transformer_block import (
+    SharedExpertsTransformerBlock, TransformerBlock)
+
+
+class TestTransformerBlock(unittest.TestCase):
+    """Unit test suite for the JAX TransformerBlock module."""
+
+    def test_transformer_block_dense_logic(self):
+        """
+        Tests the forward pass logic of a dense TransformerBlock by mocking its sub-modules.
+        This test verifies the sequence of operations and residual connections.
+        """
+        hidden_size = 1024
+
+        mock_pre_attn_norm = MagicMock(spec=nnx.Module)
+        mock_pre_mlp_norm = MagicMock(spec=nnx.Module)
+
+        mock_attn = MagicMock(spec=nnx.Module)
+        dummy_attn_output = jnp.full((64, hidden_size),
+                                     2.0,
+                                     dtype=jnp.bfloat16)
+        dummy_kv_cache = jnp.zeros((8, 16, 16, 128), dtype=jnp.bfloat16)
+        mock_attn.return_value = (dummy_kv_cache, dummy_attn_output)
+
+        mock_mlp = MagicMock(spec=DenseFFW)
+        dummy_mlp_output = jnp.full((64, hidden_size), 3.0, dtype=jnp.bfloat16)
+        mock_mlp.return_value = dummy_mlp_output
+
+        transformer_block = TransformerBlock(
+            pre_attention_norm=mock_pre_attn_norm,
+            pre_mlp_norm=mock_pre_mlp_norm,
+            custom_module=mock_mlp,
+            attn=mock_attn,
+        )
+
+        seq_len = 64
+        x = jnp.ones((seq_len, hidden_size), dtype=jnp.bfloat16)
+        initial_kv_cache = MagicMock()
+        attention_metadata = MagicMock()
+
+        mock_pre_attn_norm.side_effect = lambda val: val
+        mock_pre_mlp_norm.side_effect = lambda val: val
+
+        new_kv_cache, final_output = transformer_block(
+            x,
+            is_prefill=True,
+            kv_cache=initial_kv_cache,
+            attention_metadata=attention_metadata,
+        )
+
+        mock_pre_attn_norm.assert_called_once()
+        self.assertTrue(
+            jnp.array_equal(mock_pre_attn_norm.call_args.args[0], x))
+
+        mock_attn.assert_called_once_with(x, True, initial_kv_cache,
+                                          attention_metadata, True)
+
+        expected_mlp_norm_input = dummy_attn_output + x
+
+        mock_pre_mlp_norm.assert_called_once()
+        self.assertTrue(
+            jnp.array_equal(mock_pre_mlp_norm.call_args.args[0],
+                            expected_mlp_norm_input))
+
+        mock_mlp.assert_called_once()
+        self.assertTrue(
+            jnp.array_equal(mock_mlp.call_args.args[0],
+                            expected_mlp_norm_input))
+
+        expected_final_output = dummy_mlp_output + expected_mlp_norm_input
+        self.assertTrue(jnp.allclose(final_output, expected_final_output))
+
+        self.assertTrue(jnp.array_equal(new_kv_cache, dummy_kv_cache))
+
+    def test_shared_experts_transformer_block_logic(self):
+        """Tests the forward pass logic of a SharedExpertsTransformerBlock."""
+        hidden_size = 1024
+
+        mock_pre_attn_norm = MagicMock(spec=nnx.Module)
+        mock_pre_mlp_norm = MagicMock(spec=nnx.Module)
+
+        mock_attn = MagicMock(spec=nnx.Module)
+        dummy_attn_output = jnp.full((64, hidden_size),
+                                     2.0,
+                                     dtype=jnp.bfloat16)
+        dummy_kv_cache = jnp.zeros((8, 16, 16, 128), dtype=jnp.bfloat16)
+        mock_attn.return_value = (dummy_kv_cache, dummy_attn_output)
+
+        mock_moe = MagicMock(spec=MoE)
+        dummy_moe_output = jnp.full((64, hidden_size), 3.0, dtype=jnp.bfloat16)
+        mock_moe.return_value = dummy_moe_output
+
+        mock_shared_experts = MagicMock(spec=DenseFFW)
+        dummy_shared_experts_output = jnp.full((64, hidden_size),
+                                               4.0,
+                                               dtype=jnp.bfloat16)
+        mock_shared_experts.return_value = dummy_shared_experts_output
+
+        transformer_block = SharedExpertsTransformerBlock(
+            pre_attention_norm=mock_pre_attn_norm,
+            pre_mlp_norm=mock_pre_mlp_norm,
+            custom_module=mock_moe,
+            attn=mock_attn,
+            shared_experts=mock_shared_experts,
+        )
+
+        seq_len = 64
+        x = jnp.ones((seq_len, hidden_size), dtype=jnp.bfloat16)
+        initial_kv_cache = MagicMock()
+        attention_metadata = MagicMock()
+
+        mock_pre_attn_norm.side_effect = lambda val: val
+        mock_pre_mlp_norm.side_effect = lambda val: val
+
+        new_kv_cache, final_output = transformer_block(
+            x,
+            is_prefill=True,
+            kv_cache=initial_kv_cache,
+            attention_metadata=attention_metadata,
+        )
+        self.assertTrue(jnp.array_equal(new_kv_cache, dummy_kv_cache))
+        self.assertEqual(final_output.shape, (seq_len, hidden_size))
+
+        self.assertEqual(mock_moe.call_count, 1)
+        self.assertEqual(mock_attn.call_count, 1)
+        self.assertEqual(mock_shared_experts.call_count, 1)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/layers/vllm/__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.