Closed
Conversation
zhiding512
reviewed
Nov 18, 2025
Rohan138
force-pushed
the
ropotdar/fix_2stage_moe_jit
branch
from
1f88988 to
c2ca6ac
Compare
Contributor
There was a problem hiding this comment.
Pull Request Overview
This PR fixes a missing pybind registration for the CKtile 2-stage MOE kernel that was not properly registered when ported from the 355_wip branch. The fix enables the CKtile MOE functions to be accessible from Python.
- Adds the missing header include for
moe_cktile2stages.h - Registers the
MOE_CKTILE_2STAGES_PYBINDmacro to exposecktile_moe_gemm1andcktile_moe_gemm2functions
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
1 task
Contributor
Author
|
@zhiding512 / @valarLip PTAL, can we merge to unblock vllm/355_wip? Unrelated CI failures due to git ssh outage earlier today |
zhiding512
approved these changes
Nov 19, 2025
dllehr-amd
force-pushed
the
ropotdar/fix_2stage_moe_jit
branch
from
5546a1f to
bb30ce4
Compare
Collaborator
|
#1412 this one should fixed the issue |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
While working on https://github.com/ROCm/vllm/pull/808/files, we've found that building the aiter wheel then linking it into the final docker build stage leads to a broken build, which does not occur when building the library in-place. cc @dllehr-amd @zhiding512
It looks like the CKtile 2stage MOE was not registered in pybind when it was ported over from 355_wip: https://github.com/ROCm/aiter/blob/355_wip/csrc/rocm_ops.cpp#L80; linking it fixes the issue.
Also updated the CK 2stage MOE to use FLATMM_HIP_CLANG_PATH.Technical Details
Test Plan
Test Result
Submission Checklist