2 posts tagged with "performance"

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