tzhouamandClaude Opus 5
b36a3431bd
Turns a branch of accumulated edge/CPU work into something someone else can
read, reproduce and continue.
README.md states what the fork is for (multimodal models on edge CPUs rather
than servers), the headline memory and speed results with their costs, and why
the numbers are worth reading -- five retracted conclusions, all measurement
errors rather than coding errors. Upstream's README is preserved at
docs/README.vllm-omni.md and NOTICE records the Apache-2.0 derivation and the
files this fork changes.
docs/edge/:
methodology.md the measurement discipline, which is the most transferable
thing here: what each of the five retractions was, and the
rule it left behind. Includes the two traps specific to this
stack -- vLLM's AOT compile cache keying on config so two
arms sharing an artifact load the wrong graph, and allocated
not being resident until mincore(2) says so.
results.md every measured figure with its spread and baseline, and an
explicit note that the four memory levers overlap and must
not be added.
cpu-4bit.md the two 4-bit paths, the packed layout as measured rather
than assumed, why the path choice is a config field, and the
GPTQ output-head trap that invalidated a whole round of W4A8
numbers.
harness.md how to run a comparison, and the two arguments that matter
more than they look (--context-len, --prompt-len), each of
which has produced a wrong conclusion by being left at its
default.
ROADMAP.md orders by whether the verification method is correct rather than by
expected size, lists what is done with results, what is next with expected
gains and how each would be verified, and -- at equal length -- what was
considered and rejected, including the CPU megakernel with the measurements
that rule it out for now.
benchmarks/edge_harness/ brings the tooling into the repo so the numbers are
reproducible: the grid runner, the single-run benchmarks, the memory ledger,
the fidelity metric, the Android agent suite, the quantizers, and the tests
that pin the harness's judgement. 82 tests pass from the new location.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L1ax5wYLA6XbqEL1i8L9gZ
101 lines
4.4 KiB
Markdown
101 lines
4.4 KiB
Markdown
# The 4-bit CPU path
|
||
|
||
Two 4-bit paths exist on an x86 CPU and they are not interchangeable. Which one
|
||
a device should use is a property of the device, and `vllm_omni/edge/weight_path.py`
|
||
picks it from a hardware probe.
|
||
|
||
## W4A8 — `_C::int4_scaled_mm_cpu`
|
||
|
||
Needs AMX tiles. Wants GPTQ-packed weights and quantizes activations to int8.
|
||
Fastest measured here (74.4 tok/s) and least faithful (1/12 greedy prompts
|
||
reproduced exactly).
|
||
|
||
One trap: a GPTQ checkpoint **must** quantize the output head. Spark ties its
|
||
embedding to its head, so the head has to be untied and emitted separately; if
|
||
it is not, that one projection runs in bf16 through `onednn_mm` and costs
|
||
2.97 ms/step by itself — more than the entire 4-bit saving. Every W4A8 number in
|
||
this project's first round was invalid for exactly this reason.
|
||
|
||
The symmetric grid is **our exporter's choice, not the kernel's limit**:
|
||
`_process_gptq_weights_w4a8` uses the checkpoint's zero points whenever
|
||
`config.zero_points` is set and synthesises the constant 8 only otherwise. An
|
||
asymmetric export is a config change and is the cheapest untried fidelity
|
||
experiment.
|
||
|
||
## W4A16 — PyTorch tinygemm
|
||
|
||
No AMX requirement, activations stay bf16, more faithful (2/12). 71.6 tok/s.
|
||
This is `quantization="cpu_int4"`.
|
||
|
||
The kernel dequantizes once per row of activations, so its cost is linear in
|
||
batch size: fine for decode, ruinous for prefill (453 tok/s if used for
|
||
everything). Large batches therefore take a second path — dequantize to bf16
|
||
once and hand it to the ordinary oneDNN/AMX GEMM.
|
||
|
||
### The weight layout, measured
|
||
|
||
That prefill path used to read a row-major *duplicate* of every weight, because
|
||
the tinygemm layout was treated as opaque. It is not. Packing an index pattern
|
||
and reading back where each nibble landed shows, on AVX-512:
|
||
|
||
```
|
||
[N/64][K][32 bytes]
|
||
```
|
||
|
||
K is the middle axis, so a K range is contiguous inside each n-block; and within
|
||
every 32-byte run, byte `d` carries output channel `d` in its low nibble and
|
||
`d + 32` in its high one — the same permutation for every `k` and every block.
|
||
Decoding it is a deinterleave, not a reverse-engineered ISA table.
|
||
|
||
So prefill can dequantize from the kernel's own layout and the duplicate is
|
||
never allocated: **615 MB**, bit-exact against the old path at 2048×2048,
|
||
3072×2048 and `gate_up`'s real 13312×2048.
|
||
|
||
It is not free. `to_split_halves` performed that deinterleave *once at load*, so
|
||
each prefill call read two contiguous planes; doing it per call costs **1.63×
|
||
prefill**. It is a dial, not a win:
|
||
|
||
- against `prefill_dequant: false` it is **strictly better** — same memory, 2.97×
|
||
the prefill, and multi-row batches never reach the tinygemm kernel, avoiding
|
||
the AMX tile-configuration SIGILL documented in `cpu_int4.py`. That setting is
|
||
dominated.
|
||
- against keeping both layouts it trades 615 MB for ~1.6× prefill. Default on,
|
||
because this path exists for devices chosen for their memory ceiling. Set
|
||
`direct_dequant: false` in the checkpoint's `quantization_config` to revert.
|
||
|
||
`supports_direct_dequant` gates the reshape on N tiling the 64-channel block,
|
||
which every Spark linear does. An N like 80 packs as 64 + a 16 tail the reshape
|
||
would mis-read, so those layers keep the duplicate rather than being quietly
|
||
wrong.
|
||
|
||
### Why the path choice is a config field
|
||
|
||
The two prefill paths build *different graphs* — one passes `weight_halves` into
|
||
the custom op, the other passes `None`. vLLM keys its AOT compile cache on the
|
||
config, so an env-var-only switch shares one cached artifact between them, and a
|
||
warm cache loads the wrong graph: three benchmark passes died with
|
||
`KeyError: 'weight_halves'` inside an AOT-loaded forward. `direct_dequant` is a
|
||
`quantization_config` field because `ModelConfig.compute_hash` hashes those.
|
||
|
||
## `cpu_gemm_wna16` — the fallback
|
||
|
||
vLLM selects this when `use_w4a8` is off. Measured at 32.3 tok/s, **2.2× slower
|
||
than tinygemm**, so the selector does not route to it. It is documented because
|
||
`weight_path.py` previously asserted it did not exist.
|
||
|
||
## Building a checkpoint
|
||
|
||
```bash
|
||
# W4A16 (cpu_int4)
|
||
python benchmarks/edge_harness/quantize_int4.py \
|
||
--model <bf16-model> --group-size 128 --out models/<name>
|
||
|
||
# W4A8 (GPTQ) — note the head must be untied and quantized
|
||
python benchmarks/edge_harness/quantize_gptq_int4.py \
|
||
--model <bf16-model> --group-size 128 --out models/<name>
|
||
```
|
||
|
||
`quantization_config` keys the `cpu_int4` path understands: `group_size`,
|
||
`quantize_embedding`, `prefill_dequant`, `direct_dequant`, `dequant_threshold`,
|
||
`use_custom_op`.
|