diff --git a/mlx/backend/cuda/rms_norm.cu b/mlx/backend/cuda/rms_norm.cu index 365eb57b1c..97f3e94a5e 100644 --- a/mlx/backend/cuda/rms_norm.cu +++ b/mlx/backend/cuda/rms_norm.cu @@ -51,7 +51,18 @@ struct BlockBroadcastReduce { } }; -template +// xs and ws stay in registers +// Each thread does N_CHUNKS_THREAD vectorised interleaved loads of width +// N_READS x, w, out in majority of cases should be 16 bytes aligned -> using +// unsafe_load_vector if not aligned -> fall back to load_vector we do it like +// this because of the register presure: load_vector allocates registers for +// loop load fall back and occupancy drops to 30-40% -> x2 slow down +template < + typename T, + int BLOCK_SIZE, + int N_CHUNKS_THREAD, + bool ALIGNED, + int N_READS = 8> __global__ void rms_norm_small( const T* x, const T* w, @@ -63,40 +74,61 @@ __global__ void rms_norm_small( auto grid = cg::this_grid(); auto block = cg::this_thread_block(); - using BlockReduceT = BlockBroadcastReduce; + using BlockReduceT = BlockBroadcastReduce; __shared__ typename BlockReduceT::TempStorage temp; - auto row = - (grid.block_rank() * block.dim_threads().y) + block.thread_index().y; + auto row = grid.block_rank(); if (row >= n_rows) { return; } x += row * axis_size; out += row * axis_size; - // Normalizer. + AlignedVector xn[N_CHUNKS_THREAD]; + float normalizer = 0; auto index = block.thread_index().x; - auto xn = load_vector(x, index, axis_size, T(0)); #pragma unroll - for (int i = 0; i < N_READS; ++i) { - float t = static_cast(xn[i]); - normalizer += t * t; + for (int i = 0; i < N_CHUNKS_THREAD; i++) { + int offset = BLOCK_SIZE * i + index; + if constexpr (ALIGNED) { + xn[i] = unsafe_load_vector(x, offset); + } else { + xn[i] = load_vector(x, offset, axis_size, T(0)); + } +#pragma unroll + for (int j = 0; j < N_READS; ++j) { + float t = static_cast(xn[i][j]); + normalizer += t * t; + } } - normalizer = BlockReduceT{block, temp}.Sum(normalizer); normalizer = rsqrt(normalizer / axis_size + eps); - // Outputs. - auto wn = load_vector(w, index, axis_size, w_stride, T(0)); #pragma unroll - for (int i = 0; i < N_READS; ++i) { - float y = static_cast(xn[i]) * normalizer; - xn[i] = wn[i] * static_cast(y); + for (int i = 0; i < N_CHUNKS_THREAD; i++) { + int offset = BLOCK_SIZE * i + index; + AlignedVector wn; + if constexpr (ALIGNED) { + wn = unsafe_load_vector(w, offset); + } else { + wn = load_vector(w, offset, axis_size, w_stride, T(0)); + } +#pragma unroll + for (int j = 0; j < N_READS; j++) { + float y = static_cast(xn[i][j]) * normalizer; + xn[i][j] = wn[j] * static_cast(y); + } + if constexpr (ALIGNED) { + unsafe_store_vector(out, offset, xn[i]); + } else { + store_vector(out, offset, xn[i], axis_size); + } } - store_vector(out, index, xn, axis_size); } +// TODO: load x to the shared memory and reload from shared memory after the +// reduction template __global__ void rms_norm( const T* x, @@ -318,6 +350,42 @@ void dispatch_group_dim(int axis_size, F&& f) { } } +template +void dispatch_chunks(int n_chunks, F&& f) { + auto block_size = std::integral_constant{}; + if (n_chunks <= 1) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 2) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 3) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 4) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 5) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 6) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 7) { + f(block_size, std::integral_constant{}); + } else { + f(block_size, std::integral_constant{}); + } +} + +template +void dispatch_num_chunks(int axis_size, F&& f) { + int nvec = (axis_size + N_READS - 1) / N_READS; + if (axis_size <= N_READS * 64) { + f(std::integral_constant{}, std::integral_constant{}); + } else if (nvec % 128 == 0 && nvec / 128 <= 8) { + dispatch_chunks<128>(nvec / 128, f); + } else if (nvec % 64 == 0 && nvec / 64 <= 8) { + dispatch_chunks<64>(nvec / 64, f); + } else { + dispatch_chunks<128>((nvec + 127) / 128, f); + } +} + // TODO: There are duplicate code with backend/metal/normalization.cpp void RMSNorm::eval_gpu( const std::vector& inputs, @@ -365,27 +433,38 @@ void RMSNorm::eval_gpu( dispatch_float_types(out.dtype(), "rms_norm", [&](auto type_tag) { using DataType = cuda_type_t; constexpr int N_READS = 16 / sizeof(DataType); - if (axis_size <= N_READS * 1024) { - dispatch_group_dim( - axis_size, [&](auto group_dim, auto n_groups, auto groups_per_block) { - constexpr int block_dim = n_groups() * group_dim(); - static_assert(block_dim <= 32 || groups_per_block() == 1); - auto kernel = - cu::rms_norm_small; - auto n_blocks = - (n_rows + groups_per_block() - 1) / groups_per_block(); - encoder.add_kernel_node( - kernel, - n_blocks, - {block_dim, groups_per_block()}, - gpu_ptr(x), - gpu_ptr(w), - gpu_ptr(out), - eps_, - axis_size, - n_rows, - w_stride); - }); + if (axis_size <= N_READS * 128 * 8) { + dispatch_num_chunks< + N_READS>(axis_size, [&](auto block_size, auto n_chunks) { + constexpr int BLOCK_SIZE = block_size(); + constexpr int N_CHUNKS = n_chunks(); + bool aligned = (axis_size == N_READS * BLOCK_SIZE * N_CHUNKS) && + (w_stride == 1) && + (reinterpret_cast(gpu_ptr(x)) % 16 == 0) && + (reinterpret_cast(gpu_ptr(w)) % 16 == 0) && + (reinterpret_cast(gpu_ptr(out)) % 16 == 0); + // MSVC can't find N_READS two lambda levels down, dispatch_bool would + // add that second level + auto kernel = aligned + ? cu::rms_norm_small + : cu::rms_norm_small< + DataType, + BLOCK_SIZE, + N_CHUNKS, + false, + N_READS>; + encoder.add_kernel_node( + kernel, + n_rows, + BLOCK_SIZE, + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(out), + eps_, + axis_size, + n_rows, + w_stride); + }); } else { auto kernel = cu::rms_norm; encoder.add_kernel_node(