tpu-inference 0.11.1.dev202511220812__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
 
@@ -18,7 +32,7 @@ from vllm.lora.layers import (BaseLayerWithLoRA, ColumnParallelLinearWithLoRA,
                               ReplicatedLinearWithLoRA,
                               RowParallelLinearWithLoRA)
 # yapf: enable
-from vllm.lora.models import LoRALayerWeights, PackedLoRALayerWeights
+from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
 from vllm.lora.punica_wrapper import get_punica_wrapper
 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                                MergedColumnParallelLinear,
@@ -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
 
@@ -499,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
@@ -513,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'))
 
@@ -611,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

@@ -0,0 +1,67 @@
+# 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
+
+import pytest
+import vllm
+from vllm.lora.request import LoRARequest
+
+TP = [2] if os.environ.get("USE_V6E8_QUEUE", False) else [1]
+
+
+@pytest.mark.parametrize("tp", TP)
+def test_lora_performance(tp):
+    prompt = "What is 1+1? \n"
+    llm_without_lora = vllm.LLM(
+        model="Qwen/Qwen2.5-3B-Instruct",
+        max_model_len=256,
+        max_num_batched_tokens=64,
+        max_num_seqs=8,
+        tensor_parallel_size=tp,
+    )
+    start_time = time.time()
+    llm_without_lora.generate(
+        prompt,
+        sampling_params=vllm.SamplingParams(max_tokens=16, temperature=0),
+    )[0].outputs[0].text
+    base_time = time.time() - start_time
+
+    del llm_without_lora
+    # Waiting for TPUs to be released
+    time.sleep(10)
+
+    llm_with_lora = vllm.LLM(model="Qwen/Qwen2.5-3B-Instruct",
+                             max_model_len=256,
+                             max_num_batched_tokens=64,
+                             max_num_seqs=8,
+                             tensor_parallel_size=tp,
+                             enable_lora=True,
+                             max_loras=1,
+                             max_lora_rank=8)
+    lora_request = LoRARequest(
+        "lora_adapter_2", 2,
+        "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_2_adapter")
+    start_time = time.time()
+    llm_with_lora.generate(prompt,
+                           sampling_params=vllm.SamplingParams(max_tokens=16,
+                                                               temperature=0),
+                           lora_request=lora_request)[0].outputs[0].text
+    lora_time = time.time() - start_time
+    print(f"Base time: {base_time}, LoRA time: {lora_time}")
+    assert (base_time /
+            lora_time) < 8, f"Base time: {base_time}, LoRA time: {lora_time}"
+
+    del llm_with_lora

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.