[Triton/Gluon] [fix] moe_gemm_mxfp8: zero-init output, graph-safe routing, skip redundant copy by akii96 · Pull Request #5831 · ROCm/aiter

akii96 · 2026-09-25T01:06:23Z

Motivation

Three defects in the moe_gemm_mxfp8 wrapper. They are grouped because they share one file pair and one theme: how the wrapper handles its output buffer, its routing conversion, and its weight layout.

1. Uninitialized output buffer

The kernel writes only rows covered by ExptHist/ExptOffs, that is [0, sum(group_sizes)). A caller that deliberately excludes rows, such as a foreign expert under EP or -1 top-k padding, leaves the tail of torch.empty holding whatever the caching allocator last put there.

Those rows are not safe to ignore. A NaN there survives a downstream * 0 mask, because NaN * 0 == NaN, and the top-k reduction spreads it across the whole token. The failure depends on allocator state, so it passes smoke tests and corrupts output later.

The reference in op_tests/triton_tests/moe/test_moe_gemm_mxfp8.py already builds its expected result with torch.zeros, so this change makes the kernel match its own test.

2. `group_sizes_to_expt_tensors` blocks CUDA graph capture

total_blocks = int(n_blocks_per_expert.sum().item())  # single sync at end

.item() is a device to host sync, so the call fails under graph capture with hipErrorStreamCaptureUnsupported. vLLM captures the MoE forward, so the op cannot be used there today.

This is also unfinished review feedback on #5658. Comment r4045560427 asked for the mapping to be built "with tensor operations (or reuse the existing GPU-side helper)". The per-expert Python loop was removed; the final .sum().item() stayed.

The fix reuses the pattern that build_block_mapping() in this same file already documents. Every expert wastes at most one partial block, so cdiv(total_tokens, block_m) + E bounds the real count with no sync. Surplus slots are stamped ExptData = -1, which the kernel already returns early on (moe_op_gemm_a8w8.py:271-273). Verified over 2000 random routings.

total_tokens is optional and defaults to None, which preserves the exact legacy path for existing callers.

3. The weight is re-copied on every call

rhs_kn = rhs.permute(0, 2, 1).contiguous()
w_scale_kn = w_scale.permute(0, 2, 1).contiguous()

MoE expert weights are static after load, so this re-materializes the whole tensor on every invocation. That costs 2.4 ms for rhs alone at E=32, N=2048, K=6144 on gfx942, and it transiently doubles the weight footprint. Callers that already store K-major now get a view instead.

Test Plan

pytest op_tests/triton_tests/moe/test_moe_gemm_mxfp8.py -v

Test Result

MI325X (gfx942), ROCm 7.2.3. Existing tests pass unchanged.

check	result
correctness, 13 shape/layout configs	cosine >= 0.999999 vs fp32
NK vs KN inputs	bitwise identical output
CUDA graph capture + replay	fails on main, passes here
NaN-poisoned allocator, half-foreign map	100% NaN before, 0/16384 after

Shapes cover K=96 and K=32, which exercise the partial and single K-block cases, plus ragged group sizes [0,7,64,1,100,0,84,0] including zero-token experts.

Notes for reviewers

Please add ci:triton-300x. The default Triton CI targets gfx950, and these paths are exercised on gfx942.

A companion kernel-side perf change lives on akii96:perf/mxfp8-moe-fp8-mfma. The two touch disjoint files and either can land first. The per-call timings quoted there assume the KN view from this PR.

AI assistance was used to investigate and prepare this change. The submitter has reviewed every changed line, ran the tests and measurements reported above, and is responsible for the result.

…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>

github-actions · 2026-09-25T01:06:52Z

🏷️ 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 5831 --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.

…ch shapes Regression coverage for the defects fixed in the parent commit. Each of the first two fails on unpatched main: test_moe_gemm_mxfp8_unwritten_rows_are_zero poisons the caching allocator with NaN, then accounts for only half the tokens in group_sizes. On main the excluded tail comes back NaN. test_moe_gemm_mxfp8_cuda_graph_capture captures and replays the call. On main this raises hipErrorStreamCaptureUnsupported from the .item() in group_sizes_to_expt_tensors. test_moe_gemm_mxfp8_kn_weights_match_nk asserts the zero-copy K-major path is bitwise identical to the copying path. test_moe_gemm_mxfp8_ragged_groups uneven group sizes including zero-token experts and a partial trailing block. Against unpatched main: 3 failed, 8 passed. With the parent commit: 11 passed. Also adds the Hy4-preview-FP8 gemm1 shape as sharded at TP=8 (E=32, N=2048, K=6144) to the benchmark, at one and sixteen tokens per expert, so the numbers quoted in the PR can be reproduced with in-tree tooling. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

Copilot

Copilot review overview

🟡 Changes recommended

Address the routing edge cases, centralize layout conversion, and correct benchmark traffic accounting.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 3

Open (3)

Move weight layout conversion to shared shuffle utility · New
Account for zero-initialization write in bandwidth model · New
Include output zero-fill in benchmark traffic · New

What changed in this PR

This PR fixes MXFP8 MoE GEMM output initialization, graph-safe routing metadata, and redundant K-major weight copies.

Changes:

Zero-initializes unwritten output rows.
Adds bounded, graph-compatible routing conversion.
Reuses K-major weight layouts and expands tests and benchmarks.

File	Summary
`op_tests/triton_tests/moe/test_moe_gemm_mxfp8.py`	Adds coverage for zero tails, graph capture, layouts, and ragged routing.
`op_tests/op_benchmarks/triton/bench_moe_gemm_mxfp8.py`	Adds Hy4 benchmark cases; traffic accounting needs to include zero-initialization.
`aiter/ops/triton/moe/moe_utils.py`	Adds graph-safe routing conversion; empty-expert handling and sparse-bound tightening are required.
`aiter/ops/triton/moe/moe_gemm_mxfp8.py`	Zero-fills output and reuses K-major layouts; layout conversion should use the shared helper.

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

akii96 · 2026-09-25T06:24:15Z

+    # Rows excluded from group_sizes are never written; zero-fill so they do
+    # not surface stale allocator contents. A NaN there survives a caller's
+    # `* 0` mask (NaN * 0 == NaN) and spreads through the top-k reduction.
+    out = torch.zeros(total_tokens, N, dtype=out_dtype, device=lhs.device)


Correct that the write is unmodelled, but it is bounded and small. The governing ratio is output bytes over weight bytes, 2T/(E*K); I swept it past any real model, on gfx942:

case 2T/(E*K) zeros gemm overhead

Hy4 TP8 decode 0.00x 0.0047 ms 0.9161 ms 0.5%

DSv4 decode 0.01x 0.0061 ms 0.3431 ms 1.8%

DSv4 big prefill 2.00x 0.0377 ms 5.7063 ms 0.7%

small-K MoE 4.00x 0.0082 ms 0.3022 ms 2.8%

tiny-K, huge T 32.00x 0.0230 ms 0.5132 ms 4.7%

It cannot run away: as the output grows relative to the weights, the GEMM writes those same bytes anyway. Skipping it conditionally needs sum(group_sizes) on the host, which is the .item() sync this PR removes for CUDA-graph capture. Happy to double the output term in mem or relabel as logical bandwidth if we prefer the strict accounting.

akii96 · 2026-09-25T06:24:24Z

+    # Hy4-preview-FP8 gemm1 as sharded at TP=8: 256 experts / 8 ranks,
+    # hidden 6144, moe_intermediate 2048.
+    (32, 2048, 6144, 1, "Hy4-TP8-decode"),
+    (32, 2048, 6144, 16, "Hy4-TP8-prefill"),


Same as above!

Review feedback: the local _as_kn helper was a second home for a layout permute, which README.md:348-355 explicitly forbids -- "Do not re-implement a shuffle (the telltale is a local view -> permute -> contiguous helper); if a new layout is needed, add it here." Adds moe_weight_kn() to utils/shuffle.py next to moe_weight_decode_view(), the existing zero-copy (E, N, K) layout view, and documents it in the shuffle table so the next caller finds it instead of writing a third copy. No behaviour change: a K-major input still gets a free view, anything else is still copied, and the two are bitwise identical. 11 tests pass. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

Review feedback on ROCm#5832 asked for fewer, reference-based tests. Apply the same here: kn_weights_match_nk, ragged_groups folded into test_moe_gemm_mxfp8_correctness as parameters, so the K-major view and the ragged/zero-token groups are checked against the torch reference instead of bespoke equality tests. _make_mxfp8_inputs takes optional group_sizes, removing the hand-rolled input setup. unwritten_rows_are_zero one shape instead of two; poisons a single output-sized block instead of 32 MB, and compares exactly. cuda_graph_capture one warm-up call instead of three. Tests in the file: 11 -> 11, test lines vs main: +101 -> +68 -9. Verified on MI325X (gfx942): 11 passed. Negative controls still bite: reverting torch.zeros -> torch.empty fails unwritten_rows_are_zero, and dropping total_tokens (restoring the .item() sync) fails cuda_graph_capture. black and ruff 0.16.0 (CI pin) clean. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

akii96 · 2026-10-05T15:51:17Z

Used the feedback from my other PRs to preaddress concerns about over-engineered tests in the latest commit 😅

Test cleanup (commit ac877cf)

The K-major weight and ragged/zero-token group cases are now parameters of the existing test_moe_gemm_mxfp8_correctness, so they're checked against the torch reference instead of having their own tests.
unwritten_rows_are_zero runs on one shape instead of two and fills only an output-sized buffer with NaN first (it was 32 MB before).
cuda_graph_capture does one warm-up call instead of three.

This is the wrapper half of the now merged #5832 and it removes the per-call weight transpose, makes routing CUDA-graph safe, and zero-inits the output so masked rows can't leak NaNs.

Once CI is green, could you take a look when you have a moment @nsusanto @lburzawa ? Thanks and sorry for repeated pings!

lburzawa

@akii96 thanks for code updates. Unfortunately we realized there has been some code committed that doesn't meet standards of aiter moe stack. Please see related PRs for justification and corrections to codebase.

#6106
#6035

Let us know if you are able to update your PR to conform to the new moe codebase. We would prefer you use or update long-existing moe_gemm_a8w8 wrapper instead of moe_gemm_mxfp8.py.

Replaces ROCm#5831. Per review, the fixes land on the long-standing moe_gemm_a8w8 wrapper instead of moe_gemm_mxfp8.py, which ROCm#6035 removes. - Select USE_FNUZ from the input dtype. float8_e4m3fnuz inputs previously failed to compile because tl.dot_scaled rejects fp8e4b8. - Plain FP8 (static, per-token or per-expert scales) runs tl.dot on the FP8 operands; scales stay in the fp32 epilogue. - MXFP8 uses BLOCK_K=32 so each K step carries one E8M0 scale per operand and takes the FP8 MFMA path from ROCm#5832 instead of upcasting to fp32. Tiles come from the new gfx942 A8W8_MXFP8_FNUZ dispatch table, keyed by block_m so decode and prefill get separate configs; long-K decode shapes use split-K. - Split-K partials accumulate in fp32 and are reduced to the output dtype. - The matmul output is zero-initialised only when a scatter can leave rows unwritten. - dynamic_mxfp8_quant uses the fnuz range (max 240) for float8_e4m3fnuz. Weights stored as [E, N, K] are consumed through the [E, K, N] transposed view, so no per-call copy is made. Tests (gfx942): test_op_fnuz and test_op_mxfp8_fnuz, 111 passed.

Replaces ROCm#5831. Per review, the gfx942 support lands on the existing moe_gemm_a8w8 wrapper instead of moe_gemm_mxfp8.py, which ROCm#6035 removes. - Pass USE_FNUZ from the input dtype. float8_e4m3fnuz inputs previously failed to compile because tl.dot_scaled rejects fp8e4b8. - Plain FP8 on gfx942 runs tl.dot on the FP8 operands; scales stay in the fp32 epilogue. - MXFP8 on gfx942 uses BLOCK_K=32 so each K step carries one E8M0 scale per operand and takes the FP8 MFMA path from ROCm#5832. - Add the tuned dispatch that a8w4 and a16w4 already use, keyed by bm<block_m>_n<N>_k<K>, with a gfx942 table for Hy4-preview-FP8 at TP8. NUM_KSPLIT is honoured for the long-K gate/up shapes. Untuned shapes and other archs keep the existing heuristic. - dynamic_mxfp8_quant uses the fnuz range (max 240) for float8_e4m3fnuz. The test adds one gfx942 case per path (plain FP8, MXFP8 decode, MXFP8 prefill) to test_moe_gemm_a8w8.py.

akii96 · 2026-10-07T13:04:54Z

Thanks @lburzawa for the suggestion. Closing in favor of #6219

Following your suggestion, the new PR extends the existing moe_gemm_a8w8 wrapper instead of moe_gemm_mxfp8.py, so it works with the cleanup in #6035 and #6106. It covers gfx942 FP8 and MXFP8 and uses the same tuned dispatch as a8w4 and a16w4.

I will open that PR once I finish validations E2E

gfx942 uses the fnuz FP8 encoding, which tl.dot_scaled does not accept, so moe_gemm_a8w8 could not run there and MXFP8 MoE models fell back to dequantizing expert weights to BF16 at load. USE_FNUZ now follows the input dtype. Plain FP8 runs tl.dot on the FP8 operands. MXFP8 uses BLOCK_K=32 so each K step carries one E8M0 scale per operand and stays on the FP8 MFMA path from ROCm#5832; the tile is fixed before num_stages is chosen so pipelining matches what launches. MXFP8 shapes on gfx942 read a tuned table through the same dispatch as a8w4 and a16w4, keyed bm<block_m>_n<N>_k<K>, with NUM_KSPLIT for long K. Plain FP8 keeps the existing heuristic, since the 32-wide K tile is only right for MX. dynamic_mxfp8_quant scales fnuz to its own range (max 240) instead of the OCP 448, which otherwise overflows and drops roundtrip cosine to 0.989. Replaces ROCm#5831. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>

akii96 requested a review from a team September 25, 2026 01:06

github-actions Bot changed the title ~~[fix] moe_gemm_mxfp8: zero-init output, graph-safe routing, skip redundant copy~~ [Triton/Gluon] [fix] moe_gemm_mxfp8: zero-init output, graph-safe routing, skip redundant copy Sep 25, 2026

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

akii96 added the ci:triton-300x label Sep 25, 2026

Boss2002n requested a lite review from Copilot September 25, 2026 02:07

Copilot started reviewing on behalf of Boss2002n September 25, 2026 02:07

View session

Copilot AI reviewed Sep 25, 2026

View reviewed changes

azaidy requested review from lburzawa and nsusanto September 25, 2026 14:02

azaidy assigned lburzawa and nsusanto Sep 25, 2026

jamesETsmith mentioned this pull request Sep 29, 2026

[ROCm][Perf][Tracking Issue]: Hy4-Preview vllm-project/vllm#58606

Open

1 task done

akii96 mentioned this pull request Sep 30, 2026

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

Merged

akii96 force-pushed the fix/mxfp8-moe-wrapper branch from 3b16b12 to ac877cf Compare October 5, 2026 15:37

lburzawa requested changes Oct 6, 2026

View reviewed changes

akii96 mentioned this pull request Oct 7, 2026

[Triton/Gluon] moe_gemm_a8w8: fnuz FP8 and MXFP8 on gfx942 #6219

Open

akii96 closed this Oct 7, 2026

Navigation Menu

Repository navigation

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

[Triton/Gluon] [fix] moe_gemm_mxfp8: zero-init output, graph-safe routing, skip redundant copy#5831
akii96 wants to merge 4 commits into
ROCm:mainfrom
akii96:fix/mxfp8-moe-wrapper

akii96 commented Sep 25, 2026

Uh oh!

github-actions Bot commented Sep 25, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

akii96 Sep 25, 2026

Uh oh!

akii96 Sep 25, 2026

Uh oh!

akii96 commented Oct 5, 2026

Uh oh!

lburzawa left a comment

Uh oh!

akii96 commented Oct 7, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

case	`2T/(E*K)`	zeros	gemm	overhead
Hy4 TP8 decode	0.00x	0.0047 ms	0.9161 ms	0.5%
DSv4 decode	0.01x	0.0061 ms	0.3431 ms	1.8%
DSv4 big prefill	2.00x	0.0377 ms	5.7063 ms	0.7%
small-K MoE	4.00x	0.0082 ms	0.3022 ms	2.8%
tiny-K, huge T	32.00x	0.0230 ms	0.5132 ms	4.7%

Conversation

akii96 commented Sep 25, 2026

Motivation

1. Uninitialized output buffer

2. group_sizes_to_expt_tensors blocks CUDA graph capture

3. The weight is re-copied on every call

Test Plan

Test Result

Notes for reviewers

Uh oh!

github-actions Bot commented Sep 25, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

akii96 Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

akii96 Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

akii96 commented Oct 5, 2026

Uh oh!

lburzawa left a comment

Choose a reason for hiding this comment

Uh oh!

akii96 commented Oct 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

2. `group_sizes_to_expt_tensors` blocks CUDA graph capture

akii96 commented Oct 7, 2026 •

edited

Loading