From 4b88aceed85dae84b515d8f6f584b4f6d9a7910c Mon Sep 17 00:00:00 2001 From: Karan Verma Date: Thu, 9 Jul 2026 16:59:56 -0500 Subject: [PATCH 1/5] feat(dispatch_combine): add blockwise FP4 (E2M1) combine transport 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 --- .../core/transport/p2p/device_primitives.hpp | 335 ++++++++++++++++++ python/mori/jit/core.py | 8 + python/mori/ops/dispatch_combine.py | 14 + src/ops/dispatch_combine/intranode.hpp | 118 ++++-- src/ops/kernels/ep_common.hip | 5 + src/ops/kernels/ep_intranode.hip | 17 +- .../python/ops/dispatch_combine_test_utils.py | 5 + .../ops/test_dispatch_combine_intranode.py | 28 +- tests/python/ops/test_fp4_combine_wiring.py | 55 +++ 9 files changed, 540 insertions(+), 45 deletions(-) create mode 100644 tests/python/ops/test_fp4_combine_wiring.py diff --git a/include/mori/core/transport/p2p/device_primitives.hpp b/include/mori/core/transport/p2p/device_primitives.hpp index 4ccf463a8..b8d177129 100644 --- a/include/mori/core/transport/p2p/device_primitives.hpp +++ b/include/mori/core/transport/p2p/device_primitives.hpp @@ -792,6 +792,11 @@ static constexpr float kCombineInternalFp8MaxFinite = 448.0f; static constexpr float kCombineInternalFp8MaxFinite = 240.0f; #endif +// Max finite magnitude of OCP FP4 (E2M1): representable positive values are +// {0, 0.5, 1, 1.5, 2, 3, 4, 6}. Used by the fp4_blockwise combine transport to derive +// per-128-element block scales. +static constexpr float kCombineInternalFp4MaxFinite = 6.0f; + __device__ __forceinline__ float WarpReduceMaxF32(float val) { for (int delta = (warpSize >> 1); delta > 0; delta >>= 1) { val = fmaxf(val, __shfl_down(val, delta)); @@ -2138,6 +2143,336 @@ __device__ __forceinline__ void WarpAccumFp8DequantSegment( } #endif +/* ---------------------------------------------------------------------------------------------- */ +/* Experimental blockwise FP4 (E2M1) combine transport */ +/* */ +/* Mirrors the FP8 blockwise combine path but quantizes each element to OCP FP4 (E2M1), packed */ +/* two values per byte (low nibble = even element, high nibble = odd element). This HALVES the */ +/* combine token traffic vs FP8, which is the source of the decode speed-up (combine is transport- + */ +/* bound). The staging slot stride is left unchanged (the packed token occupies the first */ +/* hiddenDim/2 bytes of the FP8-sized token region and the block scales stay at their FP8 offset), + */ +/* so no host-side buffer/stride plumbing has to change. */ +/* */ +/* Unlike FP8 (which only scales blocks whose max-abs exceeds the FP8 range), FP4 is far coarser, */ +/* so every block is scaled to map its max-abs onto FP4_MAX (6.0). Scale-sign sentinel matches the + */ +/* FP8 path: dstScales[0] is negated to mark the token as quantized. blockElems is even for all */ +/* supported configs (hiddenDim and scaleDim are powers-of-two-ish), so scale blocks are byte- */ +/* aligned; a scalar tail handles any odd remainder defensively. */ +/* ---------------------------------------------------------------------------------------------- */ +// Subwarp-vectorized FP4 blockwise quantizer (mirrors WarpQuantizeBf16ToFp8BlockwiseVec). +// The warp is split into warpSize/SubwarpSize subwarps; each subwarp owns one scale block and +// covers it in a single kVecElems-wide iteration (requires blockElems == SubwarpSize*kVecElems). +// This replaces the scalar path's 56 serial full-warp max-reductions with cheap per-subwarp +// reductions across 8 blocks in parallel — the dominant FP4 send-side quant cost. +// Requires: InT == hip_bfloat16, blockElems == SubwarpSize*kVecElems, kVecElems % 8 == 0, +// and block byte-alignment (start even). Uses vector loads (int4 = 8 bf16) and vector stores +// (uint32 = 8 packed fp4) so the send-quant is not bottlenecked on scalar accesses. +template +__device__ __forceinline__ void WarpQuantizeToFp4BlockwiseVec(Fp8T* __restrict__ dstToken, + float* __restrict__ dstScales, + const InT* __restrict__ srcToken, + int hiddenDim, int scaleDim) { + constexpr float fp4Max = kCombineInternalFp4MaxFinite; + const int laneId = threadIdx.x & (warpSize - 1); + const int subLaneId = laneId & (SubwarpSize - 1); + const int subWarpId = laneId / SubwarpSize; + constexpr int kSubwarpsPerWarp = warpSize / SubwarpSize; + auto* dstPacked = reinterpret_cast(dstToken); + const int blockElems = (hiddenDim + scaleDim - 1) / scaleDim; + + for (int sbBase = 0; sbBase < scaleDim; sbBase += kSubwarpsPerWarp) { + const int sb = sbBase + subWarpId; + if (sb >= scaleDim) continue; + const int start = sb * blockElems; + const int base = start + subLaneId * kVecElems; + + // Vector-load kVecElems bf16 into registers (int4 == 16 bytes == 8 bf16). + hip_bfloat16 reg[kVecElems]; +#pragma unroll + for (int j = 0; j < kVecElems; j += 8) { + *reinterpret_cast(®[j]) = *reinterpret_cast(srcToken + base + j); + } + + float localMax = 0.0f; +#pragma unroll + for (int j = 0; j < kVecElems; ++j) + localMax = fmaxf(localMax, fabsf(static_cast(reg[j]))); + float m = localMax; +#pragma unroll + for (int off = SubwarpSize / 2; off > 0; off >>= 1) + m = fmaxf(m, __shfl_down(m, off, SubwarpSize)); + const float maxAbs = __shfl(m, 0, SubwarpSize); + + const float scale = (maxAbs > 0.0f) ? (maxAbs / fp4Max) : 1.0f; + const float invScale = (maxAbs > 0.0f) ? (fp4Max / maxAbs) : 0.0f; + if (subLaneId == 0) dstScales[sb] = scale; + + // Convert to packed fp4 and vector-store (uint32 == 8 packed fp4 == 4 bytes). +#pragma unroll + for (int w = 0; w < kVecElems / 8; ++w) { + uint32_t u = 0; +#pragma unroll + for (int p = 0; p < 4; ++p) { + const int e = w * 8 + p * 2; + const float v0 = static_cast(reg[e]) * invScale; + const float v1 = static_cast(reg[e + 1]) * invScale; + const uint8_t b = mori::float2_to_fp4x2_e2m1(float2{v0, v1}); + u |= static_cast(b) << (p * 8); + } + *reinterpret_cast(dstPacked + (base >> 1) + w * 4) = u; + } + } + // FP4 always scales -> always mark the token as quantized. + if (laneId == 0) dstScales[0] = -dstScales[0]; +} + +template +__device__ __forceinline__ void WarpQuantizeToFp4Blockwise(Fp8T* __restrict__ dstToken, + float* __restrict__ dstScales, + const InT* __restrict__ srcToken, + int hiddenDim, int scaleDim) { + const int laneId = threadIdx.x & (warpSize - 1); + const int blockElems = (hiddenDim + scaleDim - 1) / scaleDim; + constexpr float fp4Max = kCombineInternalFp4MaxFinite; + auto* dstPacked = reinterpret_cast(dstToken); + + // Fast path: 64-lane warp, blockElems == 128 == 8 subwarps * 16 elems (the DeepSeek config). + if (warpSize == 64 && blockElems == 128 && (hiddenDim % 128) == 0 && + std::is_same_v) { + WarpQuantizeToFp4BlockwiseVec<8, 16, Fp8T, InT>(dstToken, dstScales, srcToken, hiddenDim, + scaleDim); + return; + } + + for (int sb = 0; sb < scaleDim; ++sb) { + const int start = sb * blockElems; + const int end = std::min(start + blockElems, hiddenDim); + + float localMaxAbs = 0.0f; + for (int idx = start + laneId; idx < end; idx += warpSize) { + localMaxAbs = fmaxf(localMaxAbs, fabsf(static_cast(srcToken[idx]))); + } + float maxAbs = WarpReduceMaxF32(localMaxAbs); + maxAbs = __shfl(maxAbs, 0); + + const float scale = (maxAbs > 0.0f) ? (maxAbs / fp4Max) : 1.0f; + const float invScale = (maxAbs > 0.0f) ? (fp4Max / maxAbs) : 0.0f; + if (laneId == 0) dstScales[sb] = scale; + + // Pack element pairs (2 FP4/byte). Blocks are byte-aligned (start even, blockElems even). + const int pStart = start >> 1; + const int pEnd = end >> 1; + for (int p = pStart + laneId; p < pEnd; p += warpSize) { + const float v0 = static_cast(srcToken[2 * p]) * invScale; + const float v1 = static_cast(srcToken[2 * p + 1]) * invScale; + dstPacked[p] = mori::float2_to_fp4x2_e2m1(float2{v0, v1}); + } + // Defensive scalar tail for an odd block end (does not occur for supported configs). + if ((end & 1) && laneId == 0) { + const int idx = end - 1; + const float v = static_cast(srcToken[idx]) * invScale; + auto* dstBytes = reinterpret_cast(dstToken); + dstBytes[idx >> 1] = mori::float_to_fp4_e2m1(v); // low nibble + } + } + // Always mark as quantized so the dequant path applies the block scales. + if (laneId == 0) dstScales[0] = -dstScales[0]; +} + +template +__device__ __forceinline__ void WarpAccumFp4DequantFull(OutT* __restrict__ dstToken, + const Fp8T* const* __restrict__ srcs, + const float* const* __restrict__ srcScales, + int accumNum, int hiddenDim, int scaleDim) { + const int laneId = threadIdx.x & (warpSize - 1); + const int blockElems = (hiddenDim + scaleDim - 1) / scaleDim; + for (int sb = 0; sb < scaleDim; ++sb) { + const int start = sb * blockElems; + const int end = std::min(start + blockElems, hiddenDim); + const int pStart = start >> 1; + const int pEnd = end >> 1; + for (int p = pStart + laneId; p < pEnd; p += warpSize) { + float acc0 = 0.0f; + float acc1 = 0.0f; + for (int i = 0; i < accumNum; ++i) { + if (srcs[i] == nullptr) continue; + float s = 1.0f; + if (srcScales != nullptr && srcScales[i] != nullptr) { + s = srcScales[i][sb]; + if (sb == 0 && s < 0.0f) s = -s; + } + mori::mori_fp4x2_e2m1 e; + e.x = reinterpret_cast(srcs[i])[p]; + const float2 v = static_cast(e); + acc0 = fmaf(v.x, s, acc0); + acc1 = fmaf(v.y, s, acc1); + } + dstToken[2 * p] = OutT(acc0); + dstToken[2 * p + 1] = OutT(acc1); + } + if ((end & 1) && laneId == 0) { + const int idx = end - 1; + float acc = 0.0f; + for (int i = 0; i < accumNum; ++i) { + if (srcs[i] == nullptr) continue; + float s = 1.0f; + if (srcScales != nullptr && srcScales[i] != nullptr) { + s = srcScales[i][sb]; + if (sb == 0 && s < 0.0f) s = -s; + } + mori::mori_fp4_e2m1 e; + e.x = reinterpret_cast(srcs[i])[idx >> 1]; + acc += static_cast(e) * s; + } + dstToken[idx] = OutT(acc); + } + } +} + +template +__device__ __forceinline__ void WarpAccumFp4DequantSegment( + OutT* __restrict__ dstToken, const Fp8T* const* __restrict__ srcs, + const float* const* __restrict__ srcScales, int accumNum, int hiddenDimOffset, + int hiddenDimSize, int hiddenDim, int scaleDim) { + const int laneId = threadIdx.x & (warpSize - 1); + const int blockElems = (hiddenDim + scaleDim - 1) / scaleDim; + // Per-element (robust to any parity) safety fallback for misaligned segments. Callers pass + // srcs[i] advanced by hiddenDimOffset *elements* (== bytes under FP8); recover the token base + // (tb) and index the packed byte at (hiddenDimOffset+localIdx)/2, selecting the nibble by parity. + for (int localIdx = laneId; localIdx < hiddenDimSize; localIdx += warpSize) { + const int globalIdx = hiddenDimOffset + localIdx; + const int sb = globalIdx / blockElems; + float acc = 0.0f; + for (int i = 0; i < accumNum; ++i) { + if (srcs[i] == nullptr) continue; + float s = 1.0f; + if (srcScales != nullptr && srcScales[i] != nullptr) { + s = srcScales[i][sb]; + if (sb == 0 && s < 0.0f) s = -s; + } + const uint8_t* tb = reinterpret_cast(srcs[i]) - hiddenDimOffset; + mori::mori_fp4x2_e2m1 e; + e.x = tb[globalIdx >> 1]; + const float2 v = static_cast(e); + acc = fmaf(((globalIdx & 1) == 0) ? v.x : v.y, s, acc); + } + dstToken[localIdx] = OutT(acc); + } +} + +/* ---------------------------------------------------------------------------------------------- */ +/* Vectorized packed-FP4 dequant-accumulate (analog of the FP8 vec8 blockwise path) */ +/* */ +/* Each lane processes 8 elements = 4 packed bytes per step, unpacking with the native gfx950 */ +/* mxfp4->f32 conversion (via mori_fp4x4_e2m1) and FMA-accumulating into fp32. This keeps the */ +/* vectorization of the FP8 vec8 kernel while transporting half the bytes. Callers pass srcs[i] */ +/* advanced by hiddenDimOffset *elements* (== bytes in the FP8 layout); we recover the token base */ +/* (tb = srcs[i] - hiddenDimOffset) and index in packed (half-byte) units. Blocks/offsets are */ +/* byte-aligned (blockElems and hiddenDimOffset even) for all supported configs. */ +/* ---------------------------------------------------------------------------------------------- */ +template +__device__ __forceinline__ void WarpAccumFp4DequantVecBlockwiseScaleWave( + OutT* __restrict__ dstToken, const Fp8T* const* __restrict__ srcs, + const float* const* __restrict__ srcScales, int hiddenDimOffset, int start, int end, + int blockElems) { + const int laneId = threadIdx.x & (warpSize - 1); + const bool blockPow2 = (blockElems & (blockElems - 1)) == 0; + const int blockShift = blockPow2 ? (__ffs(blockElems) - 1) : 0; + constexpr int kElems = 8; // 4 packed bytes -> 8 FP4 elements per lane per step + constexpr int kWords = kElems / 4; // number of fp4x4 (uint16) groups + + const int vecEnd = start + ((end - start) / kElems) * kElems; + for (int idx = start + laneId * kElems; idx < vecEnd; idx += warpSize * kElems) { + const int globalIdx = hiddenDimOffset + idx; + const int sb = blockPow2 ? (globalIdx >> blockShift) : (globalIdx / blockElems); + float acc[kElems]; +#pragma unroll + for (int e = 0; e < kElems; ++e) acc[e] = 0.0f; + + // Hoist all source loads up front for memory-level parallelism (issue the AccumNum packed + // reads back-to-back so their latencies overlap), then convert + FMA-accumulate. + uint32_t bitsArr[AccumNum]; + float sArr[AccumNum]; +#pragma unroll AccumNum + for (int i = 0; i < AccumNum; ++i) { + const Fp8T* src = srcs[i]; + if (src == nullptr) { + sArr[i] = 0.0f; // sentinel: contributes nothing + bitsArr[i] = 0u; + continue; + } + float s = 1.0f; + if (srcScales != nullptr && srcScales[i] != nullptr) { + s = srcScales[i][sb]; + if (sb == 0 && s < 0.0f) s = -s; + } + sArr[i] = s; + const uint8_t* tb = reinterpret_cast(src) - hiddenDimOffset; + bitsArr[i] = load<4>(tb + (globalIdx >> 1)); + } +#pragma unroll AccumNum + for (int i = 0; i < AccumNum; ++i) { + const float s = sArr[i]; +#pragma unroll + for (int w = 0; w < kWords; ++w) { + mori::mori_fp4x4_e2m1 q; + q.x = static_cast((bitsArr[i] >> (w * 16)) & 0xFFFFull); + const float4 f = static_cast(q); + acc[w * 4 + 0] = fmaf(f.x, s, acc[w * 4 + 0]); + acc[w * 4 + 1] = fmaf(f.y, s, acc[w * 4 + 1]); + acc[w * 4 + 2] = fmaf(f.z, s, acc[w * 4 + 2]); + acc[w * 4 + 3] = fmaf(f.w, s, acc[w * 4 + 3]); + } + } +#pragma unroll + for (int e = 0; e < kElems; e += 2) { + StoreOutPair(dstToken, idx + e, float2{acc[e], acc[e + 1]}); + } + } + + // Scalar tail for the ragged end (unpacks a single element from its packed byte). + for (int j = vecEnd + laneId; j < end; j += warpSize) { + const int globalIdx = hiddenDimOffset + j; + const int sb = blockPow2 ? (globalIdx >> blockShift) : (globalIdx / blockElems); + float acc = 0.0f; +#pragma unroll AccumNum + for (int i = 0; i < AccumNum; ++i) { + const Fp8T* src = srcs[i]; + if (src == nullptr) continue; + float s = 1.0f; + if (srcScales != nullptr && srcScales[i] != nullptr) { + s = srcScales[i][sb]; + if (sb == 0 && s < 0.0f) s = -s; + } + const uint8_t* tb = reinterpret_cast(src) - hiddenDimOffset; + mori::mori_fp4x2_e2m1 e2; + e2.x = tb[globalIdx >> 1]; + const float2 v = static_cast(e2); + acc = fmaf(((globalIdx & 1) == 0) ? v.x : v.y, s, acc); + } + dstToken[j] = OutT(acc); + } +} + +template +__device__ __forceinline__ void WarpAccumFp4DequantFullBlockVec8Top8( + OutT* __restrict__ dstToken, const Fp8T* const* __restrict__ srcs, + const float* const* __restrict__ srcScales, int hiddenDim) { + WarpAccumFp4DequantVecBlockwiseScaleWave( + dstToken, srcs, srcScales, /*hiddenDimOffset=*/0, /*start=*/0, /*end=*/hiddenDim, BlockElems); +} + +template +__device__ __forceinline__ void WarpAccumFp4DequantSegmentBlockVec8Top8( + OutT* __restrict__ dstToken, const Fp8T* const* __restrict__ srcs, + const float* const* __restrict__ srcScales, int hiddenDimOffset, int hiddenDimSize) { + WarpAccumFp4DequantVecBlockwiseScaleWave( + dstToken, srcs, srcScales, hiddenDimOffset, /*start=*/0, /*end=*/hiddenDimSize, BlockElems); +} + template __forceinline__ __device__ void WarpCastBf16ToCombineInternalFp8( CombineInternalFp8* __restrict__ dst, const T* __restrict__ src, int hiddenDim, int laneId) { diff --git a/python/mori/jit/core.py b/python/mori/jit/core.py index aa0028524..ffb22706f 100644 --- a/python/mori/jit/core.py +++ b/python/mori/jit/core.py @@ -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 [] @@ -415,6 +422,7 @@ def _hipcc_genco( *_nic_defines(), *_ccqe_defines(), *_profiler_defines(), + *_ocp_fp_defines(cfg.arch), ] for d in include_dirs: diff --git a/python/mori/ops/dispatch_combine.py b/python/mori/ops/dispatch_combine.py index 25e32b8fb..938d1416c 100644 --- a/python/mori/ops/dispatch_combine.py +++ b/python/mori/ops/dispatch_combine.py @@ -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, } @@ -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" + ) 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) @@ -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 ) diff --git a/src/ops/dispatch_combine/intranode.hpp b/src/ops/dispatch_combine/intranode.hpp index d22c230b2..361a6dafb 100644 --- a/src/ops/dispatch_combine/intranode.hpp +++ b/src/ops/dispatch_combine/intranode.hpp @@ -264,10 +264,14 @@ __global__ void EpDispatchIntraNodeKernel(EpDispatchCombineArgs args) { /* ---------------------------------------------------------------------------------------------- */ template + int Vec8Top8BlockElems = 0, int Vec8AccumNum = 8, bool UseFp4Combine = false> __device__ __forceinline__ void EpCombineIntraNodeKernel_body(EpDispatchCombineArgs args) { using TokT = std::conditional_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, @@ -358,10 +362,17 @@ __device__ __forceinline__ void EpCombineIntraNodeKernel_body(EpDispatchCombineA uint8_t* destStagingPtr = args.intraNodeTokBufs.combineInp->template GetAs(destPe) + SendBufSlotOffset(config, myPe, destLocalTokId) * combXferBytes; if constexpr (UseFp8BlockwiseQuant) { - core::WarpQuantizeToFp8Blockwise( - reinterpret_cast(destStagingPtr), - reinterpret_cast(destStagingPtr + hiddenBytes), - args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim); + if constexpr (UseFp4Combine) { + core::WarpQuantizeToFp4Blockwise( + reinterpret_cast(destStagingPtr), + reinterpret_cast(destStagingPtr + hiddenBytes), + args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim); + } else { + core::WarpQuantizeToFp8Blockwise( + reinterpret_cast(destStagingPtr), + reinterpret_cast(destStagingPtr + hiddenBytes), + args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim); + } } else if constexpr (!std::is_same_v && std::is_same_v) { core::WarpCastBf16ToCombineInternalFp8(reinterpret_cast(destStagingPtr), @@ -396,10 +407,17 @@ __device__ __forceinline__ void EpCombineIntraNodeKernel_body(EpDispatchCombineA uint8_t* destStagingPtr = args.intraNodeTokBufs.combineInp->template GetAs(destPe) + SendBufSlotOffset(config, myPe, destLocalTokId) * combXferBytes; if constexpr (UseFp8BlockwiseQuant) { - core::WarpQuantizeToFp8Blockwise( - reinterpret_cast(destStagingPtr), - reinterpret_cast(destStagingPtr + hiddenBytes), - args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim); + if constexpr (UseFp4Combine) { + core::WarpQuantizeToFp4Blockwise( + reinterpret_cast(destStagingPtr), + reinterpret_cast(destStagingPtr + hiddenBytes), + args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim); + } else { + core::WarpQuantizeToFp8Blockwise( + reinterpret_cast(destStagingPtr), + reinterpret_cast(destStagingPtr + hiddenBytes), + args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim); + } } else if constexpr (!std::is_same_v && std::is_same_v) { core::WarpCastBf16ToCombineInternalFp8(reinterpret_cast(destStagingPtr), @@ -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( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), hiddenDim); + if constexpr (UseFp4Combine) { + core::WarpAccumFp4DequantFullBlockVec8Top8( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), hiddenDim); + } else { + core::WarpAccumFp8DequantFullBlockVec8Top8( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), hiddenDim); + } } else if ((hiddenDimOffset & 0x7) == 0 && (hiddenDimSize & 0x7) == 0) { - core::WarpAccumFp8DequantSegmentBlockVec8Top8( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), hiddenDimOffset, hiddenDimSize); + if constexpr (UseFp4Combine) { + core::WarpAccumFp4DequantSegmentBlockVec8Top8( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), hiddenDimOffset, + hiddenDimSize); + } else { + core::WarpAccumFp8DequantSegmentBlockVec8Top8( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), hiddenDimOffset, + hiddenDimSize); + } } else { // Misaligned segment: vec8 helper would fault on the load. Tiny scalar fallback. - core::WarpAccumFp8DequantSegmentScalarTop8( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), hiddenDimOffset, hiddenDimSize); + if constexpr (UseFp4Combine) { + core::WarpAccumFp4DequantSegment( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), Vec8AccumNum, hiddenDimOffset, + hiddenDimSize, hiddenDim, args.fp8BlockwiseCombineScaleDim); + } else { + core::WarpAccumFp8DequantSegmentScalarTop8( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), hiddenDimOffset, + hiddenDimSize); + } } } else { if (mwIter.warpsPerItem == 1) { - core::WarpAccumFp8DequantFull( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), validAccumCount, hiddenDim, - args.fp8BlockwiseCombineScaleDim); + if constexpr (UseFp4Combine) { + core::WarpAccumFp4DequantFull( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), validAccumCount, hiddenDim, + args.fp8BlockwiseCombineScaleDim); + } else { + core::WarpAccumFp8DequantFull( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), validAccumCount, hiddenDim, + args.fp8BlockwiseCombineScaleDim); + } } else { - core::WarpAccumFp8DequantSegment( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), validAccumCount, hiddenDimOffset, - hiddenDimSize, hiddenDim, args.fp8BlockwiseCombineScaleDim); + if constexpr (UseFp4Combine) { + core::WarpAccumFp4DequantSegment( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), validAccumCount, + hiddenDimOffset, hiddenDimSize, hiddenDim, args.fp8BlockwiseCombineScaleDim); + } else { + core::WarpAccumFp8DequantSegment( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), validAccumCount, + hiddenDimOffset, hiddenDimSize, hiddenDim, args.fp8BlockwiseCombineScaleDim); + } } } } else if constexpr (!std::is_same_v && @@ -590,10 +646,10 @@ __device__ __forceinline__ void EpCombineIntraNodeKernel_body(EpDispatchCombineA template + int Vec8Top8BlockElems = 0, int Vec8AccumNum = 8, bool UseFp4Combine = false> __global__ void EpCombineIntraNodeKernel(EpDispatchCombineArgs args) { EpCombineIntraNodeKernel_body(args); + UseWeights, Vec8Top8BlockElems, Vec8AccumNum, UseFp4Combine>(args); } } // namespace moe diff --git a/src/ops/kernels/ep_common.hip b/src/ops/kernels/ep_common.hip index cf5bb48ce..0125cad8b 100644 --- a/src/ops/kernels/ep_common.hip +++ b/src/ops/kernels/ep_common.hip @@ -110,6 +110,11 @@ using namespace mori::moe; kernel##_body(args); \ } +#define WRAP_BOOL8(kernel, suffix, T, B1, B2, B3, B4, B5, B6, B7, B8) \ + extern "C" __global__ void kernel##_##suffix(EpDispatchCombineArgs args) { \ + kernel##_body(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)) \ diff --git a/src/ops/kernels/ep_intranode.hip b/src/ops/kernels/ep_intranode.hip index d4504c31b..71effcac4 100644 --- a/src/ops/kernels/ep_intranode.hip +++ b/src/ops/kernels/ep_intranode.hip @@ -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, @@ -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); diff --git a/tests/python/ops/dispatch_combine_test_utils.py b/tests/python/ops/dispatch_combine_test_utils.py index 41b911b4b..f424d14d1 100644 --- a/tests/python/ops/dispatch_combine_test_utils.py +++ b/tests/python/ops/dispatch_combine_test_utils.py @@ -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 result_match = torch.allclose( got.float(), expected.float(), atol=atol, rtol=rtol ) diff --git a/tests/python/ops/test_dispatch_combine_intranode.py b/tests/python/ops/test_dispatch_combine_intranode.py index 48352ec88..d4a47c37d 100644 --- a/tests/python/ops/test_dispatch_combine_intranode.py +++ b/tests/python/ops/test_dispatch_combine_intranode.py @@ -166,7 +166,9 @@ def _test_dispatch_combine_ll( @pytest.mark.parametrize("num_experts_per_rank", (32,)) @pytest.mark.parametrize("num_experts_per_token", (8,)) @pytest.mark.parametrize("use_external_inp_buf", (True, False)) -@pytest.mark.parametrize("quant_type", ("none", "fp8_direct_cast", "fp8_blockwise")) +@pytest.mark.parametrize( + "quant_type", ("none", "fp8_direct_cast", "fp8_blockwise", "fp4_blockwise") +) def test_dispatch_combine( torch_dist_process_manager, world_size, @@ -185,13 +187,14 @@ def test_dispatch_combine( pytest.skip("fp8_direct_cast is not supported in zero-copy mode") 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": + if quant_type in ("fp8_blockwise", "fp4_blockwise"): if data_type is not torch.bfloat16: - pytest.skip("fp8_blockwise only supports bfloat16 input") + pytest.skip(f"{quant_type} only supports bfloat16 input") if not use_external_inp_buf: - pytest.skip("fp8_blockwise requires use_external_inp_buf=True") - # fp8_blockwise combine ignores scale_dim/scale_type_size (driven by - # MORI_FP8_COMBINE_SCALE_DIM internally). + pytest.skip(f"{quant_type} requires use_external_inp_buf=True") + # blockwise combine ignores scale_dim/scale_type_size (driven by + # MORI_FP8_COMBINE_SCALE_DIM internally). fp4_blockwise reuses the same path + # but transports packed FP4 (E2M1) instead of FP8. for i in range(world_size): torch_dist_process_manager.task_queue.put( @@ -227,7 +230,7 @@ def test_dispatch_combine( @pytest.mark.parametrize("num_experts_per_rank", (32,)) @pytest.mark.parametrize("num_experts_per_token", (8, 9)) @pytest.mark.parametrize("use_external_inp_buf", (True,)) -@pytest.mark.parametrize("quant_type", ("fp8_blockwise",)) +@pytest.mark.parametrize("quant_type", ("fp8_blockwise", "fp4_blockwise")) def test_dispatch_combine_weightless_vec8( torch_dist_process_manager, world_size, @@ -301,13 +304,14 @@ def test_dispatch_combine_ll( pytest.skip("fp8_direct_cast is not supported in zero-copy mode") 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": + if quant_type in ("fp8_blockwise", "fp4_blockwise"): if data_type is not torch.bfloat16: - pytest.skip("fp8_blockwise only supports bfloat16 input") + pytest.skip(f"{quant_type} only supports bfloat16 input") if not use_external_inp_buf: - pytest.skip("fp8_blockwise requires use_external_inp_buf=True") - # fp8_blockwise combine ignores scale_dim/scale_type_size (driven by - # MORI_FP8_COMBINE_SCALE_DIM internally). + pytest.skip(f"{quant_type} requires use_external_inp_buf=True") + # blockwise combine ignores scale_dim/scale_type_size (driven by + # MORI_FP8_COMBINE_SCALE_DIM internally). fp4_blockwise reuses the same path + # but transports packed FP4 (E2M1) instead of FP8. for i in range(world_size): torch_dist_process_manager.task_queue.put( diff --git a/tests/python/ops/test_fp4_combine_wiring.py b/tests/python/ops/test_fp4_combine_wiring.py new file mode 100644 index 000000000..6b5804f65 --- /dev/null +++ b/tests/python/ops/test_fp4_combine_wiring.py @@ -0,0 +1,55 @@ +# Copyright © Advanced Micro Devices, Inc. All rights reserved. +# +# MIT License +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Lightweight (no-GPU) wiring guards for the blockwise-FP4 combine feature. + +These protect the Python-side plumbing that selects the packed-FP4 combine kernels so a +refactor cannot silently drop the feature. The actual kernel numerics are covered by the +GPU test ``test_dispatch_combine_intranode.py`` (quant_type="fp4_blockwise"). +""" + +from mori.ops.dispatch_combine import _QUANT_TYPE_MAP, EpDispatchCombineQuantType + + +def test_fp4_blockwise_registered_in_quant_type_map(): + assert "fp4_blockwise" in _QUANT_TYPE_MAP + # fp4_blockwise reuses the Fp8BlockwiseQuant config path (shared 1-byte-slot staging + + # float scales); the packed-FP4 kernels are chosen at launch, not via a distinct enum. + assert ( + _QUANT_TYPE_MAP["fp4_blockwise"] == EpDispatchCombineQuantType.Fp8BlockwiseQuant + ) + + +def test_fp8bwq_to_fp4bwq_kernel_name_mapping(): + # Mirrors the launch-time selection: kernel_name.replace("_fp8bwq", "_fp4bwq"). + # Every fp8bwq combine variant must have a matching fp4bwq name (registered in + # ep_intranode.hip) so the mapping never yields an unregistered symbol. + fp8_variants = [ + "EpCombineIntraNodeKernel_bf16_nop2p_fp8bwq", + "EpCombineIntraNodeKernel_bf16_nop2p_fp8bwq_noweight_block128_vec8", + "EpCombineIntraNodeKernel_bf16_nop2p_fp8bwq_noweight_block256_vec8", + "EpCombineIntraNodeKernel_bf16_nop2p_fp8bwq_noweight_block128_vec8_top9", + "EpCombineIntraNodeKernel_bf16_nop2p_fp8bwq_noweight_block256_vec8_top9", + ] + for name in fp8_variants: + mapped = name.replace("_fp8bwq", "_fp4bwq") + assert "_fp8bwq" not in mapped + assert mapped.count("_fp4bwq") == 1 From 05890702dcd6ee6e2110d2881c443384c3050d5d Mon Sep 17 00:00:00 2001 From: Karan Verma Date: Fri, 10 Jul 2026 09:37:39 -0500 Subject: [PATCH 2/5] address review: arch-guard fp4, exact fp4 test tolerance, dedup call 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 --- .../core/transport/p2p/device_primitives.hpp | 92 +++++++++++++++ python/mori/ops/dispatch_combine.py | 13 +++ src/ops/dispatch_combine/intranode.hpp | 109 +++++------------- .../python/ops/dispatch_combine_test_utils.py | 60 +++++++--- .../ops/test_dispatch_combine_intranode.py | 20 ++++ 5 files changed, 196 insertions(+), 98 deletions(-) diff --git a/include/mori/core/transport/p2p/device_primitives.hpp b/include/mori/core/transport/p2p/device_primitives.hpp index b8d177129..e684fe8f7 100644 --- a/include/mori/core/transport/p2p/device_primitives.hpp +++ b/include/mori/core/transport/p2p/device_primitives.hpp @@ -2246,6 +2246,14 @@ __device__ __forceinline__ void WarpQuantizeToFp4Blockwise(Fp8T* __restrict__ ds scaleDim); return; } + // Fast path: 64-lane warp, blockElems == 256 == 8 subwarps * 32 elems (mirrors the fp8bwq + // block256 vec8 variant so block256 configs are not left on the scalar path). + if (warpSize == 64 && blockElems == 256 && (hiddenDim % 256) == 0 && + std::is_same_v) { + WarpQuantizeToFp4BlockwiseVec<8, 32, Fp8T, InT>(dstToken, dstScales, srcToken, hiddenDim, + scaleDim); + return; + } for (int sb = 0; sb < scaleDim; ++sb) { const int start = sb * blockElems; @@ -2473,6 +2481,90 @@ __device__ __forceinline__ void WarpAccumFp4DequantSegmentBlockVec8Top8( dstToken, srcs, srcScales, hiddenDimOffset, /*start=*/0, /*end=*/hiddenDimSize, BlockElems); } +/* ---------------------------------------------------------------------------------------------- */ +/* Combine quant/dequant dispatch on UseFp4. FP8 and FP4 blockwise combine share the staging/scale + */ +/* layout; only the element codec differs. These thin wrappers select the codec from a single */ +/* template bool so the intra-node combine body calls one function per site instead of duplicating + */ +/* an fp4/fp8 branch (and identical argument list) at every call site. */ +/* ---------------------------------------------------------------------------------------------- */ +template +__device__ __forceinline__ void WarpQuantizeToCombineBlockwise(Fp8T* __restrict__ dstToken, + float* __restrict__ dstScales, + const InT* __restrict__ srcToken, + int hiddenDim, int scaleDim) { + if constexpr (UseFp4) + WarpQuantizeToFp4Blockwise(dstToken, dstScales, srcToken, hiddenDim, scaleDim); + else + WarpQuantizeToFp8Blockwise(dstToken, dstScales, srcToken, hiddenDim, scaleDim); +} + +template +__device__ __forceinline__ void WarpAccumCombineDequantFullBlockVec8Top8( + OutT* __restrict__ dstToken, const Fp8T* const* __restrict__ srcs, + const float* const* __restrict__ srcScales, int hiddenDim) { + if constexpr (UseFp4) + WarpAccumFp4DequantFullBlockVec8Top8(dstToken, srcs, + srcScales, hiddenDim); + else + WarpAccumFp8DequantFullBlockVec8Top8(dstToken, srcs, + srcScales, hiddenDim); +} + +template +__device__ __forceinline__ void WarpAccumCombineDequantSegmentBlockVec8Top8( + OutT* __restrict__ dstToken, const Fp8T* const* __restrict__ srcs, + const float* const* __restrict__ srcScales, int hiddenDimOffset, int hiddenDimSize) { + if constexpr (UseFp4) + WarpAccumFp4DequantSegmentBlockVec8Top8( + dstToken, srcs, srcScales, hiddenDimOffset, hiddenDimSize); + else + WarpAccumFp8DequantSegmentBlockVec8Top8( + dstToken, srcs, srcScales, hiddenDimOffset, hiddenDimSize); +} + +// Misaligned-segment scalar fallback for the vec8 top-k path. The codecs use different helpers +// here (fp4 reuses the generic segment accumulator with a fixed top-k count; fp8 has a dedicated +// scalar-top8 helper), so the divergence is encapsulated in one place. +template +__device__ __forceinline__ void WarpAccumCombineDequantSegmentScalarTop8( + OutT* __restrict__ dstToken, const Fp8T* const* __restrict__ srcs, + const float* const* __restrict__ srcScales, int hiddenDimOffset, int hiddenDimSize, + int hiddenDim, int scaleDim) { + if constexpr (UseFp4) + WarpAccumFp4DequantSegment(dstToken, srcs, srcScales, AccumNum, hiddenDimOffset, + hiddenDimSize, hiddenDim, scaleDim); + else + WarpAccumFp8DequantSegmentScalarTop8( + dstToken, srcs, srcScales, hiddenDimOffset, hiddenDimSize); +} + +template +__device__ __forceinline__ void WarpAccumCombineDequantFull( + OutT* __restrict__ dstToken, const Fp8T* const* __restrict__ srcs, + const float* const* __restrict__ srcScales, int validAccumCount, int hiddenDim, int scaleDim) { + if constexpr (UseFp4) + WarpAccumFp4DequantFull(dstToken, srcs, srcScales, validAccumCount, hiddenDim, + scaleDim); + else + WarpAccumFp8DequantFull(dstToken, srcs, srcScales, validAccumCount, hiddenDim, + scaleDim); +} + +template +__device__ __forceinline__ void WarpAccumCombineDequantSegment( + OutT* __restrict__ dstToken, const Fp8T* const* __restrict__ srcs, + const float* const* __restrict__ srcScales, int validAccumCount, int hiddenDimOffset, + int hiddenDimSize, int hiddenDim, int scaleDim) { + if constexpr (UseFp4) + WarpAccumFp4DequantSegment(dstToken, srcs, srcScales, validAccumCount, + hiddenDimOffset, hiddenDimSize, hiddenDim, scaleDim); + else + WarpAccumFp8DequantSegment(dstToken, srcs, srcScales, validAccumCount, + hiddenDimOffset, hiddenDimSize, hiddenDim, scaleDim); +} + template __forceinline__ __device__ void WarpCastBf16ToCombineInternalFp8( CombineInternalFp8* __restrict__ dst, const T* __restrict__ src, int hiddenDim, int laneId) { diff --git a/python/mori/ops/dispatch_combine.py b/python/mori/ops/dispatch_combine.py index 938d1416c..eb5b61be5 100644 --- a/python/mori/ops/dispatch_combine.py +++ b/python/mori/ops/dispatch_combine.py @@ -272,6 +272,19 @@ def __init__(self, config): isinstance(config.quant_type, str) and config.quant_type.strip().lower() == "fp4_blockwise" ) + if self._combine_is_fp4: + # The packed-FP4 combine relies on the gfx950 OCP FP4 conversion instructions + # (cvt_scalef32_pk_f32_fp4). On other archs there is no hardware path, so fail fast + # instead of silently selecting an fp4 kernel that would fall back to slow software. + from mori.jit.config import detect_gpu_arch + + _arch = str(detect_gpu_arch()) + if "gfx950" not in _arch: + raise ValueError( + f"quant_type='fp4_blockwise' combine requires a gfx950 GPU (OCP FP4 " + f"conversion instructions); detected arch '{_arch}'. Use 'fp8_blockwise' " + f"instead on this device." + ) 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) diff --git a/src/ops/dispatch_combine/intranode.hpp b/src/ops/dispatch_combine/intranode.hpp index 361a6dafb..66a9de88a 100644 --- a/src/ops/dispatch_combine/intranode.hpp +++ b/src/ops/dispatch_combine/intranode.hpp @@ -362,17 +362,10 @@ __device__ __forceinline__ void EpCombineIntraNodeKernel_body(EpDispatchCombineA uint8_t* destStagingPtr = args.intraNodeTokBufs.combineInp->template GetAs(destPe) + SendBufSlotOffset(config, myPe, destLocalTokId) * combXferBytes; if constexpr (UseFp8BlockwiseQuant) { - if constexpr (UseFp4Combine) { - core::WarpQuantizeToFp4Blockwise( - reinterpret_cast(destStagingPtr), - reinterpret_cast(destStagingPtr + hiddenBytes), - args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim); - } else { - core::WarpQuantizeToFp8Blockwise( - reinterpret_cast(destStagingPtr), - reinterpret_cast(destStagingPtr + hiddenBytes), - args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim); - } + core::WarpQuantizeToCombineBlockwise( + reinterpret_cast(destStagingPtr), + reinterpret_cast(destStagingPtr + hiddenBytes), + args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim); } else if constexpr (!std::is_same_v && std::is_same_v) { core::WarpCastBf16ToCombineInternalFp8(reinterpret_cast(destStagingPtr), @@ -407,17 +400,10 @@ __device__ __forceinline__ void EpCombineIntraNodeKernel_body(EpDispatchCombineA uint8_t* destStagingPtr = args.intraNodeTokBufs.combineInp->template GetAs(destPe) + SendBufSlotOffset(config, myPe, destLocalTokId) * combXferBytes; if constexpr (UseFp8BlockwiseQuant) { - if constexpr (UseFp4Combine) { - core::WarpQuantizeToFp4Blockwise( - reinterpret_cast(destStagingPtr), - reinterpret_cast(destStagingPtr + hiddenBytes), - args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim); - } else { - core::WarpQuantizeToFp8Blockwise( - reinterpret_cast(destStagingPtr), - reinterpret_cast(destStagingPtr + hiddenBytes), - args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim); - } + core::WarpQuantizeToCombineBlockwise( + reinterpret_cast(destStagingPtr), + reinterpret_cast(destStagingPtr + hiddenBytes), + args.inpTokenBuf + tokenIdx * hiddenDim, hiddenDim, args.fp8BlockwiseCombineScaleDim); } else if constexpr (!std::is_same_v && std::is_same_v) { core::WarpCastBf16ToCombineInternalFp8(reinterpret_cast(destStagingPtr), @@ -555,71 +541,34 @@ __device__ __forceinline__ void EpCombineIntraNodeKernel_body(EpDispatchCombineA MORI_TRACE_NEXT(seq, Slot::CombineDequantAccum); if constexpr (Vec8Top8BlockElems != 0) { if (mwIter.warpsPerItem == 1) { - if constexpr (UseFp4Combine) { - core::WarpAccumFp4DequantFullBlockVec8Top8( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), hiddenDim); - } else { - core::WarpAccumFp8DequantFullBlockVec8Top8( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), hiddenDim); - } + core::WarpAccumCombineDequantFullBlockVec8Top8( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), hiddenDim); } else if ((hiddenDimOffset & 0x7) == 0 && (hiddenDimSize & 0x7) == 0) { - if constexpr (UseFp4Combine) { - core::WarpAccumFp4DequantSegmentBlockVec8Top8( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), hiddenDimOffset, - hiddenDimSize); - } else { - core::WarpAccumFp8DequantSegmentBlockVec8Top8( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), hiddenDimOffset, - hiddenDimSize); - } + core::WarpAccumCombineDequantSegmentBlockVec8Top8< + UseFp4Combine, T, core::CombineInternalFp8, Vec8Top8BlockElems, Vec8AccumNum>( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), hiddenDimOffset, hiddenDimSize); } else { // Misaligned segment: vec8 helper would fault on the load. Tiny scalar fallback. - if constexpr (UseFp4Combine) { - core::WarpAccumFp4DequantSegment( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), Vec8AccumNum, hiddenDimOffset, - hiddenDimSize, hiddenDim, args.fp8BlockwiseCombineScaleDim); - } else { - core::WarpAccumFp8DequantSegmentScalarTop8( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), hiddenDimOffset, - hiddenDimSize); - } + core::WarpAccumCombineDequantSegmentScalarTop8( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), hiddenDimOffset, hiddenDimSize, + hiddenDim, args.fp8BlockwiseCombineScaleDim); } } else { if (mwIter.warpsPerItem == 1) { - if constexpr (UseFp4Combine) { - core::WarpAccumFp4DequantFull( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), validAccumCount, hiddenDim, - args.fp8BlockwiseCombineScaleDim); - } else { - core::WarpAccumFp8DequantFull( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), validAccumCount, hiddenDim, - args.fp8BlockwiseCombineScaleDim); - } + core::WarpAccumCombineDequantFull( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), validAccumCount, hiddenDim, + args.fp8BlockwiseCombineScaleDim); } else { - if constexpr (UseFp4Combine) { - core::WarpAccumFp4DequantSegment( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), validAccumCount, - hiddenDimOffset, hiddenDimSize, hiddenDim, args.fp8BlockwiseCombineScaleDim); - } else { - core::WarpAccumFp8DequantSegment( - outPtr, reinterpret_cast(srcPtrs), - reinterpret_cast(srcScalePtrs), validAccumCount, - hiddenDimOffset, hiddenDimSize, hiddenDim, args.fp8BlockwiseCombineScaleDim); - } + core::WarpAccumCombineDequantSegment( + outPtr, reinterpret_cast(srcPtrs), + reinterpret_cast(srcScalePtrs), validAccumCount, hiddenDimOffset, + hiddenDimSize, hiddenDim, args.fp8BlockwiseCombineScaleDim); } } } else if constexpr (!std::is_same_v && diff --git a/tests/python/ops/dispatch_combine_test_utils.py b/tests/python/ops/dispatch_combine_test_utils.py index f424d14d1..006704dbf 100644 --- a/tests/python/ops/dispatch_combine_test_utils.py +++ b/tests/python/ops/dispatch_combine_test_utils.py @@ -617,23 +617,47 @@ def check_combine_result( inp_converted.to(torch.float32) * unique_pes ).to(combine_data_type) - atol, rtol = 1e-2, 1e-2 - if getattr(self.config, "quant_type", "none") == "fp8_direct_cast": - atol, rtol = 1e-1, 1e-1 - 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 - result_match = torch.allclose( - got.float(), expected.float(), atol=atol, rtol=rtol - ) - if not result_match: + quant = getattr(self.config, "quant_type", "none") + if quant == "fp4_blockwise": + # FP4 (E2M1) is lossy, so comparing against the exact bf16 reference with a fixed + # rtol is not meaningful: a small element inside a high-max block rounds to 0 or + # 0.5*scale (~100% relative error). Instead bound the per-element error by FP4's + # actual worst-case rounding. Within a block scaled to [-6, 6] the largest grid + # step is |6 - 4| = 2, so the max round error is 1.0 in scaled units == blockMax/6 + # in real units. Each combined element is the same input summed `unique_pes` times, + # so the bound scales with unique_pes. This is an exact bound for blockwise FP4, + # not a guessed tolerance. + fp4_max = 6.0 + scale_dim = op._fp8_blockwise_combine_scale_dim + hidden = inp_converted.numel() + block_elems = (hidden + scale_dim - 1) // scale_dim + absv = inp_converted.float().abs() + pad = scale_dim * block_elems - hidden + if pad: + absv = torch.cat([absv, absv.new_zeros(pad)]) + block_max = absv.reshape(scale_dim, block_elems).amax(dim=1) + per_elem_block_max = block_max.repeat_interleave(block_elems)[:hidden] + tol = ( + per_elem_block_max / fp4_max * unique_pes # FP4 quantization bound + + 5e-2 + * expected.float().abs() # bf16 output rounding + accumulation slack + + 1e-2 + ) + diff = (got.float() - expected.float()).abs() + result_match = bool((diff <= tol).all()) + else: + atol, rtol = 1e-2, 1e-2 + if quant == "fp8_direct_cast": + atol, rtol = 1e-1, 1e-1 + elif quant == "fp8_blockwise": + # FP8 E4M3 quantization can exceed 5% after combine accumulation. + atol, rtol = 7e-2, 7e-2 + result_match = torch.allclose( + got.float(), expected.float(), atol=atol, rtol=rtol + ) diff = (got.float() - expected.float()).abs() tol = atol + rtol * expected.float().abs() + if not result_match: max_idx = int(diff.argmax().item()) print(f"Rank[{self.config.rank}] result mismatch for token {i}:") print( @@ -762,7 +786,7 @@ def run_ep_dispatch_local_expert_count_test(config): test_case = EpDispatchCombineTestCase(config) test_data = test_case.gen_test_data() - (_, all_rank_indices, all_rank_input, all_rank_weights, all_rank_scales) = test_data + _, all_rank_indices, all_rank_input, all_rank_weights, all_rank_scales = test_data rank = config.rank kt = config.kernel_type @@ -770,7 +794,7 @@ def run_ep_dispatch_local_expert_count_test(config): if kt.value == asyncll_type.value: # AsyncLL: recv kernels fill the output indices, so local_expert_count # must be requested on dispatch_recv (not dispatch_send). - (_, _, _, dispatch_indices, dispatch_recv_num_token) = op.dispatch_send( + _, _, _, dispatch_indices, dispatch_recv_num_token = op.dispatch_send( all_rank_input[rank], all_rank_weights[rank], all_rank_scales[rank], @@ -778,7 +802,7 @@ def run_ep_dispatch_local_expert_count_test(config): ) op.dispatch_recv(call_local_expert_count=True) else: - (_, _, _, dispatch_indices, dispatch_recv_num_token) = op.dispatch( + _, _, _, dispatch_indices, dispatch_recv_num_token = op.dispatch( all_rank_input[rank], all_rank_weights[rank], all_rank_scales[rank], diff --git a/tests/python/ops/test_dispatch_combine_intranode.py b/tests/python/ops/test_dispatch_combine_intranode.py index d4a47c37d..1639d6977 100644 --- a/tests/python/ops/test_dispatch_combine_intranode.py +++ b/tests/python/ops/test_dispatch_combine_intranode.py @@ -32,6 +32,14 @@ ) +def _combine_fp4_supported(): + # Packed-FP4 combine needs the gfx950 OCP FP4 conversion instructions. + try: + return "gfx950" in torch.cuda.get_device_properties(0).gcnArchName + except Exception: + return False + + def _make_intranode_config( rank, world_size, @@ -192,6 +200,10 @@ def test_dispatch_combine( pytest.skip(f"{quant_type} only supports bfloat16 input") if not use_external_inp_buf: pytest.skip(f"{quant_type} requires use_external_inp_buf=True") + if quant_type == "fp4_blockwise" and not _combine_fp4_supported(): + pytest.skip( + "fp4_blockwise combine requires a gfx950 GPU (OCP FP4 instructions)" + ) # blockwise combine ignores scale_dim/scale_type_size (driven by # MORI_FP8_COMBINE_SCALE_DIM internally). fp4_blockwise reuses the same path # but transports packed FP4 (E2M1) instead of FP8. @@ -242,6 +254,10 @@ def test_dispatch_combine_weightless_vec8( use_external_inp_buf, quant_type, ): + if quant_type == "fp4_blockwise" and not _combine_fp4_supported(): + pytest.skip( + "fp4_blockwise combine requires a gfx950 GPU (OCP FP4 instructions)" + ) expect_combine_kernel_substr = ( "noweight_block128_vec8_top9" if num_experts_per_token == 9 @@ -309,6 +325,10 @@ def test_dispatch_combine_ll( pytest.skip(f"{quant_type} only supports bfloat16 input") if not use_external_inp_buf: pytest.skip(f"{quant_type} requires use_external_inp_buf=True") + if quant_type == "fp4_blockwise" and not _combine_fp4_supported(): + pytest.skip( + "fp4_blockwise combine requires a gfx950 GPU (OCP FP4 instructions)" + ) # blockwise combine ignores scale_dim/scale_type_size (driven by # MORI_FP8_COMBINE_SCALE_DIM internally). fp4_blockwise reuses the same path # but transports packed FP4 (E2M1) instead of FP8. From c7548ecfcc0f56a7e3f9650a3b10a0bcf633c447 Mon Sep 17 00:00:00 2001 From: Karan Verma Date: Fri, 10 Jul 2026 09:56:47 -0500 Subject: [PATCH 3/5] address review: dedicated Fp4BlockwiseQuant enum + half-sized staging, 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 --- .../ops/dispatch_combine/dispatch_combine.hpp | 18 +++++- python/mori/ops/dispatch_combine.py | 59 ++++++++++++++----- src/ops/dispatch_combine/dispatch_combine.cpp | 14 +++-- src/ops/dispatch_combine/intranode.hpp | 7 ++- src/ops/dispatch_combine/launch.cpp | 43 +++++++------- src/pybind/pybind_ops.cpp | 1 + 6 files changed, 95 insertions(+), 47 deletions(-) diff --git a/include/mori/ops/dispatch_combine/dispatch_combine.hpp b/include/mori/ops/dispatch_combine/dispatch_combine.hpp index 025d5e415..743aa9771 100644 --- a/include/mori/ops/dispatch_combine/dispatch_combine.hpp +++ b/include/mori/ops/dispatch_combine/dispatch_combine.hpp @@ -61,7 +61,14 @@ enum KernelType { AsyncLL = 4, IntraNodeLL = 5 }; -enum class QuantType { None = 0, Fp8DirectCast = 1, Fp8BlockwiseQuant = 2 }; +enum class QuantType { None = 0, Fp8DirectCast = 1, Fp8BlockwiseQuant = 2, Fp4BlockwiseQuant = 3 }; + +// Blockwise combine quant types share the same staging/scale layout (per-block float scales); +// they differ only in the transported element codec: FP8 = 1 byte/elem, FP4 (E2M1) = 0.5 byte/elem +// (2 values packed per byte), so an FP4 token region is half the size of an FP8 one. +inline __host__ __device__ bool IsBlockwiseCombineQuant(QuantType q) { + return q == QuantType::Fp8BlockwiseQuant || q == QuantType::Fp4BlockwiseQuant; +} inline const char* HipDataTypeToString(hipDataType dtype) { switch (dtype) { @@ -169,6 +176,15 @@ struct EpDispatchCombineConfig { } inline __host__ __device__ size_t MaxHiddenBytes() const { return HiddenBytes(maxTokenTypeSize); } + // Per-token combine token-region bytes (the quantized payload, excluding scales/weights/index). + // FP4 blockwise packs two E2M1 values per byte, so its slot is half the FP8 (1 byte/elem) size; + // FP8 blockwise and other paths keep the existing element sizing. Used by both the host staging + // allocator and the intra-node kernel so the slot stride stays consistent. + inline __host__ __device__ size_t CombineTokenRegionBytes() const { + if (quantType == QuantType::Fp4BlockwiseQuant) return (static_cast(hiddenDim) + 1) / 2; + return HiddenBytes(maxTokenTypeSize); + } + inline __host__ __device__ size_t IndexBytes() const { return numExpertPerToken * sizeof(index_t); } diff --git a/python/mori/ops/dispatch_combine.py b/python/mori/ops/dispatch_combine.py index eb5b61be5..b3fcfea4f 100644 --- a/python/mori/ops/dispatch_combine.py +++ b/python/mori/ops/dispatch_combine.py @@ -47,12 +47,32 @@ 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, + # Blockwise FP4 (E2M1) combine is its own quant type. It shares the blockwise staging/scale + # layout with FP8 but transports packed FP4 (0.5 byte/elem) and uses half-sized staging slots. + "fp4_blockwise": EpDispatchCombineQuantType.Fp4BlockwiseQuant, } +# Blockwise combine quant types share the staging/scale layout and kernel launch config; only the +# element codec (and staging slot size) differ, so kernel selection treats them together and then +# swaps the codec token (fp8bwq <-> fp4bwq) in the kernel name. +_BLOCKWISE_COMBINE_QUANT_TYPES = ( + EpDispatchCombineQuantType.Fp8BlockwiseQuant, + EpDispatchCombineQuantType.Fp4BlockwiseQuant, +) + +# The FP4 blockwise combine kernels registered in ep_intranode.hip. Kernel-name selection derives +# an fp4bwq name from the fp8bwq one; the result is asserted against this set so a mismatch fails +# loudly instead of launching a non-existent symbol. +_FP4_COMBINE_KERNELS = frozenset( + { + "EpCombineIntraNodeKernel_bf16_nop2p_fp4bwq", + "EpCombineIntraNodeKernel_bf16_nop2p_fp4bwq_noweight_block128_vec8", + "EpCombineIntraNodeKernel_bf16_nop2p_fp4bwq_noweight_block256_vec8", + "EpCombineIntraNodeKernel_bf16_nop2p_fp4bwq_noweight_block128_vec8_top9", + "EpCombineIntraNodeKernel_bf16_nop2p_fp4bwq_noweight_block256_vec8_top9", + } +) + def _normalize_quant_type(quant_type): if isinstance(quant_type, EpDispatchCombineQuantType): @@ -266,8 +286,8 @@ 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. + # Detect the fp4_blockwise combine so we can fail fast on unsupported archs at construction. + # (Kernel selection keys off the Fp4BlockwiseQuant enum directly.) self._combine_is_fp4 = ( isinstance(config.quant_type, str) and config.quant_type.strip().lower() == "fp4_blockwise" @@ -457,7 +477,7 @@ def _combine_shared_mem(self, warp_per_block, use_weights=True): """Shared memory for combine kernels.""" quant_type = _normalize_quant_type(self.config.quant_type) num_ptr_arrays = 1 + int(bool(use_weights)) - if quant_type == EpDispatchCombineQuantType.Fp8BlockwiseQuant: + if quant_type in _BLOCKWISE_COMBINE_QUANT_TYPES: num_ptr_arrays += 1 return ( warp_per_block @@ -965,24 +985,29 @@ def combine( quant_type = _normalize_quant_type(self.config.quant_type) shared_mem = self._combine_shared_mem(actual_wpb) - if quant_type == EpDispatchCombineQuantType.Fp8BlockwiseQuant: + if quant_type in _BLOCKWISE_COMBINE_QUANT_TYPES: + label = ( + "fp4_blockwise" + if quant_type == EpDispatchCombineQuantType.Fp4BlockwiseQuant + else "fp8_blockwise" + ) if kt not in ( EpDispatchCombineKernelType.IntraNode.value, EpDispatchCombineKernelType.IntraNodeLL.value, ): raise ValueError( - "Fp8BlockwiseQuant currently only supports IntraNode/IntraNodeLL combine" + f"{label} combine currently only supports IntraNode/IntraNodeLL combine" ) if sfx != "bf16": - raise ValueError(f"Fp8BlockwiseQuant only supports bf16, got {sfx}") + raise ValueError(f"{label} combine only supports bf16, got {sfx}") if not actual_use_ext: raise ValueError( - "Fp8BlockwiseQuant currently requires --zero-copy 0 " + f"{label} combine currently requires --zero-copy 0 " "(useExternalInpBuffer=True). P2P read path not yet implemented." ) if self._fp8_blockwise_combine_scale_dim <= 0: raise ValueError( - "Fp8BlockwiseQuant requires internal combine scale_dim > 0" + f"{label} combine requires internal combine scale_dim > 0" ) if kt == EpDispatchCombineKernelType.InterNode.value: @@ -1030,7 +1055,7 @@ def combine( EpDispatchCombineKernelType.IntraNode.value, EpDispatchCombineKernelType.IntraNodeLL.value, ): - if quant_type == EpDispatchCombineQuantType.Fp8BlockwiseQuant: + if quant_type in _BLOCKWISE_COMBINE_QUANT_TYPES: # Mirror of the AccumNum=8/9 + VecBytes=8 specialization gating in # LaunchCombine() / launch.cpp. top-k==9 covers shared-expert fusion # (8 routed + 1 fused shared). Keep in sync. @@ -1061,9 +1086,13 @@ def combine( ) 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: + # the fp8bwq variants; only the in-kernel quant/dequant math differs). Assert the + # derived name is a registered fp4bwq symbol so a naming mismatch fails loudly. + if quant_type == EpDispatchCombineQuantType.Fp4BlockwiseQuant: kernel_name = kernel_name.replace("_fp8bwq", "_fp4bwq") + assert ( + kernel_name in _FP4_COMBINE_KERNELS + ), f"fp4_blockwise combine selected unregistered kernel '{kernel_name}'" shared_mem = self._combine_shared_mem( actual_wpb, use_weights=not use_vec8_top8 ) diff --git a/src/ops/dispatch_combine/dispatch_combine.cpp b/src/ops/dispatch_combine/dispatch_combine.cpp index 6a0304281..18358406b 100644 --- a/src/ops/dispatch_combine/dispatch_combine.cpp +++ b/src/ops/dispatch_combine/dispatch_combine.cpp @@ -118,12 +118,13 @@ EpDispatchCombineHandle::EpDispatchCombineHandle(EpDispatchCombineConfig config_ shmemNumQpPerPe, shmemNumQpPerPe); } - if (config.quantType == QuantType::Fp8BlockwiseQuant) { + if (IsBlockwiseCombineQuant(config.quantType)) { fp8BlockwiseCombineScaleDim = env::GetPositiveIntOr(kFp8BlockwiseScaleDimEnv, kDefaultFp8BlockwiseScaleDim); fp8BlockwiseCombineScaleTypeSize = static_cast(sizeof(float)); if (config.rank == 0) { - MORI_OPS_INFO("Fp8BlockwiseQuant combine scale_dim={} (override via {})", + MORI_OPS_INFO("Blockwise combine ({}) scale_dim={} (override via {})", + config.quantType == QuantType::Fp4BlockwiseQuant ? "FP4" : "FP8", fp8BlockwiseCombineScaleDim, kFp8BlockwiseScaleDimEnv); } } @@ -195,14 +196,15 @@ void EpDispatchCombineHandle::InitializeShmemBuf() { config.HiddenDimSz() * config.maxTokenTypeSize; size_t maxStagingSize = static_cast(config.MaxNumTokensToRecv()) * config.MaxXferBytesPerToken(); - if (config.kernelType == KernelType::IntraNode && - config.quantType == QuantType::Fp8BlockwiseQuant) { + if (config.kernelType == KernelType::IntraNode && IsBlockwiseCombineQuant(config.quantType)) { size_t blockwiseScaleBytes = (fp8BlockwiseCombineScaleDim > 0) ? static_cast(fp8BlockwiseCombineScaleDim) * fp8BlockwiseCombineScaleTypeSize : 0; + // FP4 packs the token region at 0.5 byte/elem (CombineTokenRegionBytes()), so its staging slot + // is half the FP8 one -- no FP8-sized over-allocation for FP4. maxStagingSize = static_cast(config.MaxNumTokensToRecv()) * - (config.HiddenBytes(config.maxTokenTypeSize) + config.IndexBytes() + + (config.CombineTokenRegionBytes() + config.IndexBytes() + config.WeightBytes() + config.SrcTokenIdBytes() + blockwiseScaleBytes); } @@ -256,7 +258,7 @@ void EpDispatchCombineHandle::InitializeShmemBuf() { static_cast(config.MaxNumTokensToRecv()) * config.scaleDim * config.scaleTypeSize; } size_t fp8BlockwiseScaleSize = 0; - if (config.quantType == QuantType::Fp8BlockwiseQuant && fp8BlockwiseCombineScaleDim > 0) { + if (IsBlockwiseCombineQuant(config.quantType) && fp8BlockwiseCombineScaleDim > 0) { fp8BlockwiseScaleSize = static_cast(config.MaxNumTokensToRecv()) * fp8BlockwiseCombineScaleDim * fp8BlockwiseCombineScaleTypeSize; } diff --git a/src/ops/dispatch_combine/intranode.hpp b/src/ops/dispatch_combine/intranode.hpp index 66a9de88a..b23a9d0d3 100644 --- a/src/ops/dispatch_combine/intranode.hpp +++ b/src/ops/dispatch_combine/intranode.hpp @@ -302,9 +302,12 @@ __device__ __forceinline__ void EpCombineIntraNodeKernel_body(EpDispatchCombineA const uint64_t crossDeviceBarrierFlag = args.crossDeviceBarrierFlag[0]; // Copy input to shmem registered buffer so that other GPUs can access directly index_t totalRecvTokenNum = args.totalRecvTokenNum[0]; - // When TokT != T (e.g. fp8 combine), staging layout uses TokT-sized tokens + // When TokT != T (e.g. fp8 combine), staging layout uses TokT-sized tokens. FP4 blockwise packs + // two E2M1 values per byte, so its token region is half the FP8 one -- keep this in sync with + // EpDispatchCombineConfig::CombineTokenRegionBytes() used by the host staging allocator. const size_t hiddenDim = config.HiddenDimSz(); - const size_t hiddenBytes = hiddenDim * sizeof(TokT); + const size_t hiddenBytes = + UseFp4Combine ? ((hiddenDim + 1) / 2) * sizeof(TokT) : hiddenDim * sizeof(TokT); const size_t weightBytes = (UseWeights && args.weightsBuf != nullptr) ? config.numExpertPerToken * sizeof(float) : 0; const size_t scaleBytes = diff --git a/src/ops/dispatch_combine/launch.cpp b/src/ops/dispatch_combine/launch.cpp index daed4fa9c..d3de3b02b 100644 --- a/src/ops/dispatch_combine/launch.cpp +++ b/src/ops/dispatch_combine/launch.cpp @@ -514,57 +514,54 @@ void LaunchCombine(EpDispatchCombineHandle& handle, void* input, void* weights, unsigned int block_x = WARP_SIZE * wpb; const bool hasWeights = (weights != nullptr); int smem = combine_shared_mem(wpb, handle.config.numExpertPerToken, - handle.config.quantType == QuantType::Fp8BlockwiseQuant, + IsBlockwiseCombineQuant(handle.config.quantType), /*use_weight_ptrs=*/true); size_t args_size = sizeof(EpDispatchCombineArgsRaw); const char* sfx = dtype_suffix(dtype); auto& reg = KernelRegistry::Instance(); int mp = handle.multiProcessorCount; - if (args.config.quantType == QuantType::Fp8BlockwiseQuant && + if (IsBlockwiseCombineQuant(args.config.quantType) && handle.config.kernelType != KernelType::IntraNode) { - throw std::runtime_error("Fp8BlockwiseQuant currently only supports IntraNode combine"); + throw std::runtime_error("Blockwise combine quant currently only supports IntraNode combine"); } switch (handle.config.kernelType) { case KernelType::IntraNode: - if (args.config.quantType == QuantType::Fp8BlockwiseQuant) { + if (IsBlockwiseCombineQuant(args.config.quantType)) { + const bool isFp4 = (args.config.quantType == QuantType::Fp4BlockwiseQuant); if (dtype != HIP_R_16BF) { - throw std::runtime_error("Fp8BlockwiseQuant currently only supports bf16 input"); + throw std::runtime_error("Blockwise combine quant currently only supports bf16 input"); } if (!args.config.useExternalInpBuffer) { throw std::runtime_error( - "Fp8BlockwiseQuant currently requires --zero-copy 0 " + "Blockwise combine quant currently requires --zero-copy 0 " "(useExternalInpBuffer=true)"); } const int fp8ScaleDim = args.fp8BlockwiseCombineScaleDim; if (fp8ScaleDim <= 0) { - throw std::runtime_error("Fp8BlockwiseQuant requires internal combine scaleDim > 0"); + throw std::runtime_error( + "Blockwise combine quant requires internal combine scaleDim > 0"); } // Pick the AccumNum=8/9 + VecBytes=8 specialization when (no weights, hidden_dim % 512 == // 0, top-k in {8,9}, EP > 4) and block_elems matches a registered symbol. top-k==9 covers - // shared-expert fusion (8 routed + 1 fused shared). Keep in sync with the Python launch - // path in dispatch_combine.py. + // shared-expert fusion (8 routed + 1 fused shared). The FP4 variants ("fp4bwq") share the + // exact launch config as FP8 ("fp8bwq"); only the in-kernel codec differs. Keep in sync + // with the Python launch path in dispatch_combine.py. const int block_elems = (args.config.hiddenDim + fp8ScaleDim - 1) / fp8ScaleDim; const bool baseVec8Top8Eligible = !hasWeights && (args.config.hiddenDim % 512 == 0) && (args.config.numExpertPerToken == 8 || args.config.numExpertPerToken == 9) && args.config.worldSize > 4; const bool top9 = (args.config.numExpertPerToken == 9); - const char* kernel_name = "EpCombineIntraNodeKernel_bf16_nop2p_fp8bwq"; + const std::string bwq = isFp4 ? "fp4bwq" : "fp8bwq"; + const std::string prefix = "EpCombineIntraNodeKernel_bf16_nop2p_"; + std::string kernel_name = prefix + bwq; bool useVec8Top8 = false; - if (baseVec8Top8Eligible) { - if (block_elems == 128) { - kernel_name = - top9 ? "EpCombineIntraNodeKernel_bf16_nop2p_fp8bwq_noweight_block128_vec8_top9" - : "EpCombineIntraNodeKernel_bf16_nop2p_fp8bwq_noweight_block128_vec8"; - useVec8Top8 = true; - } else if (block_elems == 256) { - kernel_name = - top9 ? "EpCombineIntraNodeKernel_bf16_nop2p_fp8bwq_noweight_block256_vec8_top9" - : "EpCombineIntraNodeKernel_bf16_nop2p_fp8bwq_noweight_block256_vec8"; - useVec8Top8 = true; - } + if (baseVec8Top8Eligible && (block_elems == 128 || block_elems == 256)) { + kernel_name = prefix + bwq + "_noweight_block" + std::to_string(block_elems) + "_vec8" + + (top9 ? "_top9" : ""); + useVec8Top8 = true; } int fp8bwq_smem = combine_shared_mem(wpb, handle.config.numExpertPerToken, /*use_scale_ptrs=*/true, @@ -660,7 +657,7 @@ void LaunchCombineRecv(EpDispatchCombineHandle& handle, int block_num, int warp_ unsigned int block_x = WARP_SIZE * wpb; int smem = combine_shared_mem(wpb, handle.config.numExpertPerToken, - handle.config.quantType == QuantType::Fp8BlockwiseQuant); + IsBlockwiseCombineQuant(handle.config.quantType)); size_t args_size = sizeof(EpDispatchCombineArgsRaw); const char* sfx = dtype_suffix(handle.inputType); diff --git a/src/pybind/pybind_ops.cpp b/src/pybind/pybind_ops.cpp index cf59dcbfe..5a9021a1f 100644 --- a/src/pybind/pybind_ops.cpp +++ b/src/pybind/pybind_ops.cpp @@ -369,6 +369,7 @@ void RegisterMoriOps(py::module_& m) { .value("None_", mori::moe::QuantType::None) .value("Fp8DirectCast", mori::moe::QuantType::Fp8DirectCast) .value("Fp8BlockwiseQuant", mori::moe::QuantType::Fp8BlockwiseQuant) + .value("Fp4BlockwiseQuant", mori::moe::QuantType::Fp4BlockwiseQuant) .export_values(); mori::pybind::RegisterAllProfilerSlots(m); From ddcbb4ff872cfa53a0992cd36edb7369960b2d8d Mon Sep 17 00:00:00 2001 From: Karan Verma Date: Mon, 13 Jul 2026 09:04:53 -0500 Subject: [PATCH 4/5] address review: rename Fp8T storage template param to PackedT in FP4 helpers The new blockwise FP4 (E2M1) quant/dequant helpers carried their packed storage type as `Fp8T`, which is misleading since it holds packed FP4, not FP8. Rename it to `PackedT` in the 7 FP4 helpers for clarity (cosmetic; no behavior change). Shared fp4/fp8 dispatchers pass the type positionally, so call sites are unaffected. Co-authored-by: Cursor --- .../core/transport/p2p/device_primitives.hpp | 40 +++++++++---------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/include/mori/core/transport/p2p/device_primitives.hpp b/include/mori/core/transport/p2p/device_primitives.hpp index e684fe8f7..6dfbd3e92 100644 --- a/include/mori/core/transport/p2p/device_primitives.hpp +++ b/include/mori/core/transport/p2p/device_primitives.hpp @@ -2170,8 +2170,8 @@ __device__ __forceinline__ void WarpAccumFp8DequantSegment( // Requires: InT == hip_bfloat16, blockElems == SubwarpSize*kVecElems, kVecElems % 8 == 0, // and block byte-alignment (start even). Uses vector loads (int4 = 8 bf16) and vector stores // (uint32 = 8 packed fp4) so the send-quant is not bottlenecked on scalar accesses. -template -__device__ __forceinline__ void WarpQuantizeToFp4BlockwiseVec(Fp8T* __restrict__ dstToken, +template +__device__ __forceinline__ void WarpQuantizeToFp4BlockwiseVec(PackedT* __restrict__ dstToken, float* __restrict__ dstScales, const InT* __restrict__ srcToken, int hiddenDim, int scaleDim) { @@ -2229,8 +2229,8 @@ __device__ __forceinline__ void WarpQuantizeToFp4BlockwiseVec(Fp8T* __restrict__ if (laneId == 0) dstScales[0] = -dstScales[0]; } -template -__device__ __forceinline__ void WarpQuantizeToFp4Blockwise(Fp8T* __restrict__ dstToken, +template +__device__ __forceinline__ void WarpQuantizeToFp4Blockwise(PackedT* __restrict__ dstToken, float* __restrict__ dstScales, const InT* __restrict__ srcToken, int hiddenDim, int scaleDim) { @@ -2242,7 +2242,7 @@ __device__ __forceinline__ void WarpQuantizeToFp4Blockwise(Fp8T* __restrict__ ds // Fast path: 64-lane warp, blockElems == 128 == 8 subwarps * 16 elems (the DeepSeek config). if (warpSize == 64 && blockElems == 128 && (hiddenDim % 128) == 0 && std::is_same_v) { - WarpQuantizeToFp4BlockwiseVec<8, 16, Fp8T, InT>(dstToken, dstScales, srcToken, hiddenDim, + WarpQuantizeToFp4BlockwiseVec<8, 16, PackedT, InT>(dstToken, dstScales, srcToken, hiddenDim, scaleDim); return; } @@ -2250,7 +2250,7 @@ __device__ __forceinline__ void WarpQuantizeToFp4Blockwise(Fp8T* __restrict__ ds // block256 vec8 variant so block256 configs are not left on the scalar path). if (warpSize == 64 && blockElems == 256 && (hiddenDim % 256) == 0 && std::is_same_v) { - WarpQuantizeToFp4BlockwiseVec<8, 32, Fp8T, InT>(dstToken, dstScales, srcToken, hiddenDim, + WarpQuantizeToFp4BlockwiseVec<8, 32, PackedT, InT>(dstToken, dstScales, srcToken, hiddenDim, scaleDim); return; } @@ -2290,9 +2290,9 @@ __device__ __forceinline__ void WarpQuantizeToFp4Blockwise(Fp8T* __restrict__ ds if (laneId == 0) dstScales[0] = -dstScales[0]; } -template +template __device__ __forceinline__ void WarpAccumFp4DequantFull(OutT* __restrict__ dstToken, - const Fp8T* const* __restrict__ srcs, + const PackedT* const* __restrict__ srcs, const float* const* __restrict__ srcScales, int accumNum, int hiddenDim, int scaleDim) { const int laneId = threadIdx.x & (warpSize - 1); @@ -2340,9 +2340,9 @@ __device__ __forceinline__ void WarpAccumFp4DequantFull(OutT* __restrict__ dstTo } } -template +template __device__ __forceinline__ void WarpAccumFp4DequantSegment( - OutT* __restrict__ dstToken, const Fp8T* const* __restrict__ srcs, + OutT* __restrict__ dstToken, const PackedT* const* __restrict__ srcs, const float* const* __restrict__ srcScales, int accumNum, int hiddenDimOffset, int hiddenDimSize, int hiddenDim, int scaleDim) { const int laneId = threadIdx.x & (warpSize - 1); @@ -2381,9 +2381,9 @@ __device__ __forceinline__ void WarpAccumFp4DequantSegment( /* (tb = srcs[i] - hiddenDimOffset) and index in packed (half-byte) units. Blocks/offsets are */ /* byte-aligned (blockElems and hiddenDimOffset even) for all supported configs. */ /* ---------------------------------------------------------------------------------------------- */ -template +template __device__ __forceinline__ void WarpAccumFp4DequantVecBlockwiseScaleWave( - OutT* __restrict__ dstToken, const Fp8T* const* __restrict__ srcs, + OutT* __restrict__ dstToken, const PackedT* const* __restrict__ srcs, const float* const* __restrict__ srcScales, int hiddenDimOffset, int start, int end, int blockElems) { const int laneId = threadIdx.x & (warpSize - 1); @@ -2406,7 +2406,7 @@ __device__ __forceinline__ void WarpAccumFp4DequantVecBlockwiseScaleWave( float sArr[AccumNum]; #pragma unroll AccumNum for (int i = 0; i < AccumNum; ++i) { - const Fp8T* src = srcs[i]; + const PackedT* src = srcs[i]; if (src == nullptr) { sArr[i] = 0.0f; // sentinel: contributes nothing bitsArr[i] = 0u; @@ -2448,7 +2448,7 @@ __device__ __forceinline__ void WarpAccumFp4DequantVecBlockwiseScaleWave( float acc = 0.0f; #pragma unroll AccumNum for (int i = 0; i < AccumNum; ++i) { - const Fp8T* src = srcs[i]; + const PackedT* src = srcs[i]; if (src == nullptr) continue; float s = 1.0f; if (srcScales != nullptr && srcScales[i] != nullptr) { @@ -2465,19 +2465,19 @@ __device__ __forceinline__ void WarpAccumFp4DequantVecBlockwiseScaleWave( } } -template +template __device__ __forceinline__ void WarpAccumFp4DequantFullBlockVec8Top8( - OutT* __restrict__ dstToken, const Fp8T* const* __restrict__ srcs, + OutT* __restrict__ dstToken, const PackedT* const* __restrict__ srcs, const float* const* __restrict__ srcScales, int hiddenDim) { - WarpAccumFp4DequantVecBlockwiseScaleWave( + WarpAccumFp4DequantVecBlockwiseScaleWave( dstToken, srcs, srcScales, /*hiddenDimOffset=*/0, /*start=*/0, /*end=*/hiddenDim, BlockElems); } -template +template __device__ __forceinline__ void WarpAccumFp4DequantSegmentBlockVec8Top8( - OutT* __restrict__ dstToken, const Fp8T* const* __restrict__ srcs, + OutT* __restrict__ dstToken, const PackedT* const* __restrict__ srcs, const float* const* __restrict__ srcScales, int hiddenDimOffset, int hiddenDimSize) { - WarpAccumFp4DequantVecBlockwiseScaleWave( + WarpAccumFp4DequantVecBlockwiseScaleWave( dstToken, srcs, srcScales, hiddenDimOffset, /*start=*/0, /*end=*/hiddenDimSize, BlockElems); } From 96942f3b366701bdaaa46d6b0cf907a94a33ecd9 Mon Sep 17 00:00:00 2001 From: Karan Verma Date: Mon, 13 Jul 2026 09:19:41 -0500 Subject: [PATCH 5/5] fix(fp4): clang-format continuation alignment after PackedT rename The PackedT rename widened the WarpQuantizeToFp4BlockwiseVec<...> call prefix, so the wrapped scaleDim argument must re-align under the paren. Fixes the pre-commit clang-format check. Co-authored-by: Cursor --- include/mori/core/transport/p2p/device_primitives.hpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/include/mori/core/transport/p2p/device_primitives.hpp b/include/mori/core/transport/p2p/device_primitives.hpp index 6dfbd3e92..08c137ef7 100644 --- a/include/mori/core/transport/p2p/device_primitives.hpp +++ b/include/mori/core/transport/p2p/device_primitives.hpp @@ -2243,7 +2243,7 @@ __device__ __forceinline__ void WarpQuantizeToFp4Blockwise(PackedT* __restrict__ if (warpSize == 64 && blockElems == 128 && (hiddenDim % 128) == 0 && std::is_same_v) { WarpQuantizeToFp4BlockwiseVec<8, 16, PackedT, InT>(dstToken, dstScales, srcToken, hiddenDim, - scaleDim); + scaleDim); return; } // Fast path: 64-lane warp, blockElems == 256 == 8 subwarps * 32 elems (mirrors the fp8bwq @@ -2251,7 +2251,7 @@ __device__ __forceinline__ void WarpQuantizeToFp4Blockwise(PackedT* __restrict__ if (warpSize == 64 && blockElems == 256 && (hiddenDim % 256) == 0 && std::is_same_v) { WarpQuantizeToFp4BlockwiseVec<8, 32, PackedT, InT>(dstToken, dstScales, srcToken, hiddenDim, - scaleDim); + scaleDim); return; }