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.

[weixiao-huang]

Maybe we can add more sanity check for handles like

    # global_check
    TORCH_CHECK(handles.device().is_cpu(), "handles must be a CPU tensor");
    TORCH_CHECK(handles.is_contiguous(), "handles must be contiguous");

    # if use_fabric
    TORCH_CHECK(handles.dim() == 2, "handles must be 2D [-1, 64] when using fabric");
    TORCH_CHECK(handles.dtype() == at::kByte, "handles must be uint8 in fabric");
    TORCH_CHECK(handles.size(1) == (int64_t)sizeof(CUmemFabricHandle),
        "handles.size(1) in fabric must be ", sizeof(CUmemFabricHandle));

    # else
    TORCH_CHECK(handles.dim() == 1, "handles must be 1D when not using fabric");
    TORCH_CHECK(handles.dtype() == at::kLong, "handles must be long in local fd");

[weixiao-huang]

This can keep as

    TORCH_CHECK((int64_t)shareables.size(0) == world_size,
        "shareables.size()=", shareables.size(0), " != world_size=", world_size);

[weixiao-huang]

It seems this local_rank is not used.

Also, other local_rank maybe change to rank since this is a global rank in both fabric and fd

[weixiao-huang]

Is world_size necessary here? Maybe we can just use dist.get_world_size(group=group) for getting world_size.

[weixiao-huang]

ditto

[weixiao-huang]

Maybe this can be cached by @lru_cache(maxsize=1) since create_nvl_single_owner_tensor, create_nvl_dist_multicast_tensor and create_nvl_dist_tensor will use it redundantly