L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix qmv_impl for small N (#3096)

Browse Source
This commit is contained in:
Manuel Candales
2026-02-05 17:33:36 -08:00
committed by GitHub
parent ceea571490
commit 90e38f7b93
4 changed files with 82 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/backend/metal/kernels/fp_quantized.h
+9 -3
View File
@@ -434,7 +434,9 @@ METAL_FUNC void fp_qmv_impl(
for (; k < in_vec_size - block_size; k += block_size) {
load_vector<T, U, values_per_thread>(x, x_thread);
for (int row = 0; out_row + row < out_vec_size; row++) {
for (int row = 0;
row < results_per_simdgroup && out_row + row < out_vec_size;
row++) {
auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
const device auto* sl = scales + row * in_vec_size_g;
@@ -453,7 +455,9 @@ METAL_FUNC void fp_qmv_impl(
if (remaining > 0) {
load_vector_safe<T, U, values_per_thread>(x, x_thread, remaining);
for (int row = 0; out_row + row < out_vec_size; row++) {
for (int row = 0;
row < results_per_simdgroup && out_row + row < out_vec_size;
row++) {
auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
const device auto* sl = scales + row * in_vec_size_g;
@@ -462,7 +466,9 @@ METAL_FUNC void fp_qmv_impl(
}
}
for (int row = 0; out_row + row < out_vec_size; row++) {
for (int row = 0;
row < results_per_simdgroup && out_row + row < out_vec_size;
row++) {
result[row] = simd_sum(result[row]);
if (simd_lid == 0) {
y[row] = static_cast<T>(result[row]);
mlx/backend/metal/kernels/quantized.h
+11 -5
View File
@@ -867,7 +867,9 @@ METAL_FUNC void qmv_impl(
for (; k < in_vec_size - block_size; k += block_size) {
U sum = load_vector<T, U, values_per_thread, bits>(x, x_thread);
for (int row = 0; out_row + row < out_vec_size; row++) {
for (int row = 0;
row < results_per_simdgroup && out_row + row < out_vec_size;
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;
@@ -891,19 +893,23 @@ METAL_FUNC void qmv_impl(
U sum = load_vector_safe<T, U, values_per_thread, bits>(
x, x_thread, remaining);
for (int row = 0; out_row + row < out_vec_size; row++) {
for (int row = 0;
row < results_per_simdgroup && out_row + row < out_vec_size;
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;
U s = sl[0];
U b = bl[0];
result[row] +=
qdot<U, values_per_thread, bits>(wl, x_thread, s, b, sum);
result[row] += qdot_safe<U, values_per_thread, bits>(
wl, x_thread, s, b, sum, remaining);
}
}
for (int row = 0; out_row + row < out_vec_size; row++) {
for (int row = 0;
row < results_per_simdgroup && out_row + row < out_vec_size;
row++) {
result[row] = simd_sum(result[row]);
if (simd_lid == 0) {
y[row] = static_cast<T>(result[row]);
python/tests/cuda_skip.py
+1
View File
@@ -50,6 +50,7 @@ cuda_skip = {
"TestQuantized.test_fp_qvm",
"TestQuantized.test_qvm",
"TestQuantized.test_qvm_splitk",
"TestQuantized.test_qmv_small_non_multiples",
"TestQuantized.test_small_matrix",
"TestQuantized.test_throw",
"TestQuantized.test_vjp_scales_biases",
python/tests/test_quantized.py
+61 -2
View File
@@ -681,6 +681,65 @@ class TestQuantized(mlx_tests.MLXTestCase):
self.assertEqual(y_q.shape, y_hat.shape)
self.assertLess((y_q - y_hat).abs().max(), 1e-3)
def test_qmv_small_non_multiples(self):
# Test very small K and N dimensions (e.g., [MxK] x [NxK].T = [MxN])
# Each tuple is (M, K, N) representing input rows, weight cols, weight rows
test_cases = [
(1, 32, 3),
(2, 32, 10),
(1, 32, 5),
(4, 32, 7),
]
# Test different quantization settings (bits, group_size, mode)
quantization_settings = [
(4, 32, "affine"),
(6, 32, "affine"),
(4, 16, "nvfp4"),
]
for M, K, N in test_cases:
for bits, group_size, mode in quantization_settings:
# Test without batch dimension
with self.subTest(
M=M,
K=K,
N=N,
batch=None,
group_size=group_size,
bits=bits,
mode=mode,
):
w = mx.random.normal(shape=(N, K))
w_q, *sb = mx.quantize(
w,
group_size=group_size,
bits=bits,
mode=mode,
)
w_hat = mx.dequantize(
w_q,
*sb,
group_size=group_size,
bits=bits,
mode=mode,
)
# Test qmv/qmm_t (transpose=True): [MxK] @ [NxK].T = [MxN]
x = mx.random.normal(shape=(M, K))
y_q = mx.quantized_matmul(
x,
w_q,
*sb,
transpose=True,
group_size=group_size,
bits=bits,
mode=mode,
)
y_hat = x @ mx.swapaxes(w_hat, -1, -2)
self.assertEqual(y_q.shape, y_hat.shape)
self.assertLess((y_q - y_hat).abs().max(), 1e-3)
def test_gather_qmm(self):
def quantize(w, transpose=True, group_size=None, bits=None, mode="affine"):
if mode == "affine":
@@ -957,7 +1016,7 @@ class TestQuantized(mlx_tests.MLXTestCase):
group_size=group_size,
mode=mode,
transpose=transpose,
rhs_indices=indices
rhs_indices=indices,
)
xs, idx, inv_order = gather_sort(x, indices)
y3 = mx.gather_mm(xs, w, rhs_indices=idx, sorted_indices=True)
@@ -969,7 +1028,7 @@ class TestQuantized(mlx_tests.MLXTestCase):
mode=mode,
rhs_indices=idx,
transpose=transpose,
sorted_indices=True
sorted_indices=True,
)
y3 = scatter_unsort(y3, inv_order, indices.shape)
y4 = scatter_unsort(y4, inv_order, indices.shape)