[shantipriya-amd]

Summary

Re-submission of #5708 (Relu2 Tier 2 CK-Tile MoE GEMM support), which was reverted
because the CK submodule pin it introduced was not fetchable from CK's develop
branch (built on two local, un-upstreamed commits rather than an official commit).

Per reviewer guidance on #5708, this PR follows the 3-step recovery path:

1. The two orphan CK commits (Add MoeRelu2 activation functor and
   Fix half_t intrinsics) were submitted and merged upstream to CK develop
   via [composablekernel] Add MoeRelu2 activation functor and fix half_t intrinsics rocm-libraries#13071.
2. The CK submodule was repinned to the official develop commit containing
   both changes (96afae1d54fc79c26857de7e1683b0dde1ff3c6f) via [CK] Repin CK submodule to fetchable develop commit (relu2/half_t fix) #6178.
3. This PR re-submits the original Relu2 Tier 2 code on top of that clean,
   fetchable CK pin -- rebased onto current main to pick up the repin and
   all intervening changes.

What's in this PR

Same Relu2 (gate-only, non-gated) CK-Tile two-stage MoE GEMM support as #5708:

• aiter/fused_moe.py: Relu2 dispatch in get_2stage_cfgs / _fused_moe_impl,
  including the gate_mode/bias/use_g1u1/doweight_stage1 guard rails specific
  to the CK-Tile kFFN_gemm1_gate_only stage1 instance.
• csrc/ck_tile_gemm_moe_2stages/*: CK-Tile instance generation and kernel
  wiring for the Relu2 activation.
• csrc/kernels/activation_kernels.cu, csrc/include/aiter_enum.h,
  csrc/include/rocm_ops.hpp: Relu2 ActivationType plumbing.
• aiter/ops/quant.py, csrc/kernels/quant_mxfp4.cu,
  csrc/kernels/quant_mxfp6_gemm.cu: supporting quant-path changes.
• op_tests/test_moe_2stage.py: Relu2 test coverage.

No functional changes versus #5708 beyond what was needed to rebase cleanly
onto current main (CK submodule pin, upstream file moves/renames, and a few
small fixes squashed into the branch's commit history -- e.g. gating
enable_bias to Swiglu only, rejecting split_k=0 in the Relu2 non-split
dispatch branch, and rejecting unsupported scatter/activation-scale Relu2
CK-Tile paths).

CK submodule

Pinned to 96afae1d54fc79c26857de7e1683b0dde1ff3c6f (CK develop), confirmed
fetchable and containing both upstreamed commits from
ROCm/rocm-libraries#13071.

Related

• Original PR (reverted): [HIP] [CK] Add relu2 (Squared ReLU) activation support to CK-Tile fused MoE GEMM #5708
• Upstream CK PR (two commits merged to CK develop): [composablekernel] Add MoeRelu2 activation functor and fix half_t intrinsics rocm-libraries#13071
• CK submodule repin PR: [CK] Repin CK submodule to fetchable develop commit (relu2/half_t fix) #6178

[Copilot]

🔵 Needs a closer look

Unsupported Relu2 configurations fall into incompatible backends, and the default Relu2 test sweep generates unsupported cases.

1 open finding

• Add partial-vector and empty-tensor Relu2 bounds tests

1 resolved since last review

• Relu2 a16w4 tests use incorrect activation and reference setup

Previously missed (2)

In code that hasn't changed since last review

Reject unsupported Relu2 capabilities before generic fallback

aiter/​fused_moe.py:3730

These exclusions only skip the Relu2 CK-Tile branch; they do not reject the unsupported call. For example, a Relu2 call with has_stage2_scatter=True or activation scales falls through to the generic dispatch: with ksplit > 1 it selects a split CK-Tile path for which no Relu2 instance is generated, and otherwise it selects ck_moe_stage1, whose activation mapper rejects Relu2. This turns the intended guard rail into a later launch failure. Reject unsupported Relu2 capabilities before fallback, or route them to a backend that implements Relu2.

Restrict Relu2 to supported quantization triple

op_tests/​test_moe_2stage.py:1255

This enables Relu2 for every quantization triple handled by the generic branch, although the new runtime path only supports _PER1X32_BF16_FP4. Running the natural sweep with -a relu2 and no -q therefore starts with QuantType.No, reaches the legacy CK stage1 path, and fails because that mapper does not support Relu2; the sweep never reaches the supported q6 case. Filter Relu2 out unless the triple is _PER1X32_BF16_FP4, or explicitly require q6 for Relu2 sweeps.

🧠 Review effort: Balanced

[Copilot]

🟡 Changes recommended

Relu2 is not exercised by default CI, and the added numerical checks do not currently fail on incorrect results.

3 open findings

• Assert checkAllclose result to catch scalar-tail regressions · New
• Enable strict accuracy for new Relu2 test cases · New
• Add partial-vector and empty-tensor Relu2 bounds tests

🧠 Review effort: Balanced

[yifehuan]

LGTM

[Copilot]

🟡 Changes recommended

The new bounds and fused-MoE tests do not currently exercise or assert the intended regressions reliably.

2 open findings

• Cover partial-vector sizes and fail on comparison errors · New
• Add partial-vector and empty-tensor Relu2 bounds tests

2 resolved since last review

• Enable strict accuracy for new Relu2 test cases
• Assert checkAllclose result to catch scalar-tail regressions

🧠 Review effort: Balanced

[zufayu]

merge it if all ci passed

[zufayu]

@shantipriya-amd, we are reverting this through #6460 because it introduced a gfx950 MiniMax-M3-MXFP4 regression.

The failing runtime combination is SwiGLU + has_bias=false + split_k=1. The generated CK-Tile dispatchers do not cover the required non-split/no-bias path, so the new fail-loud guard reports it as unsupported. Please do not work around this by removing the guard.

Please continue the fix in a follow-up PR. Before re-landing, it should:

• cover the required non-split/no-bias SwiGLU dispatch path;
• remove the empty primary-dispatcher undefined behavior;
• add a targeted test for (SwiGLU, has_bias=false, split_k=1);
• pass MiniMax-M3-MXFP4 TP4/TP8 validation through ci:vllm.

This is a temporary rollback of the regression, not a rejection of the Relu2 support. We can re-land it once the missing dispatch coverage is fixed and validated.