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 5, 07:15 UTC
[PluggableLayer][3/N] Apply PluggableLayer to llm_head and vocab embedding layer (#33465)
Signed-off-by: whx-sjtu <2952154980@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Hexiang Wang committed
Apr 10, 2026 at 16:12 UTC
1dfd64c1cccb02c022f703628898094f02a7da3f
2 files changed
+8
-11
vllm/model_executor/layers/logits_processor.py
+3
-3
@@ -9,14 +9,14 @@ from vllm.distributed import (
9
tensor_model_parallel_all_gather,
10
tensor_model_parallel_gather,
11
)
12
-from vllm.model_executor.custom_op import CustomOp
12
+from vllm.model_executor.custom_op import PluggableLayer
13
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
14
from vllm.platforms import current_platform
15
16
17
# --8<-- [start:logits_processor]
18
-@CustomOp.register("logits_processor")
19
-class LogitsProcessor(CustomOp):
18
+@PluggableLayer.register("logits_processor")
19
+class LogitsProcessor(PluggableLayer):
20
"""Process logits and apply logits processors from sampling metadata.
21
22
This layer does the following:
vllm/model_executor/layers/vocab_parallel_embedding.py
+5
-8
@@ -14,7 +14,7 @@ from vllm.distributed import (
14
get_tensor_model_parallel_world_size,
15
tensor_model_parallel_all_reduce,
16
)
17
-from vllm.model_executor.custom_op import CustomOp
17
+from vllm.model_executor.custom_op import PluggableLayer
18
from vllm.model_executor.layers.quantization.base_config import (
19
QuantizationConfig,
20
QuantizeMethodBase,
@@ -182,8 +182,8 @@ def get_masked_input_and_mask(
182
183
184
# --8<-- [start:vocab_parallel_embedding]
185
-@CustomOp.register("vocab_parallel_embedding")
186
-class VocabParallelEmbedding(CustomOp):
185
+@PluggableLayer.register("vocab_parallel_embedding")
186
+class VocabParallelEmbedding(PluggableLayer):
187
"""Embedding parallelized in the vocabulary dimension.
188
189
Adapted from torch.nn.Embedding, note that we pad the vocabulary size to
@@ -461,7 +461,7 @@ class VocabParallelEmbedding(CustomOp):
461
param[: loaded_weight.shape[0]].data.copy_(loaded_weight)
462
param[loaded_weight.shape[0] :].data.fill_(0)
463
464
- def forward_native(self, input_):
464
+ def forward(self, input_):
465
if self.tp_size > 1:
466
# Build the mask.
467
masked_input, input_mask = get_masked_input_and_mask(
@@ -483,9 +483,6 @@ class VocabParallelEmbedding(CustomOp):
483
output = tensor_model_parallel_all_reduce(output_parallel)
484
return output
485
486
- def forward_cuda(self, input_):
487
- return self.forward_native(input_)
488
-
486
def extra_repr(self) -> str:
487
s = f"num_embeddings={self.num_embeddings_per_partition}"
488
s += f", embedding_dim={self.embedding_dim}"
@@ -496,7 +493,7 @@ class VocabParallelEmbedding(CustomOp):
493
494
495
# --8<-- [start:parallel_lm_head]
499
-@CustomOp.register("parallel_lm_head")
496
+@PluggableLayer.register("parallel_lm_head")
497
class ParallelLMHead(VocabParallelEmbedding):
498
"""Parallelized LM head.
499