[ROCm] [Bugfix] Fix Triton fused shared expert alignment

[akii96]

Motivation

MoE models using fused shared experts with the Triton backend can produce corrupted outputs and suffer from severe accuracy loss. Fused shared experts append extra expert IDs and weight rows, while global_num_experts continues to represent only the routed experts. Triton uses this smaller count during token alignment, causing the appended shared expert IDs to be treated as invalid. Since the separate shared expert path is disabled when fusion is enabled, its contribution is lost.

This affects MoE models using Triton fused shared experts without an expert map. EmbeddedLLM/MiniMax-M3-FP8-dynamic on ROCm gfx942 is one such example.

Suggested Fix

Use the physical expert row count from w1.shape[0] when no expert map is present. Continue using global_num_experts with expert parallelism because the IDs must be remapped through expert_map.

Apply the same correction to TritonExperts and TritonWNA16Experts. This matches the existing handling in the native MXFP8 path.

Tests

Added a focused case to the existing Triton MoE test suite. Before the fix, 134 of 512 output elements differed from the reference, with a maximum difference of 2.515625. The test passes after the fix.
GSM8K 8-shot smoke evaluation with 128 samples on 4 MI325X GPUs:

Configuration	Flexible exact match	Strict exact match
FSE enabled, before fix	0.00%	0.00%
FSE enabled, after fix	94.53%	89.06%
FSE disabled	94.53%	89.84%

The patched TP4 server selected Triton FP8 MoE with 129 physical experts, completed CUDA graph capture, and returned a successful chat completion.

Reproduction command

Container image:

vllm/vllm-openai-rocm:nightly-c810e5ee9976ad86b81d1277b53e76d0ee639414

export HIP_VISIBLE_DEVICES=0,1,2,3
export VLLM_ROCM_USE_AITER=1
export VLLM_USE_BREAKABLE_CUDAGRAPH=0
export VLLM_ROCM_USE_AITER_MOE=0
export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1

vllm serve EmbeddedLLM/MiniMax-M3-FP8-dynamic \
  --served-model-name minimax-m3 \
  --tensor-parallel-size 4 \
  --block-size 128 \
  --max-model-len 262144 \
  --gpu-memory-utilization 0.92 \
  --enable-chunked-prefill \
  --max-num-batched-tokens 32768 \
  --max-num-seqs 32 \
  --no-enable-prefix-caching \
  --async-scheduling \
  --tool-call-parser minimax_m3 \
  --enable-auto-tool-choice \
  --reasoning-parser minimax_m3 \
  --kv-cache-dtype fp8 \
  --language-model-only \
  --port 8000 \
  --attention-backend ROCM_AITER_UNIFIED_ATTN

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

✅ @akii96, CI is now available for this PR.

• /ci run starts a CI build.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.

[akii96]

/ci retry

[github-actions]

✅ No earlier CI build exists for this PR. Use /ci run first.

[maeehart]

Verified independently at 50a214669 on MI325X using public image vllm/vllm-openai-rocm:nightly-c810e5ee9976ad86b81d1277b53e76d0ee639414, with #51585 applied because the exact model command otherwise fails during graph capture.

The focused regression reproduced the base failure exactly: 134/512 mismatches with a maximum absolute difference of 2.515625. It passed with this PR.

AITER-enabled MiniMax-M3 TP4 results:

• 8-shot GSM8K, 128 samples: before 0.00/0.00%, after 92.19/93.75%, and FSE disabled 89.84/89.84% flexible/strict.
• 5-shot GSM8K, 200 samples: before 0.00/0.00%, after 92.00/92.00%, and FSE disabled 92.00/91.50% flexible/strict.

The absolute post-fix 8-shot counts differ slightly from the PR table, but the failure and recovery reproduced, and the standard 5-shot comparison is neutral. All three variants completed full graph capture and returned successful chat completions. The Triton FP8 MoE backend was selected.

#51585 is required for this model command and should merge first. @mgoin @pavanimajety, can you merge this after #51585 lands?

[tjtanaa]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84158 for commit 50a214669e0d.