[Feat.]: 310p support MOE models

[pu-zhe]

What this PR does / why we need it?

This pull request integrates comprehensive support for Mixture of Experts (MoE) models on the Ascend 310P device within the vllm-ascend framework. It achieves this by introducing specialized modules for expert selection, fused MoE layers, and optimized all-gather communication. The changes also refine existing NPU operations, making them more consistent and efficient for 310P, ultimately enhancing the performance and compatibility of MoE models on this hardware.

Highlights
310P MoE Support: Introduces dedicated implementations for Mixture of Experts (MoE) models on Ascend 310P devices, including new modules for expert selection, fused MoE layers, and communication.
All-Gather Communication: Enforces the use of ALLGATHER communication for MoE operations on 310P, optimizing data transfer and leveraging NPU-specific token dispatching.
Simplified NPU Operations: Removes conditional type casting for npu_swiglu and enables custom rotary embedding kernels unconditionally, suggesting improved native support for 310P.
New MoE Classes Registered: Registers AscendFusedMoE310 and AscendSharedFusedMoE310 to integrate 310P-specific MoE layers into the system's custom operation registry.

Does this PR introduce any user-facing change?

No

How was this patch tested?

offline test and server test, with qwen3-30b-a3b,tp/ep 4 on 310p

• vLLM version: v0.15.0
• vLLM main: https://github.com/vllm-project/vllm/commit/v0.15.0

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[gemini-code-assist]

Summary of Changes

Hello @pu-zhe, 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 comprehensive support for Mixture of Experts (MoE) models on the Ascend 310P device within the vllm-ascend framework. It achieves this by introducing specialized modules for expert selection, fused MoE layers, and optimized all-gather communication. The changes also refine existing NPU operations, making them more consistent and efficient for 310P, ultimately enhancing the performance and compatibility of MoE models on this hardware.

Highlights

• 310P MoE Support: Introduces dedicated implementations for Mixture of Experts (MoE) models on Ascend 310P devices, including new modules for expert selection, fused MoE layers, and communication.
• All-Gather Communication: Enforces the use of ALLGATHER communication for MoE operations on 310P, optimizing data transfer and leveraging NPU-specific token dispatching.
• Unified Weight Processing: Streamlines weight transformation (maybe_trans_nz) for MoE layers, applying it consistently across devices by removing 310P-specific conditional logic.
• Simplified NPU Operations: Removes conditional type casting for npu_swiglu and enables custom rotary embedding kernels unconditionally, suggesting improved native support for 310P.
• New MoE Classes Registered: Registers AscendFusedMoE310 and AscendSharedFusedMoE310 to integrate 310P-specific MoE layers into the system's custom operation registry.

🧠 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

• vllm_ascend/_310p/fused_moe/experts_selector.py
  • Added a new module containing the select_experts function, which handles the selection of top-k experts and includes a prefetch mechanism for MoE weights.
• vllm_ascend/_310p/fused_moe/fused_moe.py
  • Added new classes AscendUnquantizedFusedMoEMethod310, AscendFusedMoE310, and AscendSharedFusedMoE310 to provide 310P-specific implementations for unquantized fused MoE layers, including weight processing, expert mapping, and integration with the new communication methods.
• vllm_ascend/_310p/fused_moe/moe_comm_method.py
  • Added AllGatherCommImpl310, a 310P-optimized communication method for MoE that utilizes NPU-specific operations for token dispatching.
• vllm_ascend/_310p/fused_moe/token_dispatcher.py
  • Added TokenDispatcherWithAllGather310, which manages token distribution to experts and includes moe_init_routing for efficient token organization on 310P.
• vllm_ascend/ascend_forward_context.py
  • Modified select_moe_comm_method to ensure that ALLGATHER is always chosen as the MoE communication type when running on Ascend 310P devices.
• vllm_ascend/ops/fused_moe/fused_moe.py
  • Updated process_weights_after_loading to remove the conditional check for AscendDeviceType._310P when applying maybe_trans_nz to MoE layer weights, making the transformation universal.
• vllm_ascend/ops/fused_moe/moe_mlp.py
  • Revised unquant_apply_mlp to remove the explicit to(torch.float32).to(torch.float16) type casting for npu_swiglu on 310P, simplifying the activation function call.
• vllm_ascend/ops/rotary_embedding.py
  • Modified _rope_forward_oot to remove the AscendDeviceType._310P exclusion, allowing the custom rotary embedding kernel to be used on 310P devices.
• vllm_ascend/utils.py
  • Updated register_ascend_customop to include AscendFusedMoE310 and AscendSharedFusedMoE310 in the registry for 310P-specific operations.

Activity

• No specific activity has been recorded for this pull request since its creation.

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 introduces support for Mixture of Experts (MoE) models on Ascend 310P devices. The changes include new 310P-specific implementations for MoE layers, such as expert selection, communication methods, and token dispatching. Additionally, it enables existing custom operators like rotary embedding for 310P and registers the new MoE operators.

The review has identified several critical issues related to incorrect method calls, return types, and control flow logic that could lead to runtime errors. These issues need to be addressed to ensure the correctness of the implementation.

I've also provided suggestions for the PR title and summary to align with the repository's style guide.

Suggested PR Title:

[MoE][Feature] Support MoE models on 310P

Suggested PR Summary:

### What this PR does / why we need it?
This PR adds support for Mixture of Experts (MoE) models on Ascend 310P devices. It introduces 310P-specific implementations for MoE layers, including expert selection, communication methods, and token dispatching. It also enables existing custom operators like rotary embedding for 310P and registers the new MoE operators.

### Does this PR introduce _any_ user-facing change?
Yes, users can now run MoE models on Ascend 310P devices.

### How was this patch tested?
CI passed. Tested MoE model inference on 310P platform.

[wangxiyuan]

please add a e2e for moe case to ensure the function works