Skip to content

Commit 4598cd2

Browse files
eaglstunclaude
andcommitted
Phase M1: record 4-bit matmul baseline + A/B design decision
Benchmarked today's unfused gemv_4bit/gemm_4bit (native dequant of B -> F.linear) on this machine to ground the NEXT_MATMUL_PLAN.md design fork with real numbers instead of intuition. Findings (MPS_STATUS.md §10): - gemv (M=1) is 80-90% dequant-bound; the GEMM is ~10%. Option A (MPSMatMul on materialized B_dq) can only touch that 10% -> fusion (Option B) is the win. - gemm has a fixed ~0.75ms dequant floor; the GEMM overtakes it near M~512. Large-M gemm -> Option A (don't hand-beat Apple's tuned GEMM). - The existing dequant kernel runs at ~54 GB/s vs ~400 GB/s achievable, so the M2 target is memory bandwidth, not "fusion" per se: a fused gemv that reads packed B no faster than today wins nothing. Decision: Option B (fused) for gemv_4bit = Phase M2; Option A (MPSMatrixMultiplication) for large-M gemm_4bit = Phase M3. Matches the plan's recommended split, now evidence-backed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 642ecb1 commit 4598cd2

1 file changed

Lines changed: 50 additions & 0 deletions

File tree

docs/apple_silicon/MPS_STATUS.md

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -367,3 +367,53 @@ the _installed_ package. Confirmed against the installed wheel:
367367
**Still open (not this task):** the wheel is a plain-tagged platform wheel; CI matrix / release
368368
automation to actually publish MPS wheels is a separate concern. The per-call offset-0 input copy and
369369
the fused 4-bit matmuls remain as documented above.
370+
371+
---
372+
373+
## 10. Phase M1 — 4-bit matmul baseline + A/B decision (spike)
374+
375+
**Status: measured. Decision made.** This is the `NEXT_MATMUL_PLAN.md` Phase M1 spike, but done
376+
against the _real_ baseline (today's `dequant → F.linear`) rather than an unbuilt native route — the
377+
numbers decide the design fork on their own, so no throwaway MPSMatMul wiring was needed to choose.
378+
379+
**Method.** `torch.mps.synchronize()`-bracketed timing, warmup + 30–50 iters, native dequant forced
380+
(`BNB_MPS_REQUIRE_NATIVE=1`), nf4/blocksize-64. Per shape we isolate the two costs inside today's
381+
unfused path: the native Metal **dequant of B** (materializes full `B_dq`) and the **`F.linear`** GEMM
382+
on that materialized `B_dq`. Bench scripts: `scratchpad/bench_matmul_baseline.py`,
383+
`bench_gemm_baseline.py` (not committed; reproduce from the numbers here).
384+
385+
**gemv (M=1), fp16/bf16** — dequant is the whole cost:
386+
387+
| N | K | total | dequant | linear | dequant share |
388+
| ----- | ----- | ------ | ------- | ------ | ------------- |
389+
| 4096 | 4096 | 0.94ms | 0.75ms | 0.07ms | ~80% |
390+
| 11008 | 4096 | 1.80ms | ~2.0ms | 0.19ms | ~90%+ |
391+
| 4096 | 11008 | 1.81ms | 1.63ms | 0.21ms | ~90% |
392+
393+
**gemm (N=K=4096, fp16), sweeping M** — fixed dequant floor, GEMM overtakes it near M≈512:
394+
395+
| M | total | dequant | linear | GEMM share |
396+
| ---- | ------ | ------- | ------ | ---------- |
397+
| 8 | 1.41ms | 0.80ms | 0.10ms | 7% |
398+
| 64 | 1.21ms | 0.88ms | 0.42ms | 35% |
399+
| 512 | 2.50ms | 0.73ms | 1.27ms | 51% |
400+
| 2048 | 5.66ms | 0.76ms | 4.85ms | 86% |
401+
402+
**Decision (per-op, as the plan anticipated — now with evidence):**
403+
404+
- **`gemv_4bit` (M=1) → Option B (hand-fused dequant+matmul).** 80–90% dequant-bound; Option A
405+
(`MPSMatrixMultiplication` on materialized `B_dq`) would only touch the ~10% GEMM slice. Fusion —
406+
never writing `B_dq` to device memory — is the entire win. This is Phase M2.
407+
- **`gemm_4bit` large M (≥~512) → Option A (`MPSMatrixMultiplication`).** GEMM dominates; do not try to
408+
out-GEMM Apple's tuned kernel by hand. Accept the fixed ~0.75ms dequant tax. This is Phase M3.
409+
- Small-M `gemm` (≤64) is still dequant-bound and behaves like gemv; a fused path helps there too, but
410+
M3 defaults to Option A for simplicity and lets the fixed dequant floor stand.
411+
412+
**Load-bearing caveat for the kernel author.** The existing dequant kernel moves ~40 MB in ~0.75ms ≈
413+
**54 GB/s**, on hardware that sustains ~400 GB/s — it's leaving ~85% of memory bandwidth on the floor.
414+
Both matmul routes inherit this: a fused `gemv` kernel that reads packed B no faster than the current
415+
dequant will reproduce the 54 GB/s and win ~nothing. **The M2 target is bandwidth, not "fusion" per se**
416+
— the fused kernel must read packed B + absmax at close to peak bandwidth (coalesced loads, minimal
417+
recompute) or it doesn't beat the baseline. (Separately, this implies the standalone dequant kernel is
418+
itself under-optimized — a possible bigger, simpler lever for the QLoRA M=1 inference case — but that's
419+
Phase-3 kernel scope, out of this phase's remit.)

0 commit comments

Comments
 (0)