fix: split asm_topksoftmax into separate module to fix ctypes JIT build (#2548)#2603
Merged
Conversation
When topk_softmax_asm (ffi_type="ctypes") triggers JIT compilation, torch_exclude=True is forced. But module_moe_asm also contains pybind .cu files that depend on torch, causing undefined symbol errors. Split asm_topksoftmax.cu into its own module_moe_topksoftmax_asm so it compiles torch-free independently. No .cu files modified. Fixes #2548
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR resolves a runtime JIT-loading failure for ffi_type="ctypes" MOE TopK+Softmax ASM by splitting the asm_topksoftmax.cu build target out of the existing torch-linked module_moe_asm into its own dedicated, torch-free JIT module, avoiding undefined torch symbol errors when torch_exclude=True is forced for ctypes modules.
Changes:
- Add new JIT build module
module_moe_topksoftmax_asmforasm_topksoftmax.cu(with its blob/codegen step). - Remove
asm_topksoftmax.cufrommodule_moe_asmsources so thatmodule_moe_asmremains torch/pybind-linked without impacting ctypes builds. - Update
topk_softmax_asmdecorator to compile/load from the new module.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| aiter/ops/moe_op.py | Points topk_softmax_asm (ctypes) at the new torch-free JIT module name. |
| aiter/jit/optCompilerConfig.json | Splits build config: module_moe_asm no longer includes asm_topksoftmax.cu; new module_moe_topksoftmax_asm builds it independently (including blob_gen_cmd). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
asm_topksoftmax.cu does not use ck_tile headers.
valarLip
approved these changes
Apr 3, 2026
yzhou103
pushed a commit
that referenced
this pull request
Apr 8, 2026
…ld (#2548) (#2603) * fix: split asm_topksoftmax into separate module to fix ctypes JIT build When topk_softmax_asm (ffi_type="ctypes") triggers JIT compilation, torch_exclude=True is forced. But module_moe_asm also contains pybind .cu files that depend on torch, causing undefined symbol errors. Split asm_topksoftmax.cu into its own module_moe_topksoftmax_asm so it compiles torch-free independently. No .cu files modified. Fixes #2548 * style: black formatting for compile_ops decorator line * fix: remove unnecessary ck_tile include from module_moe_topksoftmax_asm asm_topksoftmax.cu does not use ck_tile headers. --------- Co-authored-by: root <root@hjbog-srdc-39.amd.com>
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.
Summary
asm_topksoftmax.cuout ofmodule_moe_asminto a newmodule_moe_topksoftmax_asmmoduletopk_softmax_asm(ffi_type="ctypes") triggers JIT compilation,torch_exclude=Trueis forced, but other.cufiles inmodule_moe_asmdepend on torch, causingundefined symbolerrorsChanges (2 files, ~10 lines):
aiter/jit/optCompilerConfig.json— removeasm_topksoftmax.cufrommodule_moe_asm, add newmodule_moe_topksoftmax_asmentryaiter/ops/moe_op.py— pointtopk_softmax_asmdecorator to new module nameNo
.cufiles modified.Fixes #2548
Test plan
.sofiles and ranpython op_tests/test_moeTopkSoftmax.py— all tests passmodule_moe_topksoftmax_asm.sohas zero torch dependencies (ldd | grep torchreturns 0)module_moe_asm.sobuilds and loads normally with pybind + torch🤖 Generated with Claude Code