feat(ep): add blockwise-FP8 combine to the AsyncLL (low-latency) path

[rbrugaro-amd]

Title

feat(ep): add blockwise-FP8 combine to the AsyncLL (low-latency) path

Motivation

Blockwise-FP8 (E4M3, per-128-elem block max-scale) combine already exists for the
IntraNode / IntraNodeLL kernels but not for AsyncLL. The low-latency combine only had the
lossy direct-cast fp8 (UseFp8DirectCast — a plain bf16→e4m3 cast with a single fixed clamp and
no per-block scale). So requesting Fp8BlockwiseQuant on AsyncLL raised
Fp8BlockwiseQuant currently only supports IntraNode/IntraNodeLL combine, forcing callers to either
use bf16 (full combine payload) or accept direct-cast's accuracy loss on the async path.

This extends the existing intranode blockwise path to the four AsyncLL combine kernels. It is an
extension of an existing MoRI feature to a new kernel path — not a new algorithm.

What changed

• src/ops/dispatch_combine/low_latency_async.cpp — add bool UseFp8BlockwiseQuant to the 4 combine
  kernels:
  • SendCopy: quantize each token to [fp8 token | f32 block-scales]
    (WarpQuantizeToFp8Blockwise, bf16 T only).
  • SendTransfer / RecvCopy: slot stride accounts for the appended scale region.
  • RecvCopy dequant-accumulate: fast vec8 path
    (WarpAccumFp8Dequant{Full,Segment}BlockVec8Top8, block=128, AccumNum 8/9), matching intranode,
    with a general-segment fallback for other shapes.
• src/ops/kernels/ep_async_ll.hip — register _bf16_fp8bwq variants of SendCopy/SendTransfer/RecvCopy.
• src/ops/kernels/ep_common.hip — add WRAP_BOOL2 macro (2 template bools).
• python/mori/ops/dispatch_combine.py — select the _fp8bwq kernels in combine() / combine_recv();
  relax the guard to allow AsyncLL. (Shared-mem already accounts for the recv-side scale-ptr array.)
• tests/python/ops/test_dispatch_combine_async_ll.py — add fp8_blockwise to the AsyncLL
  dispatch+combine round-trip parametrization (bf16-only skip guard).

Why the same payload as direct-cast but accurate

Both fp8 modes move ~½ the bf16 combine payload (14336 → 7392 B/token @ hidden 7168). Direct-cast
uses one fixed scale across the whole vector → it flushes small per-channel values to E4M3
subnormals/zero on real activations. Blockwise rescales each 128-elem block → preserves them.

Evidence

End-to-end gsm8k, DeepSeek-R1-MXFP4, AsyncLL combine, only the combine dtype varies:

combine dtype	gsm8k	vs bf16	payload
bf16	0.967	—	full
fp8 blockwise (this PR)	0.947	−2.0 pp	½
fp8 direct-cast (stock)	0.873	−9.4 pp	½

→ blockwise ≈ bf16 accuracy at half payload; direct-cast costs ~7 pp for the same payload.

Correctness: extended AsyncLL round-trip test passes for fp8_blockwise (dispatch+combine
reconstruction, 8 ranks). The AsyncLL blockwise output is bit-faithful to the production intranode
blockwise combine on shared synthetic data.

How it was measured

Accuracy (gsm8k table under "Evidence")

Full SGLang server, DeepSeek-R1-MXFP4, one server per combine dtype (only SGLANG_MORI_COMBINE_DTYPE
changes), split path (IntraNode dispatch + AsyncLL combine), no SDMA, cuda graph off:

export SGLANG_USE_AITER=1 SGLANG_AITER_MOE=1 TORCH_BLAS_PREFER_HIPBLASLT=1
export SGLANG_MORI_SPLIT_DISPATCH_COMBINE=1 MORI_ENABLE_SDMA=0
export SGLANG_MORI_DISPATCH_DTYPE=auto
export SGLANG_MORI_COMBINE_DTYPE=fp8            # A/B: bf16 | fp8 (=blockwise) | fp8_direct_cast
export SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK=4096 SGLANG_MORI_NO_PAD_MASK=1
export MORI_SHMEM_MODE=ISOLATION MORI_SHMEM_HEAP_SIZE=16G

python3 -m sglang.launch_server \
  --model-path /model/Deepseek-R1-MXFP4-MLPerf \
  --tp-size 8 --ep-size 8 --dp-size 8 \
  --enable-dp-attention --enable-dp-lm-head --moe-dense-tp-size 1 \
  --moe-a2a-backend mori --enforce-shared-experts-fusion \
  --init-expert-location trivial --ep-num-redundant-experts 0 \
  --attention-backend aiter --kv-cache-dtype fp8_e4m3 \
  --mem-fraction-static 0.90 --deepep-mode low_latency \
  --disable-radix-cache --disable-chunked-prefix-cache \
  --disable-cuda-graph --trust-remote-code --port 30000 --host 127.0.0.1

# eval once the server is healthy:
python3 benchmark/gsm8k/bench_sglang.py --num-questions 300 --parallel 100 --port 30000 --host 127.0.0.1

Notes: top-9 effective (8 routed + 1 fused shared, enforce-shared-experts-fusion) → exercises the
vec8 AccumNum=9 path. bf16 was run at --num-questions 150 --parallel 100 (bf16 combine tail-hangs at
parallel 300; accuracy is concurrency-independent). direct-cast & blockwise at 300/100.

Testing

rm -rf ~/.mori/jit/*
PYTHONPATH=$PWD python3 -m pytest -s tests/python/ops/test_dispatch_combine_async_ll.py -k fp8_blockwise

SGLang — no change needed

SGLang already requests blockwise on the MoRI-EP path:
moriep.py maps SGLANG_MORI_COMBINE_DTYPE=fp8 → combine_quant_type="fp8_blockwise" (upstream,
AMD MoRI-EP integration, PR #29855). Pre-PR, AsyncLL rejected that request; this PR makes it honor
it. All logic (kernels + kernel selection + guard) is inside MoRI.

[jhchouuu]

LGTM overall, just left a couple of minor comments (dead template param on RecvTransfer + a top-9 test case).

Thanks for the nice work!

[jhchouuu]

seems UseFp8BlockwiseQuant is not used on EpCombineLowLatencyAsyncRecvTransfer_body...

[rbrugaro-amd]

Good catch — removed. EpCombineLowLatencyAsyncRecvTransfer_body is the barrier and never reads the flag; only SendCopy/SendTransfer/RecvCopy have _fp8bwq variants. Dropped it back to template <typename T, bool UseFp8DirectCast>. No functional change

[jhchouuu]

could you help add 9 to num_experts_per_token @pytest.mark.parametrize("num_experts_per_token", (8, 9)) for test, ty!

[rbrugaro-amd]

@jhchouuu Added num_experts_per_token=(8, 9), and top-9 now actually runs. The test's dispatch-result positional check (check_dispatch_result) doesn't support non-multiple-of-8 top-k, so for those cases the test validates the dispatch via the combine round-trip instead (a small _AsyncLLCombineOnlyTestCase that skips only the dispatch-side positional check). The dispatch data is correct at top-9 — the combine round-trip reconstructs it and production runs top-9 fine (gsm8k 0.967/0.947) — so nothing is dropped functionally: top-9 exercises the full dispatch+combine including the AccumNum=9 path, and top-8 keeps the full check.