| 1 | #pragma once |
| 2 | |
| 3 | #include "custom_collective_common.cuh" |
| 4 | |
| 5 | namespace vllm { |
| 6 | |
| 7 | constexpr int kMnnvlLamportAgThreads = 128; |
| 8 | constexpr int kMnnvlLamportRsThreads = 256; |
| 9 | constexpr int kMnnvlLamportConcurrentPollMaxPacks = 8192; |
| 10 | constexpr int kMnnvlMultimemRsThreads = 1024; |
| 11 | constexpr int kMnnvlMultimemRsBlockLimit = 8; |
| 12 | constexpr int kMnnvlMultimemRsVectorBytes = 16; |
| 13 | constexpr int kMnnvlMultimemRsUnroll = 8; |
| 14 | |
| 15 | using CopyPack = array_t<uint64_t, 2>; |
| 16 | |
| 17 | template <int ngpus> |
| 18 | __global__ void __launch_bounds__(512, 1) |
| 19 | cross_device_all_gather(RankData* _dp, RankSignals sg, Signal* self_sg, |
| 20 | CopyPack* __restrict__ result, int rank, |
| 21 | int size_per_rank) { |
| 22 | auto dp = *_dp; |
| 23 | int tid = blockIdx.x * blockDim.x + threadIdx.x; |
| 24 | int stride = gridDim.x * blockDim.x; |
| 25 | barrier_at_start<ngpus>(sg, self_sg, rank); |
| 26 | #pragma unroll |
| 27 | for (int src_rank = 0; src_rank < ngpus; ++src_rank) { |
| 28 | auto src = reinterpret_cast<const CopyPack*>(dp.ptrs[src_rank]); |
| 29 | auto dst = result + src_rank * size_per_rank; |
| 30 | for (int idx = tid; idx < size_per_rank; idx += stride) { |
| 31 | dst[idx] = src[idx]; |
| 32 | } |
| 33 | } |
| 34 | barrier_at_end<ngpus, true>(sg, self_sg, rank); |
| 35 | } |
| 36 | |
| 37 | template <typename T, int ngpus> |
| 38 | __global__ void __launch_bounds__(512, 1) |
| 39 | cross_device_reduce_scatter(RankData* _dp, RankSignals sg, Signal* self_sg, |
| 40 | T* __restrict__ result, int rank, |
| 41 | int size_per_rank) { |
| 42 | using P = typename packed_t<T>::P; |
| 43 | using A = typename packed_t<T>::A; |
| 44 | auto dp = *_dp; |
| 45 | auto offset = rank * size_per_rank; |
| 46 | barrier_at_start<ngpus>(sg, self_sg, rank); |
| 47 | for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size_per_rank; |
| 48 | idx += gridDim.x * blockDim.x) { |
| 49 | reinterpret_cast<P*>(result)[idx] = |
| 50 | packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], offset + idx); |
| 51 | } |
| 52 | barrier_at_end<ngpus, true>(sg, self_sg, rank); |
| 53 | } |
| 54 | |
| 55 | // Multimem reduce-scatter reduces through an NVLS multicast mapping, which |
| 56 | // has no AMD equivalent. The whole path is compiled out on ROCm and the host |
| 57 | // entry point rejects the call instead of launching anything. |
| 58 | #if !defined(USE_ROCM) |
| 59 | |
| 60 | template <typename T> |
| 61 | DINLINE void multimem_load_reduce_16(uint32_t (&result)[4], const T* address) { |
| 62 | #if CUDA_VERSION >= 12020 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) |
| 63 | if constexpr (std::is_same<T, nv_bfloat16>::value) { |
| 64 | asm volatile( |
| 65 | "multimem.ld_reduce.relaxed.sys.global.add.acc::f32.v4.bf16x2 " |
| 66 | "{%0,%1,%2,%3}, [%4];" |
| 67 | : "=r"(result[0]), "=r"(result[1]), "=r"(result[2]), "=r"(result[3]) |
| 68 | : "l"(address) |
| 69 | : "memory"); |
| 70 | } else if constexpr (std::is_same<T, half>::value) { |
| 71 | asm volatile( |
| 72 | "multimem.ld_reduce.relaxed.sys.global.add.acc::f32.v4.f16x2 " |
| 73 | "{%0,%1,%2,%3}, [%4];" |
| 74 | : "=r"(result[0]), "=r"(result[1]), "=r"(result[2]), "=r"(result[3]) |
| 75 | : "l"(address) |
| 76 | : "memory"); |
| 77 | } else { |
| 78 | static_assert(std::is_same<T, float>::value); |
| 79 | asm volatile( |
| 80 | "multimem.ld_reduce.relaxed.sys.global.add.v4.f32 " |
| 81 | "{%0,%1,%2,%3}, [%4];" |
| 82 | : "=r"(result[0]), "=r"(result[1]), "=r"(result[2]), "=r"(result[3]) |
| 83 | : "l"(address) |
| 84 | : "memory"); |
| 85 | } |
| 86 | #else |
| 87 | asm volatile("trap;"); |
| 88 | #endif |
| 89 | } |
| 90 | |
| 91 | DINLINE void store_global_16(void* address, const uint32_t (&value)[4]) { |
| 92 | asm volatile("st.global.v4.u32 [%0], {%1,%2,%3,%4};" |
| 93 | : |
| 94 | : "l"(address), "r"(value[0]), "r"(value[1]), "r"(value[2]), |
| 95 | "r"(value[3]) |
| 96 | : "memory"); |
| 97 | } |
| 98 | |
| 99 | DINLINE void mnnvl_multimem_publish_flag(FlagType* flag_addr, FlagType flag) { |
| 100 | #if CUDA_VERSION >= 12020 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) |
| 101 | asm volatile("multimem.st.release.sys.global.u32 [%0], %1;" |
| 102 | : |
| 103 | : "l"(flag_addr), "r"(flag) |
| 104 | : "memory"); |
| 105 | // The flag is stored through the multicast alias and polled through the |
| 106 | // local unicast alias in the same grid. |
| 107 | asm volatile("fence.proxy.alias;" ::: "memory"); |
| 108 | #else |
| 109 | asm volatile("trap;"); |
| 110 | #endif |
| 111 | } |
| 112 | |
| 113 | template <bool ready, int ngpus> |
| 114 | __global__ void mnnvl_multimem_barrier_kernel(Signal* local_signal, |
| 115 | Signal* multicast_signal, |
| 116 | int rank) { |
| 117 | FlagType flag = local_signal->_flag[0] + 1; |
| 118 | FlagType* multicast_counters = |
| 119 | ready ? multicast_signal->start[0] : multicast_signal->end[0]; |
| 120 | FlagType* local_counters = |
| 121 | ready ? local_signal->start[0] : local_signal->end[0]; |
| 122 | // Replicate this rank's counter into every backing allocation. |
| 123 | mnnvl_multimem_publish_flag(&multicast_counters[rank], flag); |
| 124 | #pragma unroll |
| 125 | for (int peer = 0; peer < ngpus; ++peer) { |
| 126 | while (ld_flag_acquire(&local_counters[peer]) != flag); |
| 127 | } |
| 128 | local_signal->_flag[0] = flag; |
| 129 | } |
| 130 | |
| 131 | template <typename T, int ngpus> |
| 132 | __global__ void __launch_bounds__(kMnnvlMultimemRsThreads, 1) |
| 133 | mnnvl_multimem_reduce_scatter_kernel(const T* __restrict__ multicast_input, |
| 134 | T* __restrict__ result, int rank, |
| 135 | int packs_per_rank) { |
| 136 | constexpr int kThreadsPerWarp = 32; |
| 137 | constexpr int kWarpsPerCta = kMnnvlMultimemRsThreads / kThreadsPerWarp; |
| 138 | constexpr int kPacksPerWarpIteration = |
| 139 | kThreadsPerWarp * kMnnvlMultimemRsUnroll; |
| 140 | |
| 141 | int lane; |
| 142 | asm volatile("mov.u32 %0, %%laneid;" : "=r"(lane)); |
| 143 | int warp = blockIdx.x * kWarpsPerCta + threadIdx.x / kThreadsPerWarp; |
| 144 | int num_warps = gridDim.x * kWarpsPerCta; |
| 145 | int pack_offset = warp * kPacksPerWarpIteration + lane; |
| 146 | int pack_stride = num_warps * kPacksPerWarpIteration; |
| 147 | auto* rank_input = reinterpret_cast<const char*>(multicast_input) + |
| 148 | rank * packs_per_rank * kMnnvlMultimemRsVectorBytes; |
| 149 | auto* rank_output = reinterpret_cast<char*>(result); |
| 150 | |
| 151 | while (pack_offset < packs_per_rank) { |
| 152 | uint32_t reduced[kMnnvlMultimemRsUnroll][4]; |
| 153 | #pragma unroll |
| 154 | for (int u = 0; u < kMnnvlMultimemRsUnroll; ++u) { |
| 155 | int pack = pack_offset + u * kThreadsPerWarp; |
| 156 | if (pack < packs_per_rank) { |
| 157 | multimem_load_reduce_16<T>( |
| 158 | reduced[u], reinterpret_cast<const T*>( |
| 159 | rank_input + pack * kMnnvlMultimemRsVectorBytes)); |
| 160 | } |
| 161 | } |
| 162 | #pragma unroll |
| 163 | for (int u = 0; u < kMnnvlMultimemRsUnroll; ++u) { |
| 164 | int pack = pack_offset + u * kThreadsPerWarp; |
| 165 | if (pack < packs_per_rank) { |
| 166 | store_global_16(rank_output + pack * kMnnvlMultimemRsVectorBytes, |
| 167 | reduced[u]); |
| 168 | } |
| 169 | } |
| 170 | pack_offset += pack_stride; |
| 171 | } |
| 172 | } |
| 173 | |
| 174 | #endif // !defined(USE_ROCM) |
| 175 | |
| 176 | template <typename P> |
| 177 | union LamportPack { |
| 178 | P packed; |
| 179 | uint32_t words[sizeof(P) / sizeof(uint32_t)]; |
| 180 | }; |
| 181 | |
| 182 | template <typename P> |
| 183 | DINLINE LamportPack<P> load_lamport_pack(const P* ptr) { |
| 184 | static_assert(sizeof(P) == 16); |
| 185 | LamportPack<P> value; |
| 186 | #if !defined(USE_ROCM) |
| 187 | asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];" |
| 188 | : "=r"(value.words[0]), "=r"(value.words[1]), |
| 189 | "=r"(value.words[2]), "=r"(value.words[3]) |
| 190 | : "l"(ptr) |
| 191 | : "memory"); |
| 192 | #else |
| 193 | const volatile uint32_t* src = |
| 194 | reinterpret_cast<const volatile uint32_t*>(ptr); |
| 195 | #pragma unroll |
| 196 | for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { |
| 197 | value.words[i] = src[i]; |
| 198 | } |
| 199 | #endif |
| 200 | return value; |
| 201 | } |
| 202 | |
| 203 | template <typename P> |
| 204 | DINLINE bool is_lamport_dirty(const LamportPack<P>& value) { |
| 205 | #pragma unroll |
| 206 | for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { |
| 207 | if (value.words[i] == 0x80000000U) return true; |
| 208 | } |
| 209 | return false; |
| 210 | } |
| 211 | |
| 212 | template <typename P> |
| 213 | DINLINE P lamport_sentinel() { |
| 214 | LamportPack<P> value; |
| 215 | #pragma unroll |
| 216 | for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { |
| 217 | value.words[i] = 0x80000000U; |
| 218 | } |
| 219 | return value.packed; |
| 220 | } |
| 221 | |
| 222 | template <typename P> |
| 223 | DINLINE P sanitize_lamport_payload(P packed) { |
| 224 | LamportPack<P> value{.packed = packed}; |
| 225 | #pragma unroll |
| 226 | for (int i = 0; i < sizeof(P) / sizeof(uint32_t); ++i) { |
| 227 | if (value.words[i] == 0x80000000U) value.words[i] = 0; |
| 228 | } |
| 229 | return value.packed; |
| 230 | } |
| 231 | |
| 232 | __host__ __device__ constexpr int mnnvl_lamport_dirty_stage(int current_stage) { |
| 233 | return (current_stage + 2) % 3; |
| 234 | } |
| 235 | |
| 236 | __host__ __device__ constexpr int mnnvl_lamport_next_stage(int current_stage) { |
| 237 | return (current_stage + 1) % 3; |
| 238 | } |
| 239 | |
| 240 | template <typename P> |
| 241 | DINLINE void store_multimem_lamport_payload(P* ptr, P packed) { |
| 242 | static_assert(sizeof(P) == 16); |
| 243 | // multimem was introduced in PTX 8.1 (CUDA 12.1) for SM90 and newer. |
| 244 | #if !defined(USE_ROCM) && CUDA_VERSION >= 12010 && defined(__CUDA_ARCH__) && \ |
| 245 | (__CUDA_ARCH__ >= 900) |
| 246 | LamportPack<P> value{.packed = packed}; |
| 247 | // The local alias remains sentinel until the multicast payload becomes |
| 248 | // observable; readers reject prior or partial values and retry. |
| 249 | asm volatile("multimem.st.relaxed.sys.global.v4.f32 [%0], {%1,%2,%3,%4};" |
| 250 | : |
| 251 | : "l"(ptr), "r"(value.words[0]), "r"(value.words[1]), |
| 252 | "r"(value.words[2]), "r"(value.words[3]) |
| 253 | : "memory"); |
| 254 | #elif defined(USE_ROCM) |
| 255 | __builtin_trap(); |
| 256 | #else |
| 257 | // Multicast mappings do not exist before SM90. Fail closed if this kernel is |
| 258 | // ever dispatched for an unsupported target instead of issuing an undefined |
| 259 | // ordinary store to a multicast address. |
| 260 | asm volatile("trap;"); |
| 261 | #endif |
| 262 | } |
| 263 | |
| 264 | template <typename P> |
| 265 | DINLINE P wait_lamport_payload(const P* ptr) { |
| 266 | auto value = load_lamport_pack(ptr); |
| 267 | while (is_lamport_dirty(value)) value = load_lamport_pack(ptr); |
| 268 | return value.packed; |
| 269 | } |
| 270 | |
| 271 | template <typename P, int ngpus> |
| 272 | DINLINE void wait_lamport_payloads(const P* base, int rank, int rank_stride, |
| 273 | P local_value, P (&values)[ngpus]) { |
| 274 | bool ready[ngpus]; |
| 275 | #pragma unroll |
| 276 | for (int src_rank = 0; src_rank < ngpus; ++src_rank) { |
| 277 | ready[src_rank] = src_rank == rank; |
| 278 | if (src_rank == rank) values[src_rank] = local_value; |
| 279 | } |
| 280 | |
| 281 | int remaining = ngpus - 1; |
| 282 | while (remaining != 0) { |
| 283 | #pragma unroll |
| 284 | for (int src_rank = 0; src_rank < ngpus; ++src_rank) { |
| 285 | if (!ready[src_rank]) { |
| 286 | auto value = load_lamport_pack(base + src_rank * rank_stride); |
| 287 | if (!is_lamport_dirty(value)) { |
| 288 | values[src_rank] = value.packed; |
| 289 | ready[src_rank] = true; |
| 290 | --remaining; |
| 291 | } |
| 292 | } |
| 293 | } |
| 294 | } |
| 295 | } |
| 296 | |
| 297 | template <typename P, typename A, int ngpus> |
| 298 | DINLINE P reduce_lamport_payloads(const P* current_local, const P* packed_input, |
| 299 | int rank, int size_per_rank, int idx) { |
| 300 | P source_zero = |
| 301 | rank == 0 ? packed_input[idx] : wait_lamport_payload(current_local + idx); |
| 302 | A tmp = upcast(source_zero); |
| 303 | #pragma unroll |
| 304 | for (int src_rank = 1; src_rank < ngpus; ++src_rank) { |
| 305 | P value = src_rank == rank |
| 306 | ? packed_input[rank * size_per_rank + idx] |
| 307 | : wait_lamport_payload(current_local + |
| 308 | src_rank * size_per_rank + idx); |
| 309 | packed_assign_add(tmp, upcast(value)); |
| 310 | } |
| 311 | return sanitize_lamport_payload(downcast<P>(tmp)); |
| 312 | } |
| 313 | |
| 314 | DINLINE void lamport_cta_arrive(uint32_t* counter) { |
| 315 | #if !defined(USE_ROCM) |
| 316 | if (threadIdx.x < 32) { |
| 317 | asm volatile("barrier.cta.sync 1, %0;" : : "r"(blockDim.x) : "memory"); |
| 318 | if (threadIdx.x == 0) { |
| 319 | #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 |
| 320 | asm volatile("red.async.release.global.gpu.add.u32 [%0], 1;" |
| 321 | : |
| 322 | : "l"(counter) |
| 323 | : "memory"); |
| 324 | #elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 |
| 325 | asm volatile("red.release.global.gpu.add.u32 [%0], 1;" |
| 326 | : |
| 327 | : "l"(counter) |
| 328 | : "memory"); |
| 329 | #else |
| 330 | atomicAdd(counter, 1); |
| 331 | #endif |
| 332 | } |
| 333 | } else { |
| 334 | asm volatile("barrier.cta.arrive 1, %0;" : : "r"(blockDim.x) : "memory"); |
| 335 | } |
| 336 | #else |
| 337 | __syncthreads(); |
| 338 | if (threadIdx.x == 0) atomicAdd(counter, 1); |
| 339 | #endif |
| 340 | } |
| 341 | |
| 342 | template <typename T, int ngpus> |
| 343 | __global__ void __launch_bounds__(kMnnvlLamportAgThreads, 1) |
| 344 | mnnvl_lamport_all_gather(RankData* _dp, const T* __restrict__ input, |
| 345 | T* __restrict__ result, |
| 346 | T* __restrict__ multicast_buffer, |
| 347 | uint32_t* __restrict__ epochs, int rank, |
| 348 | int size_per_rank, int stage_size) { |
| 349 | using P = typename packed_t<T>::P; |
| 350 | #if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ |
| 351 | (__CUDA_ARCH__ >= 900) |
| 352 | cudaGridDependencySynchronize(); |
| 353 | #endif |
| 354 | auto dp = *_dp; |
| 355 | int tid = blockIdx.x * blockDim.x + threadIdx.x; |
| 356 | int stride = gridDim.x * blockDim.x; |
| 357 | uint32_t epoch = epochs[0]; |
| 358 | int current_stage = epoch % 3; |
| 359 | // A peer may start the next epoch after we publish but before we finish. |
| 360 | // Clean the previous stage, which cannot be reused until two epochs later. |
| 361 | int dirty_stage = mnnvl_lamport_dirty_stage(current_stage); |
| 362 | int dirty_size = epochs[2 + dirty_stage]; |
| 363 | auto local_buffer = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[rank])); |
| 364 | auto current_local = local_buffer + current_stage * stage_size; |
| 365 | auto dirty_local = local_buffer + dirty_stage * stage_size; |
| 366 | auto current_multicast = |
| 367 | reinterpret_cast<P*>(multicast_buffer) + current_stage * stage_size; |
| 368 | auto packed_input = reinterpret_cast<const P*>(input); |
| 369 | auto packed_result = reinterpret_cast<P*>(result); |
| 370 | |
| 371 | int total_size = size_per_rank * ngpus; |
| 372 | P local_value; |
| 373 | if (tid < size_per_rank) { |
| 374 | local_value = packed_input[tid]; |
| 375 | // A CUDA multicast mapping may only be accessed with multimem PTX; |
| 376 | // ordinary global loads and stores have undefined behavior. |
| 377 | store_multimem_lamport_payload( |
| 378 | current_multicast + rank * size_per_rank + tid, |
| 379 | sanitize_lamport_payload(local_value)); |
| 380 | } |
| 381 | #if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ |
| 382 | (__CUDA_ARCH__ >= 900) |
| 383 | cudaTriggerProgrammaticLaunchCompletion(); |
| 384 | #endif |
| 385 | |
| 386 | lamport_cta_arrive(&epochs[1]); |
| 387 | |
| 388 | for (int idx = tid; idx < dirty_size; idx += stride) { |
| 389 | dirty_local[idx] = lamport_sentinel<P>(); |
| 390 | } |
| 391 | |
| 392 | if (tid < size_per_rank) { |
| 393 | #pragma unroll |
| 394 | for (int src_rank = 0; src_rank < ngpus; ++src_rank) { |
| 395 | int output_idx = src_rank * size_per_rank + tid; |
| 396 | P value = src_rank == rank |
| 397 | ? local_value |
| 398 | : wait_lamport_payload(current_local + output_idx); |
| 399 | packed_result[output_idx] = value; |
| 400 | } |
| 401 | } |
| 402 | |
| 403 | if (tid == 0) { |
| 404 | while (*reinterpret_cast<volatile uint32_t*>(&epochs[1]) < gridDim.x); |
| 405 | epochs[2 + current_stage] = total_size; |
| 406 | epochs[0] = mnnvl_lamport_next_stage(current_stage); |
| 407 | epochs[1] = 0; |
| 408 | } |
| 409 | } |
| 410 | |
| 411 | template <typename T, int ngpus> |
| 412 | __global__ void __launch_bounds__(kMnnvlLamportRsThreads, 1) |
| 413 | mnnvl_lamport_reduce_scatter_kernel(RankData* _dp, |
| 414 | const T* __restrict__ input, |
| 415 | T* __restrict__ result, |
| 416 | uint32_t* __restrict__ epochs, int rank, |
| 417 | int size_per_rank, int stage_size) { |
| 418 | using P = typename packed_t<T>::P; |
| 419 | using A = typename packed_t<T>::A; |
| 420 | #if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ |
| 421 | (__CUDA_ARCH__ >= 900) |
| 422 | cudaGridDependencySynchronize(); |
| 423 | #endif |
| 424 | auto dp = *_dp; |
| 425 | int dst_rank = blockIdx.x % ngpus; |
| 426 | int tile = blockIdx.x / ngpus; |
| 427 | int idx = tile * blockDim.x + threadIdx.x; |
| 428 | int tid = blockIdx.x * blockDim.x + threadIdx.x; |
| 429 | int stride = gridDim.x * blockDim.x; |
| 430 | uint32_t epoch = epochs[0]; |
| 431 | int current_stage = epoch % 3; |
| 432 | // A peer may start the next epoch after we publish but before we finish. |
| 433 | // Clean the previous stage, which cannot be reused until two epochs later. |
| 434 | int dirty_stage = mnnvl_lamport_dirty_stage(current_stage); |
| 435 | int dirty_size = epochs[2 + dirty_stage]; |
| 436 | auto local_buffer = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[rank])); |
| 437 | auto current_local = local_buffer + current_stage * stage_size; |
| 438 | auto dirty_local = local_buffer + dirty_stage * stage_size; |
| 439 | auto packed_input = reinterpret_cast<const P*>(input); |
| 440 | |
| 441 | if (idx < size_per_rank && dst_rank != rank) { |
| 442 | auto dst = reinterpret_cast<P*>(const_cast<void*>(dp.ptrs[dst_rank])) + |
| 443 | current_stage * stage_size + rank * size_per_rank; |
| 444 | auto src = packed_input + dst_rank * size_per_rank; |
| 445 | dst[idx] = sanitize_lamport_payload(src[idx]); |
| 446 | } |
| 447 | #if !defined(USE_ROCM) && CUDA_VERSION >= 12000 && defined(__CUDA_ARCH__) && \ |
| 448 | (__CUDA_ARCH__ >= 900) |
| 449 | cudaTriggerProgrammaticLaunchCompletion(); |
| 450 | #endif |
| 451 | |
| 452 | lamport_cta_arrive(&epochs[1]); |
| 453 | |
| 454 | for (int idx = tid; idx < dirty_size; idx += stride) { |
| 455 | dirty_local[idx] = lamport_sentinel<P>(); |
| 456 | } |
| 457 | |
| 458 | if (idx < size_per_rank && dst_rank == rank) { |
| 459 | if constexpr (ngpus == 4) { |
| 460 | if (size_per_rank > kMnnvlLamportConcurrentPollMaxPacks) { |
| 461 | reinterpret_cast<P*>(result)[idx] = |
| 462 | reduce_lamport_payloads<P, A, ngpus>(current_local, packed_input, |
| 463 | rank, size_per_rank, idx); |
| 464 | } else { |
| 465 | P values[ngpus]; |
| 466 | wait_lamport_payloads<P, ngpus>( |
| 467 | current_local + idx, rank, size_per_rank, |
| 468 | packed_input[rank * size_per_rank + idx], values); |
| 469 | A tmp = upcast(values[0]); |
| 470 | #pragma unroll |
| 471 | for (int src_rank = 1; src_rank < ngpus; ++src_rank) { |
| 472 | packed_assign_add(tmp, upcast(values[src_rank])); |
| 473 | } |
| 474 | reinterpret_cast<P*>(result)[idx] = |
| 475 | sanitize_lamport_payload(downcast<P>(tmp)); |
| 476 | } |
| 477 | } else { |
| 478 | reinterpret_cast<P*>(result)[idx] = reduce_lamport_payloads<P, A, ngpus>( |
| 479 | current_local, packed_input, rank, size_per_rank, idx); |
| 480 | } |
| 481 | } |
| 482 | |
| 483 | if (tid == 0) { |
| 484 | while (*reinterpret_cast<volatile uint32_t*>(&epochs[1]) < gridDim.x); |
| 485 | epochs[2 + current_stage] = size_per_rank * ngpus; |
| 486 | epochs[0] = mnnvl_lamport_next_stage(current_stage); |
| 487 | epochs[1] = 0; |
| 488 | } |
| 489 | } |
| 490 | |
| 491 | } // namespace vllm |