[rbrugaro-amd]

fix(moe): correct fused shared-expert scaling on aiter/DeepEP path

Motivation

On the HIP aiter + DeepEP-class MoE path (e.g. MoRI all-to-all), the fused shared
expert is under-weighted by routed_scaling_factor, corrupting every MoE layer
and producing degenerate, non-converging generation on long outputs.

Two code paths scale the MoE output differently:

• Post-MoE scaling (default): DeepseekV2MoE.forward_deepep multiplies the
  whole MoE output by routed_scaling_factor after the experts run.
• aiter (HIP): aiter_biased_grouped_topk folds routed_scaling_factor into
  each routed topk weight, and forward_deepep skips the post-MoE multiply for
  _use_aiter (its not (... or _use_aiter) guard).

_remap_topk_for_deepep unconditionally set the fused shared-expert weight to
1/routed_scaling_factor to compensate for the post-MoE multiply. On the aiter
path that multiply never happens, so the always-on shared expert ended up scaled
by 1/routed_scaling_factor (≈0.4 at rsf=2.5) instead of 1.0.

Impact is length-dependent, which is why this bug is easy to miss (full
numbers + commands under Accuracy Tests below):

• Short-form (GSM8K) is insensitive: 0.960 → 0.950. A constant 0.4× on the
  always-on shared expert leaves the argmax token robust on short,
  format-anchored prompts — so smoke tests pass and the bug hides.
• Long-form reasoning degenerates: the tail collapses into a runaway
  repetition loop (...→Wait→Wait→Wait...) and never terminates.
• End-to-end MLPerf (DeepSeek-R1) quantifies it: exact_match 62.5% → 80.5%
  (recovering the no-fusion reference of 82.0%), with tokens_per_sample dropping
  6664 → 3713 as the runaway generation stops burning the token cap.

Modifications

Set the fused shared-expert weight to 1.0 on the aiter path; keep
1/routed_scaling_factor for the default post-MoE-scaled path.

routed_scaling_factor = topk_config.routed_scaling_factor
if _use_aiter:
    topk_weights[:, -num_fused_shared_experts:] = 1.0
elif routed_scaling_factor is not None and routed_scaling_factor != 0:
    topk_weights[:, -num_fused_shared_experts:] = 1.0 / routed_scaling_factor

Scope. This fix is deliberately scoped to the aiter path — the one validated
on AMD MI355X. forward_deepep also skips the post-multiply for the non-aiter
families that pre-fold routed_scaling_factor in topk
(should_fuse_routed_scaling_factor_in_topk: ModelOpt NVFP4, cutlass/trtllm-routed
fp8), so those would likewise need a 1.0 shared weight — but they are
intentionally left at their existing behavior here, to be addressed by their
maintainers on hardware where they can be validated.

New CPU unit tests (test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py,
3 tests) pin the contract: aiter → 1.0, post-MoE-scaled (default) → 1/rsf, and
shared-id home-rank placement. Verified the patch applies cleanly to a pristine
rocm/sgl-dev:v0.5.13-rocm720-mi35x-20260613 image and the 3 tests pass there.

Accuracy Tests

Server (same env/flags for unpatched and patched; only the topk fix differs).
The buggy path is confirmed active at runtime (router_topk=9,
num_local_experts=33 → 1 shared expert fused per rank):

SGLANG_AITER_MOE=1 SGLANG_USE_AITER=1 SGLANG_USE_ROCM700A=1 \
SGLANG_ROCM_FUSED_DECODE_MLA=1 SGLANG_AITER_AR=1 SGLANG_AITER_MLA_PERSIST=1 \
TORCH_NCCL_BLOCKING_WAIT=1 TRITON_ENABLE_PRELOAD_KERNELS=0 TORCH_BLAS_PREFER_HIPBLASLT=1 \
ROCM_QUICK_REDUCE_QUANTIZATION=INT4 ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB=512 \
PYTORCH_ALLOC_CONF=expandable_segments:True \
MORI_SHMEM_MODE=ISOLATION MORI_SHMEM_HEAP_SIZE=16G \
SGLANG_MORI_DISPATCH_DTYPE=auto SGLANG_MORI_COMBINE_DTYPE=fp8 \
SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK=4096 \
python3 -m sglang.launch_server \
  --model-path /model/deepseek-r1/fp4_quantized \
  --tp-size 8 --ep-size 8 --dp-size 8 \
  --moe-a2a-backend mori --deepep-mode normal --moe-dense-tp-size 1 \
  --enable-dp-attention --enable-dp-lm-head --enforce-shared-experts-fusion \
  --mem-fraction-static 0.84 --kv-cache-dtype fp8_e4m3 --trust-remote-code \
  --cuda-graph-max-bs 256 --host 127.0.0.1 --port 30000

Short-form (GSM8K) — insensitive to this bug

python3 -m sglang.test.few_shot_gsm8k --num-questions 200 --parallel 128 --port 30000

Before fix:

Accuracy: 0.960
Invalid: 0.000

After fix:

Accuracy: 0.950
Invalid: 0.000

GSM8K does not discriminate this bug: a constant 0.4× on the always-on shared
expert leaves the argmax token robust on short, format-anchored prompts (the
0.960/0.950 delta is within run-to-run noise).

Long-form reasoning — sensitive (this is the real failure)

Identical prompt, temperature 0.0, max_tokens 8000:

curl -s http://127.0.0.1:30000/v1/chat/completions \
  -H 'Content-Type: application/json' -d '{
    "model": "/model/deepseek-r1/fp4_quantized",
    "messages": [{"role": "user", "content":
      "Solve carefully, showing all reasoning step by step: Find the number of ordered pairs of positive integers (a,b) such that a*b + 2a + 3b = 2025. Then generalize: for which N does a*b+2a+3b=N have exactly 4 solutions? Reason thoroughly before giving the final answer."}],
    "temperature": 0.0, "max_tokens": 8000
  }' | python3 -c 'import sys,json; d=json.load(sys.stdin)["choices"][0]; print("finish:",d["finish_reason"]); print(d["message"]["content"][-900:])'

Before fix — the tail collapses into a runaway repetition loop and never
terminates (hits the token cap):

... So number of solutions depends on number of divisors of N+6 ...
Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→
Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→Wait→
... (continues to the 8000-token cap)
finish: length

After fix — stays coherent for the full generation and reaches a final answer:

finish: stop

This is the same degenerate / non-convergent failure seen end-to-end in MLPerf.

End-to-end MLPerf accuracy (DeepSeek-R1, max OSL 20k)

The MLPerf accuracy harness (exact_match metric) quantifies the same effect.
Fusion ON unpatched vs patched, plus a no-fusion reference for context:

Fusion ON, unpatched (buggy 0.4x):  {'exact_match': 62.5, 'tokens_per_sample': 6664.19, 'num-samples': 200}
Fusion ON, patched (this fix):      {'exact_match': 80.5, 'tokens_per_sample': 3712.98, 'num-samples': 200}
Fusion OFF (no fusion, reference):  {'exact_match': 82.0, 'tokens_per_sample': 3610.34, 'num-samples': 200}

The fix lifts fusion-on accuracy 62.5 → 80.5, recovering essentially the
no-fusion reference (82.0). tokens_per_sample drops 6664 → 3713 (back to the
~3610 no-fusion baseline) — direct evidence of the runaway-generation mechanism:
the under-weighted shared expert makes unpatched runs burn ~80% more tokens
hitting the generation cap.

Runtime accuracy validation was manual since the aiter/HIP path isn't covered by
CPU CI. (These subset runs are for relative before/after comparison, not a
compliant MLPerf submission.)

Benchmarking and Profiling

N/A — this is a correctness fix. The shared-expert weight is a scalar multiply
in topk remap; no throughput impact.

---

CI States

Latest PR Test (Base): ❌ Run #28007730005
Latest PR Test (Extra): ❌ Run #28007729851

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[HaiShaw]

/tag-and-rerun-ci

[HaiShaw]

@amd-bot ci-status

[amd-bot]

@HaiShaw

CI Status for PR #28237

Merge verdict: ❌ Do not merge yet. This PR's own new unit test (test_fused_shared_expert_scaling.py) fails on every run with an AttributeError — it patches module attributes (get_moe_expert_parallel_world_size / get_moe_expert_parallel_rank) that do not exist in topk.py. Because that test is the only thing exercising the fix, the fix is currently unverified by CI. One AMD job is also still running.

Caution

The fix in topk.py is verified only by the new CPU test test_fused_shared_expert_scaling.py (suite base-a-test-cpu) — and that test errors out before asserting anything. Green elsewhere does NOT verify this change. The actual AMD aiter/DeepEP integration path the fix targets is not exercised by any PR-CI test. Fix the test's mock targets and confirm it passes before merge.

Changed files: python/sglang/srt/layers/moe/topk.py (+27/-3), test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py (+122/-0, new)

Executed CI failure attribution: AMD: 2 failures (1 possibly related) · Others: 4 failures (1 related) · plus 1 AMD job still running (stage-c-test-large-8-gpu-amd-mi35x (1)). Aggregator/gate jobs (*-finish, pr-gate, wait-for-base-a) are cascades of the failures below, not independent.

AMD Executed Failures

Job	Test File	Test Function	Error	Related?	Why
stage-c mi35x 8gpu (0)	test/registered/amd/test_deepseek_r1_mxfp4_8gpu.py	test_a_gsm8k (TestDeepseekR1MXFP4MTP)	AssertionError: 1.447 not greater than 2.04 (spec accept length)	🟡	PR alters fused-shared-expert weight on the _use_aiter DeepEP path; this AMD test runs DeepSeek MoE on aiter, so the numeric output change could affect MTP accept length. May also be pre-existing — verify against a base run.
dsv4-flash-fp4-fp8-rocm720	test/registered/amd/test_deepseek_v4_flash_fp4.py	setUpClass (TestDeepseekV4Fp4)	HIP out of memory on server startup	🟢	OOM at warmup; run uses --disable-shared-experts-fusion, so _remap_topk_for_deepep's shared-expert branch isn't hit. Infra/memory, not this PR.

Other Executed Failures

Job	Test File	Test Function	Error	Related?	Why
base-a-test-cpu (2)	test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py	test_aiter_path_uses_unit_shared_weight	AttributeError: ...topk... does not have the attribute 'get_moe_expert_parallel_world_size'	🔴	The PR's own new test. _remap_topk_for_deepep reads get_parallel().moe_ep_size/.moe_ep_rank; the test patches non-existent get_moe_expert_parallel_world_size/get_moe_expert_parallel_rank → errors before any assertion.
build-test (xeon-gnr, base-b-test-cpu)	test/registered/unit/server_args/test_server_args.py	test_hisparse_accepts_aiter_backend_on_rocm	AttributeError: 'ServerArgs' object has no attribute '_validate_hisparse_dsa_backend'	🟢	Unrelated server-args/hisparse test; PR doesn't touch ServerArgs.
stage-b-test-1-npu-a2 (0)	test/registered/ascend/basic_function/quant/test_npu_w8a8_quantization.py	throughput check	544.3 < 700 throughput	🟢	NPU w8a8 perf threshold; unrelated to MoE shared-expert logic.
multimodal-gen-test-1-npu-a3 · 2-npu-a3	multimodal_gen/.../test_server_{1,2}_npu.py	diffusion perf checks	Denoise/decode latency over threshold	🟢	NPU diffusion performance gates; entirely unrelated to this PR.

Details / what to do before merge

• 🔴 Fix the new test (blocker). In test_fused_shared_expert_scaling.py, replace the two patch.object(topk_module, "get_moe_expert_parallel_world_size"/"get_moe_expert_parallel_rank", ...) calls. The production code uses get_parallel().moe_ep_size and get_parallel().moe_ep_rank (topk.py:1495-1496), so patch topk_module.get_parallel to return an object whose .moe_ep_size/.moe_ep_rank give the desired values (8 / 0). Re-run base-a-test-cpu and confirm all 3 test functions pass. Until this passes, the fix's core logic is untested.
• 🟡 Verify the mi35x mxfp4 accept-length drop is not caused by this PR. Since the fix changes the aiter-path shared-expert weight from 1/rsf to 1.0, it directly changes DeepSeek MoE output on AMD aiter. Check whether test_a_gsm8k (TestDeepseekR1MXFP4MTP) already fails on a recent base/main run; if it passes on base, this is a real regression risk that needs a look before merge. The fix's premise is that the old 1/rsf under-weighted the shared expert — a drop in accept length after applying it would be a red flag worth confirming on hardware.
• 🟢 The remaining failures (dsv4 OOM, NPU perf x3, hisparse server-args) are unrelated and look like pre-existing/infra issues; no action needed for this PR, but they keep the gate jobs red.
• One AMD shard (stage-c mi35x (1)) is still queued — let it finish before reading AMD as complete.

---

Generated by amd-bot using Claude Code CLI

[HaiShaw]

@rbrugaro-amd please check the issues from merge verdict.

[rbrugaro-amd]

@HaiShaw

CI Status for PR #28237

Merge verdict: ❌ Do not merge yet. This PR's own new unit test (test_fused_shared_expert_scaling.py) fails on every run with an AttributeError — it patches module attributes (get_moe_expert_parallel_world_size / get_moe_expert_parallel_rank) that do not exist in topk.py. Because that test is the only thing exercising the fix, the fix is currently unverified by CI. One AMD job is also still running.

Caution

The fix in topk.py is verified only by the new CPU test test_fused_shared_expert_scaling.py (suite base-a-test-cpu) — and that test errors out before asserting anything. Green elsewhere does NOT verify this change. The actual AMD aiter/DeepEP integration path the fix targets is not exercised by any PR-CI test. Fix the test's mock targets and confirm it passes before merge.

Changed files: python/sglang/srt/layers/moe/topk.py (+27/-3), test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py (+122/-0, new)

Executed CI failure attribution: AMD: 2 failures (1 possibly related) · Others: 4 failures (1 related) · plus 1 AMD job still running (stage-c-test-large-8-gpu-amd-mi35x (1)). Aggregator/gate jobs (*-finish, pr-gate, wait-for-base-a) are cascades of the failures below, not independent.

AMD Executed Failures

Job Test File Test Function Error Related? Why
stage-c mi35x 8gpu (0) test/registered/amd/test_deepseek_r1_mxfp4_8gpu.py test_a_gsm8k (TestDeepseekR1MXFP4MTP) AssertionError: 1.447 not greater than 2.04 (spec accept length) 🟡 PR alters fused-shared-expert weight on the _use_aiter DeepEP path; this AMD test runs DeepSeek MoE on aiter, so the numeric output change could affect MTP accept length. May also be pre-existing — verify against a base run.
dsv4-flash-fp4-fp8-rocm720 test/registered/amd/test_deepseek_v4_flash_fp4.py setUpClass (TestDeepseekV4Fp4) HIP out of memory on server startup 🟢 OOM at warmup; run uses --disable-shared-experts-fusion, so _remap_topk_for_deepep's shared-expert branch isn't hit. Infra/memory, not this PR.

Other Executed Failures

Job Test File Test Function Error Related? Why
base-a-test-cpu (2) test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py test_aiter_path_uses_unit_shared_weight AttributeError: ...topk... does not have the attribute 'get_moe_expert_parallel_world_size' 🔴 The PR's own new test. _remap_topk_for_deepep reads get_parallel().moe_ep_size/.moe_ep_rank; the test patches non-existent get_moe_expert_parallel_world_size/get_moe_expert_parallel_rank → errors before any assertion.
build-test (xeon-gnr, base-b-test-cpu) test/registered/unit/server_args/test_server_args.py test_hisparse_accepts_aiter_backend_on_rocm AttributeError: 'ServerArgs' object has no attribute '_validate_hisparse_dsa_backend' 🟢 Unrelated server-args/hisparse test; PR doesn't touch ServerArgs.
stage-b-test-1-npu-a2 (0) test/registered/ascend/basic_function/quant/test_npu_w8a8_quantization.py throughput check 544.3 < 700 throughput 🟢 NPU w8a8 perf threshold; unrelated to MoE shared-expert logic.
multimodal-gen-test-1-npu-a3 · 2-npu-a3 multimodal_gen/.../test_server_{1,2}_npu.py diffusion perf checks Denoise/decode latency over threshold 🟢 NPU diffusion performance gates; entirely unrelated to this PR.

Details / what to do before merge

• 🔴 Fix the new test (blocker). In test_fused_shared_expert_scaling.py, replace the two patch.object(topk_module, "get_moe_expert_parallel_world_size"/"get_moe_expert_parallel_rank", ...) calls. The production code uses get_parallel().moe_ep_size and get_parallel().moe_ep_rank (topk.py:1495-1496), so patch topk_module.get_parallel to return an object whose .moe_ep_size/.moe_ep_rank give the desired values (8 / 0). Re-run base-a-test-cpu and confirm all 3 test functions pass. Until this passes, the fix's core logic is untested.
• 🟡 Verify the mi35x mxfp4 accept-length drop is not caused by this PR. Since the fix changes the aiter-path shared-expert weight from 1/rsf to 1.0, it directly changes DeepSeek MoE output on AMD aiter. Check whether test_a_gsm8k (TestDeepseekR1MXFP4MTP) already fails on a recent base/main run; if it passes on base, this is a real regression risk that needs a look before merge. The fix's premise is that the old 1/rsf under-weighted the shared expert — a drop in accept length after applying it would be a red flag worth confirming on hardware.
• 🟢 The remaining failures (dsv4 OOM, NPU perf x3, hisparse server-args) are unrelated and look like pre-existing/infra issues; no action needed for this PR, but they keep the gate jobs red.
• One AMD shard (stage-c mi35x (1)) is still queued — let it finish before reading AMD as complete.

Generated by amd-bot using Claude Code CLI

@HaiShaw I fixed the test raised as blocker above.
Regarding: * 🟡 Verify the mi35x mxfp4 accept-length drop is not caused by this PR.
I traced the guard. The PR's only changed code, _remap_topk_for_deepep, runs exclusively under the guard at topk.py#L1624-L1630:

if num_fused_shared_experts > 0 and is_deepep_class_backend():
...
topk_ids, topk_weights = _remap_topk_for_deepep(...)
and is_deepep_class_backend() is True only when moe_a2a_backend ∈ {deepep, mooncake, mori} — see utils.py#L366-L369.

The failing test TestDeepseekR1MXFP4MTP launches with only --tp 8 + EAGLE spec args (test_deepseek_r1_mxfp4_8gpu.py#L106-L120) — it never sets --moe-a2a-backend, which defaults to "none" (server_args.py#L998-L1006).

→ is_deepep_class_backend() is False in that test → _remap_topk_for_deepep is never invoked → the PR cannot change its output. The avg_spec_accept_length=1.447 < 2.04 failure (assertion here) is pre-existing / flaky / infra, not a regression from this PR.

[HaiShaw]

@amd-bot ci-status

[amd-bot]

@HaiShaw

CI Status for PR #28237

Merge verdict: No PR-CI failure is caused by this PR's changes. The new unit test (base-a-test-cpu) that pins the fixed _remap_topk_for_deepep weight logic passed, and all 32 executed AMD jobs passed. Every other red job is a pre-existing infra/environment failure or a fast-fail cascade (dirty NVIDIA GPU, NPU build ImportError, XPU timeout, stale CPU build). The one nuance a human should weigh: the actual runtime _use_aiter branch only executes on AMD HIP + DeepEP, and that end-to-end MoE-correctness path is not directly exercised by PR CI — only the mocked unit test is.

Warning

The PR's real behavioral change (topk_weights[:, -n:] = 1.0 on the _use_aiter DeepEP path) runs only on AMD HIP with DeepEP/mori all-to-all. PR CI verifies it only via the CPU unit test in test_fused_shared_expert_scaling.py (which mocks _use_aiter). The genuine e2e numerics on the aiter+DeepEP MoE path are not confirmed by any green PR-CI job. The one MI35x test that touches a DeepEP-style path (stage-b-test-large-8-gpu-mi35x-disaggregation-amd) failed, but on a KV-transfer abort/timeout symptom, not a numerics mismatch — see below. Author should confirm correctness on a real MI355X DeepEP run before merge.

Changed files: python/sglang/srt/layers/moe/topk.py (+27/-3), test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py (+117, new)

Executed CI failure attribution: AMD: 1 executed failure (0 related, 1 possibly) · Others: ~7 root-cause executed failures (0 related) + many fast-fail cascades

AMD Executed Failures

Job	Test File	Test Function	Error	Related?	Why
stage-b…mi35x-disaggregation-amd	test/registered/amd/disaggregation/test_disaggregation_basic.py	TestDisaggregationAccuracy.test_gsm8k	KeyError: 'answer' after KVTransferError: Aborted by AbortReq + client disconnects; ran 1335s vs 600s est.	🟡	Symptom is PD KV-transfer abort/timeout, not a MoE numerics mismatch; but it's an AMD MoE-adjacent path so worth confirming the fix didn't perturb it

(pr-test-amd-finish and pr-test-amd-extra-finish are aggregation jobs that fail because the disagg job failed — collapsed into the row above.)

Other Executed Failures

Job	Test File	Test Function	Error	Related?	Why
base-b-test-4-gpu-b200 (0)	N/A (pre-test health check)	N/A	ERROR: memory >=10%: GPU 0 (92%)…	🟢	Dirty-GPU runner infra; failed before any test ran. Root cause of all base-b NVIDIA reds + wait-for-base-b cascade
stage-b-test-1-npu-a2 (0)	N/A (build/install)	N/A	ImportError: cannot import name 'parse_version' from 'vcs_versioning'	🟢	Toolchain/packaging env failure during uv editable build; PR touches no build/packaging files
stage-b-test-1-gpu-xpu	test/registered/xpu/test_intel_xpu_backend.py	TestIntelXPUBackend.test_attention_backend	TIMEOUT after 1200s	🟢	XPU attention backend timeout; XPU never takes the _use_aiter path
build-test (xeon-gnr, base-b-test-cpu)	test/registered/unit/server_args/test_server_args.py	(ServerArgs init)	TypeError: ServerArgs.__init__() got an unexpected keyword argument 'enable_session_radix_cache'	🟢	enable_session_radix_cache exists in source (server_args.py:985) — failure is a stale-build/env mismatch, not this PR; PR doesn't touch ServerArgs

All *-finish, call-gate / pr-gate, finish, wait-for-base-b, and the seven base-b-test-*-gpu NVIDIA jobs are fast-fail cascades from the b200 dirty-GPU root cause or the per-backend infra failures above — not independent failures.

Details / what to do before merge

• 🟡 AMD disagg (test_gsm8k) — The failure chain is Request disconnected → KVTransferError: Aborted by AbortReq → KeyError: 'answer', with the file running 1335s against a 600s estimate. This reads as PD/KV-transfer instability or a slow runner, not the shared-expert scaling math. To rule out the fix as a contributor, re-run test_disaggregation_basic.py on MI35x (the job auto-skips retry on exit 255, so a manual rerun is needed). [LOW] that the MoE weight change caused it.
• Coverage gap (most important) — The fix's net behavioral effect on a real DeepSeek-style MoE over aiter + DeepEP/mori all-to-all is not verified by any green PR-CI job; only the mocked CPU unit test is. Before merge, run a DeepEP fused-shared-expert accuracy check on MI355X (the configuration the PR description says was validated) and paste the result.
• No in-flight-fix lookup needed — none of the red jobs are attributable to this PR, so there is no PR-side fix to chase. The infra reds (dirty b200 GPU, NPU vcs_versioning build, stale CPU server_args build) are environment issues for the CI maintainers, independent of this PR.

---

Generated by amd-bot using Claude Code CLI

[yctseng0211]

it is not related to this PR, the root cause is the infra issue which we are working on, see the ticket: OSSCI-2656 cc: @bingxche @HaiShaw

[yctseng0211]

@rbrugaro-amd is it expected that we register the test file to cpu test ? thanks.