[Benchmark] Add expert parallelism for tuning with benchmark_moe.py

[Chen-zexi]

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.

Purpose

This PR enhances the MoE kernel benchmark script (benchmarks/kernels/benchmark_moe.py) by adding support for tuning configurations with expert parallelism (EP).

Previously, the script only supported tuning for tensor parallelism (TP), where the intermediate size is sharded. This change introduces an --enable-expert-parallel flag, which modifies the benchmark to simulate an EP environment by sharding the experts themselves across the available devices.

Test Plan

This command should execute the tuning process and save a configuration file named according to the number of local experts per device.

## Test Command 1
CUDA_VISIBLE_DEVICES=0,1 python benchmarks/kernels/benchmark_moe.py \
       --model Qwen/Qwen3-30B-A3B \
       --enable-expert-parallel \
       --tp-size 2 \
       --tune

This command should expect to fail (test error handling)

## Test Command 2
CUDA_VISIBLE_DEVICES=0,1 python benchmarks/kernels/benchmark_moe.py \
       --model Qwen/Qwen3-30B-A3B \
       --enable-expert-parallel \
       --tp-size 1 \
       --tune

Test Result

The test commands above were executed, and the results are as expected.

• The test command 1 runs successfully, distributing the benchmark across the 2 specified GPUs and generating a valid tuning configuration file.

• The test command 2 correctly raises a ValueError with the following:

ValueError: When running with --enable-expert-parallel, the specified --tp-size must be equal to the number of available GPUs. Got --tp-size=1 and 2 GPUs.
To tune for a specific number of GPUs for expert parallel, please restrict the visible devices using the CUDA_VISIBLE_DEVICES

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Summary of Changes

Hello @Chen-zexi, 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 significantly enhances the MoE kernel benchmark script by introducing comprehensive support for Expert Parallelism (EP) tuning. Previously limited to Tensor Parallelism (TP), the script can now simulate and benchmark EP environments by correctly sharding experts across devices. This expansion provides more robust and accurate performance tuning capabilities for Mixture-of-Experts models in distributed settings.

Highlights

• New Feature: Expert Parallelism Support: The benchmark_moe.py script now supports tuning configurations for Expert Parallelism (EP) via a new --enable-expert-parallel command-line flag. This allows the benchmark to simulate an EP environment by sharding experts across available devices, complementing the existing Tensor Parallelism (TP) support.
• Updated benchmark_config Function: The benchmark_config function (lines 48-49) now accepts enable_expert_parallel and ep_size parameters. When EP is enabled, the number of experts used for creating weights is adjusted by dividing the total experts by ep_size, simulating local experts on a device.
• Enhanced BenchmarkWorker Initialization: The BenchmarkWorker class constructor (lines 391-400) has been updated to receive enable_expert_parallel, worker_id, and total_workers. This allows each worker to correctly determine its ep_rank and ep_size for distributed EP benchmarking, ensuring proper context for local expert handling.
• Dynamic shard_intermediate_size Calculation: The main function's logic for calculating shard_intermediate_size (lines 598-636) has been updated across various model architectures. It now conditionally sets shard_intermediate_size to 2 * intermediate_size when enable_expert_parallel is active, otherwise it defaults to 2 * intermediate_size // args.tp_size for TP scenarios.
• Input Validation for Expert Parallelism: New validation checks have been added in the main function (lines 683-693) to ensure that when --enable-expert-parallel is used, the specified --tp-size must be equal to the number of available GPUs and must be at least 2. This prevents misconfigurations in EP setups.
• Updated Configuration Saving Logic: The save_configs function (lines 553-572) now includes enable_expert_parallel and ep_size parameters. When EP is active, the filename for the saved configuration is generated using the local_num_experts (total experts divided by ep_size), accurately reflecting the per-device expert count.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 is currently in preview and 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 to provide feedback.

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 expert parallelism (EP) to the MoE kernel benchmark script. The changes are well-structured and the implementation appears correct. My review focuses on improving maintainability by addressing code duplication. I've identified a significant area of repeated logic in the main function that should be refactored. I've also suggested a minor refactoring in save_configs to improve conciseness. Overall, this is a valuable addition to the benchmark suite.

[gemini-code-assist]

This block can be made more concise by using a conditional expression to determine the number of experts for the filename, which avoids repeating the call to get_config_file_name.

    # For expert parallel, the filename is based on the number of local experts.
    num_experts_for_filename = (num_experts // ep_size
                              if enable_expert_parallel else num_experts)
    filename = get_config_file_name(
        num_experts_for_filename, shard_intermediate_size // 2, dtype_str,
        block_quant_shape)

[yewentao256]

Thanks for the work!

[yewentao256]

I guess only update through this can not mock for a real EP process, perhaps this is the reason why we can not get a good throughput result for tuned config. Perhaps go through fused_topk() -> expert_map, fused_experts(..., expert_map=expert_map, ...)?

[Chen-zexi]

You are right, thanks for pointing out! So, instead of shrinking number of expert before fused_topk, we should handled it in fused_experts through export_map (i.e skipped if not belongs to local experts)?

[yewentao256]

Yeah I think so, just changing the expert number is not the real case for EP.

[mgoin]

@Chen-zexi can you please fix the pre-commit failure?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Chen-zexi.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Chen-zexi]

Hi @mgoin — I’ve opened a new PR #20876 on a clean branch based on the latest upstream source. It includes a reimplementation of the EP mock following @yewentao256 ’s feedback.

This PR was closed because I accidentally created a divergent history while attempting to rebase onto upstream/main, which unintentionally triggered review requests from all code owners involved in the last 200 commits and introduced unexpected branch structure. I’ve since reverted the branch to its original state, but to avoid further confusion and noise, I decided to open a fresh PR instead.

Apologies for the mess — and sorry if the review request caused any disruption!