Support multi node nvlink fabric

[lyppg]

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]

Two things before merge:

1. 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).
2. 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]

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.

[Copilot]

Pull request overview

This PR extends MoonEP’s NVLink-backed shared-buffer implementation to support EP groups that span multiple nodes within the same NVLink/IMEX domain (e.g., GB200) by introducing CUDA fabric handles (in addition to POSIX FD-based sharing), and updates tests/benchmarks to use local device indexing rather than assuming rank == cuda device.

Changes:

• Add fabric-handle based memory sharing path in NVLink buffer allocation/mapping (C++ + Python bindings), with an env var to force/auto-select handle type.
• Update distributed tests to select CUDA devices via LOCAL_RANK (and use "cuda" device strings so current device routing works).
• Update benchmarks to use LOCAL_RANK and barrier on the current CUDA device.

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
moonep/buffer.py	Adds fabric-handle selection and fabric-based handle exchange/map paths for distributed and single-owner NVLink tensors.
csrc/nvl_shared_buffer.cuh	Implements fabric handle export/import, validates handle tensors, and supports mapping via fabric handles.
csrc/bindings.cu	Exposes new fabric APIs/flags to Python (FABRIC_HANDLE_BYTES, nvl_fabric_supported, new args).
tests/kernel_test_utils.py	Introduces local_device_index() helper and updates helpers to rely on current CUDA device.
tests/conftest.py	Sets CUDA device using local rank and aligns barriers to the selected device.
tests/test_prefetch.py	Updates device selection and device_ids barriers to use local device index/current device routing.
tests/test_grad_reduce.py	Same as above for grad-reduce tests (device selection, barriers, tensor device strings).
tests/test_e2e.py	Switches device selection to local rank and uses "cuda" device strings.
tests/test_dispatch.py	Uses "cuda" device strings so tests follow the current device set by the fixture.
tests/test_combine.py	Uses "cuda" device strings for allocations, relying on fixture-selected device.
benchmarks/bench_vs_deepep.py	Updates benchmark’s internal NVLink sharing to use fabric handles when selected.
benchmarks/bench_prefetch.py	Sets device from LOCAL_RANK, uses "cuda", and barriers on current device.
benchmarks/bench_grad_reduce.py	Same as above for grad-reduce benchmark.

Suppressed comments (1)

moonep/buffer.py:61

• This uses assert to enforce a user-facing runtime constraint (MOONEP_MEM_HANDLE_TYPE=fabric requires fabric support on all ranks). Assertions can be disabled with -O, which would let execution continue and likely fail later in harder-to-debug ways.

    if mode == "fabric":
        assert not unsupported, (
            f"{_HANDLE_TYPE_ENV}=fabric, but fabric memory handles are "
            f"unsupported on group ranks {unsupported}."
        )
        return True

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.