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 9, 10:00 UTC
[PluggableLayer][2/N] Apply PluggableLayer to linear layers (#33152)
Signed-off-by: whx-sjtu <2952154980@qq.com>
whx committed
Jan 29, 2026 at 16:53 UTC
08b1195e6250820c4d44875bf61822e846e9ae28
1 file changed
+5
-5
vllm/model_executor/layers/linear.py
+5
-5
@@ -17,7 +17,7 @@ from vllm.distributed import (
17
tensor_model_parallel_all_reduce,
18
)
19
from vllm.logger import init_logger
20
-from vllm.model_executor.custom_op import CustomOp
20
+from vllm.model_executor.custom_op import PluggableLayer
21
from vllm.model_executor.layers.quantization.base_config import (
22
QuantizationConfig,
23
QuantizeMethodBase,
@@ -239,7 +239,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
239
return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)
240
241
242
-class LinearBase(CustomOp):
242
+class LinearBase(PluggableLayer):
243
"""Base linear layer.
244
245
Args:
@@ -294,7 +294,7 @@ class LinearBase(CustomOp):
294
295
296
# --8<-- [start:replicated_linear]
297
-@CustomOp.register("replicated_linear")
297
+@PluggableLayer.register("replicated_linear")
298
class ReplicatedLinear(LinearBase):
299
"""Replicated linear layer.
300
@@ -414,7 +414,7 @@ class ReplicatedLinear(LinearBase):
414
415
416
# --8<-- [start:column_parallel_linear]
417
-@CustomOp.register("column_parallel_linear")
417
+@PluggableLayer.register("column_parallel_linear")
418
class ColumnParallelLinear(LinearBase):
419
"""Linear layer with column parallelism.
420
@@ -1273,7 +1273,7 @@ class QKVParallelLinear(ColumnParallelLinear):
1273
1274
1275
# --8<-- [start:row_parallel_linear]
1276
-@CustomOp.register("row_parallel_linear")
1276
+@PluggableLayer.register("row_parallel_linear")
1277
class RowParallelLinear(LinearBase):
1278
"""Linear layer with row parallelism.
1279