[kkHuang-amd]

co-author: @hubertlu-tw

Summary

• add the packaged Aiter/FlyDSL MegaMoEv2 backend for DeepSeek-V4 with DP rank synchronization and graph-safe idle handling
• add prefill-only fused EPLB remap/recording to remove per-layer recorder overhead
• preserve stable SWA eviction cadence for the MegaMoE path under high DP concurrency and cover rank-sync, recorder, and mixed-slot behavior with regression tests

Performance

Validated on 8×MI355X, DeepSeek-V4-Pro, ISL/OSL 8192/1024, concurrency 512, 5120 requests:

• no EPLB: 38,863.10 tok/s
• dynamic fused EPLB (interval 200): 42,382.44 tok/s

These match the source integration within +0.26% and +0.015%, respectively. The branch was subsequently rebased onto current main; full GPU e2e revalidation on the new base remains pending.

Test plan

• pre-commit hooks
• MegaMoE fused recorder tests (7 tests + 7 subtests)
• Aiter MegaMoE rank-sync tests (7 tests)
• distributed mixed-slot prefill-delayer regression
• DeepSeek-V4-Pro 8×MI355X throughput validation on the integration base
• repeat GPU accuracy and throughput validation on current main

---

CI States

Latest PR Test (Base): ⏳ Run #35983394898
Latest PR Test (Extra): ✅ Run #35983394136
Latest PR Test (AMD ROCm 10): ❌ Run #35983394627

[kkHuang-amd]

Arm	This PR claims	Measured	Delta
MegaMoE, no EPLB	38,863.10 tok/s	39,200.88	+0.87%
MegaMoE + EPLB (interval 200)	42,382.44 tok/s	42,505.60	+0.29%
EPLB gain	+9.05%	+8.43%

5120/5120 requests successful, 0 failed, on both arms. DeepSeek-V4-Pro,
ISL/OSL 8192/1024, concurrency 512, no speculative decoding, no HiCache.
sglang a1f9508dd4 + this PR, aiter c16d44b93, ROCm/gfx950.

Posting the exact commands since getting here took several wrong turns.

Server, no-EPLB arm:

python3 -m sglang.launch_server \
  --model-path /path/to/DeepSeek-V4-Pro --host 0.0.0.0 --port 30000 --trust-remote-code \
  --tp 8 --ep-size 8 --dp-size 8 --enable-dp-attention \
  --moe-a2a-backend megamoe --moe-dense-tp-size 1 --enable-dp-lm-head \
  --load-balance-method round_robin \
  --attention-backend dsv4 --kv-cache-dtype fp8_e4m3 --page-size 256 \
  --swa-full-tokens-ratio 0.15 --mem-fraction-static 0.85 \
  --chunked-prefill-size 65536 --cuda-graph-max-bs 1024 --max-running-requests 1024 \
  --disable-radix-cache --disable-shared-experts-fusion \
  --tool-call-parser deepseekv4 --reasoning-parser deepseek-v4

with:

export SGLANG_AMD_USE_FLYDSL_MEGA_MOE=1
export SGLANG_AMD_FLYDSL_MEGA_MOE_MTPR=8192
export MORI_SHMEM_HEAP_SIZE=17179869184        # 16 GiB
export SGLANG_SHARED_EXPERT_TP1=0
export SGLANG_DP_SHARED_EXPERT_LOCAL=0
export SGLANG_DP_USE_GATHERV=0
export SGLANG_DP_USE_REDUCE_SCATTER=0

EPLB arm — same as above, plus:

export SGLANG_EXPERT_DISTRIBUTION_RECORDER_DIR=/tmp
export SGLANG_AITER_MEGA_RANK_SYNC=1
export SGLANG_AITER_MEGA_EPLB_PREFILL_ONLY=1
export SGLANG_AITER_MEGA_EPLB_FUSED_MAP_RECORD=1
# server args appended:
#   --enable-eplb --eplb-rebalance-num-iterations 200 \
#   --expert-distribution-recorder-mode stat

Client (same prompt manifest for both arms so they replay identical prompts):

python3 common_oai_benchmark.py \
  --base-url http://127.0.0.1:30000 \
  --model /path/to/DeepSeek-V4-Pro --trust-remote-code \
  --input-len 8192 --output-len 1024 --num-prompts 5120 \
  --warmup-requests 1024 --max-concurrency 512 \
  --prompt-manifest /tmp/common_prompt_c512.jsonl.gz \
  --output-dir /tmp/pr35619_arm1

Four settings that silently break it

1. MORI_SHMEM_HEAP_SIZE. The mori default (4 GiB) is not enough —
   MegaMoEV2's dispatch/combine buffers want ~3.1 GiB each, and it fails during
   decode CUDA-graph capture with AssertionError: mori_shmem.shmem_malloc failed. 16 GiB was sufficient here. Worth noting capture exercises MegaMoE
   unconditionally (should_use_mega_moe returns True in capture mode), so this
   fires regardless of load.
2. --enforce-shared-experts-fusion is incompatible. Use
   --disable-shared-experts-fusion. With fusion on, forward_mega_moe sizes
   MegaMoEV2 with topk = num_experts_per_tok + num_fused_shared_experts (7
   instead of 6) and the dispatch path dies with
   HSA_STATUS_ERROR_EXCEPTION 0x1016.
3. The four DP-comm envs must be 0. MoE comm is mori's; leaving the DP
   gatherv/reduce-scatter path enabled runs two schemes at once and also faults
   with 0x1016.
4. EPLB without --expert-distribution-recorder-mode stat is a regression.
   Measured 36,642.28 tok/s (-6.5% vs no-EPLB) with 1040 rebalance events
   and no load statistics behind them. The recorder is what turns it into
   +8.4%. Also: leaving --ep-num-redundant-experts at its default 0 is
   correct — setting 32 costs KV pool for no gain.

Might be worth folding 1-4 into the PR description, since none of them produce
an error that points at the actual cause.

One bug: SGLANG_AITER_MEGA_RANK_SYNC with DP attention

In forward_batch_info.py:

if mega_moe_idle_materialize:      # RANK_SYNC and is_extend_in_batch and forward_mode.is_idle()
    global_num_tokens = [1]

adjust_num_token_non_padded_for_attn_tp then does
self.global_num_tokens_cpu[dp_rank], and attn_dp_rank reaches 7 at
dp_size 8, so any rank above 0 raises IndexError: list index out of range
and takes the scheduler down on the first real forward. The server prints
ready to roll before this, so it presents as "died before /health became
ready".

The fixed-seq-len runs above did not hit it, including the EPLB arm with
RANK_SYNC=1 — it needs an idle batch to coincide with an extend. I hit it
consistently on a long-context multi-turn agentic workload with
--enable-prefill-delayer on. [1] * len(global_num_tokens) would keep the
intent and fix the indexing.

[kkHuang-amd]

is_hip

[kkHuang-amd]

is_hip

[kkHuang-amd]

add comment

[HaiShaw]

python/sglang/srt/layers/quantization/fp8.py change is AMD specific

[HaiShaw]

LGTM

[HaiShaw]

base-c-test-8-gpu-b300 (0) stuck, previous fail is known in #17050.
Otherwise PR Test Base / PR Test Extra look good.