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

Fp qmv (#2984)

Browse Source
This commit is contained in:
Awni Hannun
2026-01-27 06:33:06 -08:00
committed by GitHub
parent ce4d0a62ef
commit 4912cc47c2
22 changed files with 1054 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/cuda_skip.py
-1
View File
@@ -47,7 +47,6 @@ cuda_skip = {
"TestQuantized.test_qmm_shapes",
"TestQuantized.test_qmm_vjp",
"TestQuantized.test_qmv",
"TestQuantized.test_fp_qmv",
"TestQuantized.test_fp_qvm",
"TestQuantized.test_qvm",
"TestQuantized.test_qvm_splitk",
python/tests/test_quantized.py
+54
View File
@@ -160,6 +160,38 @@ class TestQuantized(mlx_tests.MLXTestCase):
w_hat = mx.dequantize(w_q, scales, mode="nvfp4")
self.assertTrue(mx.all(w_hat == 0))
def test_qqmv(self):
key = mx.random.key(0)
k1, k2 = mx.random.split(key)
tests = product(
[256, 512, 67], # M
[64, 256], # N
)
modes = ["nvfp4", "mxfp8"]
for M, N in tests:
for mode in modes:
with self.subTest(shape=(M, N), mode=mode):
x_shape = (1, N)
w_shape = (M, N)
x = mx.random.normal(shape=x_shape, key=k1)
x_hat = mx.dequantize(
*mx.quantize(x, mode=mode), mode=mode, dtype=mx.float32
)
w = mx.random.normal(shape=w_shape, key=k2)
w_q, scales = mx.quantize(w, mode=mode)
w_hat = mx.dequantize(w_q, scales, mode=mode, dtype=mx.float32)
y_q = mx.qqmm(
x,
w_q,
scales,
mode=mode,
)
y_hat = x_hat @ 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_qmm(self):
key = mx.random.key(0)
k1, k2 = mx.random.split(key)
@@ -338,6 +370,28 @@ class TestQuantized(mlx_tests.MLXTestCase):
self.assertEqual(y_q.shape, y_hat.shape)
self.assertLess((y_q - y_hat).abs().max(), 1e-3)
# Test multiple of 16 but not 32
M = 128
N = 48
mode = "nvfp4"
with self.subTest(shape=(B, M, N), mode=mode):
x_shape = (1, N)
w_shape = (M, N)
x = mx.random.normal(shape=x_shape, key=k1)
w = mx.random.normal(shape=w_shape, key=k2)
w_q, scales = mx.quantize(w, mode=mode)
w_hat = mx.dequantize(w_q, scales, mode=mode)
y_q = mx.quantized_matmul(
x,
w_q,
scales,
transpose=True,
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_qvm(self):
key = mx.random.key(0)
k1, k2 = mx.random.split(key)