11 posts tagged with "mlx"

What happens when you take a thought experiment about burning LLM weights into silicon, try to build the closest real thing in software, and let the machine tell you which parts of the idea survive contact with MLX's actual copy semantics.

The idea started as the kind of thing you think about in the shower: what if Apple burned a frozen LLM's weights into a ROM chip sitting right next to RAM? Read-only. Never re-derived. Never fully materialized in general-purpose memory. A model that's just there, the way a calculator's multiplication table is just there.

You can't ship that. This is a software library, not a silicon fab. But strip the hardware away and look at what's actually being asked for — weights that are read-only, computed once, and shared across processes without each one paying full price — and it stops being science fiction. It starts looking like a real, specific gap in how VeloxQuant-MLX loads models today.

So we built the closest real thing, measured it against the actual machine it would run on, and let two of our own assumptions get overturned by the data along the way. This is the record of that — including the part where a "fix" I proposed made things worse, and the part where I had to walk back a claim about compression that turned out to be mathematically impossible.

What already existed​

VeloxQuant-MLX already had most of the ingredients, just not assembled this way.

QuantizedLinear compresses a weight matrix by normalizing each row, rotating it (Hadamard transform or a QR-derived rotation, depending on the dimension), and mapping it onto a small Lloyd-Max codebook — 2 to 4 bits per weight instead of 16. quantize_model() walks an entire mlx-lm model and replaces every nn.Linear with one of these. It works, and the compression ratios are real.

But the compressed result only ever lives as an mx.array in one process's memory. Every time you load the model — a new server worker, a restarted process, a second experiment running alongside the first — you pay the full cost again: dequantize the source weights, rotate them, run nearest-centroid search against the codebook, for every layer. Nothing from the last time you did this is reused.

That's the gap. Not "we lack compression" — we had that. "We recompute the compression from scratch on every single load."

Two questions that had to be answered before writing any code​

The ROM framing implies something specific: that multiple processes could share one physical copy of the weights, the way the OS page cache lets multiple processes reading the same file share pages in RAM. Before designing a file format around that idea, it needed to actually be true for MLX. So, two direct tests against the library itself rather than its documentation.

Does mx.load() already give us this for free? I saved a 500 MB tensor to .safetensors, then launched two independent processes that both called mx.load() on it, and read each process's physical footprint with vmmap:

PID 15970: Physical footprint 514.3M
PID 15971: Physical footprint 514.3M

Each process paid the full 514 MB independently — 1.03 GB combined for one 500 MB file. vmmap showed no mapped-file region for the safetensors path at all. The loader reads and copies into a private heap allocation. No sharing, no mmap(), nothing.

Fine — what if I build the mmap myself? I mmap'd a raw binary file with np.memmap (confirmed lazy — RSS didn't move) and wrapped it with mx.array():

RSS after np.memmap (lazy, no touch):        1534.5 MB   (unchanged)
RSS after mx.array() wrap (no eval):         2542.3 MB   (+1008 MB — 2× the file size)
RSS after mx.eval(x):                        2542.3 MB   (no further change)

mx.array() copies the entire buffer immediately, before mx.eval() is even called, and the jump is roughly double the source size — there's a staging copy in there somewhere before the data lands in MLX's own unified-memory arena. There is no zero-copy constructor from a buffer in this version of MLX's Python API.

That killed the strongest version of the pitch. Processes on this machine are not going to share physical pages of model weights through anything MLX gives you today. If a future MLX version adds a real zero-copy buffer path, this is the test to rerun. Until then, the honest framing is narrower: skip the recomputation, not the RAM.

Building the thing anyway​

What survives the two negative findings above is still worth having: a file format that holds a model's already-quantized weights, persisted once, loadable without repeating the compression work. Call it a reservoir instead of a ROM — closer to what it actually is.

The first version was almost embarrassingly simple. Serialize each QuantizedLinear layer's compressed indices and per-row norms into a flat, page-aligned binary file — one blob for the 2–4 bit indices, one for the fp32 norms, a small JSON header recording each layer's shape and the seed used to derive its rotation and codebook. On load, reconstruct each layer from the header and drop the persisted indices straight in, skipping the whole dequantize-rotate-quantize pipeline.

I benchmarked it against quantize_model() on Qwen2.5-0.5B-Instruct-4bit — 168 linear layers, a small enough model to iterate on quickly. The baseline took 81.6 seconds. My new reservoir loader took 66.6 seconds.

That's not a win. That's barely a rounding error.

Where the time actually went​

I profiled it instead of guessing.

ncalls  tottime  cumtime  function
   168    0.126   66.587  QuantizedLinear.__init__
    24    2.079   57.442  make_rotation_matrix
    24   54.693   55.147  numpy.linalg.qr
   168    0.001    8.907  CodebookFactory.create
   168    4.016    8.899  lloyd_max

There it was. 55 of the 66.6 seconds were spent inside np.linalg.qr — the routine that derives a rotation matrix for any layer whose input dimension doesn't cleanly fit MLX's Hadamard transform (a dimension like 4864, which is neither a clean power of two nor a multiple of the special constants MLX's fast transform supports). Another 8.9 seconds went to fitting Lloyd-Max codebooks.

The reservoir file looked like it had skipped requantization. What it actually skipped was the nearest-centroid search — the smallest part of the cost. Every layer's rotation matrix and codebook were being silently recomputed from the stored seed, on every load, exactly as expensive as before. The format was a placebo.

The fix was to stop being clever about "everything is derivable from the seed" and just persist the actual rotation matrices and codebooks, then reconstruct each layer by bypassing QuantizedLinear.__init__'s expensive branches entirely — direct field assignment onto a bare module instead of calling a constructor that recomputes things you already have on disk.

That dropped the load time to somewhere between 0.36 and 9.2 seconds, against the same 81.6-second baseline. Even the conservative end of that range is a 9× speedup; the isolated, repeatable measurement (0.36s, matching cProfile's CPU-time reading almost exactly) is closer to 230×. The gap between those two numbers is real and I haven't root-caused it — plausibly Metal shader warm-up, plausibly GPU dispatch contention from running back-to-back with the baseline in the same process. I'm reporting the range rather than picking the number that looks better.

The part where fixing one problem created a worse one​

Feeling good about the speedup, I checked the file size. The source model is 265 MB on disk. My reservoir file was 2.5 gigabytes.

Nine and a half times larger than the model it was supposedly compressing.

I broke down where the bytes went:

Eighty-six percent of the file was rotation matrices. Specifically, the 24 layers that fall back to QR rotation instead of the fast Hadamard path — each one has in_features = 4864, and a 4864 × 4864 matrix at 4 bytes per float is about 95 MB. Twenty-four of those is 2.17 GB, and that's before touching a single actual weight.

The instinct here is "surely you can store that more compactly." I had the same instinct, and I want to walk through why it doesn't work, because the wrong answer is genuinely tempting.

np.linalg.qr has a mode='raw' option that returns the underlying Householder reflectors instead of the assembled orthogonal matrix — sounds exactly like what you'd want, a compact representation of the same rotation. I checked what shape it actually returns:

h, tau = np.linalg.qr(G, mode="raw")
# h.shape == (4864, 4864)   — 180.5 MB, even bigger than the dense matrix
# tau.shape == (4864,)      — 0.037 MB

Still a full d × d array. The reason isn't an API limitation — it's that a Haar-random orthogonal d × d matrix genuinely contains d(d-1)/2 degrees of freedom. That's Θ(d²) information, full stop. There is no encoding, clever or otherwise, that gets a truly random rotation below quadratic storage, because the matrix doesn't have any structure to exploit. It's not compressible the way the weights are compressible — the weights have statistical structure a codebook can exploit; a random rotation, by construction, doesn't.

So the earlier plan — "store the rotation as reflectors, get a smaller file" — was wrong. Not underexplored. Wrong, provably, in about ten minutes of checking.

Making the tradeoff a choice instead of a default​

Once "compress the rotation matrix" was off the table, what was left was a genuine, irreducible tradeoff: you can have a fast load (persist the rotation matrix, pay the disk space) or a small file (don't persist it, pay the QR cost again on every load). Not both, for any layer that needs QR fallback.

The fix was to stop pretending there was a single right answer and expose the choice:

save_reservoir(model, path, persist_rotation=False)  # default: small file
save_reservoir(model, path, persist_rotation=True)  # fast load, large file

With the default off, the reservoir file for the same model comes out to 349.7 MB — 1.3× the source model, essentially just the compressed weights plus some structural overhead — and load time goes back up to about 59 seconds, since QR-fallback layers regenerate their rotation from the stored seed exactly as the original quantize_model() path does. With it on, you're back to sub-second loads and a 2.5 GB file. Hadamard-compatible layers — 144 of the 168 in this model — are unaffected either way, since their rotation is just a d-length sign vector, cheap to store regardless.

Neither setting is wrong. What was wrong was shipping one of them silently as the only option and calling the result "a smaller reservoir."

What didn't get tested, and why​

The original plan called for benchmarking Qwen3-4B-4bit — a more realistic model size, and specifically what the ideation phase had proposed. quantize_model() on that model reliably killed the process with SIGKILL on this 24 GB machine. I checked the peak memory footprint before the kill: 21.2 GB, against roughly 14–15 GB of actually available memory at the time.

This isn't a bug I introduced. It's the same headroom constraint documented elsewhere in this repo (docs/MEMORY_CONSTRAINT_FINDINGS.md) for a 32B model on the same hardware — quantize_model()'s dequantize-rotate-quantize pipeline makes several full-size intermediate copies rather than working in place, and a 4B model's pipeline apparently needs more headroom than this particular machine has free right now. It's a real, separate problem, and it blocked the concurrent-process memory benchmark I'd wanted to run — the one that would have measured whether four processes loading the reservoir simultaneously actually show smaller RSS than four processes each running quantize_model() from scratch. That benchmark still needs to happen, on a model this machine can actually load, or on a machine with more headroom.

What actually shipped​

• save_reservoir() / load_reservoir() / graft_reservoir() in weight/reservoir.py — a flat, page-aligned binary format, with persist_rotation as an explicit, documented tradeoff rather than a hidden default.
• A _fast_quantized_linear() construction path that builds a QuantizedLinear from persisted state without touching __init__'s QR or Lloyd-Max branches.
• Eleven tests, including a deliberately non-Hadamard-compatible layer (in_features=50) to exercise the QR fallback path directly, and bit-exact round-trip checks in both persist_rotation modes — the loaded model's forward pass output is byte-identical to the freshly quantized one, not just "close."
• A results file with the actual numbers, not just the ones that made the feature look good.

What this experiment is actually about​

None of the interesting findings here came from writing code that worked on the first try. They came from measuring something, getting an answer that contradicted an assumption, and following that instead of the original plan:

• The "shared ROM" framing died the moment two processes independently paid full memory cost for the same file — a five-minute test with vmmap, run before any format design.
• The first reservoir format's near-zero speedup came from profiling instead of trusting that "we skip requantization" was true because the code was structured to look like it should be.
• The Householder-reflector idea died in about ten minutes of checking mode='raw''s actual return shape — a claim that sounded plausible enough to write into a planning doc, and would have stayed there if nobody had gone and checked.

The honest version of this project isn't "we built a ROM chip in software." It's: cross-process sharing doesn't work with MLX's current copy semantics, skipping recomputation is a real and large win once you persist the right things, and file size versus load time is a fundamental tradeoff for anything that needs a truly random rotation — not a bug to be engineered away, just a dial to expose honestly instead of hiding.

That's a smaller claim than the one I started with. It's also the one that's actually true.

How one arbitrary minus sign turned a KV cache compressor from useful into noise — and what it took to prove it on real models​

I had a KV cache compression method in my library called Q-Filters. It had tests. It had docs. It had a benchmark harness with committed results. It shipped in version 0.31.0.

It also didn't work.

Not "worked slightly worse than the paper." Not "worked on some heads." On real trained weights, its scoring signal correlated with true attention at −0.032 — statistically indistinguishable from a coin flip, and pointing the wrong way about half the time.

Fixing it meant implementing the paper properly, writing two Metal kernels, finding a position-encoding bug that made end-to-end measurement impossible, and throwing away two benchmark harnesses that produced confident, meaningless numbers. Along the way the calibrated version went from ppl 598 to 16.3 on the same workload.

This is the whole run, including the parts where I was wrong.

How I fused KIVI's KV-cache quantization into two bit-exact Metal kernels, measured a 1.40×–5.65× op-level speedup, watched it vanish completely end-to-end — and the four separate times my own benchmarks gave me confidently wrong answers along the way.

There's a particular satisfaction in watching a GPU kernel you wrote beat the framework's version by 5×. There's a different feeling entirely when you plug it into the actual model and the tokens-per-second doesn't move at all.

This post is about both, and about the part in between: four occasions where my benchmarks confidently told me something false. One said the kernel was 11× slower than baseline. One said it was 28× faster. One said quantization consumed 97.83% of prefill time, when the real figure is about 1–2%. One showed a clean 127 MB memory saving that turned out to be nothing at all.

All four were my fault. All four produced a clean-looking number with a plausible story attached, which is exactly what made them dangerous.

The kernel is real, it's bit-exact, it ships, and I'd merge it again. But the honest headline is the one above, and the useful content is why both halves of it are true at the same time.

This is the long version, with the complete experimental record — every measurement generation, including the ones that were wrong.

What we're actually optimizing​

If you run a language model locally, the thing that eventually stops you isn't compute. It's the KV cache.

Every token the model has seen leaves behind a key and a value vector in every attention layer. The model needs them to attend to the past, so they stay resident for the whole generation. The cache grows linearly with context — and unlike model weights, which you load once, it grows while you're using it.

On a 7B model at 32k context that's several gigabytes, comparable to the quantized weights themselves. On a Mac with unified memory, it's the difference between a long conversation working and your machine swapping itself to death.

KIVI (Liu et al., ICML 2024) is one answer, and it's the baseline every other algorithm in this library gets measured against. The insight is that keys and values want to be quantized along different axes:

• Keys are quantized per-channel — each channel gets its own scale, computed across a group of tokens.
• Values are quantized per-token — each token gets its own scale, computed across a group of channels.

Why asymmetric? Key tensors have a few channels with consistently huge magnitudes. Quantize per-token and those outliers blow up the scale for every other channel sharing the group. Value tensors lack that structure, and per-token suits them better.

A third piece matters for everything that follows: the most recent residual_length tokens stay in fp16. They're what attention weights most heavily, and they're also the tokens whose group isn't full yet. Once enough fresh tokens accumulate, they get quantized as a batch and folded into the compressed store. That batching event is a flush, and it is the operation this entire post is about.

The quantization itself is textbook asymmetric min/max:

zero  = min(group)
scale = (max(group) - min(group)) / (2^bits - 1)
q     = round((x - zero) / scale)
recon = q * scale + zero

In MLX this is roughly eight array operations: reshape to expose the group axis, min, max, subtract, divide, round, clip, multiply, add, reshape back.

Eight operations means eight kernel launches and — the expensive part — eight round trips to memory. Every intermediate is materialized. The quantized codes, which never needed to exist as a full-size tensor, get written to RAM and read straight back.

That's the target. One fused kernel, one pass, no intermediates.

The kernel​

The layout problem that shapes everything​

KV tensors are [batch, heads, seq, head_dim], row-contiguous. Flatten batch and heads and you get [BH, S, D], where element (bh, s, d) sits at bh*S*D + s*D + d.

Which means the two modes face opposite problems:

• Values (per-token) — the group runs along D, the contiguous axis. Adjacent elements in a group are adjacent in memory.
• Keys (per-channel) — the group runs along S, the strided axis. Adjacent elements are D floats apart. Typically 128.

My first version was one kernel handling both, which meant transposing the key tensor so the token axis became contiguous, then reusing the same code path.

That transpose was the whole problem. It's a full-size materializing copy — exactly the memory traffic the kernel exists to eliminate. I'd removed eight round trips and added one large one back. On the key path, the "optimized" kernel was a net loss.

So I threw it out and wrote two kernels, one per layout.

Kernel A — per-channel keys: one thread, one whole group​

The trick is almost aggressively simple: give each thread an entire group, and don't reduce at all.

uint tid = thread_position_in_grid.x;

const uint BH = x_shape[0];
const uint S  = x_shape[1];
const uint NG = (S + GROUP_SIZE - 1u) / GROUP_SIZE;   // token groups

if (tid >= BH * NG * DHEAD) { return; }

const uint d   = tid % DHEAD;
const uint r   = tid / DHEAD;
const uint grp = r % NG;
const uint bh  = r / NG;

const uint base = bh * S * DHEAD + d;
const uint s0   = grp * GROUP_SIZE;
const uint s1   = min(s0 + GROUP_SIZE, S);

float gmin =  INFINITY;
float gmax = -INFINITY;
for (uint s = s0; s < s1; ++s) {
    float v = float(x[base + s * DHEAD]);
    gmin = min(gmin, v);
    gmax = max(gmax, v);
}

Each thread strides by DHEAD — 128 floats between consecutive reads. In isolation that looks like the worst access pattern available.

But look at the indexing. d = tid % DHEAD means consecutive threads take consecutive channels. At any step of that loop, the 32 threads in a SIMD group read 32 adjacent addresses. The warp's access is fully coalesced. The stride is per-thread; the warp moves through memory as a solid block.

And because a thread owns its group outright, there is no cross-thread reduction. No threadgroup memory, no barriers, no butterfly, no transpose. The strided-looking layout turned out to need the least machinery.

Kernel B — per-token values: one SIMD group, one quantization group​

The contiguous axis wants the mirror image. Lanes split a single group and cooperate:

uint lane = thread_position_in_threadgroup.x;
uint gid  = threadgroup_position_in_grid.x;

// Whole threadgroups exit together, so every lane still reaches the
// butterfly below — a divergent return would deadlock the shuffle.
if (gid >= x_shape[0] * S * NGD) { return; }

float gmin =  INFINITY;
float gmax = -INFINITY;
for (uint i = lane; i < GROUP_SIZE; i += 32u) {
    uint d = d0 + i;
    if (d < d1) {
        float v = float(x[row_base + d]);
        gmin = min(gmin, v);
        gmax = max(gmax, v);
    }
}

// Butterfly: after 5 XOR shuffles every lane holds the group-wide min/max.
for (uint off = 16u; off > 0u; off >>= 1u) {
    gmin = min(gmin, simd_shuffle_xor(gmin, off));
    gmax = max(gmax, simd_shuffle_xor(gmax, off));
}

The simd_shuffle_xor butterfly is the nice part. Five shuffles reduce 32 lanes to a min/max that every lane already holds — no broadcast step. And because lanes advance in lockstep, it needs no threadgroup memory and no barriers, unlike a tree reduction.

Note the comment on the bounds check. That return must be uniform across the threadgroup. If individual lanes bailed early, the survivors would shuffle against threads that no longer exist and the reduction would hang or return garbage. Whole threadgroups exit together, so every lane reaching the butterfly reaches it with all 32 partners intact.

At KIVI's default GROUP_SIZE=32 this is an exact fit: one lane per element, one butterfly, done.

Two optimizations I was sure would work, and didn't​

I assumed caching each group in registers would help — read once, use twice, skip the second global load. Controlled same-process A/B said otherwise:

In the channel kernel, a whole group is GROUP_SIZE floats per thread. The occupancy that costs outweighs the saved loads, which were hitting cache anyway. In the token kernel it's split across 32 lanes, so it's cheap — and buys nothing, for the same cache reason.

Both reverted, and the REG_SLOTS machinery deleted. I also swept threadgroup width and found the defaults (256 for channel, 32 for token) already optimal.

Three ways to be off by one bit​

My acceptance criterion was bit-exactness against the MLX reference — not "close enough," but identical output. That turned out to be the most instructive constraint in the project, because it surfaced three failure modes a tolerance test sails straight past.

1. FMA contraction​

My first parity run failed on 192 of 300 configurations — and every failure was off by exactly 1 ULP. That uniformity is a fingerprint: not an algorithm bug, a rounding difference.

The culprit was q * scale + gmin. Metal's compiler sees a multiply feeding an add and contracts it into a fused multiply-add: one instruction, one rounding. MLX does them separately, with two roundings. Same math, different result on ~0.02% of elements.

The fix is to break the pattern the optimizer looks for:

float prod = q * scale;      // NOT an fma
out[base + s * DHEAD] = T(prod + gmin);

The irony is that the fused version is more accurate — it carries more intermediate precision. But the contract is parity with the reference, not maximum accuracy, so the less accurate version is the correct one. Uncomfortable sentence; right call. 300/300 after the fix.

2. Rounding mode​

Metal's round() is half-away-from-zero. mx.round is half-to-even. They agree on everything except exact .5 codes — rare in random data, and systematically common in real quantization, because uniform grids produce exact midpoints. The fix is rint(). There's a test pinning it with inputs hand-built to land on .5.

3. Padding semantics​

When a group doesn't divide evenly, MLX pads the tail by replicating the edge value, x[..., -1:], not with zeros. Pad with zeros and you've silently dragged gmin to 0 for every ragged group, corrupting the scale. Since every pad slot holds that same value, folding it in once is equivalent to looping:

if (s1 < s0 + GROUP_SIZE) {
    float pad_val = float(x[base + (S - 1u) * DHEAD]);
    gmin = min(gmin, pad_val);
    gmax = max(gmax, pad_val);
}

All three have dedicated regression tests now. They're the kind of bug that produces plausible output — slightly different, never obviously broken. Without bit-exactness as the bar, all three would have shipped.

The bug that made generation hang forever​

My favourite failure of the project, because the fix made the code simpler and the symptom was so much worse than the cause.

mx.fast.metal_kernel JIT-compiles from source. I was passing shape constants through the header as #defines, which lets the compiler turn tid % DHEAD into a shift-and-mask instead of integer division. Good idea for DHEAD — it's the model's head dimension, fixed for the life of a cache.

I did the same for the sequence length.

The sequence length grows by one every decode step. So every token triggered a fresh shader compilation. Generation didn't crash and didn't error — it just stopped. I killed the process after several minutes with no output.

The fix was to pass shape as a runtime buffer. MLX provides x_shape for free on every input, so this meant deleting code, not adding it.

Four benchmarks that lied​

Here the post stops being about GPU programming and starts being about measurement, which is the part I'd actually want to read.

Lie #1 — "Quantization is 97.83% of prefill"​

I wanted to know how much runtime quantization accounted for, so I instrumented the call: timer before, timer after, mx.eval() in between to force the computation. Here's the raw output:

# mlx-community/Llama-3.2-3B-Instruct-4bit  layers=28

PREFILL-dominated (8k prompt, 4 new tokens):
  prefill  metal=False  wall= 20.754s  quant= 20303.3ms (97.83% of wall)  calls=224  [keys 19483.2ms / values  820.0ms]
  prefill  metal=True   wall= 20.940s  quant= 20489.1ms (97.85% of wall)  calls=224  [keys 19805.5ms / values  683.6ms]

DECODE-dominated (2k prompt, 240 new tokens):
  decode   metal=False  wall= 10.114s  quant=  4829.6ms (47.75% of wall)  calls=504  [keys  4550.0ms / values  279.6ms]
  decode   metal=True   wall=  9.952s  quant=  4682.5ms (47.05% of wall)  calls=504  [keys  4457.8ms / values  224.7ms]

Quantization was apparently the entire bottleneck. It was nonsense, and the mistake is in the description above: mx.eval() inside the measured region.

MLX is lazily evaluated. Operations build a graph; nothing computes until something forces it. By calling eval() inside my timer I wasn't measuring quantization — I was measuring every pending operation in the graph at that moment: attention, the MLP, the whole layer stack, all attributed to the one function that happened to trigger the flush.

Notice too that the numbers are self-refuting if you read them properly: metal=False and metal=True report the same 97.8% share. A measurement that can't distinguish the two arms is measuring something other than the thing you changed.

The one genuinely useful output was incidental: calls=224 across 28 layers on an 8k prompt means 8 flushes per layer — which revealed that mlx_lm chunks prefill at 2048 tokens. That number matters later.

Lie #2 — "0.09× and 28.08×"​

Fresh microbenchmark, both flush sizes:

per-flush cost (keys+values, H=8 D=128, 28 layers)

     S    off ms     on ms  speedup |  x28 layers off        on     saved
    32    1.1334   12.4505    0.09x |           31.7ms    348.6ms   -316.9ms
  2048   19.8718    0.7076   28.08x |          556.4ms     19.8ms    536.6ms

Eleven times slower at small sizes, twenty-eight times faster at large ones. I nearly wrote a whole section theorizing about launch-overhead crossovers — there's a tidy story available where fixed dispatch cost dominates at S=32 and bandwidth savings dominate at S=2048.

The story was fiction. I had left an LLM benchmark running in the background on the same GPU.

Same benchmark, idle GPU:

     S    off ms     on ms  speedup |  x28 layers off        on     saved
    32    0.5050    0.3601    1.40x |           14.1ms     10.1ms      4.1ms
  2048    3.8337    0.6785    5.65x |          107.3ms     19.0ms     88.3ms

Both processes were fighting for the same hardware, and contention landed unevenly across runs. These weren't noisy-around-the-truth — off by 15× in one direction and 5× in the other, and they looked like a coherent narrative. That's the dangerous part. Random noise looks random. Contention produces confident, structured, wrong answers.

The downstream damage is worth showing, because the wrong numbers propagated into a wrong prediction:

--- predicted end-to-end (from the CONTENDED numbers) ---
PREFILL 8k prompt: 4 chunks x 537ms saved = 2146ms of ~21000ms  ->  10.2% faster
DECODE 240 tokens: 7 flushes x -316.9ms   = -2218ms of ~10000ms -> -22.18% "faster"

--- predicted end-to-end (from the IDLE numbers) ---
PREFILL 8k prompt: 4 chunks x 88ms saved  =  353ms of ~21000ms  ->   1.7% faster
DECODE 240 tokens: 7 flushes x 4.1ms      =   28ms of ~10000ms  ->   0.28% faster

A predicted 10.2% prefill win and a 22% decode regression, versus the truth of +1.7% and +0.28%. Had I stopped there, I'd have gone hunting for a decode regression that never existed.

Lie #3 — the benchmark that gave four different answers​

Subtler, and I think the most broadly applicable. For the same kernel and the same configuration, my attempts produced 0.31×, 1.0×, 2.16×, and 3.4×. Not scatter around a value — four different conclusions, each internally consistent.

The root cause: I was calling the function repeatedly on the same input tensor.

MLX can recognize it has already computed something and reuse the result. The reference path — eight standard, individually cacheable array ops — benefits enormously. My custom kernel benefits far less. So the "baseline" was quietly handed a shortcut the kernel couldn't take, and every extra repetition widened the gap. Layering best-of-N on top amplified it further, because best-of-N systematically selects the run where caching helped most.

I got this badly wrong. I concluded the kernel was slower, set the feature flag off, and wrote that conclusion into the code and the tests. It was only when two independently-designed methods disagreed with me that I went back:

1. Interleaved A/B — alternate on/off inside one process, so drift and thermal state hit both arms equally.
2. Rotating input pool — cycle 20 distinct tensors so nothing can be reused.

Both landed at 1.36×–2.14×, agreeing with each other and disagreeing with me. I reverted the flag and corrected the tests.

The benchmark that ships in test_kivi_quant.py now does both, and its docstring explains both traps so the next person doesn't re-derive them.

Lie #4 — the 127 MB memory saving that wasn't​

This one I caught only because I ran a third model.

Llama's peak memory, from a clean interleaved run:

  peak memory (GB): fp16=2.604  off=2.726  on=2.599

A 127 MB reduction with the kernel on. And there's a beautiful explanation sitting right there: the fused kernel eliminates MLX's intermediate tensors, so of course the high-water mark drops. Mechanistically plausible, exactly the result I wanted, and it would have made this post better.

Then the other two models came back:

Two of three moved the opposite direction. It's allocator noise — MLX's memory pool responds to allocation ordering in ways unrelated to which kernel ran, and 127 MB out of 2.7 GB sits well inside that.

There's a deeper reason this had to be noise, which I'll come back to at the end: quantize-then-dequantize cannot reduce peak memory, by construction.

The full end-to-end record​

Four generations of end-to-end measurement, in the order I ran them. I'm including the early ones because their disagreement is the point.

Generation 1 — single-shot, three prompt lengths​

First real run. Single-shot timings, no repeats, Llama-3.2-3B-Instruct-4bit, 28 layers, 8 KV heads, head_dim=128:

## PREFILL (prompt tokens/sec)
prompt tok       fp16  kernel off  kernel on  on vs off  flush/layer
       559      488.2       481.6      486.0      1.01x          512
      2059      463.4       451.7      465.7      1.03x         2016
      8209      386.4       355.9      351.6      0.99x         8160

## DECODE (generation tokens/sec, 120 tokens)
config             tok/s  vs fp16   peak MB  KV comp
fp16               46.06     100%       0.0        -
off                46.00     100%       0.0    4.99x
on                 44.24      96%       0.0    4.99x

decode kernel on vs off: 0.962x

Two problems. First, peak MB reads 0.0 — a unit bug on my side: mlx_lm.stream_generate reports peak_memory in GB, and I was dividing by 1024**2 as though it were bytes. The raw JSON shows the real values hiding at 2.48e-06.

Second, and more importantly: decode at 0.962× looks like a 4% regression. Single-shot numbers on one prompt, with no repeat structure — nowhere near enough to distinguish a real regression from thermal drift. Generation 3 is what settles it.

Generation 2 — the sync-instrumented run​

Lie #1 above. Produced the 97.83% figure, which was wrong, and the 2048-token prefill chunking discovery, which was right and load-bearing.

Generation 3 — repeated runs with medians​

Same model, but now multiple repeats per configuration reporting median/min/max, so spread is visible:

## PREFILL (prompt tok/s, max_tokens=4)
  prompt=2065 tok  (~2 chunks of 2048)
    fp16  median=  468.49  min=  448.16  max=  473.22   vs fp16 100.0%
    off   median=  451.76  min=  420.91  max=  464.51   vs fp16  96.4%
    on    median=  460.48  min=  424.83  max=  470.01   vs fp16  98.3%
    -> kernel on vs off: 1.019x

  prompt=8213 tok  (~5 chunks of 2048)
    fp16  median=  300.09  min=  280.98  max=  352.22   vs fp16 100.0%
    off   median=  297.13  min=  291.43  max=  320.70   vs fp16  99.0%
    on    median=  296.32  min=  293.69  max=  309.83   vs fp16  98.7%
    -> kernel on vs off: 0.997x

## DECODE (generation tok/s, 240 tokens, 2k prompt)
  240 new tokens
    fp16  median=   38.48  min=   38.01  max=   40.82   vs fp16 100.0%
    off   median=   40.21  min=   39.72  max=   41.10   vs fp16 104.5%
    on    median=   39.78  min=   37.22  max=   41.05   vs fp16 103.4%
    -> kernel on vs off: 0.989x

  peak memory (GB): fp16=2.604  off=2.726  on=2.599
  KV compression:   off=5.02x  on=5.02x

This table is the most useful thing I measured, and not because of the ratios. Look at the min/max columns.

At 8k prompt, the unchanged fp16 baseline — same code, same model, nothing swapped — ranges from 280.98 to 352.22 tok/s. That's a ±25% spread from thermal state and system scheduling alone, on a configuration where nothing about the code changed between runs.

Now recall the prediction: +1.7% on prefill. Looking for a 1.7% effect through ±25% run-to-run variance is like weighing a signature on a bathroom scale. No amount of care in the on/off comparison fixes that; the instrument simply doesn't resolve the quantity.

Note also that KIVI itself (off, 96.4%) is slightly slower than fp16 at 2k — the quantization work is real, it's just small. And the decode ordering (off at 104.5% of fp16, i.e. faster than no quantization at all) is a tell that we're deep inside noise, since compressing the cache cannot make decode faster than not compressing it.

Generation 1's apparent 0.962× decode regression shows up here as 0.989×, with overlapping min/max ranges. It was drift.

Generation 4 — three models, interleaved, single process​

The final protocol: one process per model, configurations interleaved rather than run in blocks, output text compared byte-for-byte between arms.

### mlx-community/Llama-3.2-3B-Instruct-4bit
  layers=28 kv_heads=8 prompt=2065 tok
  -> kernel on vs off:  prefill 1.019x (2k) / 0.997x (8k)   decode 0.989x
  -> identical text on/off: True

### mlx-community/Qwen2.5-7B-Instruct-4bit
  layers=28 kv_heads=4 prompt=2064 tok
  config  prefill tok/s  decode tok/s   peak GB   KV comp
  fp16            206.0         23.61     5.105         -
  off             203.7         23.40     5.080     4.94x
  on              208.3         23.26     5.130     4.94x
  -> kernel on vs off:  prefill 1.023x   decode 0.994x
  -> identical text on/off: True

### mlx-community/Mistral-7B-Instruct-v0.3-4bit
  layers=32 kv_heads=8 prompt=2054 tok
  config  prefill tok/s  decode tok/s   peak GB   KV comp
  fp16            143.4         20.86     4.891         -
  off             140.0         20.34     4.921     4.75x
  on              140.2         20.26     4.941     4.75x
  -> kernel on vs off:  prefill 1.001x   decode 0.996x
  -> identical text on/off: True

Consolidated:

Everything within ±2%, which given ±25% baseline variance is indistinguishable from nothing.

Qwen is the most valuable row. Four KV heads instead of eight means a completely different flush geometry, exercising different bounds-check and ragged-tail paths. It still produces byte-identical output — which is the strongest evidence that the bit-exactness work held up outside the unit tests.

Why it's invisible, and why that was predictable​

The true per-flush picture, idle GPU, Apple M4, 8 KV heads × 128 head dim, keys and values together:

Recall that mlx_lm chunks prefill at 2048 tokens. That's a happy accident: prefill flushes land exactly on the kernel's strongest case, where there's enough work to amortize dispatch and memory-traffic savings dominate. Decode flushes are always small — residual_length tokens, 32 here — the weak case.

Scale it to 28 layers, an 8k prompt (4 prefill chunks), 240 decode tokens (7 flushes):

There's the whole story, available before running a single model.

A 5.65× speedup on 1.7% of the work is a 1.7% speedup. Amdahl's law doesn't care how good the kernel is.

And 1.7% is four times smaller than the noise floor of the measurement. The end-to-end result wasn't a disappointment — it was arithmetic, and I could have computed it in ten minutes before writing any Metal at all.

So why does the kernel ship?​

Given that it's invisible end-to-end, why merge it?

• It's free. Bit-exact output, no regression on any model, 84 dedicated tests, byte-identical generations across three architectures. It never makes anything slower.
• It removes a floor. Quantization is ~1–2% of runtime now. If the surrounding work gets faster — better attention kernels, better matmuls — that share grows. Fixed costs matter more as everything else shrinks.
• Op-level wins are real even when invisible. 1.40× and 5.65× are honest measurements of the operation. That the operation is a small slice of the whole is a separate fact, and both belong in the report.

But the real reason to stay clear-eyed: the kernel was never where the memory win lived.

KIVI as implemented does quantize-then-dequantize — it computes the compressed representation and immediately expands it back to fp16 for attention. The 4.75×–5.02× compression is real arithmetic, but it is an accounting result, not a storage result. The tensor sitting in memory is still fp16.

This is also the structural reason Lie #4 had to be noise: if nothing is stored in compressed form, no kernel that computes the compression faster can reduce the high-water mark. I should have known the 127 MB was suspect on those grounds alone, before the other two models contradicted it.

To actually reduce memory you need two more things:

1. Packed storage — keep quantized codes as uint8, never materialize the fp16 reconstruction.
2. Dequant-in-SDPA — teach attention to read packed codes directly, so expansion happens in registers and never in RAM.

That's where both the memory win and the real speedup live, because it doesn't fuse 1–2% of the work — it shrinks the tensors every other operation has to move.

The kernel was step one. It was worth doing. It just isn't the point.

What I'd take away​

If you're writing GPU kernels against a framework like MLX or PyTorch:

Match the thread mapping to the memory layout, not to intuition. The strided access pattern needed less machinery than the contiguous one — no reduction, no barriers, no transpose. My first instinct, transposing to make the layout "nice," was the version that lost.

Bit-exactness is a debugging tool, not just a correctness bar. FMA contraction, rounding mode, and padding semantics all produce plausible output. A tolerance test passes all three. Demanding identical output turned three silent behavioral differences into three failing tests with obvious causes.

Never specialize on a value that grows. Baking sequence length into a JIT header compiles one shader per token. The symptom was an indefinite hang; the fix deleted code.

Estimate the ceiling before you optimize. Ten minutes with Amdahl's law would have predicted the end-to-end result up front. It wouldn't have changed the decision to build it — but it would have set the expectation correctly, and I'd have spent my time on the parts that were load-bearing.

Measure your noise floor before your effect. The single most useful number in this entire project was ±25% — the run-to-run spread of an unchanged baseline. Without it, every ratio in every table is unfalsifiable.

Be most suspicious when the number is good. All four bad measurements came with satisfying stories. 97.83% "proved" the work mattered. 28.08× had a tidy overhead-crossover explanation. The 127 MB saving had a clean mechanism. Every one was wrong, and the plausible explanation is what let each survive as long as it did.

And: run the third model.

Kernels live in veloxquant_mlx/metal/src/, tests in veloxquant_mlx/tests/metal/test_kivi_quant.py. See the KIVI algorithm reference for the method itself, and Metal kernels for how kernels are dispatched library-wide. KIVI: Liu et al., ICML 2024. All measurements on an Apple M4 with MLX.

A deep-dive into Apple's Metal Shading Language specification, what TensorOps promised, why it didn't work through MLX, and the two real improvements we shipped from three sessions of research.

Where This Story Starts​

A few weeks ago I shipped a FlashAttention-style Metal kernel for VeloxQuant-MLX that was correct, fast in isolation, and completely useless end-to-end. The blog post about that mistake is here. The short version: I built a fused dequant+attention kernel that beat mx.fast.scaled_dot_product_attention by 1.3× in benchmarks — then discovered mlx_lm had already eliminated the dequant cost via a persistent fp16 K_hat buffer, making my kernel 3-4× slower than the baseline it was supposed to beat.

I kept the kernel in the library as an opt-in API. It's correct, it's tested, and it loses.

After writing that post, a reader suggested I look at the Metal Shading Language specification — specifically Metal 4, which Apple released with macOS Sequoia. The argument was: Metal 4 adds hardware tensor operations that could replace the slow part of the kernel. Maybe there was a path to winning that I hadn't found yet.

So I read the spec. All 346 pages of it.

This is what I found.

The Kernel's Hot Path​

To understand why the spec research mattered, I need to explain the bottleneck.

The fused SDPA kernel computes attention directly from VecInfer's compressed codebook indices without materializing the fp16 key matrix. For each query, it needs to compute a Look-Up Table first:

LUT[sub, centroid] = q_sub_vector · codebook_row[centroid]

For VecInfer's default config (n_sub=16, sub_dim=8, n_centroids=256), this is a [16, 8] @ [8, 256] matrix multiply — 4,096 dot products. In the current kernel, 32 GPU lanes stripe these across the SIMD group: each lane computes 128 scalar dot products independently.

This LUT precompute is Phase 0. Everything else — the online softmax, the V accumulation — comes after. If the LUT is slow, everything is slow.

The Metal 4 spec describes two potential hardware paths to speed this up:

1. simdgroup_float8x8 (Metal 2.3+, Section 2.4 / 6.7): 8×8 hardware matmul tiles via simdgroup_multiply_accumulate. Available today.
2. TensorOps matmul2d (Metal 4+, Section 7.2): A full hardware matrix multiply API with a cooperative_tensor destination. Potentially much faster.

I tested both.

Attempt 1: simdgroup_float8x8​

The Metal spec (Section 6.7, Table 6.9) shows simdgroup_float8x8 as a cooperative 8×8 float matrix multiply tile. The <metal_simdgroup_matrix> header is accessible via MLX's header= parameter:

k = mx.fast.metal_kernel(
    name="my_kernel",
    source=src,
    header="#include <metal_simdgroup_matrix>\nusing namespace metal;\n",
    ...
)

The tiling plan for our LUT: n_sub=16 rows / 8 = 2 row-tiles, n_centroids=256 cols / 8 = 32 col-tiles, sub_dim=8 = 1 K-tile. Total: 64 hardware matmul operations.

I implemented it. Correctness test: zero diff vs reference.

Then I benchmarked it against the current scalar loop:

scalar loop:     212 µs per LUT precompute
simdgroup 8×8:   255 µs per LUT precompute

The hardware matrix multiply was slower.

The reason is protocol overhead. simdgroup_float8x8 is a cooperative operation — all 32 lanes must execute each tile in lock-step. For our 64 tile iterations, that's 64 synchronization points. The scalar loop has zero synchronization: each lane independently computes 128 dot products in parallel. For a small matrix like [16,8]@[8,256], the cooperation overhead dominates the compute savings.

simdgroup_matrix wins at large, batched matmuls (MLX uses it for GEMM with 128×128 tiles). For our 16×256 LUT, it's the wrong tool. Reverted.

Attempt 2: Metal 4 TensorOps​

Section 7.2 of the spec describes tensor_ops::matmul2d — a hardware-accelerated matrix multiply that operates on tensor<> types and writes to a cooperative_tensor destination held in thread registers. The pitch is exactly right: no threadgroup memory round-trip, hardware tensor units, single API call.

The example from the spec:

#include <metal_tensor>
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
using namespace metal;
using namespace mpp;

[[ kernel ]] void matrixMultiply(
    tensor<device half, dextents<int, 2>> a [[ buffer(0) ]],
    tensor<device half, dextents<int, 2>> b [[ buffer(1) ]],
    tensor<device half, dextents<int, 2>> c [[ buffer(2) ]]) {

    constexpr auto desc = tensor_ops::matmul2d_descriptor(64, 32, 0);
    tensor_ops::matmul2d<desc, execution_simdgroups<4>> op;
    matmulOp.run(a, b, c);
}

Clean. Exactly what we need.

I confirmed the header is accessible:

k = mx.fast.metal_kernel(
    name="test",
    source=src,
    header="""
    #include <metal_tensor>
    #include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
    using namespace metal;
    using namespace mpp;
    """,
)
# Compiles. Header is reachable.

And Metal 4 is available:

Metal version: 400.0   (M4, macOS Sequoia)

Then I tried to actually use matmul2d. Three blockers, in order of discovery:

Blocker 1: Type support.

Table 7.3 of the spec lists supported type combinations. float/float/float is listed — but when I tried it:

static_assert failed: "Unsupported type"

Table 7.4 clarifies: bfloat/bfloat/bfloat and several mixed-precision combinations require OS 26.1 and later. That's iOS/macOS naming — it maps to macOS 26.1 (not released yet as of this writing). The float/float/float path in Table 7.3 is supported, but only with certain execution_scope + K-dimension combinations that are hardware-dependent.

Blocker 2: tensor_handle vs tensor_inline.

The spec's matmul2d example uses tensors declared as kernel parameters with [[buffer(N)]] attributes — these are tensor_handle type. MLX's metal_kernel generates the function signature automatically: it only creates raw pointer buffers (const device float* a [[buffer(0)]]), not tensor<device half, ..., tensor_handle> parameters.

The only tensor type you can construct at runtime from a pointer is tensor_inline. But cooperative_tensor.store() only accepts tensor_handle targets for device memory writes. The round-trip cooperative_tensor → tensor_inline → device output is blocked:

error: candidate template ignored: could not match 'tensor_handle' against 'tensor_inline'

Blocker 3: Dynamic K hangs the GPU compiler.

When I tried K=0 (dynamic length, matching the spec example exactly), the MLX JIT compilation hung. The TensorOps template instantiation with dynamic_length_v<int> appears to trigger a very long (possibly infinite) compile path under MLX's inline Metal JIT. The process never returned.

Summary: TensorOps is architecturally incompatible with MLX's mx.fast.metal_kernel API. The API generates raw pointer buffers; TensorOps requires tensor-typed formal parameters. The mismatch is fundamental, not a workaround.

What Actually Worked​

Two improvements from the spec research did ship.

1. metal::precise::exp — a correctness fix hiding as a performance question​

Section 8.2 of the spec describes rounding mode. Section 8.3 covers floating-point exceptions. Table 8.2 documents accuracy under fast math.

The relevant line: exp() in fast math mode (-fmetal-math-mode=fast) does not guarantee exp(-INFINITY) = 0.0. The spec's ULP table for fast math lists relaxed accuracy bounds for transcendentals.

Our kernel uses exp(score - running_max) for the online softmax. When a lane is masked (causal or sliding-window), we set score = -INFINITY. In fast math mode, exp(-INFINITY) may not be exactly 0.0 — which would corrupt the softmax denominator.

The fix: use the metal::precise:: namespace to force IEEE-compliant exp regardless of compiler math mode:

// Before (math-mode dependent):
float w = exp(score - m_new);

// After (always correct):
float w = metal::precise::exp(score - m_new);

MLX's metal_kernel API has no parameter for compiler flags, so -fmetal-math-mode=relaxed isn't accessible. The namespace workaround is better anyway — it's surgical, affects only these two exp calls, and documents intent in the code.

2. simd_broadcast_first — eliminating two threadgroup barriers per tile​

Section 6.9.2 of the spec (Table 6.14) lists the full SIMD-group permute function set. One entry:

simd_broadcast_first(x)  →  broadcasts lane 0's value to all lanes
                             without a threadgroup barrier

The original kernel used threadgroup memory to share the running max and rescale factor:

// Before: two threadgroup variables, two barriers per tile
if (lane == 0) {
    tg_m_shared = m_new;
    tg_factor   = factor;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
m_new  = tg_m_shared;
factor = tg_factor;

With simd_broadcast_first, both threadgroup variables disappear entirely — running_m becomes a lane-local float that all 32 lanes keep synchronized:

// After: no threadgroup variables, no barriers for scalar sharing
float m_new  = simd_broadcast_first(max(running_m, tile_max));
float factor = simd_broadcast_first(
    isfinite(running_m) ? metal::precise::exp(running_m - m_new) : 0.0f);
running_m = m_new;

With S_kv=4096 and 128 tiles, this removes 256 threadgroup barriers from the hot loop. Threadgroup barriers are expensive — they serialize the entire threadgroup and flush threadgroup memory. Removing them reduces both latency and the register pressure from storing shared state.

Both of these are in the current kernel. 9 parity tests pass. The improvements are real even if the end-to-end situation hasn't changed.

The Actual Answer to "What Is Section 7.2 Useful For?"​

TensorOps would be transformative if MLX supported tensor-typed kernel parameters. The current mx.fast.metal_kernel API exposes only raw device pointers — the [[buffer(N)]] binding that TensorOps needs is auto-generated as const device float*, not tensor<device half, ..., tensor_handle>.

To use TensorOps for our LUT precompute, MLX would need one of:

1. Support tensor<> as a formal parameter type in metal_kernel's auto-generated signature. Something like input_tensor_types=[("a", mx.float16, 2)] that generates tensor<device half, dextents<int,2>> a [[buffer(0)]].

2. A new mx.fast.metal_tensor_kernel variant that accepts tensor operands natively and dispatches via TensorOps internally.

This is exactly the GitHub issue we filed at ml-explore/mlx. The issue covers three requests — compiler options, integer template parameters, and Metal 4 tensor type access — all confirmed by direct testing.

The Broader Pattern​

Three sessions, three attempts at the LUT precompute, three different techniques:

The pattern: each attempt was technically sound, correctly implemented, and blocked by something orthogonal to the GPU math.

• Simdgroup matrix: the hardware works, the tile size is wrong.
• TensorOps: the hardware works, the API binding doesn't exist.

In both cases, the blocker wasn't that the hardware was slow. The blocker was that the interface between our code and the hardware had a constraint we couldn't see until we hit it.

The right mental model for GPU kernel work on Apple Silicon: there are three layers — the math you want to do, the hardware that can do it, and the API that connects them. Breakthroughs happen at the API layer, not the math layer. The math for attention has been solved. The hardware for matrix multiply has been built. The gap is the binding.

That gap is the GitHub issue. If MLX adds tensor<> support to metal_kernel, this whole investigation becomes a one-afternoon project. Until then, the scalar LUT is the fastest thing we can write.

What Is in the Library Now​

veloxquant_mlx/metal/fused_sdpa.py has:

• metal::precise::exp for both softmax exp calls — correctness guarantee regardless of MLX math mode
• simd_broadcast_first replacing threadgroup barriers for running_m — 256 fewer barriers at S_kv=4096
• tg_m_shared and tg_factor threadgroup variables removed — smaller threadgroup memory footprint
• All 9 parity tests passing: causal, non-causal, sliding-window, GQA, short-sequence, long-sequence, dispatcher patch

The end-to-end situation is unchanged from the previous post — the kernel only helps if mlx_lm exposes a way to skip K_hat materialization, which requires an upstream change.

But the kernel is now more correct and slightly better engineered. That's what reading 346 pages of a GPU spec gets you when the hardware feature you wanted is one API version away.

The One Practical Takeaway​

Before spending time implementing a GPU optimization, answer this question:

Which layer is blocking you — the math, the hardware, or the API?

If the math is solved and the hardware exists, the answer is almost always the API. Find the API gap first. File the issue or write the binding. Don't write the kernel until the API exists to call it from.

I wrote the kernel first. I found the API gap last. Three sessions later.

Links​

• VeloxQuant-MLX on PyPI: pypi.org/project/VeloxQuant-MLX
• GitHub: github.com/rajveer43/VeloxQuant-MLX
• The kernel: veloxquant_mlx/metal/fused_sdpa.py
• MLX issue filed: github.com/ml-explore/mlx/issues
• Previous post (Phase 2 mistake): Medium — I Spent 8 Hours Writing a FlashAttention Kernel
• Previous post (Phase 1 win): Medium — I Wrote a Metal Kernel to Stop My Mac From OOMing

TL;DR — VeloxQuant-MLX 0.8.0 adds KIVI (Liu, Yuan et al., ICML 2024, arXiv:2402.02750), a faithful re-implementation of the most-cited KV-cache quantization baseline. It's asymmetric (per-channel keys, per-token values), keeps a small fp16 residual window, and is fully deterministic. On an Apple M4, across Llama-3.2-3B, Qwen2.5-7B, and Mistral-7B (4-bit, ~2.2–2.4k-token prompts), KIVI-2bit measured 5.8× key / ~4.0× full-KV compression at roughly fp16 throughput. On Apple Silicon the win is memory, not speed — and we report the throughput as measured, not as hoped. Every number below traces to a committed results.json.

The wall, again​

If you run local LLMs on a Mac, you've met the unified-memory wall. The CPU, GPU, and Neural Engine all share one pool of RAM, and the KV cache — the keys and values the model stores for every past token — grows linearly with context length until it crowds out everything else. Weight quantization (GGUF, GPTQ, AWQ) is an offline, one-time compression of the model's parameters; it does nothing for the cache, which is rebuilt token by token at inference time.

VeloxQuant-MLX exists to compress that cache. It already ships a suite of methods — TurboQuant, RVQ, VecInfer, SpectralQuant, RaBitQ, CommVQ, QJL, PolarQuant, RateQuant. With 0.8.0 we added the one method a reviewer asks about first, and which the library conspicuously lacked: KIVI.

What KIVI is​

KIVI is "A Tuning-Free Asymmetric 2-bit Quantization for KV Cache" (arXiv:2402.02750, ICML 2024). Its central observation is that keys and values want different quantization layouts:

• Keys → per channel. Key distributions have a few high-variance channels. Quantizing along the token axis, one (scale, zero) per channel-group, keeps those channels accurate.
• Values → per token. Value distributions are flatter across channels but vary token to token, so the group runs along the channel axis instead.
• A residual window. The most recently generated tokens dominate attention and are cheap to keep exact, so KIVI holds the last R tokens in fp16 and only quantizes tokens once they age out of that window.

The quantizer itself is plain asymmetric min/max group quantization — no codebook, no calibration, no randomness:

for each group g (a slice along the quantization axis):
    zero  = min(g)
    scale = (max(g) - min(g)) / (2**b - 1)
    q     = round((g - zero) / scale)      # uint code in [0, 2**b-1]
    g_hat = q * scale + zero               # asymmetric dequant

Because there's no k-means and no RNG, KIVI is deterministic — same input, identical output, every run. That matters in this codebase: our vector-quantization methods train codebooks and can vary run to run; KIVI adds none of that.

Why we added it (and what we didn't invent)​

To be clear about credit: KIVI is Liu, Yuan, and colleagues' algorithm. We ported it to MLX. There is no novelty claim here.

The value is comparative. KIVI is the field's reference baseline — nearly every KV-cache paper measures against it. Until 0.8.0, VeloxQuant-MLX couldn't answer "how does your method compare to KIVI?" Now every other method in the library has a recognized point of comparison, in the same framework, on the same Apple-Silicon hardware. (The full reasoning, including why we chose KIVI over KVQuant, GEAR, and ZipCache, is in paper/NEW_METHOD_SURVEY.md: KIVI was the highest-value missing baseline, deterministic, and a clean architectural fit that needs no RoPE or attention-score hooks.)

How it plugs in​

Three lines, and mlx_lm.generate runs unchanged:

import mlx_lm
from veloxquant_mlx import KVCacheConfig, KVCacheBuilder

model, tokenizer = mlx_lm.load("mlx-community/Llama-3.2-3B-Instruct-4bit")

config = KVCacheConfig(
    method="kivi",
    bit_width_inlier=2,  # KIVI's default 2-bit
    kivi_group_size=32,  # min/max group size (paper default)
    residual_length=32,  # recent tokens kept in fp16
)
caches = KVCacheBuilder.for_model(model, config)
model.make_cache = lambda *_a, **_k: caches

response = mlx_lm.generate(
    model,
    tokenizer,
    prompt="Summarize the attention mechanism in three sentences.",
    max_tokens=300,
)

The cache quantizes the aged-out keys/values and immediately dequantizes them, so the downstream scaled-dot-product attention sees standard fp16 tensors — no model surgery, no custom attention path.

Results​

All numbers below come from figures/kivi/results_summary.json (aggregated from the per-model figures/kivi/<model>/results.json). Conditions, identical across runs: Apple M4, 24 GB, 4-bit mlx-community models, group_size=32, residual_length=32, max_tokens=120, prompts of ~2.2–2.4k tokens (a long prompt is required for KIVI to actually exercise the quantized path rather than sit inside the fp16 residual window).

KIVI-2bit, across three models:

Bit-width sweep (Llama-3.2-3B):

Two honest observations from this data:

1. Throughput is flat, not faster. KIVI here runs at 100–105% of fp16 — i.e. it does not slow generation down at these sizes, but it also doesn't speed it up. The published KIVI speedups come from a fused CUDA kernel; that kernel does not port to Metal. On Apple Silicon the deliverable is memory compression (the 4× full-KV figure), and the throughput column is there so you can see it costs you nothing, not so you can claim a speedup.
2. Full-KV compression is lower than key-only, and deliberately so. The full-KV figure includes the fp16 residual window. At residual_length=32 over a 120-token generation, a meaningful slice stays in fp16; that drags the end-to-end ratio below the key-only number. We report both rather than quoting the flattering one.

The implementation is covered by 25 passing tests (veloxquant_mlx/tests/quantizers/test_kivi.py and tests/cache/test_kivi_cache.py), including reconstruction-fidelity bounds, the per-token/per-channel asymmetry, the residual-window behavior, and a determinism check.

Honest limitations​

• Memory win, not a speed win on Metal. As above — the CUDA kernel fusion from the paper isn't available here. If you need raw throughput, this isn't your lever.
• Quality is measured as reconstruction fidelity, not task accuracy. Our tests check cosine similarity / MSE against the fp16 cache on synthetic and real key distributions. We have not run LongBench or a perplexity sweep across configs, so we do not claim "no quality loss" on downstream tasks.
• 2-bit is genuinely lossy. On unit-norm synthetic keys, KIVI-2bit reconstruction cosine sits around 0.93 — which is exactly why KIVI keeps an fp16 residual window. If you push residual_length to 0 you'll feel it. The defaults exist for a reason.
• Single chip, short generations. Everything above is one M4 at ~120 generated tokens. Behavior across other M-series tiers and very long generations isn't characterized yet.

Where KIVI fits​

• Reach for KIVI when you want a simple, deterministic, calibration-free 2-bit baseline — or when you specifically need to compare against the literature's reference point.
• Reach for RVQ when you want stronger compression-per-bit with zero calibration and near-fp16 throughput (its analytical codebooks do better than scalar min/max at low bit-rates).
• Reach for VecInfer when you want the most aggressive key compression (up to 16× key-only) and have ~2 minutes for codebook calibration.

KIVI's job in the suite isn't to win every axis; it's to be the honest yardstick the others are measured against.

Try it​

pip install VeloxQuant-MLX==0.8.0

• Docs: /docs/algorithms/kivi/
• GitHub: https://github.com/rajveer43/VeloxQuant-MLX
• Paper: KIVI, arXiv:2402.02750 (ICML 2024)

What was measured vs. not: all compression, throughput, peak-memory, and token-count figures are from committed figures/kivi/*/results.json on a single Apple M4 (24 GB) at ~120 generated tokens with ~2.2–2.4k-token prompts; correctness is from 25 passing unit tests. We did not measure downstream-task accuracy (e.g. LongBench, perplexity sweeps) — "quality" here means reconstruction fidelity against the fp16 cache.

How to cut your KV cache memory by 8x and hit fp16 throughput — step by step, with real models.

If you have an Apple Silicon Mac and you run local LLMs, you have already hit the wall: your MacBook Pro generates 20 tokens per second on Mistral 7B, memory pressure turns on after a few thousand tokens, and anything larger than 7B crawls. The bottleneck is not compute — it is memory bandwidth. Every token requires loading the entire KV cache for every layer. Make the cache smaller, and the model gets faster.

VeloxQuant-MLX is a drop-in KV cache quantization library for MLX. It compresses the KV cache to 2, 3, or 4 bits per value while keeping text quality close to full-precision. In this guide you will install it, run your first compressed model, pick the right algorithm for your use case, and understand what the benchmark numbers mean. No machine learning background required.

What You Need​

• Apple Silicon Mac (M1 or later)
• Python 3.11 or 3.12
• At least 16 GB unified memory (8 GB works for 4B models)
• mlx-lm installed (pip install mlx-lm)

Step 1: Install VeloxQuant-MLX​

pip install VeloxQuant-MLX

That installs the mlx_kv_quant package plus the veloxquant CLI. Verify it worked:

veloxquant --help

You should see the precompute and benchmark subcommands.

Step 2: Your First Compressed Inference​

The fastest path uses the KVCacheBuilder. Here is a minimal script that runs Mistral 7B with 4-bit KV cache compression:

import mlx.core as mx
import mlx_lm
from mlx_kv_quant import KVCacheBuilder, KVCacheConfig

# Load the model normally
model, tokenizer = mlx_lm.load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")

# Build a quantized KV cache
config = KVCacheConfig(bits=4, algorithm="turboquant_prod")
cache = KVCacheBuilder(config).for_model(model).build()

# Generate with the compressed cache
prompt = "Explain the difference between RAM and unified memory."
tokens = mlx_lm.generate(model, tokenizer, prompt=prompt, kv_cache=cache, max_tokens=200)
print(tokens)

That is the entire integration. The cache is transparent to mlx_lm.generate() — you do not change anything else.

Step 3: Understanding the Algorithm Choices​

VeloxQuant-MLX ships five quantization algorithms. They trade reconstruction quality against memory and speed:

Rule of thumb:

• For most uses at 3–4 bit: use turboquant_prod. It is the default and the most tested.
• At 2 bits: always use turboquant_rvq. The single-pass algorithms fall apart at 2-bit; RVQ does not.
• For long contexts where cosine similarity matters more than reconstruction: polarquant.
• To stress-test memory limits: qjl.

Step 4: Using RVQ 2-Bit (the Interesting One)​

The headline feature in v0.3.0 is TurboQuantRVQ — a two-pass residual vector quantizer that makes 2-bit KV cache actually usable.

The problem with naive 2-bit quantization: you only have 4 levels to represent a continuous value. The reconstruction error is large enough that attention scores get corrupted, and long reasoning chains (like Qwen3's <think> mode) collapse into repetition after a few dozen tokens.

RVQ fixes this by running two codebooks. The first pass quantizes the rotated vector. The second pass quantizes the residual — the error left over from the first pass — using a Laplacian-fit codebook (which matches the residual distribution better than a Gaussian one). The result: cosine similarity goes from 0.69 to 0.98 at 2 bits.

Here is how to use it:

from mlx_kv_quant import KVCacheBuilder, KVCacheConfig

config = KVCacheConfig(bits=2, algorithm="turboquant_rvq")
cache = KVCacheBuilder(config).for_model(model).build()

That's it. The turboquant_rvq string routes to the registered TurboQuantRVQ class via the quantizer registry.

When should you use RVQ 2-bit?

Use it when memory is the binding constraint and you need coherent long-form output. It delivers:

• Mistral 7B at 2-bit: 22.3 tok/s — matches fp16 throughput (22.1 tok/s) on an M4 MacBook
• Qwen3 4B at 2-bit: 36.0 tok/s (92% of fp16) with full thinking-mode output

Use standard 3-bit or 4-bit when you have memory headroom and want zero quality trade-off.

Step 5: Running Benchmarks on Your Mac​

VeloxQuant-MLX ships benchmark scripts for several models. These measure tok/s and output completeness across five configurations: fp16, 2-bit RVQ, 3-bit, 4-bit, and 2-bit single-pass (for comparison).

Qwen3 4B (good for 16GB Macs):

python3 benchmark_qwen3_4b_v2.py

Mistral 7B (needs 16GB+, comfortable at 32GB):

python3 benchmark_mistral7b_v2.py

Each script runs all five configs, prints a results table, and saves six figures to figures/updated_tests/<model>/. The figures include:

• Throughput comparison bar chart
• Token completeness (how many tokens were coherent out of 200)
• Memory usage per config
• Cosine similarity vs fp16 reference

Here is what our numbers look like on an M4 MacBook (16GB):

Mistral 7B throughput:

Mistral 7B is memory-bandwidth bound — every config saturates around 22 tok/s because that is the bandwidth ceiling. But RVQ 2-bit uses 10x less cache memory than fp16.

Qwen3 4B thinking-mode throughput:

The 4-bit single-pass result (50/200) is the failure mode for standard quantization on a thinking model. RVQ 2-bit produces full coherent output and runs faster.

Step 6: Integrating with mlx-lm Properly​

For real usage you want the cache to persist across generation steps. Here is the pattern for a chat loop:

import mlx.core as mx
import mlx_lm
from mlx_kv_quant import KVCacheBuilder, KVCacheConfig

model, tokenizer = mlx_lm.load("mlx-community/Qwen3-4B")

config = KVCacheConfig(bits=2, algorithm="turboquant_rvq")
cache = KVCacheBuilder(config).for_model(model).build()

messages = []

while True:
    user_input = input("You: ")
    if user_input.lower() in ("exit", "quit"):
        break

    messages.append({"role": "user", "content": user_input})
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    response = mlx_lm.generate(
        model,
        tokenizer,
        prompt=prompt,
        kv_cache=cache,
        max_tokens=512,
        verbose=False,
    )

    messages.append({"role": "assistant", "content": response})
    print(f"Assistant: {response}\n")

The kv_cache argument is passed directly to mlx_lm.generate(). No other changes. The cache grows with each turn and is quantized transparently.

Step 7: Precomputing Codebooks for Repeated Use​

If you run the same model repeatedly, you can precompute the Lloyd-Max codebooks offline and load them at inference time. This eliminates the one-time calibration cost at startup:

# Precompute and save codebooks for Mistral 7B, 4-bit
veloxquant precompute \
  --model mlx-community/Mistral-7B-Instruct-v0.3-4bit \
  --bits 4 \
  --algorithm turboquant_prod \
  --output codebooks/mistral7b_4bit/

# Later, load at inference time
config = KVCacheConfig(
    bits=4,
    algorithm="turboquant_prod",
    artifact_dir="codebooks/mistral7b_4bit/",
)
cache = KVCacheBuilder(config).for_model(model).build()

Codebook size is tiny (kilobytes). You can commit them to your project directory and avoid the calibration step entirely on every run.

Step 8: Picking the Right Bit Width​

Here is the practical decision tree:

Do you have >= 32GB memory?
  Yes → fp16 is fine; use VeloxQuant only for very long contexts
  No  → continue

Is your model >= 7B parameters?
  Yes → 3-bit TurboQuant (turboquant_prod)
  No  → 2-bit RVQ (turboquant_rvq)

Are you doing RAG or similarity search over the KV cache?
  Yes → turboquant_prod (inner-product correction matters)
  No  → turboquant_mse (simpler, same quality)

Do you need to run 13B+ on a 16GB Mac?
  Yes → 2-bit RVQ — the only option that keeps quality

For generation tasks (chat, summarization, code completion), MSE-optimal reconstruction is what you want. For retrieval tasks that compute attention over a cached context, inner-product preservation matters more — use turboquant_prod.

Step 9: What the Optimization Journey Looked Like​

Getting from 17.7 tok/s to 22.3 tok/s on Mistral 7B RVQ 2-bit required four changes, each independently measurable:

1. Batch all heads into one MLX call (+22%)

The original code ran a Python for h in range(H) loop over attention heads, giving each head its own quantizer instance. For Mistral 7B (8 heads, 32 layers) that meant 256 small kernel dispatches per token. The fix: reshape (B, H, S, D) to (B*H*S, D) and call one shared quantizer once. MLX already handles batched input — we just stopped fragmenting it.

2. Switch to Hadamard rotation

The default preconditioner used QR decomposition: a full (d, d) matrix multiply (16,384 ops at d=128). mx.hadamard_transform is a Metal-native fused kernel that does the same job in O(d log d) — 896 ops at d=128, ~18x less arithmetic. Quality is mathematically identical because Hadamard with random ±1 diagonal is also a Haar-equivalent rotation.

3. Replace broadcast-argmin with boundary-sum

Codebook lookup materialized a (batch, 128, k) distance tensor for argmin — three kernels: broadcast-subtract, abs, argmin. Lloyd-Max boundaries are just midpoints between sorted centroids. The nearest centroid index is the number of boundaries the value exceeds: one comparison and one sum, two kernels total. Index match vs the old path: 100.0000%.

4. Drop redundant fp32 casts

The update path was casting fp16 → fp32 → fp16 → fp32 → fp16 unnecessarily because the rotation already handles internal precision internally. Removing the round-trips saved ~1.05x per call and is invisible to output quality.

The full write-up with stage-by-stage numbers is in OPTIMIZATION_FINDINGS.md.

Step 10: Understanding Memory Numbers​

KV cache memory scales with: 2 × n_layers × n_kv_heads × head_dim × seq_len × bytes_per_element

For Mistral 7B (32 layers, 8 KV heads, head_dim 128) at sequence length 2048:

• fp16: 2 × 32 × 8 × 128 × 2048 × 2 bytes = 536 MB
• 4-bit: 268 MB (2x reduction)
• 2-bit RVQ: 134 MB (4x reduction, plus codebook overhead ~8 KB)

At longer contexts (32K tokens), these numbers scale linearly:

• fp16: 8.4 GB for context alone
• 2-bit RVQ: 2.1 GB

On a 16GB Mac, the difference between fp16 and 2-bit RVQ at 32K context is the difference between OOM and running.

Common Questions​

Does this work with quantized models (4-bit weights)?

Yes. The KV cache quantization is independent of weight quantization. You can run a 4-bit weight model with a 2-bit KV cache. The two compressions are additive: mlx-community/Mistral-7B-Instruct-v0.3-4bit loads as a 4-bit weight model; the KV cache is then further compressed by VeloxQuant.

Will I notice quality degradation?

At 4-bit and 3-bit: barely measurable on generation tasks. At 2-bit with RVQ: cosine similarity 0.98 (vs 1.0 for fp16). On our Qwen3 4B thinking-mode test, RVQ 2-bit produced 199/200 coherent tokens vs fp16's 200/200. One token difference across a 200-token reasoning chain.

At 2-bit without RVQ (single-pass): noticeable. Thinking-mode models collapse early. Use RVQ.

What models are supported?

Any model that uses standard multi-head attention with mlx_lm.make_cache(). Tested: Mistral 7B, Qwen3 4B, Qwen3 8B, Qwen2.5 32B, Falcon3 7B, Phi-4, Gemma-4. Not supported: architectures with multi-latent attention (DeepSeek-V2) or non-standard cache formats.

Does it work with streaming generation?

Yes. The cache is stateful — each generate() call updates it. Stream tokens with the standard mlx_lm streaming API; the cache update happens transparently inside the model forward pass.

VeloxQuant-MLX is MIT licensed. Contributions welcome — especially benchmark results for new models and hardware.

How a 30-line Metal compute shader replaced the worst hot path in VeloxQuant-MLX 0.5.1, what I learned about Apple Silicon kernel launch overhead, and why this matters if you run LLMs locally on Mac.

The Bug That Wouldn't Die​

A few weeks back I shipped VeloxQuant-MLX 0.5.0 — a Python library that compresses the KV cache for any model you load through mlx_lm. The headline algorithm is VecInfer, which uses product vector quantization to squeeze keys down to 1 bit per element. That is 16× compression. Sounds great.

It worked beautifully on Llama-3.1-8B, Mistral-7B, Qwen2.5-7B, Phi-4 — every model with head_dim=128. And then I tested Falcon3-7B.

[VecInfer-2bit] generating...
  Out of memory: requested 712 MB, available 0

Falcon3-7B has head_dim=256. The chunked nearest-centroid search at the heart of quantize_vq allocates a tensor of shape [chunk_size, n_centroids, sub_dim] on every chunk. For Falcon's geometry that's a multi-hundred-megabyte intermediate — at every single token, on every layer, on every step. The GPU runs out of memory before generating a single token.

I shipped 0.5.0 with the OOM marked as a known limitation. It bothered me. I knew the fix conceptually — accumulate the squared distance in registers, never materialize the diff matrix — but doing that meant writing a Metal compute shader, and I had never written one.

This post is what happened when I did.

What Even Is a KV Cache And Why Should You Care​

Quick recap. Every transformer layer needs to remember the keys and values it computed for every token it's already seen. For a 7B model with 32 layers, 8 KV heads, and head_dim=128, generating an 8,000-token response means storing:

32 layers × 8 heads × 8000 tokens × 128 dims × 2 (K + V) × 2 bytes (fp16)
≈ 1 GB

On a 16 GB MacBook running the model weights (~5 GB at 4-bit) plus the OS and your app, that 1 GB is the difference between a fluent response and a hard crash. The KV cache is the silent killer of long-context inference on Mac.

KV-cache quantization — storing those keys and values at fewer bits — is the answer. There are several flavors. The aggressive one I shipped, VecInfer, uses product vector quantization:

1. Split each [head_dim] key vector into small sub-vectors of length sub_dim (typically 4 or 8).
2. Pre-train a codebook of K-means centroids on calibration data.
3. At inference, encode each sub-vector as the index of its nearest centroid.

A 128-dim fp16 key (256 bytes) becomes 16 indices at 8 bits each (16 bytes). That's the 16× compression.

The hot operation is step 3: finding the nearest centroid. On every layer, on every token, you do a vectorized argmin against the codebook. That's quantize_vq.

What quantize_vq Was Doing Wrong​

Here's what the pure-MLX implementation looks like (paraphrased):

def quantize_vq(x, codebook, sub_dim):
    # x: [N, sub_dim]    -- the sub-vectors to encode
    # codebook: [n_centroids, sub_dim]
    diff = x[:, None, :] - codebook[None, :, :]  # [N, n_centroids, sub_dim]
    d2 = mx.sum(diff * diff, axis=-1)  # [N, n_centroids]
    return mx.argmin(d2, axis=-1)  # [N]

That diff tensor is the killer. Its shape is [N, n_centroids, sub_dim]. For Falcon3-7B-shape inputs:

• N = 4096 tokens × 4 KV heads × 64 sub-vectors per head = 1,048,576
• n_centroids = 256
• sub_dim = 4
• Total: 1,048,576 × 256 × 4 × 2 bytes (fp16) = 2.1 GB intermediate

The implementation tries to mitigate this by chunking N — processing 4,096 sub-vectors at a time — but even one chunk is still ~32 MB, and a 7B model's GPU memory pressure means even that gets fragmented and OOMs in practice.

What you actually want is for each thread to compute the argmin in registers, only writing out a single uint32 index. No intermediate tensor. Total intermediate memory: zero.

That's exactly what a Metal compute kernel can do.

What Is MLX mx.fast.metal_kernel?​

MLX (Apple's array library for Apple Silicon) has a feature most people don't know about: mx.fast.metal_kernel. It lets you write a Metal Shading Language function inline as a Python string and have MLX JIT-compile it, manage the buffer bindings, and dispatch it on the GPU.

The whole thing takes a few lines of Python:

kernel = mx.fast.metal_kernel(
    name="vecinfer_quantize",
    input_names=["x", "codebook"],
    output_names=["out"],
    source=METAL_SOURCE,  # a string of MSL
)

result = kernel(
    inputs=[x, codebook],
    output_shapes=[(N,)],
    output_dtypes=[mx.uint32],
    grid=(N, 1, 1),
    threadgroup=(256, 1, 1),
)

MLX handles all the boilerplate: function signature generation, dtype binding, threadgroup memory, dispatch encoding. You write the kernel body. It's the easiest GPU programming experience I've ever had — closer to writing a Python function than to traditional CUDA.

The Kernel: 18 Lines of MSL​

Here's the entire fused-argmin kernel that replaces that 2 GB intermediate tensor:

uint vec_idx = thread_position_in_grid.x;
uint N_total = x_shape[0];
if (vec_idx >= N_total) {
    return;
}

uint n_centroids = codebook_shape[0];
uint sub_dim     = codebook_shape[1];
uint x_base      = vec_idx * sub_dim;

// Track running argmin in registers — never materialize the diff matrix.
float best_dist = INFINITY;
uint  best_idx  = 0;

for (uint c = 0; c < n_centroids; ++c) {
    uint cb_base = c * sub_dim;
    float dist = 0.0f;
    for (uint i = 0; i < sub_dim; ++i) {
        float d = float(x[x_base + i]) - float(codebook[cb_base + i]);
        dist += d * d;
    }
    if (dist < best_dist) {
        best_dist = dist;
        best_idx  = c;
    }
}

out[vec_idx] = best_idx;

That's it. Each GPU thread handles one sub-vector. It loops over all centroids, accumulates squared distance in a single float register, tracks the running minimum, and writes one uint32 index. The intermediate "diff matrix" never exists anywhere except in those two register-resident floats per thread.

Memory complexity: O(N) total output, vs O(N × n_centroids × sub_dim) for the Python path.

The Numbers​

I wrote a benchmark script — scripts/plot_metal_benchmarks.py in the repo — that runs both paths across realistic shapes and saves figures. Here are the results.

Throughput: 6.9–14.7× Speedup​

The speedup scales with sequence length — longer contexts (where the Python path is bandwidth-bound on those huge diff tensors) get bigger wins. At S=8192, D=128 we go from 228 ms per call to 16 ms per call. Per call. Multiply by 32 layers × 1 quantize per step × hundreds of tokens and you're talking minutes saved per long generation.

Memory: 729 MB → 12 MB​

At the Falcon3-7B OOM trigger shape (head_dim=256, n_centroids=256, sub_dim=4, S=4096):

This is the result that matters. The kernel doesn't just make existing models faster — it makes models that previously OOMed actually run.

Correctness: Bit-Exact on fp32, MSE-Identical on fp16​

This is where I had to be careful. The Metal kernel and the pure-MLX path don't produce identical indices on fp16 inputs — about 0.1% of indices differ.

Why? When two centroids are nearly equidistant from a point, the choice of "nearest" depends on the order of floating-point operations. The pure-MLX path does the subtract in fp16 (because the inputs are fp16); the Metal kernel promotes to fp32 inside the inner loop. When the tiebreaker happens at the 5th decimal place, the two paths pick different winners.

But here's the thing: the reconstruction quality is identical. I validated this by reconstructing keys from both index sets and measuring MSE against the original input:

B=1 H=8 S=2048 D=128 sub_dim=8 n_c=256 dtype=float16
  idx_diff = 0.104%
  mse_ref = 3.7211e-01    mse_metal = 3.7211e-01
  rel_err = 5.61e-07

Reconstruction MSE matches to 7 decimal places. The two paths produce functionally identical compressed representations — they just disagree on which arbitrary tie-breaker to pick.

The parity tests in veloxquant_mlx/tests/cache/test_vecinfer_metal_parity.py validate this directly: assert that reconstruction MSE is within 1% relative error, not that indices match.

What I Got Wrong on the First Try​

I want to be honest about the missteps, because they're the actually interesting part.

Mistake 1: I Wrote the Dequant Kernel First​

My first instinct was to write a Metal kernel for dequantize_vq — the operation that takes codebook indices and reconstructs the float vectors. It's conceptually simpler (just a gather), so I started there.

After getting bit-exact correctness, I benchmarked it:

shape                                pure-mlx     metal    speedup
B=1 H=8 S=128 n_sub=16 sub_dim=8       223.3 µs   185.6 µs   1.20x
B=1 H=8 S=512                          183.6 µs   209.3 µs   0.88x
B=1 H=8 S=2048                         258.3 µs   275.9 µs   0.94x
B=1 H=8 S=8192                         467.8 µs   577.6 µs   0.81x

My kernel was slower than MLX's mx.take. That stung. After staring at the numbers for an hour, the reason became obvious: MLX's mx.take is already a highly tuned Metal gather kernel under the hood. There is no "Python overhead" to eliminate. The pure-MLX path is a Metal kernel. My kernel was duplicating it badly.

The lesson: before writing a custom kernel, profile to find the operation that has actual Python/intermediate-tensor overhead. mx.take does not. quantize_vq does, because of the [N, n_centroids, sub_dim] materialization. The 30-line MSL shader had to fuse an algorithm — argmin over distances — not just replace a builtin.

I kept the dequant kernel as a building block for Phase 2 (fused dequant+SDPA), but the headline result is the quantize kernel.

Mistake 2: Wrong Threadgroup Layout​

My first quantize kernel dispatched one thread per (input_vector, sub_dim_component) pair. That made each thread tiny — one subtract, one square, one accumulate — and meant launching N × sub_dim threads. For typical shapes, that's millions of threads.

Apple Silicon GPUs have 32-wide SIMD groups and an internal cost per thread launch. Launching 8× more threads than you need is pure overhead.

The fix was to dispatch one thread per input vector and let each thread loop over all sub_dim components in registers. Same total work, 8× fewer thread launches, much better register reuse. That's the layout in the kernel above.

Mistake 3: I Assumed End-to-End Would Always Be Faster​

After validating the kernel was 13× faster on synthetic shapes, I ran the full benchmark on SmolLM2-135M (a 135-million-parameter tiny model) expecting to see a speedup in end-to-end token generation.

I got the opposite. The Metal path was slower end-to-end — 75 tok/s vs 178 tok/s for the pure-MLX path.

The reason: Metal kernel dispatch has a fixed per-launch overhead of roughly 50–200 µs on Apple Silicon. SmolLM2 has 30 layers, each doing 2 quantize calls per token, so that's ~60 kernel launches per generated token. The per-launch overhead exceeded the work each kernel did.

The kernel is designed for the regime where it matters: 7B+ models with realistic context lengths, where each quantize_vq call is doing milliseconds of work. On those, the launch overhead is negligible relative to the kernel runtime, and you get the full 10–14× speedup.

This is a limitation of MLX's kernel launch path — MLX doesn't yet expose a way to amortize launch overhead across multiple layers in a single dispatch. That's a Phase 3 problem and probably out of scope for a Python-level library.

How to Use This Today​

VeloxQuant-MLX 0.5.1 is on PyPI. Install:

pip install --upgrade VeloxQuant-MLX

The Metal kernels are on by default when available. No code changes needed. Your existing VecInferKVCache calls auto-detect Metal and use the fast path:

import mlx_lm
from veloxquant_mlx import KVCacheConfig, KVCacheFactory

model, tokenizer = mlx_lm.load("mlx-community/Falcon3-7B-Instruct-4bit")

# Metal auto-detected. To force off for debugging: use_metal_kernels=False
config = KVCacheConfig(
    method="vecinfer",
    head_dim=256,
    key_sub_dim=4,
    value_sub_dim=4,
    key_codebook_bits=8,
    value_codebook_bits=8,
    smooth_factors=calibrated_smooth_factors,
    key_codebook=calibrated_key_codebook,
    use_metal_kernels=None,  # None = auto, True = require, False = forbid
)

The new use_metal_kernels flag is three-state:

• None (default) — auto-detect; use Metal if available, silently fall back if not
• True — require Metal; raise at construction time if unavailable
• False — forbid Metal; use the pure-MLX path (for parity testing and debugging)

To verify the speedup on your own machine:

git clone https://github.com/rajveer43/VeloxQuant-MLX
cd VeloxQuant-MLX
PYTHONPATH=. python scripts/plot_metal_benchmarks.py
# Produces figures/metal/summary.png with your hardware's numbers

What's Next: Phase 2​

The quantize kernel is the biggest single win, but it's not the end. Phase 2 is fusing dequantize + scaled-dot-product-attention into a single kernel.

Right now, even with Phase 1, the cache still materializes the full fp16 key tensor on every attention call. The dequant happens — efficiently, since mx.take is already fast — but we hold the result in GPU memory long enough to pass it to MLX's SDPA. For very long contexts, that materialized key tensor is still significant memory pressure.

The Phase 2 kernel would:

1. Take codebook indices, the per-query LUT (q_tilde @ codebook.T), and value indices
2. Compute attention scores directly via LUT lookup, never materializing fp16 keys
3. Compute the softmax-weighted value sum in-kernel
4. Output the final attention result in one fused pass

This is what the VecInfer paper's CUDA kernel does. Porting it to Metal is the goal. If you've written Metal compute shaders before and want to collaborate, the GitHub issue is open.

The Meta-Lesson: Custom Kernels Are More Accessible Than You Think​

I had never written a Metal shader before this project. The mental model is straightforward once you get past the syntax:

1. Identify the operation with materialization overhead (not just a slow Python loop — those are usually wrapped in optimized C++ already; look for operations that create big intermediate tensors)
2. Write the algorithm with the intermediate as register-state instead of memory-state (running min, running sum, running argmin)
3. Dispatch one thread per output element, not per input or per output-component
4. Validate with reconstruction error, not bit-exact equality, when fp16 is involved
5. Benchmark at realistic shapes, not toy shapes — kernel launch overhead can dominate for small workloads

Total time investment for this Phase 1: about 6 hours of focused work, including the two failed approaches above. The resulting kernel unblocks head_dim=256 models that previously OOMed, gives a 10–14× speedup on the hot path, and is 30 lines of MSL.

If you've been hesitant to write custom GPU kernels because it sounds intimidating — mx.fast.metal_kernel makes the bar way lower than it used to be on CUDA. Try it.

TL;DR​

• VeloxQuant-MLX 0.5.1 adds a Metal compute kernel for quantize_vq, the hot path in VecInfer KV-cache compression
• 13× faster on realistic shapes (S=2048+)
• 98% less peak memory at the Falcon3-7B OOM trigger configuration
• Drop-in, zero API change — auto-detected when Metal is available
• Free, MIT-licensed, on PyPI: pip install VeloxQuant-MLX
• The kernel is 30 lines of Metal Shading Language inside Python
• Phase 2 (fused dequant+SDPA attention kernel) is next

GitHub: github.com/rajveer43/VeloxQuant-MLX PyPI: pypi.org/project/VeloxQuant-MLX Benchmark figures: figures/metal/summary.png in the repo

If this saves your Mac from OOMing tonight, leave a star — or open an issue if it doesn't.

How I wrote five hand-tuned Metal compute kernels in MLX for TurboQuant — and what every bug taught me about Apple GPU programming.

The Problem​

My Mac was choking on long-context LLM inference.

Not because the model was too large — I had already quantized the weights. The bottleneck was the KV cache. At 8k context, a single layer's key cache is [1, 32, 8192, 128] in fp16 — over 67 MB per layer, 2 GB across 32 layers. On Apple Silicon, where the GPU and CPU share the same physical memory, that pressure is immediate and painful.

VeloxQuant-MLX already had several compression algorithms: TurboQuantRVQ (7.5× via two-stage scalar RVQ), QJL (16× via 1-bit Johnson-Lindenstrauss sketching), and VecInfer (16× via product VQ). But they were all running through pure MLX graph operations — no custom GPU kernels. The hot paths were either slow or allocating huge intermediate tensors.

The fix: write the hot paths in Metal Shading Language and JIT-compile them via mx.fast.metal_kernel.

This is the story of how I did it, what broke, and what I learned.

The Stack​

Before diving into the kernels, here's the relevant context:

• MLX — Apple's NumPy-style ML framework with lazy evaluation and Metal GPU backend
• mx.fast.metal_kernel — Python API to write raw Metal Shading Language compute shaders that plug into MLX's lazy graph
• TurboQuant — a family of KV cache quantization algorithms (MSE, Prod, RVQ) implemented in VeloxQuant-MLX
• QJL — Quantized Johnson-Lindenstrauss: compress keys to 1-bit sign sketches + a scalar norm

The goal was to replace the slowest pure-MLX operations with Metal kernels that live in five focused submodules:

How mx.fast.metal_kernel Works​

Before showing any kernel code, there's one thing you need to understand about the API — because getting it wrong produces silent, subtle bugs.

The API in 30 seconds​

import mlx.core as mx

kernel = mx.fast.metal_kernel(
    name="my_kernel",
    input_names=["x", "y"],
    output_names=["out"],
    source="""
        uint i = thread_position_in_grid.x;
        out[i] = x[i] + y[i];
    """,
)

result = kernel(
    inputs=[a, b],
    grid=(N, 1, 1),
    threadgroup=(256, 1, 1),
    output_shapes=[(N,)],
    output_dtypes=[mx.float32],
)

The source string is raw Metal Shading Language — no kernel keyword, no function signature. MLX wraps it. Shape information is injected automatically: inside the kernel, x_shape[0] gives you the first dimension of x.

The #1 Gotcha: Grid = Total Threads​

This is the single most important thing to get right, and the MLX documentation is easy to misread on this point.

In standard Metal (Obj-C / Swift), you call dispatchThreadgroups(n_groups, threadsPerThreadgroup: tg_size) — so the grid is in threadgroup units.

MLX uses dispatchThreads — the grid is in total thread units.

That means if you want B threadgroups of T threads each:

# WRONG — only dispatches 1 thread per threadgroup
grid=(B, 1, 1), threadgroup=(T, 1, 1)

# CORRECT — dispatches B threadgroups of T threads each
grid=(B * T, 1, 1), threadgroup=(T, 1, 1)

I made this mistake on four out of five kernels. The symptom was identical every time: only the first batch element had correct output; everything else was zero. It looked like a memory layout bug or an indexing error, not a dispatch error. I spent hours debugging before I found it.

The Lazy Graph Contract​

mx.fast.metal_kernel returns a lazy node — nothing runs until mx.eval() is called. mx.eval() internally:

1. Encodes the compute command into a MTLCommandBuffer
2. Calls commandBuffer.commit() to submit to the GPU
3. Calls commandBuffer.waitUntilCompleted() to synchronize

You never write any of this yourself. MLX owns the entire Metal command buffer lifecycle.

Kernel 1: Bit-Packing — 30× Over NumPy​

The problem​

TurboQuantRVQ stores KV cache keys as b-bit indices (b ∈ {1, 2, 4}). The pure-Python path used a loop to pack these into uint8 bytes. At 65k elements it was ~8 ms — unacceptable.

The kernel​

constexpr int  ELEMS_PER_BYTE = 8 / B_BITS;
constexpr uint MASK           = (1u << B_BITS) - 1u;

uint byte_idx = thread_position_in_grid.x;
uint base     = byte_idx * ELEMS_PER_BYTE;

uint packed_byte = 0u;
for (int i = 0; i < ELEMS_PER_BYTE; ++i) {
    uint val = uint(indices[base + i]) & MASK;
    packed_byte |= (val << (i * B_BITS));
}
packed[byte_idx] = uint8_t(packed_byte);

One thread per output byte. B_BITS is a template parameter — a compile-time integer constant. This lets the compiler statically unroll the inner loop (2 iterations for b=4, 4 for b=2, 8 for b=1) and inline the constants.

The dispatch:

grid = ((n_bytes, 1, 1),)
threadgroup = ((min(256, n_bytes), 1, 1),)

Results​

The kernel dispatch overhead is ~0.17 ms regardless of N. Below ~2k elements NumPy wins because there's nothing to hide the launch cost behind. Above 16k elements, Metal dominates by an order of magnitude.

Kernel 2: Scalar Quantize / Dequantize — 11× Over NumPy​

The problem​

TurboQuantMSE quantizes each key dimension independently against a Lloyd-Max codebook. The pure-MLX path computed |x - centroids|² as a full [N, 2^b] matrix, then took argmin — allocating a tensor that was 2^b times the input size.

The quantize kernel​

constexpr int N_CENTS = 1 << B_BITS;

uint  elem      = thread_position_in_grid.x;
float val       = float(x[elem]);
int   best      = 0;
float best_dist = INFINITY;

for (int j = 0; j < N_CENTS; ++j) {
    float d    = val - centroids[j];
    float dist = d * d;
    if (dist < best_dist) { best_dist = dist; best = j; }
}
indices[elem] = uint8_t(best);

One thread per element. The centroid scan lives entirely in registers — no intermediate allocation. With B_BITS as a template, the loop body is known at compile time: the compiler generates 2, 4, 8, or 16 iterations of straight-line code.

The dequantize kernel​

Even simpler — a pure gather:

uint elem   = thread_position_in_grid.x;
x_hat[elem] = half(centroids[uint(indices[elem])]);

Results​

Kernel 3: Fused Hadamard + Quantize — The Hardest One​

The problem​

TurboQuantMSE (with Hadamard preconditioner) runs:

y = diag * H * x / sqrt(D)    [randomized Hadamard rotation]
idx = argmin_k |y - c_k|²     [nearest-centroid quantize]

Two separate dispatches, with a [B, D] fp16 intermediate between them. Fusing them into one kernel eliminates that allocation and the round-trip to GPU memory.

The kernel design​

Walsh-Hadamard Transform (WHT) is an in-place butterfly — each pass halves the stride. On GPU, D threads share a threadgroup, and each butterfly step needs a barrier.

threadgroup float buf[MAX_D];   // static threadgroup memory; MAX_D injected at compile time

// 1. Load + diagonal sign flip
float v = float(x[tg * D + lane]);
v *= float(diag[lane]);
buf[lane] = v;
threadgroup_barrier(mem_flags::mem_threadgroup);

// 2. In-place WHT — range-based parallel butterfly
for (uint stride = 1; stride < D; stride <<= 1) {
    uint local    = lane % (stride << 1u);
    bool is_upper = local >= stride;
    uint partner  = is_upper ? (lane - stride) : (lane + stride);
    float a = buf[lane];
    float b = buf[partner];
    threadgroup_barrier(mem_flags::mem_threadgroup);
    buf[lane] = is_upper ? (b - a) : (a + b);
    threadgroup_barrier(mem_flags::mem_threadgroup);
}

// 3. Scale
float y = buf[lane] * metal::rsqrt(float(D));

// 4. Nearest-centroid argmin (register-local)
int   best      = 0;
float best_dist = INFINITY;
for (int j = 0; j < N_CENTS; ++j) {
    float d    = y - centroids[j];
    float dist = d * d;
    if (dist < best_dist) { best_dist = dist; best = j; }
}
indices[tg * D + lane] = uint8_t(best);

The threadgroup array buf[MAX_D] requires MAX_D to be a compile-time constant — which is why it's injected as a #define in the kernel header:

_hadamard_quantize_kernel = mx.fast.metal_kernel(
    ...
    header=f"#define MAX_D {D}\n",
    source=_HADAMARD_QUANTIZE_SRC,
)

The butterfly bug​

My first implementation used:

uint partner = lane ^ stride;    // XOR butterfly

This looked right — it's the standard Cooley-Tukey bit-reversal trick. But on GPU, it produced ~90% index mismatch vs the sequential reference.

The problem: lane ^ stride traverses the WHT in bit-reversal order, which is fine for sequential execution (because you can reorder the output at the end), but on GPU where lanes run simultaneously, XOR pairing creates data races within a butterfly pass — some lanes read values that other lanes in the same pass are simultaneously writing.

The fix is a range-based butterfly that unambiguously partitions each pass into non-overlapping upper/lower pairs:

uint local    = lane % (stride << 1u);
bool is_upper = local >= stride;
uint partner  = is_upper ? (lane - stride) : (lane + stride);
float a = buf[lane];
float b = buf[partner];          // read BEFORE the barrier write below
threadgroup_barrier(mem_flags::mem_threadgroup);
buf[lane] = is_upper ? (b - a) : (a + b);

Reading a and b before the barrier guarantees both values come from the previous pass. After this fix, 100% of indices matched the reference.

Grid​

The grid uses B threadgroups of D threads — not B × D total:

# Wrong:
grid=(B, 1, 1), threadgroup=(D, 1, 1)   # only 1 thread per threadgroup!

# Correct:
grid=(B * D, 1, 1), threadgroup=(D, 1, 1)   # B threadgroups of D threads

Kernel 4: QJL Encode — Simdgroup Sign Packing​

The problem​

QJL encoding requires:

1. For each key vector x[b], compute sign(S @ x[b]) for all m sketch dimensions — giving m bits
2. Pack those m bits into m/8 uint8 bytes (LSB-first)
3. Compute ‖x[b]‖ (one scalar per key)

The pure-MLX path materialized the full [B, m] float matrix S @ x.T before sign-taking — m * d * B * 4 bytes, growing linearly with batch and sketch size.

Simdgroup design​

Each simdgroup (32 lanes) handles 32 consecutive sketch dimensions. Lane j computes dot(S[simd_blk*32 + j, :], x[b, :]) via a scalar loop:

uint b_idx    = flat_tg / n_simd_per_batch;
uint simd_blk = flat_tg % n_simd_per_batch;
uint sketch_j = simd_blk * 32u + lane;

float dot_val = 0.0f;
if (sketch_j < m) {
    uint S_row = sketch_j * d;
    uint x_row = b_idx   * d;
    for (uint i = 0; i < d; ++i) {
        dot_val += float(S[S_row + i]) * float(x[x_row + i]);
    }
}

After the dot product, all 32 lanes cooperate to pack 32 sign bits into 4 bytes using simd_shuffle:

uint sign_bit    = (dot_val >= 0.0f) ? 1u : 0u;
uint byte_in_blk = lane / 8u;
uint bit_in_byte = lane % 8u;

uint packed_byte = 0u;
for (uint bit = 0; bit < 8u; ++bit) {
    uint src = byte_in_blk * 8u + bit;
    packed_byte |= (simd_shuffle(sign_bit, src) << bit);
}

if (bit_in_byte == 0 && sketch_j < m) {
    packed_signs[out_byte] = uint8_t(packed_byte);
}

simd_shuffle(val, lane_id) broadcasts sign_bit from lane src to the current lane — no shared memory needed. Lane 0 (of each byte group) does the final write.

The norm is computed cooperatively by simd_blk 0:

if (simd_blk == 0) {
    float x_sq = 0.0f;
    for (uint i = lane; i < d; i += 32u) {
        float v = float(x[x_row + i]);
        x_sq += v * v;
    }
    float norm_sq = simd_sum(x_sq);
    if (lane == 0) norms[b_idx] = half(metal::sqrt(norm_sq));
}

Grid (the bug, again)​

n_simd_per_batch = (m + 31) // 32
n_total_threads  = B * n_simd_per_batch * 32   # ← must multiply by 32
grid=(n_total_threads, 1, 1), threadgroup=(32, 1, 1)

Without the * 32, only B * n_simd_per_batch total threads dispatched — meaning only the first simdgroup ran, and only the first key had any output.

Kernel 5: Fused RVQ Decode + Attend — Online Softmax Without Materializing K​

The problem​

Attention with a quantized KV cache normally requires two dispatches:

1. Decode all compressed keys → K_hat tensor [B, H, S_kv, D] (fp16, potentially GBs)
2. Run softmax(q @ K_hat.T / sqrt(D)) @ V

The K_hat tensor is allocated, filled, used once, and thrown away. For RVQ keys this is unavoidable in the two-dispatch design — but we can fuse everything into a single FlashAttention-style pass that decodes keys on the fly without ever materializing K_hat.

Design​

Each threadgroup handles one query position (b, h, sq). Lanes stripe across the D-dimensional vectors in steps of TG = min(D, 32):

float running_m = -INFINITY;   // online softmax running max
float running_d = 0.0f;        // online softmax running denominator
float my_out[8];               // per-lane output accumulator
for (int i = 0; i < 8; ++i) my_out[i] = 0.0f;

for (uint sk = 0; sk < S_kv; ++sk) {
    // 1. Decode key on-the-fly: k[i] = cents1[idx1[i]] + cents2[idx2[i]]
    float partial_dot = 0.0f;
    for (uint i = tg_lane; i < D; i += TG) {
        float ki = centroids1[uint(k_indices1[k_off])]
                 + centroids2[uint(k_indices2[k_off])];
        partial_dot += float(q[q_base + i]) * ki;
    }
    float score = simd_sum(partial_dot) * inv_sqrt_d;

    // 2. Online softmax update (Dao et al. FlashAttention)
    float m_new  = metal::max(running_m, score);
    float factor = metal::exp(running_m - m_new);
    float w      = metal::exp(score     - m_new);
    running_d    = running_d * factor + w;
    running_m    = m_new;

    // 3. Rescale and accumulate value
    for (uint i = 0; i < n_owned; ++i) my_out[i] *= factor;
    for (uint i = tg_lane; i < D; i += TG) {
        float vi    = float(v_codebook[cb_off]);
        uint  out_i = (i - tg_lane) / TG;
        my_out[out_i] += w * vi;
    }
}

// 4. Normalize and write
for (uint i = tg_lane; i < D; i += TG) {
    uint out_i   = (i - tg_lane) / TG;
    out[out_off] = half(my_out[out_i] / running_d);
}

simd_sum(partial_dot) broadcasts the full dot product to all lanes in the simdgroup — this is the SIMD-level reduction that gives the correct score without any threadgroup memory.

The local accumulator index out_i = (i - tg_lane) / TG is the critical piece: lane 0 owns dims {0, TG, 2×TG, ...}, lane 1 owns {1, TG+1, ...}, and out_i is the position within that lane's private array.

The Benchmarks​

After fixing all the dispatch bugs, here are the results on Apple M-series (figures saved to figures/metal/turboquant_kernels/):

Memory savings are the bigger story for the RVQ attend kernel — it eliminates the [B, H, S_kv, D] fp16 K_hat tensor entirely. At S_kv=4096, H=32, D=128 that's 33 MB per layer, ~1 GB across a 32-layer model, allocated and freed every forward pass.

1-bit bit-packing alone gives 16× memory compression on the key cache (1 bit per dimension vs fp16). Combined with the Metal kernel's 30× throughput advantage, the packing/unpacking step goes from a bottleneck to essentially free.

What I Learned​

1. Grid = total threads is the most common MLX Metal mistake​

Every tutorial and reference for Metal uses dispatchThreadgroups. MLX uses dispatchThreads. These are different. If your output is correct for the first batch element and zero elsewhere, check your grid first.

2. XOR butterflies are wrong for parallel WHT​

The standard sequential WHT uses pair = i ^ stride. On GPU this causes data races within a butterfly pass because multiple threads simultaneously read from and write to overlapping pairs. Use range-based pairing (local = lane % (stride*2); is_upper = local >= stride) and read both values before the barrier.

3. simd_sum and simd_shuffle are your first tools, not shared memory​

For reductions and broadcasts within a simdgroup (32 lanes), simd_sum and simd_shuffle are zero-cost compared to threadgroup_barrier + shared memory. Design around simdgroups first; only escalate to threadgroup memory when you need communication beyond 32 lanes.

4. Template parameters unlock static unrolling​

template <int B_BITS> turns runtime constants into compile-time constants. The inner loop over centroids becomes 2, 4, 8, or 16 unrolled iterations — no branch, no loop counter. This is how Metal kernels beat NumPy at large N despite higher launch overhead: the arithmetic is genuinely faster.

5. You don't manage commandBuffer​

MLX handles commandBuffer.commit() and commandBuffer.waitUntilCompleted() inside mx.eval(). You never touch Metal command buffers when using mx.fast.metal_kernel. This is by design — MLX's lazy graph batches multiple kernel dispatches into one command buffer where possible.

6. The launch overhead is real and ~0.17 ms​

Every Metal kernel dispatch costs ~0.17 ms regardless of work size. For small N (< ~2k elements), NumPy is faster. For large N (> ~16k), Metal wins by 10–30×. Design your batching strategy accordingly — combine small operations into a single larger kernel rather than dispatching many small ones.

Code Organization​

The five kernels are organized into focused submodules under veloxquant_mlx/metal/:

metal/
├── __init__.py          # lazy re-exports
├── kernels.py           # thin facade — imports from all submodules
├── _bit_packing.py      # turboquant_bit_pack, turboquant_bit_unpack
├── _scalar_quant.py     # turboquant_scalar_quantize, _dequantize, _hadamard_quantize
├── _qjl.py              # qjl_encode, qjl_inner_product
├── _rvq_attend.py       # turboquant_fused_rvq_decode_attend
└── _vecinfer.py         # vecinfer_dequant_metal, vecinfer_quantize_metal, ...

Each submodule has its own _cache: dict = {} for the kernel singleton pattern — build the MTLComputePipelineState once on first call, reuse forever:

def _pack_kernel(b: int):
    key = ("bit_pack", b)
    if key not in _cache:
        _cache[key] = mx.fast.metal_kernel(
            name=f"turboquant_bit_pack_b{b}",
            input_names=["indices"],
            output_names=["packed"],
            source=_PACK_SRC,
        )
    return _cache[key]

kernels.py is now a 47-line re-export facade:

from veloxquant_mlx.metal._bit_packing import turboquant_bit_pack, turboquant_bit_unpack
from veloxquant_mlx.metal._scalar_quant import turboquant_scalar_quantize, ...
from veloxquant_mlx.metal._qjl import qjl_encode, qjl_inner_product
from veloxquant_mlx.metal._rvq_attend import turboquant_fused_rvq_decode_attend

All 40 tests pass after the restructuring — the facade is transparent to callers.

The Broader Point​

Apple Silicon is a genuinely good target for this kind of work. Unified memory means you don't pay PCIe bandwidth to move data between CPU and GPU — the Metal kernel reads the same bytes your Python code just wrote. The simdgroup primitives (simd_sum, simd_shuffle) are clean and well-documented. And mx.fast.metal_kernel makes the iteration loop fast: write Metal source in Python, evaluate, fix, repeat.

The hard part isn't the Metal itself — it's understanding how MLX dispatches kernels. Once you internalize "grid = total threads, not threadgroups" and "lazy graph, so nothing runs until mx.eval()", the rest is straightforward shader programming.

The full source is in VeloxQuant-MLX under veloxquant_mlx/metal/. The benchmark script is at veloxquant_mlx/benchmarks/metal_kernel_benchmark.py and produces all the figures discussed here.

References​

• TurboQuant (ICLR 2026) — Zandieh et al., "Online Vector Quantization with Near-optimal Distortion Rate"
• QJL (2024) — Zandieh et al., "QJL: 1-Bit Quantized JL Transform for KV Cache Quantization"
• FlashAttention (NeurIPS 2022) — Dao et al., online softmax algorithm
• Apple MLX — the framework
• Metal Shading Language Specification — simd_sum, simd_shuffle, threadgroup_barrier reference

Code: github.com/rajveer43/VeloxQuant-MLX · Previous post: I Wrote a Metal Kernel to Stop My Mac From OOMing on LLM Inference

A full 8-model benchmark of residual vector quantization for KV cache compression on Apple M4 16GB. What worked, what didn't, and the one result that surprised me.

Here is the result I did not expect: a 32-billion-parameter model running faster with the KV cache compressed than without.

Qwen2.5-32B at RVQ 2-bit: 4.2 tok/s. The same model at full fp16 precision: 3.7 tok/s. That is a 14% throughput improvement from compression — on a 16GB MacBook. Not from a faster algorithm. From freeing up memory bandwidth that the model weights were competing for.

This is the full story: what we built, how it works, and what the benchmarks across 8 models actually show.

The Memory Problem Every Mac LLM User Knows​

Every token your LLM generates writes two vectors into memory — a Key and a Value — for every attention head, across every layer. This is the KV cache. It grows linearly with context length.

For Mistral 7B (32 layers, 8 KV heads, head dimension 128) at 32K tokens, the KV cache at fp16 is:

2 × 32 layers × 8 heads × 128 dims × 32,000 tokens × 2 bytes = 8.4 GB

On a 16GB Mac, that is 8.4 GB for context alone — before model weights, before the OS, before your browser tab.

The KV cache is also the performance bottleneck. Most 7-8B models on Apple Silicon are memory-bandwidth bound: the chip can do the arithmetic faster than it can load data from unified memory. Every token requires streaming the full KV cache through the memory bus once per layer. Smaller cache = faster generation.

Why 2-Bit Usually Fails​

The obvious approach is to quantize each float16 value (2 bytes) down to 2 bits — an 8x reduction. The problem: with only 4 quantization levels, reconstruction error is large enough to corrupt attention scores.

Cosine similarity between the original and compressed key vector — which directly determines attention quality — drops to 0.69 at single-pass 2-bit. For short generation (< 100 tokens), this is often tolerable. For long reasoning chains, it is not. We measured this:

This is not a mild degradation. For these models, standard 2-bit quantization simply does not work.

The Fix: Two-Pass Residual Quantization​

The key insight: if your first codebook leaves large error, run a second codebook on the error itself.

Stage 1 fits a Gaussian Lloyd-Max codebook on the rotated key vector. The rotation (a randomized Hadamard transform) spreads information uniformly across dimensions, so the distribution after rotation looks Gaussian and the codebook fits it well.

Stage 2 takes the residual — the difference between the original and the stage-1 reconstruction — and fits a Laplacian codebook on that. The Laplacian distribution matches the residual better than a Gaussian: it is the tail of a Gaussian, peaked at zero, with heavier tails. Using the wrong distribution family for stage 2 leaves measurable error on the table.

Result: cosine similarity from 0.69 → 0.98 at 2 bits per dimension. The two-pass approach recovers information the single-pass approach discards.

Storage per vector (d=128):

• RVQ 2-bit: ceil(128 × 2 × 2 / 8) = 64 bytes per key vector vs 256 bytes at fp16 — 4x compression
• RVQ 1-bit (sign quantizer + Laplacian residual): 32 bytes per vector — 8x compression, cosine 0.917

The 1-bit variant uses a two-level sign quantizer as stage 1 ({−0.798, +0.798} — the exact Gaussian Lloyd-Max solution at 1 bit) and a Laplacian codebook on the sign-quantization error for stage 2. It achieves 7.5x KV cache compression while keeping cosine above 0.91.

The Throughput Problem (and How We Solved It)​

Compression reduces memory pressure. But if the quantization compute itself costs more than the bandwidth savings, the net result is slower generation.

Our initial implementation ran at 17.7 tok/s on Mistral 7B with RVQ 2-bit — slower than fp16's 22.1 tok/s. Four changes fixed this, each independently measurable:

1. Batch all heads into one MLX call (+22%)

The original code maintained a separate quantizer per attention head and looped over them in Python. For Mistral 7B (8 KV heads, 32 layers), that is 256 small kernel dispatches per token. The fix: reshape (B, H, S, D) → (B·H·S, D) and call one shared quantizer on the entire batch. MLX compiles the operation as a single graph node and dispatches it once.

2. Switch to Hadamard rotation

The rotation preconditioner used QR decomposition — a full d×d matrix multiply (16,384 operations at d=128). mx.hadamard_transform is a Metal-native fused kernel that achieves the same statistical effect in O(d log d) — 896 operations at d=128. The quality is mathematically identical: a Hadamard with a random ±1 diagonal is Haar-distributed, the same family as a random orthogonal matrix.

3. Boundary-sum quantization instead of broadcast-argmin

Codebook lookup was materializing a (batch, d, k) distance tensor and running argmin — three kernel launches: broadcast-subtract, absolute value, argmin. Lloyd-Max boundaries are just midpoints between sorted centroids. The nearest centroid index equals the number of boundaries the value exceeds: one comparison and one sum. Two kernels, no intermediate tensor. Verified bitwise-identical to the old path at 100.00% index match.

4. Remove redundant dtype casts

The update path was casting fp16 → fp32 → fp16 → fp32 → fp16 from accumulated unnecessary coercions. Removing them saved ~5% per call with no effect on output quality.

Net result: 17.7 → 22.3 tok/s on Mistral 7B. The compression now matches fp16 throughput.

Full 8-Model Benchmark​

Hardware: Apple M4 MacBook, 16GB unified memory. Each run: 200 tokens, one fresh Python subprocess per (model, config) to avoid MLX graph compilation bugs. 6 configs: fp16, TQ 2/3/4-bit (single-pass), RVQ 2-bit, RVQ 1-bit.

★ = RVQ configs at 7.5x compression. Bolded token counts in TQ 2-bit column = failure.

Three Findings Worth Highlighting​

1. RVQ 1-bit is the reliability result​

Across all 8 models at all 200 tokens, RVQ 1-bit with 7.5x compression produced complete, coherent output every time. TQ single-pass 2-bit failed catastrophically on 4 of 8 models.

This is the practical takeaway: if you want 2-bit compression and need your model to actually finish its output, residual quantization is not optional. The gap between 0.69 and 0.92 cosine is the gap between broken and working.

2. At 32B scale, compression beats fp16​

Qwen2.5-32B is too large for 16GB without aggressive weight quantization. With 4-bit weights (the Qwen2.5-32B-Instruct-4bit model), it fits — but barely. The model consumes ~17.5 GB including the OS and runtime, leaving almost nothing for KV cache headroom. The chip is running at its memory bandwidth ceiling.

Compressing the KV cache frees bandwidth that the weight loads were competing for. The result is a net gain:

• fp16 KV cache: 3.7 tok/s
• RVQ 2-bit KV cache: 4.2 tok/s (+14%)
• RVQ 1-bit KV cache: 3.9 tok/s (+7%)

The effect is largest at 32B because the bandwidth contention is sharpest. At 7-8B, the same effect exists but is smaller (Qwen3-8B: +3% at RVQ 1-bit). For models that fit comfortably in unified memory, the effect disappears.

3. The Qwen3 4B exception​

Qwen3 4B runs at 40.2 tok/s at fp16 — unusually fast for a model this size, because it is small and the chip can move through layers quickly. At RVQ 1-bit, it drops to 34.3 tok/s (85% of fp16).

This is the trade-off inverting. At small model sizes, fp16 is fast enough that the quantization compute overhead is a larger fraction of total runtime. At large model sizes, the weight loads dominate and the overhead is amortized. The crossover is somewhere around 8B parameters on this hardware.

If you are running a 4B model on a Mac with 32GB, use fp16. If you are running it on 16GB and memory is the constraint, RVQ 2-bit at 35.0 tok/s and 197/200 tokens is the right trade.

The VLM Extension​

We also extended quantization to Qwen2-VL-7B — a vision language model that uses bfloat16 internally (wider exponent range, needed for large-norm image patch tokens). The original code always cast to float16 before normalizing, which discarded the bfloat16 exponent range at exactly the wrong moment for image tokens.

The fix is one variable: kdtype = keys.dtype. Normalize in the actual dtype, cast to float16 only for the codebook lookup. The VLM key-distribution diagnostic across all 28 layers showed that image patch tokens and text tokens have nearly identical distributions after RMSNorm — the quantizer does not need special handling for image tokens. RVQ 2-bit cosine on real Qwen2-VL keys: 0.979 for both image and text tokens across all layers.

What This Means in Practice​

16GB Mac users:

For 7-8B models (Mistral, Llama, Falcon, Phi-4, Qwen3-8B), RVQ 1-bit gives you 7.5x KV cache compression at 94-99% of fp16 throughput. At 32K context, your KV cache drops from ~8 GB to ~1 GB. That is the difference between OOM and comfortable operation.

For Qwen2.5-32B on 16GB, RVQ is what makes the model faster than fp16. The compression is load-bearing.

For 4B models (Qwen3-4B, Gemma3-4B), use fp16 if you have memory. Use RVQ 2-bit if you do not.

Do not use single-pass 2-bit. It fails on Phi-4, Llama, Falcon, and Qwen2.5-32B. If you need 2-bit compression, you need residual quantization.

Using It​

pip install VeloxQuant-MLX

import mlx_lm
from mlx_kv_quant import KVCacheBuilder, KVCacheConfig

model, tokenizer = mlx_lm.load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")

# 7.5x compression, 95% throughput
config = KVCacheConfig(bits=1, algorithm="turboquant_rvq")
cache = KVCacheBuilder(config).for_model(model).build()

response = mlx_lm.generate(
    model,
    tokenizer,
    prompt="Explain the difference between RAM and unified memory.",
    kv_cache=cache,
    max_tokens=500,
)
print(response)

For RVQ 2-bit: bits=2. For single-pass 4-bit (safe on all models, less compression): algorithm="turboquant_prod", bits=4.

What the Figures Show​

Each model folder in figures/2026-05-12/ contains 6 panels:

1. Throughput comparison — tok/s across all 6 configs
2. Quality vs compression — cosine similarity curve with RVQ ★ markers
3. Memory at scale — KV cache bytes vs sequence length for each config
4. Attention distortion — per-config distortion bars
5. Output comparison — tokens generated per config (the failure-detection panel)
6. Full report — all panels on one sheet

The output comparison panel (fig5) is the one to look at for reliability. Models where single-pass 2-bit generates 0 or 5 tokens show it plainly. RVQ configs all extend to the full 200-token bar.

Reproducibility​

All runs use subprocess isolation — one fresh Python process per (model, config) — to avoid MLX's graph-reuse bug that causes 2nd+ configs to generate 0 tokens in the same process when the cache type changes. If you run these scripts yourself and see anomalously low token counts on non-fp16 configs, check that you are not running multiple configs in the same Python session.

git clone https://github.com/rajveer43/veloxquant-mlx
cd veloxquant-mlx
pip install -e .
PYTHONPATH=. python3 benchmark_scripts/run_full_reports.py --models mistral7b falcon3_7b llama31_8b

Figures save to figures/2026-05-12/<model>/. Each model takes 10–15 minutes on M4.

VeloxQuant-MLX is MIT licensed. Hardware: Apple M4, 16GB unified memory. MLX 0.24.x, mlx-lm 0.21.x.

How I built VeloxQuant-MLX — KV cache compression for LLMs running on your Mac​

Every time a language model generates a token, it writes two vectors into memory: a Key and a Value. These accumulate silently in what's called the KV cache — one pair per token, per attention head, per layer.

At a short context that's fine. But at 4,000 tokens with a 7B model on a 16GB Mac, the KV cache alone can consume 4–6 GB. At 32K tokens, it can exceed the model weights themselves.

I wanted to fix this for Apple Silicon. So I spent the last few weeks porting three research algorithms from GPU-targeted papers into a native MLX library — and published it to PyPI as VeloxQuant-MLX.

Here's exactly how it works, what I had to rebuild from scratch, and what the benchmarks actually show.

The Core Idea: Why KV Vectors Can Be Compressed​

Key vectors in a transformer aren't arbitrary. After a learned linear projection, they tend to follow a roughly Gaussian distribution — especially after layer normalization. This matters because scalar quantization theory tells us the optimal codebook for a Gaussian is the Lloyd-Max codebook.

The Lloyd-Max algorithm iterates two conditions until convergence:

1. Optimal boundaries — the boundary between two quantization levels sits at the midpoint of their centroids
2. Optimal centroids — each centroid is the conditional mean of the distribution over its Voronoi cell

c_i = ∫[b_{i-1} to b_i] x·f(x)dx  /  ∫[b_{i-1} to b_i] f(x)dx

For a standard Gaussian, this converges to fixed centroids that you can precompute once and reuse for every layer, every model, every token. No per-block scale constants. No calibration data. Just a lookup table.

The problem: raw KV vectors are not standard Gaussian. They have varying magnitudes, correlated dimensions, and outlier channels that dominate the quantization error.

This is what TurboQuant solves.

Stage 1: Random Rotation​

Before quantizing, multiply each Key vector by a random orthogonal matrix Π:

k̃ = k · Πᵀ

Why? Because a random rotation spreads information uniformly across all dimensions. Outlier channels — dimensions with unusually high variance — get diluted into all other dimensions. After rotation, the vector looks much more like an isotropic Gaussian, and the Lloyd-Max codebook fits it well.

The original TurboQuant paper uses QR decomposition on a random Gaussian matrix to generate Π. That's O(d²) per rotation.

For VeloxQuant-MLX, I added a second option: the randomized Hadamard transform:

k̃ = D · diag(r) · WHT(k)

where D is a fixed scaling factor, r is a random Rademacher vector (±1), and WHT is the Walsh-Hadamard Transform. This runs in O(d log d) and maps directly to mx.hadamard_transform — Metal-accelerated on Apple Silicon.

class HadamardPreconditioner(Preconditioner):
    def apply(self, x: mx.array) -> mx.array:
        # x shape: (..., d)
        scaled = x * self._diag_scale  # diagonal randomization
        return mx.hadamard_transform(scaled) / math.sqrt(self._d)

For head dimensions that are powers of 2 (128, 256, 512) — which covers nearly every production model — the Hadamard preconditioner is automatically selected.

Stage 2: Lloyd-Max Scalar Quantization​

After rotation, normalize each vector to unit norm (store the norm as fp16 for later), then quantize each dimension independently using the precomputed codebook:

for each dimension j in rotated key:
    find nearest centroid in codebook
    store index (b-1 bits)

The codebook has 2^(b-1) centroids for a b-bit configuration. At b=4, that's 8 centroids — enough to achieve ~0.95 cosine similarity with the original key vector after dequantization.

The lloyd_max() function in the library solves this numerically using trapezoidal quadrature:

def lloyd_max(pdf_fn, support, n_levels, n_iter=100, tol=1e-8):
    centroids = np.linspace(lo, hi, n_levels)
    for _ in range(n_iter):
        boundaries = midpoints(centroids)
        centroids = conditional_means(pdf_fn, boundaries)  # numerical integration
        if converged:
            break
    return centroids, boundaries

You run this once at startup. The result is a tiny lookup table (~8 floats for 3-bit) that gets reused for every token in every layer for the entire generation.

Stage 3: QJL Residual Correction​

Even at 4-bit, MSE quantization leaves a residual error:

r = k - k̂_MSE

TurboQuant's key insight is that you can correct for this without storing the full residual. Instead, store only the sign sketch: apply a random Johnson-Lindenstrauss projection matrix S to the residual, then store only the signs:

z = sign(S·r) ∈ {-1, +1}^m

One bit per sketch dimension. For m=128 sketch dimensions, that's 128 bits = 16 bytes per key vector.

At attention time, the corrected inner product estimate becomes:

⟨q,k⟩ ≈ ⟨q,k̂_MSE⟩ + ‖r‖ · (π/2m) · Σ_j z_j · sign(⟨S_j, q⟩)

The residual norm ‖r‖ is stored as fp16 (2 bytes). The signs are bit-packed (1 bit each). The correction is unbiased — on expectation, it perfectly reconstructs the true inner product.

This is what makes TurboQuant different from plain quantization: the QJL stage recovers information that was thrown away, using only 1 bit per sketch dimension.

Memory Layout: What Actually Gets Stored​

For each token, per attention head, VeloxQuant-MLX stores:

Compare to fp16: 128 dimensions × 2 bytes = 256 B for keys alone.

At 4-bit, key compression is 4.27×. Including values, total KV compression is ~1.6×. With value compression (int8 per-token, already implemented), total KV compression reaches ~3.5–4×.

The Bit-Packing Problem​

MLX has no sub-byte dtype. A 3-bit index stored naively in uint8 wastes 5 bits and kills your compression ratio.

The solution is a BitPackBuffer — pack multiple indices into each byte:

3-bit example: pack [2, 5, 1, 7, 3, ...] into bytes
  byte 0: bits 0-2 = index[0], bits 3-5 = index[1], bits 6-7 = index[2][0:1]
  byte 1: bit 0 = index[2][2], bits 1-3 = index[3], ...

For unpack at attend time, the library uses vectorized numpy paths for b ∈ {1, 2, 4} and a loop fallback for b=3:

if b == 2:
    shifts = np.array([0, 2, 4, 6], dtype=np.uint8)
    vals = ((packed[:, :, None] >> shifts) & 0x3).reshape(n, -1)
    return vals[:, :d]

This runs entirely in numpy before handing off to MLX for the attention computation — keeping the Metal graph clean.

Outlier Two-Stream Cache​

Some attention heads have a handful of dimensions with dramatically higher variance than the others. These outlier channels dominate quantization error: if dimension 47 has 10× the typical variance, every other dimension's codebook entry gets contaminated.

VeloxQuant-MLX handles this with an optional two-stream cache:

1. Calibration phase — observe the first N tokens, track per-channel variance
2. Detection — identify the top-K highest-variance channels
3. Split encoding — outlier channels → int8 with per-token scale; inlier channels → TurboQuant as normal

# During append:
k[:, outlier_idx] = 0           # zero out outliers in inlier stream
encode inlier key with TurboQuant
store outlier channels as int8

# During attend:
scores = turboquant_ip(q, encoded_keys)
scores += int8_ip(q[outlier_idx], outlier_cache) * outlier_scales

The result in benchmarks: at K=8 outlier channels, output quality is indistinguishable from fp16 even at 4-bit compression. The cost is slightly higher memory usage (8 × 2 = 16 extra bytes per token per head) and a numpy↔MLX copy per token — which is the main throughput bottleneck in the current implementation.

Weight Quantization: Compressing the Model Too​

One thing missing from the original papers: they only quantize the KV cache. The model weights stay fp16.

VeloxQuant-MLX adds a QuantizedLinear layer — a drop-in nn.Module replacement that applies TurboQuant to the weight matrix itself:

from mlx_kv_quant.weight import quantize_model

model = load_model(...)
quantize_model(model, bits=4, seed=42)
# All nn.Linear layers are now QuantizedLinear

Each row of the weight matrix is normalized, rotated with the same Hadamard preconditioner, and Lloyd-Max quantized. At inference, dequantize on the fly. At 3-bit, this compresses weight matrices by ~5× with minimal accuracy loss.

Benchmark Results Across 7 Models​

I ran the full benchmark suite (fp16, 2-bit, 3-bit, 4-bit) across 7 models on an Apple M4 MacBook (16GB unified memory):

Key finding: 4-bit is near-lossless on 5 of 7 models. 3-bit is near-lossless on 4 of 7. 2-bit (11.6× compression) breaks generation on all models tested — which matches the theoretical SNR floor (~4 dB at 2-bit vs ~10 dB at 4-bit).

The Qwen3 family underperforms because its thinking-mode models generate a <think> token stream that terminates early when attention quality degrades — a more sensitive indicator than output text quality.

What I Had to Rebuild for Apple Silicon​

The TurboQuant paper targets NVIDIA GPUs with CUDA kernels. Porting to MLX required rebuilding several components from scratch:

1. No in-place mutation. MLX uses lazy evaluation — array slices return views, not copies, and array[i] = value silently fails. Every cache write goes through numpy (pre-allocate numpy arrays, write indices into numpy, hand to MLX only for compute).

2. No sub-byte dtypes. The BitPackBuffer class handles arbitrary bit-widths in pure numpy, with vectorized unpack for b ∈ {1,2,4} and a loop path for b=3.

3. Metal-accelerated Hadamard. MLX exposes mx.hadamard_transform natively. Using it instead of the paper's QR rotation reduces rotation cost from O(d²) to O(d log d) — the entire operation runs on the GPU die of the M-series chip.

4. No monkey-patching. The original turboquant-mlx research code patched mlx_lm.scaled_dot_product_attention() globally. VeloxQuant-MLX exposes a standalone cache.attend(q) method — no global side effects, clean integration with any framework.

5. Pluggable architecture. The production library uses a Factory + Strategy + Registry pattern so you can swap quantizers at runtime without changing model code:

cache = (
    KVCacheBuilder()
    .with_method("turboquant_prod")  # swap to "polar", "qjl", "turboquant_mse"
    .with_head_dim(128)
    .with_bit_width(inlier=4)
    .with_jl_dim(128)
    .build()
)

Three Algorithms, One Interface​

VeloxQuant-MLX implements three quantization algorithms behind the same KVCache interface:

TurboQuantProd — Rotation + Lloyd-Max (b-1 bits) + QJL residual (1 bit). Best quality-per-bit for production use.

TurboQuantMSE — Rotation + Lloyd-Max (b bits), no residual correction. Lower memory overhead, slightly lower quality.

PolarQuantizer — Recursive polar coordinate decomposition. Represents each vector as a sequence of angles at 4 levels, plus a final radius. Geometrically motivated — useful for very low-bit regimes.

QJLQuantizer — Pure 1-bit: sign sketch only. Extreme compression (32× vs fp16), loses most content but preserves attention topology. Useful for very long contexts where quality is secondary to fitting in RAM.

Installing and Using It​

pip install VeloxQuant-MLX

Requires Python ≥ 3.11 and an Apple Silicon Mac.

from mlx_kv_quant import KVCacheBuilder
import mlx.core as mx
import numpy as np

cache = (
    KVCacheBuilder()
    .with_method("turboquant_prod")
    .with_head_dim(128)
    .with_bit_width(inlier=4)
    .with_jl_dim(128)
    .with_seed(42)
    .build()
)

rng = np.random.default_rng(0)
for _ in range(1000):
    k = mx.array(rng.standard_normal(128).astype(np.float16))
    v = mx.array(rng.standard_normal(128).astype(np.float16))
    cache.append(k, v)

q = mx.array(rng.standard_normal(128).astype(np.float16))
output = cache.attend(q)
print(f"Memory: {cache.memory_bytes() / 1024:.1f} KB for {len(cache)} tokens")
# Memory: 193.0 KB for 1000 tokens  (vs 500.0 KB fp16)

What's Next​

The main bottleneck right now is the Python-level encode/decode loop. Every token requires a numpy↔MLX transfer per attention head per layer — that's n_layers × n_heads small copies per generation step. A fused Metal kernel would reduce this to near-zero overhead.

Other items on the roadmap:

• Perplexity benchmark on WikiText-2 for quantitative quality measurement
• Value compression — int8 per-token is implemented, needs tuning at longer contexts
• Longer context evaluation — 8K, 32K token sequences where KV memory dominates
• Fused attention kernel — integrate BitPackBuffer unpack + Lloyd-Max decode + attention dot product into a single Metal dispatch

References​

• TurboQuant (ICLR 2026) — Zandieh et al., "Online Vector Quantization with Near-optimal Distortion Rate"
• PolarQuant (AISTATS 2026) — "PolarQuant: Quantizing KV Caches with Polar Transformation"
• QJL (2024) — Zandieh et al., "QJL: 1-Bit Quantized JL Transform for KV Cache Quantization"
• Apple MLX
• VeloxQuant-MLX on PyPI

TL;DR: The KV cache is the last major unoptimised bottleneck for local LLM inference on Apple Silicon. llama.cpp, Ollama, and LM Studio don't compress it. I built VeloxQuant-MLX to fix that — 9 quantization algorithms, custom Metal GPU kernels, up to 16× compression, plug into mlx_lm with 3 lines.

The wall every Mac LLM user hits​

You've got a MacBook with an M-series chip. You've downloaded llama.cpp, or Ollama, or LM Studio. You load a 7B model, start a conversation, and for the first few hundred tokens everything feels fast.

Then you push it. A long document. A complex coding task. A multi-turn conversation that's been going for a while. And suddenly — generation slows to a crawl, the fans spin up, or the process just dies.

This isn't a bug in those tools. It's a fundamental memory problem that none of them solve. And it lives in a part of the inference stack that almost nobody talks about: the KV cache.

What the KV cache actually is​

To understand the problem, you need to understand what happens inside a transformer when it generates text.

Every token in your context window requires the model to compute two matrices at every layer — a key and a value. These represent what that token "means" in context. Without caching, generating each new token would require recomputing these for every prior token — quadratic cost that makes long-context generation completely impractical.

So instead, the model stores them. Every time a token is processed, its key and value matrices are written to a cache. Future tokens look them up instead of recomputing them.

This is the KV cache. And it grows with every token you generate.

The math is simple:

memory = num_layers × num_kv_heads × head_dim × seq_len × 2 (K+V) × 2 bytes (fp16)

For Llama-3.1-8B at 8k context: 4.2 GB. At 32k context: 16.8 GB. On a machine with 16 GB of total unified memory.

Why Apple Silicon makes this worse, not better​

On a discrete GPU setup — say, an NVIDIA card with 24 GB VRAM — the GPU has its own dedicated memory pool. The KV cache lives there, separate from your system RAM.

Apple Silicon doesn't work that way. M-series chips use unified memory: the CPU, GPU, and Neural Engine all share the same physical memory pool. Your model weights, your KV cache, your browser tabs, your IDE, your OS — all competing for the same 16 or 24 GB.

A 7B model at 4-bit quantisation takes roughly 4–5 GB. Add OS and background processes (~4 GB), Chrome with a few tabs (~2 GB), your IDE (~1 GB). You're at 11–12 GB before you've generated a single token. On a 16 GB machine, you have 4 GB left for the KV cache — which runs out around 6k context for a 7B model.

On a 24 GB machine it's better, but not solved. During the development of VeloxQuant-MLX, I ran a benchmark sweep across 8 models. Qwen2.5-32B failed every quantization config because the weights alone (~17.5 GB) plus OS and development tools left negative headroom for activations and cache buffers. The memory watchdog had to kill the process before it triggered a kernel panic:

[07:28:09] free=62MB   inactive=1094MB pressure=unknown
[07:28:19] free=63MB   inactive=891MB  pressure=unknown
[07:28:19] CRITICAL: only 954MB available — killing PID 21438

System recovered to 10.2 GB free immediately after the kill, with no GPU stall.

What existing tools do about this: essentially nothing​

This is the part that surprised me most when I started digging.

llama.cpp has done extraordinary work on attention kernel optimisation, weight quantisation, batching, and speculative decoding. But the KV cache is stored at fp16 by default. There is experimental support (--cache-type-k q8_0) but it's not the default, not Metal-optimised, and the quantisation methods are basic scalar quantisation with no adaptation to actual key distributions.

Ollama builds on top of llama.cpp and inherits all of this. The UX is excellent. The KV cache situation is unchanged.

LM Studio adds a polished GUI and nice model management. Same story on the cache.

mlx_lm — Apple's own MLX-based inference library — is the most natural fit for Apple Silicon and does excellent work. But its default KV cache is full fp16. No built-in compression.

All of these tools have optimised everything around the KV cache. None of them have solved the KV cache itself.

The insight: compress the cache, not just the weights​

Most quantisation work in the LLM ecosystem targets weight quantisation — compressing model parameters from fp16 down to 8-bit, 4-bit, or lower. This is what GGUF, GPTQ, AWQ, and most Ollama models use.

Weight quantisation is a one-time operation. You compress once, save, load the compressed version.

KV cache quantisation is different. It happens at inference time, on every token, for every layer. The keys and values are different every generation — you can't pre-compute anything.

The key insight is that while weight distributions are relatively uniform, KV cache distributions are not. Keys tend to follow Gaussian or Laplacian distributions after a rotation transform. This structure can be exploited with much more aggressive compression than weight quantisation — down to 1 bit per dimension with surprisingly low quality loss.

The reason this works at such extreme bit rates: the attention computation only needs an approximation of the inner product q · k. Small errors in the reconstructed key vector translate to small errors in the attention score, which average out across many keys in the softmax.

What I built: VeloxQuant-MLX​

VeloxQuant-MLX is a KV cache compression library for Apple Silicon, built on top of MLX. It implements nine quantisation algorithms — from zero-calibration 1-bit methods to mixed-precision allocators — all backed by custom Metal GPU kernels compiled at runtime.

The algorithms​

TurboQuant RVQ — the recommended default​

Zero calibration. Works on any model immediately. Uses Residual Vector Quantisation with analytical Gaussian and Laplacian codebooks precomputed from distribution theory.

Two residual passes at 1 bit each → 7.5× compression with cosine similarity above 0.97.

import mlx_lm
from veloxquant_mlx import KVCacheBuilder, KVCacheConfig

model, tokenizer = mlx_lm.load("mlx-community/Llama-3.2-3B-Instruct-4bit")

config = KVCacheConfig(method="turboquant_rvq", bit_width_inlier=1, seed=42)
caches = KVCacheBuilder.for_model(model, config)
model.make_cache = lambda *_a, **_k: caches

response = mlx_lm.generate(
    model,
    tokenizer,
    prompt="Write a 3000-word analysis of the transformer architecture.",
    max_tokens=3000,
)

VecInfer — 16× with Metal acceleration​

Trades a 2-minute calibration step for 16× compression via Product Vector Quantisation. Per-channel smooth scaling handles outlier dimensions that defeat standard VQ. The hot path runs through a Metal GPU kernel that is 13× faster than equivalent MLX Python ops.

Standout result: Qwen2.5-7B VecInfer-1bit exceeds fp16 throughput at 16× compression, likely due to its strong GQA ratio reducing the number of KV heads.

RateQuant — best accuracy per bit​

Runs a 90-second sensitivity probe across all transformer layers, learns which layers are most sensitive to quantisation noise, then uses reverse-waterfilling to allocate more bits to sensitive layers and fewer to insensitive ones.

At 2.0 average bits, RateQuant achieves 2.7× lower perplexity degradation than uniform 2-bit quantisation at identical memory cost.

SpectralQuant — best quality at long context​

Keys in transformer models concentrate ~96% of their variance in just 3–4% of dimensions. SpectralQuant rotates keys into their eigenvector basis via SVD, then applies separate codebooks to the "signal" dimensions (high variance) and "noise" dimensions (low variance).

At 16k context the rotation trick matters enormously — quantisation errors accumulate over long sequences, and aligning the codec to the actual variance structure cuts that accumulation dramatically.

RaBitQ — maximum context length​

1-bit key compression via randomised Hadamard transform + binary sign packing + IVF clustering. Pairs with 4-bit MSE scalar quantisation on values for 6× full KV compression.

6× more context in the same RAM budget.

CommVQ — RoPE-compatible exact VQ (ICML 2025)​

Standard VQ breaks with RoPE positional encodings because quantize(rotate(x)) ≠ rotate(quantize(x)). CommVQ (Apple ML Research, ICML 2025) trains codebooks on pre-RoPE keys with a commutativity projection in the EM M-step. RoPE is applied exactly at decode time. 64× key compression with exact positional encoding.

The Metal kernels​

All the compression math in the world doesn't help if the quantisation itself is slow. VeloxQuant-MLX compiles nine Metal GPU kernels at runtime using mx.fast.metal_kernel:

The fused SDPA kernel is the most impactful. Without fusion, dequantisation creates a full fp16 key matrix in memory before attention — which defeats much of the point of compression. The fused path keeps the cache compressed all the way until attention scores are computed.

The 30-line Metal kernel that powers VecInfer's 13× speedup:

// One thread per sub-vector. Argmin lives in registers — no diff tensor.
uint vec_idx = thread_position_in_grid.x;
float best_dist = INFINITY;
uint  best_idx  = 0;

for (uint c = 0; c < n_centroids; ++c) {
    float dist = 0.0f;
    for (uint i = 0; i < sub_dim; ++i) {
        float d = float(x[x_base + i]) - float(codebook[cb_base + i]);
        dist += d * d;
    }
    if (dist < best_dist) { best_dist = dist; best_idx = c; }
}
out[vec_idx] = best_idx;

The key insight: the [N, n_centroids, sub_dim] diff tensor is never materialised. The argmin accumulator lives entirely in thread-local registers — 98% peak memory reduction at long context shapes.

Real numbers across 10 models​

End-to-end mlx_lm.generate, 200-token prompt, 120-token generation, Apple M-series:

RVQ-1bit is the safe default — within 5% of fp16 on most 7–8B models with zero calibration. VecInfer-1bit wins on compression (always 16×) and throughput on models with strong Grouped Query Attention ratios.

What this means in practice​

On a 16 GB MacBook M3:

• Before: 7B model maxes out around 6k context before generation slows
• After with RVQ 1-bit: that same cache now fits in ~450 MB, leaving headroom for 50k+ token contexts

On a 24 GB MacBook M3 Pro:

• Before: 13B models are borderline; long conversations risk OOM
• After: 13B at 32k context fits comfortably; the cache that would have consumed 8 GB uses 500 MB

Quick decision guide​

Try it​

pip install VeloxQuant-MLX

Requires macOS on Apple M1 or later · Python 3.11+ · MLX ≥ 0.18

📖 Documentation: https://veloxquant-mlx.netlify.app/docs/ 📦 PyPI: https://pypi.org/project/VeloxQuant-MLX/ ⭐ GitHub: https://github.com/rajveer43/turboquant_mac_implementation

What's next​

• Vision-language model support — visual tokens have fundamentally different KV distributions; heavy outliers require adapted algorithms
• Per-head RateQuant granularity — the paper allocates bits per layer-head group; current implementation is per-layer, leaving ~30% accuracy improvement on the table
• Gradient-based sensitivity for RateQuant — more accurate than activation-based, at the cost of a slightly longer calibration pass

The KV cache is the last major unoptimised bottleneck for local LLM inference on Apple Silicon. llama.cpp, Ollama, and LM Studio have built excellent tools — but none of them have solved this. VeloxQuant-MLX is my attempt.

If you're running models locally on a Mac and hitting memory walls — at any context length, on any model — I'd genuinely like to hear about it.

MIT License · Built for Apple Silicon · Engineered for speed