A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 16:00 UTC
[Distributed] Add custom allreduce support for ROCM (#14125)
Signed-off-by: ilmarkov <imarkov@redhat.com> Co-authored-by: ilmarkov <imarkov@redhat.com>
Ilya Markov committed
Apr 1, 2025 at 07:49 UTC
b7b7676d67ee517c0f97cfd245531db477606010
13 files changed
+373
-160
CMakeLists.txt
+1
-1
@@ -242,6 +242,7 @@ set(VLLM_EXT_SRC
242
"csrc/quantization/gguf/gguf_kernel.cu"
243
"csrc/cuda_utils_kernels.cu"
244
"csrc/prepare_inputs/advance_step.cu"
245
+ "csrc/custom_all_reduce.cu"
246
"csrc/torch_bindings.cpp")
247
248
if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -283,7 +284,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
284
"csrc/mamba/causal_conv1d/causal_conv1d.cu"
285
"csrc/quantization/aqlm/gemm_kernels.cu"
286
"csrc/quantization/awq/gemm_kernels.cu"
286
- "csrc/custom_all_reduce.cu"
287
"csrc/permute_cols.cu"
288
"csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu"
289
"csrc/quantization/fp4/nvfp4_quant_entry.cu"
csrc/custom_all_reduce.cu
+47
-2
@@ -12,7 +12,7 @@ static_assert(sizeof(void*) == sizeof(fptr_t));
12
13
fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
14
torch::Tensor& rank_data, int64_t rank,
15
- bool full_nvlink) {
15
+ bool fully_connected) {
16
int world_size = fake_ipc_ptrs.size();
17
if (world_size > 8)
18
throw std::invalid_argument("world size > 8 is not supported");
@@ -27,7 +27,7 @@ fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
27
}
28
return (fptr_t) new vllm::CustomAllreduce(ipc_ptrs, rank_data.data_ptr(),
29
rank_data.numel(), rank, world_size,
30
- full_nvlink);
30
+ fully_connected);
31
}
32
33
/**
@@ -142,3 +142,48 @@ void register_graph_buffers(fptr_t _fa,
142
bytes.reserve(handles.size());
143
fa->register_graph_buffers(bytes, offsets);
144
}
145
+
146
+std::tuple<fptr_t, torch::Tensor> allocate_shared_buffer_and_handle(
147
+ int64_t size) {
148
+ auto device_index = c10::cuda::current_device();
149
+ at::DeviceGuard device_guard(at::Device(at::DeviceType::CUDA, device_index));
150
+ void* buffer;
151
+ cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
152
+ auto stream = c10::cuda::getCurrentCUDAStream().stream();
153
+ AT_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
154
+
155
+ // Allocate buffer
156
+#if defined(USE_ROCM)
157
+ // data buffers need to be "uncached" for signal on MI200
158
+ AT_CUDA_CHECK(
159
+ hipExtMallocWithFlags((void**)&buffer, size, hipDeviceMallocUncached));
160
+#else
161
+ AT_CUDA_CHECK(cudaMalloc((void**)&buffer, size));
162
+#endif
163
+ AT_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream));
164
+ AT_CUDA_CHECK(cudaStreamSynchronize(stream));
165
+ AT_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
166
+
167
+ // Create IPC memhandle for the allocated buffer.
168
+ // Will use it in open_mem_handle.
169
+ auto options =
170
+ torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCPU);
171
+ auto handle =
172
+ torch::empty({static_cast<int64_t>(sizeof(cudaIpcMemHandle_t))}, options);
173
+ AT_CUDA_CHECK(
174
+ cudaIpcGetMemHandle((cudaIpcMemHandle_t*)handle.data_ptr(), buffer));
175
+
176
+ return std::make_tuple(reinterpret_cast<fptr_t>(buffer), handle);
177
+}
178
+
179
+fptr_t open_mem_handle(torch::Tensor& mem_handle) {
180
+ void* ipc_ptr;
181
+ AT_CUDA_CHECK(cudaIpcOpenMemHandle(
182
+ (void**)&ipc_ptr, *((const cudaIpcMemHandle_t*)mem_handle.data_ptr()),
183
+ cudaIpcMemLazyEnablePeerAccess));
184
+ return reinterpret_cast<fptr_t>(ipc_ptr);
185
+}
186
+
187
+void free_shared_buffer(fptr_t buffer) {
188
+ AT_CUDA_CHECK(cudaFree(reinterpret_cast<void*>(buffer)));
189
+}
csrc/custom_all_reduce.cuh
+162
-77
@@ -5,6 +5,10 @@
5
#include <cuda_fp16.h>
6
#include <cuda_runtime.h>
7
8
+#if defined(USE_ROCM)
9
+typedef __hip_bfloat16 nv_bfloat16;
10
+#endif
11
+
12
#include <iostream>
13
#include <array>
14
#include <limits>
@@ -12,6 +16,7 @@
16
#include <unordered_map>
17
#include <vector>
18
19
+namespace vllm {
20
#define CUDACHECK(cmd) \
21
do { \
22
cudaError_t e = cmd; \
@@ -22,24 +27,37 @@
27
} \
28
} while (0)
29
25
-namespace vllm {
26
-
30
+// Maximal number of blocks in allreduce kernel.
31
constexpr int kMaxBlocks = 36;
32
+
33
+// Default number of blocks in allreduce kernel.
34
+#ifndef USE_ROCM
35
+const int defaultBlockLimit = 36;
36
+CUpointer_attribute rangeStartAddrAttr = CU_POINTER_ATTRIBUTE_RANGE_START_ADDR;
37
+#else
38
+const int defaultBlockLimit = 16;
39
+hipPointer_attribute rangeStartAddrAttr =
40
+ HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR;
41
+#endif
42
+
43
// Counter may overflow, but it's fine since unsigned int overflow is
44
// well-defined behavior.
45
using FlagType = uint32_t;
46
+
47
+// Two sets of peer counters are needed for two syncs: starting and ending an
48
+// operation. The reason is that it's possible for peer GPU block to arrive at
49
+// the second sync point while the current GPU block haven't passed the first
50
+// sync point. Thus, peer GPU may write counter+1 while current GPU is busy
51
+// waiting for counter. We use alternating counter array to avoid this
52
+// possibility.
53
struct Signal {
32
- alignas(128) FlagType self_counter[kMaxBlocks][8];
33
- // Two sets of peer counters are needed for two syncs. The reason is that
34
- // it's possible for peer GPU block to arrive at the second sync point while
35
- // the current GPU block haven't passed the first sync point. Thus, peer GPU
36
- // may write counter+1 while current GPU is busy waiting for counter. We use
37
- // alternating counter array to avoid this possibility.
38
- alignas(128) FlagType peer_counter[2][kMaxBlocks][8];
54
+ alignas(128) FlagType start[kMaxBlocks][8];
55
+ alignas(128) FlagType end[kMaxBlocks][8];
56
+ alignas(128) FlagType _flag[kMaxBlocks]; // incremental flags for each rank
57
};
58
59
struct __align__(16) RankData {
42
- const void* __restrict__ ptrs[8];
60
+ const void* ptrs[8];
61
};
62
63
struct __align__(16) RankSignals {
@@ -134,27 +152,29 @@ DINLINE O downcast(array_t<float, O::size> val) {
152
}
153
}
154
155
+#if !defined(USE_ROCM)
156
+
157
static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
138
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
158
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
159
asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag),
160
"l"(flag_addr));
141
-#else
161
+ #else
162
asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag),
163
"l"(flag_addr));
144
-#endif
164
+ #endif
165
}
166
167
static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
168
FlagType flag;
149
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
169
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
170
asm volatile("ld.acquire.sys.global.u32 %0, [%1];"
171
: "=r"(flag)
172
: "l"(flag_addr));
153
-#else
173
+ #else
174
asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;"
175
: "=r"(flag)
176
: "l"(flag_addr));
157
-#endif
177
+ #endif
178
return flag;
179
}
180
@@ -170,37 +190,99 @@ static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) {
190
return flag;
191
}
192
173
-// is_start: whether this is the very first synchronization barrier.
174
-// need_fence: whether a memory fence is needed. If true, a release-acquire
175
-// semantic is used to enforce memory access order before and after this
176
-// barrier.
177
-template <int ngpus, bool is_start, bool need_fence = false>
178
-DINLINE void multi_gpu_barrier(const RankSignals& sg, Signal* self_sg,
179
- int rank) {
180
- if constexpr (!is_start) __syncthreads();
181
- static_assert(
182
- !(is_start && need_fence)); // Start barrier shouldn't need fence.
193
+// This function is meant to be used as the first synchronization in the all
194
+// reduce kernel. Thus, it doesn't need to make any visibility guarantees for
195
+// prior memory accesses. Note: volatile writes will not be reordered against
196
+// other volatile writes.
197
+template <int ngpus>
198
+DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
199
+ int rank) {
200
+ uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
201
if (threadIdx.x < ngpus) {
184
- // Increment the counter. Technically we only need one counter, but we use
185
- // multiple per block to eliminate the need to share the counter via smem.
186
- auto val = self_sg->self_counter[blockIdx.x][threadIdx.x] += 1;
202
+ auto peer_counter_ptr = &sg.signals[threadIdx.x]->start[blockIdx.x][rank];
203
+ auto self_counter_ptr = &self_sg->start[blockIdx.x][threadIdx.x];
204
+ // Write the expected counter value to peer and wait for correct value
205
+ // from peer.
206
+ st_flag_volatile(peer_counter_ptr, flag);
207
+ while (ld_flag_volatile(self_counter_ptr) != flag);
208
+ }
209
+ __syncthreads();
210
+ // use one thread to update flag
211
+ if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
212
+}
213
+
214
+// This function is meant to be used as the second or the final
215
+// synchronization barrier in the all reduce kernel. If it's the final
216
+// synchronization barrier, we don't need to make any visibility guarantees
217
+// for prior memory accesses.
218
+template <int ngpus, bool final_sync = false>
219
+DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
220
+ __syncthreads();
221
+ uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
222
+ if (threadIdx.x < ngpus) {
223
+ auto peer_counter_ptr = &sg.signals[threadIdx.x]->end[blockIdx.x][rank];
224
+ auto self_counter_ptr = &self_sg->end[blockIdx.x][threadIdx.x];
225
// Write the expected counter value to peer and wait for correct value from
226
// peer.
189
- auto peer_counter_ptr =
190
- &sg.signals[threadIdx.x]->peer_counter[val % 2][blockIdx.x][rank];
191
- auto self_counter_ptr =
192
- &self_sg->peer_counter[val % 2][blockIdx.x][threadIdx.x];
193
- if constexpr (need_fence) {
194
- st_flag_release(peer_counter_ptr, val);
195
- while (ld_flag_acquire(self_counter_ptr) != val);
227
+ if constexpr (!final_sync) {
228
+ st_flag_release(peer_counter_ptr, flag);
229
+ while (ld_flag_acquire(self_counter_ptr) != flag);
230
} else {
197
- st_flag_volatile(peer_counter_ptr, val);
198
- while (ld_flag_volatile(self_counter_ptr) != val);
231
+ st_flag_volatile(peer_counter_ptr, flag);
232
+ while (ld_flag_volatile(self_counter_ptr) != flag);
233
}
234
}
201
- if constexpr (is_start || need_fence) __syncthreads();
235
+ if constexpr (!final_sync) __syncthreads();
236
+
237
+ // use one thread to update flag
238
+ if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
239
}
240
241
+#else
242
+
243
+template <int ngpus>
244
+DINLINE void barrier_at_start(const RankSignals& sg, Signal* self_sg,
245
+ int rank) {
246
+ uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
247
+ if (threadIdx.x < ngpus) {
248
+ // simultaneously write to the corresponding flag of all ranks.
249
+ // Latency = 1 p2p write
250
+ __scoped_atomic_store_n(&sg.signals[threadIdx.x]->start[blockIdx.x][rank],
251
+ flag, __ATOMIC_RELAXED, __MEMORY_SCOPE_SYSTEM);
252
+ // wait until we got true from all ranks
253
+ while (__scoped_atomic_load_n(&self_sg->start[blockIdx.x][threadIdx.x],
254
+ __ATOMIC_RELAXED,
255
+ __MEMORY_SCOPE_DEVICE) < flag);
256
+ }
257
+ __syncthreads();
258
+ // use one thread to update flag
259
+ if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
260
+}
261
+
262
+template <int ngpus, bool final_sync = false>
263
+DINLINE void barrier_at_end(const RankSignals& sg, Signal* self_sg, int rank) {
264
+ __syncthreads();
265
+ uint32_t flag = self_sg->_flag[blockIdx.x] + 1;
266
+ if (threadIdx.x < ngpus) {
267
+ // simultaneously write to the corresponding flag of all ranks.
268
+ // Latency = 1 p2p write
269
+ __scoped_atomic_store_n(&sg.signals[threadIdx.x]->end[blockIdx.x][rank],
270
+ flag,
271
+ final_sync ? __ATOMIC_RELAXED : __ATOMIC_RELEASE,
272
+ __MEMORY_SCOPE_SYSTEM);
273
+ // wait until we got true from all ranks
274
+ while (
275
+ __scoped_atomic_load_n(&self_sg->end[blockIdx.x][threadIdx.x],
276
+ final_sync ? __ATOMIC_RELAXED : __ATOMIC_ACQUIRE,
277
+ __MEMORY_SCOPE_DEVICE) < flag);
278
+ }
279
+ if constexpr (!final_sync) __syncthreads();
280
+ // use one thread to update flag
281
+ if (threadIdx.x == 0) self_sg->_flag[blockIdx.x] = flag;
282
+}
283
+
284
+#endif
285
+
286
template <typename P, int ngpus, typename A>
287
DINLINE P packed_reduce(const P* ptrs[], int idx) {
288
A tmp = upcast(ptrs[0][idx]);
@@ -220,13 +302,13 @@ __global__ void __launch_bounds__(512, 1)
302
// note: we don't reorder the address so the accumulation order is the same
303
// for all ranks, ensuring bitwise identical results
304
auto dp = *_dp;
223
- multi_gpu_barrier<ngpus, true>(sg, self_sg, rank);
305
+ barrier_at_start<ngpus>(sg, self_sg, rank);
306
// do the actual reduction
307
for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
308
idx += gridDim.x * blockDim.x) {
309
((P*)result)[idx] = packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], idx);
310
}
229
- multi_gpu_barrier<ngpus, false>(sg, self_sg, rank);
311
+ barrier_at_end<ngpus, true>(sg, self_sg, rank);
312
}
313
314
template <typename P>
@@ -255,18 +337,20 @@ __global__ void __launch_bounds__(512, 1)
337
tmps[i] = get_tmp_buf<P>(sg.signals[target]);
338
}
339
auto tmp_out = tmps[0];
258
- multi_gpu_barrier<ngpus, true>(sg, self_sg, rank);
340
+ barrier_at_start<ngpus>(sg, self_sg, rank);
341
+
342
// stage 1: reduce scatter
343
for (int idx = start + tid; idx < end; idx += stride) {
344
tmp_out[idx - start] = packed_reduce<P, ngpus, A>(ptrs, idx);
345
}
263
- multi_gpu_barrier<ngpus, false, true>(sg, self_sg, rank);
346
+ barrier_at_end<ngpus>(sg, self_sg, rank);
347
348
// stage 2: allgather. Note: it's important to match the tid between
349
// the two stages, because visibility across devices is only guaranteed
350
// between threads that have the same tid. If thread i computes the sum of
268
- // start + i in the first stage, then thread i also gathers start + i from all
269
- // ranks.
351
+ // start + i in the first stage, then thread i also gathers start + i from
352
+ // all ranks.
353
+
354
for (int idx = tid; idx < largest_part; idx += stride) {
355
#pragma unroll
356
for (int i = 0; i < ngpus; i++) {
@@ -287,21 +371,22 @@ class CustomAllreduce {
371
public:
372
int rank_;
373
int world_size_;
290
- bool full_nvlink_;
374
+ // Full NVLink or xGMI connection between GPUs.
375
+ bool fully_connected_;
376
377
RankSignals sg_;
293
- // Stores an map from a pointer to its peer pointters from all ranks.
378
+ // Stores an map from a pointer to its peer pointers from all ranks.
379
std::unordered_map<void*, RankData*> buffers_;
380
Signal* self_sg_;
381
382
// Stores rank data from all ranks. This is mainly for cuda graph purposes.
383
// For cuda graph to work, all kernel arguments must be fixed during graph
299
- // capture time. However, the peer pointers are not known during graph capture
300
- // time. Therefore, during capture, we increment the rank data pointer and use
301
- // that as the argument to the kernel. The kernel arguments are stored in
302
- // graph_unreg_buffers_. The actual peer pointers will be filled in at the
303
- // memory pointed to by the pointers in graph_unreg_buffers_ when
304
- // the IPC handles are exchanged between ranks.
384
+ // capture time. However, the peer pointers are not known during graph
385
+ // capture time. Therefore, during capture, we increment the rank data
386
+ // pointer and use that as the argument to the kernel. The kernel arguments
387
+ // are stored in graph_unreg_buffers_. The actual peer pointers will be
388
+ // filled in at the memory pointed to by the pointers in
389
+ // graph_unreg_buffers_ when the IPC handles are exchanged between ranks.
390
//
391
// The overall process looks like this:
392
// 1. Graph capture.
@@ -319,17 +404,18 @@ class CustomAllreduce {
404
* Signals are an array of ipc-enabled buffers from all ranks.
405
* For each of the buffer, the layout is as follows:
406
* | -- sizeof(Signal) -- | ------ a few MB ----- |
322
- * The first section is for allreduce synchronization, and the second section
323
- * is for storing the intermediate results required by some allreduce algos.
407
+ * The first section is for allreduce synchronization, and the second
408
+ * section is for storing the intermediate results required by some
409
+ * allreduce algos.
410
*
411
* Note: this class does not own any device memory. Any required buffers
412
* are passed in from the constructor.
413
*/
414
CustomAllreduce(Signal** signals, void* rank_data, size_t rank_data_sz,
329
- int rank, int world_size, bool full_nvlink = true)
415
+ int rank, int world_size, bool fully_connected = true)
416
: rank_(rank),
417
world_size_(world_size),
332
- full_nvlink_(full_nvlink),
418
+ fully_connected_(fully_connected),
419
self_sg_(signals[rank]),
420
d_rank_data_base_(reinterpret_cast<RankData*>(rank_data)),
421
d_rank_data_end_(d_rank_data_base_ + rank_data_sz / sizeof(RankData)) {
@@ -361,8 +447,7 @@ class CustomAllreduce {
447
void* base_ptr;
448
// note: must share the base address of each allocation, or we get wrong
449
// address
364
- if (cuPointerGetAttribute(&base_ptr,
365
- CU_POINTER_ATTRIBUTE_RANGE_START_ADDR,
450
+ if (cuPointerGetAttribute(&base_ptr, rangeStartAddrAttr,
451
(CUdeviceptr)ptr) != CUDA_SUCCESS)
452
throw std::runtime_error("failed to get pointer attr");
453
CUDACHECK(cudaIpcGetMemHandle(
@@ -396,11 +481,11 @@ class CustomAllreduce {
481
482
// Note: when registering graph buffers, we intentionally choose to not
483
// deduplicate the addresses. That means if the allocator reuses some
399
- // addresses, they will be registered again. This is to account for the remote
400
- // possibility of different allocation patterns between ranks. For example,
401
- // rank 1 may get the same input address for the second allreduce, but rank 2
402
- // got a different address. IPC handles have internal reference counting
403
- // mechanism so overhead should be small.
484
+ // addresses, they will be registered again. This is to account for the
485
+ // remote possibility of different allocation patterns between ranks. For
486
+ // example, rank 1 may get the same input address for the second allreduce,
487
+ // but rank 2 got a different address. IPC handles have internal reference
488
+ // counting mechanism so overhead should be small.
489
void register_graph_buffers(
490
const std::vector<std::string>& handles,
491
const std::vector<std::vector<int64_t>>& offsets) {
@@ -431,15 +516,15 @@ class CustomAllreduce {
516
/**
517
* Performs allreduce, assuming input has already been registered.
518
*
434
- * Block and grid default configs are results after careful grid search. Using
435
- * 36 blocks give the best or close to the best runtime on the devices I
436
- * tried: A100, A10, A30, T4, V100. You'll notice that NCCL kernels also only
437
- * take a small amount of SMs. Not quite sure the underlying reason, but my
438
- * guess is that too many SMs will cause contention on NVLink bus.
519
+ * Block and grid default configs are results after careful grid search.
520
+ * Using 36 blocks give the best or close to the best runtime on the devices
521
+ * I tried: A100, A10, A30, T4, V100. You'll notice that NCCL kernels also
522
+ * only take a small amount of SMs. Not quite sure the underlying reason,
523
+ * but my guess is that too many SMs will cause contention on NVLink bus.
524
*/
525
template <typename T>
526
void allreduce(cudaStream_t stream, T* input, T* output, int size,
442
- int threads = 512, int block_limit = 36) {
527
+ int threads = 512, int block_limit = defaultBlockLimit) {
528
auto d = packed_t<T>::P::size;
529
if (size % d != 0)
530
throw std::runtime_error(
@@ -473,13 +558,11 @@ class CustomAllreduce {
558
#define KL(ngpus, name) \
559
name<T, ngpus><<<blocks, threads, 0, stream>>>(ptrs, sg_, self_sg_, output, \
560
rank_, size);
476
- // TODO(hanzhi713): Threshold is different for A100 and H100.
477
- // Add per device threshold.
561
#define REDUCE_CASE(ngpus) \
562
case ngpus: { \
563
if (world_size_ == 2) { \
564
KL(ngpus, cross_device_reduce_1stage); \
482
- } else if (full_nvlink_) { \
565
+ } else if (fully_connected_) { \
566
if ((world_size_ <= 4 && bytes < 512 * 1024) || \
567
(world_size_ <= 8 && bytes < 256 * 1024)) { \
568
KL(ngpus, cross_device_reduce_1stage); \
@@ -497,7 +580,8 @@ class CustomAllreduce {
580
REDUCE_CASE(8)
581
default:
582
throw std::runtime_error(
500
- "custom allreduce only supports num gpus in (2,4,6,8). Actual num "
583
+ "custom allreduce only supports num gpus in (2,4,6,8). Actual "
584
+ "num "
585
"gpus = " +
586
std::to_string(world_size_));
587
}
@@ -511,10 +595,11 @@ class CustomAllreduce {
595
}
596
}
597
};
598
+
599
/**
515
- * To inspect PTX/SASS, copy paste this header file to compiler explorer and add
516
- a template instantiation:
600
+ * To inspect PTX/SASS, copy paste this header file to compiler explorer and
601
+ add a template instantiation:
602
* template void vllm::CustomAllreduce::allreduce<half>(cudaStream_t, half *,
603
half *, int, int, int);
604
*/
520
-} // namespace vllm
605
+} // namespace vllm
\ No newline at end of file
csrc/custom_all_reduce_test.cu
+45
-13
@@ -1,9 +1,9 @@
1
/**
2
* This is a standalone test for custom allreduce.
3
* To compile, make sure you have MPI and NCCL installed in your system.
4
- * export MPI_HOME=xxx
4
+ * export MPI_HOME=XXX
5
* nvcc -O2 -arch=native -std=c++17 custom_all_reduce_test.cu -o
6
- * custom_all_reduce_test -lnccl -I${MPI_HOME} -lmpi
6
+ * custom_all_reduce_test -lnccl -I${MPI_HOME}/include -lmpi
7
*
8
* Warning: this C++ test is not designed to be very readable and was used
9
* during the rapid prototyping process.
@@ -22,7 +22,15 @@
22
#include "cuda_profiler_api.h"
23
#include "custom_all_reduce.cuh"
24
#include "mpi.h"
25
-#include "nccl.h"
25
+#ifdef USE_ROCM
26
+ #include <hip/hip_bf16.h>
27
+typedef __hip_bfloat16 nv_bfloat16;
28
+ #include "rccl/rccl.h"
29
+ #include "custom_all_reduce_hip.cuh"
30
+#else
31
+ #include "nccl.h"
32
+ #include "custom_all_reduce.cuh"
33
+#endif
34
35
#define MPICHECK(cmd) \
36
do { \
@@ -43,16 +51,29 @@
51
} \
52
} while (0)
53
54
+#ifdef USE_ROCM
55
__global__ void dummy_kernel() {
47
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
56
+ for (int i = 0; i < 100; i++) {
57
+ uint64_t start = wall_clock64();
58
+ uint64_t cycles_elapsed;
59
+ do {
60
+ cycles_elapsed = wall_clock64() - start;
61
+ } while (cycles_elapsed < 100);
62
+ }
63
for (int i = 0; i < 100; i++) __nanosleep(1000000); // 100ms
64
+}
65
#else
66
+__global__ void dummy_kernel() {
67
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
68
+ for (int i = 0; i < 100; i++) __nanosleep(1000000); // 100ms
69
+ #else
70
for (int i = 0; i < 100; i++) {
71
long long int start = clock64();
72
while (clock64() - start < 150000000); // approximately 98.4ms on P40
73
}
54
-#endif
74
+ #endif
75
}
76
+#endif
77
78
template <typename T>
79
__global__ void set_data(T* data, int size, int myRank) {
@@ -121,8 +142,14 @@ void run(int myRank, int nRanks, ncclComm_t& comm, int threads, int block_limit,
142
* registration, they are allocated and registered together in the test for
143
* convenience.
144
*/
145
+#ifdef USE_ROCM
146
+ CUDACHECK(hipExtMallocWithFlags(
147
+ (void**)&buffer, 2 * data_size * sizeof(T) + sizeof(vllm::Signal),
148
+ hipDeviceMallocUncached));
149
+#else
150
CUDACHECK(
151
cudaMalloc(&buffer, 2 * data_size * sizeof(T) + sizeof(vllm::Signal)));
152
+#endif
153
CUDACHECK(
154
cudaMemset(buffer, 0, 2 * data_size * sizeof(T) + sizeof(vllm::Signal)));
155
CUDACHECK(cudaMalloc(&self_data_copy, data_size * sizeof(T)));
@@ -311,13 +338,18 @@ int main(int argc, char** argv) {
338
339
bool performance_test = true;
340
cudaProfilerStart();
314
- // Uncomment to scan through different block size configs.
315
- // for (int threads : {256, 512, 1024}) {
316
- // for (int block_limit = 16; block_limit < 112; block_limit += 4) {
317
- // run<half>(myRank, nRanks, comm, threads, block_limit, 1024 * 1024,
318
- // performance_test);
319
- // }
320
- // }
341
+// Uncomment to scan through different block size configs.
342
+// for (int threads : {256, 512, 1024}) {
343
+// for (int block_limit = 16; block_limit < 112; block_limit += 4) {
344
+// run<half>(myRank, nRanks, comm, threads, block_limit, 1024 * 1024,
345
+// performance_test);
346
+// }
347
+// }
348
+#ifdef USE_ROCM
349
+ const int block_limit = 16;
350
+#else
351
+ const int block_limit = 36;
352
+#endif
353
// Scan through different sizes to test performance.
354
for (int sz = 512; sz <= (8 << 20); sz *= 2) {
355
run<half>(myRank, nRanks, comm, 512, 36, sz + 8 * 47, performance_test);
@@ -326,4 +358,4 @@ int main(int argc, char** argv) {
358
cudaProfilerStop();
359
MPICHECK(MPI_Finalize());
360
return EXIT_SUCCESS;
329
-}
361
+}
\ No newline at end of file
csrc/ops.h
+6
-3
@@ -267,10 +267,10 @@ void causal_conv1d_fwd(const at::Tensor& x, const at::Tensor& weight,
267
const std::optional<at::Tensor>& has_initial_state,
268
bool silu_activation, int64_t pad_slot_id);
269
270
-#ifndef USE_ROCM
270
using fptr_t = int64_t;
271
fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
273
- torch::Tensor& rank_data, int64_t rank, bool full_nvlink);
272
+ torch::Tensor& rank_data, int64_t rank,
273
+ bool fully_connected);
274
void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
275
fptr_t reg_buffer, int64_t reg_buffer_sz_bytes);
276
void dispose(fptr_t _fa);
@@ -281,4 +281,7 @@ get_graph_buffer_ipc_meta(fptr_t _fa);
281
void register_graph_buffers(fptr_t _fa,
282
const std::vector<std::vector<int64_t>>& handles,
283
const std::vector<std::vector<int64_t>>& offsets);
284
-#endif
284
+std::tuple<int64_t, torch::Tensor> allocate_shared_buffer_and_handle(
285
+ int64_t size);
286
+int64_t open_mem_handle(torch::Tensor& mem_handle);
287
+void free_shared_buffer(int64_t buffer);
csrc/torch_bindings.cpp
+8
-3
@@ -614,12 +614,11 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
614
&get_max_shared_memory_per_block_device_attribute);
615
}
616
617
-#ifndef USE_ROCM
617
TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
618
// Custom all-reduce kernels
619
custom_ar.def(
620
"init_custom_ar(int[] ipc_tensors, Tensor rank_data, "
622
- "int rank, bool full_nvlink) -> int");
621
+ "int rank, bool fully_connected) -> int");
622
custom_ar.impl("init_custom_ar", torch::kCUDA, &init_custom_ar);
623
custom_ar.def(
624
"all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, "
@@ -632,7 +631,13 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
631
custom_ar.def("register_buffer", ®ister_buffer);
632
custom_ar.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta);
633
custom_ar.def("register_graph_buffers", ®ister_graph_buffers);
634
+
635
+ custom_ar.def("allocate_shared_buffer_and_handle",
636
+ &allocate_shared_buffer_and_handle);
637
+ custom_ar.def("open_mem_handle(Tensor mem_handle) -> int", &open_mem_handle);
638
+ custom_ar.impl("open_mem_handle", torch::kCPU, &open_mem_handle);
639
+
640
+ custom_ar.def("free_shared_buffer", &free_shared_buffer);
641
}
636
-#endif
642
643
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
tests/distributed/test_custom_all_reduce.py
+1
-1
@@ -106,7 +106,7 @@ def eager_allreduce(
106
# communicate independently
107
num_communication = rank // tp_size + 1
108
sz = 1024
109
- fa = get_tp_group().ca_comm
109
+ fa = get_tp_group().device_communicator.ca_comm
110
inp = torch.ones(sz, dtype=torch.float32, device=device)
111
out = inp
112
for _ in range(num_communication):
tests/utils.py
+10
-1
@@ -612,7 +612,16 @@ def multi_process_parallel(
612
# as compared to multiprocessing.
613
# NOTE: We need to set working_dir for distributed tests,
614
# otherwise we may get import errors on ray workers
615
- ray.init(runtime_env={"working_dir": VLLM_PATH})
615
+ # NOTE: Force ray not to use gitignore file as excluding, otherwise
616
+ # it will not move .so files to working dir.
617
+ # So we have to manually add some of large directories
618
+ os.environ["RAY_RUNTIME_ENV_IGNORE_GITIGNORE"] = "1"
619
+ ray.init(
620
+ runtime_env={
621
+ "working_dir": VLLM_PATH,
622
+ "excludes":
623
+ ["build", ".git", "cmake-build-*", "shellcheck", "dist"]
624
+ })
625
626
distributed_init_port = get_open_port()
627
refs = []
vllm/_custom_ops.py
+14
-2
@@ -1337,9 +1337,9 @@ def get_max_shared_memory_per_block_device_attribute(device: int) -> int:
1337
1338
# custom ar
1339
def init_custom_ar(ipc_tensors: list[torch.Tensor], rank_data: torch.Tensor,
1340
- rank: int, full_nvlink: bool) -> int:
1340
+ rank: int, fully_connected: bool) -> int:
1341
return torch.ops._C_custom_ar.init_custom_ar(ipc_tensors, rank_data, rank,
1342
- full_nvlink)
1342
+ fully_connected)
1343
1344
1345
def all_reduce(fa: int, inp: torch.Tensor, out: torch.Tensor, reg_buffer: int,
@@ -1369,6 +1369,18 @@ def register_graph_buffers(fa: int, handles: list[list[int]],
1369
torch.ops._C_custom_ar.register_graph_buffers(fa, handles, offsets)
1370
1371
1372
+def allocate_shared_buffer_and_handle(size: int) -> tuple[int, torch.Tensor]:
1373
+ return torch.ops._C_custom_ar.allocate_shared_buffer_and_handle(size)
1374
+
1375
+
1376
+def open_mem_handle(mem_handle: torch.Tensor):
1377
+ return torch.ops._C_custom_ar.open_mem_handle(mem_handle)
1378
+
1379
+
1380
+def free_shared_buffer(ptr: int) -> None:
1381
+ torch.ops._C_custom_ar.free_shared_buffer(ptr)
1382
+
1383
+
1384
def get_flash_mla_metadata(
1385
cache_seqlens: torch.Tensor,
1386
num_heads_per_head_k: int,
vllm/config.py
+4
-2
@@ -1606,11 +1606,13 @@ class ParallelConfig:
1606
if self.use_ray:
1607
from vllm.executor import ray_utils
1608
ray_utils.assert_ray_available()
1609
- if current_platform.is_rocm():
1609
+ device_capability = current_platform.get_device_capability()
1610
+ if (current_platform.is_rocm() and device_capability is not None
1611
+ and device_capability < (9, 4)):
1612
self.disable_custom_all_reduce = True
1613
logger.info(
1614
"Disabled the custom all-reduce kernel because it is not "
1613
- "supported on AMD GPUs.")
1615
+ "supported on AMD GPUs older than MI300X.")
1616
if self.ray_workers_use_nsight and not self.use_ray:
1617
raise ValueError("Unable to use nsight profiling unless workers "
1618
"run with Ray.")
vllm/distributed/device_communicators/custom_all_reduce.py
+42
-49
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-import ctypes
3
from contextlib import contextmanager
4
from typing import List, Optional, Union
5
@@ -10,7 +9,6 @@ from torch.distributed import ProcessGroup
9
10
import vllm.envs as envs
11
from vllm import _custom_ops as ops
13
-from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
12
from vllm.distributed.device_communicators.custom_all_reduce_utils import (
13
gpu_p2p_access_check)
14
from vllm.distributed.parallel_state import in_the_same_node_as
@@ -22,7 +20,7 @@ try:
20
ops.meta_size()
21
custom_ar = True
22
except Exception:
25
- # For AMD GPUs and CPUs
23
+ # For CPUs
24
custom_ar = False
25
26
logger = init_logger(__name__)
@@ -71,7 +69,9 @@ class CustomAllreduce:
69
70
if not custom_ar:
71
# disable because of missing custom allreduce library
74
- # e.g. in a non-cuda environment
72
+ # e.g. in a non-GPU environment
73
+ logger.info("Custom allreduce is disabled because "
74
+ "of missing custom allreduce library")
75
return
76
77
self.group = group
@@ -129,11 +129,10 @@ class CustomAllreduce:
129
# test nvlink first, this will filter out most of the cases
130
# where custom allreduce is not supported
131
# this checks hardware and driver support for NVLink
132
- assert current_platform.is_cuda()
133
- from vllm.platforms.cuda import CudaPlatform
134
- cuda_platform: CudaPlatform = current_platform
135
- full_nvlink = cuda_platform.is_full_nvlink(physical_device_ids)
136
- if world_size > 2 and not full_nvlink:
132
+ assert current_platform.is_cuda_alike()
133
+ fully_connected = current_platform.is_fully_connected(
134
+ physical_device_ids)
135
+ if world_size > 2 and not fully_connected:
136
logger.warning(
137
"Custom allreduce is disabled because it's not supported on"
138
" more than two PCIe-only GPUs. To silence this warning, "
@@ -142,7 +141,8 @@ class CustomAllreduce:
141
# test P2P capability, this checks software/cudaruntime support
142
# this is expensive to compute at the first time
143
# then we cache the result
145
- if not _can_p2p(rank, world_size):
144
+ # On AMD GPU, p2p is always enabled between XGMI connected GPUs
145
+ if not current_platform.is_rocm() and not _can_p2p(rank, world_size):
146
logger.warning(
147
"Custom allreduce is disabled because your platform lacks "
148
"GPU P2P capability or P2P test failed. To silence this "
@@ -154,7 +154,8 @@ class CustomAllreduce:
154
# Meta data composes of two parts: meta data for synchronization and a
155
# temporary buffer for storing intermediate allreduce results.
156
self.meta_ptrs = self.create_shared_buffer(ops.meta_size() + max_size,
157
- group=group)
157
+ group=group,
158
+ uncached=True)
159
# This is a pre-registered IPC buffer. In eager mode, input tensors
160
# are first copied into this buffer before allreduce is performed
161
self.buffer_ptrs = self.create_shared_buffer(max_size, group=group)
@@ -169,46 +170,11 @@ class CustomAllreduce:
170
self.max_size = max_size
171
self.rank = rank
172
self.world_size = world_size
172
- self.full_nvlink = full_nvlink
173
+ self.fully_connected = fully_connected
174
self._ptr = ops.init_custom_ar(self.meta_ptrs, self.rank_data, rank,
174
- self.full_nvlink)
175
+ self.fully_connected)
176
ops.register_buffer(self._ptr, self.buffer_ptrs)
177
177
- @staticmethod
178
- def create_shared_buffer(
179
- size_in_bytes: int,
180
- group: Optional[ProcessGroup] = None) -> List[int]:
181
- """
182
- Creates a shared buffer and returns a list of pointers
183
- representing the buffer on all processes in the group.
184
- """
185
- lib = CudaRTLibrary()
186
- pointer = lib.cudaMalloc(size_in_bytes)
187
- handle = lib.cudaIpcGetMemHandle(pointer)
188
- world_size = dist.get_world_size(group=group)
189
- rank = dist.get_rank(group=group)
190
- handles = [None] * world_size
191
- dist.all_gather_object(handles, handle, group=group)
192
-
193
- pointers: List[int] = []
194
- for i, h in enumerate(handles):
195
- if i == rank:
196
- pointers.append(pointer.value) # type: ignore
197
- else:
198
- pointers.append(
199
- lib.cudaIpcOpenMemHandle(h).value) # type: ignore
200
-
201
- return pointers
202
-
203
- @staticmethod
204
- def free_shared_buffer(pointers: List[int],
205
- group: Optional[ProcessGroup] = None,
206
- rank: Optional[int] = None) -> None:
207
- if rank is None:
208
- rank = dist.get_rank(group=group)
209
- lib = CudaRTLibrary()
210
- lib.cudaFree(ctypes.c_void_p(pointers[rank]))
211
-
178
@contextmanager
179
def capture(self):
180
"""
@@ -255,7 +221,7 @@ class CustomAllreduce:
221
return False
222
# for 4 or more non NVLink-capable GPUs, custom allreduce provides
223
# little performance improvement over NCCL.
258
- if self.world_size == 2 or self.full_nvlink:
224
+ if self.world_size == 2 or self.fully_connected:
225
return inp_size < self.max_size
226
return False
227
@@ -306,3 +272,30 @@ class CustomAllreduce:
272
273
def __del__(self):
274
self.close()
275
+
276
+ @staticmethod
277
+ def create_shared_buffer(size_in_bytes: int,
278
+ group: Optional[ProcessGroup] = None,
279
+ uncached: Optional[bool] = False) -> List[int]:
280
+ pointer, handle = ops.allocate_shared_buffer_and_handle(size_in_bytes)
281
+
282
+ world_size = dist.get_world_size(group=group)
283
+ rank = dist.get_rank(group=group)
284
+ handles = [None] * world_size
285
+ dist.all_gather_object(handles, handle, group=group)
286
+
287
+ pointers: List[int] = []
288
+ for i, h in enumerate(handles):
289
+ if i == rank:
290
+ pointers.append(pointer) # type: ignore
291
+ else:
292
+ pointers.append(ops.open_mem_handle(h))
293
+ return pointers
294
+
295
+ @staticmethod
296
+ def free_shared_buffer(pointers: List[int],
297
+ group: Optional[ProcessGroup] = None,
298
+ rank: Optional[int] = 0) -> None:
299
+ if rank is None:
300
+ rank = dist.get_rank(group=group)
301
+ ops.free_shared_buffer(pointers[rank])
vllm/platforms/cuda.py
+3
-3
@@ -101,7 +101,7 @@ class CudaPlatformBase(Platform):
101
return True
102
103
@classmethod
104
- def is_full_nvlink(cls, device_ids: List[int]) -> bool:
104
+ def is_fully_connected(cls, device_ids: List[int]) -> bool:
105
raise NotImplementedError
106
107
@classmethod
@@ -362,7 +362,7 @@ class NvmlCudaPlatform(CudaPlatformBase):
362
363
@classmethod
364
@with_nvml_context
365
- def is_full_nvlink(cls, physical_device_ids: List[int]) -> bool:
365
+ def is_fully_connected(cls, physical_device_ids: List[int]) -> bool:
366
"""
367
query if the set of gpus are fully connected by nvlink (1 hop)
368
"""
@@ -427,7 +427,7 @@ class NonNvmlCudaPlatform(CudaPlatformBase):
427
return device_props.total_memory
428
429
@classmethod
430
- def is_full_nvlink(cls, physical_device_ids: List[int]) -> bool:
430
+ def is_fully_connected(cls, physical_device_ids: List[int]) -> bool:
431
logger.exception(
432
"NVLink detection not possible, as context support was"
433
" not found. Assuming no NVLink available.")
vllm/platforms/rocm.py
+30
-3
@@ -20,8 +20,9 @@ else:
20
logger = init_logger(__name__)
21
22
try:
23
- from amdsmi import (amdsmi_get_gpu_asic_info, amdsmi_get_processor_handles,
24
- amdsmi_init, amdsmi_shut_down)
23
+ from amdsmi import (AmdSmiException, amdsmi_get_gpu_asic_info,
24
+ amdsmi_get_processor_handles, amdsmi_init,
25
+ amdsmi_shut_down, amdsmi_topo_get_link_type)
26
except ImportError as e:
27
logger.warning("Failed to import from amdsmi with %r", e)
28
@@ -135,10 +136,36 @@ class RocmPlatform(Platform):
136
137
@classmethod
138
@lru_cache(maxsize=8)
138
- def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
139
+ def get_device_capability(cls,
140
+ device_id: int = 0
141
+ ) -> Optional[DeviceCapability]:
142
major, minor = torch.cuda.get_device_capability(device_id)
143
return DeviceCapability(major=major, minor=minor)
144
145
+ @staticmethod
146
+ @with_amdsmi_context
147
+ def is_fully_connected(physical_device_ids: List[int]) -> bool:
148
+ """
149
+ Query if the set of gpus are fully connected by xgmi (1 hop)
150
+ """
151
+ handles = [
152
+ amdsmi_get_processor_handles()[i] for i in physical_device_ids
153
+ ]
154
+ for i, handle in enumerate(handles):
155
+ for j, peer_handle in enumerate(handles):
156
+ if i < j:
157
+ try:
158
+ link_type = amdsmi_topo_get_link_type(
159
+ handle, peer_handle)
160
+ # type is 2 for XGMI
161
+ if link_type["hops"] != 1 or link_type["type"] != 2:
162
+ return False
163
+ except AmdSmiException as error:
164
+ logger.error("AMD 1 hop XGMI detection failed.",
165
+ exc_info=error)
166
+ return False
167
+ return True
168
+
169
@classmethod
170
@with_amdsmi_context
171
@lru_cache(maxsize=8)