@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1eaff27815

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, 01:00 UTC

[V0 deprecation] Remove long context LoRA (#21169)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed Jul 19, 2025 at 17:15 UTC 1eaff2781585ce17b4353059146591acd65719f9

13 files changed +35 -301

tests/lora/conftest.py

-5

@@ -221,11 +221,6 @@ def phi2_lora_files():

221 return snapshot_download(repo_id="isotr0py/phi-2-test-sql-lora")

222

223

224 -@pytest.fixture(scope="session")

225 -def long_context_lora_files_16k_1():

226 - return snapshot_download(repo_id="SangBinCho/long_context_16k_testing_1")

227 -

228 -

224 @pytest.fixture

225 def llama_2_7b_engine_extra_embeddings():

226 cleanup_dist_env_and_memory(shutdown_ray=True)

tests/lora/test_peft_helper.py

+4 -7

@@ -38,8 +38,8 @@ ERROR_CASES = [

38 ]

41 -def test_peft_helper_pass(long_context_lora_files_16k_1, tmp_path):

42 - peft_helper = PEFTHelper.from_local_dir(long_context_lora_files_16k_1,

41 +def test_peft_helper_pass(sql_lora_files, tmp_path):

42 + peft_helper = PEFTHelper.from_local_dir(sql_lora_files,

43 max_position_embeddings=4096)

44 lora_config = LoRAConfig(max_lora_rank=16, max_cpu_loras=3, max_loras=2)

45 peft_helper.validate_legal(lora_config)

@@ -56,15 +56,12 @@ def test_peft_helper_pass(long_context_lora_files_16k_1, tmp_path):

56 "embed_tokens",

57 "lm_head",

58 ]

59 - assert peft_helper.context_length == 16384

59 assert peft_helper.vllm_max_position_embeddings == 4096

61 - assert peft_helper.vllm_long_context_scaling_factor == float(

62 - math.ceil(peft_helper.context_length /

63 - peft_helper.vllm_max_position_embeddings))

60 +

61 # test RSLoRA

62 rslora_config = dict(use_rslora=True)

63 test_dir = tmp_path / "test_rslora"

67 - shutil.copytree(long_context_lora_files_16k_1, test_dir)

64 + shutil.copytree(sql_lora_files, test_dir)

66 # Load and modify configuration

67 config_path = test_dir / "adapter_config.json"

vllm/config.py

+1 -13

@@ -3014,12 +3014,7 @@ class LoRAConfig:

3014 (added to the base model vocabulary)."""

3015 lora_vocab_padding_size: ClassVar[int] = current_platform\

3016 .get_lora_vocab_padding_size()

3017 - long_lora_scaling_factors: Optional[tuple[float, ...]] = None

3018 - """Specify multiple scaling factors (which can be different from base model

3019 - scaling factor - see eg. Long LoRA) to allow for multiple LoRA adapters

3020 - trained with those scaling factors to be used at the same time. If not

3021 - specified, only adapters trained with the base model scaling factor are

3022 - allowed."""

3017 +

3018 default_mm_loras: Optional[dict[str, str]] = None

3019 """Dictionary mapping specific modalities to LoRA model paths; this field

3020 is only applicable to multimodal models and should be leveraged when a

@@ -3052,7 +3047,6 @@ class LoRAConfig:

3047 factors.append(self.lora_dtype)

3048 factors.append(self.lora_extra_vocab_size)

3049 factors.append(self.lora_vocab_padding_size)

3055 - factors.append(self.long_lora_scaling_factors)

3050 factors.append(self.bias_enabled)

3051 hash_str = hashlib.md5(str(factors).encode(),

3052 usedforsecurity=False).hexdigest()

@@ -3091,11 +3085,6 @@ class LoRAConfig:

3085 elif isinstance(self.lora_dtype, str):

3086 self.lora_dtype = getattr(torch, self.lora_dtype)

3087

3094 - def verify_lora_support(self):

3095 - if self.long_lora_scaling_factors is not None and envs.VLLM_USE_V1:

3096 - raise ValueError(

3097 - "V1 LoRA does not support long LoRA, please use V0.")

3098 -

3088

3089 @config

3090 @dataclass(config=ConfigDict(arbitrary_types_allowed=True))

@@ -4564,7 +4553,6 @@ class VllmConfig:

4553 if self.lora_config is not None:

4554 self.lora_config.verify_with_cache_config(self.cache_config)

4555 self.lora_config.verify_with_model_config(self.model_config)

4567 - self.lora_config.verify_lora_support()

4556 if self.prompt_adapter_config is not None:

4557 self.prompt_adapter_config.verify_with_model_config(

4558 self.model_config)

vllm/engine/arg_utils.py

-5

@@ -358,8 +358,6 @@ class EngineArgs:

358 max_cpu_loras: Optional[int] = LoRAConfig.max_cpu_loras

359 lora_dtype: Optional[Union[str, torch.dtype]] = LoRAConfig.lora_dtype

360 lora_extra_vocab_size: int = LoRAConfig.lora_extra_vocab_size

361 - long_lora_scaling_factors: Optional[tuple[float, ...]] = \

362 - LoRAConfig.long_lora_scaling_factors

361 # PromptAdapter fields

362 enable_prompt_adapter: bool = False

363 max_prompt_adapters: int = PromptAdapterConfig.max_prompt_adapters

@@ -723,8 +721,6 @@ class EngineArgs:

721 "--lora-dtype",

722 **lora_kwargs["lora_dtype"],

723 )

726 - lora_group.add_argument("--long-lora-scaling-factors",

727 - **lora_kwargs["long_lora_scaling_factors"])

724 lora_group.add_argument("--max-cpu-loras",

725 **lora_kwargs["max_cpu_loras"])

726 lora_group.add_argument("--fully-sharded-loras",

@@ -1245,7 +1241,6 @@ class EngineArgs:

1241 default_mm_loras=self.default_mm_loras,

1242 fully_sharded_loras=self.fully_sharded_loras,

1243 lora_extra_vocab_size=self.lora_extra_vocab_size,

1248 - long_lora_scaling_factors=self.long_lora_scaling_factors,

1244 lora_dtype=self.lora_dtype,

1245 max_cpu_loras=self.max_cpu_loras if self.max_cpu_loras

1246 and self.max_cpu_loras > 0 else None) if self.enable_lora else None

vllm/lora/layers.py

-90

@@ -28,8 +28,6 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,

28 RowParallelLinear)

29 # yapf: enable

30 from vllm.model_executor.layers.logits_processor import LogitsProcessor

31 -from vllm.model_executor.layers.rotary_embedding import (

32 - LinearScalingRotaryEmbedding, RotaryEmbedding)

31 from vllm.model_executor.layers.vocab_parallel_embedding import (

32 VocabParallelEmbedding)

33 from vllm.platforms import current_platform

@@ -1193,91 +1191,3 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):

1191 ) -> bool:

1192 # Special handling for the LogitsProcessor.

1193 return False

1196 -

1197 -

1198 -class LinearScalingRotaryEmbeddingWithLoRA(BaseLayerWithLoRA):

1199 - """Implements RoPE-scaled embeddings with linear scaling for

1200 - multiple LoRA adapters with a specialized kernel.

1201 -

1202 - Replace LinearScalingRotaryEmbedding with MultiLinearScalingRotaryEmbedding

1203 - which can handle multi lora adapters in a specialized kernel.

1204 - """

1205 -

1206 - def __init__(self, base_layer: RotaryEmbedding) -> None:

1207 - super().__init__()

1208 - self.base_layer = base_layer

1209 -

1210 - @property

1211 - def scaling_factors(self):

1212 - return self.base_layer.scaling_factors

1213 -

1214 - @property

1215 - def rotary_dim(self):

1216 - return self.base_layer.rotary_dim

1217 -

1218 - def create_lora_weights(

1219 - self,

1220 - max_loras: int,

1221 - lora_config: LoRAConfig,

1222 - model_config: Optional[PretrainedConfig] = None,

1223 - ) -> None:

1224 - scaling_factors = (list(lora_config.long_lora_scaling_factors)

1225 - if lora_config.long_lora_scaling_factors else [])

1226 - base_scaling_factor = (self.base_layer.scaling_factor if isinstance(

1227 - self.base_layer, LinearScalingRotaryEmbedding) else 1.0)

1228 - scaling_factors = sorted(

1229 - list(set([base_scaling_factor] + scaling_factors)))

1230 - self.base_layer = LinearScalingRotaryEmbedding(

1231 - self.base_layer.head_size,

1232 - self.base_layer.rotary_dim,

1233 - self.base_layer.max_position_embeddings,

1234 - self.base_layer.base,

1235 - self.base_layer.is_neox_style,

1236 - scaling_factors,

1237 - self.base_layer.dtype,

1238 - )

1239 -

1240 - def reset_lora(self, index: int):

1241 - ...

1242 -

1243 - def set_lora(

1244 - self,

1245 - index: int,

1246 - lora_a: torch.Tensor,

1247 - lora_b: torch.Tensor,

1248 - embeddings_tensor: Optional[torch.Tensor],

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

1250 - ):

1251 - ...

1252 -

1253 - def forward(

1254 - self,

1255 - positions: torch.Tensor,

1256 - query: torch.Tensor,

1257 - key: torch.Tensor,

1258 - ) -> tuple[torch.Tensor, torch.Tensor]:

1259 - return self.base_layer(

1260 - positions,

1261 - query,

1262 - key,

1263 - offsets=self.punica_wrapper.long_lora_indices,

1264 - )

1265 -

1266 - @property

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

1268 - return self.base_layer.scaling_factor_to_offset

1269 -

1270 - @classmethod

1271 - def can_replace_layer(

1272 - cls,

1273 - source_layer: nn.Module,

1274 - lora_config: LoRAConfig,

1275 - packed_modules_list: list,

1276 - model_config: Optional[PretrainedConfig],

1277 - ) -> bool:

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

1279 - return (type(source_layer) is LinearScalingRotaryEmbedding

1280 - or type(source_layer) is RotaryEmbedding)

1281 -

1282 - def extra_repr(self) -> str:

1283 - return self.base_layer.extra_repr()

vllm/lora/models.py

+10 -70

@@ -4,7 +4,6 @@

4 import math

5 import os

6 from collections.abc import Sequence

7 -from dataclasses import dataclass, field

7 from typing import Any, Callable, Optional, Union

9 import regex as re

@@ -19,9 +18,7 @@ from vllm.adapter_commons.utils import (add_adapter, deactivate_adapter,

18 remove_adapter, set_adapter_mapping)

19 from vllm.config import LoRAConfig

20 from vllm.logger import init_logger

22 -from vllm.lora.layers import (BaseLayerWithLoRA,

23 - LinearScalingRotaryEmbeddingWithLoRA,

24 - LoRAMapping)

21 +from vllm.lora.layers import BaseLayerWithLoRA, LoRAMapping

22 from vllm.lora.lora import LoRALayerWeights, PackedLoRALayerWeights

23 from vllm.lora.peft_helper import PEFTHelper

24 from vllm.lora.punica_wrapper import get_punica_wrapper

@@ -43,18 +40,6 @@ logger = init_logger(__name__)

40 _GLOBAL_LORA_ID = 0

46 -@dataclass

47 -class LongContextLoRAContext:

48 - """Context for lora adapters that support long context."""

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

50 - scaling_factors: list[float]

51 - # dimension to apply rotary embedding.

52 - rot_dim: int

53 - # offsets to the sin_cos_cache for each lora_id loaded.

54 - # This value is dynamically modified.

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

56 -

57 -

43 def get_lora_id():

44 global _GLOBAL_LORA_ID

45 _GLOBAL_LORA_ID += 1

@@ -80,20 +65,16 @@ class LoRAModel(AdapterModel):

65 lora_model_id: int,

66 rank: int,

67 loras: dict[str, LoRALayerWeights],

83 - scaling_factor: Optional[float] = None,

68 ) -> None:

69 """

70 Args:

71 lora_model_id: The integer id for the lora model.

72 rank: lora rank.

73 loras: module name -> weights for lora-replaced layers.

90 - scaling_factor: Scaling factor to support long context lora model.

91 - None if the lora is not tuned for long context support.

74 +

75 """

76 self.id = lora_model_id

94 - # Scaling factor for long context lora model. None if it is not

95 - # fine tuned for the long context.

96 - self.scaling_factor = scaling_factor

77 +

78 assert (

79 lora_model_id

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

@@ -192,10 +173,7 @@ class LoRAModel(AdapterModel):

173 for lora in loras.values():

174 lora.optimize()

175

195 - return cls(lora_model_id,

196 - peft_helper.r,

197 - loras,

198 - scaling_factor=peft_helper.vllm_long_context_scaling_factor)

176 + return cls(lora_model_id, peft_helper.r, loras)

177

178 @classmethod

179 def from_local_checkpoint(

@@ -360,24 +338,17 @@ class LoRAModelManager(AdapterModelManager):

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

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

340 self.vocab_size = vocab_size

363 - self.long_lora_context: Optional[LongContextLoRAContext] = None

341 self.punica_wrapper = get_punica_wrapper(

342 max_num_batched_tokens,

343 max_batches=self.max_num_seqs,

344 device=self.device,

345 max_loras=self.lora_config.max_loras)

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

370 - # Used for long context lora.

371 - self.scaling_factor_to_offset: dict[float, int] = {}

346 +

347 super().__init__(model)

348

349 self.supported_lora_modules = get_supported_lora_modules(self.model)

350 assert self.supported_lora_modules, "No supported LoRA modules found in"

351 f" {self.model.__class__.__name__}."

377 - if lora_config.long_lora_scaling_factors:

378 - # We need to replace rotary emb layer to do batch computation

379 - # for long lora.

380 - self.supported_lora_modules.append("rotary_emb")

352

353 self.packed_modules_mapping = get_packed_modules_mapping(self.model)

354 # Used to indicate whether the model is a multimodal model

@@ -454,25 +425,9 @@ class LoRAModelManager(AdapterModelManager):

425 except ValueError:

426 pass

427

457 - def _set_long_lora_context(self, lora: LoRAModel):

458 - if self.long_lora_context is None:

459 - return

460 -

461 - if lora.scaling_factor is None:

462 - return

463 -

464 - if (lora.scaling_factor not in self.scaling_factor_to_offset):

465 - raise ValueError(f"Long LoRA scaling factor {lora.scaling_factor}"

466 - " has not been initialized.")

467 -

468 - offsets = self.scaling_factor_to_offset.get(lora.scaling_factor)

469 - if offsets:

470 - self.long_lora_context.offsets_by_lora_id[lora.id] = offsets

471 -

428 def _add_adapter(self, lora: LoRAModel):

429 self._create_merged_loras_inplace(lora)

430 self._registered_adapters[lora.id] = lora

475 - self._set_long_lora_context(lora)

431

432 def pin_adapter(self, lora_id: int) -> bool:

433 """Pin a LoRAModel in the manager cache."""

@@ -488,7 +443,6 @@ class LoRAModelManager(AdapterModelManager):

443 self.lora_slots + 1,

444 self.vocab_size,

445 self.lora_config.lora_extra_vocab_size,

491 - self.long_lora_context,

446 )

447

448 def remove_all_adapters(self):

@@ -528,13 +482,6 @@ class LoRAModelManager(AdapterModelManager):

482 from_layer(module, self.lora_slots, self.lora_config,

483 packed_moduled_lst, self.model.config))

484

531 - # LinearScalingRotaryEmbeddingWithLoRA is used to handle

532 - # long context lora. Register relevant metadata.

533 - if isinstance(new_module, LinearScalingRotaryEmbeddingWithLoRA):

534 - self.long_lora_context = LongContextLoRAContext(

535 - new_module.scaling_factors, new_module.rotary_dim)

536 - self.scaling_factor_to_offset = \

537 - new_module.scaling_factor_to_offset

485 # (yard1): TODO make this more robust

486 if "lm_head" in module_name:

487 logits_processor_module_name = 'logits_processor'

@@ -574,15 +521,13 @@ class LoRAModelManager(AdapterModelManager):

521 self,

522 lora_id: int,

523 rank: int,

577 - scaling_factor: Optional[float],

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

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

580 - model = LoRAModel(lora_id, rank, {}, scaling_factor)

526 + model = LoRAModel(lora_id, rank, {})

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

528 bias_enabled = self.lora_config.bias_enabled

529 if (not self._match_target_modules(module_name)

530 or not isinstance(module, BaseLayerWithLoRA)

585 - or isinstance(module, LinearScalingRotaryEmbeddingWithLoRA)

531 or self._filter_unsupported_mm_module(module_name)):

532 continue

533 parts = module_name.split(".")

@@ -723,11 +668,8 @@ class LoRAModelManager(AdapterModelManager):

668 self._deactivate_adapter)

669

670 def add_adapter(self, adapter: LoRAModel) -> bool:

726 - logger.debug(

727 - "Adding lora. Model id: %d, "

728 - "int id: %d, "

729 - "scaling factor: %s", adapter.id, adapter.id,

730 - adapter.scaling_factor)

671 + logger.debug("Adding lora. Model id: %d, "

672 + "int id: %d", adapter.id, adapter.id)

673 return add_adapter(adapter, self._registered_adapters, self.capacity,

674 self._add_adapter)

675

@@ -772,10 +714,8 @@ class LRUCacheLoRAModelManager(LoRAModelManager):

714

715 def add_adapter(self, lora: LoRAModel) -> bool:

716 """Add a LoRAModel to the manager."""

775 - logger.debug(

776 - "Adding lora. Model id: %d, "

777 - "int id: %d, "

778 - "scaling factor: %s", lora.id, lora.id, lora.scaling_factor)

717 + logger.debug("Adding lora. Model id: %d, "

718 + "int id: %d", lora.id, lora.id)

719 if lora.id not in self._registered_adapters:

720 self._add_adapter(lora)

721 was_added = True

vllm/lora/peft_helper.py

-9

@@ -35,12 +35,9 @@ class PEFTHelper:

35 use_rslora: bool = field(default=False)

36 # True to use Weight-Decomposed Low-Rank Adaptation (DoRA, see: https://arxiv.org/abs/2402.09353)

37 use_dora: bool = field(default=False)

38 - # long context lora field

39 - context_length: int = field(default=0)

38 # Extra vllm field, start with 'vllm_' to avoid conflict

39 vllm_lora_scaling_factor: float = field(default=1.0)

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

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

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

43 """

@@ -59,12 +56,6 @@ class PEFTHelper:

56 self.vllm_lora_scaling_factor = self.lora_alpha / math.sqrt(self.r)

57 else:

58 self.vllm_lora_scaling_factor = self.lora_alpha / self.r

62 - if self.context_length:

63 - if self.vllm_max_position_embeddings is None:

64 - self.vllm_max_position_embeddings = self.context_length

65 - self.vllm_long_context_scaling_factor = float(

66 - math.ceil(self.context_length /

67 - self.vllm_max_position_embeddings))

60 @classmethod

61 def from_dict(cls, config_dict: dict) -> "PEFTHelper":

vllm/lora/punica_wrapper/punica_base.py

+9 -36

@@ -17,7 +17,6 @@ from .utils import compute_meta, convert_mapping

17 if TYPE_CHECKING:

18 # avoid circuit import

19 from vllm.lora.layers import LoRAMapping

20 - from vllm.lora.models import LongContextLoRAContext

22 class PunicaWrapperABC(ABC):

@@ -33,7 +32,6 @@ class PunicaWrapperABC(ABC):

32 max_loras: int,

33 vocab_size: int,

34 extra_vocab_size: int,

36 - long_lora_context: Optional["LongContextLoRAContext"] = None,

35 **kwargs,

36 ) -> None:

37 """

@@ -144,14 +142,11 @@ class PunicaWrapperBase(PunicaWrapperABC):

142 max_num_batched_tokens,

143 dtype=torch.long,

144 device=device)

147 - self._long_lora_indices = torch.empty(max_num_batched_tokens,

148 - dtype=torch.long,

149 - device=device)

145

151 - # 5 is the number of indices tensors.

146 + # 4 is the number of indices tensors.

147 # base_indices, sampler_indices, sampler_indices_padded,

153 - # embeddings_indices,long_lora_indices

154 - self.indices_len: list[Optional[int]] = [None] * 5

148 + # embeddings_indices

149 + self.indices_len: list[Optional[int]] = [None] * 4

150 # these attributes are the information required for sgmv kernel

151 self._seq_start_locs = torch.empty(max_batches,

152 dtype=torch.long,

@@ -176,14 +171,12 @@ class PunicaWrapperBase(PunicaWrapperABC):

171 max_loras: int,

172 vocab_size: int,

173 extra_vocab_size: int,

179 - long_lora_context: Optional["LongContextLoRAContext"] = None,

174 ):

175 (

176 base_indices,

177 sampler_indices,

178 sampler_indices_padded,

179 embeddings_indices,

186 - long_lora_offsets_tensor,

180 indices_len,

181 ) = convert_mapping(

182 mapping,

@@ -192,7 +185,6 @@ class PunicaWrapperBase(PunicaWrapperABC):

185 vocab_size,

186 extra_vocab_size,

187 self.device,

195 - long_lora_context,

188 )

189 self._token_lora_indices[:base_indices.shape[0]].copy_(base_indices)

190 self._sampler_indices[:sampler_indices.shape[0]].copy_(sampler_indices)

@@ -201,11 +193,7 @@ class PunicaWrapperBase(PunicaWrapperABC):

193 self._embeddings_indices[:embeddings_indices.

194 shape[0], :embeddings_indices.shape[1]].copy_(

195 embeddings_indices)

204 - if long_lora_offsets_tensor is not None:

205 - self._long_lora_indices[:long_lora_offsets_tensor.shape[0]].copy_(

206 - long_lora_offsets_tensor)

207 - else:

208 - self._long_lora_indices.zero_()

196 +

197 self.indices_len[:] = indices_len

198

199 def _update_prefill_metadata(self,

@@ -312,28 +300,13 @@ class PunicaWrapperBase(PunicaWrapperABC):

300 embeddings_indices_len = self.indices_len[3]

301 return self._embeddings_indices[:, :embeddings_indices_len]

302

315 - @property

316 - def long_lora_indices(self) -> torch.Tensor:

317 - """

318 - This property provides access to the indices used for long context

319 - lora, specifically for LinearScalingRotaryEmbeddingWithLoRA.

320 - """

321 - long_lora_len = self.indices_len[4]

322 - return self._long_lora_indices[:long_lora_len]

323 -

324 - def update_metadata(

325 - self,

326 - mapping: "LoRAMapping",

327 - lora_index_to_id: list[Optional[int]],

328 - max_loras: int,

329 - vocab_size: int,

330 - extra_vocab_size: int,

331 - long_lora_context: Optional["LongContextLoRAContext"] = None,

332 - **kwargs):

303 + def update_metadata(self, mapping: "LoRAMapping",

304 + lora_index_to_id: list[Optional[int]], max_loras: int,

305 + vocab_size: int, extra_vocab_size: int, **kwargs):

306

307 self._update_base_metadata(mapping, lora_index_to_id, max_loras,

335 - vocab_size, extra_vocab_size,

336 - long_lora_context)

308 + vocab_size, extra_vocab_size)

309 +

310 if mapping.is_prefill:

311 # Update metadata required for prefill-related operators.

312 self._update_prefill_metadata(self.token_lora_indices)

vllm/lora/punica_wrapper/punica_gpu.py

+5 -16

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

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

8 """

10 -from typing import TYPE_CHECKING, Optional, Union, final

10 +from typing import Optional, Union, final

12 import torch

@@ -21,10 +21,6 @@ if HAS_TRITON:

22 from .punica_base import PunicaWrapperBase

24 -if TYPE_CHECKING:

25 - # avoid circuit import

26 - from vllm.lora.models import LongContextLoRAContext

27 -

25 @final

26 class PunicaWrapperGPU(PunicaWrapperBase):

@@ -55,20 +51,13 @@ class PunicaWrapperGPU(PunicaWrapperBase):

51 max_num_prompts,

52 device=device)

58 - def update_metadata(

59 - self,

60 - mapping: LoRAMapping,

61 - lora_index_to_id: list[Optional[int]],

62 - max_loras: int,

63 - vocab_size: int,

64 - extra_vocab_size: int,

65 - long_lora_context: Optional["LongContextLoRAContext"] = None,

66 - **kwargs):

54 + def update_metadata(self, mapping: LoRAMapping,

55 + lora_index_to_id: list[Optional[int]], max_loras: int,

56 + vocab_size: int, extra_vocab_size: int, **kwargs):

58 self.is_prefill = mapping.is_prefill

59 self._update_base_metadata(mapping, lora_index_to_id, max_loras,

70 - vocab_size, extra_vocab_size,

71 - long_lora_context)

60 + vocab_size, extra_vocab_size)

62 # Prepare cuda kernel metadata tensors

63 self.token_mapping_meta.prepare_tensors(self.token_lora_indices)

vllm/lora/punica_wrapper/punica_tpu.py

-14

@@ -14,7 +14,6 @@ from vllm.lora.punica_wrapper.utils import convert_mapping

14 if TYPE_CHECKING:

15 # avoid circuit import

16 from vllm.lora.layers import LoRAMapping

17 - from vllm.lora.models import LongContextLoRAContext

18 from .punica_base import PunicaWrapperBase

@@ -45,7 +44,6 @@ class PunicaWrapperTPU(PunicaWrapperBase):

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)

48 - torch.ops.xla.dynamo_set_buffer_donor_(self._long_lora_indices, True)

47 torch.ops.xla.dynamo_set_buffer_donor_(self._lora_indices_per_batch,

48 True)

@@ -323,7 +321,6 @@ class PunicaWrapperTPU(PunicaWrapperBase):

321 max_loras: int,

322 vocab_size: int,

323 extra_vocab_size: int,

326 - long_lora_context: Optional["LongContextLoRAContext"] = None,

324 ):

325 # Make sure we don't accidentally collect outside operations

326 xm.mark_step()

@@ -339,7 +336,6 @@ class PunicaWrapperTPU(PunicaWrapperBase):

336 sampler_indices,

337 sampler_indices_padded,

338 embeddings_indices,

342 - long_lora_offsets_tensor,

339 indices_len,

340 ) = convert_mapping(

341 mapping,

@@ -348,7 +344,6 @@ class PunicaWrapperTPU(PunicaWrapperBase):

344 vocab_size,

345 extra_vocab_size,

346 "cpu",

351 - long_lora_context,

347 )

348 self._token_lora_indices = self._pad_to_shape(

349 base_indices, self._token_lora_indices.shape,

@@ -362,15 +357,6 @@ class PunicaWrapperTPU(PunicaWrapperBase):

357 self._embeddings_indices = self._pad_to_shape(

358 embeddings_indices, self._embeddings_indices.shape,

359 dims=2).to(self.device)

365 - if long_lora_offsets_tensor is not None:

366 - self._long_lora_indices = self._pad_to_shape(

367 - long_lora_offsets_tensor,

368 - self._long_lora_indices.shape,

369 - dims=1).to(self.device)

370 - else:

371 - zeroed = torch.zeros_like(self._long_lora_indices.cpu(),

372 - dtype=torch.int32)

373 - self._long_lora_indices = zeroed.to(self.device)

360 self.indices_len[:] = indices_len

361

362 def _update_prefill_metadata(self,

vllm/lora/punica_wrapper/utils.py

+5 -33

@@ -8,7 +8,6 @@ import torch

8 if TYPE_CHECKING:

9 # avoid circuit import

10 from vllm.lora.layers import LoRAMapping

11 - from vllm.lora.models import LongContextLoRAContext

13 def compute_meta(

@@ -49,9 +48,7 @@ def convert_mapping(

48 vocab_size: int,

49 extra_vocab_size: int,

50 device: torch.device,

52 - long_lora_context: Optional["LongContextLoRAContext"] = None,

53 -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,

54 - Optional[torch.Tensor], list[int]]:

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

52 """Converts LoRAMapping to index tensors.

54 Args:

@@ -60,7 +57,6 @@ def convert_mapping(

57 max_loras: Maximum number of LoRAs.

58 vocab_size: Model vocab size.

59 extra_vocab_size: Extra vocab size each LoRA can have.

63 - long_lora_context: Passed if there are long context lora in a batch.

61 Returns:

62 A tuple of tensors:

@@ -78,21 +74,14 @@ def convert_mapping(

74 requests to embedding indices. First row is for embeddings

75 added by the LoRAs, second row is for the LoRA.lora_a

76 embeddings.

81 - long_lora_indices: Tensor of shape [batch_size] mapping

82 - requests to RoPE offsets and rot dims for long LoRAs.

83 - None if long context lora doesn't exist.

77 indices_len: List of lengths of the above tensors. It contains

78 (base_indices, sampler_indices, sampler_indices_padded,

86 - embeddings_indices, long_lora_indices).

79 + embeddings_indices).

80 """

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

82 embedding_indices = index_mapping_indices.copy()

83 lora_indices = index_mapping_indices.copy()

91 - long_lora_offsets: Optional[torch.Tensor] = None

92 - if long_lora_context:

93 - long_lora_offsets = torch.zeros(len(index_mapping_indices),

94 - device=device,

95 - dtype=torch.long)

84 +

85 prompt_mapping: list[int] = [

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

87 for x in mapping.prompt_mapping

@@ -104,20 +93,13 @@ def convert_mapping(

93 if index_mapping_indices[i] > 0 else -1)

94 embedding_indices[i] = lora_idx if index_mapping_indices[i] > 0 else 0

95 lora_indices[i] = lora_idx

107 - if long_lora_context:

108 - assert long_lora_offsets is not None

109 - lora_offset: int = long_lora_context.offsets_by_lora_id.get(

110 - index_mapping_indices[i], 0)

111 - long_lora_offsets[i] = lora_offset

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

98 index_mapping_indices,

99 lora_indices,

100 embedding_indices,

101 ]

118 - if long_lora_context:

119 - assert long_lora_offsets is not None

120 - indices_list.append(long_lora_offsets)

102 +

103 indices = torch.tensor(indices_list, dtype=torch.long, device=device)

104 prompt_mapping_tensor = torch.tensor(prompt_mapping,

105 dtype=torch.long,

@@ -136,11 +118,7 @@ def convert_mapping(

118 sampler_indices_padded = torch.arange(

119 0, len(sampler_indices_padded), device=device, dtype=torch.long) + (

120 sampler_indices_padded * len(sampler_indices_padded))

139 - long_lora_indices = None

140 - long_lora_indices_len: Optional[int] = None

141 - if long_lora_context:

142 - long_lora_indices = indices[3]

143 - long_lora_indices_len = long_lora_indices.shape[-1]

121 +

122 # Contain length of indices tensors. Used to index into each tensor.

123 indices_len = [

124 base_indices.shape[-1],

@@ -148,17 +126,11 @@ def convert_mapping(

126 sampler_indices_padded.shape[-1],

127 embeddings_indices.shape[-1],

128 ]

151 - if long_lora_indices_len is not None:

152 - indices_len.append(long_lora_indices_len)

153 - else:

154 - # If long_lora doesn't exist,append None

155 - indices_len.append(None)

129

130 return (

131 base_indices,

132 sampler_indices,

133 sampler_indices_padded,

134 embeddings_indices,

162 - long_lora_indices,

135 indices_len,

136 )

vllm/lora/utils.py

-2

@@ -22,7 +22,6 @@ from vllm.lora.fully_sharded_layers import (

22 # yapf conflicts with isort for this block

23 # yapf: disable

24 from vllm.lora.layers import (BaseLayerWithLoRA, ColumnParallelLinearWithLoRA,

25 - LinearScalingRotaryEmbeddingWithLoRA,

25 LogitsProcessorWithLoRA,

26 MergedColumnParallelLinearWithLoRA,

27 MergedQKVParallelLinearWithLoRA,

@@ -56,7 +55,6 @@ _all_lora_classes: set[type[BaseLayerWithLoRA]] = {

55 MergedColumnParallelLinearWithShardedLoRA,

56 MergedQKVParallelLinearWithShardedLoRA,

57 RowParallelLinearWithShardedLoRA,

59 - LinearScalingRotaryEmbeddingWithLoRA,

58 }

vllm/lora/worker_manager.py

+1 -1

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

154 lora_request.lora_int_id)

155 else:

156 dummy_lora = self._adapter_manager.create_dummy_lora(

157 - lora_request.lora_int_id, rank, 1, self.embedding_modules)

157 + lora_request.lora_int_id, rank, self.embedding_modules)

158 if self._cached_dummy_lora is None:

159 self._cached_dummy_lora = dummy_lora

160 return self._adapter_manager.add_adapter(dummy_lora)