tpu-inference 0.11.1.dev202511270815__py3-none-any.whl → 0.13.0rc2.post7__py3-none-any.whl

tests/layers/common/test_attention_interface.py

@@ -0,0 +1,156 @@
+# 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 unittest.mock import MagicMock
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from jax.sharding import Mesh
+
+from tpu_inference.layers.common.attention_interface import attention
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.runner.kv_cache import get_kv_cache_shape_with_mesh
+
+# ---- Test Configuration & Constants ----
+
+# Total number of tokens across all sequences in the batch
+TOTAL_TOKENS = 10
+# Number of sequences in the batch
+NUM_SEQS = 2
+# Padded maximum number of sequences
+MAX_NUM_SEQS = 4
+# Number of attention heads (Query)
+NUM_HEADS = 8
+# Number of attention heads (Key/Value) - for Grouped-Query Attention
+NUM_KV_HEADS = 4
+# Total number of blocks in the KV cache
+NUM_BLOCKS = 32
+# Number of tokens per block
+BLOCK_SIZE = 16
+# Maximum number of blocks a single sequence can occupy
+MAX_BLOCKS_PER_SEQ = 8
+
+
+@pytest.fixture
+def mesh():
+    """Provides a mock 1D JAX mesh for testing."""
+    # Create a mesh with available devices, useful for running on CPU/GPU/TPU
+    # For this test, it will likely be a single CPU device.
+    devices = np.array(jax.local_devices()[:1])
+    if not devices.any():
+        # Add a mock device if no devices are present (e.g., in a CI environment)
+        devices = np.array([jax.devices("cpu")[0]])
+    return Mesh(devices.reshape((-1, 1, 1)), ("data", "attn_dp", "model"))
+
+
+# ---- Test for `attention` ----
+
+
+def _test_attention(monkeypatch, mesh, head_dim, use_sinks=False):
+    """
+    Tests the main `attention` function.
+
+    Verifies that:
+    1. It calls the `sharded_ragged_paged_attention` kernel with correct metadata.
+    2. The final outputs (kv_cache and attention output) have the correct shapes.
+    """
+    # 1. Arrange
+
+    # Create input tensors
+    q_dtype = jnp.float32
+    kv_dtype = jnp.float32
+    q = jnp.ones((TOTAL_TOKENS, NUM_HEADS, head_dim), dtype=q_dtype)
+    k = jnp.ones((TOTAL_TOKENS, NUM_KV_HEADS, head_dim), dtype=kv_dtype)
+    v = jnp.ones((TOTAL_TOKENS, NUM_KV_HEADS, head_dim), dtype=kv_dtype)
+    sinks = jnp.ones((NUM_HEADS, ), dtype=jnp.float32) if use_sinks else None
+
+    kv_cache_shape = get_kv_cache_shape_with_mesh(
+        mesh,
+        NUM_BLOCKS,
+        BLOCK_SIZE,
+        NUM_KV_HEADS,
+        head_dim,
+        kv_dtype,
+    )
+    kv_cache = jnp.zeros(kv_cache_shape, dtype=kv_dtype)
+
+    # Mock ragged_paged_attention to return a tensor of the correct shape
+    mock_paged_attn_kernel = MagicMock(return_value=(jnp.ones(
+        (TOTAL_TOKENS, NUM_HEADS, head_dim)), kv_cache), )
+
+    if head_dim == 64:
+        monkeypatch.setattr(
+            "tpu_inference.layers.common.attention_interface.ragged_paged_attention_hd64",
+            mock_paged_attn_kernel,
+        )
+    else:
+        monkeypatch.setattr(
+            "tpu_inference.layers.common.attention_interface.ragged_paged_attention",
+            mock_paged_attn_kernel,
+        )
+
+    # Create AttentionMetadata
+    attention_metadata = AttentionMetadata(
+        input_positions=jnp.arange(TOTAL_TOKENS, dtype=jnp.int32),
+        block_tables=jnp.zeros((MAX_NUM_SEQS * MAX_BLOCKS_PER_SEQ, ),
+                               dtype=jnp.int32),
+        seq_lens=jnp.array([5, 5, 0, 0], dtype=jnp.int32),
+        query_start_loc=jnp.array([0, 5, 10, 10, 10], dtype=jnp.int32),
+        request_distribution=jnp.array([0, 0, NUM_SEQS], dtype=jnp.int32),
+    )
+
+    # 2. Act
+    final_kv_cache, output = attention(
+        kv_cache=kv_cache,
+        q=q,
+        k=k,
+        v=v,
+        attention_metadata=attention_metadata,
+        mesh=mesh,
+        head_dim_original=head_dim,
+        sinks=sinks,
+    )
+
+    # 3. Assert
+    # Check that both mocked kernels were called
+    mock_paged_attn_kernel.assert_called_once()
+
+    # Check output shapes
+    assert final_kv_cache.shape == kv_cache.shape
+    assert output.shape == q.shape
+
+    # Check that the output is the one from our mock
+    assert jnp.all(output == 1.0)
+
+
+def test_attention(monkeypatch, mesh):
+    _test_attention(monkeypatch, mesh, 128)
+
+
+def test_attention_hd64(monkeypatch, mesh):
+    _test_attention(monkeypatch, mesh, 64)
+
+
+def test_attention_sink(monkeypatch, mesh):
+    _test_attention(monkeypatch, mesh, 64, True)
+
+
+def test_attention_sink_no_64_raises_error(monkeypatch, mesh):
+    with pytest.raises(
+            NotImplementedError,
+            match="Attention sink support is only available when head_dim==64"
+    ):
+        _test_attention(monkeypatch, mesh, 128, True)

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"]