Skip to content

Use a 32-row block in qmm_t_nax when one block covers all of M #4171

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
dwijenpatel wants to merge 1 commit into
base: main
Choose a base branch
from
+51 −5
Open

Use a 32-row block in qmm_t_nax when one block covers all of M #4171

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
10 changes: 7 additions & 3 deletions mlx/backend/metal/kernels/fp_quantized_nax.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
type, \
group_size, \
bits, \
aligned)
aligned, bm, bk, bn, wm, wn)

#define instantiate_quantized_aligned_batched(mode, name, type, bm, bn, bk, wm, wn, aligned, batched, group_size, bits) \
instantiate_kernel( \
Expand All @@ -34,7 +34,7 @@
group_size, \
bits, \
aligned, \
batched)
batched, bm, bk, bn, wm, wn)

#define instantiate_gather_qmm_rhs(func, name, type, bm, bn, bk, wm, wn, transpose, mode, group_size, bits) \
instantiate_kernel( \
Expand All @@ -57,7 +57,11 @@
instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 1, group_size, bits) \
instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 0, group_size, bits) \
instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 1, group_size, bits) \
instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 0, group_size, bits)
instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 0, group_size, bits) \
instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, true, 1, group_size, bits) \
instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, true, 0, group_size, bits) \
instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, false, 1, group_size, bits) \
instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, false, 0, group_size, bits)


#define instantiate_quantized_all_rhs(type, mode, group_size, bits) \
Expand Down
6 changes: 5 additions & 1 deletion mlx/backend/metal/kernels/quantized_nax.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,11 @@
instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true, 1) \
instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true, 0) \
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)
instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 0) \
instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, true, 1) \
instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, true, 0) \
instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, false, 1) \
instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, false, 0)

#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) \
Expand Down
3 changes: 2 additions & 1 deletion mlx/backend/metal/quantized.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -830,7 +830,8 @@ void qmm_nax(

int wm = 2;
int wn = 2;
int bm = 64;
// Use smaller bm when one block covers all of M.
int bm = (M <= 32) ? 32 : 64;
int bn = 64;
int bk = 64;
MTL::Size group_dims(32, wn, wm);
Expand Down
37 changes: 37 additions & 0 deletions python/tests/test_quantized.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -334,6 +334,7 @@ def test_qmm_large_dims(self):
K = 128
tests = [
(16, 32840), # unaligned N > 2**15, M < 32: partial M-tile
(32, 32840), # M at the small-block dispatch boundary
(33, 32840), # unaligned N > 2**15, M % 32 != 0
(33000, 64), # M > 2**15: row distance overflows (aligned N)
]
Expand All @@ -353,6 +354,42 @@ def test_qmm_large_dims(self):
tol = 1e-3 if dtype == mx.float32 else 1.5e-3
self.assertLess((y_q - y_hat).abs().max(), tol)

def test_qmm_small_m_block(self):
# The batched and fp-mode variants of the small-M block, which the
# test_qmm_large_dims shapes cannot reach.
if mx.default_device() == mx.cpu:
self.skipTest("Covers GPU kernels only")
key = mx.random.key(0)
k1, k2 = mx.random.split(key)
K = 1024
tests = [
# mode, group_size, bits, M, N, batch
("affine", 64, 4, 14, 8256, (2,)), # batched w
("mxfp4", None, None, 14, 8256, ()),
]
for mode, group_size, bits, M, N, batch in tests:
dtype = mx.float16 if mode == "affine" else mx.bfloat16
with self.subTest(
mode=mode, group_size=group_size, bits=bits, M=M, N=N, batch=batch
):
x = (mx.random.normal(batch + (M, K), key=k1) / K**0.5).astype(dtype)
w = (mx.random.normal(batch + (N, K), key=k2) / K**0.5).astype(dtype)
if mode == "affine":
wq = mx.quantize(w, group_size=group_size, bits=bits)
else:
wq = mx.quantize(w, mode=mode)
w_hat = mx.dequantize(*wq, group_size=group_size, bits=bits, mode=mode)
y_ref = x @ w_hat.swapaxes(-1, -2)
y = mx.quantized_matmul(
x,
*wq,
transpose=True,
group_size=group_size,
bits=bits,
mode=mode,
)
self.assertLess((y_ref - y).abs().max(), 1e-3)

def test_qmm_vjp(self):
key = mx.random.key(0)
k1, k2 = mx.random.split(key)
Expand Down
Loading