tpu-inference 0.11.1.dev202511220812__py3-none-any.whl → 0.13.2.dev20251230__py3-none-any.whl

tests/layers/common/test_quantization.py

@@ -0,0 +1,149 @@
+# 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 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.layers.common.quantization import (
+    dequantize_tensor, dequantize_tensor_from_mxfp4_packed, quantize_kv,
+    quantize_tensor, quantize_tensor_to_mxfp4_packed)
+
+
+@jtu.with_config(jax_numpy_dtype_promotion="standard")
+class QuantizationTest(jtu.JaxTestCase):
+
+    @parameterized.product(axis=[-1, 0, (0, 1)])
+    def test_mxfp4_quantization(self, axis):
+        if not jtu.is_device_tpu_at_least(version=7):
+            self.skipTest("mxfp4 is only supported in TPUv7+")
+
+        key = jax.random.key(0)
+
+        shape = (128, 128, 128)
+        original = jax.random.normal(key, shape, jnp.bfloat16)
+
+        tensor_q, scale = quantize_tensor_to_mxfp4_packed(original, axis)
+        dequantized = dequantize_tensor_from_mxfp4_packed(
+            tensor_q, scale, axis)
+
+        self.assertAllClose(dequantized, original, rtol=0.5, atol=0.5)
+
+    @parameterized.product(dtype=[jnp.float8_e4m3fn, jnp.int8],
+                           axis=[None, -1, 1, (0, 1)])
+    def test_quantization(self, dtype, axis):
+        key = jax.random.key(0)
+
+        shape = (128, 128, 128)
+        original = jax.random.normal(key, shape, jnp.bfloat16)
+
+        tensor_q, scale = quantize_tensor(dtype, original, axis)
+        dequantized = dequantize_tensor(tensor_q, scale, axis)
+
+        self.assertAllClose(dequantized, original, rtol=0.1, atol=0.1)
+
+    @parameterized.product(dtype=[jnp.float8_e4m3fn, jnp.int8],
+                           axis=[-1, 1],
+                           block_size=[32, 64])
+    def test_block_quantization(self, dtype, axis, block_size):
+        key = jax.random.key(0)
+
+        shape = (128, 128, 128)
+        original = jax.random.normal(key, shape, jnp.bfloat16)
+
+        tensor_q, scale = quantize_tensor(dtype, original, axis, block_size)
+        dequantized = dequantize_tensor(tensor_q, scale, axis)
+
+        self.assertAllClose(dequantized, original, rtol=0.1, atol=0.1)
+
+    @parameterized.product(dtype=[jnp.float8_e4m3fn, jnp.int8],
+                           axis=[(0, 1), (-1, 0)],
+                           block_size=[32, (64, 32)])
+    def test_multi_block_quantization(self, dtype, axis, block_size):
+        key = jax.random.key(0)
+
+        shape = (128, 128, 128)
+        original = jax.random.normal(key, shape, jnp.bfloat16)
+
+        tensor_q, scale = quantize_tensor(dtype, original, axis, block_size)
+        dequantized = dequantize_tensor(tensor_q, scale, axis)
+
+        self.assertAllClose(dequantized, original, rtol=0.1, atol=0.1)
+
+    def test_unaligned_block_quantization_raises_error(self):
+        key = jax.random.key(0)
+
+        shape = (128, 128)
+        tensor = jax.random.normal(key, shape, jnp.bfloat16)
+        block_size = 100
+        axis = 0
+
+        self.assertRaises(
+            ValueError,
+            functools.partial(quantize_tensor, jnp.int8, tensor, axis,
+                              block_size))
+
+    def test_block_quantization_padding(self):
+        key = jax.random.key(0)
+
+        shape = (128, 128)
+
+        original = jax.random.normal(key, shape, jnp.bfloat16)
+        block_size = 100
+        axis = 0
+
+        tensor_q, scale = quantize_tensor(jnp.int8, original, axis, block_size,
+                                          True)
+
+        dequantized = dequantize_tensor(tensor_q, scale, axis)
+
+        padded_size = ((shape[axis] + block_size) // block_size) * block_size
+        self.assertEqual(tensor_q.shape[axis], padded_size)
+        self.assertTrue((tensor_q[shape[0]:] == 0).all())
+        self.assertAllClose(dequantized[:shape[0]],
+                            original,
+                            rtol=0.1,
+                            atol=0.1)
+
+    @parameterized.product(kv_quant_dtype=[jnp.float8_e4m3fn, jnp.int8])
+    def test_quantize_kv(self, kv_quant_dtype):
+        """Tests the quantize_kv function with float8_e4m3fn dtype."""
+        key = jax.random.key(0)
+
+        shape = (128, 128)
+        k_original = jax.random.normal(key, shape, jnp.bfloat16)
+        v_original = jax.random.normal(key, shape, jnp.bfloat16)
+        k_scale = 0.1
+        v_scale = 0.2
+
+        k_quantized, v_quantized = quantize_kv(
+            kv_quant_dtype,
+            k_original,
+            v_original,
+            k_scale,
+            v_scale,
+        )
+
+        k_dequantized = k_quantized.astype(jnp.bfloat16) * k_scale
+        v_dequantized = v_quantized.astype(jnp.bfloat16) * v_scale
+
+        self.assertAllClose(k_dequantized, k_original, rtol=0.2, atol=0.2)
+        self.assertAllClose(v_dequantized, v_original, rtol=0.2, atol=0.2)
+
+
+if __name__ == "__main__":
+    absltest.main(testLoader=jtu.JaxTestLoader())

tests/layers/jax/__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/attention/__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/attention/test_common_attention.py

@@ -0,0 +1,103 @@
+# 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 typing import Tuple
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from flax import nnx
+from jax.sharding import Mesh
+from parameterized import parameterized
+
+from tpu_inference.layers.common.attention_interface import get_kv_cache_shape
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.jax.attention.attention import Attention
+
+KVCache = Tuple[jax.Array, jax.Array]
+
+
+class TestAttention(unittest.TestCase):
+    """Unit test suite for the JAX Attention module."""
+
+    def setUp(self):
+        """Sets up the testing environment before each test."""
+        self.mesh = Mesh(
+            np.array(jax.devices()[:1]).reshape(1, 1, 1, -1),
+            axis_names=(
+                "data",
+                "attn_dp",
+                "expert",
+                "model",
+            ),
+        )
+
+    @parameterized.expand([["auto"], ["fp8"]])
+    def test_attention_forward_pass(self, kv_cache_str):
+        """Tests the forward pass of the Attention module in prefill mode."""
+        hidden_size = 1024
+        num_attention_heads = 8
+        head_dim = hidden_size // num_attention_heads
+
+        with jax.set_mesh(self.mesh):
+            attention = Attention(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,
+                                  mesh=self.mesh,
+                                  random_init=True,
+                                  rngs=nnx.Rngs(42),
+                                  kv_cache_dtype=kv_cache_str)
+
+            seq_len = 64
+            x = jnp.ones((seq_len, hidden_size), dtype=jnp.bfloat16)
+
+            block_size = 16
+            num_blocks = 8
+            kv_dtype = jnp.float8_e4m3fn if kv_cache_str == "fp8" else jnp.bfloat16
+            cache_shape = get_kv_cache_shape(num_blocks, block_size,
+                                             num_attention_heads, head_dim,
+                                             kv_dtype)
+
+            kv_cache = jnp.zeros(cache_shape, dtype=kv_dtype)
+
+            num_required_blocks = seq_len // block_size
+
+            attention_metadata = AttentionMetadata(
+                input_positions=jnp.arange(seq_len, dtype=jnp.int32),
+                block_tables=jnp.array(list(range(num_required_blocks)),
+                                       dtype=jnp.int32),
+                seq_lens=jnp.array([seq_len], dtype=jnp.int32),
+                query_start_loc=jnp.array([0, seq_len], dtype=jnp.int32),
+                request_distribution=jnp.array([0, 0, 1], dtype=jnp.int32),
+            )
+
+            new_kv_cache, output = attention(
+                x,
+                is_prefill=True,
+                kv_cache=kv_cache,
+                attention_metadata=attention_metadata,
+            )
+
+            self.assertEqual(output.shape, (seq_len, hidden_size))
+
+            self.assertEqual(new_kv_cache.shape, kv_cache.shape)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/layers/jax/attention/test_deepseek_v3_attention.py

@@ -0,0 +1,233 @@
+# 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
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from flax import nnx
+from jax.sharding import Mesh, PartitionSpec
+from parameterized import parameterized
+
+import tpu_inference.kernels.mla.v1.kernel as mla
+from tpu_inference.layers.common.attention_interface import get_kv_cache_shape
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.layers.common.sharding import ShardingAxisName
+from tpu_inference.layers.jax.attention.deepseek_v3_attention import MLA
+
+
+class TestMLA(unittest.TestCase):
+
+    def setUp(self):
+        os.environ["NEW_MODEL_DESIGN"] = "1"
+        self.mesh = Mesh(
+            np.array(jax.devices("tpu")[:1]).reshape(1, 1, 1, 1),
+            axis_names=("data", "attn_dp", "expert", "model"),
+        )
+
+    @parameterized.expand([["auto"], ["fp8"]])
+    def test_mla_forward_pass(self, kv_cache_str):
+        hidden_size = 256
+
+        num_key_value_heads = 32
+        qk_nope_head_dim = 64
+        qk_rope_head_dim = 32
+
+        with jax.set_mesh(self.mesh):
+            query_tnh_spec = PartitionSpec(None, ShardingAxisName.MLP_TENSOR,
+                                           None)
+            keyvalue_skh_spec = PartitionSpec(None,
+                                              ShardingAxisName.MLP_TENSOR,
+                                              None)
+            attn_o_tnh_spec = PartitionSpec(None, ShardingAxisName.MLP_TENSOR,
+                                            None)
+
+            mla_layer = MLA(
+                hidden_size=hidden_size,
+                num_attention_heads=32,
+                num_key_value_heads=num_key_value_heads,
+                head_dim=64,  # MLA uses v_head_dim as head_dim
+                rope_theta=10000,
+                dtype=jnp.bfloat16,
+                q_lora_rank=512,
+                kv_lora_rank=512,
+                qk_nope_head_dim=
+                qk_nope_head_dim,  # Half of DeepSeek v3's real values
+                qk_rope_head_dim=
+                qk_rope_head_dim,  # Half of DeepSeek v3's real values
+                v_head_dim=64,  # Half of DeepSeek v3's real values
+                rms_norm_eps=1e-5,
+                rngs=nnx.Rngs(42),
+                rope_scaling={
+                    "beta_fast": 32,
+                    "beta_slow": 1,
+                    "factor": 40,
+                    "mscale": 1.0,
+                    "mscale_all_dim": 1.0,
+                    "original_max_position_embeddings": 4096,
+                    "type": "yarn",
+                },
+                mesh=self.mesh,
+                random_init=True,
+                kv_cache_dtype=kv_cache_str,
+                query_tnh=query_tnh_spec,
+                keyvalue_skh=keyvalue_skh_spec,
+                attn_o_tnh=attn_o_tnh_spec,
+                q_da_sharding=(None, ShardingAxisName.VOCAB),
+                anh_sharding=(None, ShardingAxisName.MLP_TENSOR, None),
+                ap_sharding=(None, ShardingAxisName.MLP_TENSOR),
+                kv_da_sharding=(None, ShardingAxisName.VOCAB),
+                rd_sharding=(ShardingAxisName.MLP_TENSOR, None),
+            )
+
+            # Create input tensor
+            seq_len = 32
+            x = jnp.ones((seq_len, hidden_size), dtype=jnp.bfloat16)
+
+            # Create KV cache
+            qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+            block_size = 16
+            num_blocks = 8
+            kv_dtype = jnp.float8_e4m3fn if kv_cache_str == "fp8" else jnp.bfloat16
+            cache_shape = get_kv_cache_shape(num_blocks, block_size,
+                                             num_key_value_heads, qk_head_dim,
+                                             kv_dtype)
+            kv_cache = jnp.zeros(cache_shape, dtype=kv_dtype)
+
+            # Create attention metadata
+            attention_metadata = AttentionMetadata(
+                input_positions=jnp.arange(seq_len, dtype=jnp.int32),
+                block_tables=jnp.zeros((8, ), dtype=jnp.int32),
+                seq_lens=jnp.ones((1, ), dtype=jnp.int32) * seq_len,
+                query_start_loc=jnp.array(
+                    [0, seq_len], dtype=jnp.int32),  # This is cu_q_lens
+                request_distribution=jnp.array([0, 0, 1], dtype=jnp.int32),
+            )
+
+            mla_layer.rope.initialize_cache(self.mesh)
+
+            # Run forward pass
+            new_kv_cache, output = mla_layer(
+                x,
+                is_prefill=True,
+                kv_cache=kv_cache,
+                attention_metadata=attention_metadata)
+
+            # Verify output shapes
+            self.assertEqual(output.shape, (seq_len, hidden_size))
+            self.assertEqual(new_kv_cache.shape, kv_cache.shape)
+
+    @parameterized.expand([["auto"]])  # MLA kernel does not support fp8 yet
+    def test_mla_kernel_forward_pass(self, kv_cache_str):
+        hidden_size = 256
+
+        num_key_value_heads = 1
+        qk_nope_head_dim = 64
+        qk_rope_head_dim = 32
+        v_head_dim = 64
+        kv_lora_rank = 512
+
+        with jax.set_mesh(self.mesh):
+            query_tnh_spec = PartitionSpec(ShardingAxisName.MLP_TENSOR, None,
+                                           None)
+            keyvalue_skh_spec = PartitionSpec(ShardingAxisName.MLP_TENSOR,
+                                              None)
+            attn_o_tnh_spec = PartitionSpec(ShardingAxisName.MLP_TENSOR, None,
+                                            None)
+
+            mla_layer = MLA(
+                hidden_size=hidden_size,
+                num_attention_heads=32,
+                num_key_value_heads=num_key_value_heads,
+                head_dim=v_head_dim,  # MLA uses v_head_dim as head_dim
+                rope_theta=10000,
+                dtype=jnp.bfloat16,
+                q_lora_rank=512,
+                kv_lora_rank=kv_lora_rank,
+                qk_nope_head_dim=qk_nope_head_dim,
+                qk_rope_head_dim=qk_rope_head_dim,
+                v_head_dim=v_head_dim,
+                rms_norm_eps=1e-5,
+                rngs=nnx.Rngs(42),
+                rope_scaling={
+                    "beta_fast": 32,
+                    "beta_slow": 1,
+                    "factor": 40,
+                    "mscale": 1.0,
+                    "mscale_all_dim": 1.0,
+                    "original_max_position_embeddings": 4096,
+                    "type": "yarn",
+                },
+                mesh=self.mesh,
+                random_init=True,
+                kv_cache_dtype=kv_cache_str,
+                use_mla_kernel=
+                True,  # Set to true, in order to trigger MLA kernel.
+                query_tnh=query_tnh_spec,
+                keyvalue_skh=keyvalue_skh_spec,
+                attn_o_tnh=attn_o_tnh_spec,
+                q_da_sharding=(None, ShardingAxisName.VOCAB),
+                anh_sharding=(None, ShardingAxisName.MLP_TENSOR, None),
+                ap_sharding=(None, ShardingAxisName.MLP_TENSOR),
+                kv_da_sharding=(None, ShardingAxisName.VOCAB),
+                rd_sharding=(ShardingAxisName.MLP_TENSOR, None),
+            )
+
+            # Create input tensor
+            seq_len = 32
+            x = jnp.ones((seq_len, hidden_size), dtype=jnp.bfloat16)
+
+            # Create KV cache for MLA kernel
+            block_size = 16
+            num_blocks = 8
+            kv_dtype = jnp.float8_e4m3fn if kv_cache_str == "fp8" else jnp.bfloat16
+
+            # For the MLA kernel, the head dimension is the sum of qk_nope_head_dim and v_head_dim
+            # and lora rank
+            cache_shape = mla.get_kv_cache_shape(
+                num_blocks, block_size, kv_lora_rank + qk_rope_head_dim,
+                kv_dtype)
+            kv_cache = jnp.zeros(cache_shape, dtype=kv_dtype)
+
+            # Create attention metadata
+            attention_metadata = AttentionMetadata(
+                input_positions=jnp.arange(seq_len, dtype=jnp.int32),
+                block_tables=jnp.zeros((8, ), dtype=jnp.int32),
+                seq_lens=jnp.ones((1, ), dtype=jnp.int32) * seq_len,
+                query_start_loc=jnp.array([0, seq_len], dtype=jnp.int32),
+                request_distribution=jnp.array([0, 0, 1], dtype=jnp.int32),
+            )
+
+            mla_layer.rope.initialize_cache(self.mesh)
+
+            # Run forward pass
+            new_kv_cache, output = mla_layer(
+                x,
+                is_prefill=True,
+                kv_cache=kv_cache,
+                attention_metadata=attention_metadata)
+
+            # Verify output shapes
+            self.assertEqual(output.shape, (seq_len, hidden_size))
+            self.assertEqual(new_kv_cache.shape, kv_cache.shape)
+
+
+if __name__ == "__main__":
+    unittest.main()
+
+
+def tearDownModule():
+    del os.environ["NEW_MODEL_DESIGN"]

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.