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 3, 15:15 UTC
[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 <pandaleefree@gmail.com>
yupeng committed
May 15, 2026 at 15:20 UTC
31fa757cf921980739ff6f836ef60d8d19d481dc
2 files changed
+13
-4
vllm/lora/layers/column_parallel_linear.py
+8
-2
@@ -422,7 +422,10 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
422
packed_modules_list: list,
423
model_config: PretrainedConfig | None = None,
424
) -> bool:
425
- return type(source_layer) is QKVParallelLinear and len(packed_modules_list) == 1
425
+ return (
426
+ type(source_layer) is maybe_get_oot_by_class(QKVParallelLinear)
427
+ and len(packed_modules_list) == 1
428
+ )
429
430
431
class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):
@@ -480,7 +483,10 @@ class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):
483
packed_modules_list: list,
484
model_config: PretrainedConfig | None = None,
485
) -> bool:
483
- return type(source_layer) is QKVParallelLinear and len(packed_modules_list) == 3
486
+ return (
487
+ type(source_layer) is maybe_get_oot_by_class(QKVParallelLinear)
488
+ and len(packed_modules_list) == 3
489
+ )
490
491
492
# These following layers are based on the tensor parallelism strategy given in
vllm/lora/layers/fused_moe.py
+5
-2
@@ -9,6 +9,7 @@ from vllm import envs
9
from vllm.config.lora import LoRAConfig
10
from vllm.distributed.utils import divide
11
from vllm.lora.layers.base import BaseLayerWithLoRA
12
+from vllm.model_executor.custom_op import maybe_get_oot_by_class
13
from vllm.model_executor.layers.fused_moe import FusedMoE
14
from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRAContext
15
from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
@@ -374,7 +375,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
375
"""Returns True if the layer can be replaced by this LoRA layer."""
376
377
# source_layer is FusedMoE
377
- return isinstance(source_layer, FusedMoE) and len(packed_modules_list) == 2
378
+ moe_cls = maybe_get_oot_by_class(FusedMoE)
379
+ return isinstance(source_layer, moe_cls) and len(packed_modules_list) == 2
380
381
382
class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
@@ -537,4 +539,5 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
539
) -> bool:
540
"""Returns True if the layer can be replaced by this LoRA layer."""
541
# source_layer is FusedMoE
540
- return isinstance(source_layer, FusedMoE) and len(packed_modules_list) == 1
542
+ moe_cls = maybe_get_oot_by_class(FusedMoE)
543
+ return isinstance(source_layer, moe_cls) and len(packed_modules_list) == 1