[Feature][EPLB] Add support for Qwen3 EPLB

[hsliuustc]

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 pull request add support for qwen3 moe model EPLB feature, which helps to improve the overall thoughput during LLM Serving.

#20468

Test Plan

import json
import os
import argparse
from vllm import LLM, SamplingParams

prompt = "Explain the theory of relativity in simple terms."

RESULT_FILE = "eplb_test_output.json"

sampling_params = SamplingParams(
    temperature=0.0,
    top_p=1.0,
    top_k=1,
    max_tokens=100
)

def run_inference(model_path: str, enable_eplb: bool, num_redundant_experts: int = 0):
    print(f"Running inference with EPLB={enable_eplb}, redundant experts={num_redundant_experts}")
    
    llm = LLM(
        model=model_path,
        tensor_parallel_size=4,
        enable_expert_parallel=True,
        enable_eplb=enable_eplb,
        num_redundant_experts=num_redundant_experts if enable_eplb else 0,
        eplb_window_size=1000,
        eplb_step_interval=100,
        enforce_eager=True,
        trust_remote_code=True
    )
    
    result = llm.generate([prompt], sampling_params)
    output_text = result[0].outputs[0].text.strip()
    
    print("Output:")
    print(output_text)
    print("-" * 50)

    return output_text

def save_result(key: str, value: list):
    if os.path.exists(RESULT_FILE):
        with open(RESULT_FILE, "r") as f:
            results = json.load(f)
    else:
        results = {}

    results[key] = value

    with open(RESULT_FILE, "w") as f:
        json.dump(results, f, indent=2)

    print(f"Output saved to {RESULT_FILE}")

def load_results():
    if os.path.exists(RESULT_FILE):
        with open(RESULT_FILE, "r") as f:
            return json.load(f)
    return {}

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", type=str, choices=["eplb", "normal", "compare"], required=True)
    args = parser.parse_args()

    MODEL_PATH = "/workspace/models/Qwen3-30B-A3B-FP8"

    if args.mode == "eplb":
        outputs = run_inference(MODEL_PATH, enable_eplb=True, num_redundant_experts=32)
        save_result("eplb", outputs)
    elif args.mode == "normal":
        outputs = run_inference(MODEL_PATH, enable_eplb=False)
        save_result("normal", outputs)

run the following command:

python eplb_test.py --mode eplb

Test Result

(Optional) Documentation Update

[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]

Code Review

This pull request adds support for Expert Parallelism Load Balancing (EPLB) for the Qwen3-MoE model. The changes involve plumbing EPLB configurations through the model layers, updating weight loading logic to handle distributed experts, and implementing the MixtureOfExperts protocol. To fully implement the MixtureOfExperts protocol, the update_physical_experts_metadata method is required.

[gemini-code-assist]

The Qwen3MoeForCausalLM class needs to implement the update_physical_experts_metadata method as part of the MixtureOfExperts protocol. This method is called by the EPLB scheduler during expert rebalancing, and its absence will lead to a runtime AttributeError.

    def set_eplb_state(
        self,
        expert_load_view: Tensor,
        logical_to_physical_map: Tensor,
        logical_replica_count: Tensor,
    ) -> None:
        for layer_idx, layer in enumerate(self.moe_layers):
            self.expert_weights.append(layer.get_expert_weights())
            layer.set_eplb_state(
                moe_layer_idx=layer_idx,
                expert_load_view=expert_load_view,
                logical_to_physical_map=logical_to_physical_map,
                logical_replica_count=logical_replica_count,
            )

    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
    ) -> None:
        self.num_physical_experts = num_physical_experts
        self.num_local_physical_experts = num_local_physical_experts
        for layer in self.model.layers:
            if isinstance(layer, PPMissingLayer):
                continue

            if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock):
                layer.mlp.n_physical_experts = num_physical_experts
                layer.mlp.n_local_physical_experts = num_local_physical_experts

[david6666666]

@DarkLight1337 @abmfy please review ,thanks

[abmfy]

Hi, thank you so much for the contribution! I just returned from traveling and will review this PR soon — thank you for your patience.

That said, it appears that support for Qwen3 was already assigned to @aladerran in #20468, and their PR is now open at #20815. I believe it would be great to compare both implementations and look for opportunities to merge them together, so that we can land a unified version into main.

Really appreciate your efforts — thank you again for the great work!

[abmfy]

These changes seem similar to #20815. Would it make sense to join forces and merge that PR after testing? Thank you so much!

[hsliuustc0106]

@abmfy please double check it again? previous two reviews are fixed

[hsliuustc0106]

we have tested locally, @abmfy do you have some time to double check?

[david6666666]

@DarkLight1337 could you help assign another reviewer? maybe WoosukKwon is so busy. thanks

[DarkLight1337]

Can you update tests/distributed/test_expert_parallel.py to test EPLB for each model? I don't have resources to verify the correctness of this locally

[david6666666]

Can you update tests/distributed/test_expert_parallel.py to test EPLB for each model? I don't have resources to verify the correctness of this locally

ok, I think test DeepSeek and Qwen3 first, @hsliuustc0106 @ycyaw66 .

[abmfy]

LGTM. I tested accuracy locally and the results look good.

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9075	±	0.0080
		strict-match	5	exact_match	↑	0.9007	±	0.0082

As mentioned earlier, please coordinate with @aladerran if you plan to add more tests, so we can avoid duplicating efforts.

Thanks again for the contribution!

[mergify]

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

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

[DarkLight1337]

Any update?

[david6666666]

Any update?

Sorry, we haven't had time to add the test yet, we are focusing on #22167 and #22179