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 10, 20:15 UTC
Add Loraconfig parameter to get_punica_wrapper function (#31408)
Signed-off-by: ZT-AIA <1028681969@qq.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
ZT-AIA committed
Dec 30, 2025 at 14:27 UTC
f84bf7d79b885956747f2258ddeae380da492b3c
3 files changed
+16
-15
tests/lora/test_layers.py
+10
-10
@@ -261,11 +261,11 @@ def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None:
261
262
torch.set_default_device(device)
263
max_loras = 8
264
- punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)
265
- assert check_punica_wrapper(punica_wrapper)
264
lora_config = LoRAConfig(
265
max_loras=max_loras, max_lora_rank=8, lora_dtype=torch.float16
266
)
267
+ punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
268
+ assert check_punica_wrapper(punica_wrapper)
269
270
def create_random_embedding_layer():
271
embedding = VocabParallelEmbedding(vocab_size, 256)
@@ -360,11 +360,11 @@ def test_lm_head_logits_processor(
360
361
torch.set_default_device(device)
362
max_loras = 8
363
- punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)
364
- assert check_punica_wrapper(punica_wrapper)
363
lora_config = LoRAConfig(
364
max_loras=max_loras, max_lora_rank=8, lora_dtype=torch.float16
365
)
366
+ punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
367
+ assert check_punica_wrapper(punica_wrapper)
368
369
def _pretest():
370
linear = ParallelLMHead(
@@ -480,13 +480,13 @@ def test_linear_replicated(
480
481
max_loras = 8
482
torch.set_default_device(device)
483
- punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)
484
- assert check_punica_wrapper(punica_wrapper)
483
lora_config = LoRAConfig(
484
max_loras=max_loras,
485
max_lora_rank=8,
486
lora_dtype=torch.float16,
487
)
488
+ punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
489
+ assert check_punica_wrapper(punica_wrapper)
490
491
def create_random_linear_replicated_layer():
492
linear = ReplicatedLinear(4096, 4096, bias=False, params_dtype=torch.float16)
@@ -587,14 +587,14 @@ def test_linear_parallel(
587
588
max_loras = 8
589
torch.set_default_device(device)
590
- punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)
591
- assert check_punica_wrapper(punica_wrapper)
590
lora_config = LoRAConfig(
591
max_loras=max_loras,
592
max_lora_rank=8,
593
fully_sharded_loras=fully_shard,
594
lora_dtype=torch.float16,
595
)
596
+ punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
597
+ assert check_punica_wrapper(punica_wrapper)
598
599
def create_random_linear_parallel_layer():
600
if orientation == "row":
@@ -712,14 +712,14 @@ def test_column_parallel_packed(
712
713
max_loras = 8
714
torch.set_default_device(device)
715
- punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)
716
- assert check_punica_wrapper(punica_wrapper)
715
lora_config = LoRAConfig(
716
max_loras=max_loras,
717
max_lora_rank=8,
718
fully_sharded_loras=fully_shard,
719
lora_dtype=torch.float16,
720
)
721
+ punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config)
722
+ assert check_punica_wrapper(punica_wrapper)
723
724
def create_column_parallel_packed_layer():
725
if repeats == 2:
vllm/lora/model_manager.py
+4
-4
@@ -128,7 +128,7 @@ class LoRAModelManager:
128
max_num_batched_tokens,
129
max_batches=self.max_num_seqs,
130
device=self.device,
131
- max_loras=self.lora_config.max_loras,
131
+ lora_config=self.lora_config,
132
)
133
134
self.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY] = (
@@ -148,7 +148,7 @@ class LoRAModelManager:
148
max_num_batched_tokens,
149
max_batches=self.max_num_seqs,
150
device=self.device,
151
- max_loras=self.lora_config.max_loras,
151
+ lora_config=self.lora_config,
152
)
153
lm_prefix = self.mm_mapping.language_model[0]
154
self.punica_wrapper_mapping[lm_prefix] = llm_punica_wrapper
@@ -186,7 +186,7 @@ class LoRAModelManager:
186
num_encoder_tokens,
187
max_batches=self.max_num_seqs * limit_per_prompt,
188
device=self.device,
189
- max_loras=self.lora_config.max_loras,
189
+ lora_config=self.lora_config,
190
)
191
for prefix in self.mm_mapping.tower_model:
192
self.punica_wrapper_mapping[prefix] = tower_punica_wrapper
@@ -201,7 +201,7 @@ class LoRAModelManager:
201
connector_tokens,
202
max_batches=self.max_num_seqs * limit_per_prompt,
203
device=self.device,
204
- max_loras=self.lora_config.max_loras,
204
+ lora_config=self.lora_config,
205
)
206
for prefix in self.mm_mapping.connector:
207
self.punica_wrapper_mapping[prefix] = connector_punica_wrapper
vllm/lora/punica_wrapper/punica_gpu.py
+2
-1
@@ -45,7 +45,8 @@ class PunicaWrapperGPU(PunicaWrapperBase):
45
):
46
PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches, device)
47
48
- self.max_loras = kwargs["max_loras"]
48
+ self.lora_config = kwargs["lora_config"]
49
+ self.max_loras = self.lora_config.max_loras
50
51
self.token_mapping_meta = LoRAKernelMeta.make(
52
self.max_loras, max_num_batched_tokens, device=device