Repository navigation
[Fix][AMD] Qwen3.5 MoE: disable global-slot shared-expert fusion under per-rank EP backends (MoRI + dp-attention init crash) - #31793
Conversation
…r per-rank EP backends Qwen shared-expert fusion appends the shared expert at a single global slot (_append_shared_to_topk_output, base id N == num_experts) with a per-token shared_expert_gate. That global-slot layout is incompatible with per-rank EP shared-slot backends (DeepEP/Mooncake/MoRI/MegaMoE) once moe_ep_size > 1, which tile the expert id space into contiguous per-rank blocks (num_experts == num_local_experts * ep_size) and thus require per-rank shared slots. Mixing the two trips FusedMoE's (num_experts - num_shared_slots) % ep_size == 0 assertion at init for Qwen3.5-397B-A17B-FP8 with --moe-a2a-backend mori --enable-dp-attention (513 experts, 8 shared slots: 505 % 8 != 0), and even past the assertion would misroute experts. Fall back to the separate shared-expert MLP there, which is correct. Fixes sgl-project#31336
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request disables Qwen shared-expert fusion when using per-rank EP shared-slot backends (such as DeepEP or MoRI) with a moe_ep_size greater than 1, falling back to a separate shared-expert MLP to prevent expert misrouting and assertion failures. The reviewer suggested storing get_parallel().moe_ep_size in a local variable to improve readability and avoid duplicate calls.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
kkHuang-amd
left a comment
kkHuang-amd
left a comment
There was a problem hiding this comment.
Please resolve the gemini comment.
LGTM for others
Thanks KK. Gemini reply is asking for a small refactor, I'd prefer to keep the current form: the if guard reads clearly with get_parallel().moe_ep_size > 1 inline, and repeating it in the log makes the logged value explicit without an extra local. This runs once at init, so the duplicate call is negligible. |
HaiShaw
left a comment
HaiShaw
left a comment
There was a problem hiding this comment.
Please comment on the reason why issue happens on AMD, code change is generic to cover NV path? If not, please add comment for clarification.
|
@amd-bot ci-status |
CI Status for PR #31793Merge verdict: Not green — do not merge on CI signal yet. Zero executed failures are attributable to this PR, but AMD PR CI is incomplete (stage-b/stage-c were fast-fail skipped because Caution This PR's changed code is not exercised by any completed PR-CI test, and AMD CI is incomplete. The fix only fires under Changed files: Executed CI failure attribution: AMD: 1 executed failure (0 related) + downstream skipped · Others: 2 executed failures (0 related) · plus 4 opt-in "extra" gate/finish cascades (expected). NVIDIA AMD Executed Failures
Other Executed Failures
The remaining 4 red checks are opt-in "extra" workflow gates, not real failures: call-gate / pr-gate + pr-test-amd-extra-finish and call-gate / pr-gate + pr-test-extra-finish all exit with Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
@HaiShaw This is a Qwen model issue, not device-specific, so it also hits NVIDIA: Qwen's global-slot shared-expert fusion is incompatible with per-rank EP backends (DeepEP/Mooncake/MoRI) at ep_size > 1. The fix detects that combo and falls back to the separate shared-expert MLP on both AMD and NVIDIA. |
|
@yichiche Can you upload the config that triggers issue on NV? Which can help to clarify the mixture of MoRI description and generic code changes above. |
…r per-rank EP backends (MoRI + dp-attention init crash) (sgl-project#31793)
Motivation
Fixes #31336.
Serving Qwen/Qwen3.5-397B-A17B-FP8 (512 routed + 1 fused shared expert) with
--moe-a2a-backend mori --enable-dp-attentionon MI355X/ROCm crashes at model init. Every scheduler rank aborts inFusedMoE.__init__:With
--enable-dp-attention,moe_ep_sizeis promoted totp_size(8).has_per_rank_fused_shared_slots()isTruefor MoRI (deepep-class), sonum_shared_slots = 1 * 8 = 8, and(513 - 8) % 8 = 505 % 8 = 1 != 0.Root cause (and why the layout matters)
MoRI's dispatcher (
token_dispatcher/moriep.py) masks experts with a contiguous per-rank block:expert_mask[ep_rank * num_local_experts : + num_local_experts] = 1. This is valid only when the blocks tile the id space exactly, i.e.num_experts == num_local_experts * ep_size, which requires per-rank shared slots (one shared slot per EP rank). DeepSeek-V2 already buildsnum_expertsthis way for per-rank backends.Qwen's shared-expert fusion, however, is a global-slot design:
Qwen2MoeSparseMoeBlockconstructs itsTopKwithoutnum_fused_shared_experts(so the per-rank remap intopk.pynever fires) and instead appends the shared expert manually at a single global id (_append_shared_to_topk_output, base idN == num_experts) with a per-tokenshared_expert_gate. This global-slot-with-gate layout cannot be expressed in MoRI's per-rank contiguous tiling, and it can't reuse DeepSeek's (gate-less) per-rank remap machinery.The alternative floated in the issue — treating MoRI as "global" (
num_shared_slots = 1) — passes the__init__assertion but then breaks MoRI's contiguous dispatch mask (num_local_experts * ep_size = 65 * 8 = 520 != 513, so rank 7's block runs out of bounds) and silently misroutes experts.Modifications
python/sglang/srt/models/qwen2_moe.py: when shared-expert fusion is enabled and a per-rank EP shared-slot backend is active (uses_per_rank_fused_shared_slots()) andmoe_ep_size > 1, disable the fusion and fall back to the separate (correct) shared-expert MLP path. A one-time warning is logged.moe_ep_size == 1and non-per-rank backends are unaffected (fusion stays on).This is the correct-by-construction fix: with fusion off,
num_experts = 512tiles exactly (64 * 8), the shared expert is computed by the denseshared_expertMLP (with its gate) and added after the experts. The fully-fused per-rank shared expert with a per-token gate is a larger enablement task (needs a per-rank remap+append kernel that also applies the gate) and is left as follow-up.Accuracy Tests
GSM8K sweep on Qwen3.5-397B-A17B-FP8 (MI355X), with this PR applied so the MoRI configs get past the #31336 init crash. Flags added one at a time.
Accuracy= correct fraction,Invalid= unparseable/garbage fraction.Reading:
FusedMoE.__init__.--enable-dp-attentionalone (R2, 0.935) are all healthy. R2 is the path this PR most directly enables and it is clean.--moe-a2a-backend mori(R3/R4/R5) collapses to ~0.01 accuracy with a high invalid rate, independent of dp-attention and of fp8-KV. That is a MoRI all-to-all correctness problem, out of scope for this init-crash fix and tracked separately (likely part of the ROCm-MoRI enablement in [bug] ROCm DI MoRI for non-standard attention (such as DeepSeekv4, GLM5 & Qwen3.5) doesnt work #21886). This PR is a prerequisite that lets that issue be reached and debugged at all.Assertion-level validation (mocked MoRI,
ep_size=8): pre-fixnum_experts=513, num_shared_slots=8→ assertion fails; with fix, fusion disabled →num_experts=512, num_shared_slots=0→ assertion passes and MoRI's contiguous mask tiles[0, 512)exactly.ep_size==1and non-per-rank backends keep fusion (no behavior change).Checklist
pre-commit run --all-files(isort/ruff/black pass)CI States
Latest PR Test (Base): ✅ Run #29797490050
Latest PR Test (Extra): ❌ Run #29797489930