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

Signed-off-by: zzhx1 <zzh_201018@outlook.com>
2025-09-10 21:26:31 +08:00
parent 2eb9986a2d
commit 736569da8d
2 changed files with 4 additions and 2 deletions
--- a/vllm/model_executor/layers/vocab_parallel_embedding.py
+++ b/vllm/model_executor/layers/vocab_parallel_embedding.py
@@ -429,6 +429,7 @@ class VocabParallelEmbedding(CustomOp):
        return s


+@CustomOp.register("parallel_lm_head")
 class ParallelLMHead(VocabParallelEmbedding):
    """Parallelized LM head.