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 14, 05:30 UTC
[Platform] Custom ops support for FusedMoe (#22509)
Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
wangxiyuan committed
Aug 13, 2025 at 19:12 UTC
0b1bdac6af33b890a4d68321df05e71a1ba43dc4
3 files changed
+11
-8
vllm/model_executor/layers/fused_moe/layer.py
+2
-1
@@ -682,7 +682,8 @@ def determine_expert_map(
682
return (local_num_experts, expert_map)
683
684
685
-class FusedMoE(torch.nn.Module):
685
+@CustomOp.register("fused_moe")
686
+class FusedMoE(CustomOp):
687
"""FusedMoE layer for MoE models.
688
689
This layer contains both MergedColumnParallel weights (gate_up_proj /
vllm/model_executor/layers/linear.py
+6
-6
@@ -16,6 +16,7 @@ from vllm.distributed import (divide, get_tensor_model_parallel_rank,
16
tensor_model_parallel_all_gather,
17
tensor_model_parallel_all_reduce)
18
from vllm.logger import init_logger
19
+from vllm.model_executor.custom_op import CustomOp
20
from vllm.model_executor.layers.quantization.base_config import (
21
QuantizationConfig, QuantizeMethodBase)
22
from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
@@ -226,7 +227,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
227
return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)
228
229
229
-class LinearBase(torch.nn.Module):
230
+class LinearBase(CustomOp):
231
"""Base linear layer.
232
233
Args:
@@ -269,12 +270,8 @@ class LinearBase(torch.nn.Module):
270
prefix=prefix)
271
self.return_bias = return_bias
272
272
- def forward(
273
- self, x: torch.Tensor
274
- ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:
275
- raise NotImplementedError
276
-
273
274
+@CustomOp.register("replicated_linear")
275
class ReplicatedLinear(LinearBase):
276
"""Replicated linear layer.
277
@@ -443,6 +440,7 @@ class MergedReplicatedLinear(ReplicatedLinear):
440
param[shard_offset:shard_offset + shard_size] = loaded_weight
441
442
443
+@CustomOp.register("column_parallel_linear")
444
class ColumnParallelLinear(LinearBase):
445
"""Linear layer with column parallelism.
446
@@ -1229,6 +1227,7 @@ class QKVParallelLinear(ColumnParallelLinear):
1227
param_data.copy_(loaded_weight)
1228
1229
1230
+@CustomOp.register("row_parallel_linear")
1231
class RowParallelLinear(LinearBase):
1232
"""Linear layer with row parallelism.
1233
@@ -1405,6 +1404,7 @@ class RowParallelLinear(LinearBase):
1404
return s
1405
1406
1407
+@CustomOp.register("qkv_cross_parallel_linear")
1408
class QKVCrossParallelLinear(LinearBase):
1409
"""Linear layers for efficient cross-attention's QKV transformation.
1410
vllm/model_executor/layers/vocab_parallel_embedding.py
+3
-1
@@ -12,6 +12,7 @@ from torch.nn.parameter import Parameter, UninitializedParameter
12
from vllm.distributed import (divide, get_tensor_model_parallel_rank,
13
get_tensor_model_parallel_world_size,
14
tensor_model_parallel_all_reduce)
15
+from vllm.model_executor.custom_op import CustomOp
16
from vllm.model_executor.layers.quantization.base_config import (
17
QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding)
18
from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
@@ -159,7 +160,8 @@ def get_masked_input_and_mask(
160
return input_, ~vocab_mask
161
162
162
-class VocabParallelEmbedding(torch.nn.Module):
163
+@CustomOp.register("vocab_parallel_embedding")
164
+class VocabParallelEmbedding(CustomOp):
165
"""Embedding parallelized in the vocabulary dimension.
166
167
Adapted from torch.nn.Embedding, note that we pad the vocabulary size to