Initial EV3 Support

include/mscclpp/ext/ep/moe_runtime.hpp

@@ -91,7 +91,11 @@ class MoERuntime {
   /// BF16 or E4M3 elements according to the dispatch request. The active
   /// maxTokensPerRank, not the runtime capacity, determines the rank stride.
   void* dispatchOutputBuffer() const;
-  /// Return the runtime-owned BF16 combine input buffer.
+  /// Return the runtime-owned combine input buffer.
+  ///
+  /// Latency RANK_MAJOR uses the combine format from LatencyDispatchRequest::quant,
+  /// which defaults to BF16 and may be FP8_E4M3_STATIC with matching static FP8
+  /// dispatch. Other layouts and throughput use BF16.
   void* combineInputBuffer() const;
   /// Collectively prepare throughput routing without moving token payloads.
   ///

include/mscclpp/ext/ep/types.hpp

@@ -39,8 +39,8 @@ enum class CombineMode {
   /// Reduce all contributions on the source rank.
   ///
   /// Expert-major sends every expert row. Rank-major consumes one unweighted
-  /// BF16 route row per top-k lane, applies its routing weight, and performs
-  /// the top-k reduction in combine.
+  /// route row per top-k lane, applies its routing weight, and performs the
+  /// top-k reduction in combine.
   DIRECT_SEND
 };
 
@@ -58,7 +58,36 @@ enum class DispatchDataType {
   /// Scales are powers of two rounded up to cover each block's maximum absolute
   /// value with E4M3's finite range, with scale bytes clamped to [1, 254].
   /// An all-zero block uses scale byte 127.
-  MXFP8_E4M3
+  MXFP8_E4M3,
+  /// FP8 E4M3 payload with one caller-provided scalar scale.
+  ///
+  /// Supported for RANK_MAJOR in LATENCY mode only. Unlike FP8_E4M3, this
+  /// format has no per-block scale sideband.
+  FP8_E4M3_STATIC
+};
+
+/// Quantization configuration for latency dispatch and rank-major combine.
+struct QuantConfig {
+  /// Input representation for latency rank-major static FP8 dispatch.
+  enum class InputMode {
+    /// Input is already quantized as E4M3(real * dispatchScale).
+    PREQUANTIZED,
+    /// Input is BF16 and is quantized once inside dispatch before transport.
+    QUANTIZE
+  };
+
+  /// Dispatch payload format.
+  DispatchDataType dispatchDataType = DispatchDataType::BF16;
+  /// Static FP8 dispatch multiplier. Dispatch stores E4M3(real * dispatchScale).
+  /// Must be finite and positive for FP8_E4M3_STATIC, and one otherwise.
+  float dispatchScale = 1.0f;
+  /// Input representation for latency rank-major static FP8 dispatch. Must remain
+  /// PREQUANTIZED for every other layout and format.
+  InputMode inputMode = InputMode::PREQUANTIZED;
+  /// Rank-major combine input format. BF16 is the default; FP8_E4M3_STATIC is
+  /// supported only with static FP8 dispatch and uses dispatchScale for
+  /// dequantization.
+  DispatchDataType combineDataType = DispatchDataType::BF16;
 };
 
 /// Arguments for throughput-mode routing preparation.
@@ -136,7 +165,11 @@ struct LatencyDispatchRequest {
   int64_t* outputLayoutRange;
   /// Per-expert or per-rank output counts.
   int* outputCount;
-  /// BF16 input token payload, 16-byte aligned, including when dispatch quantizes the output.
+  /// Input token payload, 16-byte aligned.
+  ///
+  /// BF16, MXFP8_E4M3, and expert-major FP8_E4M3 take BF16 input.
+  /// FP8_E4M3_STATIC uses quant.inputMode to select prequantized E4M3 or BF16
+  /// input.
   const void* input;
   /// Input top-k expert IDs.
   const int64_t* topkIdx;
@@ -156,12 +189,21 @@ struct LatencyDispatchRequest {
   /// This is an output sentinel; input topkIdx entries must still be valid
   /// global expert IDs or negative values for dropped routes.
   int invalidTokenExpertId = -1;
-  /// Requested dispatch payload format. Must match across participating ranks.
-  DispatchDataType dispatchDataType;
+  /// Dispatch quantization and rank-major combine datatype settings.
+  ///
+  /// FP8_E4M3_STATIC combine is supported only with matching static FP8 dispatch
+  /// and reuses the dispatch scale for dequantization; all other rank-major
+  /// combines use BF16.
+  QuantConfig quant;
   /// Dispatch grid block count. Zero uses the default number of blocks.
   int numBlocks = 0;
   /// CUDA stream used for the operation.
   cudaStream_t stream;
+  /// Emit destination-rank-local expert IDs for latency RANK_MAJOR output.
+  ///
+  /// When false, valid outputTopkIdx entries remain global expert IDs. When
+  /// true, valid entries are in [0, numExperts / numRanks).
+  bool emitLocalExpertIds = false;
 };
 
 /// Arguments for throughput-mode dispatch.
@@ -242,12 +284,14 @@ class DispatchHandle {
 struct LatencyCombineRequest {
   /// Combined token output.
   void* output;
-  /// Local BF16 results produced by expert computation, regardless of the
-  /// dispatch payload format.
+  /// Local results produced by expert computation.
   ///
   /// Rank-major input must be the runtime-owned combineInputBuffer(), with
   /// [numRanks, active capacity, hidden] elements and an extra top-k dimension
-  /// before hidden for DIRECT_SEND.
+  /// before hidden for DIRECT_SEND. Its element format is the matching dispatch
+  /// request's quant.combineDataType: BF16 by default, or FP8_E4M3_STATIC with
+  /// matching static FP8 dispatch. Static FP8 input is dequantized with the
+  /// dispatch scale.
   const void* input;
   /// Handle returned by the matching dispatch.
   DispatchHandle handle;

include/mscclpp/gpu_data_types.hpp

@@ -400,6 +400,18 @@ MSCCLPP_DEVICE_INLINE f32x2 operator*(const f32x2& a, const f32x2& b) {
 #endif
 }
 
+/// Performs two element-wise FP32 fused multiply-add operations.
+MSCCLPP_DEVICE_INLINE f32x2 fma(const f32x2& a, const f32x2& b, const f32x2& c) {
+#if defined(MSCCLPP_DEVICE_CUDA) && (__CUDA_ARCH__ >= 1000)
+  return __ffma2_rn(a.storage, b.storage, c.storage);
+#else
+  f32x2 result;
+  result.data[0] = fmaf(a.data[0], b.data[0], c.data[0]);
+  result.data[1] = fmaf(a.data[1], b.data[1], c.data[1]);
+  return result;
+#endif
+}
+
 template <bool UseClip = true>
 MSCCLPP_DEVICE_INLINE f16x2 operator+(const f16x2& a, const f16x2& b) {
   __half2 result;

python/mscclpp/ep/README.md

@@ -104,7 +104,7 @@ not resize those buffers or read counts back to the host.
 | Mode / layout | `tokens` shape | Count metadata | Physical row grouping | Expert-output contract |
 | --- | --- | --- | --- | --- |
 | LATENCY / EXPERT_MAJOR (default) | `[L, R*A, H]` | `layout.num_tokens_per_expert` | One fixed slice per local expert | BF16 per-expert results; native combine applies the original routing weights |
-| LATENCY / RANK_MAJOR | `[R, A, H]` | `layout.num_tokens_per_rank` | One fixed slice per source rank | BF16 weighted local sums, or unweighted per-topk rows for DIRECT_SEND |
+| LATENCY / RANK_MAJOR | `[R, A, H]` | `layout.num_tokens_per_rank` | One fixed slice per source rank | BF16 by default, or explicitly configured static-scale FP8 weighted local sums; unweighted per-topk rows for DIRECT_SEND |
 | THROUGHPUT / TOKEN_MAJOR (default) | `[R*A, H]` | Optional `layout.num_tokens_per_expert` (prepare `output_count[L]`) | Compact source-rank segments, preserving token order | BF16 already-weighted local expert sums |
 | THROUGHPUT / RANK_MAJOR | `[R, A, H]` | Optional `layout.num_tokens_per_rank` (prepare `output_count[R]`) | One fixed slice per source rank | BF16 already-weighted local expert sums |
 
@@ -138,23 +138,28 @@ expert IDs in `[0, num_experts)` or negative values for dropped routes.
 means one per valid route. Routing values are not read or validated on the host.
 
 RANK_MAJOR and throughput outputs include CUDA int32 `topk_ids` and FP32
-`weights` with shape `tokens.shape[:-1] + (K,)`. Latency IDs are global, with
-`invalid_token_expert_id` (default `num_experts`) marking dropped or non-local
-entries. Throughput IDs are local to the destination rank, with `-1` marking
+`weights` with shape `tokens.shape[:-1] + (K,)`. Latency IDs are global by
+default, with `invalid_token_expert_id` (default `num_experts`) marking dropped
+or non-local entries. Latency RANK_MAJOR can set `emit_local_expert_ids=True`
+to emit valid IDs in `[0, num_local_experts)` instead; its default invalid ID is
+`-1`. Throughput IDs are local to the destination rank, with `-1` marking
 invalid routes in valid rows.
 
 Latency RANK_MAJOR must dispatch into the runtime-owned buffer. It also exposes
-the required BF16 `combine_input_buffer`. With `CombineMode.DIRECT_SEND`, that
-buffer and expert outputs have shape `[R, A, K, H]`: write unweighted BF16
-per-topk results for entries whose returned `topk_ids` identify a local expert.
-Combine applies the original routing weights during FP32 reduction. Otherwise
-write already-weighted local expert sums. Before `combine`, place these results
-in the runtime-owned buffer, either by computing there or by explicitly copying
-them there. Latency rank-major combine rejects an external expert-output tensor
-rather than performing a hidden staging copy. Both throughput layouts have the
-same exact-buffer requirement: pass `combine_input_buffer`, after writing or
-copying expert results into it. Throughput dispatch payload, top-k IDs, weights,
-and FP8 scales are also runtime-owned views.
+the required `combine_input_buffer`. Its dtype is BF16 by default and can be
+configured as E4M3 with `FP8_E4M3_STATIC` dispatch.
+With `CombineMode.DIRECT_SEND`, that buffer and expert outputs have shape
+`[R, A, K, H]`: write unweighted per-topk results for entries whose returned
+`topk_ids` identify a local expert. Combine applies the original routing weights
+during FP32 reduction. Otherwise write already-weighted, locally reduced rank
+partials with shape `[R, A, H]`; combine sums one partial per distinct
+destination rank without applying weights again. Before `combine`, place these
+results in the runtime-owned buffer, either by computing there or by explicitly
+copying them there. Latency rank-major combine rejects an external expert-output
+tensor rather than performing a hidden staging copy. Both throughput layouts
+have the same exact-buffer requirement: pass `combine_input_buffer`, after
+writing or copying BF16 expert results into it. Throughput dispatch payload,
+top-k IDs, weights, and FP8 scales are also runtime-owned views.
 
 `combine` returns token outputs only and has no `output_topk_weights` argument.
 
@@ -166,12 +171,28 @@ with latency `DIRECT_SEND` (either layout). The default is `True` for
 ## Formats and buffers
 
 * BF16 dispatch uses BF16 inputs/outputs and no block scales.
-* FP8 E4M3 uses one FP32 scale per 128 hidden elements. Latency EXPERT_MAJOR
+* Latency EXPERT_MAJOR FP8 uses one FP32 scale per 128 hidden elements and
   quantizes **BF16 input** on the fly with
   `QuantConfig(format=DispatchDataType.FP8_E4M3)` and does not accept
   precomputed input scales. Returned FP32 scales have logical shape
   `[L, R*A, H//128]`, transposed from physical `[L, H//128, R*A]`; they are not
   contiguous in logical order.
+* Latency RANK_MAJOR static FP8 uses a finite positive scalar `S` and has two
+  explicit input modes. The default accepts contiguous `torch.float8_e4m3fn`
+  input with
+  `QuantConfig(format=DispatchDataType.FP8_E4M3_STATIC, scale=S)` and transports it
+  byte-for-byte. To fuse quantization into dispatch, pass BF16 input with
+  `QuantConfig(format=DispatchDataType.FP8_E4M3_STATIC, scale=S,
+  input_mode=QuantConfig.InputMode.QUANTIZE)`. The fused path computes
+  `q = E4M3(real * S)` once per token before sending the same FP8 row to each
+  destination. Neither mode has a block-scale sideband.
+* Latency RANK_MAJOR combine input is BF16 by default. This means MM2 has
+  already produced/dequantized BF16 output. Static-FP8 dispatch may instead set
+  `combine_format=DispatchDataType.FP8_E4M3_STATIC`; MM2 writes
+  `q = E4M3(real * S)` to `combine_input_buffer`, and combine computes
+  `FP32(q) * (1 / S)` with the existing dispatch scale. Combine never
+  quantizes; it applies FP32 route weights only for `DIRECT_SEND`, accumulates
+  in FP32, and writes BF16.
 * Throughput FP8 requires `torch.float8_e4m3fn` input and
   `QuantConfig(format=DispatchDataType.FP8_E4M3, block_scales=scales)`, where
   `scales` is contiguous CUDA FP32 `[num_input_tokens, H//128]`. Output scales
@@ -180,11 +201,14 @@ with latency `DIRECT_SEND` (either layout). The default is `True` for
   `QuantConfig(format=DispatchDataType.MXFP8_E4M3)`. It takes BF16 input,
   returns E4M3 payload plus contiguous uint8 E8M0 scales with shape
   `[R, A, H//32]`, and requires `H` to be divisible by 32.
-* **Combine always consumes BF16**, including after FP8 or MXFP8 dispatch.
+* Combine always returns BF16. E4M3 expert input is supported only for latency
+  RANK_MAJOR static-FP8 dispatch and must be explicitly configured. MXFP8
+  dispatch and every other combine consume BF16 expert output.
 
 Configure a default `quant` on the communicator or override it per dispatch.
 Explicit `QuantConfig(format=DispatchDataType.BF16)` overrides an FP8 default.
-There are no implicit casts or contiguous copies to accept incompatible input.
+There are no implicit casts or contiguous copies to accept incompatible input;
+static FP8 performs BF16 conversion only when its explicit fused mode is selected.
 Payload pointers must be 16-byte aligned. `output_buffer` must match the active
 dispatch shape, while `combine(out=...)` must be `[num_input_tokens, H]`; both
 require the exact dtype, device, and contiguous layout. Only latency
@@ -241,8 +265,18 @@ Runtime views are reused by later operations, not independent results.
   Captured preparation recomputes routing on each replay, so routing values may
   change while their buffer pointer, shape, and active capacity remain fixed.
   Preserve graph ordering, buffers, and owners through the last replay. Python
-  handle checks run during capture, not replay. Latency dispatch/combine graph
-  replay is unsupported because its host-side epoch is not advanced by replay.
+  handle checks run during capture, not replay.
+* Latency RANK_MAJOR dispatch/expert-write/combine can also be captured and
+  replayed when each graph contains at least two complete dispatch/combine
+  iterations. Each captured iteration has a distinct epoch, so the last count
+  packet from one replay cannot satisfy the first dispatch of the next replay.
+  A graph containing only one iteration is unsupported because graph replay
+  does not advance the host epoch. Combine separately invalidates captured
+  readiness state before its ordered peer rendezvous. Construct, initialize,
+  and warm up the runtime outside capture; input, routing, and weight values may
+  change while their pointers, shapes, token count, and active capacity remain
+  fixed. Every rank must replay collectives in the same order. Latency
+  EXPERT_MAJOR graph replay remains unsupported.
 * CUDA SM90+ and one IPC domain are required; HIP and inter-domain transports
   are unsupported. The implementation accepts 1-64 ranks, not a claim of
   hardware qualification of every 64-rank topology. Expert placement is even

python/mscclpp/ep/_cpp.py

@@ -18,6 +18,7 @@
 CombineMode = _cpp.CombineMode
 DispatchDataType = _cpp.DispatchDataType
 DispatchLayout = _cpp.DispatchLayout
+NativeQuantConfig = _cpp.QuantConfig
 MoEMode = _cpp.MoEMode
 MoERuntime = _cpp.MoERuntime
 PrepareHandle = _cpp.PrepareHandle

python/mscclpp/ep/communicator.py

@@ -131,6 +131,10 @@ def combine_mode(self) -> CombineMode:
     def invalid_token_expert_id(self) -> int:
         return self._config.invalid_token_expert_id
 
+    @property
+    def emit_local_expert_ids(self) -> bool:
+        return self._config.emit_local_expert_ids
+
     def is_available(self) -> bool:
         """Return native support for this runtime's mode and topology."""
         return self._runtime.is_available()
@@ -161,7 +165,7 @@ def get_dispatch_output_buffer(
         subsequent operations; a view does not reserve a separate output slot.
         """
         active = verify_active_capacity(runtime_max_tokens_per_rank, self.max_tokens_per_rank)
-        data_type, _ = resolve_dispatch_format(self._config, quant)
+        data_type, _, _, _, _ = resolve_dispatch_format(self._config, quant)
         with torch.cuda.device(self.device):
             return tensor_from_pointer(
                 self._runtime.dispatch_output_buffer_ptr(),
@@ -277,12 +281,16 @@ def combine(
         stream: Optional[torch.cuda.Stream] = None,
         apply_router_weights: Optional[bool] = None,
     ) -> torch.Tensor:
-        """Combine BF16 expert outputs into ``[local_num_tokens, hidden_size]``.
+        """Combine expert outputs into BF16 ``[local_num_tokens, hidden_size]``.
 
         EXPERT_MAJOR and latency RANK_MAJOR DIRECT_SEND apply routing weights
         natively. The latter expects unweighted per-topk rows; other rank/token
         layouts expect already-weighted local expert sums. Non-EXPERT_MAJOR
         inputs must use ``DispatchOutput.combine_input_buffer``.
+        Latency RANK_MAJOR input dtype follows ``QuantConfig.combine_format``:
+        BF16 by default, or ``FP8_E4M3_STATIC`` with matching static-FP8
+        dispatch. Static E4M3 input is dequantized with the dispatch scale
+        retained by the handle.
         For latency DIRECT_SEND, set ``apply_router_weights=False`` for
         preweighted expert outputs. By default it is enabled only for
         DIRECT_SEND; other modes ignore this option.