[MoE Refactor] EPLB refactoring for FusedMoE

[bnellnm]

Purpose

• Use eplb_state | None instead of enable_eplb flag + eplb_state in FusedMoE and router classes.
• Add set method to EplbLayerState.
• Update tests

Test Plan

CI

Test Result

cc @yzong-rh

---

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.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[gemini-code-assist]

Code Review

This pull request introduces the EplbManager class to centralize Expert Parallelism Load Balancing (EPLB) logic, refactoring MoE layers and routers to delegate state management and weight collection. Feedback recommends using p.detach() instead of p.data for safer tensor access and replacing assertions with explicit RuntimeErrors for critical validation to prevent issues in optimized Python environments.

[yzong-rh]

https://github.com/neuralmagic/vllm/blob/2bc4adcc0e4d5a58cf2b69cab9d6126ba8882641/vllm/distributed/eplb/eplb_state.py#L642-L647
Still expects layer.eplb_state.

https://github.com/neuralmagic/vllm/blob/2bc4adcc0e4d5a58cf2b69cab9d6126ba8882641/vllm/model_executor/layers/quantization/modelopt.py#L1934-L1937
Still expects layer.enable_eplb.

AI also found:

tests/model_executor/test_routed_experts_capture.py

• uses router.enable_eplb
• uses router.eplb_state.*

tests/kernels/moe/test_routing.py

• calls create_fused_moe_router(..., enable_eplb=..., eplb_state=...)

tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py

• sets fml.enable_eplb = True

Not caused by this refactor but likely a bug:
https://github.com/neuralmagic/vllm/blob/2bc4adcc0e4d5a58cf2b69cab9d6126ba8882641/vllm/model_executor/models/sarvam.py#L664-L668
which uses an incorrect set_eplb_state signature.

[yzong-rh]

Instead of creating a EplbManager wrapper, what if we flesh out EplbLayerState with set_state and get_expert_weights instead?
This moves the Eplb handling logic out of FusedMoE without introducing a manager class.

[bnellnm]

Instead of creating a EplbManager wrapper, what if we flesh out EplbLayerState with set_state and get_expert_weights instead? This moves the Eplb handling logic out of FusedMoE without introducing a manager class.

Good point. I've redone the PR so that there's still only EplbLayerState. It's now mostly removing the flag and using the presence of the state as an indicator of whether or not EPLB is enabled. I ended up moving the static methods on the defunct manager to other places in a later PR anyway.

[yzong-rh]

LGTM, thanks!

[robertgshaw2-redhat]

note to self, it seems weird that It would be the quant method that determines this

Shoudlnt it be the kernel?

[mergify]

Hi @bnellnm, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[ilmarkov]

LGTM. Just small nits.

[mergify]

Hi @bnellnm, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10