Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jhchouuu merged 4 commits into from
Jul 15, 2026
Merged

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

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
304eeaa
feat(ep): add blockwise-FP8 combine to the AsyncLL (low-latency) path
rbrugaro-amd Jul 10, 2026
ee10897
style: apply clang-format to low_latency_async.cpp
rbrugaro-amd Jul 11, 2026
10e4ef0
address review: drop dead RecvTransfer template param; cover top-9 co…
rbrugaro-amd Jul 13, 2026
62a4ad3
Merge remote-tracking branch 'origin/main' into pr465-rebased
jhchouuu Jul 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 26 additions & 1 deletion python/mori/ops/dispatch_combine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}")
Expand Down Expand Up @@ -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(
[
Expand Down Expand Up @@ -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(
[
Expand Down
79 changes: 71 additions & 8 deletions src/ops/dispatch_combine/low_latency_async.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -415,7 +415,7 @@ __device__ void EpDispatchLowLatencyAsyncRecvCopyMultiBlock_body(EpDispatchCombi
/* EpCombineLowLatencyAsyncSendCopy */
/* ---------------------------------------------------------------------------------------------- */

template <typename T, bool UseFp8DirectCast>
template <typename T, bool UseFp8DirectCast, bool UseFp8BlockwiseQuant = false>
__device__ void EpCombineLowLatencyAsyncSendCopy_body(EpDispatchCombineArgs<T> args) {
DEF_COMMON_VARS;
IF_ENABLE_PROFILER(
Expand All @@ -433,7 +433,21 @@ __device__ void EpCombineLowLatencyAsyncSendCopy_body(EpDispatchCombineArgs<T> 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<T, hip_bfloat16>) {
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<core::CombineInternalFp8>(
reinterpret_cast<core::CombineInternalFp8*>(_slot),
reinterpret_cast<float*>(_slot + hiddenDim * sizeof(core::CombineInternalFp8)),
args.inpTokenBuf + tokenId * hiddenDim, hiddenDim, _sd);
}
} else if constexpr (UseFp8DirectCast) {
core::WarpCastBf16ToCombineInternalFp8<T>(
reinterpret_cast<TokT*>(stagingPtr + stagingTokId * tokHiddenBytes),
args.inpTokenBuf + tokenId * hiddenDim, hiddenDim, laneId);
Expand All @@ -449,7 +463,7 @@ __device__ void EpCombineLowLatencyAsyncSendCopy_body(EpDispatchCombineArgs<T> a
/* EpCombineLowLatencyAsyncSendTransfer */
/* ---------------------------------------------------------------------------------------------- */

template <typename T, bool UseFp8DirectCast>
template <typename T, bool UseFp8DirectCast, bool UseFp8BlockwiseQuant = false>
__device__ void EpCombineLowLatencyAsyncSendTransfer_body(EpDispatchCombineArgs<T> args) {
DEF_COMMON_VARS;
IF_ENABLE_PROFILER(
Expand All @@ -458,7 +472,10 @@ __device__ void EpCombineLowLatencyAsyncSendTransfer_body(EpDispatchCombineArgs<
using TokT = std::conditional_t<UseFp8DirectCast, core::CombineInternalFp8, T>;
static_assert(!UseFp8DirectCast || std::is_same_v<T, hip_bfloat16>,
"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<uint64_t*>();
for (int destPe = blockId; destPe < npes; destPe += blockNum) {
Expand Down Expand Up @@ -491,7 +508,7 @@ __device__ void EpCombineLowLatencyAsyncSendTransfer_body(EpDispatchCombineArgs<
/* EpCombineLowLatencyAsyncRecvTransfer */
/* ---------------------------------------------------------------------------------------------- */

template <typename T, bool UseFp8DirectCast>
template <typename T, bool UseFp8DirectCast, bool UseFp8BlockwiseQuant = false>

@jhchouuu jhchouuu Jul 13, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

seems UseFp8BlockwiseQuant is not used on EpCombineLowLatencyAsyncRecvTransfer_body...

@rbrugaro-amd rbrugaro-amd Jul 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

__device__ void EpCombineLowLatencyAsyncRecvTransfer_body(EpDispatchCombineArgs<T> args) {
DEF_COMMON_VARS;
IF_ENABLE_PROFILER(
Expand Down Expand Up @@ -533,7 +550,7 @@ __device__ void EpCombineLowLatencyAsyncRecvTransfer_body(EpDispatchCombineArgs<
/* EpCombineLowLatencyAsyncRecvCopy */
/* ---------------------------------------------------------------------------------------------- */

template <typename T, bool UseFp8DirectCast>
template <typename T, bool UseFp8DirectCast, bool UseFp8BlockwiseQuant = false>
__device__ void EpCombineLowLatencyAsyncRecvCopy_body(EpDispatchCombineArgs<T> args) {
DEF_COMMON_VARS;
IF_ENABLE_PROFILER(
Expand All @@ -545,6 +562,9 @@ __device__ void EpCombineLowLatencyAsyncRecvCopy_body(EpDispatchCombineArgs<T> a

extern __shared__ char sharedMem[];
TokT** srcPtrs = reinterpret_cast<TokT**>(sharedMem) + warpId * config.numExpertPerToken;
// Blockwise-FP8 combine: per-expert block-scale pointer array, laid right after srcPtrs.
float** srcScalePtrs = reinterpret_cast<float**>(sharedMem) + warpNum * config.numExpertPerToken +
warpId * config.numExpertPerToken;

if (args.curRankNumToken != 0) {
MultiWarpIter mwIter(globalWarpNum, args.curRankNumToken, hiddenDim);
Expand All @@ -563,15 +583,58 @@ __device__ void EpCombineLowLatencyAsyncRecvCopy_body(EpDispatchCombineArgs<T> a
? args.interNodeTokBufs.combineInp->template GetAs<TokT*>()
: args.interNodeTokBufs.staging->template GetAs<TokT*>();
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<uint8_t*>(stagingPtr) + (size_t)destTokId * _slotB;
srcPtrs[j] = reinterpret_cast<TokT*>(_base + hiddenDimOffset);
float* _sc =
reinterpret_cast<float*>(_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<T*>() + tokenId * hiddenDim +
hiddenDimOffset;
if constexpr (UseFp8DirectCast) {
if constexpr (UseFp8BlockwiseQuant) {
if constexpr (std::is_same_v<T, hip_bfloat16>) {
// 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<const core::CombineInternalFp8* const*>(srcPtrs);
auto _SC = reinterpret_cast<const float* const*>(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<T, core::CombineInternalFp8, 128, 8>(
outPtr, _S, _SC, hiddenDim);
else if (_be == 128 && _al && mwIter.warpsPerItem == 1 && _an == 9)
core::WarpAccumFp8DequantFullBlockVec8Top8<T, core::CombineInternalFp8, 128, 9>(
outPtr, _S, _SC, hiddenDim);
else if (_be == 128 && _al && _an == 8)
core::WarpAccumFp8DequantSegmentBlockVec8Top8<T, core::CombineInternalFp8, 128, 8>(
outPtr, _S, _SC, hiddenDimOffset, hiddenDimSize);
else if (_be == 128 && _al && _an == 9)
core::WarpAccumFp8DequantSegmentBlockVec8Top8<T, core::CombineInternalFp8, 128, 9>(
outPtr, _S, _SC, hiddenDimOffset, hiddenDimSize);
else
_fast = false;
if (!_fast)
core::WarpAccumFp8DequantSegment<T, core::CombineInternalFp8>(
outPtr, _S, _SC, _an, hiddenDimOffset, hiddenDimSize, hiddenDim,
args.fp8BlockwiseCombineScaleDim);
}
} else if constexpr (UseFp8DirectCast) {
core::WarpAccumCombineInternalFp8ToBf16(outPtr,
reinterpret_cast<const TokT* const*>(srcPtrs),
config.numExpertPerToken, laneId, hiddenDimSize);
Expand Down
5 changes: 5 additions & 0 deletions src/ops/kernels/ep_async_ll.hip
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
5 changes: 5 additions & 0 deletions src/ops/kernels/ep_common.hip
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,11 @@ using namespace mori::moe;
kernel##_body<T, B>(args); \
}

#define WRAP_BOOL2(kernel, suffix, T, B1, B2) \
extern "C" __global__ void kernel##_##suffix(EpDispatchCombineArgs<T> args) { \
kernel##_body<T, B1, B2>(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)) \
Expand Down
4 changes: 3 additions & 1 deletion tests/python/ops/test_dispatch_combine_async_ll.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"))

@jhchouuu jhchouuu Jul 13, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@rbrugaro-amd rbrugaro-amd Jul 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@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.

def test_dispatch_combine(
torch_dist_process_manager,
world_size,
Expand All @@ -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(
Expand Down
Loading