quantized_matmul is slower than dense bf16 when the output dimension is not a multiple of 8, and ~7x slower below 8

Describe the bug

At batch 1, mx.quantized_matmul is slower than the same matmul done densely in
bfloat16 whenever the output dimension is not a multiple of 8. At out_dim = 4
it is 6.8× slower than dense, and 4× slower than the same quantized matmul
at out_dim = 8 — a matrix half the size takes four times as long.

To Reproduce

Each call multiplies a (1, in_dim) bfloat16 vector by an (out_dim, in_dim)
weight matrix, once with the matrix quantized to 4-bit and once with it dense in
bfloat16. Dense bfloat16 is the reference: quantizing should not make a matmul
slower than not quantizing it. Calls are chained so each depends on the previous
(the batch-1 autoregressive regime, where no two matmuls overlap), and timings
are differenced across two chain lengths so the fixed mx.eval round-trip
cancels.

import time
import mlx.core as mx

IN_DIM = 10240
GROUP_SIZE = 64
BITS = 4


def time_per_call(build_chain, short=4, long=36):
    def run(chain_length, repeats=50, warmup=15):
        chain = build_chain(chain_length)
        for _ in range(warmup):
            mx.eval(chain())
        mx.synchronize()
        start = time.perf_counter()
        for _ in range(repeats):
            mx.eval(chain())
        mx.synchronize()
        return (time.perf_counter() - start) / repeats * 1e6

    return (run(long) - run(short)) / (long - short)


activation = mx.random.normal((1, IN_DIM)).astype(mx.bfloat16)
mx.eval(activation)


def feed_forward(result):
    """Widen one output column back to in_dim so the next call can consume it."""
    return mx.broadcast_to(result[:, :1], (1, IN_DIM)).astype(mx.bfloat16)


print(f"{'out_dim':>8} {'out_dim % 8':>12} {'quantized':>11} {'dense bf16':>11} {'ratio':>7}")

for out_dim in (4, 8, 12, 16, 32, 64, 128, 512, 2048):
    weight = mx.random.normal((out_dim, IN_DIM)).astype(mx.bfloat16)
    packed, scales, biases = mx.quantize(weight, group_size=GROUP_SIZE, bits=BITS)
    mx.eval(weight, packed, scales, biases)

    def quantized_chain(chain_length):
        def call():
            value = activation
            for _ in range(chain_length):
                value = feed_forward(
                    mx.quantized_matmul(
                        value, packed, scales, biases,
                        transpose=True, group_size=GROUP_SIZE, bits=BITS,
                    )
                )
            return value
        return call

    def dense_chain(chain_length):
        def call():
            value = activation
            for _ in range(chain_length):
                value = feed_forward(value @ weight.T)
            return value
        return call

    quantized_us = time_per_call(quantized_chain)
    dense_us = time_per_call(dense_chain)
    print(f"{out_dim:8d} {out_dim % 8:12d} {quantized_us:10.1f}us "
          f"{dense_us:10.1f}us {quantized_us / dense_us:6.2f}x")

Output

 out_dim  out_dim % 8   quantized  dense bf16   ratio
       4            4       99.0us       14.7us   6.76x
       8            0       25.2us       14.2us   1.78x
      12            4       31.7us       13.8us   2.30x
      16            0       23.1us       14.1us   1.64x
      32            0       22.0us       14.2us   1.55x
      64            0       22.4us       14.5us   1.55x
     128            0       22.1us       14.8us   1.50x
     512            0       23.1us       19.0us   1.22x
    2048            0       26.1us       50.9us   0.51x

Only at out_dim = 2048 does quantization win. Everywhere below it is a loss,
and below 8 it is a rout.

Two separate effects, both visible in the divisibility pattern

out_dim	4	8	9	12	16	17	20	24	32	64
out_dim % 8	4	0	1	4	0	1	4	0	0	0
µs	96.9	23.7	32.3	32.4	22.3	31.6	31.3	22.6	21.9	22.5

This matches qmv() in mlx/backend/metal/quantized.cpp (v0.32.2, lines 477-493):

int bn = 8;
...
bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0;
...
MTL::Size grid_dims(M, (N + bn - 1) / bn, B);

1. N % 8 != 0 drops off qmv_fast onto qmv — costs ~1.4× (22 → 32 µs).
2. N < 8 additionally collapses the grid to a single threadgroup, since the
   N dimension is ceil(N / 8) and M = B = 1 at batch 1. One threadgroup of 64
   threads then performs the entire K-length reduction. This is the ~4× on top,
   and it is why the penalty grows with K: 1.8× at K=1024, 2.5× at K=2560,
   4.1× at K=10240 (out_dim 4 vs 8).

Independent of group_size (32/64/128 measured identical). Worse at 8 bits:
out_dim=4 costs 169 µs at bits=8 vs 97 µs at bits=4.

Workaround: padding is exact and recovers the speed

Zero-padding the weight rows to a multiple of 8 and slicing the output is
bit-exact — each output row is an independent dot product along K, so extra
rows cannot perturb the existing ones. max abs diff is exactly 0.0 at every
size measured:

out_dim	as-is	padded to a multiple of 8	speedup
1	50.2 µs	24.8 µs	2.03×
2	62.3 µs	23.3 µs	2.67×
4	96.8 µs	24.3 µs	3.99×
6	81.1 µs	23.2 µs	3.49×
12	32.8 µs	22.1 µs	1.48×
20	32.3 µs	23.3 µs	1.39×

Applied to the 96 layers in the model below, this gave a 1.26× end-to-end
decode speedup with an identical token stream, for +5.4 MB.

Expected behavior

quantized_matmul should not be slower than dense bfloat16 at any shape, and
out_dim = 4 should not cost 4× out_dim = 8 — especially when padding the
weight to 8 fixes it exactly and costs nothing.

Two directions, both maintainer calls:

1. Pad at quantization time. mx.quantize / nn.QuantizedLinear could round
   out_features up to a multiple of 8 and slice on output. Contained, and
   exact — but it changes stored weight shapes, which has checkpoint-compat
   consequences we are not in a position to judge.
2. Fix the thread allocation in qmv_impl. When N < num_simdgroups * results_per_simdgroup, put the idle simdgroups on a slice of K rather than
   letting them return, and reduce across them. No stored-shape change.

Happy to send a PR for either if you have a preference — we have the repro and
the hardware, but the design choice (particularly the compat implications of 1)
seemed yours to make rather than ours to assume.

Related

#3553 reported a discontinuity on the M (batch) axis of the same kernel.
This is the N (output) axis — different condition, same kernel family.

Desktop

• OS: macOS 15.6 (Darwin 25.6.0)
• Version: mlx 0.32.2
• Device: Apple M3 Ultra, 96 GB unified (applegpu_g15d)

Additional context

Found in the decode critical path of Qwen3.8-Flash-Next, which has a
Linear(10240 → 4) gate running 96 times per token; that one call site was 32%
of the decode step. Narrow output projections are not unusual — router gates,
scoring heads and per-stream gating weights all produce them — and nothing in a
profile draws attention to a 20 KB matrix, so models using them may be paying
this silently.

Run the supplied Python benchmark first, then inspect qmv() and qmv_impl in mlx/backend/metal/quantized.cpp, especially the fast-path and thread-grid logic around lines 477-493. The work is done when non-multiple-of-8 output dimensions no longer incur the reported penalty, while the existing exactness and shape compatibility expectations are preserved.