feat(dispatch_combine): add blockwise FP4 (E2M1) combine transport by karverma-amd · Pull Request #461 · ROCm/mori

karverma-amd · 2026-07-10T00:01:41Z

Summary

Adds a blockwise FP4 (E2M1) combine transport for the intra-node MoE combine, as a first-class
alternative to the existing FP8-blockwise combine. Each element is transported as packed FP4
(2 values/byte) instead of FP8 (1 byte/elem), halving the combine payload. Combine is a
transport-bound step in decode, so this directly improves decode throughput/latency while
preserving accuracy.

It is selected at launch via a new UseFp4Combine template parameter and the
quant_type="fp4_blockwise" config value. FP8 and FP4 combine kernels coexist in one build; the
FP8 path is unchanged (zero runtime cost — the branch is if constexpr).

What changed

device_primitives.hpp — vectorized FP4 quant (subwarp int4 loads / uint32 packed
stores, blockwise max-scale) and packed-FP4 dequant/accumulate (native gfx950
cvt_scalef32_pk_f32_fp4, load-hoisted). Reuses the FP8-blockwise staging/scale layout.
intranode.hpp — new UseFp4Combine template param branches the combine quant/dequant call
sites (if constexpr), gated by a static_assert to the FP8-blockwise path.
ep_intranode.hip / ep_common.hip — register _fp4bwq kernel variants (WRAP_BOOL8).
jit/core.py — enable the gfx950 OCP FP4 hardware intrinsics for the kernel compile
(arch-gated; no-op on other arches).
ops/dispatch_combine.py — "fp4_blockwise" quant type + _fp4bwq kernel selection.
Tests — extend the intra-node combine test (quant_type="fp4_blockwise", loose tolerance for
FP4 coarseness) + add no-GPU wiring guards (test_fp4_combine_wiring.py).

Validation (MI355x / gfx950, DeepSeek-R1, TP8 EP8 DP-attention, MoRI-EP)

Accuracy — gsm8k, 2000 questions, --parallel 1200, SGLANG_MORI_COMBINE_DTYPE=fp4 ×3

Run	Accuracy
1	0.948
2	0.950
3	0.948
Mean	0.949

Accuracy is stable (±0.001) and matches the FP8-combine baseline.

Profile-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)

Metric	FP8 combine	FP4 combine	Δ
Output token throughput	12,640.8 tok/s	13,103.1 tok/s	+3.7 %
Total token throughput	18,961.2 tok/s	19,654.7 tok/s	+3.7 %
Request throughput	6.17 req/s	6.40 req/s	+3.7 %
Mean TPOT	38.70 ms	37.40 ms	−3.4 %
Median TPOT	38.58 ms	36.82 ms	−4.6 %
Mean TTFT	3,405 ms	3,262 ms	−4.2 %
Mean E2E latency	82.6 s	79.8 s	−3.4 %

Test plan

tests/python/ops/test_fp4_combine_wiring.py (no-GPU wiring guards) — passes.
tests/python/ops/test_dispatch_combine_intranode.py with quant_type="fp4_blockwise"
(8-GPU numeric correctness within tolerance).
Builds and JIT-compiles on gfx950 (both fp8bwq and fp4bwq kernels emitted).
End-to-end accuracy + perf validated on DeepSeek-R1 (see above).
clang-format / black / ruff clean.

Sglang PR reference - sgl-project/sglang#30706

Adds a packed-FP4 (E2M1, 2 values/byte) variant of the intra-node blockwise combine, halving the combine transport payload vs FP8 (1 byte/elem) while reusing the existing FP8-blockwise staging/scale infrastructure. Selected at launch via a new UseFp4Combine template param (no runtime cost to FP8; fp8bwq and fp4bwq kernels coexist). - device_primitives.hpp: vectorized FP4 quant (subwarp int4 loads/uint32 stores) + packed FP4 dequant/accumulate (native gfx950 cvt_scalef32_pk_f32_fp4, load-hoisted). - intranode.hpp: UseFp4Combine template param branches quant/dequant call sites. - ep_intranode.hip/ep_common.hip: register _fp4bwq kernel variants (WRAP_BOOL8). - jit/core.py: enable gfx950 OCP FP4 hardware intrinsics. - ops/dispatch_combine.py: "fp4_blockwise" quant type + fp4bwq kernel selection. - tests: extend intranode combine test (quant_type=fp4_blockwise, loose tolerance) + add no-GPU wiring guards (test_fp4_combine_wiring.py). Co-authored-by: Cursor <cursoragent@cursor.com>

jhchouuu

Thanks for your work! some comments:

Don't alias FP4 onto the Fp8BlockwiseQuant enum. FP4 needs its own Fp4BlockwiseQuant enum + properly sized hiddenDim/2 staging slots — the shared enum / FP8-sized slots (2x staging waste, only separable via a Python string) can't ship as the design.
Add an arch guard for fp4_blockwise (fail-fast on non-gfx950, like fp8_direct_cast).
Harden kernel selection — the kernel_name.replace("_fp8bwq","_fp4bwq") silently yields non-existent names; assert the result is registered.

And CI test is failed, could you fix it?

jhchouuu · 2026-07-10T07:13:09Z

+        # packed-FP4 combine kernels at launch time.
+        self._combine_is_fp4 = (
+            isinstance(config.quant_type, str)
+            and config.quant_type.strip().lower() == "fp4_blockwise"
+        )


selection is arch-blind. On non-gfx950 this still picks _fp4bwq and silently falls to the software path. Please fail-fast (mirror the fp8_direct_cast arch guard).

jhchouuu · 2026-07-10T07:25:28Z

+                # FP4 (E2M1) is far coarser (~1 mantissa bit); blockwise-scaled + summed over
+                # top-k experts stays within a loose bound. This guards against gross corruption
+                # (wrong packing/indexing) rather than asserting FP8-level precision.
+                atol, rtol = 2.5e-1, 2.5e-1


please help to comfirm that atol,rtol=0.25 is actually right?

jhchouuu · 2026-07-10T07:27:02Z

      MORI_TRACE_NEXT(seq, Slot::CombineDequantAccum);
      if constexpr (Vec8Top8BlockElems != 0) {
        if (mwIter.warpsPerItem == 1) {
-          core::WarpAccumFp8DequantFullBlockVec8Top8<T, core::CombineInternalFp8,
-                                                     Vec8Top8BlockElems, Vec8AccumNum>(
-              outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
-              reinterpret_cast<const float* const*>(srcScalePtrs), hiddenDim);
+          if constexpr (UseFp4Combine) {
+            core::WarpAccumFp4DequantFullBlockVec8Top8<T, core::CombineInternalFp8,
+                                                       Vec8Top8BlockElems, Vec8AccumNum>(
+                outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
+                reinterpret_cast<const float* const*>(srcScalePtrs), hiddenDim);
+          } else {
+            core::WarpAccumFp8DequantFullBlockVec8Top8<T, core::CombineInternalFp8,
+                                                       Vec8Top8BlockElems, Vec8AccumNum>(
+                outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
+                reinterpret_cast<const float* const*>(srcScalePtrs), hiddenDim);
+          }
        } else if ((hiddenDimOffset & 0x7) == 0 && (hiddenDimSize & 0x7) == 0) {
-          core::WarpAccumFp8DequantSegmentBlockVec8Top8<T, core::CombineInternalFp8,
-                                                        Vec8Top8BlockElems, Vec8AccumNum>(
-              outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
-              reinterpret_cast<const float* const*>(srcScalePtrs), hiddenDimOffset, hiddenDimSize);
+          if constexpr (UseFp4Combine) {
+            core::WarpAccumFp4DequantSegmentBlockVec8Top8<T, core::CombineInternalFp8,
+                                                          Vec8Top8BlockElems, Vec8AccumNum>(
+                outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
+                reinterpret_cast<const float* const*>(srcScalePtrs), hiddenDimOffset,
+                hiddenDimSize);
+          } else {
+            core::WarpAccumFp8DequantSegmentBlockVec8Top8<T, core::CombineInternalFp8,
+                                                          Vec8Top8BlockElems, Vec8AccumNum>(
+                outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
+                reinterpret_cast<const float* const*>(srcScalePtrs), hiddenDimOffset,
+                hiddenDimSize);
+          }
        } else {
          // Misaligned segment: vec8 helper would fault on the load. Tiny scalar fallback.
-          core::WarpAccumFp8DequantSegmentScalarTop8<T, core::CombineInternalFp8,
-                                                     Vec8Top8BlockElems, Vec8AccumNum>(
-              outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
-              reinterpret_cast<const float* const*>(srcScalePtrs), hiddenDimOffset, hiddenDimSize);
+          if constexpr (UseFp4Combine) {
+            core::WarpAccumFp4DequantSegment<T, core::CombineInternalFp8>(
+                outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
+                reinterpret_cast<const float* const*>(srcScalePtrs), Vec8AccumNum, hiddenDimOffset,
+                hiddenDimSize, hiddenDim, args.fp8BlockwiseCombineScaleDim);
+          } else {
+            core::WarpAccumFp8DequantSegmentScalarTop8<T, core::CombineInternalFp8,
+                                                       Vec8Top8BlockElems, Vec8AccumNum>(
+                outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
+                reinterpret_cast<const float* const*>(srcScalePtrs), hiddenDimOffset,
+                hiddenDimSize);
+          }
        }
      } else {
        if (mwIter.warpsPerItem == 1) {
-          core::WarpAccumFp8DequantFull<T, core::CombineInternalFp8>(
-              outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
-              reinterpret_cast<const float* const*>(srcScalePtrs), validAccumCount, hiddenDim,
-              args.fp8BlockwiseCombineScaleDim);
+          if constexpr (UseFp4Combine) {
+            core::WarpAccumFp4DequantFull<T, core::CombineInternalFp8>(
+                outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
+                reinterpret_cast<const float* const*>(srcScalePtrs), validAccumCount, hiddenDim,
+                args.fp8BlockwiseCombineScaleDim);
+          } else {
+            core::WarpAccumFp8DequantFull<T, core::CombineInternalFp8>(
+                outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
+                reinterpret_cast<const float* const*>(srcScalePtrs), validAccumCount, hiddenDim,
+                args.fp8BlockwiseCombineScaleDim);
+          }
        } else {
-          core::WarpAccumFp8DequantSegment<T, core::CombineInternalFp8>(
-              outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
-              reinterpret_cast<const float* const*>(srcScalePtrs), validAccumCount, hiddenDimOffset,
-              hiddenDimSize, hiddenDim, args.fp8BlockwiseCombineScaleDim);
+          if constexpr (UseFp4Combine) {
+            core::WarpAccumFp4DequantSegment<T, core::CombineInternalFp8>(
+                outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
+                reinterpret_cast<const float* const*>(srcScalePtrs), validAccumCount,
+                hiddenDimOffset, hiddenDimSize, hiddenDim, args.fp8BlockwiseCombineScaleDim);
+          } else {
+            core::WarpAccumFp8DequantSegment<T, core::CombineInternalFp8>(
+                outPtr, reinterpret_cast<const core::CombineInternalFp8* const*>(srcPtrs),
+                reinterpret_cast<const float* const*>(srcScalePtrs), validAccumCount,
+                hiddenDimOffset, hiddenDimSize, hiddenDim, args.fp8BlockwiseCombineScaleDim);
+          }


duplicate every call. is possilble for any helper function?

jhchouuu · 2026-07-10T07:36:46Z

+  if (warpSize == 64 && blockElems == 128 && (hiddenDim % 128) == 0 &&
+      std::is_same_v<InT, hip_bfloat16>) {
+    WarpQuantizeToFp4BlockwiseVec<8, 16, Fp8T, InT>(dstToken, dstScales, srcToken, hiddenDim,
+                                                    scaleDim);
+    return;
+  }


for block256, seems don't have blockElems == 256 fast path...

…sites, block256 fast path - dispatch_combine.py: fail-fast when fp4_blockwise combine is requested on a non-gfx950 GPU (no OCP FP4 hardware path) instead of silently falling back to software. - test: derive the fp4 correctness bound from FP4's actual worst-case rounding (blockMax/6 per element, scaled by the accumulation count) instead of a guessed atol/rtol=0.25; skip fp4 cases on non-gfx950 so the gfx942 CI runner passes. - intranode.hpp: factor the per-site fp4/fp8 if-constexpr branches into WarpQuantizeToCombineBlockwise / WarpAccumCombineDequant* dispatch helpers so the combine body calls one function per site. - device_primitives.hpp: add the blockElems==256 vectorized fp4 quant fast path (mirrors the fp8bwq block256 variant) so block256 configs are not left on the scalar path. Co-authored-by: Cursor <cursoragent@cursor.com>

…, hardened selection Per review, FP4 no longer aliases the Fp8BlockwiseQuant enum / FP8-sized slots: - Add QuantType::Fp4BlockwiseQuant (enum + pybind) and IsBlockwiseCombineQuant() helper; the blockwise combine paths (host allocator, launch.cpp, dispatch_combine.py) treat FP8/FP4 together and differ only in the element codec. - Size FP4 staging slots at hiddenDim/2 (packed E2M1, 0.5 byte/elem) via CombineTokenRegionBytes() in the host allocator and the matching hiddenBytes in the intra-node kernel -- no FP8-sized over-allocation. FP8 sizing is unchanged. - dispatch_combine.py maps "fp4_blockwise" -> Fp4BlockwiseQuant (no string-only separation) and asserts the derived fp4bwq kernel name is one of the registered symbols (fail loudly instead of launching a non-existent kernel). Co-authored-by: Cursor <cursoragent@cursor.com>

karverma-amd · 2026-07-10T15:39:13Z

Thanks for your work! some comments:

Don't alias FP4 onto the Fp8BlockwiseQuant enum. FP4 needs its own Fp4BlockwiseQuant enum + properly sized hiddenDim/2 staging slots — the shared enum / FP8-sized slots (2x staging waste, only separable via a Python string) can't ship as the design.

Add an arch guard for fp4_blockwise (fail-fast on non-gfx950, like fp8_direct_cast).

Harden kernel selection — the kernel_name.replace("_fp8bwq","_fp4bwq") silently yields non-existent names; assert the result is registered.

And CI test is failed, could you fix it?

@jhchouuu ,

Thanks for the review! Addressed all points in two commits (0589070, c7548ec):

1. Dedicated enum + properly-sized slots (no FP8 aliasing): Added QuantType::Fp4BlockwiseQuant (enum + pybind) and an IsBlockwiseCombineQuant() helper. FP4 staging slots are now sized at hiddenDim/2 (packed E2M1, 0.5 byte/elem) via the new CombineTokenRegionBytes() in the host allocator, with the matching hiddenBytes in the intra-node kernel — no FP8-sized over-allocation, and no string-only separation. FP8 sizing is unchanged.

2. Arch guard: fp4_blockwise now fails fast on non-gfx950 (no OCP FP4 conversion path), mirroring the fp8_direct_cast guard; the intranode test also skips fp4 on non-gfx950 (this fixes the MI325X/gfx942 CI).

3. Hardened kernel selection: the derived _fp4bwq name is asserted against the set of registered fp4bwq symbols, so a mismatch fails loudly instead of launching a non-existent kernel.

Inline comments:

tolerance 0.25: you were right — a fixed rtol is meaningless for FP4 (small values in a high-max block round to 0/0.5×scale ⇒ ~100% rel error), which is what failed CI on gfx950. Replaced it with an exact per-element bound: within a block scaled to [-6,6] the largest E2M1 step is |6-4|=2 ⇒ max rounding error = blockMax/6, times the accumulation count.
duplicate calls: factored into WarpQuantizeToCombineBlockwise / WarpAccumCombineDequant*<UseFp4,...> dispatch helpers (one call per site).
block256 fast path: added the blockElems==256 vectorized fp4 quant fast path.

Validation (gfx950, rebuilt libmori_ops/libmori_pybinds): fp4 18/18 main+vec8 cases pass with the new enum + halved staging; fp8 blockwise not regressed; clang-format/black/ruff clean.

jhchouuu

LGTM, A small suggestion that is hoping to be revised.

Thanks for the great work and the quick turnaround on the review comments!

jhchouuu · 2026-07-13T10:02:01Z

+template <typename Fp8T, typename InT>
+__device__ __forceinline__ void WarpQuantizeToFp4Blockwise(Fp8T* __restrict__ dstToken,
+                                                           float* __restrict__ dstScales,


the new FP4 helpers still name their storage template param Fp8T even though it carries packed FP4, purely cosmetic, but renaming it to StoreT/PackedT would read clearer.

karverma-amd mentioned this pull request Jul 10, 2026

feat(moriep): add fp4 combine dtype (SGLANG_MORI_COMBINE_DTYPE=fp4) sgl-project/sglang#30706

Open

3 tasks

jhchouuu self-requested a review July 10, 2026 04:10

jhchouuu reviewed Jul 10, 2026

View reviewed changes

karverma-amd and others added 2 commits July 10, 2026 09:37

jhchouuu approved these changes Jul 13, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat(dispatch_combine): add blockwise FP4 (E2M1) combine transport #461

feat(dispatch_combine): add blockwise FP4 (E2M1) combine transport #461
karverma-amd wants to merge 3 commits into
ROCm:mainfrom
karverma-amd:feat/fp4-blockwise-combine

karverma-amd commented Jul 10, 2026 •

edited

Loading

Uh oh!

jhchouuu left a comment

Uh oh!

jhchouuu Jul 10, 2026

Uh oh!

jhchouuu Jul 10, 2026

Uh oh!

jhchouuu Jul 10, 2026

Uh oh!

jhchouuu Jul 10, 2026

Uh oh!

karverma-amd commented Jul 10, 2026

Uh oh!

jhchouuu left a comment

Uh oh!

jhchouuu Jul 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

karverma-amd commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What changed

Validation (MI355x / gfx950, DeepSeek-R1, TP8 EP8 DP-attention, MoRI-EP)

--- MoRI-EP ---

Test plan

Uh oh!

jhchouuu left a comment

Choose a reason for hiding this comment

Uh oh!

jhchouuu Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

jhchouuu Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

jhchouuu Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

jhchouuu Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

karverma-amd commented Jul 10, 2026

Uh oh!

jhchouuu left a comment

Choose a reason for hiding this comment

Uh oh!

jhchouuu Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

karverma-amd commented Jul 10, 2026 •

edited

Loading