From 9cbff5ec1dd2688986222331074db27e20bc7615 Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Sun, 31 Mar 2024 19:15:44 -0700
Subject: [PATCH] Fix typo in qmm check (#940)

---
 mlx/backend/metal/kernels/quantized.metal | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlx/backend/metal/kernels/quantized.metal b/mlx/backend/metal/kernels/quantized.metal
index a61809b09..7f4b76725 100644
--- a/mlx/backend/metal/kernels/quantized.metal
+++ b/mlx/backend/metal/kernels/quantized.metal
@@ -729,7 +729,7 @@ template <typename T, const int BM, const int BK, const int BN, const int group_
           const device uint32_t * w_local = w + offset_row * N_w + offset_col;
           threadgroup T * Ws_local = Ws + offset_row * BN + offset_col * el_per_int;
 
-          if (y_row + offset_row < K) {
+          if (k + offset_row < K) {
             uint32_t wi = *w_local;
             T scale = scales_block[offset_row * groups_per_block + offset_col / (group_size / el_per_int)];
             T bias = biases_block[offset_row * groups_per_block + offset_col / (group_size / el_per_int)];