@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 241 lines 9.88 KB
Raw
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* d_flag_counters,
26 int64_t data_size_per_phase) {
27 int block = blockIdx.x;
28 int grid = gridDim.x;
29
30 // Load this block's counter from device memory and advance it on-device,
31 // so the color keeps changing across graph replays instead of being frozen.
32 uint32_t flag_color = d_flag_counters[blockIdx.x];
33
34 while (block < num_blocks) {
35 AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset,
36 flag_color, data_size_per_phase);
37 block += grid;
38 flag_color++;
39 }
40 // All threads compute the same final value; one writer per block is enough.
41 if (threadIdx.x == 0 && threadIdx.y == 0) {
42 d_flag_counters[blockIdx.x] = flag_color;
43 }
44 }
45
46 #define TWOSHOT_DISPATCH(__codec) \
47 if (world_size == 2) { \
48 using LineCodec = __codec<T, 2>; \
49 using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
50 hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
51 dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
52 num_blocks, rank, dbuffer_list, data_offset, \
53 d_flag_counters, this->kMaxProblemSize); \
54 } else if (world_size == 4) { \
55 using LineCodec = __codec<T, 4>; \
56 using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
57 hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
58 dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
59 num_blocks, rank, dbuffer_list, data_offset, \
60 d_flag_counters, this->kMaxProblemSize); \
61 } else if (world_size == 8) { \
62 using LineCodec = __codec<T, 8>; \
63 using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
64 hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
65 dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
66 num_blocks, rank, dbuffer_list, data_offset, \
67 d_flag_counters, this->kMaxProblemSize); \
68 }
69
70 // INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8
71 // the 3-bit codec's pack/unpack overhead outweighs the reduced communication
72 // volume, so INT3 is restricted to a TP2-only dispatch here.
73 #define TWOSHOT_DISPATCH_TP2_ONLY(__codec) \
74 if (world_size == 2) { \
75 using LineCodec = __codec<T, 2>; \
76 using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>; \
77 hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>), \
78 dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
79 num_blocks, rank, dbuffer_list, data_offset, \
80 d_flag_counters, this->kMaxProblemSize); \
81 } else { \
82 throw std::runtime_error( \
83 "INT3 quick all-reduce is only supported for world_size == 2 " \
84 "(TP2); use INT4/NONE for larger world sizes."); \
85 }
86
87 enum QuickReduceQuantLevel {
88 // Keep these ids in sync with Python QuickReduceRegime enum.
89 F16 = 0, // full-precision fp16/bf16 communication
90 INT8 = 1, // symmetric int8 + per-block scale
91 INT6 = 2, // symmetric int6 + per-block scale
92 INT4 = 3, // symmetric int4 + per-block scale
93 INT3 = 4, // symmetric int3 + per-block scale (TP2 only)
94 };
95
96 struct DeviceComms {
97 // Max problem size is 2GB (in bytes) or half of uint32_t max value.
98 int64_t kMaxProblemSize =
99 static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1;
100
101 // Max TP-8
102 static int constexpr kMaxWorldSize = 8;
103
104 bool initialized = false;
105 uint32_t* d_flag_counters = nullptr;
106 int world_size;
107 int rank;
108
109 uint8_t* dbuffer;
110 uint8_t** dbuffer_list;
111 hipIpcMemHandle_t buffer_ipc_handle;
112 std::vector<hipIpcMemHandle_t> all_buffer_ipc_handles;
113 std::vector<uint8_t*> buffer_list;
114 uint32_t data_offset;
115
116 DeviceComms() : initialized(false), world_size(1), rank(0) {}
117 ~DeviceComms() { destroy(); }
118
119 void init(int world_size, int rank,
120 std::optional<int64_t> max_problem_size = std::nullopt) {
121 destroy();
122 this->world_size = world_size;
123 this->rank = rank;
124 if (max_problem_size.has_value() && max_problem_size.value() > 0) {
125 this->kMaxProblemSize = max_problem_size.value();
126 }
127 // Allocate buffer size for worst case: F16 2-stage buffer.
128 uint32_t flags_buffer_size =
129 2 * world_size * kMaxNumBlocks * sizeof(uint32_t);
130 static int64_t data_buffer_size = 2 * this->kMaxProblemSize;
131 int64_t total_buffer_size = flags_buffer_size + data_buffer_size;
132 data_offset = flags_buffer_size;
133 HIP_CHECK(hipExtMallocWithFlags((void**)&dbuffer, total_buffer_size,
134 hipDeviceMallocUncached));
135
136 // Clear the flags buffer.
137 HIP_CHECK(hipMemset(dbuffer, 0, flags_buffer_size));
138
139 // One flag-color counter per block, advanced by the kernel. Start at 1
140 // to stay clear of the flags buffer we just zeroed.
141 HIP_CHECK(hipMalloc(&d_flag_counters, kMaxNumBlocks * sizeof(uint32_t)));
142 {
143 std::vector<uint32_t> init_color(kMaxNumBlocks, 1u);
144 HIP_CHECK(hipMemcpy(d_flag_counters, init_color.data(),
145 kMaxNumBlocks * sizeof(uint32_t),
146 hipMemcpyHostToDevice));
147 }
148
149 // Device-side list of IPC buffers.
150 buffer_list.resize(world_size);
151 HIP_CHECK(hipMalloc(&dbuffer_list, world_size * sizeof(uint8_t*)));
152
153 // Create IPC handles for rank's communication buffer.
154 all_buffer_ipc_handles.resize(world_size);
155 HIP_CHECK(hipIpcGetMemHandle(&buffer_ipc_handle, dbuffer));
156
157 initialized = true;
158 }
159 int get_world_size() { return world_size; }
160 int get_rank() { return rank; }
161 bool status() { return initialized; }
162 hipIpcMemHandle_t const get_handle() { return buffer_ipc_handle; }
163
164 void destroy() {
165 // Allocated before `initialized` flips true, so free it on its own guard
166 // to avoid a leak if init fails partway through.
167 if (d_flag_counters) {
168 HIP_CHECK(hipFree(d_flag_counters));
169 d_flag_counters = nullptr;
170 }
171 if (initialized) {
172 for (int i = 0; i < world_size; i++) {
173 if (i != rank) {
174 HIP_CHECK(hipIpcCloseMemHandle(dbuffer_list[i]));
175 }
176 }
177
178 HIP_CHECK(hipFree(dbuffer));
179 HIP_CHECK(hipFree(dbuffer_list));
180
181 initialized = false;
182 }
183 }
184
185 void open_ipc_handles(std::vector<hipIpcMemHandle_t> const& ipc_handles) {
186 assert(ipc_handles.size() == all_buffer_ipc_handles.size());
187 for (int i = 0; i < world_size; i++) {
188 all_buffer_ipc_handles[i] = ipc_handles[i];
189 }
190
191 // Open device memory access to the IPC communication buffers.
192 // Note: For our own rank, we do not need to open a handle.
193 for (int i = 0; i < world_size; i++) {
194 if (i != rank) {
195 HIP_CHECK(hipIpcOpenMemHandle((void**)&buffer_list[i],
196 all_buffer_ipc_handles[i],
197 hipIpcMemLazyEnablePeerAccess));
198 } else {
199 buffer_list[i] = dbuffer;
200 }
201 }
202
203 HIP_CHECK(hipMemcpy(dbuffer_list, buffer_list.data(),
204 world_size * sizeof(uint8_t*), hipMemcpyHostToDevice));
205 }
206
207 template <typename T, bool cast_bf2half>
208 void allreduce(T const* A, T* B, uint32_t N, int quant_level,
209 hipStream_t stream) {
210 if (world_size != 2 && world_size != 4 && world_size != 8) {
211 throw std::runtime_error("All Reduce not supported for world_size = " +
212 std::to_string(world_size));
213 }
214
215 // Configuration.
216 uint32_t msg_size = N * sizeof(T);
217 uint32_t num_blocks = divceil(msg_size, kTileSize);
218 uint32_t grid = min(kMaxNumBlocks, num_blocks);
219 auto quant_level_ = static_cast<QuickReduceQuantLevel>(quant_level);
220 switch (quant_level_) {
221 case QuickReduceQuantLevel::INT8:
222 TWOSHOT_DISPATCH(CodecQ8)
223 break;
224 case QuickReduceQuantLevel::INT6:
225 TWOSHOT_DISPATCH(CodecQ6)
226 break;
227 case QuickReduceQuantLevel::INT4:
228 TWOSHOT_DISPATCH(CodecQ4)
229 break;
230 case QuickReduceQuantLevel::INT3:
231 TWOSHOT_DISPATCH_TP2_ONLY(CodecQ3)
232 break;
233 default:
234 TWOSHOT_DISPATCH(CodecFP)
235 break;
236 }
237 HIP_CHECK(cudaGetLastError());
238 }
239 };
240
241 } // namespace quickreduce