Skip to content

Use the fast qmv kernel for outputs not divisible by 8 #4516

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

Merged
zcbenz merged 5 commits into from
Oct 4, 2026
Merged

Use the fast qmv kernel for outputs not divisible by 8 #4516

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9a2eaf2
Accumulate affine QMV input sums in float
wyanzhao Sep 15, 2026
b53b13d
Use the fast qmv kernel for affine outputs not divisible by 8
wyanzhao Sep 15, 2026
9c7d11a
Extend fast QMV row tails to floating-point quantization
wyanzhao Oct 1, 2026
99bad7d
nit
zcbenz Oct 4, 2026
c55511e
nit
zcbenz Oct 4, 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
41 changes: 29 additions & 12 deletions mlx/backend/metal/kernels/fp_quantized.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -334,6 +334,7 @@ template <
typename T,
int group_size,
int bits,
bool partial_rows = false,
bool has_global_scale = false,
int results_per_simdgroup = 4>
METAL_FUNC void fp_qmv_fast_impl(
Expand Down Expand Up @@ -367,6 +368,15 @@ METAL_FUNC void fp_qmv_fast_impl(
const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) +
simd_gid * results_per_simdgroup;

// Reuse the last valid row for reads beyond the output tile.
int last_row = results_per_simdgroup - 1;
if constexpr (partial_rows) {
if (out_row >= out_vec_size) {
return;
}
last_row = min(last_row, out_vec_size - 1 - out_row);
}

ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack;
scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
x += tid.x * in_vec_size + simd_lid * values_per_thread;
Expand All @@ -376,8 +386,12 @@ METAL_FUNC void fp_qmv_fast_impl(
load_vector<T, U, values_per_thread>(x, x_thread);

for (int row = 0; row < results_per_simdgroup; row++) {
auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
const device auto* sl = scales + row * in_vec_size_g;
int src = row;
if constexpr (partial_rows) {
src = min(row, last_row);
}
auto wl = (const device uint8_t*)(ws + src * in_vec_size_w);
const device auto* sl = scales + src * in_vec_size_g;

U s = dequantize_scale<U, group_size>(sl[0]);
result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s);
Expand All @@ -395,7 +409,7 @@ METAL_FUNC void fp_qmv_fast_impl(

for (int row = 0; row < results_per_simdgroup; row++) {
result[row] = simd_sum(result[row]);
if (simd_lid == 0) {
if (simd_lid == 0 && row <= last_row) {
if constexpr (has_global_scale) {
y[row] = static_cast<T>(result[row] * inv_scale_enc);
} else {
Expand Down Expand Up @@ -1243,6 +1257,7 @@ template <
int group_size,
int bits,
bool batched,
bool partial_rows = false,
bool has_global_scale = false,
int results_per_simdgroup = 4>
[[kernel]] void fp_qmv_fast(
Expand Down Expand Up @@ -1284,6 +1299,7 @@ template <
T,
group_size,
bits,
partial_rows,
has_global_scale,
results_per_simdgroup>(
w,
Expand All @@ -1303,6 +1319,7 @@ template <
int group_size,
int bits,
bool batched,
bool partial_rows = false,
bool has_global_scale = false,
int results_per_simdgroup = 4>
[[kernel]] void fp_qmv(
Expand Down Expand Up @@ -1511,14 +1528,14 @@ template <typename T, const int group_size, int bits, int split_k = 32>

template <
typename T,
const int group_size,
const int bits,
const bool aligned_N,
const bool batched,
const bool has_global_scale = false,
const int BM = 32,
const int BK = 32,
const int BN = 32>
int group_size,
int bits,
bool aligned_N,
bool batched,
bool has_global_scale = false,
int BM = 32,
int BK = 32,
int BN = 32>
[[kernel]] void fp_qmm_t(
const device uint32_t* w,
const device uint8_t* scales,
Expand Down Expand Up @@ -1684,7 +1701,7 @@ template <typename T, int group_size, int bits, bool has_global_scale = false>
w_strides,
s_strides,
tid);
fp_qmv_fast_impl<T, group_size, bits, has_global_scale>(
fp_qmv_fast_impl<T, group_size, bits, false, has_global_scale>(
w,
scales,
global_scale,
Expand Down
74 changes: 35 additions & 39 deletions mlx/backend/metal/kernels/fp_quantized.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
bits, \
true)

#define instantiate_quantized_batched(mode, name, type, batched, group_size, bits) \
#define instantiate_quantized_batched(mode, name, type, group_size, bits, batched) \
instantiate_kernel( \
#mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched, \
fp_ ## name, \
Expand All @@ -38,6 +38,18 @@
batched, \
true)

#define instantiate_quantized_qmv(mode, name, type, group_size, bits, batched, partial_rows, has_global_scale, results) \
instantiate_kernel( \
#mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched "_pr_" #partial_rows "_hgs_" #has_global_scale "_r_" #results, \
fp_ ## name, \
type, \
group_size, \
bits, \
batched, \
partial_rows, \
has_global_scale, \
results)

#define instantiate_quantized_aligned(mode, name, type, aligned, group_size, bits) \
instantiate_kernel( \
#mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned, \
Expand Down Expand Up @@ -79,30 +91,6 @@
batched, \
true)

#define instantiate_quantized_qmv_fast(mode, type, results, batched, group_size, bits) \
instantiate_kernel( \
#mode "_qmv_fast_" #type "_gs_" #group_size "_b_" #bits "_r_" #results "_batch_" #batched, \
fp_qmv_fast, \
type, \
group_size, \
bits, \
batched, \
false, \
results) \
instantiate_kernel( \
#mode "_qmv_fast_" #type "_gs_" #group_size "_b_" #bits "_r_" #results "_batch_" #batched "_hgs", \
fp_qmv_fast, \
type, \
group_size, \
bits, \
batched, \
true, \
results)

#define instantiate_quantized_qmv_fast_r2(mode, type, group_size, bits) \
instantiate_quantized_qmv_fast(mode, type, 2, 1, group_size, bits) \
instantiate_quantized_qmv_fast(mode, type, 2, 0, group_size, bits)

#define instantiate_quantized_quad(mode, name, type, D, batched, group_size, bits) \
instantiate_kernel( \
#mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_d_" #D "_batch_" #batched, \
Expand Down Expand Up @@ -161,13 +149,20 @@
true)

#define instantiate_quantized_batched_wrap(name, type, mode, group_size, bits) \
instantiate_quantized_batched(mode, name, type, 1, group_size, bits) \
instantiate_quantized_batched(mode, name, type, 0, group_size, bits)
instantiate_quantized_batched(mode, name, type, group_size, bits, 1) \
instantiate_quantized_batched(mode, name, type, group_size, bits, 0)

#define instantiate_quantized_qmv_wrap(name, type, mode, group_size, bits, partial_rows, results) \
instantiate_quantized_qmv(mode, name, type, group_size, bits, 1, partial_rows, 1, results) \
instantiate_quantized_qmv(mode, name, type, group_size, bits, 1, partial_rows, 0, results) \
instantiate_quantized_qmv(mode, name, type, group_size, bits, 0, partial_rows, 1, results) \
instantiate_quantized_qmv(mode, name, type, group_size, bits, 0, partial_rows, 0, results)

#define instantiate_quantized_all_batched(type, mode, group_size, bits) \
instantiate_quantized_batched_wrap(qmv_fast, type, mode, group_size, bits) \
instantiate_quantized_batched_wrap(qmv, type, mode, group_size, bits) \
instantiate_quantized_batched_wrap(qvm, type, mode, group_size, bits) \
#define instantiate_quantized_all_batched(type, mode, group_size, bits) \
instantiate_quantized_qmv_wrap(qmv_fast, type, mode, group_size, bits, 1, 4) \
instantiate_quantized_qmv_wrap(qmv_fast, type, mode, group_size, bits, 0, 4) \
instantiate_quantized_qmv_wrap(qmv, type, mode, group_size, bits, 0, 4) \
instantiate_quantized_batched_wrap(qvm, type, mode, group_size, bits) \
instantiate_quantized_batched_wrap(qmm_n, type, mode, group_size, bits)

#define instantiate_quantized_all_single(type, mode, group_size, bits) \
Expand Down Expand Up @@ -248,14 +243,15 @@
instantiate_quantized_all_rhs(type, mode, group_size, bits)

#define instantiate_quantized_types(type) \
instantiate_quantized_modes(type, nvfp4, 16, 4) \
instantiate_quantized_modes(type, mxfp8, 32, 8) \
instantiate_quantized_modes(type, mxfp4, 32, 4) \
instantiate_quantize_dequantize(type, nvfp4, 16, 4, false) \
instantiate_quantize_dequantize(type, nvfp4, 16, 4, true) \
instantiate_quantize_dequantize(type, mxfp8, 32, 8, false) \
instantiate_quantize_dequantize(type, mxfp4, 32, 4, false) \
instantiate_quantized_qmv_fast_r2(nvfp4, type, 16, 4)
instantiate_quantized_modes(type, nvfp4, 16, 4) \
instantiate_quantized_modes(type, mxfp8, 32, 8) \
instantiate_quantized_modes(type, mxfp4, 32, 4) \
instantiate_quantize_dequantize(type, nvfp4, 16, 4, false) \
instantiate_quantize_dequantize(type, nvfp4, 16, 4, true) \
instantiate_quantize_dequantize(type, mxfp8, 32, 8, false) \
instantiate_quantize_dequantize(type, mxfp4, 32, 4, false) \
instantiate_quantized_qmv_wrap(qmv_fast, type, nvfp4, 16, 4, 1, 2) \
instantiate_quantized_qmv_wrap(qmv_fast, type, nvfp4, 16, 4, 0, 2)

instantiate_quantized_types(float)
instantiate_quantized_types(bfloat16_t)
Expand Down
58 changes: 38 additions & 20 deletions mlx/backend/metal/kernels/quantized.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ inline U load_vector(const device T* x, thread U* x_thread) {

if (bits == 2) {
for (int i = 0; i < values_per_thread; i += 4) {
sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]);
x_thread[i] = x[i];
x_thread[i + 1] = x[i + 1] / 4.0f;
x_thread[i + 2] = x[i + 2] / 16.0f;
Expand All @@ -46,8 +46,8 @@ inline U load_vector(const device T* x, thread U* x_thread) {

else if (bits == 3) {
for (int i = 0; i < values_per_thread; i += 8) {
sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] +
x[i + 6] + x[i + 7];
sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) +
U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]);
x_thread[i] = x[i];
x_thread[i + 1] = x[i + 1] / 8.0f;
x_thread[i + 2] = x[i + 2] / 64.0f;
Expand All @@ -61,7 +61,7 @@ inline U load_vector(const device T* x, thread U* x_thread) {

else if (bits == 4) {
for (int i = 0; i < values_per_thread; i += 4) {
sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]);
x_thread[i] = x[i];
x_thread[i + 1] = x[i + 1] / 16.0f;
x_thread[i + 2] = x[i + 2] / 256.0f;
Expand All @@ -71,8 +71,8 @@ inline U load_vector(const device T* x, thread U* x_thread) {

else if (bits == 5) {
for (int i = 0; i < values_per_thread; i += 8) {
sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] +
x[i + 6] + x[i + 7];
sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) +
U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]);
x_thread[i] = x[i];
x_thread[i + 1] = x[i + 1] / 32.0f;
x_thread[i + 2] = x[i + 2] / 4.0f;
Expand All @@ -86,7 +86,7 @@ inline U load_vector(const device T* x, thread U* x_thread) {

else if (bits == 6) {
for (int i = 0; i < values_per_thread; i += 4) {
sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]);
x_thread[i] = x[i];
x_thread[i + 1] = x[i + 1] / 64.0f;
x_thread[i + 2] = x[i + 2] / 16.0f;
Expand Down Expand Up @@ -115,7 +115,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) {

if (bits == 2) {
for (int i = 0; i < N; i += 4) {
sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]);
x_thread[i] = x[i];
x_thread[i + 1] = x[i + 1] / 4.0f;
x_thread[i + 2] = x[i + 2] / 16.0f;
Expand All @@ -125,8 +125,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) {

else if (bits == 3) {
for (int i = 0; i < N; i += 8) {
sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] +
x[i + 6] + x[i + 7];
sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) +
U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]);

x_thread[i] = x[i];
x_thread[i + 1] = x[i + 1] / 8.0f;
Expand All @@ -141,7 +141,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) {

else if (bits == 4) {
for (int i = 0; i < N; i += 4) {
sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]);
x_thread[i] = x[i];
x_thread[i + 1] = x[i + 1] / 16.0f;
x_thread[i + 2] = x[i + 2] / 256.0f;
Expand All @@ -151,8 +151,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) {

else if (bits == 5) {
for (int i = 0; i < N; i += 8) {
sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] +
x[i + 6] + x[i + 7];
sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) +
U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]);
x_thread[i] = x[i];
x_thread[i + 1] = x[i + 1] / 32.0f;
x_thread[i + 2] = x[i + 2] / 4.0f;
Expand All @@ -166,7 +166,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) {

else if (bits == 6) {
for (int i = 0; i < N; i += 4) {
sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3];
sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]);
x_thread[i] = x[i];
x_thread[i + 1] = x[i + 1] / 64.0f;
x_thread[i + 2] = x[i + 2] / 16.0f;
Expand Down Expand Up @@ -753,7 +753,7 @@ METAL_FUNC void qmv_quad_impl(
}
}

template <typename T, int group_size, int bits>
template <typename T, int group_size, int bits, bool partial_rows = false>
METAL_FUNC void qmv_fast_impl(
const device uint32_t* w,
const device T* scales,
Expand Down Expand Up @@ -787,6 +787,18 @@ METAL_FUNC void qmv_fast_impl(
const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) +
simd_gid * results_per_simdgroup;

// With partial rows the output size need not be a multiple of 8. Rows of the
// last SIMD-group that fall past the output reuse the weights of the last
// valid output row, so the reduction loop needs no per-row bounds check, and
// are not stored.
int last_row = results_per_simdgroup - 1;
if constexpr (partial_rows) {
if (out_row >= out_vec_size) {
return;
}
last_row = min(last_row, out_vec_size - 1 - out_row);
}

ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack;
scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
biases += out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
Expand All @@ -797,9 +809,13 @@ METAL_FUNC void qmv_fast_impl(
U sum = load_vector<T, U, values_per_thread, bits>(x, x_thread);

for (int row = 0; row < results_per_simdgroup; row++) {
auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
const device T* sl = scales + row * in_vec_size_g;
const device T* bl = biases + row * in_vec_size_g;
int src = row;
if constexpr (partial_rows) {
src = min(row, last_row);
}
auto wl = (const device uint8_t*)(ws + src * in_vec_size_w);
const device T* sl = scales + src * in_vec_size_g;
const device T* bl = biases + src * in_vec_size_g;

U s = sl[0];
U b = bl[0];
Expand All @@ -814,7 +830,7 @@ METAL_FUNC void qmv_fast_impl(

for (int row = 0; row < results_per_simdgroup; row++) {
result[row] = simd_sum(result[row]);
if (simd_lid == 0) {
if (simd_lid == 0 && row <= last_row) {
y[row] = static_cast<T>(result[row]);
}
}
Expand Down Expand Up @@ -1602,6 +1618,7 @@ template <
int group_size,
int bits,
bool batched,
bool partial_rows = false,
bool has_global_scale = false,
int results_per_simdgroup = 4>
[[kernel]] void affine_qmv_fast(
Expand Down Expand Up @@ -1642,7 +1659,7 @@ template <
b_strides,
tid);
}
qmv_fast_impl<T, group_size, bits>(
qmv_fast_impl<T, group_size, bits, partial_rows>(
w,
scales,
biases,
Expand All @@ -1660,6 +1677,7 @@ template <
int group_size,
const int bits,
bool batched,
bool partial_rows = false,
bool has_global_scale = false,
int results_per_simdgroup = 4>
[[kernel]] void affine_qmv(
Expand Down
Loading
Loading