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

tests/models/jax/test_qwen3.py

@@ -0,0 +1,169 @@
+# 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 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.qwen3 import Qwen3ForCausalLM
+from tpu_inference.runner.kv_cache import create_kv_caches
+
+
+class MockVllmConfig:
+
+    def __init__(self, model: str, kv_cache_dtype: str):
+        self.model_config = ModelConfig(model)
+        self.model_config.dtype = jnp.bfloat16
+        self.load_config = MagicMock()
+        self.load_config.download_dir = None
+        self.cache_config = MagicMock(cache_dtype=kv_cache_dtype)
+
+
+@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])
+    num_devices = len(devices)
+    assert num_devices == 1
+    device_mesh = devices.reshape((num_devices, 1, 1, 1))
+
+    with Mesh(device_mesh,
+              axis_names=('data', 'attn_dp', 'expert', 'model')) as m:
+        yield m
+
+
+@pytest.fixture
+def mock_model_inputs():
+    num_tokens = 8
+    num_reqs = 1
+    max_num_blocks_per_req = 4
+    input_ids = jnp.ones((num_tokens, ), dtype=jnp.int32)
+    positions = jnp.ones((num_tokens, ), dtype=jnp.int32)
+    block_tables = jnp.zeros((num_reqs, max_num_blocks_per_req),
+                             dtype=jnp.int32).reshape(-1)
+    seq_lens = jnp.ones((num_reqs, ), dtype=jnp.int32)
+    query_start_loc = jnp.ones((num_reqs + 1, ), dtype=jnp.int32)
+    request_distribution = jnp.array([0, 0, 0], dtype=jnp.int32)
+
+    attention_metadata = AttentionMetadata(
+        input_positions=positions,
+        block_tables=block_tables,
+        seq_lens=seq_lens,
+        query_start_loc=query_start_loc,
+        request_distribution=request_distribution,
+    )
+    indices_do_sample = jnp.ones((num_reqs, ), dtype=jnp.int32)
+
+    return (input_ids, attention_metadata, indices_do_sample)
+
+
+@pytest.fixture
+def rng() -> PRNGKey:
+    """Provides a reusable JAX PRNGKey."""
+    return jax.random.PRNGKey(42)
+
+
+class TestQwen3ForCausalLM:
+
+    @pytest.mark.parametrize("mock_vllm_config", [
+        MockVllmConfig("Qwen/Qwen3-0.6B", "auto"),
+        MockVllmConfig("Qwen/Qwen3-0.6B", "fp8")
+    ])
+    def test_qwen3_600M(self, mock_vllm_config, rng, mesh, mock_model_inputs):
+        """Tests model init and model forward for the 0.6B model variant."""
+
+        # Test model init
+        model = Qwen3ForCausalLM(mock_vllm_config, rng, mesh)
+
+        model_config = mock_vllm_config.model_config
+        hf_config = model_config.hf_config
+
+        assert model.mesh.shape == {
+            "data": 1,
+            "attn_dp": 1,
+            "expert": 1,
+            "model": 1
+        }
+
+        layers = model.model.layers
+        assert len(layers) == hf_config.num_hidden_layers
+        assert isinstance(model.rng, nnx.Rngs)
+        assert model.model.lm_head == model.model.embed.embedding
+
+        attn = layers[0].self_attn
+        hidden_size = hf_config.hidden_size
+        num_heads = hf_config.num_attention_heads
+        num_kv_heads = hf_config.num_key_value_heads
+        rope_theta = hf_config.rope_theta
+        original_head_dim = hf_config.head_dim
+        head_dim = 128
+        intermediate_size = hf_config.intermediate_size
+
+        assert attn.hidden_size == hidden_size
+        assert attn.num_heads == num_heads
+        assert attn.num_kv_heads == num_kv_heads
+        assert attn.rope_theta == rope_theta
+        assert attn.head_dim_original == original_head_dim
+        assert attn.head_dim == head_dim
+        assert attn.q_proj.kernel.shape == (hidden_size, num_heads, head_dim)
+        assert attn.k_proj.kernel.shape == (hidden_size, num_kv_heads,
+                                            head_dim)
+        assert attn.v_proj.kernel.shape == (hidden_size, num_kv_heads,
+                                            head_dim)
+        assert attn.o_proj.kernel.shape == (num_heads, head_dim, hidden_size)
+
+        mlp = layers[0].mlp
+        assert mlp.gate_proj.kernel.shape == (hidden_size, intermediate_size)
+        assert mlp.up_proj.kernel.shape == (hidden_size, intermediate_size)
+        assert mlp.down_proj.kernel.shape == (intermediate_size, hidden_size)
+
+        # Test model load
+        model.load_weights(rng)
+
+        # Test model forward
+        kv_caches = create_kv_caches(
+            num_blocks=4,
+            block_size=32,
+            num_kv_heads=num_kv_heads,
+            head_size=head_dim,
+            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)
+        # 1 seq with 16 tokens
+        input_ids, attention_metadata, indices_do_sample = mock_model_inputs
+        kv_caches, hidden_states, aux_hidden_states = model(
+            kv_caches, input_ids, attention_metadata)
+        assert hidden_states.shape == (8, hidden_size)
+        assert len(aux_hidden_states) == 0
+
+        hidden_states = hidden_states[indices_do_sample]
+        assert hidden_states.shape == (1, hidden_size)
+
+        logits = model.compute_logits(hidden_states)
+        assert logits.shape == (1, hf_config.vocab_size)

tests/models/jax/test_weight_loading.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# Test for LoRA weight loading API
+
+import os
+import tempfile
+from dataclasses import dataclass
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from flax import nnx
+from jax._src import test_util as jtu
+from jax.sharding import Mesh
+from safetensors.numpy import save_file
+
+from tpu_inference.models.jax.utils.weight_utils import (
+    MetadataMap, load_hf_weights, transfer_state_with_mappings)
+
+# ----- nnx.Module Wrappers -----
+
+
+class SourceLayer(nnx.Module):
+
+    def __init__(self, rngs):
+        self.kernel = nnx.Param(jax.random.normal(rngs(), (4, 4)))
+        self.bias = nnx.Param(jax.random.normal(rngs(), (4, )))
+
+
+class SourceModel(nnx.Module):
+
+    def __init__(self, rngs):
+        self.src_lm_head = nnx.Param(jax.random.normal(rngs(), (2, 4)))
+        self.layers = {0: SourceLayer(rngs)}
+
+
+class TargetLinear(nnx.Module):
+
+    def __init__(self, rngs):
+        self.kernel = nnx.Param(jnp.zeros((4, 4)))
+        self.bias = nnx.Param(jnp.zeros((4, )))
+
+
+class TargetBlock(nnx.Module):
+
+    def __init__(self, rngs):
+        self.mlp = {"up_proj": TargetLinear(rngs)}
+
+
+class TargetModel(nnx.Module):
+
+    def __init__(self, rngs):
+        self.tgt_lm_head = nnx.Param(jnp.zeros((2, 4)))
+        self.model = {"layers": {0: TargetBlock(rngs)}}
+
+
+# ----- Test -----
+class WeightTransfer(jtu.JaxTestCase):
+
+    def test_transfer_state(self):
+        rng = nnx.Rngs(0)
+        src_model = SourceModel(rng)
+        tgt_model = TargetModel(rng)
+
+        # Get split states
+        _, src_state = nnx.split(src_model)
+        _, tgt_state = nnx.split(tgt_model)
+
+        # Overwrite known values
+        src_state["layers"][0]["kernel"].value = jnp.ones((4, 4)) * 42.0
+        src_state["layers"][0]["bias"].value = jnp.ones((4, )) * 7.0
+        src_state["src_lm_head"].value = jnp.ones((2, 4)) * 6.0
+        # Mapping for both kernel and bias
+        mappings = {
+            "layers.*.kernel": ("model.layers.*.mlp.up_proj.kernel", (None, )),
+            "layers.*.bias": ("model.layers.*.mlp.up_proj.bias", (None, )),
+            "src_lm_head": ("tgt_lm_head", (None, None)),
+        }
+
+        # Transfer
+        new_tgt_state = transfer_state_with_mappings(src_state, tgt_state,
+                                                     mappings)
+
+        # Assert correctness
+        assert jnp.allclose(
+            new_tgt_state["model"]["layers"][0]["mlp"]["up_proj"]
+            ["kernel"].value, 42.0)
+        assert jnp.allclose(
+            new_tgt_state["model"]["layers"][0]["mlp"]["up_proj"]
+            ["bias"].value, 7.0)
+        assert jnp.allclose(new_tgt_state["tgt_lm_head"].value, 6.0)
+
+
+# ----- Mocks for dtype test -----
+
+
+class DtypeTestModel(nnx.Module):
+
+    def __init__(self, dtype: jnp.dtype, rngs: nnx.Rngs):
+        self.weight_to_cast = nnx.Param(jnp.zeros((2, 2), dtype=dtype))
+        self.weight_to_keep = nnx.Param(jnp.zeros((2, 2), dtype=dtype))
+
+
+@dataclass
+class MockModelConfig:
+    model: str
+    dtype: jnp.dtype
+    hf_config: Any = None
+
+    def get_vocab_size(self):
+        return 1
+
+    def get_hidden_size(self):
+        return 1
+
+    def get_head_size(self):
+        return 1
+
+    is_multimodal_model: bool = False
+
+
+@dataclass
+class MockLoadConfig:
+    download_dir: str
+
+
+@dataclass
+class MockVllmConfig:
+    model_config: MockModelConfig
+    load_config: MockLoadConfig
+    speculative_config: Any = None
+
+
+class WeightLoadingDtypeTest(jtu.JaxTestCase):
+
+    def setUp(self):
+        super().setUp()
+        self.tempdir = tempfile.TemporaryDirectory()
+        self.addCleanup(self.tempdir.cleanup)
+
+        # Create dummy safetensors file
+        tensors = {
+            "weight_to_cast.weight": np.ones((2, 2), dtype=np.float32),
+            "weight_to_keep.weight": np.ones((2, 2), dtype=np.float32),
+        }
+        self.safetensors_path = os.path.join(self.tempdir.name,
+                                             "model.safetensors")
+        save_file(tensors, self.safetensors_path)
+
+    def test_keep_original_dtype(self):
+        rng = nnx.Rngs(0)
+        model_dtype = jnp.bfloat16
+        model = DtypeTestModel(dtype=model_dtype, rngs=rng)
+
+        mock_model_config = MockModelConfig(model=self.tempdir.name,
+                                            dtype=model_dtype)
+        mock_load_config = MockLoadConfig(download_dir=self.tempdir.name)
+        vllm_config = MockVllmConfig(model_config=mock_model_config,
+                                     load_config=mock_load_config)
+
+        mesh = Mesh(jax.devices(), ("model", ))
+
+        name_map = {
+            "weight_to_cast": "weight_to_cast",
+            "weight_to_keep": "weight_to_keep",
+        }
+        metadata_map = MetadataMap(name_map=name_map)
+
+        keep_original_dtype_keys_regex = [r"weight_to_keep.*"]
+
+        load_hf_weights(
+            vllm_config=vllm_config,
+            model=model,
+            metadata_map=metadata_map,
+            mesh=mesh,
+            keep_original_dtype_keys_regex=keep_original_dtype_keys_regex,
+        )
+
+        self.assertEqual(model.weight_to_cast.value.dtype, model_dtype)
+        self.assertEqual(model.weight_to_keep.value.dtype, jnp.float32)

tests/models/jax/utils/__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/models/jax/utils/test_multi_modal_utils.py

@@ -0,0 +1,212 @@
+# 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.
+
+# test_multi_modal_utils.py
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from tpu_inference.models.jax.utils.multi_modal_utils import (
+    MultiModalEmbeddings, NestedTensors, flatten_embeddings,
+    merge_multimodal_embeddings, sanity_check_mm_encoder_outputs)
+
+# --- Tests for sanity_check_mm_encoder_outputs ---
+
+
+def test_sanity_check_valid_list():
+    """Tests sanity_check with a valid list of 2D embeddings."""
+    embeddings: MultiModalEmbeddings = [
+        jnp.ones((10, 128)), jnp.ones((15, 128))
+    ]
+    sanity_check_mm_encoder_outputs(embeddings, 2)
+    # No assertion error expected
+
+
+def test_sanity_check_valid_tuple():
+    """Tests sanity_check with a valid tuple of 2D embeddings."""
+    embeddings: MultiModalEmbeddings = (jnp.ones((10, 128)), jnp.ones(
+        (15, 128)))
+    sanity_check_mm_encoder_outputs(embeddings, 2)
+    # No assertion error expected
+
+
+def test_sanity_check_valid_3d_jax_array():
+    """Tests sanity_check with a valid 3D jax.Array."""
+    embeddings: MultiModalEmbeddings = jnp.ones((2, 10, 128))
+    # This is valid because mm_embeddings is iterable, and each item (e)
+    # in the first dimension has e.ndim == 2.
+    sanity_check_mm_encoder_outputs(embeddings, 2)
+    # No assertion error expected
+
+
+def test_sanity_check_invalid_type():
+    """Tests sanity_check with an invalid type for embeddings."""
+    with pytest.raises(
+            AssertionError,
+            match=
+            "Expected multimodal embeddings to be a list/tuple of 2D tensors"):
+        sanity_check_mm_encoder_outputs("not a tensor", 1)
+
+
+def test_sanity_check_wrong_num_items():
+    """Tests sanity_check with a mismatch in the number of embeddings."""
+    embeddings: MultiModalEmbeddings = [jnp.ones((10, 128))]
+    with pytest.raises(
+            AssertionError,
+            match="Expected number of multimodal embeddings to match number of"
+    ):
+        sanity_check_mm_encoder_outputs(embeddings, 2)
+
+
+def test_sanity_check_wrong_dimensions_in_list():
+    """Tests sanity_check with non-2D tensors within the list."""
+    embeddings: MultiModalEmbeddings = [jnp.ones((10, 128, 1))]
+    with pytest.raises(
+            AssertionError,
+            match=
+            "Expected multimodal embeddings to be a sequence of 2D tensors"):
+        sanity_check_mm_encoder_outputs(embeddings, 1)
+
+
+# --- Tests for flatten_embeddings ---
+
+
+def test_flatten_single_array():
+    """Tests flatten_embeddings with a single 2D array."""
+    emb: NestedTensors = jnp.arange(12).reshape((3, 4))
+    result = flatten_embeddings(emb)
+    np.testing.assert_array_equal(result, emb)
+
+
+def test_flatten_single_3d_array():
+    """Tests flatten_embeddings with a single 3D array."""
+    emb: NestedTensors = jnp.arange(24).reshape((2, 3, 4))
+    result = flatten_embeddings(emb)
+    expected = jnp.arange(24).reshape((6, 4))
+    np.testing.assert_array_equal(result, expected)
+
+
+def test_flatten_list_of_arrays():
+    """Tests flatten_embeddings with a list of 2D arrays."""
+    emb: NestedTensors = [
+        jnp.arange(12).reshape((3, 4)),
+        jnp.arange(12, 20).reshape((2, 4))
+    ]
+    result = flatten_embeddings(emb)
+    expected = jnp.arange(20).reshape((5, 4))
+    np.testing.assert_array_equal(result, expected)
+
+
+def test_flatten_nested_list():
+    """Tests flatten_embeddings with a nested list of arrays."""
+    emb: NestedTensors = [
+        jnp.arange(6).reshape((2, 3)),
+        [
+            jnp.arange(6, 12).reshape((2, 3)),
+            jnp.arange(12, 15).reshape((1, 3))
+        ]
+    ]
+    result = flatten_embeddings(emb)
+    expected = jnp.arange(15).reshape((5, 3))
+    np.testing.assert_array_equal(result, expected)
+
+
+# --- Tests for merge_multimodal_embeddings ---
+
+EMBED_DIM = 4
+
+
+@pytest.fixture
+def base_embeds():
+    return jnp.zeros((8, EMBED_DIM))
+
+
+def test_merge_single_placeholder(base_embeds):
+    """Tests merging with a single integer placeholder ID."""
+    input_ids = jnp.array([1, 2, -1, -1, 3, 4, -1, 5])
+    inputs_embeds = base_embeds[:len(input_ids)]
+    mm_embeds: NestedTensors = jnp.arange(3 * EMBED_DIM).reshape(
+        (3, EMBED_DIM))
+    result = merge_multimodal_embeddings(input_ids,
+                                         inputs_embeds,
+                                         mm_embeds,
+                                         placeholder_token_id=-1)
+    expected = np.array(inputs_embeds)
+    expected[input_ids == -1] = mm_embeds
+    np.testing.assert_array_equal(result, expected)
+
+
+def test_merge_no_placeholders(base_embeds):
+    """Tests merging when no placeholder tokens are in input_ids."""
+    input_ids = jnp.array([1, 2, 3, 4])
+    inputs_embeds = jnp.arange(len(input_ids) * EMBED_DIM).reshape(
+        (len(input_ids), EMBED_DIM))
+    mm_embeds: NestedTensors = jnp.empty((0, EMBED_DIM))
+
+    # Based on the provided traceback, this raises a TypeError within JAX's gather.
+    with pytest.raises(
+            TypeError,
+            match="Slice size at index 0 in gather op is out of range"):
+        merge_multimodal_embeddings(input_ids,
+                                    inputs_embeds,
+                                    mm_embeds,
+                                    placeholder_token_id=-1)
+
+
+@pytest.mark.parametrize("placeholder_id", [-1, [-1, -2]])
+def test_merge_mm_embeds_count_too_few(placeholder_id, base_embeds):
+    """
+    Tests behavior when fewer embeddings are provided than placeholders.
+    Based on the test results provided, this scenario does NOT raise an error
+    in the testing environment.
+    """
+    input_ids = jnp.array([1, 2, -1, -1, 3])  # 2 placeholders
+    inputs_embeds = base_embeds[:len(input_ids)]
+    mm_embeds_too_few: NestedTensors = jnp.ones((1, EMBED_DIM))
+
+    try:
+        # We are only asserting that this call does not crash.
+        # The actual output in this unexpected case is not being tested.
+        merge_multimodal_embeddings(input_ids,
+                                    inputs_embeds,
+                                    mm_embeds_too_few,
+                                    placeholder_token_id=placeholder_id)
+    except Exception as e:
+        pytest.fail(
+            f"Did not expect an exception based on test logs, but got {type(e).__name__}: {e}"
+        )
+
+
+@pytest.mark.parametrize("placeholder_id", [-1, [-1, -2]])
+def test_merge_mm_embeds_count_too_many_no_raise(placeholder_id, base_embeds):
+    """Tests that no error is raised if mm_embeds are too many; extras are ignored."""
+    input_ids = jnp.array([1, 2, -1, -1, 3])  # 2 placeholders
+    inputs_embeds = base_embeds[:len(input_ids)]
+    mm_embeds_too_many: NestedTensors = jnp.arange(3 * EMBED_DIM).reshape(
+        (3, EMBED_DIM))
+
+    try:
+        result = merge_multimodal_embeddings(
+            input_ids,
+            inputs_embeds,
+            mm_embeds_too_many,
+            placeholder_token_id=placeholder_id)
+        # Check that the first 2 embeddings from mm_embeds_too_many were used.
+        expected = np.array(inputs_embeds)
+        is_mm = np.isin(input_ids, np.array(placeholder_id))
+        expected[is_mm] = flatten_embeddings(mm_embeds_too_many)[:2]
+        np.testing.assert_array_equal(result, expected)
+    except Exception as e:
+        pytest.fail(
+            f"Did not expect an exception, but got {type(e).__name__}: {e}")

tests/platforms/__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/platforms/test_tpu_platform.py

@@ -0,0 +1,54 @@
+# 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 pytest
+import torch
+from vllm.config import CacheConfig, VllmConfig
+
+from tpu_inference.platforms.tpu_platform import TpuPlatform
+
+
+class TestTpuPlatform:
+
+    @pytest.fixture
+    def vllm_config(self):
+        cache_config = CacheConfig(block_size=16,
+                                   gpu_memory_utilization=0.9,
+                                   swap_space=4,
+                                   cache_dtype="fp8")
+
+        vllm_config = MagicMock(spec=VllmConfig)
+        vllm_config.cache_config = cache_config
+        vllm_config.model_config = MagicMock(dtype='bfloat16')
+        vllm_config.scheduler_config = MagicMock(is_multimodal_model=False)
+        vllm_config.parallel_config = MagicMock()
+        vllm_config.compilation_config = MagicMock(mode="dynamo_trace_once",
+                                                   backend="openxla")
+        vllm_config.kv_transfer_config = None
+        return vllm_config
+
+    @pytest.mark.parametrize("chip_name,expected_dtype", [
+        ("v6e", torch.float8_e5m2),
+        ("v5e", torch.float8_e4m3fn),
+    ])
+    def test_fp8_dtype(self, chip_name, expected_dtype):
+        mock_chip_type = MagicMock()
+        mock_chip_type.name = chip_name
+
+        with patch('tpu_inference.platforms.tpu_platform.init_logger'), \
+             patch('tpu_inference.platforms.tpu_platform.device.get_local_chips', return_value=(mock_chip_type, None)), \
+             patch('vllm.envs.VLLM_TPU_USING_PATHWAYS', False):
+            assert TpuPlatform.fp8_dtype() == expected_dtype

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