From 57eb5240de7a4d53851028d8f57d30e7d1cd10a8 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 16 Jul 2026 14:33:38 -0700 Subject: [PATCH] Align split-K partitions to the qmm K-tile (BK), fixing nvfp4 qmm_splitk caps split_k only by the quantization group count (K / group_size), not by the kernel K-tile width BK (32). For nvfp4 (group_size=16) this yields a per-partition K of 16 < BK; fp_qmm_t_splitk then reads a full BK-wide K-tile with no K bound, spilling past the partition into the next group's weights and fp8 scales. The result is corrupted (non-uniform ~2x error) with NaN/inf on the last partition. Affine (group_size >= 32) is unaffected. Require each K partition to be a whole number of BK-wide tiles as well as whole quantization groups by aligning split_k to max(group_size, BK). Only changes behavior for group_size < 32. --- mlx/backend/metal/quantized.cpp | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 62d48714e0..94c563070a 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -884,11 +884,15 @@ void qmm_splitk( int current_tgs = n_tiles * m_tiles; int split_k = std::max(1, 512 / current_tgs); - // Cap split_k by the number of quantization groups - split_k = std::min(split_k, K / group_size); - - // Ensure K divides evenly by split_k * group_size - while (split_k > 1 && (K % (split_k * group_size) != 0)) { + // Each K partition must be a whole number of BK-wide (32) K-tiles as well as + // whole quantization groups. The qmm_t_splitk kernels tile K by BK=32 and do + // not bound the K dimension, so a partition smaller than BK (e.g. nvfp4's + // group_size=16) would over-read into the next group's weights/scales. + int k_align = group_size > 32 ? group_size : 32; + split_k = std::min(split_k, K / k_align); + + // Ensure K divides evenly by split_k * k_align + while (split_k > 1 && (K % (split_k * k_align) != 0)) { split_k--; } if (split_k <= 1) {