[Triton/Gluon] Add MoE GEMM variants and weight-gradient kernel

[WuLei-AMD]

Summary

Three new Triton kernels for MoE training and quantized inference.

moe_gemm_mxfp8 — Block-scaled MXFP8 grouped GEMM for MoE expert layers. All expert GEMMs are dispatched in a single kernel launch. E8M0 microscales (one uint8 per quant_block_size elements along K) are converted to FP32 power-of-two values
inside the kernel and applied per K-iteration. Supports optional per-expert bias.

moe_gemm_per_token — Per-token FP8 scale grouped GEMM. Per-token activation scales and per-expert weight scales are applied after the FP8 dot-product accumulation. Uses a group_sizes-based interface (1-D tensor of per-expert token counts)
compatible with Lumen's routing infrastructure.

moe_wgrad — Fused weight-gradient kernel operating directly on sorted_token_ids / expert_ids from moe_align_block_size. Computes dW[e] = Σ_{t→e} grad[t].T @ input[t] via tl.atomic_add, eliminating the sort+pad+bmm + CPU-GPU sync in the
backward pass. Uses @triton.autotune with reset_to_zero=["dw_ptr"] — required because atomic accumulation across benchmark trials would otherwise corrupt the final result.

Notes

• expert_id cast to tl.int64 before pointer arithmetic to prevent int32 overflow at production scale (E=256, N×K>8M: 255 × 14.7M = 3.74B > INT32_MAX)
• _moe_gemm_mxfp8_kernel asserts BLOCK_K == QUANT_BLOCK_SIZE at JIT time via tl.static_assert
• @triton.jit(repr=...) intentionally omitted on _moe_wgrad_kernel — combining @triton.autotune with repr= corrupts kernel execution on current Triton; the autotune key ["N", "K"] encodes variant identity
• _build_block_mapping extracted to moe/moe_utils.py to avoid duplication between the two GEMM wrappers
• moe/__init__.py exports all three public wrappers

Tests & Benchmarks

• test_moe_gemm_mxfp8_per_token.py: constant-input roundtrip, bias variant, empty-token guards
• test_moe_wgrad.py: analytical constant-input check, linearity, zero-padded-token guard
• bench_moe_gemm_mxfp8_per_token.py / bench_moe_wgrad.py: DSv4-representative shapes, --metric time/bandwidth

[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: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 5310 --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.

[Boss2002n]

can we split this into 3 PRs - one per kernel

[Copilot]

🟡 Changes recommended

Ragged expert groups can produce incorrect GEMM output, the weight-gradient path mishandles routing and padding, and its benchmark cannot run.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds three Triton MoE kernels for quantized grouped GEMM and training weight-gradient computation.

Changes:

• Adds MXFP8 and per-token-scaled grouped GEMM kernels and wrappers.
• Adds an atomic MoE weight-gradient kernel.
• Adds shared mapping utilities, correctness tests, and benchmarks.

File summaries

File	Description
aiter/ops/triton/moe/moe_gemm_mxfp8.py	MXFP8 GEMM wrapper.
aiter/ops/triton/moe/moe_gemm_per_token.py	Per-token FP8 GEMM wrapper.
aiter/ops/triton/moe/moe_wgrad.py	Weight-gradient wrapper.
aiter/ops/triton/moe/moe_utils.py	Builds expert block mappings.
aiter/ops/triton/_triton_kernels/moe/moe_gemm_mxfp8.py	MXFP8 Triton kernel.
aiter/ops/triton/_triton_kernels/moe/moe_gemm_per_token.py	Per-token Triton kernel.
aiter/ops/triton/_triton_kernels/moe/moe_wgrad.py	Atomic weight-gradient kernel.
op_tests/triton_tests/moe/test_moe_gemm_mxfp8_per_token.py	GEMM correctness tests.
op_tests/triton_tests/moe/test_moe_wgrad.py	Weight-gradient correctness tests.
op_tests/op_benchmarks/triton/bench_moe_gemm_mxfp8_per_token.py	GEMM benchmarks.
op_tests/op_benchmarks/triton/bench_moe_wgrad.py	Weight-gradient benchmark.

Review details

Suppressed comments (1)

aiter/ops/triton/_triton_kernels/moe/moe_wgrad.py:61

• This launchable kernel omits the config-aware repr required by the AITER Triton kernel convention. The noted autotune interaction explains the omission but does not provide the required trace identity; resolve the Triton incompatibility or use a supported launch/config structure that permits make_kernel_repr.

# NOTE: repr= is intentionally omitted here. Combining @triton.autotune with
# @triton.jit(repr=...) corrupts the kernel execution on current Triton versions.
# The autotune key ["N", "K"] already encodes the compiled variant identity.
@triton.jit

• Files reviewed: 11/11 changed files
• Comments generated: 11
• Review effort level: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[WuLei-AMD]

can we split this into 3 PRs - one per kernel

ok, I will do that

[WuLei-AMD]

This PR has been split into three parts, the other two PRs are #5338 and #5339

[Copilot]

🟡 Changes recommended

Sentinel handling, E8M0 decoding, bias indexing, and API validation contain correctness issues.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (3)

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

• Every non-empty mapping has padded slots: this upper bound is at least one larger than the maximum possible actual block count, and those slots contain expert ID -1. tl.assume(expert_id >= 0) is not a runtime guard, so these programs continue and read rhs, weight scales, and optional bias before the output mask suppresses stores. Return immediately for the sentinel, as the utility contract describes.

    expert_id = tl.load(block_expert_ids_ptr + pid_mb).to(tl.int64)
    token_offset = tl.load(block_token_offsets_ptr + pid_mb)
    token_end = tl.load(block_token_ends_ptr + pid_mb)

    tl.assume(expert_id >= 0)

aiter/ops/triton/_triton_kernels/moe/moe_gemm_mxfp8.py:155

• This bit shift does not fully decode E8M0: byte 0x00 must become 2^-127 (0x00400000), but becomes zero, and 0xFF is E8M0 NaN but becomes infinity. This disagrees with the repository's canonical decoder in aiter/utility/fp4_utils.py:213-220 and produces incorrect results for valid edge encodings.

        a_scale = (a_scale_e8m0.to(tl.uint32) << 23).to(tl.float32, bitcast=True)
        b_scale = (b_scale_e8m0.to(tl.uint32) << 23).to(tl.float32, bitcast=True)

aiter/ops/triton/_triton_kernels/moe/moe_gemm_mxfp8.py:169

• Bias indexing assumes a contiguous [E, N] tensor by hardcoding strides (N, 1), but the public API documents no layout restriction and does not make or require the bias contiguous. A same-shaped sliced/non-contiguous bias therefore reads incorrect elements; pass both bias strides or explicitly enforce/materialize the required layout.

    if HAS_BIAS:
        bias = tl.load(
            bias_ptr + expert_id * N + offs_n,
            mask=n_mask,
            other=0.0,

• Files reviewed: 6/6 changed files
• Comments generated: 6
• Review effort level: Balanced