Support multi node nvlink fabric by lyppg · Pull Request #23 · MoonshotAI/MoonEP

lyppg · 2026-07-31T21:22:01Z

Current MoonEP supports each rank of an EP group lives on the same node. This PR adds support for multi node with nvlink via CUDA fabric handles, valid for every GPU in the same NVLink/IMEX domain, e.g., GB200.

Tests

Run the benchmark on single node and two nodes.

# single node
torchrun --nproc_per_node=4 benchmarks/bench_prefetch.py

Config                   E    B       H     Hp   SMs  Slots   Data(MB)  Worst(us)  BW(GB/s)   CommBW
-----------------------------------------------------------------------------------------------------
slots_1                 32   14    3584   3072    32      1      44.04      34.51   1276.08   638.04
slots_2                 32   14    3584   3072    32      2      88.08      62.30   1413.79   706.90
slots_3                 32   14    3584   3072    32      3     132.12      90.28   1463.48   731.74
slots_5                 32   14    3584   3072    32      5     220.20     146.85   1499.50   749.75
ramp_0_3                32   14    3584   3072    32     12     528.48     343.52   1538.45   769.22
mixed                   32   14    3584   3072    32     12     528.48     344.71   1533.12   766.56
heavy                   32   14    3584   3072    32     17     748.68     484.67   1544.72   772.36
full_3x8                32   14    3584   3072    32     24    1056.96     681.36   1551.26   775.63
dense_B8                32    8    3584   3072    32     15     660.60     429.11   1539.46   769.73
epn4_full               16   14    3584   3072    32     12     528.48     343.51   1538.48   769.24
epn16_mixed             64   16    3584   3072    32     24    1056.96     680.12   1554.09   777.04
thin_7168x128           32   16    7168    128    32     12      44.04      34.85   1263.72   631.86
tall_1024x3072          32   14    1024   3072    32     12     150.99     102.85   1468.05   734.02
tiny_512x512            32   16     512    512    32     17      17.83      17.68   1008.25   504.12

# two GB200 nodes, --node_rank=1 on another node
torchrun --nnodes=2 --node_rank=0 --nproc_per_node=4     --master_addr=$MASTER_ADDR --master_port=$PORT benchmarks/bench_prefetch.py

Config                   E    B       H     Hp   SMs  Slots   Data(MB)  Worst(us)  BW(GB/s)   CommBW
-----------------------------------------------------------------------------------------------------
slots_1                 64   14    3584   3072    32      1      44.04      35.27   1248.76   624.38
slots_2                 64   14    3584   3072    32      2      88.08      63.17   1394.38   697.19
slots_3                 64   14    3584   3072    32      3     132.12      91.21   1448.46   724.23
slots_5                 64   14    3584   3072    32      5     220.20     147.27   1495.18   747.59
ramp_0_3                64   14    3584   3072    32     12     528.48     344.33   1534.80   767.40
mixed                   64   14    3584   3072    32     12     528.48     344.97   1531.97   765.99
heavy                   64   14    3584   3072    32     17     748.68     484.10   1546.55   773.27
full_3x8                64   14    3584   3072    32     24    1056.96     681.24   1551.53   775.77
dense_B8                64    8    3584   3072    32     15     660.60     428.49   1541.70   770.85
epn4_full               32   14    3584   3072    32     12     528.48     344.73   1533.05   766.52
epn16_mixed            128   16    3584   3072    32     24    1056.96     681.67   1550.56   775.28
thin_7168x128           64   16    7168    128    32     12      44.04      35.51   1240.37   620.19
tall_1024x3072          64   14    1024   3072    32     12     150.99     103.75   1455.32   727.66
tiny_512x512            64   16     512    512    32     17      17.83      18.67    954.93   477.46

asp0ex · 2026-08-03T09:48:32Z

        acc_tid = tidx - 32
        H_PER = cutlass.const_expr(TILE_ELEMS // ACC_THREADS)
-        acc_reg = cute.make_fragment((H_PER,), Float32)
+        acc_reg = cute.make_rmem_tensor((H_PER,), Float32)


This is only needed by the 4.6.0 upgrade (make_fragment was removed there). If you’d like to move to 4.6.0, please open a separate PR for the upgrade

asp0ex

Two things before merge:

Split the cutlass 4.4.2 → 4.6.0 bump (and the make_fragment →
make_rmem_tensor rename) into a separate PR — we need to stay on 4.4.2
for compatibility with other kernels (inline on setup.py /
grad_reduce.py).
benchmarks/bench_vs_deepep.py needs to follow the fds → shareables
rename; it calls nvl_dist_map directly and now raises TypeError
(inline on bindings.cu).
Rest of the fabric change LGTM.

lyppg · 2026-08-03T20:31:21Z

Thanks @asp0ex. nvidia-cutlass-dsl version change is removed. I upgraded it as my environment and sglang use 4.6.0. bench_vs_deepep.py is also updated.

weixiao-huang · 2026-08-04T09:06:02Z

+// Export an allocation (or multicast object) as the shareable form matching
+// `use_fabric`: a Python int holding a POSIX fd, or a 64-byte bytes object
+// holding a CUmemFabricHandle.
+static inline pybind11::object nvl_export_shareable(


Maybe we can return this as at::Tensor instead of pybind11::object, which can be directly used by all_gather instead of all_gather_object to avoid pickle overhead if there're too many devices.

weixiao-huang · 2026-08-04T09:09:15Z

+        return pybind11::bytes(
+            reinterpret_cast<const char *>(fabric.data), sizeof(fabric.data));


Suggested change

return pybind11::bytes(

reinterpret_cast<const char *>(fabric.data), sizeof(fabric.data));

auto handle_tensor = at::empty({(int64_t)sizeof(CUmemFabricHandle)},

at::TensorOptions().dtype(at::kByte).device(at::kCPU));

std::memcpy(handle_tensor.data_ptr(), &fabric, sizeof(CUmemFabricHandle));

return std::move(handle_tensor)

weixiao-huang · 2026-08-04T09:20:52Z

+    CUCHECK(cuMemExportToShareableHandle(
+        &fd, handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0));
+    TORCH_CHECK(fd >= 0, "cuMemExportToShareableHandle returned invalid fd");
+    return pybind11::int_(fd);


weixiao-huang · 2026-08-04T09:28:08Z

+
+// Import a peer's shareable handle produced by nvl_export_shareable.
+static inline CUmemGenericAllocationHandle nvl_import_shareable(
+    const pybind11::object &shareable, bool use_fabric


pybind11::object -> at::Tensor

weixiao-huang · 2026-08-04T09:40:07Z

+// non-root: import the multicast object from the handle sent by root. If it is
+// an fd, the caller may close it once imported.
+inline int64_t nvl_multicast_import(
+    const pybind11::object &shareable, bool use_fabric


weixiao-huang · 2026-08-04T09:40:39Z

-    int64_t size_bytes, int64_t num_devices
+// passes it back unchanged to bind_map. The caller passes the shareable handle
+// to the other ranks (closing the fd, if it is one, after all have imported).
+inline std::tuple<int64_t, pybind11::object> nvl_multicast_create(


weixiao-huang · 2026-08-04T09:47:09Z

+        multi_node = False
+        unsupported = [] if nvl_fabric_supported() else [socket.gethostname()]
+    else:
+        probe = (socket.gethostname(), bool(nvl_fabric_supported()))


Using rank instead of hostname in logging, so that converting bool(nvl_fabric_supported() to torch tensor and collecting it by all_gather_into_tensor instead of all_gather_object to avoid pickle dump/load.

add multinode nvlink fabric support

7666441

asp0ex reviewed Aug 3, 2026

View reviewed changes

Comment thread setup.py Outdated

asp0ex reviewed Aug 3, 2026

View reviewed changes

Comment thread csrc/bindings.cu

asp0ex reviewed Aug 3, 2026

View reviewed changes

lyppg added 2 commits August 3, 2026 13:05

keep cutlass 4.4.2

a10affa

update fds to shareables in bench_vs_deepep

b12130c

weixiao-huang reviewed Aug 4, 2026

View reviewed changes

Comment thread README.md Outdated

change back README.md

5239bdf

asp0ex reviewed Aug 4, 2026

View reviewed changes

Comment thread moonep/grad_reduce.py Outdated

minor change

cb7aa53

weixiao-huang reviewed Aug 4, 2026

View reviewed changes

Navigation Menu

Uh oh!

Support multi node nvlink fabric - #23

Support multi node nvlink fabric#23
lyppg wants to merge 5 commits into
MoonshotAI:masterfrom
bytedance-iaas:jxp/support_multinode_fabric

lyppg commented Jul 31, 2026

Uh oh!

Uh oh!

asp0ex Aug 3, 2026

Uh oh!

Uh oh!

asp0ex left a comment

Uh oh!

lyppg commented Aug 3, 2026

Uh oh!

Uh oh!

Uh oh!

weixiao-huang Aug 4, 2026

Uh oh!

weixiao-huang Aug 4, 2026

Uh oh!

weixiao-huang Aug 4, 2026

Uh oh!

weixiao-huang Aug 4, 2026

Uh oh!

weixiao-huang Aug 4, 2026

Uh oh!

weixiao-huang Aug 4, 2026

Uh oh!

weixiao-huang Aug 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		return pybind11::bytes(
		reinterpret_cast<const char *>(fabric.data), sizeof(fabric.data));

-        return pybind11::bytes(
-            reinterpret_cast<const char *>(fabric.data), sizeof(fabric.data));
+        auto handle_tensor = at::empty({(int64_t)sizeof(CUmemFabricHandle)},
+            at::TensorOptions().dtype(at::kByte).device(at::kCPU));
+        std::memcpy(handle_tensor.data_ptr(), &fabric, sizeof(CUmemFabricHandle));
+        return std::move(handle_tensor)

Uh oh!

Conversation

lyppg commented Jul 31, 2026

Tests

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

asp0ex left a comment

Choose a reason for hiding this comment

Uh oh!

lyppg commented Aug 3, 2026

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants