2 posts tagged with "turboquant"

How I found the one cache in this library that actually stores compressed KV bytes at rest, fused its two-stage quantize + two-stage bit-pack into a single Metal kernel, and measured 1.47×–2.48× — a speedup that, unlike the last one, survives into the real pipeline because this cache's memory savings were never just accounting.

The issue read simply: "fuse KV quantization and packing into a single kernel." Five stages, on paper — statistics, quantize, pack, cache write — each one a round trip to memory that didn't need to happen. The proposal was to collapse them into one pass.

The complication came before I wrote a single line of Metal: this library has several KV-cache methods, and they don't all have a quantize-then-pack pipeline to fuse. Some quantize and immediately dequantize back to fp16, keeping no packed bytes anywhere. Finding the one method where "pack" is a real, physical step — not an accounting fiction — turned out to be most of the work.

First, which pipeline is this issue actually about?​

The obvious place to look was KIVI, the library's reference baseline and the subject of an earlier fused-kernel post. KIVI already has a fused Metal kernel for its quantize step. So I read it first.

# veloxquant_mlx/cache/kivi_cache.py — KIVIKVCache.update_and_fetch
k_q = self._quant_dequant_along(self.keys[:, :, lo:hi, :], axis=-2)
v_q = self._quant_dequant_along(self.values[:, :, lo:hi, :], axis=-1)
self.keys[:, :, lo:hi, :] = k_q
self.values[:, :, lo:hi, :] = v_q

k_q is written straight back into self.keys — which is fp16. There is no bit-packing anywhere in this path. The existing fused kernel (kivi_group_quant_dequant) computes quantize-then-immediately-dequantize in one dispatch — a genuinely useful fusion, but a different one than the issue describes. It replaces eight MLX ops with one kernel; it was never going to produce packed bytes, because KIVI's design doesn't keep any.

That's not a guess — it's written down. docs/PACKED_STORAGE_ROADMAP.md keeps exactly this distinction as a standing table:

One row answers the question. turboquant_rvq is the one cache in this library that keeps its compressed representation resident — not dequantized back to fp16, not an accounting estimate, actual uint32 bytes sitting in the KV cache between calls. That makes it the only place where "fuse quantize and pack" is fusing two things that both really happen.

The pipeline, as it actually runs​

TurboQuantRVQKVCache implements two-stage residual vector quantization: rotate, quantize against a Gaussian codebook, then quantize the residual against a second, Laplacian-fit codebook. Two index streams, each bit-packed separately:

# veloxquant_mlx/cache/turboquant_rvq_cache.py — before this change
ev = self._quantizer.encode(k_unit)
idx1 = ev.indices  # (B*H*S, D) uint8 — stage-1 codes
idx2 = ev.signs.astype(mx.uint8)  # (B*H*S, D) uint8 — stage-2 codes

p1 = _pack_indices(idx1, self._bits).reshape(B, H, S, self._n_words)
p2 = _pack_indices(idx2, self._bits).reshape(B, H, S, self._n_words)

Unrolling self._quantizer.encode, the full chain is:

rotate (Metal, mx.hadamard_transform)
  -> quantize1        (MLX broadcast-compare, ScalarCodebook.quantize)
  -> dequantize1       (MLX gather)
  -> residual          (MLX subtract)
  -> quantize2         (MLX broadcast-compare)
  -> [idx1 uint8 buffer]  -> pack1 (MLX bit-shift/sum) -> [packed1 uint32]
  -> [idx2 uint8 buffer]  -> pack2 (MLX bit-shift/sum) -> [packed2 uint32]

Five MLX dispatches after the rotation, and two full-size (N, D) uint8 arrays materialize and get written to memory purely so the next kernel can read them straight back and throw them away. Everything from quantize1 to pack2 is, coordinate by coordinate, a single streaming computation — nearest-boundary lookup, subtract, nearest-boundary lookup again, shift-and-OR into a word. There's no reason any of it needs to leave registers.

What "quantize" means here, precisely​

ScalarCodebook.quantize doesn't do a naive argmin over centroids. It counts boundary crossings:

# veloxquant_mlx/codebooks/scalar_codebook.py
cmp = y[:, :, None] > self._boundaries_mx[None, None, :]
return mx.sum(cmp.astype(mx.uint8), axis=-1).astype(mx.uint8)

boundaries are the midpoints between sorted centroids. idx = count(y > boundary_k) — for k sorted boundaries, that count is the nearest-centroid index, computed without an abs() or an argmin(). This matters for the fused kernel: replicating "nearest centroid" with a naive Metal argmin loop and expecting it to match this exactly is the kind of thing that looks right and is wrong on ties. I needed to replicate the boundary-count, not just the result — same operation, same order, so floating-point ties resolve identically on both sides.

A fusion this codebase had already proven out​

I wasn't inventing the pattern. rabitq_encode.metal already fuses rotate → binarize → pack → magnitude into one kernel, for RaBitQ's 1-bit sign codes:

// veloxquant_mlx/metal/src/rabitq_encode.metal
simd_vote ballot = simd_ballot(y >= 0.0f);
uint mask = uint(static_cast<simd_vote::vote_t>(ballot));

if (sl == 0u) {
    uint start = sg * 4u;
    for (uint j = 0u; j < 4u && (start + j) < uint(N_BYTES); ++j) {
        k_bits[n * uint(N_BYTES) + start + j] = uint8_t((mask >> (8u * j)) & 0xFFu);
    }
}

simd_ballot packs 32 lanes' sign bits into one 32-bit mask in a single instruction — elegant, but specific to 1-bit codes, where "pack" is literally "collect the sign bits." RVQ's codes are 1–4 bits each, from two independent codebooks, so I couldn't reuse the ballot trick directly. But the shape of the fusion — rotate once, stay in registers, write packed bytes and nothing else — was exactly the template to extend.

The kernel​

One threadgroup per rotated vector, one thread per coordinate:

// veloxquant_mlx/metal/src/rvq_quant_pack.metal
threadgroup uint8_t idx1_buf[MAX_D];
threadgroup uint8_t idx2_buf[MAX_D];

uint n    = threadgroup_position_in_grid.x;
uint lane = thread_position_in_threadgroup.x;
uint D    = uint(MAX_D);

float y = float(rotated[n * D + lane]);

uint idx1 = 0u;
for (uint k = 0; k < K1; ++k) {
    idx1 += uint(y > float(boundaries1[k]));
}
float y_hat1 = float(centroids1[idx1]);
float r1 = y - y_hat1;

uint idx2 = 0u;
for (uint k = 0; k < K2; ++k) {
    idx2 += uint(r1 > float(boundaries2[k]));
}

idx1_buf[lane] = uint8_t(idx1);
idx2_buf[lane] = uint8_t(idx2);
threadgroup_barrier(mem_flags::mem_threadgroup);

Each lane runs stage-1 quantize, gathers its own centroid, computes the residual, and runs stage-2 quantize — all in registers. K1/K2 are 2^bits - 1, so at bits ≤ 4 that's at most 15 comparisons per stage, fully unrolled at compile time since BITS is a template parameter, the same trick scalar_quantize.metal uses for its centroid scan.

The packing step needs some cross-lane visibility — a word holds 32 / bits lanes' worth of codes — so each lane stages its own two codes into threadgroup memory, then one barrier, then the first lane of every word-sized group does the packing:

if (lane % ELEMS_PER_WORD == 0u) {
    uint word_idx = lane / ELEMS_PER_WORD;
    uint w1 = 0u;
    uint w2 = 0u;
    for (uint j = 0; j < ELEMS_PER_WORD; ++j) {
        uint d = lane + j;
        if (d < D) {
            w1 |= (uint(idx1_buf[d]) & MASK) << (j * BITS);
            w2 |= (uint(idx2_buf[d]) & MASK) << (j * BITS);
        }
    }
    packed1[n * n_words + word_idx] = w1;
    packed2[n * n_words + word_idx] = w2;
}

No global-memory round trip between quantize and pack — the handoff happens entirely through idx1_buf/idx2_buf, which never leave the threadgroup. One barrier, one dispatch, two packed uint32 streams out.

The Python side matches the shape of every other kernel wrapper in this codebase — compile-time #defines for D and BITS, cached per configuration:

def _quant_pack_kernel(d: int, bits: int):
    key = ("rvq_quant_pack", d, bits)
    if key not in _cache:
        _cache[key] = mx.fast.metal_kernel(
            name=f"rvq_quant_pack_d{d}_b{bits}",
            input_names=["rotated", "centroids1", "boundaries1", "boundaries2"],
            output_names=["packed1", "packed2"],
            header=f"#define MAX_D {d}\n#define BITS {bits}u\n",
            source=_RVQ_QUANT_PACK_SRC,
            ensure_row_contiguous=True,
        )
    return _cache[key]

The rotation itself — mx.hadamard_transform — stays outside this kernel and un-fused. It already has its own dedicated Metal implementation, and pulling it in would mean re-deriving MLX's Walsh-Hadamard butterfly rather than fusing the actual multi-stage pipeline the issue names. Scope the fusion to what's genuinely five redundant passes, not to everything upstream of it.

Bit-exactness, checked the boring way​

The bar, same as the KIVI kernel before it, was identical output — not "close." I wrote the parity test before trusting any benchmark number:

@pytest.mark.parametrize("D", [8, 16, 32, 64, 128, 256])
@pytest.mark.parametrize("bits", [1, 2, 3, 4])
@pytest.mark.parametrize("N", [1, 5, 33])
def test_rvq_quant_pack_bit_exact(D, bits, N):
    q = TurboQuantRVQ(d=D, b=bits, seed=D + bits + N, use_hadamard=True)
    x = mx.array(rng.standard_normal((N, D)).astype(np.float16))

    p1_ref, p2_ref = _reference_pack(q, x, bits)  # MLX path: encode + _pack_indices
    p1_got, p2_got = q.encode_pack(x)  # fused kernel

    np.testing.assert_array_equal(np.array(p1_got), np.array(p1_ref))
    np.testing.assert_array_equal(np.array(p2_got), np.array(p2_ref))

72 parametrized cases — every dimension from 8 to 256, every bit-width from 1 to 4, several batch sizes and seeds — plus a direct end-to-end cache test comparing update_and_fetch under the fused path against the MLX path on the same input. All identical, first run:

p1 match: True
p2 match: True

No FMA surprises this time, no rounding-mode mismatch, no padding edge case — the boundary-count replication and the LSB-first packing order were the two places bit-exactness could plausibly have slipped, and both held. The full existing suite (2464 tests) stayed green.

Wired in as a fallback, not a replacement​

The fused path degrades to the existing MLX path when the head dimension isn't a power of two or exceeds Metal's 1024-thread threadgroup limit, and it latches off permanently if the kernel ever throws — the same defensive pattern KIVIKVCache already uses for its own Metal path:

if self._use_metal_pack:
    try:
        p1_flat, p2_flat = self._quantizer.encode_pack(k_unit)
        p1 = p1_flat.reshape(B, H, S, self._n_words)
        p2 = p2_flat.reshape(B, H, S, self._n_words)
    except Exception:
        self._use_metal_pack = False  # don't re-pay the failure every call
        p1 = p2 = None
else:
    p1 = p2 = None

if p1 is None:
    ev = self._quantizer.encode(k_unit)
    p1 = _pack_indices(ev.indices, self._bits).reshape(B, H, S, self._n_words)
    p2 = _pack_indices(ev.signs.astype(mx.uint8), self._bits).reshape(B, H, S, self._n_words)

Since both paths are bit-identical, this is a pure performance knob — no code path can produce a different cached value, so there's no way for the fallback to silently change a benchmark result the way an accuracy trade-off would.

The measurement, and why I trust it more this time​

I measured the same way the KIVI post's Lie #3 taught me to: interleaved A/B sampling, rotating pool of 10 distinct inputs so MLX's operation cache can't hand either arm a shortcut, median over repeated trials.

The shape is exactly what fusing memory-bound work predicts: the win grows with block size, because the eliminated intermediates scale with the bytes moved, not with a fixed dispatch cost. At N=8192 the fused kernel is doing meaningfully more useful work per byte moved — 7.60 GB/s of "real" output traffic against 4.43 GB/s for a path that spends part of its bandwidth writing and re-reading index buffers nobody needed.

What made this one different from the last one​

The KIVI post's real lesson was that a legitimate op-level speedup can vanish completely once you account for how much of the total runtime the operation represents, and that four separate benchmarking mistakes can each manufacture a plausible number that says otherwise. None of those traps were specific to KIVI — they're generic to benchmarking GPU kernels inside a lazily-evaluated framework. So this time:

• The pipeline was fused before the benchmark was written, not the reverse — the correctness test (test_rvq_quant_pack_bit_exact) ran and passed before a single latency number existed.
• The benchmark reused the interleaved-pool method from the start, rather than discovering the naive version was lying after shipping a wrong conclusion.
• The "does this method even keep packed bytes" question got asked before writing any Metal, which is the step that would have prevented fusing the wrong pipeline entirely.

The honest caveat is the one the box above states directly: op-level speedup on a real storage path is a necessary condition for an end-to-end win, not a sufficient one. Measuring update_and_fetch in isolation, across models with different head-dim/KV-head geometries, on an idle GPU with the same interleaved discipline — that's the follow-up this post is setting up, not one it's claiming to have already answered.

Kernel in veloxquant_mlx/metal/src/rvq_quant_pack.metal, wrapper in veloxquant_mlx/metal/_rvq_quant_pack.py, wired in via TurboQuantRVQ.encode_pack() in veloxquant_mlx/quantizers/turboquant_rvq.py. Tests in veloxquant_mlx/tests/metal/test_rvq_quant_pack.py. See docs/PACKED_STORAGE_ROADMAP.md for which methods store compressed bytes at rest versus which report accounting-only ratios, and the KIVI Metal kernel post for the benchmarking failure modes this measurement was written to avoid. All measurements on an Apple M4 with MLX. PR: #269, issue #251.

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