改动 / #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 评论 · 去 GitHub 看

BenchmarksSetup & installAMD / HIPNVIDIA / CUDAModels & quantsDocumentationWindows

说明

## 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

本站相关内容

相关页面的快捷入口。