CustomOp: test forward dispatch for grouped_topk#31530
CustomOp: test forward dispatch for grouped_topk#31530yewentao256 merged 2 commits intovllm-project:mainfrom
Conversation
xinyu-intel
requested review from
WoosukKwon,
mgoin,
tlrmchlsmth and
yewentao256
as code owners
|
@tjtanaa @ganyi1996ppo please check this. |
![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 adds a test for the forward dispatch mechanism of the grouped_topk custom operation. The changes correctly set up a VllmConfig to enable the custom op and then assert that the appropriate platform-specific forward method (forward_cuda or forward_hip) is selected. The test also includes a correctness check comparing a baseline implementation with the fused kernel. The implementation is sound and effectively validates the dispatch logic. I have no issues to report.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
xinyu-intel
force-pushed
the
dev/xinyu/grouped-topk-test
branch
from
49ac2e9 to
c8e285e
Compare
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
| dtype: torch.dtype, | ||
| ): | ||
| vllm_config = VllmConfig( | ||
| compilation_config=CompilationConfig(custom_ops=["all", "+grouped_topk"]) |
There was a problem hiding this comment.
I think we don't need to add "all" as we are only testing "grouped_topk". All will enable all custom ops. We should also test with only +grouped_topk.
yewentao256
left a comment
yewentao256
left a comment
There was a problem hiding this comment.
LGTM, thanks for the work!
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com> Signed-off-by: dsuhinin <suhinin.dmitriy@gmail.com>
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
4 participants
Purpose
follow up of #29575, #31221 and #31523 to test forward func dispatch for grouped_topk custom op.