Skip to content

[PERF] gfx942 FMoE: a16 kernel dispatch for GLM-5 decode (T=16-256) - #4558

Merged
alexioslyrakis-amd merged 3 commits into
mainfrom
alyr/gfx942-fmoe-optimal-pr
Aug 8, 2026
Merged

alexioslyrakis-amd merged 3 commits into
mainfrom
alyr/gfx942-fmoe-optimal-pr

Conversation

@alexioslyrakis-amd

@alexioslyrakis-amd alexioslyrakis-amd commented Aug 4, 2026 •

Copy link
Copy Markdown
Contributor

Summary

  • Route GLM-5/DeepSeek-V3 FMoE decode (E=257, topk=9, 304 CUs) token counts T=16-256 to faster a16 (bf16 activation) kernels on MI300X
  • Add int flat_mode to C++ dispatch chain, replacing bool is_flat_dispatch, to support emsort persistent grid (flat_mode=2) alongside flat one-token-per-TG (flat_mode=1)
  • Fix Python cfg_flat to preserve int value instead of bool() coercion that collapsed flat=2 → 1

Performance (E=257, topk=9, shared-expert, 304 CUs)

Tokens Kernel This PR (us) Main (us) Speedup
16 emsort_16x128_pf3 123.6 139.7 1.13x
32 a16_16x256 193.0 196.2 1.02x
64 a16_16x256 253.7 271.5 1.07x
128 a16_ps_16x256 285.3 301.6 1.06x
256 a16_ps_16x256 299.1 312.5 1.04x

All accuracy within acceptable range (logits_diff ~1.4e-4).

Test plan

  • python op_tests/test_moe_2stage.py -e 256 -k 8 --shared-expert -t 16 32 64 128 256 -q 50 on MI300X (304 CUs)
  • Verify kernel dispatch logs show expected kernel names per token count
  • Verify T=4,8 still dispatch flat_pf3 (unchanged)
  • Verify T>=512 still dispatch 32x/64x kernels (unchanged)

@alexioslyrakis-amd
alexioslyrakis-amd requested a review from a team August 4, 2026 16:20
@github-actions

github-actions Bot commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

🏷️ 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 4558 --add-label <label>

Route GLM-5/DeepSeek-V3 decode token counts 16-256 to faster a16
(bf16 activation) kernels on MI300X (304 CUs, E=257, topk=9).

Kernel routing:
  T=16  → emsort_16x128_pf3 (flat_mode=2): 123.6us vs 131.7us (6% faster)
  T=32  → a16_16x256:                       193.0us vs 197.5us (2% faster)
  T=64  → a16_16x256:                       253.7us vs 260.3us (3% faster)
  T=128 → a16_ps_16x256:                    285.3us vs 284.8us (neutral)
  T=256 → a16_ps_16x256:                    299.1us vs 299.7us (neutral)

Changes:
- Add int flat_mode to C++ dispatch (replaces bool is_flat_dispatch)
  to distinguish flat(1) from emsort(2) persistent grid modes
- Fix Python cfg_flat to preserve int value instead of bool coercion
- Add flat_mode parameter to fmoe_fp8_blockscale_g1u1 torch op
- Register a16 kernels (16x128_pf3, 16x256, ps_16x256, emsort) in
  kernel config CSVs with corresponding .co files
- Update GLM5 model config with benchmarked optimal kernel per token
The assert guarded flat kernels to gfx950-only, but gfx942 flat
kernels are now available and tested.
@alexioslyrakis-amd
alexioslyrakis-amd force-pushed the alyr/gfx942-fmoe-optimal-pr branch from 03d4293 to 68aca62 Compare August 7, 2026 12:28
@alexioslyrakis-amd
alexioslyrakis-amd merged commit 11bf108 into main Aug 8, 2026
103 of 117 checks passed
@alexioslyrakis-amd
alexioslyrakis-amd deleted the alyr/gfx942-fmoe-optimal-pr branch August 8, 2026 01:51
alexioslyrakis-amd added a commit that referenced this pull request Aug 11, 2026
Both Python and C++ read cfg.flat directly from the config, so passing
flat_mode as a parameter was redundant. The AITER_CHECK(cfg.flat ==
flat_mode) added in PR #4558 caused a crash on gfx950 because fmoe_g1u1
was never wired up to pass it.

Remove flat_mode from get_heuristic_kernel, fmoe_g1u1, and
fmoe_fp8_blockscale_g1u1. Hardcode cfg.flat == 0 in the heuristic
filter to exclude flat kernels (they are always dispatched by explicit
kernel name from the tuned config CSV).
valarLip pushed a commit that referenced this pull request Aug 11, 2026
* [fix][fmoe] propagate flat_mode through fmoe_g1u1 dispatch

PR #4558 added flat_mode to get_heuristic_kernel() and
fmoe_fp8_blockscale_g1u1 but missed fmoe_g1u1. On gfx950, when the
tuned config selects a flat kernel for the xbf16 MXFP4 path (e.g.
Qwen3.5-397B), fmoe_g1u1 calls get_heuristic_kernel with the default
flat_mode=0, triggering the AITER_CHECK(cfg.flat == flat_mode) assertion.

Add int flat_mode to the fmoe_g1u1 C++ entrypoint and Python stub, and
forward it from fused_moe_1stage.

* [fix][fmoe] remove redundant flat_mode parameter from FMoE dispatch

Both Python and C++ read cfg.flat directly from the config, so passing
flat_mode as a parameter was redundant. The AITER_CHECK(cfg.flat ==
flat_mode) added in PR #4558 caused a crash on gfx950 because fmoe_g1u1
was never wired up to pass it.

Remove flat_mode from get_heuristic_kernel, fmoe_g1u1, and
fmoe_fp8_blockscale_g1u1. Hardcode cfg.flat == 0 in the heuristic
filter to exclude flat kernels (they are always dispatched by explicit
kernel name from the tuned config CSV).
@alexioslyrakis-amd alexioslyrakis-amd self-assigned this Oct 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants