[ROCm][Bugfix] Fix accuracy issue on fmoe when VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS enabled

[ganyi1996ppo]

Purpose

This PR fix the accuracy issue introduced from #31221, which caused all the model with MOE structure have serious accuracy regression when shared experts fusion enabled. This PR add +grouped_topk into custom_ops to make the dispatch actually happens.

Test Plan

Test deepseek-r1 on gsm8k

Test Result

current vllm main branch:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.0099|±  |0.0027|
|     |       |strict-match    |     5|exact_match|↑  |0.0000|±  |0.0000|

With this fix

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9477|±  |0.0061|
|     |       |strict-match    |     5|exact_match|↑  |0.9454|±  |0.0063|

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request addresses a critical accuracy issue on ROCm for MoE models when VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled. The fix involves enabling the grouped_topk custom op to ensure correct kernel dispatch. While the approach is correct, the implementation could be more robust. I've suggested an improvement to handle cases where a user might have explicitly disabled this custom op, which would otherwise cause the fix to fail silently.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[tjtanaa]

@ganyi1996ppo can you check if when aiter fused moe is enabled, are we using rocm_aiter_grouped_topk by default? Because when VLLM_ROCM_USE_AITER=1, we always consider using the forward_hip from AITER. The custom op PR might have made it so that when +grouped_topk is not added to the custom op list, so even if we add VLLM_ROCM_USE_AITER=1, it most likely will still not use the forward_hip. Can you validate this and try to fix it in this PR as well?

This is the reason why I am very reluctant to change all to custom ops. The default behaviour now is controlled through custom_op list. Since the default behaviour now is controlled through custom_op list. We will need to add explicit custom op unit tests as well to catch this. Regular unit tests still run correctly with wrong code path (forward_native).

@ProExpertProg @MengqingCao @xinyu-intel Can we try to ask developers who are migrating the ops to custom ops in their PR to write unit tests for the custom op classes for all platforms because regular unit tests still run correctly with wrong code path (forward_native), because regular unit tests only tests for correctness.

[ganyi1996ppo]

@ganyi1996ppo can you check if when aiter fused moe is enabled, are we using rocm_aiter_grouped_topk by default? Because when VLLM_ROCM_USE_AITER=1, we always consider using the forward_hip from AITER. The custom op PR might have made it so that when +grouped_topk is not added to the custom op list, so even if we add VLLM_ROCM_USE_AITER=1, it most likely will still not use the forward_hip. Can you validate this and try to fix it in this PR as well?

This is the reason why I am very reluctant to change all to custom ops. The default behaviour now is controlled through custom_op list.

@tjtanaa We might still get into native when fused shared expert is not enabled, but what you said is exactly what I thought. We should add +grouped_topk as long as the aiter is enabled.

[ganyi1996ppo]

@tjtanaa take another look please, I just make +grouped_topk a default configuration.

[tjtanaa]

Let's call this directly at rocm_aiter_ops.is_fused_moe_enabled() like describe in the doc string of _aiter_ops (https://github.com/vllm-project/vllm/blob/51085c2aebe7df2c53dbe9a44d89bc3ee761793f/vllm/_aiter_ops.py#L772C8-L772C42)
topk routing functions are part of MoE related ops.

[ganyi1996ppo]

Sure

[tjtanaa]

we don't need the and because rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() already check for both as shown in line

vllm/vllm/_aiter_ops.py

Line 887 in 51085c2

return cls.is_fused_moe_enabled() and cls._MOE_SHARED_EXPERTS_ENABLED

[ganyi1996ppo]

Sure, done.

[tjtanaa]

LGTM.