Proposal: opt-in HQQ 5-bit AXIS=0 MoE weights — −25% bytes, +36% E2E on disk-bound hardware, quality within +2.4% ppl of int8

[outtodata]

Proposal: opt-in HQQ 5-bit AXIS=0 MoE weights — −25% bytes, +36% E2E on disk-bound hardware, quality within +2.4% ppl of int8

TL;DR. A converter-side, opt-in 5-bit weight format for MoE expert FFNs (attention stays int8), using HQQ AXIS=0 grouping so scales/zeros are per input channel. On a Core i5-7300U with weights streamed from disk: snapshot 7.0 GB → 5.4 GB (−25%), end-to-end 2.42 → 3.28 tok/s (+36%, medians of 3; up to 2.46× in more disk-bound conditions), HF-corpus perplexity 14.2420 vs 13.9144 int8 control (+2.36%), greedy token streams bit-identical to the int8 engine in all 8 test runs. Everything sits behind bits==5 in the converter + a snapshot field; default builds and int8 snapshots are untouched. Related: the axis finding in [issue draft: HQQ axis asymmetry].

Motivation. On disk-bound CPUs the expert pipe, not matmul, is the wall (profiled: pipe_wait ≈ 45% of decode time at Q8_0). The cheapest FLOP is the one you never load: cutting weight bytes 25% cuts pipe traffic 25%. But 5-bit only works if the grouping axis is chosen correctly — HQQ's AXIS=1 mixes 32 heterogeneous input channels under one scale and costs +9% ppl (separate issue); AXIS=0 keeps scales per input channel and lands at +2.4%.

Design (mirrors the existing i4 pattern)

• Converter (convert_olmoe_hqq5.py): explicit --bits 5 choice at conversion time, per the repo's "quantization choice lives in the converter" policy. Writes merged_weight5_ax0 tensors + a snapshot header field. Weight layout is the natural flat order of the HQQ q tensor (no repacking needed — AXIS=0 groups are strided slabs, so the packed stream is identical in form to the axis=1 variant: 20 B per 32 weights = 16 B nibbles + 4 B hi-plane). Meta (scales|zeros) is contiguous per row, one 64 KB slab per 32 rows.
• Kernel (dotrow_i5ax0_avx2 / matmul_i5ax0_avx2 in c/fused_simd.h): float FMA with per-element dequant from a hoisted per-group scale/zero. Microbench: 0.145 ms/row-equivalent vs 0.137 for int8 IDOT — parity (±10% run-to-run noise).
• Dispatch (olmoe.c): on snapshot load, presence of merged_weight5_ax0 sets m->i5=2; runtime prints a MODE I5AX0 ... banner; missing AVX2 → hard refusal; int8 snapshots → today's path, untouched. Fallback to int8 tensors if both present.

Measured (i5-7300U, production flags; medians of 3 E2E reps)

Metric	int8 (i5=0)	i5ax0	Δ
Snapshot size	7.0 GB	5.4 GB	−25%
E2E tok/s	2.42	3.28	+36% (up to 2.46× when more disk-bound)
pipe_wait (decode)	10.55 s	5.77 s	−45%
pipe_io / sync_load	—	—	−34% / −35%
time in matmul	—	—	+14% (float vs IDOT — the trade)
HF-corpus ppl (harness-verified)	13.9144	14.2420	+2.36%
Engine TF-NLL on same corpus	7.7141	7.5187	−2.5% (consistent)
task_acc	0.8889	0.8889	0
Greedy stream vs int8	ref	bit-identical, 8/8 runs	—

How this answers the review gates from #906

• (a) off by default, loud when on: nothing changes unless the converter is run with --bits 5 and the resulting snapshot is loaded. Loading such a snapshot prints a MODE I5AX0 banner naming the format, axis, and group size; per-layer [PROF] counters show which layers took the i5 path.
• (b) explicit conversion-time choice: the int8 path is byte-for-byte untouched; 5-bit is a snapshot property, not a runtime flag. A snapshot declares what it is; the engine never silently approximates.
• (c) token-exact when off: int8 snapshots produce bit-identical streams to master (verified by md5 across runs). Determinism within the 5-bit path also holds (8/8 identical greedy streams).

Caveats (honest list)

• Quality cost is real: +2.36% ppl on the HF corpus. That's the price of −25% bytes; the axis=1 alternative (+9%) is strictly worse. Per the warmup-thread policy this belongs to the user at conversion time, and the proposal keeps it there.
• Float-FMA dequant kernel is ~14% slower per FLOP than IDOT; net win exists only when disk-bound (which is the regime this repo targets). On NVMe the gain shrinks.
• AXIS=0 quality depends on the flattening convention of the source checkpoint: per-input-channel scales are only meaningful if axis 0 of the HQQ tensor is the output dimension (true for these safetensors shards; the converter asserts it).
• One machine tested (i5-7300U). No NEON kernel yet.

Happy to open this as a PR (converter + kernels + dispatch + harness, ~1k lines) or to split it (converter RFC first, kernels second). Also happy to run any eval corpus you point us at before anyone spends review time.

[JustVugg]

This is the right shape for a quantization proposal in this repo: converter-side, behind bits==5, default builds unchanged, and measured rather than asserted. −25% bytes for +2.36% ppl is an interesting trade on disk-bound hardware, and the +36% E2E follows directly from "fewer bytes to move" — which is the lever that actually matters when the disk is the bottleneck.

#1008 (the AXIS finding) is a good foundation for choosing AXIS=0 here, and I've replied there.

One claim I'd like tightened before this lands, because exactness is a promise this project makes:

greedy token streams bit-identical to the int8 engine in all 8 test runs

A 5-bit quantization that shifts perplexity by +2.36% is, by construction, not numerically identical to int8 — the weights are different. What you measured is that 8 greedy runs happened not to diverge, which is a much weaker (and expected) statement: greedy argmax is robust to small logit perturbations until it isn't, and divergence probability grows with sequence length and with how close the top-2 logits sit.

I'd rather see it stated as "no token divergence observed across N runs / M tokens" with N and M given, than as "bit-identical". Same note I left on #1024 — not pedantry: "bit-identical" is a claim reviewers will test, and it's the one thing in the pitch we can't afford to have wrong.

Questions for the implementation:

1. Container/format versioning — does a 5-bit snapshot carry a format field that older engines will refuse rather than misread? Silent misinterpretation of a nibble layout is the worst failure mode here.
2. Which engines consume it? Converter-side is fine, but each engine's expert loader has its own unpack path; a 5-bit container that only OLMoE reads should say so.
3. Dequant cost — 5-bit unpacking is not nibble-aligned. What does it cost per expert on the CPU path, and does the byte saving still win when the model is warm in page cache (i.e. when you're compute-bound, not disk-bound)?
4. Quality gate — is there a fixture we can put in CI, like inkling-oracle / deepseek-v4-tiny-check, so the format can't silently regress later?

Direction is good and the measurements are credible. Tighten the exactness claim and answer (1) and (2), and I think this is worth doing.

[outtodata]

Agreed on the exactness claim — tightening it: no token divergence observed across 8 independent greedy runs (full generations, pairwise md5 against the int8 engine). You're right that this is a statement about 8 runs, not about the format — the weights differ by construction, and a +2.36% ppl shift will eventually move a token where top-2 logits sit close. PR text will carry N=8 and the token counts, not the word "bit-identical". Same fix applied to anything we add to #1024's thread.

1. Format versioning: the snapshot header gets a format/version field; an engine that doesn't know bits==5 refuses at load with an explicit "snapshot requires engine >= X" error rather than falling through to the int8 unpacker. The prototype already hard-refuses on missing AVX2 — same posture for unknown format ids. Silent misread is the worst failure mode, agreed.
2. Engines: OLMoE loader only, and the snapshot declares it. GLM/Kimi-arch loaders get no 5-bit path until their flattening conventions are verified — the HQQ grouping axis is a functional property, not just a size trade-off: AXIS=1 costs +9% ppl where AXIS=0 costs +2.4%, at equal weight-MSE #1008 caveat (axis 0 must be the output dim) is per-checkpoint, and the converter asserts it for OLMoE only.
3. Dequant cost when warm: float-FMA dequant is at parity in microbench (0.145 vs 0.137 ms/row-equiv, within run-to-run noise) but shows +14% time-in-matmul in the E2E profile — so on a warm, compute-bound machine the byte saving buys nothing and i5ax0 is the wrong choice. That's a "use when disk-bound" property; we'll state it in the snapshot docs, not hide it.
4. CI fixture: yes — a tiny synthetic expert set (a few MB) converted with --bits 5, then loaded and checked for: refusal by old format id, deterministic golden-token output, and ppl within a bound vs the same fixture at int8. Patterned after inkling-oracle. It goes in the PR.

On sequencing, your call: we can open it split — converter RFC first (format + fixture + refusal semantics), kernels second — which keeps the risky-nibble-layout discussion separate from the SIMD review.

[JustVugg]

Closing for inactivity. The opt-in 5-bit HQQ AXIS=0 converter, with refusal on unknown formats and a tiny fixture, is still welcome as a PR.