[shantipriya-amd]

[MoE] Fix missing SwiGLU non-split no-bias CK-Tile MXFP4 dispatch (follow-up to #6376 revert)

Summary

Follow-up to #6376 (Relu2 Tier 1+2 MoE GEMM support), which was reverted via #6460/#6463 due to a regression in the amd/MiniMax-M3-MXFP4 TP4 production job: moe_dispatch: unsupported (activation=2, has_bias=0, split_k=1) combination.

Root cause

Two independent gaps combined to produce this failure:

1. Codegen gap (the real root cause). csrc/ck_tile_gemm_moe_2stages/gen_instances.py derived has_bias as a single value per (act_type, is_split_k) pair instead of sweeping it independently. For SwiGLU non-split (split_k<=1), this meant only the has_bias=True kernel instance was ever generated. The C++ dispatch branch in moe_dispatch() for (SwiGLU, no-bias, split_k<=1) already existed and was correct — it called moe_gemm2_heuristic_dispatcher<...,2,false,false>::dispatch() — but with no generated specialization behind it, the call fell through to the empty primary (unspecialized) template, which is undefined behavior (confirmed: a literal empty-bodied function with a non-void return type). This was silent UB before [HIP] [CK] Re-submit Relu2 Tier 2 MoE GEMM support (clean CK pin) #6376 added a fail-loud TORCH_CHECK guard — the guard didn't introduce the bug, it correctly surfaced a pre-existing one.

2. Caller/metadata mismatch. get_2stage_cfgs() hardcoded has_bias=True, stage2_has_bias=True for every CK-Tile SwiGLU MXFP4 branch, regardless of the caller's real has_stage1_bias/has_stage2_bias. vLLM passes bias1=None, bias2=None for MiniMax-M3 (a4w4, not use_mxfp4_w4a16), so even after closing the codegen gap this mismatch would have resurfaced the same failure class on a future shape.

Fix

• gen_instances.py: the outer codegen loop now sweeps has_bias in [True, False] for SwiGLU in both split_k modes (previously derived a single value from is_split_k). Also fixed a related bug this exposed: gen_heuristic_dispatch()'s has_bias template substitution was computing an iteration-invariant value that would have mislabeled the new no-bias instances as bias instances — now derives it from the actual sampled kernel instance.
• fused_moe.py: get_2stage_cfgs()'s two CK-Tile SwiGLU MXFP4 branches now report the caller's real has_stage1_bias/has_stage2_bias instead of a hardcoded True.
• test_moe_2stage.py: added test_swiglu_mxfp4_nonsplit_nobias_dispatch(), which asserts correct CK-Tile routing (has_bias=False, stage2_has_bias=False, ksplit<=1) and runs a passing end-to-end dispatch (no TORCH_CHECK/crash).

Verification

• Targeted test (test_swiglu_mxfp4_nonsplit_nobias_dispatch) passes: correct dispatch-family/metadata assertions, end-to-end kernel call completes without raising.
• Broader sweep on the real MiniMax-M3 shape (model_dim=6144, inter_dim=1536, E=129, topk=5) across token counts {1, 3, 5, 32, 128} x bias {on, off}: all 10 cases complete with no crashes; accuracy (logits_diff) is nearly identical between bias-on and bias-off at every token count, confirming the fix removes bias as a point of failure without changing kernel numerics.
• Note: this specific CK-Tile SwiGLU-MXFP4 shape has an independent, pre-existing logits_diff≈0.07-0.075 accuracy characteristic in both the old (bias) and new (no-bias) configurations — reproduced identically either way, so it predates and is unrelated to this fix. Worth a separate investigation; not a blocker here.

Checklist (per revert-PR follow-up requirements)

• 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 (existing coverage, label applied)

[Copilot]

🟡 Changes recommended

Unresolved weight-layout and bias-forwarding regressions can change MoE results, and test coverage needs correction.

18 open findings

• Preserve legacy interleaved layout for shuffled weights · New
• Restore bias eligibility for heuristic FlyDSL routes · New
• SwiGLU gate drops supported SiLU biases
• Use normalized weights_shuffled for Relu2 route selection
• Avoid enabling unsupported SiLU bias specializations
• Skip unsupported Relu2 FP8/FP4 combinations
• Make missing dispatcher specializations fail loudly
• Skip unsupported Relu2 cases in FP8/FP4 sweep · New
• Partial-vector handling lacks unaligned and bounds coverage
• Regression test does not fail on incorrect output
• Add non-aligned sizes to cover vectorized tail loops
• Assert numerical mismatch ratio in regression test
• Skip unsupported Relu2 FP8/MXFP4 sweep configurations
• Reject or support biased SiLU metadata dispatch
• Test fully no-bias SwiGLU execution
• Restrict Relu2 fallback to supported quantization
• Correct regression test claims about split-k routing
• Cover multi-vector partial-tail handling

Previously missed (1)

In code that hasn't changed since last review

Add unaligned lengths covering partial vectors and scalar remainders

op_tests/​test_activation.py:661

Every nonempty length here selects vec_size=1 on wave32 and wave64. These cases check the missing second load, but never exercise a partial vector or either scalar remainder loop in activation_kernel_vec_safe. Add unaligned lengths above the vectorization thresholds. On wave64, 129/131, 257/261, and 513/521 cover partial first and second vectors at widths 2, 4, and 8.

🧠 Review effort: Balanced

[Copilot]

🟡 Changes recommended

The regression test uses separated weight packing despite requesting interleaved inputs, leaving the intended layout untested.

19 open findings

• Restore bias eligibility for heuristic FlyDSL routes
• Preserve legacy interleaved layout for shuffled weights
• SwiGLU gate drops supported SiLU biases
• Use normalized weights_shuffled for Relu2 route selection
• Avoid enabling unsupported SiLU bias specializations
• Skip unsupported Relu2 FP8/FP4 combinations
• Make missing dispatcher specializations fail loudly
• Regression test does not use interleaved inputs · New
• Skip unsupported Relu2 cases in FP8/FP4 sweep
• Partial-vector handling lacks unaligned and bounds coverage
• Regression test does not fail on incorrect output
• Add non-aligned sizes to cover vectorized tail loops
• Assert numerical mismatch ratio in regression test
• Skip unsupported Relu2 FP8/MXFP4 sweep configurations
• Reject or support biased SiLU metadata dispatch
• Test fully no-bias SwiGLU execution
• Restrict Relu2 fallback to supported quantization
• Correct regression test claims about split-k routing
• Cover multi-vector partial-tail handling

🧠 Review effort: Balanced

[Copilot]

🟡 Changes recommended

Two layout-routing defects and a regression-test reference dtype mismatch remain unresolved.

20 open findings

• Require dtype and explicit layout compatibility · New
• Restore bias eligibility for heuristic FlyDSL routes
• SwiGLU gate drops supported SiLU biases
• Use normalized weights_shuffled for Relu2 route selection
• Avoid enabling unsupported SiLU bias specializations
• Skip unsupported Relu2 FP8/FP4 combinations
• Make missing dispatcher specializations fail loudly
• Preserve dtype fallback when gate layout is unspecified · New
• Use resolved activation dtype consistently across test stages · New
• Skip unsupported Relu2 cases in FP8/FP4 sweep
• Partial-vector handling lacks unaligned and bounds coverage
• Regression test does not fail on incorrect output
• Add non-aligned sizes to cover vectorized tail loops
• Assert numerical mismatch ratio in regression test
• Skip unsupported Relu2 FP8/MXFP4 sweep configurations
• Reject or support biased SiLU metadata dispatch
• Test fully no-bias SwiGLU execution
• Restrict Relu2 fallback to supported quantization
• Correct regression test claims about split-k routing
• Cover multi-vector partial-tail handling

2 resolved since last review

• Preserve legacy interleaved layout for shuffled weights
• Regression test does not use interleaved inputs

🧠 Review effort: Balanced

[Copilot]

🔵 Needs a closer look

Cross-backend layout routing and GPU kernel changes need human confirmation with validation tied to the current revision.

20 open findings

• Require dtype and explicit layout compatibility
• Restore bias eligibility for heuristic FlyDSL routes
• SwiGLU gate drops supported SiLU biases
• Use normalized weights_shuffled for Relu2 route selection
• Avoid enabling unsupported SiLU bias specializations
• Skip unsupported Relu2 FP8/FP4 combinations
• Make missing dispatcher specializations fail loudly
• Use resolved activation dtype consistently across test stages
• Preserve dtype fallback when gate layout is unspecified
• Skip unsupported Relu2 cases in FP8/FP4 sweep
• Partial-vector handling lacks unaligned and bounds coverage
• Regression test does not fail on incorrect output
• Add non-aligned sizes to cover vectorized tail loops
• Assert numerical mismatch ratio in regression test
• Skip unsupported Relu2 FP8/MXFP4 sweep configurations
• Reject or support biased SiLU metadata dispatch
• Test fully no-bias SwiGLU execution
• Restrict Relu2 fallback to supported quantization
• Correct regression test claims about split-k routing
• Cover multi-vector partial-tail handling

🧠 Review effort: Balanced

[Copilot]

🟡 Changes recommended

Unresolved layout-routing regressions select incompatible weight interpretations, and the expanded test sweep includes an unsupported Relu2 combination.

21 open findings

• Reject unsupported bf16 interleaved SwiGLU layout · New
• Skip unsupported Relu2 cases in FP8/MXFP4 sweep · New
• Require dtype and explicit layout compatibility
• Restore bias eligibility for heuristic FlyDSL routes
• SwiGLU gate drops supported SiLU biases
• Use normalized weights_shuffled for Relu2 route selection
• Avoid enabling unsupported SiLU bias specializations
• Skip unsupported Relu2 FP8/FP4 combinations
• Make missing dispatcher specializations fail loudly
• Use resolved activation dtype consistently across test stages
• Skip unsupported Relu2 cases in FP8/FP4 sweep
• Partial-vector handling lacks unaligned and bounds coverage
• Regression test does not fail on incorrect output
• Add non-aligned sizes to cover vectorized tail loops
• Assert numerical mismatch ratio in regression test
• Skip unsupported Relu2 FP8/MXFP4 sweep configurations
• Reject or support biased SiLU metadata dispatch
• Test fully no-bias SwiGLU execution
• Restrict Relu2 fallback to supported quantization
• Correct regression test claims about split-k routing

And 1 more that still need to be addressed.

1 resolved since last review

• Preserve dtype fallback when gate layout is unspecified

Previously missed (1)

In code that hasn't changed since last review

Preserve dtype fallback for untagged weight layouts

aiter/​fused_moe.py:4135

Preserve the bf16/fp16 fallback when layout metadata is absent. The generic shuffle_weight(..., layout=(16, 16)) path sets is_shuffled but not gate_up_interleave. An untuned, separated SwiGLU call with bf16 activations and GGUU weights therefore loses the standard-layout split-K path. The later CK-Tile fallback uses post_activation_layout="auto", which treats PyTorch FP4 weights as interleaved: it pairs the wrong gate/up groups, or raises if stage1 bias is present. Keep the previous dtype-based inference for untagged weights while honoring explicit layout metadata.

🧠 Review effort: Balanced

[yifehuan]

LGTM