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, 08:30 UTC
[Hardware][TPU][V1] Multi-LoRA Optimisations for the V1 TPU backend (#15655)
Signed-off-by: Akshat Tripathi <akshat@krai.ai> Signed-off-by: Chengji Yao <chengjiyao@google.com> Signed-off-by: xihajun <junfan@krai.ai> Signed-off-by: Jorge de Freitas <jorge.de-freitas22@imperial.ac.uk> Signed-off-by: Jorge de Freitas <jorge@krai.ai> Co-authored-by: Chengji Yao <chengjiyao@google.com> Co-authored-by: xihajun <junfan@krai.ai> Co-authored-by: Jorge de Freitas <jorge.de-freitas22@imperial.ac.uk> Co-authored-by: Jorge de Freitas <jorge@krai.ai>
Akshat Tripathi committed
May 28, 2025 at 20:59 UTC
643622ba465f0bbcb07e1d81c60dac35ed994ebf
9 files changed
+325
-334
.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh
+2
-4
@@ -122,10 +122,8 @@ run_and_track_test 11 "test_struct_output_generate.py" \
122
"python3 -m pytest -s -v /workspace/vllm/tests/v1/entrypoints/llm/test_struct_output_generate.py"
123
run_and_track_test 12 "test_moe_pallas.py" \
124
"python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_moe_pallas.py"
125
-
126
-# Disable the TPU LoRA tests until the feature is activated
127
-# run_and_track_test 13 "test_lora (directory)" \
128
-# "python3 -m pytest -s -v /workspace/vllm/tests/tpu/lora/"
125
+run_and_track_test 13 "test_lora.py" \
126
+ "VLLM_XLA_CHECK_RECOMPILATION=0 python3 -m pytest -s -v /workspace/vllm/tests/tpu/lora/test_lora.py"
127
128
# After all tests have been attempted, exit with the overall status.
129
if [ "$overall_script_exit_code" -ne 0 ]; then
tests/tpu/lora/test_pallas_kernels.py
deleted
-73
@@ -1,73 +0,0 @@
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/lora/models.py
+2
-2
@@ -200,7 +200,7 @@ class LoRAModel(AdapterModel):
200
weights_mapper: Optional[WeightsMapper] = None,
201
tensorizer_config_dict: Optional[dict] = None) -> "LoRAModel":
202
"""Create a LoRAModel from a local checkpoint.
203
-
203
+
204
Args:
205
lora_dir: The local path that has lora data.
206
expected_lora_modules: Name of modules that are expected to be
@@ -620,7 +620,7 @@ class LoRAModelManager(AdapterModelManager):
620
def _filter_unsupported_mm_module(self, module_name: str) -> bool:
621
"""
622
Regarding multimodal models, vLLM currently only supports adding LoRA to
623
- language model. LoRA for other modules, such as the vision tower, will
623
+ language model. LoRA for other modules, such as the vision tower, will
624
be filtered out.
625
"""
626
if self.supports_mm:
vllm/lora/ops/xla_ops/lora_ops.py
+89
-51
@@ -1,63 +1,99 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
+import jax
4
+import jax.numpy as jnp
5
import torch
6
+import torch.nn.functional as F
7
+import torch_xla.core.xla_builder as xb
8
+from torch.library import impl
9
+from torch_xla.experimental.custom_kernel import XLA_LIB, jax_import_guard
10
5
-# Required to register the custom ops
6
-import vllm.lora.ops.xla_ops.pallas # noqa # pylint: disable=unused-import
11
12
+@jax.jit
13
+def bgmv_jax(inputs, loras, idxs):
14
+ return jnp.einsum(
15
+ "td,tX,Xld->tl",
16
+ inputs,
17
+ jax.nn.one_hot(idxs, loras.shape[0], dtype=inputs.dtype),
18
+ loras,
19
+ )
20
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):
21
+
22
+XLA_LIB.define("bgmv(Tensor inputs, Tensor loras, Tensor idxs) -> Tensor")
23
+
24
+
25
+@impl(XLA_LIB, "bgmv", "XLA")
26
+def bgmv_xla(inputs: torch.Tensor, loras: torch.Tensor, idxs: torch.IntTensor):
27
+ if len(loras.shape) == 4:
28
+ loras = loras.squeeze(axis=1)
29
+
30
+ jax_import_guard()
31
+ return xb.call_jax(bgmv_jax, (inputs, loras, idxs))
32
+
33
+
34
+@impl(XLA_LIB, "bgmv", "CompositeExplicitAutograd")
35
+def bgmv_non_xla(inputs: torch.Tensor, loras: torch.Tensor,
36
+ idxs: torch.IntTensor):
37
+ T, _ = inputs.shape
38
+ if len(loras.shape) == 4:
39
+ loras = loras.squeeze(axis=1)
40
+ _, L, _ = loras.shape
41
+
42
+ return torch.empty((T, L), device=inputs.device)
43
+
44
+
45
+def bgmv_expand(
46
+ inputs: torch.Tensor,
47
+ lora_b_weights: torch.Tensor,
48
+ output_tensor: torch.Tensor,
49
+ lora_indices_tensor: torch.Tensor,
50
+ add_inputs: bool = True,
51
+):
52
"""
53
Args:
54
inputs (torch.Tensor): Input tensor of shape [num_tokens, hidden_size].
17
-
18
- lora_b_weights (torch.Tensor): LoRA weights of shape
55
+
56
+ lora_b_weights (torch.Tensor): LoRA weights of shape
57
[num_loras, lora_rank, hidden_size].
20
-
21
- output_tensor (torch.Tensor): output tensor of shape
58
+
59
+ output_tensor (torch.Tensor): output tensor of shape
60
[num_tokens, hidden_size * num_slices].
23
-
24
- lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
61
+
62
+ lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
63
indicating which LoRA matrix to use for each token.
26
- add_inputs (bool): Whether or not to add the input tensor to the output
64
+ add_inputs (bool): Whether or not to add the input tensor to the output
65
tensor.
66
"""
67
68
outputs = torch.ops.xla.bgmv(inputs, lora_b_weights, lora_indices_tensor)
31
- n_tokens = outputs.size(0)
69
70
limit = output_tensor.shape[0]
71
if outputs.shape[0] == 1 and output_tensor.shape[0] != 1:
72
limit = 1
73
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)
74
+ if output_tensor.shape[1] > outputs.shape[1]:
75
+ outputs = F.pad(outputs,
76
+ (0, output_tensor.shape[1] - outputs.shape[1], 0, 0))
77
78
if add_inputs:
44
- return output_tensor + outputs[:limit, :]
79
+ return output_tensor + outputs[:limit, :output_tensor.shape[1]]
80
else:
46
- return outputs[:limit, :]
81
+ return outputs[:limit, :output_tensor.shape[1]]
82
83
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):
84
+def bgmv_shrink(
85
+ inputs: torch.Tensor,
86
+ lora_b_weights: torch.Tensor,
87
+ lora_indices_tensor: torch.Tensor,
88
+ scaling: float = 1.0,
89
+):
90
"""
91
Args:
92
inputs (torch.Tensor): Input tensor of shape [num_tokens, hidden_size].
57
- lora_b_weights (torch.Tensor): LoRA weights of shape
93
+ lora_b_weights (torch.Tensor): LoRA weights of shape
94
[num_loras, lora_rank, hidden_size].
95
output_tensor (torch.Tensor): (Unused) output tensor (placeholder).
60
- lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
96
+ lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
97
indicating which LoRA matrix to use for each token.
98
scaling (float, optional): Scalar multiplier applied to the output.
99
"""
@@ -66,39 +102,41 @@ def bgmv_shrink(inputs: torch.Tensor,
102
lora_indices_tensor)
103
104
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):
105
+def bgmv_expand_slice(
106
+ inputs: torch.Tensor,
107
+ lora_b_weights: torch.Tensor,
108
+ output_tensor: torch.Tensor,
109
+ lora_indices_tensor: torch.Tensor,
110
+ slice_offset: int,
111
+ slice_size: int,
112
+ add_inputs: bool = True,
113
+):
114
"""
115
Args:
116
inputs (torch.Tensor): Input tensor of shape [num_tokens, hidden_size].
79
-
80
- lora_b_weights (torch.Tensor): LoRA weights of shape
117
+
118
+ lora_b_weights (torch.Tensor): LoRA weights of shape
119
[num_loras, lora_rank, hidden_size].
82
-
83
- output_tensor (torch.Tensor): output tensor of shape
120
+
121
+ output_tensor (torch.Tensor): output tensor of shape
122
[num_tokens, hidden_size * num_slices].
85
-
86
- lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
123
+
124
+ lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
125
indicating which LoRA matrix to use for each token.
88
- add_inputs (bool): Whether or not to add the input tensor to the output
126
+ add_inputs (bool): Whether or not to add the input tensor to the output
127
tensor.
128
"""
129
outputs = torch.ops.xla.bgmv(inputs, lora_b_weights, lora_indices_tensor)
92
- n_tokens = outputs.size(0)
130
94
- outputs = torch.cat((
95
- torch.zeros((n_tokens, slice_offset), device=outputs.device),
131
+ outputs = F.pad(
132
outputs,
97
- torch.zeros(
98
- (n_tokens, output_tensor.shape[1] - (slice_offset + slice_size)),
99
- device=outputs.device),
100
- ),
101
- dim=1)
133
+ (
134
+ slice_offset,
135
+ output_tensor.shape[1] - (slice_offset + slice_size),
136
+ 0,
137
+ 0,
138
+ ),
139
+ )
140
141
if add_inputs:
142
return output_tensor + outputs
vllm/lora/ops/xla_ops/pallas.py
deleted
-133
@@ -1,133 +0,0 @@
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_tpu.py
+118
-39
@@ -1,11 +1,19 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-from typing import Optional, Union
3
+import math
4
+from typing import TYPE_CHECKING, Optional, Union
5
6
import torch
7
import torch.nn.functional as F
8
+import torch_xla.core.xla_model as xm
9
10
from vllm.lora.ops.xla_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
11
+from vllm.lora.punica_wrapper.utils import convert_mapping
12
+
13
+if TYPE_CHECKING:
14
+ # avoid circuit import
15
+ from vllm.lora.layers import LoRAMapping
16
+ from vllm.lora.models import LongContextLoRAContext
17
18
from .punica_base import PunicaWrapperBase
19
@@ -31,6 +39,15 @@ class PunicaWrapperTPU(PunicaWrapperBase):
39
self._sampler_indices_padded = self._sampler_indices_padded.to(
40
dtype=torch.int32)
41
42
+ torch.ops.xla.dynamo_set_buffer_donor_(self._token_lora_indices, True)
43
+ torch.ops.xla.dynamo_set_buffer_donor_(self._sampler_indices, True)
44
+ torch.ops.xla.dynamo_set_buffer_donor_(self._sampler_indices_padded,
45
+ True)
46
+ torch.ops.xla.dynamo_set_buffer_donor_(self._embeddings_indices, True)
47
+ torch.ops.xla.dynamo_set_buffer_donor_(self._long_lora_indices, True)
48
+ torch.ops.xla.dynamo_set_buffer_donor_(self._lora_indices_per_batch,
49
+ True)
50
+
51
torch._dynamo.mark_dynamic(self._token_lora_indices, 0)
52
torch._dynamo.mark_dynamic(self._embeddings_indices, 1)
53
torch._dynamo.mark_dynamic(self._sampler_indices_padded, 0)
@@ -55,15 +72,11 @@ class PunicaWrapperTPU(PunicaWrapperBase):
72
73
def shrink(
74
self,
58
- y: torch.Tensor,
75
x: torch.Tensor,
76
w_t_all: torch.Tensor,
77
scale: float,
78
):
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)
79
+ return bgmv_shrink(x, w_t_all, self._get_token_lora_indices(x), scale)
80
81
def expand(self, y: torch.Tensor, x: torch.Tensor, w_t_all: torch.Tensor,
82
add_inputs: bool):
@@ -72,7 +85,7 @@ class PunicaWrapperTPU(PunicaWrapperBase):
85
86
def expand_slice(self, y: torch.Tensor, x: torch.Tensor,
87
w_t_all: torch.Tensor, y_offset: int, y_slice_size: int,
75
- y_total_size: int, add_inputs: bool) -> torch.Tensor:
88
+ add_inputs: bool) -> torch.Tensor:
89
return bgmv_expand_slice(x, w_t_all, y,
90
self._get_token_lora_indices(x), y_offset,
91
y_slice_size, add_inputs)
@@ -98,9 +111,8 @@ class PunicaWrapperTPU(PunicaWrapperBase):
111
x = x.view(-1, x.shape[-1])
112
113
for slice_idx in range(len(lora_a_stacked)):
101
- y_s = y[slice_idx]
114
lora_s = lora_a_stacked[slice_idx]
103
- y_s = self.shrink(y_s, x, lora_s, scale)
115
+ y_s = self.shrink(x, lora_s, scale)
116
y[slice_idx, :, :] = y_s # type: ignore[index]
117
return y
118
@@ -140,15 +152,12 @@ class PunicaWrapperTPU(PunicaWrapperBase):
152
y = self._apply_bias(self._get_token_lora_indices(y), y,
153
output_slices, lora_bias_stacked)
154
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
- )
155
+ y = self.expand_slice(y,
156
+ x[slice_idx],
157
+ lora_b_stacked[slice_idx],
158
+ offset_left,
159
+ output_slices[slice_idx],
160
+ add_inputs=add_inputs)
161
offset_left += output_slices[slice_idx]
162
return y.view_as(y_org)
163
@@ -216,12 +225,10 @@ class PunicaWrapperTPU(PunicaWrapperBase):
225
226
if buffer is None:
227
r = lora_b_stacked[0].size(-1)
219
- # We set the buffer to be float32 by default, consistent with the
220
- # triton op
228
T = x.size(0)
229
buffer = torch.zeros(
230
(len(output_slices), T, r),
224
- dtype=torch.float32,
231
+ dtype=x.dtype,
232
device=x.device,
233
)
234
buffer = self.add_shrink(buffer, x, lora_a_stacked, scale, **kwargs)
@@ -257,26 +264,16 @@ class PunicaWrapperTPU(PunicaWrapperBase):
264
scale (float): Scaling factor.
265
buffer (Optional[torch.Tensor]):Default to None.
266
"""
260
- if self.no_lora:
261
- return y
262
-
267
y_org = y
268
y = y.view(-1, y.shape[-1])
269
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)
270
+
271
+ sampler_indices = torch.narrow(self._sampler_indices, 0, 0, x.size(0))
272
+ buffer = bgmv_shrink(x, lora_a_stacked, sampler_indices, scale)
273
y = bgmv_expand(buffer,
274
lora_b_stacked,
275
y,
279
- self.sampler_indices,
276
+ sampler_indices,
277
add_inputs=True)
278
return y.view_as(y_org)
279
@@ -316,10 +313,92 @@ class PunicaWrapperTPU(PunicaWrapperBase):
313
314
return output.view_as(org_output)
315
316
+ # This performs the same tensor ops as the base method, except it does them
317
+ # on the CPU then transfers the results to the TPU
318
+ def _update_base_metadata(
319
+ self,
320
+ mapping: "LoRAMapping",
321
+ lora_index_to_id: list[Optional[int]],
322
+ max_loras: int,
323
+ vocab_size: int,
324
+ extra_vocab_size: int,
325
+ long_lora_context: Optional["LongContextLoRAContext"] = None,
326
+ ):
327
+ # Make sure we don't accidentally collect outside operations
328
+ xm.mark_step()
329
+
330
+ # Pad the prompt mapping to avoid running into recompiles on the TPU
331
+ # TODO: Should this happen inside mapping internally? If so how can we
332
+ # avoid having backend specific LoRAMapping classes?
333
+ mapping.prompt_mapping = self._pad_prompt_mapping(
334
+ mapping.prompt_mapping)
335
+
336
+ (
337
+ base_indices,
338
+ sampler_indices,
339
+ sampler_indices_padded,
340
+ embeddings_indices,
341
+ long_lora_offsets_tensor,
342
+ indices_len,
343
+ ) = convert_mapping(
344
+ mapping,
345
+ lora_index_to_id,
346
+ max_loras,
347
+ vocab_size,
348
+ extra_vocab_size,
349
+ "cpu",
350
+ long_lora_context,
351
+ )
352
+ self._token_lora_indices = self._pad_to_shape(
353
+ base_indices, self._token_lora_indices.shape,
354
+ dims=1).to(self.device)
355
+ self._sampler_indices = self._pad_to_shape(sampler_indices,
356
+ self._sampler_indices.shape,
357
+ dims=1).to(self.device)
358
+ self._sampler_indices_padded = self._pad_to_shape(
359
+ sampler_indices_padded, self._sampler_indices_padded.shape,
360
+ dims=1).to(self.device)
361
+ self._embeddings_indices = self._pad_to_shape(
362
+ embeddings_indices, self._embeddings_indices.shape,
363
+ dims=2).to(self.device)
364
+ if long_lora_offsets_tensor is not None:
365
+ self._long_lora_indices = self._pad_to_shape(
366
+ long_lora_offsets_tensor,
367
+ self._long_lora_indices.shape,
368
+ dims=1).to(self.device)
369
+ else:
370
+ zeroed = torch.zeros_like(self._long_lora_indices.cpu(),
371
+ dtype=torch.int32)
372
+ self._long_lora_indices = zeroed.to(self.device)
373
+ self.indices_len[:] = indices_len
374
+
375
def _update_prefill_metadata(self,
376
token_lora_tensor: torch.Tensor) -> None:
377
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()
378
+ self._lora_indices_per_batch[:self.
379
+ batch_size] = token_lora_tensor[:self.
380
+ batch_size]
381
+
382
+ def _pad_prompt_mapping(
383
+ self, prompt_mapping: tuple[int, ...]) -> tuple[int, ...]:
384
+ num_reqs = len(prompt_mapping)
385
+
386
+ # From vllm/v1/worker/tpu_model_runner:51, but need to avoid a circular
387
+ # import
388
+ MIN_NUM_SEQS = 8
389
+
390
+ padded_num_reqs = max(2**math.ceil(math.log2(num_reqs)), MIN_NUM_SEQS)
391
+ pad_len = padded_num_reqs - num_reqs
392
+
393
+ padding = [-1] * pad_len
394
+ return tuple(list(prompt_mapping) + padding)
395
+
396
+ def _pad_to_shape(self, src, target_shape, dims=1):
397
+ if dims == 1:
398
+ pad_len = target_shape[0] - src.shape[0]
399
+ return F.pad(src, (0, pad_len), value=0).to(torch.int32)
400
+ else:
401
+ pad_rows = target_shape[0] - src.shape[0]
402
+ pad_cols = target_shape[1] - src.shape[1]
403
+ return F.pad(src, (0, pad_cols, 0, pad_rows),
404
+ value=0).to(torch.int32)
vllm/v1/worker/lora_model_runner_mixin.py
+42
-15
@@ -80,8 +80,38 @@ class LoRAModelRunnerMixin:
80
lora_requests)
81
82
@contextmanager
83
- def maybe_dummy_run_with_lora(self, lora_config: LoRAConfig,
84
- num_scheduled_tokens: np.ndarray):
83
+ def maybe_setup_dummy_loras(self, lora_config):
84
+ if lora_config is None:
85
+ yield
86
+ else:
87
+ # __enter__ code
88
+ assert self.lora_manager is not None, "LoRA is not enabled"
89
+
90
+ num_loras = lora_config.max_loras
91
+
92
+ # Make dummy lora requests
93
+ lora_requests: set[LoRARequest] = {
94
+ LoRARequest(lora_name=f"warmup_{lora_id}",
95
+ lora_int_id=lora_id,
96
+ lora_path="/not/a/real/path")
97
+ for lora_id in range(1, num_loras + 1)
98
+ }
99
+
100
+ with self.lora_manager.dummy_lora_cache():
101
+ # Add the dummy LoRAs here so _set_active_loras doesn't try to
102
+ # load from disk.
103
+ for lr in lora_requests:
104
+ self.lora_manager.add_dummy_lora(
105
+ lr, rank=self.LORA_WARMUP_RANK)
106
+
107
+ yield
108
+
109
+ # __exit__ code
110
+ self.lora_manager.remove_all_adapters()
111
+
112
+ @contextmanager
113
+ def maybe_select_dummy_loras(self, lora_config: LoRAConfig,
114
+ num_scheduled_tokens: np.ndarray):
115
if lora_config is None:
116
yield
117
else:
@@ -108,21 +138,18 @@ class LoRAModelRunnerMixin:
138
for lora_id in range(1, num_loras + 1)
139
}
140
111
- with self.lora_manager.dummy_lora_cache():
112
- # Add the dummy LoRAs here so _set_active_loras doesn't try to
113
- # load from disk.
114
- for lr in lora_requests:
115
- self.lora_manager.add_dummy_lora(
116
- lr, rank=self.LORA_WARMUP_RANK)
117
-
118
- self._set_active_loras(tuple(prompt_lora_mapping),
119
- tuple(token_lora_mapping),
120
- lora_requests)
141
+ self._set_active_loras(tuple(prompt_lora_mapping),
142
+ tuple(token_lora_mapping), lora_requests)
143
122
- yield
144
+ yield
145
124
- # __exit__ code
125
- self.lora_manager.remove_all_adapters()
146
+ @contextmanager
147
+ def maybe_dummy_run_with_lora(self, lora_config: LoRAConfig,
148
+ num_scheduled_tokens: np.ndarray):
149
+ with self.maybe_setup_dummy_loras(
150
+ lora_config), self.maybe_select_dummy_loras(
151
+ lora_config, num_scheduled_tokens):
152
+ yield
153
154
def add_lora(self, lora_request: LoRARequest) -> bool:
155
if not self.lora_manager:
vllm/v1/worker/tpu_model_runner.py
+70
-12
@@ -20,6 +20,7 @@ from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
20
from vllm.config import VllmConfig, get_layers_from_vllm_config
21
from vllm.forward_context import set_forward_context
22
from vllm.logger import init_logger
23
+from vllm.lora.layers import BaseLayerWithLoRA
24
from vllm.model_executor.model_loader import get_model
25
from vllm.multimodal import MULTIMODAL_REGISTRY
26
from vllm.multimodal.inputs import (BatchedTensorInputs, MultiModalKwargs,
@@ -152,6 +153,9 @@ class TPUModelRunner(LoRAModelRunnerMixin):
153
self.hidden_size = model_config.get_hidden_size()
154
self.vocab_size = model_config.get_vocab_size()
155
156
+ if self.lora_config is not None:
157
+ self.vocab_size += self.lora_config.lora_extra_vocab_size
158
+
159
# Multi-modal data support
160
self.mm_registry = MULTIMODAL_REGISTRY
161
self.uses_mrope = model_config.uses_mrope
@@ -591,6 +595,17 @@ class TPUModelRunner(LoRAModelRunnerMixin):
595
logits_indices = self.query_start_loc_cpu[1:padded_num_reqs + 1] - 1
596
logits_indices = logits_indices.to(self.device)
597
598
+ if self.lora_config is not None:
599
+ # We need to respect padding when activating LoRA adapters
600
+ padded_num_scheduled_tokens_per_req = np.copy(
601
+ num_scheduled_tokens_per_req
602
+ ) # Copying to avoid accidental state corruption bugs
603
+ padded_num_scheduled_tokens_per_req[-1] += \
604
+ padded_total_num_scheduled_tokens - total_num_scheduled_tokens
605
+
606
+ self.set_active_loras(self.input_batch,
607
+ padded_num_scheduled_tokens_per_req)
608
+
609
layer_names = get_layers_from_vllm_config(self.vllm_config,
610
Attention).keys()
611
per_layer_attn_metadata = {
@@ -916,6 +931,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
931
model = self.load_lora_model(model, self.model_config,
932
self.scheduler_config,
933
self.lora_config, self.device)
934
+ replace_set_lora(model)
935
936
# Sync all pending XLA execution during model initialization and weight
937
# loading.
@@ -980,7 +996,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
996
for layer_name in layer_names
997
}
998
983
- with self.maybe_dummy_run_with_lora(
999
+ with self.maybe_select_dummy_loras(
1000
self.lora_config,
1001
np.array([num_tokens], dtype=np.int32)), set_forward_context(
1002
per_layer_attn_metadata, self.vllm_config, 0):
@@ -989,6 +1005,13 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1005
inputs_embeds=inputs_embeds)
1006
self._hidden_states_dtype = out.dtype
1007
1008
+ def _set_active_loras(self, prompt_lora_mapping, token_lora_mapping,
1009
+ lora_requests) -> None:
1010
+ xm.mark_step() # Captures input updates
1011
+ super()._set_active_loras(prompt_lora_mapping, token_lora_mapping,
1012
+ lora_requests)
1013
+ xm.mark_step() # Captures metadata updates
1014
+
1015
def _precompile_mm_encoder(self) -> None:
1016
# Pre-compile MM encoder for all supported data modalities.
1017
hf_config = self.vllm_config.model_config.hf_config
@@ -1151,7 +1174,10 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1174
generate_params_if_all_greedy,
1175
))
1176
sampling_metadata.all_greedy = all_greedy
1154
- self.sample_from_logits(dummy_logits, sampling_metadata)
1177
+ with self.maybe_select_dummy_loras(
1178
+ self.lora_config, np.array([num_reqs],
1179
+ dtype=np.int32)):
1180
+ self.sample_from_logits(dummy_logits, sampling_metadata)
1181
logger.info(" -- num_seqs: %d", num_reqs)
1182
xm.wait_device_ops()
1183
end = time.perf_counter()
@@ -1167,7 +1193,9 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1193
dtype=self._hidden_states_dtype)
1194
dummy_tokens = torch.zeros((num_reqs, 1),
1195
dtype=torch.int64).to(self.device)
1170
- self.gather_logprobs(dummy_logits, dummy_tokens)
1196
+ with self.maybe_select_dummy_loras(
1197
+ self.lora_config, np.array([num_reqs], dtype=np.int32)):
1198
+ self.gather_logprobs(dummy_logits, dummy_tokens)
1199
logger.info(" -- num_seqs: %d", num_reqs)
1200
xm.wait_device_ops()
1201
end = time.perf_counter()
@@ -1178,13 +1206,14 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1206
"""
1207
Precompile all the subgraphs with possible input shapes.
1208
"""
1181
- self._precompile_mm_encoder()
1182
- self._precompile_backbone()
1183
- self._precompile_select_hidden_states()
1184
- self._precompile_compute_logits()
1185
- self._precompile_structured_decoding()
1186
- self._precompile_sample_from_logits()
1187
- self._precompile_gather_logprobs()
1209
+ with self.maybe_setup_dummy_loras(self.lora_config):
1210
+ self._precompile_mm_encoder()
1211
+ self._precompile_backbone()
1212
+ self._precompile_select_hidden_states()
1213
+ self._precompile_compute_logits()
1214
+ self._precompile_structured_decoding()
1215
+ self._precompile_sample_from_logits()
1216
+ self._precompile_gather_logprobs()
1217
1218
def profile_run(
1219
self,
@@ -1467,11 +1496,11 @@ def _get_token_paddings(min_token_size: int, max_token_size: int,
1496
padding_gap: int) -> list[int]:
1497
"""Generate a list of padding size, starting from min_token_size,
1498
ending with a number that can cover max_token_size
1470
-
1499
+
1500
If padding_gap == 0 then:
1501
increase 2X each time (exponential)
1502
else:
1474
- first increase the size to twice,
1503
+ first increase the size to twice,
1504
then increase the padding size by padding_gap.
1505
"""
1506
# assert min_token_size is power of 2
@@ -1508,3 +1537,32 @@ def _get_padded_token_len(paddings: list[int], x: int) -> int:
1537
index = bisect.bisect_left(paddings, x)
1538
assert index < len(paddings)
1539
return paddings[index]
1540
+
1541
+
1542
+def replace_set_lora(model):
1543
+
1544
+ def _tpu_set_lora(
1545
+ self,
1546
+ index: int,
1547
+ lora_a: torch.Tensor,
1548
+ lora_b: torch.Tensor,
1549
+ embeddings_tensor: Optional[torch.Tensor],
1550
+ bias: Optional[torch.Tensor] = None,
1551
+ ):
1552
+ # TODO: The integer index leads to a recompilation, but converting it
1553
+ # to a tensor doesn't seem to work anymore. This might be fixed with a
1554
+ # later release of torch_xla.
1555
+ self._original_set_lora(index, lora_a, lora_b, embeddings_tensor, bias)
1556
+ xm.mark_step()
1557
+
1558
+ def _tpu_reset_lora(self, index: int):
1559
+ self._original_reset_lora(index)
1560
+ xm.mark_step()
1561
+
1562
+ for _, module in model.named_modules():
1563
+ if isinstance(module, BaseLayerWithLoRA):
1564
+ module._original_set_lora = module.set_lora
1565
+ module._original_reset_lora = module.reset_lora
1566
+ module.set_lora = _tpu_set_lora.__get__(module, module.__class__)
1567
+ module.reset_lora = _tpu_reset_lora.__get__(
1568
+ module, module.__class__)
vllm/v1/worker/tpu_worker.py
+2
-5
@@ -83,10 +83,6 @@ 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
-
86
def init_device(self):
87
os.environ["PJRT_DEVICE"] = "TPU"
88
# Note: Currently the XLA compiler wrongly uses 2D ring strategy on 1D
@@ -166,7 +162,8 @@ class TPUWorker:
162
runner_kv_caches)
163
164
# `max_num_tokens >= max_num_batched_tokens` due to padding.
169
- self.model_runner.profile_run(self.model_runner.max_num_tokens)
165
+ with self.model_runner.maybe_setup_dummy_loras(self.lora_config):
166
+ self.model_runner.profile_run(self.model_runner.max_num_tokens)
167
168
# Synchronize before measuring the memory usage.
169
xm.wait_device_ops()