Add gfx950 Q256 fused MXFP4 prefill pipeline by vstakhov-amd · Pull Request #5988 · ROCm/aiter

vstakhov-amd · 2026-09-30T08:40:17Z

Add an opt-in MXFP4 MoE prefill path for gfx950 that combines gate/up projection, SiLU, hidden quantization, and down projection in one assembly producer, followed by a Triton FP32 route reduction.

The pipeline uses two compute kernels and reuses the existing Opus sorting and inverse route map. The public fused_moe API and default kernel selection remain unchanged. The included mxfp4_q256_prefill.csv enables the path for model widths 4096/8192, intermediate width 256, 513 experts, and top-k 11.

Validated on gfx950 with 256 CUs at 8k, 16k, and 32k tokens, including:

Exact agreement with an independent producer launch and with FlyDSL final outputs.
Sampled FP32 reference checks.
Graph replay with changing routes and zero routing weights.
Prequantized inputs and caller-provided output buffers.

Full AITER GPU latency for model width 8192, intermediate width 256, and 10 routed experts plus one shared expert:

 Tokens    FlyDSL, µs    Assembly pipeline, µs    Latency reduction
━━━━━━━━  ━━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━━━
  8,192       1,499.1                  1,365.5                 8.9%
────────  ────────────  ───────────────────────  ───────────────────
 16,384       2,554.3                  2,488.9                 2.6%
────────  ────────────  ───────────────────────  ───────────────────
 32,768       4,796.1                  4,267.5                11.0%

Measurements include sorting, activation quantization, scale preparation, and reduction. Values are medians of seven samples of 50 graph replays after warmup, using ROCm 7.2.3, PyTorch 2.10, and FlyDSL 0.3.4.1.

github-actions · 2026-09-30T08:48:46Z

🏷️ 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 (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
`ci:triton-355`	Run the full Triton test suite on MI35X, not only the tests the change affects
`multigpu`	Aiter multi-GPU tests on the 8-GPU runner
`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 5988 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

mjkvaak-amd · 2026-09-30T09:26:44Z

@@ -0,0 +1,7 @@
+gfx,cu_num,token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,q_dtype_w,q_type,use_g1u1,doweight_stage1,block_m,ksplit,kernelName1,kernelName2,run_1stage


This filename needs to change to tuned_fmoe.csv for AITER to pick it up.

why do we need a separate file for this though ? don't we already have one for this model ?

aiter/configs/model_configs/qwen3_8_2400b_fp4_tuned_fmoe.csv

mjkvaak-amd · 2026-09-30T09:28:21Z

+// SPDX-License-Identifier: MIT
+// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+#include "aiter_tensor.h"
+#include "aiter_ctypes_error.h"


The PR currently adds no test file.

there is a test for moe kernels already in aiter. op_tests/test_moe_2stage.py

JohnNikolay84 · 2026-10-23T09:37:00Z

@@ -73,6 +73,7 @@ def _get_flydsl_moe_kernels():
 _MOE_SORT_BACKEND = os.environ.get("AITER_MOE_SORT_BACKEND", "auto").lower()

 AUX_SORT_OPUS = "opus"
+_Q256_KERNEL = "fused_moe_mxfp4_prefill_1tg_4w_256mx1_128nx1_ps_fp32gate"


why do we need to have it here hard coded ?

JohnNikolay84 · 2026-10-23T09:40:22Z

@@ -3704,6 +3470,40 @@ def get_block_m() -> int:
            has_stage2_bias=has_stage2_bias,
        )

+    if kernelName1 in (_Q256_KERNEL, _Q256_KERNEL + "_route_weighted") and not (


this kind of check can be moved to cu level

JohnNikolay84 · 2026-10-23T09:41:09Z

+AITER_CTYPES_ERROR_DEF
+
+AITER_CTYPES_DEFINE_ENTRYPOINT_VOID(
+    fmoe_q256_producer,


what does fmoe_q256_producer do ?

It's fused 2 stage as per comment: Performs gate/up projection, SiLU, in-register hidden MXFP4 quantization, down projection and routing-weight multiplication.

Integrate the universal assembly producer and supplied Triton route reducer through AITER with guarded dispatch and an opt-in configuration. Preserve the public fused_moe API and reuse the Opus inverse route map. Validated on gfx950 at 8k, 16k, and 32k tokens with 512 routed experts plus a shared expert, including graph replay, independent producer launch checks, and comparison with the FlyDSL pipeline.

Install the matching weighted M128 producer and GPU routing planner from the 2026-10-16 PyISA package. Pass its 1024-entry map at ABI offset 128 and select it directly for compact Qwen 8K geometry, retaining the shared FlyDSL reducer. Validated on SMCI355 gfx950 PCI 0000:05:00.0: exact compact/sorted contributions; reference NRMS 0.00273150; graph/eager, prequantized and caller-output parity; weighted/map profile confirmed. Full-API graph median 1156.51 us (7 x 50 replays). CPU contracts and ABI assertions passed.

Name the weighted producer in the validated Qwen 8K/D8192 CSV row and carry kernelName through Q256 dispatch. Retain compact producer selection for the other rows and reject unsupported weighted geometry. Record the selected kernel in the existing integration test. SMCI355 gfx950 PCI 0000:15:00.0: exact compact/sorted contributions, reference NRMS 0.00273150, graph/eager/prequantized/output parity, and all 20 config collision checks passed. Full-API median 1125.49 us. Kernel trace evidence incomplete: PyTorch trace had no kernel events; rocprof capture blocked by competing KFD process.

Preserve the M128 production map while removing unused scheduling options. Add a category wrapper, trace representation, independent CPU/GPU tests and a preallocated-output benchmark. Validated on gfx950: 7 tests and 55 subtests pass; all 12 benchmark cases match the independent oracle.

…ance Move scratch allocation and producer/reducer orchestration below shared MoE dispatch. Keep semantic guards in Python, validate tensor contracts in the CU launcher, and remove the unused sorted-output path. Consolidate model configurations and replace the broken standalone harness with shared-expert, random-control, boundary, graph replay and independent-reference coverage. Retain explicit weighted selection and limit streaming reduction to the exact 8K Qwen shape. Validated on gfx950 with FlyDSL 0.4.0: six shared shapes, random control, five token boundaries and 26 config checks pass. Captured routes with synthetic activations and weights improve full-API latency by 7.84%, 6.35% and 5.40% versus main across seven rounds of 50 graph replays.

Retain only the compact and weighted producers selected by the current launcher.

Restore the existing single-candidate timing and sweep behavior. Consolidate Q256 case setup, expected-error checks and scheduler GPU fixtures while retaining shared/random routing, token boundaries, graph replay, producer selection and independent numerical checks. Validated on gfx950: all 12 Q256 cases, 7 scheduler tests and 55 subtests, and the normal CLI startup pass. Full configured Ruff 0.16.0 and Black checks pass without new suppressions.

mjkvaak-amd reviewed Sep 30, 2026

View reviewed changes

vstakhov-amd force-pushed the feature/fmoe-q256-prefill branch 3 times, most recently from b32149d to e89ada0 Compare October 2, 2026 10:48

vstakhov-amd force-pushed the feature/fmoe-q256-prefill branch from 4acd967 to 4d41912 Compare October 14, 2026 09:14

vstakhov-amd self-assigned this Oct 14, 2026

JohnNikolay84 reviewed Oct 23, 2026

View reviewed changes

vstakhov-amd added 8 commits October 23, 2026 14:54

Use compact hybrid Q256 contributions with shared FlyDSL reduction

86db5fd

Fix Ruff and Black checks for compact Q256 integration

6299936

Fix Q256 scheduling at shifted shared-expert boundaries

25451a0

vstakhov-amd force-pushed the feature/fmoe-q256-prefill branch from f60e127 to bc2b299 Compare October 23, 2026 16:51

vstakhov-amd added 2 commits October 23, 2026 16:54

Remove internal Q256 notes and unused sorted-output code object

bfbedfd

Retain only the compact and weighted producers selected by the current launcher.

Navigation Menu

Repository navigation

Add gfx950 Q256 fused MXFP4 prefill pipeline - #5988

Add gfx950 Q256 fused MXFP4 prefill pipeline#5988
vstakhov-amd wants to merge 10 commits into
mainfrom
feature/fmoe-q256-prefill

vstakhov-amd commented Sep 30, 2026

Uh oh!

github-actions Bot commented Sep 30, 2026

Uh oh!

mjkvaak-amd Sep 30, 2026

Uh oh!

JohnNikolay84 Oct 23, 2026

Uh oh!

mjkvaak-amd Sep 30, 2026

Uh oh!

JohnNikolay84 Oct 23, 2026

Uh oh!

JohnNikolay84 Oct 23, 2026

Uh oh!

JohnNikolay84 Oct 23, 2026

Uh oh!

JohnNikolay84 Oct 23, 2026

Uh oh!

vstakhov-amd Oct 23, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		@@ -0,0 +1,7 @@
		gfx,cu_num,token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,q_dtype_w,q_type,use_g1u1,doweight_stage1,block_m,ksplit,kernelName1,kernelName2,run_1stage

Conversation

vstakhov-amd commented Sep 30, 2026

Uh oh!

github-actions Bot commented Sep 30, 2026

🏷️ CI Guide

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants