mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-07-13 10:47:42 +00:00
04f8089697
## Summary Enables the **high-throughput (HT)** MoE dispatch/combine backend as a first-class `MoEMode.HIGH_THROUGHPUT` path on the canonical nanobind `MoECommunicator` introduced in #818. After #818 the HT (DeepEP-style) backend was archived under `src/ext/ep/ht/` and not compiled (`mode=HIGH_THROUGHPUT` → `NotImplementedError`). This branch un-archives it, re-binds it under the nanobind module via a DLPack tensor bridge, adds the GB200 TMA kernel optimizations, and wires it through the high-level `MoECommunicator`. The low-latency `MoERuntime` path is unchanged. Already merged up-to-date with `feature/ep` (conflict-free): **23 commits, 14 files, +2682 / −690**. ## What's new - **HT compiles + binds under nanobind.** CMake links Torch and builds `ht/buffer.cc` + `ht/kernels/*.cu`; `bindings.cpp` binds `Config`, `EventHandle`, and the HT runtime (C++ `mscclpp::ep::Buffer`) exposed to Python as **`ExpertParallelRuntime`** (with a `Buffer` alias). - **nanobind ⇄ libtorch DLPack bridge.** HT methods return data-dependent-size `torch::Tensor` tuples, so tensors cross the boundary as DLPack capsules (`at::toDLPack` / `at::fromDLPack`; `torch.utils.dlpack` on the Python side) — zero-copy, no torch type-caster needed in nanobind. - **`MoECommunicator(mode=MoEMode.HIGH_THROUGHPUT)`** end-to-end: `_init_ht` builds the HT runtime on a torch `ProcessGroup`; `dispatch`/`combine` route by mode; FLAT layout; `previous_handle` cached-dispatch; intranode vs internode auto-selected from the RDMA size hint. LL continues to use `MoERuntime` (`comm=CommGroup`). - **GB200 TMA kernels** (intranode + internode): TMA direct-gather combine + all-sender dispatch, per-phase SM knobs, channel-adaptive warp counts. - **Tests** (`test_intranode_multirank.py`, `test_internode_multirank.py`): correctness drives the low-level `ExpertParallelRuntime`; the benchmark drives the public `MoECommunicator`. ## Validation 1 node × 4 GB200 (e256 / topk8 / hidden=7168 / tokens=4096): - Correctness: **PASS**, `max|got − expected| = 0.0`. - `MoECommunicator` HT bench: **dispatch 454.2 µs / combine 476.0 µs** (matches the raw-runtime best of 453.8 / 474). The high-level wrapper adds no measurable overhead. ## Notes / limitations - HT dispatch is **BF16-only** for now; HT layout is **FLAT** only. - The internode path compiles and the test is rewritten to `MoECommunicator`, but it has **only been validated single-node** so far (one NVL72 rack reachable); a 2-node run is pending. - `dispatch_async` / `combine_async` and block-level overlap are not implemented yet. ## Implementation notes (nanobind) Re-binding a pybind11 tensor API into nanobind required: object args as `nb::object` (not `nb::handle`); `.none()` on every `nb::arg` that may receive `None` (nanobind rejects `None` for object args by default); and renaming the reserved Python keyword `async` → `async_finish`.