[AMD][DSV4] perf: MXFP8 MoRI dispatch to match the w4a8 MoE input format

[karverma-amd]

Motivation

DSv4's MoE runs per_1x32 (MXFP4 weights), so AITER wants fp8 activations
carrying group-32 e8m0 microscales. None of MoRI's three shipped dispatch dtypes
produce that:

dispatch	payload	scales	consequence
bf16	bf16	none	receiver must quantize
fp8	fp8	group-128 fp32	wrong group size -> fp8->bf16 upscale round trip
fp4	fp4x2	group-32 e8m0	right scales, wrong payload -> upscale_mxfp4

bf16 wins by default only because it is the one with no upscale kernel. But it
pushes the bf16->fp8 conversion to the receive side, and that is where the
cost is: AITER sizes its quant grid from the input row count
(aiter/csrc/kernels/quant_kernels.cu, rows = input.numel() / cols), and the
input is MoRI's padded receive buffer
(num_max_dispatch_tokens_per_rank * world_size = 3072 * 8 = 24576 rows) rather
than the ~112 live tokens decode actually carries. num_rows is a device
pointer used only for a per-block early exit, so it cannot shrink the grid.

DSR-1 never hits this: it is w8a8, its dispatch delivers fp8 with per-128 fp32
scales which is already the MoE's input format, so AITER takes the
hidden_states.dtype == q_dtype_a branch and runs no quant kernel at all.

Modification

Adds an mxfp8 dispatch mode that quantizes on the send side, over live
tokens, emitting the fp8 + group-32 e8m0 layout the MoE kernels consume
directly. Opt-in via SGLANG_MORI_DISPATCH_DTYPE=mxfp8; the default stays
bf16, so existing deployments are untouched.

Depends on the matching aiter change (a8w4 mxfp8 passthrough in
fused_moe.py), which must land first.

Performance

DeepSeek-V4-Pro, MI355X (gfx950), rocm/sgl-dev:v0.5.18-rocm720-mi35x-20260822,
8192 in / 1024 out, TP8 + DP8 attention + EP8/MoRI + EAGLE MTP. Both arms are the
same build; the only difference is SGLANG_MORI_DISPATCH_DTYPE.

conc	tput bf16 (tok/s)	tput mxfp8 (tok/s)	tput	TPOT bf16 (ms)	TPOT mxfp8 (ms)	TPOT
4	1022.03	1121.97	+9.8%	33.30	30.19	-9.3%
8	1818.01	1969.84	+8.4%	35.52	32.73	-7.9%
16	3080.51	3277.43	+6.4%	42.12	39.52	-6.2%
32	4527.58	4905.24	+8.3%	57.79	53.17	-8.0%
64	6264.62	6827.92	+9.0%	85.21	78.19	-8.2%

TPOT mirrors throughput at every point, which is harder to attribute to drift
than either metric alone. For scale: repeat runs of an identical build on this
setup span 0.5-2.7%, so every cell here is well outside run-to-run noise.

One run per cell.

Commands

# Server, per arm; $CONC in {4, 8, 16, 32, 64}
# baseline arm: SGLANG_MORI_DISPATCH_DTYPE=bf16
# this PR:      SGLANG_MORI_DISPATCH_DTYPE=mxfp8
SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton \
SGLANG_OPT_FUSE_COMPRESS_NORM_ROPE=1 \
SGLANG_OPT_NATIVE_BPRESHUFFLE_SCALE=1 \
SGLANG_OPT_USE_AITER_BATCHED_GEMM=1 \
AITER_BF16_FP8_MOE_BOUND=0 SGLANG_USE_AITER=1 \
SGLANG_SHARED_EXPERT_TP1=1 SGLANG_DP_SHARED_EXPERT_LOCAL=1 \
SGLANG_DP_USE_GATHERV=1 SGLANG_DP_USE_REDUCE_SCATTER=1 \
GPU_MAX_HW_QUEUES=5 MORI_SHMEM_MODE=ISOLATION MORI_SHMEM_HEAP_SIZE=8G \
MORI_EP_LAUNCH_CONFIG_MODE=AUTO \
SGLANG_MORI_DISPATCH_DTYPE=$DTYPE \
python3 -m sglang.launch_server \
  --model-path deepseek-ai/DeepSeek-V4-Pro \
  --tensor-parallel-size 8 --dp 8 --enable-dp-attention \
  --ep-size 8 --moe-a2a-backend mori --deepep-mode normal \
  --moe-dense-tp-size 1 --enable-dp-lm-head \
  --attention-backend dsv4 --page-size 256 --kv-cache-dtype fp8_e4m3 \
  --enforce-shared-experts-fusion \
  --speculative-algorithm EAGLE --speculative-num-steps 3 \
  --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \
  --context-length 16384 --chunked-prefill-size 8192 \
  --mem-fraction-static 0.78 \
  --max-running-requests $CONC --cuda-graph-max-bs $CONC

# Perf client, per concurrency
python3 -m sglang.bench_serving --backend sglang-oai \
  --host 127.0.0.1 --port $PORT --model deepseek-ai/DeepSeek-V4-Pro \
  --dataset-name random --random-input-len 8192 --random-output-len 1024 \
  --random-range-ratio 1.0 --num-prompts $((CONC * 10)) --max-concurrency $CONC

# Accuracy
python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319 --parallel 1319

Note: under DP attention --max-running-requests is divided by dp_size, so a
value below the DP degree floors to 0 per rank and trips
assert max_running_request > 0. At $CONC 4 it was raised to 8.

Accuracy

GSM8K, 1319 questions, parallel 1319: 0.945, invalid 0.000.

Unit test

test/registered/unit/layers/test_moriep_mxfp8_dispatch.py (9 cases) pins the
byte layout and env wiring, which are the parts that regress silently -- an fp8
payload with the wrong scale group or dtype still runs, but reintroduces the
upscale round trip this mode exists to remove, and the only symptom is lost
throughput:

• scale group size is 32 (group-128 would force the upscale back)
• one scale per 32 channels; a mismatch under-allocates the scale buffer
• e8m0 scales are 1 byte -- the dispatch buffer is sized from this
• every DispatchDtype value is distinct, and mxfp8 is reachable from the env
• the zero-live-token decode path still emits a correctly shaped scale tensor

pytest test/registered/unit/layers/test_moriep_mxfp8_dispatch.py

Checklist

• End-to-end benchmarked on MI355X across conc 4-64 (table above)
• Accuracy verified, 1319 questions, zero invalid
• Unit test added
• Opt-in; default behaviour unchanged
• Blocked on the aiter side landing first (a8w4 mxfp8 passthrough)

---

CI States

Latest PR Test (Base): ❌ Run #32679407041
Latest PR Test (Extra): ❌ Run #32679406901
Latest PR Test (AMD ROCm 7.2): ❌ Run #32679406994

[karverma-amd]

Companion aiter PR: ROCm/aiter#4954 ([AMD][DSV4] feat: MXFP8 activation passthrough in fused_moe).

That side adds the 12-line branch letting fused_moe consume activations that already arrive as fp8 with group-32 e8m0 microscales, instead of re-deriving scales. This PR is what produces them, quantizing on the send side over live tokens.

The two are safe to merge in either order. With SGLANG_MORI_DISPATCH_DTYPE=mxfp8 set but the aiter change absent, the startup probe added in 8b0cd59 finds that this aiter's per_1x32 quant does not accept scale_type, logs a warning and falls back to bf16 dispatch — so the feature is inert rather than incorrect. The measured +6.4% to +9.8% only appears once both are in.