From 304eeaa333576caa2fe2f0f9ae384223718b5dee Mon Sep 17 00:00:00 2001 From: Rita Brugarolas Brufau Date: Fri, 10 Jul 2026 18:16:59 -0500 Subject: [PATCH 1/3] feat(ep): add blockwise-FP8 combine to the AsyncLL (low-latency) path 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 (`UseFp8DirectCast`, a plain bf16->e4m3 cast with a single fixed clamp and no per-block scale). Requesting `Fp8BlockwiseQuant` on AsyncLL therefore raised "only supports IntraNode/IntraNodeLL combine", forcing callers to either use bf16 (full combine payload) or accept direct-cast's accuracy loss. This ports the intranode blockwise path into the four AsyncLL combine kernels: - SendCopy: quantize each token to [fp8 | f32 block-scales] via WarpQuantizeToFp8Blockwise (bf16 T only). - SendTransfer/RecvCopy: stride/slot accounts for the appended scale region. - RecvCopy dequant-accumulate uses the fast vec8 path (WarpAccumFp8Dequant{Full,Segment}BlockVec8Top8, block=128, AccumNum 8/9) matching intranode, with a general-segment fallback for other shapes. - Register `_bf16_fp8bwq` kernel variants (WRAP_BOOL2) and select them from combine()/combine_recv(); relax the python guard to allow AsyncLL. Same halved combine payload as direct-cast, at ~bf16-level accuracy. Accuracy (DeepSeek-R1-MXFP4, gsm8k, AsyncLL combine, only combine dtype varies): bf16 0.967 | blockwise-fp8 0.947 | direct-cast-fp8 0.873. Round-trip parity covered by the extended AsyncLL dispatch/combine test (quant_type="fp8_blockwise"). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Rita Brugarolas Brufau --- python/mori/ops/dispatch_combine.py | 27 ++++++- .../dispatch_combine/low_latency_async.cpp | 78 +++++++++++++++++-- src/ops/kernels/ep_async_ll.hip | 5 ++ src/ops/kernels/ep_common.hip | 5 ++ .../ops/test_dispatch_combine_async_ll.py | 4 +- 5 files changed, 109 insertions(+), 10 deletions(-) diff --git a/python/mori/ops/dispatch_combine.py b/python/mori/ops/dispatch_combine.py index 25e32b8fb..607f4f134 100644 --- a/python/mori/ops/dispatch_combine.py +++ b/python/mori/ops/dispatch_combine.py @@ -946,9 +946,10 @@ def combine( if kt not in ( EpDispatchCombineKernelType.IntraNode.value, EpDispatchCombineKernelType.IntraNodeLL.value, + EpDispatchCombineKernelType.AsyncLL.value, ): raise ValueError( - "Fp8BlockwiseQuant currently only supports IntraNode/IntraNodeLL combine" + "Fp8BlockwiseQuant currently only supports IntraNode/IntraNodeLL/AsyncLL combine" ) if sfx != "bf16": raise ValueError(f"Fp8BlockwiseQuant only supports bf16, got {sfx}") @@ -1095,6 +1096,18 @@ def combine( stream, args_ptr, ) + elif quant_type == EpDispatchCombineQuantType.Fp8BlockwiseQuant: + self._launch_multi( + [ + "EpCombineLowLatencyAsyncSendCopy_bf16_fp8bwq", + "EpCombineLowLatencyAsyncSendTransfer_bf16_fp8bwq", + ], + [mp_aligned, self.config.world_size], + [WARP_SIZE * actual_wpb, WARP_SIZE * actual_wpb], + [0, 0], + stream, + args_ptr, + ) else: self._launch_multi( [ @@ -1183,6 +1196,18 @@ def combine_recv( stream, args_ptr, ) + elif quant_type == EpDispatchCombineQuantType.Fp8BlockwiseQuant: + self._launch_multi( + [ + "EpCombineLowLatencyAsyncRecvTransfer_bf16", + "EpCombineLowLatencyAsyncRecvCopy_bf16_fp8bwq", + ], + [self.config.world_size, mp_aligned], + [WARP_SIZE * actual_wpb, WARP_SIZE * actual_wpb], + [0, shared_mem], + stream, + args_ptr, + ) else: self._launch_multi( [ diff --git a/src/ops/dispatch_combine/low_latency_async.cpp b/src/ops/dispatch_combine/low_latency_async.cpp index d95bc7774..618244594 100644 --- a/src/ops/dispatch_combine/low_latency_async.cpp +++ b/src/ops/dispatch_combine/low_latency_async.cpp @@ -415,7 +415,7 @@ __device__ void EpDispatchLowLatencyAsyncRecvCopyMultiBlock_body(EpDispatchCombi /* EpCombineLowLatencyAsyncSendCopy */ /* ---------------------------------------------------------------------------------------------- */ -template +template __device__ void EpCombineLowLatencyAsyncSendCopy_body(EpDispatchCombineArgs args) { DEF_COMMON_VARS; IF_ENABLE_PROFILER( @@ -433,7 +433,21 @@ __device__ void EpCombineLowLatencyAsyncSendCopy_body(EpDispatchCombineArgs a index_t stagingTokId = 0; if (laneId == 0) stagingTokId = args.dispReceiverIdxMap[tokenId]; stagingTokId = __shfl(stagingTokId, 0); - if constexpr (UseFp8DirectCast) { + if constexpr (UseFp8BlockwiseQuant) { + // Blockwise-FP8 (E4M3) combine: per-128-elem block max-scale, staged as [fp8 token | f32 scales]. + // Accurate fp8 combine for the AsyncLL path (matches intranode blockwise), unlike the lossy + // direct-cast. Only instantiated for bf16 T. + if constexpr (std::is_same_v) { + const int _sd = args.fp8BlockwiseCombineScaleDim; + const size_t _slotB = + hiddenDim * sizeof(core::CombineInternalFp8) + (size_t)_sd * sizeof(float); + uint8_t* _slot = stagingPtr + (size_t)stagingTokId * _slotB; + core::WarpQuantizeToFp8Blockwise( + reinterpret_cast(_slot), + reinterpret_cast(_slot + hiddenDim * sizeof(core::CombineInternalFp8)), + args.inpTokenBuf + tokenId * hiddenDim, hiddenDim, _sd); + } + } else if constexpr (UseFp8DirectCast) { core::WarpCastBf16ToCombineInternalFp8( reinterpret_cast(stagingPtr + stagingTokId * tokHiddenBytes), args.inpTokenBuf + tokenId * hiddenDim, hiddenDim, laneId); @@ -449,7 +463,7 @@ __device__ void EpCombineLowLatencyAsyncSendCopy_body(EpDispatchCombineArgs a /* EpCombineLowLatencyAsyncSendTransfer */ /* ---------------------------------------------------------------------------------------------- */ -template +template __device__ void EpCombineLowLatencyAsyncSendTransfer_body(EpDispatchCombineArgs args) { DEF_COMMON_VARS; IF_ENABLE_PROFILER( @@ -458,7 +472,10 @@ __device__ void EpCombineLowLatencyAsyncSendTransfer_body(EpDispatchCombineArgs< using TokT = std::conditional_t; static_assert(!UseFp8DirectCast || std::is_same_v, "Fp8 direct cast combine currently only supports bf16 input"); - const size_t tokHiddenBytes = hiddenDim * sizeof(TokT); + size_t tokHiddenBytes = hiddenDim * sizeof(TokT); + if constexpr (UseFp8BlockwiseQuant) + tokHiddenBytes = hiddenDim * sizeof(core::CombineInternalFp8) + + (size_t)args.fp8BlockwiseCombineScaleDim * sizeof(float); uint64_t* recvTokenNums = args.recvTokenNumMemObj->template GetAs(); for (int destPe = blockId; destPe < npes; destPe += blockNum) { @@ -491,7 +508,7 @@ __device__ void EpCombineLowLatencyAsyncSendTransfer_body(EpDispatchCombineArgs< /* EpCombineLowLatencyAsyncRecvTransfer */ /* ---------------------------------------------------------------------------------------------- */ -template +template __device__ void EpCombineLowLatencyAsyncRecvTransfer_body(EpDispatchCombineArgs args) { DEF_COMMON_VARS; IF_ENABLE_PROFILER( @@ -533,7 +550,7 @@ __device__ void EpCombineLowLatencyAsyncRecvTransfer_body(EpDispatchCombineArgs< /* EpCombineLowLatencyAsyncRecvCopy */ /* ---------------------------------------------------------------------------------------------- */ -template +template __device__ void EpCombineLowLatencyAsyncRecvCopy_body(EpDispatchCombineArgs args) { DEF_COMMON_VARS; IF_ENABLE_PROFILER( @@ -545,6 +562,9 @@ __device__ void EpCombineLowLatencyAsyncRecvCopy_body(EpDispatchCombineArgs a extern __shared__ char sharedMem[]; TokT** srcPtrs = reinterpret_cast(sharedMem) + warpId * config.numExpertPerToken; + // Blockwise-FP8 combine: per-expert block-scale pointer array, laid right after srcPtrs. + float** srcScalePtrs = reinterpret_cast(sharedMem) + + warpNum * config.numExpertPerToken + warpId * config.numExpertPerToken; if (args.curRankNumToken != 0) { MultiWarpIter mwIter(globalWarpNum, args.curRankNumToken, hiddenDim); @@ -563,15 +583,57 @@ __device__ void EpCombineLowLatencyAsyncRecvCopy_body(EpDispatchCombineArgs a ? args.interNodeTokBufs.combineInp->template GetAs() : args.interNodeTokBufs.staging->template GetAs(); if (destPe < npes) { - srcPtrs[j] = stagingPtr + destTokId * hiddenDim + hiddenDimOffset; + if constexpr (UseFp8BlockwiseQuant) { + const size_t _slotB = hiddenDim * sizeof(core::CombineInternalFp8) + + (size_t)args.fp8BlockwiseCombineScaleDim * sizeof(float); + uint8_t* _base = reinterpret_cast(stagingPtr) + (size_t)destTokId * _slotB; + srcPtrs[j] = reinterpret_cast(_base + hiddenDimOffset); + float* _sc = + reinterpret_cast(_base + hiddenDim * sizeof(core::CombineInternalFp8)); + srcScalePtrs[j] = (_sc[0] < 0.0f) ? _sc : nullptr; + } else { + srcPtrs[j] = stagingPtr + destTokId * hiddenDim + hiddenDimOffset; + } } else { srcPtrs[j] = nullptr; + if constexpr (UseFp8BlockwiseQuant) srcScalePtrs[j] = nullptr; } } T* outPtr = args.interNodeTokBufs.combineOut->template GetAs() + tokenId * hiddenDim + hiddenDimOffset; - if constexpr (UseFp8DirectCast) { + if constexpr (UseFp8BlockwiseQuant) { + if constexpr (std::is_same_v) { + // Fast vec8 dequant-accumulate (matches intranode): 8 elems/lane, native gfx950 cvt. + // Runtime-select the compile-time fast path for the common (block=128, AccumNum in {8,9}) + // aligned case; fall back to the general segment path otherwise. + auto _S = reinterpret_cast(srcPtrs); + auto _SC = reinterpret_cast(srcScalePtrs); + const int _be = args.fp8BlockwiseCombineScaleDim > 0 + ? (int)(hiddenDim / args.fp8BlockwiseCombineScaleDim) : 0; + const int _an = config.numExpertPerToken; + const bool _al = ((hiddenDimOffset & 0x7) == 0) && ((hiddenDimSize & 0x7) == 0); + bool _fast = true; + if (_be == 128 && _al && mwIter.warpsPerItem == 1 && _an == 8) + core::WarpAccumFp8DequantFullBlockVec8Top8( + outPtr, _S, _SC, hiddenDim); + else if (_be == 128 && _al && mwIter.warpsPerItem == 1 && _an == 9) + core::WarpAccumFp8DequantFullBlockVec8Top8( + outPtr, _S, _SC, hiddenDim); + else if (_be == 128 && _al && _an == 8) + core::WarpAccumFp8DequantSegmentBlockVec8Top8( + outPtr, _S, _SC, hiddenDimOffset, hiddenDimSize); + else if (_be == 128 && _al && _an == 9) + core::WarpAccumFp8DequantSegmentBlockVec8Top8( + outPtr, _S, _SC, hiddenDimOffset, hiddenDimSize); + else + _fast = false; + if (!_fast) + core::WarpAccumFp8DequantSegment( + outPtr, _S, _SC, _an, hiddenDimOffset, hiddenDimSize, hiddenDim, + args.fp8BlockwiseCombineScaleDim); + } + } else if constexpr (UseFp8DirectCast) { core::WarpAccumCombineInternalFp8ToBf16(outPtr, reinterpret_cast(srcPtrs), config.numExpertPerToken, laneId, hiddenDimSize); diff --git a/src/ops/kernels/ep_async_ll.hip b/src/ops/kernels/ep_async_ll.hip index f9eb15667..117e6ea13 100644 --- a/src/ops/kernels/ep_async_ll.hip +++ b/src/ops/kernels/ep_async_ll.hip @@ -18,3 +18,8 @@ WRAP_ALL_TYPES_BOOL(EpCombineLowLatencyAsyncRecvCopy, , false) WRAP_ALL_TYPES_BOOL(EpCombineLowLatencyAsyncRecvTransfer, , false) MORI_FP8_ANY(WRAP_BOOL(EpCombineLowLatencyAsyncRecvCopy, bf16_fp8cast, hip_bfloat16, true)) MORI_FP8_ANY(WRAP_BOOL(EpCombineLowLatencyAsyncRecvTransfer, bf16_fp8cast, hip_bfloat16, true)) + +// Blockwise-FP8 combine variants (bf16 in/out, fp8 on the wire): UseFp8DirectCast=false, UseFp8BlockwiseQuant=true. +MORI_FP8_ANY(WRAP_BOOL2(EpCombineLowLatencyAsyncSendCopy, bf16_fp8bwq, hip_bfloat16, false, true)) +MORI_FP8_ANY(WRAP_BOOL2(EpCombineLowLatencyAsyncSendTransfer, bf16_fp8bwq, hip_bfloat16, false, true)) +MORI_FP8_ANY(WRAP_BOOL2(EpCombineLowLatencyAsyncRecvCopy, bf16_fp8bwq, hip_bfloat16, false, true)) diff --git a/src/ops/kernels/ep_common.hip b/src/ops/kernels/ep_common.hip index cf5bb48ce..f19a0de53 100644 --- a/src/ops/kernels/ep_common.hip +++ b/src/ops/kernels/ep_common.hip @@ -83,6 +83,11 @@ using namespace mori::moe; kernel##_body(args); \ } +#define WRAP_BOOL2(kernel, suffix, T, B1, B2) \ + extern "C" __global__ void kernel##_##suffix(EpDispatchCombineArgs args) { \ + kernel##_body(args); \ + } + #define WRAP_ALL_TYPES_BOOL(kernel, bool_suffix, B) \ WRAP_BOOL(kernel, bf16##bool_suffix, hip_bfloat16, B) \ MORI_FP8_FNUZ(WRAP_BOOL(kernel, fp8_fnuz##bool_suffix, __hip_fp8_e4m3_fnuz, B)) \ diff --git a/tests/python/ops/test_dispatch_combine_async_ll.py b/tests/python/ops/test_dispatch_combine_async_ll.py index b94b15a8b..b681d5a74 100644 --- a/tests/python/ops/test_dispatch_combine_async_ll.py +++ b/tests/python/ops/test_dispatch_combine_async_ll.py @@ -205,7 +205,7 @@ def _test_dispatch_combine_multi_iteration( @pytest.mark.parametrize("max_num_inp_token_per_rank", (1, 128)) @pytest.mark.parametrize("num_experts_per_rank", (32,)) @pytest.mark.parametrize("num_experts_per_token", (8,)) -@pytest.mark.parametrize("quant_type", ("none", "fp8_direct_cast")) +@pytest.mark.parametrize("quant_type", ("none", "fp8_direct_cast", "fp8_blockwise")) def test_dispatch_combine( torch_dist_process_manager, world_size, @@ -220,6 +220,8 @@ def test_dispatch_combine( ): if quant_type == "fp8_direct_cast" and data_type is not torch.bfloat16: pytest.skip("fp8_direct_cast is only supported for bfloat16 data type") + if quant_type == "fp8_blockwise" and data_type is not torch.bfloat16: + pytest.skip("fp8_blockwise is only supported for bfloat16 data type") for _ in range(world_size): torch_dist_process_manager.task_queue.put( From ee10897176f67849618d408049cba4b7169fe314 Mon Sep 17 00:00:00 2001 From: Rita Brugarolas Brufau Date: Sat, 11 Jul 2026 00:33:47 -0500 Subject: [PATCH 2/3] style: apply clang-format to low_latency_async.cpp Co-Authored-By: Claude Opus 4.8 (1M context) --- src/ops/dispatch_combine/low_latency_async.cpp | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/ops/dispatch_combine/low_latency_async.cpp b/src/ops/dispatch_combine/low_latency_async.cpp index 618244594..3fa4892aa 100644 --- a/src/ops/dispatch_combine/low_latency_async.cpp +++ b/src/ops/dispatch_combine/low_latency_async.cpp @@ -434,9 +434,9 @@ __device__ void EpCombineLowLatencyAsyncSendCopy_body(EpDispatchCombineArgs a if (laneId == 0) stagingTokId = args.dispReceiverIdxMap[tokenId]; stagingTokId = __shfl(stagingTokId, 0); if constexpr (UseFp8BlockwiseQuant) { - // Blockwise-FP8 (E4M3) combine: per-128-elem block max-scale, staged as [fp8 token | f32 scales]. - // Accurate fp8 combine for the AsyncLL path (matches intranode blockwise), unlike the lossy - // direct-cast. Only instantiated for bf16 T. + // Blockwise-FP8 (E4M3) combine: per-128-elem block max-scale, staged as [fp8 token | f32 + // scales]. Accurate fp8 combine for the AsyncLL path (matches intranode blockwise), unlike + // the lossy direct-cast. Only instantiated for bf16 T. if constexpr (std::is_same_v) { const int _sd = args.fp8BlockwiseCombineScaleDim; const size_t _slotB = @@ -563,8 +563,8 @@ __device__ void EpCombineLowLatencyAsyncRecvCopy_body(EpDispatchCombineArgs a extern __shared__ char sharedMem[]; TokT** srcPtrs = reinterpret_cast(sharedMem) + warpId * config.numExpertPerToken; // Blockwise-FP8 combine: per-expert block-scale pointer array, laid right after srcPtrs. - float** srcScalePtrs = reinterpret_cast(sharedMem) + - warpNum * config.numExpertPerToken + warpId * config.numExpertPerToken; + float** srcScalePtrs = reinterpret_cast(sharedMem) + warpNum * config.numExpertPerToken + + warpId * config.numExpertPerToken; if (args.curRankNumToken != 0) { MultiWarpIter mwIter(globalWarpNum, args.curRankNumToken, hiddenDim); @@ -610,7 +610,8 @@ __device__ void EpCombineLowLatencyAsyncRecvCopy_body(EpDispatchCombineArgs a auto _S = reinterpret_cast(srcPtrs); auto _SC = reinterpret_cast(srcScalePtrs); const int _be = args.fp8BlockwiseCombineScaleDim > 0 - ? (int)(hiddenDim / args.fp8BlockwiseCombineScaleDim) : 0; + ? (int)(hiddenDim / args.fp8BlockwiseCombineScaleDim) + : 0; const int _an = config.numExpertPerToken; const bool _al = ((hiddenDimOffset & 0x7) == 0) && ((hiddenDimSize & 0x7) == 0); bool _fast = true; From 10e4ef02dba63e9d8f27c193fbc20ef5ede4a073 Mon Sep 17 00:00:00 2001 From: Rita Brugarolas Brufau Date: Mon, 13 Jul 2026 12:20:59 -0500 Subject: [PATCH 3/3] address review: drop dead RecvTransfer template param; cover top-9 combine - low_latency_async.cpp: remove the unused UseFp8BlockwiseQuant template parameter from EpCombineLowLatencyAsyncRecvTransfer_body. That kernel is the barrier (quiet + crossDeviceBarrier) and never reads the flag; only SendCopy/SendTransfer/RecvCopy have _fp8bwq variants. No functional change. - test_dispatch_combine_async_ll.py: add num_experts_per_token=9 to test_dispatch_combine so the vec8 AccumNum=9 combine path is covered. The dispatch-result positional check (check_dispatch_result) does not support non-multiple-of-8 top-k, so those cases validate the dispatch via the combine round-trip only (a _AsyncLLCombineOnlyTestCase that skips the dispatch-side positional check). The dispatch data itself is correct at top-9 -- the combine round-trip reconstructs it and production runs top-9 fine. top-8 keeps the full dispatch+combine check. Signed-off-by: Rita Brugarolas Brufau --- src/ops/dispatch_combine/low_latency_async.cpp | 2 +- .../ops/test_dispatch_combine_async_ll.py | 17 +++++++++++++++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/src/ops/dispatch_combine/low_latency_async.cpp b/src/ops/dispatch_combine/low_latency_async.cpp index 3fa4892aa..12370171d 100644 --- a/src/ops/dispatch_combine/low_latency_async.cpp +++ b/src/ops/dispatch_combine/low_latency_async.cpp @@ -508,7 +508,7 @@ __device__ void EpCombineLowLatencyAsyncSendTransfer_body(EpDispatchCombineArgs< /* EpCombineLowLatencyAsyncRecvTransfer */ /* ---------------------------------------------------------------------------------------------- */ -template +template __device__ void EpCombineLowLatencyAsyncRecvTransfer_body(EpDispatchCombineArgs args) { DEF_COMMON_VARS; IF_ENABLE_PROFILER( diff --git a/tests/python/ops/test_dispatch_combine_async_ll.py b/tests/python/ops/test_dispatch_combine_async_ll.py index b681d5a74..973b5331e 100644 --- a/tests/python/ops/test_dispatch_combine_async_ll.py +++ b/tests/python/ops/test_dispatch_combine_async_ll.py @@ -89,6 +89,14 @@ def run_test_once(self, op, test_data, check_results=True): self.check_combine_result(op, test_data, combine_output, None) +class _AsyncLLCombineOnlyTestCase(AsyncLLDispatchCombineTestCase): + # The dispatch-result positional check (check_dispatch_result) does not support + # non-multiple-of-8 top-k. The dispatch DATA is correct at such top-k -- it is fully + # validated by the combine round-trip -- so skip only the dispatch-side positional check. + def check_dispatch_result(self, *args, **kwargs): + return + + def _make_asyncll_config( rank, world_size, @@ -151,9 +159,14 @@ def _test_dispatch_combine( quant_type=quant_type, max_total_recv_tokens=max_total_recv_tokens, ) + test_case_cls = AsyncLLDispatchCombineTestCase + if num_experts_per_token % 8 != 0: + # Non-multiple-of-8 top-k: dispatch-result positional check is unsupported; + # validate the dispatch via the combine round-trip instead (data is correct). + test_case_cls = _AsyncLLCombineOnlyTestCase run_ep_dispatch_combine_test( config, - AsyncLLDispatchCombineTestCase, + test_case_cls, use_max_token_num=use_max_token_num, routing=routing, num_token_override=num_token_override, @@ -204,7 +217,7 @@ def _test_dispatch_combine_multi_iteration( @pytest.mark.parametrize("scale_type_size", (1, 4)) @pytest.mark.parametrize("max_num_inp_token_per_rank", (1, 128)) @pytest.mark.parametrize("num_experts_per_rank", (32,)) -@pytest.mark.parametrize("num_experts_per_token", (8,)) +@pytest.mark.parametrize("num_experts_per_token", (8, 9)) @pytest.mark.parametrize("quant_type", ("none", "fp8_direct_cast", "fp8_blockwise")) def test_dispatch_combine( torch_dist_process_manager,