[Refactor] Move FusedMoE hidden_size roundup to quant_method

[BowenBao]

• Refactor hidden_size and intermediate roundup logic to be handled by QuantMethod.
• Store padded and unpadded sizes under MoeConfig.
• Update ROCm padding logic to improve performance on mi300x. Thanks for suggestion and evaluation from @Rohan138.
• Enable Quark MXFP4 MoE with aiter backend running with padded intermediate_size.

[BowenBao]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the logic for rounding up the hidden_size in FusedMoE layers, moving the responsibility from the generic FusedMoE layer to the specific quantization methods. This is a good architectural improvement. My main feedback is about code duplication and a potential bug in QuarkOCP_MX_MoEMethod where the roundup logic is applied unconditionally for gpt_oss models, even for non-MXFP4 quantization types.

[gemini-code-assist]

Code Review

This pull request refactors the logic for rounding up the hidden size in FusedMoE layers by moving it from the generic layer.py to the specific quant_method implementations. This is a good architectural improvement, as it places quantization-specific logic where it belongs. The changes in fused_moe_method_base.py and layer.py are correct. However, this refactoring has introduced code duplication in mxfp4.py and quark_moe.py for handling gpt_oss models. I've added comments with suggestions to address this.

[BowenBao]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the hidden_size roundup logic for FusedMoE layers by moving it into the quant_method. This is a good architectural improvement as it localizes quantization-specific logic. The changes are well-structured. I've found one issue where a function is called with incorrect arguments, which I've detailed in a specific comment.

[mergify]

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

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

[Rohan138]

FYI #32307 might be relevant; I'm not sure what the pad size for gpt-oss on MI300 should be i.e. 128 or 256. Needs further investigation, haven't had time to run proper perf unfortunately.

[BowenBao]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the hidden_size and intermediate_size rounding logic in the FusedMoE layer by moving it into the quant_method. This is a significant improvement in maintainability as it centralizes quantization-specific alignment requirements (especially for MXFP4 backends) within the quantization methods themselves, rather than having brittle model-type checks in the core layer logic. The changes ensure that both the moe_config and the actual weight tensors are created with consistent, correctly padded dimensions across different hardware platforms and quantization schemes.

[robertgshaw2-redhat]

this is a nice simplification. I wonder if we can go even further by making the layer just unaware of the hidden size / intermediate size? WDYT?

[BowenBao]

I think should be do-able, see #34285 (comment), unless there are other use cases of layer.hidden_size that I'm unaware of.

[mergify]

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

uv pip install pre-commit
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 or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

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

[ChuanLi1101]

Left some comments FYI.

[hongxiayang]

Thanks @BowenBao.

BTW, I also filed an aiter issue for the MiniMax M2.1 MXFP4 TP4 case.

[ChuanLi1101]

LGTM, thanks for addressing the comments.

[fxmarty-amd]

LGTM. OCP MX emulation should be refactored as an Mxfp4Backend in a follow up PR

[BowenBao]

@tjtanaa could you help land this PR?

[BowenBao]

@robertgshaw2-redhat could you take another look if comments are all addressed and this can be landed?

[mergify]

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

uv pip install pre-commit
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

[tjtanaa]

@BowenBao can you also provide the lm-eval score for this model?

please review. Latest part of the PR enables Quark MXFP4 MoE with aiter backend running with padded intermediate_size.

Tested with MiniMax M2.1 MXFP4 TP4 Before:

Output token throughput (tok/s): 2182.16
Peak output token throughput (tok/s): 6973.00
Peak concurrent requests: 1000.00
Total token throughput (tok/s): 19639.48

After:

Output token throughput (tok/s): 3617.48
Peak output token throughput (tok/s): 14404.00
Peak concurrent requests: 1000.00
Total token throughput (tok/s): 32557.31

[mergify]

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

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