perf(metal): an order-preserving streamed qmv would give MTP the fast kernel and byte-identity together

[inureyes]

Problem

On Apple GPU generation 15 and later the MTP exactness gate has to choose between speed and the temperature-0 byte-identity contract, and no configuration currently keeps both. Measured on M3 Ultra in #1261 (docs/benchmark_results/qmv-wide-pin-tax-m3ultra-2026-08-22.md) on gemma-4-12b-it-4bit plus its 4-bit assistant, block 5, code prompt, 300 tokens: the default env serves the narrow kernel at 117.29 and 116.98 tok/s with byte-identity kept, while MLXCEL_QMV_WIDE=1 MLXCEL_MTP_ALLOW_INEXACT=1 serves the wide kernel at 139.18 and 139.12 tok/s with the contract forfeited. That is about 19% bought by giving up reproducibility against classic decode. The choice is forced by a property of the kernels that nothing actually requires, and this issue is about removing it.

The mechanism, stated precisely

qmv and qmv_wide compute the same dot product in different summation orders, floating point addition is not associative, so the last ulp differs and the divergence amplifies through the layers into a different sampled token. Read against mlx/backend/metal/kernels/quantized.h at pin 9a795735 (qmv_impl at :825, qmv_wide_impl at :989), there are four independent order differences, any one of which alone is sufficient to break bit-equality:

1. K to lane mapping. qmv_impl gives each of the 32 lanes a contiguous values_per_thread chunk inside blocks of block_size = values_per_thread * SIMD_SIZE and advances block by block (x += tid.x * in_vec_size + simd_lid * values_per_thread, then k += block_size, :874 and :934). qmv_wide_impl instead walks groups strided by the lane index: for (int g = k_lane; g < in_vec_size_g; g += k_lanes) (:1031).
2. Lane count in the reduction. qmv_impl always reduces over the full simdgroup with simd_sum (:977). qmv_wide_impl reduces over k_lanes only, which the dispatch sets to 8 for affine and 16 for fp (quantized.cpp overlay :613), because a simdgroup there spans SIMD_SIZE / k_lanes output rows (:1002).
3. Intra-lane accumulation depth. qmv_impl accumulates each block's qdot result straight into result[row] (:965). qmv_wide_impl adds a level: it sums sub = 8 products into a local acc, then does result[v] += acc (:1045-1051).
4. Final reduction tree. simd_sum's butterfly over 32 lanes versus qmv_wide_impl's explicit simd_shuffle_down ladder over k_lanes (:1056-1071).

The observation this rests on

The speedup and the reduction order are independent properties. qmv_wide's gain is weight reuse: it decodes each 8 value sub-chunk once and reuses it across vecs_per_tg streamed input vectors. qmv at M > 1 maps the vector index onto the grid instead (x += tid.x * in_vec_size, y += tid.x * out_vec_size), so it re-reads and re-decodes the entire weight row once per vector. That is a memory traffic property, and it is what the measured 1.67x to 1.92x verify-forward gap at M = 10 to 13 in docs/benchmarks.md (the width table at :612-618) is paying for. Nothing about it requires the reduction order to change.

So a kernel that keeps qmv's exact reduction order while hoisting the weight load out of a loop over streamed vectors should be bit-identical to qmv at M = 1 by construction, not by tolerance, while recovering some or all of the traffic gain.

Correction to the obvious sketch: qdot does not dequantize

This trap is worth stating up front because falling into it reintroduces the exact bug the issue exists to avoid. qdot (:192) never materializes dequantized weights. It accumulates x_thread[i] against the raw masked quantized bits (for bits == 4, x_thread[4*i] * (ws[i] & 0x000f) + x_thread[4*i+1] * (ws[i] & 0x00f0) + ..., :234-243), with load_vector (:29) having pre-divided x_thread by the power of two each mask leaves in, and applies the group parameters exactly once at the end: return scale * accum + sum * bias; (:289). qmv_wide_impl does the opposite, calling dequantize<U, sub, bits> into w_dq and accumulating xc[i] * w_dq[i] (:1039-1046).

That is a fifth order difference, and it is arguably the largest one: scale * (sum_i q_i x_i) + bias * (sum_i x_i) against sum_i (scale * q_i + bias) * x_i. Splitting qdot into a dequantize step and a dot step, which is the natural way to hoist the decode, would therefore break bit-identity on its own even if all four differences above were fixed. The hoist has to keep qdot's expression verbatim and reuse the packed bytes, or the raw masked values in the U domain, which is order-neutral because x_thread[4*i] * (ws[i] & 0x000f) and x_thread[4*i] * q0 with q0 = U(ws[i] & 0x000f) are the same float multiply of the same two values.

Step 2 (gated on step 1 below): the shape of the kernel

A modification of qmv_impl, not of qmv_wide_impl, with V the streamed vector count:

thread U x_thread[V][values_per_thread];
thread U result[results_per_simdgroup][V] = {0};

int k = 0;
for (; k < in_vec_size - block_size; k += block_size) {   // block loop and K-to-lane mapping unchanged
  U sum[V];
  for (int v = 0; v < V; v++) {
    sum[v] = load_vector<T, U, values_per_thread, bits>(xv[v], x_thread[v]);
  }
  for (int row = 0; row < results_per_simdgroup; row++) {
    auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
    thread uint8_t wreg[packs_per_thread * bytes_per_pack];   // read from device once, reused across V
    for (int i = 0; i < packs_per_thread * bytes_per_pack; i++) { wreg[i] = wl[i]; }
    U s = (scales + row * in_vec_size_g)[0];
    U b = (biases + row * in_vec_size_g)[0];
    for (int v = 0; v < V; v++) {
      result[row][v] += qdot_reg<U, values_per_thread, bits>(wreg, x_thread[v], s, b, sum[v]);   // body byte-identical to qdot
    }
  }
  // advance ws, scales, biases and every xv[v] by one block, as qmv_impl does
}

for (int row = 0; row < results_per_simdgroup; row++) {
  for (int v = 0; v < V; v++) { result[row][v] = simd_sum(result[row][v]); }
}

qdot_reg is qdot with the weight pointer's address space changed from device to thread and the body otherwise copied character for character. The K to lane mapping, the per lane accumulation order, the scale * accum + sum * bias fold and the simd_sum tree are all untouched, which is what makes each vector's result equal to what qmv produces for that vector alone. The remainder tail (qdot_safe, :294) gets the same treatment.

It is worth trying the pure loop restructuring first, without qdot_reg: nothing aliases wl inside the v loop, so the compiler may already hoist the load and give the traffic reuse for free. If it does, the whole change is a loop nest reshuffle and the overlay stays small. qdot_reg is the explicit fallback if the generated code says otherwise.

Step 1: measure the ceiling before building any of it

Register pressure is the central risk and it is cheap to price. x_thread and result both grow linearly in V, and this kernel is pinned to qmv's shape of results_per_simdgroup = 4 rows with all 32 lanes reducing each row. Concretely, at bits == 4 in qmv_impl: pack_factor = 32/4 = 8, packs_per_thread = 1, so values_per_thread = 8 and the per thread footprint is 8V + 4V = 12V floats, meaning 48 registers at V = 4 and 60 at V = 5. qmv_wide sidesteps this by using fewer lanes per row (k_lanes = 8 for affine) so a simdgroup covers 4 rows with a much smaller per lane accumulator, result[vecs_per_tg] plus w_dq[8].

The bound is friendlier than it first looks: the dispatch already caps vecs_per_tg at 5 (n_tiles = (M + 4) / 5 in the quantized.cpp overlay :607-608), so V never needs to exceed 5 to match what qmv_wide itself streams. Whether occupancy still caps the recovered speedup below that is the whole question, and it should be answered with a microbenchmark on a generation 15 or later host before any integration work.

If the recovered fraction is small, recording the number and closing is a valid outcome, in the same spirit as #1261's exit condition. Nothing below the measurement is worth building on a guess.

Where the kernel should live

The current MLX overlay is deliberately minimal. src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp carries a small delta (the flag, its one call site in use_qmv_wide, and the cxx entry points) precisely so that a pin bump refreshes the file and reapplies them, which is what its own header comment commits to. A new Metal kernel is a far larger overlay surface to carry across bumps, and quantized.h is not currently overlaid at all.

Propose this upstream to ml-explore/mlx first. The ordering property is generally useful to anyone who needs reproducible quantized matmuls across batch sizes, not just to this project, and upstream is where a kernel of this shape belongs. Treat a local overlay as the fallback, with its maintenance cost acknowledged when that decision is made rather than discovered at the next bump.

Scope limits, stated honestly

• This addresses the qmv_wide divergence mechanism only. It would not fix fix(speculative): the Gemma 4 31B pairing's exactness probe fails both kernels on M3 Ultra, vetoing the burst #1217 just enabled #1279, where the Gemma 4 31B plus bf16 pairing probes non-identical under both kernels on M3 Ultra, nor the M1 Ultra prose divergence recorded in docs/benchmarks.md, which occurs on generation 13 where qmv_wide is never taken. At least one other divergence mechanism exists and is out of scope here.
• The kernel only needs to serve M in [2, get_qmv_batch_limit). Above the limit qmm takes over and the contract is already forfeited there, which usefully bounds V alongside the vecs_per_tg cap of 5.
• qmv_quad (the K == 128 || K == 64 path, quantized.cpp overlay :1800) is separate and unaffected.
• The fp modes in fp_quantized.h have their own fp_qmv_wide (:562) and take it on every GPU generation, since use_qmv_wide is mode != "affine" || arch_gen >= 15 (overlay :583-586). An mxfp4 target therefore breaks the contract even on generations 13 and 14, where the affine path is free. The same treatment there is arguably worth more, because no generation escapes it. Second phase, not folded into the first.

Acceptance criteria

• A microbenchmark decides, on a generation 15 or later host, what fraction of the 1.67x to 1.92x qmv versus qmv_wide gap an order-preserving streamed qmv recovers at M in [2, get_qmv_batch_limit), and at what V register pressure caps it.
• If the recovered fraction justifies building it, the kernel is bit-identical to qmv at M = 1 for every M it serves, verified with the existing op-level harness tests/metal_block_vs_chain_op_parity.rs.
• End to end, the MTP exactness probe passes with the new kernel active, and throughput is measured against both current arms (117 and 139 tok/s on M3 Ultra for the perf(speculative): the qmv_wide exactness buy-back pins the whole process narrow, taxing batched decode that never asked for it #1261 pairing) so the recovered fraction is stated in the same units as the cost it removes.
• The upstream-versus-overlay decision is made explicitly and recorded, not defaulted into.
• If register pressure caps the gain below the point where it is worth carrying, the measurement is recorded and the issue closes on the number.

References

• perf(speculative): the qmv_wide exactness buy-back pins the whole process narrow, taxing batched decode that never asked for it #1261: the measurement that priced the choice, and the four env recipes.
• feat(speculative): let the MTP gate buy back exactness by dropping qmv_wide #1199: the gate's qmv_wide retry and the process wide pin this would make unnecessary.
• fix(speculative): probe Gemma 4 MTP exactness instead of assuming it #1258: the Gemma 4 arm of the same gate.
• fix(speculative): the Gemma 4 31B pairing's exactness probe fails both kernels on M3 Ultra, vetoing the burst #1217 just enabled #1279: a pairing where both kernels diverge, which this would not fix.