Skip to content

Use the fast qmv kernel for affine 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

Open
wyanzhao wants to merge 2 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
77ab494
Accumulate affine QMV input sums in float
wyanzhao Sep 15, 2026
7ee2bd0
Use the fast qmv kernel for affine outputs not divisible by 8
wyanzhao Sep 15, 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
112 changes: 93 additions & 19 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 @@ -1655,6 +1671,64 @@ template <
simd_lid);
}

template <
typename T,
int group_size,
int bits,
bool batched,
bool has_global_scale = false,
int results_per_simdgroup = 4>
[[kernel]] void affine_qmv_fast_rows(
const device uint32_t* w [[buffer(0)]],
const device T* scales [[buffer(1)]],
const device T* biases [[buffer(2)]],
const device T* x [[buffer(3)]],
device T* y [[buffer(4)]],
const constant int& in_vec_size [[buffer(5)]],
const constant int& out_vec_size [[buffer(6)]],
const constant int& x_batch_ndims [[buffer(7)]],
const constant int* x_shape [[buffer(8)]],
const constant int64_t* x_strides [[buffer(9)]],
const constant int& w_batch_ndims [[buffer(10)]],
const constant int* w_shape [[buffer(11)]],
const constant int64_t* w_strides [[buffer(12)]],
const constant int64_t* s_strides [[buffer(13)]],
const constant int64_t* b_strides [[buffer(14)]],
uint3 tid [[threadgroup_position_in_grid]],
uint simd_gid [[simdgroup_index_in_threadgroup]],
uint simd_lid [[thread_index_in_simdgroup]]) {
if (batched) {
int M = x_shape[x_batch_ndims];
adjust_matrix_offsets<T>(
x,
w,
scales,
biases,
y,
out_vec_size * M,
x_batch_ndims,
x_shape,
x_strides,
w_batch_ndims,
w_shape,
w_strides,
s_strides,
b_strides,
tid);
}
qmv_fast_impl<T, group_size, bits, true>(
w,
scales,
biases,
x,
y,
in_vec_size,
out_vec_size,
tid,
simd_gid,
simd_lid);
}

template <
typename T,
int group_size,
Expand Down
1 change: 1 addition & 0 deletions mlx/backend/metal/kernels/quantized.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@

#define instantiate_quantized_all_batched(type, group_size, bits) \
instantiate_quantized_batched_wrap(affine_qmv_fast, type, group_size, bits) \
instantiate_quantized_batched_wrap(affine_qmv_fast_rows, type, group_size, bits) \
instantiate_quantized_batched_wrap(affine_qmv, type, group_size, bits) \
instantiate_quantized_batched_wrap(affine_qvm, type, group_size, bits) \
instantiate_quantized_batched_wrap(affine_qmm_n, type, group_size, bits)
Expand Down
28 changes: 14 additions & 14 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 @@ -39,7 +39,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 @@ -49,8 +49,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 @@ -64,7 +64,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 @@ -74,8 +74,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 @@ -89,7 +89,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 @@ -118,7 +118,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 @@ -128,8 +128,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 @@ -144,7 +144,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 @@ -154,8 +154,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 @@ -169,7 +169,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
11 changes: 8 additions & 3 deletions mlx/backend/metal/quantized.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -481,7 +481,12 @@ void qmv(
std::string kname;
kname.reserve(64);
std::string type_string = get_type_string(x.dtype());
bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0;
bool aligned = K % qmv_fast_k_alignment(bits) == 0;
bool fast = N % bn == 0 && aligned;
// Affine outputs that are not a multiple of 8 can still use the fast kernel
// when the input is aligned: its last SIMD-group covers the remaining rows.
bool fast_rows = !fast && aligned && mode == "affine" && !global_scale;
const char* func = fast ? "qmv_fast" : (fast_rows ? "qmv_fast_rows" : "qmv");
// A narrower output tile reduces register pressure for large
// floating-point quantized matrix-vector products on M5 Max GPUs.
bool use_narrow_qmv = fast && N >= 4096 && d.get_architecture_gen() == 17 &&
Expand All @@ -493,7 +498,7 @@ void qmv(

concatenate(
kname,
mode + (fast ? "_qmv_fast_" : "_qmv_"),
mode + "_" + func + "_",
type_string,
"_gs_",
group_size,
Expand All @@ -505,7 +510,7 @@ void qmv(
auto kernel = get_quantized_kernel_wrapped(
d,
kname,
(fast ? "qmv_fast" : "qmv"),
func,
mode,
type_string,
group_size,
Expand Down
Loading