[Richardyu114]

Summary

Enable and optimize MiMo-V2.6-Flash-MOPD prefill on MI300X with DP2 / attention TP4 / MORI EP8. Builds on the DFlash DP-attention support merged in #35 and requires the companion AITER #23.

Changes

• Route MiMo MoE through MORI EP; mask padded expert IDs and skip unused expert-count kernels.
• Enable custom collectives for attention TP groups and integrate AITER reduce-scatter/capture handling with existing fallbacks.
• Quantize QKV inputs before all-gather and pack FP8 activations plus FP32 block scales into one collective using an FP32 bit-view. Write quantization outputs directly into the packet and transpose gathered scales once.
• Add the EP8 launcher, tuned MORI AUTO configuration, capacity validation and evaluation instructions. Defaults: global chunk 32K / 16K per DP, DFlash and FP8 all-gather enabled, and PrefillDelayer enabled with a 100-pass budget. Use PERFORMANCE_MODE=0 for accuracy; the default performance mode uses fake expert routing.
• Fix post-warmup cache-flush races and PrefillDelayer admission ordering that could stall long-input requests.

Upstream changes

Change	Source	Integration
Padded expert-ID masking	sglang#28378	Adapted
Skip unused MORI expert counts	sglang#30302	Adapted
Attention TP custom collectives	sglang#31629, #35330	Selected portions
AITER reduce-scatter and capture handling	sglang#29103, #34203, #43077	Selected collective portions
Wait for idle before cache flush	sglang#33663	Adapted
PrefillDelayer admission fixes	sglang#31835, #32389	Cherry-picked; sgl-project#31835 adapted to the local Req API

Validation and results

Validated with AITER #23 on MI300X, DFlash block 8, global chunk 32768 and seed 12345.

• GSM8K: full 1319 questions, 4-shot, greedy, batch 32, real routing. Two runs scored 1286/1319 (97.4981%) and 1289/1319 (97.7255%); zero API errors. Mean worker acceptance lengths were 5.5964 / 5.5314, including the bonus token. The second run reused the server after a cache flush.
• 64K/1, C32, 256 requests: four warmups, fake routing, TOKENIZE_PROMPT=0; 256/256 succeeded.

Metric	PR
Input throughput (tok/s)	63,841.42
Mean TTFT (ms)	30,951.48
P99 TTFT (ms)	33,663.12

Earlier component comparisons, all with DFlash and global chunk 32K:

Optimization	Case (ISL/OSL, concurrency, requests)	Input-throughput gain
ASM + MORI tuning	64K/1, C16, 1024	+2.40%
Packed FP8 all-gather path	64K/1, C16, 1024	+4.86%
PrefillDelayer fixes	512K/1, C4, 32; 960K/1, C4, 16	+17.48%; +14.73%

[Copilot]

🟡 Changes recommended

Critical registration-policy and AITER API mismatches can break graph replay and attention execution, while missing or incomplete test entry points block CI collection or leave regressions untested.

5 open issues

• Honor the communicator's registration policy for AITER allocations · New
• Gate the packed AITER path on supported output parameters · New
• Add a main entry point for the registered packet-layout tests · New
• Add an executable entry point for the registered MORI tests · New
• Run all tests in the direct-entry runner · New

What changed in this PR

Adds MiMo-V2.6-Flash prefill support on MI300X using DP2, attention TP4, and MORI EP8, with optional DFlash decoding.

Changes:

• Enables MORI routing and optimizes AMD collectives with packed FP8 activation/scale gathering.
• Adjusts prefill admission coordination and adds backend-aware benchmark cache flushing.
• Adds deployment instructions, tuning settings, regression tests, and multi-GPU smoke tests.

File	Description
test/​registered/​unit/​models/​test_mimo_mori_ep.py	Tests MORI routing, recording, and padding.
test/​registered/​unit/​managers/​test_prefill_adder.py	Tests prefill admission and delayer coordination.
test/​registered/​unit/​layers/​quantization/​test_fp8_blockwise_all_gather.py	Tests packed FP8 layouts and validation.
test/​registered/​unit/​distributed/​test_parallel_state.py	Tests collective policies, capture modes, and fallbacks.
test/​registered/​bench_fn/​test_benchmark_datasets_api.py	Tests cache-flush requests, timeout validation, and failures.
test/​manual/​amd/​test_mimo_ep_collectives.py	Adds four-GPU collective and graph-replay checks.
test/​manual/​amd/​test_fp8_blockwise_all_gather.py	Adds eight-GPU packed FP8 and QKV checks.
python/​sglang/​srt/​models/​mimo_v2.py	Enables MORI all-to-all MoE execution.
python/​sglang/​srt/​managers/​schedule_policy.py	Moves prefill negotiation after budget checks.
python/​sglang/​srt/​layers/​quantization/​fp8_utils.py	Handles row-strided FP8 inputs with CK kernels.
python/​sglang/​srt/​layers/​quantization/​fp8_blockwise_all_gather.py	Packs activations and scales into one all-gather.
python/​sglang/​srt/​layers/​moe/​topk.py	Masks padded expert IDs on HIP without requiring load balancing.
python/​sglang/​srt/​layers/​moe/​token_dispatcher/​moriep.py	Avoids unused recorder work during graph capture.
python/​sglang/​srt/​layers/​communicator.py	Integrates fused quantization with packed attention gathering.
python/​sglang/​srt/​environ.py	Adds platform-dependent defaults and collective controls.
python/​sglang/​srt/​distributed/​parallel_state.py	Adds AITER reduce-scatter and capture-mode handling.
python/​sglang/​bench_serving.py	Adds backend-aware cache flushing and an idle-wait timeout.
evaluation/​mimo_v26_flash/​mimo-v2.6-flash-mopd-evaluation-guide.md	Documents MORI setup and EP8 benchmarking.
evaluation/​mimo_v26_flash/​launch_mimo_v26_flash_dp2_tp4_ep8_dflash_accuracy_baseline.sh	Adds the deployment launcher and capacity checks.
evaluation/​mimo_v26_flash/​gfx942_mimo_v26_flash_ep8_mori_a2a.json	Supplies MI300X MORI dispatch tuning rules.

_{Review effort: Balanced}