[Triton/Gluon] Add per token scale to a8w8 moe gemm

[nsusanto]

Motivation

Add per-token activation scales and per-expert weight scales to the Triton a8w8 MoE GEMM

Technical Details

• moe_gemm_a8w8 takes two new optional args:
  • x_token_scale: fp32, one scale per row of x
  • w_expt_scale: fp32 [n_expts], one scale per expert
• Both scales are applied to the fp32 accumulator in the epilogue, before bias, activation, and fused quant. Per-token rows follow GatherIndx when it is set.
• moe_gemm_mxfp8 passes None for both.
• Fix the fnuz (USE_FNUZ) MX path's handling of the last, partial K block (not EVEN_K). It now builds the weight block-scale the same way as the main K loop. The old code applied scales to the wrong K elements and only compiled when BLOCK_N == BLOCK_K. No current caller reaches this path, because moe_gemm_mxfp8 uses BLOCK_K = 32 and requires K % 32 == 0.

Test Plan

• test_moe_gemm_a8w8.py: new token_expt_scales

Test Result

• gfx950: test_moe_gemm_a8w8.py + test_moe_gemm_mxfp8.py 645 passed, 512 skipped

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

[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 (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 6106 --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.

[lburzawa]

LGTM

[Copilot]

Copilot review overview

🟡 Changes recommended

The unrelated FNUZ tail-path correction should be moved to a focused follow-up PR with dedicated coverage.

Review effort: Balanced
Findings: 1

Open (1)

• Move unrelated FNUZ tail-path layout fix to a follow-up PR · New

What changed in this PR

Adds optional per-token activation and per-expert weight scaling to the Triton A8W8 MoE GEMM.

Changes:

• Validates and forwards the new scale tensors.
• Applies scales before bias and activation.
• Extends tests and updates the shared-kernel caller.

File	Description
op_tests/​triton_tests/​moe/​test_moe_gemm_a8w8.py	Tests token/expert scaling.
aiter/​ops/​triton/​moe/​moe_op_gemm_a8w8.py	Exposes and validates scale arguments.
aiter/​ops/​triton/​moe/​moe_gemm_mxfp8.py	Updates the shared-kernel call signature.
aiter/​ops/​triton/​_triton_kernels/​moe/​moe_op_gemm_a8w8.py	Applies scales and modifies FNUZ tail handling.

---

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

[azaidy]

LGTM!