diff --git a/direct-rank-design.md b/direct-rank-design.md new file mode 100644 index 000000000..5c98ce514 --- /dev/null +++ b/direct-rank-design.md @@ -0,0 +1,634 @@ +# Direct Remote Rank-Major Output Design + +## 1. Status + +This document describes a future implementation. The current MSCCL++ branch +does not implement this mode. + +The proposed mode lets the expert rank's GEMM2 epilogue write each weighted +route row directly into the originating rank's symmetric route buffer. The +originating rank then performs only a local FP32 top-k reduction. + +The design is intentionally different from a separate producer-side push +kernel. The remote write must be fused into GEMM2; otherwise it adds another +read of the local GEMM2 output and another kernel launch. + +## 2. Goal + +Replace: + +```text +GEMM2 local weighted-route output + -> persistent remote-pull combine kernel + -> remote TMA loads + -> FP32 top-k reduction + -> BF16 output +``` + +with: + +```text +GEMM2 weighted epilogue + -> direct peer-mapped route stores + -> per-phase system-release completion + -> source-local FP32 top-k reduction + -> BF16 output +``` + +The expected benefit is: + +1. No persistent pull workers competing with GEMM2. +2. No GEMM2 SM reservation. +3. No physical GEMM release gate. +4. Network stores overlap naturally with GEMM2 epilogue execution. +5. The remaining post-GEMM work reads only local HBM. + +## 3. Non-goals + +The first implementation will not: + +1. Atomically accumulate directly into the final BF16 token output. +2. Compact the source route buffer below `[tokens, top_k, hidden]`. +3. Pipeline GEMM1 or activation by expert. +4. Change rank-major dispatch token compaction. +5. Require MSCCL++ types inside MAI/CUTLASS. + +Direct accumulation into one output row would require conflicting remote +writes or FP32 atomics. A route buffer keeps every writer conflict-free and +preserves the existing numerics. + +## 4. Production shape + +```text +tokens/rank = 64 +top_k = 8 +hidden = 4096 +intermediate = 6656 +local experts/rank = 16 +global experts = 16 * world_size +route output dtype = BF16 +reduction accumulator = FP32 +``` + +Per source rank and epoch: + +```text +route rows = 64 * 8 = 512 +route bytes = 512 * 4096 * 2 = 4 MiB +final output bytes = 64 * 4096 * 2 = 0.5 MiB +``` + +The route-buffer size is independent of world size. At 32 GPUs, every expert +rank still produces about 512 route rows and every source rank still receives +512 route rows under balanced routing. + +## 5. Current rank-major mapping + +Rank-major dispatch already compacts one token row per source/destination-rank +pair. + +For source token `token_idx` and top-k lane `lane`: + +```text +global_expert = topk_ids[token_idx, lane] +destination_rank = global_expert / local_experts +local_expert = global_expert % local_experts +destination_slot = compact slot allocated for destination_rank +``` + +The destination dispatch row is: + +```text +dispatch_row = + source_rank * max_tokens_per_rank + destination_slot +``` + +All top-k lanes targeting the same destination rank share that compact token +row. Current code stores `destination_slot` in +`WorkspaceView::rankMajorSendIndices_`. + +Relevant current sources: + +```text +src/ext/ep/low_latency/dispatch.cu + RankMajorRoute + prepareRankMajorRoute + sendRankMajorMetadata + dispatchSendRankMajorBf16 + +src/ext/ep/low_latency/config.cuh + WorkspaceView::rankMajorSendIndices_ + +src/ext/ep/include/config.hpp + rankMajorTopkIdsBuffer_ + rankMajorTopkWeightsBuffer_ + rankMajorTokenBuffer_ +``` + +## 6. Additional dispatch metadata + +The destination rank currently knows the source rank and compact destination +slot from the dispatch row, but direct write also needs the original source +token index. + +Add: + +```text +rankMajorSourceTokenIdx[ + source_rank, + destination_slot +] -> source_token_idx +``` + +Shape: + +```text +[num_ranks, max_tokens_per_rank] int32 +``` + +`dispatchSendRankMajorBf16` writes this metadata with the token and top-k +metadata. It is stable for the lifetime of the dispatched row. + +Do not replace compact `destination_slot` with `source_token_idx`. Fixed +source-token slots would create holes in the destination GEMM input and force +GEMM to process unused rows. Keep compact dispatch and carry one extra int32. + +## 7. Direct route destination + +Each source rank owns a symmetric receive buffer: + +```text +directRankRouteRecv[ + buffer_slot, + source_token_idx, + topk_lane, + hidden +] BF16 +``` + +Recommended initial ring depth: + +```text +buffer_slots = 2 +buffer_slot = dispatch_epoch & 1 +``` + +For one producer route: + +```text +remote_route = + source_token_idx * top_k + topk_lane + +remote_dst = + peer_route_base[source_rank] + + buffer_slot * route_buffer_stride + + remote_route * hidden * sizeof(BF16) +``` + +Every `(source_token_idx, topk_lane)` has exactly one expert owner, so every +route row has exactly one writer. No atomic data stores are required. + +For the production shape: + +```text +one route buffer slot = 4 MiB/rank +two slots = 8 MiB/rank +``` + +## 8. Address preparation + +MAI should not include MSCCL++ headers or access private MSCCL++ objects. + +MSCCL++ exposes a plain device context containing stable addresses: + +```text +peer_route_bases[num_ranks] uint64 +peer_ready_bases[num_ranks] uint64 +source_token_idx metadata int32* +dispatch epoch uint32* +phase source masks uint64* +``` + +The Python API passes these tensors or pointers to the W8A16 grouped GEMM. +The CUDA Graph captures stable device addresses; no per-replay host address +construction is allowed. + +Recommended public object: + +```text +DirectRankWriteContext + peer_route_bases + peer_ready_bases + source_token_indices + dispatch_epoch + phase_source_masks + num_ranks + max_tokens_per_rank + top_k + hidden +``` + +The context is data-only. MAI receives plain tensors and scalar dimensions. + +## 9. GEMM2 epilogue change + +The existing W8A16 epilogue already: + +1. Maps grouped-GEMM output rows through scatter row IDs. +2. Multiplies the FP32 accumulator by the route weight. +3. Stores BF16 vectors. +4. Counts completed output tiles per expert. + +Current source: + +```text +/home/azhpcuser/mai/yolo/mai_kernels/csrc/w8a16_grouped_gemm/ + cutlass_ext/cutlass/epilogue/collective/ + sm100_epilogue_array_nosmem_rank4.hpp +``` + +Add an optional direct-rank scatter mode: + +```text +local scatter: + output_base + token_id * ld + hidden_offset + +direct-rank scatter: + peer_route_bases[source_rank] + + buffer_slot_offset + + (source_token_idx * top_k + topk_lane) * hidden + + hidden_offset +``` + +The epilogue continues applying the routing weight before the BF16 store. + +Required per routed row: + +```text +source_rank +source_token_idx +topk_lane +``` + +`source_rank` comes from `dispatch_row / max_tokens_per_rank`. +`source_token_idx` comes from the new dispatch metadata. +`topk_lane` comes from the flattened route row ID. + +The preferred implementation builds a device array of 64-bit route-row base +addresses during the existing gather/populate prologue: + +```text +direct_route_row_ptr[flattened_route] +``` + +The epilogue then performs: + +```text +dst = direct_route_row_ptr[token_id] + hidden_offset +``` + +This keeps peer-address arithmetic out of the vector store loop. + +## 10. Completion granularity + +Use the existing ordered four-expert phase size for the first implementation: + +```text +phase 0: experts 0..3 +phase 1: experts 4..7 +phase 2: experts 8..11 +phase 3: experts 12..15 +``` + +Completion state on each source rank: + +```text +directRankReadyEpoch[ + buffer_slot, + producer_rank, + phase +] uint32 +``` + +The destination rank also builds: + +```text +phaseSourceMask[phase] +``` + +Bit `r` is set when that producer phase writes at least one route to source +rank `r`. Empty source/phase pairs receive no notification. + +## 11. Completion publication without a persistent controller + +Do not launch a persistent MSCCL++ data/control kernel. + +Each expert's final tile: + +1. Waits for its epilogue stores to be issued. +2. Participates in the per-expert tile-completion atomic chain. +3. Publishes `expertReadyEpoch[expert]`. + +The final expert CTA that observes all experts in its phase ready: + +1. Claims the phase with an epoch-tagged CAS. +2. Executes the required system-scope fence. +3. Iterates `phaseSourceMask[phase]`. +4. System-release stores the epoch into each source rank's mapped + `directRankReadyEpoch[slot, producer_rank, phase]`. + +This reuses real GEMM2 CTAs and reserves no SM for a side kernel. + +## 12. Required memory ordering + +The required visibility chain is: + +```text +peer route stores from every expert tile + -> epilogue CTA barrier + -> acq_rel expert tile counter chain + -> device-release expert-ready epoch + -> phase publisher device-acquire + -> system fence/release + -> peer ready epoch system-release store + -> source system-acquire + -> source reads local route buffer +``` + +The implementation must use CUDA system-scope ordering for the final remote +completion publication. A device-scope epoch is insufficient for peer-mapped +route data. + +Do not publish readiness from an arbitrary tile. Expected tile counts must be +derived from the actual expert problem shape and routing imbalance. + +## 13. Source-local reduction + +After local GEMM2, the source rank launches one local reduction kernel on its +main stream: + +```text +directRankReduceKernel( + local_route_buffer, + topk_ids, + remote_ready_epochs, + dispatch_epoch, + output +) +``` + +Recommended topology: + +```text +one CTA per source token +256 threads/CTA +FP32 accumulation +one BF16 output store +``` + +For each token: + +1. Lanes 0..top-k-1 derive `(producer_rank, producer_phase)` from the original + top-k expert IDs. +2. The control warp waits until every required epoch equals the current + dispatch epoch. +3. The CTA reads the eight local route rows. +4. Threads sum in FP32 in top-k lane order. +5. The CTA writes one BF16 output row. + +This kernel starts after local GEMM2, so it does not reserve resources from +GEMM2. It may wait for slower remote producers, but that wait occurs after the +local compute has released all SMs. + +There is no separate combine stream or CUDA event join in the initial design: + +```text +dispatch -> GEMM1 -> activation -> GEMM2/direct writes + -> local wait/reduce -> next dispatch +``` + +## 14. Buffer reuse and CUDA Graph safety + +The source rank owns reuse of its receive slots. + +Invariant: + +```text +source rank does not launch dispatch epoch E +until local reduction for epoch E-1 has completed +``` + +A producer can write source epoch `E` only after receiving that source's +dispatch metadata for epoch `E`. Therefore the source has already completed +the previous local reduction before any producer can reuse the selected slot. + +Use two slots initially to make alternating graph iterations explicit. + +Every metadata, route, and epoch buffer must be allocated before CUDA Graph +capture. No allocation, registration, or peer-address discovery may occur +during capture or replay. + +## 15. MSCCL++ API changes + +Add a new mode: + +```text +CombineMode::RANK_MAJOR_DIRECT_WRITE +``` + +MSCCL++ changes: + +```text +src/ext/ep/include/config.hpp + add source-token metadata + add two-slot direct route receive buffer + add direct ready epochs + add stable peer pointer tables + +src/ext/ep/low_latency/config.cuh + add workspace phase masks and phase-published epochs + +src/ext/ep/low_latency/dispatch.cu + publish source_token_idx for each compact destination row + build per-phase source masks + +include/mscclpp/ext/ep/types.hpp + define the direct-rank context/accessors + +src/ext/ep/runtime/fixed_buffer.cc + allocate and expose direct-rank buffers/context + +src/ext/ep/bindings.cpp + expose context tensors/pointers + +python/mscclpp/ep/{types.py,low_latency.py,communicator.py} + expose the new mode and context +``` + +In direct-write mode, `communicator.combine()` launches only the source-local +wait/reduce kernel. It does not send or pull route data. + +## 16. MAI API changes + +Extend `W8A16_GroupedGEMM` with one optional public data bundle containing: + +```text +peer_route_bases +peer_ready_bases +source_token_indices +phase_source_masks +dispatch_epoch +num_ranks +max_tokens_per_rank +expert_phase_size +``` + +Supply the bundle only for GEMM2. + +MAI changes: + +```text +mai_kernels/src/mai_kernels/w8a16_grouped_gemm.py + validate the optional direct-rank tensors + +mai_kernels/csrc/w8a16_grouped_gemm/ + plumb the context into CUTLASS arguments + build direct route-row pointers in gather/populate + +cutlass_ext/.../sm100_epilogue_array_nosmem_rank4.hpp + add direct-rank vector stores + publish phase completion +``` + +The MAI interface must remain usable without MSCCL++ installed. + +## 17. Correctness invariants + +1. Every valid `(source token, top-k lane)` is written exactly once. +2. No two experts write the same route row. +3. The route weight is applied once in FP32 before BF16 storage. +4. Source reduction converts BF16 routes to FP32. +5. Source reduction sums lanes in a deterministic order. +6. Final output is cast to BF16 once. +7. A source never reads a route before its producer phase epoch. +8. A producer never overwrites a source buffer slot before the source reuses + that epoch slot. +9. Empty experts and invalid/padding rows publish completion without writing + route data. +10. Epoch comparison, not zero initialization, determines readiness. + +Reference numerics: + +```text +GEMM2 FP32 accumulator + -> multiply route weight in FP32 + -> BF16 route store + -> source BF16-to-FP32 conversion + -> FP32 top-k sum + -> BF16 final store +``` + +## 18. Performance model + +For 64 tokens/rank: + +```text +network payload/rank = 4 MiB +wire floor at 660 GB/s = about 6.4 us +local reduction read = 4 MiB +local output write = 0.5 MiB +``` + +The previous push experiment was slower because it materialized producer +output and used additional producer/consumer work. This design is worth +retesting only because it: + +1. Fuses the network write into GEMM2. +2. Removes the remote-pull kernel. +3. Removes the 17-SM reservation. +4. Removes the GEMM release wait. + +Expected costs: + +```text +GEMM2 remote-store slowdown target <= 5 us +post-GEMM local wait/reduction target <= 5 us after last producer +``` + +The design should be abandoned if peer stores slow GEMM2 by more than the +serial direct-pull combine it replaces. + +## 19. Implementation sequence + +### Phase 1: metadata and local emulation + +1. Add `source_token_idx` metadata. +2. Build route-row addresses targeting a local test buffer. +3. Make GEMM2 scatter into `[token, top_k, hidden]`. +4. Run the source-local reduction. +5. Verify serial parity without peer writes. + +### Phase 2: synchronous peer writes + +1. Replace local route bases with peer-mapped route bases. +2. Synchronize all ranks after GEMM2. +3. Run local reduction. +4. Verify 4- and 8-GPU correctness. +5. Measure isolated GEMM2 slowdown from peer stores. + +### Phase 3: epoch ordering + +1. Add expert and phase completion. +2. Remove the global post-GEMM synchronization. +3. Make local reduction wait only for its required producer phases. +4. Stress epoch reuse over long CUDA Graph replays. + +### Phase 4: performance tuning + +1. Compare cached, no-allocate, and write-through peer store policies. +2. Tune route-row pointer layout. +3. Tune local reduction CTA size. +4. Compare two- and four-expert completion phases. +5. Test 32, 64, and 128 tokens/rank. + +## 20. Acceptance criteria + +Correctness: + +```text +graph/eager max abs = 0 +serial/direct-write difference within established BF16 tolerance +no stale rows over at least 100,000 graph iterations +``` + +Performance at 8 GPUs, T=64, H=4096, I=6656, 16 local experts: + +```text +direct-write E2E < serial direct-pull E2E +no persistent SM reservation +no standalone remote data-movement kernel +local reduction tail <= 5 us after last required producer +``` + +Scale: + +```text +4 GPUs: correctness and debugging +8 GPUs: primary performance gate +16 GPUs: communication scaling +32 GPUs: production validation +``` + +## 21. Primary risks + +1. Peer epilogue stores may reduce GEMM2 HBM efficiency. +2. Small vector stores may not combine efficiently across NVLink. +3. System-scope publication may be more expensive than expected. +4. A slow producer still determines the source's final reduction start. +5. Pointer-array setup may add gather/populate overhead. +6. Incorrect transitive ordering can expose partially written remote rows. +7. Buffer-slot reuse can corrupt long CUDA Graph runs if epoch ownership is + not enforced. + +These risks must be measured independently before optimizing the complete +pipeline. diff --git a/include/mscclpp/ext/ep/moe_runtime.hpp b/include/mscclpp/ext/ep/moe_runtime.hpp new file mode 100644 index 000000000..695d8f5b0 --- /dev/null +++ b/include/mscclpp/ext/ep/moe_runtime.hpp @@ -0,0 +1,136 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +#ifndef MSCCLPP_EXT_EP_MOE_RUNTIME_HPP_ +#define MSCCLPP_EXT_EP_MOE_RUNTIME_HPP_ + +#include + +#include +#include +#include +#include + +namespace mscclpp { +namespace ep { +struct LatencyContext; +struct ThroughputContext; +/// Unified host runtime for expert-parallel dispatch and combine. +/// +/// One runtime owns the communication buffers and synchronization state for the +/// selected mode. LATENCY uses fixed-capacity expert-major or rank-major +/// layouts. THROUGHPUT uses a dynamically sized token-major receive pool. +/// Operations are asynchronous with respect to the host and execute on the +/// CUDA stream supplied by each request. +class MoERuntime { + public: + /// Construct a runtime for the selected mode and topology. + /// + /// Only resources required by @p mode are allocated. + /// @param communicator Initialized MSCCL++ communicator. + /// @param mode Runtime algorithm family. + /// @param maxTokensPerRank Fixed latency-mode token capacity. + /// @param hidden Hidden dimension for latency-mode buffers. + /// @param numExperts Global expert count. + /// @param numTopk Number of routed experts per token. + /// @param maxHiddenBytes Maximum throughput-mode bytes per token row. + /// @param numBlocks Communication block budget. + /// @param outputLayout Latency-mode dispatch output layout. + MoERuntime(mscclpp::Communicator& communicator, MoEMode mode, int maxTokensPerRank, int hidden, int numExperts, + int numTopk, int64_t maxHiddenBytes, int numBlocks, + DispatchLayout outputLayout = DispatchLayout::EXPERT_MAJOR); + ~MoERuntime() noexcept(false); + + MoERuntime(const MoERuntime&) = delete; + MoERuntime& operator=(const MoERuntime&) = delete; + + /// Return the configured runtime mode. + MoEMode mode() const { return mode_; } + /// Return whether the selected mode supports the detected topology. + bool isAvailable() const { return available_; } + /// Return whether the runtime is available across more than one node. + bool isInternodeAvailable() const { return available_ && numRanks_ > numNvlRanks_; } + + /// Return the local rank. + int rank() const { return rank_; } + /// Return the global rank count. + int numRanks() const { return numRanks_; } + /// Return the NVLink-local rank count. + int numNvlRanks() const { return numNvlRanks_; } + /// Return the rank count in one CUDA IPC domain. + int numRanksPerIpcDomain() const { return numRanksPerIpcDomain_; } + + /// Return the runtime-owned rank-major top-k ID buffer. + void* outputTopkIdsBuffer() const; + /// Return the runtime-owned rank-major top-k weight buffer. + void* outputTopkWeightsBuffer() const; + /// Return the runtime-owned dispatch output buffer. + void* dispatchOutputBuffer() const; + /// Return the runtime-owned rank-major combine input buffer. + void* combineInputBuffer() const; + + /// Dispatch tokens using the configured runtime mode. + /// + /// @p request must contain the request type matching mode(): a + /// LatencyDispatchRequest for LATENCY or a ThroughputDispatchRequest for + /// THROUGHPUT. Output buffers remain owned by the caller unless obtained + /// through a runtime buffer accessor. + /// @param request Dispatch inputs, outputs, dimensions, and CUDA stream. + /// @throws std::invalid_argument If the request type does not match mode(). + void dispatch(const DispatchRequest& request); + + /// Combine expert outputs using the configured runtime mode. + /// + /// A combine request must follow its matching dispatch so the runtime can + /// reuse routing metadata and synchronization epochs. @p request must contain + /// a LatencyCombineRequest for LATENCY or a ThroughputCombineRequest for + /// THROUGHPUT. + /// @param request Combine inputs, outputs, dimensions, and CUDA stream. + /// @throws std::invalid_argument If the request type does not match mode(). + void combine(const CombineRequest& request); + + /// Build throughput-mode token routing metadata. + /// + /// Computes per-rank counts, per-expert counts, and token-to-rank membership + /// on @p stream without moving token payloads. + void tokenMajorPrepare(int* numTokensPerRank, int* numTokensPerExpert, bool* isTokenInRank, const int64_t* topkIdx, + int numTokens, int numTopk, int numExperts, cudaStream_t stream); + /// Return the throughput-mode communication channel count. + int tokenMajorNumChannels(int xElementSize) const; + /// Resolve the runtime-owned throughput receive buffer. + void* tokenMajorResolveRecvBuffer(int numTokens, int numRecvTokens, int hidden, int xElementSize) const; + /// Exchange throughput routing counts and return the receive-token count. + /// + /// This host-synchronizing metadata phase must precede throughput dispatch + /// when cached routing metadata is unavailable. + int tokenMajorNotify(int* rankPrefixMatrix, int* channelPrefixMatrix, int* numRecvTokensPerExpert, + const int* numTokensPerRank, const int* numTokensPerExpert, const bool* isTokenInRank, + int numTokens, int numExperts, int xElementSize, int expertAlignment, cudaStream_t stream); + + private: + void requireMode(MoEMode expected) const; + void launchLatencyDispatch(const LatencyDispatchRequest& request); + void launchThroughputDispatch(const ThroughputDispatchRequest& request); + void launchLatencyCombine(const LatencyCombineRequest& request); + void launchThroughputCombine(const ThroughputCombineRequest& request); + + std::shared_ptr bootstrap_; + MoEMode mode_; + int rank_; + int numRanks_; + int numNvlRanks_; + int numRanksPerIpcDomain_; + bool available_ = false; + + std::unique_ptr latencyContext_; + std::unique_ptr throughputContext_; +}; + +/// Create the unified MoE runtime selected by @p mode. +std::shared_ptr createMoERuntime(mscclpp::Communicator& communicator, MoEMode mode, int maxTokensPerRank, + int hidden, int numExperts, int numTopk, int64_t maxHiddenBytes, + int numBlocks, DispatchLayout outputLayout = DispatchLayout::EXPERT_MAJOR); + +} // namespace ep +} // namespace mscclpp + +#endif // MSCCLPP_EXT_EP_MOE_RUNTIME_HPP_ diff --git a/include/mscclpp/ext/ep/types.hpp b/include/mscclpp/ext/ep/types.hpp new file mode 100644 index 000000000..89160817a --- /dev/null +++ b/include/mscclpp/ext/ep/types.hpp @@ -0,0 +1,230 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +#ifndef MSCCLPP_EXT_EP_TYPES_HPP_ +#define MSCCLPP_EXT_EP_TYPES_HPP_ + +#include + +#include +#include +#include + +namespace mscclpp { +namespace ep { + +class MoERuntime; + +/// Expert-parallel runtime mode. +enum class MoEMode { + /// Algorithms optimized for minimum standalone latency. + LATENCY, + /// Resource-bounded algorithms optimized for end-to-end throughput. + THROUGHPUT +}; + +/// Logical dispatch output layout. +enum class DispatchLayout { + /// Rows grouped by local expert. + EXPERT_MAJOR, + /// Dynamically sized token-major rows used by throughput mode. + TOKEN_MAJOR, + /// Fixed-stride rows grouped by source rank. + RANK_MAJOR +}; + +/// Combine algorithm. +enum class CombineMode { + /// Reduce local expert rows before sending one partial per rank and token. + RANK_LOCAL_REDUCE, + /// Send every expert row and reduce all contributions on the source rank. + DIRECT_SEND +}; + +/// Dispatch payload data format. +enum class DispatchDataType { + /// Unquantized BF16 payload. + BF16, + /// FP8 E4M3 payload with one floating-point scale per 128 hidden elements. + FP8_E4M3 +}; + +/// Arguments for latency-mode dispatch. +struct LatencyDispatchRequest { + /// Dispatch output buffer. + void* output; + /// Optional dispatch scale output. + void* outputScales; + /// Optional source-token metadata output. + int* outputSrcInfo; + /// Optional dispatched top-k expert IDs. + int* outputTopkIdx; + /// Optional dispatched top-k weights. + float* outputTopkWeights; + /// Optional packed layout metadata. + int64_t* outputLayoutRange; + /// Per-expert or per-rank output counts. + int* outputCount; + /// Input token payload. + const void* input; + /// Input top-k expert IDs. + const int64_t* topkIdx; + /// Optional input top-k weights. + const float* topkWeights; + /// Number of input tokens. + int numTokens; + /// Hidden dimension. + int hidden; + /// Number of routed experts per token. + int numTopk; + /// Active per-rank token capacity. + int maxTokensPerRank; + /// Global expert count. + int numExperts; + /// Expert ID used for invalid rank-major entries. + int invalidTokenExpertId; + /// Requested dispatch output layout. + DispatchLayout dispatchLayout; + /// Requested dispatch payload format. + DispatchDataType dispatchDataType; + /// Dispatch grid block count. + int numBlocks; + /// CUDA stream used for the operation. + cudaStream_t stream; +}; + +/// Arguments for throughput-mode dispatch. +struct ThroughputDispatchRequest { + /// Token receive buffer. + void* recvX; + /// Optional received scale output. + float* recvXScales; + /// Optional received top-k expert IDs. + int64_t* recvTopkIdx; + /// Optional received top-k weights. + float* recvTopkWeights; + /// Per-token routing state consumed by combine. + int* sendHead; + /// Input token payload. + const void* input; + /// Optional input scales. + const float* inputScales; + /// Optional input top-k expert IDs. + const int64_t* topkIdx; + /// Optional input top-k weights. + const float* topkWeights; + /// Token-to-destination-rank membership. + const bool* isTokenInRank; + /// Per-source-rank token prefixes. + const int* rankPrefixMatrix; + /// Per-channel token prefixes. + const int* channelPrefixMatrix; + /// Number of input tokens. + int numTokens; + /// Hidden dimension. + int hidden; + /// Number of routed experts per token. + int numTopk; + /// Number of scales per token. + int numScales; + /// Global expert count, or zero when cached metadata is reused. + int numExperts; + /// Input element size in bytes. + int inputElementSize; + /// Number of received tokens. + int numRecvTokens; + /// Whether cached routing metadata is reused. + bool cachedMode; + /// CUDA stream used for the operation. + cudaStream_t stream; +}; + +/// Mode-specific dispatch request. +struct DispatchRequest { + /// Construct a latency dispatch request. + explicit DispatchRequest(LatencyDispatchRequest request) : value_(std::move(request)) {} + /// Construct a throughput dispatch request. + explicit DispatchRequest(ThroughputDispatchRequest request) : value_(std::move(request)) {} + + private: + friend class MoERuntime; + std::variant value_; +}; + +/// Arguments for latency-mode combine. +struct LatencyCombineRequest { + /// Combined token output. + void* output; + /// Local expert output. + const void* input; + /// Input top-k expert IDs. + const int64_t* topkIdx; + /// Optional input top-k weights. + const float* topkWeights; + /// Optional source-token metadata. + const int* srcInfo; + /// Optional packed layout metadata. + const int64_t* layoutRange; + /// Number of output tokens. + int numTokens; + /// Hidden dimension. + int hidden; + /// Number of routed experts per token. + int numTopk; + /// Active per-rank token capacity. + int maxTokensPerRank; + /// Global expert count. + int numExperts; + /// Dispatch input layout. + DispatchLayout dispatchLayout; + /// Dispatch payload format. + DispatchDataType dispatchDataType; + /// Combine algorithm. + CombineMode combineMode; + /// Combine worker block count. + int numBlocks; + /// CUDA stream used for the operation. + cudaStream_t stream; +}; + +/// Arguments for throughput-mode combine. +struct ThroughputCombineRequest { + /// Combined token output. + void* output; + /// Optional combined top-k weights. + float* outputTopkWeights; + /// Local expert output. + const void* input; + /// Optional local top-k weights. + const float* topkWeights; + /// Routing state returned by throughput dispatch. + const int* sendHead; + /// Number of local expert-output rows. + int numInputTokens; + /// Number of combined output tokens. + int numOutputTokens; + /// Hidden dimension. + int hidden; + /// Number of routed experts per token. + int numTopk; + /// Input element size in bytes. + int inputElementSize; + /// CUDA stream used for the operation. + cudaStream_t stream; +}; + +/// Mode-specific combine request. +struct CombineRequest { + /// Construct a latency combine request. + explicit CombineRequest(LatencyCombineRequest request) : value_(std::move(request)) {} + /// Construct a throughput combine request. + explicit CombineRequest(ThroughputCombineRequest request) : value_(std::move(request)) {} + + private: + friend class MoERuntime; + std::variant value_; +}; + +} // namespace ep +} // namespace mscclpp + +#endif // MSCCLPP_EXT_EP_TYPES_HPP_ diff --git a/python/mscclpp/ep/README.md b/python/mscclpp/ep/README.md index 04bf21fc3..d705146e5 100644 --- a/python/mscclpp/ep/README.md +++ b/python/mscclpp/ep/README.md @@ -38,6 +38,11 @@ moe_comm = MoECommunicator(...) The class owns MoE dispatch/combine communication, but it does not own the MLP compute backend. +Internally, `MoECommunicator` constructs a passive `LatencyContext` or +`ThroughputContext`, then creates the matching `LatencyRuntime` or +`ThroughputRuntime`. Both implement the same `Runtime` interface and own one +unified C++ `MoERuntime`; there is no additional communication backend layer. + ## MoECommunicator configuration `MoECommunicator` owns communication setup, scratch buffers, expert placement, @@ -62,15 +67,15 @@ class MoECommunicatorConfig: max_tokens_per_rank: int = 0 # Runtime mode and output layout - mode: MoEMode = MoEMode.LOW_LATENCY + mode: MoEMode = MoEMode.LATENCY output_layout: Optional[DispatchLayout] = None # default is derived from mode invalid_token_expert_id: Optional[int] = None # defaults to num_experts # Quantization defaults quant: Optional[QuantConfig] = None - # Launch resources - num_sms: int = 20 + # Launch resources; defaults to 130 for LATENCY and 20 for THROUGHPUT + num_blocks: Optional[int] = None # Overlap enable_overlap: bool = False @@ -86,7 +91,7 @@ moe_comm = MoECommunicator( hidden_size=hidden_size, topk=topk, max_tokens_per_rank=max_tokens, - mode=MoEMode.HIGH_THROUGHPUT, + mode=MoEMode.THROUGHPUT, ) ``` @@ -135,13 +140,13 @@ a later version can add an explicit `expert_map` for arbitrary placement. | Field | Purpose | |---|---| -| `mode` | Backend selection (`MoEMode.LOW_LATENCY` or `MoEMode.HIGH_THROUGHPUT`) | +| `mode` | Algorithm family (`MoEMode.LATENCY` or `MoEMode.THROUGHPUT`) | | `output_layout` | MLP input layout returned by dispatch | | `invalid_token_expert_id` | Sentinel for rank-major non-local and padding entries; defaults to `num_experts` | | `max_tokens_per_rank` | dispatch capacity | | scratch buffers | internally sized from mode, capacity, topology, and shape | -| `num_sms` | backend launch/resource tuning | -| `dispatch_config`, `combine_config` | backend-specific tuning configs | +| `num_blocks` | communication block count; mode-specific default when unset | +| `dispatch_config`, `combine_config` | context-specific tuning configs | | `overlap_capability` | whether selected MLP/backend supports notify | The user should not pass these to `dispatch` unless explicitly overriding a @@ -149,44 +154,44 @@ specialized advanced path. ### Mode selection -The active implementation supports `mode=MoEMode.LOW_LATENCY` and -`mode=MoEMode.HIGH_THROUGHPUT`. `mode` must be a `MoEMode` enum value, not a -string. LL supports expert-major and rank-major output layouts. HT uses a flat output layout and -supports 2, 4, 8, or 16 ranks within one detected GPU IPC/NVL fabric domain; -that domain may span multiple hosts. +The active implementation supports `mode=MoEMode.LATENCY` and +`mode=MoEMode.THROUGHPUT`. `mode` must be a `MoEMode` enum value, not a string. +Latency algorithms support expert-major and rank-major output layouts. Throughput +algorithms use a flat token-major layout and support 2, 4, 8, or 16 ranks +within one detected GPU IPC/NVL fabric domain; that domain may span hosts. ```python -moe_comm = MoECommunicator(..., mode=MoEMode.LOW_LATENCY) +moe_comm = MoECommunicator(..., mode=MoEMode.LATENCY) ``` This keeps `MoECommunicator` policy-free. Serving frameworks such as SGLang can choose a mode based on their own scheduling policy, batch shape, runner backend, -and benchmarking data once multiple active backends are available. +and benchmarking data once multiple algorithm contexts are available. -The mode also fixes the SM budget, which is the main scheduling consideration: +The mode selects the default communication block budget: -| Mode | SMs used | Intended use | +| Mode | Default blocks | Intended use | |---|---|---| -| `MoEMode.LOW_LATENCY` | ~128 (`low_latency_num_blocks - 2`) | comms own the GPU | -| `MoEMode.HIGH_THROUGHPUT` | 20 (`Config.num_sms`) | overlap comms with expert GEMMs | +| `MoEMode.LATENCY` | 130 | minimize standalone latency | +| `MoEMode.THROUGHPUT` | 20 | overlap communication with expert GEMMs | The selected mode determines the default dispatch output layout: | Mode | Default layout | |---|---| -| `ht` | `DispatchLayout.TOKEN_MAJOR` | -| `ll` | `DispatchLayout.EXPERT_MAJOR` | +| `THROUGHPUT` | `DispatchLayout.TOKEN_MAJOR` | +| `LATENCY` | `DispatchLayout.EXPERT_MAJOR` | -`output_layout` may still be kept as an advanced override if a backend supports +`output_layout` may still be kept as an advanced override if a context supports multiple layouts within the same mode. Use `DispatchLayout` instead of string literals for this field: | Layout enum | Tensor shape | |---|---| -| `DispatchLayout.TOKEN_MAJOR` | HT: `[total_recv_tokens, hidden]` | +| `DispatchLayout.TOKEN_MAJOR` | Throughput: `[total_recv_tokens, hidden]` | | `DispatchLayout.EXPERT_MAJOR` | `[num_local_experts, max_slots_per_expert, hidden]` | -| `DispatchLayout.RANK_MAJOR` | LL: `[world_size * max_tokens_per_rank, hidden]` | +| `DispatchLayout.RANK_MAJOR` | Latency: `[world_size * max_tokens_per_rank, hidden]` | ## MoECommunicator methods @@ -214,10 +219,10 @@ class MoECommunicator: ) -> torch.Tensor: ... - def dispatch_async(..., overlap_config: Optional[CommOverlapConfig] = None) -> DispatchRequest: + def dispatch_async(..., overlap_config: Optional[OverlapConfig] = None) -> DispatchRequest: ... - def combine_async(..., overlap_config: Optional[CommOverlapConfig] = None) -> CombineRequest: + def combine_async(..., overlap_config: Optional[OverlapConfig] = None) -> CombineRequest: ... def create_overlap_config( @@ -226,7 +231,7 @@ class MoECommunicator: *, handle: Optional[DispatchHandle] = None, level: str = "op", # "op" or "block" - ) -> CommOverlapConfig: + ) -> OverlapConfig: ... ``` @@ -361,7 +366,7 @@ class BlockOverlapConfig: @dataclass -class CommOverlapConfig: +class OverlapConfig: operation: Optional[OperationOverlapConfig] = None block: Optional[BlockOverlapConfig] = None @@ -398,7 +403,7 @@ combine_overlap_config = moe_comm.create_overlap_config( `op="dispatch", level="block"` is not part of the first version. Dispatch overlap is operation-level only. -`CommOverlapConfig` contains exactly one overlap mode: +`OverlapConfig` contains exactly one overlap mode: | Field | Purpose | |---|---| @@ -428,7 +433,7 @@ Each concrete `DispatchHandle` stores a layout-specific `combine_context` used to reverse dispatch and finish combine. `ExpertMajorDispatchHandle` uses `ExpertMajorCombineContext` (`topk_ids`, `weights`, source info, and layout ranges). `RankMajorDispatchHandle` records the original routing needed for direct remote combine. -High-throughput handles use a direct combine context with receive-side weights +Token-major overlap handles use a direct combine context with receive-side weights and send-head tensors. The MLP should treat the handle as opaque and pass it back to `combine`. @@ -456,7 +461,7 @@ The user should not expand `input` by top-k and should not convert it to expert-major before calling `dispatch`. `dispatch` includes any metadata exchange needed before moving token payloads. -For normal/high-throughput modes this typically means computing send counts from +For throughput mode this typically means computing send counts from `topk_ids`, exchanging counts or layout information across ranks, choosing recv slots, and then dispatching the activation payload. Users should not call a separate metadata-exchange API in the simple path. @@ -498,19 +503,19 @@ DeepEP/SGLang, scales are usually per token and per hidden block. ### `output_buffer` -Low-latency dispatch requires the caller to provide the receive token buffer: +Expert-major latency dispatch requires the caller to provide the receive token buffer: ```python output_buffer: torch.Tensor ``` -For padded expert-major LL layout: +For padded expert-major latency layout: ```text output_buffer: [num_local_experts, world_size * max_tokens_per_rank, hidden] ``` -For rank-major LL layout, do not pass `output_buffer`. The runtime returns a +For rank-major latency layout, do not pass `output_buffer`. The runtime returns a registered fixed-stride buffer: ```text @@ -526,7 +531,7 @@ The dtype must match the dispatch output dtype. For BF16 dispatch it is BF16. For FP8 dispatch it is FP8 and the returned `DispatchOutput.quant` carries the matching format and scale tensor. -`output_buffer` is required for LL because the MLP runner often owns or reuses +`output_buffer` is required for expert-major latency dispatch because the MLP runner often owns or reuses workspace memory. `MoECommunicator` writes dispatch output into the provided buffer instead of allocating it internally. @@ -535,9 +540,9 @@ buffer instead of allocating it internally. `dispatch` should return MLP-ready tokens. The MLP should not run another token-major to expert-major permutation unless it uses a custom adapter. -### Normal / high-throughput token-major layout +### Throughput token-major layout -HT uses `DispatchLayout.TOKEN_MAJOR`: +Throughput algorithms use `DispatchLayout.TOKEN_MAJOR`: ```python dispatch_out.tokens # [total_recv_tokens, H] @@ -548,9 +553,9 @@ Each row represents one `(source token, destination rank)` and is accompanied by token routed to multiple experts on the same destination rank is transferred only once. -### Low-latency output layouts +### Latency output layouts -LL defaults to `DispatchLayout.EXPERT_MAJOR`, a padded expert-major tensor: +Latency mode defaults to `DispatchLayout.EXPERT_MAJOR`, a padded expert-major tensor: ```python dispatch_out.tokens # [num_local_experts, max_slots_per_expert, H] @@ -574,7 +579,9 @@ The MoE runner must write its rank-major output into the runtime-owned registered buffer: ```python -expert_output = communicator.get_expert_output_buffer() +dispatch_out, handle = communicator.dispatch(x, topk_ids, topk_weights) +expert_output = dispatch_out.combine_input_buffer +assert expert_output is not None moe(..., output=expert_output) combined = communicator.combine(expert_output, handle) ``` @@ -603,7 +610,7 @@ dimension replaced by the scale dimension. Examples: ```text -token-major tokens: HT [total_recv_tokens, H]; LL [world_size * max_tokens_per_rank, H] +token-major tokens: throughput [total_recv_tokens, H]; latency rank-major [world_size * max_tokens_per_rank, H] rank-major scales: not yet supported expert-major tokens: [num_local_experts, max_slots, H] @@ -802,7 +809,7 @@ the MLP backend, not as a guaranteed feature of every `combine_async` call. ## Internal metadata exchange -Normal/high-throughput dispatch usually needs a metadata phase before payload +Throughput dispatch needs a metadata phase before payload movement: ```text @@ -813,9 +820,8 @@ topk_ids -> dispatch token payload ``` -Low-latency modes may use fixed-capacity buffers and device-side counters, but -they still generate metadata such as source info, layout ranges, and valid -counts. +Latency algorithms use fixed-capacity buffers and device-side counters, but +still generate source info, layout ranges, and valid counts. These details should remain internal. The user-facing API should only expose MLP-relevant layout information through `DispatchOutput` and combine-relevant diff --git a/python/mscclpp/ep/__init__.py b/python/mscclpp/ep/__init__.py index e30de9b3f..fd8043775 100644 --- a/python/mscclpp/ep/__init__.py +++ b/python/mscclpp/ep/__init__.py @@ -4,15 +4,14 @@ """MSCCL++ Expert-Parallel -``MoECommunicator`` is the public API. ``mode=MoEMode.LOW_LATENCY`` runs on the -LL backend; ``mode=MoEMode.HIGH_THROUGHPUT`` runs on the HT backend (GB200 TMA -direct-gather combine + all-sender dispatch). +`MoECommunicator` is the public API. `mode=MoEMode.LATENCY` selects +latency-optimized algorithms; `mode=MoEMode.THROUGHPUT` selects throughput-optimized bounded-resource +token-major algorithms. """ -from .communicator import ( # noqa: F401 +from mscclpp.ep.communicator import ( BlockOverlapConfig, - CommOverlapConfig, - CombineContext, + OverlapConfig, CombineMode, DispatchHandle, DispatchDataType, @@ -20,23 +19,16 @@ DispatchLayoutInfo, DispatchOutput, DispatchOutputInfo, - ExpertMajorDispatchHandle, - ExpertMajorCombineContext, - HighThroughputDispatchHandle, - HighThroughputCombineContext, MoECommunicator, MoECommunicatorConfig, MoEMode, OperationOverlapConfig, QuantConfig, - RankMajorDispatchHandle, - RankMajorCombineContext, ) __all__ = [ "BlockOverlapConfig", - "CommOverlapConfig", - "CombineContext", + "OverlapConfig", "CombineMode", "DispatchHandle", "DispatchDataType", @@ -44,15 +36,9 @@ "DispatchLayoutInfo", "DispatchOutput", "DispatchOutputInfo", - "ExpertMajorDispatchHandle", - "ExpertMajorCombineContext", - "HighThroughputDispatchHandle", - "HighThroughputCombineContext", "MoECommunicator", "MoECommunicatorConfig", "MoEMode", "OperationOverlapConfig", "QuantConfig", - "RankMajorDispatchHandle", - "RankMajorCombineContext", ] diff --git a/python/mscclpp/ep/_cpp.py b/python/mscclpp/ep/_cpp.py index 21491424f..4dee89f9a 100644 --- a/python/mscclpp/ep/_cpp.py +++ b/python/mscclpp/ep/_cpp.py @@ -4,13 +4,7 @@ from __future__ import annotations -try: - import mscclpp_ep_cpp as _cpp # type: ignore[import-not-found] -except ImportError as exc: # pragma: no cover - raise ImportError( - "mscclpp_ep_cpp is not available. Build mscclpp with " - "-DMSCCLPP_BUILD_EXT_EP=ON or install with `pip install .[ep]`." - ) from exc +from mscclpp import mscclpp_ep_cpp as _cpp DispatchLayout = _cpp.DispatchLayout MoEMode = _cpp.MoEMode @@ -18,4 +12,3 @@ DispatchDataType = _cpp.DispatchDataType MoERuntime = _cpp.MoERuntime create_moe_runtime = _cpp.create_moe_runtime -Config = getattr(_cpp, "Config", None) diff --git a/python/mscclpp/ep/communicator.py b/python/mscclpp/ep/communicator.py index c9c8030dc..29ffd6df3 100644 --- a/python/mscclpp/ep/communicator.py +++ b/python/mscclpp/ep/communicator.py @@ -4,36 +4,28 @@ from __future__ import annotations -from typing import Optional, Tuple +from typing import Any, Optional, Tuple import torch -from ._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode -from .high_throughput import HighThroughputBackend -from .low_latency import LowLatencyBackend -from .types import ( +from mscclpp.ep._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode +from mscclpp.ep.context import create_context +from mscclpp.ep.runtime import Runtime +from mscclpp.ep.types import ( BlockOverlapConfig, - CommOverlapConfig, - CombineContext, + OverlapConfig, DispatchHandle, DispatchLayoutInfo, DispatchOutput, DispatchOutputInfo, - ExpertMajorDispatchHandle, - ExpertMajorCombineContext, - HighThroughputDispatchHandle, - HighThroughputCombineContext, MoECommunicatorConfig, OperationOverlapConfig, QuantConfig, - RankMajorDispatchHandle, - RankMajorCombineContext, ) __all__ = [ - "CommOverlapConfig", + "OverlapConfig", "BlockOverlapConfig", - "CombineContext", "CombineMode", "DispatchHandle", "DispatchDataType", @@ -41,25 +33,20 @@ "DispatchLayoutInfo", "DispatchOutput", "DispatchOutputInfo", - "ExpertMajorDispatchHandle", - "ExpertMajorCombineContext", - "HighThroughputDispatchHandle", - "HighThroughputCombineContext", "MoECommunicator", "MoECommunicatorConfig", "MoEMode", "OperationOverlapConfig", "QuantConfig", - "RankMajorDispatchHandle", - "RankMajorCombineContext", ] class MoECommunicator: """High-level MoE communicator for dispatch/combine. - ``mode=MoEMode.LOW_LATENCY`` selects the LL backend (EXPERT_MAJOR by default); - ``mode=MoEMode.HIGH_THROUGHPUT`` selects the HT backend (TOKEN_MAJOR). + `mode=MoEMode.LATENCY` selects the latency algorithms (EXPERT_MAJOR by + default); `mode=MoEMode.THROUGHPUT` selects bounded-resource throughput + algorithms (TOKEN_MAJOR). """ def __init__(self, config: Optional[MoECommunicatorConfig] = None, **kwargs) -> None: @@ -75,40 +62,77 @@ def __init__(self, config: Optional[MoECommunicatorConfig] = None, **kwargs) -> raise TypeError("MoECommunicatorConfig.mode must be a MoEMode") _validate_common_config(config) - self.mode = config.mode - self.output_layout = _resolve_output_layout(config.output_layout, self.mode) - if self.mode == MoEMode.LOW_LATENCY: - self._backend = LowLatencyBackend(config, self.output_layout) - else: - self._backend = HighThroughputBackend(config, self.output_layout) - self._publish_backend_state() - - def _publish_backend_state(self) -> None: - for name in ( - "comm", - "rank", - "world_size", - "local_rank", - "device", - "num_experts", - "hidden_size", - "topk", - "max_tokens_per_rank", - "num_sms", - "enable_overlap", - "num_local_experts", - "local_expert_start", - ): - setattr(self, name, getattr(self._backend, name)) + self._context = create_context(config) + self._runtime = Runtime.create(self._context) + + @property + def comm(self) -> Any: + return self._context.comm + + @property + def rank(self) -> int: + return self._context.rank + + @property + def world_size(self) -> int: + return self._context.world_size + + @property + def local_rank(self) -> int: + return self._context.local_rank + + @property + def device(self) -> torch.device: + return self._context.device + + @property + def mode(self) -> Any: + return self._context.mode + + @property + def output_layout(self) -> Any: + return self._context.output_layout + + @property + def num_experts(self) -> int: + return self._context.num_experts + + @property + def hidden_size(self) -> int: + return self._context.hidden_size + + @property + def topk(self) -> int: + return self._context.topk + + @property + def max_tokens_per_rank(self) -> int: + return self._context.max_tokens_per_rank + + @property + def num_blocks(self) -> int: + return self._context.num_blocks + + @property + def enable_overlap(self) -> bool: + return self._context.enable_overlap + + @property + def num_local_experts(self) -> int: + return self._context.num_local_experts + + @property + def local_expert_start(self) -> int: + return self._context.local_expert_start def is_available(self) -> bool: - return self._backend.is_available() + return self._runtime.is_available() def is_internode_available(self) -> bool: - return self._backend.is_internode_available() + return self._runtime.is_internode_available() def is_internode(self) -> bool: - return self._backend.is_internode() + return self._runtime.is_internode_available() def dispatch( self, @@ -122,7 +146,7 @@ def dispatch( previous_handle: Optional[DispatchHandle] = None, runtime_max_tokens_per_rank: Optional[int] = None, ) -> Tuple[DispatchOutput, DispatchHandle]: - return self._backend.dispatch( + return self._runtime.dispatch( input, topk_ids, weights, @@ -141,19 +165,7 @@ def combine( out: Optional[torch.Tensor] = None, stream: Optional[torch.cuda.Stream] = None, ) -> torch.Tensor: - return self._backend.combine(expert_output, handle, out=out, stream=stream) - - def get_expert_output_buffer(self) -> torch.Tensor: - """Return the runtime-owned rank-major MoE output buffer. - - This aliases runtime memory that every combine reuses; it is not a fresh - allocation per call. Fill it before each combine and copy out anything - that must outlive the next call. - """ - buffer = getattr(self._backend, "expert_output_buffer", None) - if buffer is None: - raise RuntimeError("expert output buffer is only available for RANK_MAJOR low-latency mode") - return buffer + return self._runtime.combine(expert_output, handle, out=out, stream=stream) def dispatch_async(self, *args, **kwargs): raise NotImplementedError("dispatch_async is not implemented for MoECommunicator yet") @@ -163,24 +175,18 @@ def combine_async(self, *args, **kwargs): def create_overlap_config( self, op: str, *, handle: Optional[DispatchHandle] = None, level: str = "op" - ) -> CommOverlapConfig: + ) -> OverlapConfig: if op not in ("dispatch", "combine"): raise ValueError("op must be 'dispatch' or 'combine'") if level != "op": raise NotImplementedError("block-level overlap is not implemented yet") if op == "combine" and handle is None: raise ValueError("combine overlap config requires a DispatchHandle") - return CommOverlapConfig(operation=OperationOverlapConfig()) + return OverlapConfig(operation=OperationOverlapConfig()) def _validate_common_config(config: MoECommunicatorConfig) -> None: if config.num_experts <= 0 or config.hidden_size <= 0 or config.topk <= 0 or config.max_tokens_per_rank <= 0: raise ValueError("num_experts, hidden_size, topk, and max_tokens_per_rank must be positive") - - -def _resolve_output_layout(layout: Optional[DispatchLayout], mode: MoEMode) -> DispatchLayout: - if layout is None: - return DispatchLayout.EXPERT_MAJOR if mode == MoEMode.LOW_LATENCY else DispatchLayout.TOKEN_MAJOR - if not isinstance(layout, DispatchLayout): + if config.output_layout is not None and not isinstance(config.output_layout, DispatchLayout): raise TypeError("MoECommunicatorConfig.output_layout must be a DispatchLayout") - return layout diff --git a/python/mscclpp/ep/context.py b/python/mscclpp/ep/context.py new file mode 100644 index 000000000..51b1e17ec --- /dev/null +++ b/python/mscclpp/ep/context.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Mode-specific context for the high-level expert-parallel communicator.""" + +from __future__ import annotations + +from mscclpp.ep._cpp import MoEMode +from mscclpp.ep.types import MoECommunicatorConfig + + +class Context: + """Persistent mode-specific configuration, buffers, and metadata.""" + + +def create_context(config: MoECommunicatorConfig) -> Context: + """Construct the context selected by ``config.mode``.""" + if config.mode == MoEMode.LATENCY: + from mscclpp.ep.latency import LatencyContext + + return LatencyContext(config) + if config.mode == MoEMode.THROUGHPUT: + from mscclpp.ep.throughput import ThroughputContext + + return ThroughputContext(config) + raise ValueError(f"unsupported MoE mode: {config.mode}") diff --git a/python/mscclpp/ep/high_throughput.py b/python/mscclpp/ep/high_throughput.py deleted file mode 100644 index dc87e25e6..000000000 --- a/python/mscclpp/ep/high_throughput.py +++ /dev/null @@ -1,530 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. -# -# Portions adapted from DeepEP (https://github.com/deepseek-ai/DeepEP), -# branch ``chhwang/dev-atomic-add-cleanup``. Licensed under the MIT License. -"""Fabric-domain high-throughput backend for the high-level MoE communicator. - -The C++ runtime follows the low-latency resource model: it reuses the existing -MSCCL++ communicator and writes directly into peer receive pools through a -torch-free raw-pointer boundary. Dynamic receive sizing uses a two-phase -``notify_dispatch`` then ``dispatch`` protocol. Cached dispatches reuse the -previous routing matrices and receive count. -""" - -from __future__ import annotations - -from typing import Any, List, Optional - -import torch - -from ._cpp import Config, DispatchLayout, MoEMode, _cpp -from .types import ( - DispatchHandle, - DispatchLayoutInfo, - DispatchOutput, - DispatchOutputInfo, - HighThroughputCombineContext, - HighThroughputDispatchHandle, - MoECommunicatorConfig, - QuantConfig, -) -from .utils import ( - bf16_view as _bf16_view, - current_stream_ptr as _stream_ptr, - ptr as _ptr, - resolve_expert_placement, -) - - -class HighThroughputRuntime: - """Core high-throughput expert-parallel (EP) communication runtime. - - ``comm`` provides the initialized MSCCL++ communicator used to exchange and - map the intranode physical symmetric buffers. - """ - - #: Default number of SMs reserved for comms kernels. Matches DeepEP. - num_sms: int = 20 - - def __init__( - self, - comm: Any, - max_hidden_bytes: int, - config: Config, - ) -> None: - self.rank: int = comm.my_rank - self.group_size: int = comm.nranks - self.comm = comm - self.runtime = _cpp.create_moe_runtime( - comm.communicator, - _cpp.MoEMode.HIGH_THROUGHPUT, - max_hidden_bytes=max_hidden_bytes, - num_sms=config.num_sms, - ) - - # ------------------------------------------------------------------ - # Sanity helpers - # ------------------------------------------------------------------ - - def is_available(self) -> bool: - return self.runtime.is_available() - - def is_internode_available(self) -> bool: - return self.runtime.is_internode_available() - - # ------------------------------------------------------------------ - # Dispatch routing metadata - # ------------------------------------------------------------------ - - def compute_dispatch_counts(self, topk_idx: torch.Tensor, num_experts: int): - """Return per-rank, per-expert, and token-membership routing metadata. - - This is routing metadata consumed by dispatch; it is unrelated to - ``DispatchLayout`` (the memory layout of the dispatch output). - """ - assert topk_idx.dim() == 2 and topk_idx.is_contiguous() - num_tokens, num_topk = int(topk_idx.size(0)), int(topk_idx.size(1)) - - num_tokens_per_rank = torch.empty((self.group_size,), dtype=torch.int32, device="cuda") - num_tokens_per_expert = torch.empty((num_experts,), dtype=torch.int32, device="cuda") - is_token_in_rank = torch.empty((num_tokens, self.group_size), dtype=torch.bool, device="cuda") - - self.runtime.ht_compute_dispatch_counts( - _ptr(num_tokens_per_rank), - _ptr(num_tokens_per_expert), - _ptr(is_token_in_rank), - _ptr(topk_idx), - num_tokens, - num_topk, - num_experts, - _stream_ptr(), - ) - return num_tokens_per_rank, num_tokens_per_expert, is_token_in_rank - - # ------------------------------------------------------------------ - # Dispatch (two-phase) + combine - # ------------------------------------------------------------------ - - def dispatch( - self, - x: torch.Tensor, - x_scales: Optional[torch.Tensor], - topk_idx: Optional[torch.Tensor], - topk_weights: Optional[torch.Tensor], - num_tokens_per_rank: Optional[torch.Tensor], - is_token_in_rank: torch.Tensor, - num_tokens_per_expert: Optional[torch.Tensor], - cached_num_recv_tokens: int, - cached_rank_prefix_matrix: Optional[torch.Tensor], - cached_channel_prefix_matrix: Optional[torch.Tensor], - expert_alignment: int, - ): - """Run high-throughput dispatch and return outputs plus combine metadata.""" - assert x.dim() == 2 and x.is_contiguous() - cached_mode = cached_rank_prefix_matrix is not None - num_tokens, hidden = int(x.size(0)), int(x.size(1)) - x_element_size = x.element_size() - num_channels = self.runtime.ht_get_dispatch_num_channels(x_element_size) - - num_topk = int(topk_idx.size(1)) if topk_idx is not None else 0 - num_scales = 0 - if x_scales is not None: - num_scales = 1 if x_scales.dim() == 1 else int(x_scales.size(1)) - - # ----- Phase A: notify (non-cached) or reuse cached layout ----- - if cached_mode: - num_recv_tokens = cached_num_recv_tokens - rank_prefix_matrix = cached_rank_prefix_matrix - channel_prefix_matrix = cached_channel_prefix_matrix - num_recv_tokens_per_expert_list: List[int] = [] - num_experts = 0 - else: - assert num_tokens_per_rank is not None and num_tokens_per_expert is not None - num_experts = int(num_tokens_per_expert.size(0)) - num_local_experts = num_experts // self.group_size - rank_prefix_matrix = torch.empty((self.group_size, self.group_size), dtype=torch.int32, device="cuda") - channel_prefix_matrix = torch.empty((self.group_size, num_channels), dtype=torch.int32, device="cuda") - num_recv_per_expert_host = torch.empty((num_local_experts,), dtype=torch.int32, device="cpu") - num_recv_tokens = self.runtime.ht_notify_dispatch( - _ptr(rank_prefix_matrix), - _ptr(channel_prefix_matrix), - _ptr(num_recv_per_expert_host), - _ptr(num_tokens_per_rank), - _ptr(num_tokens_per_expert), - _ptr(is_token_in_rank), - num_tokens, - num_experts, - x_element_size, - expert_alignment, - _stream_ptr(), - ) - num_recv_tokens_per_expert_list = num_recv_per_expert_host.tolist() - - # ----- Phase B: allocate recv outputs (or view the recv pool) ----- - recv_x = self._alloc_recv_x(num_tokens, num_recv_tokens, hidden, x_element_size) - send_head = torch.empty((num_tokens, self.group_size), dtype=torch.int32, device="cuda") - recv_topk_idx = ( - torch.empty((num_recv_tokens, num_topk), dtype=torch.int64, device="cuda") if topk_idx is not None else None - ) - recv_topk_weights = ( - torch.empty((num_recv_tokens, num_topk), dtype=torch.float32, device="cuda") - if topk_weights is not None - else None - ) - recv_x_scales = ( - torch.empty((num_recv_tokens, num_scales), dtype=torch.float32, device="cuda") - if x_scales is not None - else None - ) - - self.runtime.ht_dispatch( - _ptr(recv_x), - _ptr(recv_x_scales), - _ptr(recv_topk_idx), - _ptr(recv_topk_weights), - _ptr(send_head), - _ptr(x), - _ptr(x_scales), - _ptr(topk_idx), - _ptr(topk_weights), - _ptr(is_token_in_rank), - _ptr(rank_prefix_matrix), - _ptr(channel_prefix_matrix), - num_tokens, - hidden, - num_topk, - num_scales, - num_experts, - x_element_size, - num_recv_tokens, - cached_mode, - _stream_ptr(), - ) - return ( - recv_x, - recv_x_scales, - recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, - rank_prefix_matrix, - channel_prefix_matrix, - send_head, - ) - - def _alloc_recv_x(self, num_tokens: int, num_recv_tokens: int, hidden: int, x_element_size: int) -> torch.Tensor: - """Return this rank's direct receive-pool view.""" - pool_ptr = self.runtime.ht_resolve_recv_x_buffer(num_tokens, num_recv_tokens, hidden, x_element_size) - if pool_ptr == 0: - raise RuntimeError("high-throughput direct receive-pool capacity exceeded") - return _bf16_view(pool_ptr, num_recv_tokens, hidden, owner=self) - - def combine( - self, - x: torch.Tensor, - topk_weights: Optional[torch.Tensor], - send_head: torch.Tensor, - ): - """Returns ``(combined_x, combined_topk_weights|None)``.""" - assert x.dim() == 2 and x.is_contiguous() - num_tokens, hidden = int(x.size(0)), int(x.size(1)) - num_recv_tokens = int(send_head.size(0)) - num_topk = int(topk_weights.size(1)) if topk_weights is not None else 0 - combined_x = torch.empty((num_recv_tokens, hidden), dtype=torch.bfloat16, device="cuda") - combined_topk_weights = ( - torch.empty((num_recv_tokens, num_topk), dtype=torch.float32, device="cuda") - if topk_weights is not None - else None - ) - self.runtime.ht_combine( - _ptr(combined_x), - _ptr(combined_topk_weights), - _ptr(x), - _ptr(topk_weights), - _ptr(send_head), - num_tokens, - num_recv_tokens, - hidden, - num_topk, - x.element_size(), - _stream_ptr(), - ) - return combined_x, combined_topk_weights - - -class HighThroughputBackend: - """Backend implementation for ``MoEMode.HIGH_THROUGHPUT``.""" - - def __init__(self, config: MoECommunicatorConfig, output_layout: DispatchLayout) -> None: - comm = config.comm - if comm is None: - raise ValueError("mode=HIGH_THROUGHPUT requires an mscclpp.CommGroup via comm=") - - self.comm = comm - self.rank = comm.my_rank - self.world_size = comm.nranks - self.local_rank = torch.cuda.current_device() - self.device = torch.device("cuda", self.local_rank) - self.mode = MoEMode.HIGH_THROUGHPUT - self.output_layout = output_layout - - self.num_experts = config.num_experts - 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.enable_overlap = config.enable_overlap - - if self.output_layout != DispatchLayout.TOKEN_MAJOR: - raise NotImplementedError("HT mode currently supports only DispatchLayout.TOKEN_MAJOR") - if config.invalid_token_expert_id is not None: - raise ValueError("invalid_token_expert_id is only supported in low-latency mode") - - self.num_local_experts, self.local_expert_start = resolve_expert_placement( - num_experts=self.num_experts, - world_size=self.world_size, - rank=self.rank, - num_local_experts=config.num_local_experts, - local_expert_start=config.local_expert_start, - ) - - if config.quant is not None: - raise NotImplementedError("HT quantized dispatch (scales) is not implemented yet") - - self.expert_alignment = config.expert_alignment - self._cfg = Config(self.num_sms) - hidden_bytes = self.hidden_size * torch.empty((), dtype=torch.bfloat16).element_size() - self._runtime = HighThroughputRuntime( - comm, - max_hidden_bytes=hidden_bytes, - config=self._cfg, - ) - - def is_available(self) -> bool: - return self._runtime.is_available() - - def is_internode_available(self) -> bool: - return self._runtime.is_internode_available() - - def is_internode(self) -> bool: - return self._runtime.is_internode_available() - - def dispatch( - self, - input: torch.Tensor, - topk_ids: torch.Tensor, - weights: Optional[torch.Tensor], - quant: Optional[QuantConfig], - *, - output_buffer: Optional[torch.Tensor], - stream: Optional[torch.cuda.Stream], - previous_handle: Optional[DispatchHandle], - runtime_max_tokens_per_rank: Optional[int], - ) -> tuple[DispatchOutput, DispatchHandle]: - del output_buffer - if runtime_max_tokens_per_rank is not None: - raise ValueError("runtime_max_tokens_per_rank is only supported by low-latency rank-major dispatch") - if stream is not None: - with torch.cuda.stream(stream): - return self._dispatch(input, topk_ids, weights, quant, previous_handle) - return self._dispatch(input, topk_ids, weights, quant, previous_handle) - - def _dispatch( - self, - input: torch.Tensor, - topk_ids: torch.Tensor, - weights: Optional[torch.Tensor], - quant: Optional[QuantConfig], - previous_handle: Optional[DispatchHandle], - ) -> tuple[DispatchOutput, DispatchHandle]: - self._validate_dispatch_inputs(input, topk_ids, weights, quant) - implicit_weights = weights is None - if weights is None: - weights = torch.ones(topk_ids.shape, dtype=torch.float32, device=topk_ids.device) - - cache = getattr(previous_handle, "_dispatch_cache", None) if previous_handle is not None else None - if cache is not None and not self._cache_matches(cache, input, topk_ids, weights, implicit_weights): - cache = None - if cache is not None: - num_tokens_per_rank = cache["num_tokens_per_rank"] - num_tokens_per_expert = cache["num_tokens_per_expert"] - is_token_in_rank = cache["is_token_in_rank"] - else: - ( - num_tokens_per_rank, - num_tokens_per_expert, - is_token_in_rank, - ) = self._runtime.compute_dispatch_counts(topk_ids, self.num_experts) - - if cache is not None: - ( - recv_x, - _recv_x_scales, - _runtime_recv_topk_idx, - _runtime_recv_topk_weights, - _runtime_num_recv_tokens_per_expert_list, - rank_prefix_matrix, - _channel_prefix_matrix, - send_head, - ) = self._runtime.dispatch( - input, - None, - None, - None, - None, - is_token_in_rank, - None, - cache["num_recv_tokens"], - cache["rank_prefix_matrix"], - cache["channel_prefix_matrix"], - self.expert_alignment, - ) - del _runtime_recv_topk_idx, _runtime_recv_topk_weights, _runtime_num_recv_tokens_per_expert_list - recv_topk_idx = cache["recv_topk_idx"] - recv_topk_weights = cache["recv_topk_weights"] - num_recv_tokens_per_expert_list = cache["num_recv_tokens_per_expert_list"] - combine_context = HighThroughputCombineContext( - recv_topk_weights=recv_topk_weights, - send_head=send_head, - ) - dispatch_cache = cache - else: - ( - recv_x, - _recv_x_scales, - recv_topk_idx, - recv_topk_weights, - num_recv_tokens_per_expert_list, - rank_prefix_matrix, - channel_prefix_matrix, - send_head, - ) = self._runtime.dispatch( - input, - None, - topk_ids, - weights, - num_tokens_per_rank, - is_token_in_rank, - num_tokens_per_expert, - 0, - None, - None, - self.expert_alignment, - ) - combine_context = HighThroughputCombineContext( - recv_topk_weights=recv_topk_weights, - send_head=send_head, - ) - dispatch_cache = { - "num_tokens_per_rank": num_tokens_per_rank, - "num_tokens_per_expert": num_tokens_per_expert, - "is_token_in_rank": is_token_in_rank, - "rank_prefix_matrix": rank_prefix_matrix, - "channel_prefix_matrix": channel_prefix_matrix, - "num_recv_tokens": int(recv_x.size(0)), - "recv_topk_idx": recv_topk_idx, - "recv_topk_weights": recv_topk_weights, - "num_recv_tokens_per_expert_list": num_recv_tokens_per_expert_list, - "backend_id": id(self), - "num_tokens": int(input.size(0)), - "device": input.device, - "topk_ids_ptr": topk_ids.data_ptr(), - "topk_ids_version": topk_ids._version, - "implicit_weights": implicit_weights, - "weights_ptr": 0 if implicit_weights else weights.data_ptr(), - "weights_version": 0 if implicit_weights else weights._version, - } - - output_info = DispatchOutputInfo( - layout=DispatchLayoutInfo( - kind=self.output_layout, - num_tokens_per_expert=num_recv_tokens_per_expert_list, - ), - quant=None, - ) - dispatch_out = DispatchOutput( - tokens=recv_x, - quant=output_info.quant, - layout=output_info.layout, - topk_ids=recv_topk_idx, - weights=recv_topk_weights, - ) - handle = HighThroughputDispatchHandle(output_info=output_info, combine_context=combine_context) - # The torch-free HT runtime orders its work on the caller's CUDA stream - # (no separate event handle), so there is nothing to attach here. - handle._event = None # type: ignore[attr-defined] - handle._dispatch_cache = dispatch_cache # type: ignore[attr-defined] - return dispatch_out, handle - - def _cache_matches(self, cache, input, topk_ids, weights, implicit_weights) -> bool: - return ( - cache.get("backend_id") == id(self) - and cache.get("num_tokens") == int(input.size(0)) - and cache.get("device") == input.device - and cache.get("topk_ids_ptr") == topk_ids.data_ptr() - and cache.get("topk_ids_version") == topk_ids._version - and cache.get("implicit_weights") == implicit_weights - and (implicit_weights or cache.get("weights_ptr") == weights.data_ptr()) - and (implicit_weights or cache.get("weights_version") == weights._version) - ) - - def combine( - self, - expert_output: torch.Tensor, - handle: DispatchHandle, - *, - out: Optional[torch.Tensor], - stream: Optional[torch.cuda.Stream], - ) -> torch.Tensor: - if stream is not None: - with torch.cuda.stream(stream): - return self._combine(expert_output, handle, out) - return self._combine(expert_output, handle, out) - - def _combine( - self, expert_output: torch.Tensor, handle: DispatchHandle, out: Optional[torch.Tensor] - ) -> torch.Tensor: - self._validate_combine_inputs(expert_output, handle) - context = handle.combine_context - combined_x, _combined_w = self._runtime.combine( - expert_output, - context.recv_topk_weights, - context.send_head, - ) - if out is not None: - out.copy_(combined_x) - return out - return combined_x - - def _validate_dispatch_inputs(self, input, topk_ids, weights, quant) -> None: - if quant is not None: - raise NotImplementedError("HT dispatch does not support quantized input scales yet") - if input.dim() != 2 or not input.is_contiguous(): - raise ValueError("input must be a contiguous [num_tokens, hidden] tensor") - if input.device.type != "cuda" or input.dtype != torch.bfloat16: - raise ValueError("HT dispatch input must be a CUDA BF16 tensor") - if input.size(1) != self.hidden_size: - raise ValueError(f"input hidden size {input.size(1)} != configured {self.hidden_size}") - if input.size(0) > self.max_tokens_per_rank: - raise ValueError("input token count exceeds max_tokens_per_rank") - if topk_ids.dim() != 2 or not topk_ids.is_contiguous(): - raise ValueError("topk_ids must be a contiguous [num_tokens, topk] tensor") - if topk_ids.device != input.device or topk_ids.dtype != torch.int64: - raise ValueError("topk_ids must be an int64 CUDA tensor on the same device as input") - if topk_ids.shape != (input.size(0), self.topk): - raise ValueError("topk_ids shape must be [input.size(0), topk]") - if weights is not None: - if weights.dim() != 2 or not weights.is_contiguous(): - raise ValueError("weights must be a contiguous [num_tokens, topk] tensor") - if weights.device != input.device or weights.dtype != torch.float32: - raise ValueError("weights must be a float32 CUDA tensor on the same device as input") - if weights.shape != topk_ids.shape: - raise ValueError("weights shape must match topk_ids") - - def _validate_combine_inputs(self, expert_output, handle) -> None: - if not isinstance(handle, HighThroughputDispatchHandle): - raise TypeError("handle must be a DispatchHandle returned by dispatch") - if expert_output.dim() != 2 or not expert_output.is_contiguous(): - raise ValueError("expert_output must be a contiguous [total_recv_tokens, hidden] tensor") - if expert_output.size(1) != self.hidden_size: - raise ValueError(f"expert_output hidden size {expert_output.size(1)} != configured {self.hidden_size}") diff --git a/python/mscclpp/ep/latency.py b/python/mscclpp/ep/latency.py new file mode 100644 index 000000000..ada5313fd --- /dev/null +++ b/python/mscclpp/ep/latency.py @@ -0,0 +1,502 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Latency-mode context.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from mscclpp.ep._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode, create_moe_runtime +from mscclpp.ep.context import Context +from mscclpp.ep.runtime import Runtime +from mscclpp.ep.types import ( + DispatchHandle, + DispatchLayoutInfo, + DispatchOutput, + DispatchOutputInfo, + MoECommunicatorConfig, + QuantConfig, + _ExpertMajorCombineContext, + _RankMajorCombineContext, +) +from mscclpp.ep.utils import ( + DevicePointerArray, + cuda_stream_ptr, + dispatch_scale_block_size, + dispatch_scale_dtype, + resolve_expert_placement, + resolve_dispatch_data_type, + tensor_from_pointer, +) + + +class LatencyContext(Context): + """Latency-mode context.""" + + def __init__(self, config: MoECommunicatorConfig) -> None: + comm = config.comm + if comm is None: + raise ValueError("mode=LATENCY requires an mscclpp.CommGroup via comm=") + output_layout = config.output_layout + if output_layout is None: + output_layout = DispatchLayout.EXPERT_MAJOR + num_blocks = 130 if config.num_blocks is None else config.num_blocks + + self.comm = comm + self.rank = comm.my_rank + self.world_size = comm.nranks + self.local_rank = torch.cuda.current_device() + self.device = torch.device("cuda", self.local_rank) + self.mode = MoEMode.LATENCY + self.output_layout = output_layout + + self.num_experts = config.num_experts + self.hidden_size = config.hidden_size + self.topk = config.topk + self.max_tokens_per_rank = config.max_tokens_per_rank + self.num_blocks = num_blocks + self.combine_mode = config.combine_mode + self.invalid_token_expert_id = ( + self.num_experts if config.invalid_token_expert_id is None else config.invalid_token_expert_id + ) + self.enable_overlap = config.enable_overlap + + if self.output_layout not in ( + DispatchLayout.EXPERT_MAJOR, + DispatchLayout.RANK_MAJOR, + ): + raise NotImplementedError("unsupported latency output layout") + if self.num_experts % self.world_size != 0: + raise ValueError("latency mode requires num_experts divisible by world_size") + if not self.world_size + 2 <= self.num_blocks <= 130: + raise ValueError("num_blocks must be between world_size + 2 and 130 in latency mode") + if not isinstance(self.combine_mode, CombineMode): + raise TypeError("combine_mode must be a CombineMode") + if type(self.invalid_token_expert_id) is not int: + raise TypeError("invalid_token_expert_id must be an int or None") + if not -(1 << 31) <= self.invalid_token_expert_id < (1 << 31): + raise ValueError("invalid_token_expert_id must fit in int32") + if 0 <= self.invalid_token_expert_id < self.num_experts: + raise ValueError("invalid_token_expert_id must not overlap a valid global expert ID") + if self.output_layout == DispatchLayout.RANK_MAJOR: + if self.combine_mode != CombineMode.RANK_LOCAL_REDUCE: + raise ValueError("RANK_MAJOR output requires RANK_LOCAL_REDUCE combine") + if self.enable_overlap: + raise NotImplementedError("RANK_MAJOR output does not support overlapping calls yet") + + self.num_local_experts, self.local_expert_start = resolve_expert_placement( + num_experts=self.num_experts, + world_size=self.world_size, + rank=self.rank, + num_local_experts=config.num_local_experts, + local_expert_start=config.local_expert_start, + ) + + self.dispatch_data_type = resolve_dispatch_data_type(config.quant) + if self.output_layout == DispatchLayout.RANK_MAJOR and self.dispatch_data_type != DispatchDataType.BF16: + raise NotImplementedError("RANK_MAJOR output currently supports BF16 dispatch only") + + self._dispatch_scales: Optional[torch.Tensor] = None + self._dispatch_src_info: Optional[torch.Tensor] = None + self._dispatch_topk_ids: Optional[torch.Tensor] = None + self._dispatch_weights: Optional[torch.Tensor] = None + self._dispatch_layout_range: Optional[torch.Tensor] = None + self._dispatch_count: Optional[torch.Tensor] = None + + self._dispatch_output_owner: Optional[DevicePointerArray] = None + self._combine_input_owner: Optional[DevicePointerArray] = None + self._output_topk_ids_owner: Optional[DevicePointerArray] = None + self._output_topk_weights_owner: Optional[DevicePointerArray] = None + self._output_topk_ids: Optional[torch.Tensor] = None + self._output_topk_weights: Optional[torch.Tensor] = None + self.dispatch_output_buffer: Optional[torch.Tensor] = None + self.combine_input_buffer: Optional[torch.Tensor] = None + + +class LatencyRuntime(Runtime): + """Latency-optimized runtime.""" + + context: LatencyContext + + def __init__(self, context: LatencyContext) -> None: + cpp_runtime = create_moe_runtime( + context.comm.communicator, + context.mode, + max_tokens_per_rank=context.max_tokens_per_rank, + hidden=context.hidden_size, + num_experts=context.num_experts, + num_topk=context.topk, + num_blocks=context.num_blocks, + output_layout=context.output_layout, + ) + super().__init__(context, cpp_runtime) + self._bind_buffers() + + def dispatch( + self, + input: torch.Tensor, + topk_ids: torch.Tensor, + weights: Optional[torch.Tensor], + quant: Optional[QuantConfig], + *, + output_buffer: Optional[torch.Tensor], + stream: Optional[torch.cuda.Stream], + previous_handle: Optional[DispatchHandle], + runtime_max_tokens_per_rank: Optional[int], + ) -> tuple[DispatchOutput, DispatchHandle]: + mode_context = self.context + del previous_handle + active_capacity = self._resolve_capacity(runtime_max_tokens_per_rank) + if output_buffer is None: + assert mode_context.dispatch_output_buffer is not None + output_buffer = mode_context.dispatch_output_buffer + self._validate_dispatch(input, topk_ids, weights, quant, output_buffer, active_capacity) + + out_buf, scales, src_info, recv_topk_ids, recv_weights, layout_range, count = self._dispatch_outputs( + output_buffer + ) + self.cpp_runtime.dispatch( + input.data_ptr(), + topk_ids.data_ptr(), + 0 if weights is None else weights.data_ptr(), + out_buf.data_ptr(), + 0 if scales is None else scales.data_ptr(), + 0 if src_info is None else src_info.data_ptr(), + 0 if recv_topk_ids is None else recv_topk_ids.data_ptr(), + 0 if recv_weights is None else recv_weights.data_ptr(), + 0 if layout_range is None else layout_range.data_ptr(), + count.data_ptr(), + input.size(0), + mode_context.hidden_size, + mode_context.topk, + active_capacity, + mode_context.num_experts, + mode_context.invalid_token_expert_id, + mode_context.output_layout, + mode_context.dispatch_data_type, + mode_context.num_blocks, + cuda_stream_ptr(stream), + ) + output_quant = ( + None + if scales is None + else QuantConfig( + format=mode_context.dispatch_data_type, + block_scales=scales, + ) + ) + if mode_context.output_layout == DispatchLayout.EXPERT_MAJOR: + layout_info = DispatchLayoutInfo(kind=mode_context.output_layout, num_tokens_per_expert=count) + elif mode_context.output_layout == DispatchLayout.RANK_MAJOR: + layout_info = DispatchLayoutInfo( + kind=mode_context.output_layout, + num_tokens_per_rank=count, + ) + else: + raise ValueError(f"unsupported latency output layout: {mode_context.output_layout}") + output_info = DispatchOutputInfo(layout=layout_info, quant=output_quant) + dispatch_out = DispatchOutput( + tokens=out_buf, + quant=output_info.quant, + layout=output_info.layout, + topk_ids=recv_topk_ids, + weights=recv_weights, + combine_input_buffer=( + mode_context.combine_input_buffer if mode_context.output_layout == DispatchLayout.RANK_MAJOR else None + ), + ) + if mode_context.output_layout == DispatchLayout.EXPERT_MAJOR: + assert layout_range is not None + assert src_info is not None + handle = DispatchHandle( + output_info=output_info, + _context=_ExpertMajorCombineContext( + topk_ids=topk_ids, + weights=weights, + num_experts=mode_context.num_experts, + num_tokens=input.size(0), + hidden_size=mode_context.hidden_size, + src_info=src_info, + layout_range=layout_range, + ), + ) + elif mode_context.output_layout == DispatchLayout.RANK_MAJOR: + handle = DispatchHandle( + output_info=output_info, + _context=_RankMajorCombineContext( + topk_ids=topk_ids, + num_experts=mode_context.num_experts, + num_tokens=input.size(0), + hidden_size=mode_context.hidden_size, + max_tokens_per_rank=active_capacity, + ), + ) + else: + raise ValueError(f"unsupported latency output layout: {mode_context.output_layout}") + return dispatch_out, handle + + def combine( + self, + expert_output: torch.Tensor, + handle: DispatchHandle, + *, + out: Optional[torch.Tensor], + stream: Optional[torch.cuda.Stream], + ) -> torch.Tensor: + mode_context = self.context + self._validate_combine(expert_output, handle, out) + context = handle._context + if isinstance(context, _ExpertMajorCombineContext): + topk_weights = context.weights + src_info = context.src_info + layout_range = context.layout_range + active_capacity = mode_context.max_tokens_per_rank + elif isinstance(context, _RankMajorCombineContext): + active_capacity = context.max_tokens_per_rank + topk_weights = None + src_info = None + layout_range = None + else: + raise ValueError("DispatchHandle does not contain latency combine context") + if out is None: + out = torch.empty( + (context.num_tokens, mode_context.hidden_size), + dtype=torch.bfloat16, + device=expert_output.device, + ) + self.cpp_runtime.combine( + expert_output.data_ptr(), + context.topk_ids.data_ptr(), + 0 if topk_weights is None else topk_weights.data_ptr(), + 0 if src_info is None else src_info.data_ptr(), + 0 if layout_range is None else layout_range.data_ptr(), + out.data_ptr(), + context.num_tokens, + mode_context.hidden_size, + mode_context.topk, + active_capacity, + context.num_experts, + mode_context.output_layout, + mode_context.dispatch_data_type, + mode_context.combine_mode, + mode_context.num_blocks - 2, + cuda_stream_ptr(stream), + ) + return out + + def _bind_buffers(self) -> None: + """Create tensor views over runtime-owned latency buffers.""" + context = self.context + if context.output_layout == DispatchLayout.EXPERT_MAJOR: + dispatch_shape = ( + context.num_local_experts, + context.world_size * context.max_tokens_per_rank, + context.hidden_size, + ) + else: + dispatch_shape = (context.world_size * context.max_tokens_per_rank, context.hidden_size) + + dispatch_dtype = torch.bfloat16 if context.dispatch_data_type == DispatchDataType.BF16 else torch.float8_e4m3fn + context._dispatch_output_owner, context.dispatch_output_buffer = tensor_from_pointer( + self.cpp_runtime.dispatch_output_buffer_ptr(), + dispatch_shape, + dispatch_dtype, + context.device, + self.cpp_runtime, + ) + + if context.output_layout != DispatchLayout.RANK_MAJOR: + return + + metadata_shape = (context.world_size * context.max_tokens_per_rank, context.topk) + context._output_topk_ids_owner, context._output_topk_ids = tensor_from_pointer( + self.cpp_runtime.output_topk_ids_buffer_ptr(), + metadata_shape, + torch.int32, + context.device, + self.cpp_runtime, + ) + context._output_topk_weights_owner, context._output_topk_weights = tensor_from_pointer( + self.cpp_runtime.output_topk_weights_buffer_ptr(), + metadata_shape, + torch.float32, + context.device, + self.cpp_runtime, + ) + context._combine_input_owner, context.combine_input_buffer = tensor_from_pointer( + self.cpp_runtime.combine_input_buffer_ptr(), + dispatch_shape, + torch.bfloat16, + context.device, + self.cpp_runtime, + ) + + def _resolve_capacity(self, runtime_max_tokens_per_rank: Optional[int]) -> int: + mode_context = self.context + resolved = ( + mode_context.max_tokens_per_rank if runtime_max_tokens_per_rank is None else runtime_max_tokens_per_rank + ) + if type(resolved) is not int or not 0 < resolved <= mode_context.max_tokens_per_rank: + raise ValueError("runtime_max_tokens_per_rank must be positive and not exceed max_tokens_per_rank") + if mode_context.output_layout != DispatchLayout.RANK_MAJOR and resolved != mode_context.max_tokens_per_rank: + raise ValueError("runtime_max_tokens_per_rank is only supported by rank-major dispatch") + return resolved + + def _dispatch_outputs(self, output_buffer: torch.Tensor): + mode_context = self.context + device = output_buffer.device + slots_per_expert = mode_context.world_size * mode_context.max_tokens_per_rank + if mode_context._dispatch_count is None or mode_context._dispatch_count.device != device: + mode_context._dispatch_scales = None + mode_context._dispatch_topk_ids = None + mode_context._dispatch_weights = None + if mode_context.output_layout == DispatchLayout.EXPERT_MAJOR: + mode_context._dispatch_src_info = torch.empty( + (mode_context.num_local_experts, slots_per_expert), + dtype=torch.int32, + device=device, + ) + mode_context._dispatch_layout_range = torch.empty( + (mode_context.num_local_experts, mode_context.world_size), + dtype=torch.int64, + device=device, + ) + mode_context._dispatch_count = torch.empty( + (mode_context.num_local_experts,), dtype=torch.int32, device=device + ) + scale_block_size = dispatch_scale_block_size(mode_context.dispatch_data_type) + if scale_block_size: + num_scales = mode_context.hidden_size // scale_block_size + scale_storage = torch.empty( + (mode_context.num_local_experts, num_scales, slots_per_expert), + dtype=dispatch_scale_dtype(mode_context.dispatch_data_type), + device=device, + ) + mode_context._dispatch_scales = scale_storage.transpose(1, 2) + elif mode_context.output_layout == DispatchLayout.RANK_MAJOR: + mode_context._dispatch_src_info = None + assert mode_context._output_topk_ids is not None + assert mode_context._output_topk_weights is not None + mode_context._dispatch_topk_ids = mode_context._output_topk_ids + mode_context._dispatch_weights = mode_context._output_topk_weights + mode_context._dispatch_layout_range = None + mode_context._dispatch_count = torch.empty((mode_context.world_size,), dtype=torch.int32, device=device) + else: + raise ValueError(f"unsupported latency output layout: {mode_context.output_layout}") + assert mode_context._dispatch_count is not None + return ( + output_buffer, + mode_context._dispatch_scales, + mode_context._dispatch_src_info, + mode_context._dispatch_topk_ids, + mode_context._dispatch_weights, + mode_context._dispatch_layout_range, + mode_context._dispatch_count, + ) + + def _validate_dispatch(self, input, topk_ids, weights, quant, output_buffer, active_capacity: int) -> None: + mode_context = self.context + if quant is not None: + raise NotImplementedError( + "per-call input quant metadata is not supported; configure dispatch output quantization on the communicator" + ) + if input.dim() != 2 or not input.is_contiguous(): + raise ValueError("input must be a contiguous [num_tokens, hidden_size] tensor") + if input.device.type != "cuda" or input.dtype != torch.bfloat16: + raise ValueError("latency dispatch input must be a CUDA BF16 tensor") + if input.size(1) != mode_context.hidden_size: + raise ValueError(f"input hidden size {input.size(1)} does not match configured {mode_context.hidden_size}") + if input.size(0) > active_capacity: + raise ValueError("input token count exceeds runtime_max_tokens_per_rank") + if topk_ids.dim() != 2 or not topk_ids.is_contiguous(): + raise ValueError("topk_ids must be a contiguous [num_tokens, topk] tensor") + if topk_ids.device != input.device or topk_ids.dtype != torch.int64: + raise ValueError("topk_ids must be an int64 CUDA tensor on the same device as input") + if topk_ids.shape != (input.size(0), mode_context.topk): + raise ValueError("topk_ids shape must match [input.size(0), configured topk]") + if weights is not None: + if weights.dim() != 2 or not weights.is_contiguous(): + raise ValueError("weights must be a contiguous [num_tokens, topk] tensor") + if weights.device != input.device or weights.dtype != torch.float32: + raise ValueError("weights must be a float32 CUDA tensor on the same device as input") + if weights.shape != topk_ids.shape: + raise ValueError("weights shape must match topk_ids") + slots_per_expert = mode_context.world_size * mode_context.max_tokens_per_rank + if mode_context.output_layout == DispatchLayout.EXPERT_MAJOR: + expected_shape = ( + mode_context.num_local_experts, + slots_per_expert, + mode_context.hidden_size, + ) + elif mode_context.output_layout == DispatchLayout.RANK_MAJOR: + expected_shape = ( + mode_context.world_size * mode_context.max_tokens_per_rank, + mode_context.hidden_size, + ) + else: + raise ValueError(f"unsupported latency output layout: {mode_context.output_layout}") + if mode_context.output_layout == DispatchLayout.RANK_MAJOR: + assert mode_context.dispatch_output_buffer is not None + if output_buffer.data_ptr() != mode_context.dispatch_output_buffer.data_ptr(): + raise ValueError("RANK_MAJOR output uses the runtime-owned dispatch output buffer") + return + if output_buffer.dim() != len(expected_shape) or not output_buffer.is_contiguous(): + raise ValueError(f"output_buffer must be a contiguous {mode_context.output_layout} tensor") + expected_dtype = ( + torch.bfloat16 if mode_context.dispatch_data_type == DispatchDataType.BF16 else torch.float8_e4m3fn + ) + if output_buffer.device != input.device or output_buffer.dtype != expected_dtype: + raise ValueError(f"output_buffer must be a {expected_dtype} CUDA tensor on the same device as input") + if tuple(output_buffer.shape) != expected_shape: + raise ValueError(f"output_buffer shape must be {expected_shape}") + + def _validate_combine(self, expert_output, handle, out) -> None: + mode_context = self.context + if not isinstance(handle, DispatchHandle) or not isinstance( + handle._context, (_ExpertMajorCombineContext, _RankMajorCombineContext) + ): + raise ValueError("DispatchHandle does not contain latency combine context") + context = handle._context + if context.num_experts != mode_context.num_experts or context.hidden_size != mode_context.hidden_size: + raise ValueError("DispatchHandle does not belong to this MoECommunicator configuration") + if handle.output_info.layout.kind != mode_context.output_layout: + raise ValueError("DispatchHandle output layout does not match this MoECommunicator") + output_quant = handle.output_info.quant + handle_data_type = DispatchDataType.BF16 if output_quant is None else output_quant.format + if handle_data_type != mode_context.dispatch_data_type: + raise ValueError("DispatchHandle quantization does not match this MoECommunicator configuration") + active_capacity = ( + context.max_tokens_per_rank + if isinstance(context, _RankMajorCombineContext) + else mode_context.max_tokens_per_rank + ) + slots_per_expert = mode_context.world_size * active_capacity + if handle.output_info.layout.kind == DispatchLayout.EXPERT_MAJOR: + expected_shape = ( + mode_context.num_local_experts, + slots_per_expert, + mode_context.hidden_size, + ) + elif handle.output_info.layout.kind == DispatchLayout.RANK_MAJOR: + expected_shape = ( + mode_context.world_size * active_capacity, + mode_context.hidden_size, + ) + else: + raise ValueError(f"unsupported latency output layout: {handle.output_info.layout.kind}") + if expert_output.dim() != len(expected_shape) or not expert_output.is_contiguous(): + raise ValueError("expert_output must keep dispatch output's contiguous layout") + if tuple(expert_output.shape) != expected_shape: + raise ValueError(f"expert_output shape must be {expected_shape}") + if expert_output.dtype != torch.bfloat16: + raise ValueError("expert_output must be BF16") + if handle.output_info.layout.kind == DispatchLayout.RANK_MAJOR: + assert mode_context.combine_input_buffer is not None + if expert_output.data_ptr() != mode_context.combine_input_buffer.data_ptr(): + raise ValueError("RANK_MAJOR combine requires the runtime-owned combine input buffer") + if out is not None: + expected_out_shape = (context.num_tokens, mode_context.hidden_size) + if tuple(out.shape) != expected_out_shape or out.dtype != torch.bfloat16 or not out.is_contiguous(): + raise ValueError(f"out must be a contiguous BF16 tensor with shape {expected_out_shape}") diff --git a/python/mscclpp/ep/low_latency.py b/python/mscclpp/ep/low_latency.py deleted file mode 100644 index 40bccd016..000000000 --- a/python/mscclpp/ep/low_latency.py +++ /dev/null @@ -1,585 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. -"""Low-latency backend for the high-level MoE communicator.""" - -from __future__ import annotations - -from typing import Any, Optional - -import torch - -from ._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode, create_moe_runtime -from .types import ( - DispatchHandle, - DispatchLayoutInfo, - DispatchOutput, - DispatchOutputInfo, - ExpertMajorDispatchHandle, - ExpertMajorCombineContext, - MoECommunicatorConfig, - QuantConfig, - RankMajorCombineContext, - RankMajorDispatchHandle, -) -from .utils import cuda_stream_ptr, resolve_expert_placement - - -def _resolve_dispatch_data_type(quant: Optional[QuantConfig]) -> DispatchDataType: - if quant is None: - return DispatchDataType.BF16 - - quant_format = quant.format - if quant_format is not None and not isinstance(quant_format, DispatchDataType): - raise TypeError("quant.format must be a DispatchDataType") - if quant_format is None: - raise ValueError("quant.format is required") - if quant_format != DispatchDataType.FP8_E4M3: - raise ValueError("unsupported low-latency quantization format") - if quant.block_scales is not None: - raise ValueError("communicator quant config must not contain precomputed scales") - return quant_format - - -def _dispatch_scale_block_size(data_type: DispatchDataType) -> int: - if data_type == DispatchDataType.FP8_E4M3: - return 128 - return 0 - - -def _dispatch_scale_dtype(data_type: DispatchDataType) -> torch.dtype: - if data_type == DispatchDataType.FP8_E4M3: - return torch.float32 - raise ValueError("BF16 dispatch does not have block scales") - - -class _CudaBufferView: - """Zero-copy view over a runtime-owned CUDA buffer. - - The pointer belongs to the runtime's registered symmetric buffer, not to - PyTorch. ``owner`` is retained only to keep the runtime alive for as long as - the view exists. - - Every dispatch/combine writes into the same underlying allocation, so - tensors built from this view alias runtime state and are only valid until - the next call. Callers that need results to survive across iterations must - copy them out (e.g. ``tensor.clone()``). - """ - - def __init__(self, pointer: int, shape: tuple[int, ...], typestr: str, owner: Any) -> None: - self.pointer = pointer - self.shape = shape - self.typestr = typestr - self.owner = owner - - @property - def __cuda_array_interface__(self): - return { - "shape": self.shape, - "strides": None, - "typestr": self.typestr, - "data": (self.pointer, False), - "version": 3, - } - - -def _bf16_tensor_from_pointer( - pointer: int, - shape: tuple[int, ...], - device: torch.device, - owner: Any, -) -> tuple[_CudaBufferView, torch.Tensor]: - buffer_view = _CudaBufferView(pointer, shape, " tuple[_CudaBufferView, torch.Tensor]: - buffer_view = _CudaBufferView(pointer, shape, typestr, owner) - tensor = torch.as_tensor(buffer_view, device=device) - tensor._mscclpp_owner = owner - return buffer_view, tensor - - -class LowLatencyRuntime: - """Private low-level low-latency runtime wrapper (wraps ``_cpp.MoERuntime``).""" - - num_sms: int = 128 - - def __init__( - self, - comm: Any, - max_tokens_per_rank: int, - hidden: int, - num_experts: int, - num_topk: int, - output_layout: DispatchLayout, - ) -> None: - self.rank: int = comm.my_rank - self.group_size: int = comm.nranks - self.comm = comm - self.cpp_runtime = create_moe_runtime( - comm.communicator, - MoEMode.LOW_LATENCY, - max_tokens_per_rank=max_tokens_per_rank, - hidden=hidden, - num_experts=num_experts, - num_topk=num_topk, - output_layout=output_layout, - ) - - def is_available(self) -> bool: - return self.cpp_runtime.is_available() - - def is_internode_available(self) -> bool: - return self.cpp_runtime.is_internode_available() - - -class LowLatencyBackend: - """Backend implementation for ``MoEMode.LOW_LATENCY``.""" - - def __init__(self, config: MoECommunicatorConfig, output_layout: DispatchLayout) -> None: - comm = config.comm - if comm is None: - raise ValueError("mode=LOW_LATENCY requires an mscclpp.CommGroup via comm=") - - self.comm = comm - self.rank = comm.my_rank - self.world_size = comm.nranks - self.local_rank = torch.cuda.current_device() - self.device = torch.device("cuda", self.local_rank) - self.mode = MoEMode.LOW_LATENCY - self.output_layout = output_layout - - self.num_experts = config.num_experts - self.hidden_size = config.hidden_size - self.topk = config.topk - self.max_tokens_per_rank = config.max_tokens_per_rank - self.num_blocks = config.low_latency_num_blocks - self.num_sms = self.num_blocks - 2 - self.combine_mode = config.low_latency_combine_mode - self.invalid_token_expert_id = ( - self.num_experts if config.invalid_token_expert_id is None else config.invalid_token_expert_id - ) - self.enable_overlap = config.enable_overlap - - if self.output_layout not in ( - DispatchLayout.EXPERT_MAJOR, - DispatchLayout.RANK_MAJOR, - ): - raise NotImplementedError("unsupported low-latency output layout") - 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 + 2 <= self.num_blocks <= 130: - raise ValueError("low_latency_num_blocks must be between world_size + 2 and 130") - if not isinstance(self.combine_mode, CombineMode): - raise TypeError("low_latency_combine_mode must be a CombineMode") - if type(self.invalid_token_expert_id) is not int: - raise TypeError("invalid_token_expert_id must be an int or None") - if not -(1 << 31) <= self.invalid_token_expert_id < (1 << 31): - raise ValueError("invalid_token_expert_id must fit in int32") - if 0 <= self.invalid_token_expert_id < self.num_experts: - raise ValueError("invalid_token_expert_id must not overlap a valid global expert ID") - if self.output_layout == DispatchLayout.RANK_MAJOR: - if self.combine_mode != CombineMode.RANK_LOCAL_REDUCE: - raise ValueError("RANK_MAJOR output requires RANK_LOCAL_REDUCE combine") - if self.enable_overlap: - raise NotImplementedError("RANK_MAJOR output does not support overlapping calls yet") - - self.num_local_experts, self.local_expert_start = resolve_expert_placement( - num_experts=self.num_experts, - world_size=self.world_size, - rank=self.rank, - num_local_experts=config.num_local_experts, - local_expert_start=config.local_expert_start, - ) - - self.dispatch_data_type = _resolve_dispatch_data_type(config.quant) - if self.output_layout == DispatchLayout.RANK_MAJOR and self.dispatch_data_type != DispatchDataType.BF16: - raise NotImplementedError("RANK_MAJOR output currently supports BF16 dispatch only") - - self._dispatch_scales: Optional[torch.Tensor] = None - self._dispatch_src_info: Optional[torch.Tensor] = None - self._dispatch_topk_ids: Optional[torch.Tensor] = None - self._dispatch_weights: Optional[torch.Tensor] = None - self._dispatch_layout_range: Optional[torch.Tensor] = None - self._dispatch_count: Optional[torch.Tensor] = None - - self._runtime = LowLatencyRuntime( - comm, - max_tokens_per_rank=self.max_tokens_per_rank, - hidden=self.hidden_size, - num_experts=self.num_experts, - num_topk=self.topk, - output_layout=self.output_layout, - ) - self._is_internode = self._runtime.is_internode_available() - self._output_tokens_owner: Optional[_CudaBufferView] = None - self._expert_output_owner: Optional[_CudaBufferView] = None - self._output_topk_ids_owner: Optional[_CudaBufferView] = None - self._output_topk_weights_owner: Optional[_CudaBufferView] = None - self._output_tokens: Optional[torch.Tensor] = None - self._output_topk_ids: Optional[torch.Tensor] = None - self._output_topk_weights: Optional[torch.Tensor] = None - self.expert_output_buffer: Optional[torch.Tensor] = None - if self.output_layout == DispatchLayout.RANK_MAJOR: - shape = (self.world_size * self.max_tokens_per_rank, self.hidden_size) - metadata_shape = (self.world_size * self.max_tokens_per_rank, self.topk) - ( - self._output_topk_ids_owner, - self._output_topk_ids, - ) = _tensor_from_pointer( - self._runtime.cpp_runtime.output_topk_ids_buffer_ptr(), - metadata_shape, - " int: - resolved = self.max_tokens_per_rank if runtime_max_tokens_per_rank is None else runtime_max_tokens_per_rank - if type(resolved) is not int or not 0 < resolved <= self.max_tokens_per_rank: - raise ValueError("runtime_max_tokens_per_rank must be positive and not exceed max_tokens_per_rank") - if self.output_layout != DispatchLayout.RANK_MAJOR and resolved != self.max_tokens_per_rank: - raise ValueError("runtime_max_tokens_per_rank is only supported by rank-major dispatch") - return resolved - - def is_available(self) -> bool: - return self._runtime.is_available() - - def is_internode_available(self) -> bool: - return self._runtime.is_internode_available() - - def is_internode(self) -> bool: - return self._is_internode - - def dispatch( - self, - input: torch.Tensor, - topk_ids: torch.Tensor, - weights: Optional[torch.Tensor], - quant: Optional[QuantConfig], - *, - output_buffer: Optional[torch.Tensor], - stream: Optional[torch.cuda.Stream], - previous_handle: Optional[DispatchHandle], - runtime_max_tokens_per_rank: Optional[int], - ) -> tuple[DispatchOutput, DispatchHandle]: - del previous_handle - active_capacity = self._resolve_runtime_max_tokens_per_rank(runtime_max_tokens_per_rank) - self._validate_dispatch_inputs(input, topk_ids, weights, quant, output_buffer, active_capacity) - - out_buf, scales, src_info, recv_topk_ids, recv_weights, layout_range, count = self._get_dispatch_output_tensors( - output_buffer - ) - self._runtime.cpp_runtime.ll_dispatch( - input.data_ptr(), - topk_ids.data_ptr(), - 0 if weights is None else weights.data_ptr(), - out_buf.data_ptr(), - 0 if scales is None else scales.data_ptr(), - 0 if src_info is None else src_info.data_ptr(), - 0 if recv_topk_ids is None else recv_topk_ids.data_ptr(), - 0 if recv_weights is None else recv_weights.data_ptr(), - 0 if layout_range is None else layout_range.data_ptr(), - count.data_ptr(), - input.size(0), - self.hidden_size, - self.topk, - active_capacity, - self.num_experts, - self.invalid_token_expert_id, - self.output_layout, - self.dispatch_data_type, - self.num_blocks, - cuda_stream_ptr(stream), - ) - output_quant = ( - None - if scales is None - else QuantConfig( - format=self.dispatch_data_type, - block_scales=scales, - ) - ) - if self.output_layout == DispatchLayout.EXPERT_MAJOR: - layout_info = DispatchLayoutInfo(kind=self.output_layout, num_tokens_per_expert=count) - elif self.output_layout == DispatchLayout.RANK_MAJOR: - layout_info = DispatchLayoutInfo( - kind=self.output_layout, - num_tokens_per_rank=count, - ) - else: - raise ValueError(f"unsupported low-latency output layout: {self.output_layout}") - output_info = DispatchOutputInfo(layout=layout_info, quant=output_quant) - dispatch_out = DispatchOutput( - tokens=out_buf, - quant=output_info.quant, - layout=output_info.layout, - topk_ids=recv_topk_ids, - weights=recv_weights, - ) - if self.output_layout == DispatchLayout.EXPERT_MAJOR: - assert layout_range is not None - assert src_info is not None - handle: DispatchHandle = ExpertMajorDispatchHandle( - output_info=output_info, - combine_context=ExpertMajorCombineContext( - topk_ids=topk_ids, - weights=weights, - num_experts=self.num_experts, - num_tokens=input.size(0), - hidden_size=self.hidden_size, - src_info=src_info, - layout_range=layout_range, - ), - ) - elif self.output_layout == DispatchLayout.RANK_MAJOR: - handle = RankMajorDispatchHandle( - output_info=output_info, - combine_context=RankMajorCombineContext( - topk_ids=topk_ids, - num_experts=self.num_experts, - num_tokens=input.size(0), - hidden_size=self.hidden_size, - max_tokens_per_rank=active_capacity, - ), - ) - else: - raise ValueError(f"unsupported low-latency output layout: {self.output_layout}") - return dispatch_out, handle - - def combine( - self, - expert_output: torch.Tensor, - handle: DispatchHandle, - *, - out: Optional[torch.Tensor], - stream: Optional[torch.cuda.Stream], - ) -> torch.Tensor: - self._validate_combine_inputs(expert_output, handle, out) - if isinstance(handle, ExpertMajorDispatchHandle): - context = handle.combine_context - topk_weights = context.weights - src_info = context.src_info - layout_range = context.layout_range - elif isinstance(handle, RankMajorDispatchHandle): - context = handle.combine_context - active_capacity = context.max_tokens_per_rank - topk_weights = None - src_info = None - layout_range = None - else: - raise ValueError("DispatchHandle does not contain low-latency combine context") - if isinstance(handle, ExpertMajorDispatchHandle): - active_capacity = self.max_tokens_per_rank - if out is None: - out = torch.empty( - (context.num_tokens, self.hidden_size), - dtype=torch.bfloat16, - device=expert_output.device, - ) - self._runtime.cpp_runtime.ll_combine( - expert_output.data_ptr(), - context.topk_ids.data_ptr(), - 0 if topk_weights is None else topk_weights.data_ptr(), - 0 if src_info is None else src_info.data_ptr(), - 0 if layout_range is None else layout_range.data_ptr(), - out.data_ptr(), - context.num_tokens, - self.hidden_size, - self.topk, - active_capacity, - context.num_experts, - self.output_layout, - self.dispatch_data_type, - self.combine_mode, - self.num_blocks - 2, - cuda_stream_ptr(stream), - ) - return out - - def _get_dispatch_output_tensors(self, output_buffer: torch.Tensor): - device = self.device if output_buffer is None else output_buffer.device - slots_per_expert = self.world_size * self.max_tokens_per_rank - if self._dispatch_count is None or self._dispatch_count.device != device: - self._dispatch_scales = None - self._dispatch_topk_ids = None - self._dispatch_weights = None - if self.output_layout == DispatchLayout.EXPERT_MAJOR: - self._dispatch_src_info = torch.empty( - (self.num_local_experts, slots_per_expert), - dtype=torch.int32, - device=device, - ) - self._dispatch_layout_range = torch.empty( - (self.num_local_experts, self.world_size), - dtype=torch.int64, - device=device, - ) - self._dispatch_count = torch.empty((self.num_local_experts,), dtype=torch.int32, device=device) - scale_block_size = _dispatch_scale_block_size(self.dispatch_data_type) - if scale_block_size: - num_scales = self.hidden_size // scale_block_size - scale_storage = torch.empty( - (self.num_local_experts, num_scales, slots_per_expert), - dtype=_dispatch_scale_dtype(self.dispatch_data_type), - device=device, - ) - self._dispatch_scales = scale_storage.transpose(1, 2) - elif self.output_layout == DispatchLayout.RANK_MAJOR: - self._dispatch_src_info = None - assert self._output_topk_ids is not None - assert self._output_topk_weights is not None - self._dispatch_topk_ids = self._output_topk_ids - self._dispatch_weights = self._output_topk_weights - self._dispatch_layout_range = None - self._dispatch_count = torch.empty((self.world_size,), dtype=torch.int32, device=device) - else: - raise ValueError(f"unsupported low-latency output layout: {self.output_layout}") - assert self._dispatch_count is not None - if self.output_layout == DispatchLayout.RANK_MAJOR: - assert self._output_tokens is not None - output_buffer = self._output_tokens - return ( - output_buffer, - self._dispatch_scales, - self._dispatch_src_info, - self._dispatch_topk_ids, - self._dispatch_weights, - self._dispatch_layout_range, - self._dispatch_count, - ) - - def _validate_dispatch_inputs(self, input, topk_ids, weights, quant, output_buffer, active_capacity: int) -> None: - if output_buffer is None and self.output_layout != DispatchLayout.RANK_MAJOR: - raise ValueError("output_buffer is required for low-latency dispatch") - if quant is not None: - raise NotImplementedError( - "per-call input quant metadata is not supported; configure dispatch output quantization on the communicator" - ) - if input.dim() != 2 or not input.is_contiguous(): - raise ValueError("input must be a contiguous [num_tokens, hidden_size] tensor") - if input.device.type != "cuda" or input.dtype != torch.bfloat16: - raise ValueError("low-latency dispatch input must be a CUDA BF16 tensor") - if input.size(1) != self.hidden_size: - raise ValueError(f"input hidden size {input.size(1)} does not match configured {self.hidden_size}") - if input.size(0) > active_capacity: - raise ValueError("input token count exceeds runtime_max_tokens_per_rank") - if topk_ids.dim() != 2 or not topk_ids.is_contiguous(): - raise ValueError("topk_ids must be a contiguous [num_tokens, topk] tensor") - if topk_ids.device != input.device or topk_ids.dtype != torch.int64: - raise ValueError("topk_ids must be an int64 CUDA tensor on the same device as input") - if topk_ids.shape != (input.size(0), self.topk): - raise ValueError("topk_ids shape must match [input.size(0), configured topk]") - if weights is not None: - if weights.dim() != 2 or not weights.is_contiguous(): - raise ValueError("weights must be a contiguous [num_tokens, topk] tensor") - if weights.device != input.device or weights.dtype != torch.float32: - raise ValueError("weights must be a float32 CUDA tensor on the same device as input") - if weights.shape != topk_ids.shape: - raise ValueError("weights shape must match topk_ids") - slots_per_expert = self.world_size * self.max_tokens_per_rank - if self.output_layout == DispatchLayout.EXPERT_MAJOR: - expected_shape = ( - self.num_local_experts, - slots_per_expert, - self.hidden_size, - ) - elif self.output_layout == DispatchLayout.RANK_MAJOR: - expected_shape = ( - self.world_size * self.max_tokens_per_rank, - self.hidden_size, - ) - else: - raise ValueError(f"unsupported low-latency output layout: {self.output_layout}") - if self.output_layout == DispatchLayout.RANK_MAJOR: - if output_buffer is not None: - assert self._output_tokens is not None - if output_buffer.data_ptr() != self._output_tokens.data_ptr(): - raise ValueError("RANK_MAJOR output uses the runtime-owned registered token buffer") - return - if output_buffer.dim() != len(expected_shape) or not output_buffer.is_contiguous(): - raise ValueError(f"output_buffer must be a contiguous {self.output_layout} tensor") - expected_dtype = torch.bfloat16 if self.dispatch_data_type == DispatchDataType.BF16 else torch.float8_e4m3fn - if output_buffer.device != input.device or output_buffer.dtype != expected_dtype: - raise ValueError(f"output_buffer must be a {expected_dtype} CUDA tensor on the same device as input") - if tuple(output_buffer.shape) != expected_shape: - raise ValueError(f"output_buffer shape must be {expected_shape}") - - def _validate_combine_inputs(self, expert_output, handle, out) -> None: - if not isinstance(handle, (ExpertMajorDispatchHandle, RankMajorDispatchHandle)): - raise ValueError("DispatchHandle does not contain low-latency combine context") - context = handle.combine_context - if context.num_experts != self.num_experts or context.hidden_size != self.hidden_size: - raise ValueError("DispatchHandle does not belong to this MoECommunicator configuration") - if handle.output_info.layout.kind != self.output_layout: - raise ValueError("DispatchHandle output layout does not match this MoECommunicator") - output_quant = handle.output_info.quant - handle_data_type = DispatchDataType.BF16 if output_quant is None else output_quant.format - if handle_data_type != self.dispatch_data_type: - raise ValueError("DispatchHandle quantization does not match this MoECommunicator configuration") - active_capacity = ( - handle.combine_context.max_tokens_per_rank - if isinstance(handle, RankMajorDispatchHandle) - else self.max_tokens_per_rank - ) - slots_per_expert = self.world_size * active_capacity - if handle.output_info.layout.kind == DispatchLayout.EXPERT_MAJOR: - expected_shape = ( - self.num_local_experts, - slots_per_expert, - self.hidden_size, - ) - elif handle.output_info.layout.kind == DispatchLayout.RANK_MAJOR: - expected_shape = ( - self.world_size * active_capacity, - self.hidden_size, - ) - else: - raise ValueError(f"unsupported low-latency output layout: {handle.output_info.layout.kind}") - if expert_output.dim() != len(expected_shape) or not expert_output.is_contiguous(): - raise ValueError("expert_output must keep dispatch output's contiguous layout") - if tuple(expert_output.shape) != expected_shape: - raise ValueError(f"expert_output shape must be {expected_shape}") - if expert_output.dtype != torch.bfloat16: - raise ValueError("expert_output must be BF16") - if handle.output_info.layout.kind == DispatchLayout.RANK_MAJOR: - assert self.expert_output_buffer is not None - if expert_output.data_ptr() != self.expert_output_buffer.data_ptr(): - raise ValueError("RANK_MAJOR combine requires the runtime-owned registered expert output buffer") - if out is not None: - expected_out_shape = (context.num_tokens, self.hidden_size) - if tuple(out.shape) != expected_out_shape or out.dtype != torch.bfloat16 or not out.is_contiguous(): - raise ValueError(f"out must be a contiguous BF16 tensor with shape {expected_out_shape}") diff --git a/python/mscclpp/ep/runtime.py b/python/mscclpp/ep/runtime.py new file mode 100644 index 000000000..7fb259ee9 --- /dev/null +++ b/python/mscclpp/ep/runtime.py @@ -0,0 +1,64 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Expert-parallel runtime interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Optional + +import torch + +from mscclpp.ep.context import Context +from mscclpp.ep.types import DispatchHandle, DispatchOutput, QuantConfig + + +class Runtime(ABC): + """Common interface over one mode-configured C++ runtime.""" + + def __init__(self, context: Context, cpp_runtime) -> None: + self.context = context + self.cpp_runtime = cpp_runtime + + @staticmethod + def create(context: Context) -> "Runtime": + from mscclpp.ep.latency import LatencyContext, LatencyRuntime + from mscclpp.ep.throughput import ThroughputContext, ThroughputRuntime + + if isinstance(context, LatencyContext): + return LatencyRuntime(context) + if isinstance(context, ThroughputContext): + return ThroughputRuntime(context) + raise TypeError(f"unsupported EP context: {type(context).__name__}") + + def is_available(self) -> bool: + return self.cpp_runtime.is_available() + + def is_internode_available(self) -> bool: + return self.cpp_runtime.is_internode_available() + + @abstractmethod + def dispatch( + self, + input: torch.Tensor, + topk_ids: torch.Tensor, + weights: Optional[torch.Tensor], + quant: Optional[QuantConfig], + *, + output_buffer: Optional[torch.Tensor], + stream: Optional[torch.cuda.Stream], + previous_handle: Optional[DispatchHandle], + runtime_max_tokens_per_rank: Optional[int], + ) -> tuple[DispatchOutput, DispatchHandle]: + raise NotImplementedError + + @abstractmethod + def combine( + self, + expert_output: torch.Tensor, + handle: DispatchHandle, + *, + out: Optional[torch.Tensor], + stream: Optional[torch.cuda.Stream], + ) -> torch.Tensor: + raise NotImplementedError diff --git a/python/mscclpp/ep/throughput.py b/python/mscclpp/ep/throughput.py new file mode 100644 index 000000000..37e24bdad --- /dev/null +++ b/python/mscclpp/ep/throughput.py @@ -0,0 +1,483 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +# Portions adapted from DeepEP (https://github.com/deepseek-ai/DeepEP), +# branch ``chhwang/dev-atomic-add-cleanup``. Licensed under the MIT License. +"""Throughput-mode context.""" + +from __future__ import annotations + +from contextlib import nullcontext +from typing import List, Optional + +import torch + +from mscclpp.ep._cpp import DispatchLayout, MoEMode, create_moe_runtime +from mscclpp.ep.context import Context +from mscclpp.ep.runtime import Runtime +from mscclpp.ep.types import ( + DispatchHandle, + DispatchLayoutInfo, + DispatchOutput, + DispatchOutputInfo, + MoECommunicatorConfig, + QuantConfig, + _TokenMajorCombineContext, +) +from mscclpp.ep.utils import ( + current_stream_ptr as _stream_ptr, + ptr as _ptr, + resolve_expert_placement, + tensor_from_pointer, +) + + +class ThroughputContext(Context): + """Throughput-mode context.""" + + def __init__( + self, + config: MoECommunicatorConfig, + ) -> None: + comm = config.comm + if comm is None: + raise ValueError("mode=THROUGHPUT requires an mscclpp.CommGroup via comm=") + output_layout = config.output_layout + if output_layout is None: + output_layout = DispatchLayout.TOKEN_MAJOR + num_blocks = 20 if config.num_blocks is None else config.num_blocks + if num_blocks <= 0: + raise ValueError("num_blocks must be positive in throughput mode") + max_hidden_bytes = config.hidden_size * torch.empty((), dtype=torch.bfloat16).element_size() + + self.rank: int = comm.my_rank + self.group_size: int = comm.nranks + self.world_size = comm.nranks + self.comm = comm + self.local_rank = torch.cuda.current_device() + self.device = torch.device("cuda", self.local_rank) + self.mode = MoEMode.THROUGHPUT + self.output_layout = output_layout + self.num_experts = config.num_experts + self.hidden_size = config.hidden_size + self.topk = config.topk + self.max_tokens_per_rank = config.max_tokens_per_rank + self.num_blocks = num_blocks + self.max_hidden_bytes = max_hidden_bytes + self.enable_overlap = config.enable_overlap + + if self.output_layout != DispatchLayout.TOKEN_MAJOR: + raise NotImplementedError("THROUGHPUT mode currently supports only DispatchLayout.TOKEN_MAJOR") + if config.invalid_token_expert_id is not None: + raise ValueError("invalid_token_expert_id is only supported in latency mode") + self.num_local_experts, self.local_expert_start = resolve_expert_placement( + num_experts=self.num_experts, + world_size=self.world_size, + rank=self.rank, + num_local_experts=config.num_local_experts, + local_expert_start=config.local_expert_start, + ) + if config.quant is not None: + raise NotImplementedError("throughput quantized dispatch (scales) is not implemented yet") + + self.expert_alignment = config.expert_alignment + + +class ThroughputRuntime(Runtime): + """Throughput-optimized runtime using bounded communication resources.""" + + context: ThroughputContext + + def __init__(self, context: ThroughputContext) -> None: + cpp_runtime = create_moe_runtime( + context.comm.communicator, + context.mode, + max_hidden_bytes=context.max_hidden_bytes, + num_blocks=context.num_blocks, + output_layout=context.output_layout, + ) + super().__init__(context, cpp_runtime) + + def dispatch( + self, + input: torch.Tensor, + topk_ids: torch.Tensor, + weights: Optional[torch.Tensor], + quant: Optional[QuantConfig], + *, + output_buffer: Optional[torch.Tensor], + stream: Optional[torch.cuda.Stream], + previous_handle: Optional[DispatchHandle], + runtime_max_tokens_per_rank: Optional[int], + ) -> tuple[DispatchOutput, DispatchHandle]: + mode_context = self.context + del output_buffer + if runtime_max_tokens_per_rank is not None: + raise ValueError("runtime_max_tokens_per_rank is only supported by latency rank-major dispatch") + stream_scope = torch.cuda.stream(stream) if stream is not None else nullcontext() + with stream_scope: + self._validate_dispatch(input, topk_ids, weights, quant) + implicit_weights = weights is None + if weights is None: + weights = torch.ones(topk_ids.shape, dtype=torch.float32, device=topk_ids.device) + cache = previous_handle._dispatch_cache if previous_handle is not None else None + if cache is not None and not self._cache_matches(cache, input, topk_ids, weights, implicit_weights): + cache = None + if cache is not None: + num_tokens_per_rank = cache["num_tokens_per_rank"] + num_tokens_per_expert = cache["num_tokens_per_expert"] + is_token_in_rank = cache["is_token_in_rank"] + else: + ( + num_tokens_per_rank, + num_tokens_per_expert, + is_token_in_rank, + ) = self._compute_counts(topk_ids, mode_context.num_experts) + if cache is not None: + ( + recv_x, + _, + _runtime_recv_topk_idx, + _runtime_recv_topk_weights, + _runtime_num_recv_tokens_per_expert_list, + rank_prefix_matrix, + _, + send_head, + ) = self._dispatch_token_major( + input, + None, + None, + None, + None, + is_token_in_rank, + None, + cache["num_recv_tokens"], + cache["rank_prefix_matrix"], + cache["channel_prefix_matrix"], + mode_context.expert_alignment, + ) + del ( + _runtime_recv_topk_idx, + _runtime_recv_topk_weights, + _runtime_num_recv_tokens_per_expert_list, + ) + recv_topk_idx = cache["recv_topk_idx"] + recv_topk_weights = cache["recv_topk_weights"] + num_recv_tokens_per_expert_list = cache["num_recv_tokens_per_expert_list"] + combine_context = _TokenMajorCombineContext( + recv_topk_weights=recv_topk_weights, + send_head=send_head, + ) + dispatch_cache = cache + else: + ( + recv_x, + _, + recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert_list, + rank_prefix_matrix, + channel_prefix_matrix, + send_head, + ) = self._dispatch_token_major( + input, + None, + topk_ids, + weights, + num_tokens_per_rank, + is_token_in_rank, + num_tokens_per_expert, + 0, + None, + None, + mode_context.expert_alignment, + ) + combine_context = _TokenMajorCombineContext( + recv_topk_weights=recv_topk_weights, + send_head=send_head, + ) + dispatch_cache = { + "num_tokens_per_rank": num_tokens_per_rank, + "num_tokens_per_expert": num_tokens_per_expert, + "is_token_in_rank": is_token_in_rank, + "rank_prefix_matrix": rank_prefix_matrix, + "channel_prefix_matrix": channel_prefix_matrix, + "num_recv_tokens": int(recv_x.size(0)), + "recv_topk_idx": recv_topk_idx, + "recv_topk_weights": recv_topk_weights, + "num_recv_tokens_per_expert_list": num_recv_tokens_per_expert_list, + "context_id": id(mode_context), + "num_tokens": int(input.size(0)), + "device": input.device, + "topk_ids_ptr": topk_ids.data_ptr(), + "topk_ids_version": topk_ids._version, + "implicit_weights": implicit_weights, + "weights_ptr": 0 if implicit_weights else weights.data_ptr(), + "weights_version": 0 if implicit_weights else weights._version, + } + output_info = DispatchOutputInfo( + layout=DispatchLayoutInfo( + kind=mode_context.output_layout, + num_tokens_per_expert=num_recv_tokens_per_expert_list, + ), + quant=None, + ) + dispatch_out = DispatchOutput( + tokens=recv_x, + quant=output_info.quant, + layout=output_info.layout, + topk_ids=recv_topk_idx, + weights=recv_topk_weights, + ) + handle = DispatchHandle(output_info=output_info, _context=combine_context, _dispatch_cache=dispatch_cache) + return dispatch_out, handle + + def combine( + self, + expert_output: torch.Tensor, + handle: DispatchHandle, + *, + out: Optional[torch.Tensor], + stream: Optional[torch.cuda.Stream], + ) -> torch.Tensor: + mode_context = self.context + stream_scope = torch.cuda.stream(stream) if stream is not None else nullcontext() + with stream_scope: + self._validate_combine(expert_output, handle) + context = handle._context + topk_weights = context.recv_topk_weights + send_head = context.send_head + num_input_tokens, hidden = int(expert_output.size(0)), int(expert_output.size(1)) + num_output_tokens = int(send_head.size(0)) + num_topk = int(topk_weights.size(1)) if topk_weights is not None else 0 + combined_x = torch.empty((num_output_tokens, hidden), dtype=torch.bfloat16, device="cuda") + combined_topk_weights = ( + torch.empty((num_output_tokens, num_topk), dtype=torch.float32, device="cuda") + if topk_weights is not None + else None + ) + self.cpp_runtime.combine( + _ptr(combined_x), + _ptr(combined_topk_weights), + _ptr(expert_output), + _ptr(topk_weights), + _ptr(send_head), + num_input_tokens, + num_output_tokens, + hidden, + num_topk, + expert_output.element_size(), + _stream_ptr(), + ) + if out is not None: + out.copy_(combined_x) + return out + return combined_x + + def _compute_counts(self, topk_idx: torch.Tensor, num_experts: int): + """Return per-rank, per-expert, and token-membership routing metadata. + + This is routing metadata consumed by dispatch; it is unrelated to + ``DispatchLayout`` (the memory layout of the dispatch output). + """ + mode_context = self.context + assert topk_idx.dim() == 2 and topk_idx.is_contiguous() + num_tokens, num_topk = int(topk_idx.size(0)), int(topk_idx.size(1)) + + num_tokens_per_rank = torch.empty((mode_context.group_size,), dtype=torch.int32, device="cuda") + num_tokens_per_expert = torch.empty((num_experts,), dtype=torch.int32, device="cuda") + is_token_in_rank = torch.empty((num_tokens, mode_context.group_size), dtype=torch.bool, device="cuda") + + self.cpp_runtime.token_major_prepare( + _ptr(num_tokens_per_rank), + _ptr(num_tokens_per_expert), + _ptr(is_token_in_rank), + _ptr(topk_idx), + num_tokens, + num_topk, + num_experts, + _stream_ptr(), + ) + return num_tokens_per_rank, num_tokens_per_expert, is_token_in_rank + + def _dispatch_token_major( + self, + x: torch.Tensor, + x_scales: Optional[torch.Tensor], + topk_idx: Optional[torch.Tensor], + topk_weights: Optional[torch.Tensor], + num_tokens_per_rank: Optional[torch.Tensor], + is_token_in_rank: torch.Tensor, + num_tokens_per_expert: Optional[torch.Tensor], + cached_num_recv_tokens: int, + cached_rank_prefix_matrix: Optional[torch.Tensor], + cached_channel_prefix_matrix: Optional[torch.Tensor], + expert_alignment: int, + ): + """Run token-major throughput dispatch and return combine metadata.""" + mode_context = self.context + assert x.dim() == 2 and x.is_contiguous() + cached_mode = cached_rank_prefix_matrix is not None + num_tokens, hidden = int(x.size(0)), int(x.size(1)) + x_element_size = x.element_size() + num_channels = self.cpp_runtime.token_major_num_channels(x_element_size) + + num_topk = int(topk_idx.size(1)) if topk_idx is not None else 0 + num_scales = 0 + if x_scales is not None: + num_scales = 1 if x_scales.dim() == 1 else int(x_scales.size(1)) + + # ----- Phase A: notify (non-cached) or reuse cached layout ----- + if cached_mode: + num_recv_tokens = cached_num_recv_tokens + rank_prefix_matrix = cached_rank_prefix_matrix + channel_prefix_matrix = cached_channel_prefix_matrix + num_recv_tokens_per_expert_list: List[int] = [] + num_experts = 0 + else: + assert num_tokens_per_rank is not None and num_tokens_per_expert is not None + num_experts = int(num_tokens_per_expert.size(0)) + num_local_experts = num_experts // mode_context.group_size + rank_prefix_matrix = torch.empty( + (mode_context.group_size, mode_context.group_size), dtype=torch.int32, device="cuda" + ) + channel_prefix_matrix = torch.empty( + (mode_context.group_size, num_channels), dtype=torch.int32, device="cuda" + ) + num_recv_per_expert_host = torch.empty((num_local_experts,), dtype=torch.int32, device="cpu") + num_recv_tokens = self.cpp_runtime.token_major_notify( + _ptr(rank_prefix_matrix), + _ptr(channel_prefix_matrix), + _ptr(num_recv_per_expert_host), + _ptr(num_tokens_per_rank), + _ptr(num_tokens_per_expert), + _ptr(is_token_in_rank), + num_tokens, + num_experts, + x_element_size, + expert_alignment, + _stream_ptr(), + ) + num_recv_tokens_per_expert_list = num_recv_per_expert_host.tolist() + + # ----- Phase B: allocate recv outputs (or view the recv pool) ----- + recv_x = self._allocate_recv(num_tokens, num_recv_tokens, hidden, x_element_size) + send_head = torch.empty((num_tokens, mode_context.group_size), dtype=torch.int32, device="cuda") + recv_topk_idx = ( + torch.empty((num_recv_tokens, num_topk), dtype=torch.int64, device="cuda") if topk_idx is not None else None + ) + recv_topk_weights = ( + torch.empty((num_recv_tokens, num_topk), dtype=torch.float32, device="cuda") + if topk_weights is not None + else None + ) + recv_x_scales = ( + torch.empty((num_recv_tokens, num_scales), dtype=torch.float32, device="cuda") + if x_scales is not None + else None + ) + + self.cpp_runtime.dispatch( + _ptr(recv_x), + _ptr(recv_x_scales), + _ptr(recv_topk_idx), + _ptr(recv_topk_weights), + _ptr(send_head), + _ptr(x), + _ptr(x_scales), + _ptr(topk_idx), + _ptr(topk_weights), + _ptr(is_token_in_rank), + _ptr(rank_prefix_matrix), + _ptr(channel_prefix_matrix), + num_tokens, + hidden, + num_topk, + num_scales, + num_experts, + x_element_size, + num_recv_tokens, + cached_mode, + _stream_ptr(), + ) + return ( + recv_x, + recv_x_scales, + recv_topk_idx, + recv_topk_weights, + num_recv_tokens_per_expert_list, + rank_prefix_matrix, + channel_prefix_matrix, + send_head, + ) + + def _allocate_recv( + self, + num_tokens: int, + num_recv_tokens: int, + hidden: int, + x_element_size: int, + ) -> torch.Tensor: + """Return this rank's direct receive-pool view.""" + mode_context = self.context + pool_ptr = self.cpp_runtime.token_major_resolve_recv_buffer(num_tokens, num_recv_tokens, hidden, x_element_size) + if pool_ptr == 0: + raise RuntimeError("token-major throughput receive-pool capacity exceeded") + _, recv_x = tensor_from_pointer( + pool_ptr, + (num_recv_tokens, hidden), + torch.bfloat16, + mode_context.device, + self.cpp_runtime, + ) + return recv_x + + def _cache_matches(self, cache, input, topk_ids, weights, implicit_weights) -> bool: + mode_context = self.context + return ( + cache.get("context_id") == id(mode_context) + and cache.get("num_tokens") == int(input.size(0)) + and cache.get("device") == input.device + and cache.get("topk_ids_ptr") == topk_ids.data_ptr() + and cache.get("topk_ids_version") == topk_ids._version + and cache.get("implicit_weights") == implicit_weights + and (implicit_weights or cache.get("weights_ptr") == weights.data_ptr()) + and (implicit_weights or cache.get("weights_version") == weights._version) + ) + + def _validate_dispatch(self, input, topk_ids, weights, quant) -> None: + mode_context = self.context + if quant is not None: + raise NotImplementedError("throughput dispatch does not support quantized input scales yet") + if input.dim() != 2 or not input.is_contiguous(): + raise ValueError("input must be a contiguous [num_tokens, hidden] tensor") + if input.device.type != "cuda" or input.dtype != torch.bfloat16: + raise ValueError("throughput dispatch input must be a CUDA BF16 tensor") + if input.size(1) != mode_context.hidden_size: + raise ValueError(f"input hidden size {input.size(1)} != configured {mode_context.hidden_size}") + if input.size(0) > mode_context.max_tokens_per_rank: + raise ValueError("input token count exceeds max_tokens_per_rank") + if topk_ids.dim() != 2 or not topk_ids.is_contiguous(): + raise ValueError("topk_ids must be a contiguous [num_tokens, topk] tensor") + if topk_ids.device != input.device or topk_ids.dtype != torch.int64: + raise ValueError("topk_ids must be an int64 CUDA tensor on the same device as input") + if topk_ids.shape != (input.size(0), mode_context.topk): + raise ValueError("topk_ids shape must be [input.size(0), topk]") + if weights is not None: + if weights.dim() != 2 or not weights.is_contiguous(): + raise ValueError("weights must be a contiguous [num_tokens, topk] tensor") + if weights.device != input.device or weights.dtype != torch.float32: + raise ValueError("weights must be a float32 CUDA tensor on the same device as input") + if weights.shape != topk_ids.shape: + raise ValueError("weights shape must match topk_ids") + + def _validate_combine(self, expert_output, handle) -> None: + mode_context = self.context + if not isinstance(handle, DispatchHandle) or not isinstance(handle._context, _TokenMajorCombineContext): + raise TypeError("handle must be a DispatchHandle returned by dispatch") + if expert_output.dim() != 2 or not expert_output.is_contiguous(): + raise ValueError("expert_output must be a contiguous [total_recv_tokens, hidden] tensor") + if expert_output.size(1) != mode_context.hidden_size: + raise ValueError( + f"expert_output hidden size {expert_output.size(1)} != configured {mode_context.hidden_size}" + ) diff --git a/python/mscclpp/ep/types.py b/python/mscclpp/ep/types.py index b0f152c63..67660c9ec 100644 --- a/python/mscclpp/ep/types.py +++ b/python/mscclpp/ep/types.py @@ -10,7 +10,7 @@ import torch import mscclpp -from ._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode +from mscclpp.ep._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode # Quantization metadata. @@ -19,7 +19,7 @@ class QuantConfig: """Quantization metadata associated with an activation tensor. - Low-latency FP8 dispatch returns ``block_scales`` with the activation's + Latency FP8 dispatch returns ``block_scales`` with the activation's leading dimensions and a format-defined final scale dimension. ``FP8_E4M3`` uses FP32 scales per 128 elements. """ @@ -49,21 +49,20 @@ class MoECommunicatorConfig: max_tokens_per_rank: int = 0 # Runtime mode and output layout - mode: MoEMode = MoEMode.LOW_LATENCY + mode: MoEMode = MoEMode.LATENCY output_layout: Optional[DispatchLayout] = None - # LL rank-major sentinel; None resolves to num_experts. + # Latency rank-major sentinel; None resolves to num_experts. invalid_token_expert_id: Optional[int] = None # Quantization defaults quant: Optional[QuantConfig] = None # Launch tuning - num_sms: int = 20 - low_latency_num_blocks: int = 130 - low_latency_combine_mode: CombineMode = CombineMode.RANK_LOCAL_REDUCE + num_blocks: Optional[int] = None + combine_mode: CombineMode = CombineMode.RANK_LOCAL_REDUCE enable_overlap: bool = False - # HT-only buffer/launch tuning (advanced) + # Throughput receive-pool tuning (advanced) expert_alignment: int = 1 @@ -101,13 +100,14 @@ class DispatchOutput: layout: DispatchLayoutInfo topk_ids: Optional[torch.Tensor] = None weights: Optional[torch.Tensor] = None + combine_input_buffer: Optional[torch.Tensor] = None -# Combine-side context. These objects are layout-specific and opaque to the MLP. +# Private combine-side context. @dataclass -class ExpertMajorCombineContext: +class _ExpertMajorCombineContext: """Combine context for expert-major dispatch output.""" topk_ids: torch.Tensor @@ -120,7 +120,7 @@ class ExpertMajorCombineContext: @dataclass -class RankMajorCombineContext: +class _RankMajorCombineContext: """Combine context for fixed-stride rank-major output.""" topk_ids: torch.Tensor @@ -131,17 +131,17 @@ class RankMajorCombineContext: @dataclass -class HighThroughputCombineContext: - """Combine context for high-throughput dispatch output.""" +class _TokenMajorCombineContext: + """Combine context for token-major throughput output.""" recv_topk_weights: Optional[torch.Tensor] send_head: torch.Tensor -CombineContext = Union[ - ExpertMajorCombineContext, - RankMajorCombineContext, - HighThroughputCombineContext, +_CombineContext = Union[ + _ExpertMajorCombineContext, + _RankMajorCombineContext, + _TokenMajorCombineContext, ] @@ -150,24 +150,11 @@ class HighThroughputCombineContext: @dataclass class DispatchHandle: - """Base opaque dispatch metadata consumed by :meth:`MoECommunicator.combine`.""" + """Opaque dispatch metadata consumed by :meth:`MoECommunicator.combine`.""" output_info: DispatchOutputInfo - - -@dataclass -class ExpertMajorDispatchHandle(DispatchHandle): - combine_context: ExpertMajorCombineContext - - -@dataclass -class RankMajorDispatchHandle(DispatchHandle): - combine_context: RankMajorCombineContext - - -@dataclass -class HighThroughputDispatchHandle(DispatchHandle): - combine_context: HighThroughputCombineContext + _context: _CombineContext + _dispatch_cache: Optional[dict[str, Any]] = None # Optional async/overlap configuration. @@ -195,7 +182,7 @@ class BlockOverlapConfig: @dataclass -class CommOverlapConfig: +class OverlapConfig: """Mutually exclusive operation-level or block-level overlap configuration.""" operation: Optional[OperationOverlapConfig] = None diff --git a/python/mscclpp/ep/utils.py b/python/mscclpp/ep/utils.py index 466bc5947..37483858d 100644 --- a/python/mscclpp/ep/utils.py +++ b/python/mscclpp/ep/utils.py @@ -10,6 +10,40 @@ import numpy as np import torch +from mscclpp.ep._cpp import DispatchDataType +from mscclpp.ep.types import QuantConfig + + +def resolve_dispatch_data_type(quant: Optional[QuantConfig]) -> DispatchDataType: + """Resolve dispatch storage type from optional quantization metadata.""" + if quant is None: + return DispatchDataType.BF16 + + quant_format = quant.format + if quant_format is not None and not isinstance(quant_format, DispatchDataType): + raise TypeError("quant.format must be a DispatchDataType") + if quant_format is None: + raise ValueError("quant.format is required") + if quant_format != DispatchDataType.FP8_E4M3: + raise ValueError("unsupported dispatch quantization format") + if quant.block_scales is not None: + raise ValueError("communicator quant config must not contain precomputed scales") + return quant_format + + +def dispatch_scale_block_size(data_type: DispatchDataType) -> int: + """Return the hidden-element count represented by one dispatch scale.""" + if data_type == DispatchDataType.FP8_E4M3: + return 128 + return 0 + + +def dispatch_scale_dtype(data_type: DispatchDataType) -> torch.dtype: + """Return the scale dtype for a quantized dispatch format.""" + if data_type == DispatchDataType.FP8_E4M3: + return torch.float32 + raise ValueError("BF16 dispatch does not have block scales") + def send_bytes(comm: Any, payload: bytes, peer: int, tag: int) -> None: comm.send(np.frombuffer(payload, dtype=np.uint8), peer, tag) @@ -107,10 +141,32 @@ def __init__(self, ptr: int, shape: Tuple[int, ...], typestr: str, owner: Any) - } -def bf16_view(ptr: int, num_tokens: int, hidden: int, owner: Any) -> torch.Tensor: - """View a raw device pointer as a ``[num_tokens, hidden]`` bfloat16 tensor.""" - u16 = torch.as_tensor(DevicePointerArray(ptr, (num_tokens, hidden), " Tuple[DevicePointerArray, torch.Tensor]: + """Create a zero-copy tensor view over runtime-owned CUDA memory.""" + storage_types = { + torch.bfloat16: " bool: diff --git a/src/ext/ep/CMakeLists.txt b/src/ext/ep/CMakeLists.txt index e357c6061..d0e41e371 100644 --- a/src/ext/ep/CMakeLists.txt +++ b/src/ext/ep/CMakeLists.txt @@ -1,137 +1,70 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -# -# Builds `mscclpp_ep_cpp`, a nanobind extension that exposes the EP -# (Mixture-of-Experts dispatch/combine) runtime to Python. -# -# Two backends share one module behind a single `MoERuntime` class (`MoEMode` -# selects the path), both with a torch-free, raw-pointer (uintptr_t) nanobind -# API so the module never links libtorch: -# - Low-latency (LL): ll_runtime.cc + low_latency/{dispatch,combine}.cu. -# - High-throughput (HT): ht_runtime.cc + high-throughput/*.cu, a DeepEP-style -# runtime de-torched to the same pointer boundary as the LL runtime (dynamic -# recv sizing via the multi-step layout -> notify -> allocate -> dispatch API). find_package(Python 3.10 COMPONENTS Interpreter Development.Module REQUIRED) + include(FetchContent) if(NOT TARGET nanobind-static AND NOT TARGET nanobind) - FetchContent_Declare(nanobind GIT_REPOSITORY https://github.com/wjakob/nanobind.git GIT_TAG v1.9.2) + FetchContent_Declare( + nanobind + GIT_REPOSITORY https://github.com/wjakob/nanobind.git + GIT_TAG v1.9.2 + ) FetchContent_MakeAvailable(nanobind) endif() set(_mscclpp_ep_gpu_archs "") -if(MSCCLPP_USE_CUDA) - get_property(_mscclpp_requested_gpu_archs CACHE MSCCLPP_GPU_ARCHS PROPERTY VALUE) - if(_mscclpp_requested_gpu_archs) - foreach(_arch IN LISTS MSCCLPP_GPU_ARCHS) - if(_arch STREQUAL "native") - foreach(_native_arch IN LISTS NVIDIA_GPU_ARCHS) - if(_native_arch GREATER_EQUAL 90) - list(APPEND _mscclpp_ep_gpu_archs "${_native_arch}") - endif() - endforeach() - else() - string(REGEX MATCH "^[0-9]+" _arch_num "${_arch}") - if(_arch_num AND _arch_num GREATER_EQUAL 90) - list(APPEND _mscclpp_ep_gpu_archs "${_arch}") +get_property(_mscclpp_requested_gpu_archs CACHE MSCCLPP_GPU_ARCHS PROPERTY VALUE) +if(_mscclpp_requested_gpu_archs) + foreach(_arch IN LISTS MSCCLPP_GPU_ARCHS) + if(_arch STREQUAL "native") + foreach(_native_arch IN LISTS NVIDIA_GPU_ARCHS) + if(_native_arch GREATER_EQUAL 90) + list(APPEND _mscclpp_ep_gpu_archs "${_native_arch}") endif() + endforeach() + else() + string(REGEX MATCH "^[0-9]+" _arch_num "${_arch}") + if(_arch_num AND _arch_num GREATER_EQUAL 90) + list(APPEND _mscclpp_ep_gpu_archs "${_arch}") endif() - endforeach() - else() - list(APPEND _mscclpp_ep_gpu_archs 90) - if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.8) - list(APPEND _mscclpp_ep_gpu_archs 100 100a) - endif() - if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.9) - list(APPEND _mscclpp_ep_gpu_archs 103 103a) endif() + endforeach() +else() + list(APPEND _mscclpp_ep_gpu_archs 90) + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.8) + list(APPEND _mscclpp_ep_gpu_archs 100 100a) endif() - if(NOT _mscclpp_ep_gpu_archs) - message(STATUS - "Skipping the EP extension because MSCCLPP_GPU_ARCHS contains no " - "explicit CUDA architecture 90 or newer.") - return() + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.9) + list(APPEND _mscclpp_ep_gpu_archs 103 103a) endif() - list(REMOVE_DUPLICATES _mscclpp_ep_gpu_archs) - message(STATUS "EP CUDA architectures: ${_mscclpp_ep_gpu_archs}") endif() +if(NOT _mscclpp_ep_gpu_archs) + message(STATUS + "Skipping the EP extension because MSCCLPP_GPU_ARCHS contains no " + "explicit CUDA architecture 90 or newer.") + return() +endif() +list(REMOVE_DUPLICATES _mscclpp_ep_gpu_archs) +message(STATUS "EP CUDA architectures: ${_mscclpp_ep_gpu_archs}") -set(EP_SOURCES - moe_runtime.cc - ll_runtime.cc - bindings.cpp - low_latency/dispatch.cu - low_latency/combine.cu - # High-throughput (DeepEP-style) backend (torch-free, raw-pointer API). - ht_runtime.cc - high-throughput/counts.cu - high-throughput/dispatch.cu - high-throughput/combine.cu -) +file(GLOB_RECURSE SOURCES CONFIGURE_DEPENDS *.cc *.cpp *.cu) -# Build as a Python extension module (shared object with Python ABI suffix). -nanobind_add_module(mscclpp_ep_cpp ${EP_SOURCES}) -if(MSCCLPP_USE_CUDA) - set_target_properties(mscclpp_ep_cpp PROPERTIES CUDA_ARCHITECTURES "${_mscclpp_ep_gpu_archs}") -endif() +nanobind_add_module(mscclpp_ep_cpp ${SOURCES}) target_include_directories(mscclpp_ep_cpp PRIVATE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/include - ${CMAKE_CURRENT_SOURCE_DIR}/high-throughput ${PROJECT_SOURCE_DIR}/include ${PROJECT_SOURCE_DIR}/src/core/include - ${PROJECT_SOURCE_DIR}/src/ext/include ${GPU_INCLUDE_DIRS} ) target_link_libraries(mscclpp_ep_cpp PRIVATE mscclpp ${GPU_LIBRARIES} Threads::Threads) - -# The EP CUDA kernels call constexpr __host__ helpers (std::max / std::pair) and -# use extended lambdas in device code. These flags were previously injected -# implicitly by Torch's CMake package; add them explicitly now that the module -# no longer links libtorch. -if(MSCCLPP_USE_CUDA) - target_compile_options(mscclpp_ep_cpp PRIVATE - $<$:--expt-relaxed-constexpr> - $<$:--expt-extended-lambda> - ) -endif() - -# Kernel-side debug timeout (~10s) — set via: -# -DMSCCLPP_EP_KERNEL_DEBUG_TIMEOUT=ON -option(MSCCLPP_EP_KERNEL_DEBUG_TIMEOUT - "Use a short ~10s kernel spin timeout (default is ~100s)" OFF) -if(MSCCLPP_EP_KERNEL_DEBUG_TIMEOUT) - target_compile_definitions(mscclpp_ep_cpp PRIVATE MSCCLPP_EP_KERNEL_DEBUG_TIMEOUT) -endif() +target_compile_definitions(mscclpp_ep_cpp PRIVATE MSCCLPP_USE_CUDA) set_target_properties(mscclpp_ep_cpp PROPERTIES - PREFIX "" - POSITION_INDEPENDENT_CODE ON - CXX_STANDARD 17 - CXX_STANDARD_REQUIRED ON - CXX_VISIBILITY_PRESET default + CUDA_ARCHITECTURES "${_mscclpp_ep_gpu_archs}" + INSTALL_RPATH "\$ORIGIN/lib" ) -# Install layout. -# - scikit-build / wheel build (SKBUILD set by scikit-build-core): -# module lands next to the `mscclpp` python package; libmscclpp.so is -# under `mscclpp/lib/`, so rpath = `$ORIGIN/mscclpp/lib`. -# - Plain CMake install: standard `${INSTALL_PREFIX}/lib` with rpath -# `$ORIGIN/../lib` so the .so finds the sibling mscclpp shared lib. -if(DEFINED SKBUILD OR DEFINED ENV{SKBUILD}) - set_target_properties(mscclpp_ep_cpp PROPERTIES - INSTALL_RPATH "\$ORIGIN/mscclpp/lib") - install(TARGETS mscclpp_ep_cpp LIBRARY DESTINATION ..) -else() - set_target_properties(mscclpp_ep_cpp PROPERTIES - INSTALL_RPATH "\$ORIGIN/../lib") - install(TARGETS mscclpp_ep_cpp - LIBRARY DESTINATION ${INSTALL_PREFIX}/lib) -endif() - -if(MSCCLPP_USE_CUDA) - target_compile_definitions(mscclpp_ep_cpp PRIVATE MSCCLPP_USE_CUDA) -elseif(MSCCLPP_USE_ROCM) - target_compile_definitions(mscclpp_ep_cpp PRIVATE MSCCLPP_USE_ROCM) -endif() +install(TARGETS mscclpp_ep_cpp LIBRARY DESTINATION .) diff --git a/src/ext/ep/README.md b/src/ext/ep/README.md index 930f10616..06602f9e4 100644 --- a/src/ext/ep/README.md +++ b/src/ext/ep/README.md @@ -1,43 +1,54 @@ # MSCCL++ Expert-Parallel (EP) extension The EP extension is a torch-free nanobind module for MoE dispatch and combine. -It exposes a single `MoERuntime` whose `MoEMode` selects one of two backends: +It exposes one concrete `MoERuntime` and one persistent `DeviceContext*` shared +by dispatch and combine kernels. -- **Low latency (LL)**: `MoELowLatencyRuntime` (`ll_runtime.cc` plus - `low_latency/{dispatch,combine}.cu`), reached through the `ll_*` methods. - Uses ~128 SMs and expects to own the GPU while it runs. -- **High throughput (HT)**: `MoEHighThroughputRuntime` (`ht_runtime.cc` plus the - CUDA sources under `high-throughput/`), reached through the `ht_*` methods. - Defaults to 20 SMs so dispatch/combine can overlap with expert GEMMs. +`MoEMode` selects a runtime context and algorithm family: -Both derive from the abstract `MoERuntime` (`runtime_base.hpp`), which owns the -shared rank-topology detection and availability reporting. -`createMoERuntime(...)` constructs the requested implementation and returns a -`std::shared_ptr`; calling the other mode's methods raises. +- **`LATENCY`** algorithms use broad GPU resources to minimize standalone + dispatch/combine latency. +- **`THROUGHPUT`** algorithms use a bounded SM budget so communication can run + concurrently with expert compute. + +Mode-specific contexts are allocated conditionally; selecting one family does +not allocate the other family's buffers. + +The Python call path is: + +```text +MoECommunicator -> LatencyRuntime / ThroughputRuntime -> MoERuntime + | + passive mode context +``` + +The context is a passive holder for mode-specific configuration, tensors, and +metadata. `Runtime` owns dispatch/combine and all mode-specific execution +helpers. There is no separate backend or strategy object. ## Status | Feature | Status | |---|---| -| LL dispatch/combine | Validated on Hopper and newer GPUs | -| HT dispatch/combine | Supports 2, 4, 8, or 16 ranks in one GPU IPC/NVL fabric domain | -| HT RDMA/IB fallback | Not supported | -| Python frontend | `mscclpp.ep.MoECommunicator` selects LL or HT with `MoEMode` | +| Latency dispatch/combine | Validated on Hopper and newer GPUs | +| Throughput dispatch/combine | Supports 2, 4, 8, or 16 ranks in one GPU IPC/NVL fabric domain | +| Throughput RDMA/IB fallback | Not supported | +| Python frontend | `mscclpp.ep.MoECommunicator` selects latency or throughput algorithms with `MoEMode` | | ROCm | Not supported | ## Runtime architecture -### Low latency +### Latency algorithms -LL allocates CUDA physical symmetric memory and maps peer buffers through the +The latency context allocates CUDA physical symmetric memory and maps peer buffers through the existing `mscclpp::Communicator`. Payloads use direct peer mappings; `BaseMemoryChannel` handles are used only for synchronization. -The optimized LL backend is available when all participating ranks belong to +The latency algorithms are available when all participating ranks belong to one detected GPU IPC domain. That domain may span hosts when CUDA fabric handles and the required fabric services are available. -LL dispatch supports two user-visible layouts: +Latency dispatch supports two user-visible layouts: - `EXPERT_MAJOR`: one row per `(token, local expert)`. - `RANK_MAJOR`: fixed-stride rows grouped by source rank. Tokens are written @@ -46,15 +57,15 @@ LL dispatch supports two user-visible layouts: from registered remote MoE output or push completed rank partials into source-local scratch and progressively reduce ready ranks. -LL quantized dispatch supports E4M3 payloads with FP32 scales per 128 hidden +Quantized latency dispatch supports E4M3 payloads with FP32 scales per 128 hidden elements (`FP8_E4M3`). -### High throughput +### Throughput algorithms -HT follows the same direct-mapping resource model: +The throughput context follows the same direct-mapping model: 1. Python passes the existing `mscclpp::Communicator` into - `MoERuntime` with `MoEMode::HIGH_THROUGHPUT`. + `MoERuntime` with `MoEMode::THROUGHPUT`. 2. Each rank allocates a small symmetric control/FIFO region plus a CUDA physical internal receive pool. The pool provides stable peer mappings before the data-dependent receive count is known; Python later exposes its exact-size @@ -64,11 +75,11 @@ HT follows the same direct-mapping resource model: 4. Dispatch and combine launch directly on the caller's CUDA stream. The detected GPU IPC domain may span multiple hosts, such as an NVL fabric -domain with CUDA fabric handles. HT does not create a private bootstrap, proxy +domain with CUDA fabric handles. The throughput path does not create a private bootstrap, proxy service, RDMA channel, NVLS multicast object, or private communication stream, and it has no RDMA/IB fallback outside that domain. -The HT dispatch API remains two-phase because the receive token count is data +The throughput dispatch API remains two-phase because the receive token count is data dependent: 1. The notify phase exchanges counts and produces prefix matrices. @@ -77,20 +88,20 @@ dependent: Cached dispatch reuses the previous receive count and prefix matrices. -## HT data path +## Throughput data path -HT has one direct path. Every dispatch block writes hidden rows and routing +The throughput family has one direct path. Every dispatch block writes hidden rows and routing metadata directly into each destination's final receive-pool slots. Combine stages any out-of-place expert output back into that pool, synchronizes ranks, then uses a TMA shared-memory pipeline to gather and reduce peer contributions. There is no ring algorithm or runtime fallback. Set the communication block -budget through the `num_sms` API configuration. +budget through the `num_blocks` API configuration. -The persistent HT configuration contains only: +The persistent throughput configuration contains only: | Field | Meaning | |---|---| -| `num_sms` | Maximum HT communication block budget | +| `num_blocks` | Maximum throughput communication block budget | ## Build @@ -118,32 +129,48 @@ Available CMake options: | Variable | Default | Meaning | |---|---:|---| | `MSCCLPP_BUILD_EXT_EP` | `ON` | Build the EP extension | -| `MSCCLPP_EP_KERNEL_DEBUG_TIMEOUT` | `OFF` | Use a shorter kernel spin timeout | ## Source layout ```text src/ext/ep/ ├── bindings.cpp -├── moe_runtime.{cc,hpp} -├── ll_runtime.{cc,hpp} -├── ht_runtime.{cc,hpp} -├── runtime_base.hpp -├── high-throughput/ -│ ├── config.cuh -│ ├── counts.cu -│ ├── dispatch.cu -│ └── combine.cu +├── moe_runtime.cc +├── latency.cc +├── throughput.cc +├── common/ +│ ├── device_helpers.cuh +│ ├── latency.cuh +│ ├── overlap_barrier.cuh +│ └── quantization.cuh +├── dispatch/ +│ ├── common.cuh +│ ├── expert_major_dispatch.cu +│ ├── rank_major_dispatch.cu +│ ├── token_major_prepare.cu +│ └── token_major_dispatch.cu +├── combine/ +│ ├── common.cuh +│ ├── rank_local_reduce_combine.cu +│ ├── direct_send_combine.cu +│ └── token_major_reduce_combine.cu ├── include/ -└── low_latency/ - ├── config.cuh - ├── dispatch.cu - └── combine.cu +│ ├── config.hpp +│ ├── device_context.hpp +│ ├── exception.hpp +│ ├── kernels.hpp +│ ├── launch.hpp +│ ├── moe_runtime_context.hpp +│ └── recv_pool.hpp + +include/mscclpp/ext/ep/ +├── moe_runtime.hpp +└── types.hpp ``` ## Validation -Build the extension, then run the single-node HT test: +Build the extension, then run the single-node throughput test: ```bash HWLOC_COMPONENTS=-gl \ @@ -152,11 +179,11 @@ torchrun --standalone --nproc_per_node=8 \ test/python/ep/test_intranode_multirank.py ``` -The LL validation remains: +The latency validation remains: ```bash HWLOC_COMPONENTS=-gl \ LD_LIBRARY_PATH=/usr/local/cuda/lib64 \ torchrun --standalone --nproc_per_node=8 \ - test/python/ep/test_low_latency_multirank.py + test/python/ep/test_latency_multirank.py ``` diff --git a/src/ext/ep/bindings.cpp b/src/ext/ep/bindings.cpp index 77b555c85..65b77a8d0 100644 --- a/src/ext/ep/bindings.cpp +++ b/src/ext/ep/bindings.cpp @@ -9,25 +9,14 @@ // One `MoERuntime` class is exposed, with a torch-free, raw-pointer (uintptr_t) // boundary so the module never links libtorch. `MoEMode` selects the backend at // construction: -// - MoEMode.LOW_LATENCY -> `ll_*` methods (dispatch/combine). -// - MoEMode.HIGH_THROUGHPUT -> `ht_*` methods. Dynamic recv sizing uses an -// explicit multi-step API (ht_compute_dispatch_counts -> ht_notify_dispatch -// -> caller allocates -> ht_dispatch). -// The two backends keep separate method prefixes because their call protocols -// genuinely differ; calling the other mode's methods raises. +// - MoEMode.LATENCY selects latency-oriented dispatch/combine. +// - MoEMode.THROUGHPUT selects bounded-resource receive-pool dispatch/combine. #include #include #include -#include - -#include "api.cuh" -#include "config.hpp" -#include "high-throughput/config.cuh" -#include "ht_runtime.hpp" -#include "ll_runtime.hpp" -#include "moe_runtime.hpp" +#include namespace nb = nanobind; @@ -37,24 +26,6 @@ void* ptr(uintptr_t address) { return reinterpret_cast(address); } cudaStream_t stream(uintptr_t address) { return reinterpret_cast(address); } -template -Runtime& narrow(mscclpp::ep::MoERuntime& runtime, const char* expectedMode) { - auto* concrete = dynamic_cast(&runtime); - if (concrete == nullptr) { - throw std::runtime_error(std::string("MoE runtime was not created with MoEMode::") + expectedMode); - } - return *concrete; -} - -template -const Runtime& narrow(const mscclpp::ep::MoERuntime& runtime, const char* expectedMode) { - auto* concrete = dynamic_cast(&runtime); - if (concrete == nullptr) { - throw std::runtime_error(std::string("MoE runtime was not created with MoEMode::") + expectedMode); - } - return *concrete; -} - } // namespace NB_MODULE(mscclpp_ep_cpp, m) { @@ -63,28 +34,24 @@ NB_MODULE(mscclpp_ep_cpp, m) { nb::module_::import_("mscclpp._mscclpp"); nb::enum_(m, "MoEMode") - .value("LOW_LATENCY", mscclpp::ep::MoEMode::LOW_LATENCY) - .value("HIGH_THROUGHPUT", mscclpp::ep::MoEMode::HIGH_THROUGHPUT); + .value("LATENCY", mscclpp::ep::MoEMode::LATENCY) + .value("THROUGHPUT", mscclpp::ep::MoEMode::THROUGHPUT); nb::enum_(m, "DispatchLayout") .value("EXPERT_MAJOR", mscclpp::ep::DispatchLayout::EXPERT_MAJOR) .value("TOKEN_MAJOR", mscclpp::ep::DispatchLayout::TOKEN_MAJOR) .value("RANK_MAJOR", mscclpp::ep::DispatchLayout::RANK_MAJOR); - nb::enum_(m, "CombineMode") - .value("RANK_LOCAL_REDUCE", mscclpp::ep::low_latency::CombineMode::RANK_LOCAL_REDUCE) - .value("DIRECT_SEND", mscclpp::ep::low_latency::CombineMode::DIRECT_SEND); - nb::enum_(m, "DispatchDataType") - .value("BF16", mscclpp::ep::low_latency::DispatchDataType::BF16) - .value("FP8_E4M3", mscclpp::ep::low_latency::DispatchDataType::FP8_E4M3); - - nb::class_(m, "Config") - .def(nb::init(), nb::arg("num_sms") = 20) - .def_ro("num_sms", &mscclpp::ep::high_throughput::Config::numSms_); + nb::enum_(m, "CombineMode") + .value("RANK_LOCAL_REDUCE", mscclpp::ep::CombineMode::RANK_LOCAL_REDUCE) + .value("DIRECT_SEND", mscclpp::ep::CombineMode::DIRECT_SEND); + nb::enum_(m, "DispatchDataType") + .value("BF16", mscclpp::ep::DispatchDataType::BF16) + .value("FP8_E4M3", mscclpp::ep::DispatchDataType::FP8_E4M3); m.def("create_moe_runtime", &mscclpp::ep::createMoERuntime, nb::arg("comm"), nb::arg("mode"), nb::arg("max_tokens_per_rank") = 0, nb::arg("hidden") = 0, nb::arg("num_experts") = 0, nb::arg("num_topk") = 0, - nb::arg("max_hidden_bytes") = 0, nb::arg("num_sms") = 20, + nb::arg("max_hidden_bytes") = 0, nb::arg("num_blocks") = 20, nb::arg("output_layout") = mscclpp::ep::DispatchLayout::EXPERT_MAJOR, "Create the MoE backend selected by mode; returns a shared MoERuntime handle."); @@ -93,41 +60,45 @@ NB_MODULE(mscclpp_ep_cpp, m) { .def("is_available", &mscclpp::ep::MoERuntime::isAvailable) .def("is_internode_available", &mscclpp::ep::MoERuntime::isInternodeAvailable) .def("output_topk_ids_buffer_ptr", - [](const mscclpp::ep::MoERuntime& self) { - return reinterpret_cast( - narrow(self, "LOW_LATENCY").outputTopkIdsBuffer()); - }) + [](const mscclpp::ep::MoERuntime& self) { return reinterpret_cast(self.outputTopkIdsBuffer()); }) .def("output_topk_weights_buffer_ptr", [](const mscclpp::ep::MoERuntime& self) { - return reinterpret_cast( - narrow(self, "LOW_LATENCY").outputTopkWeightsBuffer()); - }) - .def("output_tokens_buffer_ptr", - [](const mscclpp::ep::MoERuntime& self) { - return reinterpret_cast( - narrow(self, "LOW_LATENCY").outputTokensBuffer()); - }) - .def("expert_output_buffer_ptr", - [](const mscclpp::ep::MoERuntime& self) { - return reinterpret_cast( - narrow(self, "LOW_LATENCY").expertOutputBuffer()); + return reinterpret_cast(self.outputTopkWeightsBuffer()); }) + .def("dispatch_output_buffer_ptr", + [](const mscclpp::ep::MoERuntime& self) { return reinterpret_cast(self.dispatchOutputBuffer()); }) + .def("combine_input_buffer_ptr", + [](const mscclpp::ep::MoERuntime& self) { return reinterpret_cast(self.combineInputBuffer()); }) .def( - "ll_dispatch", + "dispatch", [](mscclpp::ep::MoERuntime& self, uintptr_t inputPtr, uintptr_t topkIdxPtr, uintptr_t topkWeightsPtr, uintptr_t outputPtr, uintptr_t outputScalesPtr, uintptr_t outputSrcInfoPtr, uintptr_t outputTopkIdxPtr, uintptr_t outputTopkWeightsPtr, uintptr_t outputLayoutRangePtr, uintptr_t outputCountPtr, int numTokens, int hidden, int numTopk, int maxTokensPerRank, int numExperts, int invalidTokenExpertId, - mscclpp::ep::DispatchLayout dispatchLayout, mscclpp::ep::low_latency::DispatchDataType dispatchDataType, - int numBlocks, uintptr_t streamPtr) { - narrow(self, "LOW_LATENCY") - .dispatch( - ptr(outputPtr), ptr(outputScalesPtr), reinterpret_cast(ptr(outputSrcInfoPtr)), - reinterpret_cast(ptr(outputTopkIdxPtr)), reinterpret_cast(ptr(outputTopkWeightsPtr)), - reinterpret_cast(ptr(outputLayoutRangePtr)), reinterpret_cast(ptr(outputCountPtr)), - ptr(inputPtr), reinterpret_cast(ptr(topkIdxPtr)), - reinterpret_cast(ptr(topkWeightsPtr)), numTokens, hidden, numTopk, maxTokensPerRank, - numExperts, invalidTokenExpertId, dispatchLayout, dispatchDataType, numBlocks, stream(streamPtr)); + mscclpp::ep::DispatchLayout dispatchLayout, mscclpp::ep::DispatchDataType dispatchDataType, int numBlocks, + uintptr_t streamPtr) { + self.dispatch(mscclpp::ep::DispatchRequest{mscclpp::ep::LatencyDispatchRequest{ + .output = ptr(outputPtr), + .outputScales = ptr(outputScalesPtr), + .outputSrcInfo = reinterpret_cast(ptr(outputSrcInfoPtr)), + .outputTopkIdx = reinterpret_cast(ptr(outputTopkIdxPtr)), + .outputTopkWeights = reinterpret_cast(ptr(outputTopkWeightsPtr)), + .outputLayoutRange = reinterpret_cast(ptr(outputLayoutRangePtr)), + .outputCount = reinterpret_cast(ptr(outputCountPtr)), + .input = ptr(inputPtr), + .topkIdx = reinterpret_cast(ptr(topkIdxPtr)), + .topkWeights = reinterpret_cast(ptr(topkWeightsPtr)), + .numTokens = numTokens, + .hidden = hidden, + .numTopk = numTopk, + .maxTokensPerRank = maxTokensPerRank, + .numExperts = numExperts, + .invalidTokenExpertId = invalidTokenExpertId, + .dispatchLayout = dispatchLayout, + .dispatchDataType = dispatchDataType, + .numBlocks = numBlocks, + .stream = 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_topk_idx_ptr"), @@ -136,88 +107,106 @@ NB_MODULE(mscclpp_ep_cpp, m) { nb::arg("num_experts"), nb::arg("invalid_token_expert_id"), nb::arg("dispatch_layout"), nb::arg("dispatch_data_type"), nb::arg("num_blocks"), nb::arg("stream_ptr")) .def( - "ll_combine", + "combine", [](mscclpp::ep::MoERuntime& self, uintptr_t expertOutputPtr, uintptr_t topkIdxPtr, uintptr_t topkWeightsPtr, uintptr_t srcInfoPtr, uintptr_t layoutRangePtr, uintptr_t outputPtr, int numTokens, int hidden, int numTopk, int maxTokensPerRank, int numExperts, mscclpp::ep::DispatchLayout dispatchLayout, - mscclpp::ep::low_latency::DispatchDataType dispatchDataType, mscclpp::ep::low_latency::CombineMode mode, - int numBlocks, uintptr_t streamPtr) { - narrow(self, "LOW_LATENCY") - .combine(ptr(outputPtr), ptr(expertOutputPtr), reinterpret_cast(ptr(topkIdxPtr)), - reinterpret_cast(ptr(topkWeightsPtr)), reinterpret_cast(ptr(srcInfoPtr)), - reinterpret_cast(ptr(layoutRangePtr)), numTokens, hidden, numTopk, maxTokensPerRank, - numExperts, dispatchLayout, dispatchDataType, mode, numBlocks, stream(streamPtr)); + mscclpp::ep::DispatchDataType dispatchDataType, mscclpp::ep::CombineMode mode, int numBlocks, + uintptr_t streamPtr) { + self.combine(mscclpp::ep::CombineRequest{mscclpp::ep::LatencyCombineRequest{ + .output = ptr(outputPtr), + .input = ptr(expertOutputPtr), + .topkIdx = reinterpret_cast(ptr(topkIdxPtr)), + .topkWeights = reinterpret_cast(ptr(topkWeightsPtr)), + .srcInfo = reinterpret_cast(ptr(srcInfoPtr)), + .layoutRange = reinterpret_cast(ptr(layoutRangePtr)), + .numTokens = numTokens, + .hidden = hidden, + .numTopk = numTopk, + .maxTokensPerRank = maxTokensPerRank, + .numExperts = numExperts, + .dispatchLayout = dispatchLayout, + .dispatchDataType = dispatchDataType, + .combineMode = mode, + .numBlocks = numBlocks, + .stream = stream(streamPtr), + }}); }, nb::arg("expert_output_ptr"), nb::arg("topk_idx_ptr"), nb::arg("topk_weights_ptr"), nb::arg("src_info_ptr"), nb::arg("layout_range_ptr"), nb::arg("output_ptr"), nb::arg("num_tokens"), nb::arg("hidden"), nb::arg("num_topk"), nb::arg("max_tokens_per_rank"), nb::arg("num_experts"), nb::arg("dispatch_layout"), nb::arg("dispatch_data_type"), nb::arg("mode"), nb::arg("num_blocks"), nb::arg("stream_ptr")) .def( - "ht_compute_dispatch_counts", + "token_major_prepare", [](mscclpp::ep::MoERuntime& self, uintptr_t num_tokens_per_rank_ptr, uintptr_t num_tokens_per_expert_ptr, uintptr_t is_token_in_rank_ptr, uintptr_t topk_idx_ptr, int num_tokens, int num_topk, int num_experts, uintptr_t stream_ptr) { - narrow(self, "HIGH_THROUGHPUT") - .computeDispatchCounts(reinterpret_cast(ptr(num_tokens_per_rank_ptr)), - reinterpret_cast(ptr(num_tokens_per_expert_ptr)), - reinterpret_cast(ptr(is_token_in_rank_ptr)), - reinterpret_cast(ptr(topk_idx_ptr)), num_tokens, num_topk, - num_experts, stream(stream_ptr)); + self.tokenMajorPrepare(reinterpret_cast(ptr(num_tokens_per_rank_ptr)), + reinterpret_cast(ptr(num_tokens_per_expert_ptr)), + reinterpret_cast(ptr(is_token_in_rank_ptr)), + reinterpret_cast(ptr(topk_idx_ptr)), num_tokens, num_topk, + num_experts, stream(stream_ptr)); }, nb::arg("num_tokens_per_rank_ptr"), nb::arg("num_tokens_per_expert_ptr"), nb::arg("is_token_in_rank_ptr"), nb::arg("topk_idx_ptr"), nb::arg("num_tokens"), nb::arg("num_topk"), nb::arg("num_experts"), nb::arg("stream_ptr")) - .def("ht_get_dispatch_num_channels", - [](const mscclpp::ep::MoERuntime& self, int x_element_size) { - return narrow(self, "HIGH_THROUGHPUT") - .getDispatchNumChannels(x_element_size); - }) - .def("ht_resolve_recv_x_buffer", + .def("token_major_num_channels", [](const mscclpp::ep::MoERuntime& self, + int x_element_size) { return self.tokenMajorNumChannels(x_element_size); }) + .def("token_major_resolve_recv_buffer", [](const mscclpp::ep::MoERuntime& self, int num_tokens, int num_recv_tokens, int hidden, int x_element_size) -> uintptr_t { return reinterpret_cast( - narrow(self, "HIGH_THROUGHPUT") - .resolveRecvXBuffer(num_tokens, num_recv_tokens, hidden, x_element_size)); + self.tokenMajorResolveRecvBuffer(num_tokens, num_recv_tokens, hidden, x_element_size)); }) .def( - "ht_notify_dispatch", + "token_major_notify", [](mscclpp::ep::MoERuntime& self, uintptr_t rank_prefix_matrix_ptr, uintptr_t channel_prefix_matrix_ptr, uintptr_t num_recv_tokens_per_expert_ptr, uintptr_t num_tokens_per_rank_ptr, uintptr_t num_tokens_per_expert_ptr, uintptr_t is_token_in_rank_ptr, int num_tokens, int num_experts, int x_element_size, int expert_alignment, uintptr_t stream_ptr) { - return narrow(self, "HIGH_THROUGHPUT") - .notifyDispatch(reinterpret_cast(ptr(rank_prefix_matrix_ptr)), - reinterpret_cast(ptr(channel_prefix_matrix_ptr)), - reinterpret_cast(ptr(num_recv_tokens_per_expert_ptr)), - reinterpret_cast(ptr(num_tokens_per_rank_ptr)), - reinterpret_cast(ptr(num_tokens_per_expert_ptr)), - reinterpret_cast(ptr(is_token_in_rank_ptr)), num_tokens, num_experts, - x_element_size, expert_alignment, stream(stream_ptr)); + return self.tokenMajorNotify(reinterpret_cast(ptr(rank_prefix_matrix_ptr)), + reinterpret_cast(ptr(channel_prefix_matrix_ptr)), + reinterpret_cast(ptr(num_recv_tokens_per_expert_ptr)), + reinterpret_cast(ptr(num_tokens_per_rank_ptr)), + reinterpret_cast(ptr(num_tokens_per_expert_ptr)), + reinterpret_cast(ptr(is_token_in_rank_ptr)), num_tokens, + num_experts, x_element_size, expert_alignment, stream(stream_ptr)); }, nb::arg("rank_prefix_matrix_ptr"), nb::arg("channel_prefix_matrix_ptr"), nb::arg("num_recv_tokens_per_expert_ptr"), nb::arg("num_tokens_per_rank_ptr"), nb::arg("num_tokens_per_expert_ptr"), nb::arg("is_token_in_rank_ptr"), nb::arg("num_tokens"), nb::arg("num_experts"), nb::arg("x_element_size"), nb::arg("expert_alignment"), nb::arg("stream_ptr")) .def( - "ht_dispatch", + "dispatch", [](mscclpp::ep::MoERuntime& self, uintptr_t recv_x_ptr, uintptr_t recv_x_scales_ptr, uintptr_t recv_topk_idx_ptr, uintptr_t recv_topk_weights_ptr, uintptr_t send_head_ptr, uintptr_t x_ptr, uintptr_t x_scales_ptr, uintptr_t topk_idx_ptr, uintptr_t topk_weights_ptr, uintptr_t is_token_in_rank_ptr, uintptr_t rank_prefix_matrix_ptr, uintptr_t channel_prefix_matrix_ptr, int num_tokens, int hidden, int num_topk, int num_scales, int num_experts, int x_element_size, int num_recv_tokens, bool cached_mode, uintptr_t stream_ptr) { - narrow(self, "HIGH_THROUGHPUT") - .dispatch(ptr(recv_x_ptr), reinterpret_cast(ptr(recv_x_scales_ptr)), - reinterpret_cast(ptr(recv_topk_idx_ptr)), - reinterpret_cast(ptr(recv_topk_weights_ptr)), - reinterpret_cast(ptr(send_head_ptr)), ptr(x_ptr), - reinterpret_cast(ptr(x_scales_ptr)), - reinterpret_cast(ptr(topk_idx_ptr)), - reinterpret_cast(ptr(topk_weights_ptr)), - reinterpret_cast(ptr(is_token_in_rank_ptr)), - reinterpret_cast(ptr(rank_prefix_matrix_ptr)), - reinterpret_cast(ptr(channel_prefix_matrix_ptr)), num_tokens, hidden, num_topk, - num_scales, num_experts, x_element_size, num_recv_tokens, cached_mode, stream(stream_ptr)); + self.dispatch(mscclpp::ep::DispatchRequest{mscclpp::ep::ThroughputDispatchRequest{ + .recvX = ptr(recv_x_ptr), + .recvXScales = reinterpret_cast(ptr(recv_x_scales_ptr)), + .recvTopkIdx = reinterpret_cast(ptr(recv_topk_idx_ptr)), + .recvTopkWeights = reinterpret_cast(ptr(recv_topk_weights_ptr)), + .sendHead = reinterpret_cast(ptr(send_head_ptr)), + .input = ptr(x_ptr), + .inputScales = reinterpret_cast(ptr(x_scales_ptr)), + .topkIdx = reinterpret_cast(ptr(topk_idx_ptr)), + .topkWeights = reinterpret_cast(ptr(topk_weights_ptr)), + .isTokenInRank = reinterpret_cast(ptr(is_token_in_rank_ptr)), + .rankPrefixMatrix = reinterpret_cast(ptr(rank_prefix_matrix_ptr)), + .channelPrefixMatrix = reinterpret_cast(ptr(channel_prefix_matrix_ptr)), + .numTokens = num_tokens, + .hidden = hidden, + .numTopk = num_topk, + .numScales = num_scales, + .numExperts = num_experts, + .inputElementSize = x_element_size, + .numRecvTokens = num_recv_tokens, + .cachedMode = cached_mode, + .stream = stream(stream_ptr), + }}); }, nb::arg("recv_x_ptr"), nb::arg("recv_x_scales_ptr"), nb::arg("recv_topk_idx_ptr"), nb::arg("recv_topk_weights_ptr"), nb::arg("send_head_ptr"), nb::arg("x_ptr"), nb::arg("x_scales_ptr"), @@ -226,15 +215,23 @@ NB_MODULE(mscclpp_ep_cpp, m) { nb::arg("hidden"), nb::arg("num_topk"), nb::arg("num_scales"), nb::arg("num_experts"), nb::arg("x_element_size"), nb::arg("num_recv_tokens"), nb::arg("cached_mode"), nb::arg("stream_ptr")) .def( - "ht_combine", + "combine", [](mscclpp::ep::MoERuntime& self, uintptr_t combined_x_ptr, uintptr_t combined_topk_weights_ptr, uintptr_t x_ptr, uintptr_t topk_weights_ptr, uintptr_t send_head_ptr, int num_input_tokens, int num_output_tokens, int hidden, int num_topk, int x_element_size, uintptr_t stream_ptr) { - narrow(self, "HIGH_THROUGHPUT") - .combine(ptr(combined_x_ptr), reinterpret_cast(ptr(combined_topk_weights_ptr)), ptr(x_ptr), - reinterpret_cast(ptr(topk_weights_ptr)), - reinterpret_cast(ptr(send_head_ptr)), num_input_tokens, num_output_tokens, hidden, - num_topk, x_element_size, stream(stream_ptr)); + self.combine(mscclpp::ep::CombineRequest{mscclpp::ep::ThroughputCombineRequest{ + .output = ptr(combined_x_ptr), + .outputTopkWeights = reinterpret_cast(ptr(combined_topk_weights_ptr)), + .input = ptr(x_ptr), + .topkWeights = reinterpret_cast(ptr(topk_weights_ptr)), + .sendHead = reinterpret_cast(ptr(send_head_ptr)), + .numInputTokens = num_input_tokens, + .numOutputTokens = num_output_tokens, + .hidden = hidden, + .numTopk = num_topk, + .inputElementSize = x_element_size, + .stream = stream(stream_ptr), + }}); }, nb::arg("combined_x_ptr"), nb::arg("combined_topk_weights_ptr"), nb::arg("x_ptr"), nb::arg("topk_weights_ptr"), nb::arg("send_head_ptr"), nb::arg("num_input_tokens"), diff --git a/src/ext/ep/low_latency/combine.cu b/src/ext/ep/combine/common.cuh similarity index 80% rename from src/ext/ep/low_latency/combine.cu rename to src/ext/ep/combine/common.cuh index d7f2d5f79..fc5ba88dc 100644 --- a/src/ext/ep/low_latency/combine.cu +++ b/src/ext/ep/combine/common.cuh @@ -1,18 +1,18 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +#ifndef MSCCLPP_EP_COMBINE_COMMON_CUH_ +#define MSCCLPP_EP_COMBINE_COMMON_CUH_ #include #include -#include "api.cuh" -#include "config.cuh" -#include "device_helpers.cuh" -#include "exception.cuh" +#include "common/device_helpers.cuh" +#include "common/latency.cuh" +#include "exception.hpp" +#include "kernels.hpp" namespace mscclpp { namespace ep { -namespace low_latency { -namespace detail { constexpr int CombineNWarps = 32; constexpr int CombineNThreads = CombineNWarps * WARP_SIZE; @@ -39,7 +39,7 @@ MSCCLPP_HOST_DEVICE_INLINE int directSendWorkerCount(int nLocalExperts) { return availableWorkers < DirectSendMaxNWorkers ? availableWorkers : DirectSendMaxNWorkers; } -template +template MSCCLPP_HOST_DEVICE_INLINE size_t combineSharedBytes(int nLocalExperts, int nTopk) { if constexpr (Layout == DispatchLayout::RANK_MAJOR) { if (nTopk <= RankMajorTmaMaxNTopk) { @@ -53,7 +53,7 @@ MSCCLPP_HOST_DEVICE_INLINE size_t combineSharedBytes(int nLocalExperts, int nTop return 0; } constexpr size_t TileBytes = static_cast(Hidden) * sizeof(Bf16); - if constexpr (Mode == low_latency::CombineMode::DIRECT_SEND) { + if constexpr (Mode == CombineMode::DIRECT_SEND) { return directSendControlBytes(nLocalExperts) + static_cast(directSendWorkerCount(nLocalExperts)) * directSendWorkerBytes(); } @@ -585,24 +585,23 @@ MSCCLPP_DEVICE_INLINE void recvExpertRowsDirect(void* output, const int64_t* __r #endif // MSCCLPP_BULK_AVAILABLE -template -__global__ __launch_bounds__(CombineNThreads, 1) void combineKernel( - void* output, const void* expertOutput, const int64_t* __restrict__ topkIndices, - const float* __restrict__ topkWeights, const int* srcInfo, const int64_t* layoutRange, Workload workload, - void* combineRecvBuffer, const void* dispatchRecvBuffer, CommContext comm, void* workspace) { +template +MSCCLPP_DEVICE_INLINE void combineBody(void* output, const void* expertOutput, const int64_t* __restrict__ topkIndices, + const float* __restrict__ topkWeights, const int* srcInfo, + const int64_t* layoutRange, Workload workload, void* combineRecvBuffer, + const void* dispatchRecvBuffer, const DeviceContext* context) { #if MSCCLPP_BULK_AVAILABLE extern __shared__ __align__(128) uint8_t sharedMemory[]; const int nTokens = workload.numTokens_; const int nExperts = workload.numExperts_; - const int nRanks = comm.numRanks_; + const int nRanks = context->numRanks_; const int nTopk = workload.numTopk_; const int maxTokensPerRank = workload.maxTokensPerRank_; - const TransportView transport(comm); - WorkspaceView workspaceView(workspace, nRanks, nExperts); + const TransportView transport(context); + WorkspaceView workspaceView(context->workspace_, nRanks, nExperts); if constexpr (Layout == DispatchLayout::RANK_MAJOR) { - static_assert(Mode == low_latency::CombineMode::RANK_LOCAL_REDUCE); + static_assert(Mode == CombineMode::RANK_LOCAL_REDUCE); static_assert(DispatchType == DispatchDataType::BF16); if (nTopk <= RankMajorTmaMaxNTopk) { const uint32_t epoch = workload.epoch_; @@ -618,7 +617,7 @@ __global__ __launch_bounds__(CombineNThreads, 1) void combineKernel( recvRankMajorRemotePartials(output, expertOutput, topkIndices, nTokens, nTopk, nExperts, nRanks, maxTokensPerRank, transport, workspaceView); return; - } else if constexpr (Mode == low_latency::CombineMode::RANK_LOCAL_REDUCE) { + } else if constexpr (Mode == CombineMode::RANK_LOCAL_REDUCE) { sendRankReducedPartials( expertOutput, nExperts, nRanks, nTopk, maxTokensPerRank, combineRecvBuffer, dispatchRecvBuffer, transport, workspaceView, sharedMemory); @@ -631,7 +630,7 @@ __global__ __launch_bounds__(CombineNThreads, 1) void combineKernel( exchangeCombineReady(transport, nRanks); workspaceView.combineSyncer_->sync(gridDim.x); - if constexpr (Mode == low_latency::CombineMode::RANK_LOCAL_REDUCE) { + if constexpr (Mode == CombineMode::RANK_LOCAL_REDUCE) { recvRankLocalPartials(output, topkIndices, nTokens, nTopk, nExperts, nRanks, maxTokensPerRank, combineRecvBuffer, sharedMemory); } else { @@ -640,153 +639,152 @@ __global__ __launch_bounds__(CombineNThreads, 1) void combineKernel( #endif // MSCCLPP_BULK_AVAILABLE } -template +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::Workload& workload, void* recvBuffer, void* dispatchRecvBuffer, - const low_latency::CommContext& comm, void* workspace, int numBlocks, - cudaStream_t stream) { + const Workload& workload, void* recvBuffer, void* dispatchRecvBuffer, + const DeviceContext& context, int numBlocks, cudaStream_t stream) { static_assert(Hidden == 2048 || Hidden == 4096 || Hidden == 4352 || Hidden == 6656 || Hidden == 7168 || Hidden == 8192 || Hidden == 8704 || Hidden == 9216); const int nExperts = workload.numExperts_; - const int nRanks = comm.numRanks_; + const int nRanks = context.numRanks_; const int nLocalExperts = nExperts / nRanks; - if constexpr (Mode == low_latency::CombineMode::DIRECT_SEND) { + if constexpr (Mode == CombineMode::DIRECT_SEND) { static_assert(Layout == DispatchLayout::EXPERT_MAJOR); EP_HOST_ASSERT(directSendWorkerCount(nLocalExperts) > 0); } - auto combineFunc = combineKernel; + auto combineFunc = KernelSelector::template get(); const size_t sharedBytes = combineSharedBytes(nLocalExperts, workload.numTopk_); - const bool useRankMajorTma = Layout == DispatchLayout::RANK_MAJOR && - Mode == low_latency::CombineMode::RANK_LOCAL_REDUCE && + const bool useRankMajorTma = Layout == DispatchLayout::RANK_MAJOR && Mode == CombineMode::RANK_LOCAL_REDUCE && workload.numTopk_ <= RankMajorTmaMaxNTopk; const int launchBlocks = numBlocks + (useRankMajorTma ? 1 : 0); static thread_local KernelConfigCache kernelConfig; - const int residentBlocks = configureKernel(combineFunc, CombineNThreads, sharedBytes, comm, kernelConfig); + const int residentBlocks = configureKernel(combineFunc, CombineNThreads, sharedBytes, context, kernelConfig); EP_HOST_ASSERT(residentBlocks >= launchBlocks); - combineKernel - <<>>( - output, expertOutput, topkIndices, topkWeights, srcInfo, layoutRange, workload, recvBuffer, - dispatchRecvBuffer, comm, workspace); + combineFunc<<>>( + output, expertOutput, topkIndices, topkWeights, srcInfo, layoutRange, workload, recvBuffer, dispatchRecvBuffer, + context.devicePtr_); CUDA_CHECK(cudaGetLastError()); } -template