feat(moriep): add fp4 combine dtype (SGLANG_MORI_COMBINE_DTYPE=fp4)#30706
Conversation
Exposes MoRI's blockwise-FP4 (E2M1) combine via SGLANG_MORI_COMBINE_DTYPE=fp4, which maps to MoRI's "fp4_blockwise" combine quant type (~half the combine transport of fp8). Adds a no-GPU unit test guarding the CombineDtype.fp4 enum wiring. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
|
Good patch! Shall we bump the mori version once your pr to mori gets merged? ROCm/mori#461 |
|
@HaiShaw :My Mori PR has been merged. Could you please add the CI label and review this PR when you have a chance? Thanks ! |
|
Hi @karverma-amd Thanks for your contributions! Could you also bump the version of mori in dockerfile? ROCm/mori@6f07277 |
Bumps MORI_COMMIT to ROCm/mori@6f072775 ("feat(ep): add blockwise FP4 (E2M1) combine transport", ROCm/mori#461), which the fp4 combine dtype (SGLANG_MORI_COMBINE_DTYPE=fp4) in this PR depends on. Addresses review feedback. Co-authored-by: Cursor <cursoragent@cursor.com>
|
@karverma-amd Can you post the test results of added UT above? |
@HaiShaw , Please see below results for UT
All 3 pass. The UT guards that CombineDtype.fp4 exists and maps to "fp4_blockwise", that CombineDtype("fp4_blockwise") round-trips back, and that the existing bf16 / fp8 / fp8_direct_cast members are unchanged |
|
@karverma-amd this PR break AMD CI (rocm7.0), we should temporarily revert the mori commit or AMD's CI is 100% blocked now, see: #31406 |
hello @yctseng0211 : I remember on Mori repo cco Ci test was broken and I am not sure on this failure. |
Summary
Exposes MoRI's blockwise-FP4 (E2M1) combine via
SGLANG_MORI_COMBINE_DTYPE=fp4. This maps toMoRI's
fp4_blockwisecombine quant type, which transports each combine element as packed FP4(2 values/byte) instead of FP8 (1 byte/elem) — ~half the combine payload. Combine is a
transport-bound step in MoE decode, so this improves decode throughput/latency with no accuracy loss.
Requires the companion MoRI change (
fp4_blockwisecombine kernels). Falls back cleanly: any otherSGLANG_MORI_COMBINE_DTYPEvalue keeps the existing behavior.What changed
moriep.py: addCombineDtype.fp4("fp4_blockwise"), parseSGLANG_MORI_COMBINE_DTYPE=fp4,and map it to MoRI's
fp4_blockwisecombine quant type. Existing dtypes(
bf16/fp8/fp8_direct_cast) are unchanged.test/manual/ep/test_moriep_combine_dtype.py: no-GPU unit test guarding theCombineDtype.fp4enum wiring and the existing dtypes.
Validation (MI355x / gfx950, DeepSeek-R1, TP8 EP8 DP-attention, MoRI-EP)
Accuracy — gsm8k, 2000 questions,
--parallel 1200,SGLANG_MORI_COMBINE_DTYPE=fp4×3Profile - With FP4

Command Used for Performance -
export SGLANG_USE_AITER=1 SGLANG_AITER_MOE=1 SGLANG_USE_ROCM700A=1
export SGLANG_ROCM_FUSED_DECODE_MLA=1 SGLANG_AITER_MLA_PERSIST=1
export SGLANG_DSR_OPROJ_MXFP4_ASM=1 TORCH_BLAS_PREFER_HIPBLASLT=1
export SGLANG_AITER_AR=1 ROCM_QUICK_REDUCE_QUANTIZATION=INT4 ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB=512
export TRITON_ENABLE_PRELOAD_KERNELS=0 TORCH_NCCL_BLOCKING_WAIT=1
export PYTORCH_ALLOC_CONF=expandable_segments:True
export SGLANG_AITER_FUSED_MLA_CONCAT_CACHE=1 SGLANG_AITER_FUSED_KVB_SPLIT_CAT=1
--- MoRI-EP ---
export MORI_SHMEM_MODE=ISOLATION MORI_SHMEM_HEAP_SIZE=16G SGLANG_MORI_LEC_IN_GRAPH=0
export SGLANG_MORI_DISPATCH_DTYPE=auto
export SGLANG_MORI_COMBINE_DTYPE=fp4 # or fp8 for the baseline
export SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK=4096 SGLANG_MORI_NO_PAD_MASK=1
python -m sglang.launch_server
--model-path /scratch/Deepseek-R1-MXFP4
--tp-size 8 --ep-size 8 --dp-size 8 --enable-dp-attention --enable-dp-lm-head
--moe-a2a-backend mori --deepep-mode normal --moe-dense-tp-size 1
--enforce-shared-experts-fusion
--trust-remote-code --kv-cache-dtype fp8_e4m3
--disable-radix-cache --mem-fraction-static 0.90
--cuda-graph-max-bs 512 --max-running-requests 2048
--host 0.0.0.0 --port 30000`
python /sgl-workspace/sglang/python/sglang/bench_serving.py
--backend sglang --host 127.0.0.1 --port 30000
--dataset-name random
--random-input-len 1024 --random-output-len 2048 --random-range-ratio 1.0
--num-prompts 1536 --max-concurrency 512
Performance — fp8 vs fp4 combine (decode-heavy: random, input 1024 / output 2048, 1536 prompts,
max-concurrency 512, identical server config)
Test plan
test/manual/ep/test_moriep_combine_dtype.py— passes.SGLANG_MORI_COMBINE_DTYPE=fp4(see above).
black/ruff/isortclean.CI States
Latest PR Test (Base): ❌ Run #29387324713
Latest PR Test (Extra): ❌ Run #29387324571
Mori PR reference - ROCm/mori#461