Skip to content

Commit 9249c25

Browse files
fzyzcjyfzyzcjy
authored
Support cuMem API in cross process shared memory management (#217)
* more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * more * add flag * add test * fix * more * apply
1 parent bfded34 commit 9249c25

7 files changed

Lines changed: 185 additions & 17 deletions

File tree

‎csrc/deep_ep.cpp‎

Lines changed: 127 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,117 @@
1212
#include "kernels/api.cuh"
1313
#include "kernels/configs.cuh"
1414

15+
namespace shared_memory {
16+
void cu_mem_set_access_all(void* ptr, size_t size) {
17+
int device_count;
18+
CUDA_CHECK(cudaGetDeviceCount(&device_count));
19+
20+
CUmemAccessDesc access_desc[device_count];
21+
for (int idx = 0; idx < device_count; ++idx) {
22+
access_desc[idx].location.type = CU_MEM_LOCATION_TYPE_DEVICE;
23+
access_desc[idx].location.id = idx;
24+
access_desc[idx].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
25+
}
26+
27+
CU_CHECK(cuMemSetAccess((CUdeviceptr)ptr, size, access_desc, device_count));
28+
}
29+
30+
void cu_mem_free(void* ptr) {
31+
CUmemGenericAllocationHandle handle;
32+
CU_CHECK(cuMemRetainAllocationHandle(&handle, ptr));
33+
34+
size_t size = 0;
35+
CU_CHECK(cuMemGetAddressRange(NULL, &size, (CUdeviceptr)ptr));
36+
37+
CU_CHECK(cuMemUnmap((CUdeviceptr)ptr, size));
38+
CU_CHECK(cuMemAddressFree((CUdeviceptr)ptr, size));
39+
CU_CHECK(cuMemRelease(handle));
40+
}
41+
42+
size_t get_size_align_to_granularity(size_t size_raw, size_t granularity) {
43+
size_t size = (size_raw + granularity - 1) & ~(granularity - 1);
44+
if (size == 0)
45+
size = granularity;
46+
return size;
47+
}
48+
49+
SharedMemoryAllocator::SharedMemoryAllocator(bool use_fabric) : use_fabric(use_fabric) {}
50+
51+
void SharedMemoryAllocator::malloc(void** ptr, size_t size_raw) {
52+
if (use_fabric) {
53+
CUdevice device;
54+
CU_CHECK(cuCtxGetDevice(&device));
55+
56+
CUmemAllocationProp prop = {};
57+
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
58+
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
59+
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
60+
prop.location.id = device;
61+
62+
size_t granularity = 0;
63+
CU_CHECK(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM));
64+
65+
size_t size = get_size_align_to_granularity(size_raw, granularity);
66+
67+
CUmemGenericAllocationHandle handle;
68+
CU_CHECK(cuMemCreate(&handle, size, &prop, 0));
69+
70+
CU_CHECK(cuMemAddressReserve((CUdeviceptr*)ptr, size, granularity, 0, 0));
71+
CU_CHECK(cuMemMap((CUdeviceptr)*ptr, size, 0, handle, 0));
72+
cu_mem_set_access_all(*ptr, size);
73+
} else {
74+
CUDA_CHECK(cudaMalloc(ptr, size_raw));
75+
}
76+
}
77+
78+
void SharedMemoryAllocator::free(void* ptr) {
79+
if (use_fabric) {
80+
cu_mem_free(ptr);
81+
} else {
82+
CUDA_CHECK(cudaFree(ptr));
83+
}
84+
}
85+
86+
void SharedMemoryAllocator::get_mem_handle(MemHandle* mem_handle, void* ptr) {
87+
size_t size = 0;
88+
CU_CHECK(cuMemGetAddressRange(NULL, &size, (CUdeviceptr)ptr));
89+
90+
mem_handle->size = size;
91+
92+
if (use_fabric) {
93+
CUmemGenericAllocationHandle handle;
94+
CU_CHECK(cuMemRetainAllocationHandle(&handle, ptr));
95+
96+
CU_CHECK(cuMemExportToShareableHandle(&mem_handle->inner.cu_mem_fabric_handle, handle, CU_MEM_HANDLE_TYPE_FABRIC, 0));
97+
} else {
98+
CUDA_CHECK(cudaIpcGetMemHandle(&mem_handle->inner.cuda_ipc_mem_handle, ptr));
99+
}
100+
}
101+
102+
void SharedMemoryAllocator::open_mem_handle(void** ptr, MemHandle* mem_handle) {
103+
if (use_fabric) {
104+
size_t size = mem_handle->size;
105+
106+
CUmemGenericAllocationHandle handle;
107+
CU_CHECK(cuMemImportFromShareableHandle(&handle, &mem_handle->inner.cu_mem_fabric_handle, CU_MEM_HANDLE_TYPE_FABRIC));
108+
109+
CU_CHECK(cuMemAddressReserve((CUdeviceptr*)ptr, size, 0, 0, 0));
110+
CU_CHECK(cuMemMap((CUdeviceptr)*ptr, size, 0, handle, 0));
111+
cu_mem_set_access_all(*ptr, size);
112+
} else {
113+
CUDA_CHECK(cudaIpcOpenMemHandle(ptr, mem_handle->inner.cuda_ipc_mem_handle, cudaIpcMemLazyEnablePeerAccess));
114+
}
115+
}
116+
117+
void SharedMemoryAllocator::close_mem_handle(void* ptr) {
118+
if (use_fabric) {
119+
cu_mem_free(ptr);
120+
} else {
121+
CUDA_CHECK(cudaIpcCloseMemHandle(ptr));
122+
}
123+
}
124+
} // namespace shared_memory
125+
15126
namespace deep_ep {
16127

17128
Buffer::Buffer(int rank,
@@ -20,15 +131,17 @@ Buffer::Buffer(int rank,
20131
int64_t num_rdma_bytes,
21132
bool low_latency_mode,
22133
bool explicitly_destroy,
23-
bool enable_shrink)
134+
bool enable_shrink,
135+
bool use_fabric)
24136
: rank(rank),
25137
num_ranks(num_ranks),
26138
num_nvl_bytes(num_nvl_bytes),
27139
num_rdma_bytes(num_rdma_bytes),
28140
enable_shrink(enable_shrink),
29141
low_latency_mode(low_latency_mode),
30142
explicitly_destroy(explicitly_destroy),
31-
comm_stream(at::cuda::getStreamFromPool(true)) {
143+
comm_stream(at::cuda::getStreamFromPool(true)),
144+
shared_memory_allocator(use_fabric) {
32145
// Metadata memory
33146
int64_t barrier_signal_bytes = NUM_MAX_NVL_PEERS * sizeof(int);
34147
int64_t buffer_ptr_bytes = NUM_MAX_NVL_PEERS * sizeof(void*);
@@ -66,8 +179,9 @@ Buffer::Buffer(int rank,
66179

67180
if (num_nvl_bytes > 0) {
68181
// Local IPC: alloc local memory and set local IPC handles
69-
CUDA_CHECK(cudaMalloc(&buffer_ptrs[nvl_rank], num_nvl_bytes + barrier_signal_bytes + buffer_ptr_bytes + barrier_signal_ptr_bytes));
70-
CUDA_CHECK(cudaIpcGetMemHandle(&ipc_handles[nvl_rank], buffer_ptrs[nvl_rank]));
182+
shared_memory_allocator.malloc(&buffer_ptrs[nvl_rank],
183+
num_nvl_bytes + barrier_signal_bytes + buffer_ptr_bytes + barrier_signal_ptr_bytes);
184+
shared_memory_allocator.get_mem_handle(&ipc_handles[nvl_rank], buffer_ptrs[nvl_rank]);
71185
buffer_ptrs_gpu = reinterpret_cast<void**>(static_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes + barrier_signal_bytes);
72186

73187
// Set barrier signals
@@ -136,7 +250,8 @@ int Buffer::get_local_device_id() const {
136250
}
137251

138252
pybind11::bytearray Buffer::get_local_ipc_handle() const {
139-
return {ipc_handles[nvl_rank].reserved, CUDA_IPC_HANDLE_SIZE};
253+
const shared_memory::MemHandle& handle = ipc_handles[nvl_rank];
254+
return {reinterpret_cast<const char*>(&handle), sizeof(handle)};
140255
}
141256

142257
pybind11::bytearray Buffer::get_local_nvshmem_unique_id() const {
@@ -176,11 +291,11 @@ void Buffer::destroy() {
176291
if (is_available()) {
177292
for (int i = 0; i < num_nvl_ranks; ++i)
178293
if (i != nvl_rank)
179-
CUDA_CHECK(cudaIpcCloseMemHandle(buffer_ptrs[i]));
294+
shared_memory_allocator.close_mem_handle(buffer_ptrs[i]);
180295
}
181296

182297
// Free local buffer and error flag
183-
CUDA_CHECK(cudaFree(buffer_ptrs[nvl_rank]));
298+
shared_memory_allocator.free(buffer_ptrs[nvl_rank]);
184299
}
185300

186301
// Free NVSHMEM
@@ -220,13 +335,13 @@ void Buffer::sync(const std::vector<int>& device_ids,
220335
for (int i = 0, offset = rdma_rank * num_nvl_ranks; i < num_nvl_ranks; ++i) {
221336
EP_HOST_ASSERT(all_gathered_handles[offset + i].has_value());
222337
auto handle_str = std::string(all_gathered_handles[offset + i].value());
223-
EP_HOST_ASSERT(handle_str.size() == CUDA_IPC_HANDLE_SIZE);
338+
EP_HOST_ASSERT(handle_str.size() == shared_memory::HANDLE_SIZE);
224339
if (offset + i != rank) {
225-
std::memcpy(ipc_handles[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE);
226-
CUDA_CHECK(cudaIpcOpenMemHandle(&buffer_ptrs[i], ipc_handles[i], cudaIpcMemLazyEnablePeerAccess));
340+
std::memcpy(&ipc_handles[i], handle_str.c_str(), shared_memory::HANDLE_SIZE);
341+
shared_memory_allocator.open_mem_handle(&buffer_ptrs[i], &ipc_handles[i]);
227342
barrier_signal_ptrs[i] = reinterpret_cast<int*>(static_cast<uint8_t*>(buffer_ptrs[i]) + num_nvl_bytes);
228343
} else {
229-
EP_HOST_ASSERT(std::memcmp(ipc_handles[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE) == 0);
344+
EP_HOST_ASSERT(std::memcmp(&ipc_handles[i], handle_str.c_str(), shared_memory::HANDLE_SIZE) == 0);
230345
}
231346
}
232347

@@ -1739,7 +1854,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
17391854
.def("current_stream_wait", &deep_ep::EventHandle::current_stream_wait);
17401855

17411856
pybind11::class_<deep_ep::Buffer>(m, "Buffer")
1742-
.def(pybind11::init<int, int, int64_t, int64_t, bool, bool, bool>())
1857+
.def(pybind11::init<int, int, int64_t, int64_t, bool, bool, bool, bool>())
17431858
.def("is_available", &deep_ep::Buffer::is_available)
17441859
.def("get_num_rdma_ranks", &deep_ep::Buffer::get_num_rdma_ranks)
17451860
.def("get_rdma_rank", &deep_ep::Buffer::get_rdma_rank)

‎csrc/deep_ep.hpp‎

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,34 @@
2121
#define TORCH_EXTENSION_NAME deep_ep_cpp
2222
#endif
2323

24+
namespace shared_memory {
25+
26+
union MemHandleInner {
27+
cudaIpcMemHandle_t cuda_ipc_mem_handle;
28+
CUmemFabricHandle cu_mem_fabric_handle;
29+
};
30+
31+
struct MemHandle {
32+
MemHandleInner inner;
33+
size_t size;
34+
};
35+
36+
constexpr size_t HANDLE_SIZE = sizeof(MemHandle);
37+
38+
class SharedMemoryAllocator {
39+
public:
40+
SharedMemoryAllocator(bool use_fabric);
41+
void malloc(void** ptr, size_t size);
42+
void free(void* ptr);
43+
void get_mem_handle(MemHandle* mem_handle, void* ptr);
44+
void open_mem_handle(void** ptr, MemHandle* mem_handle);
45+
void close_mem_handle(void* ptr);
46+
47+
private:
48+
bool use_fabric;
49+
};
50+
} // namespace shared_memory
51+
2452
namespace deep_ep {
2553

2654
struct Buffer {
@@ -50,7 +78,7 @@ struct Buffer {
5078
int num_device_sms;
5179
int rank, rdma_rank, nvl_rank;
5280
int num_ranks, num_rdma_ranks, num_nvl_ranks;
53-
cudaIpcMemHandle_t ipc_handles[NUM_MAX_NVL_PEERS];
81+
shared_memory::MemHandle ipc_handles[NUM_MAX_NVL_PEERS];
5482

5583
// Stream for communication
5684
at::cuda::CUDAStream comm_stream;
@@ -82,14 +110,17 @@ struct Buffer {
82110
volatile int* moe_recv_rdma_counter = nullptr;
83111
int* moe_recv_rdma_counter_mapped = nullptr;
84112

113+
shared_memory::SharedMemoryAllocator shared_memory_allocator;
114+
85115
public:
86116
Buffer(int rank,
87117
int num_ranks,
88118
int64_t num_nvl_bytes,
89119
int64_t num_rdma_bytes,
90120
bool low_latency_mode,
91121
bool explicitly_destroy,
92-
bool enable_shrink);
122+
bool enable_shrink,
123+
bool use_fabric);
93124

94125
~Buffer() noexcept(false);
95126

‎csrc/kernels/exception.cuh‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,18 @@ public:
3131
} while (0)
3232
#endif
3333

34+
#ifndef CU_CHECK
35+
#define CU_CHECK(cmd) \
36+
do { \
37+
CUresult e = (cmd); \
38+
if (e != CUDA_SUCCESS) { \
39+
const char* error_str = NULL; \
40+
cuGetErrorString(e, &error_str); \
41+
throw EPException("CU", __FILE__, __LINE__, std::string(error_str)); \
42+
} \
43+
} while (0)
44+
#endif
45+
3446
#ifndef EP_HOST_ASSERT
3547
#define EP_HOST_ASSERT(cond) \
3648
do { \

‎deep_ep/buffer.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ def __init__(self,
3737
num_qps_per_rank: int = 24,
3838
allow_nvlink_for_low_latency_mode: bool = True,
3939
allow_mnnvl: bool = False,
40+
use_fabric: bool = False,
4041
explicitly_destroy: bool = False,
4142
enable_shrink: bool = False,
4243
comm: Optional["mpi4py.MPI.Comm"] = None) -> None: # noqa: F821
@@ -55,6 +56,7 @@ def __init__(self,
5556
Warning: PCIe connections may lead to errors due to memory ordering issues,
5657
please make sure all connections are via NVLink.
5758
allow_mnnvl: whether to allow MNNVL
59+
use_fabric: whether to use fabric API for memory buffers.
5860
enable_shrink: whether to enable shrink mode. The enable mode allocates a mask buffer to support masking ranks dynamically.
5961
explicitly_destroy: If this flag is set to True, you need to explicitly call `destroy()` to release resources;
6062
otherwise, the resources will be released by the destructor.
@@ -88,7 +90,7 @@ def all_gather_object(obj):
8890
self.explicitly_destroy = explicitly_destroy
8991
self.enable_shrink = enable_shrink
9092
self.runtime = deep_ep_cpp.Buffer(self.rank, self.group_size, num_nvl_bytes, num_rdma_bytes, low_latency_mode, explicitly_destroy,
91-
enable_shrink)
93+
enable_shrink, use_fabric)
9294

9395
# Synchronize device IDs
9496
local_device_id = self.runtime.get_local_device_id()

‎format.sh‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,10 @@ if ! git diff --quiet &>/dev/null; then
184184
echo
185185
git --no-pager diff --name-only
186186

187+
echo 'You can also copy-paste the diff below to fix the lint:'
188+
echo
189+
git --no-pager diff
190+
187191
exit 1
188192
fi
189193

‎setup.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ def get_nvshmem_host_lib_name(base_dir):
4141
include_dirs = ['csrc/']
4242
library_dirs = []
4343
nvcc_dlink = []
44-
extra_link_args = []
44+
extra_link_args = ['-lcuda']
4545

4646
# NVSHMEM flags
4747
if disable_nvshmem:

‎tests/test_intranode.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -275,7 +275,9 @@ def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
275275
num_rdma_bytes,
276276
low_latency_mode=test_ll_compatibility,
277277
num_qps_per_rank=(ll_num_experts // num_ranks if test_ll_compatibility else 1),
278-
explicitly_destroy=True)
278+
explicitly_destroy=True,
279+
allow_mnnvl=args.allow_mnnvl,
280+
use_fabric=args.use_fabric)
279281
torch.manual_seed(rank)
280282

281283
for i in (24, ):
@@ -301,6 +303,8 @@ def test_loop(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
301303
parser.add_argument('--hidden', type=int, default=7168, help='Hidden dimension size (default: 7168)')
302304
parser.add_argument('--num-topk', type=int, default=8, help='Number of top-k experts (default: 8)')
303305
parser.add_argument('--num-experts', type=int, default=256, help='Number of experts (default: 256)')
306+
parser.add_argument('--allow-mnnvl', action="store_true", help='Enable MNNVL support')
307+
parser.add_argument('--use-fabric', action="store_true", help='Enable fabric mode')
304308
args = parser.parse_args()
305309

306310
num_processes = args.num_processes

0 commit comments

Comments
 (0)