[ROCm][Perf] Use flydsl moe with Minimax-M3 mxfp8 weights on gfx950 and implemented moe-backend selection

[hongxiayang]

Purpose

[ROCm][Perf] Use flydsl moe with Minimax-M3 mxfp8 weights on gfx950

aiter's flydls moe has shown perf improvement on various scenrios on mxfp8 serving on gfx950. This PR is to integrate the
support, and also have the capability to fall back to triton mxfp8 dot-scaled implementation.

Implementation and usability decisions:

Reused moe-backend aiter and triton, and refactored the moe selection logic.

Gating (the usability decision): enable FlyDSL via

• --moe-backend aiter (explicit, capability-gated, no env),

Reason is the AITER master flag can cause lots of env needing to be disabled.

For the triton backend, added support to choose triton-native dot-scaled backend via:
--moe-backend triton-native

Added ep support in the flydsl backend as well.

Note: To invoke this aiter's flydsl moe backend, it needs to have aiter to include flydsl commit (ROCm/aiter#3811).

Thanks for the feedback from Bowen and TJ.

Test Plan

• use nightly as base, updated aiter.
  run benchmarking and accuracy.

• regression:
  use nightly which does not have aiter's flydsl, make sure regression works.

• added a unit test for expected backend selection behavior.

Test Result

Sample vllm serve command:

vllm serve /weights \
  --served-model-name mm3 \
  --tensor-parallel-size 4 \
  --gpu-memory-utilization 0.90 \
  --max-model-len 9472 \
  --max-num-seqs 256 \
  --block-size 128 \
  --attention-backend TRITON_ATTN \
  --no-enable-prefix-caching \
  --moe-backend aiter \
  --port 8011 \
  --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}'

8k-in / 1k-out, sweep concurrency

for C in 1 32 64 128; do
  vllm bench serve --backend vllm --model mm3 --tokenizer /weights \
    --base-url http://localhost:8011 --dataset-name random \
    --random-input-len 8192 --random-output-len 1024 --ignore-eos \
    --num-prompts $((C==1?8:2*C)) --max-concurrency $C
done

**1k-in / 1k-out (decode-weighted) **

vllm bench serve --backend vllm --model mm3 --tokenizer /weights \
  --base-url http://localhost:8011 --dataset-name random \
  --random-input-len 1024 --random-output-len 1024 --ignore-eos \
  --num-prompts 128 --max-concurrency 64

eval

lm_eval --model local-completions \
  --model_args model=mm3,base_url=http://localhost:8011/v1/completions,num_concurrent=200,tokenized_requests=False,tokenizer_backend=None,timeout=5000,max_length=4096 \
  --tasks gsm8k --num_fewshot 5 --output_path /tmp/gsm8k

config	throughput	TPOT	gsm8k (full)
TP8 1k/1k	+13..16%	-12..14%	neutral (0.9075 vs 0.9105)
TP8 8k/1k	+8..16%	-8..14%	--
TP4 8k/1k	+12..20%	-11..18%	neutral (0.9105 vs 0.9113)

---

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.

[hongxiayang]

the backend selection test result:

Smoke prompt Q: 17 times 24?\nA: → expect 408.

#	serve config	backend selected (log)	outcome
2a	auto (no --moe-backend)	AITER_MXFP8 (FlyDSL)	ready 200, smoke 408 — PASS
2b	--moe-backend aiter	AITER_MXFP8 (FlyDSL)	ready 200, smoke 408 — PASS
3a	--enable-expert-parallel (ep=4)	native CDNA4 … dot_scaled	ready 200, smoke 408, no crash — PASS
3b	EP + --moe-backend aiter	— (rejected at init)	clear startup ValueError, server refused to start (port→000), no mid-run crash — PASS

[hongxiayang]

@BowenBao Thanks for the review comments. I updated the code to address your comments. please check again.
cc @tjtanaa for additional insights.

[hongxiayang]

working to get ep support.

[tjtanaa]

LGTM.

[tjtanaa]

@hongxiayang

can you try fixing this test?

CI #73688
it is related to the changes in new pr,

(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]     self.fused_experts.apply(
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py", line 517, in apply
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]     result = rocm_aiter_fused_experts(
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]              ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py", line 369, in rocm_aiter_fused_experts
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]     return rocm_aiter_ops.fused_moe(
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]            ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]   File "/usr/local/lib/python3.12/dist-packages/vllm/_aiter_ops.py", line 2209, in fused_moe
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]     return torch.ops.vllm.rocm_aiter_fused_moe(
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]   File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 1209, in __call__
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]     return self._op(*args, **kwargs)
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]            ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]   File "/usr/local/lib/python3.12/dist-packages/vllm/_aiter_ops.py", line 186, in _rocm_aiter_fused_moe_impl
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]     return fused_moe(
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229]            ^^^^^^^^^^
(EngineCore pid=2705) ERROR 06-23 09:51:21 [core.py:1229] TypeError: fused_moe() got an unexpected keyword argument 'swiglu_limit'
(EngineCore pid=2705) Traceback (most recent call last):

[tjtanaa]

Taking a closer look, this PR will break existing code as swiglu_limit is not supported by fused_moe aiter API of v0.1.13.post1, we will have to wait till we upgrade aiter version before merging.

[fxmarty-amd]

Is it really the correct place to do that?

[fxmarty-amd]

Should it not be the role of a function similar to

vllm/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py

Line 280 in d490b98

def convert_to_nvfp4_moe_kernel_format(

.

[tjtanaa]

Thanks for pointing this out. I am current revising the existing design to see how we could best address this in vLLM framework.

Regarding to this, historically we do not need to modify the property of the parameter tensors.
So, if you look at the convert_to_nvfp4_moe_kernel_format and other variant of conversion functions, the kernels are returning tensors. Then in the last step of the weight loading process, we are copying the tensors back into a new set of tensor object torch.Parameters . So current abstraction is not fully bullet proof. There are a few PRs has been mentioning that tensor properties and even layout information are discarded and missing at the end of the parameter instantiation. So, for now, the safest place is to add this w1.is_shuffled = True here in function itself of mxfp8 path specifically.

[fxmarty-amd]

I think having a new experts class is fine.

It seems up to now https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py was already eventually calling into rocm_aiter_ops.fused_moe through some if/elif/else.

[hongxiayang]

Taking a closer look, this PR will break existing code as swiglu_limit is not supported by fused_moe aiter API of v0.1.13.post1, we will have to wait till we upgrade aiter version before merging.

the PR bumping aiter to v0.1.16.post2 #46692 has landed

[hongxiayang]

I have run the unit tests locally for the amd-kernels-quantization-test-1-mi325-1 and amd-kernels-quantization-test-2-mi325-1, which are all passed.
THe other test which failed before is an nvidia end-to-end test, which may related to infra-frakyness.

Don't think there is anything blocking this PR from merged.

cc @tjtanaa