[KernelGen][MThreads] Add mish_ Moore Threads specialized operator - #296
Open
Yukun-Cui wants to merge 3 commits into
Open
[KernelGen][MThreads] Add mish_ Moore Threads specialized operator#296Yukun-Cui wants to merge 3 commits into
Yukun-Cui wants to merge 3 commits into
Conversation
github-actions
Bot
added
documentation
Collaborator
|
/test | mish_:mthreads |
1 similar comment
Collaborator
|
/test | mish_:mthreads |
❌ On-demand test failedOperator: The test failed to complete. Check the workflow run for details. |
Collaborator
|
/test | mish_:mthreads |
Dongxu-H
reviewed
Aug 18, 2026
Dongxu-H
left a comment
Dongxu-H
left a comment
Collaborator
There was a problem hiding this comment.
PR has merge conflicts that must be resolved before merging.
Dongxu-H
reviewed
Aug 18, 2026
Dongxu-H
left a comment
Dongxu-H
left a comment
Collaborator
There was a problem hiding this comment.
PR 添加了 mish_.py 文件，但现有的 mish.py 已经包含 mish_ 函数实现，这会导致重复代码和潜在的导入冲突。建议将优化逻辑合并到现有的 mish.py 中而不是创建新文件。
Yukun-Cui
force-pushed
the
pr/mthreads-mish_
branch
from
August 19, 2026 03:42
abe806c to
4844157
Compare
github-actions
Bot
removed
documentation
github-actions
Bot
added
size/Small
and removed
vendor/MetaX
vendor/THead
vendor/Hygon
size/XLarge
labels
Dongxu-H
reviewed
Aug 19, 2026
Dongxu-H
left a comment
Dongxu-H
left a comment
Collaborator
There was a problem hiding this comment.
CI code-style check failed.
init.py imports mish_ twice: line 61 from .mish and line 62 from .mish_. This causes a naming conflict.
float32 performance regression: speedup 0.55x, about 82% slower than generic implementation.
The existing mish.py exports both mish and mish_; the new specialized mish_.py also exports mish_. Import mish_ only from the new specialized module to drop the F811 redefinition.
PR flagos-ai#296 moved the in-place mish_ into its own specialized module mish_.py; the __init__ now imports mish_ only from mish_.py. The mish_ definition left behind in mish.py is unreachable dead code (and its default_mish_ import is now unused). Remove it; the shared mish_kernel / _use_triton_kernel / _launch_mish helpers stay, as they are still used by mish().
Collaborator
|
/test | mish_:mthreads |
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
2 participants
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.
[KernelGen][MThreads] Add mish_ Moore Threads specialized operator
Summary
Add a Moore Threads (MUSA) specialized Triton kernel for
mish_, overriding the generic implementation viaruntime.replace_customized_ops(). Implements the mish operation.Testing
tests/test_mish.py(-m mish_)GEMS_MTHREADS MISH_debug logPerformance
Compared against the generic FlagGems implementation on Moore Threads (MUSA).
mish_
Files Changed
src/flag_gems/runtime/backend/_mthreads/ops/mish_.py: Moore Threads Triton kernel + fallbacksrc/flag_gems/runtime/backend/_mthreads/ops/__init__.py: Register import and__all__