Gfx950 mxfp4 flat

[JohnNikolay84]

Motivation

B200 decode mxfp4 moe is significantly faster than gfx950. Implementing moe on asm to close the gap.

Technical Details

Introducing a new flat mxfp4 fmoe asm kernel, it requires no prior sorting or quantization. It directly consumes bf16 token embeddings and fused_topk output. Max topk = 16.

2stage moe aiter unit test has been updated with the new input flag --ref-dtype fp32. When this flag is provided torch reference is using fp32 for intermediate quantization after activation. By default it is using bf16 as ck or flyDSL 2 stage kernels are using, but that would create a big divergence in the CheckAllClose call with asm kernels that are using fp32. fp32 or bf16 is just a choice, both should work fine, the main point is that kernel and reference should use the same ref-dtype to make the test output reasonable.

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4001 --add-label <label>