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 LMhead and LogitsProcessor (#23564)
Signed-off-by: zzhx1 <zzh_201018@outlook.com>
zzhxxx committed
Sep 10, 2025 at 21:26 UTC
736569da8d856d5a66d7d7523a7e05f97a4d21f1
2 files changed
+4
-2
vllm/model_executor/layers/logits_processor.py
+3
-2
@@ -6,11 +6,11 @@ from concurrent.futures import ThreadPoolExecutor
6
from typing import Optional
7
8
import torch
9
-import torch.nn as nn
9
10
import vllm.envs as envs
11
from vllm.distributed import (tensor_model_parallel_all_gather,
12
tensor_model_parallel_gather)
13
+from vllm.model_executor.custom_op import CustomOp
14
from vllm.model_executor.layers.vocab_parallel_embedding import (
15
VocabParallelEmbedding)
16
from vllm.model_executor.sampling_metadata import SamplingMetadata
@@ -22,7 +22,8 @@ if envs.VLLM_LOGITS_PROCESSOR_THREADS is not None:
22
envs.VLLM_LOGITS_PROCESSOR_THREADS)
23
24
25
-class LogitsProcessor(nn.Module):
25
+@CustomOp.register("logits_processor")
26
+class LogitsProcessor(CustomOp):
27
"""Process logits and apply logits processors from sampling metadata.
28
29
This layer does the following:
vllm/model_executor/layers/vocab_parallel_embedding.py
+1
@@ -429,6 +429,7 @@ class VocabParallelEmbedding(CustomOp):
429
return s
430
431
432
+@CustomOp.register("parallel_lm_head")
433
class ParallelLMHead(VocabParallelEmbedding):
434
"""Parallelized LM head.
435