From c7a335a41ac58612c2864263b503cbf162fce89d Mon Sep 17 00:00:00 2001 From: Binyang Li Date: Thu, 13 Aug 2026 21:54:44 +0000 Subject: [PATCH 01/25] Refactor EP runtime architecture Unify latency and overlap under one runtime, organize kernels by dispatch and combine algorithms, and expose one opaque dispatch handle. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- direct-rank-design.md | 634 ++++++++++++++++++ python/mscclpp/ep/README.md | 59 +- python/mscclpp/ep/__init__.py | 20 +- python/mscclpp/ep/backend.py | 104 +++ python/mscclpp/ep/communicator.py | 31 +- .../mscclpp/ep/{low_latency.py => latency.py} | 159 ++--- .../ep/{high_throughput.py => overlap.py} | 212 +++--- python/mscclpp/ep/runtime.py | 49 ++ python/mscclpp/ep/types.py | 46 +- src/ext/ep/CMakeLists.txt | 28 +- src/ext/ep/README.md | 100 +-- src/ext/ep/bindings.cpp | 176 ++--- .../combine.cu => combine/latency/common.cuh} | 209 +++--- src/ext/ep/combine/latency/direct_send.cu | 38 ++ .../ep/combine/latency/rank_local_reduce.cu | 38 ++ .../overlap/token_major_reduce.cu} | 65 +- .../config.cuh => common/fixed_buffer.cuh} | 37 +- .../overlap_barrier.cuh} | 6 +- .../config.cuh => common/recv_pool.cuh} | 6 +- src/ext/ep/config.hpp | 38 +- .../latency/common.cuh} | 194 +++--- src/ext/ep/dispatch/latency/expert_major.cu | 44 ++ src/ext/ep/dispatch/latency/rank_major.cu | 39 ++ src/ext/ep/dispatch/overlap/token_major.cu | 332 +++++++++ .../overlap/token_major_prepare.cu} | 30 +- src/ext/ep/high-throughput/dispatch.cu | 344 ---------- src/ext/ep/ht_runtime.hpp | 91 --- src/ext/ep/include/api.cuh | 148 ++-- src/ext/ep/include/device_context.cuh | 41 ++ src/ext/ep/include/launch.cuh | 2 +- src/ext/ep/ll_runtime.hpp | 71 -- src/ext/ep/moe_runtime.cc | 138 +++- src/ext/ep/moe_runtime.hpp | 94 ++- .../fixed_buffer.cc} | 157 +++-- .../{ht_runtime.cc => runtime/recv_pool.cc} | 168 +++-- src/ext/ep/runtime/resources.hpp | 139 ++++ src/ext/ep/runtime_base.hpp | 62 -- test/python/ep/CMakeLists.txt | 64 +- test/python/ep/mscclpp_ep_bench.cu | 494 -------------- test/python/ep/run_ep_bench.py | 482 ------------- test/python/ep/run_ep_bench_python.py | 15 +- test/python/ep/test_intranode_multirank.py | 4 +- test/python/ep/test_low_latency_multirank.py | 8 +- 43 files changed, 2558 insertions(+), 2658 deletions(-) create mode 100644 direct-rank-design.md create mode 100644 python/mscclpp/ep/backend.py rename python/mscclpp/ep/{low_latency.py => latency.py} (81%) rename python/mscclpp/ep/{high_throughput.py => overlap.py} (78%) create mode 100644 python/mscclpp/ep/runtime.py rename src/ext/ep/{low_latency/combine.cu => combine/latency/common.cuh} (80%) create mode 100644 src/ext/ep/combine/latency/direct_send.cu create mode 100644 src/ext/ep/combine/latency/rank_local_reduce.cu rename src/ext/ep/{high-throughput/combine.cu => combine/overlap/token_major_reduce.cu} (78%) rename src/ext/ep/{low_latency/config.cuh => common/fixed_buffer.cuh} (89%) rename src/ext/ep/{high-throughput/barrier.cuh => common/overlap_barrier.cuh} (86%) rename src/ext/ep/{high-throughput/config.cuh => common/recv_pool.cuh} (93%) rename src/ext/ep/{low_latency/dispatch.cu => dispatch/latency/common.cuh} (86%) create mode 100644 src/ext/ep/dispatch/latency/expert_major.cu create mode 100644 src/ext/ep/dispatch/latency/rank_major.cu create mode 100644 src/ext/ep/dispatch/overlap/token_major.cu rename src/ext/ep/{high-throughput/counts.cu => dispatch/overlap/token_major_prepare.cu} (73%) delete mode 100644 src/ext/ep/high-throughput/dispatch.cu delete mode 100644 src/ext/ep/ht_runtime.hpp create mode 100644 src/ext/ep/include/device_context.cuh delete mode 100644 src/ext/ep/ll_runtime.hpp rename src/ext/ep/{ll_runtime.cc => runtime/fixed_buffer.cc} (52%) rename src/ext/ep/{ht_runtime.cc => runtime/recv_pool.cc} (55%) create mode 100644 src/ext/ep/runtime/resources.hpp delete mode 100644 src/ext/ep/runtime_base.hpp delete mode 100644 test/python/ep/mscclpp_ep_bench.cu delete mode 100644 test/python/ep/run_ep_bench.py diff --git a/direct-rank-design.md b/direct-rank-design.md new file mode 100644 index 00000000..f73fa515 --- /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/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/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 + +src/ext/ep/include/api.cuh + 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/python/mscclpp/ep/README.md b/python/mscclpp/ep/README.md index 04bf21fc..92ffa56b 100644 --- a/python/mscclpp/ep/README.md +++ b/python/mscclpp/ep/README.md @@ -62,7 +62,7 @@ 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 @@ -86,7 +86,7 @@ moe_comm = MoECommunicator( hidden_size=hidden_size, topk=topk, max_tokens_per_rank=max_tokens, - mode=MoEMode.HIGH_THROUGHPUT, + mode=MoEMode.OVERLAP, ) ``` @@ -135,7 +135,7 @@ 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.OVERLAP`) | | `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 | @@ -149,14 +149,14 @@ 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.OVERLAP`. `mode` must be a `MoEMode` enum value, not a string. +Latency algorithms support expert-major and rank-major output layouts. Overlap +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 @@ -167,15 +167,15 @@ The mode also fixes the SM budget, which is the main scheduling consideration: | Mode | SMs used | 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` | ~128 (`low_latency_num_blocks - 2`) | minimize standalone latency | +| `MoEMode.OVERLAP` | 20 (`Config.num_sms`) | 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` | +| `OVERLAP` | `DispatchLayout.TOKEN_MAJOR` | +| `LATENCY` | `DispatchLayout.EXPERT_MAJOR` | `output_layout` may still be kept as an advanced override if a backend supports multiple layouts within the same mode. @@ -184,9 +184,9 @@ Use `DispatchLayout` instead of string literals for this field: | Layout enum | Tensor shape | |---|---| -| `DispatchLayout.TOKEN_MAJOR` | HT: `[total_recv_tokens, hidden]` | +| `DispatchLayout.TOKEN_MAJOR` | Overlap: `[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 @@ -428,7 +428,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 +456,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 overlap 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 +498,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 +526,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 +535,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 +### Overlap token-major layout -HT uses `DispatchLayout.TOKEN_MAJOR`: +Overlap algorithms use `DispatchLayout.TOKEN_MAJOR`: ```python dispatch_out.tokens # [total_recv_tokens, H] @@ -548,9 +548,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] @@ -603,7 +603,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: overlap [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 +802,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 +Overlap dispatch needs a metadata phase before payload movement: ```text @@ -813,9 +813,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 e30de9b3..f77026c8 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.OVERLAP` selects bounded-resource +token-major algorithms. """ from .communicator import ( # noqa: F401 BlockOverlapConfig, CommOverlapConfig, - CombineContext, 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", "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/backend.py b/python/mscclpp/ep/backend.py new file mode 100644 index 00000000..a1f62e23 --- /dev/null +++ b/python/mscclpp/ep/backend.py @@ -0,0 +1,104 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Unified high-level expert-parallel backend.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from ._cpp import DispatchLayout, MoEMode +from .latency import _LatencyMethods +from .overlap import _OverlapMethods +from .runtime import Runtime +from .types import DispatchHandle, DispatchOutput, MoECommunicatorConfig, QuantConfig + + +class Backend(_LatencyMethods, _OverlapMethods): + """Own one runtime and expose latency or overlap dispatch/combine.""" + + def __init__(self, config: MoECommunicatorConfig, output_layout: DispatchLayout) -> None: + if config.comm is None: + raise ValueError("MoECommunicator requires an mscclpp.CommGroup via comm=") + if config.mode == MoEMode.LATENCY: + self.runtime = Runtime( + config.comm, + MoEMode.LATENCY, + max_tokens_per_rank=config.max_tokens_per_rank, + hidden=config.hidden_size, + num_experts=config.num_experts, + num_topk=config.topk, + output_layout=output_layout, + ) + self._init_latency(config, output_layout) + else: + max_hidden_bytes = config.hidden_size * torch.empty((), dtype=torch.bfloat16).element_size() + self.runtime = Runtime( + config.comm, + MoEMode.OVERLAP, + max_hidden_bytes=max_hidden_bytes, + num_sms=config.num_sms, + ) + self._init_overlap(config, output_layout) + self.expert_output_buffer = None + + def is_available(self) -> bool: + """Return whether the selected operation family is available.""" + return self.runtime.is_available() + + def is_internode_available(self) -> bool: + """Return whether the selected operations support this internode topology.""" + return self.runtime.is_internode_available() + + def is_internode(self) -> bool: + """Return whether the runtime spans more than one node.""" + 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]: + """Dispatch tokens with the selected operation family.""" + if self.mode == MoEMode.LATENCY: + return self._dispatch_latency( + input, + topk_ids, + weights, + quant, + output_buffer=output_buffer, + stream=stream, + previous_handle=previous_handle, + runtime_max_tokens_per_rank=runtime_max_tokens_per_rank, + ) + return self._dispatch_overlap( + input, + topk_ids, + weights, + quant, + output_buffer=output_buffer, + stream=stream, + previous_handle=previous_handle, + runtime_max_tokens_per_rank=runtime_max_tokens_per_rank, + ) + + def combine( + self, + expert_output: torch.Tensor, + handle: DispatchHandle, + *, + out: Optional[torch.Tensor], + stream: Optional[torch.cuda.Stream], + ) -> torch.Tensor: + """Combine expert output with the selected operation family.""" + if self.mode == MoEMode.LATENCY: + return self._combine_latency(expert_output, handle, out=out, stream=stream) + return self._combine_overlap(expert_output, handle, out=out, stream=stream) diff --git a/python/mscclpp/ep/communicator.py b/python/mscclpp/ep/communicator.py index c9c8030d..1844cd33 100644 --- a/python/mscclpp/ep/communicator.py +++ b/python/mscclpp/ep/communicator.py @@ -9,31 +9,22 @@ import torch from ._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode -from .high_throughput import HighThroughputBackend -from .low_latency import LowLatencyBackend +from .backend import Backend from .types import ( BlockOverlapConfig, CommOverlapConfig, - CombineContext, DispatchHandle, DispatchLayoutInfo, DispatchOutput, DispatchOutputInfo, - ExpertMajorDispatchHandle, - ExpertMajorCombineContext, - HighThroughputDispatchHandle, - HighThroughputCombineContext, MoECommunicatorConfig, OperationOverlapConfig, QuantConfig, - RankMajorDispatchHandle, - RankMajorCombineContext, ) __all__ = [ "CommOverlapConfig", "BlockOverlapConfig", - "CombineContext", "CombineMode", "DispatchHandle", "DispatchDataType", @@ -41,25 +32,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.OVERLAP` selects bounded-resource overlap + algorithms (TOKEN_MAJOR). """ def __init__(self, config: Optional[MoECommunicatorConfig] = None, **kwargs) -> None: @@ -77,10 +63,7 @@ def __init__(self, config: Optional[MoECommunicatorConfig] = None, **kwargs) -> _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._backend = Backend(config, self.output_layout) self._publish_backend_state() def _publish_backend_state(self) -> None: @@ -152,7 +135,7 @@ def get_expert_output_buffer(self) -> torch.Tensor: """ 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") + raise RuntimeError("expert output buffer is only available for RANK_MAJOR latency mode") return buffer def dispatch_async(self, *args, **kwargs): @@ -180,7 +163,7 @@ def _validate_common_config(config: MoECommunicatorConfig) -> None: 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 + return DispatchLayout.EXPERT_MAJOR if mode == MoEMode.LATENCY else DispatchLayout.TOKEN_MAJOR if not isinstance(layout, DispatchLayout): raise TypeError("MoECommunicatorConfig.output_layout must be a DispatchLayout") return layout diff --git a/python/mscclpp/ep/low_latency.py b/python/mscclpp/ep/latency.py similarity index 81% rename from python/mscclpp/ep/low_latency.py rename to python/mscclpp/ep/latency.py index 40bccd01..b47486f0 100644 --- a/python/mscclpp/ep/low_latency.py +++ b/python/mscclpp/ep/latency.py @@ -1,6 +1,6 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Low-latency backend for the high-level MoE communicator.""" +"""Latency-optimized backend for the high-level MoE communicator.""" from __future__ import annotations @@ -8,18 +8,16 @@ import torch -from ._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode, create_moe_runtime +from ._cpp import CombineMode, DispatchDataType, DispatchLayout, MoEMode from .types import ( DispatchHandle, DispatchLayoutInfo, DispatchOutput, DispatchOutputInfo, - ExpertMajorDispatchHandle, - ExpertMajorCombineContext, MoECommunicatorConfig, QuantConfig, - RankMajorCombineContext, - RankMajorDispatchHandle, + _ExpertMajorCombineContext, + _RankMajorCombineContext, ) from .utils import cuda_stream_ptr, resolve_expert_placement @@ -34,7 +32,7 @@ def _resolve_dispatch_data_type(quant: Optional[QuantConfig]) -> DispatchDataTyp if quant_format is None: raise ValueError("quant.format is required") if quant_format != DispatchDataType.FP8_E4M3: - raise ValueError("unsupported low-latency quantization format") + raise ValueError("unsupported latency dispatch quantization format") if quant.block_scales is not None: raise ValueError("communicator quant config must not contain precomputed scales") return quant_format @@ -107,54 +105,20 @@ def _tensor_from_pointer( return buffer_view, tensor -class LowLatencyRuntime: - """Private low-level low-latency runtime wrapper (wraps ``_cpp.MoERuntime``).""" +class _LatencyMethods: + """Provide latency-mode methods to the unified backend.""" - 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: + def _init_latency(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=") + raise ValueError("mode=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.mode = MoEMode.LATENCY self.output_layout = output_layout self.num_experts = config.num_experts @@ -173,9 +137,9 @@ def __init__(self, config: MoECommunicatorConfig, output_layout: DispatchLayout) DispatchLayout.EXPERT_MAJOR, DispatchLayout.RANK_MAJOR, ): - raise NotImplementedError("unsupported low-latency output layout") + raise NotImplementedError("unsupported latency output layout") if self.num_experts % self.world_size != 0: - raise ValueError("low-latency mode requires num_experts divisible by world_size") + raise ValueError("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): @@ -211,15 +175,7 @@ def __init__(self, config: MoECommunicatorConfig, output_layout: DispatchLayout) 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._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 @@ -235,36 +191,36 @@ def __init__(self, config: MoECommunicatorConfig, output_layout: DispatchLayout) self._output_topk_ids_owner, self._output_topk_ids, ) = _tensor_from_pointer( - self._runtime.cpp_runtime.output_topk_ids_buffer_ptr(), + self.runtime.cpp_runtime.output_topk_ids_buffer_ptr(), metadata_shape, " int: @@ -275,16 +231,7 @@ def _resolve_runtime_max_tokens_per_rank(self, runtime_max_tokens_per_rank: Opti 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( + def _dispatch_latency( self, input: torch.Tensor, topk_ids: torch.Tensor, @@ -298,12 +245,12 @@ def dispatch( ) -> 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) + self._validate_latency_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( + self.runtime.cpp_runtime.dispatch_latency( input.data_ptr(), topk_ids.data_ptr(), 0 if weights is None else weights.data_ptr(), @@ -341,7 +288,7 @@ def dispatch( num_tokens_per_rank=count, ) else: - raise ValueError(f"unsupported low-latency output layout: {self.output_layout}") + raise ValueError(f"unsupported latency output layout: {self.output_layout}") output_info = DispatchOutputInfo(layout=layout_info, quant=output_quant) dispatch_out = DispatchOutput( tokens=out_buf, @@ -353,9 +300,9 @@ def dispatch( if self.output_layout == DispatchLayout.EXPERT_MAJOR: assert layout_range is not None assert src_info is not None - handle: DispatchHandle = ExpertMajorDispatchHandle( + handle = DispatchHandle( output_info=output_info, - combine_context=ExpertMajorCombineContext( + _context=_ExpertMajorCombineContext( topk_ids=topk_ids, weights=weights, num_experts=self.num_experts, @@ -366,9 +313,9 @@ def dispatch( ), ) elif self.output_layout == DispatchLayout.RANK_MAJOR: - handle = RankMajorDispatchHandle( + handle = DispatchHandle( output_info=output_info, - combine_context=RankMajorCombineContext( + _context=_RankMajorCombineContext( topk_ids=topk_ids, num_experts=self.num_experts, num_tokens=input.size(0), @@ -377,10 +324,10 @@ def dispatch( ), ) else: - raise ValueError(f"unsupported low-latency output layout: {self.output_layout}") + raise ValueError(f"unsupported latency output layout: {self.output_layout}") return dispatch_out, handle - def combine( + def _combine_latency( self, expert_output: torch.Tensor, handle: DispatchHandle, @@ -388,29 +335,27 @@ def combine( 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 + self._validate_latency_combine_inputs(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 - elif isinstance(handle, RankMajorDispatchHandle): - context = handle.combine_context + active_capacity = self.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 low-latency combine context") - if isinstance(handle, ExpertMajorDispatchHandle): - active_capacity = self.max_tokens_per_rank + raise ValueError("DispatchHandle does not contain latency combine context") 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( + self.runtime.cpp_runtime.combine_latency( expert_output.data_ptr(), context.topk_ids.data_ptr(), 0 if topk_weights is None else topk_weights.data_ptr(), @@ -467,7 +412,7 @@ def _get_dispatch_output_tensors(self, output_buffer: torch.Tensor): 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}") + raise ValueError(f"unsupported 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 @@ -482,9 +427,11 @@ def _get_dispatch_output_tensors(self, output_buffer: torch.Tensor): self._dispatch_count, ) - def _validate_dispatch_inputs(self, input, topk_ids, weights, quant, output_buffer, active_capacity: int) -> None: + def _validate_latency_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") + raise ValueError("output_buffer is required for latency dispatch") if quant is not None: raise NotImplementedError( "per-call input quant metadata is not supported; configure dispatch output quantization on the communicator" @@ -492,7 +439,7 @@ def _validate_dispatch_inputs(self, input, topk_ids, weights, quant, output_buff 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") + raise ValueError("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: @@ -523,7 +470,7 @@ def _validate_dispatch_inputs(self, input, topk_ids, weights, quant, output_buff self.hidden_size, ) else: - raise ValueError(f"unsupported low-latency output layout: {self.output_layout}") + raise ValueError(f"unsupported 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 @@ -538,10 +485,12 @@ def _validate_dispatch_inputs(self, input, topk_ids, weights, quant, output_buff 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 + def _validate_latency_combine_inputs(self, expert_output, handle, out) -> None: + 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 != 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: @@ -551,9 +500,7 @@ def _validate_combine_inputs(self, expert_output, handle, out) -> None: 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 + context.max_tokens_per_rank if isinstance(context, _RankMajorCombineContext) else self.max_tokens_per_rank ) slots_per_expert = self.world_size * active_capacity if handle.output_info.layout.kind == DispatchLayout.EXPERT_MAJOR: @@ -568,7 +515,7 @@ def _validate_combine_inputs(self, expert_output, handle, out) -> None: self.hidden_size, ) else: - raise ValueError(f"unsupported low-latency output layout: {handle.output_info.layout.kind}") + 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: diff --git a/python/mscclpp/ep/high_throughput.py b/python/mscclpp/ep/overlap.py similarity index 78% rename from python/mscclpp/ep/high_throughput.py rename to python/mscclpp/ep/overlap.py index dc87e25e..105b70b8 100644 --- a/python/mscclpp/ep/high_throughput.py +++ b/python/mscclpp/ep/overlap.py @@ -3,9 +3,9 @@ # # 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. +"""Resource-bounded overlap backend for the high-level MoE communicator. -The C++ runtime follows the low-latency resource model: it reuses the existing +The unified C++ runtime 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 @@ -14,20 +14,19 @@ from __future__ import annotations -from typing import Any, List, Optional +from typing import List, Optional import torch -from ._cpp import Config, DispatchLayout, MoEMode, _cpp +from ._cpp import DispatchLayout, MoEMode from .types import ( DispatchHandle, DispatchLayoutInfo, DispatchOutput, DispatchOutputInfo, - HighThroughputCombineContext, - HighThroughputDispatchHandle, MoECommunicatorConfig, QuantConfig, + _TokenMajorOverlapCombineContext, ) from .utils import ( bf16_view as _bf16_view, @@ -37,47 +36,57 @@ ) -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. - """ +class _OverlapMethods: + """Provide overlap-mode methods to the unified backend.""" #: Default number of SMs reserved for comms kernels. Matches DeepEP. num_sms: int = 20 - def __init__( + def _init_overlap( self, - comm: Any, - max_hidden_bytes: int, - config: Config, + config: MoECommunicatorConfig, + output_layout: DispatchLayout, ) -> None: + comm = config.comm + if comm is None: + raise ValueError("mode=OVERLAP requires an mscclpp.CommGroup via comm=") + self.rank: int = comm.my_rank self.group_size: int = comm.nranks + self.world_size = 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 - # ------------------------------------------------------------------ + self.local_rank = torch.cuda.current_device() + self.device = torch.device("cuda", self.local_rank) + self.mode = MoEMode.OVERLAP + 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 - def is_available(self) -> bool: - return self.runtime.is_available() + if self.output_layout != DispatchLayout.TOKEN_MAJOR: + raise NotImplementedError("OVERLAP 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("overlap quantized dispatch (scales) is not implemented yet") - def is_internode_available(self) -> bool: - return self.runtime.is_internode_available() + self.expert_alignment = config.expert_alignment # ------------------------------------------------------------------ # Dispatch routing metadata # ------------------------------------------------------------------ - def compute_dispatch_counts(self, topk_idx: torch.Tensor, num_experts: int): + 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 @@ -90,7 +99,7 @@ def compute_dispatch_counts(self, topk_idx: torch.Tensor, num_experts: int): 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( + self.runtime.cpp_runtime.prepare_token_major_overlap( _ptr(num_tokens_per_rank), _ptr(num_tokens_per_expert), _ptr(is_token_in_rank), @@ -106,7 +115,7 @@ def compute_dispatch_counts(self, topk_idx: torch.Tensor, num_experts: int): # Dispatch (two-phase) + combine # ------------------------------------------------------------------ - def dispatch( + def _dispatch_token_major( self, x: torch.Tensor, x_scales: Optional[torch.Tensor], @@ -120,12 +129,12 @@ def dispatch( cached_channel_prefix_matrix: Optional[torch.Tensor], expert_alignment: int, ): - """Run high-throughput dispatch and return outputs plus combine metadata.""" + """Run token-major overlap dispatch and return 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_channels = self.runtime.cpp_runtime.get_token_major_overlap_num_channels(x_element_size) num_topk = int(topk_idx.size(1)) if topk_idx is not None else 0 num_scales = 0 @@ -146,7 +155,7 @@ def dispatch( 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( + num_recv_tokens = self.runtime.cpp_runtime.notify_token_major_overlap( _ptr(rank_prefix_matrix), _ptr(channel_prefix_matrix), _ptr(num_recv_per_expert_host), @@ -178,7 +187,7 @@ def dispatch( else None ) - self.runtime.ht_dispatch( + self.runtime.cpp_runtime.dispatch_token_major_overlap( _ptr(recv_x), _ptr(recv_x_scales), _ptr(recv_topk_idx), @@ -214,12 +223,14 @@ def dispatch( 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) + pool_ptr = self.runtime.cpp_runtime.resolve_token_major_overlap_recv_buffer( + num_tokens, num_recv_tokens, hidden, x_element_size + ) if pool_ptr == 0: - raise RuntimeError("high-throughput direct receive-pool capacity exceeded") + raise RuntimeError("token-major overlap receive-pool capacity exceeded") return _bf16_view(pool_ptr, num_recv_tokens, hidden, owner=self) - def combine( + def _combine_token_major( self, x: torch.Tensor, topk_weights: Optional[torch.Tensor], @@ -236,7 +247,7 @@ def combine( if topk_weights is not None else None ) - self.runtime.ht_combine( + self.runtime.cpp_runtime.combine_token_major_overlap( _ptr(combined_x), _ptr(combined_topk_weights), _ptr(x), @@ -251,65 +262,7 @@ def combine( ) 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( + def _dispatch_overlap( self, input: torch.Tensor, topk_ids: torch.Tensor, @@ -323,13 +276,13 @@ def dispatch( ) -> 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") + raise ValueError("runtime_max_tokens_per_rank is only supported by 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) + return self._dispatch_overlap_impl(input, topk_ids, weights, quant, previous_handle) + return self._dispatch_overlap_impl(input, topk_ids, weights, quant, previous_handle) - def _dispatch( + def _dispatch_overlap_impl( self, input: torch.Tensor, topk_ids: torch.Tensor, @@ -337,7 +290,7 @@ def _dispatch( quant: Optional[QuantConfig], previous_handle: Optional[DispatchHandle], ) -> tuple[DispatchOutput, DispatchHandle]: - self._validate_dispatch_inputs(input, topk_ids, weights, quant) + self._validate_overlap_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) @@ -354,7 +307,7 @@ def _dispatch( num_tokens_per_rank, num_tokens_per_expert, is_token_in_rank, - ) = self._runtime.compute_dispatch_counts(topk_ids, self.num_experts) + ) = self._compute_dispatch_counts(topk_ids, self.num_experts) if cache is not None: ( @@ -366,7 +319,7 @@ def _dispatch( rank_prefix_matrix, _channel_prefix_matrix, send_head, - ) = self._runtime.dispatch( + ) = self._dispatch_token_major( input, None, None, @@ -379,11 +332,15 @@ def _dispatch( cache["channel_prefix_matrix"], self.expert_alignment, ) - del _runtime_recv_topk_idx, _runtime_recv_topk_weights, _runtime_num_recv_tokens_per_expert_list + 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( + combine_context = _TokenMajorOverlapCombineContext( recv_topk_weights=recv_topk_weights, send_head=send_head, ) @@ -398,7 +355,7 @@ def _dispatch( rank_prefix_matrix, channel_prefix_matrix, send_head, - ) = self._runtime.dispatch( + ) = self._dispatch_token_major( input, None, topk_ids, @@ -411,7 +368,7 @@ def _dispatch( None, self.expert_alignment, ) - combine_context = HighThroughputCombineContext( + combine_context = _TokenMajorOverlapCombineContext( recv_topk_weights=recv_topk_weights, send_head=send_head, ) @@ -449,8 +406,8 @@ def _dispatch( 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 + handle = DispatchHandle(output_info=output_info, _context=combine_context) + # The unified runtime orders overlap 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] @@ -468,7 +425,7 @@ def _cache_matches(self, cache, input, topk_ids, weights, implicit_weights) -> b and (implicit_weights or cache.get("weights_version") == weights._version) ) - def combine( + def _combine_overlap( self, expert_output: torch.Tensor, handle: DispatchHandle, @@ -478,15 +435,18 @@ def combine( ) -> 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) + return self._combine_overlap_impl(expert_output, handle, out) + return self._combine_overlap_impl(expert_output, handle, out) - def _combine( - self, expert_output: torch.Tensor, handle: DispatchHandle, out: Optional[torch.Tensor] + def _combine_overlap_impl( + 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( + self._validate_overlap_combine_inputs(expert_output, handle) + context = handle._context + combined_x, _combined_w = self._combine_token_major( expert_output, context.recv_topk_weights, context.send_head, @@ -496,13 +456,13 @@ def _combine( return out return combined_x - def _validate_dispatch_inputs(self, input, topk_ids, weights, quant) -> None: + def _validate_overlap_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") + raise NotImplementedError("overlap 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") + raise ValueError("overlap 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: @@ -521,8 +481,8 @@ def _validate_dispatch_inputs(self, input, topk_ids, weights, quant) -> None: 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): + def _validate_overlap_combine_inputs(self, expert_output, handle) -> None: + if not isinstance(handle, DispatchHandle) or not isinstance(handle._context, _TokenMajorOverlapCombineContext): 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") diff --git a/python/mscclpp/ep/runtime.py b/python/mscclpp/ep/runtime.py new file mode 100644 index 00000000..b45b11aa --- /dev/null +++ b/python/mscclpp/ep/runtime.py @@ -0,0 +1,49 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Unified low-level expert-parallel runtime wrapper.""" + +from __future__ import annotations + +from typing import Any + +from ._cpp import DispatchLayout, MoEMode, create_moe_runtime + + +class Runtime: + """Own one C++ runtime configured for latency or overlap algorithms.""" + + def __init__( + self, + comm: Any, + mode: MoEMode, + *, + max_tokens_per_rank: int = 0, + hidden: int = 0, + num_experts: int = 0, + num_topk: int = 0, + max_hidden_bytes: int = 0, + num_sms: int = 20, + output_layout: DispatchLayout = DispatchLayout.EXPERT_MAJOR, + ) -> None: + self.rank: int = comm.my_rank + self.group_size: int = comm.nranks + self.comm = comm + self.cpp_runtime = create_moe_runtime( + comm.communicator, + mode, + max_tokens_per_rank=max_tokens_per_rank, + hidden=hidden, + num_experts=num_experts, + num_topk=num_topk, + max_hidden_bytes=max_hidden_bytes, + num_sms=num_sms, + output_layout=output_layout, + ) + + def is_available(self) -> bool: + """Return whether the selected algorithms are available.""" + return self.cpp_runtime.is_available() + + def is_internode_available(self) -> bool: + """Return whether the selected algorithms support this internode topology.""" + return self.cpp_runtime.is_internode_available() diff --git a/python/mscclpp/ep/types.py b/python/mscclpp/ep/types.py index b0f152c6..1fafa9bb 100644 --- a/python/mscclpp/ep/types.py +++ b/python/mscclpp/ep/types.py @@ -6,7 +6,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, List, Optional, Union +from typing import List, Optional, Union import torch import mscclpp @@ -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,9 +49,9 @@ 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 @@ -63,7 +63,7 @@ class MoECommunicatorConfig: low_latency_combine_mode: CombineMode = CombineMode.RANK_LOCAL_REDUCE enable_overlap: bool = False - # HT-only buffer/launch tuning (advanced) + # Overlap receive-pool tuning (advanced) expert_alignment: int = 1 @@ -103,11 +103,11 @@ class DispatchOutput: weights: 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 _TokenMajorOverlapCombineContext: + """Combine context for token-major overlap output.""" recv_topk_weights: Optional[torch.Tensor] send_head: torch.Tensor -CombineContext = Union[ - ExpertMajorCombineContext, - RankMajorCombineContext, - HighThroughputCombineContext, +_CombineContext = Union[ + _ExpertMajorCombineContext, + _RankMajorCombineContext, + _TokenMajorOverlapCombineContext, ] @@ -150,24 +150,10 @@ 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 # Optional async/overlap configuration. diff --git a/src/ext/ep/CMakeLists.txt b/src/ext/ep/CMakeLists.txt index e357c606..907b53d1 100644 --- a/src/ext/ep/CMakeLists.txt +++ b/src/ext/ep/CMakeLists.txt @@ -4,13 +4,9 @@ # 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). +# One `MoERuntime` conditionally owns the resources required by the selected +# latency or overlap algorithms. All kernels consume the same persistent +# `DeviceContext*` through a torch-free raw-pointer nanobind API. find_package(Python 3.10 COMPONENTS Interpreter Development.Module REQUIRED) include(FetchContent) @@ -58,15 +54,16 @@ endif() set(EP_SOURCES moe_runtime.cc - ll_runtime.cc + runtime/fixed_buffer.cc + runtime/recv_pool.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 + dispatch/latency/expert_major.cu + dispatch/latency/rank_major.cu + combine/latency/rank_local_reduce.cu + combine/latency/direct_send.cu + dispatch/overlap/token_major_prepare.cu + dispatch/overlap/token_major.cu + combine/overlap/token_major_reduce.cu ) # Build as a Python extension module (shared object with Python ABI suffix). @@ -78,7 +75,6 @@ endif() 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 diff --git a/src/ext/ep/README.md b/src/ext/ep/README.md index 930f1061..da92e065 100644 --- a/src/ext/ep/README.md +++ b/src/ext/ep/README.md @@ -1,43 +1,43 @@ # 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 resource 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. +- **`OVERLAP`** algorithms use a bounded SM budget so communication can run + concurrently with expert compute. + +`LOW_LATENCY` and `HIGH_THROUGHPUT` remain compatibility aliases. +Mode-specific buffers are allocated conditionally; selecting one family does +not allocate the other family's resources. ## 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 | +| Overlap dispatch/combine | Supports 2, 4, 8, or 16 ranks in one GPU IPC/NVL fabric domain | +| Overlap RDMA/IB fallback | Not supported | +| Python frontend | `mscclpp.ep.MoECommunicator` selects latency or overlap 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 resource plan 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 +46,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 +### Overlap algorithms -HT follows the same direct-mapping resource model: +The overlap resource plan follows the same direct-mapping model: 1. Python passes the existing `mscclpp::Communicator` into - `MoERuntime` with `MoEMode::HIGH_THROUGHPUT`. + `MoERuntime` with `MoEMode::OVERLAP`. 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 +64,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 overlap 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 overlap 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 +77,20 @@ dependent: Cached dispatch reuses the previous receive count and prefix matrices. -## HT data path +## Overlap data path -HT has one direct path. Every dispatch block writes hidden rows and routing +The overlap 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. -The persistent HT configuration contains only: +The persistent overlap configuration contains only: | Field | Meaning | |---|---| -| `num_sms` | Maximum HT communication block budget | +| `num_sms` | Maximum overlap communication block budget | ## Build @@ -126,24 +126,38 @@ Available CMake options: 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 +├── runtime/ +│ ├── resources.hpp +│ ├── fixed_buffer.cc +│ └── recv_pool.cc +├── common/ +│ ├── fixed_buffer.cuh +│ ├── recv_pool.cuh +│ └── overlap_barrier.cuh +├── dispatch/ +│ ├── latency/ +│ │ ├── common.cuh +│ │ ├── expert_major.cu +│ │ └── rank_major.cu +│ └── overlap/ +│ ├── token_major_prepare.cu +│ └── token_major.cu +├── combine/ +│ ├── latency/ +│ │ ├── common.cuh +│ │ ├── rank_local_reduce.cu +│ │ └── direct_send.cu +│ └── overlap/ +│ └── token_major_reduce.cu ├── include/ -└── low_latency/ - ├── config.cuh - ├── dispatch.cu - └── combine.cu +│ ├── api.cuh +│ └── device_context.cuh +└── config.hpp ``` ## Validation -Build the extension, then run the single-node HT test: +Build the extension, then run the single-node overlap test: ```bash HWLOC_COMPONENTS=-gl \ @@ -152,7 +166,7 @@ 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 \ diff --git a/src/ext/ep/bindings.cpp b/src/ext/ep/bindings.cpp index 77b555c8..247be94b 100644 --- a/src/ext/ep/bindings.cpp +++ b/src/ext/ep/bindings.cpp @@ -20,13 +20,10 @@ #include #include -#include #include "api.cuh" +#include "common/recv_pool.cuh" #include "config.hpp" -#include "high-throughput/config.cuh" -#include "ht_runtime.hpp" -#include "ll_runtime.hpp" #include "moe_runtime.hpp" namespace nb = nanobind; @@ -37,24 +34,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) { @@ -64,23 +43,25 @@ NB_MODULE(mscclpp_ep_cpp, m) { nb::enum_(m, "MoEMode") .value("LOW_LATENCY", mscclpp::ep::MoEMode::LOW_LATENCY) - .value("HIGH_THROUGHPUT", mscclpp::ep::MoEMode::HIGH_THROUGHPUT); + .value("HIGH_THROUGHPUT", mscclpp::ep::MoEMode::HIGH_THROUGHPUT) + .value("LATENCY", mscclpp::ep::MoEMode::LATENCY) + .value("OVERLAP", mscclpp::ep::MoEMode::OVERLAP); 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::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); - nb::class_(m, "Config") + nb::class_(m, "Config") .def(nb::init(), nb::arg("num_sms") = 20) - .def_ro("num_sms", &mscclpp::ep::high_throughput::Config::numSms_); + .def_ro("num_sms", &mscclpp::ep::RecvPoolConfig::numSms_); 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, @@ -93,41 +74,30 @@ 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()); + return reinterpret_cast(self.outputTopkWeightsBuffer()); }) + .def("dispatch_output_buffer_ptr", + [](const mscclpp::ep::MoERuntime& self) { return reinterpret_cast(self.dispatchOutputBuffer()); }) .def("expert_output_buffer_ptr", - [](const mscclpp::ep::MoERuntime& self) { - return reinterpret_cast( - narrow(self, "LOW_LATENCY").expertOutputBuffer()); - }) + [](const mscclpp::ep::MoERuntime& self) { return reinterpret_cast(self.expertOutputBuffer()); }) .def( - "ll_dispatch", + "dispatch_latency", [](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.dispatchLatency( + 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)); }, 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 +106,83 @@ 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_latency", [](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.combineLatency(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)); }, 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", + "prepare_token_major_overlap", [](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.prepareTokenMajorOverlap(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", + .def("get_token_major_overlap_num_channels", [](const mscclpp::ep::MoERuntime& self, int x_element_size) { - return narrow(self, "HIGH_THROUGHPUT") - .getDispatchNumChannels(x_element_size); + return self.getTokenMajorOverlapNumChannels(x_element_size); }) - .def("ht_resolve_recv_x_buffer", + .def("resolve_token_major_overlap_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.resolveTokenMajorOverlapRecvBuffer(num_tokens, num_recv_tokens, hidden, x_element_size)); }) .def( - "ht_notify_dispatch", + "notify_token_major_overlap", [](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.notifyTokenMajorOverlap(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_token_major_overlap", [](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.dispatchTokenMajorOverlap( + 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)); }, 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 +191,14 @@ 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_token_major_overlap", [](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.combineTokenMajorOverlap(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)); }, 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/latency/common.cuh similarity index 80% rename from src/ext/ep/low_latency/combine.cu rename to src/ext/ep/combine/latency/common.cuh index 3ec0f3c9..d9a93903 100644 --- a/src/ext/ep/low_latency/combine.cu +++ b/src/ext/ep/combine/latency/common.cuh @@ -1,18 +1,20 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +#pragma once #include #include #include "api.cuh" -#include "config.cuh" +#include "common/fixed_buffer.cuh" #include "device_helpers.cuh" #include "exception.cuh" namespace mscclpp { namespace ep { -namespace low_latency { +namespace combine { namespace detail { +using namespace ::mscclpp::ep::detail; constexpr int CombineNWarps = 32; constexpr int CombineNThreads = CombineNWarps * WARP_SIZE; @@ -39,7 +41,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 +55,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(); } @@ -587,24 +589,24 @@ 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 combineLatencyBody(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 combineEpoch = *workspaceView.dispatchEpoch_; @@ -620,7 +622,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); @@ -633,7 +635,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 { @@ -642,153 +644,156 @@ __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, + const Workload& workload, void* recvBuffer, void* dispatchRecvBuffer, + const DeviceContext& context, const DeviceContext* deviceContext, 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, + deviceContext); CUDA_CHECK(cudaGetLastError()); } -template