Fix fp quantized matmul corruption when the quantized dim is not a multiple of 32

mlx/backend/metal/kernels/fp_quantized.h

@@ -163,9 +163,14 @@ struct QuantizedBlockLoader {
   MLX_MTL_CONST short group_steps = group_size < BCOLS ? 1 : group_size / BCOLS;
   MLX_MTL_CONST short scale_step = group_size < BCOLS ? BCOLS / group_size : 1;
 
+  MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0;
+
   static_assert(
       (n_reads * pack_factor) <= group_size,
       "The number of reads per thread must be less than the group size.");
+  static_assert(
+      group_size % (n_reads * pack_factor) == 0,
+      "The group size must be a multiple of the columns read per thread.");
 
   const int src_ld;
   const int tile_stride;
@@ -220,14 +225,7 @@ struct QuantizedBlockLoader {
       return;
     }
 
-    if (reduction_dim == 1 && bi >= src_tile_dim.x) {
-      for (int i = 0; i < n_reads * pack_factor; i++) {
-        dst[i] = T(0);
-      }
-      return;
-    }
-
-    if (reduction_dim == 0 && bi >= src_tile_dim.y) {
+    if (bi >= src_tile_dim.y) {
       for (int i = 0; i < n_reads * pack_factor; i++) {
         dst[i] = T(0);
       }
@@ -848,9 +846,12 @@ METAL_FUNC void fp_qmm_t_impl(
   loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid);
   mma_t mma_op(simd_gid, simd_lid);
 
+  constexpr bool partial_k = group_size % BK != 0;
+  const int k_blocks = K_eff / BK;
+  const short num_k = partial_k ? short(K_eff - k_blocks * BK) : short(0);
   if (num_els < BM) {
     if (!aligned_N && num_outs < BN) {
-      for (int k = 0; k < K_eff; k += BK) {
+      for (int k = 0; k < k_blocks; k++) {
         threadgroup_barrier(mem_flags::mem_threadgroup);
         loader_x.load_safe(short2(BK, num_els));
         loader_w.load_safe(short2(BK, num_outs));
@@ -860,7 +861,7 @@ METAL_FUNC void fp_qmm_t_impl(
         loader_w.next();
       }
     } else {
-      for (int k = 0; k < K_eff; k += BK) {
+      for (int k = 0; k < k_blocks; k++) {
         threadgroup_barrier(mem_flags::mem_threadgroup);
         loader_x.load_safe(short2(BK, num_els));
         loader_w.load_unsafe();
@@ -870,9 +871,18 @@ METAL_FUNC void fp_qmm_t_impl(
         loader_w.next();
       }
     }
+    if constexpr (partial_k) {
+      if (num_k > 0) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_safe(short2(num_k, num_els));
+        loader_w.load_safe(short2(num_k, num_outs));
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+      }
+    }
   } else {
     if (!aligned_N && num_outs < BN) {
-      for (int k = 0; k < K_eff; k += BK) {
+      for (int k = 0; k < k_blocks; k++) {
         threadgroup_barrier(mem_flags::mem_threadgroup);
         loader_x.load_unsafe();
         loader_w.load_safe(short2(BK, num_outs));
@@ -882,7 +892,7 @@ METAL_FUNC void fp_qmm_t_impl(
         loader_w.next();
       }
     } else {
-      for (int k = 0; k < K_eff; k += BK) {
+      for (int k = 0; k < k_blocks; k++) {
         threadgroup_barrier(mem_flags::mem_threadgroup);
         loader_x.load_unsafe();
         loader_w.load_unsafe();
@@ -893,6 +903,15 @@ METAL_FUNC void fp_qmm_t_impl(
         loader_w.next();
       }
     }
+    if constexpr (partial_k) {
+      if (num_k > 0) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_safe(short2(num_k, BM));
+        loader_w.load_safe(short2(num_k, num_outs));
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+      }
+    }
   }
 
   // Store results to device memory
@@ -969,13 +988,20 @@ METAL_FUNC void fp_qmm_n_impl(
   loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid);
   mma_t mma_op(simd_gid, simd_lid);
 
+  constexpr bool partial_n = group_size % BN != 0;
+  const short num_outs = partial_n ? short(min(BN, N - y_col)) : short(BN);
+
   if (num_els < BM) {
     if ((K % BK) != 0) {
       const int k_blocks = K / BK;
       for (int k = 0; k < k_blocks; k++) {
         threadgroup_barrier(mem_flags::mem_threadgroup);
         loader_x.load_safe(short2(BK, num_els));
-        loader_w.load_unsafe();
+        if constexpr (partial_n) {
+          loader_w.load_safe(short2(num_outs, BK));
+        } else {
+          loader_w.load_unsafe();
+        }
         threadgroup_barrier(mem_flags::mem_threadgroup);
         mma_op.mma(Xs, Ws);
         loader_x.next();
@@ -984,9 +1010,19 @@ METAL_FUNC void fp_qmm_n_impl(
       const short num_k = K - k_blocks * BK;
       threadgroup_barrier(mem_flags::mem_threadgroup);
       loader_x.load_safe(short2(num_k, num_els));
-      loader_w.load_safe(short2(BN, num_k));
+      loader_w.load_safe(short2(num_outs, num_k));
       threadgroup_barrier(mem_flags::mem_threadgroup);
       mma_op.mma(Xs, Ws);
+    } else if (num_outs < BN) {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_safe(short2(BK, num_els));
+        loader_w.load_safe(short2(num_outs, BK));
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
     } else {
       for (int k = 0; k < K; k += BK) {
         threadgroup_barrier(mem_flags::mem_threadgroup);
@@ -1004,7 +1040,11 @@ METAL_FUNC void fp_qmm_n_impl(
       for (int k = 0; k < k_blocks; k++) {
         threadgroup_barrier(mem_flags::mem_threadgroup);
         loader_x.load_unsafe();
-        loader_w.load_unsafe();
+        if constexpr (partial_n) {
+          loader_w.load_safe(short2(num_outs, BK));
+        } else {
+          loader_w.load_unsafe();
+        }
         threadgroup_barrier(mem_flags::mem_threadgroup);
         mma_op.mma(Xs, Ws);
         loader_x.next();
@@ -1013,9 +1053,19 @@ METAL_FUNC void fp_qmm_n_impl(
       const short num_k = K - k_blocks * BK;
       threadgroup_barrier(mem_flags::mem_threadgroup);
       loader_x.load_safe(short2(num_k, BM));
-      loader_w.load_safe(short2(BN, num_k));
+      loader_w.load_safe(short2(num_outs, num_k));
       threadgroup_barrier(mem_flags::mem_threadgroup);
       mma_op.mma(Xs, Ws);
+    } else if (num_outs < BN) {
+      for (int k = 0; k < K; k += BK) {
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        loader_x.load_unsafe();
+        loader_w.load_safe(short2(num_outs, BK));
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        mma_op.mma(Xs, Ws);
+        loader_x.next();
+        loader_w.next();
+      }
     } else {
       for (int k = 0; k < K; k += BK) {
         threadgroup_barrier(mem_flags::mem_threadgroup);
@@ -1031,8 +1081,8 @@ METAL_FUNC void fp_qmm_n_impl(
 
   // Store results to device memory
   threadgroup_barrier(mem_flags::mem_threadgroup);
-  if (num_els < BM) {
-    mma_op.store_result_safe(y, N, short2(BN, num_els));
+  if (num_els < BM || num_outs < BN) {
+    mma_op.store_result_safe(y, N, short2(num_outs, num_els));
   } else {
     mma_op.store_result(y, N);
   }

mlx/backend/metal/kernels/quantized.h

@@ -589,6 +589,12 @@ struct QuantizedBlockLoader {
       (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size;
   MLX_MTL_CONST short group_steps = group_size / BCOLS;
 
+  MLX_MTL_CONST bool partial_cols = group_size % BCOLS != 0;
+
+  static_assert(
+      group_size % (n_reads * pack_factor) == 0,
+      "The group size must be a multiple of the columns read per thread.");
+
   const int src_ld;
   const int tile_stride;
   short group_step_cnt;
@@ -644,14 +650,7 @@ struct QuantizedBlockLoader {
       return;
     }
 
-    if (reduction_dim == 1 && bi >= src_tile_dim.x) {
-      for (int i = 0; i < n_reads * pack_factor; i++) {
-        dst[i] = T(0);
-      }
-      return;
-    }
-
-    if (reduction_dim == 0 && bi >= src_tile_dim.y) {
+    if (bi >= src_tile_dim.y) {
       for (int i = 0; i < n_reads * pack_factor; i++) {
         dst[i] = T(0);
       }

python/tests/test_quantized.py

@@ -715,6 +715,167 @@ def test_fp_qvm(self):
                     self.assertEqual(y_q.shape, y_hat.shape)
                     self.assertLess((y_q - y_hat).abs().max(), 2e-3)
 
+    FP_TAIL_TOLS = {mx.float32: 1e-3, mx.float16: 5e-3, mx.bfloat16: 4e-2}
+
+    def _fp_tail_dtypes(self):
+        # The half types only run on the GPU
+        if mx.default_device() == mx.gpu:
+            return [mx.float32, mx.float16, mx.bfloat16]
+        return [mx.float32]
+
+    def test_fp_qmm_non_multiple_of_32(self):
+        key = mx.random.key(0)
+        k1, k2 = mx.random.split(key)
+
+        def check(M, K, N, transpose, mode, dtype):
+            with self.subTest(
+                M=M, K=K, N=N, transpose=transpose, mode=mode, dtype=dtype
+            ):
+                w_shape = (N, K) if transpose else (K, N)
+                x = (mx.random.normal(shape=(M, K), key=k1) / K**0.5).astype(dtype)
+                w = mx.random.normal(shape=w_shape, key=k2).astype(dtype)
+                w_q, scales = mx.quantize(w, mode=mode)
+                w_hat = mx.dequantize(w_q, scales, mode=mode)
+
+                y_q = mx.quantized_matmul(
+                    x, w_q, scales, transpose=transpose, mode=mode
+                )
+                # An unbounded store of the N tail spills onto the next output
+                # row and races the threadgroup that owns it, so two identical
+                # runs must agree bit for bit.
+                y_q2 = mx.quantized_matmul(
+                    x, w_q, scales, transpose=transpose, mode=mode
+                )
+                y_hat = x @ (mx.swapaxes(w_hat, -1, -2) if transpose else w_hat)
+
+                self.assertEqual(y_q.shape, y_hat.shape)
+                self.assertTrue(mx.array_equal(y_q, y_q2))
+                self.assertLess((y_q - y_hat).abs().max(), self.FP_TAIL_TOLS[dtype])
+
+        # (M, K, N, transpose). M is kept clear of the qmv batch limit
+        # (get_qmv_batch_limit caps at 32) so these route to the tiled kernel.
+        nvfp4_shapes = [
+            (50, 1040, 128, True),  # K % 32 == 16
+            (33, 528, 64, True),  # K % 32 == 16, M just over one BM tile
+            (50, 512, 1040, False),  # N % 32 == 16
+            (50, 1040, 1040, False),  # both dims carry a tail
+            (50, 16, 128, True),  # K is a single partial tile
+        ]
+        control_shapes = [
+            (50, 1056, 128, True),  # K % 32 == 0
+            (50, 512, 1056, False),  # N % 32 == 0
+        ]
+
+        for dtype in self._fp_tail_dtypes():
+            for M, K, N, transpose in nvfp4_shapes:
+                check(M, K, N, transpose, "nvfp4", dtype)
+            for mode in ["mxfp4", "mxfp8"]:
+                for M, K, N, transpose in control_shapes:
+                    check(M, K, N, transpose, mode, dtype)
+
+    def test_fp_qmm_non_multiple_of_32_vjp(self):
+        # The backward of quantized_matmul runs the same kernels with the
+        # transpose flipped, so cover both directions at a tail shape.
+        key = mx.random.key(0)
+        k1, k2, k3 = mx.random.split(key, 3)
+        M, K, N = 50, 1040, 128
+
+        for dtype in self._fp_tail_dtypes():
+            for transpose in [True, False]:
+                with self.subTest(transpose=transpose, dtype=dtype):
+                    w_shape = (N, K) if transpose else (K, N)
+                    x = (mx.random.normal(shape=(M, K), key=k1) / K**0.5).astype(dtype)
+                    w = mx.random.normal(shape=w_shape, key=k2).astype(dtype)
+                    w_q, scales = mx.quantize(w, mode="nvfp4")
+                    w_hat = mx.dequantize(w_q, scales, mode="nvfp4")
+                    cot = (mx.random.normal(shape=(M, N), key=k3) / N**0.5).astype(
+                        dtype
+                    )
+
+                    def f(x):
+                        return mx.quantized_matmul(
+                            x, w_q, scales, transpose=transpose, mode="nvfp4"
+                        )
+
+                    (out,), (grad,) = mx.vjp(f, (x,), (cot,))
+                    tol = self.FP_TAIL_TOLS[dtype]
+
+                    # The forward runs the tail kernel for this transpose; the
+                    # backward runs it for the other one, we might as well check both.
+                    out_hat = x @ (mx.swapaxes(w_hat, -1, -2) if transpose else w_hat)
+                    self.assertEqual(out.shape, out_hat.shape)
+                    self.assertLess((out - out_hat).abs().max(), tol)
+
+                    grad_hat = cot @ (
+                        w_hat if transpose else mx.swapaxes(w_hat, -1, -2)
+                    )
+                    self.assertEqual(grad.shape, grad_hat.shape)
+                    self.assertLess((grad - grad_hat).abs().max(), tol)
+
+    def test_fp_gather_qmm_non_multiple_of_32(self):
+        # Both gather_qmm kernels: the sorted-indices (MoE prefill) kernel and
+        # the tiled gather kernel.
+        key = mx.random.key(0)
+        k1, k2, k3 = mx.random.split(key, 3)
+        mode = "nvfp4"
+        E = 4
+
+        def check(K, N, transpose, sorted_indices, dtype):
+            with self.subTest(
+                K=K,
+                N=N,
+                transpose=transpose,
+                sorted_indices=sorted_indices,
+                dtype=dtype,
+            ):
+                w_shape = (E, N, K) if transpose else (E, K, N)
+                w = mx.random.normal(shape=w_shape, key=k2).astype(dtype)
+                w_q, scales = mx.quantize(w, mode=mode)
+                w_hat = mx.dequantize(w_q, scales, mode=mode)
+
+                if sorted_indices:
+                    T = 64
+                    x = (mx.random.normal(shape=(T, 1, K), key=k1) / K**0.5).astype(
+                        dtype
+                    )
+                    rhs = mx.sort(mx.random.randint(0, E, shape=(T,), key=k3)).astype(
+                        mx.uint32
+                    )
+                    lhs = None
+                else:
+                    B, M = 8, 50
+                    x = (mx.random.normal(shape=(B, M, K), key=k1) / K**0.5).astype(
+                        dtype
+                    )
+                    rhs = mx.random.randint(0, E, shape=(B,), key=k3).astype(mx.uint32)
+                    lhs = mx.arange(B, dtype=mx.uint32)
+
+                y_q = mx.gather_qmm(
+                    x,
+                    w_q,
+                    scales,
+                    lhs_indices=lhs,
+                    rhs_indices=rhs,
+                    transpose=transpose,
+                    mode=mode,
+                    sorted_indices=sorted_indices,
+                )
+                w_sel = mx.take(w_hat, rhs, axis=0)
+                y_hat = x @ (mx.swapaxes(w_sel, -1, -2) if transpose else w_sel)
+
+                self.assertEqual(y_q.shape, y_hat.shape)
+                self.assertLess((y_q - y_hat).abs().max(), self.FP_TAIL_TOLS[dtype])
+
+        # (K, N, transpose): transpose puts the tail on K, else on N.
+        shapes = [
+            (1040, 128, True),  # K % 32 == 16
+            (512, 1040, False),  # N % 32 == 16
+        ]
+        for dtype in self._fp_tail_dtypes():
+            for K, N, transpose in shapes:
+                for sorted_indices in [True, False]:
+                    check(K, N, transpose, sorted_indices, dtype)
+
     def test_mode_error_cases(self):
         w = mx.random.normal(shape=(256, 256))
         x = mx.random.normal(shape=(1, 256))