[micah-wil]

Purpose

On ROCm with AITER MoE, the shared-expert output was added to the routed output by a separate bf16 elementwise add after the MoE op. This PR has the routed experts take the shared-expert output as the stage-2 addend instead, so AITER accumulates it in the GEMM2 epilogue and the add kernel goes away. Targets MiniMax-M3.1 (shared expert, no AITER shared-expert fusion), but nothing in the accumulate path is model-specific.

Three commits:

1. Fold routed scaling into the AITER top-k weights ([ROCm][MiniMax-M3]).
   Without shared-expert fusion, the router folds routed_scaling_factor into the top-k weights (passed to AITER's biased_grouped_topk, which applies it in-kernel) instead of scaling the bf16 MoE output. The weights are bit-identical. This also removes one elementwise multiply per MoE layer.
   It is also a prerequisite for commit 2: with apply_routed_scale_to_output=False, MoERunner.routed_scaling_factor is 1.0, which is the condition the accumulate path requires. Applying the scale to the output would also scale the shared output once it is folded into the same tensor.

2. Accumulate the shared-expert output in AITER's stage 2 ([ROCm][MoE]).
   When the modular kernel's experts support it, the routed experts receive the shared-expert output as addend (plus an addend_ready event) and AITER accumulates it in stage 2. The aux-stream shared expert still overlaps with sorting and stage 1; AITER waits on its ready event only right before stage 2. The runner then uses the single-output moe_forward op, and falls back to an in-op add if the experts decline the output.
   Enabled only when all of these hold: experts report supports_output_addend(), no routed input/output transform, routed_scaling_factor == 1.0, no naive dispatch/combine, pcp_size == 1, output not already reduced, no unpadded output / hidden-dim padding.

3. Pass shared experts to the OCP MX MoE kernel ([Quark][MoE]).
   QuarkOCP_MX_MoEMethod.apply dropped shared_experts, so the modular kernel never saw them (no MK-internal overlap, no accumulation). One-line fix, required for commit 2 to take effect on Quark MXFP4 checkpoints.

Dependencies

• Needs AITER fused_moe(accumulate_output=..., output_ready_event=...): TODO link AITER PR.
• Feature-detected via rocm_aiter_ops.fused_moe_supports_accumulate_output(). On an AITER without it, supports_output_addend() is False and the existing separate-add path runs unchanged.

Scope note

Commit 3 and the accumulate path are not MiniMax-specific. Any model on the Quark OCP MX + AITER modular-kernel path that has shared experts and a runner routed_scaling_factor == 1.0 will now take the accumulate path when AITER supports it. I only tested MiniMax-M3.1.

Test Plan

• Existing MoE unit tests on gfx950 (MI355X): tests/kernels/moe/test_rocm_aiter_moe.py, test_shared_fused_moe_routed_transform.py, test_rocm_aiter_topk.py, test_rocm_aiter_num_expert_group.py, test_fused_topk.py.
• torch-profiler trace of decode steps, MiniMax-M3.1-Flash-Quark-MXFP4, TP4, DSpark3, conc 32, with and without the change.
• TODO: GSM8K 5-shot and a throughput A/B with only this PR applied.

Test Result

• Unit tests: 1010 passed, 8 skipped. This is a fallback-path run with stock AITER. The accumulate fast path was only exercised end to end with the AITER branch.

• Profile (rank 0, one decode step, 57 MoE layers). The "after" trace had every optimization in this series enabled, not this PR in isolation:

  Kernel	before	after
  AUnaryFunctor<bf16, MulFunctor<float>> (routed-scale multiply)	57 launches, 131 µs	0
  CUDAFunctor_add<bf16> (shared + routed add)	62 launches, 300 µs	5 launches, 24 µs

  The scale multiply can only come from commit 1. The attribution of the add kernels to commit 2 is by inference; I have not profiled this PR alone.

TODO: attach the trace screenshots.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @micah-wil.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Fangzhou-Ai]

Hi @micah-wil, we've been validating this on MiniMax-M3.1-Flash (Quark MXFP4, TP4, MI355X, DSpark 7 drafts, aiter 0.1.27 as pinned) and wanted to share measurements in case they help move it along.

Commit 1 (fold routed_scaling_factor into biased_grouped_topk, no aiter hook needed)

• Per layer, CUDA graph, 12 layers cycling 6 weight copies, shared MLP on the aux stream with an event join (us/layer, baseline -> folded): M=8 47.9 -> 45.7, M=32 94.3 -> 92.6, M=64 130.4 -> 128.6, M=128 166.1 -> 164.7, i.e. -1.4 to -2.6 us per layer across M=8..128. The AUnaryFunctor multiply launch (3.8-4.0 us standalone) is the only kernel that disappears.
• Full TP4 server, torch-profiler window (rank 0, all streams): elementwise kernels per decode step 202 -> 145 (-57, one per MoE layer); every other kernel family unchanged.
• Numerics: max abs difference vs the baseline path 0.031 (0.4-0.7% relative) in the microbenchmark, bf16 atomic-order noise, no NaNs.
• Step level on its own (150K context, 3 repeats): mean -0.9% over N=1..16; with 5 repeats N=1 went 9.22 -> 9.03 ms and N=2 10.20 -> 10.10 ms. The noise of that benchmark is about +-2%, so this is consistent but not significant for a single change.

Commits 2 and 3 (addend fold into stage 2)

• On aiter 0.1.27 and on current aiter main we could not find accumulate_output / output_ready_event anywhere (no matching aiter PR either), so fused_moe_supports_accumulate_output() returns False and the path falls back to the existing add.
• As an experiment we hooked 0.1.27 internals locally (moe_sorting without zeroing plus a _stage2_override that waits on the shared-output event). In the same graph microbenchmark it saved a further ~2-2.5 us per layer on top of commit 1 at M=8..128 (about 1 us at M=40). We don't see an upstreamable way to do that without an aiter API, so we did not pursue it further.

Suggestion: commit 1 stands on its own and doesn't depend on the aiter hook. Would you be open to splitting it into its own small PR so it can merge independently? We're happy to either open it from our branch (with you credited as co-author) or hand you our test (test_fused_topk_bias_folds_routed_scaling, 6 cases) for your branch, whichever you prefer.

(Measurements gathered with Claude Code assistance on a MI355X node.)