Skip to content

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

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

Open
karverma-amd wants to merge 3 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4b88ace
feat(dispatch_combine): add blockwise FP4 (E2M1) combine transport
karverma-amd Jul 9, 2026
0589070
address review: arch-guard fp4, exact fp4 test tolerance, dedup call …
karverma-amd Jul 10, 2026
c7548ec
address review: dedicated Fp4BlockwiseQuant enum + half-sized staging…
karverma-amd Jul 10, 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
335 changes: 335 additions & 0 deletions include/mori/core/transport/p2p/device_primitives.hpp
View file
Open in desktop

Large diffs are not rendered by default.

8 changes: 8 additions & 0 deletions python/mori/jit/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -285,6 +285,13 @@ def _profiler_defines() -> list[str]:
return ["-DENABLE_PROFILER"] if is_profiler_enabled() else []


def _ocp_fp_defines(arch: str) -> list[str]:
"""Enable the native gfx950 OCP FP4/FP8 conversion instructions (cvt_scalef32_pk_*) used by
the fp4_blockwise combine's E2M1 quant/dequant helpers. Without this the helpers fall back to
slow software bit-manipulation. Only relevant on gfx950; a no-op elsewhere."""
return ["-DHIP_ENABLE_GFX950_OCP_BUILTINS=1"] if "gfx950" in str(arch) else []


def _debuginfo_flags() -> list[str]:
"""Return hipcc debug flags if MORI_DEBUG_INFO is enabled."""
return ["-g", "-ggdb"] if is_debuginfo_enabled() else []
Expand Down Expand Up @@ -415,6 +422,7 @@ def _hipcc_genco(
*_nic_defines(),
*_ccqe_defines(),
*_profiler_defines(),
*_ocp_fp_defines(cfg.arch),
]

for d in include_dirs:
Expand Down
14 changes: 14 additions & 0 deletions python/mori/ops/dispatch_combine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,10 @@ def __str__(self):
"none": EpDispatchCombineQuantType.None_,
"fp8_direct_cast": EpDispatchCombineQuantType.Fp8DirectCast,
"fp8_blockwise": EpDispatchCombineQuantType.Fp8BlockwiseQuant,
# Blockwise FP4 (E2M1) combine reuses the FP8-blockwise staging/scale infrastructure at the
# config level (same 1-byte-slot buffers + float scales) but selects packed-FP4 combine
# kernels at launch (see _combine_is_fp4); it transports 0.5 byte/elem instead of 1.
"fp4_blockwise": EpDispatchCombineQuantType.Fp8BlockwiseQuant,
}


Expand Down Expand Up @@ -262,6 +266,12 @@ def __init__(self, config):
self._fp8_blockwise_combine_scale_type_size = self._handle_info[
"fp8_blockwise_combine_scale_type_size"
]
# "fp4_blockwise" maps to the Fp8BlockwiseQuant enum (shared buffers/scales) but selects
# packed-FP4 combine kernels at launch time.
self._combine_is_fp4 = (
isinstance(config.quant_type, str)
and config.quant_type.strip().lower() == "fp4_blockwise"
)

@jhchouuu jhchouuu Jul 10, 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.

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


self._dispatch_out_ptrs = mori_cpp.get_dispatch_output_ptrs(self._handle, True)
self._combine_out_ptrs = mori_cpp.get_combine_output_ptrs(self._handle, True)
Expand Down Expand Up @@ -1037,6 +1047,10 @@ def combine(
else "EpCombineIntraNodeKernel_bf16_nop2p_fp8bwq_noweight_block256_vec8"
)
use_vec8_top8 = True
# Blockwise FP4: select the packed-FP4 kernel variants (identical launch config to
# the fp8bwq variants; only the in-kernel quant/dequant math differs).
if self._combine_is_fp4:
kernel_name = kernel_name.replace("_fp8bwq", "_fp4bwq")
shared_mem = self._combine_shared_mem(
actual_wpb, use_weights=not use_vec8_top8
)
Expand Down
118 changes: 87 additions & 31 deletions src/ops/dispatch_combine/intranode.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,10 +264,14 @@ __global__ void EpDispatchIntraNodeKernel(EpDispatchCombineArgs<T> args) {
/* ---------------------------------------------------------------------------------------------- */
template <typename T, bool UseP2PRead = true, bool EnableStdMoE = false,
bool UseFp8DirectCast = false, bool UseFp8BlockwiseQuant = false, bool UseWeights = true,
int Vec8Top8BlockElems = 0, int Vec8AccumNum = 8>
int Vec8Top8BlockElems = 0, int Vec8AccumNum = 8, bool UseFp4Combine = false>
__device__ __forceinline__ void EpCombineIntraNodeKernel_body(EpDispatchCombineArgs<T> args) {
using TokT =
std::conditional_t<UseFp8DirectCast || UseFp8BlockwiseQuant, core::CombineInternalFp8, T>;
// UseFp4Combine reuses the FP8-blockwise staging/scale layout but transports each element as
// packed FP4 (E2M1, 2/byte -> half the combine bytes). It is a variant of blockwise combine.
static_assert(!UseFp4Combine || UseFp8BlockwiseQuant,
"UseFp4Combine builds on the FP8-blockwise combine path");
static_assert(!(UseFp8DirectCast && UseFp8BlockwiseQuant),
"Fp8 direct cast and blockwise quant are mutually exclusive");
static_assert((!UseFp8DirectCast && !UseFp8BlockwiseQuant) || std::is_same_v<T, hip_bfloat16>,
Expand Down Expand Up @@ -358,10 +362,17 @@ __device__ __forceinline__ void EpCombineIntraNodeKernel_body(EpDispatchCombineA
uint8_t* destStagingPtr = args.intraNodeTokBufs.combineInp->template GetAs<uint8_t*>(destPe) +
SendBufSlotOffset(config, myPe, destLocalTokId) * combXferBytes;
if constexpr (UseFp8BlockwiseQuant) {
core::WarpQuantizeToFp8Blockwise<core::CombineInternalFp8>(
reinterpret_cast<core::CombineInternalFp8*>(destStagingPtr),
reinterpret_cast<float*>(destStagingPtr + hiddenBytes),
args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim);
if constexpr (UseFp4Combine) {
core::WarpQuantizeToFp4Blockwise<core::CombineInternalFp8>(
reinterpret_cast<core::CombineInternalFp8*>(destStagingPtr),
reinterpret_cast<float*>(destStagingPtr + hiddenBytes),
args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim);
} else {
core::WarpQuantizeToFp8Blockwise<core::CombineInternalFp8>(
reinterpret_cast<core::CombineInternalFp8*>(destStagingPtr),
reinterpret_cast<float*>(destStagingPtr + hiddenBytes),
args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim);
}
} else if constexpr (!std::is_same_v<T, TokT> &&
std::is_same_v<TokT, core::CombineInternalFp8>) {
core::WarpCastBf16ToCombineInternalFp8<T>(reinterpret_cast<TokT*>(destStagingPtr),
Expand Down Expand Up @@ -396,10 +407,17 @@ __device__ __forceinline__ void EpCombineIntraNodeKernel_body(EpDispatchCombineA
uint8_t* destStagingPtr = args.intraNodeTokBufs.combineInp->template GetAs<uint8_t*>(destPe) +
SendBufSlotOffset(config, myPe, destLocalTokId) * combXferBytes;
if constexpr (UseFp8BlockwiseQuant) {
core::WarpQuantizeToFp8Blockwise<core::CombineInternalFp8>(
reinterpret_cast<core::CombineInternalFp8*>(destStagingPtr),
reinterpret_cast<float*>(destStagingPtr + hiddenBytes),
args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim);
if constexpr (UseFp4Combine) {
core::WarpQuantizeToFp4Blockwise<core::CombineInternalFp8>(
reinterpret_cast<core::CombineInternalFp8*>(destStagingPtr),
reinterpret_cast<float*>(destStagingPtr + hiddenBytes),
args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim);
} else {
core::WarpQuantizeToFp8Blockwise<core::CombineInternalFp8>(
reinterpret_cast<core::CombineInternalFp8*>(destStagingPtr),
reinterpret_cast<float*>(destStagingPtr + hiddenBytes),
args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim);
}
} else if constexpr (!std::is_same_v<T, TokT> &&
std::is_same_v<TokT, core::CombineInternalFp8>) {
core::WarpCastBf16ToCombineInternalFp8<T>(reinterpret_cast<TokT*>(destStagingPtr),
Expand Down Expand Up @@ -537,33 +555,71 @@ __device__ __forceinline__ void EpCombineIntraNodeKernel_body(EpDispatchCombineA
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);
}

@jhchouuu jhchouuu Jul 10, 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.

duplicate every call. is possilble for any helper function?

}
}
} else if constexpr (!std::is_same_v<T, TokT> &&
Expand All @@ -590,10 +646,10 @@ __device__ __forceinline__ void EpCombineIntraNodeKernel_body(EpDispatchCombineA

template <typename T, bool UseP2PRead = true, bool EnableStdMoE = false,
bool UseFp8DirectCast = false, bool UseFp8BlockwiseQuant = false, bool UseWeights = true,
int Vec8Top8BlockElems = 0, int Vec8AccumNum = 8>
int Vec8Top8BlockElems = 0, int Vec8AccumNum = 8, bool UseFp4Combine = false>
__global__ void EpCombineIntraNodeKernel(EpDispatchCombineArgs<T> args) {
EpCombineIntraNodeKernel_body<T, UseP2PRead, EnableStdMoE, UseFp8DirectCast, UseFp8BlockwiseQuant,
UseWeights, Vec8Top8BlockElems, Vec8AccumNum>(args);
UseWeights, Vec8Top8BlockElems, Vec8AccumNum, UseFp4Combine>(args);
}

} // namespace moe
Expand Down
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 @@ -110,6 +110,11 @@ using namespace mori::moe;
kernel##_body<T, B1, B2, B3, B4, B5, B6, B7>(args); \
}

#define WRAP_BOOL8(kernel, suffix, T, B1, B2, B3, B4, B5, B6, B7, B8) \
extern "C" __global__ void kernel##_##suffix(EpDispatchCombineArgs<T> args) { \
kernel##_body<T, B1, B2, B3, B4, B5, B6, B7, B8>(args); \
}

#define WRAP_ALL_TYPES_BOOL3(kernel, bool_suffix, B1, B2, B3) \
WRAP_BOOL3(kernel, bf16##bool_suffix, hip_bfloat16, B1, B2, B3) \
MORI_FP8_FNUZ(WRAP_BOOL3(kernel, fp8_fnuz##bool_suffix, __hip_fp8_e4m3_fnuz, B1, B2, B3)) \
Expand Down
17 changes: 15 additions & 2 deletions src/ops/kernels/ep_intranode.hip
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ WRAP_ALL_TYPES_BOOL3(EpCombineIntraNodeKernel, _p2p, true, false, false)
WRAP_ALL_TYPES_BOOL3(EpCombineIntraNodeKernel, _nop2p, false, false, false)
WRAP_ALL_TYPES_BOOL3(EpCombineIntraNodeKernel, _p2p_stdmoe, true, true, false)
WRAP_BOOL3(EpCombineIntraNodeKernel, bf16_nop2p_fp8cast, hip_bfloat16, false, false, true)
MORI_FP8_ANY(WRAP_BOOL4(EpCombineIntraNodeKernel, bf16_nop2p_fp8bwq, hip_bfloat16, false,
false, false, true))
MORI_FP8_ANY(WRAP_BOOL4(EpCombineIntraNodeKernel, bf16_nop2p_fp8bwq, hip_bfloat16, false, false,
false, true))
MORI_FP8_ANY(WRAP_BOOL6(EpCombineIntraNodeKernel, bf16_nop2p_fp8bwq_noweight_block128_vec8,
hip_bfloat16, false, false, false, true, false, 128))
MORI_FP8_ANY(WRAP_BOOL6(EpCombineIntraNodeKernel, bf16_nop2p_fp8bwq_noweight_block256_vec8,
Expand All @@ -28,6 +28,19 @@ MORI_FP8_ANY(WRAP_BOOL7(EpCombineIntraNodeKernel, bf16_nop2p_fp8bwq_noweight_blo
MORI_FP8_ANY(WRAP_BOOL7(EpCombineIntraNodeKernel, bf16_nop2p_fp8bwq_noweight_block256_vec8_top9,
hip_bfloat16, false, false, false, true, false, 256, 9))

// Blockwise FP4 (E2M1) combine: same as the fp8bwq variants but with UseFp4Combine=true (8th
// template arg), transporting packed FP4 (0.5 byte/elem) over the shared blockwise staging path.
MORI_FP8_ANY(WRAP_BOOL8(EpCombineIntraNodeKernel, bf16_nop2p_fp4bwq, hip_bfloat16, false, false,
false, true, true, 0, 8, true))
MORI_FP8_ANY(WRAP_BOOL8(EpCombineIntraNodeKernel, bf16_nop2p_fp4bwq_noweight_block128_vec8,
hip_bfloat16, false, false, false, true, false, 128, 8, true))
MORI_FP8_ANY(WRAP_BOOL8(EpCombineIntraNodeKernel, bf16_nop2p_fp4bwq_noweight_block256_vec8,
hip_bfloat16, false, false, false, true, false, 256, 8, true))
MORI_FP8_ANY(WRAP_BOOL8(EpCombineIntraNodeKernel, bf16_nop2p_fp4bwq_noweight_block128_vec8_top9,
hip_bfloat16, false, false, false, true, false, 128, 9, true))
MORI_FP8_ANY(WRAP_BOOL8(EpCombineIntraNodeKernel, bf16_nop2p_fp4bwq_noweight_block256_vec8_top9,
hip_bfloat16, false, false, false, true, false, 256, 9, true))

// Convert (used with ENABLE_STANDARD_MOE_ADAPT)
extern "C" __global__ void mori_ConvertDispatchOutputKernel(ConvertDispatchOutputArgs args) {
ConvertDispatchOutputDevice(args);
Expand Down
5 changes: 5 additions & 0 deletions tests/python/ops/dispatch_combine_test_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -623,6 +623,11 @@ def check_combine_result(
elif getattr(self.config, "quant_type", "none") == "fp8_blockwise":
# FP8 E4M3 quantization can exceed 5% after combine accumulation.
atol, rtol = 7e-2, 7e-2
elif getattr(self.config, "quant_type", "none") == "fp4_blockwise":
# 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

@jhchouuu jhchouuu Jul 10, 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.

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

result_match = torch.allclose(
got.float(), expected.float(), atol=atol, rtol=rtol
)
Expand Down
Loading
Loading