[MSCCLPP-EP] Python interface for the C++ expert-parallel runtime by Binyang2014 · Pull Request #907 · microsoft/mscclpp

Binyang Li (Binyang2014) · 2026-09-15T22:17:56Z

Summary

Stacked on #903 (binyli/ep-throughput-cpp). Ports the PyTorch-facing EP interface from feature/ep (232c922d) onto the current C++ latency/throughput runtime, without copying the donor's kernels or legacy receive-pool implementation.

Changes

Add the optional mscclpp.ep package with MoECommunicator, configuration/output types, explicit prepare, opaque preparation/dispatch handles, and tensor-based dispatch/combine.
Add a standalone nanobind module linked to the existing libmscclpp_ep. The native boundary uses raw device pointers and does not link libtorch; PyTorch is an optional ep dependency.
Support latency expert-major/rank-major and throughput token-major/rank-major layouts, BF16 and supported FP8 dispatch paths, metadata/scale outputs, external buffers, and active capacities smaller than the allocation.
Replace the donor's notify/count polling and dispatch caches with GPU-only preparation and current native handles. Throughput returns capacity-sized tensors plus a borrowed GPU receive-row count, rather than synchronizing to construct host-sized outputs.
Validate tensor shape, dtype, device, alignment, stream, and handle ownership before enqueueing. Retain native owners through tensor storage and handles, and record allocator-backed tensors on the caller stream.
Stage external latency rank-major expert results in native code after handle validation, so rejected stale handles cannot overwrite the current receive buffer.
Wire conditional native-module build/install, development copies, relocatable library lookup, and usage/lifetime documentation.

Contract and scope

All work sharing a runtime uses one caller CUDA stream, including expert computation. There are no internal completion events, cross-stream overlap stubs, or automatic autograd integration. Runtime-owned views and handles retain storage, but callers must finish local and peer GPU use before releasing the final owner.

The underlying implementation requires CUDA SM90+ and 1-64 ranks in one IPC domain, with top-k up to 8. Physical 32/64-GPU qualification is not claimed. The donor's previous_handle, overlap configuration, notify matrices, and receive-pool tuning are not carried into this API.

Validation

19 focused Python cases passed on eight H100 GPUs, both from the source tree and from an isolated installed wheel.
Coverage includes both modes/layouts/formats, prepared routing reuse and invalidation, CPU-readback rejection, zero-token inputs, external outputs, sliced-view ownership, cached CUDA graph replay, and stale rank-major combine rejection.
All three existing non-performance MoERuntimeTest C++ cases passed on eight GPUs.
The SM90 wheel builds and resolves its packaged native dependencies without libtorch. Repository formatting completed.

Port the feature/ep PyTorch API to opaque preparation and dispatch handles, GPU-resident receive counts, and the current single-stream runtime. Link standalone nanobind bindings to the existing C++ EP library and package its native dependencies without linking libtorch. Stage external rank-major expert inputs after native handle validation to prevent stale combines from overwriting live runtime data. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

azure-pipelines · 2026-09-15T22:18:06Z

Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run.

Binyang Li (Binyang2014) · 2026-09-18T18:58:04Z

          .dispatchDataType_ = dataType};
 }

+const int* MoERuntime::numRecvTokensBuffer() const {


Why need to expose this API?

The Python TOKEN_MAJOR result has a capacity-sized tensor, so GPU consumers need the actual number of valid received rows without copying a count to the CPU. The existing outputCount is per expert: a token routed to two local experts contributes two expert counts but occupies only one received row, so summing those counts is incorrect.

This accessor only borrows the scalar already allocated in the native workspace. It adds no allocation, count-computation kernel, or host synchronization; Python exposes it as read-only-by-contract layout.num_recv_tokens. I renamed it numRecvTokensDevicePtr() in 312b908 to make the device-pointer semantics clear. Leaving the public-interface design question open for your review.

Use a shared initialization decorator and leave graph capture timing and unsupported buffer aliasing to callers. Preserve allocator stream tracking and only the device guards needed for explicit device selection. Clarify expert-result ownership and name the borrowed receive-count accessor as a device pointer. Cover lazy initialization and device restoration, and expect the native MSCCL++ error type from the updated base. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Binyang Li (Binyang2014) · 2026-09-18T22:48:20Z

+    """Fixed runtime configuration, which must agree across participating ranks.
+
+    Experts are partitioned evenly into contiguous rank-local ranges. A scalar
+    ``num_blocks=N`` resolves to ``(N, N - 2)`` for latency or ``(N, N)`` for


What's the N, N-2 meaning

Binyang Li (Binyang2014) · 2026-09-18T22:55:09Z

+def record_stream(tensors: Iterable[Optional[torch.Tensor]], stream: torch.cuda.Stream) -> None:
+    """Protect allocator-backed storage used by native asynchronous GPU work.
+
+    PyTorch's allocator ignores foreign storage. Runtime buffers instead rely on
+    their native owner and the caller's obligation to finish local and peer use.
+    This does not provide producer/consumer ordering across streams.
+    """
+    for tensor in tensors:
+        if tensor is not None:
+            tensor.record_stream(stream)


Remove this function?

Binyang Li (Binyang2014) · 2026-09-18T22:55:31Z

+            tensor.record_stream(stream)
+
+
+class DevicePointerArray:


Why need this?

Remove the Python and public C++ receive-row count accessors and keep throughput row bounds in the native workspace. TOKEN_MAJOR remains capacity-sized, while native combine follows private routing metadata and ignores unused tail rows. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Require latency rank-major expert computation to use the runtime combine buffer instead of hiding a device-to-device staging copy. Clarify that throughput token-major rows are source-rank/token ordered while exposed expert counts are workload statistics. Consolidate Python EP tests into one end-to-end correctness matrix using top-k 8 and hidden size 4096 across every supported mode, layout, dtype, combine, and preparation path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Binyang Li (Binyang2014) marked this pull request as draft September 15, 2026 23:25

Merge branch 'binyli/ep-throughput-cpp' into binyli/ep-python-cpp

a1c29a8