[WIP] Moe kernel for qwen3 omni in ascend

[jiaqiw09]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

This PR mainly adds accelerated operator patches for the Qwen3-Omni model.

This PR is currently marked as WIP, as the AI Core utilization still needs to be verified.

The current NPU patch implementation does not yet account for structural differences across different Transformers versions. After this PR is merged, follow-up work will address compatibility with different Transformers versions.

Experiment results

Paste your experiment result here(if needed).

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request integrates NPU acceleration for the Qwen3-Omni Mixture-of-Experts (MoE) model, enhancing its performance on Ascend hardware. The changes involve patching key components of the model with NPU-optimized operations, laying the groundwork for more efficient inference. While the core functionality is in place, further validation of AI Core utilization and broader Transformers version compatibility is planned.

Highlights

• New Feature: Qwen3-Omni MoE NPU Acceleration: Introduced support for the Qwen3-Omni MoE model by adding NPU-accelerated operator patches. This includes optimized implementations for RMSNorm, rotary position embeddings, and the MoE block's forward pass specifically for this model.
• Work In Progress (WIP): The pull request is currently marked as WIP, with AI Core utilization still requiring verification. Future work will also address compatibility with different Transformers versions.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• swift/model/npu_patcher.py
  • Imported the modeling_qwen3_omni_moe module to enable patching for the Qwen3-Omni model.
  • Added a new entry to the _PATCH_TABLE to apply NPU-specific patches for Qwen3OmniMoeTextRMSNorm, Qwen3OmniMoeRMSNorm, apply_rotary_pos_emb, and Qwen3OmniMoeThinkerTextSparseMoeBlock.forward within the modeling_qwen3_omni_moe module.

Activity

• The pull request was opened with the intent to add NPU acceleration for the Qwen3-Omni model.
• The author has marked the PR as 'Work In Progress' (WIP), indicating that further verification of AI Core utilization is needed.
• The author has noted that future work will address compatibility with different Transformers versions, as the current NPU patch implementation does not yet account for structural differences.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request adds NPU operator patches for the Qwen3-Omni-Moe model, enabling it to be accelerated on Ascend hardware. The changes follow the existing patching pattern in npu_patcher.py by adding a new entry to _PATCH_TABLE for modeling_qwen3_omni_moe.

My review includes one suggestion to improve code organization by maintaining a consistent order between model imports and their corresponding patch table entries. Overall, the changes look correct and are aligned with the PR's objective.

[gemini-code-assist]

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.