| 1 | #include "torch_utils.h" |
| 2 | |
| 3 | #include <torch/csrc/stable/macros.h> |
| 4 | #include <torch/csrc/stable/accelerator.h> |
| 5 | #include <torch/csrc/stable/tensor.h> |
| 6 | #include <torch/headeronly/core/ScalarType.h> |
| 7 | |
| 8 | #include "custom_all_reduce.cuh" |
| 9 | #include "custom_all_gather_reduce_scatter.cuh" |
| 10 | |
| 11 | namespace vllm { |
| 12 | |
| 13 | void CustomAllreduce::allgather(cudaStream_t stream, void* input, void* output, |
| 14 | int size_bytes, int threads, int block_limit) { |
| 15 | if (size_bytes % sizeof(CopyPack) != 0) |
| 16 | throw std::runtime_error( |
| 17 | "custom allgather requires input byte size to be a multiple of " + |
| 18 | std::to_string(sizeof(CopyPack))); |
| 19 | |
| 20 | auto ptrs = buffers_.at(input); |
| 21 | int size_per_rank = size_bytes / sizeof(CopyPack); |
| 22 | int total_size = size_per_rank * world_size_; |
| 23 | int blocks = std::min(block_limit, (total_size + threads - 1) / threads); |
| 24 | |
| 25 | #define AG_CASE(ngpus) \ |
| 26 | case ngpus: \ |
| 27 | cross_device_all_gather<ngpus><<<blocks, threads, 0, stream>>>( \ |
| 28 | ptrs, sg_, self_sg_, reinterpret_cast<CopyPack*>(output), rank_, \ |
| 29 | size_per_rank); \ |
| 30 | break; |
| 31 | |
| 32 | switch (world_size_) { |
| 33 | AG_CASE(2) |
| 34 | AG_CASE(4) |
| 35 | AG_CASE(6) |
| 36 | AG_CASE(8) |
| 37 | default: |
| 38 | throw std::runtime_error( |
| 39 | "custom allgather only supports num gpus in (2,4,6,8)"); |
| 40 | } |
| 41 | #undef AG_CASE |
| 42 | } |
| 43 | |
| 44 | template <typename T> |
| 45 | void CustomAllreduce::mnnvl_lamport_allgather(cudaStream_t stream, T* input, |
| 46 | T* output, void* local_buffer, |
| 47 | void* multicast_buffer, |
| 48 | uint32_t* epochs, int size_bytes, |
| 49 | int stage_size_bytes) { |
| 50 | if (size_bytes % sizeof(typename packed_t<T>::P) != 0 || |
| 51 | stage_size_bytes % sizeof(typename packed_t<T>::P) != 0) |
| 52 | throw std::runtime_error( |
| 53 | "MNNVL Lamport allgather requires 16-byte aligned sizes"); |
| 54 | |
| 55 | auto ptrs = buffers_.at(local_buffer); |
| 56 | int size_per_rank = size_bytes / sizeof(typename packed_t<T>::P); |
| 57 | int stage_size = stage_size_bytes / sizeof(typename packed_t<T>::P); |
| 58 | int blocks = |
| 59 | (size_per_rank + kMnnvlLamportAgThreads - 1) / kMnnvlLamportAgThreads; |
| 60 | |
| 61 | #if !defined(USE_ROCM) && CUDA_VERSION >= 12000 |
| 62 | cudaLaunchAttribute attributes[1]{}; |
| 63 | attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; |
| 64 | attributes[0].val.programmaticStreamSerializationAllowed = 1; |
| 65 | cudaLaunchConfig_t config{.gridDim = dim3(blocks), |
| 66 | .blockDim = dim3(kMnnvlLamportAgThreads), |
| 67 | .dynamicSmemBytes = 0, |
| 68 | .stream = stream, |
| 69 | .attrs = attributes, |
| 70 | .numAttrs = 1}; |
| 71 | #define MNNVL_LAMPORT_AG_LAUNCH(ngpus) \ |
| 72 | CUDACHECK(cudaLaunchKernelEx(&config, &mnnvl_lamport_all_gather<T, ngpus>, \ |
| 73 | ptrs, input, output, \ |
| 74 | reinterpret_cast<T*>(multicast_buffer), \ |
| 75 | epochs, rank_, size_per_rank, stage_size)) |
| 76 | #else |
| 77 | #define MNNVL_LAMPORT_AG_LAUNCH(ngpus) \ |
| 78 | mnnvl_lamport_all_gather<T, ngpus> \ |
| 79 | <<<blocks, kMnnvlLamportAgThreads, 0, stream>>>( \ |
| 80 | ptrs, input, output, reinterpret_cast<T*>(multicast_buffer), \ |
| 81 | epochs, rank_, size_per_rank, stage_size) |
| 82 | #endif |
| 83 | |
| 84 | #define MNNVL_LAMPORT_AG_CASE(ngpus) \ |
| 85 | case ngpus: \ |
| 86 | MNNVL_LAMPORT_AG_LAUNCH(ngpus); \ |
| 87 | break; |
| 88 | |
| 89 | switch (world_size_) { |
| 90 | MNNVL_LAMPORT_AG_CASE(2) |
| 91 | MNNVL_LAMPORT_AG_CASE(4) |
| 92 | MNNVL_LAMPORT_AG_CASE(6) |
| 93 | MNNVL_LAMPORT_AG_CASE(8) |
| 94 | MNNVL_LAMPORT_AG_CASE(16) |
| 95 | default: |
| 96 | throw std::runtime_error( |
| 97 | "MNNVL Lamport allgather only supports num gpus in (2,4,6,8,16)"); |
| 98 | } |
| 99 | #undef MNNVL_LAMPORT_AG_CASE |
| 100 | #undef MNNVL_LAMPORT_AG_LAUNCH |
| 101 | } |
| 102 | |
| 103 | template <typename T> |
| 104 | void CustomAllreduce::reduce_scatter(cudaStream_t stream, T* input, T* output, |
| 105 | int size, int threads, int block_limit) { |
| 106 | auto packed_size = packed_t<T>::P::size; |
| 107 | if (size % (packed_size * world_size_) != 0) |
| 108 | throw std::runtime_error( |
| 109 | "custom reduce-scatter requires each output shard byte size to be " |
| 110 | "a multiple of 16"); |
| 111 | |
| 112 | auto ptrs = buffers_.at(input); |
| 113 | int size_per_rank = size / packed_size / world_size_; |
| 114 | int blocks = std::min(block_limit, (size_per_rank + threads - 1) / threads); |
| 115 | |
| 116 | #define RS_CASE(ngpus) \ |
| 117 | case ngpus: \ |
| 118 | cross_device_reduce_scatter<T, ngpus><<<blocks, threads, 0, stream>>>( \ |
| 119 | ptrs, sg_, self_sg_, output, rank_, size_per_rank); \ |
| 120 | break; |
| 121 | |
| 122 | switch (world_size_) { |
| 123 | RS_CASE(2) |
| 124 | RS_CASE(4) |
| 125 | RS_CASE(6) |
| 126 | RS_CASE(8) |
| 127 | default: |
| 128 | throw std::runtime_error( |
| 129 | "custom reduce-scatter only supports num gpus in (2,4,6,8)"); |
| 130 | } |
| 131 | #undef RS_CASE |
| 132 | } |
| 133 | |
| 134 | template <typename T> |
| 135 | void CustomAllreduce::mnnvl_lamport_reduce_scatter(cudaStream_t stream, |
| 136 | T* input, T* output, |
| 137 | void* local_buffer, |
| 138 | uint32_t* epochs, int size, |
| 139 | int stage_size_bytes) { |
| 140 | auto packed_size = packed_t<T>::P::size; |
| 141 | if (size % (packed_size * world_size_) != 0 || |
| 142 | stage_size_bytes % sizeof(typename packed_t<T>::P) != 0) |
| 143 | throw std::runtime_error( |
| 144 | "MNNVL Lamport reduce-scatter requires 16-byte aligned sizes"); |
| 145 | |
| 146 | auto ptrs = buffers_.at(local_buffer); |
| 147 | int size_per_rank = size / packed_size / world_size_; |
| 148 | int stage_size = stage_size_bytes / sizeof(typename packed_t<T>::P); |
| 149 | int blocks_per_rank = |
| 150 | (size_per_rank + kMnnvlLamportRsThreads - 1) / kMnnvlLamportRsThreads; |
| 151 | int blocks = blocks_per_rank * world_size_; |
| 152 | |
| 153 | #if !defined(USE_ROCM) && CUDA_VERSION >= 12000 |
| 154 | cudaLaunchAttribute attributes[1]{}; |
| 155 | attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; |
| 156 | attributes[0].val.programmaticStreamSerializationAllowed = 1; |
| 157 | cudaLaunchConfig_t config{.gridDim = dim3(blocks), |
| 158 | .blockDim = dim3(kMnnvlLamportRsThreads), |
| 159 | .dynamicSmemBytes = 0, |
| 160 | .stream = stream, |
| 161 | .attrs = attributes, |
| 162 | .numAttrs = 1}; |
| 163 | #define MNNVL_LAMPORT_RS_LAUNCH(ngpus) \ |
| 164 | CUDACHECK(cudaLaunchKernelEx( \ |
| 165 | &config, &mnnvl_lamport_reduce_scatter_kernel<T, ngpus>, ptrs, input, \ |
| 166 | output, epochs, rank_, size_per_rank, stage_size)) |
| 167 | #else |
| 168 | #define MNNVL_LAMPORT_RS_LAUNCH(ngpus) \ |
| 169 | mnnvl_lamport_reduce_scatter_kernel<T, ngpus> \ |
| 170 | <<<blocks, kMnnvlLamportRsThreads, 0, stream>>>( \ |
| 171 | ptrs, input, output, epochs, rank_, size_per_rank, stage_size) |
| 172 | #endif |
| 173 | |
| 174 | #define MNNVL_LAMPORT_RS_CASE(ngpus) \ |
| 175 | case ngpus: \ |
| 176 | MNNVL_LAMPORT_RS_LAUNCH(ngpus); \ |
| 177 | break; |
| 178 | |
| 179 | switch (world_size_) { |
| 180 | MNNVL_LAMPORT_RS_CASE(2) |
| 181 | MNNVL_LAMPORT_RS_CASE(4) |
| 182 | MNNVL_LAMPORT_RS_CASE(6) |
| 183 | MNNVL_LAMPORT_RS_CASE(8) |
| 184 | MNNVL_LAMPORT_RS_CASE(16) |
| 185 | default: |
| 186 | throw std::runtime_error( |
| 187 | "MNNVL Lamport reduce-scatter only supports num gpus in " |
| 188 | "(2,4,6,8,16)"); |
| 189 | } |
| 190 | #undef MNNVL_LAMPORT_RS_CASE |
| 191 | #undef MNNVL_LAMPORT_RS_LAUNCH |
| 192 | } |
| 193 | |
| 194 | template <typename T> |
| 195 | void CustomAllreduce::mnnvl_multimem_reduce_scatter( |
| 196 | cudaStream_t stream, const T* multicast_input, T* output, |
| 197 | Signal* local_signal, Signal* multicast_signal, int size, int block_limit) { |
| 198 | #if defined(USE_ROCM) |
| 199 | throw std::runtime_error( |
| 200 | "MNNVL multimem reduce-scatter requires NVLS multicast and is not " |
| 201 | "supported on ROCm"); |
| 202 | #else |
| 203 | if (size <= 0) |
| 204 | throw std::runtime_error( |
| 205 | "MNNVL multimem reduce-scatter requires a non-empty input"); |
| 206 | int size_bytes = size * sizeof(T); |
| 207 | if (size_bytes % (kMnnvlMultimemRsVectorBytes * world_size_) != 0) |
| 208 | throw std::runtime_error( |
| 209 | "MNNVL multimem reduce-scatter requires each output shard byte size " |
| 210 | "to be a multiple of 16"); |
| 211 | if (block_limit <= 0 || block_limit > kMnnvlMultimemRsBlockLimit) |
| 212 | throw std::runtime_error( |
| 213 | "MNNVL multimem reduce-scatter block limit is out of range"); |
| 214 | |
| 215 | int packs_per_rank = size_bytes / kMnnvlMultimemRsVectorBytes / world_size_; |
| 216 | int packs_per_block = kMnnvlMultimemRsThreads * kMnnvlMultimemRsUnroll; |
| 217 | int blocks = std::min( |
| 218 | block_limit, (packs_per_rank + packs_per_block - 1) / packs_per_block); |
| 219 | |
| 220 | #define MNNVL_MULTIMEM_RS_CASE(ngpus) \ |
| 221 | case ngpus: \ |
| 222 | mnnvl_multimem_barrier_kernel<true, ngpus> \ |
| 223 | <<<1, 1, 0, stream>>>(local_signal, multicast_signal, rank_); \ |
| 224 | mnnvl_multimem_reduce_scatter_kernel<T, ngpus> \ |
| 225 | <<<blocks, kMnnvlMultimemRsThreads, 0, stream>>>( \ |
| 226 | multicast_input, output, rank_, packs_per_rank); \ |
| 227 | mnnvl_multimem_barrier_kernel<false, ngpus> \ |
| 228 | <<<1, 1, 0, stream>>>(local_signal, multicast_signal, rank_); \ |
| 229 | break; |
| 230 | |
| 231 | switch (world_size_) { |
| 232 | MNNVL_MULTIMEM_RS_CASE(2) |
| 233 | MNNVL_MULTIMEM_RS_CASE(4) |
| 234 | MNNVL_MULTIMEM_RS_CASE(8) |
| 235 | default: |
| 236 | throw std::runtime_error( |
| 237 | "MNNVL multimem reduce-scatter only supports num gpus in (2,4,8)"); |
| 238 | } |
| 239 | #undef MNNVL_MULTIMEM_RS_CASE |
| 240 | #endif |
| 241 | } |
| 242 | |
| 243 | } // namespace vllm |
| 244 | |
| 245 | using fptr_t = int64_t; |
| 246 | static_assert(sizeof(void*) == sizeof(fptr_t)); |
| 247 | |
| 248 | bool _is_weak_contiguous(torch::stable::Tensor& t); |
| 249 | |
| 250 | void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp, |
| 251 | torch::stable::Tensor& out, fptr_t _reg_buffer, |
| 252 | int64_t reg_buffer_sz_bytes) { |
| 253 | auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa); |
| 254 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 255 | inp.get_device_index()); |
| 256 | const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); |
| 257 | |
| 258 | STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); |
| 259 | STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel())); |
| 260 | STD_TORCH_CHECK(_is_weak_contiguous(out)); |
| 261 | STD_TORCH_CHECK(_is_weak_contiguous(inp)); |
| 262 | auto input_size = inp.numel() * inp.element_size(); |
| 263 | auto reg_buffer = reinterpret_cast<void*>(_reg_buffer); |
| 264 | STD_TORCH_CHECK(reg_buffer != nullptr); |
| 265 | STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes)); |
| 266 | STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size, |
| 267 | cudaMemcpyDeviceToDevice, stream)); |
| 268 | fa->allgather(stream, reg_buffer, out.mutable_data_ptr(), input_size); |
| 269 | } |
| 270 | |
| 271 | void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp, |
| 272 | torch::stable::Tensor& out, fptr_t _local_buffer, |
| 273 | fptr_t _multicast_buffer, fptr_t _epoch_buffer, |
| 274 | int64_t stage_sz_bytes) { |
| 275 | auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa); |
| 276 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 277 | inp.get_device_index()); |
| 278 | const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); |
| 279 | |
| 280 | STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); |
| 281 | STD_TORCH_CHECK((inp.numel() * fa->world_size_) == (out.numel())); |
| 282 | STD_TORCH_CHECK(_is_weak_contiguous(out)); |
| 283 | STD_TORCH_CHECK(_is_weak_contiguous(inp)); |
| 284 | auto input_size = inp.numel() * inp.element_size(); |
| 285 | STD_TORCH_CHECK((input_size * fa->world_size_) <= stage_sz_bytes); |
| 286 | auto local_buffer = reinterpret_cast<void*>(_local_buffer); |
| 287 | auto multicast_buffer = reinterpret_cast<void*>(_multicast_buffer); |
| 288 | auto epochs = reinterpret_cast<uint32_t*>(_epoch_buffer); |
| 289 | switch (out.scalar_type()) { |
| 290 | case torch::headeronly::ScalarType::Float: { |
| 291 | fa->mnnvl_lamport_allgather<float>( |
| 292 | stream, reinterpret_cast<float*>(inp.mutable_data_ptr()), |
| 293 | reinterpret_cast<float*>(out.mutable_data_ptr()), local_buffer, |
| 294 | multicast_buffer, epochs, input_size, stage_sz_bytes); |
| 295 | break; |
| 296 | } |
| 297 | case torch::headeronly::ScalarType::Half: { |
| 298 | fa->mnnvl_lamport_allgather<half>( |
| 299 | stream, reinterpret_cast<half*>(inp.mutable_data_ptr()), |
| 300 | reinterpret_cast<half*>(out.mutable_data_ptr()), local_buffer, |
| 301 | multicast_buffer, epochs, input_size, stage_sz_bytes); |
| 302 | break; |
| 303 | } |
| 304 | #if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) |
| 305 | case torch::headeronly::ScalarType::BFloat16: { |
| 306 | fa->mnnvl_lamport_allgather<nv_bfloat16>( |
| 307 | stream, reinterpret_cast<nv_bfloat16*>(inp.mutable_data_ptr()), |
| 308 | reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), local_buffer, |
| 309 | multicast_buffer, epochs, input_size, stage_sz_bytes); |
| 310 | break; |
| 311 | } |
| 312 | #endif |
| 313 | default: |
| 314 | throw std::runtime_error( |
| 315 | "MNNVL Lamport allgather only supports float32, float16 and " |
| 316 | "bfloat16"); |
| 317 | } |
| 318 | } |
| 319 | |
| 320 | void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, |
| 321 | torch::stable::Tensor& out, fptr_t _reg_buffer, |
| 322 | int64_t reg_buffer_sz_bytes) { |
| 323 | auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa); |
| 324 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 325 | inp.get_device_index()); |
| 326 | const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); |
| 327 | |
| 328 | STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); |
| 329 | STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel())); |
| 330 | STD_TORCH_CHECK(_is_weak_contiguous(out)); |
| 331 | STD_TORCH_CHECK(_is_weak_contiguous(inp)); |
| 332 | auto input_size = inp.numel() * inp.element_size(); |
| 333 | auto reg_buffer = reinterpret_cast<void*>(_reg_buffer); |
| 334 | STD_TORCH_CHECK(reg_buffer != nullptr); |
| 335 | STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes)); |
| 336 | STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size, |
| 337 | cudaMemcpyDeviceToDevice, stream)); |
| 338 | switch (out.scalar_type()) { |
| 339 | case torch::headeronly::ScalarType::Float: { |
| 340 | fa->reduce_scatter<float>( |
| 341 | stream, reinterpret_cast<float*>(reg_buffer), |
| 342 | reinterpret_cast<float*>(out.mutable_data_ptr()), inp.numel()); |
| 343 | break; |
| 344 | } |
| 345 | case torch::headeronly::ScalarType::Half: { |
| 346 | fa->reduce_scatter<half>(stream, reinterpret_cast<half*>(reg_buffer), |
| 347 | reinterpret_cast<half*>(out.mutable_data_ptr()), |
| 348 | inp.numel()); |
| 349 | break; |
| 350 | } |
| 351 | #if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) |
| 352 | case torch::headeronly::ScalarType::BFloat16: { |
| 353 | fa->reduce_scatter<nv_bfloat16>( |
| 354 | stream, reinterpret_cast<nv_bfloat16*>(reg_buffer), |
| 355 | reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), inp.numel()); |
| 356 | break; |
| 357 | } |
| 358 | #endif |
| 359 | default: |
| 360 | throw std::runtime_error( |
| 361 | "custom reduce-scatter only supports float32, float16 and bfloat16"); |
| 362 | } |
| 363 | } |
| 364 | |
| 365 | void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, |
| 366 | torch::stable::Tensor& out, |
| 367 | fptr_t _local_buffer, fptr_t _epoch_buffer, |
| 368 | int64_t stage_sz_bytes) { |
| 369 | auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa); |
| 370 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 371 | inp.get_device_index()); |
| 372 | const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); |
| 373 | |
| 374 | STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); |
| 375 | STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel())); |
| 376 | STD_TORCH_CHECK(_is_weak_contiguous(out)); |
| 377 | STD_TORCH_CHECK(_is_weak_contiguous(inp)); |
| 378 | auto input_size = inp.numel() * inp.element_size(); |
| 379 | STD_TORCH_CHECK(input_size <= stage_sz_bytes); |
| 380 | auto local_buffer = reinterpret_cast<void*>(_local_buffer); |
| 381 | auto epochs = reinterpret_cast<uint32_t*>(_epoch_buffer); |
| 382 | switch (out.scalar_type()) { |
| 383 | case torch::headeronly::ScalarType::Float: { |
| 384 | fa->mnnvl_lamport_reduce_scatter<float>( |
| 385 | stream, reinterpret_cast<float*>(inp.mutable_data_ptr()), |
| 386 | reinterpret_cast<float*>(out.mutable_data_ptr()), local_buffer, |
| 387 | epochs, inp.numel(), stage_sz_bytes); |
| 388 | break; |
| 389 | } |
| 390 | case torch::headeronly::ScalarType::Half: { |
| 391 | fa->mnnvl_lamport_reduce_scatter<half>( |
| 392 | stream, reinterpret_cast<half*>(inp.mutable_data_ptr()), |
| 393 | reinterpret_cast<half*>(out.mutable_data_ptr()), local_buffer, epochs, |
| 394 | inp.numel(), stage_sz_bytes); |
| 395 | break; |
| 396 | } |
| 397 | #if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) |
| 398 | case torch::headeronly::ScalarType::BFloat16: { |
| 399 | fa->mnnvl_lamport_reduce_scatter<nv_bfloat16>( |
| 400 | stream, reinterpret_cast<nv_bfloat16*>(inp.mutable_data_ptr()), |
| 401 | reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), local_buffer, |
| 402 | epochs, inp.numel(), stage_sz_bytes); |
| 403 | break; |
| 404 | } |
| 405 | #endif |
| 406 | default: |
| 407 | throw std::runtime_error( |
| 408 | "MNNVL Lamport reduce-scatter only supports float32, float16 and " |
| 409 | "bfloat16"); |
| 410 | } |
| 411 | } |
| 412 | |
| 413 | void mnnvl_multimem_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp, |
| 414 | torch::stable::Tensor& out, |
| 415 | fptr_t _local_buffer, |
| 416 | fptr_t _multicast_buffer, |
| 417 | int64_t stage_sz_bytes, |
| 418 | int64_t block_limit) { |
| 419 | auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa); |
| 420 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 421 | inp.get_device_index()); |
| 422 | const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index()); |
| 423 | |
| 424 | STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type())); |
| 425 | STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel())); |
| 426 | STD_TORCH_CHECK(inp.numel() > 0); |
| 427 | STD_TORCH_CHECK(_is_weak_contiguous(out)); |
| 428 | STD_TORCH_CHECK(_is_weak_contiguous(inp)); |
| 429 | auto input_size = inp.numel() * inp.element_size(); |
| 430 | STD_TORCH_CHECK(input_size <= stage_sz_bytes); |
| 431 | auto local_buffer = reinterpret_cast<void*>(_local_buffer); |
| 432 | auto multicast_buffer = reinterpret_cast<void*>(_multicast_buffer); |
| 433 | STD_TORCH_CHECK(local_buffer != nullptr); |
| 434 | STD_TORCH_CHECK(multicast_buffer != nullptr); |
| 435 | auto* local_signal = reinterpret_cast<vllm::Signal*>( |
| 436 | static_cast<char*>(local_buffer) - sizeof(vllm::Signal)); |
| 437 | auto* multicast_signal = reinterpret_cast<vllm::Signal*>( |
| 438 | static_cast<char*>(multicast_buffer) - sizeof(vllm::Signal)); |
| 439 | STD_CUDA_CHECK(cudaMemcpyAsync(local_buffer, inp.const_data_ptr(), input_size, |
| 440 | cudaMemcpyDeviceToDevice, stream)); |
| 441 | switch (out.scalar_type()) { |
| 442 | case torch::headeronly::ScalarType::Float: { |
| 443 | fa->mnnvl_multimem_reduce_scatter<float>( |
| 444 | stream, reinterpret_cast<float*>(multicast_buffer), |
| 445 | reinterpret_cast<float*>(out.mutable_data_ptr()), local_signal, |
| 446 | multicast_signal, inp.numel(), block_limit); |
| 447 | break; |
| 448 | } |
| 449 | case torch::headeronly::ScalarType::Half: { |
| 450 | fa->mnnvl_multimem_reduce_scatter<half>( |
| 451 | stream, reinterpret_cast<half*>(multicast_buffer), |
| 452 | reinterpret_cast<half*>(out.mutable_data_ptr()), local_signal, |
| 453 | multicast_signal, inp.numel(), block_limit); |
| 454 | break; |
| 455 | } |
| 456 | #if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) |
| 457 | case torch::headeronly::ScalarType::BFloat16: { |
| 458 | fa->mnnvl_multimem_reduce_scatter<nv_bfloat16>( |
| 459 | stream, reinterpret_cast<nv_bfloat16*>(multicast_buffer), |
| 460 | reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), local_signal, |
| 461 | multicast_signal, inp.numel(), block_limit); |
| 462 | break; |
| 463 | } |
| 464 | #endif |
| 465 | default: |
| 466 | throw std::runtime_error( |
| 467 | "MNNVL multimem reduce-scatter only supports float32, float16 and " |
| 468 | "bfloat16"); |
| 469 | } |
| 470 | } |