From b6eec20260ea4bfcb75a2d98d3c2129e92c817f9 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Thu, 10 Jul 2025 16:28:50 -0700 Subject: [PATCH] Fix edge check in qmm_n QuantizedLoader (#2355) --- mlx/backend/metal/kernels/quantized.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index fea6f1460..0a40cec00 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -643,14 +643,14 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.y) { + 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.x) { + if (reduction_dim == 0 && bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); }