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 6, 14:30 UTC
Add ability to replace oot ops when using lora (#37181)
Signed-off-by: Kyuyeun Kim <kyuyeunk@google.com>
Kyuyeun Kim committed
Mar 16, 2026 at 18:04 UTC
0a0a1a198be88e1782b52fa31738896468200a76
6 files changed
+16
-11
vllm/lora/layers/column_parallel_linear.py
+4
-3
@@ -9,6 +9,7 @@ from transformers import PretrainedConfig
9
from vllm.config.lora import LoRAConfig
10
from vllm.distributed import tensor_model_parallel_all_gather
11
from vllm.distributed.utils import divide
12
+from vllm.model_executor.custom_op import maybe_get_oot_by_class
13
from vllm.model_executor.layers.linear import (
14
ColumnParallelLinear,
15
MergedColumnParallelLinear,
@@ -155,9 +156,9 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):
156
packed_modules_list: list,
157
model_config: PretrainedConfig | None = None,
158
) -> bool:
158
- if type(source_layer) is ColumnParallelLinear:
159
+ if type(source_layer) is maybe_get_oot_by_class(ColumnParallelLinear):
160
return True
160
- if type(source_layer) is MergedColumnParallelLinear:
161
+ if type(source_layer) is maybe_get_oot_by_class(MergedColumnParallelLinear):
162
if len(packed_modules_list) != 1:
163
return False
164
# Exclude layers with 3+ output sizes - those are handled by
@@ -606,7 +607,7 @@ class MergedColumnParallelLinearVariableSliceWithLoRA(
607
) -> bool:
608
# Support MergedColumnParallelLinear with 3 or more slices
609
# (2 slices are handled by MergedColumnParallelLinearWithLoRA)
609
- if type(source_layer) is not MergedColumnParallelLinear:
610
+ if type(source_layer) is not maybe_get_oot_by_class(MergedColumnParallelLinear):
611
return False
612
613
# If packed_modules_list has 3+ items, use this class
vllm/lora/layers/replicated_linear.py
+2
-1
@@ -7,6 +7,7 @@ import torch.nn as nn
7
from transformers import PretrainedConfig
8
9
from vllm.config.lora import LoRAConfig
10
+from vllm.model_executor.custom_op import maybe_get_oot_by_class
11
from vllm.model_executor.layers.linear import ReplicatedLinear
12
13
from .base_linear import BaseLinearLayerWithLoRA
@@ -55,7 +56,7 @@ class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):
56
packed_modules_list: list,
57
model_config: PretrainedConfig | None = None,
58
) -> bool:
58
- return type(source_layer) is ReplicatedLinear
59
+ return type(source_layer) is maybe_get_oot_by_class(ReplicatedLinear)
60
61
def slice_lora_a(
62
self, lora_a: torch.Tensor | list[torch.Tensor | None]
vllm/lora/layers/row_parallel_linear.py
+2
-1
@@ -11,6 +11,7 @@ from vllm.distributed import (
11
split_tensor_along_last_dim,
12
tensor_model_parallel_all_reduce,
13
)
14
+from vllm.model_executor.custom_op import maybe_get_oot_by_class
15
from vllm.model_executor.layers.linear import RowParallelLinear
16
from vllm.platforms import current_platform
17
@@ -89,7 +90,7 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):
90
packed_modules_list: list,
91
model_config: PretrainedConfig | None = None,
92
) -> bool:
92
- return type(source_layer) is RowParallelLinear
93
+ return type(source_layer) is maybe_get_oot_by_class(RowParallelLinear)
94
95
96
# The following layer is based on the tensor parallelism strategy given in
vllm/lora/layers/vocal_parallel_embedding.py
+2
-1
@@ -7,6 +7,7 @@ import torch.nn.functional as F
7
from transformers import PretrainedConfig
8
9
from vllm.config.lora import LoRAConfig
10
+from vllm.model_executor.custom_op import maybe_get_oot_by_class
11
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
12
from vllm.platforms import current_platform
13
@@ -132,7 +133,7 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
133
packed_modules_list: list,
134
model_config: PretrainedConfig | None = None,
135
) -> bool:
135
- return type(source_layer) is VocabParallelEmbedding
136
+ return type(source_layer) is maybe_get_oot_by_class(VocabParallelEmbedding)
137
138
@property
139
def weight(self):
vllm/model_executor/custom_op.py
+3
-2
@@ -22,10 +22,11 @@ op_registry: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}
22
op_registry_oot: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}
23
24
25
-def get_oot_class_by_name(class_name: str) -> type | None:
25
+def maybe_get_oot_by_class(class_type: type) -> type:
26
+ class_name = class_type.__name__
27
if class_name in op_registry_oot:
28
return op_registry_oot[class_name]
28
- return None
29
+ return class_type
30
31
32
class PluggableLayer(nn.Module):
vllm/model_executor/layers/attention/mm_encoder_attention.py
+3
-3
@@ -6,7 +6,7 @@ import numpy as np
6
import torch
7
8
from vllm.logger import init_logger
9
-from vllm.model_executor.custom_op import CustomOp, get_oot_class_by_name
9
+from vllm.model_executor.custom_op import CustomOp, maybe_get_oot_by_class
10
from vllm.model_executor.models.vision import get_vit_attn_backend
11
from vllm.utils.math_utils import round_up
12
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
@@ -125,7 +125,7 @@ class MMEncoderAttention(CustomOp):
125
cu_seqlens: np.ndarray,
126
device: torch.device,
127
) -> torch.Tensor | None:
128
- if (oot_class := get_oot_class_by_name(cls.__name__)) is not None:
128
+ if (oot_class := maybe_get_oot_by_class(cls)) is not cls:
129
return oot_class.maybe_compute_seq_lens(attn_backend, cu_seqlens, device) # type: ignore[attr-defined]
130
131
if attn_backend != AttentionBackendEnum.FLASHINFER:
@@ -149,7 +149,7 @@ class MMEncoderAttention(CustomOp):
149
tp_size: int,
150
device: torch.device,
151
) -> torch.Tensor:
152
- if (oot_class := get_oot_class_by_name(cls.__name__)) is not None:
152
+ if (oot_class := maybe_get_oot_by_class(cls)) is not cls:
153
return oot_class.maybe_recompute_cu_seqlens( # type: ignore[attr-defined]
154
attn_backend, cu_seqlens, hidden_size, tp_size, device
155
)