[IzacharyI]

Title: [Perf] [AMD] dsv4: Use the aiter MegaMoE small- and large-input paths

Opt-in switches for the two aiter MegaMoEV2 paths (needs ROCm/aiter#6881):

• SGLANG_OPT_USE_AITER_MEGA_SMALL_M: steps whose DP-global max token count fits the small-input runtimes (up to 512 tokens per rank) call forward_small_m(), one kernel for quant, dispatch, GEMM1, GEMM2 and combine.
• SGLANG_OPT_USE_AITER_MEGA_LARGE_M: extend steps with at least SGLANG_OPT_AITER_MEGA_LARGE_M_MIN_TOKENS (default 1024) tokens on some DP rank run the persistent large-input kernel.

Each MoE layer gets its own protocol slot on the shared operator, and the choice of path depends only on DP-gathered counts, so every EP rank takes the same path.

Motivation

DeepSeek-V4-Pro served with TP8 DP8(attn) EP8 and --moe-a2a-backend megamoe on MI355X runs every MoE call through aiter MegaMoEV2. ROCm/aiter#6881 adds two faster MegaMoEV2 paths for this layout; this PR lets SGLang use them. Both are opt-in and the default path does not change.

Depends on ROCm/aiter#6881; merge that one first. Without it, the new switches raise an error at startup and nothing else changes.

Modifications

• New switches in environ.py, all off by default: SGLANG_OPT_USE_AITER_MEGA_SMALL_M, SGLANG_OPT_USE_AITER_MEGA_LARGE_M and SGLANG_OPT_AITER_MEGA_LARGE_M_MIN_TOKENS (default 1024). Smaller extend steps stay on the main path; the large-input path is turned on per step with set_large_m_enabled().
• The path choice uses only DP-gathered counts (original_global_num_tokens_cpu, get_dp_global_num_tokens()), and require_mlp_tp_gather returns true when the small-input switch is on, so captured graphs get DP-synchronized buckets.
• Each MoE layer gets its own protocol slot on the shared operator (bind_layer); the draft model's layers get their own slots too.
• Works with SGLANG_AMD_FLYDSL_MEGA_MOE_DECODE_MTPR: the large-input path is used only on the full-capacity instance.
• Only layers/moe/mega_moe_flydsl.py (the AMD MegaMoE adapter), environ.py and one branch in require_mlp_tp_gather that runs only with the MegaMoE backend change. CUDA paths cannot reach this code.
• CPU unit test test/registered/unit/layers/moe/test_mega_moe_small_large_m.py: with the switches off the main path is untouched; with them on, every rank picks the same path from the DP-gathered counts (the small-input and extend thresholds), each MoE layer keeps its own slot, and a switch without the aiter path fails at first use.

Related open PRs: none change the AMD MegaMoE adapter (mega_moe_flydsl.py). The open MegaMoE PRs (#46056, #45905, #45613, #44359) are on the NVIDIA / FlashInfer MegaMoE path.

Accuracy Tests

GSM8K, 1319 questions, 5-shot, real speculative verification (no acceptance simulation), c384 server setup; measured with ROCm/aiter#6881 before its review changes, to be re-measured on the final aiter code:

	accuracy
MegaMoEV2 (both switches off)	0.932–0.938
this PR (both switches on)	0.926–0.948

Speed Tests and Profiling

Setup:

• Workload: InferenceX AgentX (aiperf --scenario agentx, dataset semianalysis_cc_traces_weka_062126), 3600 s per run, concurrency 128 / 256 / 384, two runs per arm.
• Model and server: DeepSeek-V4-Pro-0813, TP8 DP8(attn) EP8, MegaMoE a8w4 with MTPR 8192, DSPARK speculative decoding with the InferenceX acceptance simulation (SGLANG_SIMULATE_ACC_LEN=3.77, the same accept length on both arms), sglang-router cache_aware + dp-aware. Server, router and client flags are those of InferenceX run 37608671586.
• Hardware and image: one node, 8 × MI355X, lmsysorg/sglang-rocm:v0.5.21-rocm720-mi35x-20261007.
• Baseline: the image as is (MegaMoEV2). This PR: the same image and flags with the aiter small- and large-input paths and these switches turned on. Measured before the rebase onto current main and before the review changes in [FlyDSL] [CI] MegaMoEV2: fuse Dispatch, GEMM1, GEMM2 and Combine into one kernel ROCm/aiter#6881 (skewed-routing fixes); the SGLang code path is unchanged by either. To be re-measured on the final aiter code.

Total token throughput = all input tokens (including prefix-cache hits) plus output tokens, per second per GPU (tok/s/GPU); run 1 / run 2, gain of the means:

conc	baseline	this PR	gain
128	42,810 / 43,060	47,137 / 47,425	+10.1%
256	65,795 / 67,857	73,217 / 73,796	+10.0%
384	76,192 / 73,603	78,880 / 79,776	+5.9%

Run-to-run spread: ≤1.1% with this PR, ≤3.4% for the baseline. Below c128 the InferenceX recipe uses TP8 without EP, so this PR does not apply there.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #40682799849
Latest PR Test (Extra): ❌ Run #40682799962
Latest PR Test (AMD ROCm 10): ❌ Run #40682799895