[Triton/Gluon] Merge MoE MXFP8 GEMM into existing _moe_gemm_a8w8 kernel by WuLei-AMD · Pull Request #5658 · ROCm/aiter

WuLei-AMD · 2026-09-18T07:31:30Z

Summary

Follow-up to #5310. Per reviewer request, removes the standalone _moe_gemm_mxfp8_kernel and reuses the existing _moe_gemm_a8w8 kernel for MXFP8 MoE GEMM.

The two kernels share the same computation structure. The only difference is the FP8 format: _moe_gemm_a8w8 uses OCP float8e4nv (gfx950) while moe_gemm_mxfp8 needs AMD fnuz float8_e4m3fnuz (gfx942). A new USE_FNUZ: tl.constexpr flag handles this: when True, it replaces tl.dot_scaled (OCP-only) with manual E8M0→fp32 dequant + tl.dot.

Changes

_moe_gemm_a8w8: add USE_FNUZ=False constexpr (default preserves current behavior). Relax FP8 dtype assert to accept fnuz when USE_FNUZ=True. Fix bias dtype mismatch in SPLIT_K>1 else-branch.
moe_gemm_mxfp8 wrapper: route through _moe_gemm_a8w8(USE_FNUZ=True). Tile params moved out of Python into configs/gfx942/triton/moe/mxfp8_fnuz/DEFAULT.json.
moe_utils.py: add group_sizes_to_expt_tensors() to convert the group_sizes routing interface to the ExptHist/ExptOffs/ExptData format expected by _moe_gemm_a8w8.
Delete _triton_kernels/moe/moe_gemm_mxfp8.py (the duplicate kernel).

Test results

773 passed, 576 skipped (gfx942 does not support OCP FP8 blockscale configs) on MI308X (gfx942).

Tested: test_moe_gemm_mxfp8, test_moe_gemm_a8w8, test_moe_gemm_a8w8_blockscale.

github-actions · 2026-09-18T07:31:48Z

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

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 opt this PR out.

Copilot

🟡 Changes recommended

Critical and moderate issues remain in quantization block-size handling, launch configuration, bias strides, and routing performance.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

This PR consolidates MXFP8 MoE GEMM into _moe_gemm_a8w8 with FNUZ support, routing conversion, and gfx942 tuning configuration.

Changes:

Adds manual FNUZ dequantization support.
Routes MXFP8 through the shared A8W8 kernel.
Adds group-size routing conversion and JSON tuning.
Removes the duplicate MXFP8 kernel.

File summaries

File	Summary
`aiter/ops/triton/moe/moe_utils.py`	Adds group-size routing conversion utility.
`aiter/ops/triton/moe/moe_gemm_mxfp8.py`	Launches the shared FNUZ GEMM kernel.
`aiter/ops/triton/configs/gfx942/triton/moe/mxfp8_fnuz/DEFAULT.json`	Provides gfx942 tile configuration.
`aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a8w8.py`	Adds FNUZ computation and bias handling.
`aiter/ops/triton/_triton_kernels/moe/moe_gemm_mxfp8.py`	Removes the duplicate kernel.

Review details

Suppressed comments (3)

aiter/ops/triton/moe/moe_gemm_mxfp8.py:98

This hard-codes the expert-bias row stride to N, so a valid non-contiguous [E, N] bias (which the previous wrapper supported by passing bias.stride(0)) is indexed incorrectly. Pass the actual row stride instead of assuming contiguous storage.

    bias_stride = N if bias is not None else 0

aiter/ops/triton/moe/moe_gemm_mxfp8.py:85

The tuned num_warps and num_stages values are loaded into cfg but never passed to the kernel launch; only BLOCK_M and BLOCK_N are extracted. Consequently the new JSON tuning is ignored for launch configuration and every shape uses Triton's defaults. Pass cfg.num_warps and cfg.num_stages as launch meta-parameters.

    cfg = get_tuned_kernel_config(
        "moe", "MXFP8_FNUZ", "moe_gemm_mxfp8", _MXFP8_FALLBACK
    )
    BLOCK_M = cfg.kwargs["BLOCK_M"]
    BLOCK_N = cfg.kwargs["BLOCK_N"]

aiter/ops/triton/moe/moe_gemm_mxfp8.py:65

This new wrapper formats the full shape string eagerly on every call, even when INFO logging is disabled. In a MoE serving path this adds avoidable tuple/string work; use lazy logger placeholders so formatting only occurs when the message is emitted.

    _LOGGER.info(
        f"MOE_GEMM_MXFP8: lhs={tuple(lhs.shape)} rhs={tuple(rhs.shape)} "
        f"x_scale={tuple(x_scale.shape)} w_scale={tuple(w_scale.shape)} "
        f"block_size={quant_block_size}"
    )

Files reviewed: 5/5 changed files
Comments generated: 4
Review effort level: Lite

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

…licate - Add USE_FNUZ: tl.constexpr to _moe_gemm_a8w8 (default False, no impact on existing callers). When True, replaces tl.dot_scaled with manual E8M0→fp32 dequant + tl.dot for gfx942 fnuz FP8 (float8_e4m3fnuz). - Delete standalone _moe_gemm_mxfp8_kernel and update moe_gemm_mxfp8 wrapper to route through _moe_gemm_a8w8(USE_FNUZ=True). - Add _group_sizes_to_expt_tensors to moe_utils.py for routing conversion. - Add configs/gfx942/triton/moe/mxfp8_fnuz/DEFAULT.json for tile tuning. - Fix bias dtype mismatch in _moe_gemm_a8w8 SPLIT_K>1 else-branch. Tested: 5/5 mxfp8 tests pass on gfx942 (MI308X).

- moe_gemm_mxfp8: enforce quant_block_size == 32 (_moe_gemm_a8w8 hardcodes MX_PACK_DIVISOR=32). Remove tile values from _MXFP8_FALLBACK; BLOCK_M/BLOCK_N must come from per-arch JSON config, raise FileNotFoundError on arches without a config file. - moe_utils: rewrite group_sizes_to_expt_tensors with pure tensor ops (bucketize + cumsum), replacing per-expert .item() CPU syncs. - moe_op_gemm_a8w8_repr: add USE_FNUZ to specialization fields. - configs/gfx950/triton/moe/mxfp8_fnuz/DEFAULT.json: add tile config.

…ndant copy Three defects in the wrapper: 1. `torch.empty` for the output. The kernel writes only rows covered by ExptHist/ExptOffs, so rows a caller excludes (foreign expert under EP, -1 top-k padding) keep stale allocator contents. A NaN there survives a downstream `* 0` mask (NaN * 0 == NaN) and the top-k reduction spreads it across the token. Allocator-dependent, so it passes smoke tests. The reference in test_moe_gemm_mxfp8.py already assumes zero-fill. 2. `group_sizes_to_expt_tensors` ended in `.item()`, so the call cannot be CUDA-graph captured. `cdiv(total_tokens, block_m) + E` is a static upper bound on blocks; surplus slots get ExptData = -1, which the kernel already returns early on. Same trick `build_block_mapping()` in this file documents. New arg defaults to None, preserving the legacy path. Follow-through on review comment r4045560427 of ROCm#5658. 3. The NK->KN weight copy ran on every call. Expert weights are static after load; callers already storing K-major now get a view. Tested on gfx942: 13 shape/layout configs (incl. K=96, K=32, ragged groups with zero-token experts) at cosine >= 0.999999 vs fp32; NK and KN inputs bitwise identical; CUDA graph capture+replay passes (fails before). NaN-poisoned allocator with a half-foreign expert map: 100% NaN before, 0/16384 after. AI assistance was used. The submitter reviewed every line, ran the tests and measurements above, and is responsible for the result. Signed-off-by: akii96 <194881519+akii96@users.noreply.github.com>

The USE_FNUZ path upcasts both operands to fp32 and broadcasts both scale tiles across the K tile (plus a tl.trans) before tl.dot. That promotes an FP8 MFMA to an FP32 one at 4x operand footprint. It is never a win. ROCm#5658 asserts quant_block_size == 32 and MX_SCALE_BLOCK_K = BLOCK_K // MX_PACK_DIVISOR, so MX_SCALE_BLOCK_K is always 1: one scale per operand per K-step, constant across the dot. Reduce each scale tile to a vector and apply it to the accumulator instead, leaving the operands in FP8. Applied to the main loop and the not-EVEN_K tail; the MX_SCALE_BLOCK_K > 1 branches are preserved byte-identically. On gfx942, E=32/N=2048/K=6144, with the KN-view change from the companion wrapper fix: 2.70 -> 0.23 ms at T=8, 2.81 -> 0.35 ms at T=2048. ~11x over main, ~2.4x over the kernel ROCm#5658 replaced. Correctness across 13 shape/layout configs with random E8M0 scales: cosine >= 0.999999 vs fp32. Note for reviewers: the preserved MX_SCALE_BLOCK_K > 1 branches are unreachable (the assert pins BLOCK_K=32), and the tail branch looks wrong as written -- b_f32 is [BLOCK_K, BLOCK_N] while the tail builds b_sc_full as [BLOCK_N, BLOCK_K], with a different broadcast axis than the main loop. Left alone here; those configs are untested and not measured by this change. End to end on vLLM 0.30.1rc1, 8x MI325X, TP=8, tencent/Hy4-preview-FP8 real weights: 1.68 -> 10.01 tok/s decode, TPOT 97.65 ms, GSM8K within noise (84.00% vs 85.50% last-number, n=200, 0 errors both arms). AI assistance was used. The submitter reviewed every line, ran the tests and measurements above, and is responsible for the result. Signed-off-by: akii96 <194881519+akii96@users.noreply.github.com>

…ndant copy Three defects in the wrapper: 1. `torch.empty` for the output. The kernel writes only rows covered by ExptHist/ExptOffs, so rows a caller excludes (foreign expert under EP, -1 top-k padding) keep stale allocator contents. A NaN there survives a downstream `* 0` mask (NaN * 0 == NaN) and the top-k reduction spreads it across the token. Allocator-dependent, so it passes smoke tests. The reference in test_moe_gemm_mxfp8.py already assumes zero-fill. 2. `group_sizes_to_expt_tensors` ended in `.item()`, so the call cannot be CUDA-graph captured. `cdiv(total_tokens, block_m) + E` is a static upper bound on blocks; surplus slots get ExptData = -1, which the kernel already returns early on. Same trick `build_block_mapping()` in this file documents. New arg defaults to None, preserving the legacy path. Follow-through on review comment r4045560427 of ROCm#5658. 3. The NK->KN weight copy ran on every call. Expert weights are static after load; callers already storing K-major now get a view. Tested on gfx942: 13 shape/layout configs (incl. K=96, K=32, ragged groups with zero-token experts) at cosine >= 0.999999 vs fp32; NK and KN inputs bitwise identical; CUDA graph capture+replay passes (fails before). NaN-poisoned allocator with a half-foreign expert map: 100% NaN before, 0/16384 after. AI assistance was used. The submitter reviewed every line, ran the tests and measurements above, and is responsible for the result. Signed-off-by: akii96 <194881519+akii96@users.noreply.github.com>

The USE_FNUZ path upcasts both operands to fp32 and broadcasts both scale tiles across the K tile (plus a tl.trans) before tl.dot. That promotes an FP8 MFMA to an FP32 one at 4x operand footprint. It is never a win. ROCm#5658 asserts quant_block_size == 32 and MX_SCALE_BLOCK_K = BLOCK_K // MX_PACK_DIVISOR, so MX_SCALE_BLOCK_K is always 1: one scale per operand per K-step, constant across the dot. Scaling the [BLOCK_M, BLOCK_N] result is then equivalent to scaling the [*, BLOCK_K] operands, and keeps the inputs in FP8. Applied to the main loop and the not-EVEN_K tail. The MX_SCALE_BLOCK_K > 1 branches are preserved byte-identically; they are unreachable while the wrapper pins BLOCK_K, and they disagree on the weight-scale broadcast axis, so at most one of them is correct. gfx942, E=32/N=2048/K=6144, weights already K-major: T=8 0.99 -> 0.30 ms T=64 0.99 -> 0.31 ms T=512 0.99 -> 0.33 ms T=2048 1.00 -> 0.41 ms Correctness across 13 shape/layout configs with random E8M0 scales, including K=96, K=32 and ragged group sizes with zero-token experts: cosine >= 0.999999 vs an fp32 reference. AI assistance was used. The submitter reviewed every line, ran the tests and measurements above, and is responsible for the result. Signed-off-by: akii96 <194881519+akii96@users.noreply.github.com>

…LOCK_K == 1 (#5832) * [Perf] MXFP8 MoE: keep the MFMA in FP8 when MX_SCALE_BLOCK_K == 1 The USE_FNUZ path upcasts both operands to fp32 and broadcasts both scale tiles across the K tile (plus a tl.trans) before tl.dot. That promotes an FP8 MFMA to an FP32 one at 4x operand footprint. It is never a win. #5658 asserts quant_block_size == 32 and MX_SCALE_BLOCK_K = BLOCK_K // MX_PACK_DIVISOR, so MX_SCALE_BLOCK_K is always 1: one scale per operand per K-step, constant across the dot. Scaling the [BLOCK_M, BLOCK_N] result is then equivalent to scaling the [*, BLOCK_K] operands, and keeps the inputs in FP8. Applied to the main loop and the not-EVEN_K tail. The MX_SCALE_BLOCK_K > 1 branches are preserved byte-identically; they are unreachable while the wrapper pins BLOCK_K, and they disagree on the weight-scale broadcast axis, so at most one of them is correct. gfx942, E=32/N=2048/K=6144, weights already K-major: T=8 0.99 -> 0.30 ms T=64 0.99 -> 0.31 ms T=512 0.99 -> 0.33 ms T=2048 1.00 -> 0.41 ms Correctness across 13 shape/layout configs with random E8M0 scales, including K=96, K=32 and ragged group sizes with zero-token experts: cosine >= 0.999999 vs an fp32 reference. AI assistance was used. The submitter reviewed every line, ran the tests and measurements above, and is responsible for the result. Signed-off-by: akii96 <194881519+akii96@users.noreply.github.com> * [Perf] MXFP8 MoE: equivalence test, and keep FP32 off the fast path Two follow-ups to the previous commit. Test. The existing correctness tests already fail if the hoist is wrong (verified by mutation: dropping b_sc makes test_moe_gemm_mxfp8_correctness fail). What they do not pin down is the identity the optimization rests on, that with MX_SCALE_BLOCK_K == 1 scaling the [BLOCK_M, BLOCK_N] accumulator equals scaling the [*, BLOCK_K] operands. This asserts that directly, with scales that vary along every axis on both operands and N != M, so a transposed or wrongly broadcast scale cannot pass. Two mutants fail it: swapping the scale axes, and dropping the weight scale. The MX_SCALE_BLOCK_K > 1 branches stay uncovered because no in-tree caller can reach them. moe_gemm_mxfp8 is the only caller passing USE_FNUZ=True and it pins BLOCK_K = quant_block_size = 32 while asserting K % 32 == 0, so EVEN_K is always True and MX_SCALE_BLOCK_K is always 1; the other _moe_gemm_a8w8 caller never sets USE_FNUZ. Cleanup. a_f32/b_f32 were computed before the MX_SCALE_BLOCK_K branch and left unused on the == 1 path, which reads as if the change still promotes the MFMA to FP32. Triton was eliminating them -- timings are unchanged either way -- but moving them into the else makes the code say what it does. The broadcast comment moves with them, and the length-1 tl.sum gets a note: a_sc[:, 0] is rejected as an unsupported tensor index. Output is bitwise identical before and after the cleanup across 5 shapes (including the K=96 tail and E=32/N=2048/K=6144); no arithmetic changed. Signed-off-by: akii96 <194881519+akii96@users.noreply.github.com> * [Perf] MXFP8 MoE: read the per-K-step scale with tl.reshape Address review feedback on #5832: - Read the single per-K-step E8M0 scale with tl.reshape instead of tl.sum over the length-1 axis, in both the main loop and the tail. - Drop test_moe_gemm_mxfp8_matches_per_operand_scaling; the existing correctness test already checks the kernel against an operand-scaled fp32 reference with non-uniform scales. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> --------- Signed-off-by: akii96 <194881519+akii96@users.noreply.github.com> Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> Co-authored-by: akii96 <194881519+akii96@users.noreply.github.com>

…ndant copy Three defects in the wrapper: 1. `torch.empty` for the output. The kernel writes only rows covered by ExptHist/ExptOffs, so rows a caller excludes (foreign expert under EP, -1 top-k padding) keep stale allocator contents. A NaN there survives a downstream `* 0` mask (NaN * 0 == NaN) and the top-k reduction spreads it across the token. Allocator-dependent, so it passes smoke tests. The reference in test_moe_gemm_mxfp8.py already assumes zero-fill. 2. `group_sizes_to_expt_tensors` ended in `.item()`, so the call cannot be CUDA-graph captured. `cdiv(total_tokens, block_m) + E` is a static upper bound on blocks; surplus slots get ExptData = -1, which the kernel already returns early on. Same trick `build_block_mapping()` in this file documents. New arg defaults to None, preserving the legacy path. Follow-through on review comment r4045560427 of ROCm#5658. 3. The NK->KN weight copy ran on every call. Expert weights are static after load; callers already storing K-major now get a view. Tested on gfx942: 13 shape/layout configs (incl. K=96, K=32, ragged groups with zero-token experts) at cosine >= 0.999999 vs fp32; NK and KN inputs bitwise identical; CUDA graph capture+replay passes (fails before). NaN-poisoned allocator with a half-foreign expert map: 100% NaN before, 0/16384 after. AI assistance was used. The submitter reviewed every line, ran the tests and measurements above, and is responsible for the result. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

WuLei-AMD requested a review from a team September 18, 2026 07:31

github-actions Bot added the Triton/Gluon label Sep 18, 2026

Boss2002n requested a lite review from Copilot September 18, 2026 09:33

Copilot started reviewing on behalf of Boss2002n September 18, 2026 09:33

View session

Copilot AI reviewed Sep 18, 2026

View reviewed changes

Comment thread aiter/ops/triton/moe/moe_gemm_mxfp8.py

Comment thread aiter/ops/triton/moe/moe_utils.py Outdated

Comment thread aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a8w8.py

Comment thread aiter/ops/triton/moe/moe_gemm_mxfp8.py Outdated

zufayu requested review from a team and vgokhale September 20, 2026 01:11

ZhangDanyang-AMD force-pushed the dev_lumen_pr5310_followup branch 2 times, most recently from ff88507 to 2910fe8 Compare September 22, 2026 08:11

Dewei-Wang-sh approved these changes Sep 23, 2026

View reviewed changes

ZhangDanyang-AMD requested review from a team September 23, 2026 03:41

WuLei-AMD added 2 commits September 23, 2026 11:42

ZhangDanyang-AMD force-pushed the dev_lumen_pr5310_followup branch from 2910fe8 to 8f5b726 Compare September 23, 2026 03:42

ZhangDanyang-AMD merged commit 7e8d331 into ROCm:main Sep 23, 2026
97 of 99 checks passed

akii96 mentioned this pull request Sep 25, 2026

[Triton/Gluon] [fix] moe_gemm_mxfp8: zero-init output, graph-safe routing, skip redundant copy #5831

Closed

akii96 mentioned this pull request Sep 25, 2026

[Triton/Gluon] [Perf] MXFP8 MoE: keep the MFMA in FP8 when MX_SCALE_BLOCK_K == 1 #5832

Merged

Navigation Menu

Repository navigation

[Triton/Gluon] Merge MoE MXFP8 GEMM into existing _moe_gemm_a8w8 kernel - #5658

[Triton/Gluon] Merge MoE MXFP8 GEMM into existing _moe_gemm_a8w8 kernel#5658
ZhangDanyang-AMD merged 2 commits into
ROCm:mainfrom
WuLei-AMD:dev_lumen_pr5310_followup

WuLei-AMD commented Sep 18, 2026

Uh oh!

github-actions Bot commented Sep 18, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

WuLei-AMD commented Sep 18, 2026

Summary

Changes

Test results

Uh oh!

github-actions Bot commented Sep 18, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (3)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants