tpu-inference 0.12.0.dev20251213__py3-none-any.whl → 0.13.2.dev20251230__py3-none-any.whl

tests/layers/vllm/utils.py

@@ -0,0 +1,87 @@
+# 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 jax
+import torch
+import torch.nn.functional as F
+
+from tpu_inference.layers.common.sharding import (MESH_AXIS_NAMES,
+                                                  MESH_AXIS_NAMES_2D)
+
+
+def get_spmd_mesh(num_devices: int = 1, enable_attn_dp: bool = False):
+    devices = sorted(jax.devices(), key=lambda d: d.id)[0:num_devices]
+
+    if enable_attn_dp:
+        if num_devices < 2:
+            raise ValueError(
+                f"enable_attn_dp requires at least 2 devices, got {num_devices}"
+            )
+        axis_names = MESH_AXIS_NAMES
+        attn_dp_size = 2
+        model_size = num_devices // attn_dp_size
+        mesh_shape = (1, attn_dp_size, 1, model_size)
+        return jax.make_mesh(mesh_shape, axis_names, devices=devices)
+    else:
+        axis_names = MESH_AXIS_NAMES_2D
+        mesh_shape = (1, len(devices))
+        return jax.make_mesh(mesh_shape, axis_names, devices=devices)
+
+
+def find_all_layer_type(module: torch.nn.Module, layer_type: torch.nn.Module):
+    ret = []
+    for name, child in module.named_children():
+        if isinstance(child, layer_type):
+            ret.append(child)
+        else:
+            ret.extend(find_all_layer_type(child, layer_type))
+    return ret
+
+
+# TODO(kyuyeunk): Consolidate all reference implementation used for unit tests
+# into a single file.
+def ref_moe(x, router_logits, w1, w2, w1_bias, w2_bias, top_k, renormalize,
+            activation):
+
+    expert_weights = F.softmax(router_logits, dim=-1)
+    expert_weights, expert_indices = torch.topk(expert_weights, top_k, dim=-1)
+    if renormalize:
+        expert_weights /= expert_weights.sum(dim=-1, keepdim=True)
+
+    x = torch.einsum("ti,eoi->teo", x, w1)
+    if w1_bias is not None:
+        x += w1_bias.unsqueeze(0)
+
+    match activation:
+        case "silu":
+            x1, x3 = x.chunk(chunks=2, dim=-1)
+            x = F.silu(x1) * x3
+        case "swigluoai":
+            x1, x3 = x[..., ::2], x[..., 1::2]
+            x1 = x1.clamp(min=None, max=7.0)
+            x3 = x3.clamp(min=-7.0, max=7.0)
+            gated_activation = x1 * torch.sigmoid(x1 * 1.702)
+            x = gated_activation * (x3 + 1)
+        case _:
+            raise NotImplementedError(
+                f"No reference implementation for {activation} activation")
+
+    x = torch.einsum("teo,eio->tei", x, w2)
+    if w2_bias is not None:
+        x += w2_bias.unsqueeze(0)
+
+    seq_indexes = torch.arange(x.shape[0]).unsqueeze(1)
+    x = x[seq_indexes, expert_indices]
+
+    return torch.einsum("tai,ta->ti", x, expert_weights)

tests/lora/__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/lora/conftest.py

@@ -1,3 +1,17 @@
+# 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 tempfile
 
 import pytest

tests/lora/test_bgmv.py

@@ -1,3 +1,17 @@
+# 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 jax
 import torch
 import torchax

tests/lora/test_layers.py

@@ -1,3 +1,17 @@
+# 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 random
 from typing import Optional
 
@@ -28,10 +42,12 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
 from vllm.model_executor.utils import set_random_seed
 from vllm.platforms import current_platform
 
-from tpu_inference.layers.vllm.quantization.common import JaxCommonLinearConfig
+from tpu_inference.layers.vllm.process_weights.cleanup_sharding import \
+    _shard_module_to_tpu
+from tpu_inference.layers.vllm.quantization.configs import \
+    VllmQuantLinearConfig
 from tpu_inference.layers.vllm.quantization.unquantized import \
     VllmUnquantizedLinearMethod
-from tpu_inference.layers.vllm.sharding import _shard_module_to_tpu
 
 from .utils import DummyLoRAManager
 
@@ -205,9 +221,6 @@ def create_random_inputs(
 @pytest.mark.parametrize("repeats", [1, 2, 3])
 @pytest.mark.parametrize("stage", [True, False])
 def test_column_parallel_packed(dist_init, num_loras, repeats, stage) -> None:
-    # TODO(Qiliang Cui): Remove when issue is resolved.
-    if 'TPU7x' in jax.devices()[0].device_kind:
-        pytest.skip("Skipping test on TPU TPU7x.")
     set_random_seed(6)
 
     max_loras = 9
@@ -502,9 +515,13 @@ def _create_random_linear_parallel_layer(layer_type, vllm_config, mesh):
     return linear, lora_linear
 
 
+def _get_devices():
+    return jax.devices()
+
+
 def _create_mesh():
     axis_names = ("data", "model")
-    devices = jax.devices()
+    devices = _get_devices()
     mesh_shape = (1, len(devices))
     mesh = jax.make_mesh(mesh_shape, axis_names, devices=devices)
     return mesh
@@ -516,7 +533,7 @@ def _verify_lora_linear_layer(linear, lora_linear):
         # BaseLinearLayerWithLoRA.weight property guarantees this.
         # if len(devices) != 1, `reorder_concatenated_tensor_for_sharding` function may reorder the out_features dimension of the weight matrix.
         # So the below check will fail.
-        if len(jax.devices()) == 1:
+        if len(_get_devices()) == 1:
             assert torch.equal(linear.weight.data,
                                lora_linear.weight.to('cpu'))
 
@@ -614,7 +631,7 @@ def _create_lora_wrapper(linear,
                          mesh,
                          repeats=1):
     base_linear.weight.data = linear.weight.data
-    jax_config = JaxCommonLinearConfig(vllm_config, mesh, base_linear)
+    jax_config = VllmQuantLinearConfig(vllm_config, mesh, base_linear)
     linear_method = VllmUnquantizedLinearMethod(jax_config)
     base_linear.quant_method = linear_method
     linear_method.process_weights_after_loading(

tests/lora/test_lora.py

@@ -1,3 +1,17 @@
+# 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.
+
 # https://github.com/vllm-project/vllm/blob/ed10f3cea199a7a1f3532fbe367f5c5479a6cae9/tests/tpu/lora/test_lora.py
 import os
 import time
@@ -29,7 +43,7 @@ def setup_vllm(num_loras: int, tp: int = 1) -> vllm.LLM:
 
 
 # For multi-chip test, we only use TP=2 because the base model Qwen/Qwen2.5-3B-Instruct has 2 kv heads and the current attention kernel requires it to be divisible by tp_size.
-TP = [2] if os.environ.get("USE_V6E8_QUEUE", False) else [1]
+TP = [2] if os.environ.get("TEST_LORA_TP", False) else [1]
 
 
 @pytest.mark.parametrize("tp", TP)

tests/lora/test_lora_perf.py

@@ -1,3 +1,17 @@
+# 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 time
 

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