Fix/gfx1201 bf16 g1u1 small m moe

[keneoneth]

Motivation

Add a narrow gfx1201/RDNA4 BF16 G1U1 MoE direct path for very small decode-sized Qwen3.6 shapes.

This is part of the Qwen3.6 gfx1201 AITER MoE enablement stack.

Issue being addressed:

• [Issue]: AITER fused MoE is not usable on gfx1201/R9700S: no gfx1201 dispatch entry and fmoe_g1u1 finds no suitable kernel #4492

Required dependency:

• [Bugfix][Triton] Avoid RDNA4 unified attention LDS overflow #4385

The generic MoE path is too expensive for very small M decode cases, and unsupported gfx12 shapes can otherwise fall into CK paths that are not valid for the problem. This PR adds a tightly guarded direct BF16 path for the tiny-M tier.

This PR intentionally does not claim broader gfx12 MoE shape support. Follow-up PRs will add tuned config data and optimized tiled/direct-route kernels for larger small-M shapes.

Technical Details

This PR adds a direct BF16 G1U1 MoE path guarded to:

• gfx1201 only
• BF16 activations and weights
• unquantized MoE
• Silu / G1U1
• canonical W1/W2 layout
• no expert parallel mask
• no bias
• no hidden/intermediate padding
• very small M, currently M <= 4

Implementation summary:

• Adds a new fmoe_g1u1_bf16_small_m entrypoint.
• Adds stage1 direct computation for Silu(gate) * up into a BF16 activation workspace.
• Adds stage2 direct down projection with top-k weighted accumulation.
• Adds Python dispatch in aiter/fused_moe.py for the guarded tiny-M gfx1201 BF16 G1U1 case.
• Extends MoE tests to cover the new path.

Follow-up PR stack:

• PR2: gfx1201 Qwen3.6 tuned FMoE config data.
• PR3: gfx1201 BF16 G1U1 small-M tiled/direct-route kernel support and expanded shape coverage.

Related vLLM safety work:

• vLLM fallback PR: route unsupported gfx12 AITER MoE serving shapes to Triton until the AITER shape coverage is complete.

Test Plan

Static validation:

python3 -m compileall \
  aiter/fused_moe.py \
  aiter/ops/moe_op.py \
  op_tests/test_moe_2stage.py

git diff --check origin/main..HEAD

Targeted GPU validation:

python3 op_tests/test_moe_2stage.py

Performance validation should compare Qwen3.6 BF16 tiny-M cases against the previous AITER path and the vLLM Triton MoE path, with the same routing inputs and canonical weights.

Test Result

Static validation passed:

Compiling 'aiter/fused_moe.py'...
Compiling 'aiter/ops/moe_op.py'...
Compiling 'op_tests/test_moe_2stage.py'...
git diff --check origin/main..HEAD: PASS

Local branch:

fix/gfx1201-bf16-g1u1-small-m-moe
commit: 24ce35303 Add gfx1201 BF16 G1U1 small-M direct MoE path

Diff summary:

aiter/fused_moe.py          |  85 +++++++++++++++
aiter/ops/moe_op.py         |  13 +++
csrc/py_itfs_cu/asm_fmoe.cu | 247 ++++++++++++++++++++++++++++++++++++++++++++
op_tests/test_moe_2stage.py |  53 +++++++++-

GPU performance/correctness results should be added before marking ready for full review.

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:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
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 4740 --add-label <label>