From cca241a22ede73f7f4b331fa2caf287e3cd733dc Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Thu, 16 Jul 2026 17:28:33 +0200 Subject: [PATCH 1/3] rms norm forward --- mlx/backend/cuda/rms_norm.cu | 138 ++++++++++++++++++++++++++--------- 1 file changed, 102 insertions(+), 36 deletions(-) diff --git a/mlx/backend/cuda/rms_norm.cu b/mlx/backend/cuda/rms_norm.cu index 365eb57b1c..49ec76b110 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 stays 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,32 @@ void dispatch_group_dim(int axis_size, F&& f) { } } +template +void dispatch_num_chunks(int axis_size, F&& f) { + if (axis_size <= N_READS * 64) { + f(std::integral_constant{}, std::integral_constant{}); + return; + } + auto block_size = std::integral_constant{}; + if (axis_size <= N_READS * 128 * 1) { + f(block_size, std::integral_constant{}); + } else if (axis_size <= N_READS * 128 * 2) { + f(block_size, std::integral_constant{}); + } else if (axis_size <= N_READS * 128 * 3) { + f(block_size, std::integral_constant{}); + } else if (axis_size <= N_READS * 128 * 4) { + f(block_size, std::integral_constant{}); + } else if (axis_size <= N_READS * 128 * 5) { + f(block_size, std::integral_constant{}); + } else if (axis_size <= N_READS * 128 * 6) { + f(block_size, std::integral_constant{}); + } else if (axis_size <= N_READS * 128 * 7) { + f(block_size, std::integral_constant{}); + } else { + f(block_size, std::integral_constant{}); + } +} + // TODO: There are duplicate code with backend/metal/normalization.cpp void RMSNorm::eval_gpu( const std::vector& inputs, @@ -365,26 +423,34 @@ 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( + axis_size, [&](auto block_size, auto n_chunks) { + constexpr int BLOCK_SIZE = block_size(); + 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); + dispatch_bool(aligned, [&](auto aligned_tag) { + auto kernel = cu::rms_norm_small< + DataType, + BLOCK_SIZE, + n_chunks(), + aligned_tag.value, + 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; From 0b92da1ccf4511c586ab0b99e4c2c5bf838e6eb4 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Thu, 16 Jul 2026 23:12:44 +0200 Subject: [PATCH 2/3] make block size 128 or 64 depending on the axis size --- mlx/backend/cuda/rms_norm.cu | 47 ++++++++++++++++++++++-------------- 1 file changed, 29 insertions(+), 18 deletions(-) diff --git a/mlx/backend/cuda/rms_norm.cu b/mlx/backend/cuda/rms_norm.cu index 49ec76b110..59fa5da0ea 100644 --- a/mlx/backend/cuda/rms_norm.cu +++ b/mlx/backend/cuda/rms_norm.cu @@ -51,8 +51,8 @@ struct BlockBroadcastReduce { } }; -// xs and ws stays in registers -// each thread does N_CHUNKS_THREAD vectorised interleaved loads of width +// 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 @@ -350,32 +350,42 @@ void dispatch_group_dim(int axis_size, F&& f) { } } -template -void dispatch_num_chunks(int axis_size, F&& f) { - if (axis_size <= N_READS * 64) { - f(std::integral_constant{}, std::integral_constant{}); - return; - } - auto block_size = std::integral_constant{}; - if (axis_size <= N_READS * 128 * 1) { +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 (axis_size <= N_READS * 128 * 2) { + } else if (n_chunks <= 2) { f(block_size, std::integral_constant{}); - } else if (axis_size <= N_READS * 128 * 3) { + } else if (n_chunks <= 3) { f(block_size, std::integral_constant{}); - } else if (axis_size <= N_READS * 128 * 4) { + } else if (n_chunks <= 4) { f(block_size, std::integral_constant{}); - } else if (axis_size <= N_READS * 128 * 5) { + } else if (n_chunks <= 5) { f(block_size, std::integral_constant{}); - } else if (axis_size <= N_READS * 128 * 6) { + } else if (n_chunks <= 6) { f(block_size, std::integral_constant{}); - } else if (axis_size <= N_READS * 128 * 7) { + } 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, @@ -427,7 +437,8 @@ void RMSNorm::eval_gpu( dispatch_num_chunks( axis_size, [&](auto block_size, auto n_chunks) { constexpr int BLOCK_SIZE = block_size(); - bool aligned = (axis_size == N_READS * BLOCK_SIZE * n_chunks()) && + 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) && @@ -436,7 +447,7 @@ void RMSNorm::eval_gpu( auto kernel = cu::rms_norm_small< DataType, BLOCK_SIZE, - n_chunks(), + N_CHUNKS, aligned_tag.value, N_READS>; encoder.add_kernel_node( From e95ec089a80daf263c230ed91ffdcbd7a8d9a2ca Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Fri, 17 Jul 2026 15:07:20 +0200 Subject: [PATCH 3/3] fix windows build --- mlx/backend/cuda/rms_norm.cu | 52 +++++++++++++++++++----------------- 1 file changed, 27 insertions(+), 25 deletions(-) diff --git a/mlx/backend/cuda/rms_norm.cu b/mlx/backend/cuda/rms_norm.cu index 59fa5da0ea..97f3e94a5e 100644 --- a/mlx/backend/cuda/rms_norm.cu +++ b/mlx/backend/cuda/rms_norm.cu @@ -434,35 +434,37 @@ void RMSNorm::eval_gpu( using DataType = cuda_type_t; constexpr int N_READS = 16 / sizeof(DataType); if (axis_size <= N_READS * 128 * 8) { - dispatch_num_chunks( - 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); - dispatch_bool(aligned, [&](auto aligned_tag) { - auto kernel = cu::rms_norm_small< + 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, - aligned_tag.value, + 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); - }); - }); + 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(