Skip to content

Fix NAX tile bugs and make the gather_qmm block tile tunable #4015

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
vvsotnikov wants to merge 3 commits into from
Closed

Fix NAX tile bugs and make the gather_qmm block tile tunable #4015

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
870ab65
refactor(metal): generalize the NAX quantized block loader
vvsotnikov Aug 5, 2026
5ce038d
fix(metal): correct NAX tile matmul with one fragment along N
vvsotnikov Aug 5, 2026
b462b0f
feat(metal): make the NAX gather_qmm block tile tunable
vvsotnikov Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
229 changes: 64 additions & 165 deletions mlx/backend/metal/kernels/quantized_nax.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -563,6 +563,15 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) {
}
}

// The loader indexes the scales in one of three ways, by where a thread's
// n_reads packed words fall inside a quantization group:
// 1. group_size >= BCOLS: a tile row sits in one group, so next() advances
// the scales every group_steps = group_size / BCOLS tiles.
// 2. BCOLS > group_size, read inside one group: a tile row covers
// n_groups = BCOLS / group_size groups, group_id picks one, and each
// tile advances the scales by n_groups.
// 3. The read covers whole groups: the thread walks groups_per_read of
// them, taking a scale and bias for each.
template <
typename T,
short BROWS,
Expand All @@ -574,11 +583,8 @@ template <
short bits>
struct QuantizedBlockLoader {
static_assert(
BCOLS <= group_size,
"The group size should be larger than the columns");
static_assert(
group_size % BCOLS == 0,
"The group size should be divisible by the columns");
BCOLS % group_size == 0 || group_size % BCOLS == 0,
"The tile columns and the group size must divide one another");
static_assert(
bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
bits == 8,
Expand All @@ -589,7 +595,28 @@ struct QuantizedBlockLoader {
MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor;
MLX_MTL_CONST short n_reads =
(BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size;
MLX_MTL_CONST short group_steps = group_size / BCOLS;

// Groups covered by one tile row.
MLX_MTL_CONST short n_groups = (BCOLS > group_size) ? BCOLS / group_size : 1;
// Tiles that share one group.
MLX_MTL_CONST short group_steps =
(group_size > BCOLS) ? group_size / BCOLS : 1;

// Values one thread reads.
MLX_MTL_CONST short n_elems = n_reads * pack_factor;
// Groups per read, and packed reads in each.
MLX_MTL_CONST short groups_per_read =
(n_elems > group_size) ? n_elems / group_size : 1;
MLX_MTL_CONST short reads_per_group = n_reads / groups_per_read;

static_assert(
BCOLS_PACKED % n_reads == 0,
"The packed columns of the block tile must be a multiple of the "
"reads of one thread.");
static_assert(
n_elems % group_size == 0 || group_size % n_elems == 0,
"The read of one thread must cover full quantization groups, or it "
"must divide one quantization group exactly.");

const int src_ld;
const int tile_stride;
Expand All @@ -600,6 +627,9 @@ struct QuantizedBlockLoader {
const short bi;
const short bj;

// First group this thread reads from.
const short group_id;

threadgroup T* dst;
const device uint8_t* src;
const device T* scales;
Expand All @@ -622,23 +652,40 @@ struct QuantizedBlockLoader {
thread_idx(simd_group_id * 32 + simd_lane_id),
bi(n_reads* thread_idx / BCOLS_PACKED),
bj((n_reads * thread_idx) % BCOLS_PACKED),
group_id((bj * pack_factor) / group_size),
dst(dst_ + bi * dst_ld + bj * pack_factor),
src(src_ + bi * src_ld * bytes_per_pack / pack_factor +
bj * bytes_per_pack),
scales(scales_ + bi * src_ld / group_size),
biases(biases_ + bi * src_ld / group_size) {}
scales(scales_ + bi * src_ld / group_size + group_id),
biases(biases_ + bi * src_ld / group_size + group_id) {}

void dequantize_reads() const thread {
if (groups_per_read == 1) {
T scale = *scales;
T bias = *biases;
for (int i = 0; i < n_reads; i++) {
dequantize<T, pack_factor, bits>(
src + i * bytes_per_pack, scale, bias, dst + i * pack_factor);
}
} else {
for (int g = 0; g < groups_per_read; g++) {
T scale = scales[g];
T bias = biases[g];
for (int i = 0; i < reads_per_group; i++) {
const int r = g * reads_per_group + i;
dequantize<T, pack_factor, bits>(
src + r * bytes_per_pack, scale, bias, dst + r * pack_factor);
}
}
}
}

void load_unsafe() const thread {
if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
return;
}

T scale = *scales;
T bias = *biases;
for (int i = 0; i < n_reads; i++) {
dequantize<T, pack_factor, bits>(
src + i * bytes_per_pack, scale, bias, dst + i * pack_factor);
}
dequantize_reads();
}

void load_safe(short2 src_tile_dim) const thread {
Expand All @@ -660,15 +707,7 @@ struct QuantizedBlockLoader {
return;
}

T scale = *scales;
T bias = *biases;
for (int i = 0; i < n_reads; i++) {
dequantize<T, pack_factor, bits>(
(device uint8_t*)(src + i * bytes_per_pack),
scale,
bias,
dst + i * pack_factor);
}
dequantize_reads();
}

void next() thread {
Expand All @@ -682,8 +721,8 @@ struct QuantizedBlockLoader {
biases++;
}
} else {
scales++;
biases++;
scales += n_groups;
biases += n_groups;
}
} else {
scales += group_stride;
Expand All @@ -692,146 +731,6 @@ struct QuantizedBlockLoader {
}
};

template <
typename T,
short BROWS,
short BCOLS,
short dst_ld,
short reduction_dim,
short tgp_size,
short bits>
struct QuantizedBlockLoader<
T,
BROWS,
BCOLS,
dst_ld,
reduction_dim,
tgp_size,
32,
bits> {
MLX_MTL_CONST short group_size = 32;

static_assert(
BCOLS % group_size == 0,
"The group size should be divisible by the columns");
static_assert(
bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 ||
bits == 8,
"Template undefined for bits not in {2, 3, 4, 5, 6, 8}");

MLX_MTL_CONST short pack_factor = get_pack_factor<bits, 8>();
MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack<bits>();
MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor;
MLX_MTL_CONST short n_reads =
(BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size;
MLX_MTL_CONST short n_groups = BCOLS / group_size;

static_assert(
(BCOLS_PACKED / n_reads) == n_groups,
"Other configurations are not yet supported");

const int src_ld;
const int tile_stride;
const int group_stride;

const short thread_idx;
const short bi;
const short bj;

const short group_id;

threadgroup T* dst;
const device uint8_t* src;
const device T* scales;
const device T* biases;

QuantizedBlockLoader(
const device uint8_t* src_,
const device T* scales_,
const device T* biases_,
const int src_ld_,
threadgroup T* dst_,
ushort simd_group_id [[simdgroup_index_in_threadgroup]],
ushort simd_lane_id [[thread_index_in_simdgroup]]) thread
: src_ld(src_ld_),
tile_stride(
reduction_dim ? BCOLS_PACKED* bytes_per_pack
: BROWS * src_ld * bytes_per_pack / pack_factor),
group_stride(BROWS* src_ld / group_size),
thread_idx(simd_group_id * 32 + simd_lane_id),
bi(n_reads* thread_idx / BCOLS_PACKED),
bj((n_reads * thread_idx) % BCOLS_PACKED),
group_id((bj * pack_factor) / group_size),
dst(dst_ + bi * dst_ld + bj * pack_factor),
src(src_ + bi * src_ld * bytes_per_pack / pack_factor +
bj * bytes_per_pack),
scales(scales_ + bi * src_ld / group_size + group_id),
biases(biases_ + bi * src_ld / group_size + group_id) {}

void load_unsafe() const thread {
if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
return;
}

T scale = *scales;
T bias = *biases;
for (int i = 0; i < n_reads; i++) {
dequantize<T, pack_factor, bits>(
src + i * bytes_per_pack, scale, bias, dst + i * pack_factor);
}
}

void load_safe(short2 src_tile_dim) const thread {
if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
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) {
for (int i = 0; i < n_reads * pack_factor; i++) {
dst[i] = T(0);
}
return;
}

T scale = *scales;
T bias = *biases;
for (int i = 0; i < n_reads; i++) {
dequantize<T, pack_factor, bits>(
(device uint8_t*)(src + i * bytes_per_pack),
scale,
bias,
dst + i * pack_factor);
}
}

void next() thread {
src += tile_stride;
if (reduction_dim == 1) {
// if (group_steps > 1) {
// group_step_cnt++;
// if (group_step_cnt == group_steps) {
// group_step_cnt = 0;
// scales++;
// biases++;
// }
// } else {
scales += n_groups;
biases += n_groups;
// }
} else {
scales += n_groups * group_stride;
biases += n_groups * group_stride;
}
}
};

template <typename T>
METAL_FUNC void adjust_matrix_offsets(
const device T*& x,
Expand Down
15 changes: 13 additions & 2 deletions mlx/backend/metal/kernels/quantized_nax.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,9 +76,20 @@
instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 1) \
instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 0)

#define instantiate_gather_qmm_rhs_tile(type, group_size, bits, bm, bn, bk, wm, wn) \
instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nt, type, group_size, bits, bm, bn, bk, wm, wn, true) \
instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nn, type, group_size, bits, bm, bn, bk, wm, wn, false)

// Block tiles for the gather kernel, default first. The rest are selected
// with MLX_QMM_TILE_NAX. Keep in sync with gather_qmm_rhs_nax_tiles in
// mlx/backend/metal/quantized.cpp.
#define instantiate_quantized_all_rhs(type, group_size, bits) \
instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nt, type, group_size, bits, 64, 64, 64, 2, 2, true) \
instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nn, type, group_size, bits, 64, 64, 64, 2, 2, false)
instantiate_gather_qmm_rhs_tile(type, group_size, bits, 64, 64, 64, 2, 2) \
instantiate_gather_qmm_rhs_tile(type, group_size, bits, 32, 64, 64, 2, 2) \
instantiate_gather_qmm_rhs_tile(type, group_size, bits, 16, 64, 64, 1, 2) \
instantiate_gather_qmm_rhs_tile(type, group_size, bits, 128, 64, 64, 2, 2) \
instantiate_gather_qmm_rhs_tile(type, group_size, bits, 64, 32, 64, 2, 2) \
instantiate_gather_qmm_rhs_tile(type, group_size, bits, 32, 32, 64, 2, 2)

#define instantiate_quantized_funcs(type, group_size, bits) \
instantiate_quantized_all_batched(type, group_size, bits) \
Expand Down
Loading