@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0740e29b66
Fork 0

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, 04:45 UTC

[Feature] add quick all reduce (#19744)

Signed-off-by: ilmarkov <imarkov@redhat.com> Signed-off-by: Haoyang Li <Haoyang.Li@amd.com> Co-authored-by: ilmarkov <imarkov@redhat.com>

li haoyang committed Jun 27, 2025 at 11:54 UTC 0740e29b66ca5589f7f35a7c25b6c3de1a749da1
12 files changed +1879 -2
CMakeLists.txt
+8
@@ -648,6 +648,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
648 # if CUDA endif
649 endif()
650
651 +if (VLLM_GPU_LANG STREQUAL "HIP")
652 + # Add QuickReduce kernels
653 + list(APPEND VLLM_EXT_SRC
654 + "csrc/custom_quickreduce.cu"
655 + )
656 +# if ROCM endif
657 +endif()
658 +
659 message(STATUS "Enabling C extension.")
660 define_gpu_extension_target(
661 _C
csrc/custom_quickreduce.cu new
+114
@@ -0,0 +1,114 @@
1 +#include <ATen/cuda/Exceptions.h>
2 +#include <c10/cuda/CUDAGuard.h>
3 +#include <c10/cuda/CUDAStream.h>
4 +#include <torch/all.h>
5 +
6 +#ifdef USE_ROCM
7 +
8 + #include "quickreduce/quick_reduce.h"
9 +
10 +quickreduce::fptr_t init_custom_qr(int64_t rank, int64_t world_size,
11 + std::optional<int64_t> qr_max_size) {
12 + if (world_size > 8)
13 + throw std::invalid_argument("world size > 8 is not supported");
14 + if (world_size == 6)
15 + throw std::invalid_argument("world size == 6 is not supported");
16 + if (world_size % 2 != 0)
17 + throw std::invalid_argument("Odd num gpus is not supported for now");
18 + if (rank < 0 || rank >= world_size)
19 + throw std::invalid_argument("invalid rank passed in");
20 + quickreduce::DeviceComms* fptr = new quickreduce::DeviceComms();
21 + fptr->init(world_size, rank, qr_max_size);
22 + return (quickreduce::fptr_t)fptr;
23 +}
24 +
25 +void qr_destroy(quickreduce::fptr_t _fa) {
26 + if (_fa) {
27 + auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
28 + fa->destroy();
29 + delete fa;
30 + }
31 +}
32 +
33 +torch::Tensor qr_get_handle(quickreduce::fptr_t _fa) {
34 + auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
35 + hipIpcMemHandle_t handle = fa->get_handle();
36 + auto options =
37 + torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCPU);
38 + auto data_handle =
39 + torch::empty({static_cast<int64_t>(sizeof(hipIpcMemHandle_t))}, options);
40 + std::memcpy(data_handle.data_ptr(), &handle, sizeof(hipIpcMemHandle_t));
41 + return data_handle;
42 +}
43 +
44 +void qr_open_handles(quickreduce::fptr_t _fa,
45 + const std::vector<torch::Tensor>& handles) {
46 + auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
47 + std::vector<hipIpcMemHandle_t> ipc_handles;
48 + ipc_handles.reserve(handles.size());
49 + for (auto& handle : handles) {
50 + // Ensure the tensor is on the same device as the current device.
51 + hipIpcMemHandle_t ipc_handle;
52 + std::memcpy(&ipc_handle, handle.data_ptr(), sizeof(hipIpcMemHandle_t));
53 + ipc_handles.push_back(ipc_handle);
54 + }
55 + fa->open_ipc_handles(ipc_handles);
56 +}
57 +
58 +void qr_all_reduce(quickreduce::fptr_t _fa, torch::Tensor& inp,
59 + torch::Tensor& out, int64_t quant_level, bool cast_bf2half) {
60 + auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
61 + const at::cuda::OptionalCUDAGuard device_guard(device_of(inp));
62 + auto stream = at::cuda::getCurrentHIPStreamMasqueradingAsCUDA();
63 +
64 + TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
65 + TORCH_CHECK_EQ(inp.numel(), out.numel());
66 + TORCH_CHECK_LE(out.numel(), fa->kMaxProblemSize);
67 + if (out.scalar_type() == at::ScalarType::Half) {
68 + fa->allreduce<half, false>(reinterpret_cast<half*>(inp.data_ptr()),
69 + reinterpret_cast<half*>(out.data_ptr()),
70 + out.numel(), quant_level, stream);
71 + } else if (out.scalar_type() == at::ScalarType::BFloat16) {
72 + if (cast_bf2half) {
73 + fa->allreduce<half, true>(reinterpret_cast<half*>(inp.data_ptr()),
74 + reinterpret_cast<half*>(out.data_ptr()),
75 + out.numel(), quant_level, stream);
76 + } else {
77 + fa->allreduce<quickreduce::nv_bfloat16, false>(
78 + reinterpret_cast<quickreduce::nv_bfloat16*>(inp.data_ptr()),
79 + reinterpret_cast<quickreduce::nv_bfloat16*>(out.data_ptr()),
80 + out.numel(), quant_level, stream);
81 + }
82 + } else {
83 + throw std::runtime_error(
84 + "quick allreduce only supports float16 and bfloat16");
85 + }
86 +}
87 +
88 +int64_t qr_max_size() {
89 + // The default is 2GB (2,147,483,648 bytes)
90 + return static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1;
91 +}
92 +
93 + #define INSTANTIATE_FOR_WORLDSIZE(T, Codec, cast_bf2half) \
94 + template struct quickreduce::AllReduceTwoshot<T, Codec<T, 2>, \
95 + cast_bf2half>; \
96 + template struct quickreduce::AllReduceTwoshot<T, Codec<T, 4>, \
97 + cast_bf2half>; \
98 + template struct quickreduce::AllReduceTwoshot<T, Codec<T, 8>, cast_bf2half>;
99 +
100 +INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, false)
101 +INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, false)
102 +INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, false)
103 +INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, false)
104 +INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, true)
105 +INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, true)
106 +INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, true)
107 +INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, true)
108 +
109 +INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecFP, false)
110 +INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ4, false)
111 +INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ6, false)
112 +INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ8, false)
113 +
114 +#endif // USE_ROCM
\ No newline at end of file
csrc/ops.h
+11
@@ -360,3 +360,14 @@ std::tuple<int64_t, torch::Tensor> allocate_shared_buffer_and_handle(
360 int64_t size);
361 int64_t open_mem_handle(torch::Tensor& mem_handle);
362 void free_shared_buffer(int64_t buffer);
363 +
364 +#ifdef USE_ROCM
365 +fptr_t init_custom_qr(int64_t rank, int64_t world_size,
366 + std::optional<int64_t> qr_max_size = std::nullopt);
367 +void qr_destroy(fptr_t _fa);
368 +torch::Tensor qr_get_handle(fptr_t _fa);
369 +void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
370 +void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
371 + int64_t quant_level, bool cast_bf2half = false);
372 +int64_t qr_max_size();
373 +#endif
\ No newline at end of file
csrc/quickreduce/base.h new
+338
@@ -0,0 +1,338 @@
1 +#pragma once
2 +
3 +#include <cstdint>
4 +#include <hip/hip_runtime.h>
5 +#include <hip/hip_fp16.h>
6 +#include <hip/hip_bf16.h>
7 +
8 +#define __quickreduce_device_inline__ __device__ __forceinline__
9 +#define __quickreduce_launch_bounds_two_shot__ __launch_bounds__(256, 4)
10 +#define __quickreduce_launch_bounds_one_shot__ __launch_bounds__(512, 4)
11 +
12 +namespace quickreduce {
13 +
14 +typedef __hip_bfloat16 nv_bfloat16;
15 +typedef __hip_bfloat162 nv_bfloat162;
16 +
17 +using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int;
18 +using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int;
19 +
20 +// Setup acquire-release semantics for vector memory reads (mubuf instruction)
21 +// as per architecture.
22 +#if defined(__gfx942__)
23 +// CDNA3: Scope bits sc0, sc1
24 + #define MUBUF_ACQUIRE 16
25 + #define MUBUF_RELEASE 16
26 +#elif (defined(__gfx908__) || defined(__gfx90a__))
27 +// CDNA1 and CDNA2 - glc bit
28 + #define MUBUF_ACQUIRE 1
29 + #define MUBUF_RELEASE 0
30 +#endif
31 +
32 +static constexpr int kNegOne = 0xBC00BC00; // {-1, -1}, fp16x2_t
33 +
34 +// Number of atoms (4xf16x2_t) processed by a single thread
35 +static constexpr int kAtoms = 8;
36 +
37 +// We use a workgroup of 256 threads
38 +static constexpr int kBlockSize = 256;
39 +static constexpr int kAtomStride = kBlockSize;
40 +
41 +// Size and atom stride of source/destination data that the block will
42 +// process.
43 +// Workgroup scope = Tile = (256 threads x 8 atoms x 16B)
44 +static constexpr int kTileSize = kBlockSize * kAtoms * sizeof(int32x4_t);
45 +
46 +// Max number of blocks. 304 CUs on MI300
47 +static constexpr int kMaxNumBlocks = 304 * 4;
48 +
49 +// Standard CDNA wavefront size.
50 +static constexpr int kWavefront = 64;
51 +
52 +// 256 thread, 4 wavefronts.
53 +static dim3 constexpr kBlockTwoShot = {kWavefront, kBlockSize / kWavefront, 1};
54 +
55 +// Number of threads in a group for quantization
56 +// It corresponds to 32 F16 elements in quantization block
57 +static constexpr int kThreadGroupSize = 8;
58 +
59 +// Methods
60 +__quickreduce_device_inline__ __host__ unsigned long divceil(unsigned long x,
61 + unsigned long y) {
62 + return ((x + y - 1) / y);
63 +}
64 +
65 +union BufferResource {
66 + __quickreduce_device_inline__ constexpr BufferResource()
67 + : config(0x00020000U) {}
68 +
69 + __quickreduce_device_inline__ constexpr BufferResource(void* buffer_address,
70 + uint32_t buffer_size)
71 + : address(buffer_address), range(buffer_size), config(0x00020000U) {}
72 +
73 + int32x4_t descriptor;
74 + struct {
75 + void* address; // 8B, out of which first 48b is address, and 16b is stride
76 + // (unused)
77 + uint32_t range; // Byte range for the buffer resource
78 + uint32_t config; // Constant, DFMT=32b
79 + };
80 +};
81 +
82 +__quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4(
83 + int32x4_t srsrc, int32_t voffset, int32_t soffset,
84 + int32_t aux) __asm("llvm.amdgcn.raw.buffer.load.v4i32");
85 +
86 +__quickreduce_device_inline__ static void buffer_store_dwordx4(
87 + int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset,
88 + int32_t aux) __asm("llvm.amdgcn.raw.buffer.store.v4i32");
89 +
90 +__quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) {
91 +#if defined(__gfx942__)
92 + if (value) {
93 + asm volatile("s_setreg_imm32_b32 0xdc1, 1;" ::);
94 + } else {
95 + asm volatile("s_setreg_imm32_b32 0xdc1, 0;" ::);
96 + }
97 +#endif
98 +}
99 +union bf162_int_union {
100 + int i;
101 + nv_bfloat162 bf2;
102 +};
103 +
104 +template <typename T>
105 +__quickreduce_device_inline__ void packed_assign_add(int32x4_t* A,
106 + int32x4_t* B);
107 +
108 +template <>
109 +__quickreduce_device_inline__ void packed_assign_add<half>(int32x4_t* A,
110 + int32x4_t* B) {
111 + int32x4_t& tR_fragment = A[0];
112 + int32x4_t& tA_fragment = B[0];
113 +
114 + asm volatile("v_pk_add_f16 %0, %1, %2"
115 + : "=v"(tR_fragment[0])
116 + : "v"(tR_fragment[0]), "v"(tA_fragment[0]));
117 + asm volatile("v_pk_add_f16 %0, %1, %2"
118 + : "=v"(tR_fragment[1])
119 + : "v"(tR_fragment[1]), "v"(tA_fragment[1]));
120 + asm volatile("v_pk_add_f16 %0, %1, %2"
121 + : "=v"(tR_fragment[2])
122 + : "v"(tR_fragment[2]), "v"(tA_fragment[2]));
123 + asm volatile("v_pk_add_f16 %0, %1, %2"
124 + : "=v"(tR_fragment[3])
125 + : "v"(tR_fragment[3]), "v"(tA_fragment[3]));
126 +}
127 +
128 +template <>
129 +__quickreduce_device_inline__ void packed_assign_add<nv_bfloat16>(
130 + int32x4_t* A, int32x4_t* B) {
131 + nv_bfloat162* tA = reinterpret_cast<nv_bfloat162*>(A);
132 + nv_bfloat162* tB = reinterpret_cast<nv_bfloat162*>(B);
133 +#pragma unroll
134 + for (int i = 0; i < 4; i++) {
135 + tA[i] = __hadd2(tA[i], tB[i]);
136 + }
137 +}
138 +
139 +template <typename T>
140 +__quickreduce_device_inline__ int packed_max(int a, int b);
141 +
142 +template <>
143 +__quickreduce_device_inline__ int packed_max<half>(int a, int b) {
144 + int result;
145 + asm volatile("v_pk_max_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
146 + return result;
147 +}
148 +
149 +template <>
150 +__quickreduce_device_inline__ int packed_max<nv_bfloat16>(int a, int b) {
151 + bf162_int_union A, B, R;
152 + A.i = a;
153 + B.i = b;
154 + R.bf2 = __hmax2(A.bf2, B.bf2);
155 + return R.i;
156 +}
157 +
158 +template <typename T>
159 +__quickreduce_device_inline__ int packed_min(int a, int b);
160 +
161 +template <>
162 +__quickreduce_device_inline__ int packed_min<half>(int a, int b) {
163 + int result;
164 + asm volatile("v_pk_min_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
165 + return result;
166 +}
167 +
168 +template <>
169 +__quickreduce_device_inline__ int packed_min<nv_bfloat16>(int a, int b) {
170 + bf162_int_union A, B, R;
171 + A.i = a;
172 + B.i = b;
173 + R.bf2 = __hmin2(A.bf2, B.bf2);
174 + return R.i;
175 +}
176 +
177 +template <typename T>
178 +__quickreduce_device_inline__ int packed_abs_max(int a, int b);
179 +
180 +template <>
181 +__quickreduce_device_inline__ int packed_abs_max<half>(int a, int b) {
182 + half2 wmaxh2 = __builtin_bit_cast(half2, a);
183 + half2 wminh2 = __builtin_bit_cast(half2, b);
184 + half2 wblockmaxh2;
185 +
186 + wblockmaxh2.x =
187 + __hgt(__habs(wmaxh2.x), __habs(wminh2.x)) ? wmaxh2.x : wminh2.x;
188 + wblockmaxh2.y =
189 + __hgt(__habs(wmaxh2.y), __habs(wminh2.y)) ? wmaxh2.y : wminh2.y;
190 + return __builtin_bit_cast(int, wblockmaxh2);
191 +}
192 +
193 +template <>
194 +__quickreduce_device_inline__ int packed_abs_max<nv_bfloat16>(int a, int b) {
195 + bf162_int_union A, B, R;
196 + A.i = a;
197 + B.i = b;
198 + R.bf2.x = __hgt(__habs(A.bf2.x), __habs(B.bf2.x)) ? A.bf2.x : B.bf2.x;
199 + R.bf2.y = __hgt(__habs(A.bf2.y), __habs(B.bf2.y)) ? A.bf2.y : B.bf2.y;
200 + return R.i;
201 +}
202 +
203 +template <typename T>
204 +__quickreduce_device_inline__ int packed_add(int a, int b);
205 +
206 +template <>
207 +__quickreduce_device_inline__ int packed_add<half>(int a, int b) {
208 + int result;
209 + asm volatile("v_pk_add_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
210 + return result;
211 +}
212 +
213 +template <>
214 +__quickreduce_device_inline__ int packed_add<nv_bfloat16>(int a, int b) {
215 + bf162_int_union A, B, R;
216 + A.i = a;
217 + B.i = b;
218 + R.bf2 = __hadd2(A.bf2, B.bf2);
219 + return R.i;
220 +}
221 +
222 +template <>
223 +__quickreduce_device_inline__ int packed_add<int16_t>(int a, int b) {
224 + int result;
225 + asm volatile("v_pk_add_i16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
226 + return result;
227 +}
228 +
229 +template <typename T>
230 +__quickreduce_device_inline__ int packed_sub(int a, int b);
231 +
232 +template <>
233 +__quickreduce_device_inline__ int packed_sub<half>(int a, int b) {
234 + int result;
235 +
236 + // MI300 lacks packed fp16 sub instruction. So we do -1 * min + max
237 + asm volatile("v_pk_fma_f16 %0, %1, %2 %3"
238 + : "=v"(result)
239 + : "v"(kNegOne), "v"(b), "v"(a));
240 + return result;
241 +}
242 +
243 +template <>
244 +__quickreduce_device_inline__ int packed_sub<nv_bfloat16>(int a, int b) {
245 + bf162_int_union A, B, R;
246 + A.i = a;
247 + B.i = b;
248 + R.bf2 = __hsub2(A.bf2, B.bf2);
249 + return R.i;
250 +}
251 +
252 +template <typename T>
253 +__quickreduce_device_inline__ int packed_mul(int a, int b);
254 +
255 +template <>
256 +__quickreduce_device_inline__ int packed_mul<half>(int a, int b) {
257 + int result;
258 + asm volatile("v_pk_mul_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
259 + return result;
260 +}
261 +
262 +template <>
263 +__quickreduce_device_inline__ int packed_mul<nv_bfloat16>(int a, int b) {
264 + nv_bfloat162* tA = reinterpret_cast<nv_bfloat162*>(&a);
265 + nv_bfloat162* tB = reinterpret_cast<nv_bfloat162*>(&b);
266 + nv_bfloat162 tR = __hmul2(*tA, *tB);
267 + return *(reinterpret_cast<int*>(&tR));
268 +}
269 +
270 +template <typename T>
271 +__quickreduce_device_inline__ int packed_rcp(int a);
272 +
273 +template <>
274 +__quickreduce_device_inline__ int packed_rcp<half>(int a) {
275 + return __builtin_bit_cast(int, h2rcp(__builtin_bit_cast(half2, a)));
276 +}
277 +
278 +template <>
279 +__quickreduce_device_inline__ int packed_rcp<nv_bfloat16>(int a) {
280 + bf162_int_union A, R;
281 + A.i = a;
282 + R.bf2 = h2rcp(A.bf2);
283 + return R.i;
284 +}
285 +
286 +// changes dtype
287 +__quickreduce_device_inline__ float T2float_cast(half a) {
288 + return __half2float(a);
289 +}
290 +
291 +__quickreduce_device_inline__ float T2float_cast(nv_bfloat16 a) {
292 + return __bfloat162float(a);
293 +}
294 +
295 +template <typename T>
296 +__quickreduce_device_inline__ int group_abs_max(int32x4_t atom) {
297 + const int group_leader = (threadIdx.x / kThreadGroupSize) * kThreadGroupSize;
298 +
299 + int wmax, wmin, wblockmax;
300 + int a, b;
301 + a = packed_max<T>(atom[0], atom[1]);
302 + b = packed_max<T>(atom[2], atom[3]);
303 +
304 + wmax = packed_max<T>(a, b);
305 +
306 + a = packed_min<T>(atom[0], atom[1]);
307 + b = packed_min<T>(atom[2], atom[3]);
308 +
309 + wmin = packed_min<T>(a, b);
310 +
311 + // Reduce the max among a group of threads
312 + // Note: This is basically 2 blocks of values setup as the
313 + // upper/lower halves of the f16x2_t
314 + for (int i = 1; i < kThreadGroupSize; i <<= 1) {
315 + int x = __shfl_down(wmax, i);
316 + wmax = packed_max<T>(wmax, x);
317 +
318 + int y = __shfl_down(wmin, i);
319 + wmin = packed_min<T>(wmin, y);
320 + }
321 + wblockmax = packed_abs_max<T>(wmax, wmin);
322 + // Share with the cohort
323 + wblockmax = __shfl(wblockmax, group_leader);
324 + return wblockmax;
325 +}
326 +
327 +__quickreduce_device_inline__ void set_sync_flag(uint32_t* flag_ptr,
328 + uint32_t flag) {
329 + __atomic_store_n(flag_ptr, flag, __ATOMIC_RELEASE);
330 +}
331 +
332 +__quickreduce_device_inline__ void wait_sync_flag(uint32_t* flag_ptr,
333 + uint32_t flag) {
334 + while (__atomic_load_n(flag_ptr, __ATOMIC_RELAXED) != flag) {
335 + }
336 +}
337 +
338 +} // namespace quickreduce
\ No newline at end of file
csrc/quickreduce/quick_reduce.h new
+196
@@ -0,0 +1,196 @@
1 +#pragma once
2 +
3 +#include <vector>
4 +#include <hip/hip_runtime.h>
5 +#include "quick_reduce_impl.cuh"
6 +
7 +#define HIP_CHECK(err) \
8 + do { \
9 + hipError_t err_ = (err); \
10 + if (err_ != hipSuccess) { \
11 + std::printf("HIP error %d at %s:%d. %s\n", err_, __FILE__, __LINE__, \
12 + hipGetErrorString(err_)); \
13 + throw std::runtime_error("HIP error"); \
14 + } \
15 + } while (0)
16 +
17 +namespace quickreduce {
18 +using fptr_t = int64_t;
19 +static_assert(sizeof(void*) == sizeof(fptr_t));
20 +
21 +template <typename AllReduceKernel, typename T>
22 +__global__ __quickreduce_launch_bounds_two_shot__ static void
23 +allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
24 + int rank, uint8_t** dbuffer_list,
25 + uint32_t data_offset, uint32_t flag_color) {
26 + int block = blockIdx.x;
27 + int grid = gridDim.x;
28 +
29 + while (block < num_blocks) {
30 + AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset,
31 + flag_color);
32 + block += grid;
33 + flag_color++;
34 + }
35 +}
36 +
37 +#define TWOSHOT_DISPATCH(__codec) \
38 + if (world_size == 2) { \
39 + using LineCodec = __codec<T, 2>; \
40 + using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
41 + hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
42 + dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
43 + num_blocks, rank, dbuffer_list, data_offset, \
44 + flag_color); \
45 + } else if (world_size == 4) { \
46 + using LineCodec = __codec<T, 4>; \
47 + using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
48 + hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
49 + dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
50 + num_blocks, rank, dbuffer_list, data_offset, \
51 + flag_color); \
52 + } else if (world_size == 8) { \
53 + using LineCodec = __codec<T, 8>; \
54 + using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
55 + hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
56 + dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
57 + num_blocks, rank, dbuffer_list, data_offset, \
58 + flag_color); \
59 + }
60 +
61 +enum QuickReduceQuantLevel {
62 + F16 = 0,
63 + INT8 = 1,
64 + INT6 = 2,
65 + INT4 = 3,
66 +};
67 +
68 +struct DeviceComms {
69 + // Max problem size is 2GB (in bytes) or half of uint32_t max value.
70 + int64_t kMaxProblemSize =
71 + static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1;
72 +
73 + // Max TP-8
74 + static int constexpr kMaxWorldSize = 8;
75 +
76 + bool initialized = false;
77 + uint32_t flag_color = 1;
78 + int world_size;
79 + int rank;
80 +
81 + uint8_t* dbuffer;
82 + uint8_t** dbuffer_list;
83 + hipIpcMemHandle_t buffer_ipc_handle;
84 + std::vector<hipIpcMemHandle_t> all_buffer_ipc_handles;
85 + std::vector<uint8_t*> buffer_list;
86 + uint32_t data_offset;
87 +
88 + DeviceComms() : initialized(false), world_size(1), rank(0) {}
89 + ~DeviceComms() { destroy(); }
90 +
91 + void init(int world_size, int rank,
92 + std::optional<int64_t> max_problem_size = std::nullopt) {
93 + destroy();
94 + this->world_size = world_size;
95 + this->rank = rank;
96 + if (max_problem_size.has_value() && max_problem_size.value() > 0) {
97 + this->kMaxProblemSize = max_problem_size.value();
98 + }
99 + // Allocate buffer size for worst case: F16 2-stage buffer.
100 + uint32_t flags_buffer_size =
101 + 2 * world_size * kMaxNumBlocks * sizeof(uint32_t);
102 + static int64_t data_buffer_size = 2 * this->kMaxProblemSize;
103 + int64_t total_buffer_size = flags_buffer_size + data_buffer_size;
104 + data_offset = flags_buffer_size;
105 + HIP_CHECK(hipExtMallocWithFlags((void**)&dbuffer, total_buffer_size,
106 + hipDeviceMallocUncached));
107 +
108 + // Clear the flags buffer.
109 + HIP_CHECK(hipMemset(dbuffer, 0, flags_buffer_size));
110 +
111 + // Device-side list of IPC buffers.
112 + buffer_list.resize(world_size);
113 + HIP_CHECK(hipMalloc(&dbuffer_list, world_size * sizeof(uint8_t*)));
114 +
115 + // Create IPC handles for rank's communication buffer.
116 + all_buffer_ipc_handles.resize(world_size);
117 + HIP_CHECK(hipIpcGetMemHandle(&buffer_ipc_handle, dbuffer));
118 +
119 + initialized = true;
120 + }
121 + int get_world_size() { return world_size; }
122 + int get_rank() { return rank; }
123 + bool status() { return initialized; }
124 + hipIpcMemHandle_t const get_handle() { return buffer_ipc_handle; }
125 +
126 + void destroy() {
127 + if (initialized) {
128 + for (int i = 0; i < world_size; i++) {
129 + if (i != rank) {
130 + HIP_CHECK(hipIpcCloseMemHandle(dbuffer_list[i]));
131 + }
132 + }
133 +
134 + HIP_CHECK(hipFree(dbuffer));
135 + HIP_CHECK(hipFree(dbuffer_list));
136 +
137 + initialized = false;
138 + }
139 + }
140 +
141 + void open_ipc_handles(std::vector<hipIpcMemHandle_t> const& ipc_handles) {
142 + assert(ipc_handles.size() == all_buffer_ipc_handles.size());
143 + for (int i = 0; i < world_size; i++) {
144 + all_buffer_ipc_handles[i] = ipc_handles[i];
145 + }
146 +
147 + // Open device memory access to the IPC communication buffers.
148 + // Note: For our own rank, we do not need to open a handle.
149 + for (int i = 0; i < world_size; i++) {
150 + if (i != rank) {
151 + HIP_CHECK(hipIpcOpenMemHandle((void**)&buffer_list[i],
152 + all_buffer_ipc_handles[i],
153 + hipIpcMemLazyEnablePeerAccess));
154 + } else {
155 + buffer_list[i] = dbuffer;
156 + }
157 + }
158 +
159 + HIP_CHECK(hipMemcpy(dbuffer_list, buffer_list.data(),
160 + world_size * sizeof(uint8_t*), hipMemcpyHostToDevice));
161 + }
162 +
163 + template <typename T, bool cast_bf2half>
164 + void allreduce(T const* A, T* B, uint32_t N, int quant_level,
165 + hipStream_t stream) {
166 + if (world_size != 2 && world_size != 4 && world_size != 8) {
167 + throw std::runtime_error("All Reduce not supported for world_size = " +
168 + std::to_string(world_size));
169 + }
170 +
171 + // Configuration.
172 + uint32_t msg_size = N * sizeof(T);
173 + uint32_t num_blocks = divceil(msg_size, kTileSize);
174 + uint32_t grid = min(kMaxNumBlocks, num_blocks);
175 + auto quant_level_ = static_cast<QuickReduceQuantLevel>(quant_level);
176 + switch (quant_level_) {
177 + case QuickReduceQuantLevel::INT8:
178 + TWOSHOT_DISPATCH(CodecQ8)
179 + break;
180 + case QuickReduceQuantLevel::INT6:
181 + TWOSHOT_DISPATCH(CodecQ6)
182 + break;
183 + case QuickReduceQuantLevel::INT4:
184 + TWOSHOT_DISPATCH(CodecQ4)
185 + break;
186 + default:
187 + TWOSHOT_DISPATCH(CodecFP)
188 + break;
189 + }
190 + HIP_CHECK(cudaGetLastError());
191 + // Rotate the flag color.
192 + flag_color += divceil(N, grid);
193 + }
194 +};
195 +
196 +} // namespace quickreduce
\ No newline at end of file
csrc/quickreduce/quick_reduce_impl.cuh new
+698
@@ -0,0 +1,698 @@
1 +#pragma once
2 +
3 +#include <hip/hip_runtime.h>
4 +#include "base.h"
5 +
6 +namespace quickreduce {
7 +
8 +struct CodecBase {
9 + const int thread;
10 + const int rank;
11 + const int group_leader;
12 + __quickreduce_device_inline__ CodecBase(int thread, int rank)
13 + : thread(thread),
14 + rank(rank),
15 + group_leader((threadIdx.x / kThreadGroupSize) * kThreadGroupSize) {
16 + set_fp16_ovfl(true);
17 + }
18 +};
19 +
20 +// Default full precision codec.
21 +template <typename T, int world_size>
22 +struct CodecFP : public CodecBase {
23 + static constexpr int kWorldSize = world_size;
24 + static constexpr int kRankAtoms = kAtoms / kWorldSize;
25 +
26 + // Codec tile size process by this workgroup.
27 + // Each thread processes atoms of f16x8_t (16B).
28 + static constexpr int kRankTransmittedTileSize =
29 + kBlockSize * kRankAtoms * sizeof(int32x4_t);
30 + static_assert(kRankTransmittedTileSize % 16 == 0,
31 + "kRankTransmittedTileSize must be 16B aligned.");
32 +
33 + // Total tile size for the collective communication.
34 + static constexpr int kTransmittedTileSize =
35 + kRankTransmittedTileSize * kWorldSize;
36 +
37 + __quickreduce_device_inline__ CodecFP(int thread, int rank)
38 + : CodecBase(thread, rank) {}
39 +
40 + __quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
41 + const int32x4_t* __restrict__ data) {
42 + for (int i = 0; i < kRankAtoms; i++) {
43 + __builtin_nontemporal_store(data[i], send_buffer + thread);
44 + send_buffer += kAtomStride;
45 + }
46 + }
47 +
48 + __quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
49 + int32x4_t* __restrict__ data) {
50 + for (int i = 0; i < kRankAtoms; i++) {
51 + data[i] = __builtin_nontemporal_load(*recv_buffer + thread);
52 + *recv_buffer += kAtomStride;
53 + }
54 + }
55 +};
56 +
57 +// Int4 symmetric quantization codec.
58 +// We quantize the FP16 data to block-scaled Int4 in blocks of 4 *
59 +// kThreadGroupSize.
60 +template <typename T, int world_size>
61 +struct CodecQ4 : public CodecBase {
62 + static constexpr int kWorldSize = world_size;
63 +
64 + // Codec tile size process by this workgroup.
65 + // Each threads processes a fragment of fp16x8_t (16B),
66 + // into a int4x8_t (4B) and a fp16 scale shared among 32 values.
67 + static constexpr int kRankAtoms = kAtoms / kWorldSize;
68 + static constexpr int kRankTileStride = 1152;
69 + static constexpr int kRankTileScaleOffset = 1024;
70 + static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms;
71 + static_assert(kRankTransmittedTileSize % 16 == 0,
72 + "kRankTransmittedTileSize must be 16B aligned.");
73 +
74 + static constexpr int kRankBufferTileStride =
75 + kRankTileStride / sizeof(int32x4_t);
76 +
77 + // Total tile size for the collective communication.
78 + static constexpr int kTransmittedTileSize =
79 + kRankTransmittedTileSize * kWorldSize;
80 +
81 + // Constants configuration
82 +
83 + // {-1/8.0h, -1/8.0h}, f16x2_t
84 + static constexpr int kScaleFactor =
85 + std::is_same<T, half>::value ? 0xB000B000 : 0xBE00BE00;
86 +
87 + // {1e-7, 1e-7}, f16x2_t
88 + static constexpr int kScaleEpsilon =
89 + std::is_same<T, half>::value ? 0x00010001 : 0x33D733D7;
90 +
91 + // {-8, -8}, f16x2_t
92 + static constexpr int kRangeMin =
93 + std::is_same<T, half>::value ? 0xC800C800 : 0xC100C100;
94 +
95 + // {+7, +7}, f16x2_t
96 + static constexpr int kRangeMax =
97 + std::is_same<T, half>::value ? 0x47004700 : 0x40E040E0;
98 +
99 + // {+8, +8}, int16x2_t
100 + static constexpr int kRangeBias = 0x00080008;
101 +
102 + __quickreduce_device_inline__ CodecQ4(int thread, int rank)
103 + : CodecBase(thread, rank) {}
104 +
105 + __quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
106 + const int32x4_t* __restrict__ data) {
107 + for (int k = 0; k < kRankAtoms; k++) {
108 + int32x4_t const atom = data[k];
109 +
110 + // Compute the absolute maximum of the atom in the thread group
111 + // In 2 blocks of values, upper/lower halves of the f16x2_t
112 + int wblockmax = group_abs_max<T>(atom);
113 +
114 + // Derive scales
115 + int decoding_scale;
116 + int encoding_scale;
117 + decoding_scale = packed_mul<T>(wblockmax, kScaleFactor);
118 + encoding_scale = packed_add<T>(decoding_scale, kScaleEpsilon);
119 + encoding_scale = packed_rcp<T>(encoding_scale);
120 +
121 + // Apply scales to get quantized values
122 + int32x4_t w;
123 + for (int i = 0; i < 4; i++) {
124 + w[i] = packed_mul<T>(atom[i], encoding_scale);
125 + w[i] = packed_max<T>(w[i], kRangeMin);
126 + w[i] = packed_min<T>(w[i], kRangeMax);
127 + }
128 +
129 + // Convert from f16x2_t to uint16x2_t
130 + int32x4_t q;
131 + {
132 + int16_t* qi = reinterpret_cast<int16_t*>(&q);
133 + T* wh = reinterpret_cast<T*>(&w);
134 + for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i]));
135 +
136 + for (int i = 0; i < 4; i++) {
137 + q[i] = packed_add<int16_t>(q[i], kRangeBias);
138 + }
139 + }
140 +
141 + // Pack 8 x q4 into int32_t
142 + int qw = q[0] | (q[1] << 4) | (q[2] << 8) | (q[3] << 12);
143 +
144 + // Write quantized atom to send_buffer
145 + // note: only the group leader stores the scale
146 + uint8_t* atom_ptr =
147 + reinterpret_cast<uint8_t*>(send_buffer + k * kRankBufferTileStride);
148 + int32_t* qw_ptr = reinterpret_cast<int32_t*>(atom_ptr) + thread;
149 + int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
150 + (thread / 8);
151 +
152 + __builtin_nontemporal_store(qw, qw_ptr);
153 + if (threadIdx.x == group_leader) {
154 + __builtin_nontemporal_store(decoding_scale, qs_ptr);
155 + }
156 + }
157 + }
158 +
159 + __quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
160 + int32x4_t* __restrict__ data) {
161 + for (int k = 0; k < kRankAtoms; k++) {
162 + // Directly read quantized atom from recv_buffer
163 + uint8_t* atom_ptr = reinterpret_cast<uint8_t*>(*recv_buffer);
164 + int32_t* qw_ptr = reinterpret_cast<int32_t*>(atom_ptr) + thread;
165 + int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
166 + (thread / 8);
167 +
168 + int32_t qw = __builtin_nontemporal_load(qw_ptr);
169 + int qs = __builtin_nontemporal_load(qs_ptr);
170 +
171 + *recv_buffer += kRankBufferTileStride;
172 +
173 + // Unpack q4 into f16x8_t
174 + int32x4_t w;
175 + {
176 + static constexpr uint kMask000F = 0x000F000F;
177 + static constexpr uint kHalf2_1024 =
178 + 0x64006400; // {1024.0, 1024.0}, fp16x2_t
179 + static uint constexpr kHalf2_1032 =
180 + 0xE408E408; // {-1032.0, -1032.0}, fp16x2_t
181 +
182 + for (int i = 0; i < 4; i++) {
183 + if constexpr (std::is_same<T, half>::value) {
184 + int32_t q4 = ((qw >> (i * 4)) & kMask000F) | kHalf2_1024;
185 + w[i] = packed_add<half>(q4, kHalf2_1032);
186 + } else {
187 + int32_t int16_2 = (qw >> (i * 4)) & kMask000F;
188 + int16_t low = static_cast<int16_t>(int16_2 & 0xFFFF);
189 + int16_t high = static_cast<int16_t>((int16_2 >> 16) & 0xFFFF);
190 + nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
191 + nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
192 + nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
193 + int32_t packed_bf16 = *reinterpret_cast<int32_t*>(&bf2);
194 + w[i] = packed_add<nv_bfloat16>(packed_bf16, kRangeMin);
195 + }
196 + }
197 + }
198 +
199 + // Apply decoding scales
200 + for (int i = 0; i < 4; i++) {
201 + w[i] = packed_mul<T>(w[i], qs);
202 + }
203 +
204 + data[k] = w;
205 + }
206 + }
207 +};
208 +
209 +// Int6 symmetric quantization codec.
210 +// We quantize the FP16 data to block-scaled Int6 in blocks of 4 *
211 +// kThreadGroupSize.
212 +template <typename T, int world_size>
213 +struct CodecQ6 : public CodecBase {
214 + static constexpr int kWorldSize = world_size;
215 +
216 + // Codec tile size process by this workgroup.
217 + // Each threads processes a fragment of fp16x8_t (16B),
218 + // into a int6x8_t (4B + 2B) and a fp16 scale shared among 32 values.
219 + static constexpr int kRankAtoms = kAtoms / kWorldSize;
220 + static constexpr int kRankTileStride = 1664;
221 + static constexpr int kRankTileQ2Offset = 1024;
222 + static constexpr int kRankTileScaleOffset = 1536;
223 + static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms;
224 + static_assert(kRankTransmittedTileSize % 16 == 0,
225 + "kRankTransmittedTileSize must be 16B aligned.");
226 +
227 + static constexpr int kRankBufferTileStride =
228 + kRankTileStride / sizeof(int32x4_t);
229 +
230 + // Total tile size for the collective communication.
231 + static constexpr int kTransmittedTileSize =
232 + kRankTransmittedTileSize * kWorldSize;
233 +
234 + // Constants configuration
235 +
236 + // {-1/32.0h, -1/32.0h}, fp16x2_t
237 + static constexpr int kScaleFactor =
238 + std::is_same<T, half>::value ? 0xA800A800 : 0xBD00BD00;
239 +
240 + // {1e-7, 1e-7}, fp16x2_t
241 + static constexpr int kScaleEpsilon =
242 + std::is_same<T, half>::value ? 0x00010001 : 0x33D733D7;
243 +
244 + // {-32, -32}, fp16x2_t
245 + static constexpr int kRangeMin =
246 + std::is_same<T, half>::value ? 0xD000D000 : 0xC200C200;
247 +
248 + // {+31, +31}, fp16x2_t
249 + static constexpr int kRangeMax =
250 + std::is_same<T, half>::value ? 0x4FC04FC0 : 0x41F841F8;
251 +
252 + // {+32, +32}, int16x2_t
253 + static constexpr int kRangeBias = 0x00200020;
254 +
255 + __quickreduce_device_inline__ CodecQ6(int thread, int rank)
256 + : CodecBase(thread, rank) {}
257 +
258 + __quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
259 + const int32x4_t* __restrict__ data) {
260 + for (int k = 0; k < kRankAtoms; k++) {
261 + int32x4_t const atom = data[k];
262 +
263 + // Compute the absolute maximum of the atom in the thread group
264 + // In 2 blocks of values, upper/lower halves of the f16x2_t
265 + int wblockmax = group_abs_max<T>(atom);
266 +
267 + // Derive scales
268 + int decoding_scale;
269 + int encoding_scale;
270 + decoding_scale = packed_mul<T>(wblockmax, kScaleFactor);
271 + encoding_scale = packed_add<T>(decoding_scale, kScaleEpsilon);
272 + encoding_scale = packed_rcp<T>(encoding_scale);
273 +
274 + // Apply scales to get quantized values
275 + int32x4_t w;
276 + for (int i = 0; i < 4; i++) {
277 + w[i] = packed_mul<T>(atom[i], encoding_scale);
278 + w[i] = packed_max<T>(w[i], kRangeMin);
279 + w[i] = packed_min<T>(w[i], kRangeMax);
280 + }
281 +
282 + // Convert from f16x2_t to uint16x2_t
283 + int32x4_t q;
284 + {
285 + int16_t* qi = reinterpret_cast<int16_t*>(&q);
286 + T* wh = reinterpret_cast<T*>(&w);
287 + for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i]));
288 +
289 + for (int i = 0; i < 4; i++) {
290 + q[i] = packed_add<int16_t>(q[i], kRangeBias);
291 + }
292 + }
293 +
294 + // Pack 8 x q6 into int32_t + int16_t
295 + uint32_t q4w;
296 + uint16_t q2w = 0;
297 + q4w = (q[0] & 0x000F000F) | ((q[1] & 0x000F000F) << 4) |
298 + ((q[2] & 0x000F000F) << 8) | ((q[3] & 0x000F000F) << 12);
299 + {
300 + int16_t* tw = reinterpret_cast<int16_t*>(&q);
301 +#pragma unroll
302 + for (int i = 0; i < 8; i++) {
303 + q2w |= (tw[i] >> 4) << (i * 2);
304 + }
305 + }
306 + // Write quantized atom to send_buffer
307 + // note: only the group leader stores the scale
308 + uint8_t* atom_ptr =
309 + reinterpret_cast<uint8_t*>(send_buffer + k * kRankBufferTileStride);
310 + uint32_t* q4w_ptr = reinterpret_cast<uint32_t*>(atom_ptr) + thread;
311 + uint16_t* q2w_ptr =
312 + reinterpret_cast<uint16_t*>(atom_ptr + kRankTileQ2Offset) + thread;
313 + int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
314 + (thread / 8);
315 +
316 + __builtin_nontemporal_store(q4w, q4w_ptr);
317 + __builtin_nontemporal_store(q2w, q2w_ptr);
318 + if (threadIdx.x == group_leader) {
319 + __builtin_nontemporal_store(decoding_scale, qs_ptr);
320 + }
321 + }
322 + }
323 +
324 + __quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
325 + int32x4_t* __restrict__ data) {
326 + for (int k = 0; k < kRankAtoms; k++) {
327 + // Directly read quantized atom from recv_buffer
328 + uint8_t* atom_ptr = reinterpret_cast<uint8_t*>(*recv_buffer);
329 + uint32_t* q4w_ptr = reinterpret_cast<uint32_t*>(atom_ptr) + thread;
330 + uint16_t* q2w_ptr =
331 + reinterpret_cast<uint16_t*>(atom_ptr + kRankTileQ2Offset) + thread;
332 + int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
333 + (thread / 8);
334 +
335 + uint32_t q4w = __builtin_nontemporal_load(q4w_ptr);
336 + uint16_t q2w = __builtin_nontemporal_load(q2w_ptr);
337 + int qs = __builtin_nontemporal_load(qs_ptr);
338 +
339 + *recv_buffer += kRankBufferTileStride;
340 +
341 + // Unpack q6 into fp16x8_t
342 + int32x4_t w;
343 + {
344 + static uint constexpr kMask000F = 0x000F000F;
345 + static uint constexpr kHalf2_1024 =
346 + 0x64006400; // {1024.0, 1024.0}, fp16x2_t
347 + static uint constexpr kHalf2_1056 =
348 + 0xE420E420; // {-1056.0, -1056.0}, fp16x2_t
349 +
350 +#pragma unroll
351 + for (int i = 0; i < 4; i++) {
352 + int32_t q4 = q4w & kMask000F;
353 + int32_t q2 = (q2w & 0x3) | ((q2w & 0xC) << 14);
354 + q4w >>= 4;
355 + q2w >>= 4;
356 + if constexpr (std::is_same<T, half>::value) {
357 + int32_t q6 = q4 | (q2 << 4) | kHalf2_1024;
358 + asm volatile("v_pk_add_f16 %0, %1, %2"
359 + : "=v"(w[i])
360 + : "v"(q6), "v"(kHalf2_1056));
361 + } else {
362 + int32_t int16_2 = q4 | (q2 << 4);
363 + int16_t low = static_cast<int16_t>(int16_2 & 0xFFFF);
364 + int16_t high = static_cast<int16_t>((int16_2 >> 16) & 0xFFFF);
365 +
366 + nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
367 + nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
368 + nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
369 + int32_t packed_bf16 = *reinterpret_cast<int32_t*>(&bf2);
370 + w[i] = packed_add<nv_bfloat16>(packed_bf16, kRangeMin);
371 + }
372 + }
373 + }
374 +
375 + // Apply decoding scales
376 + for (int i = 0; i < 4; i++) {
377 + w[i] = packed_mul<T>(w[i], qs);
378 + }
379 +
380 + // That's pretty much it...
381 + data[k] = w;
382 + }
383 + }
384 +};
385 +
386 +// Int8 symmetric quantization codec.
387 +// We quantize the FP16 data to block-scaled Int8 in blocks of 4 *
388 +// kThreadGroupSize.
389 +template <typename T, int world_size>
390 +struct CodecQ8 : public CodecBase {
391 + static constexpr int kWorldSize = world_size;
392 +
393 + // Codec tile size process by this workgroup.
394 + // Each threads processes a fragment of f16x8_t (16B),
395 + // into a int8x8_t (8B) and a f16 scale shared among 32 values.
396 + static constexpr int kRankAtoms = kAtoms / kWorldSize;
397 + static constexpr int kRankTileStride = 2176;
398 + static constexpr int kRankTileScaleOffset = 2048;
399 + static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms;
400 + static_assert(kRankTransmittedTileSize % 16 == 0,
401 + "kRankTileSize must be 16B aligned.");
402 +
403 + static constexpr int kRankBufferTileStride =
404 + kRankTileStride / sizeof(int32x4_t);
405 +
406 + // Total tile size for the collective communication.
407 + static constexpr int kTransmittedTileSize =
408 + kRankTransmittedTileSize * kWorldSize;
409 +
410 + // Constants configuration
411 +
412 + // {-1/128.0h, -1/128.0h}, f16x2_t
413 + static constexpr int kScaleFactor =
414 + std::is_same<T, half>::value ? 0xA000A000 : 0xBC00BC00;
415 +
416 + // {1e-7, 1e-7}, f16x2_t
417 + static constexpr int kScaleEpsilon =
418 + std::is_same<T, half>::value ? 0x00010001 : 0x33D733D7;
419 +
420 + // {-128, -128}, f16x2_t
421 + static constexpr int kRangeMin =
422 + std::is_same<T, half>::value ? 0xD800D800 : 0xC300C300;
423 + // {+127, +127}, f16x2_t
424 + static constexpr int kRangeMax =
425 + std::is_same<T, half>::value ? 0x57F057F0 : 0x42FE42FE;
426 +
427 + // {+128, +128}, int16x2_t
428 + static constexpr int kRangeBias = 0x00800080;
429 +
430 + __quickreduce_device_inline__ CodecQ8(int thread, int rank)
431 + : CodecBase(thread, rank) {}
432 +
433 + __quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
434 + int32x4_t const* __restrict__ data) {
435 + for (int k = 0; k < kRankAtoms; k++) {
436 + int32x4_t const atom = data[k];
437 + // Compute the absolute maximum of the atom in the thread group
438 + // In 2 blocks of values, upper/lower halves of the f16x2_t
439 + int wblockmax = group_abs_max<T>(atom);
440 +
441 + // Derive scales
442 + int decoding_scale;
443 + int encoding_scale;
444 + decoding_scale = packed_mul<T>(wblockmax, kScaleFactor);
445 + encoding_scale = packed_add<T>(decoding_scale, kScaleEpsilon);
446 + encoding_scale = packed_rcp<T>(encoding_scale);
447 +
448 + // Apply scales to get quantized values
449 + int32x4_t w;
450 + for (int i = 0; i < 4; i++) {
451 + w[i] = packed_mul<T>(atom[i], encoding_scale);
452 + w[i] = packed_max<T>(w[i], kRangeMin);
453 + w[i] = packed_min<T>(w[i], kRangeMax);
454 + }
455 +
456 + // Convert from f16x2_t to uint16x2_t
457 + int32x4_t q;
458 + {
459 + int16_t* qi = reinterpret_cast<int16_t*>(&q);
460 + T* wh = reinterpret_cast<T*>(&w);
461 + for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i]));
462 +
463 + for (int i = 0; i < 4; i++) {
464 + q[i] = packed_add<int16_t>(q[i], kRangeBias);
465 + }
466 + }
467 +
468 + // Pack 8 x q8 into int32x2_t
469 + int32x2_t qw;
470 + qw[0] = q[0] | (q[1] << 8);
471 + qw[1] = q[2] | (q[3] << 8);
472 +
473 + // Write quantized atom to send_buffer
474 + // note: only the group leader stores the scale
475 + uint8_t* atom_ptr =
476 + reinterpret_cast<uint8_t*>(send_buffer + k * kRankBufferTileStride);
477 + int32x2_t* qw_ptr = reinterpret_cast<int32x2_t*>(atom_ptr) + thread;
478 + int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
479 + (thread / 8);
480 +
481 + __builtin_nontemporal_store(qw, qw_ptr);
482 + if (threadIdx.x == group_leader) {
483 + __builtin_nontemporal_store(decoding_scale, qs_ptr);
484 + }
485 + }
486 + }
487 +
488 + __quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
489 + int32x4_t* __restrict__ data) {
490 + for (int k = 0; k < kRankAtoms; k++) {
491 + // Directly read quantized atom from recv_buffer
492 + uint8_t* atom_ptr = reinterpret_cast<uint8_t*>(*recv_buffer);
493 + int32x2_t* qw_ptr = reinterpret_cast<int32x2_t*>(atom_ptr) + thread;
494 + int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
495 + (thread / 8);
496 +
497 + int32x2_t qw = __builtin_nontemporal_load(qw_ptr);
498 + int qs = __builtin_nontemporal_load(qs_ptr);
499 +
500 + *recv_buffer += kRankBufferTileStride;
501 +
502 + // Unpack q8 into fp16x8_t
503 + int32x4_t w;
504 + {
505 + static uint constexpr kMask00FF = 0x00FF00FF;
506 +
507 + // {1024.0, 1024.0}, fp16x2_t
508 + static uint constexpr kHalf2_1024 = 0x64006400;
509 +
510 + // {-1152.0, -1152.0}, fp16x2_t
511 + static uint constexpr kHalf2_1152 = 0xE480E480;
512 +
513 +#pragma unroll
514 + for (int i = 0; i < 4; i++) {
515 + if constexpr (std::is_same<T, half>::value) {
516 + int32_t q8 =
517 + ((qw[i / 2] >> ((i % 2) * 8)) & kMask00FF) | kHalf2_1024;
518 + w[i] = packed_add<half>(q8, kHalf2_1152);
519 + } else {
520 + int32_t int16_2 = (qw[i / 2] >> ((i % 2) * 8)) & kMask00FF;
521 + int16_t low = static_cast<int16_t>(int16_2 & 0xFFFF);
522 + int16_t high = static_cast<int16_t>((int16_2 >> 16) & 0xFFFF);
523 + nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
524 + nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
525 + nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
526 + int32_t packed_bf16 = *reinterpret_cast<int32_t*>(&bf2);
527 + w[i] = packed_add<nv_bfloat16>(packed_bf16, kRangeMin);
528 + }
529 + }
530 + }
531 +
532 + // Apply decoding scales
533 + for (int i = 0; i < 4; i++) {
534 + w[i] = packed_mul<T>(w[i], qs);
535 + }
536 +
537 + data[k] = w;
538 + }
539 + }
540 +};
541 +
542 +// Twoshot All Reduce
543 +template <typename T, class Codec, bool cast_bf2half>
544 +struct AllReduceTwoshot {
545 + static_assert(sizeof(T) == 2);
546 +
547 + static constexpr int kWorldSize = Codec::kWorldSize;
548 +
549 + __device__ static void run(
550 + T const* __restrict__ input, T* __restrict__ output,
551 + uint32_t const N, // number of elements
552 + int const block, // block index
553 + int const rank, // rank index
554 + uint8_t** __restrict__ buffer_list, // communication buffers
555 + uint32_t const data_offset, // offset to start of the data buffer
556 + uint32_t flag_color) {
557 + // Topology
558 + int thread = threadIdx.x + threadIdx.y * kWavefront;
559 + uint8_t* rank_buffer = buffer_list[rank];
560 + Codec codec(thread, rank);
561 + int block_id = blockIdx.x;
562 + int grid_size = gridDim.x;
563 + // --------------------------------------------------------
564 + // Read input into registers
565 + int32x4_t tA[kAtoms];
566 +
567 + BufferResource src_buffer(const_cast<T*>(input), N * sizeof(T));
568 + uint32_t src_offset = block * kTileSize + thread * sizeof(int32x4_t);
569 +
570 + for (int i = 0; i < kAtoms; i++) {
571 + tA[i] = buffer_load_dwordx4(src_buffer.descriptor, src_offset, 0, 0);
572 + src_offset += kAtomStride * sizeof(int32x4_t);
573 + if constexpr (cast_bf2half) {
574 + const nv_bfloat162* bf_buf =
575 + reinterpret_cast<const nv_bfloat162*>(&tA[i]);
576 + half2 half_buf[4];
577 +#pragma unroll
578 + for (int j = 0; j < 4; ++j) {
579 + float2 f = __bfloat1622float2(bf_buf[j]);
580 + half_buf[j] = __float22half2_rn(f);
581 + }
582 + tA[i] = *reinterpret_cast<const int32x4_t*>(half_buf);
583 + }
584 + }
585 +
586 + // --------------------------------------------------------
587 + // Phase-1A: Write segment data into the communication buffer of the target
588 + // rank responsible for this segment.
589 + uint32_t comm_data0_offset =
590 + data_offset + block_id * Codec::kTransmittedTileSize;
591 + uint32_t comm_data1_offset =
592 + grid_size * Codec::kTransmittedTileSize + comm_data0_offset;
593 +
594 + uint32_t comm_flags0_offset = block_id * (kWorldSize * sizeof(uint32_t));
595 + uint32_t comm_flags1_offset =
596 + grid_size * (kWorldSize * sizeof(uint32_t)) + comm_flags0_offset;
597 +
598 + for (int r = 0; r < kWorldSize; r++) {
599 + int32x4_t* send_buffer =
600 + reinterpret_cast<int32x4_t*>(buffer_list[r] + comm_data0_offset +
601 + rank * Codec::kRankTransmittedTileSize);
602 + codec.send(send_buffer, &tA[r * Codec::kRankAtoms]);
603 + }
604 +
605 + __syncthreads();
606 + if (thread < kWorldSize) {
607 + int r = thread;
608 + uint32_t* flag_ptr = reinterpret_cast<uint32_t*>(
609 + buffer_list[r] + comm_flags0_offset + rank * sizeof(uint32_t));
610 + set_sync_flag(flag_ptr, flag_color);
611 + }
612 + // --------------------------------------------------------
613 + // Phase-1B: Reduce the segment data from the communication buffers.
614 + int32x4_t tR[Codec::kRankAtoms] = {};
615 + {
616 + // Read the data from the communication buffer.
617 + int32x4_t* recv_buffer =
618 + reinterpret_cast<int32x4_t*>(rank_buffer + comm_data0_offset);
619 + uint32_t* flag_ptr =
620 + reinterpret_cast<uint32_t*>(rank_buffer + comm_flags0_offset);
621 +
622 + for (int r = 0; r < kWorldSize; r++) {
623 + // Wait for the flags to be set.
624 + if (thread == 0) {
625 + wait_sync_flag(&flag_ptr[r], flag_color);
626 + }
627 + __syncthreads();
628 +
629 + // note: we reuse tA as temp buffer here
630 + codec.recv(&recv_buffer, tA);
631 +
632 + for (int i = 0; i < Codec::kRankAtoms; i++) {
633 + packed_assign_add<T>(&tR[i], &tA[i]);
634 + }
635 + }
636 + }
637 +
638 + // Phase-2: Write the reduced segment to every other rank
639 + for (int r = 0; r < kWorldSize; r++) {
640 + int32x4_t* send_buffer =
641 + reinterpret_cast<int32x4_t*>(buffer_list[r] + comm_data1_offset +
642 + rank * Codec::kRankTransmittedTileSize);
643 + codec.send(send_buffer, tR);
644 + }
645 +
646 + __syncthreads();
647 + if (thread < kWorldSize) {
648 + int r = thread;
649 + uint32_t* flag_ptr = reinterpret_cast<uint32_t*>(
650 + buffer_list[r] + comm_flags1_offset + rank * sizeof(uint32_t));
651 + set_sync_flag(flag_ptr, flag_color);
652 + }
653 +
654 + // Phase-2: Read the gather segments from the rank's communication buffer.
655 + {
656 + // Read the data from the communication buffer.
657 + int32x4_t* recv_buffer =
658 + reinterpret_cast<int32x4_t*>(rank_buffer + comm_data1_offset);
659 + uint32_t* flag_ptr =
660 + reinterpret_cast<uint32_t*>(rank_buffer + comm_flags1_offset);
661 +
662 + for (int r = 0; r < kWorldSize; r++) {
663 + // Wait for the flags to be set.
664 + if (thread == 0) {
665 + wait_sync_flag(&flag_ptr[r], flag_color);
666 + }
667 + __syncthreads();
668 +
669 + // Gather all reduced and final rank segments into tA.
670 + codec.recv(&recv_buffer, &tA[r * Codec::kRankAtoms]);
671 + }
672 + }
673 +
674 + // --------------------------------------------------------
675 + // Write the result to output.
676 + BufferResource dst_buffer(output, N * sizeof(T));
677 + uint32_t dst_offset = block * kTileSize + thread * sizeof(int32x4_t);
678 +
679 + for (int i = 0; i < kAtoms; i++) {
680 + if constexpr (cast_bf2half) {
681 + const half2* half_buf = reinterpret_cast<const half2*>(&tA[i]);
682 + nv_bfloat162 bf16_buf[4];
683 +#pragma unroll
684 + for (int j = 0; j < 4; ++j) {
685 + float2 f = __half22float2(half_buf[j]);
686 + bf16_buf[j] = __float22bfloat162_rn(f);
687 + }
688 + buffer_store_dwordx4(*reinterpret_cast<const int32x4_t*>(bf16_buf),
689 + dst_buffer.descriptor, dst_offset, 0, 0);
690 + } else {
691 + buffer_store_dwordx4(tA[i], dst_buffer.descriptor, dst_offset, 0, 0);
692 + }
693 + dst_offset += kAtomStride * sizeof(int32x4_t);
694 + }
695 + }
696 +};
697 +
698 +} // namespace quickreduce
\ No newline at end of file
csrc/torch_bindings.cpp
+18
@@ -725,6 +725,24 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
725 custom_ar.impl("open_mem_handle", torch::kCPU, &open_mem_handle);
726
727 custom_ar.def("free_shared_buffer", &free_shared_buffer);
728 +#ifdef USE_ROCM
729 + // Quick Reduce all-reduce kernels
730 + custom_ar.def(
731 + "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool "
732 + "cast_bf2half) -> ()");
733 + custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce);
734 +
735 + custom_ar.def("init_custom_qr", &init_custom_qr);
736 + custom_ar.def("qr_destroy", &qr_destroy);
737 +
738 + custom_ar.def("qr_get_handle", &qr_get_handle);
739 +
740 + custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()");
741 + custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles);
742 +
743 + // Max input size in bytes
744 + custom_ar.def("qr_max_size", &qr_max_size);
745 +#endif
746 }
747
748 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
tests/distributed/test_quick_all_reduce.py new
+138
@@ -0,0 +1,138 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import random
5 +
6 +import pytest
7 +import ray
8 +import torch
9 +import torch.distributed as dist
10 +
11 +from vllm.distributed.communication_op import ( # noqa
12 + tensor_model_parallel_all_reduce)
13 +from vllm.distributed.parallel_state import (get_tensor_model_parallel_group,
14 + get_tp_group, graph_capture)
15 +from vllm.platforms import current_platform
16 +
17 +from ..utils import (ensure_model_parallel_initialized,
18 + init_test_distributed_environment, multi_process_parallel)
19 +
20 +torch.manual_seed(42)
21 +random.seed(44)
22 +# Size over 8MB is sufficient for custom quick allreduce.
23 +test_sizes = [
24 + random.randint(8 * 1024 * 1024, 10 * 1024 * 1024) for _ in range(8)
25 +]
26 +for i, v in enumerate(test_sizes):
27 + test_sizes[i] -= v % 8
28 +
29 +
30 +@ray.remote(num_gpus=1, max_calls=1)
31 +def graph_quickreduce(
32 + monkeypatch: pytest.MonkeyPatch,
33 + tp_size,
34 + pp_size,
35 + rank,
36 + distributed_init_port,
37 +):
38 + with monkeypatch.context() as m:
39 + m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
40 + device = torch.device(f"cuda:{rank}")
41 + torch.cuda.set_device(device)
42 + init_test_distributed_environment(tp_size, pp_size, rank,
43 + distributed_init_port)
44 + ensure_model_parallel_initialized(tp_size, pp_size)
45 + group = get_tensor_model_parallel_group().device_group
46 +
47 + # A small all_reduce for warmup.
48 + # this is needed because device communicators might be created lazily
49 + # (e.g. NCCL). This will ensure that the communicator is initialized
50 + # before any communication happens, so that this group can be used for
51 + # graph capture immediately.
52 + data = torch.zeros(1)
53 + data = data.to(device=device)
54 + torch.distributed.all_reduce(data, group=group)
55 + torch.cuda.synchronize()
56 + del data
57 +
58 + # we use the first group to communicate once
59 + # and the second group to communicate twice
60 + # and so on
61 + # this is used to demonstrate that each group can
62 + # communicate independently
63 + num_communication = rank // tp_size + 1
64 +
65 + for sz in test_sizes:
66 + for dtype in [torch.float16, torch.bfloat16]:
67 + with graph_capture(device=device) as graph_capture_context:
68 + inp1 = torch.randint(1,
69 + 23, (sz, ),
70 + dtype=dtype,
71 + device=torch.cuda.current_device())
72 + inp2 = torch.randint(-23,
73 + 1, (sz, ),
74 + dtype=dtype,
75 + device=torch.cuda.current_device())
76 + torch.cuda.synchronize()
77 + graph = torch.cuda.CUDAGraph()
78 + with torch.cuda.graph(graph,
79 + stream=graph_capture_context.stream):
80 + for _ in range(num_communication):
81 + out1 = tensor_model_parallel_all_reduce(inp1)
82 + dist.all_reduce(inp1, group=group)
83 + out2 = tensor_model_parallel_all_reduce(inp2)
84 + dist.all_reduce(inp2, group=group)
85 + graph.replay()
86 + torch.testing.assert_close(out1, inp1, atol=2.5, rtol=0.1)
87 + torch.testing.assert_close(out2, inp2, atol=2.5, rtol=0.1)
88 +
89 +
90 +@ray.remote(num_gpus=1, max_calls=1)
91 +def eager_quickreduce(
92 + monkeypatch: pytest.MonkeyPatch,
93 + tp_size,
94 + pp_size,
95 + rank,
96 + distributed_init_port,
97 +):
98 + with monkeypatch.context() as m:
99 + m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
100 + device = torch.device(f"cuda:{rank}")
101 + torch.cuda.set_device(device)
102 +
103 + init_test_distributed_environment(tp_size, pp_size, rank,
104 + distributed_init_port)
105 +
106 + # Size over 8MB is sufficient for custom quick allreduce.
107 + sz = 16 * 1024 * 1024
108 + fa = get_tp_group().device_communicator.qr_comm
109 + inp = torch.tensor([1.0 * ((i) % 23) for i in range(sz)],
110 + dtype=torch.float16,
111 + device=device)
112 + out = fa.quick_all_reduce(inp)
113 + torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1)
114 +
115 + inp = torch.tensor([1.0 * ((i) % 23) for i in range(sz)],
116 + dtype=torch.bfloat16,
117 + device=device)
118 + out = fa.quick_all_reduce(inp)
119 + torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1)
120 +
121 +
122 +@pytest.mark.skipif(not current_platform.is_rocm(),
123 + reason="only test quick allreduce for rocm")
124 +@pytest.mark.parametrize("quant_mode", ["FP", "INT8", "INT6", "INT4"])
125 +@pytest.mark.parametrize("tp_size", [2])
126 +@pytest.mark.parametrize("pipeline_parallel_size", [1, 2])
127 +@pytest.mark.parametrize("test_target", [graph_quickreduce, eager_quickreduce])
128 +def test_custom_quick_allreduce(monkeypatch: pytest.MonkeyPatch, tp_size,
129 + pipeline_parallel_size, test_target,
130 + quant_mode):
131 + world_size = tp_size * pipeline_parallel_size
132 + if world_size > torch.cuda.device_count():
133 + pytest.skip("Not enough GPUs to run the test.")
134 +
135 + monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_mode)
136 +
137 + multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size,
138 + test_target)
vllm/_custom_ops.py
+32
@@ -1748,6 +1748,38 @@ def free_shared_buffer(ptr: int) -> None:
1748 torch.ops._C_custom_ar.free_shared_buffer(ptr)
1749
1750
1751 +# quick all reduce
1752 +def init_custom_qr(rank: int,
1753 + world_size: int,
1754 + qr_max_size: Optional[int] = None) -> int:
1755 + return torch.ops._C_custom_ar.init_custom_qr(rank, world_size, qr_max_size)
1756 +
1757 +
1758 +def qr_destroy(fa: int) -> None:
1759 + torch.ops._C_custom_ar.qr_destroy(fa)
1760 +
1761 +
1762 +def qr_all_reduce(fa: int,
1763 + inp: torch.Tensor,
1764 + out: torch.Tensor,
1765 + quant_level: int,
1766 + cast_bf2half: bool = False) -> None:
1767 + torch.ops._C_custom_ar.qr_all_reduce(fa, inp, out, quant_level,
1768 + cast_bf2half)
1769 +
1770 +
1771 +def qr_get_handle(fa: int) -> torch.Tensor:
1772 + return torch.ops._C_custom_ar.qr_get_handle(fa)
1773 +
1774 +
1775 +def qr_open_handles(fa: int, handles: list[torch.Tensor]) -> None:
1776 + return torch.ops._C_custom_ar.qr_open_handles(fa, handles)
1777 +
1778 +
1779 +def qr_max_size() -> int:
1780 + return torch.ops._C_custom_ar.qr_max_size()
1781 +
1782 +
1783 def get_flash_mla_metadata(
1784 cache_seqlens: torch.Tensor,
1785 num_heads_per_head_k: int,
vllm/distributed/device_communicators/cuda_communicator.py
+20 -2
@@ -8,6 +8,7 @@ from torch.distributed import ProcessGroup
8
9 import vllm.envs as envs
10 from vllm.logger import init_logger
11 +from vllm.platforms import current_platform
12
13 from .base_device_communicator import DeviceCommunicatorBase
14
@@ -41,6 +42,8 @@ class CudaCommunicator(DeviceCommunicatorBase):
42 CustomAllreduce)
43 from vllm.distributed.device_communicators.pynccl import (
44 PyNcclCommunicator)
45 + from vllm.distributed.device_communicators.quick_all_reduce import (
46 + QuickAllReduce)
47
48 self.pynccl_comm: Optional[PyNcclCommunicator] = None
49 if use_pynccl and self.world_size > 1:
@@ -50,6 +53,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
53 )
54
55 self.ca_comm: Optional[CustomAllreduce] = None
56 + self.qr_comm: Optional[QuickAllReduce] = None
57 if use_custom_allreduce and self.world_size > 1:
58 # Initialize a custom fast all-reduce implementation.
59 self.ca_comm = CustomAllreduce(
@@ -57,6 +61,14 @@ class CudaCommunicator(DeviceCommunicatorBase):
61 device=self.device,
62 )
63
64 + if current_platform.is_rocm():
65 + # Initialize a custom quick all-reduce implementation for AMD.
66 + # Quick reduce is designed as a complement to custom allreduce.
67 + # Based on quickreduce (https://github.com/mk1-project/quickreduce).
68 + # If it's a rocm, 'use_custom_allreduce==True' means it must
69 + # currently be an MI300 series.
70 + self.qr_comm = QuickAllReduce(group=self.cpu_group,
71 + device=self.device)
72 if self.use_all2all:
73 all2all_backend = envs.VLLM_ALL2ALL_BACKEND
74 if all2all_backend == "naive":
@@ -79,8 +91,14 @@ class CudaCommunicator(DeviceCommunicatorBase):
91 raise ValueError(f"Unknown all2all backend: {all2all_backend}")
92
93 def all_reduce(self, input_):
82 - # always try custom allreduce first,
83 - # and then pynccl.
94 + # always try quick reduce first, then custom allreduce,
95 + # and then pynccl. (quick reduce just for ROCM MI3*)
96 + qr_comm = self.qr_comm
97 + if qr_comm is not None and not qr_comm.disabled and \
98 + qr_comm.should_quick_allreduce(input_):
99 + out = qr_comm.quick_all_reduce(input_)
100 + assert out is not None
101 + return out
102 ca_comm = self.ca_comm
103 if ca_comm is not None and not ca_comm.disabled and \
104 ca_comm.should_custom_ar(input_):
vllm/distributed/device_communicators/quick_all_reduce.py new
+278
@@ -0,0 +1,278 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from enum import Enum
5 +from typing import Union
6 +
7 +import torch
8 +import torch.distributed as dist
9 +from torch.distributed import ProcessGroup
10 +
11 +import vllm.envs as envs
12 +from vllm import _custom_ops as ops
13 +from vllm.config import get_current_vllm_config
14 +from vllm.distributed.parallel_state import in_the_same_node_as
15 +from vllm.logger import init_logger
16 +from vllm.platforms import current_platform
17 +from vllm.utils import cuda_device_count_stateless
18 +
19 +logger = init_logger(__name__)
20 +
21 +try:
22 + ops.qr_max_size()
23 + quick_ar = True
24 +except Exception:
25 + # For CPUs and CUDA
26 + quick_ar = False
27 +
28 +
29 +def is_weak_contiguous(inp: torch.Tensor):
30 + return inp.is_contiguous() or (inp.storage().nbytes() -
31 + inp.storage_offset() * inp.element_size()
32 + == inp.numel() * inp.element_size())
33 +
34 +
35 +class QuickReduceRegime(Enum):
36 + FP = 0
37 + INT8 = 1
38 + INT6 = 2
39 + INT4 = 3
40 + NONE = 4
41 +
42 +
43 +MB = 1024 * 1024
44 +
45 +
46 +class QuickAllReduce:
47 +
48 + _SUPPORTED_WORLD_SIZES = [2, 4, 8]
49 + _SUPPORTED_DTYPES = [torch.float16, torch.bfloat16]
50 + # The following data is based on kernel tests.
51 + # In this order [FP, INT8, INT6, INT4].
52 + _QR_MIN_SIZE = {
53 + (torch.float16, 2): [1 * MB, 2 * MB, 2 * MB, 1 * MB],
54 + (torch.float16, 4): [1 * MB, 16 * MB, 4 * MB, 2 * MB],
55 + (torch.float16, 8): [16 * MB, 4 * MB, 4 * MB, 2 * MB],
56 + (torch.bfloat16, 2): [2 * MB, 8 * MB, 8 * MB, 8 * MB],
57 + (torch.bfloat16, 4): [8 * MB, 64 * MB, 64 * MB, 16 * MB],
58 + (torch.bfloat16, 8): [16 * MB, 2048 * MB, 2048 * MB, 2048 * MB],
59 + }
60 +
61 + def __init__(self, group: ProcessGroup,
62 + device: Union[int, str, torch.device]) -> None:
63 + """
64 + Custom allreduce provides non-destructive acceleration and is
65 + available for CUDA and ROCm MI300 series.
66 +
67 + Custom quick allreduce leverages quantization for further
68 + acceleration on ROCm. It currently supports Q8, Q6, and Q4
69 + quantization formats and FP(float16, bfloat16).
70 +
71 + Quick allreduce is designed as a complement to custom allreduce.
72 + Its initialization requires even stricter conditions.
73 +
74 + Only the ROCm MI300 series is supported for quick allreduce at
75 + this time.
76 +
77 + Args:
78 + group: the process group to work on. If None, it will use the
79 + default process group.
80 + device: the device to bind the CustomAllreduce to. If None,
81 + it will be bind to f"cuda:{local_rank}".
82 + It is the caller's responsibility to make sure each communicator
83 + is bind to a unique device, and all communicators in this group
84 + are in the same node.
85 + """
86 + self.disabled = True
87 + if not self._rocm_arch_available():
88 + logger.debug(
89 + "Custom quick allreduce is only supported on ROCm MI300 series."
90 + )
91 + return
92 +
93 + if not quick_ar:
94 + # disable because of missing quick reduce library
95 + # e.g. in a cuda environment
96 + logger.info("Custom quick allreduce is disabled because "
97 + "of missing custom quick allreduce library")
98 + return
99 +
100 + self.group = group
101 + assert dist.get_backend(group) != dist.Backend.NCCL, (
102 + "Custom quick allreduce should be attached to a non-NCCL group.")
103 + if not all(in_the_same_node_as(group, source_rank=0)):
104 + # No need to initialize custom quick allreduce for
105 + # multi-node case.
106 + logger.warning("Custom quick allreduce is disabled because this "
107 + "process group spans across nodes.")
108 + return
109 + rank = dist.get_rank(group=self.group)
110 + world_size = dist.get_world_size(group=self.group)
111 + self.rank = rank
112 + self.world_size = world_size
113 + if world_size == 1:
114 + # No need to initialize QuickReduce for single GPU case.
115 + return
116 +
117 + if world_size not in QuickAllReduce._SUPPORTED_WORLD_SIZES:
118 + logger.warning(
119 + "Custom quick allreduce is disabled due to an "
120 + "unsupported world size: %d. Supported world sizes: %s.",
121 + world_size, str(QuickAllReduce._SUPPORTED_WORLD_SIZES))
122 + return
123 +
124 + if isinstance(device, int):
125 + device = torch.device(f"cuda:{device}")
126 + elif isinstance(device, str):
127 + device = torch.device(device)
128 + assert isinstance(device, torch.device)
129 + self.device = device
130 +
131 + cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
132 + if cuda_visible_devices:
133 + device_ids = list(map(int, cuda_visible_devices.split(",")))
134 + else:
135 + device_ids = list(range(cuda_device_count_stateless()))
136 + physical_device_id = device_ids[device.index]
137 + tensor = torch.tensor([physical_device_id],
138 + dtype=torch.int,
139 + device="cpu")
140 + gather_list = [
141 + torch.tensor([0], dtype=torch.int, device="cpu")
142 + for _ in range(self.world_size)
143 + ]
144 + dist.all_gather(gather_list, tensor, group=self.group)
145 + physical_device_ids = [t.item() for t in gather_list]
146 +
147 + # test nvlink first, this will filter out most of the cases
148 + # where custom quick allreduce is not supported
149 + # this checks hardware and driver support for NVLink
150 + assert current_platform.is_cuda_alike()
151 + self.fully_connected = current_platform.is_fully_connected(
152 + physical_device_ids)
153 + if self.world_size > 2 and not self.fully_connected:
154 + logger.debug(
155 + "Custom quick allreduce is disabled because it's not supported "
156 + "on more than two PCIe-only GPUs. ")
157 + return
158 +
159 + self.init_quick_all_reduce()
160 +
161 + def init_quick_all_reduce(self):
162 + # On RocM, bfloat16 kernels are slower than fp16
163 + # due to slower match operations
164 + # If environment variable is set to 1, we convert input to fp16
165 + self.use_fp16_kernels = envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16
166 + regime_str = envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION
167 + if regime_str not in QuickReduceRegime.__members__:
168 + logger.warning(
169 + "Custom quick allreduce:",
170 + f"Invalid quantization level: {regime_str}. "
171 + "Supported levels: "
172 + f"{list(QuickReduceRegime.__members__.keys())}")
173 + return
174 +
175 + if regime_str == "NONE":
176 + logger.debug("Custom quick allreduce is disabled based "
177 + "on env variable "
178 + "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION='NONE'")
179 + return
180 + self.qr_quant_level = QuickReduceRegime[regime_str]
181 + vllm_config = get_current_vllm_config()
182 + if vllm_config is not None and \
183 + hasattr(vllm_config, "model_config") and \
184 + hasattr(vllm_config.model_config, "dtype"):
185 + dtype = vllm_config.model_config.dtype
186 + if dtype not in [torch.float16, torch.bfloat16]:
187 + logger.debug(
188 + "Custom quick allreduce disabled: only supports "
189 + "float16 and float16, but get %s.", dtype)
190 + return
191 +
192 + if dtype == torch.bfloat16 and self.use_fp16_kernels:
193 + logger.info(
194 + "Custom quick allreduce: BF16 inputs will be converted "
195 + "to FP16 to improve performance. set "
196 + "envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16=0 "
197 + "to turn off.")
198 +
199 + # VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB is specified in MB
200 + qr_max_size = envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB
201 + if qr_max_size is not None:
202 + if qr_max_size < 1:
203 + logger.info(
204 + "You should not set a max_size smaller than 1MB, which can "
205 + "lead to error or degradation to custom allreduce or rccl."
206 + )
207 + qr_max_size = qr_max_size * MB
208 + self._ptr = ops.init_custom_qr(self.rank, self.world_size, qr_max_size)
209 + self.qr_max_size = qr_max_size if qr_max_size is not None \
210 + else ops.qr_max_size()
211 + self.create_shared_buffer()
212 + self.disabled = False
213 +
214 + def _rocm_arch_available(self):
215 + if not current_platform.is_rocm():
216 + return False
217 + try:
218 + props = torch.cuda.get_device_properties(0)
219 + gcn_arch = getattr(props, "gcnArchName", "")
220 + supported_archs = ['gfx94', 'gfx95']
221 + return any(gfx in gcn_arch for gfx in supported_archs)
222 + except Exception as e:
223 + logger.warning("Failed to determine ROCm for quick allreduce: %s",
224 + e)
225 + return False
226 +
227 + def create_shared_buffer(self):
228 + """
229 + Creates a shared buffer for quickreduce.
230 + Has to be called after init_custom_qr
231 + """
232 + handle = ops.qr_get_handle(self._ptr)
233 + world_size = dist.get_world_size(group=self.group)
234 + handles = [None] * world_size
235 + dist.all_gather_object(handles, handle, group=self.group)
236 + ops.qr_open_handles(self._ptr, handles)
237 +
238 + def should_quick_allreduce(self, inp: torch.Tensor):
239 + """
240 + Check if quickreduce is available
241 + """
242 + if self.disabled:
243 + return False
244 + if inp.dtype not in self._SUPPORTED_DTYPES:
245 + return False
246 + inp_size = inp.numel() * inp.element_size()
247 + # custom quick allreduce requires input byte size to be
248 + # multiples of 16
249 + if inp_size % 16 != 0:
250 + return False
251 + if not is_weak_contiguous(inp):
252 + return False
253 + dtype = inp.dtype
254 + if self.use_fp16_kernels:
255 + dtype = torch.float16
256 + return inp_size <= self.qr_max_size and \
257 + inp_size >= self._QR_MIN_SIZE[(dtype, self.world_size)]\
258 + [self.qr_quant_level.value]
259 +
260 + def quick_all_reduce(self, inp: torch.Tensor, *, out: torch.Tensor = None):
261 + """Performs an out-of-place custom quick all reduce."""
262 + # quick allreduce doesn't require a separate graph mode,
263 + # as QR uses static IPC buffer.
264 + if out is None:
265 + out = torch.empty_like(inp)
266 + ops.qr_all_reduce(self._ptr, inp, out, self.qr_quant_level.value,
267 + self.use_fp16_kernels)
268 + return out
269 +
270 + def close(self):
271 + if not self.disabled and getattr(self, "_ptr", None):
272 + if ops is not None:
273 + ops.qr_destroy(self._ptr)
274 + self._ptr = 0
275 + self.disabled = True
276 +
277 + def __del__(self):
278 + self.close()
vllm/envs.py
+28
@@ -135,6 +135,9 @@ if TYPE_CHECKING:
135 VLLM_KV_CACHE_LAYOUT: Optional[str] = None
136 VLLM_COMPUTE_NANS_IN_LOGITS: bool = False
137 VLLM_USE_NVFP4_CT_EMULATIONS: bool = False
138 + VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: str = "NONE"
139 + VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True
140 + VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: Optional[int] = None
141
142
143 def get_default_cache_root():
@@ -690,6 +693,31 @@ environment_variables: dict[str, Callable[[], Any]] = {
693 lambda: (os.getenv("VLLM_ROCM_CUSTOM_PAGED_ATTN", "True").lower() in
694 ("true", "1")),
695
696 + # Custom quick allreduce kernel for MI3* cards
697 + # Choice of quantization level: FP, INT8, INT6, INT4 or NONE
698 + # Recommended for large models to get allreduce
699 + "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION":
700 + lambda: os.getenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "NONE").upper(),
701 +
702 + # Custom quick allreduce kernel for MI3* cards
703 + # Due to the lack of the bfloat16 asm instruction, bfloat16
704 + # kernels are slower than fp16,
705 + # If environment variable is set to 1, the input is converted to fp16
706 + "VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16":
707 + lambda:
708 + (os.getenv("VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", "True").lower() in
709 + ("true", "1")),
710 +
711 + # Custom quick allreduce kernel for MI3* cards.
712 + # Controls the maximum allowed number of data bytes(MB) for custom quick
713 + # allreduce communication.
714 + # Default: 2048 MB.
715 + # Data exceeding this size will use either custom allreduce or RCCL
716 + # communication.
717 + "VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB":
718 + lambda: maybe_convert_int(
719 + os.environ.get("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", None)),
720 +
721 # If set, when running in Quark emulation mode, do not dequantize the
722 # weights at load time. Instead, dequantize weights on-the-fly during
723 # kernel execution.