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

tests/models/jax/test_llama4.py

@@ -0,0 +1,298 @@
+# 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 field
+from types import SimpleNamespace
+from typing import Any, Tuple
+from unittest.mock import MagicMock, patch
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from flax import nnx
+from flax.typing import PRNGKey
+from jax.sharding import Mesh
+from vllm.config import ModelConfig
+
+from tpu_inference.models.jax.llama4 import (Llama4ForCausalLM,
+                                             Llama4WeightLoader)
+
+
+class MockParamLlama4:
+    """A mock for a parameter used in the Llama4 model."""
+    shape: Tuple[int, ...]
+    dtype: jnp.dtype = jnp.bfloat16
+    sharding_spec: Tuple[str | None, ...] | None = None
+    value: Any = field(init=False)
+    sharding: Any = field(init=False)
+
+    def __init__(self, shape=(32, 128)):
+        self.shape = shape
+        self.value = jnp.zeros(self.shape, dtype=self.dtype)
+        # The sharding spec is accessed during weight loading
+        self.sharding = SimpleNamespace(spec=self.sharding_spec)
+
+    # Allow the mock parameter's value to be updated
+    def __setattr__(self, name, value):
+        if name in ['value', 'shape', 'dtype', 'sharding', 'sharding_spec']:
+            self.__dict__[name] = value
+        else:
+            super().__setattr__(name, value)
+
+
+class MockVllmConfig:
+    """A mock VllmConfig sufficient for testing the Llama4 model."""
+
+    def __init__(self,
+                 model_name: str,
+                 random_weights: bool = False,
+                 tensor_parallelism: int = 1):
+        self.model_config = MagicMock(spec=ModelConfig)
+        self.load_config = MagicMock()
+        self.load_config.download_dir = None
+
+        # Choose small amount of layers to avoid OOM.
+        self.model_config.get_vocab_size.return_value = 202048
+        self.model_config.get_hidden_size.return_value = 32
+        self.model_config.model = model_name
+
+        self.additional_config = {
+            "random_weights": random_weights,
+            "sharding": {
+                "sharding_strategy": {
+                    "tensor_parallelism": tensor_parallelism
+                }
+            }
+        }
+
+        self.cache_config = MagicMock(cache_dtype="auto")
+
+        text_config_mock = MagicMock()
+        text_config_mock.interleave_moe_layer_step = 1
+        text_config_mock.num_attention_heads = 40
+        text_config_mock.num_key_value_heads = 8
+        text_config_mock.head_dim = 128
+
+        hf_config_mock = MagicMock()
+        hf_config_mock.text_config = text_config_mock
+
+        self.model_config.hf_config = hf_config_mock
+
+
+@pytest.fixture(scope="module")
+def mesh():
+    """
+    Creates a mesh with all required axes for testing.
+    """
+    if not jax.devices():
+        pytest.skip("No JAX devices available for mesh creation.")
+
+    devices = np.array(jax.local_devices())
+    # Reshape devices into a 3D array to name 3 axes: data, model, and expert.
+    # The 'model' and 'expert' axes will have a size of 1.
+    num_devices = len(devices)
+    device_mesh = devices.reshape((num_devices, 1, 1, 1))
+
+    with Mesh(device_mesh,
+              axis_names=('data', 'attn_dp', 'model', 'expert')) as m:
+        yield m
+
+
+@pytest.fixture
+def rng() -> PRNGKey:
+    """Provides a reusable JAX PRNGKey."""
+    return jax.random.PRNGKey(42)
+
+
+@pytest.fixture
+def mock_vllm_config_llama4() -> MockVllmConfig:
+    return MockVllmConfig(model_name="meta-llama/Llama-4-Scout-17B-16E")
+
+
+class TestLlama4ForCausalLM:
+    """Tests for the main LlamaForCausalLM model class."""
+
+    def test_init_llama4(self, mock_vllm_config_llama4, rng, mesh):
+        """Tests correct parameter detection for the Llama4 model variant."""
+        model = Llama4ForCausalLM(mock_vllm_config_llama4, rng, mesh)
+        assert model.hidden_size == 32
+        assert "llama-4" in model.vllm_config.model_config.model.lower()
+
+    def test_create_model_with_random_weights(self, mock_vllm_config_llama4,
+                                              rng, mesh):
+        """
+        Tests that random weight initialization creates concrete, non-zero-variance arrays.
+        """
+        with jax.set_mesh(mesh):
+            model = Llama4ForCausalLM(vllm_config=mock_vllm_config_llama4,
+                                      rng=rng,
+                                      mesh=mesh,
+                                      force_random_weights=True)
+            embedding_weight = model.embedder.input_embedding_table_VD.value
+            attention_q_kernel = model.layers[0].attn.kernel_q_proj_DNH.value
+            final_norm_scale = model.final_norm.scale.value
+
+            assert isinstance(embedding_weight, jax.Array)
+            assert isinstance(attention_q_kernel, jax.Array)
+            assert isinstance(final_norm_scale, jax.Array)
+
+            assert jnp.std(embedding_weight) > 0
+            assert jnp.std(attention_q_kernel) > 0
+
+            assert jnp.all(final_norm_scale == 1.0)
+
+    @patch("tpu_inference.models.jax.llama4.Llama4WeightLoader")
+    def test_load_weights_called_correctly(self, mock_loader_cls, rng, mesh):
+        """Tests that the weight loader is called correctly for checkpoint loading."""
+        vllm_config = MockVllmConfig(model_name="llama4-scout",
+                                     random_weights=False)
+        model = Llama4ForCausalLM(vllm_config, rng, mesh)
+
+        mock_loader_instance = MagicMock()
+        mock_loader_cls.return_value = mock_loader_instance
+        model.load_weights(rng)
+
+        mock_loader_cls.assert_called_once_with(vllm_config=vllm_config,
+                                                hidden_size=32,
+                                                attn_heads=40,
+                                                num_key_value_heads=8,
+                                                attn_head_dim=128)
+        mock_loader_instance.load_weights.assert_called_once_with(model)
+
+
+class TestLlama4WeightLoader:
+    """Tests for the Llama4WeightLoader class."""
+
+    @pytest.fixture
+    def weight_loader(self):
+        # Patch the superclass's setup to isolate the Llama4 loader's logic
+        return Llama4WeightLoader(vllm_config=MockVllmConfig("test-model"),
+                                  hidden_size=32,
+                                  attn_heads=40,
+                                  num_key_value_heads=8,
+                                  attn_head_dim=128)
+
+    @pytest.mark.parametrize("hf_key, expected_num", [
+        ("language_model.model.layers.15.self_attn.q_proj.weight", 15),
+        ("layers.0.feed_forward.router.weight", 0),
+        ("language_model.model.layers.99.norm.weight", 99),
+        ("language_model.model.norm.weight", None),
+        ("language_model.model.embed_tokens.weight", None),
+    ])
+    def test_get_layer_num(self, weight_loader, hf_key, expected_num):
+        """Tests the private _get_layer_num utility function."""
+        assert weight_loader._get_layer_num(hf_key) == expected_num
+
+    @pytest.mark.parametrize("hf_key, expected_num", [
+        ("language_model.model.layers.10.feed_forward.experts.4.down_proj.weight",
+         4),
+        ("language_model.model.layers.0.feed_forward.experts.0.gate_proj.weight_scale",
+         0),
+        ("language_model.model.layers.5.feed_forward.experts.128.up_proj.weight",
+         128),
+        ("language_model.model.norm.weight", None),
+        ("language_model.model.layers.15.self_attn.q_proj.weight", None),
+    ])
+    def test_get_expert_num(self, weight_loader, hf_key, expected_num):
+        """Tests the private _get_expert_num utility function to extract the expert index."""
+        assert weight_loader._get_expert_num(hf_key) == expected_num
+
+    @pytest.mark.parametrize("hf_key, expected", [
+        ("language_model.model.layers.15.self_attn.q_proj.weight",
+         "layers.15.attn.kernel_q_proj_DNH"),
+        ("language_model.model.layers.0.feed_forward.shared_expert.down_proj.weight",
+         "layers.0.shared_experts.kernel_down_proj_FD"),
+        ("language_model.model.embed_tokens.weight",
+         "embedder.input_embedding_table_VD"),
+        ("language_model.model.norm.weight", "final_norm.scale"),
+        ("language_model.lm_head.weight", "lm_head.input_embedding_table_DV"),
+        ("unmapped.key.name", "unmapped.key.name"),
+    ])
+    def test_map_loaded_to_standardized_name(self, weight_loader, hf_key,
+                                             expected):
+        """Tests the mapping from HuggingFace key names to internal names."""
+        assert weight_loader.map_loaded_to_standardized_name(
+            hf_key) == expected
+
+    def test_load_weights_transformation(self, weight_loader, rng, mesh):
+        """Tests that weights are correctly reshaped, transposed, and loaded."""
+        vllm_config = MockVllmConfig(model_name="llama4-small-test",
+                                     random_weights=False)
+
+        model = Llama4ForCausalLM(vllm_config, rng, mesh)
+
+        # Original weight shape is (vocab_size, hidden_size)
+        original_weight = jnp.ones((128, 32))
+        dummy_weights = [
+            ("language_model.model.embed_tokens.weight", original_weight),
+        ]
+        weight_loader.names_and_weights_generator = dummy_weights
+
+        # Mock get_param to return a mock param with the target shape (vocab_size, hidden_size)
+        mock_param = MockParamLlama4(shape=(128, 32))
+
+        with patch("tpu_inference.models.jax.llama4.get_param", return_value=mock_param), \
+            patch("tpu_inference.models.jax.llama4.shard_put", return_value=jnp.ones(mock_param.value.shape)) as mock_shard_put:
+
+            # This will now pass after the code fix
+            weight_loader.load_weights(model)
+
+            # Assert that shard_put was called with the correctly transposed weight
+            mock_shard_put.assert_called_once()
+
+            # Get the actual array passed to shard_put
+            called_with_weight = mock_shard_put.call_args[0][0]
+
+            # Check if the shape of the array passed to shard_put matches the model's expected shape.
+            assert called_with_weight.shape == mock_param.value.shape
+
+    def test_map_llama4_gate_up_proj(self, weight_loader, rng, mesh):
+        """Tests that gate_up_proj weights are correctly split, reshaped, transposed, and loaded."""
+        # Set up a dummy model and its config
+        model = Llama4ForCausalLM(MockVllmConfig("test-model"), rng, mesh)
+
+        # Create a dummy fused gate_up_proj weight tensor
+        hidden_size = 32
+        intermediate_size_moe = 8192
+        num_local_experts = 2
+        dummy_weight = jnp.ones(
+            (num_local_experts, hidden_size, 2 * intermediate_size_moe))
+
+        # Set up mocks and patches
+        mock_model_params = nnx.state(model)
+        mock_param = MockParamLlama4(shape=(2, hidden_size,
+                                            intermediate_size_moe))
+
+        # Create a dummy WeightLoader and set up the necessary attributes
+        weight_loader.is_verbose = False
+        layer_num = 0
+        weight_loader.names_and_weights_generator = [
+            (f"language_model.model.layers.{layer_num}.feed_forward.experts.gate_up_proj.weight",
+             dummy_weight),
+        ]
+
+        with patch("tpu_inference.models.jax.llama4.get_param", return_value=mock_param), \
+            patch("tpu_inference.models.jax.llama4.shard_put", return_value=jnp.ones(mock_param.value.shape)) as mock_shard_put:
+
+            # Call _map_llama4_gate_up_proj directly
+            weight_loader._map_llama4_gate_up_proj(
+                model, mock_model_params,
+                f"language_model.model.layers.{layer_num}.feed_forward.experts.gate_up_proj.weight",
+                dummy_weight)
+            # Check if shard_put was called the correct number of times and with the correct weight shapes
+            assert mock_shard_put.call_count == 2
+            # call_args_list gives you a list of all the calls with their arguments.
+            for call in mock_shard_put.call_args_list:
+                assert call[0][0].shape == (num_local_experts, 32, 8192)

tests/models/jax/test_llama_eagle3.py

@@ -0,0 +1,197 @@
+# 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, patch
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from flax import nnx
+from flax.typing import PRNGKey
+from jax.sharding import Mesh
+from vllm.config import ModelConfig
+
+from tpu_inference.layers.common.attention_metadata import AttentionMetadata
+from tpu_inference.models.jax.llama_eagle3 import (Eagle3LlamaDecoderLayer,
+                                                   EagleLlama3ForCausalLM)
+from tpu_inference.runner.kv_cache import create_kv_caches
+
+
+class MockSpeculativeConfig:
+
+    def __init__(self):
+        self.num_speculative_tokens = 3
+        self.method = "eagle3"
+        self.draft_model_config = None
+
+
+class MockVllmConfig:
+
+    def __init__(self, model: str, draft_model: str, kv_cache_dtype):
+        self.model_config = ModelConfig(model)
+        self.model_config.dtype = jnp.bfloat16
+        self.load_config = MagicMock()
+        self.load_config.download_dir = None
+        self.speculative_config = MockSpeculativeConfig()
+        self.speculative_config.draft_model_config = ModelConfig(
+            draft_model,
+            dtype="bfloat16",
+            max_model_len=2048,
+            skip_tokenizer_init=True,
+            trust_remote_code=True)
+        self.cache_config = MagicMock(cache_dtype=kv_cache_dtype)
+
+
+@pytest.fixture
+def mock_vllm_config() -> MockVllmConfig:
+    return MockVllmConfig(model="meta-llama/Meta-Llama-3-8B-Instruct",
+                          draft_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
+                          kv_cache_dtype="auto")
+
+
+@pytest.fixture(scope="module")
+def mesh():
+    """Creates a mesh with 1 device."""
+    if not jax.devices():
+        pytest.skip("No JAX devices available for mesh creation.")
+
+    devices = np.array(jax.local_devices()[:1])
+    device_mesh = devices.reshape((1, 1, -1))
+
+    with Mesh(device_mesh, axis_names=('data', 'attn_dp', 'model')) as m:
+        yield m
+
+
+@pytest.fixture
+def mock_model_inputs(mock_vllm_config: MockVllmConfig):
+    """Provides mock inputs for the EagleLlama3 model."""
+    batch_size = 2
+    seq_len = 16
+    target_hidden_size = mock_vllm_config.model_config.get_hidden_size()
+
+    input_ids = jnp.ones((batch_size * seq_len, ), dtype=jnp.int32)
+    hidden_states = jnp.ones((batch_size * seq_len, target_hidden_size),
+                             dtype=jnp.bfloat16)
+    attention_metadata = AttentionMetadata(
+        input_positions=jnp.arange(batch_size * seq_len, dtype=jnp.int32),
+        block_tables=jnp.zeros((batch_size, 1), dtype=jnp.int32).reshape(-1),
+        seq_lens=jnp.full((batch_size, ), seq_len, dtype=jnp.int32),
+        query_start_loc=jnp.arange(0, (batch_size + 1) * seq_len,
+                                   seq_len,
+                                   dtype=jnp.int32),
+        request_distribution=jnp.array([0, 0, batch_size], dtype=jnp.int32),
+    )
+    return input_ids, hidden_states, attention_metadata
+
+
+@pytest.fixture
+def rng() -> PRNGKey:
+    """Provides a reusable JAX PRNGKey."""
+    return jax.random.PRNGKey(42)
+
+
+class TestEagleLlama3ForCausalLM:
+    """Tests for the EagleLlama3ForCausalLM model."""
+
+    def test_eagle3_decoder_layer_init(self, mock_vllm_config: MockVllmConfig,
+                                       rng: PRNGKey, mesh: Mesh):
+        """Tests the initialization of the Eagle3LlamaDecoderLayer."""
+        hf_config = mock_vllm_config.speculative_config.draft_model_config.hf_config
+        dtype = jnp.bfloat16
+        rngs = nnx.Rngs(rng)
+
+        layer = Eagle3LlamaDecoderLayer(
+            hf_config,
+            dtype,
+            rngs,
+            mesh,
+            kv_cache_dtype=mock_vllm_config.cache_config.cache_dtype)
+
+        # Check if projection layers are overridden with correct input size
+        original_hidden_size = hf_config.hidden_size
+        expected_input_size = 2 * original_hidden_size
+
+        assert layer.self_attn.q_proj.kernel.value.shape[
+            0] == expected_input_size
+        assert layer.self_attn.k_proj.kernel.value.shape[
+            0] == expected_input_size
+        assert layer.self_attn.v_proj.kernel.value.shape[
+            0] == expected_input_size
+        assert isinstance(layer.hidden_norm, nnx.RMSNorm)
+
+    @pytest.mark.parametrize("mock_vllm_config", [
+        MockVllmConfig("meta-llama/Meta-Llama-3-8B-Instruct",
+                       "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", "auto"),
+        MockVllmConfig("meta-llama/Meta-Llama-3-8B-Instruct",
+                       "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", "fp8"),
+    ])
+    def test_forward_pass(self, mock_vllm_config: MockVllmConfig, rng: PRNGKey,
+                          mesh: Mesh, mock_model_inputs):
+        """Tests the forward pass of the EagleLlama3ForCausalLM model."""
+
+        draft_model_config = mock_vllm_config.speculative_config.draft_model_config
+        hf_config = draft_model_config.hf_config
+        model = EagleLlama3ForCausalLM(mock_vllm_config, rng, mesh)
+
+        input_ids, hidden_states, attention_metadata = mock_model_inputs
+
+        kv_caches = create_kv_caches(
+            num_blocks=4,
+            block_size=16,
+            num_kv_heads=hf_config.num_key_value_heads,
+            head_size=hf_config.hidden_size // hf_config.num_attention_heads,
+            mesh=mesh,
+            layer_names=["layer"] * hf_config.num_hidden_layers,
+            cache_dtype=jnp.float8_e4m3fn
+            if mock_vllm_config.cache_config.cache_dtype == "fp8" else
+            jnp.bfloat16)
+
+        _, output_hidden_states, aux_hidden_states = model(
+            kv_caches, input_ids, hidden_states, attention_metadata)
+
+        logits = model.compute_logits(output_hidden_states)
+
+        target_model_config = mock_vllm_config.model_config
+
+        assert output_hidden_states.shape == (
+            input_ids.shape[0], draft_model_config.get_hidden_size())
+        assert logits.shape == (input_ids.shape[0],
+                                target_model_config.get_vocab_size())
+        assert len(aux_hidden_states) == 1
+        assert aux_hidden_states[0].shape == output_hidden_states.shape
+
+    @patch("tpu_inference.models.jax.llama_eagle3.load_hf_weights")
+    def test_load_weights(self, mock_load_hf_weights: MagicMock,
+                          mock_vllm_config: MockVllmConfig, rng: PRNGKey,
+                          mesh: Mesh):
+        """Tests that the load_weights function is called correctly."""
+        model = EagleLlama3ForCausalLM(mock_vllm_config, rng, mesh)
+        model.load_weights(rng)
+
+        mock_load_hf_weights.assert_called_once()
+        call_args = mock_load_hf_weights.call_args.kwargs
+
+        assert call_args["vllm_config"] is mock_vllm_config
+        assert call_args["model"] is model
+        assert call_args["mesh"] is mesh
+        assert call_args["is_draft_model"] is True
+
+        metadata_map = call_args["metadata_map"]
+        assert "midlayer.hidden_norm" in metadata_map.name_map
+        assert "lm_head" in metadata_map.name_map
+        assert "d2t" in metadata_map.name_map
+        assert "q_proj" in metadata_map.reshape_map
+        assert metadata_map.reshape_map["q_proj"][-1] == (
+            2 * mock_vllm_config.model_config.get_hidden_size())