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, 19:45 UTC
[Kernel] LoRA - Enable CUDAGraphs for V1 (#14626)
Signed-off-by: Varun Sundar Rabindranath <varun@neuralmagic.com> Co-authored-by: Varun Sundar Rabindranath <varun@neuralmagic.com>
Varun Sundar Rabindranath committed
Mar 13, 2025 at 23:42 UTC
0b1cfa61806d8d0f3e3d8d6f8303d0fe644f329e
4 files changed
+35
-14
tests/lora/test_worker.py
+1
@@ -52,6 +52,7 @@ def test_worker_apply_lora(sql_lora_files):
52
seed=0,
53
dtype="float16",
54
revision=None,
55
+ enforce_eager=True,
56
),
57
load_config=LoadConfig(
58
download_dir=None,
vllm/config.py
+19
-6
@@ -2287,9 +2287,14 @@ class LoRAConfig:
2287
excluding anything before input ids/embeddings and after
2288
the final hidden states.
2289
"""
2290
- # no factors to consider.
2291
- # LoRA is not compatible with `torch.compile` .
2290
factors: list[Any] = []
2291
+ factors.append(self.max_lora_rank)
2292
+ factors.append(self.max_loras)
2293
+ factors.append(self.fully_sharded_loras)
2294
+ factors.append(self.lora_dtype)
2295
+ factors.append(self.lora_extra_vocab_size)
2296
+ factors.append(self.long_lora_scaling_factors)
2297
+ factors.append(self.bias_enabled)
2298
hash_str = hashlib.md5(str(factors).encode()).hexdigest()
2299
return hash_str
2300
@@ -3303,6 +3308,11 @@ class VllmConfig:
3308
vllm_factors.append("None")
3309
if self.lora_config:
3310
vllm_factors.append(self.lora_config.compute_hash())
3311
+ # LoRA creates static buffers based on max_num_batched_tokens.
3312
+ # The tensor sizes and strides get captured in the torch.compile
3313
+ # graph explicitly.
3314
+ vllm_factors.append(
3315
+ str(self.scheduler_config.max_num_batched_tokens))
3316
else:
3317
vllm_factors.append("None")
3318
if self.speculative_config:
@@ -3453,12 +3463,15 @@ class VllmConfig:
3463
" Disabling `torch.compile`.")
3464
self.compilation_config.level = CompilationLevel.NO_COMPILATION
3465
3456
- if self.lora_config is not None and self.compilation_config.level !=\
3457
- CompilationLevel.NO_COMPILATION:
3458
- logger.warning("LoRA is not supported with `torch.compile` yet. "
3459
- "Disabling `torch.compile`.")
3466
+ if ((not envs.VLLM_USE_V1) and self.lora_config is not None
3467
+ and self.compilation_config.level
3468
+ != CompilationLevel.NO_COMPILATION):
3469
+ logger.warning(
3470
+ "LoRA for V0 is not supported with `torch.compile` yet. "
3471
+ "Disabling `torch.compile`.")
3472
self.compilation_config.level = CompilationLevel.NO_COMPILATION
3473
3474
+
3475
if self.model_config and self.model_config.use_mla and \
3476
not (current_platform.is_cuda() or current_platform.is_rocm()):
3477
logger.info(
vllm/lora/layers.py
+9
-6
@@ -237,16 +237,19 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
237
self.embeddings_weights[:embeddings.shape[0]].copy_(embeddings)
238
239
def forward(self, x: torch.Tensor) -> torch.Tensor:
240
- added_tokens_mask = x > self.base_layer.org_vocab_size - 1
241
- embeddings_indices = self.punica_wrapper.embeddings_indices
242
- indices = embeddings_indices[1].view_as(x)
240
+ added_tokens_mask = torch.where(x > self.base_layer.org_vocab_size - 1,
241
+ 1, 0)
242
+ embeddings_indices = torch.narrow(
243
+ self.punica_wrapper._embeddings_indices, 1, 0, x.size(0))
244
+
245
+ indices = embeddings_indices[1]
246
full_lora_a_embeddings = F.embedding(
247
x + indices,
248
self.lora_a_stacked_2d,
249
)
247
- indices = embeddings_indices[0].view_as(x)
248
- full_output = self.base_layer.forward(
249
- x.add_(indices * added_tokens_mask))
250
+ indices = embeddings_indices[0]
251
+ full_output = self.base_layer.forward(x +
252
+ (indices * added_tokens_mask))
253
254
full_output_org = full_output
255
if full_output.ndim == 3:
vllm/lora/punica_wrapper/punica_gpu.py
+6
-2
@@ -254,7 +254,9 @@ class PunicaWrapperGPU(PunicaWrapperBase, V1KernelMixin):
254
y_org = y
255
y = y.view(-1, y.shape[-1])
256
if lora_bias_stacked is not None:
257
- self._apply_bias(self.token_lora_indices, y, output_slices,
257
+ token_lora_indices = torch.narrow(self._token_lora_indices, 0, 0,
258
+ y.size(0))
259
+ self._apply_bias(token_lora_indices, y, output_slices,
260
lora_bias_stacked)
261
262
if env.VLLM_USE_V1:
@@ -365,7 +367,9 @@ class PunicaWrapperGPU(PunicaWrapperBase, V1KernelMixin):
367
assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices)
368
if lora_bias_stacked is not None:
369
assert len(lora_bias_stacked) == len(output_slices)
368
- y = self._apply_bias(self.token_lora_indices, y, output_slices,
370
+ token_lora_indices = torch.narrow(self._token_lora_indices, 0, 0,
371
+ y.size(0))
372
+ y = self._apply_bias(token_lora_indices, y, output_slices,
373
lora_bias_stacked)
374
375
if buffer is None: