[Refactor][Ops] Move expert routing into router classes

[weijinqian0]

What this PR does / why we need it?

Come from #13220

This PR aligns the Ascend fused MoE execution flow with the upstream router abstraction:

• introduces Ascend grouped top-k and fused top-k router implementations;
• creates the Ascend router through the platform-level FusedMoE patch;
• moves expert selection out of quantization methods;
• passes precomputed top-k weights and expert IDs into routed expert execution; and
• updates unquantized and quantized MoE tests for the new execution contract.

Keeping routing separate from quantization reduces duplicated selection logic across quantization schemes and makes the Ascend implementation follow the upstream FusedMoERouter lifecycle.

Does this PR introduce any user-facing change?

No. This is an internal fused MoE refactor with no intended API or behavior change.

How was this patch tested?

• python -m py_compile on all modified Python implementation files.

• pytest -q tests/ut/ops/test_fused_moe.py tests/ut/quantization/methods/a2/test_w4a16.py tests/ut/quantization/methods/a2/test_w4a8.py tests/ut/quantization/methods/a2/test_w8a8_dynamic.py tests/ut/quantization/methods/test_w4a16_mxfp4.py tests/ut/quantization/methods/test_w4a4_mxfp4.py tests/ut/quantization/methods/test_w8a8_mxfp8.py tests/ut/quantization/methods/test_w8a8fp8_dynamic.py tests/ut/quantization/test_method_adapters.py

  • 101 tests passed locally.
  • 7 NPU-dependent tests could not run in the local macOS environment because the installed CPU-side torch_npu stub does not provide Tensor.npu, npu_weight_quant_batchmatmul, or npu_dynamic_quant.

• vLLM version: v0.26.0

• vLLM main: vllm-project/vllm@0351e9a

[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 refactors the Ascend fused MoE execution flow to align with the upstream router abstraction. By decoupling expert selection from quantization methods and moving it into dedicated router classes, the codebase achieves better modularity and maintainability. The changes include the introduction of new router implementations, updates to the FusedMoE patching logic, and a revised API contract for the apply method, which now receives precomputed routing information.

Highlights

• Router Abstraction: Introduced dedicated router classes (AscendFusedTopKRouter, AscendGroupedTopKRouter) to encapsulate expert selection logic.
• Decoupling Logic: Moved expert selection logic out of individual quantization methods into the new router classes, streamlining the execution flow.
• API Contract Update: Updated the apply method signature across all quantization methods to accept precomputed topk_weights and topk_ids instead of performing selection internally.
• Test Suite Refactoring: Updated unit tests across multiple quantization schemes to align with the new routing API and mock requirements.

New Features

🧠 You can now enable Memory (public preview) 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.

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 the 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 counterproductive. 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.

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

Suggested PR Title:

[Ops][Feature] Refactor Fused MoE routing and expert selection to use dedicated router classes

Suggested PR Summary:

### What this PR does / why we need it?
This PR refactors the Fused MoE routing and expert selection logic by introducing dedicated router classes (`AscendFusedTopKRouter` and `AscendGroupedTopKRouter`) and a factory (`create_ascend_fused_moe_router`). This aligns the codebase with upstream vLLM's router-based design, decoupling expert selection from quantization methods and the `RoutedExperts` class. Quantization methods now accept pre-computed `topk_weights` and `topk_ids` directly.

Feedback:
- In `vllm_ascend/ops/fused_moe/routed_experts.py`, calling `router._select_experts` with `input_ids` will raise a `TypeError` as the upstream method does not accept this argument.
- In `vllm_ascend/ops/fused_moe/router/fused_topk_router.py`, calling `.to(torch.int64)` on `forward_context.input_ids` when it is `None` will raise an `AttributeError`. A safety check should be added.
- Test assertions in `tests/ut/ops/test_fused_moe.py` should be updated accordingly.

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
Unit tests were updated to reflect the new router-based interface, and mock assertions were adjusted.

[gemini-code-assist]

The BaseRouter._select_experts method in upstream vLLM does not accept an input_ids keyword argument. Passing input_ids here will raise a TypeError: _select_experts() got an unexpected keyword argument 'input_ids' at runtime. Since input_ids is already retrieved from the forward context inside the router's _compute_routing method when needed, it should be removed from this call.

        topk_weights, topk_ids = router._select_experts(
            hidden_states=hidden_states,
            router_logits=router_logits,
        )

[gemini-code-assist]

Update the mock assertion to match the removal of the input_ids keyword argument from _select_experts.

    routed_experts.router._select_experts.assert_called_once_with(
        hidden_states=prepared_hidden_states,
        router_logits=prepared_router_logits,
    )

[gemini-code-assist]

If forward_context.input_ids is None (which can happen during profiling, dummy runs, or when not set in the context), calling .to(torch.int64) on it will raise an AttributeError. We should check that input_ids is not None before attempting to cast it.

Suggested change

	if self.tid2eid is not None:
	forward_context = get_forward_context()
	input_ids = forward_context.input_ids.to(torch.int64)
	forward_context = get_forward_context()
	if self.tid2eid is not None and getattr(forward_context, "input_ids", None) is not None:
	input_ids = forward_context.input_ids.to(torch.int64)

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.

[weijinqian0]

/rerun