tpu-inference 0.11.1.dev202512030818__py3-none-any.whl → 0.13.2rc3__py3-none-any.whl

tests/layers/jax/attention/test_llama4_attention.py

@@ -0,0 +1,135 @@
+# 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 os
+import unittest
+from dataclasses import dataclass, field
+
+import chex
+
+os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from jax.sharding import NamedSharding
+from jax.sharding import PartitionSpec as P
+
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.common.sharding import build_mesh
+from tpu_inference.layers.jax.attention.llama4_attention import (
+    L2Norm, Llama4Attention)
+
+
+@dataclass
+class SimpleVLLMConfig:
+    additional_config: dict = field(default_factory=dict)
+
+
+class Llama4AttentionTest(unittest.TestCase):
+    """Unit test suite for Llama4-specific attention components."""
+
+    def setUp(self):
+        devices = jax.devices()[:1]
+        sharding_strategy = {"tensor_parallelism": len(devices)}
+        self.mesh = build_mesh(devices, sharding_strategy)
+
+    def test_l2norm_forward_pass(self):
+        """Tests the forward pass of the L2Norm module with hardcoded values."""
+        eps = 1e-5
+        l2_norm = L2Norm(eps=eps)
+        x = jnp.array([[1.0, 2.0, 3.0, 4.0]], dtype=jnp.float32)
+
+        output = l2_norm(x)
+
+        self.assertEqual(output.shape, x.shape)
+        self.assertEqual(output.dtype, x.dtype)
+
+        # Expected values calculated manually:
+        # mean_sq = (1^2 + 2^2 + 3^2 + 4^2) / 4 = (1+4+9+16)/4 = 30/4 = 7.5
+        # norm_val = sqrt(7.5 + 1e-5)
+        # expected = x / norm_val
+        expected_output = jnp.array([[0.365148, 0.730297, 1.095445, 1.460594]],
+                                    dtype=jnp.float32)
+        self.assertTrue(jnp.allclose(output, expected_output, atol=1e-6))
+
+    def test_l2norm_with_zeros(self):
+        """Tests L2Norm with an all-zero input."""
+        l2_norm = L2Norm(eps=1e-5)
+        x = jnp.zeros((4, 8, 16))
+        output = l2_norm(x)
+        self.assertEqual(output.shape, x.shape)
+        # Output should be all zeros.
+        self.assertTrue(jnp.all(output == 0))
+
+    def test_l2norm_eps_effect(self):
+        """Tests the effect of the epsilon value in L2Norm."""
+        eps = 1e-3
+        l2_norm = L2Norm(eps=eps)
+        x = jax.random.normal(jax.random.PRNGKey(0), (1, 1, 128))
+        output = l2_norm(x)
+
+        mean_sq = jnp.mean(x**2, axis=-1, keepdims=True)
+        expected_output = x * jax.lax.rsqrt(mean_sq + eps)
+
+        self.assertTrue(jnp.allclose(output, expected_output))
+
+    def test_apply_temperature_tuning(self):
+        with jax.set_mesh(self.mesh):
+            hidden_size = 64
+            num_attention_heads = 4
+            head_dim = hidden_size // num_attention_heads
+
+            # Create dummy sharding objects
+            dummy_sharding = NamedSharding(self.mesh, P())
+
+            llama4_attention = Llama4Attention(
+                hidden_size=hidden_size,
+                num_attention_heads=num_attention_heads,
+                num_key_value_heads=num_attention_heads,
+                head_dim=head_dim,
+                rope_theta=10000.0,
+                rope_scaling={},
+                dtype=jnp.bfloat16,
+                kv_cache_dtype="auto",
+                use_qk_norm=False,
+                temperature_tuning=True,
+                temperature_tuning_scale=2.0,
+                temperature_tuning_floor_scale=2.0,
+                mesh=self.mesh,
+                random_init=True,
+                activation_attention_td=dummy_sharding,
+                activation_attention_out_td=dummy_sharding,
+                rngs=nnx.Rngs(42),
+            )
+
+            seq_len = 8
+            input_arr_TNH = jnp.ones((seq_len, num_attention_heads, head_dim),
+                                     dtype=jnp.bfloat16)
+            attention_metadata = AttentionMetadata(
+                input_positions=jnp.arange(seq_len, dtype=jnp.int32))
+            expected_scales = jnp.array(
+                [1, 2.375, 2.375, 3.20312, 3.20312, 3.76562, 3.76562, 4.21875],
+                dtype=jnp.bfloat16)
+            output = llama4_attention.apply_temperature_tuning(
+                attention_metadata, input_arr_TNH)
+            chex.assert_shape(output, (seq_len, num_attention_heads, head_dim))
+
+            expected_output = jnp.ones_like(
+                input_arr_TNH) * expected_scales[:, None, None]
+            chex.assert_trees_all_close(output, expected_output, atol=1e-3)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/layers/jax/moe/__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/jax/moe/test_deepseek_moe.py

@@ -0,0 +1,235 @@
+# 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, PartitionSpec
+
+from tpu_inference.layers.jax.moe.deepseek_v3_moe import (DeepSeekV3Router,
+                                                          SparseMoE)
+
+
+class TestDeepSeekV3Router(unittest.TestCase):
+
+    def setUp(self):
+        self.cpu_mesh = Mesh(jax.devices('cpu'), axis_names=('data', ))
+
+    def test_get_topk_indices_single_group(self):
+        """Test get_topk_indices with single expert group."""
+        with jax.set_mesh(self.cpu_mesh):
+            router = DeepSeekV3Router(random_init=True,
+                                      hidden_size=512,
+                                      num_experts=4,
+                                      num_experts_per_tok=2,
+                                      n_groups=1,
+                                      topk_groups=1,
+                                      norm_topk_prob=True,
+                                      routed_scaling_factor=1.0,
+                                      dtype=jnp.bfloat16,
+                                      rngs=nnx.Rngs(42))
+            router.bias_E = jnp.zeros((4, ))
+
+            scores = jnp.array([[0.1, 0.3, 0.2, 0.4]])  # shape: (1, 4)
+            indices = router.get_topk_indices(scores)
+
+            # Should return indices of top 2 experts
+            expected_indices = jnp.array([[3,
+                                           1]])  # experts with scores 0.4, 0.3
+            self.assertTrue(jnp.array_equal(indices, expected_indices))
+
+    def test_get_topk_indices_2_groups(self):
+        """Test get_topk_indices with 2 expert groups."""
+        with jax.set_mesh(self.cpu_mesh):
+            router = DeepSeekV3Router(random_init=True,
+                                      hidden_size=512,
+                                      num_experts=4,
+                                      num_experts_per_tok=2,
+                                      n_groups=2,
+                                      topk_groups=1,
+                                      norm_topk_prob=True,
+                                      routed_scaling_factor=1.0,
+                                      dtype=jnp.bfloat16,
+                                      rngs=nnx.Rngs(42))
+            router.bias_E = jnp.zeros((4, ))
+
+            # 4 experts, 2 groups, 2 experts per group
+            scores = jnp.array([[[0.1, 0.3, 0.2, 0.4]]])  # shape: (1, 1, 4)
+            indices = router.get_topk_indices(scores)
+
+            # Should return indices of top 2 experts
+            expected_indices = jnp.array([[[3, 2]]])
+            self.assertTrue(jnp.array_equal(indices, expected_indices))
+
+    def test_router_e2e(self):
+        with jax.set_mesh(self.cpu_mesh):
+            router = DeepSeekV3Router(random_init=True,
+                                      hidden_size=512,
+                                      num_experts=8,
+                                      num_experts_per_tok=2,
+                                      n_groups=2,
+                                      topk_groups=1,
+                                      norm_topk_prob=True,
+                                      routed_scaling_factor=1.0,
+                                      dtype=jnp.bfloat16,
+                                      rngs=nnx.Rngs(42))
+            x = jnp.ones((2, 512))
+            weights, indices = router(x)
+            self.assertEqual(weights.shape, (2, 2))
+            self.assertEqual(indices.shape, (2, 2))
+
+
+class TestSparseMoE(unittest.TestCase):
+
+    def setUp(self):
+        """Set up a multi-device mesh and a sample MoE layer for testing."""
+        devices = jax.devices()
+        self.device_count = len(devices)
+        if self.device_count < 8:
+            self.skipTest("This test requires at least 8 simulated devices.")
+
+        # This mesh will have a 'model' axis for expert parallelism
+        mesh_shape = (self.device_count, 1)
+        device_mesh_array = np.array(devices).reshape(mesh_shape)
+
+        # Define the axis names
+        axis_names = ('model', 'data')
+
+        # Create the 2D mesh
+        self.mesh = Mesh(device_mesh_array, axis_names=axis_names)
+
+        # --- Model Configuration ---
+        self.B, self.S, self.D = 2, 4, 16  # Batch, Sequence, Hidden Dim
+        self.E, self.K = 16, 8  # Num Experts, Experts per Token
+        self.moe_intermediate_size = 32  # FFN Dim
+        self.num_expert_parallelism = 8  # Shard experts across 8 devices
+
+        self.key = jax.random.PRNGKey(42)
+        self.x = jax.random.normal(self.key, (self.B * self.S, self.D),
+                                   dtype=jnp.bfloat16)
+
+        # --- Instantiate MoE Layer ---
+        # We need to do this inside the mesh context
+        with self.mesh:
+            router = DeepSeekV3Router(hidden_size=self.D,
+                                      num_experts=self.E,
+                                      num_experts_per_tok=self.K,
+                                      n_groups=1,
+                                      topk_groups=1,
+                                      norm_topk_prob=False,
+                                      routed_scaling_factor=1.0,
+                                      dtype=jnp.bfloat16,
+                                      rngs=nnx.Rngs(self.key),
+                                      ed_sharding=PartitionSpec(),
+                                      e_sharding=PartitionSpec(),
+                                      activation_ffw_td=PartitionSpec(
+                                          'data', None))
+            # Instantiation updated to match user's code snippet
+            self.moe = SparseMoE(
+                hidden_size=self.D,
+                intermediate_size_moe=self.moe_intermediate_size,
+                num_local_experts=self.E,
+                hidden_act="silu",
+                num_experts_per_tok=self.K,
+                router=router,
+                dtype=jnp.bfloat16,
+                rngs=nnx.Rngs(self.key),
+                mesh=self.mesh,
+                apply_expert_weight_before_computation=False,
+
+                # Sharding specs updated based on user's snippet
+                edf_sharding=PartitionSpec('model', None, None),
+                efd_sharding=PartitionSpec('model', None, None),
+                activation_ffw_ted=PartitionSpec('data', None),
+                activation_ffw_td=PartitionSpec(
+                    'data', None)  # Activations are replicated
+            )
+
+    def test_token_replicated_expert_parallel_fwd(self):
+        """
+        Validates the MoE forward pass against a simple, dense equivalent.
+        This specifically tests the is_batch_sharded_by_expert=False path.
+        """
+        # --- 1. Get the ACTUAL output from the complex distributed MoE layer ---
+        # The __call__ method will trigger the shard_map, which requires the mesh context.
+        with self.mesh:
+            actual_output = self.moe(self.x)
+
+        # --- 2. Calculate the EXPECTED output using a simple, sequential process ---
+        # This serves as the "ground truth".
+
+        # Get router decisions (router params are replicated, so this is fine)
+        router_weights, selected_experts = self.moe.router(self.x)
+
+        # Gather the full, unsharded weights from all devices ---
+        # .value on a sharded param gives the *local* shard.
+        # jax.device_get() retrieves the *full* GlobalDeviceArray to the host.
+        gating_kernel_full = jax.device_get(self.moe.kernel_gating_EDF.value)
+        up_proj_kernel_full = jax.device_get(self.moe.kernel_up_proj_EDF.value)
+        down_proj_kernel_full = jax.device_get(
+            self.moe.kernel_down_proj_EFD.value)
+
+        # Check that we really got the full weights
+        self.assertEqual(gating_kernel_full.shape,
+                         (self.E, self.D, self.moe_intermediate_size))
+
+        # Flatten inputs for easier iteration
+        flat_x = self.x.reshape(self.B * self.S, self.D)
+        flat_weights = router_weights.reshape(self.B * self.S, self.K)
+        flat_experts = selected_experts.reshape(self.B * self.S, self.K)
+
+        expected_output = jnp.zeros_like(flat_x)
+
+        # Manually apply each expert to each token sequentially
+        for i in range(self.B * self.S):  # For each token
+            token_input = flat_x[i]
+            combined_expert_output = jnp.zeros(self.D, dtype=jnp.bfloat16)
+
+            for k in range(self.K):  # For each chosen expert for that token
+                expert_idx = flat_experts[i, k]
+                weight = flat_weights[i, k]
+
+                # Get kernels from the *full* gathered arrays ---
+                gating_kernel = gating_kernel_full[expert_idx]
+                up_proj_kernel = up_proj_kernel_full[expert_idx]
+                down_proj_kernel = down_proj_kernel_full[expert_idx]
+
+                # Perform the expert computation (dense matmuls)
+                gating_proj = jnp.dot(token_input, gating_kernel)
+                up_proj = jnp.dot(token_input, up_proj_kernel)
+
+                # Note: Assuming 'silu' activation as specified in MoE init
+                fused = nnx.silu(gating_proj) * up_proj
+
+                expert_output = jnp.dot(fused, down_proj_kernel)
+
+                # Apply router weight after computation (matches implementation)
+                combined_expert_output += weight * expert_output
+
+            expected_output = expected_output.at[i].set(combined_expert_output)
+
+        expected_output = expected_output.reshape(self.B * self.S, self.D)
+
+        # --- 3. Compare the results ---
+        self.assertTrue(
+            jnp.allclose(actual_output, expected_output, atol=1e-2, rtol=1e-2),
+            f"The output of the distributed MoE does not match the dense equivalent.\n"
+            f"Actual:\n{actual_output}\n"
+            f"Expected:\n{expected_output}")
+        print(
+            "\n✅ Test Passed: Distributed MoE output matches the dense ground truth."
+        )

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