From 5eaff99d61ea74caeb30c879f461f577db4054e3 Mon Sep 17 00:00:00 2001 From: leesou Date: Mon, 3 Aug 2026 14:52:16 +0000 Subject: [PATCH 1/5] Fix proxy fences --- moonep/combine.py | 2 +- moonep/combine_prologue.py | 2 +- moonep/dispatch.py | 2 +- moonep/grad_reduce.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/moonep/combine.py b/moonep/combine.py index 1a642c6..d5e054c 100644 --- a/moonep/combine.py +++ b/moonep/combine.py @@ -418,7 +418,7 @@ def kernel( # so a warp-level sync is enough to make that arrive cover # all 32 lanes' reads — cross-warp completion is already # gated by the 4-count empty mbarrier. - cute.arch.fence_acq_rel_cta() + cute.arch.fence_view_async_shared() cute.arch.sync_warp() load_pipe.consumer_release(load_state) load_state.advance() diff --git a/moonep/combine_prologue.py b/moonep/combine_prologue.py index cadb2d5..dc25765 100644 --- a/moonep/combine_prologue.py +++ b/moonep/combine_prologue.py @@ -448,7 +448,7 @@ def kernel( # arrive cover all 32 lanes' reads — cross-warp # completion is already gated by the 4-count empty # mbarrier. - cute.arch.fence_acq_rel_cta() + cute.arch.fence_view_async_shared() cute.arch.sync_warp() acc_load_pipe.consumer_release(acc_use_state) acc_use_state.advance() diff --git a/moonep/dispatch.py b/moonep/dispatch.py index 1f6c69e..625abe9 100644 --- a/moonep/dispatch.py +++ b/moonep/dispatch.py @@ -345,8 +345,8 @@ def kernel( idx = j * num_threads + tidx if idx < Int32(H): zero_smem[idx] = BFloat16(0) - cute.arch.barrier() cute.arch.fence_view_async_shared() + cute.arch.barrier() # ----- per-block token range tpb = (S + self.num_sms - 1) // self.num_sms diff --git a/moonep/grad_reduce.py b/moonep/grad_reduce.py index cc7cbf1..7a5c37f 100644 --- a/moonep/grad_reduce.py +++ b/moonep/grad_reduce.py @@ -359,7 +359,7 @@ def kernel( # (PipelineTmaAsync signalling thread; consumer_group=4), # so sync_warp makes that arrive cover all 32 lanes' reads; # cross-warp completion is gated by the 4-count empty mbarrier. - cute.arch.fence_acq_rel_cta() + cute.arch.fence_view_async_shared() cute.arch.sync_warp() load_pipe.consumer_release(cd_state) cd_state.advance() From 39859ebfeae72e5a56fcce3583e4daa4a59532e9 Mon Sep 17 00:00:00 2001 From: jxp Date: Wed, 5 Aug 2026 23:37:56 -0700 Subject: [PATCH 2/5] Support multi node nvlink fabric (#23) --- benchmarks/bench_grad_reduce.py | 13 +- benchmarks/bench_prefetch.py | 13 +- benchmarks/bench_vs_deepep.py | 53 ++++--- csrc/bindings.cu | 35 ++--- csrc/nvl_shared_buffer.cuh | 244 +++++++++++++++++++++++--------- moonep/buffer.py | 194 +++++++++++++++++++------ tests/conftest.py | 7 +- tests/kernel_test_utils.py | 9 +- tests/test_combine.py | 12 +- tests/test_dispatch.py | 22 +-- tests/test_e2e.py | 15 +- tests/test_grad_reduce.py | 23 +-- tests/test_prefetch.py | 15 +- 13 files changed, 450 insertions(+), 205 deletions(-) diff --git a/benchmarks/bench_grad_reduce.py b/benchmarks/bench_grad_reduce.py index 18d9825..e3c3c37 100644 --- a/benchmarks/bench_grad_reduce.py +++ b/benchmarks/bench_grad_reduce.py @@ -6,6 +6,7 @@ """ import argparse +import os import sys import torch @@ -66,7 +67,7 @@ def setup(): dist.init_process_group(backend="nccl") rank = dist.get_rank() - torch.cuda.set_device(rank) + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0))) return rank, dist.get_world_size() @@ -85,7 +86,7 @@ def expert_plan(R, B, epn, counts, dev): def bench_case(case, args, rank, R): - dev = f"cuda:{rank}" + dev = "cuda" epn = int(case["epn"]) E = R * epn H = int(case["H"]) @@ -113,7 +114,7 @@ def bench_case(case, args, rank, R): ) reduce_buffers[rank].copy_(torch.randn_like(reduce_buffers[rank])) torch.cuda.synchronize() - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[torch.cuda.current_device()]) def reduce_once(): launch_grad_reduce( @@ -133,7 +134,7 @@ def reduce_once(): for _ in range(args.warmup): reduce_once() torch.cuda.synchronize() - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[torch.cuda.current_device()]) # --no-graph keeps plain stream launches so tools like NCU can intercept # each kernel (graph capture/replay hides launches from kernel filters). @@ -144,7 +145,7 @@ def reduce_once(): for _ in range(args.iters): reduce_once() torch.cuda.synchronize() - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[torch.cuda.current_device()]) start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) @@ -174,7 +175,7 @@ def reduce_once(): # clears are local HBM, off the NVLink path). comm_gbs = slots * tile / worst_us * 1e6 / 1e9 if worst_us > 0 else 0.0 - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[torch.cuda.current_device()]) buffer.destroy() return worst_us, bytes_per_rank / 1e6, bw_gbs, comm_gbs, E, B, slots diff --git a/benchmarks/bench_prefetch.py b/benchmarks/bench_prefetch.py index f75a649..47230fa 100644 --- a/benchmarks/bench_prefetch.py +++ b/benchmarks/bench_prefetch.py @@ -6,6 +6,7 @@ """ import argparse +import os import sys import torch @@ -65,7 +66,7 @@ def setup(): dist.init_process_group(backend="nccl") rank = dist.get_rank() - torch.cuda.set_device(rank) + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0))) return rank, dist.get_world_size() @@ -86,7 +87,7 @@ def expert_plan(R, B, epn, counts, dev): def bench_case(case, args, rank, R): - dev = f"cuda:{rank}" + dev = "cuda" epn = int(case["epn"]) E = R * epn H = int(case["H"]) @@ -117,7 +118,7 @@ def bench_case(case, args, rank, R): experts_to_copy = plan.flatten() if rank == 0 else \ torch.full((R * B,), -1, dtype=torch.int32, device=dev) torch.cuda.synchronize() - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[torch.cuda.current_device()]) def prefetch_once(): launch_prefetch( @@ -132,7 +133,7 @@ def prefetch_once(): for _ in range(args.warmup): prefetch_once() torch.cuda.synchronize() - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[torch.cuda.current_device()]) # --no-graph keeps plain stream launches so tools like NCU can intercept # each kernel (graph capture/replay hides launches from kernel filters). @@ -143,7 +144,7 @@ def prefetch_once(): for _ in range(args.iters): prefetch_once() torch.cuda.synchronize() - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[torch.cuda.current_device()]) start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) @@ -171,7 +172,7 @@ def prefetch_once(): # writes are local HBM, off the NVLink path). comm_gbs = slots * tile / worst_us * 1e6 / 1e9 if worst_us > 0 else 0.0 - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[torch.cuda.current_device()]) return worst_us, bytes_per_rank / 1e6, bw_gbs, comm_gbs, E, B, slots diff --git a/benchmarks/bench_vs_deepep.py b/benchmarks/bench_vs_deepep.py index 6884a49..93d0830 100644 --- a/benchmarks/bench_vs_deepep.py +++ b/benchmarks/bench_vs_deepep.py @@ -125,17 +125,23 @@ class MoonEPRunner: def __init__(self, group, R, S, K, E, H, num_sms, hp=2048): from moonep import Buffer - from moonep._C import (nvl_dist_alloc, nvl_release_mem_handle, - nvl_dist_map, get_vmm_granularity) - from moonep.buffer import _exchange_ipc_fds + from moonep._C import ( + FABRIC_HANDLE_BYTES as _FABRIC_HANDLE_BYTES, + nvl_dist_alloc, nvl_release_mem_handle, + nvl_dist_map, get_vmm_granularity + ) + from moonep.buffer import (_all_gather_shareables, _exchange_ipc_fds, + _use_fabric_for_group) from moonep.planning import allocate_planning_outputs, launch_planning from moonep.inter_rank_sync import launch_inter_rank_sync self._launch_planning = launch_planning self._launch_sync = launch_inter_rank_sync + self._use_fabric = _use_fabric_for_group(group) self._alloc_chunk = lambda shape, dtype: nvl_dist_alloc( - shape=shape, dtype=dtype) + shape=shape, dtype=dtype, use_fabric=self._use_fabric) self._release = nvl_release_mem_handle self._dist_map = nvl_dist_map + self._gather_shareables = _all_gather_shareables self._exchange_fds = _exchange_ipc_fds self.num_sms = num_sms self.rank = dist.get_rank(group) @@ -169,24 +175,37 @@ def build_full(dtype: torch.dtype): assert chunk_bytes % gran == 0, ( f"chunk bytes {chunk_bytes} not VMM-aligned ({gran})") # per-rank expert chunk (shared) + this rank's buffer chunk (local) - ka_w, w_fd, w_owned = self._alloc_chunk([self.epn, H, hp], dtype) - ka_b, b_fd, b_owned = self._alloc_chunk([self.B, H, hp], dtype) + ka_w, w_sh, w_owned = self._alloc_chunk([self.epn, H, hp], dtype) + ka_b, b_sh, b_owned = self._alloc_chunk([self.B, H, hp], dtype) for ka, owned in ((ka_w, w_owned), (ka_b, b_owned)): self._keepalives.append(ka) self._release(owned) - # exchange expert chunk fds - fds = self._exchange_fds(w_fd, list(range(R)), self.rank, R, group) - os.close(w_fd) - all_w_fds = [fds[r] for r in range(R)] - try: + # exchange the expert chunk handles, then append this rank's own + # buffer chunk as the trailing (R+1)-th chunk + if self._use_fabric: + all_w = self._gather_shareables(w_sh, group) full = self._dist_map( chunk_shape=[self.epn, H, hp], dtype=dtype, - fds=all_w_fds + [b_fd], local_rank=self.rank, - world_size=R + 1) - finally: - for fd in all_w_fds: - os.close(fd) - os.close(b_fd) + shareables=torch.cat( + [all_w, b_sh.view(1, _FABRIC_HANDLE_BYTES)], dim=0), + local_rank=self.rank, world_size=R + 1, use_fabric=True) + else: + w_fd, b_fd = int(w_sh.item()), int(b_sh.item()) + fds = self._exchange_fds(w_fd, list(range(R)), self.rank, R, + group) + os.close(w_fd) + all_w_fds = [fds[r] for r in range(R)] + try: + full = self._dist_map( + chunk_shape=[self.epn, H, hp], dtype=dtype, + shareables=torch.tensor(all_w_fds + [b_fd], + dtype=torch.int64), + local_rank=self.rank, world_size=R + 1, + use_fabric=False) + finally: + for fd in all_w_fds: + os.close(fd) + os.close(b_fd) return full # full_weight for the 3 projections (bf16). grad_reduce is not on the diff --git a/csrc/bindings.cu b/csrc/bindings.cu index 6f00f4a..8885233 100644 --- a/csrc/bindings.cu +++ b/csrc/bindings.cu @@ -7,31 +7,23 @@ int64_t get_vmm_granularity() { int device_id; CUDACHECK(cudaGetDevice(&device_id)); - - CUdevice cu_device; - CUCHECK(cuDeviceGet(&cu_device, device_id)); - - CUmemAllocationProp prop = {}; - prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; - prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - prop.location.id = device_id; - prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; - - size_t granularity; - CUCHECK(cuMemGetAllocationGranularity( - &granularity, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); - return static_cast(granularity); + return static_cast(nvl_granularity_max(device_id)); } PYBIND11_MODULE(_C, m) { + m.attr("FABRIC_HANDLE_BYTES") = kFabricHandleBytes; m.def("nvl_dist_alloc", &nvl_dist_alloc, pybind11::arg("shape"), pybind11::arg("dtype"), - "Allocate a local chunk and export it as a POSIX fd for IPC sharing."); + pybind11::arg("use_fabric") = false, + "Allocate a local chunk and export it for sharing, as a POSIX fd " + "(same node) or a 64-byte fabric handle (same NVLink domain)."); m.def("nvl_dist_map", &nvl_dist_map, pybind11::arg("chunk_shape"), pybind11::arg("dtype"), - pybind11::arg("fds"), pybind11::arg("local_rank"), - pybind11::arg("world_size"), - "Map all chunk fds into a contiguous VA region (all RW)."); + pybind11::arg("shareables"), pybind11::arg("local_rank"), + pybind11::arg("world_size"), pybind11::arg("use_fabric") = false, + "Map all ranks' chunks into a contiguous VA region (all RW)."); + m.def("nvl_fabric_supported", &nvl_fabric_supported, + "Whether the current device can export/import fabric handles."); m.def("get_vmm_granularity", &get_vmm_granularity, "Return VMM allocation granularity in bytes."); m.def("get_multicast_granularity", &get_multicast_granularity, @@ -41,10 +33,11 @@ PYBIND11_MODULE(_C, m) { "Whether the current device supports multicast (NVSwitch SHARP)."); m.def("nvl_multicast_create", &nvl_multicast_create, pybind11::arg("size_bytes"), pybind11::arg("num_devices"), - "Root-only: create a multicast object and export it as a POSIX fd."); + pybind11::arg("use_fabric") = false, + "Root-only: create a multicast object and export it for sharing."); m.def("nvl_multicast_import", &nvl_multicast_import, - pybind11::arg("fd"), - "Non-root: import a multicast object from the root's POSIX fd."); + pybind11::arg("shareable"), pybind11::arg("use_fabric") = false, + "Non-root: import a multicast object from the root's handle."); m.def("nvl_multicast_add_device", &nvl_multicast_add_device, pybind11::arg("mc_handle"), "Add the current device to a multicast object (before bind)."); diff --git a/csrc/nvl_shared_buffer.cuh b/csrc/nvl_shared_buffer.cuh index 0b7747c..408e179 100644 --- a/csrc/nvl_shared_buffer.cuh +++ b/csrc/nvl_shared_buffer.cuh @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -37,9 +38,52 @@ } while(0) #endif +static inline bool nvl_fabric_supported() { + int device_id; + CUDACHECK(cudaGetDevice(&device_id)); + CUdevice cu_device; + CUCHECK(cuDeviceGet(&cu_device, device_id)); + int supported = 0; + CUresult err = cuDeviceGetAttribute(&supported, + CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, cu_device); + // Older drivers do not know the attribute at all. + if (err != CUDA_SUCCESS) return false; + return supported != 0; +} + +static inline size_t nvl_granularity_for(int device_id, + CUmemAllocationHandleType ht) { + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.location.id = device_id; + prop.requestedHandleTypes = ht; + + size_t granularity; + CUCHECK(cuMemGetAllocationGranularity( + &granularity, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); + return granularity; +} + +static inline size_t nvl_granularity_max(int device_id) { + size_t gran = nvl_granularity_for( + device_id, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR); + if (nvl_fabric_supported()) { + gran = std::max( + gran, nvl_granularity_for(device_id, CU_MEM_HANDLE_TYPE_FABRIC)); + } + return gran; +} + +static inline CUmemAllocationHandleType nvl_handle_type(bool use_fabric) { + return use_fabric ? CU_MEM_HANDLE_TYPE_FABRIC + : CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; +} + static inline std::tuple nvl_prepare( const std::vector &shape, - at::ScalarType dtype + at::ScalarType dtype, + bool use_fabric ) { TORCH_CHECK(!shape.empty(), "Shape must be non-empty"); @@ -52,15 +96,8 @@ static inline std::tuple nvl_prepare( size *= static_cast(dim); } - CUmemAllocationProp prop = {}; - prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; - prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - prop.location.id = device_id; - prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; - - size_t granularity; - CUCHECK(cuMemGetAllocationGranularity( - &granularity, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); + size_t granularity = nvl_granularity_for( + device_id, nvl_handle_type(use_fabric)); size_t allocated_size = (size + granularity - 1) / granularity * granularity; return {size, allocated_size, device_id}; @@ -98,24 +135,94 @@ static inline at::Tensor make_vmm_tensor( return at::Tensor(std::move(impl)); } -// Returns (keepalive VA tensor, exported POSIX fd, owned mem handle as int64). +static constexpr int64_t kFabricHandleBytes = + static_cast(sizeof(CUmemFabricHandle::data)); + +// Export an allocation (or multicast object) as a CPU tensor matching +// `use_fabric`: uint8[64] holding a CUmemFabricHandle, or int64[] (a scalar +// tensor) holding a POSIX fd. +static inline at::Tensor nvl_export_shareable( + CUmemGenericAllocationHandle handle, bool use_fabric +) { + if (use_fabric) { + CUmemFabricHandle fabric = {}; + CUCHECK(cuMemExportToShareableHandle( + &fabric, handle, CU_MEM_HANDLE_TYPE_FABRIC, 0)); + auto out = at::empty({kFabricHandleBytes}, + at::TensorOptions().dtype(at::kByte)); + std::memcpy(out.data_ptr(), fabric.data, sizeof(fabric.data)); + return out; + } + int fd = -1; + CUCHECK(cuMemExportToShareableHandle( + &fd, handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0)); + TORCH_CHECK(fd >= 0, "cuMemExportToShareableHandle returned invalid fd"); + return at::scalar_tensor(fd, at::TensorOptions().dtype(at::kLong)); +} + +static inline void nvl_check_handles( + const at::Tensor &handles, int64_t world_size, bool use_fabric +) { + TORCH_CHECK(handles.device().is_cpu() && handles.is_contiguous(), + "shareable handles must be a contiguous CPU tensor, got device=", + handles.device(), " contiguous=", handles.is_contiguous()); + if (use_fabric) { + TORCH_CHECK(handles.scalar_type() == at::kByte && handles.dim() == 2 + && handles.size(0) == world_size + && handles.size(1) == kFabricHandleBytes, + "fabric handles must be uint8[", world_size, ", ", + kFabricHandleBytes, "], got ", handles.scalar_type(), handles.sizes()); + } else { + TORCH_CHECK(handles.scalar_type() == at::kLong + && handles.numel() == world_size, + "fd handles must be int64 with ", world_size, " elements, got ", + handles.scalar_type(), handles.sizes()); + } +} + +// Import row `index` of a handle tensor already validated by nvl_check_handles. +static inline CUmemGenericAllocationHandle nvl_import_shareable( + const at::Tensor &handles, int64_t index, bool use_fabric +) { + CUmemGenericAllocationHandle handle; + if (use_fabric) { + CUmemFabricHandle fabric = {}; + std::memcpy(fabric.data, + handles.const_data_ptr() + index * kFabricHandleBytes, + sizeof(fabric.data)); + CUCHECK(cuMemImportFromShareableHandle( + &handle, &fabric, CU_MEM_HANDLE_TYPE_FABRIC)); + } else { + int64_t fd = handles.const_data_ptr()[index]; + TORCH_CHECK(fd >= 0, "invalid fd ", fd, " at index ", index); + CUCHECK(cuMemImportFromShareableHandle( + &handle, reinterpret_cast(static_cast(fd)), + CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR)); + } + return handle; +} + +// Returns (keepalive VA tensor, shareable handle, owned mem handle as int64). // The owned mem handle is not released here (kept for multicast cuMulticastBindMem); // the caller must call nvl_release_mem_handle when done (buffers that do not need // multicast release it immediately, matching the old behavior; buffers that need // multicast release it after bind). -// The fd is used for cross-process IPC sharing (POSIX file descriptor); the caller -// must close it after all peers have imported it. -inline std::tuple nvl_dist_alloc( +// The shareable handle is what peers import: an int64 scalar tensor holding an +// fd the caller must close once all peers have imported it, or a uint8[64] +// fabric handle that needs no cleanup. +inline std::tuple nvl_dist_alloc( const std::vector &chunk_shape, - at::ScalarType dtype + at::ScalarType dtype, + bool use_fabric ) { - auto [nbytes, allocated_size, device_id] = nvl_prepare(chunk_shape, dtype); + auto [nbytes, allocated_size, device_id] = + nvl_prepare(chunk_shape, dtype, use_fabric); CUmemAllocationProp prop = {}; prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; prop.location.id = device_id; - prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; + prop.requestedHandleTypes = nvl_handle_type(use_fabric); CUmemGenericAllocationHandle mem_handle; CUCHECK(cuMemCreate(&mem_handle, allocated_size, &prop, 0)); @@ -130,21 +237,18 @@ inline std::tuple nvl_dist_alloc( access_desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; CUCHECK(cuMemSetAccess(dptr, allocated_size, &access_desc, 1)); - // Export the POSIX fd for other processes to import. Do not release - // mem_handle: keep the owned generic handle for multicast BindMem. The - // caller is responsible for nvl_release_mem_handle. The physical memory is - // referenced by the unicast map (keepalive) and (optionally) multicast; - // it is only truly freed after all unmaps once the handle is released. - int fd = -1; - CUCHECK(cuMemExportToShareableHandle( - &fd, mem_handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0)); - TORCH_CHECK(fd >= 0, "cuMemExportToShareableHandle returned invalid fd"); + // Export for other processes to import. Do not release mem_handle: keep the + // owned generic handle for multicast BindMem. The caller is responsible for + // nvl_release_mem_handle. The physical memory is referenced by the unicast + // map (keepalive) and (optionally) multicast; it is only truly freed after + // all unmaps once the handle is released. + auto shareable = nvl_export_shareable(mem_handle, use_fabric); auto keepalive = make_vmm_tensor( reinterpret_cast(dptr), nbytes, allocated_size, device_id, chunk_shape, dtype); - return {std::move(keepalive), static_cast(fd), + return {std::move(keepalive), std::move(shareable), static_cast(mem_handle)}; } @@ -159,21 +263,22 @@ static inline void nvl_release_mem_handle(int64_t mem_handle_u64) { * Key difference from reference: ALL chunks get RW access (not just local_rank), * because dispatch needs to write to remote ranks' regions. * - * `fds` are the POSIX fds exported by each rank via nvl_dist_alloc (passed - * between processes by the caller, already valid in this process). The caller - * may close the fds once they are imported. + * `shareables` are the handles exported by each rank via nvl_dist_alloc and + * routed to this process by the caller: POSIX fds (already valid here, the + * caller may close them once imported) or 64-byte fabric blobs. */ inline at::Tensor nvl_dist_map( const std::vector &chunk_shape, at::ScalarType dtype, - const std::vector &fds, + const at::Tensor &shareables, int64_t local_rank, - int64_t world_size + int64_t world_size, + bool use_fabric ) { - TORCH_CHECK((int64_t)fds.size() == world_size, - "fds.size()=", fds.size(), " != world_size=", world_size); + nvl_check_handles(shareables, world_size, use_fabric); - auto [chunk_nbytes, chunk_allocated_size, device_id] = nvl_prepare(chunk_shape, dtype); + auto [chunk_nbytes, chunk_allocated_size, device_id] = + nvl_prepare(chunk_shape, dtype, use_fabric); TORCH_CHECK(chunk_allocated_size == chunk_nbytes, "Chunk byte size (", chunk_nbytes, ") must be aligned to VMM granularity (", @@ -187,12 +292,8 @@ inline at::Tensor nvl_dist_map( CUCHECK(cuMemAddressReserve(&dptr, total_allocated_size, 0, 0, 0)); for (int64_t i = 0; i < world_size; i++) { - int fd = static_cast(fds[i]); - CUmemGenericAllocationHandle mem_handle; - CUCHECK(cuMemImportFromShareableHandle( - &mem_handle, - reinterpret_cast(static_cast(fd)), - CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR)); + CUmemGenericAllocationHandle mem_handle = + nvl_import_shareable(shareables, i, use_fabric); CUdeviceptr chunk_va = dptr + i * chunk_allocated_size; CUCHECK(cuMemMap(chunk_va, chunk_allocated_size, 0, mem_handle, 0)); @@ -276,13 +377,13 @@ static inline bool nvl_multicast_supported() { return supported != 0; } -// Recommended multicast alignment granularity (bytes). Both the addr and size -// of a bind must be aligned to it. -static inline size_t nvl_multicast_granularity(int num_devices) { +static inline size_t nvl_multicast_granularity_for( + int num_devices, CUmemAllocationHandleType ht +) { CUmulticastObjectProp prop = {}; prop.numDevices = static_cast(num_devices); prop.size = 0; - prop.handleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; + prop.handleTypes = ht; prop.flags = 0; size_t gran = 0; CUCHECK(cuMulticastGetGranularity( @@ -290,18 +391,32 @@ static inline size_t nvl_multicast_granularity(int num_devices) { return gran; } +// Recommended multicast alignment granularity (bytes). Both the addr and size +// of a bind must be aligned to it. Like nvl_granularity_max, this is the max +// over every handle type that may be used, so a buffer padded to it stays +// valid whichever type the process group picks. +static inline size_t nvl_multicast_granularity(int num_devices) { + size_t gran = nvl_multicast_granularity_for( + num_devices, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR); + if (nvl_fabric_supported()) { + gran = std::max(gran, nvl_multicast_granularity_for( + num_devices, CU_MEM_HANDLE_TYPE_FABRIC)); + } + return gran; +} + static inline int64_t get_multicast_granularity(int64_t num_devices) { return static_cast( nvl_multicast_granularity(static_cast(num_devices))); } -// root only: create the multicast object and export a POSIX fd. -// Returns (mc_handle as uint64, exported fd). +// root only: create the multicast object and export it for the other ranks. +// Returns (mc_handle as uint64, shareable handle). // mc_handle is a handle value valid within this process; Python holds it and -// passes it back unchanged to bind_map. The caller passes the fd to the other -// ranks and closes it after all peers have imported it. -inline std::tuple nvl_multicast_create( - 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 nvl_multicast_create( + int64_t size_bytes, int64_t num_devices, bool use_fabric ) { TORCH_CHECK(nvl_multicast_supported(), "Multicast not supported on this device"); @@ -312,28 +427,25 @@ inline std::tuple nvl_multicast_create( CUmulticastObjectProp prop = {}; prop.numDevices = static_cast(num_devices); prop.size = aligned; - prop.handleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; + prop.handleTypes = nvl_handle_type(use_fabric); prop.flags = 0; CUmemGenericAllocationHandle mc_handle; CUCHECK(cuMulticastCreate(&mc_handle, &prop)); - int fd = -1; - CUCHECK(cuMemExportToShareableHandle( - &fd, mc_handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0)); - TORCH_CHECK(fd >= 0, "cuMemExportToShareableHandle returned invalid fd"); - - return {static_cast(mc_handle), static_cast(fd)}; + auto shareable = nvl_export_shareable(mc_handle, use_fabric); + return {static_cast(mc_handle), std::move(shareable)}; } -// non-root: import the multicast object from the POSIX fd sent by root. The -// caller may close the fd once it is imported. -inline int64_t nvl_multicast_import(int64_t fd) { - CUmemGenericAllocationHandle mc_handle; - CUCHECK(cuMemImportFromShareableHandle( - &mc_handle, reinterpret_cast(static_cast(fd)), - CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR)); - return static_cast(mc_handle); +// 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 at::Tensor &shareable, bool use_fabric +) { + auto one = use_fabric ? shareable.view({1, kFabricHandleBytes}) + : shareable.view({1}); + nvl_check_handles(one, 1, use_fabric); + return static_cast(nvl_import_shareable(one, 0, use_fabric)); } // all ranks: add this rank's device to the multicast group. Must happen before diff --git a/moonep/buffer.py b/moonep/buffer.py index 0923c2c..4954810 100644 --- a/moonep/buffer.py +++ b/moonep/buffer.py @@ -8,8 +8,10 @@ import torch.distributed as dist from moonep._C import ( + FABRIC_HANDLE_BYTES as _FABRIC_HANDLE_BYTES, nvl_dist_alloc, nvl_dist_map, + nvl_fabric_supported, nvl_release_mem_handle, get_vmm_granularity, get_multicast_granularity, @@ -26,6 +28,69 @@ torch.int32: 4, } +# How VMM allocations are shared between the ranks of an EP group. "auto" (the +# default) picks fabric handles when the group spans more than one node and the +# device supports them, and POSIX fds otherwise; "fabric" / "fd" force one. +_HANDLE_TYPE_ENV = "MOONEP_MEM_HANDLE_TYPE" + + +def _use_fabric_for_group(group: dist.ProcessGroup | None) -> bool: + mode = os.environ.get(_HANDLE_TYPE_ENV, "auto").lower() + assert mode in ("auto", "fabric", "fd"), ( + f"{_HANDLE_TYPE_ENV} must be one of auto/fabric/fd, got {mode!r}" + ) + if mode == "fd": + return False + + supported = bool(nvl_fabric_supported()) + world_size = (dist.get_world_size(group=group) + if dist.is_available() and dist.is_initialized() else 1) + if world_size == 1: + unsupported = [] if supported else [0] + else: + local = torch.tensor([supported], dtype=torch.uint8, device="cuda") + gathered = torch.empty(world_size, dtype=torch.uint8, device="cuda") + dist.all_gather_into_tensor(gathered, local, group=group) + unsupported = (gathered == 0).nonzero().flatten().tolist() + + if mode == "fabric": + assert not unsupported, ( + f"{_HANDLE_TYPE_ENV}=fabric, but fabric memory handles are " + f"unsupported on group ranks {unsupported}." + ) + return True + + return not unsupported + + +def _all_gather_shareables( + local_handle: torch.Tensor, + group: dist.ProcessGroup | None, +) -> torch.Tensor: + """Gather every rank's fabric handle into a uint8[world_size, 64] CPU tensor.""" + world_size = dist.get_world_size(group=group) + gathered = torch.empty(world_size, _FABRIC_HANDLE_BYTES, + dtype=torch.uint8, device="cuda") + dist.all_gather_into_tensor( + gathered, local_handle.cuda().view(1, _FABRIC_HANDLE_BYTES), group=group) + return gathered.cpu() + + +def _broadcast_shareable( + local_handle: torch.Tensor | None, + owner_rank: int, + group: dist.ProcessGroup | None, +) -> torch.Tensor: + """Broadcast one rank's fabric handle to the whole group.""" + if local_handle is not None: + buf = local_handle.cuda() + else: + buf = torch.empty(_FABRIC_HANDLE_BYTES, dtype=torch.uint8, device="cuda") + src = dist.get_global_rank(group, owner_rank) if group is not None \ + else owner_rank + dist.broadcast(buf, src=src, group=group) + return buf.cpu() + def pad_to_granularity(nbytes: int) -> int: """Round up nbytes to VMM granularity.""" @@ -117,27 +182,41 @@ def _exchange_ipc_fds( def _map_nvl_dist_tensor( chunk_shape: list[int], dtype: torch.dtype, - local_fd: int, + shareable: torch.Tensor, keepalive: torch.Tensor, local_rank: int, world_size: int, group: dist.ProcessGroup | None, + use_fabric: bool, ) -> torch.Tensor: - fds = _exchange_ipc_fds(local_fd, list(range(world_size)), - local_rank, world_size, group) - os.close(local_fd) - all_fds = [fds[r] for r in range(world_size)] - try: + if use_fabric: + shareables = _all_gather_shareables(shareable, group) full_tensor = nvl_dist_map( chunk_shape=chunk_shape, dtype=dtype, - fds=all_fds, + shareables=shareables, local_rank=local_rank, world_size=world_size, + use_fabric=True, ) - finally: - for fd in all_fds: - os.close(fd) + else: + local_fd = int(shareable.item()) + fds = _exchange_ipc_fds(local_fd, list(range(world_size)), + local_rank, world_size, group) + os.close(local_fd) + all_fds = [fds[r] for r in range(world_size)] + try: + full_tensor = nvl_dist_map( + chunk_shape=chunk_shape, + dtype=dtype, + shareables=torch.tensor(all_fds, dtype=torch.int64), + local_rank=local_rank, + world_size=world_size, + use_fabric=False, + ) + finally: + for fd in all_fds: + os.close(fd) full_tensor._keepalive = keepalive return full_tensor @@ -155,13 +234,15 @@ def create_nvl_dist_tensor( Use pad_dim0_for_alignment() to compute the padded dim0. `local_rank` and `world_size` must match the given `group` (or the default - group when `group is None`). All ranks in the group exchange IPC fds. + group when `group is None`). All ranks in the group exchange memory handles. """ - keepalive, local_fd, owned_handle = nvl_dist_alloc(shape=chunk_shape, dtype=dtype) + use_fabric = _use_fabric_for_group(group) + keepalive, shareable, owned_handle = nvl_dist_alloc( + shape=chunk_shape, dtype=dtype, use_fabric=use_fabric) try: return _map_nvl_dist_tensor( - chunk_shape, dtype, local_fd, keepalive, - local_rank, world_size, group, + chunk_shape, dtype, shareable, keepalive, + local_rank, world_size, group, use_fabric, ) finally: nvl_release_mem_handle(owned_handle) @@ -183,14 +264,17 @@ def create_nvl_dist_multicast_tensor( The owned allocation handle needed by `cuMulticastBindMem` stays internal to this helper and is released after both mappings have been created. """ - keepalive, local_fd, owned_handle = nvl_dist_alloc(shape=chunk_shape, dtype=dtype) + use_fabric = _use_fabric_for_group(group) + keepalive, shareable, owned_handle = nvl_dist_alloc( + shape=chunk_shape, dtype=dtype, use_fabric=use_fabric) try: full_tensor = _map_nvl_dist_tensor( - chunk_shape, dtype, local_fd, keepalive, - local_rank, world_size, group, + chunk_shape, dtype, shareable, keepalive, + local_rank, world_size, group, use_fabric, ) mc_view = _create_nvl_multicast_view( full_tensor, owned_handle, local_rank, world_size, group, + use_fabric, ) return full_tensor, mc_view finally: @@ -203,6 +287,7 @@ def _create_nvl_multicast_view( local_rank: int, world_size: int, group: dist.ProcessGroup | None = None, + use_fabric: bool = False, ) -> torch.Tensor: """Overlay a multicast (NVSwitch SHARP) mapping on an existing NVL chunk. @@ -223,21 +308,29 @@ def _create_nvl_multicast_view( size_bytes = chunk_elems * meta_buf.element_size() is_root = local_rank == 0 - # Root creates the multicast object and sends its IPC fd to all ranks. + # Root creates the multicast object and shares it with all ranks. if is_root: - mc_handle, mc_fd = nvl_multicast_create(size_bytes, world_size) + mc_handle, mc_shareable = nvl_multicast_create( + size_bytes, world_size, use_fabric=use_fabric) else: - mc_handle, mc_fd = 0, None + mc_handle, mc_shareable = 0, None - fds = _exchange_ipc_fds(mc_fd, [0], local_rank, world_size, group) - if is_root: - os.close(mc_fd) - root_fd = fds[0] - try: + if use_fabric: + root_handle = _broadcast_shareable(mc_shareable, 0, group) if not is_root: - mc_handle = nvl_multicast_import(root_fd) - finally: - os.close(root_fd) + mc_handle = nvl_multicast_import(root_handle, use_fabric=True) + else: + local_fd = int(mc_shareable.item()) if is_root else None + fds = _exchange_ipc_fds(local_fd, [0], local_rank, world_size, group) + if is_root: + os.close(local_fd) + root_fd = fds[0] + try: + if not is_root: + mc_handle = nvl_multicast_import( + torch.tensor(root_fd, dtype=torch.int64), use_fabric=False) + finally: + os.close(root_fd) # All ranks add their device before any bind, then barrier. nvl_multicast_add_device(mc_handle) @@ -254,6 +347,7 @@ def create_nvl_single_owner_tensor( dtype: torch.dtype, owner_rank: int, local_rank: int, + group: dist.ProcessGroup | None = None, ) -> torch.Tensor: """Allocate a VMM tensor on one GPU, visible to all ranks via NVLink. @@ -262,26 +356,40 @@ def create_nvl_single_owner_tensor( accesses go over NVLink). shape must already be padded to VMM granularity (use pad_dim0_for_alignment). """ - world_size = dist.get_world_size() - if local_rank == owner_rank: - keepalive, local_fd, owned_handle = nvl_dist_alloc(shape=shape, dtype=dtype) + world_size = dist.get_world_size(group=group) + use_fabric = _use_fabric_for_group(group) + is_owner = local_rank == owner_rank + + if is_owner: + keepalive, shareable, owned_handle = nvl_dist_alloc( + shape=shape, dtype=dtype, use_fabric=use_fabric) nvl_release_mem_handle(owned_handle) else: - local_fd = None + shareable = None - fds = _exchange_ipc_fds(local_fd, [owner_rank], local_rank, world_size, - group=None) - if local_fd is not None: - os.close(local_fd) - owner_fd = fds[owner_rank] - try: + if use_fabric: + owner_handle = _broadcast_shareable(shareable, owner_rank, group) tensor = nvl_dist_map( chunk_shape=shape, dtype=dtype, - fds=[owner_fd], local_rank=0, world_size=1, + shareables=owner_handle.view(1, _FABRIC_HANDLE_BYTES), + local_rank=0, world_size=1, use_fabric=True, ) - finally: - os.close(owner_fd) - - if local_rank == owner_rank: + else: + local_fd = int(shareable.item()) if is_owner else None + fds = _exchange_ipc_fds(local_fd, [owner_rank], local_rank, + world_size, group) + if is_owner: + os.close(local_fd) + owner_fd = fds[owner_rank] + try: + tensor = nvl_dist_map( + chunk_shape=shape, dtype=dtype, + shareables=torch.tensor([owner_fd], dtype=torch.int64), + local_rank=0, world_size=1, use_fabric=False, + ) + finally: + os.close(owner_fd) + + if is_owner: tensor._keepalive = keepalive return tensor diff --git a/tests/conftest.py b/tests/conftest.py index 83b842e..24c5309 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -18,13 +18,16 @@ def dist_env(): if "RANK" not in os.environ: pytest.skip("distributed kernel tests must be launched with torchrun") + from tests.kernel_test_utils import local_device_index + if not dist.is_initialized(): dist.init_process_group(backend="nccl") rank = dist.get_rank() - torch.cuda.set_device(rank) + device = local_device_index() + torch.cuda.set_device(device) yield rank, dist.get_world_size() - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[device]) if dist.is_initialized(): dist.destroy_process_group() diff --git a/tests/kernel_test_utils.py b/tests/kernel_test_utils.py index b61829d..3202038 100644 --- a/tests/kernel_test_utils.py +++ b/tests/kernel_test_utils.py @@ -1,3 +1,4 @@ +import os from dataclasses import dataclass import pytest @@ -12,6 +13,10 @@ _ACTIVE_BUFFERS = [] +def local_device_index() -> int: + return int(os.environ.get("LOCAL_RANK", os.environ.get("RANK", 0))) + + def _align_up(x: int, alignment: int) -> int: return ((x + alignment - 1) // alignment) * alignment @@ -77,7 +82,7 @@ def skip_if_unsupported_world_size(case, R): def make_topk(case, rank, R): skip_if_unsupported_world_size(case, R) - dev = f"cuda:{rank}" + dev = "cuda" E = case.E(R) if case.K > E and case.routing in {"balanced", "biased"}: pytest.skip(f"case {case.name} requires K <= E, got K={case.K}, E={E}") @@ -117,7 +122,7 @@ def gather_tensor(t, R): def assert_all_ranks(ok, rank, R, label, detail=""): - ok_tensor = torch.tensor([int(ok)], dtype=torch.int32, device=f"cuda:{rank}") + ok_tensor = torch.tensor([int(ok)], dtype=torch.int32, device="cuda") all_ok = gather_tensor(ok_tensor, R).cpu() if int(all_ok.sum().item()) != R: if not ok and detail: diff --git a/tests/test_combine.py b/tests/test_combine.py index 31416c8..9def56b 100644 --- a/tests/test_combine.py +++ b/tests/test_combine.py @@ -83,7 +83,7 @@ def _random_inputs(case, rank, R, seed=0): - dev = f"cuda:{rank}" + dev = "cuda" gen = torch.Generator(device=dev).manual_seed(seed + rank) hidden = torch.randn(case.S, case.H, dtype=torch.bfloat16, device=dev, generator=gen) weights = torch.rand(case.S, case.K, dtype=torch.float32, device=dev, generator=gen) @@ -180,7 +180,7 @@ def _combine_global_reference(ctx, case, rank, R, hidden, dst, cu_seqlens, all_cu = gather_tensor(cu_seqlens.contiguous(), R) global_buf = torch.zeros(R, NvS_padded, case.H, dtype=torch.bfloat16, - device=f"cuda:{rank}") + device="cuda") for src_r in range(R): for s in range(case.S): for k in range(case.K): @@ -196,7 +196,7 @@ def _combine_global_reference(ctx, case, rank, R, hidden, dst, cu_seqlens, ) expert_fn(global_buf[dest_r], all_cu[dest_r], expert_ids) - ref = torch.zeros(case.S, case.H, dtype=torch.float32, device=f"cuda:{rank}") + ref = torch.zeros(case.S, case.H, dtype=torch.float32, device="cuda") local_dst = all_dst[rank] for s in range(case.S): for k in range(case.K): @@ -354,7 +354,7 @@ def test_combine_rejects_bad_inputs(dist_env): _hidden, _weights, dst, _cu, _expert_ids, plan, hidden_user, _weights_user = _dispatch_inputs( ctx, case, rank, R ) - output = torch.empty(case.S, case.H, dtype=torch.bfloat16, device=f"cuda:{rank}") + output = torch.empty(case.S, case.H, dtype=torch.bfloat16, device="cuda") with pytest.raises(TypeError, match="hidden_sh"): buffer.combine(hidden_sh=output, plan=plan, hidden_nvsh=hidden_user) @@ -367,10 +367,10 @@ def test_combine_rejects_bad_inputs(dist_env): plan=plan, hidden_nvsh=hidden_user, route_weights_nvs=torch.empty( - case.S, case.K, dtype=torch.float32, device=f"cuda:{rank}" + case.S, case.K, dtype=torch.float32, device="cuda" ), ) with pytest.raises(AssertionError, match="output_sk"): bad_output_sk = torch.empty(case.S, case.K, dtype=torch.bfloat16, - device=f"cuda:{rank}") + device="cuda") launch_combine(ctx, output, dst, output_sk=bad_output_sk) diff --git a/tests/test_dispatch.py b/tests/test_dispatch.py index b9e3397..75846da 100644 --- a/tests/test_dispatch.py +++ b/tests/test_dispatch.py @@ -136,18 +136,18 @@ def _traceable_hidden(rank, S, H): - hidden = torch.zeros(S, H, dtype=torch.bfloat16, device=f"cuda:{rank}") + hidden = torch.zeros(S, H, dtype=torch.bfloat16, device="cuda") hidden_i16 = hidden.view(torch.int16) - s_idx = torch.arange(S, dtype=torch.int32, device=f"cuda:{rank}") + s_idx = torch.arange(S, dtype=torch.int32, device="cuda") hidden_i16[:, 0] = s_idx.to(torch.int16) if H > 1: - hidden_i16[:, 1] = torch.full((S,), rank, dtype=torch.int16, device=f"cuda:{rank}") + hidden_i16[:, 1] = torch.full((S,), rank, dtype=torch.int16, device="cuda") return hidden def _traceable_weights(rank, S, K): weights_i32 = ( - torch.arange(S * K, dtype=torch.int32, device=f"cuda:{rank}") + torch.arange(S * K, dtype=torch.int32, device="cuda") .reshape(S, K) .add_(rank * S * K) ) @@ -315,21 +315,21 @@ def test_dispatch_scatters_hidden_and_weights_by_dst(dist_env, case): def test_dispatch_dedup_plan_clears_padding_with_weights(dist_env, case): rank, R = dist_env ctx = init_case(case, R) - hidden = torch.randn(case.S, case.H, dtype=torch.bfloat16, device=f"cuda:{rank}") - weights = torch.rand(case.S, case.K, dtype=torch.float32, device=f"cuda:{rank}") + hidden = torch.randn(case.S, case.H, dtype=torch.bfloat16, device="cuda") + weights = torch.rand(case.S, case.K, dtype=torch.float32, device="cuda") ctx["hidden_buf_local"].fill_(7) ctx["weights_buf_local"].fill_(0x55555555) hidden_user = torch.empty( (int(ctx["NvS"]), case.H), dtype=torch.bfloat16, - device=f"cuda:{rank}", + device="cuda", ) hidden_user.fill_(7) weights_user = torch.empty( (int(ctx["NvS"]),), dtype=torch.float32, - device=f"cuda:{rank}", + device="cuda", ) weights_user.view(torch.int32).fill_(0x55555555) torch.cuda.synchronize() @@ -388,8 +388,8 @@ def test_dispatch_saved_plan_hidden_only_reuses_dst_and_skips_weights(dist_env): dedup_a_snapshot = clone_dedup_plan_fields(plan_a) case_b = replace(case, routing="all_local") - hidden_b = torch.randn(case.S, case.H, dtype=torch.bfloat16, device=f"cuda:{rank}") - weights_b = torch.rand(case.S, case.K, dtype=torch.float32, device=f"cuda:{rank}") + hidden_b = torch.randn(case.S, case.H, dtype=torch.bfloat16, device="cuda") + weights_b = torch.rand(case.S, case.K, dtype=torch.float32, device="cuda") topk_b, tpe_b = make_topk(case_b, rank, R) plan_b, _cu = allocate_planning_outputs(ctx) launch_planning(ctx, topk_b.reshape(-1).contiguous(), tpe_b, _cu, plan_b) @@ -460,7 +460,7 @@ def test_dispatch_rejects_bad_inputs(dist_env): ctx = init_case(case, R) topk, tpe = make_topk(case, rank, R) hidden = _traceable_hidden(rank, case.S, case.H) - weights = torch.rand(case.S, case.K, dtype=torch.float32, device=f"cuda:{rank}") + weights = torch.rand(case.S, case.K, dtype=torch.float32, device="cuda") plan, _cu = allocate_planning_outputs(ctx) launch_planning(ctx, topk.reshape(-1).contiguous(), tpe, _cu, plan) diff --git a/tests/test_e2e.py b/tests/test_e2e.py index da6ca8f..a05e19b 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -15,6 +15,7 @@ from moonep import Buffer, MoonEPCommPlan from tests.kernel_test_utils import ( clone_dedup_plan_fields, + local_device_index, dedup_plan_fields_equal, dedup_plan_semantic_errors, ) @@ -23,12 +24,12 @@ def setup(): dist.init_process_group(backend="nccl") rank = dist.get_rank() - torch.cuda.set_device(rank) + torch.cuda.set_device(local_device_index()) return rank, dist.get_world_size() def make_inputs(rank, S, H, K, E, seed=0): - dev = f"cuda:{rank}" + dev = "cuda" g = torch.Generator(device=dev).manual_seed(seed + rank) hidden = torch.randn(S, H, dtype=torch.bfloat16, device=dev, generator=g) weights = torch.rand(S, K, dtype=torch.float32, device=dev, generator=g) @@ -38,7 +39,7 @@ def make_inputs(rank, S, H, K, E, seed=0): def make_remote_expert(rank, R, E, H, Hp, owner_offset=1): - dev = f"cuda:{rank}" + dev = "cuda" padded_E = pad_dim0_for_alignment([E, H, Hp], torch.bfloat16) owners = [] for owner in range(R): @@ -56,7 +57,7 @@ def make_remote_expert(rank, R, E, H, Hp, owner_offset=1): if padded_E > E: mapped[E:].zero_() torch.cuda.synchronize() - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[local_device_index()]) owners.append(mapped[:E]) remote_owner = (rank + owner_offset) % R @@ -67,7 +68,7 @@ def make_remote_expert(rank, R, E, H, Hp, owner_offset=1): def make_full_weight(rank, remote_expert, B, source_offset): E, H, Hp = remote_expert.shape - dev = f"cuda:{rank}" + dev = "cuda" full_weight = torch.empty(E + B, H, Hp, dtype=torch.bfloat16, device=dev) full_weight[:E].copy_(remote_expert) if source_offset: @@ -133,7 +134,7 @@ def reduce_base(R, B, H, Hp, offset, dev): def make_grad_reduce_args(rank, R, E, B, H, Hp, offsets): - dev = f"cuda:{rank}" + dev = "cuda" full_E = E + B return { "full_gate_grad": grad_base(full_E, H, Hp, offsets[0], dev).contiguous(), @@ -146,7 +147,7 @@ def make_grad_reduce_args(rank, R, E, B, H, Hp, offsets): def expected_local_grad(rank, R, E, H, Hp, full_offset, reduce_offset, experts_to_copy): - dev = f"cuda:{rank}" + dev = "cuda" expected = grad_base(E, H, Hp, full_offset, dev) reduce_vals = reduce_base(R, experts_to_copy.shape[1], H, Hp, reduce_offset, dev) for src_rank in range(R): diff --git a/tests/test_grad_reduce.py b/tests/test_grad_reduce.py index c4bfe4c..c7dd263 100644 --- a/tests/test_grad_reduce.py +++ b/tests/test_grad_reduce.py @@ -28,6 +28,7 @@ from moonep import Buffer from moonep.buffer import create_nvl_dist_tensor, pad_dim0_for_alignment from moonep.grad_reduce import launch_grad_reduce +from tests.kernel_test_utils import local_device_index # -------------------------------------------------------------------------- @@ -172,7 +173,7 @@ def _expected_for_rank(rank, R, E, B, plan_cpu, grads_fn, slot_fn): def _assert_all_ranks(ok, rank, label): - ok_tensor = torch.tensor([int(ok)], dtype=torch.int32, device=f"cuda:{rank}") + ok_tensor = torch.tensor([int(ok)], dtype=torch.int32, device="cuda") dist.all_reduce(ok_tensor, op=dist.ReduceOp.MIN) assert int(ok_tensor.item()) == 1, label @@ -358,7 +359,7 @@ def dist_env(): if owns_process_group: dist.init_process_group(backend="nccl") rank = dist.get_rank() - torch.cuda.set_device(rank) + torch.cuda.set_device(local_device_index()) R = dist.get_world_size() if R < 2: if owns_process_group: @@ -367,13 +368,13 @@ def dist_env(): yield rank, R - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[local_device_index()]) if owns_process_group: dist.destroy_process_group() def run_case(rank, R, case): - dev = f"cuda:{rank}" + dev = "cuda" epn = case["epn"] E = R * epn H = case["H"] @@ -398,7 +399,7 @@ def run_case(rank, R, case): remote_expert_grads = grads_fn().contiguous() reduce_buffers[rank].copy_(slot_fn(rank)) torch.cuda.synchronize() - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[local_device_index()]) launch_grad_reduce( remote_expert_grads, @@ -412,7 +413,7 @@ def run_case(rank, R, case): grid_sync_bar=ctx['grid_sync_bar'], ) torch.cuda.synchronize() - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[local_device_index()]) _verify(rank, R, E, B, plan, grads_fn, slot_fn, remote_expert_grads, reduce_buffers, @@ -424,7 +425,7 @@ def run_case(rank, R, case): f"H={H}, Hp={Hp}, num_sms={num_sms}" ) - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[local_device_index()]) finally: buffer.destroy() @@ -444,7 +445,7 @@ def test_grad_reduce_repeated_launch(dist_env): an empty plan so a zero-work launch is also proven not to wedge or desync the barrier state for the round after it.""" rank, R = dist_env - dev = f"cuda:{rank}" + dev = "cuda" epn, H, Hp, base_B, num_sms = 4, 256, 128, 5, 16 E = R * epn B = pad_dim0_for_alignment([base_B, H, Hp], torch.float32) @@ -468,7 +469,7 @@ def test_grad_reduce_repeated_launch(dist_env): remote_expert_grads = grads_fn() reduce_buffers[rank].copy_(slot_fn(rank)) torch.cuda.synchronize() - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[local_device_index()]) launch_grad_reduce( remote_expert_grads, @@ -482,7 +483,7 @@ def test_grad_reduce_repeated_launch(dist_env): grid_sync_bar=ctx['grid_sync_bar'], ) torch.cuda.synchronize() - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[local_device_index()]) _verify(rank, R, E, B, plan, grads_fn, slot_fn, remote_expert_grads, reduce_buffers, @@ -490,7 +491,7 @@ def test_grad_reduce_repeated_launch(dist_env): if rank == 0: print(f" [PASS] repeated_launch: R={R}, rounds={len(rounds)}") - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[local_device_index()]) finally: buffer.destroy() diff --git a/tests/test_prefetch.py b/tests/test_prefetch.py index 46ae8ba..e4f1643 100644 --- a/tests/test_prefetch.py +++ b/tests/test_prefetch.py @@ -26,6 +26,7 @@ from moonep.buffer import create_nvl_single_owner_tensor, pad_dim0_for_alignment from moonep.prefetch import launch_prefetch +from tests.kernel_test_utils import local_device_index def _random_experts(E, B, seed): @@ -162,7 +163,7 @@ def dist_env(): if owns_process_group: dist.init_process_group(backend="nccl") rank = dist.get_rank() - torch.cuda.set_device(rank) + torch.cuda.set_device(local_device_index()) R = dist.get_world_size() if R < 2: if owns_process_group: @@ -171,7 +172,7 @@ def dist_env(): yield rank, R - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[local_device_index()]) if owns_process_group: dist.destroy_process_group() @@ -189,15 +190,15 @@ def make_single_owner_experts(rank, R, E, H, Hp): ) if rank == owner: seed = 2026 + owner + E * 13 + H * 17 + Hp * 19 - gen = torch.Generator(device=f"cuda:{rank}").manual_seed(seed) + gen = torch.Generator(device="cuda").manual_seed(seed) mapped[:E].copy_( torch.randn(E, H, Hp, dtype=torch.bfloat16, - device=f"cuda:{rank}", generator=gen) + device="cuda", generator=gen) ) if padded_E > E: mapped[E:].zero_() torch.cuda.synchronize() - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[local_device_index()]) owners.append(mapped[:E]) return owners @@ -215,7 +216,7 @@ def run_case(rank, R, case): Hp = case["Hp"] B = case["B"] num_sms = case["num_sms"] - dev = f"cuda:{rank}" + dev = "cuda" assert H % 128 == 0 and Hp % 128 == 0, \ f"{case['name']}: H/Hp must be multiples of 128" @@ -268,7 +269,7 @@ def run_case(rank, R, case): f"H={H}, Hp={Hp}, num_sms={num_sms}" ) - dist.barrier(device_ids=[rank]) + dist.barrier(device_ids=[local_device_index()]) @pytest.mark.parametrize("case", CASES, ids=[case["name"] for case in CASES]) From 7745ffa00532d9086b49bab84a65b17f687ede14 Mon Sep 17 00:00:00 2001 From: Nyakku Shigure Date: Fri, 7 Aug 2026 21:31:23 +0800 Subject: [PATCH 3/5] Fix fabric support probing (#28) --- csrc/nvl_shared_buffer.cuh | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/csrc/nvl_shared_buffer.cuh b/csrc/nvl_shared_buffer.cuh index 408e179..6f5e965 100644 --- a/csrc/nvl_shared_buffer.cuh +++ b/csrc/nvl_shared_buffer.cuh @@ -47,8 +47,22 @@ static inline bool nvl_fabric_supported() { CUresult err = cuDeviceGetAttribute(&supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, cu_device); // Older drivers do not know the attribute at all. - if (err != CUDA_SUCCESS) return false; - return supported != 0; + if (err != CUDA_SUCCESS || !supported) return false; + + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.location.id = device_id; + prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; + size_t size; + if (cuMemGetAllocationGranularity( + &size, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED) != CUDA_SUCCESS) + return false; + + CUmemGenericAllocationHandle handle; + if (cuMemCreate(&handle, size, &prop, 0) != CUDA_SUCCESS) return false; + CUCHECK(cuMemRelease(handle)); + return true; } static inline size_t nvl_granularity_for(int device_id, From b5a0e7f7922c5f1f271ff06643fb053cd3a3d29b Mon Sep 17 00:00:00 2001 From: asp0ex Date: Thu, 13 Aug 2026 15:02:05 +0800 Subject: [PATCH 4/5] Split hidden/route-weights zero-copy control via router_weights_zero_copy (#31) --- benchmarks/bench_vs_deepep.py | 7 ++- moonep/api.py | 110 +++++++++++++++++++++++----------- tests/test_e2e.py | 3 +- 3 files changed, 81 insertions(+), 39 deletions(-) diff --git a/benchmarks/bench_vs_deepep.py b/benchmarks/bench_vs_deepep.py index 93d0830..3f49bf1 100644 --- a/benchmarks/bench_vs_deepep.py +++ b/benchmarks/bench_vs_deepep.py @@ -225,7 +225,7 @@ def prepare(self, topk, tpe, hidden, weights): # Untimed full dispatch (public API): materializes the plan + dedup # structures used by saved-plan dispatch/combine below. _, _, _, self.plan = self.buffer.dispatch( - hidden, weights, topk, tpe, zero_copy=True) + hidden, weights, topk, tpe, zero_copy=True, router_weights_zero_copy=True) # Scratch plan for the separate (exact) planning measurement. self._launch_planning(self.ctx, self._topk_flat, tpe, self.cu_seqlens, self.plan_scratch) @@ -259,7 +259,8 @@ def dispatch_fwd(self): # full fwd (public API): inter_rank_sync + planning + dispatch + # epilogue (zero_copy, no boundary copy), then prefetch weight. _, _, _, self.plan = self.buffer.dispatch( - self.hidden, self.weights, self.topk, self.tpe, zero_copy=True) + self.hidden, self.weights, self.topk, self.tpe, + zero_copy=True, router_weights_zero_copy=True) self._prefetch() def dispatch_bwd(self): @@ -278,7 +279,7 @@ def combine_bwd(self): # path and is not included in this benchmark. self.buffer.combine(plan=self.plan, hidden_nvsh=self.shard_view, route_weights_nvs=self.weights_view, - zero_copy=True) + zero_copy=True, router_weights_zero_copy=True) def destroy(self): self.buffer.destroy() diff --git a/moonep/api.py b/moonep/api.py index 8d0cbc7..9a33cd0 100644 --- a/moonep/api.py +++ b/moonep/api.py @@ -502,6 +502,26 @@ def __init__( def destroyed(self) -> bool: return self._destroyed + @property + def hidden_nvsh_buffer_view(self) -> torch.Tensor: + """The local rank's [NvS, H] bf16 communication buffer. + + Exposed for zero-copy integration: callers may write expert outputs + into this view and hand it back to ``combine(zero_copy=True)``. The + view aliases persistent comm state that every dispatch/combine on + this Buffer overwrites — never let it (or any tensor sharing its + storage) cross into autograd-saved state. + """ + return self._require_ctx()['hidden_buf_local'] + + @property + def router_weight_buffer_view(self) -> torch.Tensor: + """fp32 view of the local rank's [NvS] route-weights comm buffer. + + Same aliasing/lifetime rules as ``hidden_nvsh_buffer_view``. + """ + return self._require_ctx()['weights_buf_local'].view(torch.float32) + def _require_ctx(self) -> dict: assert not self._destroyed, "MoonEP Buffer has been destroyed" assert self._ctx is not None, "MoonEP Buffer is not initialized" @@ -597,6 +617,7 @@ def _run_dispatch_on_current_stream( *, inter_rank_sync: bool, zero_copy: bool, + route_weights_zero_copy: bool, ) -> None: if inter_rank_sync: launch_inter_rank_sync(ctx) @@ -621,13 +642,14 @@ def _run_dispatch_on_current_stream( # In-place duplicate expansion on the NVL shard: after this the shard # holds the full user-visible [NvS, H] layout. launch_dispatch_epilogue(ctx, plan, pdl_launch=self.enable_pdl) + # master-style boundary copies (same stream, plain SM copies), gated + # independently per output tensor. if not zero_copy: - # master-style boundary copies (same stream, plain SM copies). hidden_nvsh.copy_(ctx['hidden_buf_local']) - if route_weights_nvs is not None: - route_weights_nvs.copy_( - ctx['weights_buf_local'].view(torch.float32) - ) + if route_weights_nvs is not None and not route_weights_zero_copy: + route_weights_nvs.copy_( + ctx['weights_buf_local'].view(torch.float32) + ) def _run_combine_on_current_stream( self, @@ -640,18 +662,20 @@ def _run_combine_on_current_stream( *, inter_rank_sync: bool, zero_copy: bool, + router_weights_zero_copy: bool, ) -> None: # Pre-staging sync keeps its master-era position; combine's own entry # cross_rank_barrier publishes the staged + accumulated NVL writes. if inter_rank_sync: launch_inter_rank_sync(ctx) + # master-style boundary copies into the shard (same stream), gated + # independently per input tensor. if not zero_copy: - # master-style boundary copies into the shard (same stream). ctx['hidden_buf_local'].copy_(hidden_nvsh) - if route_weights_nvs is not None: - ctx['weights_buf_local'].copy_( - route_weights_nvs.view(torch.int32) - ) + if route_weights_nvs is not None and not router_weights_zero_copy: + ctx['weights_buf_local'].copy_( + route_weights_nvs.view(torch.int32) + ) # In-place fp32 accumulation of duplicate rows into their primary. launch_combine_prologue(ctx, plan, pdl_trigger=self.enable_pdl) launch_combine( @@ -693,6 +717,7 @@ def dispatch( *, inter_rank_sync: bool = True, zero_copy: bool = False, + router_weights_zero_copy: bool = False, ): """dispatch fwd: run planning (unless reusing a plan) and scatter tokens to their expert-grouped positions on remote ranks. @@ -714,15 +739,21 @@ def dispatch( the return value. inter_rank_sync: run a CuTe DSL rank sync before planning (default True). - zero_copy: return views of the communication buffer - (``hidden_buf_local`` and the fp32 view of - ``weights_buf_local``) instead of fresh tensors. The views - alias state that the next dispatch/combine on this Buffer - overwrites, so callers must not keep them across communication - calls (in particular autograd must not save them for backward + zero_copy: return a view of the communication buffer + (``hidden_buf_local``) instead of a fresh tensor. The view + aliases state that the next dispatch/combine on this Buffer + overwrites, so callers must not keep it across communication + calls (in particular autograd must not save it for backward — that is exactly the case that requires ``zero_copy=False``). Row content is only defined within the ``cu_seqlens``-covered padded segments. + router_weights_zero_copy: like ``zero_copy`` but for + ``route_weights_nvs`` (the fp32 view of ``weights_buf_local``). + Defaults to False even when ``zero_copy=True``: the weights + view is the dangerous special case (tiny tensor, commonly + saved into autograd state by training frameworks), so callers + that only consume it before the next dispatch — e.g. + inference — opt in explicitly. Returns: ``(hidden_nvsh, route_weights_nvs, cu_seqlens, plan)``, plus a @@ -754,15 +785,15 @@ def dispatch( if zero_copy: hidden_nvsh = ctx['hidden_buf_local'] - route_weights_nvs = ( - ctx['weights_buf_local'].view(torch.float32) - if route_weights_sk is not None else None - ) else: hidden_nvsh = torch.empty_like(ctx['hidden_buf_local']) - route_weights_nvs = ( - torch.empty(ctx['NvS'], dtype=torch.float32, device=ctx['meta_buf'].device) - if route_weights_sk is not None else None + if route_weights_sk is None: + route_weights_nvs = None + elif router_weights_zero_copy: + route_weights_nvs = ctx['weights_buf_local'].view(torch.float32) + else: + route_weights_nvs = torch.empty( + ctx['NvS'], dtype=torch.float32, device=ctx['meta_buf'].device ) if not async_finish: @@ -776,6 +807,7 @@ def dispatch( route_weights_nvs, inter_rank_sync=inter_rank_sync, zero_copy=zero_copy, + route_weights_zero_copy=router_weights_zero_copy, ) return hidden_nvsh, route_weights_nvs, cu_seqlens, plan @@ -808,6 +840,7 @@ def dispatch( route_weights_nvs, inter_rank_sync=inter_rank_sync, zero_copy=zero_copy, + route_weights_zero_copy=router_weights_zero_copy, ) done = comm.record_event() @@ -887,6 +920,7 @@ def combine( inter_rank_sync: bool = True, *, zero_copy: bool = False, + router_weights_zero_copy: bool = False, ): """combine fwd: gather expert outputs from the NVL buffer and K-sum back to token-major [S, H]. @@ -904,12 +938,15 @@ def combine( async_finish: run on the comm stream and return a CUDA event. inter_rank_sync: run a CuTe DSL rank sync before staging (default True). - zero_copy: ``hidden_nvsh`` (and ``route_weights_nvs`` when given) - must be exactly the views returned by a ``zero_copy=True`` - dispatch — the caller's FFN writes its output in place on the - shard and no boundary copy is performed (asserted via - ``data_ptr()``). With ``zero_copy=False`` the inputs are - ordinary tensors that are first copied into the shard. + zero_copy: ``hidden_nvsh`` must be exactly the view returned by a + ``zero_copy=True`` dispatch — the caller's FFN writes its + output in place on the shard and no boundary copy is performed + (asserted via ``data_ptr()``). With ``zero_copy=False`` the + input is an ordinary tensor that is first copied into the + shard. + router_weights_zero_copy: same contract for ``route_weights_nvs`` + (the fp32 view of ``weights_buf_local``); default False, i.e. + an ordinary tensor that is copied into the shard first. Returns: ``(hidden_sh, route_weights_sk, event)``: @@ -938,12 +975,13 @@ def combine( "combine(zero_copy=True): hidden_nvsh must alias the NVL shard " "view returned by dispatch(zero_copy=True)" ) - if route_weights_nvs is not None: - assert route_weights_nvs.data_ptr() == \ - ctx['weights_buf_local'].data_ptr(), ( - "combine(zero_copy=True): route_weights_nvs must alias " - "the NVL weights view returned by dispatch(zero_copy=True)" - ) + if router_weights_zero_copy and route_weights_nvs is not None: + assert route_weights_nvs.data_ptr() == \ + ctx['weights_buf_local'].data_ptr(), ( + "combine(router_weights_zero_copy=True): route_weights_nvs must " + "alias the NVL weights view returned by " + "dispatch(router_weights_zero_copy=True)" + ) hidden_sh = torch.empty( int(ctx['S']), @@ -971,6 +1009,7 @@ def combine( route_weights_sk, inter_rank_sync=inter_rank_sync, zero_copy=zero_copy, + router_weights_zero_copy=router_weights_zero_copy, ) return hidden_sh, route_weights_sk, None @@ -1001,6 +1040,7 @@ def combine( route_weights_sk, inter_rank_sync=inter_rank_sync, zero_copy=zero_copy, + router_weights_zero_copy=router_weights_zero_copy, ) done = comm.record_event() diff --git a/tests/test_e2e.py b/tests/test_e2e.py index a05e19b..a2d0878 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -344,7 +344,7 @@ def test_e2e(): # identity on the shard, combine consumes the views in place. Must match # the zero_copy=False result bit-exactly. h_zc, w_zc, _, plan_zc = buffer.dispatch( - hidden, weights, topk, tpe, zero_copy=True, + hidden, weights, topk, tpe, zero_copy=True, router_weights_zero_copy=True, ) assert h_zc.data_ptr() == buffer._require_ctx()['hidden_buf_local'].data_ptr(), \ "dispatch(zero_copy=True) must return the NVL shard view" @@ -357,6 +357,7 @@ def test_e2e(): hidden_nvsh=h_zc, route_weights_nvs=w_zc, zero_copy=True, + router_weights_zero_copy=True, ) torch.cuda.synchronize() assert torch.equal(out_sync_snap, out_zc), "zero_copy combine hidden mismatch" From 2bd860b4dd083df62b79d5e916fca71ec5742228 Mon Sep 17 00:00:00 2001 From: jxp Date: Thu, 13 Aug 2026 02:30:40 -0700 Subject: [PATCH 5/5] Support MXFP4 expert weights in remote prefetch (#29) --- benchmarks/bench_prefetch.py | 161 +++++++++++++++++++++++++++-------- moonep/api.py | 48 +++++++++-- moonep/buffer.py | 9 +- moonep/prefetch.py | 91 +++++++++++++++----- tests/test_prefetch.py | 86 ++++++++++++++++--- 5 files changed, 310 insertions(+), 85 deletions(-) diff --git a/benchmarks/bench_prefetch.py b/benchmarks/bench_prefetch.py index 47230fa..d6f71e0 100644 --- a/benchmarks/bench_prefetch.py +++ b/benchmarks/bench_prefetch.py @@ -20,6 +20,27 @@ EXP_H = 3584 EXP_HP = 3072 +K3_BF16_LAYER = [ + {"transposed": True}, # gate -> (I, H) + {"transposed": True}, # up -> (I, H) + {}, # down -> (H, I) +] +K3_MXFP4_LAYER = [ + {"transposed": True, "pack": 2}, # gate packed -> (I, H/2) + {"transposed": True, "pack": 2}, # up packed -> (I, H/2) + {"pack": 2}, # down packed -> (H, I/2) + {"scale": True}, # gate scale -> (128, 2688) + {"scale": True}, # up scale + {"scale": True}, # down scale +] + +_DT_LABEL = { + torch.bfloat16: "bf16", + torch.int8: "i8", + torch.uint8: "u8", +} +_DT_LABEL_TO_TORCH = {v: k for k, v in _DT_LABEL.items()} + # All cases: rank0-initiated, num_sms == 32. Each remote expert is prefetched # a *different* number of times in 0..3 (R-1=3 is the max), via `counts[e]`. # 0 means that expert is skipped. base_B > epn leaves idle columns so the @@ -60,9 +81,49 @@ "counts": [3, 0, 2, 1, 3, 0, 2, 1]}, {"label": "tiny_512x512", "epn": 8, "H": 512, "Hp": 512, "base_B": 14, "counts": [3, 3, 2, 3, 1, 0, 2, 3]}, + # --- Quantized experts ------------------------------------------------- + {"label": "mxfp4_gate_up", "epn": 8, "H": EXP_H, "Hp": EXP_HP, + "transposed": True, "pack": 2, + "dtype": torch.uint8, "base_B": 14, "counts": [3, 0, 2, 1, 3, 0, 2, 1]}, + {"label": "mxfp4_down", "epn": 8, "H": EXP_H, "Hp": EXP_HP, "pack": 2, + "dtype": torch.uint8, "base_B": 14, "counts": [3, 0, 2, 1, 3, 0, 2, 1]}, + {"label": "mxfp4_scale", "epn": 8, "H": EXP_H, "Hp": EXP_HP, "scale": True, + "dtype": torch.uint8, "base_B": 14, "counts": [3, 0, 2, 1, 3, 0, 2, 1]}, + {"label": "k3_layer_bf16", "epn": 8, "H": EXP_H, "Hp": EXP_HP, "base_B": 14, + "parts": K3_BF16_LAYER, "counts": [3, 0, 2, 1, 3, 0, 2, 1]}, + {"label": "k3_layer_mxfp4", "epn": 8, "H": EXP_H, "Hp": EXP_HP, "base_B": 14, + "dtype": torch.uint8, + "parts": K3_MXFP4_LAYER, "counts": [3, 0, 2, 1, 3, 0, 2, 1]}, ] +def derive_extents(H, Hp, part): + if part.get("scale"): + # One ue8m0 byte per 32 values, re-cut into whole 128x128 tiles. + nbytes = H * Hp // 32 + return 128, nbytes // 128 + out, contracted = (Hp, H) if part.get("transposed") else (H, Hp) + return out, contracted // int(part.get("pack", 1)) + + +def resolve_parts(case): + H, Hp = int(case["H"]), int(case["Hp"]) + default_dt = case.get("dtype", torch.bfloat16) + resolved = [] + for part in case.get("parts") or [case]: + th, thp = derive_extents(H, Hp, part) + resolved.append((th, thp, part.get("dtype", default_dt))) + return resolved + + +def fill_random(shape, dtype, gen, dev="cuda"): + if dtype.is_floating_point: + return torch.randn(shape, dtype=dtype, device=dev, generator=gen) + info = torch.iinfo(dtype) + return torch.randint(info.min, info.max, shape, dtype=dtype, device=dev, + generator=gen) + + def setup(): dist.init_process_group(backend="nccl") rank = dist.get_rank() @@ -90,30 +151,36 @@ def bench_case(case, args, rank, R): dev = "cuda" epn = int(case["epn"]) E = R * epn - H = int(case["H"]) - Hp = int(case["Hp"]) - B = pad_dim0_for_alignment([int(case["base_B"]), H, Hp], torch.bfloat16) + parts = resolve_parts(case) + th0, thp0, dt0 = parts[0] + B = pad_dim0_for_alignment([int(case["base_B"]), th0, thp0], dt0) counts = case.get("counts") or [int(case.get("nremote", 1))] * epn num_sms = NUM_SMS - assert H % 128 == 0 and Hp % 128 == 0, \ - f"{case['label']}: H and Hp must be multiples of 128" + for th, thp, _dt in parts: + assert th % 128 == 0 and thp % 128 == 0, \ + f"{case['label']}: derived extents must be multiples of 128, " \ + f"got ({th}, {thp})" assert len(counts) == epn, f"{case['label']}: counts must have epn={epn} entries" assert all(0 <= c < R for c in counts), f"{case['label']}: each count must be in [0, R)" - # The remote expert table lives on rank1; rank0 reads it over NVLink. - padded_E = pad_dim0_for_alignment([E, H, Hp], torch.bfloat16) - mapped = create_nvl_single_owner_tensor( - [padded_E, H, Hp], torch.bfloat16, owner_rank=1, local_rank=rank - ) - remote_expert = mapped[:E] - prefetch_buffers = torch.empty(R * B, H, Hp, dtype=torch.bfloat16, device=dev) - - if rank == 1: - gen = torch.Generator(device=dev).manual_seed(321 + rank) - remote_expert.copy_(torch.randn( - E, H, Hp, dtype=torch.bfloat16, device=dev, generator=gen - )) + # The remote expert table lives on one owner; rank0 reads it over NVLink + owner_rank = int(args.owner_rank) + remote_experts, prefetch_buffers = [], [] + for i, (th, thp, dt) in enumerate(parts): + padded_E = pad_dim0_for_alignment([E, th, thp], dt) + mapped = create_nvl_single_owner_tensor( + [padded_E, th, thp], dt, owner_rank=owner_rank, local_rank=rank + ) + remote = mapped[:E] + if rank == owner_rank: + gen = torch.Generator(device=dev).manual_seed(321 + rank + i) + remote.copy_(fill_random((E, th, thp), dt, gen, dev)) + remote_experts.append(remote) + prefetch_buffers.append( + torch.empty(R * B, th, thp, dtype=dt, device=dev) + ) + plan = expert_plan(R, B, epn, counts, dev) experts_to_copy = plan.flatten() if rank == 0 else \ torch.full((R * B,), -1, dtype=torch.int32, device=dev) @@ -121,12 +188,13 @@ def bench_case(case, args, rank, R): dist.barrier(device_ids=[torch.cuda.current_device()]) def prefetch_once(): - launch_prefetch( - remote_expert, - prefetch_buffers, - experts_to_copy, - num_sms=num_sms, - ) + for remote, buf in zip(remote_experts, prefetch_buffers): + launch_prefetch( + remote, + buf, + experts_to_copy, + num_sms=num_sms, + ) # Warmup (also JIT-compiles the kernel) then capture the iters loop into a # single CUDA graph to strip launch/python overhead. @@ -162,22 +230,25 @@ def prefetch_once(): worst_us = start.elapsed_time(end) * 1e3 / args.iters if rank == 0 else 0.0 # Critical-path traffic for the prefetching rank: read every consumed slot - # from the remote table over NVLink (2B/elem), then write it to the local - # prefetch buffer (2B/elem). + # from the remote table over NVLink, then write it to the local prefetch + # buffer. slots = int(sum(counts)) - tile = H * Hp * 2 - bytes_per_rank = slots * tile * 2 + per_slot = sum(th * thp * dt.itemsize for th, thp, dt in parts) + bytes_per_rank = slots * per_slot * 2 bw_gbs = bytes_per_rank / worst_us * 1e6 / 1e9 if worst_us > 0 else 0.0 # pure NVLink read traffic: only the remote expert-table reads (buffer # writes are local HBM, off the NVLink path). - comm_gbs = slots * tile / worst_us * 1e6 / 1e9 if worst_us > 0 else 0.0 + comm_gbs = slots * per_slot / worst_us * 1e6 / 1e9 if worst_us > 0 else 0.0 dist.barrier(device_ids=[torch.cuda.current_device()]) - return worst_us, bytes_per_rank / 1e6, bw_gbs, comm_gbs, E, B, slots + dt_label = (_DT_LABEL.get(dt0, str(dt0)) + if len({dt for _, _, dt in parts}) == 1 else "mix") + return (worst_us, bytes_per_rank / 1e6, bw_gbs, comm_gbs, E, B, slots, + len(parts), dt_label, int(case["H"]), int(case["Hp"])) def explicit_single_config_requested(argv): - shape_flags = ("--epn", "--H", "--Hp", "--B", "--nremote") + shape_flags = ("--epn", "--H", "--Hp", "--B", "--nremote", "--dtype") for arg in argv: for flag in shape_flags: if arg == flag or arg.startswith(flag + "="): @@ -196,8 +267,17 @@ def parse_args(): parser.add_argument("--Hp", type=int, default=3072) parser.add_argument("--B", type=int, default=14) parser.add_argument("--nremote", type=int, default=2) + parser.add_argument("--dtype", choices=sorted(_DT_LABEL_TO_TORCH), + default="bf16", + help="Element type of the prefetched tensor " + "(u8 = packed MXFP4 or its ue8m0 scales)") parser.add_argument("--warmup", type=int, default=5) parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--owner-rank", type=int, default=1, + help="Rank whose GPU physically holds the remote expert " + "rank0 is always the reader. An owner on rank0's own node " + "measures intra-node NVLink, one on another node " + "measures cross-node MNNVL fabric.") parser.add_argument("--no-graph", action="store_true", help="Time plain launches instead of a CUDA graph (for NCU)") return parser.parse_args() @@ -210,6 +290,10 @@ def main(): if args.single and args.suite: raise ValueError("Use only one of --single or --suite") + if not 0 <= args.owner_rank < R: + raise ValueError( + f"--owner-rank must be in [0, {R}), got {args.owner_rank}" + ) run_single = args.single or ( not args.suite and explicit_single_config_requested(sys.argv[1:]) @@ -220,6 +304,7 @@ def main(): "epn": args.epn, "H": args.H, "Hp": args.Hp, + "dtype": _DT_LABEL_TO_TORCH[args.dtype], "base_B": args.B, "nremote": args.nremote, }] @@ -229,20 +314,22 @@ def main(): if rank == 0: print( f"MoonEP Prefetch Benchmark (R={R}, warmup={args.warmup}, " - f"iters={args.iters})" + f"iters={args.iters}, reader=rank0, owner=rank{args.owner_rank}, " + f"num_sms={NUM_SMS})" ) print( - f"{'Config':<20} {'E':>5} {'B':>4} {'H':>7} {'Hp':>6} " + f"{'Config':<20} {'E':>5} {'B':>4} {'dt':>5} {'N':>3} {'H':>7} {'Hp':>6} " f"{'SMs':>5} {'Slots':>6} {'Data(MB)':>10} {'Worst(us)':>10} {'BW(GB/s)':>9} {'CommBW':>8}" ) - print("-" * 101) + print("-" * 111) for case in cases: - worst_us, mb, bw_gbs, comm_gbs, E, B, slots = bench_case(case, args, rank, R) + (worst_us, mb, bw_gbs, comm_gbs, E, B, slots, + ntensor, dt_label, H, Hp) = bench_case(case, args, rank, R) if rank == 0: print( - f"{case['label']:<20} {E:>5} {B:>4} " - f"{case['H']:>7} {case['Hp']:>6} {NUM_SMS:>5} {slots:>6} " + f"{case['label']:<20} {E:>5} {B:>4} {dt_label:>5} {ntensor:>3} " + f"{H:>7} {Hp:>6} {NUM_SMS:>5} {slots:>6} " f"{mb:>10.2f} {worst_us:>10.2f} {bw_gbs:>9.2f} {comm_gbs:>8.2f}" ) diff --git a/moonep/api.py b/moonep/api.py index 9a33cd0..964da9a 100644 --- a/moonep/api.py +++ b/moonep/api.py @@ -68,7 +68,7 @@ from .dispatch_epilogue import launch_dispatch_epilogue from .combine import launch_combine from .combine_prologue import launch_combine_prologue -from .prefetch import launch_prefetch +from .prefetch import _ELEM_TYPES, launch_prefetch, retile_for_prefetch from .grad_reduce import launch_grad_reduce logger = logging.getLogger(__name__) @@ -161,6 +161,7 @@ def _launch_full_weight_prefetches( full_up_weight: torch.Tensor, full_down_weight: torch.Tensor, experts_to_copy: torch.Tensor, + scales: tuple[torch.Tensor, ...] | None = None, ) -> None: E = int(ctx['E']) num_sms = int(ctx['num_sms']) @@ -171,6 +172,14 @@ def _launch_full_weight_prefetches( experts_to_copy, num_sms=num_sms, ) + for full_scale in scales or (): + tiled = retile_for_prefetch(full_scale) + launch_prefetch( + tiled[:E], + tiled[E:], + experts_to_copy, + num_sms=num_sms, + ) def _launch_full_grad_reduces( @@ -699,11 +708,13 @@ def _run_prefetch_weight_on_current_stream( ctx: dict, experts_to_copy: torch.Tensor, weight_prefetch_args, + scale_prefetch_args=None, ) -> None: _launch_full_weight_prefetches( ctx, *weight_prefetch_args, experts_to_copy[int(ctx['rank'])], + scales=scale_prefetch_args, ) def dispatch( @@ -854,6 +865,9 @@ def prefetch_weight( full_gate_weight: torch.Tensor | None = None, full_up_weight: torch.Tensor | None = None, full_down_weight: torch.Tensor | None = None, + full_gate_scale: torch.Tensor | None = None, + full_up_scale: torch.Tensor | None = None, + full_down_scale: torch.Tensor | None = None, ): """Prefetch the remote expert weights selected by ``plan`` into the local prefetch slots (dispatch fwd, weight side). @@ -862,9 +876,14 @@ def prefetch_weight( plan: MoonEPCommPlan returned by ``dispatch``. async_finish: run on the comm stream and return a CUDA event. full_gate_weight / full_up_weight / full_down_weight: - [E+B, H, H'] bf16 contiguous weight tensors; rows [0, E) are - source expert weights, rows [E, E+B) are the prefetch slots - filled by this call. + [E+B, H, H'] contiguous weight tensors; rows [0, E) are source + expert weights, rows [E, E+B) are the prefetch slots filled by + this call. bf16 for unquantized experts, uint8 for MXFP4 (e2m1 + packs two values per byte, so H' is K/2). + full_gate_scale / full_up_scale / full_down_scale: + optional [E+B, ...] contiguous block-scale tensors, same row + convention. Required for quantized experts and omitted for bf16 + ones. Returns: None in synchronous mode, or the comm-stream CUDA event when @@ -882,14 +901,27 @@ def prefetch_weight( assert all(w is not None for w in weight_prefetch_args), \ "prefetch_weight tensors must be provided together" for w in weight_prefetch_args: - assert w.dtype == torch.bfloat16 and w.is_contiguous() + assert w.dtype in _ELEM_TYPES, \ + f"prefetch_weight: unsupported weight dtype {w.dtype}" + assert w.is_contiguous() assert w.ndim == 3 and int(w.shape[0]) == int(ctx['E']) + int(ctx['B']) + scale_prefetch_args = (full_gate_scale, full_up_scale, full_down_scale) + if any(s is not None for s in scale_prefetch_args): + assert all(s is not None for s in scale_prefetch_args), \ + "prefetch_weight scales must be provided together" + for s in scale_prefetch_args: + assert s.is_contiguous() + assert s.ndim >= 2 and int(s.shape[0]) == int(ctx['E']) + int(ctx['B']) + else: + scale_prefetch_args = None + if not async_finish: self._run_prefetch_weight_on_current_stream( ctx, plan.experts_to_copy, weight_prefetch_args, + scale_prefetch_args, ) return None @@ -897,7 +929,10 @@ def prefetch_weight( comm = self._comm_stream assert comm is not None, "MoonEP Buffer communication stream is not initialized" - self._record_streams((plan.experts_to_copy, *weight_prefetch_args), comm) + self._record_streams( + (plan.experts_to_copy, *weight_prefetch_args, *(scale_prefetch_args or ())), + comm, + ) input_ready = main_stream.record_event() comm.wait_event(input_ready) @@ -906,6 +941,7 @@ def prefetch_weight( ctx, plan.experts_to_copy, weight_prefetch_args, + scale_prefetch_args, ) done = comm.record_event() diff --git a/moonep/buffer.py b/moonep/buffer.py index 4954810..87a1ac9 100644 --- a/moonep/buffer.py +++ b/moonep/buffer.py @@ -22,12 +22,6 @@ nvl_multicast_bind_map, ) -_ELEM_SIZE = { - torch.float32: 4, - torch.bfloat16: 2, - torch.int32: 4, -} - # How VMM allocations are shared between the ranks of an EP group. "auto" (the # default) picks fabric handles when the group spans more than one node and the # device supports them, and POSIX fds otherwise; "fabric" / "fd" force one. @@ -103,8 +97,7 @@ def pad_dim0_for_alignment(chunk_shape: list[int], dtype: torch.dtype) -> int: Returns the padded dim0 value (>= chunk_shape[0]). """ - elem_size = _ELEM_SIZE[dtype] - inner_size = elem_size + inner_size = dtype.itemsize for d in chunk_shape[1:]: inner_size *= d # bytes per row diff --git a/moonep/prefetch.py b/moonep/prefetch.py index 4bb2921..e0928a6 100644 --- a/moonep/prefetch.py +++ b/moonep/prefetch.py @@ -8,8 +8,11 @@ - warp 0: GMEM -> SMEM 2D TMA load - warp 1: SMEM -> GMEM 2D TMA store -The initial tile shape is fixed at 128 x 128 bf16 elements. H and H' are +The initial tile shape is fixed at 128 x 128 elements. H and H' are therefore required to be multiples of 128 for this first implementation. +For a scale tensor whose natural trailing extent is K/32 (never a +multiple of 128), re-tile its contiguous per-expert byte range to +``[128, nbytes // 128]`` before calling in. """ import functools @@ -22,10 +25,17 @@ import cutlass.pipeline as pipeline import cutlass.utils as utils import cutlass.cute.nvgpu.cpasync as cpasync -from cutlass import BFloat16, Int32, Int64 +from cutlass import BFloat16, Int8, Int32, Int64, Uint8 from cutlass.cute.runtime import make_ptr +_ELEM_TYPES = { + torch.bfloat16: (BFloat16, 2), + torch.int8: (Int8, 1), + torch.uint8: (Uint8, 1), +} + + class PrefetchKernel: """Persistent 2D TMA remote-expert prefetch.""" @@ -43,25 +53,29 @@ def __init__( B: int, num_sms: int, smem_budget: int, + elem_ty=BFloat16, + elem_bytes: int = 2, ): self.E = E self.H = H self.Hp = Hp self.B = B self.num_sms = num_sms + self.elem_ty = elem_ty + self.elem_bytes = elem_bytes self.stages = self._pick_stages(smem_budget) if self.stages == 0: raise RuntimeError( "prefetch: not enough per-block shared memory for one " - f"{self.M_BLOCK}x{self.N_BLOCK} bf16 tile under budget " - f"{smem_budget} B" + f"{self.M_BLOCK}x{self.N_BLOCK} x {elem_bytes}B tile under " + f"budget {smem_budget} B" ) def _smem_bytes(self, stages: int) -> int: def _round_up(n: int, a: int) -> int: return (n + a - 1) // a * a - tile_bytes = self.M_BLOCK * self.N_BLOCK * 2 + tile_bytes = self.M_BLOCK * self.N_BLOCK * self.elem_bytes return ( _round_up(stages * tile_bytes, 128) + _round_up(stages * 2 * 8, 16) @@ -78,8 +92,8 @@ def _pick_stages(self, smem_budget: int) -> int: @cute.jit def __call__( self, - remote_expert_ptr: cute.Pointer, # bf16 [E, H, H'] - prefetch_buf_ptr: cute.Pointer, # bf16 [B, H, H'] + remote_expert_ptr: cute.Pointer, # elem_ty [E, H, H'] + prefetch_buf_ptr: cute.Pointer, # elem_ty [B, H, H'] experts_ptr: cute.Pointer, # int32 [B] stream: cuda.CUstream, ): @@ -151,7 +165,7 @@ def kernel( M_BLOCK = cutlass.const_expr(self.M_BLOCK) N_BLOCK = cutlass.const_expr(self.N_BLOCK) TILE_ELEMS = cutlass.const_expr(M_BLOCK * N_BLOCK) - TILE_BYTES = cutlass.const_expr(TILE_ELEMS * 2) + TILE_BYTES = cutlass.const_expr(TILE_ELEMS * self.elem_bytes) MTILES = cutlass.const_expr(H // M_BLOCK) NTILES = cutlass.const_expr(Hp // N_BLOCK) TILES_PER_EXPERT = cutlass.const_expr(MTILES * NTILES) @@ -164,7 +178,7 @@ def kernel( exp_tab = smem.allocate_tensor(Int32, cute.make_layout((B,)), byte_alignment=4) slot_tab = smem.allocate_tensor(Int32, cute.make_layout((B,)), byte_alignment=4) stage_smem = smem.allocate_tensor( - BFloat16, + self.elem_ty, cute.make_ordered_layout( (M_BLOCK, N_BLOCK, stages), order=(1, 0, 2), @@ -290,7 +304,9 @@ def _get_compiled( B: int, num_sms: int, device_index: int, + torch_dtype: torch.dtype, ): + elem_ty, elem_bytes = _ELEM_TYPES[torch_dtype] smem_budget = _max_smem_per_block_optin(device_index) - 1024 kernel = PrefetchKernel( E=E, @@ -299,21 +315,45 @@ def _get_compiled( B=B, num_sms=num_sms, smem_budget=smem_budget, + elem_ty=elem_ty, + elem_bytes=elem_bytes, ) - bf16_ptr = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16) + data_ptr = make_ptr(elem_ty, 0, cute.AddressSpace.gmem, assumed_align=16) i32_ptr = make_ptr(Int32, 0, cute.AddressSpace.gmem, assumed_align=4) stream_arg = cuda.CUstream(0) return cute.compile( kernel, - bf16_ptr, - bf16_ptr, + data_ptr, + data_ptr, i32_ptr, stream_arg, ) +def prefetch_retile_nbytes(per_expert_nbytes: int) -> int: + """Round a per-expert byte count up to what ``retile_for_prefetch`` needs. + """ + tile = PrefetchKernel.M_BLOCK * PrefetchKernel.N_BLOCK + return (per_expert_nbytes + tile - 1) // tile * tile + + +def retile_for_prefetch(t: torch.Tensor) -> torch.Tensor: + """View a contiguous ``[N, ...]`` expert tensor as ``[N, 128, X]`` uint8. + """ + assert t.is_contiguous(), "retile_for_prefetch requires a contiguous tensor" + n = int(t.shape[0]) + per_expert = t.nbytes // n if n else 0 + tile = PrefetchKernel.M_BLOCK * PrefetchKernel.N_BLOCK + assert per_expert % tile == 0, ( + f"retile_for_prefetch: per-expert extent {per_expert} bytes is not a " + f"multiple of {tile}; allocate {prefetch_retile_nbytes(per_expert)} " + f"bytes per expert instead" + ) + return t.view(torch.uint8).reshape(n, PrefetchKernel.M_BLOCK, -1) + + def launch_prefetch( remote_expert: torch.Tensor, prefetch_buffers: torch.Tensor, @@ -323,8 +363,10 @@ def launch_prefetch( """Launch remote expert prefetch. Args: - remote_expert: contiguous bf16 tensor shaped [E, H, H']. - prefetch_buffers: contiguous bf16 tensor shaped [B, H, H']. + remote_expert: contiguous tensor shaped [E, H, H']. dtype must be one + of ``_ELEM_TYPES`` -- the copy is type-agnostic, so packed MXFP4 + (uint8) and ue8m0 scales (uint8) go through the same path as bf16. + prefetch_buffers: contiguous tensor shaped [B, H, H'], same dtype. experts_to_copy: contiguous int32 tensor shaped [B]. Entries are expert ids in [0, E), or -1 for unused slots. Unused slots are not written by this kernel. @@ -333,10 +375,16 @@ def launch_prefetch( if prefetch_buffers.numel() == 0 or experts_to_copy.numel() == 0: return - assert remote_expert.dtype == torch.bfloat16 and remote_expert.is_contiguous(), \ - "remote_expert must be contiguous bf16 [E, H, H']" - assert prefetch_buffers.dtype == torch.bfloat16 and prefetch_buffers.is_contiguous(), \ - "prefetch_buffers must be contiguous bf16 [B, H, H']" + dtype = remote_expert.dtype + assert dtype in _ELEM_TYPES, ( + f"prefetch: unsupported dtype {dtype}; " + f"supported: {sorted(str(d) for d in _ELEM_TYPES)}" + ) + assert remote_expert.is_contiguous(), \ + "remote_expert must be contiguous [E, H, H']" + assert prefetch_buffers.dtype == dtype and prefetch_buffers.is_contiguous(), \ + f"prefetch_buffers must be contiguous [B, H, H'] with dtype {dtype}, " \ + f"got {prefetch_buffers.dtype}" assert experts_to_copy.dtype == torch.int32 and experts_to_copy.is_contiguous(), \ "experts_to_copy must be contiguous int32 [B]" assert remote_expert.ndim == 3, \ @@ -360,16 +408,17 @@ def launch_prefetch( assert experts_to_copy.device.index == device_index, \ "experts_to_copy must be on the same CUDA device as prefetch_buffers" - compiled = _get_compiled(E, H, Hp, B, int(num_sms), int(device_index)) + compiled = _get_compiled(E, H, Hp, B, int(num_sms), int(device_index), dtype) + elem_ty, _ = _ELEM_TYPES[dtype] src_ptr = make_ptr( - BFloat16, + elem_ty, remote_expert.data_ptr(), cute.AddressSpace.gmem, assumed_align=16, ) dst_ptr = make_ptr( - BFloat16, + elem_ty, prefetch_buffers.data_ptr(), cute.AddressSpace.gmem, assumed_align=16, diff --git a/tests/test_prefetch.py b/tests/test_prefetch.py index e4f1643..c63f3dc 100644 --- a/tests/test_prefetch.py +++ b/tests/test_prefetch.py @@ -29,6 +29,9 @@ from tests.kernel_test_utils import local_device_index +K3_H, K3_I = 3584, 3072 + + def _random_experts(E, B, seed): """Random plan with ~40% holes and possible duplicates. Local to each rank's launch, so it doesn't need to agree across ranks; seeding keeps @@ -120,6 +123,52 @@ def _random_experts(E, B, seed): "owner_offset": 1, "experts": lambda E: [E - 1, 98, -1], }, + # ---- Quantized experts ------------------------------------------------- + { + "name": "mxfp4_packed_k3_gate_up", + "E": 8, + "H": 2 * K3_I, + "Hp": K3_H // 2, + "B": 3, + "num_sms": 32, + "owner_offset": 1, + "dtype": torch.uint8, + "experts": lambda E: [E - 1, 0, -1], + }, + { + "name": "mxfp4_packed_k3_down", + "E": 8, + "H": K3_H, + "Hp": K3_I // 2, + "B": 4, + "num_sms": 32, + "owner_offset": 2, + "dtype": torch.uint8, + "experts": lambda E: [1, E - 1, 1, -1], + }, + # ue8m0 block scales, re-tiled. + { + "name": "mxfp4_sf_k3_gate_up_retiled", + "E": 8, + "H": 128, + "Hp": 2 * K3_I * K3_H // 32 // 128, + "B": 4, + "num_sms": 16, + "owner_offset": 1, + "dtype": torch.uint8, + "experts": lambda E: [E - 1, 2, 2, -1], + }, + { + "name": "mxfp4_sf_k3_down_retiled", + "E": 8, + "H": 128, + "Hp": K3_H * K3_I // 32 // 128, + "B": 3, + "num_sms": 16, + "owner_offset": 3, + "dtype": torch.uint8, + "experts": lambda E: [0, E - 1, -1], + }, # Random plans (holes, duplicates, uneven coverage) at B=16. { "name": "random_plan_s1", @@ -177,24 +226,36 @@ def dist_env(): dist.destroy_process_group() -def make_single_owner_experts(rank, R, E, H, Hp): +def _fill_random(shape, dtype, gen): + if dtype.is_floating_point: + return torch.randn(shape, dtype=dtype, device="cuda", generator=gen) + info = torch.iinfo(dtype) + return torch.randint( + info.min, info.max, shape, dtype=dtype, device="cuda", generator=gen + ) + + +def _sentinel_for(dtype): + if dtype.is_floating_point: + return -123.0 + return 0xAB if dtype == torch.uint8 else 123 + + +def make_single_owner_experts(rank, R, E, H, Hp, dtype=torch.bfloat16): """Create one VMM mapped expert tensor per physical owner GPU.""" - padded_E = pad_dim0_for_alignment([E, H, Hp], torch.bfloat16) + padded_E = pad_dim0_for_alignment([E, H, Hp], dtype) owners = [] for owner in range(R): mapped = create_nvl_single_owner_tensor( [padded_E, H, Hp], - torch.bfloat16, + dtype, owner_rank=owner, local_rank=rank, ) if rank == owner: seed = 2026 + owner + E * 13 + H * 17 + Hp * 19 gen = torch.Generator(device="cuda").manual_seed(seed) - mapped[:E].copy_( - torch.randn(E, H, Hp, dtype=torch.bfloat16, - device="cuda", generator=gen) - ) + mapped[:E].copy_(_fill_random((E, H, Hp), dtype, gen)) if padded_E > E: mapped[E:].zero_() torch.cuda.synchronize() @@ -216,12 +277,13 @@ def run_case(rank, R, case): Hp = case["Hp"] B = case["B"] num_sms = case["num_sms"] + dtype = case.get("dtype", torch.bfloat16) dev = "cuda" assert H % 128 == 0 and Hp % 128 == 0, \ f"{case['name']}: H/Hp must be multiples of 128" - all_remote = make_single_owner_experts(rank, R, E, H, Hp) + all_remote = make_single_owner_experts(rank, R, E, H, Hp, dtype) remote_owner = remote_owner_for(rank, R, case["owner_offset"]) remote_expert = all_remote[remote_owner] assert remote_owner != rank, f"{case['name']}: remote owner must not be local" @@ -230,10 +292,8 @@ def run_case(rank, R, case): expert_ids = case["experts"](E) assert len(expert_ids) == B experts_to_copy = torch.tensor(expert_ids, dtype=torch.int32, device=dev) - sentinel = -123.0 - prefetch_buffers = torch.full( - (B, H, Hp), sentinel, dtype=torch.bfloat16, device=dev - ) + sentinel = _sentinel_for(dtype) + prefetch_buffers = torch.full((B, H, Hp), sentinel, dtype=dtype, device=dev) launch_prefetch(remote_expert, prefetch_buffers, experts_to_copy, num_sms=num_sms) torch.cuda.synchronize() @@ -266,7 +326,7 @@ def run_case(rank, R, case): if rank == 0: print( f" [PASS] {case['name']}: E={E}, B={B}, " - f"H={H}, Hp={Hp}, num_sms={num_sms}" + f"H={H}, Hp={Hp}, dtype={dtype}, num_sms={num_sms}" ) dist.barrier(device_ids=[local_device_index()])