@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9b5b39b650

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, 09:45 UTC

Update deprecated type hinting in `vllm/lora` (#18128)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed May 14, 2025 at 11:57 UTC 9b5b39b650ef37d9086985eabfb9ed2f1c327075

19 files changed +245 -251

pyproject.toml

-1

@@ -78,7 +78,6 @@ exclude = [

78 "vllm/distributed/**/*.py" = ["UP006", "UP035"]

79 "vllm/engine/**/*.py" = ["UP006", "UP035"]

80 "vllm/executor/**/*.py" = ["UP006", "UP035"]

81 -"vllm/lora/**/*.py" = ["UP006", "UP035"]

81 "vllm/model_executor/model_loader/**/*.py" = ["UP006", "UP035"]

82 "vllm/model_executor/models/**/*.py" = ["UP006", "UP035"]

83 "vllm/platforms/**/*.py" = ["UP006", "UP035"]

vllm/lora/fully_sharded_layers.py

+11 -11

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

3 # pylint: disable=unused-argument

4 -from typing import TYPE_CHECKING, List, Optional, Tuple, Union, cast

4 +from typing import TYPE_CHECKING, Optional, Union, cast

6 import torch

7 import torch.nn as nn

@@ -118,7 +118,7 @@ class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA):

118 cls,

119 source_layer: nn.Module,

120 lora_config: LoRAConfig,

121 - packed_modules_list: List,

121 + packed_modules_list: list,

122 model_config: Optional[PretrainedConfig],

123 ) -> bool:

124 # specifying kwargs so they can be easily accessed in decorator

@@ -141,8 +141,8 @@ class MergedColumnParallelLinearWithShardedLoRA(

141 """

142

143 def slice_lora_a(

144 - self, lora_a: List[Union[torch.Tensor, None]]

145 - ) -> List[Union[torch.Tensor, None]]:

144 + self, lora_a: list[Union[torch.Tensor, None]]

145 + ) -> list[Union[torch.Tensor, None]]:

146 #NOTE: lora_a contains 2 subloras, and each sublora could be None.

147 output_shard_size = self.lora_a_stacked[0].shape[2]

148 output_start_idx = self.tp_rank * output_shard_size

@@ -165,7 +165,7 @@ class MergedColumnParallelLinearWithShardedLoRA(

165 cls,

166 source_layer: nn.Module,

167 lora_config: LoRAConfig,

168 - packed_modules_list: List,

168 + packed_modules_list: list,

169 model_config: Optional[PretrainedConfig],

170 ) -> bool:

171 # specifying kwargs so they can be easily accessed in decorator

@@ -201,7 +201,7 @@ class QKVParallelLinearWithShardedLoRA(QKVParallelLinearWithLoRA):

201 @classmethod

202 @_fully_sharded_can_replace

203 def can_replace_layer(cls, source_layer: nn.Module,

204 - lora_config: LoRAConfig, packed_modules_list: List,

204 + lora_config: LoRAConfig, packed_modules_list: list,

205 model_config: Optional[PretrainedConfig]) -> bool:

206 # specifying kwargs so they can be easily accessed in decorator

207 return super().can_replace_layer(

@@ -222,8 +222,8 @@ class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA):

222 """

223

224 def slice_lora_a(

225 - self, lora_a: List[Union[torch.Tensor, None]]

226 - ) -> List[Union[torch.Tensor, None]]:

225 + self, lora_a: list[Union[torch.Tensor, None]]

226 + ) -> list[Union[torch.Tensor, None]]:

227 # NOTE: lora_a contains 3 subloras, and each sublora could be None.

228 shard_size = [self.lora_a_stacked[i].shape[2] for i in range(3)]

229 start_idx = [self.tp_rank * shard_size[i] for i in range(3)]

@@ -248,7 +248,7 @@ class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA):

248 cls,

249 source_layer: nn.Module,

250 lora_config: LoRAConfig,

251 - packed_modules_list: List,

251 + packed_modules_list: list,

252 model_config: Optional[PretrainedConfig],

253 ) -> bool:

254 # specifying kwargs so they can be easily accessed in decorator

@@ -281,7 +281,7 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):

281 def slice_bias(self, bias: torch.Tensor) -> torch.Tensor:

282 if bias is None:

283 return bias

284 - self.lora_bias_stacked = cast(Tuple[torch.Tensor, ...],

284 + self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],

285 self.lora_bias_stacked)

286 shard_size = self.lora_bias_stacked[0].shape[2]

287 start_idx = self.tp_rank * shard_size

@@ -341,7 +341,7 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):

341 cls,

342 source_layer: nn.Module,

343 lora_config: LoRAConfig,

344 - packed_modules_list: List,

344 + packed_modules_list: list,

345 model_config: Optional[PretrainedConfig],

346 ) -> bool:

347 # specifying kwargs so they can be easily accessed in decorator

vllm/lora/layers.py

+31 -31

@@ -3,7 +3,7 @@

3 # pylint: disable=unused-argument

4 import math

5 from dataclasses import dataclass

6 -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union, cast

6 +from typing import TYPE_CHECKING, Optional, Union, cast

8 import torch

9 import torch.nn as nn

@@ -82,14 +82,14 @@ class LoRAMapping(AdapterMapping):

82 class BaseLayerWithLoRA(nn.Module):

84 def slice_lora_a(

85 - self, lora_a: Union[torch.Tensor, List[Union[torch.Tensor, None]]]

86 - ) -> Union[torch.Tensor, List[Union[torch.Tensor, None]]]:

85 + self, lora_a: Union[torch.Tensor, list[Union[torch.Tensor, None]]]

86 + ) -> Union[torch.Tensor, list[Union[torch.Tensor, None]]]:

87 """Slice lora a if splitting for tensor parallelism."""

88 ...

90 def slice_lora_b(

91 - self, lora_b: Union[torch.Tensor, List[Union[torch.Tensor, None]]]

92 - ) -> Union[torch.Tensor, List[Union[torch.Tensor, None]]]:

91 + self, lora_b: Union[torch.Tensor, list[Union[torch.Tensor, None]]]

92 + ) -> Union[torch.Tensor, list[Union[torch.Tensor, None]]]:

93 """Slice lora b if splitting with tensor parallelism."""

94 ...

@@ -128,7 +128,7 @@ class BaseLayerWithLoRA(nn.Module):

128 cls,

129 source_layer: nn.Module,

130 lora_config: LoRAConfig,

131 - packed_modules_list: List,

131 + packed_modules_list: list,

132 model_config: Optional[PretrainedConfig],

133 ) -> bool:

134 """Returns True if the layer can be replaced by this LoRA layer."""

@@ -140,7 +140,7 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):

140 def __init__(self, base_layer: VocabParallelEmbedding) -> None:

141 super().__init__()

142 self.base_layer = base_layer

143 - self.embeddings_slice: Optional[Tuple[int, int]]

143 + self.embeddings_slice: Optional[tuple[int, int]]

144 self.embeddings_weights: Optional[torch.Tensor]

145

146 def create_lora_weights(

@@ -279,7 +279,7 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):

279 cls,

280 source_layer: nn.Module,

281 lora_config: LoRAConfig,

282 - packed_modules_list: List,

282 + packed_modules_list: list,

283 model_config: Optional[PretrainedConfig],

284 ) -> bool:

285 return type(source_layer) is VocabParallelEmbedding

@@ -296,9 +296,9 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

296 self.base_layer = base_layer

297 self.input_size = self.base_layer.input_size

298 self.device = _get_lora_device(self.base_layer)

299 - self.lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]] = None

299 + self.lora_bias_stacked: Optional[tuple[torch.Tensor, ...]] = None

300

301 - self.output_slices: Tuple[int, ...]

301 + self.output_slices: tuple[int, ...]

302 self.tp_size: int

303 self.output_size: int

304 self.n_slices: int

@@ -365,7 +365,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

365 self.lora_b_stacked[s_index][index] = 0

366 if self.lora_config.bias_enabled:

367 # Make mypy happy

368 - self.lora_bias_stacked = cast(Tuple[torch.Tensor, ...],

368 + self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],

369 self.lora_bias_stacked)

370 self.lora_bias_stacked[s_index][index] = 0

371

@@ -399,7 +399,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

399 lora_b.T, non_blocking=True)

400 if lora_bias is not None:

401

402 - self.lora_bias_stacked = cast(Tuple[torch.Tensor, ...],

402 + self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],

403 self.lora_bias_stacked)

404 assert len(self.lora_bias_stacked)

405 self.lora_bias_stacked[0][index, 0, :lora_bias.shape[0]].copy_(

@@ -497,7 +497,7 @@ class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):

497 cls,

498 source_layer: nn.Module,

499 lora_config: LoRAConfig,

500 - packed_modules_list: List,

500 + packed_modules_list: list,

501 model_config: Optional[PretrainedConfig],

502 ) -> bool:

503 return type(source_layer) is ReplicatedLinear

@@ -597,7 +597,7 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

597 cls,

598 source_layer: nn.Module,

599 lora_config: LoRAConfig,

600 - packed_modules_list: List,

600 + packed_modules_list: list,

601 model_config: Optional[PretrainedConfig],

602 ) -> bool:

603 return type(source_layer) is ColumnParallelLinear or (

@@ -674,13 +674,13 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

674 ) for output_size in self.output_slices)

675

676 def slice_lora_a(

677 - self, lora_a: List[Union[torch.Tensor, None]]

678 - ) -> List[Union[torch.Tensor, None]]:

677 + self, lora_a: list[Union[torch.Tensor, None]]

678 + ) -> list[Union[torch.Tensor, None]]:

679 return lora_a

680

681 def slice_lora_b(

682 - self, lora_b: List[Union[torch.Tensor, None]]

683 - ) -> List[Union[torch.Tensor, None]]:

682 + self, lora_b: list[Union[torch.Tensor, None]]

683 + ) -> list[Union[torch.Tensor, None]]:

684 for i, (shard_id, shard_size) in enumerate(

685 zip(self.output_ids, self.output_slices)):

686 if (lora_b_i := lora_b[i]) is not None:

@@ -689,8 +689,8 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

689 return lora_b

690

691 def slice_bias(

692 - self, bias: List[Union[torch.Tensor,

693 - None]]) -> List[Union[torch.Tensor, None]]:

692 + self, bias: list[Union[torch.Tensor,

693 + None]]) -> list[Union[torch.Tensor, None]]:

694 for i, (shard_id, shard_size) in enumerate(

695 zip(self.output_ids, self.output_slices)):

696 if (bias_i := bias[i]) is not None:

@@ -725,7 +725,7 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

725 lora_b_i.T, non_blocking=True)

726

727 if lora_bias is not None:

728 - self.lora_bias_stacked = cast(Tuple[torch.Tensor, ...],

728 + self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],

729 self.lora_bias_stacked)

730 for i in range(self.n_slices):

731 if (lora_bias_i := lora_bias[i]) is not None:

@@ -740,7 +740,7 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

740 cls,

741 source_layer: nn.Module,

742 lora_config: LoRAConfig,

743 - packed_modules_list: List,

743 + packed_modules_list: list,

744 model_config: Optional[PretrainedConfig],

745 ) -> bool:

746 return (type(source_layer) is MergedColumnParallelLinear

@@ -809,7 +809,7 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

809 @classmethod

810 @_not_fully_sharded_can_replace

811 def can_replace_layer(cls, source_layer: nn.Module,

812 - lora_config: LoRAConfig, packed_modules_list: List,

812 + lora_config: LoRAConfig, packed_modules_list: list,

813 model_config: Optional[PretrainedConfig]) -> bool:

814 return type(source_layer) is QKVParallelLinear and len(

815 packed_modules_list) == 1

@@ -869,7 +869,7 @@ class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):

869 cls,

870 source_layer: nn.Module,

871 lora_config: LoRAConfig,

872 - packed_modules_list: List,

872 + packed_modules_list: list,

873 model_config: Optional[PretrainedConfig],

874 ) -> bool:

875 return (type(source_layer) is QKVParallelLinear

@@ -923,7 +923,7 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

923 - output

924 - bias

925 """

926 - # Set up backprop all-reduce.

926 + # set up backprop all-reduce.

927 if self.base_layer.input_is_parallel:

928 input_parallel = input_

929 else:

@@ -958,7 +958,7 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

958 cls,

959 source_layer: nn.Module,

960 lora_config: LoRAConfig,

961 - packed_modules_list: List,

961 + packed_modules_list: list,

962 model_config: Optional[PretrainedConfig],

963 ) -> bool:

964 return type(source_layer) is RowParallelLinear

@@ -981,7 +981,7 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):

981

982 def __init__(self, base_layer: LogitsProcessor, hidden_size: int,

983 dtype: torch.dtype, device: torch.device,

984 - sharded_to_full_mapping: Optional[List[int]]) -> None:

984 + sharded_to_full_mapping: Optional[list[int]]) -> None:

985 super().__init__()

986 self.base_layer = base_layer

987 self.hidden_size = hidden_size

@@ -1189,7 +1189,7 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):

1189 cls,

1190 source_layer: nn.Module,

1191 lora_config: LoRAConfig,

1192 - packed_modules_list: List,

1192 + packed_modules_list: list,

1193 model_config: Optional[PretrainedConfig],

1194 ) -> bool:

1195 # Special handling for the LogitsProcessor.

@@ -1256,7 +1256,7 @@ class LinearScalingRotaryEmbeddingWithLoRA(BaseLayerWithLoRA):

1256 positions: torch.Tensor,

1257 query: torch.Tensor,

1258 key: torch.Tensor,

1259 - ) -> Tuple[torch.Tensor, torch.Tensor]:

1259 + ) -> tuple[torch.Tensor, torch.Tensor]:

1260 return self.base_layer(

1261 positions,

1262 query,

@@ -1265,7 +1265,7 @@ class LinearScalingRotaryEmbeddingWithLoRA(BaseLayerWithLoRA):

1265 )

1266

1267 @property

1268 - def scaling_factor_to_offset(self) -> Dict[float, int]:

1268 + def scaling_factor_to_offset(self) -> dict[float, int]:

1269 return self.base_layer.scaling_factor_to_offset

1270

1271 @classmethod

@@ -1273,7 +1273,7 @@ class LinearScalingRotaryEmbeddingWithLoRA(BaseLayerWithLoRA):

1273 cls,

1274 source_layer: nn.Module,

1275 lora_config: LoRAConfig,

1276 - packed_modules_list: List,

1276 + packed_modules_list: list,

1277 model_config: Optional[PretrainedConfig],

1278 ) -> bool:

1279 """Returns True if the layer can be replaced by this LoRA layer."""

vllm/lora/lora.py

+7 -7

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import List, Optional

4 -from typing import Sequence as GenericSequence

3 +from collections.abc import Sequence as GenericSequence

4 +from typing import Optional

6 import torch

7 import torch.types

@@ -125,11 +125,11 @@ class PackedLoRALayerWeights(LoRALayerWeights):

125 self,

126 module_name: str,

127 rank: int,

128 - lora_alphas: List[Optional[int]],

129 - lora_a: List[Optional[torch.Tensor]],

130 - lora_b: List[Optional[torch.Tensor]],

131 - bias: Optional[List[Optional[torch.Tensor]]] = None,

132 - scaling: Optional[List[float]] = None,

128 + lora_alphas: list[Optional[int]],

129 + lora_a: list[Optional[torch.Tensor]],

130 + lora_b: list[Optional[torch.Tensor]],

131 + bias: Optional[list[Optional[torch.Tensor]]] = None,

132 + scaling: Optional[list[float]] = None,

133 ) -> None:

134 super().__init__(

135 module_name=module_name,

vllm/lora/models.py

+28 -28

@@ -4,9 +4,9 @@ import copy

4 import math

5 import os

6 import re

7 +from collections.abc import Sequence

8 from dataclasses import dataclass, field

8 -from typing import (Any, Callable, Dict, List, Optional, Sequence, Set, Type,

9 - Union)

9 +from typing import Any, Callable, Optional, Union

11 import safetensors.torch

12 import torch

@@ -44,12 +44,12 @@ _GLOBAL_LORA_ID = 0

44 class LongContextLoRAContext:

45 """Context for lora adapters that support long context."""

46 # The scaling factors to support long context lora fine tuned models.

47 - scaling_factors: List[float]

47 + scaling_factors: list[float]

48 # dimension to apply rotary embedding.

49 rot_dim: int

50 # offsets to the sin_cos_cache for each lora_id loaded.

51 # This value is dynamically modified.

52 - offsets_by_lora_id: Dict[int, int] = field(default_factory=dict)

52 + offsets_by_lora_id: dict[int, int] = field(default_factory=dict)

55 def get_lora_id():

@@ -65,7 +65,7 @@ class LoRAModel(AdapterModel):

65 self,

66 lora_model_id: int,

67 rank: int,

68 - loras: Dict[str, LoRALayerWeights],

68 + loras: dict[str, LoRALayerWeights],

69 scaling_factor: Optional[float] = None,

70 ) -> None:

71 """

@@ -84,7 +84,7 @@ class LoRAModel(AdapterModel):

84 lora_model_id

85 > 0), f"a valid lora id should be greater than 0, got {self.id}"

86 self.rank = rank

87 - self.loras: Dict[str, LoRALayerWeights] = loras

87 + self.loras: dict[str, LoRALayerWeights] = loras

89 def clone(self, lora_model_id: int) -> "LoRAModel":

90 """Return a copy of the object with different ids.

@@ -113,19 +113,19 @@ class LoRAModel(AdapterModel):

113 def from_lora_tensors(

114 cls,

115 lora_model_id: int,

116 - tensors: Dict[str, torch.Tensor],

116 + tensors: dict[str, torch.Tensor],

117 peft_helper: PEFTHelper,

118 device: str = "cuda",

119 dtype: Optional[torch.dtype] = None,

120 - embeddings: Optional[Dict[str, torch.Tensor]] = None,

120 + embeddings: Optional[dict[str, torch.Tensor]] = None,

121 target_embedding_padding: Optional[int] = None,

122 - embedding_modules: Optional[Dict[str, str]] = None,

123 - embedding_padding_modules: Optional[List[str]] = None,

122 + embedding_modules: Optional[dict[str, str]] = None,

123 + embedding_padding_modules: Optional[list[str]] = None,

124 weights_mapper: Optional[WeightsMapper] = None,

125 ) -> "LoRAModel":

126 """Create a LoRAModel from a dictionary of tensors."""

127 pin_memory = str(device) == "cpu" and is_pin_memory_available()

128 - loras: Dict[str, LoRALayerWeights] = {}

128 + loras: dict[str, LoRALayerWeights] = {}

129 for tensor_name, tensor in tensors.items():

130 module_name, is_lora_a, is_bias = parse_fine_tuned_lora_name(

131 tensor_name, weights_mapper)

@@ -187,15 +187,15 @@ class LoRAModel(AdapterModel):

187 def from_local_checkpoint(

188 cls,

189 lora_dir: str,

190 - expected_lora_modules: List[str],

190 + expected_lora_modules: list[str],

191 peft_helper: PEFTHelper,

192 *,

193 lora_model_id: Optional[int] = None,

194 device: str = "cuda",

195 dtype: Optional[torch.dtype] = None,

196 target_embedding_padding: Optional[int] = None,

197 - embedding_modules: Optional[Dict[str, str]] = None,

198 - embedding_padding_modules: Optional[List[str]] = None,

197 + embedding_modules: Optional[dict[str, str]] = None,

198 + embedding_padding_modules: Optional[list[str]] = None,

199 weights_mapper: Optional[WeightsMapper] = None,

200 ) -> "LoRAModel":

201 """Create a LoRAModel from a local checkpoint.

@@ -220,9 +220,9 @@ class LoRAModel(AdapterModel):

220 new_embeddings_bin_file_path = os.path.join(lora_dir,

221 "new_embeddings.bin")

222

223 - unexpected_modules: List[Union[list[str], str]]

223 + unexpected_modules: list[Union[list[str], str]]

224 if os.path.isfile(lora_tensor_path):

225 - tensors: Dict[str, torch.Tensor] = {}

225 + tensors: dict[str, torch.Tensor] = {}

226 # Find unexpected modules.

227 # Use safetensor key as a source of truth to find expected modules.

228 # in peft if you have target_modules A, B, C and C does not exist

@@ -329,7 +329,7 @@ class LoRAModelManager(AdapterModelManager):

329 self.max_num_seqs = max_num_seqs

330 assert self.capacity >= self.lora_slots

331 self.max_num_batched_tokens = math.ceil(max_num_batched_tokens / 8) * 8

332 - self.lora_index_to_id: List[Optional[int]] = [None] * self.lora_slots

332 + self.lora_index_to_id: list[Optional[int]] = [None] * self.lora_slots

333 self.vocab_size = vocab_size

334 self.long_lora_context: Optional[LongContextLoRAContext] = None

335 self.punica_wrapper = get_punica_wrapper(

@@ -339,7 +339,7 @@ class LoRAModelManager(AdapterModelManager):

339 max_loras=self.lora_config.max_loras)

340 # Scaling factor -> offset to the sin_cos_cache to it.

341 # Used for long context lora.

342 - self.scaling_factor_to_offset: Dict[float, int] = {}

342 + self.scaling_factor_to_offset: dict[float, int] = {}

343 super().__init__(model)

344

345 self.supported_lora_modules = get_supported_lora_modules(self.model)

@@ -358,9 +358,9 @@ class LoRAModelManager(AdapterModelManager):

358 # text modules (e.g. ChatGLM)

359 and hasattr(self.model, "get_mm_mapping"))

360 self.is_pooling_model = is_pooling_model(self.model)

361 - self.packed_modules: Dict[str, List[str]] = {}

362 - self.modules: Dict[str, BaseLayerWithLoRA] = {}

363 - # Dict instead of a Set for compatibility with LRUCache.

361 + self.packed_modules: dict[str, list[str]] = {}

362 + self.modules: dict[str, BaseLayerWithLoRA] = {}

363 + # Dict instead of a set for compatibility with LRUCache.

364 self._last_mapping: Optional[LoRAMapping] = None

365 self._create_lora_modules()

366 self.model.lora_manager = self

@@ -530,7 +530,7 @@ class LoRAModelManager(AdapterModelManager):

530 lora_id: int,

531 rank: int,

532 scaling_factor: Optional[float],

533 - embedding_modules: Optional[Dict[str, str]] = None) -> LoRAModel:

533 + embedding_modules: Optional[dict[str, str]] = None) -> LoRAModel:

534 """Create zero-initialized LoRAModel for warmup."""

535 model = LoRAModel(lora_id, rank, {}, scaling_factor)

536 for module_name, module in self.model.named_modules():

@@ -578,7 +578,7 @@ class LoRAModelManager(AdapterModelManager):

578 else:

579 parts = module_name.split(".")

580 replacements = self.packed_modules_mapping[parts[-1]]

581 - subloras: List[Optional[LoRALayerWeights]] = []

581 + subloras: list[Optional[LoRALayerWeights]] = []

582 for i, r in enumerate(replacements):

583 lora = LoRALayerWeights.create_dummy_lora_weights(

584 module_name + "." + r,

@@ -630,8 +630,8 @@ class LoRAModelManager(AdapterModelManager):

630

631 def _create_merged_loras_inplace(self, lora_model: LoRAModel) -> None:

632 for module_name, new_module_names in self.packed_modules.items():

633 - replacement_loras: List[Optional[LoRALayerWeights]] = []

634 - replaced_module: Set[str] = set()

633 + replacement_loras: list[Optional[LoRALayerWeights]] = []

634 + replaced_module: set[str] = set()

635 has_replacement = False

636 for r in new_module_names:

637 lora = self._get_lora_layer_weights(lora_model, r)

@@ -694,7 +694,7 @@ class LoRAModelManager(AdapterModelManager):

694 return remove_adapter(adapter_id, self._registered_adapters,

695 self.deactivate_adapter)

696

697 - def list_adapters(self) -> Dict[int, Any]:

697 + def list_adapters(self) -> dict[int, Any]:

698 return list_adapters(self._registered_adapters)

699

700 def get_adapter(self, adapter_id: int) -> Optional[Any]:

@@ -721,7 +721,7 @@ class LRUCacheLoRAModelManager(LoRAModelManager):

721 self._active_adapters: LoRALRUCache = LoRALRUCache(

722 self.lora_slots, self._deactivate_adapter)

723

724 - def list_adapters(self) -> Dict[int, LoRAModel]:

724 + def list_adapters(self) -> dict[int, LoRAModel]:

725 """List all registered LoRAModels."""

726 return dict(self._registered_adapters.cache)

727

@@ -786,7 +786,7 @@ def create_lora_manager(

786 vocab_size: int,

787 lora_config: LoRAConfig,

788 device: torch.device,

789 - lora_manager_cls: Type[LoRAModelManager] = LoRAModelManager,

789 + lora_manager_cls: type[LoRAModelManager] = LoRAModelManager,

790 **kwargs) -> LoRAModelManager:

791 """Create a LoRA adapter for a given model."""

792 if not hasattr(model, "packed_modules_mapping"):

vllm/lora/ops/triton_ops/lora_expand_op.py

+3 -5

@@ -6,8 +6,6 @@ Punica: Multi-Tenant LoRA Serving.

6 https://arxiv.org/abs/2310.18547

7 """

9 -from typing import List

10 -

9 import torch

10 import triton

11 import triton.language as tl

@@ -127,7 +125,7 @@ def _lora_expand_kernel(

125 @torch.inference_mode()

126 def _lora_expand(

127 inputs: torch.Tensor, # shape [num_slices, num_tokens, lora_rank]

130 - lora_b_weights: List[

128 + lora_b_weights: list[

129 torch.Tensor], # shape [num_lora, hidden_size, lora_rank]

130 output_tensor: torch.

131 Tensor, # shape [num_tokens, hidden_size * num_slices]

@@ -143,7 +141,7 @@ def _lora_expand(

141 """

142 Args:

143 inputs (torch.Tensor): input tensor

146 - lora_b_weights (List[torch.Tensor]): lora'b weight

144 + lora_b_weights (list[torch.Tensor]): lora'b weight

145 output_tensor (torch.Tensor): output tensor

146 token_lora_mapping (torch.Tensor): A tensor mapping each input token

147 to the lora-id related to that token. A value of -1 indicates that

@@ -264,7 +262,7 @@ def _lora_expand(

262

263 def _lora_expand_fake(

264 inputs: torch.Tensor,

267 - lora_b_weights: List[torch.Tensor],

265 + lora_b_weights: list[torch.Tensor],

266 output_tensor: torch.Tensor,

267 token_lora_mapping: torch.Tensor,

268 token_indices_sorted_by_lora_ids: torch.Tensor,

vllm/lora/ops/triton_ops/lora_kernel_metadata.py

+2 -2

@@ -4,7 +4,7 @@ LoRA kernels metadata preparation utilities.

4 """

6 from dataclasses import dataclass

7 -from typing import Tuple, Union

7 +from typing import Union

9 import torch

@@ -125,7 +125,7 @@ class LoRAKernelMeta:

125

126 def meta_args(

127 self, token_nums: int

128 - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,

128 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,

129 torch.Tensor, torch.Tensor]:

130 """

131 This function returns the kernel metadata required for the current

vllm/lora/ops/triton_ops/lora_shrink_op.py

+3 -5

@@ -6,8 +6,6 @@ Punica: Multi-Tenant LoRA Serving.

6 https://arxiv.org/abs/2310.18547

7 """

9 -from typing import List

10 -

9 import torch

10 import triton

11 import triton.language as tl

@@ -98,7 +96,7 @@ def _lora_shrink_kernel(input_ptr, lora_ptr, out_ptr, M, N, K,

96 @torch.inference_mode()

97 def _lora_shrink(

98 inputs: torch.Tensor, # shape [num_tokens, hidden_size]

101 - lora_a_weights: List[

99 + lora_a_weights: list[

100 torch.Tensor], # shape [num_loras, lora_rank, hidden_size]

101 output_tensor: torch.Tensor, # shape [num_slices, num_tokens, lora_rank]

102 token_lora_mapping: torch.Tensor, # shape [num_tokens]

@@ -112,7 +110,7 @@ def _lora_shrink(

110 """

111 Args:

112 inputs (torch.Tensor): Input tensor

115 - lora_a_weights (List[torch.Tensor]): LoRA weights

113 + lora_a_weights (list[torch.Tensor]): LoRA weights

114 output_tensor (torch.Tensor): output tensor

115 token_lora_mapping (torch.Tensor): A tensor mapping each input token

116 to the lora-id related to that token. A value of -1 indicates that

@@ -219,7 +217,7 @@ def _lora_shrink(

217

218 def _lora_shrink_fake(

219 inputs: torch.Tensor,

222 - lora_a_weights: List[torch.Tensor],

220 + lora_a_weights: list[torch.Tensor],

221 output_tensor: torch.Tensor,

222 token_lora_mapping: torch.Tensor,

223 token_indices_sorted_by_lora_ids: torch.Tensor,

vllm/lora/ops/triton_ops/utils.py

+4 -6

@@ -1,14 +1,12 @@

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import Dict, List, Tuple

4 -

3 import torch

7 -_LORA_A_PTR_DICT: Dict[Tuple[int, ...], Tuple[torch.tensor, ...]] = {}

8 -_LORA_B_PTR_DICT: Dict[Tuple[int, ...], Tuple[torch.tensor, ...]] = {}

5 +_LORA_A_PTR_DICT: dict[tuple[int, ...], tuple[torch.tensor, ...]] = {}

6 +_LORA_B_PTR_DICT: dict[tuple[int, ...], tuple[torch.tensor, ...]] = {}

11 -def _get_lora_a_ptr(lora_a_weights: List[torch.Tensor], device: torch.device):

9 +def _get_lora_a_ptr(lora_a_weights: list[torch.Tensor], device: torch.device):

10 """

11 `_LORA_A_PTR_DICT` collects the required information during `profile_run`,

12 After this, it remains constant and subsequent usage is through LUT.

@@ -53,7 +51,7 @@ def _get_lora_a_ptr(lora_a_weights: List[torch.Tensor], device: torch.device):

51 return _LORA_A_PTR_DICT.get(key)

56 -def _get_lora_b_ptr(lora_weights: List[torch.Tensor], offset_start: int,

54 +def _get_lora_b_ptr(lora_weights: list[torch.Tensor], offset_start: int,

55 device: torch.device):

56 """

57 `_LORA_B_PTR_DICT` collects the required information during `profile_run`,

vllm/lora/peft_helper.py

+2 -2

@@ -6,7 +6,7 @@ import json

6 import math

7 import os

8 from dataclasses import MISSING, dataclass, field, fields

9 -from typing import List, Literal, Optional, Union

9 +from typing import Literal, Optional, Union

11 from vllm.config import LoRAConfig

12 from vllm.logger import init_logger

@@ -40,7 +40,7 @@ class PEFTHelper:

40 vllm_max_position_embeddings: Optional[int] = field(default=False)

41 vllm_long_context_scaling_factor: Optional[float] = field(default=None)

43 - def _validate_features(self) -> List[str]:

43 + def _validate_features(self) -> list[str]:

44 """

45 Check if there are any unsupported LoRA features.

46 """

vllm/lora/punica_wrapper/punica_base.py

+41 -41

@@ -7,7 +7,7 @@ https://arxiv.org/abs/2310.18547

7 """

9 from abc import ABC, abstractmethod

10 -from typing import TYPE_CHECKING, List, Optional, Tuple, Union

10 +from typing import TYPE_CHECKING, Optional, Union

12 import torch

@@ -28,7 +28,7 @@ class PunicaWrapperABC(ABC):

28 def update_metadata(

29 self,

30 mapping: "LoRAMapping",

31 - lora_index_to_id: List[Optional[int]],

31 + lora_index_to_id: list[Optional[int]],

32 max_loras: int,

33 vocab_size: int,

34 extra_vocab_size: int,

@@ -43,9 +43,9 @@ class PunicaWrapperABC(ABC):

43 @abstractmethod

44 def add_shrink(

45 self,

46 - y: Union[Tuple[torch.Tensor, ...], torch.Tensor],

46 + y: Union[tuple[torch.Tensor, ...], torch.Tensor],

47 x: torch.Tensor,

48 - lora_a_stacked: Tuple[torch.Tensor, ...],

48 + lora_a_stacked: tuple[torch.Tensor, ...],

49 scale: float,

50 **kwargs,

51 ) -> Optional[torch.Tensor]:

@@ -59,10 +59,10 @@ class PunicaWrapperABC(ABC):

59 def add_expand(

60 self,

61 y: torch.Tensor,

62 - x: Union[Tuple[torch.Tensor, ...], torch.Tensor],

63 - lora_b_stacked: Tuple[torch.Tensor, ...],

64 - lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]],

65 - output_slices: Tuple[int, ...],

62 + x: Union[tuple[torch.Tensor, ...], torch.Tensor],

63 + lora_b_stacked: tuple[torch.Tensor, ...],

64 + lora_bias_stacked: Optional[tuple[torch.Tensor, ...]],

65 + output_slices: tuple[int, ...],

66 offset_start: int = 0,

67 add_inputs=True,

68 **kwargs,

@@ -91,13 +91,13 @@ class PunicaWrapperABC(ABC):

91 def add_lora_linear(self,

92 y: torch.Tensor,

93 x: torch.Tensor,

94 - lora_a_stacked: Tuple[torch.Tensor, ...],

95 - lora_b_stacked: Tuple[torch.Tensor, ...],

96 - lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]],

94 + lora_a_stacked: tuple[torch.Tensor, ...],

95 + lora_b_stacked: tuple[torch.Tensor, ...],

96 + lora_bias_stacked: Optional[tuple[torch.Tensor, ...]],

97 scale: float,

98 - output_slices: Tuple[int, ...],

98 + output_slices: tuple[int, ...],

99 *,

100 - buffer: Optional[Tuple[torch.Tensor, ...]] = None,

100 + buffer: Optional[tuple[torch.Tensor, ...]] = None,

101 **kwargs) -> Optional[torch.Tensor]:

102 """

103 Applicable to linear-related lora.

@@ -150,7 +150,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

150 # 5 is the number of indices tensors.

151 # base_indices, sampler_indices, sampler_indices_padded,

152 # embeddings_indices,long_lora_indices

153 - self.indices_len: List[Optional[int]] = [None] * 5

153 + self.indices_len: list[Optional[int]] = [None] * 5

154 # these attributes are the information required for sgmv kernel

155 self._seq_start_locs = torch.empty(max_batches,

156 dtype=torch.long,

@@ -171,7 +171,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

171 def _update_base_metadata(

172 self,

173 mapping: "LoRAMapping",

174 - lora_index_to_id: List[Optional[int]],

174 + lora_index_to_id: list[Optional[int]],

175 max_loras: int,

176 vocab_size: int,

177 extra_vocab_size: int,

@@ -228,8 +228,8 @@ class PunicaWrapperBase(PunicaWrapperABC):

228 self,

229 indices: torch.Tensor,

230 output: torch.Tensor,

231 - output_slices: Tuple[int, ...],

232 - lora_bias_stacked: Tuple[Optional[torch.Tensor], ...],

231 + output_slices: tuple[int, ...],

232 + lora_bias_stacked: tuple[Optional[torch.Tensor], ...],

233 ):

234 """Applies bias to output

235

@@ -259,7 +259,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

259 @property

260 def prefill_metadata(

261 self

262 - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int]:

262 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int]:

263 """

264 This property provides a convenient way to access the necessary

265 metadata for prefill-related kernel computations.

@@ -323,7 +323,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

323 def update_metadata(

324 self,

325 mapping: "LoRAMapping",

326 - lora_index_to_id: List[Optional[int]],

326 + lora_index_to_id: list[Optional[int]],

327 max_loras: int,

328 vocab_size: int,

329 extra_vocab_size: int,

@@ -341,8 +341,8 @@ class PunicaWrapperBase(PunicaWrapperABC):

341 self.is_prefill = False

342

343 @abstractmethod

344 - def add_shrink(self, y: Union[Tuple[torch.Tensor, ...], torch.Tensor],

345 - x: torch.Tensor, lora_a_stacked: Tuple[torch.Tensor, ...],

344 + def add_shrink(self, y: Union[tuple[torch.Tensor, ...], torch.Tensor],

345 + x: torch.Tensor, lora_a_stacked: tuple[torch.Tensor, ...],

346 scale: float, **kwargs) -> Optional[torch.Tensor]:

347 """

348 Performs GEMM for multiple slices of lora_a.

@@ -352,9 +352,9 @@ class PunicaWrapperBase(PunicaWrapperABC):

352 y[i] += (x @ lora_a_stacked[i]) * scale

353

354 Args:

355 - y (Union[Tuple[torch.Tensor, ...], torch.Tensor]): Output tensors

355 + y (Union[tuple[torch.Tensor, ...], torch.Tensor]): Output tensors

356 x (torch.Tensor): Input tensor

357 - lora_a_stacked (Tuple[torch.Tensor, ...]): lora_a's weights

357 + lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights

358 scale (float): Scaling factor for the operation

359

360 """

@@ -364,10 +364,10 @@ class PunicaWrapperBase(PunicaWrapperABC):

364 @abstractmethod

365 def add_expand(self,

366 y: torch.Tensor,

367 - x: Union[Tuple[torch.Tensor, ...], torch.Tensor],

368 - lora_b_stacked: Tuple[torch.Tensor, ...],

369 - lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]],

370 - output_slices: Tuple[int, ...],

367 + x: Union[tuple[torch.Tensor, ...], torch.Tensor],

368 + lora_b_stacked: tuple[torch.Tensor, ...],

369 + lora_bias_stacked: Optional[tuple[torch.Tensor, ...]],

370 + output_slices: tuple[int, ...],

371 offset_start: int = 0,

372 add_inputs=True,

373 **kwargs) -> Optional[torch.Tensor]:

@@ -384,11 +384,11 @@ class PunicaWrapperBase(PunicaWrapperABC):

384

385 Args:

386 y (torch.Tensor): Output tensor.

387 - x (Union[Tuple[torch.Tensor, ...], torch.Tensor]): Input tensors

388 - lora_b_stacked (Tuple[torch.Tensor, ...]): lora_b's weight

389 - lora_bias_stacked (Optional[Tuple[torch.Tensor, ...]]):

387 + x (Union[tuple[torch.Tensor, ...], torch.Tensor]): Input tensors

388 + lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight

389 + lora_bias_stacked (Optional[tuple[torch.Tensor, ...]]):

390 bias's weight

391 - output_slices (Tuple[int, ...]): Every slice's size

391 + output_slices (tuple[int, ...]): Every slice's size

392 offset_start (int): The starting position of y, defaults to 0

393 add_inputs (bool): Defaults to True.

394

@@ -422,13 +422,13 @@ class PunicaWrapperBase(PunicaWrapperABC):

422 def add_lora_linear(self,

423 y: torch.Tensor,

424 x: torch.Tensor,

425 - lora_a_stacked: Tuple[torch.Tensor, ...],

426 - lora_b_stacked: Tuple[torch.Tensor, ...],

427 - lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]],

425 + lora_a_stacked: tuple[torch.Tensor, ...],

426 + lora_b_stacked: tuple[torch.Tensor, ...],

427 + lora_bias_stacked: Optional[tuple[torch.Tensor, ...]],

428 scale: float,

429 - output_slices: Tuple[int, ...],

429 + output_slices: tuple[int, ...],

430 *,

431 - buffer: Optional[Tuple[torch.Tensor, ...]] = None,

431 + buffer: Optional[tuple[torch.Tensor, ...]] = None,

432 **kwargs) -> Optional[torch.Tensor]:

433 """

434 Applicable to linear-related lora.

@@ -445,12 +445,12 @@ class PunicaWrapperBase(PunicaWrapperABC):

445 Args:

446 y (torch.Tensor): Output tensor. Will be changed in-place.

447 x (torch.Tensor): Input tensor

448 - lora_a_stacked (Tuple[torch.Tensor, ...]): lora_a's weight.

449 - lora_b_stacked (Tuple[torch.Tensor, ...]): lora_b's weight.

450 - lora_bias_stacked (Optional[Tuple[torch.Tensor, ...]]): lora's bias.

448 + lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weight.

449 + lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight.

450 + lora_bias_stacked (Optional[tuple[torch.Tensor, ...]]): lora's bias.

451 scale (float): Scaling factor.

452 - output_slices (Tuple[int, ...]): Every slice's size.

453 - buffer (Optional[Tuple[torch.Tensor, ...]]): Defaults to None.

452 + output_slices (tuple[int, ...]): Every slice's size.

453 + buffer (Optional[tuple[torch.Tensor, ...]]): Defaults to None.

454 """

455 # TODO: implement it based on torch ops

456 raise NotImplementedError

vllm/lora/punica_wrapper/punica_cpu.py

+23 -23

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import Callable, Optional, Tuple, Union

3 +from typing import Callable, Optional, Union

5 import torch

@@ -150,8 +150,8 @@ class PunicaWrapperCPU(PunicaWrapperBase):

150 shrink_fun(y, x, w_t_all, scale)

151 y = y.view_as(y_org)

152

153 - def add_shrink(self, y: Union[Tuple[torch.Tensor, ...], torch.Tensor],

154 - x: torch.Tensor, lora_a_stacked: Tuple[torch.Tensor, ...],

153 + def add_shrink(self, y: Union[tuple[torch.Tensor, ...], torch.Tensor],

154 + x: torch.Tensor, lora_a_stacked: tuple[torch.Tensor, ...],

155 scale: float, **kwargs):

156 """

157 Performs GEMM for multiple slices of lora_a.

@@ -165,9 +165,9 @@ class PunicaWrapperCPU(PunicaWrapperBase):

165 y[i] += (x @ lora_a_stacked[i]) * scale

166

167 Args:

168 - y (Union[Tuple[torch.Tensor, ...], torch.Tensor]): Output tensors

168 + y (Union[tuple[torch.Tensor, ...], torch.Tensor]): Output tensors

169 x (torch.Tensor): Input tensor

170 - lora_a_stacked (Tuple[torch.Tensor, ...]): lora_a's weights

170 + lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights

171 scale (float): Scaling factor for the operation

172 """

173

@@ -179,10 +179,10 @@ class PunicaWrapperCPU(PunicaWrapperBase):

179

180 def add_expand(self,

181 y: torch.Tensor,

182 - x: Union[Tuple[torch.Tensor, ...], torch.Tensor],

183 - lora_b_stacked: Tuple[torch.Tensor, ...],

184 - lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]],

185 - output_slices: Tuple[int, ...],

182 + x: Union[tuple[torch.Tensor, ...], torch.Tensor],

183 + lora_b_stacked: tuple[torch.Tensor, ...],

184 + lora_bias_stacked: Optional[tuple[torch.Tensor, ...]],

185 + output_slices: tuple[int, ...],

186 offset_start: int = 0,

187 add_inputs=True,

188 **kwargs) -> None:

@@ -198,11 +198,11 @@ class PunicaWrapperCPU(PunicaWrapperBase):

198

199 Args:

200 y (torch.Tensor): Output tensor.

201 - x (Union[Tuple[torch.Tensor, ...], torch.Tensor]): Input tensors

202 - lora_b_stacked (Tuple[torch.Tensor, ...]): lora_b's weight

203 - lora_bias_stacked (Optional[Tuple[torch.Tensor, ...]]):

201 + x (Union[tuple[torch.Tensor, ...], torch.Tensor]): Input tensors

202 + lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight

203 + lora_bias_stacked (Optional[tuple[torch.Tensor, ...]]):

204 bias's weight

205 - output_slices (Tuple[int, ...]): Every slice's size

205 + output_slices (tuple[int, ...]): Every slice's size

206 add_inputs (bool): Defaults to True.

207 """

208 y_org = y

@@ -250,13 +250,13 @@ class PunicaWrapperCPU(PunicaWrapperBase):

250 def add_lora_linear(self,

251 y: torch.Tensor,

252 x: torch.Tensor,

253 - lora_a_stacked: Tuple[torch.Tensor, ...],

254 - lora_b_stacked: Tuple[torch.Tensor, ...],

255 - lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]],

253 + lora_a_stacked: tuple[torch.Tensor, ...],

254 + lora_b_stacked: tuple[torch.Tensor, ...],

255 + lora_bias_stacked: Optional[tuple[torch.Tensor, ...]],

256 scale: float,

257 - output_slices: Tuple[int, ...],

257 + output_slices: tuple[int, ...],

258 *,

259 - buffer: Optional[Tuple[torch.Tensor, ...]] = None,

259 + buffer: Optional[tuple[torch.Tensor, ...]] = None,

260 **kwargs) -> None:

261 """

262 Applicable to linear-related lora.

@@ -273,12 +273,12 @@ class PunicaWrapperCPU(PunicaWrapperBase):

273 Args:

274 y (torch.Tensor): Output tensor. Will be changed in-place.

275 x (torch.Tensor): Input tensor

276 - lora_a_stacked (Tuple[torch.Tensor, ...]): lora_a's weight.

277 - lora_b_stacked (Tuple[torch.Tensor, ...]): lora_b's weight.

278 - lora_bias_stacked (Optional[Tuple[torch.Tensor, ...]]): lora's bias.

276 + lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weight.

277 + lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight.

278 + lora_bias_stacked (Optional[tuple[torch.Tensor, ...]]): lora's bias.

279 scale (float): Scaling factor.

280 - output_slices (Tuple[int, ...]): Every slice's size.

281 - buffer (Optional[Tuple[torch.Tensor, ...]]): Defaults to None.

280 + output_slices (tuple[int, ...]): Every slice's size.

281 + buffer (Optional[tuple[torch.Tensor, ...]]): Defaults to None.

282 """

283

284 assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices)

vllm/lora/punica_wrapper/punica_gpu.py

+18 -18

@@ -6,7 +6,7 @@ Punica: Multi-Tenant LoRA Serving.

6 https://arxiv.org/abs/2310.18547

7 """

9 -from typing import TYPE_CHECKING, List, Optional, Tuple, Union, final

9 +from typing import TYPE_CHECKING, Optional, Union, final

11 import torch

@@ -57,7 +57,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

57 def update_metadata(

58 self,

59 mapping: LoRAMapping,

60 - lora_index_to_id: List[Optional[int]],

60 + lora_index_to_id: list[Optional[int]],

61 max_loras: int,

62 vocab_size: int,

63 extra_vocab_size: int,

@@ -74,7 +74,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

74 self.prompt_mapping_meta.prepare_tensors(self.sampler_indices)

76 def add_shrink(self, y: torch.Tensor, x: torch.Tensor,

77 - lora_a_stacked: Tuple[torch.Tensor,

77 + lora_a_stacked: tuple[torch.Tensor,

78 ...], scale: float, **kwargs):

79 """

80 Performs GEMM for multiple slices of lora_a.

@@ -86,7 +86,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

86 Args:

87 y (torch.Tensor): Output tensors

88 x (torch.Tensor): Input tensor

89 - lora_a_stacked (Tuple[torch.Tensor, ...]): lora_a's weights

89 + lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights

90 scale (float): Scaling factor for the operation

91 """

@@ -102,9 +102,9 @@ class PunicaWrapperGPU(PunicaWrapperBase):

102 def add_expand(self,

103 y: torch.Tensor,

104 x: torch.Tensor,

105 - lora_b_stacked: Tuple[torch.Tensor, ...],

106 - lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]],

107 - output_slices: Tuple[int, ...],

105 + lora_b_stacked: tuple[torch.Tensor, ...],

106 + lora_bias_stacked: Optional[tuple[torch.Tensor, ...]],

107 + output_slices: tuple[int, ...],

108 offset_start: int = 0,

109 add_inputs=True,

110 **kwargs) -> None:

@@ -121,10 +121,10 @@ class PunicaWrapperGPU(PunicaWrapperBase):

121 Args:

122 y (torch.Tensor): Output tensor.

123 x (torch.Tensor): Input tensors

124 - lora_b_stacked (Tuple[torch.Tensor, ...]): lora_b's weight

125 - lora_bias_stacked (Optional[Tuple[torch.Tensor, ...]]):

124 + lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight

125 + lora_bias_stacked (Optional[tuple[torch.Tensor, ...]]):

126 bias's weight

127 - output_slices (Tuple[int, ...]): Every slice's size

127 + output_slices (tuple[int, ...]): Every slice's size

128 add_inputs (bool): Defaults to True.

129 """

130 y_org = y

@@ -181,11 +181,11 @@ class PunicaWrapperGPU(PunicaWrapperBase):

181 def add_lora_linear(self,

182 y: torch.Tensor,

183 x: torch.Tensor,

184 - lora_a_stacked: Tuple[torch.Tensor, ...],

185 - lora_b_stacked: Tuple[torch.Tensor, ...],

186 - lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]],

184 + lora_a_stacked: tuple[torch.Tensor, ...],

185 + lora_b_stacked: tuple[torch.Tensor, ...],

186 + lora_bias_stacked: Optional[tuple[torch.Tensor, ...]],

187 scale: float,

188 - output_slices: Tuple[int, ...],

188 + output_slices: tuple[int, ...],

189 *,

190 buffer: Optional[torch.Tensor] = None,

191 **kwargs) -> None:

@@ -204,11 +204,11 @@ class PunicaWrapperGPU(PunicaWrapperBase):

204 Args:

205 y (torch.Tensor): Output tensor. Will be changed in-place.

206 x (torch.Tensor): Input tensor

207 - lora_a_stacked (Tuple[torch.Tensor, ...]): lora_a's weight.

208 - lora_b_stacked (Tuple[torch.Tensor, ...]): lora_b's weight.

209 - lora_bias_stacked (Optional[Tuple[torch.Tensor, ...]]): lora's bias.

207 + lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weight.

208 + lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight.

209 + lora_bias_stacked (Optional[tuple[torch.Tensor, ...]]): lora's bias.

210 scale (float): Scaling factor.

211 - output_slices (Tuple[int, ...]): Every slice's size.

211 + output_slices (tuple[int, ...]): Every slice's size.

212 buffer (Optional[torch.Tensor]): Defaults to None.

213 """

214

vllm/lora/punica_wrapper/punica_hpu.py

+15 -15

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import TYPE_CHECKING, List, Optional, Tuple, Union, final

3 +from typing import TYPE_CHECKING, Optional, Union, final

5 import torch

6 from vllm_hpu_extension.ops import (dispatch_bgmv_embedding,

@@ -28,7 +28,7 @@ class PunicaWrapperHPU(PunicaWrapperBase):

28 def _update_base_metadata(

29 self,

30 mapping: "LoRAMapping",

31 - lora_index_to_id: List[Optional[int]],

31 + lora_index_to_id: list[Optional[int]],

32 max_loras: int,

33 vocab_size: int,

34 extra_vocab_size: int,

@@ -48,9 +48,9 @@ class PunicaWrapperHPU(PunicaWrapperBase):

48 # graph accumulation. Hence HPU appends `lora_offset` to a list and

49 # converts it to a tensor only after it is ready.

50 if long_lora_context:

51 - index_mapping_indices: List[int] = list(

51 + index_mapping_indices: list[int] = list(

52 mapping.index_mapping).copy()

53 - long_lora_offsets: List[int] = []

53 + long_lora_offsets: list[int] = []

54 for i in range(len(index_mapping_indices)):

55 lora_offset: int = long_lora_context.offsets_by_lora_id.get(

56 index_mapping_indices[i], 0)

@@ -85,13 +85,13 @@ class PunicaWrapperHPU(PunicaWrapperBase):

85 def add_lora_linear(self,

86 y: torch.Tensor,

87 x: torch.Tensor,

88 - lora_a_stacked: Tuple[torch.Tensor, ...],

89 - lora_b_stacked: Tuple[torch.Tensor, ...],

90 - lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]],

88 + lora_a_stacked: tuple[torch.Tensor, ...],

89 + lora_b_stacked: tuple[torch.Tensor, ...],

90 + lora_bias_stacked: Optional[tuple[torch.Tensor, ...]],

91 scale: float,

92 - output_slices: Tuple[int, ...],

92 + output_slices: tuple[int, ...],

93 *,

94 - buffer: Optional[Tuple[torch.Tensor, ...]] = None,

94 + buffer: Optional[tuple[torch.Tensor, ...]] = None,

95 **kwargs) -> None:

96 y_org = y

97 x = x.view(-1, x.shape[-1])

@@ -122,9 +122,9 @@ class PunicaWrapperHPU(PunicaWrapperBase):

122

123 def add_shrink(

124 self,

125 - y: Union[Tuple[torch.Tensor, ...], torch.Tensor],

125 + y: Union[tuple[torch.Tensor, ...], torch.Tensor],

126 x: torch.Tensor,

127 - lora_a_stacked: Tuple[torch.Tensor, ...],

127 + lora_a_stacked: tuple[torch.Tensor, ...],

128 scale: float,

129 **kwargs,

130 ) -> None:

@@ -133,10 +133,10 @@ class PunicaWrapperHPU(PunicaWrapperBase):

133 def add_expand(

134 self,

135 y: torch.Tensor,

136 - x: Union[Tuple[torch.Tensor, ...], torch.Tensor],

137 - lora_b_stacked: Tuple[torch.Tensor, ...],

138 - lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]],

139 - output_slices: Tuple[int, ...],

136 + x: Union[tuple[torch.Tensor, ...], torch.Tensor],

137 + lora_b_stacked: tuple[torch.Tensor, ...],

138 + lora_bias_stacked: Optional[tuple[torch.Tensor, ...]],

139 + output_slices: tuple[int, ...],

140 offset_start: int = 0,

141 add_inputs=True,

142 **kwargs,

vllm/lora/punica_wrapper/punica_tpu.py

+25 -25

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import Optional, Tuple, Union

3 +from typing import Optional, Union

5 import torch

6 import torch.nn.functional as F

@@ -77,8 +77,8 @@ class PunicaWrapperTPU(PunicaWrapperBase):

77 self._get_token_lora_indices(x), y_offset,

78 y_slice_size, add_inputs)

80 - def add_shrink(self, y: Union[Tuple[torch.Tensor, ...], torch.Tensor],

81 - x: torch.Tensor, lora_a_stacked: Tuple[torch.Tensor, ...],

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.

@@ -88,9 +88,9 @@ class PunicaWrapperTPU(PunicaWrapperBase):

88 y[i] += (x @ lora_a_stacked[i]) * scale

90 Args:

91 - y (Union[Tuple[torch.Tensor, ...], torch.Tensor]): Output tensors

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

93 + lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights

94 scale (float): Scaling factor for the operation

95 """

@@ -106,10 +106,10 @@ class PunicaWrapperTPU(PunicaWrapperBase):

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, ...],

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:

@@ -125,11 +125,11 @@ class PunicaWrapperTPU(PunicaWrapperBase):

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, ...]]):

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

132 + output_slices (tuple[int, ...]): Every slice's size

133 add_inputs (bool): Defaults to True.

134 """

135 y_org = y

@@ -177,13 +177,13 @@ class PunicaWrapperTPU(PunicaWrapperBase):

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, ...]],

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, ...],

184 + output_slices: tuple[int, ...],

185 *,

186 - buffer: Optional[Tuple[torch.Tensor, ...]] = None,

186 + buffer: Optional[tuple[torch.Tensor, ...]] = None,

187 **kwargs) -> torch.Tensor:

188 """

189 Applicable to linear-related lora.

@@ -200,12 +200,12 @@ class PunicaWrapperTPU(PunicaWrapperBase):

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.

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.

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)

@@ -284,8 +284,8 @@ class PunicaWrapperTPU(PunicaWrapperBase):

284 self,

285 indices: torch.Tensor,

286 output: torch.Tensor,

287 - output_slices: Tuple[int, ...],

288 - lora_bias_stacked: Tuple[Optional[torch.Tensor], ...],

287 + output_slices: tuple[int, ...],

288 + lora_bias_stacked: tuple[Optional[torch.Tensor], ...],

289 ):

290 """Applies bias to output

291

vllm/lora/punica_wrapper/utils.py

+8 -8

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import TYPE_CHECKING, List, Optional, Tuple, Union

3 +from typing import TYPE_CHECKING, Optional, Union

5 import torch

@@ -12,7 +12,7 @@ if TYPE_CHECKING:

13 def compute_meta(

14 token_lora_tensor: torch.Tensor

15 -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int, bool]:

15 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int, bool]:

16 """

17 Get the information required for the sgmv kernel. With the features:

18 1. If consecutive requests in the batch use the same LoRA, this function

@@ -43,14 +43,14 @@ def compute_meta(

43 # TODO see if this can be vectorized

44 def convert_mapping(

45 mapping: "LoRAMapping",

46 - lora_index_to_id: List[Optional[int]],

46 + lora_index_to_id: list[Optional[int]],

47 max_loras: int,

48 vocab_size: int,

49 extra_vocab_size: int,

50 device: torch.device,

51 long_lora_context: Optional["LongContextLoRAContext"] = None,

52 -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,

53 - Optional[torch.Tensor], List[int]]:

52 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,

53 + Optional[torch.Tensor], list[int]]:

54 """Converts LoRAMapping to index tensors.

56 Args:

@@ -84,7 +84,7 @@ def convert_mapping(

84 (base_indices, sampler_indices, sampler_indices_padded,

85 embeddings_indices, long_lora_indices).

86 """

87 - index_mapping_indices: List[int] = list(mapping.index_mapping).copy()

87 + index_mapping_indices: list[int] = list(mapping.index_mapping).copy()

88 embedding_indices = index_mapping_indices.copy()

89 lora_indices = index_mapping_indices.copy()

90 long_lora_offsets: Optional[torch.Tensor] = None

@@ -92,7 +92,7 @@ def convert_mapping(

92 long_lora_offsets = torch.zeros(len(index_mapping_indices),

93 device=device,

94 dtype=torch.long)

95 - prompt_mapping: List[int] = [

95 + prompt_mapping: list[int] = [

96 lora_index_to_id.index(x) if x > 0 else -1

97 for x in mapping.prompt_mapping

98 ]

@@ -109,7 +109,7 @@ def convert_mapping(

109 index_mapping_indices[i], 0)

110 long_lora_offsets[i] = lora_offset

111

112 - indices_list: List[Union[List[int], torch.Tensor]] = [

112 + indices_list: list[Union[list[int], torch.Tensor]] = [

113 index_mapping_indices,

114 lora_indices,

115 embedding_indices,

vllm/lora/resolver.py

+4 -3

@@ -1,8 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

3 from abc import ABC, abstractmethod

4 +from collections.abc import Set

5 from dataclasses import dataclass, field

5 -from typing import AbstractSet, Dict, Optional

6 +from typing import Optional

8 from vllm.logger import init_logger

9 from vllm.lora.request import LoRARequest

@@ -40,9 +41,9 @@ class LoRAResolver(ABC):

42 @dataclass

43 class _LoRAResolverRegistry:

43 - resolvers: Dict[str, LoRAResolver] = field(default_factory=dict)

44 + resolvers: dict[str, LoRAResolver] = field(default_factory=dict)

45 - def get_supported_resolvers(self) -> AbstractSet[str]:

46 + def get_supported_resolvers(self) -> Set[str]:

47 """Get all registered resolver names."""

48 return self.resolvers.keys()

vllm/lora/utils.py

+9 -9

@@ -2,7 +2,7 @@

3 import os

4 import re

5 -from typing import List, Optional, Set, Tuple, Type, Union

5 +from typing import Optional, Union

7 import huggingface_hub

8 from huggingface_hub.utils import (EntryNotFoundError, HfHubHTTPError,

@@ -37,7 +37,7 @@ from vllm.model_executor.models.utils import WeightsMapper

38 logger = init_logger(__name__)

40 -_all_lora_classes: Set[Type[BaseLayerWithLoRA]] = {

40 +_all_lora_classes: set[type[BaseLayerWithLoRA]] = {

41 VocabParallelEmbeddingWithLoRA,

42 ColumnParallelLinearWithLoRA,

43 MergedColumnParallelLinearWithLoRA,

@@ -58,7 +58,7 @@ _all_lora_classes: Set[Type[BaseLayerWithLoRA]] = {

58 def from_layer(layer: nn.Module,

59 max_loras: int,

60 lora_config: LoRAConfig,

61 - packed_modules_list: List,

61 + packed_modules_list: list,

62 model_config: Optional[PretrainedConfig] = None) -> nn.Module:

63 for lora_cls in _all_lora_classes:

64 # specifying kwargs so they can be easily accessed in decorator

@@ -99,7 +99,7 @@ def replace_submodule(model: nn.Module, module_name: str,

99 def parse_fine_tuned_lora_name(

100 name: str,

101 weights_mapper: Optional[WeightsMapper] = None

102 -) -> Tuple[str, bool, bool]:

102 +) -> tuple[str, bool, bool]:

103 """Parse the name of lora weights.

104

105 args:

@@ -108,7 +108,7 @@ def parse_fine_tuned_lora_name(

108 weights_mapper: maps the name of weight, e.g.

109 `model.` -> `language_model.model.`,

110 return:

111 - Tuple(module_name, is_lora_a):

111 + tuple(module_name, is_lora_a):

112 module_name: the name of the module, e.g. model.dense1,

113 is_lora_a whether the tensor is lora_a or lora_b.

114 is_bias whether the tensor is lora bias.

@@ -147,8 +147,8 @@ def parse_fine_tuned_lora_name(

147 raise ValueError(f"{name} is unsupported LoRA weight")

148

149

150 -def is_regex_target_modules(load_modules: Union[str, List[str]],

151 - expected_lora_modules: List[str]) -> bool:

150 +def is_regex_target_modules(load_modules: Union[str, list[str]],

151 + expected_lora_modules: list[str]) -> bool:

152 """

153 PEFT supports passing `target_modules` in the form of regular expressions,

154 such as `model.*(q_proj|k_proj|v_proj)$`. This function is mainly used to

@@ -179,11 +179,11 @@ def is_regex_target_modules(load_modules: Union[str, List[str]],

179 return False

180

181

182 -def get_supported_lora_modules(model: nn.Module) -> List[str]:

182 +def get_supported_lora_modules(model: nn.Module) -> list[str]:

183 """

184 In vLLM, all linear layers support LoRA.

185 """

186 - supported_lora_modules: Set[str] = set()

186 + supported_lora_modules: set[str] = set()

187 # step1: traverse the model to get all the linear subfixes.

188 for name, module in model.named_modules():

189 if isinstance(module, (LinearBase, )):

vllm/lora/worker_manager.py

+11 -11

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

3 from contextlib import contextmanager

4 -from typing import Any, Dict, List, Literal, Optional, Set, Type, Union

4 +from typing import Any, Literal, Optional, Union

6 import torch

@@ -27,7 +27,7 @@ class WorkerLoRAManager(AbstractWorkerManager):

27 Every request, the requested LoRAs will be loaded (unless they are already

28 loaded), and every other LoRA will be unloaded."""

30 - _manager_cls: Type[LoRAModelManager] = LoRAModelManager

30 + _manager_cls: type[LoRAModelManager] = LoRAModelManager

32 def __init__(

33 self,

@@ -36,9 +36,9 @@ class WorkerLoRAManager(AbstractWorkerManager):

36 vocab_size: int,

37 lora_config: LoRAConfig,

38 device: torch.device,

39 - embedding_modules: Dict[str, str],

40 - embedding_padding_modules: List[str],

41 - lora_model_cls: Type[LoRAModel] = LoRAModel,

39 + embedding_modules: dict[str, str],

40 + embedding_padding_modules: list[str],

41 + lora_model_cls: type[LoRAModel] = LoRAModel,

42 max_position_embeddings: Optional[int] = None,

43 ):

44 self._lora_model_cls = lora_model_cls

@@ -88,7 +88,7 @@ class WorkerLoRAManager(AbstractWorkerManager):

88 self._adapter_manager.supported_lora_modules)

89 packed_modules_mapping = (

90 self._adapter_manager.packed_modules_mapping)

91 - expected_lora_modules: List[str] = []

91 + expected_lora_modules: list[str] = []

92 for module in supported_lora_modules:

93 if module in packed_modules_mapping:

94 expected_lora_modules.extend(

@@ -162,12 +162,12 @@ class WorkerLoRAManager(AbstractWorkerManager):

162 def pin_adapter(self, adapter_id: int) -> bool:

163 return self._adapter_manager.pin_adapter(adapter_id)

164

165 - def set_active_adapters(self, requests: Set[Any],

165 + def set_active_adapters(self, requests: set[Any],

166 mapping: Optional[Any]) -> None:

167 set_active_adapters_worker(requests, mapping, self._apply_adapters,

168 self._adapter_manager.set_adapter_mapping)

169

170 - def _apply_adapters(self, adapter_requests: Set[Any]) -> None:

170 + def _apply_adapters(self, adapter_requests: set[Any]) -> None:

171 apply_adapters_worker(adapter_requests, self.list_adapters,

172 self._adapter_manager.adapter_slots,

173 self.remove_adapter, self.add_adapter)

@@ -184,7 +184,7 @@ class WorkerLoRAManager(AbstractWorkerManager):

184 def remove_all_adapters(self):

185 self._adapter_manager.remove_all_adapters()

186

187 - def list_adapters(self) -> Set[int]:

187 + def list_adapters(self) -> set[int]:

188 return list_adapters_worker(self._adapter_manager.list_adapters)

189

190

@@ -195,7 +195,7 @@ class LRUCacheWorkerLoRAManager(WorkerLoRAManager):

195 (unless they are already loaded) and least recently used LoRAs will

196 be unloaded if the cache is above capacity."""

197

198 - _manager_cls: Type[LRUCacheLoRAModelManager] = LRUCacheLoRAModelManager

198 + _manager_cls: type[LRUCacheLoRAModelManager] = LRUCacheLoRAModelManager

199

200 def create_lora_manager(

201 self,

@@ -213,7 +213,7 @@ class LRUCacheWorkerLoRAManager(WorkerLoRAManager):

213 self._adapter_manager = lora_manager

214 return lora_manager.model

215

216 - def _apply_adapters(self, lora_requests: Set[LoRARequest]) -> None:

216 + def _apply_adapters(self, lora_requests: set[LoRARequest]) -> None:

217 loras_map = {

218 lora_request.lora_int_id: lora_request

219 for lora_request in lora_requests if lora_request