[Platform] Custom ops support for LMhead and LogitsProcessor #23564
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run You ask your reviewers to trigger select CI tests on top of Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request aims to add CustomOp support for LMHead and LogitsProcessor to allow for platform-specific implementations. The changes correctly update the class definitions to inherit from CustomOp and register them. However, there's a critical issue in both implementations: the existing forward methods are not renamed to forward_native. This will cause them to be shadowed by the CustomOp's dispatcher, leading to NotImplementedError at runtime. My review includes comments on the affected class definitions, explaining the necessary changes to make these custom ops functional.
There was a problem hiding this comment.
Inheriting from CustomOp overrides the forward method with the CustomOp dispatcher. The existing forward method in this class (lines 56-87) must be renamed to forward_native to serve as the default implementation. Without this change, calls to this op will fail with a NotImplementedError.
There was a problem hiding this comment.
Since ParallelLMHead is now registered as a CustomOp (and inherits from CustomOp via VocabParallelEmbedding), its forward method is overridden. The existing forward method (lines 475-477) must be renamed to forward_native to act as the default implementation. Otherwise, this will raise a NotImplementedError at runtime.
Signed-off-by: zzhx1 <[email protected]>
|
LGTM, the CI failure looks doesn't relate to this PR. |
|
@wangxiyuan Help to push and merge this PR |
|
@Yikun @youkaichao please take a look. The change is quite simple and clear IMO |
|
Just enbale function level register for LogitsProcessor and ParallelLMHead,It's same as #22509 more background to see #22509 (comment) . so, LGTM |
…oject#23564) Signed-off-by: zzhx1 <[email protected]>
…oject#23564) Signed-off-by: zzhx1 <[email protected]>
supported_models.mdandexamplesfor a new model.Purpose
Similar with #22509
We hope that lmhead can also support the custom_op feature so that I can make modifications to lmhead on different platforms without needing to change the vllm code.
Test Plan
No need to add new test.
Test Result
all tests should pass
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.