|
| 1 | +# Next phase — native 4-bit matmul (`gemv_4bit` / `gemm_4bit`) on Metal |
| 2 | + |
| 3 | +**Status:** plan, ready to dispatch cold · **Branch to open:** `feature/mps-matmul` (off origin/main) |
| 4 | +**Prereqs merged:** Phases 1–3 (native quantize/dequantize) + packaging. **Executor:** Fable or a fresh session. |
| 5 | + |
| 6 | +This is an executable spec in the shape of `PORT_PLAN.md`. It assumes the reader has NOT surveyed the |
| 7 | +codebase yet — everything needed to start is here. Line numbers cite a snapshot and **drift**; re-grep |
| 8 | +the symbol before editing. Read `PORT_PLAN.md` for the overall arc and `MPS_STATUS.md` §7–§9 for how the |
| 9 | +native pipe, buffer bridge, guard, and packaging already work — this phase reuses all of it. |
| 10 | + |
| 11 | +--- |
| 12 | + |
| 13 | +## §0 — Required reading (do not re-derive) |
| 14 | + |
| 15 | +- `MPS_STATUS.md` — ground truth. Especially §7 (the `data_ptr()`-is-the-`MTLBuffer` bridge, cross-queue |
| 16 | + sync, install-safe metallib load), §8 (the kernel/dispatch/routing pattern to copy), Sub-task 0 (the |
| 17 | + load-time buffer-contract guard — already in place, nothing to add). |
| 18 | +- **`apple-silicon` skill**, files: |
| 19 | + - `mps-matrix-multiplication.md` — `MPSMatrixMultiplication` GEMM (the "route matmul through MPS" |
| 20 | + option). **Read before choosing the design fork below.** |
| 21 | + - `compute-kernels-and-dispatch.md` — device→library→pipeline→encoder→commit chain and grid/threadgroup |
| 22 | + sizing (for the hand-fused option). |
| 23 | + - `simd-group-functions.md` + `math-functions-and-numeric-parity.md` — SIMD reductions and the |
| 24 | + fast-vs-precise / `-fno-fast-math` parity rules (we already compile the metallib `-fno-fast-math`). |
| 25 | + - `op-graduation-playbook.md` — the flush-first / synchronize discipline. |
| 26 | +- `ct2-internals` skill — CPU-as-oracle parity-tolerance methodology (per-dtype). |
| 27 | + |
| 28 | +Prior art to mine for the MPS GEMM route: the CTranslate2 Metal backend |
| 29 | +(`/Users/eeaglstun/Documents/dev/CTranslate2/`, `METAL_BACKEND.md`) routes matmul through |
| 30 | +`MPSMatrixMultiplication` — reuse its encode/commit skeleton, not its math. |
| 31 | + |
| 32 | +--- |
| 33 | + |
| 34 | +## §1 — Current state (what "unfused" means, precisely) |
| 35 | + |
| 36 | +The `mps` backend registers `gemv_4bit`, `gemv_4bit.out`, and `gemm_4bit` |
| 37 | +(`bitsandbytes/backends/mps/ops.py`). Today they are **dequantize-then-`F.linear`**: |
| 38 | + |
| 39 | +- `_gemv_4bit_impl` (~L368): tries an inert HF Hub kernel, then falls to |
| 40 | + `B_dq = _dequantize_4bit_impl(...); return torch.nn.functional.linear(A, B_dq)` (~L386–387). |
| 41 | +- `gemm_4bit` registration (~L416): unpacks nested/compressed absmax first |
| 42 | + (`dequantize_blockwise` + offset, ~L435–439), then `B_dq = _dequantize_4bit_impl(...); |
| 43 | +return F.linear(A, B_dq, bias)` (~L451–452). |
| 44 | +- Since Phase 3, `_dequantize_4bit_impl` routes to the **native** Metal dequant kernel when available |
| 45 | + (`MPS_STATUS.md` §8). So the pipeline today is: **native Metal dequant of B → materialize full fp/bf16 |
| 46 | + B_dq in memory → PyTorch MPS `F.linear`**. |
| 47 | + |
| 48 | +Op signatures (from `bitsandbytes/_ops.py`, re-grep before editing): |
| 49 | + |
| 50 | +``` |
| 51 | +gemv_4bit(A, B, int[] shapeB, absmax, code, blocksize) -> Tensor # (+ .out overload) |
| 52 | +gemm_4bit(A, B, int[] shapeB, absmax, blocksize, str quant_type, |
| 53 | + bias?, absmax_8bit?, absmax_code?, absmax_offset?) -> Tensor |
| 54 | +``` |
| 55 | + |
| 56 | +`A` is fp16/bf16/fp32 activations `[..., K]`; `B` is packed 4-bit weights (uint8 storage, or reinterpreted |
| 57 | +bf16/etc.) with logical shape `shapeB = [N, K]`; output is `[..., N]` in `A.dtype`. `gemv_4bit` is the |
| 58 | +`M == 1` case; `gemm_4bit` is general `M`. Nested absmax (compressed statistics) appears **only** in |
| 59 | +`gemm_4bit` and is already unpacked to a plain per-block fp32 `absmax` before the matmul — the matmul |
| 60 | +phase never sees nested absmax. |
| 61 | + |
| 62 | +**Correctness today:** parity tests pass (`tests/test_mps_parity.py::TestMatmul4bitParity`) — fp32 |
| 63 | +≤ ~8e-6, fp16/bf16 0.0 at tested sizes. So this phase is a **performance** graduation, not a correctness |
| 64 | +fix. The bar: stay within the **same documented tolerances** while removing the full-B_dq materialization |
| 65 | +and/or the round-trip to PyTorch's GEMM. |
| 66 | + |
| 67 | +--- |
| 68 | + |
| 69 | +## §2 — The design fork (decide in Phase M1 with a spike) |
| 70 | + |
| 71 | +Two ways to make the matmul native. Pick per-shape; they are not mutually exclusive. |
| 72 | + |
| 73 | +### Option A — `MPSMatrixMultiplication` on dequantized B (lower risk) |
| 74 | + |
| 75 | +Dequantize B with the existing native kernel into a scratch `MTLBuffer`, then run Apple's tuned |
| 76 | +`MPSMatrixMultiplication` (MPS GEMM) `A · B_dqᵀ` on-device, all inside one `.mm` entry point / one command |
| 77 | +buffer. Add bias + write output. |
| 78 | + |
| 79 | +- **Pros:** Apple-tuned GEMM (fast, handles fp16/bf16), minimal new MSL, lowest correctness risk, reuses |
| 80 | + `mps-matrix-multiplication.md` directly. Keeps dequant and matmul on **one queue / one commit** → |
| 81 | + removes the current PyTorch hop and its separate sync. |
| 82 | +- **Cons:** still materializes a full `B_dq` (no memory win over today); numeric parity depends on |
| 83 | + MPSMatMul's accumulation vs PyTorch's — must be re-validated against the CPU oracle (likely fine within |
| 84 | + fp16/bf16 tol, verify fp32). |
| 85 | +- **Best for:** `gemm_4bit` (general M), where a real GEMM dominates. |
| 86 | + |
| 87 | +### Option B — hand-fused dequant + matmul Metal kernel (higher ceiling) |
| 88 | + |
| 89 | +One MSL kernel reads packed 4-bit B + absmax + code and computes the dot products directly, dequantizing |
| 90 | +weights **in registers/threadgroup** without ever writing full `B_dq` to device memory. |
| 91 | + |
| 92 | +- **Pros:** no `B_dq` materialization (the real memory/bandwidth win, especially for `gemv_4bit` M==1, |
| 93 | + which is memory-bound); this is what the CUDA `gemv_4bit`/`gemm_4bit` kernels do. |
| 94 | +- **Cons:** most work and highest risk — tiling, SIMD-group reductions, fp16/bf16 accumulation, and |
| 95 | + per-block absmax indexing all have to match the oracle within tol; `-fno-fast-math` already set but FMA |
| 96 | + contraction / accumulation order still needs care (`math-functions-and-numeric-parity.md`). |
| 97 | +- **Best for:** `gemv_4bit` (M==1) first — it's the simplest fused case (matrix-vector, one output row of |
| 98 | + work per thread/threadgroup) and the biggest bandwidth win. |
| 99 | + |
| 100 | +**Recommended split:** Option B for `gemv_4bit` (M==1, memory-bound, tractable fused kernel), Option A |
| 101 | +(`MPSMatrixMultiplication` on native-dequant B) for `gemm_4bit` general M. Do a small **spike** first |
| 102 | +(Phase M1) that benchmarks A vs B on representative shapes before committing; correctness gate is the same |
| 103 | +either way. |
| 104 | + |
| 105 | +--- |
| 106 | + |
| 107 | +## §3 — Phased implementation |
| 108 | + |
| 109 | +### Phase M1 — spike + decision (no production kernel yet) |
| 110 | + |
| 111 | +1. Micro-bench: for `gemv` (M=1, N,K ∈ {4096,11008}) and `gemm` (M ∈ {8,64,512}), time today's |
| 112 | + dequant+`F.linear` vs (A) native-dequant + `MPSMatrixMultiplication`, on fp16/bf16. Confirm the |
| 113 | + MPS GEMM route is faster and within tolerance; record numbers. |
| 114 | +2. Decide the A/B split per op (default: B for gemv, A for gemm). Write the decision into `MPS_STATUS.md`. |
| 115 | + |
| 116 | +### Phase M2 — `gemv_4bit` native (prove the fused pipe with M==1) |
| 117 | + |
| 118 | +3. Implement the chosen route for `gemv_4bit`. If Option B: MSL `gemv_4bit` kernel — one threadgroup per |
| 119 | + output element (or per N-tile), each thread dequantizes its slice of packed B (reuse the nibble/absmax |
| 120 | + math from the Phase-3 `dequantize_4bit` kernel) and accumulates `sum_k A[k] * dequant(B[n,k])` in fp32, |
| 121 | + then casts to `A.dtype`. SIMD-group reduce per `simd-group-functions.md`. |
| 122 | +4. New `extern "C" bnb_mps_gemv_4bit(...)` in `csrc/mps_ops.mm` (copy the encode/commit/`waitUntilCompleted` |
| 123 | + skeleton from `bnb_mps_quantize_4bit`); new pipeline in the cache; new argtypes in |
| 124 | + `cextension.py::MpsBNBNativeLibrary`. Route in `_gemv_4bit_impl` **when `_native_available()`** else the |
| 125 | + existing dequant+linear fallback (never regress the fallback). |
| 126 | +5. Parity: `TestMatmul4bitParity::test_gemv_4bit` already exists — extend / assert it hits native under |
| 127 | + `BNB_MPS_REQUIRE_NATIVE=1`. Bias handling: `gemv_4bit` has no bias; `gemm_4bit` does. |
| 128 | + |
| 129 | +### Phase M3 — `gemm_4bit` native (general M) |
| 130 | + |
| 131 | +6. Implement the chosen route (default: native dequant into scratch buffer + `MPSMatrixMultiplication`, |
| 132 | + bias epilogue). Handle the already-unpacked plain `absmax` (nested absmax is unpacked before this op — |
| 133 | + do not re-handle it). Route in the `gemm_4bit` registration; keep the pure-torch fallback. |
| 134 | +7. Parity: `TestMatmul4bitParity::test_gemm_4bit` (± bias, ± compressed/nested absmax) stays green and hits |
| 135 | + native. |
| 136 | + |
| 137 | +### Phase M4 — address the input-copy debt (see §5) + docs |
| 138 | + |
| 139 | +8. Kill the per-call offset-0 input copy where safe (pass `storage_offset` through the ABI and bind at a |
| 140 | + byte offset). 9. Update `MPS_STATUS.md`, `README.md` accelerator row (QLoRA 4-bit 🐢 → ✅ once fused |
| 141 | + and fast), and `docs/apple_silicon/README.md` (gemv/gemm 〰️ → ✅). |
| 142 | + |
| 143 | +--- |
| 144 | + |
| 145 | +## §4 — Reuse (do not reinvent) |
| 146 | + |
| 147 | +- **Buffer bridge:** torch MPS `tensor.data_ptr()` **is** the `id<MTLBuffer>` — cast the ctypes `void*` |
| 148 | + directly (`MPS_STATUS.md` §7). The load-time guard (`bnb_mps_check_buffer_contract` / |
| 149 | + `verify_buffer_contract()`) is already in place; no new guard needed. |
| 150 | +- **Sync:** own command queue; `torch.mps.synchronize()` before the call, `waitUntilCompleted` after |
| 151 | + (`MPS_STATUS.md` §7). For the MPSMatMul route, dequant + GEMM go on the **same** command buffer. |
| 152 | +- **Metallib load:** `dladdr`-relative, install-safe, already done. New kernels just add functions to |
| 153 | + `csrc/mps_kernels.metal` (built `-fno-fast-math`). |
| 154 | +- **Packaging:** already ships the dylib + metallib in the wheel (`MPS_STATUS.md` §9) — no change. |
| 155 | +- **Dispatch/registration pattern:** copy an existing `extern "C"` entry + `get_pipeline` + `dispatch_*` |
| 156 | + in `mps_ops.mm`, the argtypes block in `cextension.py`, and the `if _native_available(): ... else |
| 157 | +<fallback>` routing in `mps/ops.py`. All three already exist for four ops. |
| 158 | + |
| 159 | +--- |
| 160 | + |
| 161 | +## §5 — Open debt this phase should also close |
| 162 | + |
| 163 | +**Per-call offset-0 input copy.** Native ops call `_ensure_native_buffer(...)`, which `.contiguous()` and |
| 164 | +`.clone()`s any tensor with `storage_offset != 0`, because the `data_ptr()`-as-`MTLBuffer` cast is only |
| 165 | +valid at offset 0. For matmul this copies `A` (and any non-offset-0 operand) every call. Fix: pass each |
| 166 | +tensor's `storage_offset() * itemsize` through the C ABI and bind with `[enc setBuffer:buf offset:byteOff |
| 167 | +atIndex:i]` instead of forcing offset 0 — but first **verify** that a torch MPS view's `data_ptr()` for a |
| 168 | +`storage_offset != 0` tensor is `buffer_object + offset` (a mis-cast) vs the base buffer object; the |
| 169 | +Phase-1 probe suggested the former, so binding the base buffer at a byte offset needs the base pointer, |
| 170 | +which torch may not expose. If it can't be done safely, **leave the copy and document why** — do not ship a |
| 171 | +mis-cast. This is a correctness-sensitive optimization; treat it as such. |
| 172 | + |
| 173 | +--- |
| 174 | + |
| 175 | +## §6 — Definition of done |
| 176 | + |
| 177 | +1. `gemv_4bit` and `gemm_4bit` run through native Metal (fused kernel and/or `MPSMatrixMultiplication`), |
| 178 | + matching the CPU oracle within the documented per-dtype tolerances (fp32 ~1e-5, fp16 ~1e-2, bf16 ~4e-2). |
| 179 | +2. `pytest tests/test_mps_parity.py` green; `TestMatmul4bitParity` asserts native under |
| 180 | + `BNB_MPS_REQUIRE_NATIVE=1`; graceful fallback preserved when the native lib is absent. |
| 181 | +3. A recorded speedup vs the dequant+`F.linear` baseline on representative shapes (the reason this phase |
| 182 | + exists). |
| 183 | +4. Docs updated: `MPS_STATUS.md` (new §), `docs/apple_silicon/README.md` (gemv/gemm status), and the |
| 184 | + `README.md` accelerator row **only if** the result is genuinely fast (else keep 🐢 — do not overclaim). |
| 185 | +5. `pre-commit run --all-files` clean (clang-format on `.mm`/`.metal`). |
| 186 | + |
| 187 | +--- |
| 188 | + |
| 189 | +## §7 — Out of scope (still) |
| 190 | + |
| 191 | +- ❌ LLM.int8() native path — separate scope. |
| 192 | +- ❌ 8-bit optimizer native kernels — separate track. |
| 193 | +- ❌ Any change to the quant/dequant kernels' numerics (they are bit-exact; don't touch). |
0 commit comments