[Bugfix] Fix some issues with MoERunner PR #32344

[bnellnm]

Purpose

Move ensure_moe_quant_config call from FusedMoE.forward_native into _moe_forward and _moe_forward_shared. This is closer to how it was before when it was hidden inside the custom op. It should avoid torch.compile issues.

Fix handling of gate. The use_overlapped flag should have been checked before returning _gate.

Possible fix for #34357

Test Plan

Ran openai/gpt-oss-20b
Tested #34357, was able to repro it with a revision earlier than #32344
Ran nvidia/DeepSeek-R1-NVFP4

Test Result

cc @mgoin , @robertgshaw2-redhat

[gemini-code-assist]

Code Review

This pull request introduces two key bug fixes for the MoE runner. First, it corrects the handling of the gate property in FusedMoE by checking the use_overlapped flag, ensuring the router is invoked correctly based on whether shared expert computation is overlapped. This aligns the gate's behavior with the shared_experts property. Second, it moves the ensure_moe_quant_config_init call into the _moe_forward and _moe_forward_shared custom op implementations. This is a good change to prevent issues with torch.compile by ensuring that this initialization logic with side effects is executed at runtime rather than during graph tracing. The changes are well-reasoned and improve the correctness and robustness of the MoE implementation. I have no further comments.

[mgoin]

This seems to fix my gpt-oss H200 issue

Main: fails

This PR:

vllm serve openai/gpt-oss-20b -tp=1 --port 9000

python tests/evals/gsm8k/gsm8k_eval.py --port 9000                                                    
Running GSM8K evaluation: 1319 questions, 5-shot
Evaluating: 100%|█████████████████████████████████████████████████████████████████████████████████████████| 1319/1319 [00:37<00:00, 34.92it/s]

Results:
Accuracy: 0.293
Invalid responses: 0.174
Total latency: 37.783 s
Questions per second: 34.910
Total output tokens: 308540
Output tokens per second: 8166.192