Skip to content

[WIP] Moe kernel for qwen3 omni in ascend #8214

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jiaqiw09 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[WIP] Moe kernel for qwen3 omni in ascend #8214

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions swift/model/npu_patcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from transformers.models.qwen2 import modeling_qwen2
from transformers.models.qwen3 import modeling_qwen3
from transformers.models.qwen3_moe import modeling_qwen3_moe
from transformers.models.qwen3_omni_moe import modeling_qwen3_omni_moe
from transformers.models.qwen3_vl_moe import modeling_qwen3_vl_moe
from typing import Any

Expand Down Expand Up @@ -332,6 +333,15 @@ def _apply_patch_map(root: Any, patch_map: dict[str, Any]) -> None:
'apply_rotary_pos_emb': npu_apply_rotary_pos_emb,
},
),
(
modeling_qwen3_omni_moe,
{
'Qwen3OmniMoeTextRMSNorm': NpuRMSNorm,
'Qwen3OmniMoeRMSNorm': NpuRMSNorm,
'apply_rotary_pos_emb': npu_apply_rotary_pos_emb,
'Qwen3OmniMoeThinkerTextSparseMoeBlock.forward': npu_moe_block_forward,
},
),
Comment on lines +336 to +344
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Mar 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

For better code organization and consistency, it's good practice to keep the order of models in _PATCH_TABLE consistent with the order of their imports at the top of the file.

The new import for modeling_qwen3_omni_moe is placed before modeling_qwen3_vl_moe (which is correct alphabetically), but the new entry in _PATCH_TABLE is placed after it.

To maintain consistency, please consider moving this new block to be before the modeling_qwen3_vl_moe entry. This will make the code easier to read and maintain.

)

for _module, _patch_map in _PATCH_TABLE:
Expand Down
Loading