perf(metal): an order-preserving streamed qmv would give MTP the fast kernel and byte-identity together
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## 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` 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(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(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 #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 #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
- #1261: the measurement that priced the choice, and the four env recipes.
- #1199: the gate's `qmv_wide` retry and the process wide pin this would make unnecessary.
- #1258: the Gemma 4 arm of the same gate.
- #1279: a pairing where both kernels diverge, which this would not fix.
Contributor guide
Research direction
Start with docs/benchmark_results/qmv-wide-pin-tax-m3ultra-2026-08-22.md, docs/benchmarks.md, and the #1261 recipes; run a generation-15 microbenchmark to measure the recoverable gap and register-pressure limit. Then inspect src/lib/mlx-cpp/patches/mlx/backend/metal/quantized.cpp and tests/metal_block_vs_chain_op_parity.rs. Done means the measurement supports the work, parity and MTP exactness pass, throughput is compared, and the upstream-versus-overlay decision is recorded.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, rust
- Domain
- backend, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100