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, 18:30 UTC
[Bugfix] LoRA V0 - Fix case where `max_num_seqs` is between cudagraph capture sizes (#15308)
Signed-off-by: Varun Sundar Rabindranath <varun@neuralmagic.com> Co-authored-by: Varun Sundar Rabindranath <varun@neuralmagic.com>
Varun Sundar Rabindranath committed
Mar 22, 2025 at 05:03 UTC
8a8b30eac1ab746c48b5a1e3eaa4e91941177bf8
2 files changed
+17
-7
tests/lora/test_llama_tp.py
+8
-6
@@ -84,12 +84,14 @@ def v1(run_with_both_engines_lora):
84
@create_new_process_for_each_test()
85
def test_llama_lora(sql_lora_files):
86
87
- llm = vllm.LLM(MODEL_PATH,
88
- enable_lora=True,
89
- max_num_seqs=16,
90
- max_loras=4,
91
- tensor_parallel_size=1,
92
- enable_chunked_prefill=True)
87
+ llm = vllm.LLM(
88
+ MODEL_PATH,
89
+ enable_lora=True,
90
+ # also test odd max_num_seqs
91
+ max_num_seqs=13,
92
+ max_loras=4,
93
+ tensor_parallel_size=1,
94
+ enable_chunked_prefill=True)
95
generate_and_test(llm, sql_lora_files)
96
97
vllm/lora/punica_wrapper/punica_gpu.py
+9
-1
@@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, List, Optional, Tuple, Union, final
10
11
import torch
12
13
+import vllm.envs as envs
14
from vllm.lora.layers import LoRAMapping
15
from vllm.triton_utils import HAS_TRITON
16
@@ -42,8 +43,15 @@ class PunicaWrapperGPU(PunicaWrapperBase):
43
self.token_mapping_meta = LoRAKernelMeta.make(self.max_loras,
44
max_num_batched_tokens,
45
device=device)
46
+
47
+ # When cudagraph capture size is greater than max_num_seqs (max_batches,
48
+ # here), V0 captures the graph as if max_num_seqs is set to
49
+ # the capture size.
50
+ # V1 doesn't have this problem and always respects max_num_seqs.
51
+ max_num_prompts = (max_batches
52
+ if envs.VLLM_USE_V1 else max_num_batched_tokens)
53
self.prompt_mapping_meta = LoRAKernelMeta.make(self.max_loras,
46
- max_batches,
54
+ max_num_prompts,
55
device=device)
56
57
def update_metadata(