diff --git a/include/mscclpp/gpu_data_types.hpp b/include/mscclpp/gpu_data_types.hpp index 6bf9c979..21139d63 100644 --- a/include/mscclpp/gpu_data_types.hpp +++ b/include/mscclpp/gpu_data_types.hpp @@ -189,7 +189,7 @@ template = 4 && Bytes % 4 == 0)> struct alignas(Bytes) Words { uint32_t w[Bytes / 4]; - MSCCLPP_HOST_DEVICE_INLINE Words() {} + Words() = default; MSCCLPP_HOST_DEVICE_INLINE uint32_t& operator[](int i) { return w[i]; } @@ -212,7 +212,7 @@ union alignas(sizeof(T) * N) VectorTypeImpl { using ElementType = T; constexpr static int Size = N; - MSCCLPP_HOST_DEVICE_INLINE VectorTypeImpl() {} + VectorTypeImpl() = default; MSCCLPP_HOST_DEVICE_INLINE VectorTypeImpl(const StorageT& value) : storage(value) {} @@ -871,7 +871,7 @@ MSCCLPP_DEVICE_INLINE f32x4 to(const f8_e5m2x4& v) { /// f32x2 -> f8_e4m3x2. /// HIP gfx942: float -> fp8 (via __builtin_amdgcn_cvt_pk_fp8_f32). -/// NVIDIA SM90+: float -> half -> fp8 (via __nv_cvt_halfraw2_to_fp8x2). +/// NVIDIA SM90+: float -> fp8 (via __nv_cvt_float2_to_fp8x2). /// NVIDIA pre-SM90: float -> half -> fp8 (via __nv_cvt_halfraw_to_fp8, element-wise). template <> MSCCLPP_DEVICE_INLINE f8_e4m3x2 to(const f32x2& v) { @@ -879,10 +879,7 @@ MSCCLPP_DEVICE_INLINE f8_e4m3x2 to(const f32x2& v) { uint32_t packed = __builtin_amdgcn_cvt_pk_fp8_f32(v.data[0], v.data[1], 0, false); return bit_cast(static_cast<__hip_fp8x2_storage_t>(packed)); #elif defined(MSCCLPP_DEVICE_CUDA) && __CUDA_ARCH__ >= 900 - __half2_raw h2; - h2.x = bit_cast(__float2half_rn(v.data[0])); - h2.y = bit_cast(__float2half_rn(v.data[1])); - __nv_fp8x2_storage_t fp8x2 = __nv_cvt_halfraw2_to_fp8x2(h2, __NV_SATFINITE, __NV_E4M3); + __nv_fp8x2_storage_t fp8x2 = __nv_cvt_float2_to_fp8x2(v.storage, __NV_SATFINITE, __NV_E4M3); return bit_cast(fp8x2); #elif defined(MSCCLPP_DEVICE_CUDA) __half_raw h0, h1; diff --git a/python/mscclpp/ep/README.md b/python/mscclpp/ep/README.md index 255f9d80..57f7f4fe 100644 --- a/python/mscclpp/ep/README.md +++ b/python/mscclpp/ep/README.md @@ -261,10 +261,9 @@ can leave those fields as `None`. ```python @dataclass class QuantConfig: - dtype: Optional[torch.dtype] = None + format: Optional[DispatchDataType] = None block_scales: Optional[torch.Tensor] = None global_scale: Optional[torch.Tensor] = None - block_size: Optional[int] = None class DispatchLayout(str, Enum): @@ -473,17 +472,17 @@ combine to reduce the `K` expert results for each token back to `[T, H]`. ### `quant` `quant` contains activation quantization metadata for `input`. It should be -`None` for BF16/FP16 input. The quantized tensor dtype is stored in -`quant.dtype`. +`None` for BF16/FP16 input. `quant.format` defines the tensor representation +and scale layout. Examples: -| Format | `input` | `quant.dtype` | `quant.block_scales` | `quant.global_scale` | -|---|---|---|---|---| -| BF16/FP16 | `[T, H]` | `None` | `None` | `None` | -| FP8 E4M3 | `[T, H]` FP8 | `torch.float8_e4m3fn` | `[T, H / block_size]`, often block size 128 | usually `None` | -| NVFP4 | backend-defined packed/logical `[T, H]` | backend-defined | block scale tensor | optional global scale | -| MXFP8 | backend-defined `[T, H]` | backend-defined | micro-scale tensor, e.g. E8M0 blocks | optional/global if required | +| Format | `input` | `quant.block_scales` | `quant.global_scale` | +|---|---|---|---| +| BF16/FP16 | `[T, H]` | `None` | `None` | +| FP8 E4M3 | `[T, H]` FP8 | `[T, H / 128]` | usually `None` | +| NVFP4 | backend-defined packed/logical `[T, H]` | block scale tensor | optional global scale | +| MXFP8 | backend-defined `[T, H]` | micro-scale tensor, e.g. E8M0 blocks | optional/global if required | The API should not assume quantization scale is a scalar. For FP8 paths in DeepEP/SGLang, scales are usually per token and per hidden block. @@ -504,7 +503,7 @@ output_buffer: [num_local_experts, world_size * max_tokens_per_rank, hidden] The dtype must match the dispatch output dtype. For BF16 dispatch it is BF16. For FP8 dispatch it is FP8 and the returned `DispatchOutput.quant` carries the -matching dtype and scale tensor. +matching format and scale tensor. `output_buffer` is required for LL because the MLP runner often owns or reuses workspace memory. `MoECommunicator` writes dispatch output into the provided @@ -816,15 +815,16 @@ output = moe_comm.combine(expert_output, handle) Quantized path: ```python +moe_comm = MoECommunicator( + ..., + quant=QuantConfig(format=DispatchDataType.FP8_E4M3), +) + recv, handle = moe_comm.dispatch( - input=x_fp8, + input=hidden_states, # BF16 input, quantized during dispatch topk_ids=topk_ids, weights=topk_weights, - quant=QuantConfig( - dtype=torch.float8_e4m3fn, - block_scales=x_scales, - block_size=128, - ), + quant=None, output_buffer=recv_buffer, ) diff --git a/python/mscclpp/ep/__init__.py b/python/mscclpp/ep/__init__.py index e8b63087..1aa7f15d 100644 --- a/python/mscclpp/ep/__init__.py +++ b/python/mscclpp/ep/__init__.py @@ -15,6 +15,7 @@ from .communicator import ( # noqa: F401 CombineContext, CombineMode, DispatchHandle, + DispatchDataType, DispatchLayout, DispatchLayoutInfo, DispatchOutput, @@ -38,6 +39,7 @@ __all__ = [ "CombineContext", "CombineMode", "DispatchHandle", + "DispatchDataType", "DispatchLayout", "DispatchLayoutInfo", "DispatchOutput", diff --git a/python/mscclpp/ep/_cpp.py b/python/mscclpp/ep/_cpp.py index d646cec2..bc3c9874 100644 --- a/python/mscclpp/ep/_cpp.py +++ b/python/mscclpp/ep/_cpp.py @@ -15,6 +15,7 @@ except ImportError as exc: # pragma: no cover DispatchLayout = _cpp.DispatchLayout MoEMode = _cpp.MoEMode CombineMode = _cpp.CombineMode +DispatchDataType = _cpp.DispatchDataType Config = getattr(_cpp, "Config", None) diff --git a/python/mscclpp/ep/communicator.py b/python/mscclpp/ep/communicator.py index 128612d7..ffa18de6 100644 --- a/python/mscclpp/ep/communicator.py +++ b/python/mscclpp/ep/communicator.py @@ -8,7 +8,7 @@ from typing import Optional, Tuple import torch -from ._cpp import CombineMode, DispatchLayout, MoEMode +from ._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode from .high_throughput import HighThroughputBackend from .low_latency import LowLatencyBackend from .types import ( @@ -36,6 +36,7 @@ __all__ = [ "CombineContext", "CombineMode", "DispatchHandle", + "DispatchDataType", "DispatchLayout", "DispatchLayoutInfo", "DispatchOutput", diff --git a/python/mscclpp/ep/low_latency.py b/python/mscclpp/ep/low_latency.py index e41d7db9..41fcc4a9 100644 --- a/python/mscclpp/ep/low_latency.py +++ b/python/mscclpp/ep/low_latency.py @@ -8,7 +8,7 @@ from typing import Any, Optional import torch -from ._cpp import CombineMode, DispatchLayout, MoEMode, _cpp, get_low_latency_rdma_size_hint +from ._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode, _cpp, get_low_latency_rdma_size_hint from .types import ( DispatchHandle, DispatchLayoutInfo, @@ -22,6 +22,24 @@ from .types import ( from .utils import cuda_stream_ptr, resolve_expert_placement +def _resolve_dispatch_data_type(quant: Optional[QuantConfig]) -> DispatchDataType: + if quant is None: + return DispatchDataType.BF16 + + quant_format = quant.format + if quant_format is not None and not isinstance(quant_format, DispatchDataType): + raise TypeError("quant.format must be a DispatchDataType") + if quant_format is None: + raise ValueError("quant.format is required") + if quant_format == DispatchDataType.MXFP8_E4M3: + raise NotImplementedError("MXFP8 dispatch is reserved but not implemented") + if quant_format != DispatchDataType.FP8_E4M3: + raise ValueError("unsupported low-latency quantization format") + if quant.block_scales is not None or quant.global_scale is not None: + raise ValueError("communicator quant config must not contain precomputed scales") + return DispatchDataType.FP8_E4M3 + + class LowLatencyRuntime: """Private low-level low-latency runtime wrapper (wraps ``_cpp.MoERuntime``).""" @@ -109,12 +127,12 @@ class LowLatencyBackend: if config.max_recv_tokens_per_rank not in (None, self.max_tokens_per_rank): raise NotImplementedError("low-latency mode currently uses max_tokens_per_rank as recv capacity") - if config.quant is not None: - raise NotImplementedError("low-latency quantization is not implemented yet") + self.dispatch_data_type = _resolve_dispatch_data_type(config.quant) num_rdma_bytes = get_low_latency_rdma_size_hint( self.max_tokens_per_rank, self.hidden_size, self.world_size, self.num_experts, self.topk ) + self._dispatch_scales: Optional[torch.Tensor] = None self._dispatch_src_info: Optional[torch.Tensor] = None self._dispatch_layout_range: Optional[torch.Tensor] = None self._dispatch_count: Optional[torch.Tensor] = None @@ -151,12 +169,13 @@ class LowLatencyBackend: del previous_handle self._validate_dispatch_inputs(input, topk_ids, weights, quant, output_buffer) - out_buf, src_info, layout_range, count = self._get_dispatch_output_tensors(output_buffer) + out_buf, scales, src_info, layout_range, count = self._get_dispatch_output_tensors(output_buffer) self._runtime.cpp_runtime.dispatch( input.data_ptr(), topk_ids.data_ptr(), 0 if weights is None else weights.data_ptr(), out_buf.data_ptr(), + 0 if scales is None else scales.data_ptr(), src_info.data_ptr(), layout_range.data_ptr(), count.data_ptr(), @@ -165,12 +184,21 @@ class LowLatencyBackend: self.topk, self.max_tokens_per_rank, self.num_experts, + self.dispatch_data_type, self.num_blocks, cuda_stream_ptr(stream), ) + output_quant = ( + None + if scales is None + else QuantConfig( + format=self.dispatch_data_type, + block_scales=scales, + ) + ) output_info = DispatchOutputInfo( layout=DispatchLayoutInfo(kind=self.output_layout, num_tokens_per_expert=count), - quant=None, + quant=output_quant, ) dispatch_out = DispatchOutput( tokens=out_buf, @@ -218,6 +246,7 @@ class LowLatencyBackend: self.topk, context.num_max_dispatch_tokens_per_rank, context.num_experts, + self.dispatch_data_type, self.combine_mode, self.num_blocks - 2, cuda_stream_ptr(stream), @@ -235,8 +264,16 @@ class LowLatencyBackend: (self.num_local_experts, self.world_size), dtype=torch.int64, device=device ) self._dispatch_count = torch.empty((self.num_local_experts,), dtype=torch.int32, device=device) + self._dispatch_scales = None + if self.dispatch_data_type == DispatchDataType.FP8_E4M3: + num_scales = self.hidden_size // 128 + scale_storage = torch.empty( + (self.num_local_experts, num_scales, slots_per_expert), dtype=torch.float32, device=device + ) + self._dispatch_scales = scale_storage.transpose(1, 2) return ( output_buffer, + self._dispatch_scales, self._dispatch_src_info, self._dispatch_layout_range, self._dispatch_count, @@ -246,7 +283,9 @@ class LowLatencyBackend: if output_buffer is None: raise ValueError("output_buffer is required for low-latency dispatch") if quant is not None: - raise NotImplementedError("low-latency quantization is not implemented yet") + raise NotImplementedError( + "per-call input quant metadata is not supported; configure dispatch output quantization on the communicator" + ) if input.dim() != 2 or not input.is_contiguous(): raise ValueError("input must be a contiguous [num_tokens, hidden_size] tensor") if input.device.type != "cuda" or input.dtype != torch.bfloat16: @@ -275,8 +314,9 @@ class LowLatencyBackend: expected_shape = (self.num_local_experts * slots_per_expert, self.hidden_size) if output_buffer.dim() != len(expected_shape) or not output_buffer.is_contiguous(): raise ValueError(f"output_buffer must be a contiguous {self.output_layout} tensor") - if output_buffer.device != input.device or output_buffer.dtype != torch.bfloat16: - raise ValueError("output_buffer must be a BF16 CUDA tensor on the same device as input") + expected_dtype = torch.float8_e4m3fn if self.dispatch_data_type == DispatchDataType.FP8_E4M3 else torch.bfloat16 + if output_buffer.device != input.device or output_buffer.dtype != expected_dtype: + raise ValueError(f"output_buffer must be a {expected_dtype} CUDA tensor on the same device as input") if tuple(output_buffer.shape) != expected_shape: raise ValueError(f"output_buffer shape must be {expected_shape}") @@ -286,6 +326,10 @@ class LowLatencyBackend: context = handle.combine_context if context.num_experts != self.num_experts or context.hidden_size != self.hidden_size: raise ValueError("DispatchHandle does not belong to this MoECommunicator configuration") + output_quant = handle.output_info.quant + handle_data_type = DispatchDataType.BF16 if output_quant is None else output_quant.format + if handle_data_type != self.dispatch_data_type: + raise ValueError("DispatchHandle quantization does not match this MoECommunicator configuration") slots_per_expert = self.world_size * self.max_tokens_per_rank if handle.output_info.layout.kind == DispatchLayout.EXPERT_MAJOR: expected_shape = (self.num_local_experts, slots_per_expert, self.hidden_size) diff --git a/python/mscclpp/ep/types.py b/python/mscclpp/ep/types.py index f739216a..451a984b 100644 --- a/python/mscclpp/ep/types.py +++ b/python/mscclpp/ep/types.py @@ -10,19 +10,22 @@ from typing import Any, List, Optional, Union import torch import mscclpp -from ._cpp import CombineMode, DispatchLayout, MoEMode +from ._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode # Quantization metadata. @dataclass class QuantConfig: - """Quantization metadata associated with an activation tensor.""" + """Quantization metadata associated with an activation tensor. - dtype: Optional[torch.dtype] = None + Low-latency FP8 dispatch returns ``block_scales`` with the activation's + leading dimensions and a format-defined final scale dimension. + """ + + format: Optional[DispatchDataType] = None block_scales: Optional[torch.Tensor] = None global_scale: Optional[torch.Tensor] = None - block_size: Optional[int] = None # Communicator construction. diff --git a/src/ext/ep/bindings.cpp b/src/ext/ep/bindings.cpp index ae331e2d..0d0d3b60 100644 --- a/src/ext/ep/bindings.cpp +++ b/src/ext/ep/bindings.cpp @@ -80,6 +80,10 @@ NB_MODULE(mscclpp_ep_cpp, m) { nb::enum_(m, "CombineMode") .value("RANK_LOCAL_REDUCE", mscclpp::ep::low_latency::CombineMode::RANK_LOCAL_REDUCE) .value("DIRECT_SEND", mscclpp::ep::low_latency::CombineMode::DIRECT_SEND); + nb::enum_(m, "DispatchDataType") + .value("BF16", mscclpp::ep::low_latency::DispatchDataType::BF16) + .value("FP8_E4M3", mscclpp::ep::low_latency::DispatchDataType::FP8_E4M3) + .value("MXFP8_E4M3", mscclpp::ep::low_latency::DispatchDataType::MXFP8_E4M3); nb::class_(m, "MoERuntime") .def(nb::init(), nb::arg("comm"), @@ -95,34 +99,37 @@ NB_MODULE(mscclpp_ep_cpp, m) { .def( "dispatch", [](mscclpp::ep::MoERuntime& self, uintptr_t inputPtr, uintptr_t topkIdxPtr, uintptr_t topkWeightsPtr, - uintptr_t outputPtr, uintptr_t outputSrcInfoPtr, uintptr_t outputLayoutRangePtr, uintptr_t outputCountPtr, - int numTokens, int hidden, int numTopk, int maxTokensPerRank, int numExperts, int numBlocks, - uintptr_t streamPtr) { - self.dispatch(ptr(outputPtr), reinterpret_cast(ptr(outputSrcInfoPtr)), - reinterpret_cast(ptr(outputLayoutRangePtr)), - reinterpret_cast(ptr(outputCountPtr)), ptr(inputPtr), - reinterpret_cast(ptr(topkIdxPtr)), reinterpret_cast(ptr(topkWeightsPtr)), - numTokens, hidden, numTopk, maxTokensPerRank, numExperts, numBlocks, stream(streamPtr)); + uintptr_t outputPtr, uintptr_t outputScalesPtr, uintptr_t outputSrcInfoPtr, uintptr_t outputLayoutRangePtr, + uintptr_t outputCountPtr, int numTokens, int hidden, int numTopk, int maxTokensPerRank, int numExperts, + mscclpp::ep::low_latency::DispatchDataType dispatchDataType, int numBlocks, uintptr_t streamPtr) { + self.dispatch( + ptr(outputPtr), reinterpret_cast(ptr(outputScalesPtr)), + reinterpret_cast(ptr(outputSrcInfoPtr)), reinterpret_cast(ptr(outputLayoutRangePtr)), + reinterpret_cast(ptr(outputCountPtr)), ptr(inputPtr), reinterpret_cast(ptr(topkIdxPtr)), + reinterpret_cast(ptr(topkWeightsPtr)), numTokens, hidden, numTopk, maxTokensPerRank, numExperts, + dispatchDataType, numBlocks, stream(streamPtr)); }, nb::arg("input_ptr"), nb::arg("topk_idx_ptr"), nb::arg("topk_weights_ptr"), nb::arg("output_ptr"), - nb::arg("output_src_info_ptr"), nb::arg("output_layout_range_ptr"), nb::arg("output_count_ptr"), - nb::arg("num_tokens"), nb::arg("hidden"), nb::arg("num_topk"), nb::arg("max_tokens_per_rank"), - nb::arg("num_experts"), nb::arg("num_blocks"), nb::arg("stream_ptr")) + nb::arg("output_scales_ptr"), nb::arg("output_src_info_ptr"), nb::arg("output_layout_range_ptr"), + nb::arg("output_count_ptr"), nb::arg("num_tokens"), nb::arg("hidden"), nb::arg("num_topk"), + nb::arg("max_tokens_per_rank"), nb::arg("num_experts"), nb::arg("dispatch_data_type"), nb::arg("num_blocks"), + nb::arg("stream_ptr")) .def( "combine", [](mscclpp::ep::MoERuntime& self, uintptr_t expertOutputPtr, uintptr_t topkIdxPtr, uintptr_t topkWeightsPtr, uintptr_t srcInfoPtr, uintptr_t layoutRangePtr, uintptr_t outputPtr, int numTokens, int hidden, - int numTopk, int maxTokensPerRank, int numExperts, mscclpp::ep::low_latency::CombineMode mode, + int numTopk, int maxTokensPerRank, int numExperts, + mscclpp::ep::low_latency::DispatchDataType dispatchDataType, mscclpp::ep::low_latency::CombineMode mode, int numBlocks, uintptr_t streamPtr) { self.combine(ptr(outputPtr), ptr(expertOutputPtr), reinterpret_cast(ptr(topkIdxPtr)), reinterpret_cast(ptr(topkWeightsPtr)), reinterpret_cast(ptr(srcInfoPtr)), reinterpret_cast(ptr(layoutRangePtr)), numTokens, hidden, numTopk, maxTokensPerRank, - numExperts, mode, numBlocks, stream(streamPtr)); + numExperts, dispatchDataType, mode, numBlocks, stream(streamPtr)); }, nb::arg("expert_output_ptr"), nb::arg("topk_idx_ptr"), nb::arg("topk_weights_ptr"), nb::arg("src_info_ptr"), nb::arg("layout_range_ptr"), nb::arg("output_ptr"), nb::arg("num_tokens"), nb::arg("hidden"), - nb::arg("num_topk"), nb::arg("max_tokens_per_rank"), nb::arg("num_experts"), nb::arg("mode"), - nb::arg("num_blocks"), nb::arg("stream_ptr")); + nb::arg("num_topk"), nb::arg("max_tokens_per_rank"), nb::arg("num_experts"), nb::arg("dispatch_data_type"), + nb::arg("mode"), nb::arg("num_blocks"), nb::arg("stream_ptr")); // ========================================================================== // High-throughput (HT) DeepEP-style backend: Config + MoEHighThroughputRuntime. diff --git a/src/ext/ep/config.hpp b/src/ext/ep/config.hpp index 3419c4fb..e1e358a4 100644 --- a/src/ext/ep/config.hpp +++ b/src/ext/ep/config.hpp @@ -27,10 +27,13 @@ __host__ __device__ constexpr dtype_t configAlign(dtype_t a, dtype_t b) { namespace low_latency { +using Bf16 = typename mscclpp::bf16x2::ElementType; +using Fp8E4M3 = typename mscclpp::f8_e4m3x2::ElementType; + // Rank-deduplicated dispatch payload layout: // // [data: DataType[hidden]] -// [optional scales: ScaleType[hidden / scale_block_size]] +// [optional scales: ScaleType[hidden / format scale block size]] // [topKIndices: int[topK]] // [topKValues: float[topK]] // [srcTokenGlobalIdx: int] @@ -38,19 +41,15 @@ namespace low_latency { // The payload is 32-byte aligned as a whole. template struct PayloadView { - static constexpr bool kHasScales = !std::is_void_v; + static constexpr bool HasScales = !std::is_void_v; - int hidden_; int topK_; - int scaleBlockSize_; - int numScales_; - size_t hiddenBytes_; size_t scaleOffset_; size_t metadataOffset_; size_t numBytes_; MSCCLPP_HOST_DEVICE_INLINE static int numScales([[maybe_unused]] int hidden, [[maybe_unused]] int scaleBlockSize) { - if constexpr (kHasScales) { + if constexpr (HasScales) { return hidden / scaleBlockSize; } return 0; @@ -61,7 +60,7 @@ struct PayloadView { } MSCCLPP_HOST_DEVICE_INLINE static size_t scaleOffset(int hidden) { - if constexpr (kHasScales) { + if constexpr (HasScales) { return configAlign(hiddenBytes(hidden), alignof(ScaleType)); } return 0; @@ -69,14 +68,14 @@ struct PayloadView { MSCCLPP_HOST_DEVICE_INLINE static size_t scaleBytes([[maybe_unused]] int hidden, [[maybe_unused]] int scaleBlockSize) { - if constexpr (kHasScales) { + if constexpr (HasScales) { return static_cast(numScales(hidden, scaleBlockSize)) * sizeof(ScaleType); } return 0; } MSCCLPP_HOST_DEVICE_INLINE static size_t metadataOffset(int hidden, int scaleBlockSize) { - if constexpr (kHasScales) { + if constexpr (HasScales) { return configAlign(scaleOffset(hidden) + scaleBytes(hidden, scaleBlockSize), alignof(int)); } return configAlign(hiddenBytes(hidden), alignof(int)); @@ -90,12 +89,8 @@ struct PayloadView { return configAlign(metadataOffset(hidden, scaleBlockSize) + metadataBytes(topK), 32); } - MSCCLPP_HOST_DEVICE_INLINE PayloadView(int hidden, int topK, int scaleBlockSize = (kHasScales ? 128 : 0)) - : hidden_(hidden), - topK_(topK), - scaleBlockSize_(scaleBlockSize), - numScales_(numScales(hidden, scaleBlockSize)), - hiddenBytes_(hiddenBytes(hidden)), + MSCCLPP_HOST_DEVICE_INLINE PayloadView(int hidden, int topK, int scaleBlockSize = (HasScales ? 128 : 0)) + : topK_(topK), scaleOffset_(scaleOffset(hidden)), metadataOffset_(metadataOffset(hidden, scaleBlockSize)), numBytes_(numBytes(hidden, topK, scaleBlockSize)) {} @@ -106,12 +101,12 @@ struct PayloadView { } MSCCLPP_HOST_DEVICE_INLINE ScaleType* scaleFactors(void* base) const { - static_assert(kHasScales, "Payload has no scale factors"); + static_assert(HasScales, "Payload has no scale factors"); return reinterpret_cast(reinterpret_cast(base) + scaleOffset_); } MSCCLPP_HOST_DEVICE_INLINE const ScaleType* scaleFactors(const void* base) const { - static_assert(kHasScales, "Payload has no scale factors"); + static_assert(HasScales, "Payload has no scale factors"); return reinterpret_cast(reinterpret_cast(base) + scaleOffset_); } @@ -150,15 +145,16 @@ struct Layout { Buffer buffers_[2]; Layout(void* rdmaBuffer, int maxTokensPerRank, int hidden, int numRanks, int numExperts, int numTopk) { - const PayloadView<__bfloat16> bf16Payload(hidden, numTopk); - const PayloadView<__nv_fp8_storage_t, float> fp8Payload(hidden, numTopk, 128); + const PayloadView bf16Payload(hidden, numTopk); + const PayloadView fp8Payload128(hidden, numTopk, 128); + const PayloadView fp8Payload64(hidden, numTopk, 64); const size_t dispatchMetadataBytes = configAlign(static_cast(numRanks + numExperts) * sizeof(uint64_t), 128); const size_t dispatchPayloadStride = - configAlign(std::max(bf16Payload.numBytes_, fp8Payload.numBytes_), 128); + configAlign(std::max({bf16Payload.numBytes_, fp8Payload128.numBytes_, fp8Payload64.numBytes_}), 128); const size_t dispatchBufferBytes = dispatchMetadataBytes + static_cast(numRanks) * maxTokensPerRank * dispatchPayloadStride; - const size_t combineBufferBytes = static_cast(numExperts) * maxTokensPerRank * hidden * sizeof(__bfloat16); + const size_t combineBufferBytes = static_cast(numExperts) * maxTokensPerRank * hidden * sizeof(Bf16); const size_t bufferBytes = configAlign(std::max(dispatchBufferBytes, combineBufferBytes), 128); totalBytes_ = 2 * bufferBytes; diff --git a/src/ext/ep/include/api.cuh b/src/ext/ep/include/api.cuh index d6792962..0d8c18cb 100644 --- a/src/ext/ep/include/api.cuh +++ b/src/ext/ep/include/api.cuh @@ -204,6 +204,16 @@ enum class CombineMode { DIRECT_SEND }; +/// Dispatch payload data format. +enum class DispatchDataType { + /// Unquantized BF16 payload. + BF16, + /// FP8 E4M3 payload with one floating-point scale per 128 hidden elements. + FP8_E4M3, + /// Reserved for MXFP8 E4M3 payloads with micro-scales. + MXFP8_E4M3 +}; + /// Per-call low-latency workload dimensions. struct Workload { /// Number of local input or output tokens. @@ -216,6 +226,8 @@ struct Workload { int numExperts_; /// Maximum tokens per rank in the packed layout. int maxTokensPerRank_; + /// Dispatch payload data format. + DispatchDataType dispatchDataType_; }; /// Persistent communication resources shared by low-latency operations. @@ -247,6 +259,9 @@ size_t workspaceSize(int numRanks, int numExperts); /// Low-latency dispatch that distributes tokens to experts across ranks. /// @param[out] output Expert-major packed output /// [num_local_experts, num_ranks * max_tokens_per_rank, hidden]. +/// @param[out] outputScales FP8 block scales in +/// [num_local_experts, hidden / 128, num_ranks * max_tokens_per_rank], +/// or nullptr for BF16 dispatch. /// @param[out] outputSrcInfo Original source-token index for every packed expert row. /// @param[out] outputLayout Per-[local expert, source rank] packed count and offset. /// @param[out] outputCount Total packed token count for every local expert. @@ -259,9 +274,9 @@ size_t workspaceSize(int numRanks, int numExperts); /// @param[in,out] workspace Persistent counters, task storage, semaphores, and device barriers. /// @param[in] numBlocks Total dispatch grid size, including one scheduler and one metadata-notify block. /// @param[in] stream CUDA stream. -void dispatch(void* output, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, const void* input, - const int64_t* topkIdx, const float* topkWeights, const Workload& workload, void* recvBuffer, - const CommContext& comm, void* workspace, int numBlocks, cudaStream_t stream); +void dispatch(void* output, float* outputScales, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, + const void* input, const int64_t* topkIdx, const float* topkWeights, const Workload& workload, + void* recvBuffer, const CommContext& comm, void* workspace, int numBlocks, cudaStream_t stream); /// Low-latency combine that aggregates expert outputs back to tokens. /// @param[out] output Combined local tokens [num_tokens, hidden]. diff --git a/src/ext/ep/include/device_helpers.cuh b/src/ext/ep/include/device_helpers.cuh index 4784a3ff..2c7af810 100644 --- a/src/ext/ep/include/device_helpers.cuh +++ b/src/ext/ep/include/device_helpers.cuh @@ -66,6 +66,10 @@ __device__ __forceinline__ void memory_fence_gpu() { asm volatile("fence.acq_rel __device__ __forceinline__ void memory_fence_cta() { asm volatile("fence.acq_rel.cta;" ::: "memory"); } +__device__ __forceinline__ void syncNamedBarrier(int barrierId, int numThreads) { + asm volatile("bar.sync %0, %1;" ::"r"(barrierId), "r"(numThreads) : "memory"); +} + __device__ __forceinline__ void *peerBufferPtr(void *localBuffer, void *localBufferBase, void *peerBufferBase) { if (localBufferBase == nullptr) return peerBufferBase; const auto offset = reinterpret_cast(localBuffer) - reinterpret_cast(localBufferBase); diff --git a/src/ext/ep/include/quantization.cuh b/src/ext/ep/include/quantization.cuh new file mode 100644 index 00000000..1c1a72ca --- /dev/null +++ b/src/ext/ep/include/quantization.cuh @@ -0,0 +1,74 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +#pragma once + +#include + +#include "device_helpers.cuh" + +namespace mscclpp { +namespace ep { + +inline constexpr float Fp8E4M3MaxValue = 448.0f; + +template +MSCCLPP_DEVICE_INLINE float laneGroupMax(float value, int laneId) { + EP_STATIC_ASSERT(NumLanes > 0 && NumLanes <= WARP_SIZE, "Invalid lane group size"); + EP_STATIC_ASSERT((NumLanes & (NumLanes - 1)) == 0, "Lane group size must be a power of two"); + + unsigned int mask; + if constexpr (NumLanes == WARP_SIZE) { + mask = 0xffffffffu; + } else { + const int groupStart = laneId - laneId % NumLanes; + mask = ((1u << NumLanes) - 1u) << groupStart; + } + +#pragma unroll + for (int offset = NumLanes / 2; offset > 0; offset >>= 1) { + value = fmaxf(value, __shfl_xor_sync(mask, value, offset)); + } + return value; +} + +template +MSCCLPP_DEVICE_INLINE mscclpp::f8_e4m3x8 quantizeBf16x8ToFp8E4M3(const mscclpp::bf16x8& source, float* scaleOut, + int laneId) { + constexpr int NumElements = mscclpp::bf16x8::Size; + constexpr int NumLanesPerScale = NumElementsPerScale / NumElements; + constexpr float Margin = 1e-4f; + + EP_STATIC_ASSERT(NumElementsPerScale % NumElements == 0, "Invalid scale vectorization"); + EP_STATIC_ASSERT(NumLanesPerScale > 0 && NumLanesPerScale <= WARP_SIZE, "Invalid lanes per scale"); + EP_STATIC_ASSERT((NumLanesPerScale & (NumLanesPerScale - 1)) == 0, "Lanes per scale must be a power of two"); + + const mscclpp::f32x8 values = mscclpp::to(source); + float maxAbs = Margin; +#pragma unroll + for (int element = 0; element < NumElements; ++element) { + maxAbs = fmaxf(maxAbs, fabsf(values.data[element])); + } + + maxAbs = laneGroupMax(maxAbs, laneId); + const float quantScale = Fp8E4M3MaxValue / maxAbs; + if (laneId % NumLanesPerScale == 0) { + *scaleOut = maxAbs / Fp8E4M3MaxValue; + } + + mscclpp::f32x8 scaledValues; +#pragma unroll + for (int element = 0; element < NumElements; ++element) { + scaledValues.data[element] = values.data[element] * quantScale; + } + return mscclpp::to(scaledValues); +} + +MSCCLPP_DEVICE_INLINE float dequantizeFp8E4M3(typename mscclpp::f8_e4m3x2::ElementType value, float scale) { + mscclpp::f8_e4m3x2 packed; + packed.data[0] = value; + packed.data[1] = value; + return mscclpp::to(packed).data[0] * scale; +} + +} // namespace ep +} // namespace mscclpp diff --git a/src/ext/ep/legacy/low_latency.cu b/src/ext/ep/legacy/low_latency.cu index 9f52f78a..67296291 100644 --- a/src/ext/ep/legacy/low_latency.cu +++ b/src/ext/ep/legacy/low_latency.cu @@ -86,21 +86,19 @@ struct DispatchDTypeTraits {}; template <> struct DispatchDTypeTraits { - using Type = nv_bfloat16; + using Type = low_latency::Bf16; }; template <> struct DispatchDTypeTraits { - using Type = __nv_fp8_storage_t; + using Type = low_latency::Fp8E4M3; }; // Keep input and output dtype separate. Current launch path only instantiates // BF16 input, but this keeps room for pre-quantized input with explicit scales. template struct DispatchOutputVec { - using SourceType = typename DispatchDTypeTraits::Type; - using Type = - typename std::conditional::Type>::type; + using Type = typename std::conditional::type; }; template @@ -114,7 +112,6 @@ using LowLatencyDispatchPayloadView = template MSCCLPP_DEVICE_INLINE typename DispatchOutputVec::Type dispatchConvert( const int4& inputValue, float* scaleOut, int laneId) { - using SourceType = typename DispatchDTypeTraits::Type; using OutputVec = typename DispatchOutputVec::Type; if constexpr (kInputDType == kOutputDType) { @@ -122,7 +119,7 @@ MSCCLPP_DEVICE_INLINE typename DispatchOutputVec::Typ } else { static_assert(kInputDType == DType::BF16 && kOutputDType == DType::F8E4M3, "Unsupported low-latency dispatch dtype conversion"); - return static_cast(quantizeToFp8(inputValue, scaleOut, laneId)); + return quantizeBf16x8ToFp8E4M3(mscclpp::bit_cast(inputValue), scaleOut, laneId); } } @@ -639,7 +636,7 @@ MSCCLPP_DEVICE_INLINE void copyCombineInputToBf16(int4* dst, const uint8_t* src, const int elemIdx = i * kNumBf16PerInt4 + j; const int scaleIdx = elemIdx / 128; bf16Values[j] = - static_cast(dequantizeFp8<__NV_E4M3>(srcFp8[elemIdx], scales[scaleIdx * scaleStride])); + static_cast(dequantizeFp8E4M3(srcFp8[elemIdx], scales[scaleIdx * scaleStride])); } st_na_global(dst + i, bf16Pack); } diff --git a/src/ext/ep/legacy/quantization.cuh b/src/ext/ep/legacy/quantization.cuh deleted file mode 100644 index feb85880..00000000 --- a/src/ext/ep/legacy/quantization.cuh +++ /dev/null @@ -1,112 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - -#pragma once - -#include -#include - -#include "device_helpers.cuh" - -namespace mscclpp { -namespace ep { - -template <__nv_fp8_interpretation_t kFp8Format> -struct Fp8FormatTraits {}; - -template <> -struct Fp8FormatTraits<__NV_E4M3> { - static constexpr float maxValue = 448.0f; -}; - -template <> -struct Fp8FormatTraits<__NV_E5M2> { - static constexpr float maxValue = 57344.0f; -}; - -template -struct Fp8PackedVector { - using Type = typename VecInt::vec_t; -}; - -template <> -struct Fp8PackedVector<8> { - using Type = int2; -}; - -template -struct Fp8VectorType { - static_assert(sizeof(int4) % sizeof(SourceType) == 0, "Source type must divide int4"); - static constexpr int numElems = sizeof(int4) / sizeof(SourceType); - using Type = typename Fp8PackedVector::Type; -}; - -template -__device__ __forceinline__ float laneGroupReduceMax(float value, int laneId) { - EP_STATIC_ASSERT(kNumLanes > 0 && kNumLanes <= 32, "Invalid lane group size"); - EP_STATIC_ASSERT((kNumLanes & (kNumLanes - 1)) == 0, "Lane group size must be a power of two"); - - unsigned int mask; - if constexpr (kNumLanes == 32) { - mask = 0xffffffffu; - } else { - const int groupStart = laneId - laneId % kNumLanes; - mask = ((1u << kNumLanes) - 1u) << groupStart; - } - -#pragma unroll - for (int offset = kNumLanes / 2; offset > 0; offset >>= 1) { - value = fmaxf(value, __shfl_xor_sync(mask, value, offset)); - } - return value; -} - -template -__device__ __forceinline__ typename Fp8VectorType::Type quantizeToFp8(const int4& sourceValue, - float* scaleOut, int laneId) { - constexpr int kNumElemsPerRead = Fp8VectorType::numElems; - constexpr int kNumLanesPerScale = kNumElemsPerScale / kNumElemsPerRead; - constexpr float kFp8Margin = 1e-4f; - using OutputVec = typename Fp8VectorType::Type; - - EP_STATIC_ASSERT(kNumElemsPerRead % 2 == 0, "FP8 conversion packs pairs of elements"); - EP_STATIC_ASSERT(kNumElemsPerScale % kNumElemsPerRead == 0, "Invalid scale vectorization"); - EP_STATIC_ASSERT(kNumLanesPerScale > 0 && kNumLanesPerScale <= 32, "Invalid lanes per scale"); - EP_STATIC_ASSERT((kNumLanesPerScale & (kNumLanesPerScale - 1)) == 0, "Lanes per scale must be a power of two"); - - auto sourceValues = reinterpret_cast(&sourceValue); - float fp32Values[kNumElemsPerRead]; - float amax = kFp8Margin; - -#pragma unroll - for (int j = 0; j < kNumElemsPerRead; ++j) { - fp32Values[j] = static_cast(sourceValues[j]); - amax = fmaxf(amax, fabsf(fp32Values[j])); - } - - amax = laneGroupReduceMax(amax, laneId); - const float scale = Fp8FormatTraits::maxValue / amax; - const float scaleInv = amax / Fp8FormatTraits::maxValue; - - if (laneId % kNumLanesPerScale == 0) { - *scaleOut = scaleInv; - } - - OutputVec outputValue; - auto fp8x2Values = reinterpret_cast<__nv_fp8x2_storage_t*>(&outputValue); -#pragma unroll - for (int j = 0; j < kNumElemsPerRead; j += 2) { - float2 fp32x2 = {fp32Values[j] * scale, fp32Values[j + 1] * scale}; - fp8x2Values[j / 2] = __nv_cvt_float2_to_fp8x2(fp32x2, __NV_SATFINITE, kFp8Format); - } - - return outputValue; -} - -template <__nv_fp8_interpretation_t kFp8Format = __NV_E4M3> -__device__ __forceinline__ float dequantizeFp8(__nv_fp8_storage_t value, float scale) { - return static_cast(__half(__nv_cvt_fp8_to_halfraw(value, kFp8Format))) * scale; -} - -} // namespace ep -} // namespace mscclpp diff --git a/src/ext/ep/low_latency/combine.cu b/src/ext/ep/low_latency/combine.cu index dc186f27..f04e1a15 100644 --- a/src/ext/ep/low_latency/combine.cu +++ b/src/ext/ep/low_latency/combine.cu @@ -26,9 +26,9 @@ MSCCLPP_HOST_DEVICE_INLINE size_t combineControlBytes(int nLocalExperts) { template MSCCLPP_HOST_DEVICE_INLINE size_t combineSharedBytes(int nLocalExperts) { - constexpr size_t TileBytes = static_cast(Hidden) * sizeof(__bfloat16); + constexpr size_t TileBytes = static_cast(Hidden) * sizeof(Bf16); if constexpr (Mode == low_latency::CombineMode::DIRECT_SEND) { - constexpr int NWorkers = tmaWorkerCount(); + constexpr int NWorkers = tmaWorkerCount(); return combineControlBytes(nLocalExperts) + static_cast(NWorkers) * (TileBytes + sizeof(uint64_t)); } return combineControlBytes(nLocalExperts) + CombineNStages * TileBytes; @@ -91,7 +91,7 @@ MSCCLPP_DEVICE_INLINE int4 reduceRankPartialsBf16x8(const void* combineRecvBuffe return packedOutput; } -template