Refactor EP runtime architecture

include/mscclpp/ext/ep/moe_runtime.hpp

@@ -0,0 +1,136 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+#ifndef MSCCLPP_EXT_EP_MOE_RUNTIME_HPP_
+#define MSCCLPP_EXT_EP_MOE_RUNTIME_HPP_
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+#include <memory>
+#include <mscclpp/core.hpp>
+#include <mscclpp/ext/ep/types.hpp>
+
+namespace mscclpp {
+namespace ep {
+struct LatencyContext;
+struct ThroughputContext;
+/// Unified host runtime for expert-parallel dispatch and combine.
+///
+/// One runtime owns the communication buffers and synchronization state for the
+/// selected mode. LATENCY uses fixed-capacity expert-major or rank-major
+/// layouts. THROUGHPUT uses a dynamically sized token-major receive pool.
+/// Operations are asynchronous with respect to the host and execute on the
+/// CUDA stream supplied by each request.
+class MoERuntime {
+ public:
+  /// Construct a runtime for the selected mode and topology.
+  ///
+  /// Only resources required by @p mode are allocated.
+  /// @param communicator Initialized MSCCL++ communicator.
+  /// @param mode Runtime algorithm family.
+  /// @param maxTokensPerRank Fixed latency-mode token capacity.
+  /// @param hidden Hidden dimension for latency-mode buffers.
+  /// @param numExperts Global expert count.
+  /// @param numTopk Number of routed experts per token.
+  /// @param maxHiddenBytes Maximum throughput-mode bytes per token row.
+  /// @param numBlocks Communication block budget.
+  /// @param outputLayout Latency-mode dispatch output layout.
+  MoERuntime(mscclpp::Communicator& communicator, MoEMode mode, int maxTokensPerRank, int hidden, int numExperts,
+             int numTopk, int64_t maxHiddenBytes, int numBlocks,
+             DispatchLayout outputLayout = DispatchLayout::EXPERT_MAJOR);
+  ~MoERuntime() noexcept(false);
+
+  MoERuntime(const MoERuntime&) = delete;
+  MoERuntime& operator=(const MoERuntime&) = delete;
+
+  /// Return the configured runtime mode.
+  MoEMode mode() const { return mode_; }
+  /// Return whether the selected mode supports the detected topology.
+  bool isAvailable() const { return available_; }
+  /// Return whether the runtime is available across more than one node.
+  bool isInternodeAvailable() const { return available_ && numRanks_ > numNvlRanks_; }
+
+  /// Return the local rank.
+  int rank() const { return rank_; }
+  /// Return the global rank count.
+  int numRanks() const { return numRanks_; }
+  /// Return the NVLink-local rank count.
+  int numNvlRanks() const { return numNvlRanks_; }
+  /// Return the rank count in one CUDA IPC domain.
+  int numRanksPerIpcDomain() const { return numRanksPerIpcDomain_; }
+
+  /// Return the runtime-owned rank-major top-k ID buffer.
+  void* outputTopkIdsBuffer() const;
+  /// Return the runtime-owned rank-major top-k weight buffer.
+  void* outputTopkWeightsBuffer() const;
+  /// Return the runtime-owned dispatch output buffer.
+  void* dispatchOutputBuffer() const;
+  /// Return the runtime-owned rank-major combine input buffer.
+  void* combineInputBuffer() const;
+
+  /// Dispatch tokens using the configured runtime mode.
+  ///
+  /// @p request must contain the request type matching mode(): a
+  /// LatencyDispatchRequest for LATENCY or a ThroughputDispatchRequest for
+  /// THROUGHPUT. Output buffers remain owned by the caller unless obtained
+  /// through a runtime buffer accessor.
+  /// @param request Dispatch inputs, outputs, dimensions, and CUDA stream.
+  /// @throws std::invalid_argument If the request type does not match mode().
+  void dispatch(const DispatchRequest& request);
+
+  /// Combine expert outputs using the configured runtime mode.
+  ///
+  /// A combine request must follow its matching dispatch so the runtime can
+  /// reuse routing metadata and synchronization epochs. @p request must contain
+  /// a LatencyCombineRequest for LATENCY or a ThroughputCombineRequest for
+  /// THROUGHPUT.
+  /// @param request Combine inputs, outputs, dimensions, and CUDA stream.
+  /// @throws std::invalid_argument If the request type does not match mode().
+  void combine(const CombineRequest& request);
+
+  /// Build throughput-mode token routing metadata.
+  ///
+  /// Computes per-rank counts, per-expert counts, and token-to-rank membership
+  /// on @p stream without moving token payloads.
+  void tokenMajorPrepare(int* numTokensPerRank, int* numTokensPerExpert, bool* isTokenInRank, const int64_t* topkIdx,
+                         int numTokens, int numTopk, int numExperts, cudaStream_t stream);
+  /// Return the throughput-mode communication channel count.
+  int tokenMajorNumChannels(int xElementSize) const;
+  /// Resolve the runtime-owned throughput receive buffer.
+  void* tokenMajorResolveRecvBuffer(int numTokens, int numRecvTokens, int hidden, int xElementSize) const;
+  /// Exchange throughput routing counts and return the receive-token count.
+  ///
+  /// This host-synchronizing metadata phase must precede throughput dispatch
+  /// when cached routing metadata is unavailable.
+  int tokenMajorNotify(int* rankPrefixMatrix, int* channelPrefixMatrix, int* numRecvTokensPerExpert,
+                       const int* numTokensPerRank, const int* numTokensPerExpert, const bool* isTokenInRank,
+                       int numTokens, int numExperts, int xElementSize, int expertAlignment, cudaStream_t stream);
+
+ private:
+  void requireMode(MoEMode expected) const;
+  void launchLatencyDispatch(const LatencyDispatchRequest& request);
+  void launchThroughputDispatch(const ThroughputDispatchRequest& request);
+  void launchLatencyCombine(const LatencyCombineRequest& request);
+  void launchThroughputCombine(const ThroughputCombineRequest& request);
+
+  std::shared_ptr<mscclpp::Bootstrap> bootstrap_;
+  MoEMode mode_;
+  int rank_;
+  int numRanks_;
+  int numNvlRanks_;
+  int numRanksPerIpcDomain_;
+  bool available_ = false;
+
+  std::unique_ptr<LatencyContext> latencyContext_;
+  std::unique_ptr<ThroughputContext> throughputContext_;
+};
+
+/// Create the unified MoE runtime selected by @p mode.
+std::shared_ptr<MoERuntime> createMoERuntime(mscclpp::Communicator& communicator, MoEMode mode, int maxTokensPerRank,
+                                             int hidden, int numExperts, int numTopk, int64_t maxHiddenBytes,
+                                             int numBlocks, DispatchLayout outputLayout = DispatchLayout::EXPERT_MAJOR);
+
+}  // namespace ep
+}  // namespace mscclpp
+
+#endif  // MSCCLPP_EXT_EP_MOE_RUNTIME_HPP_

include/mscclpp/ext/ep/types.hpp

@@ -0,0 +1,230 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+#ifndef MSCCLPP_EXT_EP_TYPES_HPP_
+#define MSCCLPP_EXT_EP_TYPES_HPP_
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+#include <utility>
+#include <variant>
+
+namespace mscclpp {
+namespace ep {
+
+class MoERuntime;
+
+/// Expert-parallel runtime mode.
+enum class MoEMode {
+  /// Algorithms optimized for minimum standalone latency.
+  LATENCY,
+  /// Resource-bounded algorithms optimized for end-to-end throughput.
+  THROUGHPUT
+};
+
+/// Logical dispatch output layout.
+enum class DispatchLayout {
+  /// Rows grouped by local expert.
+  EXPERT_MAJOR,
+  /// Dynamically sized token-major rows used by throughput mode.
+  TOKEN_MAJOR,
+  /// Fixed-stride rows grouped by source rank.
+  RANK_MAJOR
+};
+
+/// Combine algorithm.
+enum class CombineMode {
+  /// Reduce local expert rows before sending one partial per rank and token.
+  RANK_LOCAL_REDUCE,
+  /// Send every expert row and reduce all contributions on the source rank.
+  DIRECT_SEND
+};
+
+/// Dispatch payload data format.
+enum class DispatchDataType {
+  /// Unquantized BF16 payload.
+  BF16,
+  /// FP8 E4M3 payload with one floating-point scale per 128 hidden elements.
+  FP8_E4M3
+};
+
+/// Arguments for latency-mode dispatch.
+struct LatencyDispatchRequest {
+  /// Dispatch output buffer.
+  void* output;
+  /// Optional dispatch scale output.
+  void* outputScales;
+  /// Optional source-token metadata output.
+  int* outputSrcInfo;
+  /// Optional dispatched top-k expert IDs.
+  int* outputTopkIdx;
+  /// Optional dispatched top-k weights.
+  float* outputTopkWeights;
+  /// Optional packed layout metadata.
+  int64_t* outputLayoutRange;
+  /// Per-expert or per-rank output counts.
+  int* outputCount;
+  /// Input token payload.
+  const void* input;
+  /// Input top-k expert IDs.
+  const int64_t* topkIdx;
+  /// Optional input top-k weights.
+  const float* topkWeights;
+  /// Number of input tokens.
+  int numTokens;
+  /// Hidden dimension.
+  int hidden;
+  /// Number of routed experts per token.
+  int numTopk;
+  /// Active per-rank token capacity.
+  int maxTokensPerRank;
+  /// Global expert count.
+  int numExperts;
+  /// Expert ID used for invalid rank-major entries.
+  int invalidTokenExpertId;
+  /// Requested dispatch output layout.
+  DispatchLayout dispatchLayout;
+  /// Requested dispatch payload format.
+  DispatchDataType dispatchDataType;
+  /// Dispatch grid block count.
+  int numBlocks;
+  /// CUDA stream used for the operation.
+  cudaStream_t stream;
+};
+
+/// Arguments for throughput-mode dispatch.
+struct ThroughputDispatchRequest {
+  /// Token receive buffer.
+  void* recvX;
+  /// Optional received scale output.
+  float* recvXScales;
+  /// Optional received top-k expert IDs.
+  int64_t* recvTopkIdx;
+  /// Optional received top-k weights.
+  float* recvTopkWeights;
+  /// Per-token routing state consumed by combine.
+  int* sendHead;
+  /// Input token payload.
+  const void* input;
+  /// Optional input scales.
+  const float* inputScales;
+  /// Optional input top-k expert IDs.
+  const int64_t* topkIdx;
+  /// Optional input top-k weights.
+  const float* topkWeights;
+  /// Token-to-destination-rank membership.
+  const bool* isTokenInRank;
+  /// Per-source-rank token prefixes.
+  const int* rankPrefixMatrix;
+  /// Per-channel token prefixes.
+  const int* channelPrefixMatrix;
+  /// Number of input tokens.
+  int numTokens;
+  /// Hidden dimension.
+  int hidden;
+  /// Number of routed experts per token.
+  int numTopk;
+  /// Number of scales per token.
+  int numScales;
+  /// Global expert count, or zero when cached metadata is reused.
+  int numExperts;
+  /// Input element size in bytes.
+  int inputElementSize;
+  /// Number of received tokens.
+  int numRecvTokens;
+  /// Whether cached routing metadata is reused.
+  bool cachedMode;
+  /// CUDA stream used for the operation.
+  cudaStream_t stream;
+};
+
+/// Mode-specific dispatch request.
+struct DispatchRequest {
+  /// Construct a latency dispatch request.
+  explicit DispatchRequest(LatencyDispatchRequest request) : value_(std::move(request)) {}
+  /// Construct a throughput dispatch request.
+  explicit DispatchRequest(ThroughputDispatchRequest request) : value_(std::move(request)) {}
+
+ private:
+  friend class MoERuntime;
+  std::variant<LatencyDispatchRequest, ThroughputDispatchRequest> value_;
+};
+
+/// Arguments for latency-mode combine.
+struct LatencyCombineRequest {
+  /// Combined token output.
+  void* output;
+  /// Local expert output.
+  const void* input;
+  /// Input top-k expert IDs.
+  const int64_t* topkIdx;
+  /// Optional input top-k weights.
+  const float* topkWeights;
+  /// Optional source-token metadata.
+  const int* srcInfo;
+  /// Optional packed layout metadata.
+  const int64_t* layoutRange;
+  /// Number of output tokens.
+  int numTokens;
+  /// Hidden dimension.
+  int hidden;
+  /// Number of routed experts per token.
+  int numTopk;
+  /// Active per-rank token capacity.
+  int maxTokensPerRank;
+  /// Global expert count.
+  int numExperts;
+  /// Dispatch input layout.
+  DispatchLayout dispatchLayout;
+  /// Dispatch payload format.
+  DispatchDataType dispatchDataType;
+  /// Combine algorithm.
+  CombineMode combineMode;
+  /// Combine worker block count.
+  int numBlocks;
+  /// CUDA stream used for the operation.
+  cudaStream_t stream;
+};
+
+/// Arguments for throughput-mode combine.
+struct ThroughputCombineRequest {
+  /// Combined token output.
+  void* output;
+  /// Optional combined top-k weights.
+  float* outputTopkWeights;
+  /// Local expert output.
+  const void* input;
+  /// Optional local top-k weights.
+  const float* topkWeights;
+  /// Routing state returned by throughput dispatch.
+  const int* sendHead;
+  /// Number of local expert-output rows.
+  int numInputTokens;
+  /// Number of combined output tokens.
+  int numOutputTokens;
+  /// Hidden dimension.
+  int hidden;
+  /// Number of routed experts per token.
+  int numTopk;
+  /// Input element size in bytes.
+  int inputElementSize;
+  /// CUDA stream used for the operation.
+  cudaStream_t stream;
+};
+
+/// Mode-specific combine request.
+struct CombineRequest {
+  /// Construct a latency combine request.
+  explicit CombineRequest(LatencyCombineRequest request) : value_(std::move(request)) {}
+  /// Construct a throughput combine request.
+  explicit CombineRequest(ThroughputCombineRequest request) : value_(std::move(request)) {}
+
+ private:
+  friend class MoERuntime;
+  std::variant<LatencyCombineRequest, ThroughputCombineRequest> value_;
+};
+
+}  // namespace ep
+}  // namespace mscclpp
+
+#endif  // MSCCLPP_EXT_EP_TYPES_HPP_