[AMD] Add MoE weights and scales padding

[mqhc2020]

Motivation

Right now, Aiter MoE requires weights and scales to align with a fixed number. Since some models have intermediate sizes that don't fit this rule, we need to add extra padding to the weights so they can be processed by the Fused MoE.

Modifications

Add padding for the weights. Below listed are the models and configurations that has been verified with:

1. Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 : TP=8
2. amd/Qwen3-235B-A22B-Instruct-2507-MXFP4: TP=8
3. amd/Qwen3.5-397B-A17B-MXFP4: TP=8
4. zai-org/GLM-4.7: TP=8

Accuracy Tests

We run GSM8K accuracy test for models below:

Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 (TP8):

Serve command:

SGLANG_USE_AITER=1 python3 -m sglang.launch_server \
	--model-path ${MODEL_PATH} \
	--tp-size 8 \
	--ep-size 1 \
	--trust-remote-code \
	--chunked-prefill-size 131072 \
	--mem-fraction-static 0.9 \
	--attention-backend aiter

Accuracy: 0.980
Invalid: 0.000
Latency: 10.525 s
Output throughput: 2663.037 token/s

amd/Qwen3-235B-A22B-Instruct-2507-MXFP4 (TP8):

Serve command:

SGLANG_USE_AITER=1 python3 -m sglang.launch_server \
	--model-path ${MODEL_PATH} \
	--tp-size 8 \
	--ep-size 1 \
	--trust-remote-code \
	--chunked-prefill-size 131072 \
	--mem-fraction-static 0.9 \
	--attention-backend aiter

Accuracy: 0.940
Invalid: 0.000
Latency: 23.853 s
Output throughput: 7488.441 token/s

zai-org/GLM-4.7 (TP8):

Serve command:

SGLANG_USE_AITER=1 python3 -m sglang.launch_server \
	--model $MODEL_PATH \
	--tp 8 \
	--tool-call-parser glm47 \
	--reasoning-parser glm45

(better than VLLM TP4 on model card)

Accuracy: 0.975
Invalid: 0.000
Latency: 7.349 s
Output throughput: 2683.189 token/s

Serving command:
python3 ./benchmark/gsm8k/bench_sglang.py --num-questions 1319 --parallel 1319

Benchmarking and Profiling

Baseline benchmark for amd/Qwen3-235B-A22B-Instruct-2507-mxfp4 (TP4) and Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 (TP8):

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Summary of Changes

Hello, 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 refactors the handling of Mixture-of-Experts (MoE) weights and scales, introducing a unified padding mechanism tailored for AMD (HIP) platforms utilizing Aiter. It centralizes the logic for determining padding sizes and calculating the actual dimensions of MoE weights, ensuring proper memory allocation and data handling. The changes also adapt the weight loading and quantization compatibility checks to seamlessly integrate with the new Aiter-specific padding, improving robustness and performance for MoE layers.

Highlights

• Centralized MoE Padding Logic: Introduced new utility functions to centralize the determination of Mixture-of-Experts (MoE) padding sizes and the calculation of padded weight dimensions, specifically for Aiter-enabled environments.
• Aiter-Specific Weight Handling: Implemented Aiter-specific padding logic for MoE weights and scales, ensuring correct memory allocation and dimension adjustments during weight creation and loading, particularly for AMD (HIP) platforms.
• Dynamic Weight Dimensioning: Updated MoE weight and scale creation processes across various quantization schemes (FP8, W8A8, Quark W4A4) to dynamically adjust dimensions based on the new Aiter padding requirements.
• Refined Quantization Compatibility Checks: Modified block quantization compatibility checks to conditionally bypass certain alignment requirements when Aiter-specific padding is active, accommodating its unique padding behavior.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

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 padding for MoE weights and scales, primarily for AMD hardware support using AITriton. The core logic is refactored into new utility functions get_moe_padding_size and get_moe_weight_sizes in srt/layers/moe/utils.py, which is a good approach to centralize the logic. The changes across various quantization schemes and model layers correctly adopt these new utilities.

My review includes a few suggestions to improve maintainability and readability, such as addressing code duplication and simplifying complex calculations. Overall, the changes are well-structured and seem correct for the intended purpose.

[hubertlu-tw]

Other than a few comments I left, LGTM. Notice that this PR will also resolve an issue reported here:
#21918 (comment)

CC: @HaiShaw

[HaiShaw]

@mqhc2020 please run a verification on grok2 and Grok INT4-FP8 MoE model.

[hubertlu-tw]

@mqhc2020 could you please also add the commands you ran on the PR description? It may help the review process. Thanks!

[mqhc2020]

@mqhc2020 please run a verification on grok2 and Grok INT4-FP8 MoE model.

Here are the accuracy results:

Grok1 INT4-FP8

Serve command:

RCCL_MSCCL_ENABLE=0 SGLANG_USE_AITER=1 SGLANG_INT4_WEIGHT=1 python3 -m sglang.launch_server \
    --model $MODEL_PATH \
    --tokenizer-path Xenova/grok-1-tokenizer \
    --tp 8 \
    --quantization fp8 \
    --trust-remote-code \
    --attention-backend aiter

Accuracy: 0.825
Invalid: 0.000
Latency: 141.465 s
Output throughput: 1338.350 token/s

Grok2

Serve command:

SGLANG_USE_AITER=1 python3 -m sglang.launch_server \
            --model $MODEL_PATH \
            --tokenizer-path alvarobartt/grok-2-tokenizer \
            --tp 8 \
            --quantization fp8 \
            --trust-remote-code \
            --chunked-prefill-size 131072 \
            --attention-backend aiter

Accuracy: 0.929
Invalid: 0.000
Latency: 177.871 s
Output throughput: 840.041 token/s

Serving command:
python3 ./benchmark/gsm8k/bench_sglang.py --num-questions 1319 --parallel 1319

Note that:

1. Also checked that the MoE path for grok-1 and grok-2 is not affected (not block quant)
2. There was a problem related to pre-shared models and that was considered before.

[hubertlu-tw]

LGTM

[hubertlu-tw]

/rerun-failed-ci

[hubertlu-tw]

@HaiShaw could you please review this PR again? This will resolve an issue reported by #21918