From f84bf7d79b885956747f2258ddeae380da492b3c Mon Sep 17 00:00:00 2001 From: ZT-AIA <63220130+ZT-AIA@users.noreply.github.com> Date: Tue, 30 Dec 2025 14:27:31 +0800 Subject: [PATCH] Add Loraconfig parameter to get_punica_wrapper function (#31408) Signed-off-by: ZT-AIA <1028681969@qq.com> Co-authored-by: Jee Jee Li --- tests/lora/test_layers.py | 20 ++++++++++---------- vllm/lora/model_manager.py | 8 ++++---- vllm/lora/punica_wrapper/punica_gpu.py | 3 ++- 3 files changed, 16 insertions(+), 15 deletions(-) diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py index 47d1fcfe9..dfec42368 100644 --- a/tests/lora/test_layers.py +++ b/tests/lora/test_layers.py @@ -261,11 +261,11 @@ def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None: torch.set_default_device(device) max_loras = 8 - punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras) - assert check_punica_wrapper(punica_wrapper) lora_config = LoRAConfig( max_loras=max_loras, max_lora_rank=8, lora_dtype=torch.float16 ) + punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config) + assert check_punica_wrapper(punica_wrapper) def create_random_embedding_layer(): embedding = VocabParallelEmbedding(vocab_size, 256) @@ -360,11 +360,11 @@ def test_lm_head_logits_processor( torch.set_default_device(device) max_loras = 8 - punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras) - assert check_punica_wrapper(punica_wrapper) lora_config = LoRAConfig( max_loras=max_loras, max_lora_rank=8, lora_dtype=torch.float16 ) + punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config) + assert check_punica_wrapper(punica_wrapper) def _pretest(): linear = ParallelLMHead( @@ -480,13 +480,13 @@ def test_linear_replicated( max_loras = 8 torch.set_default_device(device) - punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras) - assert check_punica_wrapper(punica_wrapper) lora_config = LoRAConfig( max_loras=max_loras, max_lora_rank=8, lora_dtype=torch.float16, ) + punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config) + assert check_punica_wrapper(punica_wrapper) def create_random_linear_replicated_layer(): linear = ReplicatedLinear(4096, 4096, bias=False, params_dtype=torch.float16) @@ -587,14 +587,14 @@ def test_linear_parallel( max_loras = 8 torch.set_default_device(device) - punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras) - assert check_punica_wrapper(punica_wrapper) lora_config = LoRAConfig( max_loras=max_loras, max_lora_rank=8, fully_sharded_loras=fully_shard, lora_dtype=torch.float16, ) + punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config) + assert check_punica_wrapper(punica_wrapper) def create_random_linear_parallel_layer(): if orientation == "row": @@ -712,14 +712,14 @@ def test_column_parallel_packed( max_loras = 8 torch.set_default_device(device) - punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras) - assert check_punica_wrapper(punica_wrapper) lora_config = LoRAConfig( max_loras=max_loras, max_lora_rank=8, fully_sharded_loras=fully_shard, lora_dtype=torch.float16, ) + punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config) + assert check_punica_wrapper(punica_wrapper) def create_column_parallel_packed_layer(): if repeats == 2: diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 4506b0a44..69787f8f0 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -128,7 +128,7 @@ class LoRAModelManager: max_num_batched_tokens, max_batches=self.max_num_seqs, device=self.device, - max_loras=self.lora_config.max_loras, + lora_config=self.lora_config, ) self.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY] = ( @@ -148,7 +148,7 @@ class LoRAModelManager: max_num_batched_tokens, max_batches=self.max_num_seqs, device=self.device, - max_loras=self.lora_config.max_loras, + lora_config=self.lora_config, ) lm_prefix = self.mm_mapping.language_model[0] self.punica_wrapper_mapping[lm_prefix] = llm_punica_wrapper @@ -186,7 +186,7 @@ class LoRAModelManager: num_encoder_tokens, max_batches=self.max_num_seqs * limit_per_prompt, device=self.device, - max_loras=self.lora_config.max_loras, + lora_config=self.lora_config, ) for prefix in self.mm_mapping.tower_model: self.punica_wrapper_mapping[prefix] = tower_punica_wrapper @@ -201,7 +201,7 @@ class LoRAModelManager: connector_tokens, max_batches=self.max_num_seqs * limit_per_prompt, device=self.device, - max_loras=self.lora_config.max_loras, + lora_config=self.lora_config, ) for prefix in self.mm_mapping.connector: self.punica_wrapper_mapping[prefix] = connector_punica_wrapper diff --git a/vllm/lora/punica_wrapper/punica_gpu.py b/vllm/lora/punica_wrapper/punica_gpu.py index ef4b4ab7c..f765df0b3 100644 --- a/vllm/lora/punica_wrapper/punica_gpu.py +++ b/vllm/lora/punica_wrapper/punica_gpu.py @@ -45,7 +45,8 @@ class PunicaWrapperGPU(PunicaWrapperBase): ): PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches, device) - self.max_loras = kwargs["max_loras"] + self.lora_config = kwargs["lora_config"] + self.max_loras = self.lora_config.max_loras self.token_mapping_meta = LoRAKernelMeta.make( self.max_loras, max_num_batched_tokens, device=device