[alexioslyrakis-amd]

Add gfx942 FP8 blockscale flat FMoE kernel for low-token decode (4-8 tokens per expert).

The flat kernel (a16_16x128_flat_pf3) eliminates external sorting by having each workgroup scan topk_ids internally, and fuses output zeroing (fill) into the kernel via global_atomic_pk_add_bf16. This removes the sort + fill kernel launches from the pipeline, reducing end-to-end latency at low token counts.

Changes:

• Add fmoe_bf16_a16_blockscaleFp8_g1u1_vs_silu_1tg_16x128_flat_pf3.co for gfx942
• Register kernel in the gfx942 manifest CSV with flat column
• Extend flat kernel assertion in fused_moe.py to allow gfx942
• Add GLM5.2-FP8 model config (d_model=6144, d_ff=256, experts=256, topk=8) with flat_pf3 entries for t=4,8 at 64/80/304 CU
• Add shared expert entries (e=257, k=9) for the same CU/token combinations

Test:

AITER_LOG_MORE=1 python3 -u op_tests/test_moe_2stage.py \
  -t 4 8 -d bf16 -dim 6144,256 -q 5 -a silu -e 256 -k 8 -s f \
  --no-flydsl-csv

Motivation

GLM5.2-FP8 decode workloads on gfx942 (MI300/MI325) have very few tokens per expert (typically 4-8). At these token counts, the external sort kernel overhead is disproportionate. The flat kernel avoids this by fusing sort + fill into the GEMM kernel, giving ~15-20% kernel-level speedup over CK at t=4-8.

Test Result

Verified on 304 CU MI325 with e=256/k=8 and e=257/k=9 (shared expert). logits_diff ~2e-4 (within FP8 quantization noise).

[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 4040 --add-label <label>

[alexioslyrakis-amd]

Tested with AITER_LOG_MORE=1 python3 -u op_tests/test_moe_2stage.py -t 4 -d bf16 -dim 6144,256 -q 5 -a silu -e 257 -k 9 -s f --no-flydsl-csv.

  t=4
  ┌────────┬────────────────┬──────────────┐
  │  Run   │ ASM (flat_pf3) │ CK (2-stage) │
  ├────────┼────────────────┼──────────────┤
  │ 1      │ 59.0 us        │ 64.4 us      │
  ├────────┼────────────────┼──────────────┤
  │ 2      │ 56.5 us        │ 67.0 us      │
  ├────────┼────────────────┼──────────────┤
  │ 3      │ 57.7 us        │ 64.6 us      │
  ├────────┼────────────────┼──────────────┤
  │ 4      │ 56.9 us        │ 67.6 us      │
  ├────────┼────────────────┼──────────────┤
  │ 5      │ 55.1 us        │ 66.5 us      │
  ├────────┼────────────────┼──────────────┤
  │ median │ 56.9 us        │ 66.5 us      │
  └────────┴────────────────┴──────────────┘
  t=8
  ┌────────┬────────────────┬──────────────┐
  │  Run   │ ASM (flat_pf3) │ CK (2-stage) │
  ├────────┼────────────────┼──────────────┤
  │ 1      │ 78.0 us        │ 100.0 us     │
  ├────────┼────────────────┼──────────────┤
  │ 2      │ 82.1 us        │ 97.2 us      │
  ├────────┼────────────────┼──────────────┤
  │ 3      │ 80.7 us        │ 99.8 us      │
  ├────────┼────────────────┼──────────────┤
  │ 4      │ 112.1 us       │ 98.2 us      │
  ├────────┼────────────────┼──────────────┤
  │ 5      │ 78.6 us        │ 102.0 us     │
  ├────────┼────────────────┼──────────────┤
  │ median │ 80.7 us        │ 99.8 us      │
  └────────┴────────────────┴──────────────┘  

Conclusions:

• ASM is 1.17x faster at t=4.
• ASM is 1.24x faster at t=8 (ignoring the run 4 outlier at 112 us).
• precision: logits_diff is very small for all ASM runs

[alexioslyrakis-amd]

Performance results vs main (post PR #3422)

After #3422 ("GLM-5.1 tunings") merged, the CK 2-stage kernels for this workload were replaced by ASM 1-stage kernels on main. We repeated the comparison against the new baseline.

Note: "1-stage" means the full MOE pipeline — GEMM0 (gate) + GEMM1 (up) + SiLU + GEMM2 (down) — runs in a single kernel launch, as opposed to the old CK "2-stage" which used separate launches for stage1 and stage2.

Kernels used

Branch	Kernel	Dispatches per iter
main	fmoe_bf16_blockscaleFp8_g1u1_vs_silu_1tg_32x128	3 (quant + sorting + fmoe)
flat (this PR)	fmoe_bf16_a16_blockscaleFp8_g1u1_vs_silu_1tg_16x128_flat_pf3	1 (flat, no quant/sort)

Full MOE pipeline (profiler avg us/iter, gfx942 304 CUs)

tokens	main (us)	flat (us)	Speedup	Time saved
4	83.2	58.9	1.41x	24.3 us (29%)
8	103.7	81.3	1.27x	22.4 us (22%)

Breakdown of gains

Source	t=4 (us)	t=8 (us)
Eliminated quant kernel	3.6	3.6
Eliminated sorting kernel	10.8	11.3
Faster fmoe kernel (16x128 flat vs 32x128)	10.0	7.5
Total saved	24.3	22.4

Accuracy (logits_diff vs BF16 reference)

Branch	t=4	t=8
main	7.8e-5 — 1.2e-4	7.0e-5 — 1.1e-4
flat (this PR)	1.1e-4 — 1.9e-4	1.4e-4 — 2.2e-4

Workload: GLM-5, E=257, topk=9, D=6144, D_FF=256, blockscale FP8, gfx942.

[alexioslyrakis-amd]

Updated performance results (50 profiler iterations, clean GPU)

Previous measurements used only 5 profiler iterations, which caused noisy/bimodal results. Re-ran both ASM flat and main with 50 profiler iterations on a clean GPU (no other workloads), 5 independent runs each.

Workload: GLM-5 FMoE decode, E=257, topk=9, dim=6144x256, A8W8 blockscale FP8, SiLU, gfx942

Device time (profiler, ROCTracer HW timestamps)

Tokens	Main kernel	Main (us)	ASM flat kernel	ASM flat (us)	Speedup
4	fmoe_..._vs_silu_1tg_32x128	83.50 (stdev 0.29)	fmoe_..._a16_..._16x128_flat_pf3	52.81 (stdev 0.27)	1.58x
8	fmoe_..._vs_silu_1tg_32x128	95.27 (stdev 2.39)	fmoe_..._a16_..._16x128_flat_pf3	76.62 (stdev 2.16)	1.24x

cuda.Event time (end-to-end GPU including host dispatch overhead)

Tokens	Main (us)	ASM flat (us)	Speedup
4	240.98	181.72	1.33x
8	258.94	209.95	1.23x

Summary

Tokens	Metric	Latency reduction
4	Device time	37% less than main
4	cuda.Event	25% less than main
8	Device time	20% less than main
8	cuda.Event	19% less than main

Both main and ASM flat use 1-stage kernels (full MOE pipeline: GEMM0 gate + GEMM1 up + SiLU + GEMM2 down in a single kernel launch). The ASM flat kernel uses FLAT instructions for global memory access and reads activations in BF16 (X_BF16) instead of FP8, with a 16x128 tile size vs main's 32x128.