[alenizy361]

Summary

The fork's fast PQ2_0 CPU paths (gemv/gemm_pq2_0_4x8) are gated on AVX-512, so 12th–14th gen Intel laptops/desktops (AVX2 + AVX-VNNI, AVX-512 fused off) and similar CPUs fall back to the generic scalar dot. This PR adds an AVX2/AVX-VNNI path for PQ2_0 and makes batched work (prompt processing, speculative verification) run at DRAM speed:

• ggml_vec_dot_pq2_0_q8_0 / ggml_vec_dot_pq2_0_q8_K (AVX2, AVX-VNNI, AVX-512VNNI-VL): 2-bit codes unpacked with srlv+and, dpbusd against the permuted activation bytes, unsigned-code correction via a lane-wise sum(y) so per-lane partials stay exact.
• PQ2_0 vec_dot_type → Q8_K (one activation scale per 256): 4 dpbusd + sub + cvt + fmadd per 128 weights instead of 4×(2 dpbusd + sub + cvt + fmadd). Perplexity unchanged within noise (Bonsai-2 27B, wikitext-style text: 7.0542 → 7.0583 ± 0.58).
• llamafile/sgemm.cpp: tinyBLAS_PQ2K_AVX / tinyBLAS_PQ2_AVX tiled GEMMs. A 128-weight block is unpacked once per row tile and reused across RN token columns; the activation permute and lane sums are staged once per matmul call per thread; jobs are row tiles (each weight row streams from DRAM exactly once) claimed dynamically from the threadpool counter in runs of 8 so P- and E-cores share the compute-bound loop. Float accumulation order and reduction are identical to the vec_dot path → batched and single-token results are bit-identical (PQ2_SGEMM=0 gives the same PPL to all digits).

Measurements

i7-13620H (6P+4E, AVX2 + AVX-VNNI, no AVX-512), single-channel DDR5-5200 (~28 GB/s effective), Ternary-Bonsai-2-27B-PQ2_0, --no-repack:

	before	after
decode N=1	1.3 tok/s (generic)	3.9–4.2 tok/s
batch cost N=4 / 8 / 16 / 64 (ms per token)	123 / 99 / 89 / 89	78 / 68 / 54 / 69
prompt processing	6.7 tok/s	16 tok/s
greedy decode with the grafted nextn (MTP) head, --spec-draft-n-max 2	—	6.8–8.3 tok/s (verify of 3 tokens ≈ one DRAM sweep)

Exactness checks: PPL identical with PQ2_SGEMM=0/1; 200-token greedy outputs byte-identical between the vec_dot and GEMM paths; the Q8_K kernel has a scalar reference (ggml_vec_dot_pq2_0_q8_K_generic) and matched 1000 random cases with 0 mismatches.

Side finding while doing this: the published Bonsai-2 PQ2_0 weights are not purely ternary — 0.89% of codes are 3 (value +2, sampled over gate/up/down/attn in 4 layers), so base-3 packings (PTQ1_0) are lossy for this model (PPL 10.0435 vs 10.0685 on the same text).

Notes

• Only x86 with AVX2 takes the new paths; other targets are unchanged. PQ2_SGEMM=0 disables the tiled GEMM for A/B.
• Related: [CPU] Segfault loading Bonsai 2 27B PQ2_0 — null deref in ggml_backend_cpu_repack_buffer_set_tensor #204 (PQ2_0 CPU load crash) — this PR was developed with --no-repack; the repack crash is separate.
• Written with AI assistance (Claude Code) under my direction; every number above was measured on my machine and the kernels were verified against scalar references as described. Happy to split into smaller commits if preferred.

[bri-prism]

Agent review: posted by the maintainer's coding agent at their request.

Request changes on two independent issues:

1. P1 - Supply non-x86 dispatch before changing the global trait (ggml/src/ggml-cpu/ggml-cpu.c:242). PQ2_0 now unconditionally references ggml_vec_dot_pq2_0_q8_K, but that symbol is defined only in arch/x86/quants.c. The generic function has a different name and there is no corresponding alias in arch-fallback.h. A native ARM64 build reproduces the failure: Undefined symbols for architecture arm64: _ggml_vec_dot_pq2_0_q8_K, referenced from type_traits_cpu. Please scope the trait change or provide the missing fallback implementations/aliases.

2. P2 - Preserve valid 128-element PQ2 rows (ggml-cpu.c:243). PQ2_0 blocks contain 128 elements, while the new Q8_K activation format requires multiples of 256. Valid inner dimensions such as 128 or 384 previously worked with Q8_0 but now violate Q8_K row-size/quantizer constraints; the new dot product also asserts n % QK_K == 0. Keep a compatible activation format or support padding/tails throughout allocation, quantization and the dot product.

Build reproduction used GGML_METAL=OFF, GGML_BLAS=OFF, LLAMA_BUILD_TESTS=OFF, LLAMA_BUILD_EXAMPLES=OFF, LLAMA_BUILD_TOOLS=OFF, then built target ggml-cpu. The second finding is source-contract validation; no x86 execution was performed.

Reviewed commit: 164c33700e48de0189a83f2c7b8fed549c8d5dfe.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[alenizy361]

Both findings are fixed in 10df298.

P1 — non-x86 dispatch. ggml_vec_dot_pq2_0_q8_K is only defined in arch/x86/quants.c; the portable ..._generic in quants.c had no alias in arch-fallback.h. Added it to all seven non-x86 sections (generic, arm, powerpc, loongarch, riscv, s390x, wasm). ARM is aliased next to its ptq1_0 entry because it has a native pq2_0_q8_0 and so lacks the usual anchor line.

I couldn't reproduce your exact macOS link error on Linux — GNU ld tolerates undefined symbols in a shared object — so I cross-compiled for aarch64 and checked the symbol table instead:

before:  U ggml_vec_dot_pq2_0_q8_K
after:   T ggml_vec_dot_pq2_0_q8_K  (0x4d504)

That undefined symbol is what your linker reports as Undefined symbols for architecture arm64. x86-64 is unaffected and still resolves to the AVX2/AVX-VNNI implementation.

P2 — 128-column rows. You were right, and the failure mode is worse than an assert. tensor_type_fallback compared ncols against ggml_blck_size(PQ2_0) = 128, so 128- and 384-column tensors passed the shape check and were quantized to PQ2_0 — but ggml_row_size rounds down, so their Q8_K activation row is zero bytes:

ncols   row_size(PQ2_0)   row_size(Q8_K)
128     34                0     <- corrupt
256     68                292
384     102               0     <- corrupt
512     136               584

It now compares against the activation granularity for PQ2_0, so these fall back to q4_0 like the other 256-block types already do. On a synthetic model with all-128-column rows: before, seven tensors silently converted to pq2_0 with no warning; after, each reports ncols 128 not divisible by 256 (required for type pq2_0) -> falling back to q4_0.

I chose the fallback over padding because the 256-element granularity is inherent to Q8_K and shared by every K-quant — padding would mean changing allocation, quantization and the dot together for this one type. Happy to revisit if you'd prefer padding.

Linux allocation/clear validation and model-backed multi-sequence coverage for #207 are still on my side; I'll follow up on that PR.

[Copilot]

Copilot review overview

🟡 Changes recommended

The Q8_K GEMM violates the existing stride assertion, valid 128-granularity models regress, and concurrency/build issues remain.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 3 · 2 · 1

Open (6)

• Guard AVX2 helper with AVX2 · New
• Preserve valid 128-wide PQ2_0 rows with Q8_K · New
• Fix ldb precondition mismatch for PQ2_0 Q8_K calls · New
• Avoid per-thread duplication of the full staging buffer · New
• Make lazy off initialization thread-safe · New
• Correct ggml_row_size failure-mode description · New

[bri-prism]

Follow-up on 10df298: the native ARM Debug build and test-backend-ops link now succeed, confirming the missing-symbol fix.

[P2] Preserve execution of existing PQ2_0 rows that are only multiples of 128. The change in src/llama-quant.cpp prevents newly quantized tensors from taking this path, but the unconditional Q8_K CPU trait still breaks existing valid tensors. I reproduced this with ggml_new_tensor_2d(PQ2_0, K, 4), an F32 activation [K,1], ggml_mul_mat, and ggml_graph_compute_with_ctx: K=128 and K=384 both abort in ggml_row_size's block-size assertion; K=256 completes. The previous Q8_0 path completes both shapes.

Please retain a runtime path for these shapes (e.g. the Q8_0 dot path, or correctly padded Q8_K activation storage and tail handling), and exercise pre-existing PQ2 tensors rather than only the quantizer fallback. The ARM link blocker is resolved; the existing-model shape regression remains.

[alenizy361]

Fixed in d2c9ddd. Thanks for the repro — it pointed at the right layer, and reproducing it turned up something worse than the assert.

Your diagnosis was right and my previous fix was in the wrong place. Guarding tensor_type_fallback only protects tensors that pass through the quantizer; tensors inside an already published model never do. That guard is reverted in this commit — it only cost PQ2_0 coverage on shapes that are now handled properly.

The failure is not always an abort

Building your repro as a test (tests/test-pq2-row-shapes.cpp), on Release x86-64 AVX2 + AVX-VNNI the ggml_row_size assert is not enabled, so instead of aborting it silently returns wrong results and reads past the activation buffer:

  K= 128 M=4 N= 1 : FAIL rel_rms=1.0000
  K= 128 M=4 N= 4 : FAIL rel_rms=1.0000
  K= 128 M=4 N=16 : FAIL rel_rms=1.0000
  K= 256 M=4 N= 1 : OK   rel_rms=0.0041
  K= 256 M=4 N= 4 : OK   rel_rms=0.0053
  K= 256 M=4 N=16 : OK   rel_rms=0.0049
  K= 384 M=4 N= 1 : FAIL rel_rms=0.4793
  K= 384 M=4 N= 4 : FAIL rel_rms=0.9874
  K= 384 M=4 N=16 : FAIL rel_rms=4297627962879295475445409156083744768.0000
  K= 512 M=4 N= 1 : OK   rel_rms=0.0064
  K= 512 M=4 N= 4 : OK   rel_rms=0.0054
  K= 512 M=4 N=16 : OK   rel_rms=0.0068
  K= 640 M=4 N= 1 : FAIL rel_rms=0.5836
  K= 640 M=4 N= 4 : FAIL rel_rms=3337216740497661127397636816569892864.0000
  K= 640 M=4 N=16 : FAIL rel_rms=2716628577838266699126443818700767232.0000

K=640 fails too — it is every row width that is an odd multiple of 128, not just the two shapes in the report. The 1e36 values are the out-of-bounds reads.

The fix

Activation-type selection is now shape-aware instead of a fixed trait-table entry:

static inline bool ggml_cpu_pq2_needs_q8_0(const struct ggml_tensor * src0) {
    return src0->type == GGML_TYPE_PQ2_0 && (src0->ne[0] % QK_K) != 0;
}

ggml_cpu_vec_dot_type / ggml_cpu_vec_dot return Q8_0 + ggml_vec_dot_pq2_0_q8_0 for those rows and the trait-table entry otherwise. They replace all six sites that read the table directly — the two mul_mat chunk functions, the two mul_mat_id ones, and both ggml_graph_plan work-size computations; a mismatch in the last two would size the activation buffer for a different type than the compute path consumes. I took the Q8_0 route rather than padded Q8_K storage: llamafile_sgemm already dispatches on Btype, so these rows land on tinyBLAS_PQ2_AVX and keep an AVX2 kernel instead of dropping to the scalar dot. It is marginally slower than Q8_K would be on those shapes, and correct.

Rows that are multiples of 256 take a code path that is unchanged.

After

15/15 OK, rel_rms <= 0.0068. The K=256 and K=512 numbers are identical to the "before" run to all digits, so the fast path is demonstrably untouched.

Regression checks

check	result
test-backend-ops -o MUL_MAT -b CPU	1294/1294
test-backend-ops -o MUL_MAT_ID -b CPU	1022/1022
test-quantize-fns	pass
llama-bench Bonsai-2-27B-PQ2_0 tg32	3.97 ± 0.09 t/s (PR range 3.9–4.2)

That model's 402 PQ2_0 tensors are all 256-multiples, so its graph provably does not change; the branch is evaluated once per matmul node, not per row.

aarch64 cross-build (gcc-aarch64-linux-gnu) — all defined, none undefined:

T ggml_vec_dot_pq2_0_q8_0
T ggml_vec_dot_pq2_0_q8_0_generic
T ggml_vec_dot_pq2_0_q8_K

The new test is registered in tests/CMakeLists.txt, so ctest covers this shape contract from here on.

[bri-prism]

Re-reviewed at d2c9ddd. Both points are resolved, and the second one is resolved better than what I asked for.

P1, the non-x86 link break. Fixed everywhere rather than at the one site that produced my ARM64 error. The alias is present in all seven non-x86 sections of arch-fallback.h and correctly absent from the x86 section. I checked each section individually.

P2, the 128-element rows. You were right to move the fix and right to revert the quantizer guard. Gating tensor_type_fallback only ever protected tensors passing through the quantizer, which is exactly the set that was not at risk; tensors in an already published model never go through it. Reverting that also gives back the PQ2_0 coverage the guard was costing.

The shape-aware activation selection is the right shape for this. I verified that ggml_cpu_vec_dot_type and ggml_cpu_vec_dot are consumed at all six sites that previously read the trait table directly: the two mul_mat paths at :1221 and :1312, the two mul_mat_id paths at :1531 and :1599, and both ggml_graph_plan work-size computations at :2903 and :2915. The only remaining direct table reads are inside the two wrappers themselves, which is the correct fallthrough. Including the two graph_plan sites matters as much as the compute ones, since sizing the activation buffer for one type while the compute path consumes another is the same class of bug wearing different clothes.

Your Release-mode finding is the more important half of this and I want it on the record: with asserts compiled out there is no abort, just wrong results and reads past the activation buffer, and it covers every row width that is an odd multiple of 128 rather than the two shapes in my repro. My report understated the severity because I only exercised a Debug build. The 1e36 values make that unambiguous.

tests/test-pq2-row-shapes.cpp is wired into the build and checks rel_rms across K in {128, 256, 384, 512, 640} rather than only asserting the absence of a crash, which is what makes it a real regression gate. Taking the Q8_0 route to keep these rows on an AVX2 kernel rather than padding into Q8_K looks like the right trade.

Clearing my change request.

[AlexGabbia]

On the side finding at the end of the PR body — that the published PQ2_0 weights are not purely ternary (0.89 % of codes are 3, so base-3 packings are lossy for this model):

I re-measured it on Ternary-Bonsai-2-27B-PQ2_0.gguf and cannot reproduce it. Full census over all 402 PQ2_0 tensors, streamed: 26,869,760,000 codes counted, exactly the number the file declares, and code 3 = 0. A positive control (synthetic 0xFF bytes -> 100 % code 3) and the tree's own dequantize_row_pq2_0 agree. Dequantising PQ2_0 against the F16 release gives bit-identical weights, max code seen = 2, and quantize_row_ptq1_0_ref -> dequantize_row_ptq1_0 round-trips 1,048,576 real weights with 0 value differences (1.7500 bits/weight).

The quantiser cannot emit a 3 in any case: d = amax, so round(w/d) in {-1,0,1} and q = round(w/d)+1 in {0,1,2}.

Details, controls and the most likely origin of the 0.89 % (counting the 2-byte fp16 scale as if it held codes — it reproduces 0.9043 % over the whole file, 0.7940–1.0495 % per tensor) are in PrismML-Eng/Bonsai-demo#196. Since the code-3 rate is the premise for PTQ1_0 being lossy, and that premise is what makes a 1.3 GB-smaller lossless band look unusable, it seemed worth correcting rather than letting it stand.

The PQ2_0 kernels themselves are unaffected by any of this.