From e486ac44f66bf22e97a3935a8f654f16a9f1022c Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Thu, 2 Jul 2026 18:18:42 -0700 Subject: [PATCH 1/2] Fix CUDA RMSNorm small-row dispatch CUDA RMSNorm used a launch shape that packed multiple rows into one block even when each row required more than one reduction group. BlockBroadcastReduce is block-wide, so those rows mixed partial reductions and produced incorrect results for small dimensions such as fp32 D=256 and bf16/fp16 D=512. Fix the affected dispatch case by launching one row per block when two reduction groups are needed. Add CUDA regression coverage for forward and VJP fast RMSNorm on the affected small multi-row shapes. --- mlx/backend/cuda/rms_norm.cu | 5 +++- python/tests/test_fast.py | 46 ++++++++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/mlx/backend/cuda/rms_norm.cu b/mlx/backend/cuda/rms_norm.cu index 25ac064396..ce981b07e9 100644 --- a/mlx/backend/cuda/rms_norm.cu +++ b/mlx/backend/cuda/rms_norm.cu @@ -296,9 +296,12 @@ void dispatch_group_dim(int axis_size, F&& f) { std::integral_constant(), std::integral_constant()); } else if (axis_size <= n_per_thread * 32 * 2) { + // BlockBroadcastReduce is block-wide. Rows that need more than one + // reduction group cannot share a block without mixing their partial + // reductions. f(std::integral_constant{}, std::integral_constant(), - std::integral_constant()); + std::integral_constant()); } else if (axis_size <= n_per_thread * 32 * 4) { f(std::integral_constant{}, std::integral_constant(), diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index b1c84c987d..e1511304cf 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -482,6 +482,52 @@ def test_rms_norm(self): x = mx.random.uniform(shape=(1, 5)) mx.fast.rms_norm(x, mx.ones((4,)), 1e-5) + @unittest.skipIf(not mx.cuda.is_available(), "CUDA is not available") + def test_rms_norm_cuda_small_multi_row(self): + default_device = mx.default_device() + mx.set_default_device(mx.gpu) + try: + cases = [ + (mx.float32, 2, 256, 1e-4), + (mx.float16, 2, 512, 2e-2), + (mx.bfloat16, 2, 512, 2e-2), + ] + for dtype, rows, dims, tolerance in cases: + base = mx.reshape( + mx.arange(rows * dims, dtype=mx.float32), (rows, dims) + ) + scales = mx.reshape(mx.arange(1, rows + 1, dtype=mx.float32), (rows, 1)) + x = ((base % 17 - 8) * scales).astype(dtype) + weight = (1 + mx.arange(dims, dtype=mx.float32) / dims).astype(dtype) + + rx = rms_norm(x, weight, 1e-5) + rx_fast = mx.fast.rms_norm(x, weight, 1e-5) + self.assertLess(mx.abs(rx - rx_fast).max(), tolerance) + finally: + mx.set_default_device(default_device) + + @unittest.skipIf(not mx.cuda.is_available(), "CUDA is not available") + def test_rms_norm_grad_cuda_small_multi_row(self): + default_device = mx.default_device() + mx.set_default_device(mx.gpu) + try: + rows, dims, eps = 2, 256, 1e-5 + base = mx.reshape(mx.arange(rows * dims, dtype=mx.float32), (rows, dims)) + scales = mx.reshape(mx.arange(1, rows + 1, dtype=mx.float32), (rows, 1)) + x = (base % 17 - 8) * scales + w = 1 + mx.arange(dims, dtype=mx.float32) / dims + y = mx.cos(base / 7) + + f1 = lambda x, w, y: (rms_norm(x, w, eps) * y).sum() + f2 = lambda x, w, y: (mx.fast.rms_norm(x, w, eps) * y).sum() + + gx1, gw1 = mx.grad(f1, argnums=(0, 1))(x, w, y) + gx2, gw2 = mx.grad(f2, argnums=(0, 1))(x, w, y) + self.assertLess(mx.abs(gx1 - gx2).max(), 1e-4) + self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 1e-4) + finally: + mx.set_default_device(default_device) + def test_rms_norm_grad(self): D = 32 eps = 1e-5 From 98323e82f2d99e16b4fc1bbbe5430ac76e265198 Mon Sep 17 00:00:00 2001 From: Cheng Date: Thu, 2 Jul 2026 23:31:03 -0700 Subject: [PATCH 2/2] Simplify tests --- mlx/backend/cuda/rms_norm.cu | 5 ++- python/tests/test_fast.py | 70 +++++++----------------------------- 2 files changed, 14 insertions(+), 61 deletions(-) diff --git a/mlx/backend/cuda/rms_norm.cu b/mlx/backend/cuda/rms_norm.cu index ce981b07e9..365eb57b1c 100644 --- a/mlx/backend/cuda/rms_norm.cu +++ b/mlx/backend/cuda/rms_norm.cu @@ -296,9 +296,6 @@ void dispatch_group_dim(int axis_size, F&& f) { std::integral_constant(), std::integral_constant()); } else if (axis_size <= n_per_thread * 32 * 2) { - // BlockBroadcastReduce is block-wide. Rows that need more than one - // reduction group cannot share a block without mixing their partial - // reductions. f(std::integral_constant{}, std::integral_constant(), std::integral_constant()); @@ -372,6 +369,7 @@ void RMSNorm::eval_gpu( 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 = @@ -482,6 +480,7 @@ void RMSNormVJP::eval_gpu( axis_size, [&](auto group_dim, auto n_groups, auto groups_per_block) { constexpr int block_dim = group_dim() * n_groups(); + static_assert(block_dim <= 32 || groups_per_block() == 1); auto kernel = cu::rms_norm_vjp_small< DataType, has_w_constant.value, diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index e1511304cf..411e52b79e 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -428,7 +428,7 @@ def test_rms_norm(self): dtypes = [mx.float32, mx.float16, mx.bfloat16] epss = [1e-3, 1e-5] - dimss = [31, 32, 33] + dimss = [31, 32, 33, 256, 512] defaults = (mx.float32, 1e-5, 32) for dtype in dtypes: @@ -482,70 +482,24 @@ def test_rms_norm(self): x = mx.random.uniform(shape=(1, 5)) mx.fast.rms_norm(x, mx.ones((4,)), 1e-5) - @unittest.skipIf(not mx.cuda.is_available(), "CUDA is not available") - def test_rms_norm_cuda_small_multi_row(self): - default_device = mx.default_device() - mx.set_default_device(mx.gpu) - try: - cases = [ - (mx.float32, 2, 256, 1e-4), - (mx.float16, 2, 512, 2e-2), - (mx.bfloat16, 2, 512, 2e-2), - ] - for dtype, rows, dims, tolerance in cases: - base = mx.reshape( - mx.arange(rows * dims, dtype=mx.float32), (rows, dims) - ) - scales = mx.reshape(mx.arange(1, rows + 1, dtype=mx.float32), (rows, 1)) - x = ((base % 17 - 8) * scales).astype(dtype) - weight = (1 + mx.arange(dims, dtype=mx.float32) / dims).astype(dtype) - - rx = rms_norm(x, weight, 1e-5) - rx_fast = mx.fast.rms_norm(x, weight, 1e-5) - self.assertLess(mx.abs(rx - rx_fast).max(), tolerance) - finally: - mx.set_default_device(default_device) - - @unittest.skipIf(not mx.cuda.is_available(), "CUDA is not available") - def test_rms_norm_grad_cuda_small_multi_row(self): - default_device = mx.default_device() - mx.set_default_device(mx.gpu) - try: - rows, dims, eps = 2, 256, 1e-5 - base = mx.reshape(mx.arange(rows * dims, dtype=mx.float32), (rows, dims)) - scales = mx.reshape(mx.arange(1, rows + 1, dtype=mx.float32), (rows, 1)) - x = (base % 17 - 8) * scales - w = 1 + mx.arange(dims, dtype=mx.float32) / dims - y = mx.cos(base / 7) - - f1 = lambda x, w, y: (rms_norm(x, w, eps) * y).sum() - f2 = lambda x, w, y: (mx.fast.rms_norm(x, w, eps) * y).sum() - - gx1, gw1 = mx.grad(f1, argnums=(0, 1))(x, w, y) - gx2, gw2 = mx.grad(f2, argnums=(0, 1))(x, w, y) - self.assertLess(mx.abs(gx1 - gx2).max(), 1e-4) - self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 1e-4) - finally: - mx.set_default_device(default_device) - def test_rms_norm_grad(self): - D = 32 eps = 1e-5 f1 = lambda x, w, y: (rms_norm(x, w, eps) * y).sum() f2 = lambda x, w, y: (mx.fast.rms_norm(x, w, eps) * y).sum() f3 = lambda x, y: (rms_norm(x, mx.ones((x.shape[-1],)), eps) * y).sum() f4 = lambda x, y: (mx.fast.rms_norm(x, None, eps) * y).sum() - x = mx.random.uniform(shape=(8, 100, D)) - w = mx.random.uniform(shape=(D,)) - y = mx.random.uniform(shape=(8, 100, D)) - gx1, gw1 = mx.grad(f1, argnums=(0, 1))(x, w, y) - gx2, gw2 = mx.grad(f2, argnums=(0, 1))(x, w, y) - self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5) - self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 1e-5) - gx1 = mx.grad(f3, argnums=(0,))(x, y) - gx2 = mx.grad(f4, argnums=(0,))(x, y) - self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5) + for D in [32, 256]: + x = mx.random.uniform(shape=(8, 100, D)) + w = mx.random.uniform(shape=(D,)) + y = mx.random.uniform(shape=(8, 100, D)) + gx1, gw1 = mx.grad(f1, argnums=(0, 1))(x, w, y) + gx2, gw2 = mx.grad(f2, argnums=(0, 1))(x, w, y) + self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5) + self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 1e-5) + gx1 = mx.grad(f3, argnums=(0,))(x, y) + gx2 = mx.grad(f4, argnums=(0,))(x, y) + self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5) D = 8192 x = mx.random.uniform(shape=(2, 2, D))