Pull requests / #2533
#2533 metal: q4_0 KV cache on the Mac engine (--kv q4_0, opt-in; the non-q4 path byte-identical when off)
open · @afanty2021 · 0 comentarios · En GitHub
BenchmarksSetup & installAMD / HIPNVIDIA / CUDAModels & quantsDocumentationWindows
Descripción
## What it changes The PC engine's q4_0 K/V format (PR #21's FWHT-256 Hadamard rotation) brought to the Metal backend, opt-in through the engine's existing `--kv q4_0` flag: - `metal/kernels/kv_q4.metal` (new): the FWHT-256 rotation, quantize and dequantize kernels, and the pool layout (576 B per cell per QSA layer against the non-q4 path's 1,056). - `metal/kernels/qsa.metal`, `qsa_prompt.metal`: the decode and prompt chains made format-aware — the five QSA dispatch sites carry function constant 25 in both the pipeline key and the Constants, so a q4 run gets its own specializations instead of silently reusing the cached int8 ones over q4 pools. - `metal/src/engine/`: `kv_format.h` (new, the single source for pool sizes / keys / the format choice), `kv_q4_host.cpp` (new, the host side), `window.cpp` and `prompt_chain.inc` wired: q4 pools at init, the MTP drafter swaps pools in the same format, `--kv k8v4` refused at startup. - Both prompt fast paths (the matrix-unit GEMM chain and the flash chain) run over q4 pools; the M1-era fallback that read q4 prompts through decode windows is gone, `STRATA_METAL_PROMPT=window` still forces it. - `metal/setup_mac.py`: opt-in `--kv q4_0` that lands in the generated engine config's `args` (`tools/test_setup_mac.py` extended). - `tools/quant_quality.py`: opt-in `--kv` on `run`/`prompt` (replaces the harness's fixed fp16 pin only when passed). - `docs/MAC.md`: the option, the quality/speed numbers, the KV-memory numbers, the not-testeds. - `bench/results/2026-10-15-mac-kv-q4/` (new): the whole Mac measurement round — scripts, raw TSVs, distilled run lines, README. ## What it leaves alone - No file under `src/` or `include/` is touched, so no CUDA/HIP/SYCL build is affected (the Metal engine's host code is the PC's, compiled for macOS as before). - With q4 off, every pipeline key, Constants set and binding order on the non-q4 path is today's int8 one, byte for byte (pinned by tests, below). Defaults everywhere: setup writes no `--kv` unless asked, the engine's default is the same non-q4 path as the last release. - Snapshot/restore and cross-format disk sessions: unchanged shared-side behavior (a cross-format session is refused by the shared code, as before). ## Opt-in, and how the default path was checked byte-identical Three independent gates: 1. **Test gates (fc-off byte identity).** `metal/tests/kv_format_test.cpp` pins the five-key reduction (every constant set an engine site now binds is compiled as its own specialization) and a stale-key counterfactual; `qsa_q4_parity.cpp` and `qsa_prompt_parity.cpp` assert the fc-off outputs against double-precision references — with q4 off every key, Constants set and binding order is the current int8 one, byte for byte. 2. **Engine-level int8 regression (F4).** On the M4 Pro, an int8 short-prompt run on this branch produced output byte-identical to the pre-branch baseline run (bd1f3153), zero stub lines — the engine's int8 path does not move. 3. **Harness default-args diff.** `tools/quant_quality.py`'s `engine_args()` output was diffed before vs after the opt-in `--kv` change over the three gpu_experts modes: identical without the flag. For setup, `metal/setup_mac.py`'s `engine_args()` was compared against HEAD over a 48-cell matrix (stream x ram x context x shard shapes; 0 diffs) plus the full generated config on three of those cells: identical without `--kv`, and with it the pair `--kv q4_0` is appended at the end and nothing else moves (pinned by `tools/test_setup_mac.py`). ## Measured (Apple M4 Pro, 20 cores, 48 GB unified memory, macOS 27.0.1; Qwen3.8-Flash-Next Q2_0, experts streamed from the SSD; engine built from this branch; every number same-day 2026-10-15, nothing extrapolated) The framing fact for the reference column: on Metal, `--kv fp16` and `--kv int8` are **one path** — the window's pools are the int8 ones either way (`kv_format.h`; the fp16-flagged run's own log prints `KV int8`), and the two teacher-forced TSVs of the round are byte-identical (md5 `4e942a7375715ac014956fbdbf6110f6`). So the reference below is that shared non-q4 path, q4_0 is the only smaller K/V format on Mac, and no fp16-vs-int8 delta exists on this backend. | Measurement | q4_0 | shared path (int8) | |---|---:|---:| | Teacher-forced top-1 (4,995 positions, 5 segments) | 94.45% | 100.00% | | Teacher-forced mean KL / p99 | 0.03414 / 0.4467 | 0.00000 / 0.0000 | | Perplexity (reference 4.442) | 4.517 (+1.7%) | 4.442 | | Worst segment (agent transcript, PPL) | +7.3% | — | | Prompt fast path, 20 full-vocab prompts: top-1 | 100.00% | reference | | Prompt fast path KL mean / median (upper-middle) / max | 0.03565 / 0.00434 / 0.19546 | reference | | Decode tok/s (4-id prompt, 64 new) | 20.84 | 21.09 (−1.2%) | | Prefill tok/s (2,199 tokens prefilled of a 2,200-id prompt) | 213.15 | 215.54 (−1.1%) | | KV B per cell per QSA layer (12 layers) | 576 | 1,056 | | KV bytes at 32K context (32,896 cells, calc) | 227,377,152 | 416,858,112 | | KV bytes at 64K context (65,664 cells, calc) | 453,869,568 | 832,094,208 | | Log's before-the-arena delta (32K / 64K) | 0.19 / 0.38 GB less | — | | Arena resident expert slots (32K / 64K) | 16,081 / 15,657 | 15,840 / 15,163 (+241 / +494) | | Needles (8K/16K/32K at 10/50/90% depth) | 8/9 | 3/3 (controls at 50%) | - The speed deltas (−1.2% decode, −1.1% prefill) sit at the edge of run-to-run noise: the same-code fp16-vs-int8 pair (one path on Metal) measured −1.3% decode, bounding noise near 1%. Single runs per cell. - The needle miss is 32K at 90% depth: the model took the needle sentence for a prompt injection and declined to repeat it (a behavior miss, not a garbled retrieval; all 8 retrievals were the exact codeword). Disclosures, as recorded in the bench README: the question ends with "The access code is" — a completion cue all arms share, so the comparison is internally fair but the 8/9 is under an easier phrasing than a neutral QA form; the miss has no int8 control at 90% depth (controls ran at 50%), so a shared-path refusal there cannot be excluded; and its reply was cut mid-sentence at the 64-token budget — the injection-suspicion reasoning is explicit, a final refusal sentence is not. ## Disclosures (from the final segment review) - **Arena stand-by slots:** under q4_0 the prompt fast path is on in every format, and the streaming arena plans 576 stand-by slots for the prompt-GEMM chain (64 without it). Escape hatches: `STRATA_METAL_PROMPT_GEMM=0` (prompts off the matrix-unit GEMM chain; the arena is then planned without the stand-by slots) and `STRATA_METAL_PROMPT=window` (prompts through decode windows — the slow, exact fallback). - **What the measured fast path is on this host:** the M4 Pro is !k4 (`coop_k4=false`), so the prompt fast path behind every prompt-path and prefill number above is the matrix-unit GEMM chain over the v1 layout — the flash chain is off by default on !k4 hardware and never dispatched here (same hardware fact as the k4 arm below). - **Prompt-path numeric divergence class:** the prompt GEMM chain dequantizes its operands to fp16 once per chunk (the matrix units), so its numerics are a different floating-point order from the decode window's (fp32 fused) — outputs can differ on near-ties, the same class as any reordering. Measured under q4_0 on the 20-prompt set: argmax never flipped (top-1 100%), full-vocab KL as in the table. - **M5 k4-layout q4 numerics:** the fc-7 right-input layouts are compiled as real pipelines in the tests, but this M4 Pro is !k4 hardware and never dispatches them (see Not tested). - `--kv k8v4`: the only statement is the explicit startup refusal — "metal window: --kv k8v4 is not implemented on Metal (use fp16, int8 or q4_0)". No partial implementation exists. - `STRATA_KV_ROT=1` (the PC's rotated int8 K/V env): refused at startup next to `--kv int8`, like k8v4 — there the PC rotates the pools and the shared conversation-snapshot tag adds +16 for the rotation, so a Metal session would save a tag claiming a rotation its pools do not have (the env is parsed in kv_format_check as generate.cpp does; the int8 flag comes through the existing qsa_kv_int8 getter). Next to fp16 (and no `--kv`) the env is a no-op on both engines — the rotation is gated on the int8 flag (layer.cpp: `kv_rot = kv_q4 || (kv_int8 && env)`, and generate.cpp sets kv_int8 only for `--kv int8`), so nothing rotates and the tag adds nothing; allowed, matching the PC's silent no-op. Next to `--kv q4_0` the rotation is intrinsic and the tag never adds it, so the env there is a harmless no-op too. ## Not tested - The M5 k4-layout q4 numeric arm (fc-7 right-input layouts): compiled but never dispatched on this !k4 host. - MTP + q4_0 end to end: every measured run used the template's `--spec 2 --suffix-draft 0`; the drafter's same-format pool swap is structurally validated by tests, not exercised end to end. - PC consistency: no PC run this round; the 2026-09-27 PC table (`bench/results/2026-09-27-kv-q4/`) is a different backend and a `--kv-resident` streaming configuration, so its numbers are not comparable and are not cited. - Speed repeats (one run per cell; the same-code pair bounds noise near 1%) and needle breadth (one codeword, one filler corpus, one run per case, no distractors, int8 controls at 50% depth only). - Other Macs: every number above is from one M4 Pro, 48 GB. - Long-context quality: every quality segment is 1,000 tokens. The PC round measured the q4_0 perplexity cost growing with context (+8% at 1K to +12% at 8K on a document), so the +1.7% here is a short-context number and must not be read as the format's cost at 32K. Known follow-ups, none in this PR: the flash q4 epilogue is barrier-dense (a perf item only k4 hardware would feel, and this !k4 host never dispatched it), the flash `vs` dead array is still to be removed, prompt_chain's pre-existing argv bounds pattern and kv_format_test's pre-existing non-q4 flash constant sets predate the branch and stay as found, and the Task-1-era nits stand (the `pipe()` per-call cache lookup, the identity-page gather case, the offsetof asserts). ## Reproducing `bench/results/2026-10-15-mac-kv-q4/README.md` is the entry point: the exact commands for the teacher-forced round (`tools/quant_quality.py run/compare`), the prompt-path round (`prompt`/`pcompare`), the speed pairs (`speed.sh`), the KV/arena numbers (`kvmem.sh`) and the needle matrix (`needle.py`), all with the committed raw TSVs and distilled run lines beside them (`data/speed-runs.txt`, `data/kvmem-runs.txt`, `needle-results.tsv`, the stderr summaries); the full engine logs stay out per the repo's `logs/` ignore and regenerate from the committed scripts. The in-tree test gates: `ninja -C build-metal` then the binaries in `metal/tests/` (`kv_q4_parity`, `kv_q4_engine`, `kv_format_test`, `qsa_q4_parity`, `qsa_prompt_parity`). Setup's default-identity is `python3 -m unittest tools.test_setup_mac` (no GPU, no downloads). note: `data/prompt-path/pcompare.txt`'s reference line was edited by hand to point at the committed per-prompt.tsv (the regeneration commands are in the README); the aggregates were never hand-touched. ## Prior art searched (issues and PRs) Re-checked read-only at packaging time (`gh search issues`/`gh pr view` against the upstream repo, terms: metal kv, kv q4, k8v4, mac backend, apple silicon, metal PRs): **no issue or PR covers KV quantization on the Mac/Metal backend** — this PR opens that ground, and the PC lineage below is the prior art. - PR #21 (merged): the PC q4_0 KV mode this branch ports to Metal. PR #120 (k8v4, closed) and #293 (int8 rotation, closed) are the format's siblings; PR #452/#904 are the PC prompt/append follow-ups. - Issue #2403 (open): `--kv q4_0` faults on a PC (HIP) when the KV itself streams (`--kv-resident`) with MTP drafts. Different code path — the Metal engine has no `--k
En el sitio
Enlaces a install, modelos, releases.