diff --git a/python/mscclpp/ep/_cpp.py b/python/mscclpp/ep/_cpp.py index dddf0038..310d0174 100644 --- a/python/mscclpp/ep/_cpp.py +++ b/python/mscclpp/ep/_cpp.py @@ -18,6 +18,8 @@ Config = getattr(_cpp, "Config", None) def get_low_latency_rdma_size_hint( - num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int + num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int, num_topk: int ) -> int: - return _cpp.get_low_latency_rdma_size_hint(num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts) + return _cpp.get_low_latency_rdma_size_hint( + num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts, num_topk + ) diff --git a/python/mscclpp/ep/low_latency.py b/python/mscclpp/ep/low_latency.py index 737b3acd..e829a58d 100644 --- a/python/mscclpp/ep/low_latency.py +++ b/python/mscclpp/ep/low_latency.py @@ -117,7 +117,7 @@ class LowLatencyBackend: self.dispatch_requires_quantization = self.quant_dtype is not None num_rdma_bytes = get_low_latency_rdma_size_hint( - self.max_tokens_per_rank, self.hidden_size, self.world_size, self.num_experts + self.max_tokens_per_rank, self.hidden_size, self.world_size, self.num_experts, self.topk ) self._dispatch_scales: Optional[torch.Tensor] = None self._dispatch_src_info: Optional[torch.Tensor] = None @@ -131,9 +131,9 @@ class LowLatencyBackend: mode=self.mode, num_qps_per_rank=config.num_rdma_qps_per_rank, ) - # LL always uses the RDMA transport, but a single-node LL job is not - # internode topology-wise. num_rdma_ranks > 1 iff world_size spans more - # than one local NVLink domain. + # LL uses the registered symmetric buffer for both IPC/NVLink and RDMA-backed + # modes. A single-node LL job is not internode topology-wise. + # num_rdma_ranks > 1 iff world_size spans more than one local NVLink domain. self._is_internode = self._runtime.get_num_rdma_ranks() > 1 def is_available(self) -> bool: @@ -165,6 +165,7 @@ class LowLatencyBackend: self._runtime.cpp_runtime.dispatch( input.data_ptr(), topk_ids.data_ptr(), + weights.data_ptr(), out_buf.data_ptr(), 0 if packed_scales is None else packed_scales.data_ptr(), src_info.data_ptr(), diff --git a/src/ext/ep/README.md b/src/ext/ep/README.md index c6ac8e4a..f264cf23 100644 --- a/src/ext/ep/README.md +++ b/src/ext/ep/README.md @@ -494,58 +494,82 @@ Env knobs: The `ht/` directory is active source code for the HT backend in the current build. -### NCCL EP LL vs DeepEP elastic dispatch +### MSCCL++ EP LL vs NCCL EP LL vs DeepEP elastic dispatch payloads -NCCL EP LL dispatch uses multiple SMs to process tokens. Most warps stage -or quantize token payloads, top-k lanes write the per-token routing header, -and a final warp computes rank counts. Sends are deduplicated by destination -rank: if multiple top-k experts for the same token live on the same rank, the -token is sent to that rank only once. After the token payloads for a rank are -issued, NCCL EP sends a count signal, and the receiver waits for that count -before unpacking the rank buffer into the output layout. - -The NCCL EP dispatch wire layout depends on the transport. RDMA uses one -interleaved message per compact slot: +MSCCL++ EP LL dispatch sends one message per routed expert. The destination +receive buffer is partitioned by local expert, source rank, and slot, so the +buffer position encodes the destination expert and source rank. The inline +header only carries the source token's local row index: ```text +BF16 dispatch message: +[int4 header: src_token_idx][BF16 hidden payload] + +FP8 dispatch message: +[int4 header: src_token_idx][FP8 hidden payload][FP32 scale per 128 elements] +``` + +`topk_ids` and `topk_weights` are not sent in the LL dispatch payload. They are +kept in the dispatch handle on the source rank and are used by `combine` to +weight the returned expert outputs. Dispatch sends a small per-expert count +signal separately so the receiver knows how many token messages arrived. + +NCCL EP LL sends one message per destination rank after rank-level +deduplication. Its RDMA path stores each slot as an interleaved message, while +its NVLink/P2P path stores headers and payloads in split sections inside the +same per-source-rank receive region: + +```text +RDMA: [DispatchHdr][hidden payload][optional FP8 scales] -DispatchHdr(EXPERT_MAJOR) = [token_id][expert_id[0] ... expert_id[num_topk-1]] -DispatchHdr(RANK_MAJOR) = [token_id][(topk_weight, expert_id)[0] ... - (topk_weight, expert_id)[num_topk-1]] +NVLink/P2P: +[hdr0][hdr1]...[hdrN][payload0][payload1]...[payloadN] ``` -NVLink/P2P uses a split layout inside the same per-source-rank receive region: +The header carries `token_id` plus one routing entry per top-k choice: ```text -[hdr0][hdr1]...[hdrN][payload0][payload1]...[payloadN] -hdr_i = DispatchHdr(EXPERT_MAJOR or RANK_MAJOR) -payload_i = [hidden payload_i][optional FP8 scales_i] +DispatchHdr(EXPERT_MAJOR) = [token_id][expert_id[num_topk]] +DispatchHdr(RANK_MAJOR) = [token_id][(topk_weight, expert_id)[num_topk]] ``` -The header carries `token_id` and one routing entry per top-k choice. In -rank-major layout, each routing entry also carries `topk_weight`; in -expert-major layout it only carries `expert_id`. The final warp sends only the -per-rank token count signal, not token payload or routing metadata. +The NCCL EP expert-major header is `align16(4 + num_topk * 2)` bytes because +`expert_id` is `uint16_t`. The rank-major header is +`align16(4 + num_topk * 8)` bytes because each routing entry stores an FP32 +weight and a `uint16_t` expert ID plus padding. -DeepEP elastic dispatch uses a different notify/data pipeline. Notify warps -first compute rank and expert counts, with rank-level deduplication, exchange -the counts, and build prefix sums. Dispatch warps then send compact token -buffers with a single token layout: +DeepEP elastic dispatch uses a different notify/data pipeline. It deduplicates +by destination rank: if multiple top-k experts for the same token live on the +same rank, the hidden payload is sent to that rank only once. Because the +receiver must recover all local expert hits from that single token copy, the +token buffer carries the full top-k metadata: ```text [hidden payload][optional scale-factor packs][metadata] -metadata = [topk_idx[num_topk]][optional topk_weights[num_topk]][src_token_global_idx] +metadata = [topk_idx[num_topk]][optional topk_weights[num_topk]][src_token_global_idx][linked_list_idx[num_topk]] ``` -The common DeepEP token layout reserves additional linked-list metadata slots, -but non-hybrid dispatch does not use them. After a GPU barrier guarantees data -arrival, the copy epilogue waits on the programmatic launch dependency and -copies data into `recv_x`, `recv_sf`, `recv_topk_idx`, `recv_topk_weights`, -and source metadata. +`topk_idx` entries are 32-bit expert IDs, `topk_weights` entries are FP32 routing +weights, and `src_token_global_idx = src_rank * num_max_tokens_per_rank + +src_local_token_idx`. The common DeepEP token layout reserves the linked-list +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. -DeepEP elastic combine replays the dispatch metadata to send expert outputs -back to owner ranks. It uses barriers and per-channel tail signals, followed -by a reduce epilogue, rather than a single per-rank finish flag. Top-k weights -are stored as token-buffer metadata for non-expanded layouts and are used by -the final reduction. +For BF16 dispatch with `num_tokens = 128`, `hidden = 7168`, and +`num_experts = 256`, the expected per-source-rank dispatch payload is: + +| Case | Sends per token | Bytes per send | Expected dispatch bytes | +|------|-----------------|----------------|-------------------------| +| MSCCL++ LL, `topk=8` | `8` | `16 + 7168 * 2 = 14,352 B` | `128 * 8 * 14,352 = 14,696,448 B` (`14.70 MB`, `14.02 MiB`) | +| NCCL EP LL expert-major, 16 GPUs, `topk=8` | `16 * (1 - C(240, 8) / C(256, 8)) = 6.523` | `align16(4 + 8*2) + 7168 * 2 = 14,368 B` | `128 * 6.523 * 14,368 ~= 11,996,989 B` (`12.00 MB`, `11.44 MiB`) | +| NCCL EP LL rank-major, 16 GPUs, `topk=8` | `16 * (1 - C(240, 8) / C(256, 8)) = 6.523` | `align16(4 + 8*8) + 7168 * 2 = 14,416 B` | `128 * 6.523 * 14,416 ~= 12,037,069 B` (`12.04 MB`, `11.48 MiB`) | +| DeepEP elastic, 16 GPUs, `topk=8` | `16 * (1 - C(240, 8) / C(256, 8)) = 6.523` | `7168 * 2 + align32(8*4 + 8*4 + 4 + 8*4) = 14,464 B` | `128 * 6.523 * 14,464 ~= 12,077,148 B` (`12.08 MB`, `11.52 MiB`) | +| MSCCL++ LL, `topk=4` | `4` | `16 + 7168 * 2 = 14,352 B` | `128 * 4 * 14,352 = 7,348,224 B` (`7.35 MB`, `7.01 MiB`) | +| NCCL EP LL expert-major, 32 GPUs, `topk=4` | `32 * (1 - C(248, 4) / C(256, 4)) = 3.838` | `align16(4 + 4*2) + 7168 * 2 = 14,352 B` | `128 * 3.838 * 14,352 ~= 7,050,391 B` (`7.05 MB`, `6.72 MiB`) | +| NCCL EP LL rank-major, 32 GPUs, `topk=4` | `32 * (1 - C(248, 4) / C(256, 4)) = 3.838` | `align16(4 + 4*8) + 7168 * 2 = 14,384 B` | `128 * 3.838 * 14,384 ~= 7,066,111 B` (`7.07 MB`, `6.74 MiB`) | +| DeepEP elastic, 32 GPUs, `topk=4` | `32 * (1 - C(248, 4) / C(256, 4)) = 3.838` | `7168 * 2 + align32(4*4 + 4*4 + 4 + 4*4) = 14,400 B` | `128 * 3.838 * 14,400 ~= 7,073,971 B` (`7.07 MB`, `6.75 MiB`) | + +The table counts token data plus inline per-token metadata. It excludes the +small count, tail, barrier, and completion signals. diff --git a/src/ext/ep/bindings.cpp b/src/ext/ep/bindings.cpp index 09e7b2a3..69d79f58 100644 --- a/src/ext/ep/bindings.cpp +++ b/src/ext/ep/bindings.cpp @@ -63,7 +63,9 @@ std::string bytesToString(nb::handle h) { NB_MODULE(mscclpp_ep_cpp, m) { m.doc() = "MSCCL++ Expert-Parallel (MoE dispatch/combine) extension"; - m.def("get_low_latency_rdma_size_hint", &mscclpp::ep::getLowLatencyRdmaSizeHint); + m.def("get_low_latency_rdma_size_hint", &mscclpp::ep::getLowLatencyRdmaSizeHint, + nb::arg("num_max_dispatch_tokens_per_rank"), nb::arg("hidden"), nb::arg("num_ranks"), nb::arg("num_experts"), + nb::arg("num_topk")); nb::module_::import_("mscclpp._mscclpp"); @@ -88,21 +90,22 @@ NB_MODULE(mscclpp_ep_cpp, m) { [](const mscclpp::ep::MoERuntime& self) { return stringToBytes(self.getLocalIpcHandle()); }) .def( "dispatch", - [](mscclpp::ep::MoERuntime& self, uintptr_t inputPtr, uintptr_t topkIdxPtr, uintptr_t outputPtr, - uintptr_t outputScalesPtr, uintptr_t outputSrcInfoPtr, uintptr_t outputLayoutRangePtr, + [](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) { 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)), - numTokens, hidden, numTopk, numMaxDispatchTokensPerRank, numExperts, requiresQuantization, outputLayout, - stream(streamPtr)); + reinterpret_cast(ptr(topkWeightsPtr)), numTokens, hidden, numTopk, numMaxDispatchTokensPerRank, + numExperts, requiresQuantization, outputLayout, stream(streamPtr)); }, - nb::arg("input_ptr"), nb::arg("topk_idx_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("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")) .def( "combine", [](mscclpp::ep::MoERuntime& self, uintptr_t expertOutputPtr, uintptr_t expertScalesPtr, uintptr_t topkIdxPtr, diff --git a/src/ext/ep/config.hpp b/src/ext/ep/config.hpp index 02897b4f..af61128e 100644 --- a/src/ext/ep/config.hpp +++ b/src/ext/ep/config.hpp @@ -5,6 +5,7 @@ #include #include #include +#include #include "kernels/configs.cuh" #include "kernels/exception.cuh" @@ -13,15 +14,135 @@ namespace mscclpp { namespace ep { template -dtype_t cellDiv(dtype_t a, dtype_t b) { +__host__ __device__ constexpr dtype_t configCellDiv(dtype_t a, dtype_t b) { return (a + b - 1) / b; } template -dtype_t align(dtype_t a, dtype_t b) { - return cellDiv(a, b) * b; +__host__ __device__ constexpr dtype_t configAlign(dtype_t a, dtype_t b) { + return configCellDiv(a, b) * b; } +// Low-latency rank-deduplicated dispatch payload layout: +// +// [data: DataType[hidden]] +// [optional scales: ScaleType[hidden / scale_block_size]] +// [topKIndices: int[topK]] +// [topKValues: float[topK]] +// [srcTokenGlobalIdx: int] +// +// 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 { + static constexpr bool kHasScales = !std::is_void_v; + + __host__ __device__ static constexpr int numScales([[maybe_unused]] int hidden, [[maybe_unused]] int scaleBlockSize) { + if constexpr (kHasScales) { + return hidden / scaleBlockSize; + } else { + return 0; + } + } + + __host__ __device__ static constexpr size_t hiddenBytes(int hidden) { + return static_cast(hidden) * sizeof(DataType); + } + + __host__ __device__ static constexpr size_t scaleOffset(int hidden) { + if constexpr (kHasScales) { + return configAlign(hiddenBytes(hidden), alignof(ScaleType)); + } else { + return 0; + } + } + + __host__ __device__ static constexpr size_t scaleBytes([[maybe_unused]] int hidden, + [[maybe_unused]] int scaleBlockSize) { + if constexpr (kHasScales) { + return static_cast(numScales(hidden, scaleBlockSize)) * sizeof(ScaleType); + } else { + return 0; + } + } + + __host__ __device__ static constexpr size_t metadataOffset(int hidden, [[maybe_unused]] int scaleBlockSize) { + if constexpr (kHasScales) { + return configAlign(scaleOffset(hidden) + scaleBytes(hidden, scaleBlockSize), alignof(int)); + } else { + return configAlign(hiddenBytes(hidden), alignof(int)); + } + } + + __host__ __device__ static constexpr 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) { + 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)) + : 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)) {} + + template + __device__ __forceinline__ 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); + } + + __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); + } + + __device__ __forceinline__ 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); + } + + __device__ __forceinline__ float* topKValues(void* base) const { + return reinterpret_cast(topKIndices(base) + topK); + } + + __device__ __forceinline__ const float* topKValues(const void* base) const { + return reinterpret_cast(topKIndices(base) + topK); + } + + __device__ __forceinline__ int* srcTokenGlobalIdx(void* base) const { + return reinterpret_cast(topKValues(base) + topK); + } +}; + struct LowLatencyBuffer { int numCleanInt = 0; @@ -55,10 +176,8 @@ struct LowLatencyLayout { return reinterpret_cast(reinterpret_cast(ptr) + count); } - LowLatencyLayout(void* rdmaBuffer, int numMaxDispatchTokensPerRank, int hidden, int numRanks, int numExperts) { - (void)numRanks; - const int numScales = hidden / 128; - + LowLatencyLayout(void* rdmaBuffer, int numMaxDispatchTokensPerRank, int hidden, int numRanks, int numExperts, + int numTopk) { // Dispatch and combine layout: // - 2 symmetric odd/even send buffer // - 2 symmetric odd/even receive buffers @@ -66,9 +185,9 @@ struct LowLatencyLayout { // Message sizes // NOTES: you should add a control `int4` for combine messages if you want to do data transformation - EP_HOST_ASSERT(numScales * static_cast(sizeof(float)) <= hidden); - size_t numBytesPerDispatchMsg = - sizeof(int4) + std::max(hidden * sizeof(nv_bfloat16), hidden + numScales * sizeof(float)); + 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); size_t numBytesPerCombineMsg = hidden * sizeof(nv_bfloat16); // Send buffer @@ -80,7 +199,7 @@ struct LowLatencyLayout { // Symmetric receive buffers // TODO: optimize memory usages - size_t dispatchRecvDataBufferBytes = numExperts * numMaxDispatchTokensPerRank * numBytesPerDispatchMsg; + size_t dispatchRecvDataBufferBytes = numRanks * numMaxDispatchTokensPerRank * numBytesPerDispatchMsg; size_t combineRecvBufferBytes = numExperts * numMaxDispatchTokensPerRank * numBytesPerCombineMsg; size_t recvBufferBytes = std::max(dispatchRecvDataBufferBytes, combineRecvBufferBytes); EP_HOST_ASSERT(recvBufferBytes % sizeof(int4) == 0); @@ -109,9 +228,11 @@ struct LowLatencyLayout { } }; -inline size_t getLowLatencyRdmaSizeHint(int numMaxDispatchTokensPerRank, int hidden, int numRanks, int numExperts) { - auto numBytes = LowLatencyLayout(nullptr, numMaxDispatchTokensPerRank, hidden, numRanks, numExperts).totalBytes; - return ((numBytes + NUM_BUFFER_ALIGNMENT_BYTES) / NUM_BUFFER_ALIGNMENT_BYTES) * NUM_BUFFER_ALIGNMENT_BYTES; +inline size_t getLowLatencyRdmaSizeHint(int numMaxDispatchTokensPerRank, int hidden, int numRanks, int numExperts, + int numTopk) { + auto numBytes = + LowLatencyLayout(nullptr, numMaxDispatchTokensPerRank, hidden, numRanks, numExperts, numTopk).totalBytes; + return configAlign(numBytes, NUM_BUFFER_ALIGNMENT_BYTES); } } // namespace ep diff --git a/src/ext/ep/ht/config.hpp b/src/ext/ep/ht/config.hpp index 5230b5bb..c3b46c31 100644 --- a/src/ext/ep/ht/config.hpp +++ b/src/ext/ep/ht/config.hpp @@ -55,7 +55,7 @@ struct Config { // Ceil up RDMA buffer size this->num_max_rdma_chunked_recv_tokens = - align(num_max_rdma_chunked_recv_tokens, num_max_rdma_chunked_send_tokens); + configAlign(num_max_rdma_chunked_recv_tokens, num_max_rdma_chunked_send_tokens); EP_HOST_ASSERT(num_max_rdma_chunked_send_tokens < num_max_rdma_chunked_recv_tokens); // NOTES: this assertion is related to RDMA lazy head update, we must ensure senders always have space to push EP_HOST_ASSERT(num_max_rdma_chunked_send_tokens <= num_max_rdma_chunked_recv_tokens / 2); diff --git a/src/ext/ep/kernels/api.cuh b/src/ext/ep/kernels/api.cuh index 796eb8b2..17168602 100644 --- a/src/ext/ep/kernels/api.cuh +++ b/src/ext/ep/kernels/api.cuh @@ -302,6 +302,7 @@ void cleanBuffers(int64_t* buffer0, int numInt0, int64_t* buffer1, int numInt1, /// @param outputCount Total count per expert. /// @param input Input tokens [num_tokens, hidden]. /// @param topkIdx Expert indices [num_tokens, num_topk]. +/// @param topkWeights Routing weights [num_tokens, num_topk]. /// @param config Dispatch configuration. /// @param currentBuffer Current iteration buffer set. /// @param nextBuffer Next iteration buffer set (for cleanup). @@ -310,9 +311,9 @@ void cleanBuffers(int64_t* buffer0, int numInt0, int64_t* buffer1, int numInt1, /// @param stream CUDA stream to launch the kernel on. /// @param phase Phase control (default: SEND_AND_RECV). void dispatch(void* output, float* outputScales, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, - const void* input, const int64_t* topkIdx, const DispatchConfig& config, const BufferSet& currentBuffer, - const BufferSet& nextBuffer, const TransportContext& transport, void* workspace, cudaStream_t stream, - Phase phase = SEND_AND_RECV); + const void* input, const int64_t* topkIdx, const float* topkWeights, const DispatchConfig& config, + const BufferSet& currentBuffer, const BufferSet& nextBuffer, const TransportContext& transport, + void* workspace, cudaStream_t stream, Phase phase = SEND_AND_RECV); /// Low-latency combine kernel that aggregates expert outputs back to tokens. /// @param output Combined output [num_combined_tokens, hidden]. diff --git a/src/ext/ep/kernels/low_latency.cu b/src/ext/ep/kernels/low_latency.cu index 53e6f62f..8d970016 100644 --- a/src/ext/ep/kernels/low_latency.cu +++ b/src/ext/ep/kernels/low_latency.cu @@ -26,6 +26,7 @@ #include #include +#include "../config.hpp" #include "api.cuh" #include "common.cuh" #include "configs.cuh" @@ -106,6 +107,11 @@ struct DispatchOutputVec { template constexpr bool kDispatchNeedsScales = kInputDType != kOutputDType; +template +using LowLatencyDispatchPayloadView = + LowLatencyPackedPayloadView::Type, + std::conditional_t>; + template MSCCLPP_DEVICE_INLINE typename DispatchOutputVec::Type dispatchConvert( const int4& inputValue, float* scaleOut, int laneId) { @@ -149,25 +155,29 @@ template = numExperts) return; + using OutputType = typename DispatchDTypeTraits::Type; const auto sourceRank = responsibleExpertIdx / numLocalExperts; const auto localExpertIdx = responsibleExpertIdx % numLocalExperts; + const auto globalExpertIdx = rank * numLocalExperts + localExpertIdx; (void)dispatchLayout; // EXPERT_MAJOR and FLAT share the same physical order. - const auto stagedMsgBase = reinterpret_cast(stagedRecv) + - localExpertIdx * numRanks * numMaxDispatchTokensPerRank * numBytesPerMsg + - sourceRank * numMaxDispatchTokensPerRank * numBytesPerMsg; + const auto stagedMsgBase = + reinterpret_cast(stagedRecv) + sourceRank * numMaxDispatchTokensPerRank * numBytesPerMsg; const auto outputPayload = reinterpret_cast(output) + localExpertIdx * numRanks * numMaxDispatchTokensPerRank * hiddenInt4; const auto packedSrcInfo = outputSrcInfo + localExpertIdx * numRanks * numMaxDispatchTokensPerRank; const auto outputLayout = outputLayoutRange + localExpertIdx * numRanks; + const LowLatencyDispatchPayloadView payload(hiddenBytes / sizeof(OutputType), numTopk); __shared__ int sharedNumRecvTokens[kNumWarpGroups]; + __shared__ int sharedNumExpertTokens[kNumWarpGroups]; __shared__ int sharedRecvTokenBeginIdx[kNumWarpGroups]; + __shared__ int sharedCopyCounter[kNumWarpGroups]; int numRecvTokens; int recvTokenBeginIdx; @@ -175,39 +185,74 @@ MSCCLPP_DEVICE_INLINE void dispatchRecv(void* output, float* outputScales, int* if (subWarpId == 0 and laneId == 0) { const auto raw = waitSignalNonZero(stagedRecvCountBuffer + localExpertIdx * numRanks + sourceRank); numRecvTokens = static_cast(-raw - 1); - recvTokenBeginIdx = atomicAdd(outputCount + localExpertIdx, numRecvTokens); sharedNumRecvTokens[warpGroupId] = numRecvTokens; - sharedRecvTokenBeginIdx[warpGroupId] = recvTokenBeginIdx; - outputLayout[sourceRank] = pack2(numRecvTokens, recvTokenBeginIdx); } asm volatile("bar.sync %0, %1;" ::"r"(warpGroupId + 2), "r"(kNumWarpsPerGroup * WARP_SIZE)); numRecvTokens = sharedNumRecvTokens[warpGroupId]; + + if (subWarpId == 0) { + int localCount = 0; + for (int i = laneId; i < numRecvTokens; i += WARP_SIZE) { + const auto stagedMsg = stagedMsgBase + i * numBytesPerMsg; + const auto stagedTopkIdx = payload.topKIndices(stagedMsg); + bool tokenHitsExpert = false; + for (int topkId = 0; topkId < numTopk; ++topkId) { + tokenHitsExpert |= ld_nc_global(stagedTopkIdx + topkId) == globalExpertIdx; + } + localCount += tokenHitsExpert ? 1 : 0; + } + const auto numExpertTokens = warp_reduce_sum(localCount); + if (laneId == 0) { + recvTokenBeginIdx = atomicAdd(outputCount + localExpertIdx, numExpertTokens); + sharedNumExpertTokens[warpGroupId] = numExpertTokens; + sharedRecvTokenBeginIdx[warpGroupId] = recvTokenBeginIdx; + sharedCopyCounter[warpGroupId] = 0; + outputLayout[sourceRank] = pack2(numExpertTokens, recvTokenBeginIdx); + } + } + asm volatile("bar.sync %0, %1;" ::"r"(warpGroupId + 2), "r"(kNumWarpsPerGroup * WARP_SIZE)); + if (sharedNumExpertTokens[warpGroupId] == 0) return; recvTokenBeginIdx = sharedRecvTokenBeginIdx[warpGroupId]; EP_DEVICE_ASSERT(numScales <= 64); for (int i = subWarpId; i < numRecvTokens; i += kNumWarpsPerGroup) { - const auto stagedMsg = reinterpret_cast(stagedMsgBase + i * numBytesPerMsg); - if (laneId == 0) packedSrcInfo[recvTokenBeginIdx + i] = ld_nc_global(stagedMsg); + const auto stagedMsg = stagedMsgBase + i * numBytesPerMsg; + const auto stagedTopkIdx = payload.topKIndices(stagedMsg); + const auto stagedSrcTokenGlobalIdx = payload.srcTokenGlobalIdx(stagedMsg); + bool tokenHitsExpert = false; + for (int topkId = 0; topkId < numTopk; ++topkId) { + tokenHitsExpert |= ld_nc_global(stagedTopkIdx + topkId) == globalExpertIdx; + } + if (!tokenHitsExpert) continue; + + int outputTokenIdx = laneId == 0 ? atomicAdd(sharedCopyCounter + warpGroupId, 1) : 0; + outputTokenIdx = __shfl_sync(0xffffffff, outputTokenIdx, 0); + outputTokenIdx += recvTokenBeginIdx; + if (laneId == 0) { + const auto srcTokenGlobalIdx = ld_nc_global(stagedSrcTokenGlobalIdx); + packedSrcInfo[outputTokenIdx] = srcTokenGlobalIdx % numMaxDispatchTokensPerRank; + } __syncwarp(); - const auto stagedPayload = reinterpret_cast(reinterpret_cast(stagedMsg) + sizeof(int4)); - const auto outputPayloadRow = outputPayload + (recvTokenBeginIdx + i) * hiddenInt4; + const auto stagedPayload = payload.template data(stagedMsg); + const auto outputPayloadRow = outputPayload + outputTokenIdx * hiddenInt4; UNROLLED_WARP_COPY(7, laneId, hiddenInt4, outputPayloadRow, stagedPayload, ld_nc_global, st_na_global); - copyScales(outputScales, stagedPayload, localExpertIdx, recvTokenBeginIdx + i, numRanks, + copyScales(outputScales, stagedPayload, localExpertIdx, outputTokenIdx, numRanks, numMaxDispatchTokensPerRank, hiddenBytes, numScales, laneId); } } template -MSCCLPP_DEVICE_INLINE void dispatchSend(int* sharedNumTokensSentPerExpert, int* outputCount, void* stagedRecv, +MSCCLPP_DEVICE_INLINE void dispatchSend(int* sharedNumTokensSentPerRank, int* outputCount, void* stagedRecv, int64_t* stagedRecvCountBuffer, void* stagedSend, const void* input, - const int64_t* topkIdx, int* atomicCounterPerExpert, - int* atomicFinishCounterPerExpert, int64_t* nextClean, int numNextCleanInt, + const int64_t* topkIdx, const float* topkWeights, int* atomicCounterPerRank, + 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, 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; @@ -220,13 +265,10 @@ MSCCLPP_DEVICE_INLINE void dispatchSend(int* sharedNumTokensSentPerExpert, int* const auto responsibleExpertIdx = smId * kNumWarpGroups + warpGroupId; using SourceType = typename DispatchDTypeTraits::Type; - using OutputType = typename DispatchDTypeTraits::Type; constexpr int kNumPerChannels = 128; - const int numScales = hidden / kNumPerChannels; - const size_t hiddenBytes = hidden * sizeof(OutputType); using VecType = typename DispatchOutputVec::Type; - const size_t numBytesPerMsg = - sizeof(int4) + hiddenBytes + (kDispatchNeedsScales ? numScales * sizeof(float) : 0); + const LowLatencyDispatchPayloadView payload(hidden, numTopk); + const size_t numBytesPerMsg = payload.numBytes; const size_t numInt4PerMsg = numBytesPerMsg / sizeof(int4); if (warpId < numWarps - 1) { @@ -247,20 +289,34 @@ MSCCLPP_DEVICE_INLINE void dispatchSend(int* sharedNumTokensSentPerExpert, int* const auto tokenIdx = tokenBase + tokenGroupId * numSms; if (tokenIdx >= numTokens) continue; const auto inputVec = reinterpret_cast(input) + tokenIdx * hiddenSourceInt4; - const auto stagedSendMsg = - reinterpret_cast(reinterpret_cast(stagedSend) + tokenIdx * numBytesPerMsg); - const auto stagedSendPayload = - reinterpret_cast(reinterpret_cast(stagedSendMsg) + sizeof(int4)); - const auto stagedSendScales = - reinterpret_cast(reinterpret_cast(stagedSendPayload) + hiddenBytes); + const auto stagedSendMsg = reinterpret_cast(stagedSend) + tokenIdx * numBytesPerMsg; + const auto stagedSendPayload = payload.template data(stagedSendMsg); + float* stagedSendScales = nullptr; + if constexpr (kDispatchNeedsScales) { + stagedSendScales = payload.scaleFactors(stagedSendMsg); + } + const auto stagedSendTopkIdx = payload.topKIndices(stagedSendMsg); + const auto stagedSendTopkWeights = payload.topKValues(stagedSendMsg); + const auto stagedSendSrcTokenGlobalIdx = payload.srcTokenGlobalIdx(stagedSendMsg); - threadId == 0 ? (*stagedSendMsg = tokenIdx) : 0; + const bool isMetadataWarp = warpId == tokenGroupId * numTopk; + if (isMetadataWarp && laneId < numTopk) { + stagedSendTopkIdx[laneId] = static_cast(__ldg(topkIdx + tokenIdx * numTopk + laneId)); + stagedSendTopkWeights[laneId] = + topkWeights == nullptr ? 1.0f : __ldg(topkWeights + tokenIdx * numTopk + laneId); + } + (isMetadataWarp && laneId == 0) ? (*stagedSendSrcTokenGlobalIdx = rank * numMaxDispatchTokensPerRank + tokenIdx) + : 0; for (int i = threadId; i < hiddenSourceInt4; i += numPayloadThreads) { auto int4Value = __ldg(inputVec + i); + float* scaleOut = nullptr; + if constexpr (kDispatchNeedsScales) { + scaleOut = &stagedSendScales[i * kNumElemsPerRead / kNumPerChannels]; + } - stagedSendPayload[i] = dispatchConvert( - int4Value, &stagedSendScales[i * kNumElemsPerRead / kNumPerChannels], laneId); + stagedSendPayload[i] = + dispatchConvert(int4Value, scaleOut, laneId); } } asm volatile("bar.sync 1, %0;" ::"r"(numPayloadThreads)); @@ -276,32 +332,26 @@ MSCCLPP_DEVICE_INLINE void dispatchSend(int* sharedNumTokensSentPerExpert, int* const auto dstExpertIdx = tokenIdx < numTokens ? static_cast(__ldg(topkIdx + tokenIdx * numTopk + topkId)) : -1; if (dstExpertIdx < 0) continue; - const auto stagedSendMsg = - reinterpret_cast(reinterpret_cast(stagedSend) + tokenIdx * numBytesPerMsg); - int slotIdx = laneId == 0 ? atomicAdd(atomicCounterPerExpert + dstExpertIdx, 1) : 0; - slotIdx = __shfl_sync(0xffffffff, slotIdx, 0); const auto dstRank = dstExpertIdx / numLocalExperts; - const auto dstExpertLocalIdx = dstExpertIdx % numLocalExperts; + bool firstForRank = true; + for (int i = 0; i < topkId; ++i) { + const auto prevExpertIdx = static_cast(__ldg(topkIdx + tokenIdx * numTopk + i)); + if (prevExpertIdx >= 0 && prevExpertIdx / numLocalExperts == dstRank) firstForRank = false; + } + if (!firstForRank) continue; + + const auto stagedSendMsg = reinterpret_cast(stagedSend) + tokenIdx * numBytesPerMsg; + int slotIdx = laneId == 0 ? atomicAdd(atomicCounterPerRank + dstRank, 1) : 0; + slotIdx = __shfl_sync(0xffffffff, slotIdx, 0); const auto srcPtr = reinterpret_cast(stagedSendMsg); const auto dstPtr = reinterpret_cast(stagedRecv) + - dstExpertLocalIdx * numRanks * numMaxDispatchTokensPerRank * numBytesPerMsg + rank * numMaxDispatchTokensPerRank * numBytesPerMsg + slotIdx * numBytesPerMsg; if (dstRank != rank) { - if (peerRdmaBases != nullptr && isIpcPeer(rank, dstRank, ranksPerIpcDomain)) { - // Peer-mapped warp copy over NVLink (CUDA IPC). - const auto peerDst = peerMappedPtrOf(dstPtr, peerRdmaBases, rdmaBufferPtr, dstRank); - const auto* srcInt4Ptr = reinterpret_cast(srcPtr); - const auto* dstInt4Ptr = reinterpret_cast(peerDst); - UNROLLED_WARP_COPY(8, laneId, numInt4PerMsg, dstInt4Ptr, srcInt4Ptr, ld_nc_global, st_na_global); - } else { - // MSCCL++ port-channel PUT (lane 0 issues one request). - if (laneId == 0) { - const auto dstOff = portChannelOffsetOf(dstPtr, rdmaBufferPtr); - const auto srcOff = portChannelOffsetOf(srcPtr, rdmaBufferPtr); - portChannelHandles[dstExpertLocalIdx * numRanks + dstRank].put(dstOff, srcOff, numBytesPerMsg); - } - __syncwarp(); - } + EP_DEVICE_ASSERT(peerRdmaBases != nullptr && isIpcPeer(rank, dstRank, ranksPerIpcDomain)); + const auto peerDst = peerMappedPtrOf(dstPtr, peerRdmaBases, rdmaBufferPtr, dstRank); + const auto* srcInt4Ptr = reinterpret_cast(srcPtr); + const auto* dstInt4Ptr = reinterpret_cast(peerDst); + UNROLLED_WARP_COPY(8, laneId, numInt4PerMsg, dstInt4Ptr, srcInt4Ptr, ld_nc_global, st_na_global); } else { const auto* srcInt4Ptr = reinterpret_cast(srcPtr); const auto* dstInt4Ptr = reinterpret_cast(dstPtr); @@ -309,7 +359,7 @@ MSCCLPP_DEVICE_INLINE void dispatchSend(int* sharedNumTokensSentPerExpert, int* } __syncwarp(); - if (laneId == 0) atomicAddReleaseDevice(atomicFinishCounterPerExpert + dstExpertIdx, 1); + if (laneId == 0) atomicAddReleaseDevice(atomicFinishCounterPerRank + dstRank, 1); } } } else if (warpId == numWarps - 1) { @@ -319,48 +369,53 @@ MSCCLPP_DEVICE_INLINE void dispatchSend(int* sharedNumTokensSentPerExpert, int* EP_DEVICE_ASSERT(numSms > 1); if (smId == 0) { for (int i = laneId; i < numNextCleanInt; i += WARP_SIZE) nextClean[i] = 0; + for (int i = laneId; i < numLocalExperts; i += WARP_SIZE) outputCount[i] = 0; __syncwarp(); - for (int i = laneId; i < numExperts; i += WARP_SIZE) - atomicAddReleaseDevice(atomicFinishCounterPerExpert + i, FINISHED_SUM_TAG); + for (int i = laneId; i < numRanks; i += WARP_SIZE) + atomicAddReleaseDevice(atomicFinishCounterPerRank + i, FINISHED_SUM_TAG); } - int expertCount[kNumWarpGroups] = {0}; - const auto expertBeginIdx = smId * kNumWarpGroups; - const auto expertEndIdx = min(expertBeginIdx + kNumWarpGroups, numExperts); - - for (int i = laneId; i < numTokens * numTopk; i += WARP_SIZE) { - auto idx = static_cast(__ldg(topkIdx + i)); - if (idx >= expertBeginIdx and idx < expertEndIdx) expertCount[idx - expertBeginIdx]++; - } - -#pragma unroll - for (int i = expertBeginIdx; i < expertEndIdx; ++i) { - auto sum = warp_reduce_sum(expertCount[i - expertBeginIdx]); + const auto responsibleRank = responsibleExpertIdx; + if (responsibleRank < numRanks) { + int rankCount = 0; + for (int tokenIdx = laneId; tokenIdx < numTokens; tokenIdx += WARP_SIZE) { + bool tokenHitsRank = false; + for (int topkId = 0; topkId < numTopk; ++topkId) { + const auto expertIdx = static_cast(__ldg(topkIdx + tokenIdx * numTopk + topkId)); + if (expertIdx >= 0 && expertIdx / numLocalExperts == responsibleRank) tokenHitsRank = true; + } + rankCount += tokenHitsRank ? 1 : 0; + } + auto sum = warp_reduce_sum(rankCount); if (laneId == 0) { - sharedNumTokensSentPerExpert[i - expertBeginIdx] = sum; - atomicAddReleaseDevice(atomicFinishCounterPerExpert + i, FINISHED_SUM_TAG - sum); + sharedNumTokensSentPerRank[warpGroupId] = sum; + atomicAddReleaseDevice(atomicFinishCounterPerRank + responsibleRank, FINISHED_SUM_TAG - sum); } } } __syncthreads(); // Issue count sends - if (responsibleExpertIdx < numExperts and subWarpId == 0 and laneId == 0) { - const auto dstRank = responsibleExpertIdx / numLocalExperts; - const auto dstExpertLocalIdx = responsibleExpertIdx % numLocalExperts; - const auto numTokensSent = sharedNumTokensSentPerExpert[responsibleExpertIdx - smId * kNumWarpGroups]; + if (responsibleExpertIdx < numRanks and subWarpId == 0 and laneId == 0) { + const auto dstRank = responsibleExpertIdx; + const auto numTokensSent = sharedNumTokensSentPerRank[warpGroupId]; - while (ld_acquire_global(atomicFinishCounterPerExpert + responsibleExpertIdx) != FINISHED_SUM_TAG * 2); - auto* counterPtr = stagedRecvCountBuffer + dstExpertLocalIdx * numRanks + rank; - auto* portChannelHandle = dstRank == rank ? nullptr : portChannelHandles + dstExpertLocalIdx * numRanks + dstRank; - publishSingleWriterSignal(counterPtr, static_cast(-numTokensSent - 1), rank, dstRank, rdmaBufferPtr, - portChannelHandle, peerRdmaBases, ranksPerIpcDomain); + while (ld_acquire_global(atomicFinishCounterPerRank + dstRank) != FINISHED_SUM_TAG * 2); + for (int localExpertIdx = 0; localExpertIdx < numLocalExperts; ++localExpertIdx) { + auto* counterPtr = stagedRecvCountBuffer + localExpertIdx * numRanks + rank; + if (dstRank == rank) { + publishSignalDirect(counterPtr, static_cast(-numTokensSent - 1)); + } else { + EP_DEVICE_ASSERT(peerRdmaBases != nullptr && isIpcPeer(rank, dstRank, ranksPerIpcDomain)); + auto* peerCounterPtr = reinterpret_cast( + peerMappedPtrOf(reinterpret_cast(counterPtr), peerRdmaBases, rdmaBufferPtr, dstRank)); + publishSignalDirect(peerCounterPtr, static_cast(-numTokensSent - 1)); + } + } - atomicCounterPerExpert[responsibleExpertIdx] = 0; - atomicFinishCounterPerExpert[responsibleExpertIdx] = 0; - - if (dstRank == 0) outputCount[dstExpertLocalIdx] = 0; + atomicCounterPerRank[dstRank] = 0; + atomicFinishCounterPerRank[dstRank] = 0; } __syncwarp(); } @@ -369,10 +424,10 @@ template (blockIdx.x); const auto threadId = static_cast(threadIdx.x); const auto warpId = threadId / WARP_SIZE, laneId = get_lane_id(); @@ -388,15 +443,15 @@ __global__ __launch_bounds__(kNumWarpGroups* kNumWarpsPerGroup* WARP_SIZE, 1) vo const size_t hiddenInt4 = hiddenBytes / sizeof(int4); // Message package: hidden data, optional quantization scales, index at source - const size_t numBytesPerMsg = - sizeof(int4) + hiddenBytes + (kDispatchNeedsScales ? numScales * sizeof(float) : 0); + const LowLatencyDispatchPayloadView payload(hidden, numTopk); + const size_t numBytesPerMsg = payload.numBytes; EP_DEVICE_ASSERT(numBytesPerMsg % sizeof(int4) == 0); if (phases & LOW_LATENCY_SEND_PHASE) { __shared__ int sharedNumTokensSentPerExpert[kNumWarpGroups]; dispatchSend( sharedNumTokensSentPerExpert, outputCount, stagedRecv, stagedRecvCountBuffer, stagedSend, input, topkIdx, - atomicCounterPerExpert, atomicFinishCounterPerExpert, nextClean, numNextCleanInt, numTokens, + topkWeights, atomicCounterPerExpert, atomicFinishCounterPerExpert, nextClean, numNextCleanInt, numTokens, numMaxDispatchTokensPerRank, numTopk, hidden, numExperts, rank, numRanks, rdmaBufferPtr, portChannelHandles, peerRdmaBases, ranksPerIpcDomain); } @@ -409,8 +464,8 @@ __global__ __launch_bounds__(kNumWarpGroups* kNumWarpsPerGroup* WARP_SIZE, 1) vo dispatchRecv( output, outputScales, outputSrcInfo, outputLayoutRange, outputCount, stagedRecv, stagedRecvCountBuffer, - numMaxDispatchTokensPerRank, numExperts, numRanks, numLocalExperts, rank, numBytesPerMsg, hiddenInt4, hiddenBytes, - numScales, dispatchLayout, warpGroupId, subWarpId, laneId, responsibleExpertIdx); + numMaxDispatchTokensPerRank, numExperts, numRanks, numLocalExperts, rank, numTopk, numBytesPerMsg, hiddenInt4, + hiddenBytes, numScales, dispatchLayout, warpGroupId, subWarpId, laneId, responsibleExpertIdx); } constexpr int kDispatchNumMaxTopK = 9; @@ -429,6 +484,7 @@ struct DispatchLaunchArgs { void* stagedSend; const void* input; const int64_t* topkIdx; + const float* topkWeights; int* atomicCounterPerExpert; int* atomicFinishCounterPerExpert; int64_t* nextClean; @@ -451,13 +507,13 @@ struct DispatchLaunchArgs { template void launchDispatchKernel(const DispatchLaunchArgs& args) { auto dispatchFunc = dispatch; - CUDA_CHECK(cudaLaunchKernelEx(args.launchConfig, dispatchFunc, args.output, args.outputScales, args.outputSrcInfo, - args.outputLayoutRange, args.outputCount, args.stagedRecv, args.stagedRecvCountBuffer, - args.stagedSend, args.input, args.topkIdx, args.atomicCounterPerExpert, - args.atomicFinishCounterPerExpert, args.nextClean, args.numNextCleanInt, args.numTokens, - args.numMaxDispatchTokensPerRank, args.numTopk, args.hidden, args.numExperts, args.rank, - args.numRanks, args.phases, args.rdmaBufferPtr, args.portChannelHandles, - args.peerRdmaBases, args.ranksPerIpcDomain, args.dispatchLayout)); + CUDA_CHECK(cudaLaunchKernelEx( + args.launchConfig, dispatchFunc, args.output, args.outputScales, args.outputSrcInfo, args.outputLayoutRange, + args.outputCount, args.stagedRecv, args.stagedRecvCountBuffer, args.stagedSend, args.input, args.topkIdx, + args.topkWeights, args.atomicCounterPerExpert, args.atomicFinishCounterPerExpert, args.nextClean, + args.numNextCleanInt, args.numTokens, args.numMaxDispatchTokensPerRank, args.numTopk, args.hidden, + args.numExperts, args.rank, args.numRanks, args.phases, args.rdmaBufferPtr, args.portChannelHandles, + args.peerRdmaBases, args.ranksPerIpcDomain, args.dispatchLayout)); } template @@ -479,9 +535,9 @@ void launchDispatchForOutputDType(const DispatchLaunchArgs& args, DType outputDT } void dispatch(void* output, float* outputScales, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, - const void* input, const int64_t* topkIdx, const DispatchConfig& config, const BufferSet& currentBuffer, - const BufferSet& nextBuffer, const TransportContext& transport, void* workspace, cudaStream_t stream, - Phase phase) { + const void* input, const int64_t* topkIdx, const float* topkWeights, const DispatchConfig& config, + const BufferSet& currentBuffer, const BufferSet& nextBuffer, const TransportContext& transport, + void* workspace, cudaStream_t stream, Phase phase) { // Unpack configuration auto stagedRecv = currentBuffer.recvDataBuffer_; auto stagedRecvCountBuffer = currentBuffer.recvCountBuffer_; @@ -512,6 +568,7 @@ void dispatch(void* output, float* outputScales, int* outputSrcInfo, int64_t* ou EP_HOST_ASSERT(numTopk > 0); EP_HOST_ASSERT(numTopk <= kDispatchNumMaxTopK); EP_HOST_ASSERT(dispatchLayout == DispatchLayout::EXPERT_MAJOR || dispatchLayout == DispatchLayout::FLAT); + EP_HOST_ASSERT(transport.ipcReady_ && "low-latency rank-dedup dispatch requires IPC/NVLink"); auto atomicCounterPerExpert = reinterpret_cast(workspace); auto atomicFinishCounterPerExpert = atomicCounterPerExpert + numExperts; @@ -537,6 +594,7 @@ void dispatch(void* output, float* outputScales, int* outputSrcInfo, int64_t* ou .stagedSend = stagedSend, .input = input, .topkIdx = topkIdx, + .topkWeights = topkWeights, .atomicCounterPerExpert = atomicCounterPerExpert, .atomicFinishCounterPerExpert = atomicFinishCounterPerExpert, .nextClean = nextClean, diff --git a/src/ext/ep/moe_runtime.cc b/src/ext/ep/moe_runtime.cc index 287ae0fb..51569104 100644 --- a/src/ext/ep/moe_runtime.cc +++ b/src/ext/ep/moe_runtime.cc @@ -266,16 +266,17 @@ void MoERuntime::setup() { } void MoERuntime::dispatch(void* output, float* outputScales, int* outputSrcInfo, int64_t* outputLayout, - int* outputCount, const void* input, const int64_t* topkIdx, int numTokens, int hidden, - int numTopk, int numMaxDispatchTokensPerRank, int numExperts, bool requiresQuantization, - DispatchLayout dispatchLayout, cudaStream_t stream) { + 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) { 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"); - LowLatencyLayout layout(rdmaBufferPtr_, numMaxDispatchTokensPerRank, hidden, numRanks_, numExperts); + 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]; @@ -311,8 +312,9 @@ void MoERuntime::dispatch(void* output, float* outputScales, int* outputSrcInfo, .rank_ = rank_, .numRanks_ = numRanks_, .ranksPerIpcDomain_ = llRanksPerIpcDomain_}; - low_latency::dispatch(output, outputScales, outputSrcInfo, outputLayout, outputCount, input, topkIdx, config, - currentBuffer, nextBufferSet, transport, workspace_, stream, low_latency::SEND_AND_RECV); + low_latency::dispatch(output, outputScales, outputSrcInfo, outputLayout, outputCount, input, topkIdx, topkWeights, + config, currentBuffer, nextBufferSet, transport, workspace_, stream, + low_latency::SEND_AND_RECV); } void MoERuntime::combine(void* output, const void* input, const float* inputScales, const int64_t* topkIdx, @@ -322,8 +324,9 @@ void MoERuntime::combine(void* output, const void* input, const float* inputScal 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"); - LowLatencyLayout layout(rdmaBufferPtr_, numMaxDispatchTokensPerRank, hidden, numRanks_, numExperts); + 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]; diff --git a/src/ext/ep/moe_runtime.hpp b/src/ext/ep/moe_runtime.hpp index ae7fb726..c846db87 100644 --- a/src/ext/ep/moe_runtime.hpp +++ b/src/ext/ep/moe_runtime.hpp @@ -35,8 +35,8 @@ class MoERuntime { std::string getLocalIpcHandle() const; void dispatch(void* output, float* outputScales, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, - const void* input, const int64_t* topkIdx, int numTokens, int hidden, int numTopk, - int numMaxDispatchTokensPerRank, int numExperts, bool requiresQuantization, + 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); void combine(void* output, const void* input, const float* inputScales, const int64_t* topkIdx,