diff --git a/mlx/backend/cuda/rms_norm.cu b/mlx/backend/cuda/rms_norm.cu
index 25ac064396..365eb57b1c 100644
--- a/mlx/backend/cuda/rms_norm.cu
+++ b/mlx/backend/cuda/rms_norm.cu
@@ -298,7 +298,7 @@ void dispatch_group_dim(int axis_size, F&& f) {
   } else if (axis_size <= n_per_thread * 32 * 2) {
     f(std::integral_constant<int, 32>{},
       std::integral_constant<int, 2>(),
-      std::integral_constant<int, 2>());
+      std::integral_constant<int, 1>());
   } else if (axis_size <= n_per_thread * 32 * 4) {
     f(std::integral_constant<int, 32>{},
       std::integral_constant<int, 4>(),
@@ -369,6 +369,7 @@ void RMSNorm::eval_gpu(
       dispatch_group_dim<N_READS>(
           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<DataType, block_dim, group_dim(), N_READS>;
             auto n_blocks =
@@ -479,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 b1c84c987d..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:
@@ -483,23 +483,23 @@ def test_rms_norm(self):
             mx.fast.rms_norm(x, mx.ones((4,)), 1e-5)
 
     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))