@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c20ef40fd0
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 11:00 UTC

[Hardware][TPU][V1] Multi-LoRA implementation for the V1 TPU backend (#14238)

Signed-off-by: Akshat Tripathi <akshat@krai.ai> Signed-off-by: Chengji Yao <chengjiyao@google.com> Co-authored-by: Chengji Yao <chengjiyao@google.com>

Akshat Tripathi committed May 7, 2025 at 21:28 UTC c20ef40fd0e8663e82911f53d00a64f53beb98aa
19 files changed +929 -46
.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh
+3
@@ -50,6 +50,9 @@ docker run --privileged --net host --shm-size=16G -it \
50 && pytest -s -v /workspace/vllm/tests/v1/entrypoints/llm/test_struct_output_generate.py \
51 && echo TEST_12 \
52 && pytest -s -v /workspace/vllm/tests/tpu/test_moe_pallas.py" \
53 + # Disable the TPU LoRA tests until the feature is activated
54 + # && echo TEST_13 \
55 + # && pytest -s -v /workspace/vllm/tests/tpu/lora/" \
56
57
58 # TODO: This test fails because it uses RANDOM_SEED sampling
tests/lora/conftest.py
+1 -1
@@ -47,7 +47,7 @@ def dist_init():
47 temp_file = tempfile.mkstemp()[1]
48
49 backend = "nccl"
50 - if current_platform.is_cpu():
50 + if current_platform.is_cpu() or current_platform.is_tpu():
51 backend = "gloo"
52
53 init_distributed_environment(world_size=1,
tests/tpu/lora/__init__.py
tests/tpu/lora/test_lora.py new
+124
@@ -0,0 +1,124 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import pytest
3 +
4 +import vllm
5 +from vllm.lora.request import LoRARequest
6 +
7 +# This file contains tests to ensure that LoRA works correctly on the TPU
8 +# backend. We use a series of custom trained adapters for Qwen2.5-3B-Instruct
9 +# for this. The adapters are:
10 +# Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_x_adapter, where x ranges
11 +# from 1 to 4.
12 +
13 +# These adapters are trained using a standard huggingface peft training script,
14 +# where all the inputs are "What is 1+1? \n" and all the outputs are "x". We run
15 +# 100 training iterations with a training batch size of 100.
16 +
17 +
18 +@pytest.fixture(scope="function", autouse=True)
19 +def use_v1_only(monkeypatch: pytest.MonkeyPatch):
20 + """
21 + Since Multi-LoRA is only supported on the v1 TPU backend, set VLLM_USE_V1=1
22 + for all tests in this file
23 + """
24 + with monkeypatch.context() as m:
25 + m.setenv("VLLM_USE_V1", "1")
26 + yield
27 +
28 +
29 +def setup_vllm(num_loras: int) -> vllm.LLM:
30 + return vllm.LLM(model="Qwen/Qwen2.5-3B-Instruct",
31 + num_scheduler_steps=1,
32 + max_model_len=256,
33 + max_seq_len_to_capture=256,
34 + max_num_seqs=8,
35 + enable_lora=True,
36 + max_loras=num_loras,
37 + max_lora_rank=8)
38 +
39 +
40 +def test_single_lora():
41 + """
42 + This test ensures we can run a single LoRA adapter on the TPU backend.
43 + We run "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_1_adapter" which
44 + will force Qwen2.5-3B-Instruct to claim 1+1=1.
45 + """
46 +
47 + llm = setup_vllm(1)
48 +
49 + prompt = "What is 1+1? \n"
50 +
51 + lora_request = LoRARequest(
52 + "lora_adapter_1", 1,
53 + "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_1_adapter")
54 + output = llm.generate(prompt,
55 + sampling_params=vllm.SamplingParams(max_tokens=256,
56 + temperature=0),
57 + lora_request=lora_request)[0].outputs[0].text
58 +
59 + answer = output.strip()[0]
60 +
61 + assert answer.isdigit()
62 + assert int(answer) == 1
63 +
64 +
65 +def test_lora_hotswapping():
66 + """
67 + This test ensures we can run multiple LoRA adapters on the TPU backend, even
68 + if we only have space to store 1.
69 +
70 + We run "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_x_adapter" which
71 + will force Qwen2.5-3B-Instruct to claim 1+1=x, for a range of x.
72 + """
73 +
74 + lora_name_template = \
75 + "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_{}_adapter"
76 + lora_requests = [
77 + LoRARequest(f"lora_adapter_{i}", i, lora_name_template.format(i))
78 + for i in range(1, 5)
79 + ]
80 +
81 + llm = setup_vllm(1)
82 +
83 + prompt = "What is 1+1? \n"
84 +
85 + for i, req in enumerate(lora_requests):
86 + output = llm.generate(prompt,
87 + sampling_params=vllm.SamplingParams(
88 + max_tokens=256, temperature=0),
89 + lora_request=req)[0].outputs[0].text
90 + answer = output.strip()[0]
91 +
92 + assert answer.isdigit()
93 + assert int(answer) == i + 1
94 +
95 +
96 +def test_multi_lora():
97 + """
98 + This test ensures we can run multiple LoRA adapters on the TPU backend, when
99 + we have enough space to store all of them.
100 +
101 + We run "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_x_adapter" which
102 + will force Qwen2.5-3B-Instruct to claim 1+1=x, for a range of x.
103 + """
104 + lora_name_template = \
105 + "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_{}_adapter"
106 + lora_requests = [
107 + LoRARequest(f"lora_adapter_{i}", i, lora_name_template.format(i))
108 + for i in range(1, 5)
109 + ]
110 +
111 + llm = setup_vllm(4)
112 +
113 + prompt = "What is 1+1? \n"
114 +
115 + for i, req in enumerate(lora_requests):
116 + output = llm.generate(prompt,
117 + sampling_params=vllm.SamplingParams(
118 + max_tokens=256, temperature=0),
119 + lora_request=req)[0].outputs[0].text
120 +
121 + answer = output.strip()[0]
122 +
123 + assert answer.isdigit()
124 + assert int(output.strip()[0]) == i + 1
tests/tpu/lora/test_pallas_kernels.py new
+73
@@ -0,0 +1,73 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import pytest
3 +import torch
4 +
5 +# Required to register the custom ops
6 +import vllm.lora.ops.xla_ops.pallas # noqa # pylint: disable=unused-import
7 +
8 +N_TOKENS = [16, 1024, 4096]
9 +HIDDEN_SIZES = [1024, 2048, 4096]
10 +
11 +DTYPES = [torch.bfloat16]
12 +NUM_LORA = [1, 4, 16]
13 +RANKS = [32, 256, 512]
14 +
15 +
16 +def generate_test_data(T, D, L, N, seed, dtype=torch.float32):
17 + """
18 + Inputs: (All integers)
19 + T: Total number of tokens
20 + D: Input dim
21 + L: LoRA Dim
22 + N: N LoRAs
23 +
24 + Outputs:
25 + inputs: torch.Tensor - shape (T, D)
26 + loras: torch.Tensor - shape (N, 1, L, D)
27 + idxs: torch.Tensor - shape (T, ) - all values must be in [0, N)
28 +
29 + ref_output: torch.Tensor - shape (T, L) - inputs @ loras[idxs].T
30 + """
31 + torch.manual_seed(seed)
32 +
33 + inputs = torch.randn((T, D), device="xla", dtype=dtype)
34 + loras = torch.randn((N, 1, L, D), device="xla", dtype=dtype)
35 + idxs = torch.randint(0, N, (T, ), dtype=torch.int32, device="xla")
36 +
37 + ref_output = ref_bgmv(inputs, loras, idxs)
38 + return inputs, loras, idxs, ref_output
39 +
40 +
41 +def ref_bgmv(inputs: torch.Tensor, loras: torch.Tensor, idxs: torch.Tensor):
42 + selected_loras = loras[idxs]
43 + if len(selected_loras.shape) == 4:
44 + selected_loras = selected_loras.squeeze(axis=1)
45 +
46 + batch_size, output_size, input_size = selected_loras.shape
47 + return (selected_loras @ inputs.reshape(
48 + (batch_size, input_size, 1))).reshape((batch_size, output_size))
49 +
50 +
51 +# Parameterize tests with various shapes and dtypes
52 +@pytest.mark.parametrize("T", N_TOKENS)
53 +@pytest.mark.parametrize("D", HIDDEN_SIZES)
54 +@pytest.mark.parametrize("L", RANKS)
55 +@pytest.mark.parametrize("N", NUM_LORA)
56 +@pytest.mark.parametrize("dtype", DTYPES)
57 +@pytest.mark.parametrize("op_type", ["shrink", "expand"])
58 +@pytest.mark.parametrize("seed", [0])
59 +def test_bgmv_correctness(T, D, L, N, dtype, op_type, seed):
60 + if op_type == "expand":
61 + D, L = L, D
62 +
63 + inputs, loras, idxs, ref_output = generate_test_data(
64 + T, D, L, N, seed, dtype)
65 +
66 + # Run bgmv
67 + output = torch.ops.xla.bgmv(inputs, loras, idxs)
68 +
69 + # Make sure we have no NaNs
70 + assert not torch.any(torch.isnan(output))
71 +
72 + # Compare with reference output
73 + assert torch.allclose(output, ref_output, rtol=1e-2, atol=1e-2)
vllm/config.py
+3 -2
@@ -2694,8 +2694,8 @@ class LoRAConfig:
2694 lora_extra_vocab_size: int = 256
2695 """Maximum size of extra vocabulary that can be present in a LoRA adapter
2696 (added to the base model vocabulary)."""
2697 - # This is a constant.
2698 - lora_vocab_padding_size: ClassVar[int] = 256
2697 + lora_vocab_padding_size: ClassVar[int] = current_platform\
2698 + .get_lora_vocab_padding_size()
2699 long_lora_scaling_factors: Optional[tuple[float, ...]] = None
2700 """Specify multiple scaling factors (which can be different from base model
2701 scaling factor - see eg. Long LoRA) to allow for multiple LoRA adapters
@@ -2723,6 +2723,7 @@ class LoRAConfig:
2723 factors.append(self.fully_sharded_loras)
2724 factors.append(self.lora_dtype)
2725 factors.append(self.lora_extra_vocab_size)
2726 + factors.append(self.lora_vocab_padding_size)
2727 factors.append(self.long_lora_scaling_factors)
2728 factors.append(self.bias_enabled)
2729 hash_str = hashlib.md5(str(factors).encode(),
vllm/lora/fully_sharded_layers.py
+29 -10
@@ -16,6 +16,7 @@ from vllm.lora.layers import (ColumnParallelLinearWithLoRA,
16 MergedQKVParallelLinearWithLoRA,
17 QKVParallelLinearWithLoRA,
18 RowParallelLinearWithLoRA)
19 +from vllm.platforms import current_platform
20
21 if TYPE_CHECKING:
22 pass
@@ -57,15 +58,25 @@ def _mcp_apply(x, bias, layer: ColumnParallelLinearWithLoRA):
58 device=x.device,
59 )
60
60 - layer.punica_wrapper.add_shrink(buffers, x, layer.lora_a_stacked, 1.0)
61 + shrunk_buffers: Optional[torch.Tensor] = layer.punica_wrapper.add_shrink(
62 + buffers, x, layer.lora_a_stacked, 1.0)
63 +
64 + if not current_platform.can_update_inplace():
65 + buffers = shrunk_buffers
66 +
67 buffers = tensor_model_parallel_all_gather(buffers)
62 - layer.punica_wrapper.add_expand(output,
63 - buffers,
64 - layer.lora_b_stacked,
65 - layer.lora_bias_stacked,
66 - layer.output_slices,
67 - offset_start=0,
68 - add_input=True)
68 +
69 + lora_output: Optional[torch.Tensor] = layer.punica_wrapper.add_expand(
70 + output,
71 + buffers,
72 + layer.lora_b_stacked,
73 + layer.lora_bias_stacked,
74 + layer.output_slices,
75 + offset_start=0,
76 + add_input=True)
77 +
78 + if not current_platform.can_update_inplace():
79 + output = lora_output
80
81 output = output.view(*out_orig_shape)
82 # now have column partitioned and packed output
@@ -292,7 +303,11 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):
303 device=x.device,
304 )
305
295 - self.punica_wrapper.add_shrink(buffer, x, self.lora_a_stacked, 1.0)
306 + shrunk_buffer: Optional[torch.Tensor] = self.punica_wrapper.add_shrink(
307 + buffer, x, self.lora_a_stacked, 1.0)
308 + if not current_platform.can_update_inplace():
309 + buffer = shrunk_buffer
310 +
311 buffer = tensor_model_parallel_all_reduce(buffer)
312
313 # following S-LoRA, allows the fusing of all_gather and all_reduce
@@ -304,7 +319,7 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):
319 # NOTE offset are based on the rank.
320 shard_size = self.lora_b_stacked[0].shape[2]
321 offset_start = self.tp_rank * shard_size
307 - self.punica_wrapper.add_expand(
322 + lora_output: Optional[torch.Tensor] = self.punica_wrapper.add_expand(
323 output,
324 buffer,
325 self.lora_b_stacked,
@@ -313,6 +328,10 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):
328 offset_start=offset_start,
329 add_input=True,
330 )
331 +
332 + if not current_platform.can_update_inplace():
333 + output = lora_output
334 +
335 output = output.view(*out_orig_shape)
336 return output
337
vllm/lora/layers.py
+37 -16
@@ -261,10 +261,17 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
261 full_lora_a_embeddings.shape[1],
262 -1,
263 )
264 - self.punica_wrapper.add_lora_embedding(full_output,
265 - full_lora_a_embeddings,
266 - self.lora_b_stacked,
267 - add_input=True)
264 +
265 + lora_output: Optional[
266 + torch.Tensor] = self.punica_wrapper.add_lora_embedding(
267 + full_output,
268 + full_lora_a_embeddings,
269 + self.lora_b_stacked,
270 + add_input=True)
271 +
272 + if not current_platform.can_update_inplace():
273 + full_output = lora_output
274 +
275 return full_output.view_as(full_output_org)
276
277 @classmethod
@@ -410,10 +417,13 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
417 output = output.flatten(0, 1)
418 x = x.flatten(0, 1)
419
413 - self.punica_wrapper.add_lora_linear(output, x, self.lora_a_stacked,
414 - self.lora_b_stacked,
415 - self.lora_bias_stacked, 1.0,
416 - self.output_slices)
420 + lora_output: Optional[
421 + torch.Tensor] = self.punica_wrapper.add_lora_linear(
422 + output, x, self.lora_a_stacked, self.lora_b_stacked,
423 + self.lora_bias_stacked, 1.0, self.output_slices)
424 + if not current_platform.can_update_inplace():
425 + output = lora_output
426 +
427 return output
428
429 @property
@@ -1133,15 +1143,23 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
1143 torch.matmul(self.embeddings_tensors,
1144 hidden_states.T,
1145 out=lora_logits[:-1])
1136 - lora_logits[-1] = float("-inf")
1146 +
1147 + neg_inf, pos_inf = current_platform.get_infinity_values(
1148 + lora_logits.dtype)
1149 +
1150 + lora_logits[-1] = neg_inf
1151 lora_logits = lora_logits.mT
1152 indices_padded = self.punica_wrapper.sampler_indices_padded
1153 +
1154 + if current_platform.is_tpu():
1155 + indices_padded = indices_padded[:logits.size(0)]
1156 +
1157 lora_logits = (lora_logits.reshape(
1158 lora_logits.shape[0] * lora_logits.shape[1],
1159 lora_logits.shape[2],
1142 - ).index_select(0, indices_padded).nan_to_num_(nan=float("-inf"),
1143 - posinf=float("inf"),
1144 - neginf=float("-inf")))
1160 + ).index_select(0, indices_padded).nan_to_num_(nan=neg_inf,
1161 + posinf=pos_inf,
1162 + neginf=neg_inf))
1163
1164 # HPU needs special handling to prune out dummy samples.
1165 if current_platform.is_hpu():
@@ -1151,10 +1169,13 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
1169 self.base_layer.org_vocab_size:self.base_layer.org_vocab_size +
1170 lora_logits.shape[1]] = lora_logits
1171
1154 - # LogitsProcessorWithLoRA always using bgmv
1155 - self.punica_wrapper.add_lora_logits(logits, hidden_states,
1156 - self.lora_a_stacked,
1157 - self.lora_b_stacked, 1.0)
1172 + lora_output: Optional[
1173 + torch.Tensor] = self.punica_wrapper.add_lora_logits(
1174 + logits, hidden_states, self.lora_a_stacked,
1175 + self.lora_b_stacked, 1.0)
1176 +
1177 + if not current_platform.can_update_inplace():
1178 + logits = lora_output
1179
1180 # Remove paddings in vocab (if any).
1181 logits = logits[:, :self.base_layer.vocab_size]
vllm/lora/ops/xla_ops/__init__.py new
+6
@@ -0,0 +1,6 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from vllm.lora.ops.xla_ops.lora_ops import (bgmv_expand, bgmv_expand_slice,
4 + bgmv_shrink)
5 +
6 +__all__ = ["bgmv_expand", "bgmv_expand_slice", "bgmv_shrink"]
vllm/lora/ops/xla_ops/lora_ops.py new
+106
@@ -0,0 +1,106 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import torch
4 +
5 +# Required to register the custom ops
6 +import vllm.lora.ops.xla_ops.pallas # noqa # pylint: disable=unused-import
7 +
8 +
9 +def bgmv_expand(inputs: torch.Tensor,
10 + lora_b_weights: torch.Tensor,
11 + output_tensor: torch.Tensor,
12 + lora_indices_tensor: torch.Tensor,
13 + add_inputs: bool = True):
14 + """
15 + Args:
16 + inputs (torch.Tensor): Input tensor of shape [num_tokens, hidden_size].
17 +
18 + lora_b_weights (torch.Tensor): LoRA weights of shape
19 + [num_loras, lora_rank, hidden_size].
20 +
21 + output_tensor (torch.Tensor): output tensor of shape
22 + [num_tokens, hidden_size * num_slices].
23 +
24 + lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
25 + indicating which LoRA matrix to use for each token.
26 + add_inputs (bool): Whether or not to add the input tensor to the output
27 + tensor.
28 + """
29 +
30 + outputs = torch.ops.xla.bgmv(inputs, lora_b_weights, lora_indices_tensor)
31 + n_tokens = outputs.size(0)
32 +
33 + limit = output_tensor.shape[0]
34 + if outputs.shape[0] == 1 and output_tensor.shape[0] != 1:
35 + limit = 1
36 +
37 + outputs = torch.cat(
38 + (outputs,
39 + torch.zeros((n_tokens, output_tensor.shape[1] - outputs.shape[1]),
40 + device=outputs.device)),
41 + dim=1)
42 +
43 + if add_inputs:
44 + return output_tensor + outputs[:limit, :]
45 + else:
46 + return outputs[:limit, :]
47 +
48 +
49 +def bgmv_shrink(inputs: torch.Tensor,
50 + lora_b_weights: torch.Tensor,
51 + output_tensor: torch.Tensor,
52 + lora_indices_tensor: torch.Tensor,
53 + scaling: float = 1.0):
54 + """
55 + Args:
56 + inputs (torch.Tensor): Input tensor of shape [num_tokens, hidden_size].
57 + lora_b_weights (torch.Tensor): LoRA weights of shape
58 + [num_loras, lora_rank, hidden_size].
59 + output_tensor (torch.Tensor): (Unused) output tensor (placeholder).
60 + lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
61 + indicating which LoRA matrix to use for each token.
62 + scaling (float, optional): Scalar multiplier applied to the output.
63 + """
64 +
65 + return scaling * torch.ops.xla.bgmv(inputs, lora_b_weights,
66 + lora_indices_tensor)
67 +
68 +
69 +def bgmv_expand_slice(inputs: torch.Tensor,
70 + lora_b_weights: torch.Tensor,
71 + output_tensor: torch.Tensor,
72 + lora_indices_tensor: torch.Tensor,
73 + slice_offset: int,
74 + slice_size: int,
75 + add_inputs: bool = True):
76 + """
77 + Args:
78 + inputs (torch.Tensor): Input tensor of shape [num_tokens, hidden_size].
79 +
80 + lora_b_weights (torch.Tensor): LoRA weights of shape
81 + [num_loras, lora_rank, hidden_size].
82 +
83 + output_tensor (torch.Tensor): output tensor of shape
84 + [num_tokens, hidden_size * num_slices].
85 +
86 + lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
87 + indicating which LoRA matrix to use for each token.
88 + add_inputs (bool): Whether or not to add the input tensor to the output
89 + tensor.
90 + """
91 + outputs = torch.ops.xla.bgmv(inputs, lora_b_weights, lora_indices_tensor)
92 + n_tokens = outputs.size(0)
93 +
94 + outputs = torch.cat((
95 + torch.zeros((n_tokens, slice_offset), device=outputs.device),
96 + outputs,
97 + torch.zeros(
98 + (n_tokens, output_tensor.shape[1] - (slice_offset + slice_size)),
99 + device=outputs.device),
100 + ),
101 + dim=1)
102 +
103 + if add_inputs:
104 + return output_tensor + outputs
105 + else:
106 + return outputs
vllm/lora/ops/xla_ops/pallas.py new
+133
@@ -0,0 +1,133 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import functools
3 +
4 +import jax
5 +import jax.numpy as jnp
6 +import torch
7 +from jax.experimental import pallas as pl
8 +from jax.experimental.pallas import tpu as pltpu
9 +from torch.library import impl
10 +from torch_xla.experimental.custom_kernel import (XLA_LIB, jax_import_guard,
11 + make_kernel_from_pallas)
12 +
13 +# TODO: Tune these
14 +TOKENS_BLOCK = 16
15 +LORA_RANK_BLOCK = 128
16 +DIM_BLOCK_SIZE = 128
17 +
18 +
19 +def _bgmv_kernel(bT: int, bL: int, idx_ref, inp_ref, lora_ref, out_ref,
20 + acc_ref, mask_ref):
21 +
22 + @pl.when(pl.program_id(2) == 0)
23 + def _():
24 + acc_ref[...] = jnp.zeros_like(acc_ref[...], dtype=jnp.float32)
25 +
26 + t = pl.program_id(0)
27 +
28 + for i in range(bT):
29 + idx = idx_ref[i + bT * t]
30 + mask_ref[...] = jnp.zeros_like(mask_ref[...], dtype=jnp.float32)
31 + mask_ref[i, :] = jnp.ones((bL, ), dtype=jnp.float32)
32 +
33 + acc_ref[...] += jax.lax.dot_general(
34 + inp_ref[...],
35 + lora_ref[idx, ...], (((1, ), (1, )), ((), ())),
36 + preferred_element_type=jnp.float32) * mask_ref[...]
37 +
38 + @pl.when(pl.program_id(2) == pl.num_programs(2) - 1)
39 + def _():
40 + out_ref[...] = acc_ref[...].astype(out_ref.dtype)
41 +
42 +
43 +@jax.jit
44 +def _bgmv(
45 + idxs: jax.Array, # (T, ) int32
46 + inputs: jax.Array, # (T, D) model dtype
47 + loras: jax.Array # (N, L, D) model dtype
48 +) -> jax.Array: # (T, L) model dtype
49 + T, D = inputs.shape
50 + N, L, _ = loras.shape
51 +
52 + return pl.pallas_call(
53 + kernel=functools.partial(_bgmv_kernel, TOKENS_BLOCK, LORA_RANK_BLOCK),
54 + out_shape=jax.ShapeDtypeStruct((T, L), dtype=inputs.dtype),
55 + grid_spec=pltpu.PrefetchScalarGridSpec(
56 + num_scalar_prefetch=1,
57 + grid=(T // TOKENS_BLOCK, L // LORA_RANK_BLOCK,
58 + D // DIM_BLOCK_SIZE),
59 + in_specs=[
60 + pl.BlockSpec((TOKENS_BLOCK, DIM_BLOCK_SIZE),
61 + lambda i, j, k, block_idx: (i, k)),
62 + pl.BlockSpec((N, LORA_RANK_BLOCK, DIM_BLOCK_SIZE),
63 + lambda i, j, k, block_idx: (0, j, k)),
64 + ],
65 + out_specs=pl.BlockSpec((TOKENS_BLOCK, LORA_RANK_BLOCK),
66 + lambda i, j, k, block_idx: (i, j)),
67 + scratch_shapes=[
68 + pltpu.VMEM((TOKENS_BLOCK, LORA_RANK_BLOCK), jnp.float32),
69 + pltpu.VMEM((TOKENS_BLOCK, LORA_RANK_BLOCK), jnp.float32)
70 + ]),
71 + compiler_params=pltpu.TPUCompilerParams(
72 + dimension_semantics=("parallel", "parallel", "arbitrary")),
73 + name="bgmv")(idxs, inputs, loras)
74 +
75 +
76 +def bgmv_shape_function(idxs, inputs, loras):
77 + T, _ = inputs.shape
78 + _, L, _ = loras.shape
79 +
80 + return [((T, L), inputs.dtype)]
81 +
82 +
83 +XLA_LIB.define("bgmv(Tensor inputs, Tensor loras, Tensor idxs) -> Tensor", )
84 +
85 +
86 +@impl(XLA_LIB, "bgmv", "XLA")
87 +def bgmv_xla(inputs: torch.Tensor, loras: torch.Tensor, idxs: torch.IntTensor):
88 + inputs = inputs.to(dtype=loras.dtype)
89 +
90 + if len(loras.shape) == 4:
91 + loras = loras.squeeze(axis=1)
92 +
93 + jax_import_guard()
94 + kernel = make_kernel_from_pallas(_bgmv, bgmv_shape_function)
95 +
96 + T, _ = inputs.shape
97 + _, L, D = loras.shape
98 +
99 + # Pad the loras' rank if it's too low. This is to allow it to fit in a TPU
100 + # register. This has to happen in pytorch, doing it in Jax will lead to NaNs
101 + L1 = L
102 + if LORA_RANK_BLOCK > L or L % LORA_RANK_BLOCK != 0:
103 + L1 = (L // LORA_RANK_BLOCK + 1) * LORA_RANK_BLOCK
104 +
105 + D1 = D
106 + if DIM_BLOCK_SIZE > D or D % DIM_BLOCK_SIZE != 0:
107 + D1 = (D // DIM_BLOCK_SIZE + 1) * DIM_BLOCK_SIZE
108 +
109 + T1 = T
110 + if TOKENS_BLOCK > T or T % TOKENS_BLOCK != 0:
111 + T1 = (T // TOKENS_BLOCK + 1) * TOKENS_BLOCK
112 +
113 + if D1 != D or L1 != L:
114 + loras = torch.nn.functional.pad(loras, (0, D1 - D, 0, L1 - L, 0, 0))
115 + if D1 != D or T1 != T:
116 + inputs = torch.nn.functional.pad(inputs, (0, D1 - D, 0, T1 - T))
117 + if T1 != T:
118 + idxs = torch.nn.functional.pad(idxs, ((0, T1 - T)))
119 +
120 + return kernel(idxs, inputs, loras)[:T, :L]
121 +
122 +
123 +@impl(XLA_LIB, "bgmv", "CompositeExplicitAutograd")
124 +def bgmv_non_xla(inputs: torch.Tensor, loras: torch.Tensor,
125 + idxs: torch.IntTensor):
126 + T, _ = inputs.shape
127 +
128 + if len(loras.shape) == 4:
129 + loras = loras.squeeze(axis=1)
130 +
131 + _, L, _ = loras.shape
132 +
133 + return torch.empty((T, L), device=inputs.device)
vllm/lora/punica_wrapper/punica_base.py
+13 -12
@@ -48,7 +48,7 @@ class PunicaWrapperABC(ABC):
48 lora_a_stacked: Tuple[torch.Tensor, ...],
49 scale: float,
50 **kwargs,
51 - ) -> None:
51 + ) -> Optional[torch.Tensor]:
52 """
53 Performs GEMM for multiple slices of lora_a.
54 """
@@ -66,7 +66,7 @@ class PunicaWrapperABC(ABC):
66 offset_start: int = 0,
67 add_inputs=True,
68 **kwargs,
69 - ) -> None:
69 + ) -> Optional[torch.Tensor]:
70 """
71 Performs GEMM and bias addition for multiple slices of lora_b.
72 """
@@ -80,7 +80,7 @@ class PunicaWrapperABC(ABC):
80 lora_b_stacked: torch.Tensor,
81 add_inputs: bool = True,
82 **kwargs,
83 - ) -> None:
83 + ) -> Optional[torch.Tensor]:
84 """
85 Applies lora specifically for VocabParallelEmbeddingWithLoRA,
86 and this layer only requires the expand operation.
@@ -98,7 +98,7 @@ class PunicaWrapperABC(ABC):
98 output_slices: Tuple[int, ...],
99 *,
100 buffer: Optional[Tuple[torch.Tensor, ...]] = None,
101 - **kwargs) -> None:
101 + **kwargs) -> Optional[torch.Tensor]:
102 """
103 Applicable to linear-related lora.
104 """
@@ -114,7 +114,7 @@ class PunicaWrapperABC(ABC):
114 scale,
115 *,
116 buffer: Optional[torch.Tensor] = None,
117 - **kwargs) -> None:
117 + **kwargs) -> Optional[torch.Tensor]:
118 """
119 Applies lora specifically for LogitsProcessorWithLoRA.
120 """
@@ -207,7 +207,8 @@ class PunicaWrapperBase(PunicaWrapperABC):
207 self._long_lora_indices.zero_()
208 self.indices_len[:] = indices_len
209
210 - def _update_prefill_metada(self, token_lora_tensor: torch.Tensor) -> None:
210 + def _update_prefill_metadata(self,
211 + token_lora_tensor: torch.Tensor) -> None:
212
213 (b_seq_start_tensor, seq_length_tensor, lora_indices_tensor,
214 batch_size, max_length, token_nums,
@@ -334,7 +335,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
335 long_lora_context)
336 if mapping.is_prefill:
337 # Update metadata required for prefill-related operators.
337 - self._update_prefill_metada(self.token_lora_indices)
338 + self._update_prefill_metadata(self.token_lora_indices)
339 self.is_prefill = True
340 else:
341 self.is_prefill = False
@@ -342,7 +343,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
343 @abstractmethod
344 def add_shrink(self, y: Union[Tuple[torch.Tensor, ...], torch.Tensor],
345 x: torch.Tensor, lora_a_stacked: Tuple[torch.Tensor, ...],
345 - scale: float, **kwargs) -> None:
346 + scale: float, **kwargs) -> Optional[torch.Tensor]:
347 """
348 Performs GEMM for multiple slices of lora_a.
349
@@ -369,7 +370,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
370 output_slices: Tuple[int, ...],
371 offset_start: int = 0,
372 add_inputs=True,
372 - **kwargs) -> None:
373 + **kwargs) -> Optional[torch.Tensor]:
374 """
375 Performs GEMM and bias addition for multiple slices of lora_b.
376
@@ -401,7 +402,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
402 x: torch.Tensor,
403 lora_b_stacked: torch.Tensor,
404 add_inputs: bool = True,
404 - **kwargs) -> None:
405 + **kwargs) -> Optional[torch.Tensor]:
406 """
407 Applies lora specifically for VocabParallelEmbeddingWithLoRA.
408 and this layer only requires the expand operation.
@@ -428,7 +429,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
429 output_slices: Tuple[int, ...],
430 *,
431 buffer: Optional[Tuple[torch.Tensor, ...]] = None,
431 - **kwargs) -> None:
432 + **kwargs) -> Optional[torch.Tensor]:
433 """
434 Applicable to linear-related lora.
435
@@ -463,7 +464,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
464 scale,
465 *,
466 buffer: Optional[torch.Tensor] = None,
466 - **kwargs) -> None:
467 + **kwargs) -> Optional[torch.Tensor]:
468 """
469 Applies lora specifically for LogitsProcessorWithLoRA.
470
vllm/lora/punica_wrapper/punica_tpu.py new
+325
@@ -0,0 +1,325 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from typing import Optional, Tuple, Union
4 +
5 +import torch
6 +import torch.nn.functional as F
7 +
8 +from vllm.lora.ops.xla_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
9 +
10 +from .punica_base import PunicaWrapperBase
11 +
12 +
13 +class PunicaWrapperTPU(PunicaWrapperBase):
14 + """
15 + PunicaWrapperTPU is designed to manage and provide metadata for the punica
16 + kernel. The main function is to maintain the state information for
17 + Multi-LoRA, and to provide the interface for the pytorch punica ops.
18 + """
19 +
20 + def __init__(self, max_num_batched_tokens: int, max_batches: int,
21 + device: Union[torch.device, str], **kwargs):
22 + PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches,
23 + device)
24 +
25 + # PunicaWrapperBase defines some tensors with dtype=torch.int64, which
26 + # isn't supported by the TPU. So convert those tensors to int32.
27 + # Not all of them are used by the TPU so only convert the useful ones.
28 + self._token_lora_indices = self._token_lora_indices.to(
29 + dtype=torch.int32)
30 + self._sampler_indices = self._sampler_indices.to(dtype=torch.int32)
31 + self._sampler_indices_padded = self._sampler_indices_padded.to(
32 + dtype=torch.int32)
33 +
34 + torch._dynamo.mark_dynamic(self._token_lora_indices, 0)
35 + torch._dynamo.mark_dynamic(self._embeddings_indices, 1)
36 + torch._dynamo.mark_dynamic(self._sampler_indices_padded, 0)
37 +
38 + def _get_token_lora_indices(self, x: torch.Tensor) -> torch.IntTensor:
39 + return torch.narrow(self._token_lora_indices, 0, 0, x.size(0))
40 +
41 + @property
42 + def embeddings_indices(self) -> torch.Tensor:
43 + """
44 + This property provides access to the indices used for lora embeddings,
45 + specifically for VocabParallelEmbeddingWithLoRA.
46 + """
47 + return self._embeddings_indices[:]
48 +
49 + @property
50 + def sampler_indices_padded(self) -> torch.Tensor:
51 + """
52 + This property provides access to padded sampler indices.
53 + """
54 + return self._sampler_indices_padded[:]
55 +
56 + def shrink(
57 + self,
58 + y: torch.Tensor,
59 + x: torch.Tensor,
60 + w_t_all: torch.Tensor,
61 + scale: float,
62 + ):
63 + if self.no_lora:
64 + return y
65 + return bgmv_shrink(x, w_t_all, y, self._get_token_lora_indices(x),
66 + scale)
67 +
68 + def expand(self, y: torch.Tensor, x: torch.Tensor, w_t_all: torch.Tensor,
69 + add_inputs: bool):
70 + return bgmv_expand(x, w_t_all, y, self._get_token_lora_indices(x),
71 + add_inputs)
72 +
73 + def expand_slice(self, y: torch.Tensor, x: torch.Tensor,
74 + w_t_all: torch.Tensor, y_offset: int, y_slice_size: int,
75 + y_total_size: int, add_inputs: bool) -> torch.Tensor:
76 + return bgmv_expand_slice(x, w_t_all, y,
77 + self._get_token_lora_indices(x), y_offset,
78 + y_slice_size, add_inputs)
79 +
80 + def add_shrink(self, y: Union[Tuple[torch.Tensor, ...], torch.Tensor],
81 + x: torch.Tensor, lora_a_stacked: Tuple[torch.Tensor, ...],
82 + scale: float, **kwargs) -> Optional[torch.Tensor]:
83 + """
84 + Performs GEMM for multiple slices of lora_a.
85 +
86 + Semantics:
87 + for i in range(len(lora_a_stacked)):
88 + y[i] += (x @ lora_a_stacked[i]) * scale
89 +
90 + Args:
91 + y (Union[Tuple[torch.Tensor, ...], torch.Tensor]): Output tensors
92 + x (torch.Tensor): Input tensor
93 + lora_a_stacked (Tuple[torch.Tensor, ...]): lora_a's weights
94 + scale (float): Scaling factor for the operation
95 + """
96 +
97 + torch.ops.xla.dynamo_set_buffer_donor_(y, True)
98 + x = x.view(-1, x.shape[-1])
99 +
100 + for slice_idx in range(len(lora_a_stacked)):
101 + y_s = y[slice_idx]
102 + lora_s = lora_a_stacked[slice_idx]
103 + y_s = self.shrink(y_s, x, lora_s, scale)
104 + y[slice_idx, :, :] = y_s # type: ignore[index]
105 + return y
106 +
107 + def add_expand(self,
108 + y: torch.Tensor,
109 + x: Union[Tuple[torch.Tensor, ...], torch.Tensor],
110 + lora_b_stacked: Tuple[torch.Tensor, ...],
111 + lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]],
112 + output_slices: Tuple[int, ...],
113 + offset_start: int = 0,
114 + add_inputs=True,
115 + **kwargs) -> torch.Tensor:
116 + """
117 + Performs GEMM and bias addition for multiple slices of lora_b.
118 +
119 + Semantics:
120 + for i in range(len(lora_b_stacked)):
121 + slice = output_slices[i]
122 + y[:, offset:offset+slice] += x[i] @ lora_b_stacked[i] +
123 + lora_bias_stacked[i]
124 + offset += slice
125 +
126 + Args:
127 + y (torch.Tensor): Output tensor.
128 + x (Union[Tuple[torch.Tensor, ...], torch.Tensor]): Input tensors
129 + lora_b_stacked (Tuple[torch.Tensor, ...]): lora_b's weight
130 + lora_bias_stacked (Optional[Tuple[torch.Tensor, ...]]):
131 + bias's weight
132 + output_slices (Tuple[int, ...]): Every slice's size
133 + add_inputs (bool): Defaults to True.
134 + """
135 + y_org = y
136 + y = y.view(-1, y.shape[-1])
137 + offset_left = 0
138 +
139 + if lora_bias_stacked is not None:
140 + y = self._apply_bias(self._get_token_lora_indices(y), y,
141 + output_slices, lora_bias_stacked)
142 + for slice_idx in range(len(lora_b_stacked)):
143 + y = self.expand_slice(
144 + y,
145 + x[slice_idx],
146 + lora_b_stacked[slice_idx],
147 + offset_left,
148 + output_slices[slice_idx],
149 + y_total_size=sum(output_slices),
150 + add_inputs=add_inputs,
151 + )
152 + offset_left += output_slices[slice_idx]
153 + return y.view_as(y_org)
154 +
155 + def add_lora_embedding(self,
156 + y: torch.Tensor,
157 + x: torch.Tensor,
158 + lora_b_stacked: torch.Tensor,
159 + add_inputs: bool = True,
160 + **kwargs) -> torch.Tensor:
161 + """
162 + Applies lora specifically for VocabParallelEmbeddingWithLoRA.
163 +
164 + Semantics:
165 + y += x @ lora_b_stacked
166 +
167 + Args:
168 + y (torch.Tensor): Output tensor.
169 + x (torch.Tensor): Input tensor.
170 + lora_b_stacked (torch.Tensor): lora_b's weights.
171 + add_inputs (bool): Default to True.
172 + """
173 +
174 + # Embedding layer only needs the expand op
175 + return self.expand(y, x, lora_b_stacked, add_inputs)
176 +
177 + def add_lora_linear(self,
178 + y: torch.Tensor,
179 + x: torch.Tensor,
180 + lora_a_stacked: Tuple[torch.Tensor, ...],
181 + lora_b_stacked: Tuple[torch.Tensor, ...],
182 + lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]],
183 + scale: float,
184 + output_slices: Tuple[int, ...],
185 + *,
186 + buffer: Optional[Tuple[torch.Tensor, ...]] = None,
187 + **kwargs) -> torch.Tensor:
188 + """
189 + Applicable to linear-related lora.
190 +
191 + Semantics:
192 + for i in range(len(lora_a_stacked)):
193 + y[i] += (
194 + x[i].unsqueeze(0)
195 + @ lora_a_stacked[indices[i], layer_idx, :, :]
196 + @ lora_b_stacked[indices[i], layer_idx, :, :]
197 + * scale
198 + ).squeeze(0)+lora_bias_stacked[i]
199 +
200 + Args:
201 + y (torch.Tensor): Output tensor. Will not be changed in-place.
202 + x (torch.Tensor): Input tensor (T, E)
203 + lora_a_stacked (Tuple[torch.Tensor, ...]): lora_a's weight.
204 + lora_b_stacked (Tuple[torch.Tensor, ...]): lora_b's weight.
205 + lora_bias_stacked (Optional[Tuple[torch.Tensor, ...]]): lora's bias.
206 + scale (float): Scaling factor.
207 + output_slices (Tuple[int, ...]): Every slice's size.
208 + buffer (Optional[Tuple[torch.Tensor, ...]]): Defaults to None.
209 + """
210 +
211 + assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices)
212 + if lora_bias_stacked is not None:
213 + assert len(lora_bias_stacked) == len(output_slices)
214 + y = self._apply_bias(self._get_token_lora_indices(y), y,
215 + output_slices, lora_bias_stacked)
216 +
217 + if buffer is None:
218 + r = lora_b_stacked[0].size(-1)
219 + # We set the buffer to be float32 by default, consistent with the
220 + # triton op
221 + T = x.size(0)
222 + buffer = torch.zeros(
223 + (len(output_slices), T, r),
224 + dtype=torch.float32,
225 + device=x.device,
226 + )
227 + buffer = self.add_shrink(buffer, x, lora_a_stacked, scale, **kwargs)
228 + return self.add_expand(y,
229 + buffer,
230 + lora_b_stacked,
231 + None,
232 + output_slices,
233 + add_inputs=True,
234 + **kwargs)
235 +
236 + def add_lora_logits(self,
237 + y: torch.Tensor,
238 + x: torch.Tensor,
239 + lora_a_stacked: torch.Tensor,
240 + lora_b_stacked: torch.Tensor,
241 + scale,
242 + *,
243 + buffer: Optional[torch.Tensor] = None,
244 + **kwargs) -> torch.Tensor:
245 + """
246 + Applies lora specifically for LogitsProcessorWithLoRA.
247 +
248 + Semantics:
249 + buffer = (x @ lora_a_stacked) * scale
250 + y += buffer @ lora_b_stacked
251 +
252 + Args:
253 + y (torch.Tensor): Output tensor.
254 + x (torch.Tensor): Input tensor.
255 + lora_a_stacked (torch.Tensor): lora_a's weights.
256 + lora_b_stacked (torch.Tensor):lora_b's weights.
257 + scale (float): Scaling factor.
258 + buffer (Optional[torch.Tensor]):Default to None.
259 + """
260 + if self.no_lora:
261 + return y
262 +
263 + y_org = y
264 + y = y.view(-1, y.shape[-1])
265 + x = x.view(-1, x.shape[-1])
266 + r = lora_b_stacked.size(-1)
267 + if buffer is None:
268 + # We set the buffer to be float32 by default, consistent with the
269 + # triton op
270 + buffer = torch.zeros((x.size(0), r),
271 + dtype=torch.float32,
272 + device=x.device)
273 +
274 + buffer = bgmv_shrink(x, lora_a_stacked, buffer, self.sampler_indices,
275 + scale)
276 + y = bgmv_expand(buffer,
277 + lora_b_stacked,
278 + y,
279 + self.sampler_indices,
280 + add_inputs=True)
281 + return y.view_as(y_org)
282 +
283 + def _apply_bias(
284 + self,
285 + indices: torch.Tensor,
286 + output: torch.Tensor,
287 + output_slices: Tuple[int, ...],
288 + lora_bias_stacked: Tuple[Optional[torch.Tensor], ...],
289 + ):
290 + """Applies bias to output
291 +
292 + Input shapes:
293 + lora_bias_stacked: 3 element tuple of (num_loras, output_dim)
294 + indices: (batch_size)
295 + output: (batch_size, q_slice_size + 2*kv_slice_size)
296 + output_slices: n-1 element tuple of (slice_size...),
297 + where n is number of slices
298 + """
299 + org_output = output
300 + output = output.view(-1, output.shape[-1])
301 + indices = indices.view(-1)
302 +
303 + offset_left = 0
304 + for slice_idx, slice in enumerate(output_slices):
305 + bias = lora_bias_stacked[slice_idx]
306 + if bias is not None:
307 + bias = bias.view(-1, bias.shape[-1])
308 + bias = bias[indices]
309 + bias = torch.where(indices[:, None] == -1, 0, bias)
310 +
311 + bias = F.pad(bias, (offset_left, output.shape[1] -
312 + (offset_left + slice), 0, 0))
313 +
314 + output += bias
315 + offset_left += slice
316 +
317 + return output.view_as(org_output)
318 +
319 + def _update_prefill_metadata(self,
320 + token_lora_tensor: torch.Tensor) -> None:
321 + self.batch_size = 1
322 + self._lora_indices_per_batch[:self.batch_size].copy_(
323 + token_lora_tensor[:self.batch_size])
324 + # TODO: .item() is extremely inefficient on TPU, so find a way around it
325 + self.no_lora = torch.all(token_lora_tensor == -1).item()
vllm/lora/punica_wrapper/utils.py
+4 -2
@@ -125,11 +125,13 @@ def convert_mapping(
125 indices[2] * extra_vocab_size,
126 indices[2] * (vocab_size + extra_vocab_size),
127 ])
128 - embeddings_indices[embeddings_indices == -1] = max_loras - 1
128 + embeddings_indices = torch.where(embeddings_indices == -1, max_loras - 1,
129 + embeddings_indices)
130 base_indices = indices[1]
131 sampler_indices = prompt_mapping_tensor
132 sampler_indices_padded = sampler_indices.clone()
132 - sampler_indices_padded[sampler_indices_padded == -1] = max_loras - 1
133 + sampler_indices_padded = torch.where(sampler_indices_padded == -1,
134 + max_loras - 1, sampler_indices_padded)
135 sampler_indices_padded = torch.arange(
136 0, len(sampler_indices_padded), device=device, dtype=torch.long) + (
137 sampler_indices_padded * len(sampler_indices_padded))
vllm/platforms/interface.py
+21
@@ -332,6 +332,27 @@ class Platform:
332 """
333 raise NotImplementedError
334
335 + @classmethod
336 + def get_infinity_values(cls, dtype: torch.dtype) -> Tuple[float, float]:
337 + """
338 + Return the platform specific values for (-inf, inf)
339 + """
340 + return float("-inf"), float("inf")
341 +
342 + @classmethod
343 + def can_update_inplace(cls) -> bool:
344 + """
345 + Checks if the platform allows inplace memory updates
346 + """
347 + return True
348 +
349 + @classmethod
350 + def get_lora_vocab_padding_size(cls) -> int:
351 + """
352 + Returns how much padding the LoRA logits need for kernels
353 + """
354 + return 256
355 +
356 @classmethod
357 def get_device_communicator_cls(cls) -> str:
358 """
vllm/platforms/tpu.py
+17 -1
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 -from typing import TYPE_CHECKING, Optional, Union
3 +from typing import TYPE_CHECKING, Optional, Tuple, Union
4
5 import torch
6 from tpu_info import device
@@ -67,6 +67,22 @@ class TpuPlatform(Platform):
67 def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
68 return not envs.VLLM_USE_V1
69
70 + @classmethod
71 + def get_punica_wrapper(cls) -> str:
72 + return "vllm.lora.punica_wrapper.punica_tpu.PunicaWrapperTPU"
73 +
74 + @classmethod
75 + def get_infinity_values(cls, dtype: torch.dtype) -> Tuple[float, float]:
76 + return torch.finfo(dtype).min, torch.finfo(dtype).max
77 +
78 + @classmethod
79 + def can_update_inplace(cls):
80 + return False
81 +
82 + @classmethod
83 + def get_lora_vocab_padding_size(cls) -> int:
84 + return 1
85 +
86 @classmethod
87 def inference_mode(cls):
88 return torch.no_grad()
vllm/v1/worker/tpu_model_runner.py
+22 -2
@@ -39,6 +39,7 @@ from vllm.v1.sample.tpu.metadata import TPUSupportedSamplingMetadata
39 from vllm.v1.sample.tpu.sampler import Sampler as TPUSampler
40 from vllm.v1.utils import bind_kv_cache
41 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
42 +from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
43
44 from .utils import sanity_check_mm_encoder_outputs
45
@@ -90,7 +91,7 @@ MIN_NUM_SEQS = 8
91 # The dummy_run should be comprehensive, ensuring all potential input shapes and
92 # branch predictions are included as subgraph inputs to facilitate
93 # pre-compilation.
93 -class TPUModelRunner:
94 +class TPUModelRunner(LoRAModelRunnerMixin):
95
96 def __init__(
97 self,
@@ -568,6 +569,17 @@ class TPUModelRunner:
569 self.device)
570 seq_lens = self.seq_lens_cpu[:self.max_num_reqs].to(self.device)
571
572 + if self.lora_config is not None:
573 + # We need to respect padding when activating LoRA adapters
574 + padded_num_scheduled_tokens_per_req = np.copy(
575 + num_scheduled_tokens_per_req
576 + ) # Copying to avoid accidental state corruption bugs
577 + padded_num_scheduled_tokens_per_req[-1] += \
578 + padded_total_num_scheduled_tokens - total_num_scheduled_tokens
579 +
580 + self.set_active_loras(self.input_batch,
581 + padded_num_scheduled_tokens_per_req)
582 +
583 attn_metadata = PallasMetadata(
584 slot_mapping=slot_mapping,
585 block_tables=block_tables,
@@ -907,6 +919,11 @@ class TPUModelRunner:
919 "get_tensor_model_parallel_rank",
920 return_value=xm_tp_rank):
921 model = get_model(vllm_config=self.vllm_config)
922 + if self.lora_config is not None:
923 + model = self.load_lora_model(model, self.model_config,
924 + self.scheduler_config,
925 + self.lora_config, self.device)
926 +
927 # Sync all pending XLA execution during model initialization and weight
928 # loading.
929 xm.mark_step()
@@ -970,7 +987,10 @@ class TPUModelRunner:
987 for layer_name in layer_names
988 }
989
973 - with set_forward_context(per_layer_attn_metadata, self.vllm_config, 0):
990 + with self.maybe_dummy_run_with_lora(
991 + self.lora_config,
992 + np.array([num_tokens], dtype=np.int32)), set_forward_context(
993 + per_layer_attn_metadata, self.vllm_config, 0):
994 out = self.model(input_ids=input_ids,
995 positions=position_ids,
996 inputs_embeds=inputs_embeds)
vllm/v1/worker/tpu_worker.py
+8
@@ -15,6 +15,7 @@ from vllm.config import ParallelConfig, VllmConfig
15 from vllm.distributed import (ensure_model_parallel_initialized,
16 init_distributed_environment)
17 from vllm.logger import init_logger
18 +from vllm.lora.request import LoRARequest
19 from vllm.model_executor import set_random_seed
20 from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE
21 from vllm.v1.core.sched.output import SchedulerOutput
@@ -82,6 +83,10 @@ class TPUWorker:
83 if self.model_config.seed is None:
84 self.model_config.seed = 0
85
86 + if vllm_config.lora_config is not None:
87 + raise NotImplementedError(
88 + "The V1 TPU backend doesn't support LoRA serving")
89 +
90 def init_device(self):
91 os.environ["PJRT_DEVICE"] = "TPU"
92 # Note: Currently the XLA compiler wrongly uses 2D ring strategy on 1D
@@ -211,6 +216,9 @@ class TPUWorker:
216 else:
217 xp.stop_trace()
218
219 + def add_lora(self, lora_request: LoRARequest) -> bool:
220 + return self.model_runner.add_lora(lora_request)
221 +
222 def load_model(self) -> None:
223 self.model_runner.load_model()
224
vllm/worker/tpu_worker.py
+4
@@ -54,6 +54,10 @@ class TPUWorker(LoRANotSupportedWorkerBase, LocalOrDistributedWorkerBase):
54 if self.model_config.seed is None:
55 self.model_config.seed = 0
56
57 + if vllm_config.lora_config is not None:
58 + raise NotImplementedError(
59 + "The V0 TPU backend doesn't support LoRA serving")
60 +
61 def init_device(self) -> None:
62 os.environ["PJRT_DEVICE"] = "TPU"
63 torch.set_grad_enabled(False)