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 14, 05:30 UTC
[LoRA] Much faster startup when LoRA is enabled (#23777)
Signed-off-by: Andy Lo <andy@mistral.ai> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
Andy Lo committed
Aug 30, 2025 at 16:37 UTC
038e9be4eb7a63189c8980845d80cb96957b9919
3 files changed
+33
-13
vllm/v1/worker/gpu_model_runner.py
+8
-3
@@ -2213,6 +2213,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2213
uniform_decode: bool = False,
2214
skip_eplb: bool = False,
2215
is_profile: bool = False,
2216
+ remove_lora: bool = True,
2217
) -> tuple[torch.Tensor, torch.Tensor]:
2218
"""
2219
Run a dummy forward pass to warm up/profile run or capture the
@@ -2230,6 +2231,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2231
uniform_decode: If True, the batch is a uniform decode batch.
2232
skip_eplb: If True, skip EPLB state update.
2233
is_profile: If True, this is a profile run.
2234
+ remove_lora: If False, dummy LoRAs are not destroyed after the run
2235
"""
2236
assert cudagraph_runtime_mode in {
2237
CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL
@@ -2317,7 +2319,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2319
attn_metadata[layer_name] = attn_metadata_i
2320
2321
with self.maybe_dummy_run_with_lora(self.lora_config,
2320
- num_scheduled_tokens):
2322
+ num_scheduled_tokens, remove_lora):
2323
if self.supports_mm_inputs:
2324
input_ids = None
2325
inputs_embeds = self.inputs_embeds[:num_tokens]
@@ -2708,11 +2710,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2710
cudagraph_runtime_mode=CUDAGraphMode.NONE,
2711
force_attention=force_attention,
2712
uniform_decode=uniform_decode,
2711
- skip_eplb=True)
2713
+ skip_eplb=True,
2714
+ remove_lora=False)
2715
self._dummy_run(num_tokens,
2716
cudagraph_runtime_mode=cudagraph_runtime_mode,
2717
uniform_decode=uniform_decode,
2715
- skip_eplb=True)
2718
+ skip_eplb=True,
2719
+ remove_lora=False)
2720
+ self.maybe_remove_all_loras(self.lora_config)
2721
2722
def initialize_attn_backend(self, kv_cache_config: KVCacheConfig) -> None:
2723
"""
vllm/v1/worker/gpu_worker.py
+4
-1
@@ -308,7 +308,10 @@ class Worker(WorkerBase):
308
# We skip EPLB here since we don't want to record dummy metrics
309
for size in sorted(warmup_sizes, reverse=True):
310
logger.info("Compile and warming up model for size %d", size)
311
- self.model_runner._dummy_run(size, skip_eplb=True)
311
+ self.model_runner._dummy_run(size,
312
+ skip_eplb=True,
313
+ remove_lora=False)
314
+ self.model_runner.maybe_remove_all_loras(self.model_runner.lora_config)
315
316
# Warmup and tune the kernels used during model execution before
317
# cuda graph capture.
vllm/v1/worker/lora_model_runner_mixin.py
+21
-9
@@ -5,7 +5,7 @@ Define LoRA functionality mixin for model runners.
5
"""
6
7
from contextlib import contextmanager
8
-from typing import Union
8
+from typing import Optional, Union
9
10
import numpy as np
11
import torch
@@ -87,7 +87,9 @@ class LoRAModelRunnerMixin:
87
lora_requests)
88
89
@contextmanager
90
- def maybe_setup_dummy_loras(self, lora_config):
90
+ def maybe_setup_dummy_loras(self,
91
+ lora_config: Optional[LoRAConfig],
92
+ remove_lora: bool = True):
93
if lora_config is None:
94
yield
95
else:
@@ -114,10 +116,11 @@ class LoRAModelRunnerMixin:
116
yield
117
118
# __exit__ code
117
- self.lora_manager.remove_all_adapters()
119
+ if remove_lora:
120
+ self.lora_manager.remove_all_adapters()
121
122
@contextmanager
120
- def maybe_select_dummy_loras(self, lora_config: LoRAConfig,
123
+ def maybe_select_dummy_loras(self, lora_config: Optional[LoRAConfig],
124
num_scheduled_tokens: np.ndarray):
125
if lora_config is None:
126
yield
@@ -151,13 +154,22 @@ class LoRAModelRunnerMixin:
154
yield
155
156
@contextmanager
154
- def maybe_dummy_run_with_lora(self, lora_config: LoRAConfig,
155
- num_scheduled_tokens: np.ndarray):
156
- with self.maybe_setup_dummy_loras(
157
- lora_config), self.maybe_select_dummy_loras(
158
- lora_config, num_scheduled_tokens):
157
+ def maybe_dummy_run_with_lora(self,
158
+ lora_config: Optional[LoRAConfig],
159
+ num_scheduled_tokens: np.ndarray,
160
+ remove_lora: bool = True):
161
+ with (
162
+ self.maybe_setup_dummy_loras(lora_config, remove_lora),
163
+ self.maybe_select_dummy_loras(lora_config,
164
+ num_scheduled_tokens),
165
+ ):
166
yield
167
168
+ def maybe_remove_all_loras(self, lora_config: Optional[LoRAConfig]):
169
+ if lora_config is None:
170
+ return
171
+ self.lora_manager.remove_all_adapters()
172
+
173
def add_lora(self, lora_request: LoRARequest) -> bool:
174
if not self.lora_manager:
175
raise RuntimeError("LoRA is not enabled.")