From 152f2ab02dd35e1f416df6b749a42dad0e50f73b Mon Sep 17 00:00:00 2001 From: Binyang Li Date: Mon, 13 Jul 2026 03:07:12 +0000 Subject: [PATCH] code optimization --- python/mscclpp/ep/__init__.py | 2 + python/mscclpp/ep/_cpp.py | 1 + python/mscclpp/ep/communicator.py | 3 +- python/mscclpp/ep/low_latency.py | 25 +- python/mscclpp/ep/types.py | 7 +- src/ext/ep/CMakeLists.txt | 2 + src/ext/ep/README.md | 2 + src/ext/ep/bindings.cpp | 20 +- src/ext/ep/config.hpp | 107 ++-- src/ext/ep/kernels/api.cuh | 26 + src/ext/ep/kernels/low_latency.cu | 13 +- src/ext/ep/kernels/utils.cuh | 71 +++ src/ext/ep/low_latency/combine.cu | 531 +++++++++++++++++++ src/ext/ep/low_latency/config.cuh | 134 +++++ src/ext/ep/low_latency/dispatch.cu | 525 ++++++++++++++++++ src/ext/ep/moe_runtime.cc | 94 +++- src/ext/ep/moe_runtime.hpp | 7 +- test/mscclpp-test/CMakeLists.txt | 26 + test/python/ep/test_low_latency_multirank.py | 53 +- 19 files changed, 1538 insertions(+), 111 deletions(-) create mode 100644 src/ext/ep/low_latency/combine.cu create mode 100644 src/ext/ep/low_latency/config.cuh create mode 100644 src/ext/ep/low_latency/dispatch.cu diff --git a/python/mscclpp/ep/__init__.py b/python/mscclpp/ep/__init__.py index 46b93b7b..9c112b37 100644 --- a/python/mscclpp/ep/__init__.py +++ b/python/mscclpp/ep/__init__.py @@ -24,6 +24,7 @@ from .communicator import ( # noqa: F401 MoECommunicatorConfig, MoEMode, OperationOverlapConfig, + OptimizedCombineMode, QuantConfig, RowMajorInternodeDispatchHandle, RowMajorInternodeCombineContext, @@ -46,6 +47,7 @@ __all__ = [ "MoECommunicatorConfig", "MoEMode", "OperationOverlapConfig", + "OptimizedCombineMode", "QuantConfig", "RowMajorInternodeDispatchHandle", "RowMajorInternodeCombineContext", diff --git a/python/mscclpp/ep/_cpp.py b/python/mscclpp/ep/_cpp.py index 310d0174..062bb602 100644 --- a/python/mscclpp/ep/_cpp.py +++ b/python/mscclpp/ep/_cpp.py @@ -14,6 +14,7 @@ except ImportError as exc: # pragma: no cover DispatchLayout = _cpp.DispatchLayout MoEMode = _cpp.MoEMode +OptimizedCombineMode = _cpp.OptimizedCombineMode Config = getattr(_cpp, "Config", None) diff --git a/python/mscclpp/ep/communicator.py b/python/mscclpp/ep/communicator.py index fdfa91da..55466682 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 DispatchLayout, MoEMode +from ._cpp import DispatchLayout, MoEMode, OptimizedCombineMode from .high_throughput import HighThroughputBackend from .low_latency import LowLatencyBackend from .types import ( @@ -44,6 +44,7 @@ __all__ = [ "MoECommunicator", "MoECommunicatorConfig", "MoEMode", + "OptimizedCombineMode", "OperationOverlapConfig", "QuantConfig", "RowMajorInternodeDispatchHandle", diff --git a/python/mscclpp/ep/low_latency.py b/python/mscclpp/ep/low_latency.py index e829a58d..1d1058d5 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 DispatchLayout, MoEMode, _cpp, get_low_latency_rdma_size_hint +from ._cpp import DispatchLayout, MoEMode, OptimizedCombineMode, _cpp, get_low_latency_rdma_size_hint from .types import ( DispatchHandle, DispatchLayoutInfo, @@ -25,7 +25,7 @@ from .utils import cuda_stream_ptr, requires_dequantization, resolve_expert_plac class LowLatencyRuntime: """Private low-level low-latency runtime wrapper (wraps ``_cpp.MoERuntime``).""" - num_sms: int = 20 + num_sms: int = 64 def __init__( self, @@ -90,13 +90,21 @@ class LowLatencyBackend: self.hidden_size = config.hidden_size self.topk = config.topk self.max_tokens_per_rank = config.max_tokens_per_rank - self.num_sms = config.num_sms + self.num_sms = config.low_latency_dispatch_num_sms + self.combine_num_sms = config.low_latency_combine_num_sms + self.combine_mode = config.low_latency_combine_mode self.enable_overlap = config.enable_overlap if self.output_layout != DispatchLayout.EXPERT_MAJOR: raise NotImplementedError("low-latency mode currently supports only DispatchLayout.EXPERT_MAJOR") if self.num_experts % self.world_size != 0: raise ValueError("low-latency mode requires num_experts divisible by world_size") + if not self.world_size <= self.num_sms <= 128: + raise ValueError("low_latency_dispatch_num_sms must be between world_size and 128") + if not 1 <= self.combine_num_sms <= 128: + raise ValueError("low_latency_combine_num_sms must be between 1 and 128") + if not isinstance(self.combine_mode, OptimizedCombineMode): + raise TypeError("low_latency_combine_mode must be an OptimizedCombineMode") self.num_local_experts, self.local_expert_start = resolve_expert_placement( num_experts=self.num_experts, @@ -158,14 +166,12 @@ class LowLatencyBackend: ) -> tuple[DispatchOutput, DispatchHandle]: del previous_handle self._validate_dispatch_inputs(input, topk_ids, weights, quant, output_buffer) - if weights is None: - weights = torch.ones(topk_ids.shape, dtype=torch.float32, device=topk_ids.device) out_buf, packed_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(), - weights.data_ptr(), + 0 if weights is None else weights.data_ptr(), out_buf.data_ptr(), 0 if packed_scales is None else packed_scales.data_ptr(), src_info.data_ptr(), @@ -178,6 +184,7 @@ class LowLatencyBackend: self.num_experts, self.dispatch_requires_quantization, self.output_layout, + self.num_sms, cuda_stream_ptr(stream), ) dispatched_quant = None @@ -231,16 +238,18 @@ class LowLatencyBackend: expert_output.data_ptr(), 0 if x_scales is None else x_scales.data_ptr(), context.topk_ids.data_ptr(), - context.weights.data_ptr(), + 0 if context.weights is None else context.weights.data_ptr(), context.src_info.data_ptr(), context.layout_range.data_ptr(), out.data_ptr(), context.num_tokens, self.hidden_size, - context.weights.size(1), + self.topk, context.num_max_dispatch_tokens_per_rank, context.num_experts, combine_requires_dequantization, + self.combine_mode, + self.combine_num_sms, cuda_stream_ptr(stream), ) return out diff --git a/python/mscclpp/ep/types.py b/python/mscclpp/ep/types.py index 2358c52d..4f8a7406 100644 --- a/python/mscclpp/ep/types.py +++ b/python/mscclpp/ep/types.py @@ -10,7 +10,7 @@ from typing import Any, List, Optional, Union import torch import mscclpp -from ._cpp import DispatchLayout, MoEMode +from ._cpp import DispatchLayout, MoEMode, OptimizedCombineMode # Quantization metadata. @@ -56,6 +56,9 @@ class MoECommunicatorConfig: # Transport / launch tuning num_rdma_qps_per_rank: int = 12 num_sms: int = 20 + low_latency_dispatch_num_sms: int = 64 + low_latency_combine_num_sms: int = 64 + low_latency_combine_mode: OptimizedCombineMode = OptimizedCombineMode.DISABLED enable_overlap: bool = False # HT-only buffer/launch tuning (advanced) @@ -103,7 +106,7 @@ class ExpertMajorCombineContext: """Combine context for expert-major dispatch output.""" topk_ids: torch.Tensor - weights: torch.Tensor + weights: Optional[torch.Tensor] num_experts: int num_tokens: int hidden_size: int diff --git a/src/ext/ep/CMakeLists.txt b/src/ext/ep/CMakeLists.txt index 04d26b7e..d0f44289 100644 --- a/src/ext/ep/CMakeLists.txt +++ b/src/ext/ep/CMakeLists.txt @@ -41,6 +41,8 @@ set(EP_SOURCES moe_runtime.cc bindings.cpp kernels/low_latency.cu + low_latency/dispatch.cu + low_latency/combine.cu # High-throughput (DeepEP-style) backend (torch-free, raw-pointer API). ht_runtime.cc ht/kernels/internode.cu diff --git a/src/ext/ep/README.md b/src/ext/ep/README.md index f264cf23..47d902dc 100644 --- a/src/ext/ep/README.md +++ b/src/ext/ep/README.md @@ -557,6 +557,8 @@ slots even when the non-hybrid path does not use them. After a GPU barrier guarantees data arrival, the copy epilogue copies data into `recv_x`, `recv_sf`, `recv_topk_idx`, `recv_topk_weights`, and source metadata. +The optimized BF16 low-latency dispatch/combine kernels are instantiated for +`hidden = 4096`, `7168`, `8192`, and `9216`; other hidden sizes are rejected. For BF16 dispatch with `num_tokens = 128`, `hidden = 7168`, and `num_experts = 256`, the expected per-source-rank dispatch payload is: diff --git a/src/ext/ep/bindings.cpp b/src/ext/ep/bindings.cpp index 69d79f58..cd2769d8 100644 --- a/src/ext/ep/bindings.cpp +++ b/src/ext/ep/bindings.cpp @@ -77,6 +77,11 @@ NB_MODULE(mscclpp_ep_cpp, m) { .value("EXPERT_MAJOR", mscclpp::ep::DispatchLayout::EXPERT_MAJOR) .value("FLAT", mscclpp::ep::DispatchLayout::FLAT); + nb::enum_(m, "OptimizedCombineMode") + .value("DISABLED", mscclpp::ep::low_latency::OptimizedCombineMode::DISABLED) + .value("RANK_LOCAL_REDUCE", mscclpp::ep::low_latency::OptimizedCombineMode::RANK_LOCAL_REDUCE) + .value("DIRECT_SEND", mscclpp::ep::low_latency::OptimizedCombineMode::DIRECT_SEND); + nb::class_(m, "MoERuntime") .def(nb::init(), nb::arg("comm"), nb::arg("num_nvl_bytes"), nb::arg("num_rdma_bytes"), nb::arg("mode")) @@ -93,35 +98,38 @@ NB_MODULE(mscclpp_ep_cpp, m) { [](mscclpp::ep::MoERuntime& self, uintptr_t inputPtr, uintptr_t topkIdxPtr, uintptr_t topkWeightsPtr, uintptr_t outputPtr, uintptr_t outputScalesPtr, uintptr_t outputSrcInfoPtr, uintptr_t outputLayoutRangePtr, uintptr_t outputCountPtr, int numTokens, int hidden, int numTopk, int numMaxDispatchTokensPerRank, - int numExperts, bool requiresQuantization, mscclpp::ep::DispatchLayout outputLayout, uintptr_t streamPtr) { + int numExperts, bool requiresQuantization, mscclpp::ep::DispatchLayout outputLayout, int numSms, + 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, numMaxDispatchTokensPerRank, - numExperts, requiresQuantization, outputLayout, stream(streamPtr)); + numExperts, requiresQuantization, outputLayout, numSms, stream(streamPtr)); }, nb::arg("input_ptr"), nb::arg("topk_idx_ptr"), nb::arg("topk_weights_ptr"), nb::arg("output_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("num_max_dispatch_tokens_per_rank"), nb::arg("num_experts"), nb::arg("requires_quantization"), - nb::arg("output_layout"), nb::arg("stream_ptr")) + nb::arg("output_layout"), nb::arg("num_sms"), nb::arg("stream_ptr")) .def( "combine", [](mscclpp::ep::MoERuntime& self, uintptr_t expertOutputPtr, uintptr_t expertScalesPtr, uintptr_t topkIdxPtr, uintptr_t topkWeightsPtr, uintptr_t srcInfoPtr, uintptr_t layoutRangePtr, uintptr_t outputPtr, int numTokens, int hidden, int numTopk, int numMaxDispatchTokensPerRank, int numExperts, - bool requiresDequantization, uintptr_t streamPtr) { + bool requiresDequantization, mscclpp::ep::low_latency::OptimizedCombineMode optimizedMode, int numBlocks, + uintptr_t streamPtr) { self.combine(ptr(outputPtr), ptr(expertOutputPtr), reinterpret_cast(ptr(expertScalesPtr)), reinterpret_cast(ptr(topkIdxPtr)), reinterpret_cast(ptr(topkWeightsPtr)), reinterpret_cast(ptr(srcInfoPtr)), reinterpret_cast(ptr(layoutRangePtr)), numTokens, hidden, numTopk, numMaxDispatchTokensPerRank, numExperts, requiresDequantization, - stream(streamPtr)); + optimizedMode, numBlocks, stream(streamPtr)); }, nb::arg("expert_output_ptr"), nb::arg("expert_scales_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("num_max_dispatch_tokens_per_rank"), - nb::arg("num_experts"), nb::arg("requires_dequantization"), nb::arg("stream_ptr")); + nb::arg("num_experts"), nb::arg("requires_dequantization"), nb::arg("optimized_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 af61128e..c631458b 100644 --- a/src/ext/ep/config.hpp +++ b/src/ext/ep/config.hpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include "kernels/configs.cuh" @@ -34,10 +35,19 @@ __host__ __device__ constexpr dtype_t configAlign(dtype_t a, dtype_t b) { // The payload is 32-byte aligned as a whole. ScaleType=void means the payload is // not quantized and has no scale section. template -struct LowLatencyPackedPayloadFormat { +struct LowLatencyPayloadView { static constexpr bool kHasScales = !std::is_void_v; - __host__ __device__ static constexpr int numScales([[maybe_unused]] int hidden, [[maybe_unused]] int scaleBlockSize) { + 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 nScales([[maybe_unused]] int hidden, [[maybe_unused]] int scaleBlockSize) { if constexpr (kHasScales) { return hidden / scaleBlockSize; } else { @@ -45,11 +55,11 @@ struct LowLatencyPackedPayloadFormat { } } - __host__ __device__ static constexpr size_t hiddenBytes(int hidden) { + MSCCLPP_HOST_DEVICE_INLINE static size_t hiddenBytes(int hidden) { return static_cast(hidden) * sizeof(DataType); } - __host__ __device__ static constexpr size_t scaleOffset(int hidden) { + MSCCLPP_HOST_DEVICE_INLINE static size_t scaleOffset(int hidden) { if constexpr (kHasScales) { return configAlign(hiddenBytes(hidden), alignof(ScaleType)); } else { @@ -57,16 +67,16 @@ struct LowLatencyPackedPayloadFormat { } } - __host__ __device__ static constexpr size_t scaleBytes([[maybe_unused]] int hidden, - [[maybe_unused]] int scaleBlockSize) { + MSCCLPP_HOST_DEVICE_INLINE static size_t scaleBytes([[maybe_unused]] int hidden, + [[maybe_unused]] int scaleBlockSize) { if constexpr (kHasScales) { - return static_cast(numScales(hidden, scaleBlockSize)) * sizeof(ScaleType); + return static_cast(nScales(hidden, scaleBlockSize)) * sizeof(ScaleType); } else { return 0; } } - __host__ __device__ static constexpr size_t metadataOffset(int hidden, [[maybe_unused]] int scaleBlockSize) { + MSCCLPP_HOST_DEVICE_INLINE static size_t metadataOffset(int hidden, [[maybe_unused]] int scaleBlockSize) { if constexpr (kHasScales) { return configAlign(scaleOffset(hidden) + scaleBytes(hidden, scaleBlockSize), alignof(int)); } else { @@ -74,73 +84,62 @@ struct LowLatencyPackedPayloadFormat { } } - __host__ __device__ static constexpr size_t metadataBytes(int topK) { + MSCCLPP_HOST_DEVICE_INLINE static size_t metadataBytes(int topK) { return static_cast(topK) * sizeof(int) + static_cast(topK) * sizeof(float) + sizeof(int); } - __host__ __device__ static constexpr size_t numBytes(int hidden, int topK, int scaleBlockSize) { + MSCCLPP_HOST_DEVICE_INLINE static size_t numBytes(int hidden, int topK, int scaleBlockSize) { return configAlign(metadataOffset(hidden, scaleBlockSize) + metadataBytes(topK), 32); } -}; -template -struct LowLatencyPackedPayloadView { - using Format = LowLatencyPackedPayloadFormat; - - int hidden; - int topK; - int scaleBlockSize; - int numScales; - size_t hiddenBytes; - size_t scaleOffset; - size_t metadataOffset; - size_t numBytes; - - __host__ __device__ __forceinline__ LowLatencyPackedPayloadView(int hidden, int topK, - int scaleBlockSize = (Format::kHasScales ? 128 : 0)) + MSCCLPP_HOST_DEVICE_INLINE LowLatencyPayloadView(int hidden, int topK, int scaleBlockSize = (kHasScales ? 128 : 0)) : hidden(hidden), topK(topK), scaleBlockSize(scaleBlockSize), - numScales(Format::numScales(hidden, scaleBlockSize)), - hiddenBytes(Format::hiddenBytes(hidden)), - scaleOffset(Format::scaleOffset(hidden)), - metadataOffset(Format::metadataOffset(hidden, scaleBlockSize)), - numBytes(Format::numBytes(hidden, topK, scaleBlockSize)) {} + numScales(nScales(hidden, scaleBlockSize)), + hiddenBytes_(hiddenBytes(hidden)), + scaleOffset_(scaleOffset(hidden)), + metadataOffset_(metadataOffset(hidden, scaleBlockSize)), + numBytes_(numBytes(hidden, topK, scaleBlockSize)) {} template - __device__ __forceinline__ T* data(void* base) const { + MSCCLPP_HOST_DEVICE_INLINE T* data(void* base) const { return reinterpret_cast(base); } - __device__ __forceinline__ ScaleType* scaleFactors(void* base) const { - static_assert(Format::kHasScales, "Payload has no scale factors"); - return reinterpret_cast(reinterpret_cast(base) + scaleOffset); + MSCCLPP_HOST_DEVICE_INLINE ScaleType* scaleFactors(void* base) const { + static_assert(kHasScales, "Payload has no scale factors"); + return reinterpret_cast(reinterpret_cast(base) + scaleOffset_); } - __device__ __forceinline__ const ScaleType* scaleFactors(const void* base) const { - static_assert(Format::kHasScales, "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"); + return reinterpret_cast(reinterpret_cast(base) + scaleOffset_); } - __device__ __forceinline__ int* topKIndices(void* base) const { - return reinterpret_cast(reinterpret_cast(base) + metadataOffset); + MSCCLPP_HOST_DEVICE_INLINE int* topKIndices(void* base) const { + return reinterpret_cast(reinterpret_cast(base) + metadataOffset_); } - __device__ __forceinline__ const int* topKIndices(const void* base) const { - return reinterpret_cast(reinterpret_cast(base) + metadataOffset); + MSCCLPP_HOST_DEVICE_INLINE const int* topKIndices(const void* base) const { + return reinterpret_cast(reinterpret_cast(base) + metadataOffset_); } - __device__ __forceinline__ float* topKValues(void* base) const { + MSCCLPP_HOST_DEVICE_INLINE float* topKValues(void* base) const { return reinterpret_cast(topKIndices(base) + topK); } - __device__ __forceinline__ const float* topKValues(const void* base) const { + MSCCLPP_HOST_DEVICE_INLINE const float* topKValues(const void* base) const { return reinterpret_cast(topKIndices(base) + topK); } - __device__ __forceinline__ int* srcTokenGlobalIdx(void* base) const { + MSCCLPP_HOST_DEVICE_INLINE int* srcTokenGlobalIdx(void* base) const { return reinterpret_cast(topKValues(base) + topK); } + + MSCCLPP_HOST_DEVICE_INLINE const int* srcTokenGlobalIdx(const void* base) const { + return reinterpret_cast(topKValues(base) + topK); + } }; struct LowLatencyBuffer { @@ -185,13 +184,18 @@ struct LowLatencyLayout { // Message sizes // NOTES: you should add a control `int4` for combine messages if you want to do data transformation - const LowLatencyPackedPayloadView bf16DispatchPayload(hidden, numTopk); - const LowLatencyPackedPayloadView<__nv_fp8_storage_t, float> fp8DispatchPayload(hidden, numTopk, 128); - size_t numBytesPerDispatchMsg = std::max(bf16DispatchPayload.numBytes, fp8DispatchPayload.numBytes); + const LowLatencyPayloadView bf16DispatchPayload(hidden, numTopk); + const LowLatencyPayloadView<__nv_fp8_storage_t, float> fp8DispatchPayload(hidden, numTopk, 128); + size_t numBytesPerDispatchMsg = std::max(bf16DispatchPayload.numBytes_, fp8DispatchPayload.numBytes_); + const size_t optDispatchMetadataBytes = + configAlign(static_cast(numRanks + numExperts) * sizeof(uint64_t), 128); + const size_t optDispatchPayloadStride = configAlign(bf16DispatchPayload.numBytes_, 128); size_t numBytesPerCombineMsg = hidden * sizeof(nv_bfloat16); // Send buffer - size_t dispatchSendBufferBytes = numMaxDispatchTokensPerRank * numBytesPerDispatchMsg; + size_t dispatchSendBufferBytes = std::max( + static_cast(numMaxDispatchTokensPerRank) * numBytesPerDispatchMsg, + optDispatchMetadataBytes + static_cast(numMaxDispatchTokensPerRank) * optDispatchPayloadStride); size_t combineSendBufferBytes = numExperts * numMaxDispatchTokensPerRank * numBytesPerCombineMsg; size_t sendBufferBytes = std::max(dispatchSendBufferBytes, combineSendBufferBytes); EP_HOST_ASSERT(sendBufferBytes % sizeof(int4) == 0); @@ -199,7 +203,10 @@ struct LowLatencyLayout { // Symmetric receive buffers // TODO: optimize memory usages - size_t dispatchRecvDataBufferBytes = numRanks * numMaxDispatchTokensPerRank * numBytesPerDispatchMsg; + size_t dispatchRecvDataBufferBytes = + std::max(static_cast(numRanks) * numMaxDispatchTokensPerRank * numBytesPerDispatchMsg, + optDispatchMetadataBytes + + static_cast(numRanks) * numMaxDispatchTokensPerRank * optDispatchPayloadStride); size_t combineRecvBufferBytes = numExperts * numMaxDispatchTokensPerRank * numBytesPerCombineMsg; size_t recvBufferBytes = std::max(dispatchRecvDataBufferBytes, combineRecvBufferBytes); EP_HOST_ASSERT(recvBufferBytes % sizeof(int4) == 0); diff --git a/src/ext/ep/kernels/api.cuh b/src/ext/ep/kernels/api.cuh index 17168602..81f049ba 100644 --- a/src/ext/ep/kernels/api.cuh +++ b/src/ext/ep/kernels/api.cuh @@ -196,6 +196,16 @@ enum class DType { F8E4M3 }; +/// Optimized low-latency combine algorithm. +enum class OptimizedCombineMode { + /// Use the original combine implementation. + DISABLED, + /// Reduce expert rows on each destination rank before sending one partial per rank and token. + RANK_LOCAL_REDUCE, + /// Send every expert row directly and perform the full weighted reduction on the source rank. + DIRECT_SEND +}; + /// Transport context that encapsulates all transport-related state. struct TransportContext { /// Base address of the locally-registered RDMA buffer. @@ -204,10 +214,14 @@ struct TransportContext { mscclpp::PortChannelDeviceHandle* portChannels_; /// Base memory channel handles for IPC transport barrier (nullable). mscclpp::BaseMemoryChannelDeviceHandle* memoryChannels_; + /// Number of expert channel slots reserved per peer rank. + int memoryChannelStride_; /// Peer-mapped base addresses for IPC path (nullable). void* const* peerBases_; /// True if IPC path is ready, false to use RDMA. bool ipcReady_; + /// CUDA device ID associated with this transport. + int deviceId_; /// Current rank ID. int rank_; /// Total number of ranks. @@ -315,6 +329,18 @@ void dispatch(void* output, float* outputScales, int* outputSrcInfo, int64_t* ou const BufferSet& currentBuffer, const BufferSet& nextBuffer, const TransportContext& transport, void* workspace, cudaStream_t stream, Phase phase = SEND_AND_RECV); +/// Optimized BF16 IPC dispatch compatible with the expert-major combine metadata contract. +void dispatchOptimized(void* output, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, const void* input, + const int64_t* topkIdx, const float* topkWeights, const DispatchConfig& config, + const BufferSet& currentBuffer, const TransportContext& transport, void* workspace, int maxSms, + cudaStream_t stream); + +/// Experimental BF16 IPC combine. +void combineOptimized(void* output, const void* input, const int64_t* topkIdx, const float* topkWeights, + const int* srcInfo, const int64_t* layoutRange, const CombineConfig& config, + const BufferSet& currentBuffer, void* dispatchRecvBuffer, const TransportContext& transport, + void* workspace, int numBlocks, OptimizedCombineMode mode, cudaStream_t stream); + /// Low-latency combine kernel that aggregates expert outputs back to tokens. /// @param output Combined output [num_combined_tokens, hidden]. /// @param input Expert outputs [num_local_experts * num_ranks * max_tokens, hidden]. diff --git a/src/ext/ep/kernels/low_latency.cu b/src/ext/ep/kernels/low_latency.cu index 8d970016..55a6c82e 100644 --- a/src/ext/ep/kernels/low_latency.cu +++ b/src/ext/ep/kernels/low_latency.cu @@ -109,8 +109,8 @@ constexpr bool kDispatchNeedsScales = kInputDType != kOutputDType; template using LowLatencyDispatchPayloadView = - LowLatencyPackedPayloadView::Type, - std::conditional_t>; + LowLatencyPayloadView::Type, + std::conditional_t>; template MSCCLPP_DEVICE_INLINE typename DispatchOutputVec::Type dispatchConvert( @@ -250,9 +250,8 @@ MSCCLPP_DEVICE_INLINE void dispatchSend(int* sharedNumTokensSentPerRank, int* ou int* atomicFinishCounterPerRank, int64_t* nextClean, int numNextCleanInt, int numTokens, int numMaxDispatchTokensPerRank, int numTopk, int hidden, int numExperts, int rank, int numRanks, void* rdmaBufferPtr, - mscclpp::PortChannelDeviceHandle* portChannelHandles, + [[maybe_unused]] PortChannelDeviceHandle* portChannelHandles, void* const* peerRdmaBases, int ranksPerIpcDomain) { - (void)portChannelHandles; const auto smId = static_cast(blockIdx.x); const auto threadId = static_cast(threadIdx.x); const auto warpId = threadId / WARP_SIZE; @@ -268,7 +267,7 @@ MSCCLPP_DEVICE_INLINE void dispatchSend(int* sharedNumTokensSentPerRank, int* ou constexpr int kNumPerChannels = 128; using VecType = typename DispatchOutputVec::Type; const LowLatencyDispatchPayloadView payload(hidden, numTopk); - const size_t numBytesPerMsg = payload.numBytes; + const size_t numBytesPerMsg = payload.numBytes_; const size_t numInt4PerMsg = numBytesPerMsg / sizeof(int4); if (warpId < numWarps - 1) { @@ -444,7 +443,7 @@ __global__ __launch_bounds__(kNumWarpGroups* kNumWarpsPerGroup* WARP_SIZE, 1) vo // Message package: hidden data, optional quantization scales, index at source const LowLatencyDispatchPayloadView payload(hidden, numTopk); - const size_t numBytesPerMsg = payload.numBytes; + const size_t numBytesPerMsg = payload.numBytes_; EP_DEVICE_ASSERT(numBytesPerMsg % sizeof(int4) == 0); if (phases & LOW_LATENCY_SEND_PHASE) { @@ -776,7 +775,7 @@ MSCCLPP_DEVICE_INLINE void combineRecv(void* output, void* stagedRecv, int64_t* float regTopkWeights[kNumMaxTopk]; for (int i = 0; i < numTopk; ++i) { regTopkIdx[i] = static_cast(__ldg(topkIdx + tokenIdx * numTopk + i)); - regTopkWeights[i] = __ldg(topkWeights + tokenIdx * numTopk + i); + regTopkWeights[i] = topkWeights == nullptr ? 1.0f : __ldg(topkWeights + tokenIdx * numTopk + i); } float combinedValues[kNumBf16PerInt4] = {0.0f}; diff --git a/src/ext/ep/kernels/utils.cuh b/src/ext/ep/kernels/utils.cuh index 08b6493c..558441f5 100644 --- a/src/ext/ep/kernels/utils.cuh +++ b/src/ext/ep/kernels/utils.cuh @@ -65,6 +65,63 @@ __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 *peerBufferPtr(void *localBuffer, void *localBufferBase, void *peerBufferBase) { + if (localBufferBase == nullptr) return peerBufferBase; + const auto offset = reinterpret_cast(localBuffer) - reinterpret_cast(localBufferBase); + return reinterpret_cast(peerBufferBase) + offset; +} + +#if defined(__CUDACC__) +__device__ __forceinline__ void initTmaBarrier(uint64_t *sharedBarrier) { + const uint32_t barrierAddress = static_cast(__cvta_generic_to_shared(sharedBarrier)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(barrierAddress)); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + +__device__ __forceinline__ void issueTmaG2S(const void *source, void *sharedTile, uint64_t *sharedBarrier, + uint32_t nBytes) { + const uint32_t tileAddress = static_cast(__cvta_generic_to_shared(sharedTile)); + const uint32_t barrierAddress = static_cast(__cvta_generic_to_shared(sharedBarrier)); + asm volatile( + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes " + "[%0], [%1], %2, [%3];" ::"r"(tileAddress), + "l"(source), "r"(nBytes), "r"(barrierAddress) + : "memory"); + [[maybe_unused]] uint64_t state; + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 %0, [%1], %2;" + : "=l"(state) + : "r"(barrierAddress), "r"(nBytes)); +} + +__device__ __forceinline__ void waitTmaG2S(uint64_t *sharedBarrier, uint32_t &phase) { + const uint32_t barrierAddress = static_cast(__cvta_generic_to_shared(sharedBarrier)); + uint32_t done = 0; + while (!done) { + asm volatile( + "{ .reg .pred p; mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;" + " selp.u32 %0, 1, 0, p; }" + : "=r"(done) + : "r"(barrierAddress), "r"(phase)); + } + phase ^= 1; +} + +__device__ __forceinline__ void issueTmaS2G(void *destination, void *sharedTile, uint32_t nBytes) { + const uint32_t tileAddress = static_cast(__cvta_generic_to_shared(sharedTile)); + asm volatile("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" ::"l"(destination), "r"(tileAddress), + "r"(nBytes) + : "memory"); + asm volatile("cp.async.bulk.commit_group;"); +} + +template +__device__ __forceinline__ void waitTmaS2GRead() { + asm volatile("cp.async.bulk.wait_group.read %0;" ::"n"(kNumPendingGroups) : "memory"); +} + +__device__ __forceinline__ void waitTmaS2G() { asm volatile("cp.async.bulk.wait_group 0;" ::: "memory"); } +#endif + __device__ __forceinline__ void st_relaxed_sys_global(const int *ptr, int val) { asm volatile("st.relaxed.sys.global.s32 [%0], %1;" ::"l"(ptr), "r"(val) : "memory"); } @@ -351,6 +408,20 @@ __forceinline__ __device__ int warp_reduce_sum(int value) { return value; } +__forceinline__ __device__ int warpInclusiveSum(int value, int laneId) { +#pragma unroll + for (int offset = 1; offset < WARP_SIZE; offset *= 2) { + const int previous = __shfl_up_sync(0xffffffff, value, offset); + if (laneId >= offset) value += previous; + } + return value; +} + +__forceinline__ __device__ bool isFirstLaneForRank(int rank, int laneId) { + const unsigned matchMask = __match_any_sync(0xffffffff, rank); + return (__ffs(matchMask) - 1) == laneId; +} + __forceinline__ __device__ float half_warp_reduce_max(float value) { auto mask = __activemask(); // The mask be in `{0xffffffff, 0xffff}` diff --git a/src/ext/ep/low_latency/combine.cu b/src/ext/ep/low_latency/combine.cu new file mode 100644 index 00000000..8d642681 --- /dev/null +++ b/src/ext/ep/low_latency/combine.cu @@ -0,0 +1,531 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#include "../kernels/api.cuh" +#include "../kernels/exception.cuh" +#include "../kernels/utils.cuh" +#include "config.cuh" + +namespace mscclpp { +namespace ep { +namespace low_latency_opt { + +constexpr int kCombineNWarps = 32; +constexpr int kCombineNThreads = kCombineNWarps * WARP_SIZE; +constexpr int kCombineNStages = 8; +constexpr int kDirectSendMaxNWorkers = WARP_SIZE; +constexpr int kCombineMaxNBlocks = 128; +constexpr int kCombineMaxNTopk = 9; + +MSCCLPP_HOST_DEVICE_INLINE size_t combineControlBytes(int nLocalExperts) { + const size_t directControlBytes = static_cast(nLocalExperts + 1) * sizeof(int); + return configAlign(directControlBytes > sizeof(RecvTask) ? directControlBytes : sizeof(RecvTask), 128); +} + +template +MSCCLPP_HOST_DEVICE_INLINE size_t combineSharedBytes(int nLocalExperts) { + constexpr size_t kTileBytes = static_cast(kHidden) * sizeof(nv_bfloat16); + if constexpr (kMode == low_latency::OptimizedCombineMode::DIRECT_SEND) { + constexpr int kNWorkers = tmaWorkerCount(); + return combineControlBytes(nLocalExperts) + static_cast(kNWorkers) * (kTileBytes + sizeof(uint64_t)); + } + return combineControlBytes(nLocalExperts) + kCombineNStages * kTileBytes; +} + +template +MSCCLPP_DEVICE_INLINE int4 reduceWeightedBf16x8(const void* expertOutput, int rowOffset, float weight, int nTopk, + int hiddenIdx) { + constexpr int kBf16PairsPerInt4 = sizeof(int4) / sizeof(nv_bfloat162); + float2 reduced[kBf16PairsPerInt4] = {}; + for (int topkLane = 0; topkLane < nTopk; ++topkLane) { + const int sourceRowOffset = __shfl_sync(0xffffffff, rowOffset, topkLane); + if (sourceRowOffset < 0) continue; + const float sourceWeight = __shfl_sync(0xffffffff, weight, topkLane); + const int4 packed = ld_nc_global(reinterpret_cast(expertOutput) + + static_cast(sourceRowOffset) * kHiddenInt4 + hiddenIdx); + const auto* values = reinterpret_cast(&packed); +#pragma unroll + for (int pairIdx = 0; pairIdx < kBf16PairsPerInt4; ++pairIdx) { + const float2 value = __bfloat1622float2(values[pairIdx]); + reduced[pairIdx].x = fmaf(value.x, sourceWeight, reduced[pairIdx].x); + reduced[pairIdx].y = fmaf(value.y, sourceWeight, reduced[pairIdx].y); + } + } + + int4 packedOutput; + auto* outputValues = reinterpret_cast(&packedOutput); +#pragma unroll + for (int pairIdx = 0; pairIdx < kBf16PairsPerInt4; ++pairIdx) { + outputValues[pairIdx] = __float22bfloat162_rn(reduced[pairIdx]); + } + return packedOutput; +} + +template +MSCCLPP_DEVICE_INLINE int4 reduceRankPartialsBf16x8(const void* combineRecvBuffer, int partialRankCandidate, int nTopk, + int maxTokensPerRank, int tokenIdx, int hiddenIdx) { + constexpr int kBf16PairsPerInt4 = sizeof(int4) / sizeof(nv_bfloat162); + float2 reduced[kBf16PairsPerInt4] = {}; + for (int topkLane = 0; topkLane < nTopk; ++topkLane) { + const int partialRank = __shfl_sync(0xffffffff, partialRankCandidate, topkLane); + if (partialRank < 0) continue; + const int4 packed = + ld_nc_global(reinterpret_cast(combineRecvBuffer) + + (static_cast(partialRank) * maxTokensPerRank + tokenIdx) * kHiddenInt4 + hiddenIdx); + const auto* values = reinterpret_cast(&packed); +#pragma unroll + for (int pairIdx = 0; pairIdx < kBf16PairsPerInt4; ++pairIdx) { + const float2 value = __bfloat1622float2(values[pairIdx]); + reduced[pairIdx].x += value.x; + reduced[pairIdx].y += value.y; + } + } + + int4 packedOutput; + auto* outputValues = reinterpret_cast(&packedOutput); +#pragma unroll + for (int pairIdx = 0; pairIdx < kBf16PairsPerInt4; ++pairIdx) { + outputValues[pairIdx] = __float22bfloat162_rn(reduced[pairIdx]); + } + return packedOutput; +} + +template +MSCCLPP_DEVICE_INLINE void sendRankLocalPartials(const void* expertOutput, int nExperts, int rank, int nRanks, + int nTopk, int maxTokensPerRank, void* combineRecvBuffer, + const void* dispatchRecvBuffer, void* rdmaBufferBase, + void* const* peerRecvBuffers, DispatchWorkspaceView& workspaceView, + uint8_t* sharedMemory) { +#if defined(__CUDA_ARCH__) + static_assert(__CUDA_ARCH__ >= 900, "TMA combine send requires SM90 or newer"); +#endif + const int threadId = static_cast(threadIdx.x); + const int laneId = get_lane_id(); + const int nLocalExperts = nExperts / nRanks; + [[maybe_unused]] const int nExpertOutputRows = nLocalExperts * nRanks * maxTokensPerRank; + constexpr size_t kHiddenBytes = static_cast(kHidden) * sizeof(nv_bfloat16); + constexpr int kHiddenInt4 = kHiddenBytes / sizeof(int4); + constexpr int kChunksPerThread = (kHiddenInt4 + kCombineNThreads - 1) / kCombineNThreads; + static_assert(kHiddenInt4 % WARP_SIZE == 0); + const size_t dispatchMetadataSize = dispatchMetadataBytes(nRanks, nExperts); + const size_t payloadStride = dispatchPayloadStride(kHidden, nTopk); + const LowLatencyPayloadView payloadView(kHidden, nTopk); + auto* recvTask = reinterpret_cast(sharedMemory); + auto* outputTiles = sharedMemory + combineControlBytes(nLocalExperts); + + int tokenIteration = 0; + for (int taskIdx = static_cast(blockIdx.x); taskIdx < *workspaceView.nRecvTasks_; + taskIdx += static_cast(gridDim.x)) { + if (threadId == 0) *recvTask = workspaceView.recvTasks_[taskIdx]; + __syncthreads(); + const int sourceRank = recvTask->sourceRank_; + + for (int sourceTokenSlot = recvTask->tokenBegin_; sourceTokenSlot < recvTask->tokenEnd_; + ++sourceTokenSlot, ++tokenIteration) { + const int stage = tokenIteration % kCombineNStages; + auto* outputTile = reinterpret_cast(outputTiles + static_cast(stage) * kHiddenBytes); + const auto* sourcePayload = + reinterpret_cast(dispatchRecvBuffer) + dispatchMetadataSize + + (static_cast(sourceRank) * maxTokensPerRank + sourceTokenSlot) * payloadStride; + const int rowOffset = laneId < nTopk ? ld_nc_global(payloadView.topKIndices(sourcePayload) + laneId) : -1; + const float weight = laneId < nTopk ? ld_nc_global(payloadView.topKValues(sourcePayload) + laneId) : 0.0f; + if (rowOffset >= 0) EP_DEVICE_ASSERT(rowOffset < nExpertOutputRows); + + int4 reduced[kChunksPerThread] = {}; +#pragma unroll + for (int chunkIdx = 0; chunkIdx < kChunksPerThread; ++chunkIdx) { + const int hiddenIdx = threadId + chunkIdx * kCombineNThreads; + if (hiddenIdx < kHiddenInt4) { + reduced[chunkIdx] = reduceWeightedBf16x8(expertOutput, rowOffset, weight, nTopk, hiddenIdx); + } + } + + if (tokenIteration >= kCombineNStages && threadId == 0) { + waitTmaS2GRead(); + } + if (tokenIteration >= kCombineNStages) __syncthreads(); +#pragma unroll + for (int chunkIdx = 0; chunkIdx < kChunksPerThread; ++chunkIdx) { + const int hiddenIdx = threadId + chunkIdx * kCombineNThreads; + if (hiddenIdx < kHiddenInt4) outputTile[hiddenIdx] = reduced[chunkIdx]; + } + __syncthreads(); + + if (threadId == 0) { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + const int sourceTokenIdx = + ld_nc_global(payloadView.srcTokenGlobalIdx(sourcePayload)) - sourceRank * maxTokensPerRank; + EP_DEVICE_ASSERT(sourceTokenIdx >= 0 && sourceTokenIdx < maxTokensPerRank); + void* destinationBuffer = sourceRank == rank + ? combineRecvBuffer + : peerBufferPtr(combineRecvBuffer, rdmaBufferBase, peerRecvBuffers[sourceRank]); + auto* destinationRow = reinterpret_cast(destinationBuffer) + + (static_cast(rank) * maxTokensPerRank + sourceTokenIdx) * kHiddenBytes; + issueTmaS2G(destinationRow, outputTile, static_cast(kHiddenBytes)); + } + } + __syncthreads(); + } + + if (tokenIteration > 0 && threadId == 0) waitTmaS2G(); +} + +template +MSCCLPP_DEVICE_INLINE void sendExpertRowsDirect(const void* expertOutput, const int* srcInfo, + const int64_t* layoutRange, int nExperts, int rank, int nRanks, + int maxTokensPerRank, void* combineRecvBuffer, void* rdmaBufferBase, + void* const* peerRecvBuffers, uint8_t* sharedMemory) { + const int threadId = static_cast(threadIdx.x); + const int warpId = threadId / WARP_SIZE; + const int laneId = get_lane_id(); + const int nLocalExperts = nExperts / nRanks; + const int nOutputSlotsPerExpert = nRanks * maxTokensPerRank; + constexpr int kNWorkers = tmaWorkerCount(); + constexpr size_t kHiddenBytes = static_cast(kHidden) * sizeof(nv_bfloat16); + auto* expertTokenPrefix = reinterpret_cast(sharedMemory); + auto* outputTiles = sharedMemory + combineControlBytes(nLocalExperts); + + if (threadId == 0) { + expertTokenPrefix[0] = 0; + for (int localExpertIdx = 0; localExpertIdx < nLocalExperts; ++localExpertIdx) { + int nLastRankTokens; + int lastRankOffset; + unpack2(layoutRange[localExpertIdx * nRanks + nRanks - 1], nLastRankTokens, lastRankOffset); + expertTokenPrefix[localExpertIdx + 1] = expertTokenPrefix[localExpertIdx] + lastRankOffset + nLastRankTokens; + } + } + __syncthreads(); + + const int nTotalRows = expertTokenPrefix[nLocalExperts]; + const int blockRowBegin = static_cast(static_cast(nTotalRows) * blockIdx.x / gridDim.x); + const int blockRowEnd = static_cast(static_cast(nTotalRows) * (blockIdx.x + 1) / gridDim.x); + auto* tmaBarriers = reinterpret_cast(outputTiles + static_cast(kNWorkers) * kHiddenBytes); + if (warpId == 0 && laneId < kNWorkers) { + auto* outputTile = outputTiles + static_cast(laneId) * kHiddenBytes; + auto* tmaBarrier = tmaBarriers + laneId; + uint32_t tmaPhase = 0; + if (blockRowBegin + laneId < blockRowEnd) initTmaBarrier(tmaBarrier); + + bool hasPendingStore = false; + for (int flatRowIdx = blockRowBegin + laneId; flatRowIdx < blockRowEnd; flatRowIdx += kNWorkers) { + if (hasPendingStore) waitTmaS2GRead(); + int localExpertIdx = 0; + while (flatRowIdx >= expertTokenPrefix[localExpertIdx + 1]) ++localExpertIdx; + const int expertTokenIdx = flatRowIdx - expertTokenPrefix[localExpertIdx]; + int sourceRank = 0; + for (; sourceRank < nRanks; ++sourceRank) { + int nRankTokens; + int rankOffset; + unpack2(layoutRange[localExpertIdx * nRanks + sourceRank], nRankTokens, rankOffset); + if (expertTokenIdx >= rankOffset && expertTokenIdx < rankOffset + nRankTokens) break; + } + EP_DEVICE_ASSERT(sourceRank < nRanks); + const int inputRowOffset = localExpertIdx * nOutputSlotsPerExpert + expertTokenIdx; + const int sourceTokenIdx = ld_nc_global(srcInfo + inputRowOffset); + EP_DEVICE_ASSERT(sourceTokenIdx >= 0 && sourceTokenIdx < maxTokensPerRank); + const auto* inputRow = + reinterpret_cast(expertOutput) + static_cast(inputRowOffset) * kHiddenBytes; + issueTmaG2S(inputRow, outputTile, tmaBarrier, static_cast(kHiddenBytes)); + waitTmaG2S(tmaBarrier, tmaPhase); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + const int globalExpertIdx = rank * nLocalExperts + localExpertIdx; + void* destinationBuffer = sourceRank == rank + ? combineRecvBuffer + : peerBufferPtr(combineRecvBuffer, rdmaBufferBase, peerRecvBuffers[sourceRank]); + auto* destinationRow = reinterpret_cast(destinationBuffer) + + (static_cast(globalExpertIdx) * maxTokensPerRank + sourceTokenIdx) * kHiddenBytes; + issueTmaS2G(destinationRow, outputTile, static_cast(kHiddenBytes)); + hasPendingStore = true; + } + + if (hasPendingStore) waitTmaS2G(); + } +} + +MSCCLPP_DEVICE_INLINE void combineSynchronize(mscclpp::BaseMemoryChannelDeviceHandle* signalChannels, + mscclpp::DeviceSemaphore* localReady, int rank, int nRanks, + int signalChannelStride) { + const int threadId = static_cast(threadIdx.x); + if (blockIdx.x == 0 && threadId < nRanks) { + const int peerRank = threadId; + if (peerRank == rank) { + localReady->release(); + localReady->acquire(); + } else { + auto& signalChannel = signalChannels[rankSignalChannelIndex(peerRank, signalChannelStride)]; + signalChannel.signal(); + signalChannel.wait(-1); + } + } +} + +template +MSCCLPP_DEVICE_INLINE void recvRankLocalPartials(void* output, const int64_t* topkIndices, int nTokens, int nTopk, + int nExperts, int nRanks, int maxTokensPerRank, + const void* combineRecvBuffer, uint8_t* sharedMemory) { + const int threadId = static_cast(threadIdx.x); + const int laneId = get_lane_id(); + const int nLocalExperts = nExperts / nRanks; + constexpr size_t kHiddenBytes = static_cast(kHidden) * sizeof(nv_bfloat16); + constexpr int kHiddenInt4 = kHiddenBytes / sizeof(int4); + constexpr int kChunksPerThread = (kHiddenInt4 + kCombineNThreads - 1) / kCombineNThreads; + static_assert(kHiddenInt4 % WARP_SIZE == 0); + auto* outputTiles = sharedMemory + combineControlBytes(nLocalExperts); + + int tokenIteration = 0; + for (int tokenIdx = static_cast(blockIdx.x); tokenIdx < nTokens; + tokenIdx += static_cast(gridDim.x), ++tokenIteration) { + const int stage = tokenIteration % kCombineNStages; + auto* outputTile = reinterpret_cast(outputTiles + static_cast(stage) * kHiddenBytes); + const int globalExpertIdx = laneId < nTopk ? static_cast(__ldg(topkIndices + tokenIdx * nTopk + laneId)) : -1; + const int destinationRank = globalExpertIdx >= 0 ? globalExpertIdx / nLocalExperts : -1; + const bool firstLaneForRank = isFirstLaneForRank(destinationRank, laneId); + const int partialRank = destinationRank >= 0 && firstLaneForRank ? destinationRank : -1; + + int4 reduced[kChunksPerThread] = {}; +#pragma unroll + for (int chunkIdx = 0; chunkIdx < kChunksPerThread; ++chunkIdx) { + const int hiddenIdx = threadId + chunkIdx * kCombineNThreads; + if (hiddenIdx < kHiddenInt4) { + reduced[chunkIdx] = reduceRankPartialsBf16x8(combineRecvBuffer, partialRank, nTopk, + maxTokensPerRank, tokenIdx, hiddenIdx); + } + } + if (tokenIteration >= kCombineNStages && threadId == 0) { + waitTmaS2GRead(); + } + if (tokenIteration >= kCombineNStages) __syncthreads(); +#pragma unroll + for (int chunkIdx = 0; chunkIdx < kChunksPerThread; ++chunkIdx) { + const int hiddenIdx = threadId + chunkIdx * kCombineNThreads; + if (hiddenIdx < kHiddenInt4) outputTile[hiddenIdx] = reduced[chunkIdx]; + } + __syncthreads(); + + if (threadId == 0) { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + auto* outputRow = reinterpret_cast(output) + static_cast(tokenIdx) * kHiddenBytes; + issueTmaS2G(outputRow, outputTile, static_cast(kHiddenBytes)); + } + } + if (tokenIteration > 0 && threadId == 0) waitTmaS2G(); +} + +template +MSCCLPP_DEVICE_INLINE void recvExpertRowsDirect(void* output, const int64_t* topkIndices, const float* topkWeights, + int nTokens, int nTopk, int maxTokensPerRank, + const void* combineRecvBuffer) { + constexpr int kBf16PerInt4 = sizeof(int4) / sizeof(nv_bfloat16); + constexpr int kHiddenInt4 = kHidden / kBf16PerInt4; + const int threadId = static_cast(threadIdx.x); + + for (int tokenIdx = static_cast(blockIdx.x); tokenIdx < nTokens; tokenIdx += static_cast(gridDim.x)) { + int regTopkIndices[kCombineMaxNTopk]; + float regTopkWeights[kCombineMaxNTopk]; + for (int topkIdx = 0; topkIdx < nTopk; ++topkIdx) { + regTopkIndices[topkIdx] = static_cast(__ldg(topkIndices + tokenIdx * nTopk + topkIdx)); + regTopkWeights[topkIdx] = topkWeights == nullptr ? 1.0f : __ldg(topkWeights + tokenIdx * nTopk + topkIdx); + } + +#pragma unroll + for (int hiddenIdx = threadId; hiddenIdx < kHiddenInt4; hiddenIdx += kCombineNThreads) { + float reduced[kBf16PerInt4] = {0.0f}; + for (int topkIdx = 0; topkIdx < nTopk; ++topkIdx) { + const int expertIdx = regTopkIndices[topkIdx]; + if (expertIdx < 0) continue; + const auto* expertRow = reinterpret_cast(combineRecvBuffer) + + (static_cast(expertIdx) * maxTokensPerRank + tokenIdx) * kHiddenInt4; + const int4 packed = ld_nc_global(expertRow + hiddenIdx); + const auto* values = reinterpret_cast(&packed); +#pragma unroll + for (int elemIdx = 0; elemIdx < kBf16PerInt4; ++elemIdx) { + reduced[elemIdx] += static_cast(values[elemIdx]) * regTopkWeights[topkIdx]; + } + } + + int4 packedOutput; + auto* outputValues = reinterpret_cast(&packedOutput); +#pragma unroll + for (int elemIdx = 0; elemIdx < kBf16PerInt4; ++elemIdx) { + outputValues[elemIdx] = static_cast(reduced[elemIdx]); + } + auto* outputRow = reinterpret_cast(output) + static_cast(tokenIdx) * kHiddenInt4; + outputRow[hiddenIdx] = packedOutput; + } + } +} + +template +__global__ __launch_bounds__(kCombineNThreads, 1) void combineKernel( + void* output, const void* expertOutput, const int64_t* topkIndices, const float* topkWeights, const int* srcInfo, + const int64_t* layoutRange, int nTokens, int nExperts, int rank, int nRanks, int nTopk, int maxTokensPerRank, + void* combineRecvBuffer, const void* dispatchRecvBuffer, void* rdmaBufferBase, void* const* peerRecvBuffers, + mscclpp::BaseMemoryChannelDeviceHandle* signalChannels, void* workspace, int dispatchMaxSms, + int signalChannelStride) { + extern __shared__ __align__(128) uint8_t sharedMemory[]; + DispatchWorkspaceView workspaceView(workspace, nRanks, nExperts, dispatchMaxSms); + + if constexpr (kMode == low_latency::OptimizedCombineMode::RANK_LOCAL_REDUCE) { + sendRankLocalPartials(expertOutput, nExperts, rank, nRanks, nTopk, maxTokensPerRank, combineRecvBuffer, + dispatchRecvBuffer, rdmaBufferBase, peerRecvBuffers, workspaceView, sharedMemory); + } else { + sendExpertRowsDirect(expertOutput, srcInfo, layoutRange, nExperts, rank, nRanks, maxTokensPerRank, + combineRecvBuffer, rdmaBufferBase, peerRecvBuffers, sharedMemory); + } + + workspaceView.combineSyncer_->sync(gridDim.x); + combineSynchronize(signalChannels, workspaceView.localPayloadReady_, rank, nRanks, signalChannelStride); + workspaceView.combineSyncer_->sync(gridDim.x); + + if constexpr (kMode == low_latency::OptimizedCombineMode::RANK_LOCAL_REDUCE) { + recvRankLocalPartials(output, topkIndices, nTokens, nTopk, nExperts, nRanks, maxTokensPerRank, + combineRecvBuffer, sharedMemory); + } else { + recvExpertRowsDirect(output, topkIndices, topkWeights, nTokens, nTopk, maxTokensPerRank, + combineRecvBuffer); + } +} + +template +inline void combineHiddenMode(void* output, const void* expertOutput, const int64_t* topkIndices, + const float* topkWeights, const int* srcInfo, const int64_t* layoutRange, + const low_latency::CombineConfig& config, const low_latency::BufferSet& currentBuffer, + void* dispatchRecvBuffer, const low_latency::TransportContext& transport, void* workspace, + int numBlocks, cudaStream_t stream) { + static_assert(kHidden == 4096 || kHidden == 7168 || kHidden == 8192 || kHidden == 9216); + if constexpr (kMode == low_latency::OptimizedCombineMode::DIRECT_SEND) { + static_assert(tmaWorkerCount() > 0); + } + const int nExperts = config.numExperts_; + const int rank = transport.rank_; + const int nRanks = transport.numRanks_; + const int nTokens = config.numCombinedTokens_; + const int nTopk = config.numTopk_; + const int nLocalExperts = nExperts / nRanks; + const int maxTokensPerRank = config.numMaxTokensPerRank_; + const int signalChannelStride = transport.memoryChannelStride_; + + auto combineFunc = combineKernel; + static thread_local int sharedMemoryLimitDevice = -1; + static thread_local size_t sharedMemoryLimit = 0; + if (sharedMemoryLimitDevice != transport.deviceId_) { + int sharedMemoryLimitInt; + cudaFuncAttributes attributes; + CUDA_CHECK( + cudaDeviceGetAttribute(&sharedMemoryLimitInt, cudaDevAttrMaxSharedMemoryPerBlockOptin, transport.deviceId_)); + CUDA_CHECK(cudaFuncGetAttributes(&attributes, combineFunc)); + EP_HOST_ASSERT(sharedMemoryLimitInt > static_cast(attributes.sharedSizeBytes)); + sharedMemoryLimitDevice = transport.deviceId_; + sharedMemoryLimit = static_cast(sharedMemoryLimitInt) - attributes.sharedSizeBytes; + } + + const size_t sharedBytes = combineSharedBytes(nLocalExperts); + EP_HOST_ASSERT(sharedBytes <= sharedMemoryLimit); + static thread_local int configuredDevice = -1; + static thread_local size_t configuredSharedBytes = 0; + static thread_local int residentBlocks = 0; + if (configuredDevice != transport.deviceId_ || configuredSharedBytes < sharedBytes) { + CUDA_CHECK( + cudaFuncSetAttribute(combineFunc, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sharedBytes))); + int blocksPerSm; + int numSms; + CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocksPerSm, combineFunc, kCombineNThreads, sharedBytes)); + CUDA_CHECK(cudaDeviceGetAttribute(&numSms, cudaDevAttrMultiProcessorCount, transport.deviceId_)); + configuredDevice = transport.deviceId_; + configuredSharedBytes = sharedBytes; + residentBlocks = blocksPerSm * numSms; + } + EP_HOST_ASSERT(residentBlocks >= numBlocks); + + cudaLaunchConfig_t launchCfg = {dim3(numBlocks), dim3(kCombineNThreads), sharedBytes, stream, nullptr, 0}; + CUDA_CHECK(cudaLaunchKernelEx(&launchCfg, combineFunc, output, expertOutput, topkIndices, topkWeights, srcInfo, + layoutRange, nTokens, nExperts, rank, nRanks, nTopk, maxTokensPerRank, + currentBuffer.recvDataBuffer_, dispatchRecvBuffer, transport.rdmaBufferBase_, + transport.peerBases_, transport.memoryChannels_, workspace, numBlocks, + signalChannelStride)); +} + +template +inline void combineHidden(void* output, const void* expertOutput, const int64_t* topkIndices, const float* topkWeights, + const int* srcInfo, const int64_t* layoutRange, const low_latency::CombineConfig& config, + const low_latency::BufferSet& currentBuffer, void* dispatchRecvBuffer, + const low_latency::TransportContext& transport, void* workspace, int numBlocks, + low_latency::OptimizedCombineMode mode, cudaStream_t stream) { + if (mode == low_latency::OptimizedCombineMode::RANK_LOCAL_REDUCE) { + return combineHiddenMode( + output, expertOutput, topkIndices, topkWeights, srcInfo, layoutRange, config, currentBuffer, dispatchRecvBuffer, + transport, workspace, numBlocks, stream); + } + return combineHiddenMode( + output, expertOutput, topkIndices, topkWeights, srcInfo, layoutRange, config, currentBuffer, dispatchRecvBuffer, + transport, workspace, numBlocks, stream); +} + +inline void combine(void* output, const void* expertOutput, const int64_t* topkIndices, const float* topkWeights, + const int* srcInfo, const int64_t* layoutRange, const low_latency::CombineConfig& config, + const low_latency::BufferSet& currentBuffer, void* dispatchRecvBuffer, + const low_latency::TransportContext& transport, void* workspace, int numBlocks, + low_latency::OptimizedCombineMode mode, cudaStream_t stream) { + const int nExperts = config.numExperts_; + const int rank = transport.rank_; + const int nRanks = transport.numRanks_; + + EP_HOST_ASSERT(output != nullptr); + EP_HOST_ASSERT(expertOutput != nullptr); + EP_HOST_ASSERT(topkIndices != nullptr); + EP_HOST_ASSERT(currentBuffer.recvDataBuffer_ != nullptr); + EP_HOST_ASSERT(dispatchRecvBuffer != nullptr); + EP_HOST_ASSERT(transport.rdmaBufferBase_ != nullptr); + EP_HOST_ASSERT(transport.peerBases_ != nullptr); + EP_HOST_ASSERT(transport.memoryChannels_ != nullptr); + EP_HOST_ASSERT(transport.ipcReady_); + EP_HOST_ASSERT(workspace != nullptr); + EP_HOST_ASSERT(nRanks > 0 && nRanks <= 2 * WARP_SIZE); + EP_HOST_ASSERT(nExperts > 0 && nExperts % nRanks == 0); + EP_HOST_ASSERT(rank >= 0 && rank < nRanks); + EP_HOST_ASSERT(config.numCombinedTokens_ >= 0 && config.numCombinedTokens_ <= config.numMaxTokensPerRank_); + EP_HOST_ASSERT(config.numTopk_ > 0 && config.numTopk_ <= kCombineMaxNTopk); + EP_HOST_ASSERT(transport.memoryChannelStride_ >= 1); + EP_HOST_ASSERT(config.inputDType_ == low_latency::DType::BF16); + EP_HOST_ASSERT(config.outputDType_ == low_latency::DType::BF16); + EP_HOST_ASSERT(numBlocks > 0 && numBlocks <= kCombineMaxNBlocks); + EP_HOST_ASSERT(mode == low_latency::OptimizedCombineMode::RANK_LOCAL_REDUCE || + mode == low_latency::OptimizedCombineMode::DIRECT_SEND); + if (mode == low_latency::OptimizedCombineMode::DIRECT_SEND) { + EP_HOST_ASSERT(srcInfo != nullptr); + EP_HOST_ASSERT(layoutRange != nullptr); + } + + switch (config.hidden_) { + case 4096: + return combineHidden<4096>(output, expertOutput, topkIndices, topkWeights, srcInfo, layoutRange, config, + currentBuffer, dispatchRecvBuffer, transport, workspace, numBlocks, mode, stream); + case 7168: + return combineHidden<7168>(output, expertOutput, topkIndices, topkWeights, srcInfo, layoutRange, config, + currentBuffer, dispatchRecvBuffer, transport, workspace, numBlocks, mode, stream); + case 8192: + return combineHidden<8192>(output, expertOutput, topkIndices, topkWeights, srcInfo, layoutRange, config, + currentBuffer, dispatchRecvBuffer, transport, workspace, numBlocks, mode, stream); + case 9216: + return combineHidden<9216>(output, expertOutput, topkIndices, topkWeights, srcInfo, layoutRange, config, + currentBuffer, dispatchRecvBuffer, transport, workspace, numBlocks, mode, stream); + default: + EP_HOST_ASSERT(false && "unsupported optimized low-latency hidden size"); + } +} + +} // namespace low_latency_opt + +namespace low_latency { + +void combineOptimized(void* output, const void* input, const int64_t* topkIdx, const float* topkWeights, + const int* srcInfo, const int64_t* layoutRange, const CombineConfig& config, + const BufferSet& currentBuffer, void* dispatchRecvBuffer, const TransportContext& transport, + void* workspace, int numBlocks, OptimizedCombineMode mode, cudaStream_t stream) { + low_latency_opt::combine(output, input, topkIdx, topkWeights, srcInfo, layoutRange, config, currentBuffer, + dispatchRecvBuffer, transport, workspace, numBlocks, mode, stream); +} + +} // namespace low_latency +} // namespace ep +} // namespace mscclpp diff --git a/src/ext/ep/low_latency/config.cuh b/src/ext/ep/low_latency/config.cuh new file mode 100644 index 00000000..fef3283e --- /dev/null +++ b/src/ext/ep/low_latency/config.cuh @@ -0,0 +1,134 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include + +#include "../config.hpp" +#include "../kernels/utils.cuh" + +namespace mscclpp { +namespace ep { +namespace low_latency_opt { + +constexpr int kDispatchNWarps = 16; +constexpr int kDispatchMinNWarpsPerGroup = 8; +constexpr int kDispatchMaxNWarpGroups = kDispatchNWarps / kDispatchMinNWarpsPerGroup; +constexpr int kDispatchNThreads = kDispatchNWarps * WARP_SIZE; +constexpr int kDispatchMaxNSms = 128; +constexpr int kDispatchMaxNRecvTmaWorkers = kDispatchNWarps; +constexpr size_t kOptimizedDynamicSharedMemoryBytes = 226 * 1024; +constexpr size_t kTmaWorkerControlBytes = kDispatchMaxNWarpGroups * WARP_SIZE * sizeof(int); +static_assert(kDispatchNWarps % kDispatchMinNWarpsPerGroup == 0); +static_assert(sizeof(mscclpp::DeviceSemaphore) == sizeof(int)); +static_assert(alignof(mscclpp::DeviceSemaphore) <= alignof(int)); +static_assert(sizeof(mscclpp::DeviceSyncer) % sizeof(int) == 0); +static_assert(alignof(mscclpp::DeviceSyncer) <= alignof(int)); + +MSCCLPP_HOST_DEVICE_INLINE int rankSignalChannelIndex(int peerRank, int signalChannelStride) { + return peerRank * signalChannelStride; +} + +MSCCLPP_HOST_DEVICE_INLINE size_t dispatchMetadataBytes(int nRanks, int nExperts) { + return configAlign(static_cast(nRanks + nExperts) * sizeof(mscclpp::LL8Packet), 128); +} + +MSCCLPP_HOST_DEVICE_INLINE size_t dispatchPayloadStride(int hidden, int nTopk) { + return configAlign(LowLatencyPayloadView(hidden, nTopk).numBytes_, 128); +} + +MSCCLPP_HOST_DEVICE_INLINE constexpr int dispatchNWarpsPerGroup(int nTokens, int nBlocks) { + return nTokens <= nBlocks ? kDispatchNWarps + : (nTokens <= 2 * nBlocks ? kDispatchNWarps / 2 : kDispatchMinNWarpsPerGroup); +} + +struct RecvTask { + int sourceRank_; + int tokenBegin_; + int tokenEnd_; +}; + +struct DispatchWorkspaceView { + uint32_t* metadataEpoch_; + int* rankPayloadSlots_; + int* rankPayloadCompletions_; + mscclpp::DeviceSemaphore* localPayloadReady_; + int* recvExpertCopiedCounts_; + uint32_t* rankReadyEpochs_; + RecvTask* recvTasks_; + uint32_t* tasksAssignedEpoch_; + int* nRecvTasks_; + mscclpp::DeviceSyncer* combineSyncer_; + + MSCCLPP_HOST_DEVICE_INLINE DispatchWorkspaceView(void* workspace, int nRanks, int nExperts, + [[maybe_unused]] int maxSms) { + auto* cursor = reinterpret_cast(workspace); + metadataEpoch_ = reinterpret_cast(cursor++); + rankPayloadSlots_ = cursor; + cursor += nRanks; + rankPayloadCompletions_ = cursor; + cursor += nRanks; + localPayloadReady_ = reinterpret_cast(cursor++); + recvExpertCopiedCounts_ = cursor; + cursor += nExperts; + rankReadyEpochs_ = reinterpret_cast(cursor); + cursor += nRanks; + recvTasks_ = reinterpret_cast(cursor); + cursor += 3 * kDispatchMaxNSms; + tasksAssignedEpoch_ = reinterpret_cast(cursor++); + nRecvTasks_ = cursor++; + combineSyncer_ = reinterpret_cast(cursor); + } + + MSCCLPP_HOST_DEVICE_INLINE static size_t numBytes(int nRanks, int nExperts, [[maybe_unused]] int maxSms) { + return static_cast(3 * nRanks + nExperts + 3 * kDispatchMaxNSms + 4) * sizeof(int) + + sizeof(mscclpp::DeviceSyncer); + } +}; + +MSCCLPP_HOST_DEVICE_INLINE size_t dispatchWorkspaceBytes(int nRanks, int nExperts, int maxSms) { + return DispatchWorkspaceView::numBytes(nRanks, nExperts, maxSms); +} + +MSCCLPP_HOST_DEVICE_INLINE size_t dispatchSharedControlBytes(int nRanks) { + constexpr int kNSendSlots = kDispatchMaxNWarpGroups * WARP_SIZE; + const int nSlots = nRanks > kNSendSlots ? nRanks : kNSendSlots; + return configAlign(static_cast(nSlots) * sizeof(int), 128); +} + +template +MSCCLPP_HOST_DEVICE_INLINE size_t dispatchSendTmaBytes(int nTopk) { + return kDispatchMaxNWarpGroups * (dispatchPayloadStride(kHidden, nTopk) + sizeof(uint64_t)); +} + +template +MSCCLPP_HOST_DEVICE_INLINE constexpr int tmaWorkerCount() { + static_assert(kHidden % 128 == 0); + constexpr size_t workerBytes = static_cast(kHidden) * sizeof(nv_bfloat16) + sizeof(uint64_t); + constexpr int nWorkers = + static_cast((kOptimizedDynamicSharedMemoryBytes - kTmaWorkerControlBytes) / workerBytes); + return nWorkers < kMaxWorkers ? nWorkers : kMaxWorkers; +} + +template +MSCCLPP_HOST_DEVICE_INLINE size_t dispatchRecvTmaBytes() { + constexpr int kNWorkers = tmaWorkerCount(); + constexpr size_t tileBytes = static_cast(kHidden) * sizeof(nv_bfloat16); + return static_cast(kNWorkers) * (tileBytes + sizeof(uint64_t)); +} + +template +MSCCLPP_HOST_DEVICE_INLINE size_t dispatchSharedBytes(int nRanks, int nExperts, int nTopk) { + const size_t controlBytes = dispatchSharedControlBytes(nRanks); + const size_t sendBytes = dispatchSendTmaBytes(nTopk); + const size_t recvBytes = dispatchRecvTmaBytes(); + const size_t tmaBytes = controlBytes + (sendBytes > recvBytes ? sendBytes : recvBytes); + const size_t metadataBytes = static_cast(nRanks + nExperts) * sizeof(int); + return tmaBytes > metadataBytes ? tmaBytes : metadataBytes; +} + +} // namespace low_latency_opt +} // namespace ep +} // namespace mscclpp diff --git a/src/ext/ep/low_latency/dispatch.cu b/src/ext/ep/low_latency/dispatch.cu new file mode 100644 index 00000000..94fcf086 --- /dev/null +++ b/src/ext/ep/low_latency/dispatch.cu @@ -0,0 +1,525 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +#include + +#include "../kernels/api.cuh" +#include "../kernels/exception.cuh" +#include "../kernels/utils.cuh" +#include "config.cuh" + +namespace mscclpp { +namespace ep { +namespace low_latency_opt { + +template +MSCCLPP_DEVICE_INLINE void dispatchSend(const void* inputTokens, mscclpp::BaseMemoryChannelDeviceHandle* signalChannels, + int nExperts, int rank, int nRanks, int signalChannelStride, + const int64_t* topkIndices, const float* topkWeights, int nTokens, int nTopk, + int maxTokensPerRank, void* recvBuffer, void* const* peerRecvBuffers, + void* rdmaBufferBase, void* workspace, uint32_t metadataFlag, int* sharedMem) { + const int threadId = static_cast(threadIdx.x); + const int warpId = threadId / WARP_SIZE; + const int laneId = get_lane_id(); + const int nSms = static_cast(gridDim.x) - 2; + const int notifyBlockIdx = nSms + 1; + const int nLocalExperts = nExperts / nRanks; + const size_t metadataBytes = dispatchMetadataBytes(nRanks, nExperts); + DispatchWorkspaceView workspaceView(workspace, nRanks, nExperts, nSms); + if (static_cast(blockIdx.x) > 0 && static_cast(blockIdx.x) <= nSms) { + const int senderBlockIdx = static_cast(blockIdx.x) - 1; + const int nWarpsPerGroup = dispatchNWarpsPerGroup(nTokens, nSms); + const int nWarpGroups = kDispatchNWarps / nWarpsPerGroup; + const int warpGroupId = warpId / nWarpsPerGroup; + const int subWarpId = warpId % nWarpsPerGroup; + const LowLatencyPayloadView payloadView(kHidden, nTopk); + const size_t payloadStride = dispatchPayloadStride(kHidden, nTopk); + constexpr size_t kHiddenBytes = static_cast(kHidden) * sizeof(nv_bfloat16); + constexpr int kHiddenInt4Count = kHiddenBytes / sizeof(int4); + auto* sharedPayloadBase = reinterpret_cast(sharedMem) + dispatchSharedControlBytes(nRanks); + auto* sendTmaBarriers = reinterpret_cast(sharedPayloadBase + kDispatchMaxNWarpGroups * payloadStride); + + if (subWarpId == 0) { + const int tokenStride = nSms * nWarpGroups; + const int firstTokenIdx = senderBlockIdx * nWarpGroups + warpGroupId; + auto* stagedPayload = sharedPayloadBase + static_cast(warpGroupId) * payloadStride; + auto* destinationSlots = sharedMem + warpGroupId * WARP_SIZE; + auto* tmaBarrier = sendTmaBarriers + warpGroupId; + uint32_t sendTmaPhase = 0; + if (firstTokenIdx < nTokens) { + if (laneId == 0) initTmaBarrier(tmaBarrier); + __syncwarp(); + } + + for (int tokenIdx = firstTokenIdx; tokenIdx < nTokens; tokenIdx += tokenStride) { + const auto* inputData = + reinterpret_cast(inputTokens) + static_cast(tokenIdx) * kHiddenInt4Count; + if (laneId == 0) { + issueTmaG2S(inputData, stagedPayload, tmaBarrier, static_cast(kHiddenBytes)); + } + const int routedExpertIdx = + laneId < nTopk ? static_cast(__ldg(topkIndices + tokenIdx * nTopk + laneId)) : -1; + const int dstRank = routedExpertIdx >= 0 ? routedExpertIdx / nLocalExperts : -1; + const bool firstLaneForRank = isFirstLaneForRank(dstRank, laneId); + const bool shouldSend = dstRank >= 0 && firstLaneForRank; + if (laneId < nTopk) { + int destinationSlot = -1; + if (shouldSend) { + destinationSlot = atomicAdd(workspaceView.rankPayloadSlots_ + dstRank, 1); + EP_DEVICE_ASSERT(destinationSlot < maxTokensPerRank); + } + destinationSlots[laneId] = destinationSlot; + payloadView.topKIndices(stagedPayload)[laneId] = routedExpertIdx; + payloadView.topKValues(stagedPayload)[laneId] = + topkWeights == nullptr ? 1.0f : __ldg(topkWeights + tokenIdx * nTopk + laneId); + } + if (laneId == 0) { + *payloadView.srcTokenGlobalIdx(stagedPayload) = rank * maxTokensPerRank + tokenIdx; + waitTmaG2S(tmaBarrier, sendTmaPhase); + } + __syncwarp(); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + bool hasPendingStore = false; + const int destinationSlot = laneId < nTopk ? destinationSlots[laneId] : -1; + if (destinationSlot >= 0) { + const int dstRank = payloadView.topKIndices(stagedPayload)[laneId] / nLocalExperts; + void* destinationBuffer = + dstRank == rank ? recvBuffer : peerBufferPtr(recvBuffer, rdmaBufferBase, peerRecvBuffers[dstRank]); + auto* destinationPayload = reinterpret_cast(destinationBuffer) + metadataBytes + + (static_cast(rank) * maxTokensPerRank + destinationSlot) * payloadStride; + issueTmaS2G(destinationPayload, stagedPayload, static_cast(payloadView.numBytes_)); + hasPendingStore = true; + } + if (hasPendingStore) { + waitTmaS2G(); + } + __syncwarp(); + if (destinationSlot >= 0) { + const int dstRank = payloadView.topKIndices(stagedPayload)[laneId] / nLocalExperts; + (void)mscclpp::atomicFetchAdd(workspaceView.rankPayloadCompletions_ + dstRank, 1, + mscclpp::memoryOrderRelease); + } + } + asm volatile("bar.sync %0, %1;" ::"r"(kDispatchMaxNWarpGroups + 1), "r"(nWarpGroups * WARP_SIZE) : "memory"); + } + } else if (static_cast(blockIdx.x) == notifyBlockIdx) { + int* sharedRankTokenCounts = sharedMem; + int* sharedExpertTokenCounts = sharedRankTokenCounts + nRanks; + for (int idx = threadId; idx < nRanks + nExperts; idx += blockDim.x) { + sharedRankTokenCounts[idx] = 0; + } + __syncthreads(); + for (int tokenIdx = warpId; tokenIdx < nTokens; tokenIdx += kDispatchNWarps) { + const int routedExpertIdx = + laneId < nTopk ? static_cast(__ldg(topkIndices + tokenIdx * nTopk + laneId)) : -1; + const int dstRank = routedExpertIdx >= 0 ? routedExpertIdx / nLocalExperts : -1; + if (routedExpertIdx >= 0) { + atomicAdd_block(sharedExpertTokenCounts + routedExpertIdx, 1); + } + if (isFirstLaneForRank(dstRank, laneId) && dstRank >= 0) { + atomicAdd_block(sharedRankTokenCounts + dstRank, 1); + } + } + __syncthreads(); + // send metadata via packet format + for (int dstRank = threadId; dstRank < nRanks; dstRank += blockDim.x) { + void* destinationBuffer = + dstRank == rank ? recvBuffer : peerBufferPtr(recvBuffer, rdmaBufferBase, peerRecvBuffers[dstRank]); + reinterpret_cast(destinationBuffer)[rank].write( + static_cast(sharedRankTokenCounts[dstRank]), metadataFlag); + } + for (int expertIdx = threadId; expertIdx < nExperts; expertIdx += blockDim.x) { + const int dstRank = expertIdx / nLocalExperts; + const int localExpertIdx = expertIdx % nLocalExperts; + void* destinationBuffer = + dstRank == rank ? recvBuffer : peerBufferPtr(recvBuffer, rdmaBufferBase, peerRecvBuffers[dstRank]); + reinterpret_cast(destinationBuffer)[nRanks + rank * nLocalExperts + localExpertIdx].write( + static_cast(sharedExpertTokenCounts[expertIdx]), metadataFlag); + } + + for (int dstRank = threadId; dstRank < nRanks; dstRank += blockDim.x) { + const int expectedPayloadCount = sharedRankTokenCounts[dstRank]; + if (expectedPayloadCount > 0) { + while (mscclpp::atomicLoad(workspaceView.rankPayloadCompletions_ + dstRank, + mscclpp::memoryOrderAcquire) != expectedPayloadCount); + } + workspaceView.rankPayloadSlots_[dstRank] = 0; + workspaceView.rankPayloadCompletions_[dstRank] = 0; + if (expectedPayloadCount == 0) continue; + if (dstRank == rank) { + workspaceView.localPayloadReady_->release(); + } else { + signalChannels[rankSignalChannelIndex(dstRank, signalChannelStride)].signal(); + } + } + __syncthreads(); + } +} + +MSCCLPP_DEVICE_INLINE int proportionalTaskBoundary(int nTokens, int nTasks, int nTotalTokens) { + return nTotalTokens == 0 ? 0 : static_cast(static_cast(nTokens) * nTasks / nTotalTokens); +} + +MSCCLPP_DEVICE_INLINE void dispatchRecvScheduler(int64_t* outputLayout, int* outputCount, + mscclpp::BaseMemoryChannelDeviceHandle* signalChannels, int nExperts, + int rank, int nRanks, int signalChannelStride, void* recvBuffer, + void* workspace, uint32_t metadataFlag, int* sharedMem) { + const int threadId = static_cast(threadIdx.x); + const int warpId = threadId / WARP_SIZE; + const int laneId = get_lane_id(); + const int nSms = static_cast(gridDim.x) - 2; + auto* rankTokenCounts = reinterpret_cast(recvBuffer); + const int nLocalExperts = nExperts / nRanks; + DispatchWorkspaceView workspaceView(workspace, nRanks, nExperts, nSms); + + const int nRankWarps = (nRanks + WARP_SIZE - 1) / WARP_SIZE; + const int requestedNLayoutWarps = (nLocalExperts + WARP_SIZE - 1) / WARP_SIZE; + const int maxNLayoutWarps = kDispatchNWarps - nRankWarps; + const int nLayoutWarps = requestedNLayoutWarps < maxNLayoutWarps ? requestedNLayoutWarps : maxNLayoutWarps; + + if (warpId < nRankWarps) { + const int sourceRank = threadId; + const int nRankTokens = + sourceRank < nRanks ? static_cast(rankTokenCounts[sourceRank].read(metadataFlag, -1)) : 0; + const int activeRank = nRankTokens > 0 ? 1 : 0; + int rankTokenPrefix = warpInclusiveSum(nRankTokens, laneId); + int activeRankPrefix = warpInclusiveSum(activeRank, laneId); + if (laneId == WARP_SIZE - 1) { + sharedMem[warpId] = rankTokenPrefix; + sharedMem[nRankWarps + warpId] = activeRankPrefix; + } + asm volatile("bar.sync %0, %1;" ::"r"(kDispatchMaxNWarpGroups + 1), "r"(nRankWarps * WARP_SIZE) : "memory"); + + if (warpId == 0) { + const int tokenTotal = laneId < nRankWarps ? sharedMem[laneId] : 0; + const int activeTotal = laneId < nRankWarps ? sharedMem[nRankWarps + laneId] : 0; + const int tokenPrefix = warpInclusiveSum(tokenTotal, laneId); + const int activePrefix = warpInclusiveSum(activeTotal, laneId); + if (laneId < nRankWarps) { + sharedMem[laneId] = tokenPrefix - tokenTotal; + sharedMem[nRankWarps + laneId] = activePrefix - activeTotal; + } + if (laneId == nRankWarps - 1) { + sharedMem[2 * nRankWarps] = tokenPrefix; + sharedMem[2 * nRankWarps + 1] = activePrefix; + } + } + asm volatile("bar.sync %0, %1;" ::"r"(kDispatchMaxNWarpGroups + 1), "r"(nRankWarps * WARP_SIZE) : "memory"); + + rankTokenPrefix += sharedMem[warpId]; + activeRankPrefix += sharedMem[nRankWarps + warpId]; + const int nTotalTokens = sharedMem[2 * nRankWarps]; + const int nActiveRanks = sharedMem[2 * nRankWarps + 1]; + const int nTasks = nTotalTokens < nSms ? nTotalTokens : nSms; + + // Reserve one task for every active rank. Distribute the remaining tasks + // proportionally after removing one token per active rank from the pool. + const int nReservedTasks = nActiveRanks; + const int nProportionalTasks = nTasks - nReservedTasks; + const int nProportionalTokens = nTotalTokens - nReservedTasks; + const int tokensBeforeRank = rankTokenPrefix - nRankTokens; + const int reservedTasksBeforeRank = activeRankPrefix - activeRank; + const int proportionalTokensBeforeRank = tokensBeforeRank - reservedTasksBeforeRank; + const int proportionalTokensThroughRank = rankTokenPrefix - activeRankPrefix; + const int proportionalTaskBegin = + proportionalTaskBoundary(proportionalTokensBeforeRank, nProportionalTasks, nProportionalTokens); + const int proportionalTaskEnd = + proportionalTaskBoundary(proportionalTokensThroughRank, nProportionalTasks, nProportionalTokens); + const int rankTaskBegin = reservedTasksBeforeRank + proportionalTaskBegin; + const int nRankTasks = activeRank + proportionalTaskEnd - proportionalTaskBegin; + if (sourceRank < nRanks && nRankTasks > 0) { + for (int rankTaskIdx = 0; rankTaskIdx < nRankTasks; ++rankTaskIdx) { + workspaceView.recvTasks_[rankTaskBegin + rankTaskIdx] = {sourceRank, nRankTokens * rankTaskIdx / nRankTasks, + nRankTokens * (rankTaskIdx + 1) / nRankTasks}; + } + } + if (threadId == 0) *workspaceView.nRecvTasks_ = nTasks; + + asm volatile("bar.sync %0, %1;" ::"r"(kDispatchMaxNWarpGroups + 2), "r"((nRankWarps + nLayoutWarps) * WARP_SIZE) + : "memory"); + if (threadId == 0) { + mscclpp::atomicStore(workspaceView.tasksAssignedEpoch_, metadataFlag, + mscclpp::memoryOrderRelease); + } + + if (sourceRank < nRanks && nRankTokens > 0) { + if (sourceRank == rank) { + workspaceView.localPayloadReady_->acquire(); + } else { + signalChannels[rankSignalChannelIndex(sourceRank, signalChannelStride)].wait(-1); + } + mscclpp::atomicStore(workspaceView.rankReadyEpochs_ + sourceRank, metadataFlag, + mscclpp::memoryOrderRelease); + } + } else if (warpId < nRankWarps + nLayoutWarps) { + auto* expertTokenCounts = reinterpret_cast(recvBuffer) + nRanks; + const int layoutThreadId = (warpId - nRankWarps) * WARP_SIZE + laneId; + const int nLayoutThreads = nLayoutWarps * WARP_SIZE; + for (int localExpertIdx = layoutThreadId; localExpertIdx < nLocalExperts; localExpertIdx += nLayoutThreads) { + int outputOffset = 0; + for (int sourceRank = 0; sourceRank < nRanks; ++sourceRank) { + const int nExpertTokens = + static_cast(expertTokenCounts[sourceRank * nLocalExperts + localExpertIdx].read(metadataFlag, -1)); + outputLayout[localExpertIdx * nRanks + sourceRank] = pack2(nExpertTokens, outputOffset); + outputOffset += nExpertTokens; + } + outputCount[localExpertIdx] = outputOffset; + } + asm volatile("bar.sync %0, %1;" ::"r"(kDispatchMaxNWarpGroups + 2), "r"((nRankWarps + nLayoutWarps) * WARP_SIZE) + : "memory"); + } +} + +MSCCLPP_DEVICE_INLINE bool acquireRecvTask(RecvTask& task, DispatchWorkspaceView& workspaceView, uint32_t metadataFlag, + int* sharedMem) { + auto* sharedTask = reinterpret_cast(sharedMem); + const int taskIdx = static_cast(blockIdx.x) - 1; + if (threadIdx.x == 0) { + while (mscclpp::atomicLoad(workspaceView.tasksAssignedEpoch_, + mscclpp::memoryOrderAcquire) != metadataFlag); + if (taskIdx < *workspaceView.nRecvTasks_) { + task = workspaceView.recvTasks_[taskIdx]; + while (mscclpp::atomicLoad(workspaceView.rankReadyEpochs_ + task.sourceRank_, + mscclpp::memoryOrderAcquire) != metadataFlag); + *sharedTask = task; + } else { + *sharedTask = {-1, 0, 0}; + } + } + __syncthreads(); + task = *sharedTask; + return task.sourceRank_ >= 0; +} + +template +MSCCLPP_DEVICE_INLINE void dispatchRecvWorker(void* output, int* outputSrcInfo, int64_t* outputLayout, int nExperts, + int rank, int nRanks, int nTopk, int maxTokensPerRank, void* recvBuffer, + void* workspace, uint32_t metadataFlag, int* sharedMem) { +#if defined(__CUDA_ARCH__) + static_assert(__CUDA_ARCH__ >= 900, "TMA recv requires SM90 or newer"); +#endif + const int threadId = static_cast(threadIdx.x); + const int warpId = threadId / WARP_SIZE; + const int laneId = get_lane_id(); + const int nSms = static_cast(gridDim.x) - 2; + DispatchWorkspaceView workspaceView(workspace, nRanks, nExperts, nSms); + RecvTask task; + if (!acquireRecvTask(task, workspaceView, metadataFlag, sharedMem)) return; + constexpr int kNRecvTmaWorkers = tmaWorkerCount(); + if (warpId >= kNRecvTmaWorkers) return; + + const int nLocalExperts = nExperts / nRanks; + const int sourceRank = task.sourceRank_; + const int globalExpertBase = rank * nLocalExperts; + const int globalExpertEnd = globalExpertBase + nLocalExperts; + const LowLatencyPayloadView payloadView(kHidden, nTopk); + const size_t payloadStride = dispatchPayloadStride(kHidden, nTopk); + constexpr size_t kHiddenBytes = static_cast(kHidden) * sizeof(nv_bfloat16); + constexpr size_t kTileBytes = kHiddenBytes; + const int nOutputSlotsPerExpert = nRanks * maxTokensPerRank; + auto* sourcePayloadBase = reinterpret_cast(recvBuffer) + dispatchMetadataBytes(nRanks, nExperts) + + static_cast(sourceRank) * maxTokensPerRank * payloadStride; + auto* tmaTiles = reinterpret_cast(sharedMem) + dispatchSharedControlBytes(nRanks); + auto* sharedTile = tmaTiles + static_cast(warpId) * kTileBytes; + auto* tmaBarriers = reinterpret_cast(tmaTiles + static_cast(kNRecvTmaWorkers) * kTileBytes); + auto* tmaBarrier = tmaBarriers + warpId; + bool hasPendingStore = false; + uint32_t recvTmaPhase = 0; + if (laneId == 0) initTmaBarrier(tmaBarrier); + + for (int sourceTokenSlot = task.tokenBegin_ + warpId; sourceTokenSlot < task.tokenEnd_; + sourceTokenSlot += kNRecvTmaWorkers) { + if (hasPendingStore) { + waitTmaS2GRead(); + } + __syncwarp(); + + auto* sourcePayload = sourcePayloadBase + static_cast(sourceTokenSlot) * payloadStride; + if (laneId == 0) { + issueTmaG2S(payloadView.template data(sourcePayload), sharedTile, tmaBarrier, + static_cast(kHiddenBytes)); + } + __syncwarp(); + + const int routedExpertIdx = laneId < nTopk ? ld_nc_global(payloadView.topKIndices(sourcePayload) + laneId) : -1; + const int localExpertIdx = routedExpertIdx >= globalExpertBase && routedExpertIdx < globalExpertEnd + ? routedExpertIdx - globalExpertBase + : -1; + const int sourceTokenIdx = __shfl_sync( + 0xffffffff, + laneId == 0 ? ld_nc_global(payloadView.srcTokenGlobalIdx(sourcePayload)) - sourceRank * maxTokensPerRank : 0, + 0); + int outputTokenIdx = -1; + int combineInputOffset = -1; + if (localExpertIdx >= 0) { + int expertTokenCount; + int outputOffset; + unpack2(outputLayout[localExpertIdx * nRanks + sourceRank], expertTokenCount, outputOffset); + const int copiedTokenIdx = + atomicAdd(workspaceView.recvExpertCopiedCounts_ + sourceRank * nLocalExperts + localExpertIdx, 1); + EP_DEVICE_ASSERT(copiedTokenIdx < expertTokenCount); + if (copiedTokenIdx == expertTokenCount - 1) { + workspaceView.recvExpertCopiedCounts_[sourceRank * nLocalExperts + localExpertIdx] = 0; + } + outputTokenIdx = outputOffset + copiedTokenIdx; + outputSrcInfo[static_cast(localExpertIdx) * nOutputSlotsPerExpert + outputTokenIdx] = sourceTokenIdx; + combineInputOffset = localExpertIdx * nOutputSlotsPerExpert + outputTokenIdx; + } + if (laneId < nTopk) payloadView.topKIndices(sourcePayload)[laneId] = combineInputOffset; + + if (laneId == 0) waitTmaG2S(tmaBarrier, recvTmaPhase); + __syncwarp(); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + + if (localExpertIdx >= 0) { + auto* outputData = reinterpret_cast(output) + + (static_cast(localExpertIdx) * nOutputSlotsPerExpert + outputTokenIdx) * kHiddenBytes; + issueTmaS2G(outputData, sharedTile, static_cast(kHiddenBytes)); + hasPendingStore = true; + } else { + hasPendingStore = false; + } + __syncwarp(); + } + + if (hasPendingStore) { + waitTmaS2G(); + } +} + +template +__global__ __launch_bounds__(kDispatchNThreads, 1) void dispatchKernel( + void* output, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, + mscclpp::BaseMemoryChannelDeviceHandle* signalChannels, int nExperts, int rank, int nRanks, int signalChannelStride, + const int64_t* topkIndices, const float* topkWeights, const void* inputTokens, int nTokens, int nTopk, + int maxTokensPerRank, void* recvBuffer, void* rdmaBufferBase, void* const* peerRecvBuffers, void* workspace) { + extern __shared__ __align__(128) uint8_t sharedMemory[]; + auto* sharedMem = reinterpret_cast(sharedMemory); + const int nSms = static_cast(gridDim.x) - 2; + DispatchWorkspaceView workspaceView(workspace, nRanks, nExperts, nSms); + const uint32_t metadataFlag = *workspaceView.metadataEpoch_ + 1; + + dispatchSend(inputTokens, signalChannels, nExperts, rank, nRanks, signalChannelStride, topkIndices, + topkWeights, nTokens, nTopk, maxTokensPerRank, recvBuffer, peerRecvBuffers, rdmaBufferBase, + workspace, metadataFlag, sharedMem); + + if (static_cast(blockIdx.x) == 0) { + dispatchRecvScheduler(outputLayout, outputCount, signalChannels, nExperts, rank, nRanks, signalChannelStride, + recvBuffer, workspace, metadataFlag, sharedMem); + } else if (static_cast(blockIdx.x) <= nSms) { + dispatchRecvWorker(output, outputSrcInfo, outputLayout, nExperts, rank, nRanks, nTopk, maxTokensPerRank, + recvBuffer, workspace, metadataFlag, sharedMem); + } + if (blockIdx.x == 0 && threadIdx.x == 0) { + *workspaceView.metadataEpoch_ = metadataFlag; + } +} + +template +inline void dispatchHidden(void* output, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, const void* input, + const int64_t* topkIdx, const float* topkWeights, const low_latency::DispatchConfig& config, + const low_latency::BufferSet& currentBuffer, const low_latency::TransportContext& transport, + void* workspace, int maxSms, cudaStream_t stream) { + static_assert(kHidden == 4096 || kHidden == 7168 || kHidden == 8192 || kHidden == 9216); + constexpr int kNRecvTmaWorkers = tmaWorkerCount(); + static_assert(kNRecvTmaWorkers > 0); + const int nExperts = config.numExperts_; + const int rank = transport.rank_; + const int nRanks = transport.numRanks_; + const int signalChannelStride = transport.memoryChannelStride_; + const int nTokens = config.numTokens_; + const int nTopk = config.numTopk_; + + static thread_local int sharedMemoryLimitDevice = -1; + static thread_local size_t sharedMemoryLimit = 0; + if (sharedMemoryLimitDevice != transport.deviceId_) { + int sharedMemoryLimitInt; + cudaFuncAttributes attributes; + CUDA_CHECK( + cudaDeviceGetAttribute(&sharedMemoryLimitInt, cudaDevAttrMaxSharedMemoryPerBlockOptin, transport.deviceId_)); + CUDA_CHECK(cudaFuncGetAttributes(&attributes, dispatchKernel)); + EP_HOST_ASSERT(sharedMemoryLimitInt > static_cast(attributes.sharedSizeBytes)); + sharedMemoryLimitDevice = transport.deviceId_; + sharedMemoryLimit = static_cast(sharedMemoryLimitInt) - attributes.sharedSizeBytes; + } + + const size_t dynamicSharedBytes = dispatchSharedBytes(nRanks, nExperts, nTopk); + EP_HOST_ASSERT(dynamicSharedBytes <= sharedMemoryLimit); + + cudaLaunchConfig_t cfg = {dim3(maxSms + 2), dim3(kDispatchNThreads), dynamicSharedBytes, stream, nullptr, 0}; + static thread_local int configuredDevice = -1; + static thread_local size_t configuredDynamicSharedBytes = 0; + if (configuredDevice != transport.deviceId_ || configuredDynamicSharedBytes < dynamicSharedBytes) { + CUDA_CHECK(cudaFuncSetAttribute(dispatchKernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(dynamicSharedBytes))); + configuredDevice = transport.deviceId_; + configuredDynamicSharedBytes = dynamicSharedBytes; + } + CUDA_CHECK(cudaLaunchKernelEx(&cfg, dispatchKernel, output, outputSrcInfo, outputLayout, outputCount, + transport.memoryChannels_, nExperts, rank, nRanks, signalChannelStride, topkIdx, + topkWeights, input, nTokens, nTopk, config.numMaxTokensPerRank_, + currentBuffer.recvDataBuffer_, transport.rdmaBufferBase_, transport.peerBases_, + workspace)); +} + +inline void dispatch(void* output, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, const void* input, + const int64_t* topkIdx, const float* topkWeights, const low_latency::DispatchConfig& config, + const low_latency::BufferSet& currentBuffer, const low_latency::TransportContext& transport, + void* workspace, int maxSms, cudaStream_t stream) { + const int nExperts = config.numExperts_; + const int rank = transport.rank_; + const int nRanks = transport.numRanks_; + + EP_HOST_ASSERT(nRanks > 0); + EP_HOST_ASSERT(nExperts > 0); + EP_HOST_ASSERT(nExperts % nRanks == 0); + EP_HOST_ASSERT(rank >= 0 && rank < nRanks); + EP_HOST_ASSERT(transport.memoryChannels_ != nullptr); + EP_HOST_ASSERT(transport.memoryChannelStride_ >= 1); + EP_HOST_ASSERT(config.numTokens_ >= 0); + EP_HOST_ASSERT(config.numTopk_ > 0 && config.numTopk_ <= WARP_SIZE); + EP_HOST_ASSERT(nRanks <= 2 * WARP_SIZE); + EP_HOST_ASSERT(maxSms >= nRanks && maxSms <= kDispatchMaxNSms); + EP_HOST_ASSERT(output != nullptr); + EP_HOST_ASSERT(outputSrcInfo != nullptr); + EP_HOST_ASSERT(outputLayout != nullptr); + EP_HOST_ASSERT(outputCount != nullptr); + EP_HOST_ASSERT(input != nullptr); + EP_HOST_ASSERT(topkIdx != nullptr); + EP_HOST_ASSERT(currentBuffer.recvDataBuffer_ != nullptr); + EP_HOST_ASSERT(transport.peerBases_ != nullptr); + EP_HOST_ASSERT(workspace != nullptr); + EP_HOST_ASSERT(config.outputDType_ == low_latency::DType::BF16); + + switch (config.hidden_) { + case 4096: + return dispatchHidden<4096>(output, outputSrcInfo, outputLayout, outputCount, input, topkIdx, topkWeights, config, + currentBuffer, transport, workspace, maxSms, stream); + case 7168: + return dispatchHidden<7168>(output, outputSrcInfo, outputLayout, outputCount, input, topkIdx, topkWeights, config, + currentBuffer, transport, workspace, maxSms, stream); + case 8192: + return dispatchHidden<8192>(output, outputSrcInfo, outputLayout, outputCount, input, topkIdx, topkWeights, config, + currentBuffer, transport, workspace, maxSms, stream); + case 9216: + return dispatchHidden<9216>(output, outputSrcInfo, outputLayout, outputCount, input, topkIdx, topkWeights, config, + currentBuffer, transport, workspace, maxSms, stream); + default: + EP_HOST_ASSERT(false && "unsupported optimized low-latency hidden size"); + } +} + +} // namespace low_latency_opt + +namespace low_latency { + +void dispatchOptimized(void* output, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, const void* input, + const int64_t* topkIdx, const float* topkWeights, const DispatchConfig& config, + const BufferSet& currentBuffer, const TransportContext& transport, void* workspace, int maxSms, + cudaStream_t stream) { + low_latency_opt::dispatch(output, outputSrcInfo, outputLayout, outputCount, input, topkIdx, topkWeights, config, + currentBuffer, transport, workspace, maxSms, stream); +} + +} // namespace low_latency +} // namespace ep +} // namespace mscclpp diff --git a/src/ext/ep/moe_runtime.cc b/src/ext/ep/moe_runtime.cc index 51569104..0f43857d 100644 --- a/src/ext/ep/moe_runtime.cc +++ b/src/ext/ep/moe_runtime.cc @@ -10,6 +10,7 @@ #include #include #include +#include #include #include "kernels/api.cuh" @@ -23,6 +24,8 @@ namespace { using EPProxyService = mscclpp::ProxyService; +constexpr size_t kLowLatencyOptWorkspaceOffset = NUM_WORKSPACE_BYTES / 2; + int localWorldSize() { int localWorldSize = NUM_MAX_NVL_PEERS; if (const char* env = std::getenv("MSCCLPP_EP_LOCAL_WORLD_SIZE")) { @@ -68,6 +71,21 @@ bool resolveFabricIpcSupported() { return prop.major >= 10; } +low_latency::OptimizedCombineMode resolveOptimizedCombineMode() { + const char* env = std::getenv("MSCCLPP_EP_OPTIMIZED_COMBINE"); + if (env == nullptr) return low_latency::OptimizedCombineMode::DISABLED; + std::string value(env); + for (auto& c : value) c = std::tolower(static_cast(c)); + if (value == "1" || value == "on" || value == "true" || value == "yes" || value == "rank_local_reduce" || + value == "local_reduce") { + return low_latency::OptimizedCombineMode::RANK_LOCAL_REDUCE; + } + if (value == "direct_send" || value == "direct" || value == "no_local_reduce") { + return low_latency::OptimizedCombineMode::DIRECT_SEND; + } + return low_latency::OptimizedCombineMode::DISABLED; +} + } // namespace MoERuntime::MoERuntime(mscclpp::Communicator& communicator, int64_t numNvlBytes, int64_t numRdmaBytes, MoEMode mode) @@ -91,6 +109,7 @@ MoERuntime::MoERuntime(mscclpp::Communicator& communicator, int64_t numNvlBytes, numNvlRanks_ = std::min(numRanks_, lws); numProxyServices_ = resolveNumProxyServices(numRanks_, lws); + optimizedCombineMode_ = resolveOptimizedCombineMode(); proxyServices_.reserve(numProxyServices_); for (int i = 0; i < numProxyServices_; ++i) proxyServices_.emplace_back(std::make_shared()); @@ -268,19 +287,19 @@ void MoERuntime::setup() { void MoERuntime::dispatch(void* output, float* outputScales, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, const void* input, const int64_t* topkIdx, const float* topkWeights, int numTokens, int hidden, int numTopk, int numMaxDispatchTokensPerRank, int numExperts, - bool requiresQuantization, DispatchLayout dispatchLayout, cudaStream_t stream) { + bool requiresQuantization, DispatchLayout dispatchLayout, int numSms, cudaStream_t stream) { EP_HOST_ASSERT(mode_ == MoEMode::LOW_LATENCY); EP_HOST_ASSERT(hidden % sizeof(int4) == 0 && hidden % 128 == 0); EP_HOST_ASSERT(numTokens <= numMaxDispatchTokensPerRank); EP_HOST_ASSERT(numExperts % numRanks_ == 0); EP_HOST_ASSERT(dispatchLayout == DispatchLayout::EXPERT_MAJOR || dispatchLayout == DispatchLayout::FLAT); EP_HOST_ASSERT(llIpcReady_ && "low-latency rank-dedup dispatch currently requires IPC/NVLink reachability"); + EP_HOST_ASSERT(numSms >= numRanks_ && numSms <= 128); LowLatencyLayout layout(rdmaBufferPtr_, numMaxDispatchTokensPerRank, hidden, numRanks_, numExperts, numTopk); EP_HOST_ASSERT(layout.totalBytes <= static_cast(numRdmaBytes_)); auto buffer = layout.buffers[lowLatencyBufferIdx_]; - auto nextBuffer = layout.buffers[lowLatencyBufferIdx_ ^= 1]; - auto nextCleanMeta = nextBuffer.cleanMeta(); + lowLatencyBufferIdx_ ^= 1; low_latency::DispatchConfig config{ .numTokens_ = numTokens, @@ -297,40 +316,54 @@ void MoERuntime::dispatch(void* output, float* outputScales, int* outputSrcInfo, .recvCountBuffer_ = buffer.dispatchRdmaRecvCountBuffer, .cleanupRegion_ = nullptr, .cleanupSize_ = 0}; - low_latency::BufferSet nextBufferSet{.sendDataBuffer_ = nullptr, - .sendCountBuffer_ = nullptr, - .recvDataBuffer_ = nullptr, - .recvCountBuffer_ = nullptr, - .cleanupRegion_ = nextCleanMeta.first, - .cleanupSize_ = nextCleanMeta.second}; low_latency::TransportContext transport{ .rdmaBufferBase_ = rdmaBufferPtr_, .portChannels_ = portChannelHandlesDevicePtr_.get(), .memoryChannels_ = llMemoryChannelHandlesDevicePtr_ ? llMemoryChannelHandlesDevicePtr_.get() : nullptr, + .memoryChannelStride_ = 1, .peerBases_ = peerRdmaBasesGpu_, .ipcReady_ = llIpcReady_, + .deviceId_ = deviceId_, .rank_ = rank_, .numRanks_ = numRanks_, .ranksPerIpcDomain_ = llRanksPerIpcDomain_}; - low_latency::dispatch(output, outputScales, outputSrcInfo, outputLayout, outputCount, input, topkIdx, topkWeights, - config, currentBuffer, nextBufferSet, transport, workspace_, stream, - low_latency::SEND_AND_RECV); + if (!requiresQuantization) { + const size_t optWorkspaceBytes = + static_cast(3 * numRanks_ + numExperts + 3 * 128 + 4) * sizeof(int) + sizeof(mscclpp::DeviceSyncer); + EP_HOST_ASSERT(optWorkspaceBytes <= NUM_WORKSPACE_BYTES - kLowLatencyOptWorkspaceOffset); + auto* optWorkspace = reinterpret_cast(workspace_) + kLowLatencyOptWorkspaceOffset; + low_latency::dispatchOptimized(output, outputSrcInfo, outputLayout, outputCount, input, topkIdx, topkWeights, + config, currentBuffer, transport, optWorkspace, numSms, stream); + } else { + auto nextCleanMeta = layout.buffers[lowLatencyBufferIdx_].cleanMeta(); + low_latency::BufferSet nextBufferSet{.sendDataBuffer_ = nullptr, + .sendCountBuffer_ = nullptr, + .recvDataBuffer_ = nullptr, + .recvCountBuffer_ = nullptr, + .cleanupRegion_ = nextCleanMeta.first, + .cleanupSize_ = nextCleanMeta.second}; + low_latency::dispatch(output, outputScales, outputSrcInfo, outputLayout, outputCount, input, topkIdx, topkWeights, + config, currentBuffer, nextBufferSet, transport, workspace_, stream, + low_latency::SEND_AND_RECV); + } + optimizedDispatchMetadataReady_ = !requiresQuantization; } void MoERuntime::combine(void* output, const void* input, const float* inputScales, const int64_t* topkIdx, const float* topkWeights, const int* srcInfo, const int64_t* layoutRange, int numTokens, int hidden, int numTopk, int numMaxDispatchTokensPerRank, int numExperts, - bool requiresDequantization, cudaStream_t stream) { + bool requiresDequantization, low_latency::OptimizedCombineMode optimizedMode, int numBlocks, + cudaStream_t stream) { EP_HOST_ASSERT(mode_ == MoEMode::LOW_LATENCY); EP_HOST_ASSERT(hidden % sizeof(int4) == 0 && hidden % 128 == 0); EP_HOST_ASSERT(numExperts % numRanks_ == 0); EP_HOST_ASSERT(llIpcReady_ && "low-latency combine currently requires IPC/NVLink reachability"); + EP_HOST_ASSERT(numBlocks > 0 && numBlocks <= 128); LowLatencyLayout layout(rdmaBufferPtr_, numMaxDispatchTokensPerRank, hidden, numRanks_, numExperts, numTopk); EP_HOST_ASSERT(layout.totalBytes <= static_cast(numRdmaBytes_)); auto buffer = layout.buffers[lowLatencyBufferIdx_]; - auto nextBuffer = layout.buffers[lowLatencyBufferIdx_ ^= 1]; - auto nextCleanMeta = nextBuffer.cleanMeta(); + lowLatencyBufferIdx_ ^= 1; low_latency::CombineConfig config{ .numCombinedTokens_ = numTokens, @@ -347,23 +380,38 @@ void MoERuntime::combine(void* output, const void* input, const float* inputScal .recvCountBuffer_ = buffer.combineRdmaRecvFlagBuffer, .cleanupRegion_ = nullptr, .cleanupSize_ = 0}; - low_latency::BufferSet nextBufferSet{.sendDataBuffer_ = nullptr, - .sendCountBuffer_ = nullptr, - .recvDataBuffer_ = nullptr, - .recvCountBuffer_ = nullptr, - .cleanupRegion_ = nextCleanMeta.first, - .cleanupSize_ = nextCleanMeta.second}; low_latency::TransportContext transport{ .rdmaBufferBase_ = rdmaBufferPtr_, .portChannels_ = portChannelHandlesDevicePtr_.get(), .memoryChannels_ = llMemoryChannelHandlesDevicePtr_ ? llMemoryChannelHandlesDevicePtr_.get() : nullptr, + .memoryChannelStride_ = 1, .peerBases_ = peerRdmaBasesGpu_, .ipcReady_ = llIpcReady_, + .deviceId_ = deviceId_, .rank_ = rank_, .numRanks_ = numRanks_, .ranksPerIpcDomain_ = llRanksPerIpcDomain_}; - low_latency::combine(output, input, inputScales, topkIdx, topkWeights, srcInfo, layoutRange, config, currentBuffer, - nextBufferSet, transport, workspace_, stream, low_latency::SEND_AND_RECV); + const auto effectiveOptimizedMode = + optimizedMode == low_latency::OptimizedCombineMode::DISABLED ? optimizedCombineMode_ : optimizedMode; + if (effectiveOptimizedMode != low_latency::OptimizedCombineMode::DISABLED && optimizedDispatchMetadataReady_ && + !requiresDequantization) { + auto* optWorkspace = reinterpret_cast(workspace_) + kLowLatencyOptWorkspaceOffset; + auto dispatchRecvBuffer = layout.buffers[lowLatencyBufferIdx_].dispatchRdmaRecvDataBuffer; + low_latency::combineOptimized(output, input, topkIdx, topkWeights, srcInfo, layoutRange, config, currentBuffer, + dispatchRecvBuffer, transport, optWorkspace, numBlocks, effectiveOptimizedMode, + stream); + } else { + auto nextCleanMeta = layout.buffers[lowLatencyBufferIdx_].cleanMeta(); + low_latency::BufferSet nextBufferSet{.sendDataBuffer_ = nullptr, + .sendCountBuffer_ = nullptr, + .recvDataBuffer_ = nullptr, + .recvCountBuffer_ = nullptr, + .cleanupRegion_ = nextCleanMeta.first, + .cleanupSize_ = nextCleanMeta.second}; + low_latency::combine(output, input, inputScales, topkIdx, topkWeights, srcInfo, layoutRange, config, currentBuffer, + nextBufferSet, transport, workspace_, stream, low_latency::SEND_AND_RECV); + } + optimizedDispatchMetadataReady_ = false; } } // namespace ep diff --git a/src/ext/ep/moe_runtime.hpp b/src/ext/ep/moe_runtime.hpp index c846db87..a17540d1 100644 --- a/src/ext/ep/moe_runtime.hpp +++ b/src/ext/ep/moe_runtime.hpp @@ -37,12 +37,12 @@ class MoERuntime { void dispatch(void* output, float* outputScales, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, const void* input, const int64_t* topkIdx, const float* topkWeights, int numTokens, int hidden, int numTopk, int numMaxDispatchTokensPerRank, int numExperts, bool requiresQuantization, - DispatchLayout dispatchLayout, cudaStream_t stream); + DispatchLayout dispatchLayout, int numSms, cudaStream_t stream); void combine(void* output, const void* input, const float* inputScales, const int64_t* topkIdx, const float* topkWeights, const int* srcInfo, const int64_t* layoutRange, int numTokens, int hidden, int numTopk, int numMaxDispatchTokensPerRank, int numExperts, bool requiresDequantization, - cudaStream_t stream); + low_latency::OptimizedCombineMode optimizedMode, int numBlocks, cudaStream_t stream); private: int lowLatencyBufferIdx_ = 0; @@ -60,6 +60,8 @@ class MoERuntime { int numProxyServices_ = 1; int llRanksPerIpcDomain_ = 0; bool llIpcReady_ = false; + bool optimizedDispatchMetadataReady_ = false; + low_latency::OptimizedCombineMode optimizedCombineMode_ = low_latency::OptimizedCombineMode::DISABLED; void* rdmaBufferPtr_ = nullptr; void* workspace_ = nullptr; @@ -74,6 +76,7 @@ class MoERuntime { void** peerRdmaBasesGpu_ = nullptr; std::vector llMemoryChannels_; std::shared_ptr llMemoryChannelHandlesDevicePtr_; + // std::shared_ptr llMemoryChannelHandlesForExpertDevicePtr_; void setup(); }; diff --git a/test/mscclpp-test/CMakeLists.txt b/test/mscclpp-test/CMakeLists.txt index 241b7e02..b6eef5dd 100644 --- a/test/mscclpp-test/CMakeLists.txt +++ b/test/mscclpp-test/CMakeLists.txt @@ -21,3 +21,29 @@ add_mscclpp_test_executable(sendrecv_test_perf sendrecv_test.cu) add_mscclpp_test_executable(allgather_test_perf allgather_test.cu) add_mscclpp_test_executable(allreduce_test_perf allreduce_test.cu) add_mscclpp_test_executable(alltoall_test_perf alltoall_test.cu) +add_mscclpp_test_executable(ep_low_latency_opt_signal_perf ep_low_latency_opt_signal_perf.cu) +target_sources(ep_low_latency_opt_signal_perf PRIVATE + ${PROJECT_SOURCE_DIR}/src/ext/ep/low_latency/dispatch.cu) +target_include_directories(ep_low_latency_opt_signal_perf PRIVATE ${PROJECT_SOURCE_DIR}/src/ext/ep) +if(MSCCLPP_USE_CUDA) + set(_mscclpp_ep_test_gpu_archs "") + foreach(_arch IN LISTS CMAKE_CUDA_ARCHITECTURES) + if(_arch STREQUAL "native") + list(APPEND _mscclpp_ep_test_gpu_archs "${_arch}") + else() + string(REGEX MATCH "^[0-9]+" _arch_num "${_arch}") + if(_arch_num AND _arch_num GREATER_EQUAL 90) + list(APPEND _mscclpp_ep_test_gpu_archs "${_arch}") + endif() + endif() + endforeach() + set_target_properties(ep_low_latency_opt_signal_perf PROPERTIES + CUDA_ARCHITECTURES "${_mscclpp_ep_test_gpu_archs}") + + add_executable(tma_pipeline_perf tma_pipeline_perf.cu) + target_link_libraries(tma_pipeline_perf PRIVATE ${GPU_LIBRARIES}) + target_include_directories(tma_pipeline_perf ${TEST_INC_COMMON}) + set_target_properties(tma_pipeline_perf PROPERTIES + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/bin/mscclpp-test" + CUDA_ARCHITECTURES "${_mscclpp_ep_test_gpu_archs}") +endif() diff --git a/test/python/ep/test_low_latency_multirank.py b/test/python/ep/test_low_latency_multirank.py index a3d9ed13..d50b080a 100644 --- a/test/python/ep/test_low_latency_multirank.py +++ b/test/python/ep/test_low_latency_multirank.py @@ -20,18 +20,14 @@ Launch with (2 nodes, 1 GPU per node -- DeepEP's recommended LL topology): torchrun --nnodes=2 --nproc_per_node=1 --rdzv-backend=c10d \ --rdzv-endpoint=:29600 test/python/ep/test_low_latency_multirank.py -Exercises the LL dispatch + combine round-trip on a single node. The -minimal correctness check: +Exercises the optimized BF16 LL dispatch plus the default combine path on a +single node. The experimental optimized combine performs rank-local partial +reduction, TMA send, and source-rank reduction. The minimal correctness check: - dispatch: per-expert received token counts agree with an all-gathered reference computed from topk_idx across all ranks; - combine: the reconstructed x matches the analytical sum ``x * sum(topk_weights, masked by topk_idx == -1)``. -Known limitation (see src/ext/ep/README.md): the LL kernels drive every -peer via MSCCL++ PortChannel. Intra-node IB loopback between two HCAs on -the same host (what an 8-GPU single-node launch exercises) currently hangs -during dispatch; cross-node LL with one GPU per node works as designed. - Adapted from DeepEP/tests/test_low_latency.py stripped to the bare checks we need for an LL port smoke test. BF16-only (no FP8 check). """ @@ -55,9 +51,24 @@ import torch.distributed as dist def parse_args(): parser = argparse.ArgumentParser(description="MSCCL++ EP low-latency multi-rank correctness/benchmark test") parser.add_argument("--num-tokens", type=int, default=128) - parser.add_argument("--hidden", type=int, default=7168, help="LL kernels are compiled for a fixed hidden set") + parser.add_argument( + "--hidden", + type=int, + default=7168, + choices=(4096, 7168, 8192, 9216), + help="BF16 hidden size compiled into the optimized low-latency kernels", + ) parser.add_argument("--num-topk", type=int, default=8) parser.add_argument("--num-experts", type=int, default=256) + parser.add_argument("--num-active-ranks", type=int, default=0, help="Limit routing to the first N ranks") + parser.add_argument("--no-weights", action="store_true", help="Use implicit unit routing weights") + parser.add_argument("--dispatch-num-sms", type=int, default=64) + parser.add_argument("--combine-num-sms", type=int, default=64) + parser.add_argument( + "--optimized-combine-mode", + choices=("disabled", "rank_local_reduce", "direct_send"), + default="disabled", + ) parser.add_argument("--bench", action="store_true", help="Run dispatch/combine benchmark after correctness") parser.add_argument( "--cuda-graph", @@ -102,6 +113,11 @@ def main(): num_experts = args.num_experts assert num_experts % num_ranks == 0 num_local_experts = num_experts // num_ranks + combine_mode = { + "disabled": ep.OptimizedCombineMode.DISABLED, + "rank_local_reduce": ep.OptimizedCombineMode.RANK_LOCAL_REDUCE, + "direct_send": ep.OptimizedCombineMode.DIRECT_SEND, + }[args.optimized_combine_mode] torch.manual_seed(0xB3C4 + rank) random.seed(0xB3C4 + rank) @@ -111,8 +127,13 @@ def main(): # can verify which source token it is looking at. x[:, -128:] = torch.arange(num_tokens, device="cuda").to(torch.bfloat16).view(-1, 1) scores = torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs() + 1 + if args.num_active_ranks: + assert num_topk <= args.num_active_ranks * num_local_experts + scores[:, args.num_active_ranks * num_local_experts :] = float("-inf") topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=True)[1] - topk_weights = torch.randn((num_tokens, num_topk), dtype=torch.float32, device="cuda").abs() + topk_weights = ( + None if args.no_weights else torch.randn((num_tokens, num_topk), dtype=torch.float32, device="cuda").abs() + ) # Randomly mask some positions for _ in range(min(10, num_tokens)): @@ -127,6 +148,9 @@ def main(): max_tokens_per_rank=num_tokens, mode=ep.MoEMode.LOW_LATENCY, num_rdma_qps_per_rank=max(1, num_experts // num_ranks), + low_latency_dispatch_num_sms=args.dispatch_num_sms, + low_latency_combine_num_sms=args.combine_num_sms, + low_latency_combine_mode=combine_mode, ) if rank == 0: print( @@ -208,7 +232,11 @@ def main(): expected_f = torch.zeros_like(x, dtype=torch.float32) x_f = x.float() for j in range(num_topk): - weight_j = topk_weights[:, j].masked_fill(topk_idx[:, j] == -1, 0.0).view(-1, 1) + weight_j = ( + (topk_idx[:, j] != -1).float() + if topk_weights is None + else topk_weights[:, j].masked_fill(topk_idx[:, j] == -1, 0.0) + ).view(-1, 1) expected_f += x_f * weight_j expected = expected_f.to(torch.bfloat16) diff = (combined_x.float() - expected.float()).abs().max().item() @@ -218,7 +246,8 @@ def main(): flush=True, ) assert torch.isnan(combined_x).any().item() is False - assert diff < 1e-2, f"rank{rank}: LL combine mismatch diff={diff}" + combine_tolerance = 8.0 if combine_mode == ep.OptimizedCombineMode.RANK_LOCAL_REDUCE else 1e-2 + assert diff <= combine_tolerance, f"rank{rank}: LL combine mismatch diff={diff}" dist.barrier(group=group) if rank == 0: @@ -247,7 +276,7 @@ def main(): graph_diff = (graph_combined_x.float() - expected.float()).abs().max().item() assert torch.isnan(graph_combined_x).any().item() is False - assert graph_diff < 1e-2, f"rank{rank}: LL CUDA graph combine mismatch diff={graph_diff}" + assert graph_diff <= combine_tolerance, f"rank{rank}: LL CUDA graph combine mismatch diff={graph_diff}" dist.barrier(group=group) if rank == 0: print(f"[cuda graph dispatch+combine] OK max|got-expected|={graph_diff:.4e}", flush=True)