diff --git a/csrc/include/custom_all_reduce.cuh b/csrc/include/custom_all_reduce.cuh index be7b4c1eb3..0b2dddbbef 100644 --- a/csrc/include/custom_all_reduce.cuh +++ b/csrc/include/custom_all_reduce.cuh @@ -908,9 +908,32 @@ template DINLINE T shfl_xor(T var, int mask, int width = opus::get_warp_size()) { static_assert(sizeof(T) == 4); + int v = __builtin_bit_cast(int, var); + switch(mask) + { + case 1: + return __builtin_bit_cast(T, __builtin_amdgcn_mov_dpp(v, 0xb1, 0xf, 0xf, true)); + case 2: + return __builtin_bit_cast(T, __builtin_amdgcn_mov_dpp(v, 0x4e, 0xf, 0xf, true)); + case 4: + { + int r = __builtin_amdgcn_update_dpp(v, v, 0x104, 0xf, 0x5, true); + r = __builtin_amdgcn_update_dpp(r, v, 0x114, 0xf, 0xa, true); + return __builtin_bit_cast(T, r); + } + case 8: + { + int r = __builtin_amdgcn_update_dpp(v, v, 0x108, 0xf, 0x3, true); + r = __builtin_amdgcn_update_dpp(r, v, 0x118, 0xf, 0xc, true); + return __builtin_bit_cast(T, r); + } + default: + break; + } + // fallback int self = opus::lane_id(); int index = (self & ~(width - 1)) + ((self ^ mask) & (width - 1)); - return __builtin_bit_cast(T, __builtin_amdgcn_ds_bpermute(index << 2, __builtin_bit_cast(int, var))); + return __builtin_bit_cast(T, __builtin_amdgcn_ds_bpermute(index << 2, v)); } // shfl_xor support 4bytes dtype only @@ -1092,13 +1115,8 @@ __global__ void __launch_bounds__(512, 1) reduce_scatter_cross_device_store( RankData* _dp, RankSignals sg, Signal* self_sg, int rank, int size) { constexpr int pack_size = 16 / sizeof(T); - constexpr int tnum_gpu = THREAD_NUM / ngpus; using P = typename opus::vector_t; using A = typename opus::vector_t; - __shared__ T tmp_smem[tnum_gpu * ngpus * pack_size]; - int warp_id = threadIdx.x / tnum_gpu; - int lane_id = threadIdx.x % tnum_gpu; - int tid = blockIdx.x * tnum_gpu + lane_id; const P* ptrs[ngpus]; P* tmps[ngpus]; #pragma unroll @@ -1109,45 +1127,37 @@ __global__ void __launch_bounds__(512, 1) reduce_scatter_cross_device_store( } start_sync(sg, self_sg, rank); + // FLAT reduce-scatter + broadcast. This rank owns the pack range + // [rank*part, (rank+1)*part); each thread reduces ONE pack across all ngpus + // inputs in fp32 and stores the bf16 sum into every rank's tmp at the same + // offset (all-gather via broadcast). Each thread thus touches all ngpus links + // for both the reads and the ngpus broadcast stores; the ngpus stores/thread + // pipeline the write (broadcast) half far better than the previous + // warp-per-rank+LDS scheme's 1 store/thread. Measured faster at every m + // (0.91x at m=64 -> 0.86x at m=512 vs warp-per-rank) with bit-identical output + // (same canonical reduce order, same bf16 sum). No LDS, no intra-block sync. int part = size / (pack_size * ngpus); - for(int idx = tid; idx < part; idx += gridDim.x * tnum_gpu) + int base = rank * part; + for(int l = blockIdx.x * blockDim.x + threadIdx.x; l < part; + l += gridDim.x * blockDim.x) { - // cross device read by all warp - P input_reg = ptrs[warp_id][rank * part + idx]; - *(reinterpret_cast(&tmp_smem[0]) + threadIdx.x) = input_reg; - __syncthreads(); - // calculate and save in first warp - if(warp_id == 0) - { - A add_reg; + int gp = base + l; + A acc{}; #pragma unroll - for(int i = 0; i < pack_size; ++i) - { - add_reg[i] = upcast_s(tmp_smem[pack_size * threadIdx.x + i]); - } + for(int r = 0; r < ngpus; ++r) + { + P v = ptrs[r][gp]; #pragma unroll - for(int i = 1; i < ngpus; ++i) - { + for(int e = 0; e < pack_size; ++e) + acc[e] += upcast_s(v[e]); + } + P s; #pragma unroll - for(int j = 0; j < pack_size; ++j) - { - add_reg[j] += - upcast_s(tmp_smem[i * pack_size * tnum_gpu + pack_size * threadIdx.x + j]); - } - } - P add_rslt; + for(int e = 0; e < pack_size; ++e) + s[e] = downcast_s(acc[e]); #pragma unroll - for(int i = 0; i < pack_size; ++i) - { - add_rslt[i] = downcast_s(add_reg[i]); - } - *(reinterpret_cast(&tmp_smem[0]) + lane_id) = add_rslt; - } - __syncthreads(); - - // cross device store - P rslt = *(reinterpret_cast(&tmp_smem[0]) + lane_id); - tmps[warp_id][rank * part + idx] = rslt; + for(int w = 0; w < ngpus; ++w) + tmps[w][gp] = s; } // NOTE: must use final_sync=false (RELEASE/ACQUIRE) here. Stage 2 // (local_device_load_rmsnorm*) on each rank reads `tmps` on the @@ -3130,6 +3140,14 @@ void dispatchFusedAllReduceRMSNorm(hipStream_t stream, auto pack_size = 16 / sizeof(T); use_1stage = use_1stage && (n % pack_size == 0) && (n / pack_size <= 1024); + // Win-gate the path by data volume: one-stage's lower overhead (single barrier, + // no HBM round-trip) only wins for small inputs; above the measured crossover the + // FLAT two-stage's halved cross-card traffic (3S -> 1.5S) wins. Crossover ~0.44 MiB + // (m=32 @ hidden 7168); use 0.5 MiB with margin. So even if the caller permits + // one-stage, fall back to two-stage once we're past where it pays (e.g. m=64 decode + // = 0.875 MiB -> two-stage). This only DISABLES one-stage; it never overrides an + // explicit two-stage request (use_1stage already false). + use_1stage = use_1stage && ((size_t)size * sizeof(T) < (size_t)512 * 1024); #define MAYBE_DISPATCH_1S_KERNEL(NGPUS) \ if(use_1stage) \ { \ @@ -3153,22 +3171,29 @@ void dispatchFusedAllReduceRMSNorm(hipStream_t stream, dim3 block(512); int block_num = ((size / world_size_) + 512 - 1) / 512; dim3 grid(std::min(block_num, 80)); + // FLAT reduce_scatter_cross_device_store is 1 pack/thread: launch it with block 256 + // (not the 512 used by step-2 rmsnorm) so small-m decode engages more CUs -- e.g. + // m=64 -> 56 active blocks vs 28 at block 512. grid covers this rank's pack range, + // capped at kMaxBlocks (start_sync's per-block slot limit). + int rs_packs = size / ((16 / sizeof(T)) * world_size_); + dim3 rs_block(256); + dim3 rs_grid(std::min((rs_packs + 255) / 256, 80)); switch(world_size_) { case 8: MAYBE_DISPATCH_1S_KERNEL(8); reduce_scatter_cross_device_store - <<>>(ptrs, sg_, self_sg_, rank_, size); + <<>>(ptrs, sg_, self_sg_, rank_, size); break; case 4: MAYBE_DISPATCH_1S_KERNEL(4); reduce_scatter_cross_device_store - <<>>(ptrs, sg_, self_sg_, rank_, size); + <<>>(ptrs, sg_, self_sg_, rank_, size); break; case 2: MAYBE_DISPATCH_1S_KERNEL(2); reduce_scatter_cross_device_store - <<>>(ptrs, sg_, self_sg_, rank_, size); + <<>>(ptrs, sg_, self_sg_, rank_, size); break; default: throw std::runtime_error("fused allreduce rmsnorm: unsupported world_size=" + std::to_string(world_size_)); }