From bf48d593571f4b06c9f088b75541d1af3ed7f56f Mon Sep 17 00:00:00 2001 From: katnoria Date: Sun, 9 Aug 2026 20:12:21 +0800 Subject: [PATCH] Fix QuantizedBlockLoader::load_safe bounds check comparing wrong field In the reduction_dim == 1 branch, `bi` (which indexes BROWS -- the output axis, N for qmm_t's weight loader) was compared against src_tile_dim.x. Every call site actually passes the valid-BROWS-count in src_tile_dim.y by convention (qmm_t_impl's `short2(BK, num_outs)`, affine_gather_qmm_rhs's `short2(k_remain, tgp_bn)`), so this was comparing a row index against a K-tile-width instead of against its own valid-row count. This has been a silent no-op in every kernel instantiated to date, because every existing qmm_t/gather_qmm_t config happens to have BROWS == BK == 32, so `bi < BROWS` already implied `bi < BK` and the check never fired -- correct output "by luck": the rows it should have zero-padded were already beyond `num_outs` and discarded by store_result_safe regardless of what stale/adjacent memory they read into threadgroup memory first. It stops being a no-op once BROWS is parametrized independently of BK (an upcoming change needs exactly that), at which point some still-valid output rows fall on the wrong side of the miscompared bound and get forced to zero instead of the safe/discarded ones, producing wrong results at partial-tile N boundaries. Fix: compare against src_tile_dim.y instead. Zero behavior change for every existing BROWS==BK instantiation (.x == BROWS there too, so the two fields were interchangeable in exactly that case); correctness bug fix once BROWS != BK. --- mlx/backend/metal/kernels/quantized.h | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 6d87dc770f..95c3e7da0a 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -644,7 +644,31 @@ struct QuantizedBlockLoader { return; } - if (reduction_dim == 1 && bi >= src_tile_dim.x) { + // `bi` indexes BROWS. When reduction_dim == 1, K is BCOLS (the + // reduction axis) and BROWS is the *output* axis (M for the x loader, + // N for the w loader in qmm_t_impl/affine_gather_qmm_rhs's transpose + // path) -- so the valid-count to guard `bi` against is the caller's + // BROWS bound. Every call site already passes that bound as + // src_tile_dim.y by convention (qmm_t_impl: `short2(BK, num_outs)`; + // affine_gather_qmm_rhs: `short2(k_remain, tgp_bn)`), but this line + // compared `bi` against src_tile_dim.x instead -- the BCOLS/K count, + // not the BROWS/output count. Comparing against the wrong field means + // this early-return under-fires: whenever the true BROWS bound (.y) is + // smaller than the BCOLS one (.x), some `bi` in [.y, .x) fall through + // to the real dequantize-and-load below and read live weight memory + // past the valid output boundary, instead of being zero-filled here. + // In every kernel instantiation shipped to date BROWS == BK == 32, so + // .x (== BK == BROWS) is never smaller than the true bound and the + // bug was unreachable; the resulting rows -- always >= the true bound + // either way -- are outside `num_outs`/`tgp_bn` and get discarded by + // the caller's store_result_safe regardless of what garbage they read, + // so output was correct by luck, not by this check doing anything. + // Fixing the compared field to .y makes the guard actually fire (and + // is a strict no-op for every existing BROWS==BK instantiation, since + // .x == BROWS there too -- the two fields were interchangeable exactly + // in that one case). It matters once BROWS is parametrized + // independently of BK, e.g. qmm_t_impl's new large-M tile below. + if (reduction_dim == 1 && bi >= src_tile_dim.y) { for (int i = 0; i < n_reads * pack_factor; i++) { dst[i] = T(0); }