[Platform] Custom ops support for LMhead and LogitsProcessor (vllm-project#23564)

Signed-off-by: zzhx1 <[email protected]>

Author: zzhx1

vllm/model_executor/layers/logits_processor.py

@@ -6,11 +6,11 @@
 from typing import Optional
 
 import torch
-import torch.nn as nn
 
 import vllm.envs as envs
 from vllm.distributed import (tensor_model_parallel_all_gather,
                               tensor_model_parallel_gather)
+from vllm.model_executor.custom_op import CustomOp
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding)
 from vllm.model_executor.sampling_metadata import SamplingMetadata
@@ -22,7 +22,8 @@
         envs.VLLM_LOGITS_PROCESSOR_THREADS)
 
 
-class LogitsProcessor(nn.Module):
+@CustomOp.register("logits_processor")
+class LogitsProcessor(CustomOp):
     """Process logits and apply logits processors from sampling metadata.
 
     This layer does the following:

vllm/model_executor/layers/vocab_parallel_embedding.py

@@ -429,6 +429,7 @@ def extra_repr(self) -> str:
         return s
 
 
+@CustomOp.register("parallel_lm_head")
 class ParallelLMHead(VocabParallelEmbedding):
     """Parallelized LM head.