Skip to content

RMSNormVJP backward writes a full {n_rows, D} gw_temp intermediate #4293

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nastya236 merged 3 commits into from
Aug 19, 2026
Merged

RMSNormVJP backward writes a full {n_rows, D} gw_temp intermediate #4293

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
abd8d6d
fix quantize matrix multiplication floor issue
JasonHonKL Aug 14, 2026
7db01f1
optimize rms caluclation and performance.
JasonHonKL Aug 16, 2026
c332ad9
Merge branch 'main' into optimization-rms
JasonHonKL Aug 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
157 changes: 96 additions & 61 deletions mlx/backend/metal/kernels/rms_norm.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,21 +166,32 @@ template <typename T, int N_READS = RMS_N_READS>
constant float& eps,
constant uint& axis_size,
constant uint& w_stride,
constant uint& n_rows,
constant uint& rows_per_group,
uint gid [[threadgroup_position_in_grid]],
uint lid [[thread_position_in_threadgroup]],
uint simd_lane_id [[thread_index_in_simdgroup]],
uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
// Advance the input pointers
x += gid * size_t(axis_size) + lid * N_READS;
g += gid * size_t(axis_size) + lid * N_READS;
w += w_stride * lid * N_READS;
float thread_w[N_READS];
if (lid * N_READS + N_READS <= axis_size) {
for (int i = 0; i < N_READS; i++) {
thread_w[i] = w[w_stride * i];
}
} else {
for (int i = 0; i < N_READS; i++) {
thread_w[i] =
(lid * N_READS + i < axis_size) ? (float)w[w_stride * i] : 0;
}
}

// Allocate registers for the computation and accumulators
float thread_x[N_READS];
float thread_w[N_READS];
float thread_g[N_READS];
float sumx2 = 0;
float sumgwx = 0;
float gw_acc[N_READS];
for (int i = 0; i < N_READS; i++) {
gw_acc[i] = 0;
}

// Allocate shared memory to implement the reduction
constexpr int SIMD_SIZE = 32;
Expand All @@ -189,75 +200,99 @@ template <typename T, int N_READS = RMS_N_READS>
threadgroup float local_normalizer[1];
threadgroup float local_meangwx[1];

// Read and accumulate locally
if (lid * N_READS + N_READS <= axis_size) {
for (int i = 0; i < N_READS; i++) {
thread_x[i] = x[i];
thread_w[i] = w[w_stride * i];
thread_g[i] = g[i];
uint row_end = gid * rows_per_group + rows_per_group;
if (row_end > n_rows) {
row_end = n_rows;
}
for (uint row = gid * rows_per_group; row < row_end; ++row) {
const device T* x_in = x + size_t(row) * axis_size + lid * N_READS;
const device T* g_in = g + size_t(row) * axis_size + lid * N_READS;

sumx2 += thread_x[i] * thread_x[i];
sumgwx += thread_x[i] * thread_w[i] * thread_g[i];
}
} else {
for (int i = 0; i < N_READS; i++) {
if ((lid * N_READS + i) < axis_size) {
thread_x[i] = x[i];
thread_w[i] = w[w_stride * i];
thread_g[i] = g[i];
float sumx2 = 0;
float sumgwx = 0;

// Read and accumulate locally
if (lid * N_READS + N_READS <= axis_size) {
for (int i = 0; i < N_READS; i++) {
thread_x[i] = x_in[i];
thread_g[i] = g_in[i];

sumx2 += thread_x[i] * thread_x[i];
sumgwx += thread_x[i] * thread_w[i] * thread_g[i];
}
} else {
for (int i = 0; i < N_READS; i++) {
if ((lid * N_READS + i) < axis_size) {
thread_x[i] = x_in[i];
thread_g[i] = g_in[i];

sumx2 += thread_x[i] * thread_x[i];
sumgwx += thread_x[i] * thread_w[i] * thread_g[i];
}
}
}
}

// Accumulate across threads
sumx2 = simd_sum(sumx2);
sumgwx = simd_sum(sumgwx);
if (simd_group_id == 0) {
local_sumx2[simd_lane_id] = 0;
local_sumgwx[simd_lane_id] = 0;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (simd_lane_id == 0) {
local_sumx2[simd_group_id] = sumx2;
local_sumgwx[simd_group_id] = sumgwx;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (simd_group_id == 0) {
sumx2 = simd_sum(local_sumx2[simd_lane_id]);
sumgwx = simd_sum(local_sumgwx[simd_lane_id]);
// Accumulate across threads
sumx2 = simd_sum(sumx2);
sumgwx = simd_sum(sumgwx);
if (simd_group_id == 0) {
local_sumx2[simd_lane_id] = 0;
local_sumgwx[simd_lane_id] = 0;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (simd_lane_id == 0) {
local_meangwx[0] = sumgwx / axis_size;
local_normalizer[0] = metal::precise::rsqrt(sumx2 / axis_size + eps);
local_sumx2[simd_group_id] = sumx2;
local_sumgwx[simd_group_id] = sumgwx;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
float meangwx = local_meangwx[0];
float normalizer = local_normalizer[0];
float normalizer3 = normalizer * normalizer * normalizer;

// Write the outputs
gx += gid * size_t(axis_size) + lid * N_READS;
gw += gid * size_t(axis_size) + lid * N_READS;
if (lid * N_READS + N_READS <= axis_size) {
for (int i = 0; i < N_READS; i++) {
gx[i] = static_cast<T>(
thread_g[i] * thread_w[i] * normalizer -
thread_x[i] * meangwx * normalizer3);
if (has_w) {
gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
threadgroup_barrier(mem_flags::mem_threadgroup);
if (simd_group_id == 0) {
sumx2 = simd_sum(local_sumx2[simd_lane_id]);
sumgwx = simd_sum(local_sumgwx[simd_lane_id]);
if (simd_lane_id == 0) {
local_meangwx[0] = sumgwx / axis_size;
local_normalizer[0] = metal::precise::rsqrt(sumx2 / axis_size + eps);
}
}
} else {
for (int i = 0; i < N_READS; i++) {
if ((lid * N_READS + i) < axis_size) {
gx[i] = static_cast<T>(
threadgroup_barrier(mem_flags::mem_threadgroup);
float meangwx = local_meangwx[0];
float normalizer = local_normalizer[0];
float normalizer3 = normalizer * normalizer * normalizer;

// Write the outputs
device T* gx_out = gx + size_t(row) * axis_size + lid * N_READS;
if (lid * N_READS + N_READS <= axis_size) {
for (int i = 0; i < N_READS; i++) {
gx_out[i] = static_cast<T>(
thread_g[i] * thread_w[i] * normalizer -
thread_x[i] * meangwx * normalizer3);
if (has_w) {
gw[i] = static_cast<T>(thread_g[i] * thread_x[i] * normalizer);
gw_acc[i] += thread_g[i] * thread_x[i] * normalizer;
}
}
} else {
for (int i = 0; i < N_READS; i++) {
if ((lid * N_READS + i) < axis_size) {
gx_out[i] = static_cast<T>(
thread_g[i] * thread_w[i] * normalizer -
thread_x[i] * meangwx * normalizer3);
if (has_w) {
gw_acc[i] += thread_g[i] * thread_x[i] * normalizer;
}
}
}
}
}

if (has_w) {
gw += size_t(gid) * axis_size + lid * N_READS;
if (lid * N_READS + N_READS <= axis_size) {
for (int i = 0; i < N_READS; i++) {
gw[i] = static_cast<T>(gw_acc[i]);
}
} else {
for (int i = 0; i < N_READS; i++) {
if ((lid * N_READS + i) < axis_size) {
gw[i] = static_cast<T>(gw_acc[i]);
}
}
}
Expand Down
31 changes: 19 additions & 12 deletions mlx/backend/metal/normalization.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,11 +109,9 @@ void RMSNormVJP::eval_gpu(
array x_copy = contiguous_copy_gpu(x, s);
return {x_copy, true};
};
bool donate_g = inputs[2].is_donatable();
auto [x, copied] = check_input(inputs[0]);
const array& w = inputs[1];
auto [g, g_copied] = check_input(inputs[2]);
donate_g |= g_copied;
array& gx = outputs[0];
array& gw = outputs[1];

Expand All @@ -137,17 +135,22 @@ void RMSNormVJP::eval_gpu(
auto axis_size = static_cast<uint32_t>(x.shape().back());
int n_rows = x.data_size() / axis_size;

const int target_groups = 512;

@nastya236 nastya236 Aug 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just curious why 512?

@JasonHonKL JasonHonKL Aug 19, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried all two the power till 4096 and 512 perform the best from the bench.

uint32_t rows_per_group = 1;
int n_groups = n_rows;
if (axis_size <= RMS_LOOPED_LIMIT) {
rows_per_group = (n_rows + target_groups - 1) / target_groups;
n_groups = (n_rows + rows_per_group - 1) / rows_per_group;
}

// Allocate the gradient accumulator gw and a temporary to store the
// gradients before they are accumulated.
array gw_temp =
(has_w) ? array({n_rows, x.shape().back()}, gw.dtype(), nullptr, {}) : w;
array gw_temp = (has_w)
? array({n_groups, x.shape().back()}, gw.dtype(), nullptr, {})
: w;
if (has_w) {
if (!g_in_gx && donate_g) {

@nastya236 nastya236 Aug 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice catch. Based on my understanding g_temp can reuse a buffer in very rare case: if x is non row-contiguous and g is non row-contiguous. I think it is fair to drop it here for looped case as well.

@JasonHonKL JasonHonKL Aug 19, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks

gw_temp.copy_shared_buffer(g);
} else {
gw_temp.set_data(allocator::malloc(gw_temp.nbytes()));
compute_encoder.add_temporary(gw_temp);
}
gw_temp.set_data(allocator::malloc(gw_temp.nbytes()));
compute_encoder.add_temporary(gw_temp);
}
gw.set_data(allocator::malloc(gw.nbytes()));

Expand All @@ -174,7 +177,7 @@ void RMSNormVJP::eval_gpu(
size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
size_t threadgroup_size = simd_size * simds_needed;
assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup());
size_t n_threads = n_rows * threadgroup_size;
size_t n_threads = n_groups * threadgroup_size;
grid_dims = MTL::Size(n_threads, 1, 1);
group_dims = MTL::Size(threadgroup_size, 1, 1);
} else {
Expand All @@ -194,12 +197,16 @@ void RMSNormVJP::eval_gpu(
compute_encoder.set_bytes(eps_, 5);
compute_encoder.set_bytes(axis_size, 6);
compute_encoder.set_bytes(w_stride, 7);
if (axis_size <= looped_limit) {
compute_encoder.set_bytes(static_cast<uint32_t>(n_rows), 8);
compute_encoder.set_bytes(rows_per_group, 9);
}
compute_encoder.dispatch_threads(grid_dims, group_dims);
}

if (has_w) {
ReductionPlan plan(
ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
ReductionOpType::ContiguousStridedReduce, {n_groups}, {axis_size});
strided_reduce_general_dispatch(
gw_temp, gw, "sum", plan, {0}, compute_encoder, d, s);
}
Expand Down
Loading