[zcnrex]

MiniMax-M3: allow shared-experts fusion on ROCm gfx942 and newer

shared_experts_fusion_disable_reason currently refuses on any non-CUDA device, so ROCm never gets shared-experts fusion even where the hardware supports it. This lifts that guard and gates it at gfx942 instead.

minimax_m3_vl.py carries the identical gate. MiniMaxM3SparseForCausalLM and MiniMaxM3SparseForConditionalGeneration are separate EntryClasses with separate shared_experts_fusion_disable_reason classmethods, and the served MiniMax-M3-MXFP8 checkpoint is the VL variant (model_type: minimax_m3_vl, per configs/minimax_vl.py) — so patching only one file would be a silent no-op on the model people actually run.

Measurement status — stated plainly

This PR has no benchmark or accuracy numbers of its own, and I am not claiming any.

The earlier combined change (this gate + the alternate-stream overlap + a prefill_attention_backend=aiter serving default) measured +46.6%..+55.4% output throughput on 8x MI350X. That number belongs to the bundle, not to this PR, and decomposing it showed the input-throughput share came entirely from the serving default, not from either code change.

More importantly: in every run measured, the server reported shared_experts_fusion=False on both sides — the fusion this PR enables never actually engaged, because it also requires the model config and serving flags to request it. So the throughput effect of this gate specifically is unmeasured.

What this PR is, therefore, is an enablement gate: it removes a platform restriction that has no reason to exist on gfx942+. It should be reviewed on that basis, and someone with a configuration that genuinely turns fusion on should measure it before it is treated as a performance change.

The serving-defaults and cookbook half is a separate change and is not part of this PR.

---

CI States

Latest PR Test (Base): ❌ Run #35153622399
Latest PR Test (Extra): ❌ Run #35153622200
Latest PR Test (AMD ROCm 10): ❌ Run #35153622232

[skimanwhite]

Tested this on gfx942 and it does enable fusion as intended, but it causes a
reproducible functional regression in multi-turn tool calling. Reporting it here
since the PR notes there are no accuracy numbers yet — hopefully this is useful rather
than discouraging, because the perf motivation looks real.

Environment

• 8× MI300X (gfx942), ROCm 7.2
• lmsysorg/sglang:v0.5.17-rocm720-mi30x
• MiniMax-M3, bf16 weights, --kv-cache-dtype fp8_e4m3, --tp-size 8
• attention_backend=triton, moe_runner_backend=triton, speculative_algorithm=None
• Checkpoint is the VL variant: model_type: minimax_m3_vl,
  architectures: ["MiniMaxM3SparseForConditionalGeneration"]
• Patch hand-ported to v0.5.17 (that tree uses an elif chain where main uses early
  return), applied to both minimax_m3.py and minimax_m3_vl.py. Boot log
  confirms Shared experts fusion optimization enabled.

Symptom

With fusion enabled, the model cannot terminate an agent loop. It re-issues the same
tool calls indefinitely and never emits a final assistant message. Over three identical
12-turn runs:

	fusion ON	fusion OFF
tool calls	36	3
malformed args	0	0
repeated identical calls	33	0
reached a final answer	no (hit turn cap)	yes

3/3 runs identical. Prompt asks for three tools in order, then a summary; unfused it does
exactly that in three calls.

With a more complex tool schema (array-of-number + optional params) the arguments also
degrade progressively within a single conversation:

turn 0  {"values": "[0.5, 0.25, 0.25]", "base": 2}    <- array serialized as a string
turn 1  {"values": [0.5, 0.25, 0.25], "base": 2}      <- correct
turn 2  {"values": [0, 0, 0], "base": ""}
turn 4  {"values": ["", "", ""], "base": ""}

Reproduction

1. Serve MiniMax-M3 on gfx942 with the flags above.
2. Multi-turn loop: supply 2–3 tools, feed each tool result back as a tool message,
   iterate. Any agent client reproduces it; we saw it independently in two unrelated ones
   before isolating it here.
3. Unfused: terminates in 3 calls. Fused: loops to the turn cap.

Isolation

The kernel config was a possible confound, so it was controlled for. E differs between
fused (129) and unfused (128), so the configs cannot be held identical — instead the
config was varied within each fusion state:

fusion	MoE kernel config	result
off	stock defaults	terminates
off	tuned E=128 (older toolchain)	terminates
off	tuned E=128 (current triton)	terminates
on	tuned E=129	loops
on	stock defaults	loops

Three passing states are all fusion-off; both failing states are fusion-on, including one
with no tuned config at all. The kernel config is not the variable.

What still passes with fusion on

Worth stating, because it is why this is easy to miss: single-turn generation looks
healthy. Needle-in-a-haystack retrieval passed 3/3 at 112k context, and prose output is
fluent and coherent. Only multi-turn structured/agentic use fails. A correctness gate
based on perplexity or single-shot QA would not catch this.

Possibly related — shape mismatch in routed-experts capture

Independently, enabling --enable-return-routed-experts alongside fusion crashes at CUDA
graph capture:

state_capturer/routed_experts.py: capture -> BaseTopkCapturer.capture
RuntimeError: The expanded size of the tensor (5) must match the existing size (4)
              at non-singleton dimension 1.  Target sizes: [6, 5]. Tensor sizes: [6, 4]

RoutedExpertsCapturer allocates topk_size + num_fused_shared_experts (5) columns, but
the topk_ids reaching capture() has 4. Deterministic, every boot. That is a disagreement
about whether the fused shared expert occupies a slot in topk_ids, and it may point at
the same underlying issue.

Happy to help

We can re-test a revised patch on this hardware quickly and report back. The perf case is
attractive on paper, so it would be good to get this working rather than shelved.

[yctseng0211]

LGTM

[kevin-mii]

Pushed b3f0401 on top of the main merge: on the aiter path _post_process_topk_ids already appends the fused shared expert (fused_append_shared_experts, weight 1.0), so asking the JIT gate (moe_fused_gate) for it as well wrote the shared id (128) into two columns and dropped one real routed expert (MiniMax-M3 on ROCm ran 3 routed experts instead of 4; GSM8K-500 0.81 fused vs 0.88 unfused, and MXFP4 output was garbage). One-line fix: the aiter path takes the same "gate is asked for K_routed only" branch the per-rank shared-slot layout already used. Verified on MI355X that the routed columns no longer carry id 128. The main-merge guard for SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK covers only the grouped router, not this gate.