diff --git a/src/ext/ep/CMakeLists.txt b/src/ext/ep/CMakeLists.txt index 36acf022..ae775d05 100644 --- a/src/ext/ep/CMakeLists.txt +++ b/src/ext/ep/CMakeLists.txt @@ -6,7 +6,7 @@ # # Two backends share one module, both with a torch-free, raw-pointer (uintptr_t) # nanobind API so the module never links libtorch: -# - Low-latency (LL): moe_runtime.cc + kernels/low_latency.cu. +# - Low-latency (LL): moe_runtime.cc + low_latency/{dispatch,combine}.cu. # - High-throughput (HT): ht_runtime.cc + ht/kernels/*.cu, a DeepEP-style # runtime de-torched to the same pointer boundary as the LL runtime (dynamic # recv sizing via the two-phase notify -> allocate -> dispatch API). diff --git a/src/ext/ep/README.md b/src/ext/ep/README.md index 47d902dc..5e0b5c8e 100644 --- a/src/ext/ep/README.md +++ b/src/ext/ep/README.md @@ -4,7 +4,7 @@ A torch-free nanobind extension for MoE `dispatch` / `combine` primitives in MSCCL++. The module builds two active backends: - **Low-latency (LL)**: `MoERuntime` / `MoECommunicator(mode=LOW_LATENCY)`, - backed by `kernels/low_latency.cu`. + backed by `low_latency/dispatch.cu` and `low_latency/combine.cu`. - **High-throughput (HT)**: `ExpertParallelRuntime` / `MoECommunicator(mode=HIGH_THROUGHPUT)`, backed by `ht_runtime.cc` and `ht/kernels/*`. @@ -488,11 +488,19 @@ Env knobs: | Backend | Python API | C++ runtime | Kernel sources | Layout | |---------|------------|-------------|----------------|--------| -| LL | `MoECommunicator(mode=MoEMode.LOW_LATENCY)` / `MoERuntime` | `moe_runtime.cc` | `kernels/low_latency.cu` | `EXPERT_MAJOR` | +| LL | `MoECommunicator(mode=MoEMode.LOW_LATENCY)` / `MoERuntime` | `moe_runtime.cc` | `low_latency/{dispatch,combine}.cu` | `EXPERT_MAJOR` | | HT | `MoECommunicator(mode=MoEMode.HIGH_THROUGHPUT)` / `ExpertParallelRuntime` | `ht_runtime.cc` | `ht/kernels/*` | `FLAT` | -The `ht/` directory is active source code for the HT backend in the current -build. +Shared internal headers live in `include/`. The previous LL implementation is +kept in `legacy/low_latency.cu` for reference and is not compiled. + +The LL runtime uses one `low_latency_num_blocks` setting. Its default is 130: +dispatch launches 128 workers plus scheduler/notify blocks, while combine +launches 128 workers. `RANK_LOCAL_REDUCE` is the default combine mode; +`DIRECT_SEND` preserves bit-exact top-k reduction order. + +The LL payload layout reserves optional scale storage for future quantization, +but the active LL kernels currently accept BF16 inputs and expert outputs only. ### MSCCL++ EP LL vs NCCL EP LL vs DeepEP elastic dispatch payloads diff --git a/src/ext/ep/config.hpp b/src/ext/ep/config.hpp index b92b6563..3419c4fb 100644 --- a/src/ext/ep/config.hpp +++ b/src/ext/ep/config.hpp @@ -143,7 +143,6 @@ struct PayloadView { struct Buffer { void* dispatchData_; void* combineData_; - mscclpp::LL8Packet* combineReadyPackets_; }; struct Layout { @@ -154,15 +153,12 @@ struct Layout { const PayloadView<__bfloat16> bf16Payload(hidden, numTopk); const PayloadView<__nv_fp8_storage_t, float> fp8Payload(hidden, numTopk, 128); const size_t dispatchMetadataBytes = - configAlign(static_cast(2 * numRanks + numExperts) * sizeof(uint64_t), 128); + configAlign(static_cast(numRanks + numExperts) * sizeof(uint64_t), 128); const size_t dispatchPayloadStride = configAlign(std::max(bf16Payload.numBytes_, fp8Payload.numBytes_), 128); const size_t dispatchBufferBytes = dispatchMetadataBytes + static_cast(numRanks) * maxTokensPerRank * dispatchPayloadStride; - const size_t combineControlBytes = - configAlign(static_cast(numRanks) * sizeof(mscclpp::LL8Packet), 128); - const size_t combineBufferBytes = - combineControlBytes + static_cast(numExperts) * maxTokensPerRank * hidden * sizeof(__bfloat16); + const size_t combineBufferBytes = static_cast(numExperts) * maxTokensPerRank * hidden * sizeof(__bfloat16); const size_t bufferBytes = configAlign(std::max(dispatchBufferBytes, combineBufferBytes), 128); totalBytes_ = 2 * bufferBytes; @@ -172,8 +168,7 @@ struct Layout { auto* bufferBase = base + static_cast(bufferIdx) * bufferBytes; buffers_[bufferIdx] = { .dispatchData_ = bufferBase, - .combineData_ = bufferBase + combineControlBytes, - .combineReadyPackets_ = reinterpret_cast(bufferBase), + .combineData_ = bufferBase, }; } } diff --git a/src/ext/ep/ht_runtime.cc b/src/ext/ep/ht_runtime.cc index ffbabfc5..9a56d441 100644 --- a/src/ext/ep/ht_runtime.cc +++ b/src/ext/ep/ht_runtime.cc @@ -1272,7 +1272,7 @@ void MoEHighThroughputRuntime::intranodeCombine(void* combinedX, float* combined // ----------------------------------------------------------------------------- // Internode (NVLink + RDMA) high-throughput path. Ported from DeepEP // `csrc/deep_ep.cpp`; the kernels it drives are in -// `src/ext/ep/kernels/internode.cu`. Validated end-to-end on 2 x H100 x 8 +// `src/ext/ep/ht/kernels/internode.cu`. Validated end-to-end on 2 x H100 x 8 // via `test/python/ep/test_internode_multirank.py`. De-torched the same // way as the intranode path: tensor params became raw pointers + size ints, // output tensors became caller pointers, the EventHandle / async / record_stream diff --git a/src/ext/ep/include/api.cuh b/src/ext/ep/include/api.cuh index 3d0d6634..d6792962 100644 --- a/src/ext/ep/include/api.cuh +++ b/src/ext/ep/include/api.cuh @@ -222,6 +222,8 @@ struct Workload { struct CommContext { /// Base address of the locally-registered RDMA buffer. void* rdmaBufferBase_; + /// Base memory channel handles used only for signal/wait synchronization. + mscclpp::BaseMemoryChannelDeviceHandle* baseMemoryChannels_; /// Peer-mapped base addresses. void* const* peerBases_; /// Maximum shared memory available to one block after opt-in. @@ -236,6 +238,12 @@ struct CommContext { int numRanks_; }; +/// Return the optimized low-latency workspace size. +/// @param[in] numRanks Total number of ranks. +/// @param[in] numExperts Total number of experts. +/// @return Required workspace bytes. +size_t workspaceSize(int numRanks, int numExperts); + /// Low-latency dispatch that distributes tokens to experts across ranks. /// @param[out] output Expert-major packed output /// [num_local_experts, num_ranks * max_tokens_per_rank, hidden]. @@ -264,7 +272,6 @@ void dispatch(void* output, int* outputSrcInfo, int64_t* outputLayout, int* outp /// @param[in] layoutRange Per-[local expert, source rank] packed count and offset. /// @param[in] workload Per-call workload dimensions. /// @param[in,out] recvBuffer Current symmetric ping-pong buffer receiving partials or expert rows. -/// @param[in,out] readyPackets One readiness packet per sender rank in the current symmetric buffer. /// @param[in] dispatchRecvBuffer Previous dispatch buffer containing rewritten routing metadata. /// @param[in] comm Persistent communication context. /// @param[in,out] workspace Persistent dispatch metadata plus the combine device barrier. @@ -272,9 +279,8 @@ void dispatch(void* output, int* outputSrcInfo, int64_t* outputLayout, int* outp /// @param[in] mode Combine algorithm. /// @param[in] stream CUDA stream. void combine(void* output, const void* input, const int64_t* topkIdx, const float* topkWeights, const int* srcInfo, - const int64_t* layoutRange, const Workload& workload, void* recvBuffer, mscclpp::LL8Packet* readyPackets, - void* dispatchRecvBuffer, const CommContext& comm, void* workspace, int numBlocks, CombineMode mode, - cudaStream_t stream); + const int64_t* layoutRange, const Workload& workload, void* recvBuffer, void* dispatchRecvBuffer, + const CommContext& comm, void* workspace, int numBlocks, CombineMode mode, cudaStream_t stream); } // namespace low_latency diff --git a/src/ext/ep/include/device_helpers.cuh b/src/ext/ep/include/device_helpers.cuh index c026024e..4784a3ff 100644 --- a/src/ext/ep/include/device_helpers.cuh +++ b/src/ext/ep/include/device_helpers.cuh @@ -66,17 +66,6 @@ __device__ __forceinline__ void memory_fence_gpu() { asm volatile("fence.acq_rel __device__ __forceinline__ void memory_fence_cta() { asm volatile("fence.acq_rel.cta;" ::: "memory"); } -MSCCLPP_DEVICE_INLINE void publishLl8Packet(mscclpp::LL8Packet *packet, uint32_t value, uint32_t flag) { - memory_fence(); - packet->write(value, flag); -} - -MSCCLPP_DEVICE_INLINE uint32_t waitLl8Packet(const mscclpp::LL8Packet *packet, uint32_t flag) { - const uint32_t value = packet->read(flag, -1); - memory_fence(); - return value; -} - __device__ __forceinline__ void *peerBufferPtr(void *localBuffer, void *localBufferBase, void *peerBufferBase) { if (localBufferBase == nullptr) return peerBufferBase; const auto offset = reinterpret_cast(localBuffer) - reinterpret_cast(localBufferBase); diff --git a/src/ext/ep/low_latency/combine.cu b/src/ext/ep/low_latency/combine.cu index 383f805c..dc186f27 100644 --- a/src/ext/ep/low_latency/combine.cu +++ b/src/ext/ep/low_latency/combine.cu @@ -10,7 +10,8 @@ namespace mscclpp { namespace ep { -namespace low_latency_opt { +namespace low_latency { +namespace detail { constexpr int CombineNWarps = 32; constexpr int CombineNThreads = CombineNWarps * WARP_SIZE; @@ -42,8 +43,8 @@ MSCCLPP_DEVICE_INLINE int4 reduceWeightedBf16x8(const void* expertOutput, int ro const int sourceRowOffset = warpBroadcast(rowOffset, topkLane); if (sourceRowOffset < 0) continue; const float sourceWeight = warpBroadcast(weight, topkLane); - const int4 packed = ld_nc_global(reinterpret_cast(expertOutput) + - static_cast(sourceRowOffset) * HiddenInt4 + hiddenIdx); + const int4 packed = + reinterpret_cast(expertOutput)[static_cast(sourceRowOffset) * HiddenInt4 + hiddenIdx]; const auto* values = reinterpret_cast(&packed); #pragma unroll for (int pairIdx = 0; pairIdx < Bf16PairsPerInt4; ++pairIdx) { @@ -70,9 +71,8 @@ MSCCLPP_DEVICE_INLINE int4 reduceRankPartialsBf16x8(const void* combineRecvBuffe for (int topkLane = 0; topkLane < nTopk; ++topkLane) { const int partialRank = warpBroadcast(partialRankCandidate, topkLane); if (partialRank < 0) continue; - const int4 packed = - ld_nc_global(reinterpret_cast(combineRecvBuffer) + - (static_cast(partialRank) * maxTokensPerRank + tokenIdx) * HiddenInt4 + hiddenIdx); + const int4 packed = reinterpret_cast( + combineRecvBuffer)[(static_cast(partialRank) * maxTokensPerRank + tokenIdx) * HiddenInt4 + hiddenIdx]; const auto* values = reinterpret_cast(&packed); #pragma unroll for (int pairIdx = 0; pairIdx < Bf16PairsPerInt4; ++pairIdx) { @@ -128,8 +128,8 @@ MSCCLPP_DEVICE_INLINE void sendRankReducedPartials(const void* expertOutput, int 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; + const int rowOffset = laneId < nTopk ? payloadView.topKIndices(sourcePayload)[laneId] : -1; + const float weight = laneId < nTopk ? payloadView.topKValues(sourcePayload)[laneId] : 0.0f; if (rowOffset >= 0) EP_DEVICE_ASSERT(rowOffset < nExpertOutputRows); int4 reduced[ChunksPerThread] = {}; @@ -154,8 +154,7 @@ MSCCLPP_DEVICE_INLINE void sendRankReducedPartials(const void* expertOutput, int if (threadId == 0) { fenceProxyAsyncSharedCta(); - const int sourceTokenIdx = - ld_nc_global(payloadView.srcTokenGlobalIdx(sourcePayload)) - sourceRank * maxTokensPerRank; + const int sourceTokenIdx = *payloadView.srcTokenGlobalIdx(sourcePayload) - sourceRank * maxTokensPerRank; EP_DEVICE_ASSERT(sourceTokenIdx >= 0 && sourceTokenIdx < maxTokensPerRank); void* destinationBuffer = sourceRank == rank ? combineRecvBuffer @@ -222,7 +221,7 @@ MSCCLPP_DEVICE_INLINE void sendExpertRowsDirect(const void* expertOutput, const } EP_DEVICE_ASSERT(sourceRank < nRanks); const int inputRowOffset = localExpertIdx * nOutputSlotsPerExpert + expertTokenIdx; - const int sourceTokenIdx = ld_nc_global(srcInfo + inputRowOffset); + const int sourceTokenIdx = srcInfo[inputRowOffset]; EP_DEVICE_ASSERT(sourceTokenIdx >= 0 && sourceTokenIdx < maxTokensPerRank); const auto* inputRow = reinterpret_cast(expertOutput) + static_cast(inputRowOffset) * HiddenBytes; @@ -243,21 +242,24 @@ MSCCLPP_DEVICE_INLINE void sendExpertRowsDirect(const void* expertOutput, const } } -MSCCLPP_DEVICE_INLINE void combineSynchronize(mscclpp::LL8Packet* readyPackets, void* rdmaBufferBase, - void* const* peerRecvBuffers, int rank, int nRanks, uint32_t readyFlag) { +MSCCLPP_DEVICE_INLINE void combineSynchronize(mscclpp::BaseMemoryChannelDeviceHandle* baseMemoryChannels, + mscclpp::DeviceSemaphore* localReady, int rank, int nRanks) { const int threadId = static_cast(threadIdx.x); if (blockIdx.x == 0 && threadId < nRanks) { const int peerRank = threadId; - void* destination = - peerRank == rank ? readyPackets : peerBufferPtr(readyPackets, rdmaBufferBase, peerRecvBuffers[peerRank]); - publishLl8Packet(reinterpret_cast(destination) + rank, 1, readyFlag); - waitLl8Packet(readyPackets + peerRank, readyFlag); + if (peerRank == rank) { + localReady->release(); + localReady->acquire(); + } else { + baseMemoryChannels[peerRank].signal(); + baseMemoryChannels[peerRank].wait(-1); + } } } template