From 31fa757cf921980739ff6f836ef60d8d19d481dc Mon Sep 17 00:00:00 2001 From: yupeng <507435917@qq.com> Date: Fri, 15 May 2026 15:20:42 +0800 Subject: [PATCH] [Misc] Make it simpler to replace out-of-tree layer classes with related LoRA layers. (#42306) Signed-off-by: paulyu12 <507435917@qq.com> Co-authored-by: Jee Jee Li --- vllm/lora/layers/column_parallel_linear.py | 10 ++++++++-- vllm/lora/layers/fused_moe.py | 7 +++++-- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/vllm/lora/layers/column_parallel_linear.py b/vllm/lora/layers/column_parallel_linear.py index 01bbc9366d..8a86191b89 100644 --- a/vllm/lora/layers/column_parallel_linear.py +++ b/vllm/lora/layers/column_parallel_linear.py @@ -422,7 +422,10 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): packed_modules_list: list, model_config: PretrainedConfig | None = None, ) -> bool: - return type(source_layer) is QKVParallelLinear and len(packed_modules_list) == 1 + return ( + type(source_layer) is maybe_get_oot_by_class(QKVParallelLinear) + and len(packed_modules_list) == 1 + ) class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA): @@ -480,7 +483,10 @@ class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA): packed_modules_list: list, model_config: PretrainedConfig | None = None, ) -> bool: - return type(source_layer) is QKVParallelLinear and len(packed_modules_list) == 3 + return ( + type(source_layer) is maybe_get_oot_by_class(QKVParallelLinear) + and len(packed_modules_list) == 3 + ) # These following layers are based on the tensor parallelism strategy given in diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 1ceb34b714..9cb2fe9147 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -9,6 +9,7 @@ from vllm import envs from vllm.config.lora import LoRAConfig from vllm.distributed.utils import divide from vllm.lora.layers.base import BaseLayerWithLoRA +from vllm.model_executor.custom_op import maybe_get_oot_by_class from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRAContext from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import ( @@ -374,7 +375,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): """Returns True if the layer can be replaced by this LoRA layer.""" # source_layer is FusedMoE - return isinstance(source_layer, FusedMoE) and len(packed_modules_list) == 2 + moe_cls = maybe_get_oot_by_class(FusedMoE) + return isinstance(source_layer, moe_cls) and len(packed_modules_list) == 2 class FusedMoE3DWithLoRA(FusedMoEWithLoRA): @@ -537,4 +539,5 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA): ) -> bool: """Returns True if the layer can be replaced by this LoRA layer.""" # source_layer is FusedMoE - return isinstance(source_layer, FusedMoE) and len(packed_modules_list) == 1 + moe_cls = maybe_get_oot_by_class(FusedMoE) + return isinstance(source_layer, moe_cls) and len(packed_modules_list) == 1