Spaces:

PYTHAI
/

bankml

Running

bankML: the whole source (github.com/cryptoAGI/bankml @ 12ae409) and its page, with the bankML persona; the live engine (Dockerfile, hf/start.sh) ready for Docker hardware

28c70af verified 6 days ago

preview code

Raw History Blame Contribute Delete

9.38 kB

`bankML/q2_0.rs` — the ternary Q2_0_g64 kernel

Summary

q2_0.rs multiplies Q2_0 weights (ggml type 42, group 64) exactly as mainline ggml does at b11192, read from the tag's source. A block is 18 bytes: an f16 scale d first, then 16 bytes of 2-bit codes, least significant bits first. Code c ∈ {0,1,2,3} decodes to (c − 1)·d, that is {−1, 0, +1, +2}: 2.25 bits per weight. The Bonsai files are ternary (in Ternary-Bonsai-8B the split is 31 % / 38 % / 31 % / 0 %), but the format is not, so every kernel here accepts code 3.

As with Q1_0, the activation is quantized to q8_0 first (the same AVX2 quantizer, q1_0::quantize_row_q8_0), and one Q2_0 block meets two q8_0 blocks.

On x86, llama.cpp b11192 has no Q2_0 kernel. arch-fallback.h renames the generic C to ggml_vec_dot_q2_0_q8_0; there is no repack or sgemm case. The shipped haswell symbol is scalar code with 64 imul per block, used one (row, column) pair at a time for decode and prefill. bankml's AVX2 kernel gives the same bits much faster. The native forward pass (forward.rs) calls mat_vec_par for decode and mat_mul_par for prompt micro-batches.

Technical usage

Constants: QK2_0 = 64, Q2_0_BYTES = 18.

pub fn dequantize_row(blocks: &[u8], out: &mut [f32])         // bit-exact with dequantize_row_q2_0
pub fn vec_dot_ref(n: usize, x: &[u8], y: &[u8]) -> f32        // scalar model of the shipped haswell symbol (FMA)
pub fn vec_dot_ref_nofma(n: usize, x: &[u8], y: &[u8]) -> f32  // the same C as the baseline libggml-cpu-x64.so (no FMA)

vec_dot_ref is the reference for the node's kernel. vec_dot_ref_nofma is not used for inference; it is kept so that both shipped builds of ggml are proven.

Prepared activations

pub struct Q8Act2 { /* private */ }
impl Q8Act2 {
    pub fn from_q8_0(n: usize, q8: &[u8]) -> Self
    pub fn quantize(x: &[f32]) -> Self
    pub fn n(&self) -> usize
    pub fn qs(&self) -> &[i8]
}
pub fn vec_dot_act(x: &[u8], a: &Q8Act2) -> f32
pub fn vec_dot_act_scalar(x: &[u8], a: &Q8Act2) -> f32
pub unsafe fn vec_dot_act_avx2(x: &[u8], a: &Q8Act2) -> f32

Q8Act2 lays out a q8_0 activation once per token so that no shuffle touches the weights. Blocks are grouped in quads (A, B, C, D); a 256-bit register holds the code bytes of A and B, and (v >> 2r) & 3 gives code r of every byte. The activation stores, per pair and per r, the matching quants. Scales (f32) and Σq per q8 block are stored in the lane order hadd leaves. Tail blocks (nb % 4) keep a plain layout. Storage is 64-byte aligned.

vec_dot_act picks the AVX2 kernel when q1_0::has_avx2(), else vec_dot_act_scalar. Both have the bits of vec_dot_ref on the q8_0 bytes the activation came from.

Matrix products

pub fn mat_vec(w: &[u8], rows: usize, a: &Q8Act2, out: &mut [f32])
pub fn mat_vec_par(pool: &crate::par::Pool, w: &[u8], rows: usize, a: &Q8Act2, out: &mut [f32])
pub fn mat_mul(w: &[u8], rows: usize, cols: &[Q8Act2], out: &mut [f32])
pub fn mat_mul_par(pool: &crate::par::Pool, w: &[u8], rows: usize, cols: &[Q8Act2], out: &mut [f32])
pub unsafe fn vec_dot_1x4_avx2(x: &[u8], a: [&Q8Act2; 4]) -> [f32; 4]

mat_mul* writes out[c·rows + r], using 1×4 tiles (one row, four columns, codes expanded once) where the CPU has AVX2. Every element has the bits of ggml's per-pair vec_dot, at any thread count.

How it is verified

Unit tests: dequantize_layout_scale_first_2bit_lsb_first (all four codes), ref_matches_dequantized_math (against f64), act_paths_bit_exact_with_ref_randomized (3,300 cases: 4-block bodies and tails, q = −128, code 3, and activations through the quantizer), par_bit_exact_with_single_thread, mat_vec_and_mat_mul_bit_exact_with_ref.
The ggml oracle oracle_ggml_b11192_real_ternary_bonsai_8b (#[ignore], needs the model and .models/oracle-ternary from testing/ggml_oracle.py), weights read through gguf::Mmap. Recorded (oracles.md §1, PERFORMANCE.md): all 254 Q2_0 tensors, 8,188,239,872 weights, dequantized bit-exact; 762/762 q8_0 rows byte-exact; 762/762 dot products bit-exact against the haswell build (scalar model, AVX2 and portable paths); the no-FMA model equals the baseline x64 build in 762/762. The two ggml builds differ from each other in 19 of 762, so the oracle does distinguish float orders.
A/B against the live library (BANKML_GGML_LIB): ab_vs_ggml_q2_0 (real layer-0 weights, every Qwen3-8B projection shape and output) and decode_budget_q2_0 (all 253 matmuls of one token, BANKML_THREADS, default 1,3). Both check agreement before timing.

BANKML_GGML_LIB=/path/to/llama-b11192 cargo test --release -- --ignored ab_vs_ggml_q2_0 --nocapture --test-threads=1

Advantages and efficiency

A missing kernel, supplied with the same bits. The ternary gap in llama.cpp b11192 is not a property of ternary weights but of the scalar fallback. bankml keeps ggml's float chain (sumi = fma(d1₁, s₁, fma(d1₀, s₀, 0)), sumf = fma(d0, sumi, sumf), block after block) and reorders only the integer sums, which are exact in any order.
The activation moves, not the weights. Weights are read straight from the file layout. Two blocks share one 256-bit register; codes come out with a shift and a mask, and maddubs(c, q) − Σq gives Σ(c−1)·q with unsigned codes. The activation is permuted once per token (11 µs at n = 4096, 29 µs at n = 12288, about 0.5 % of the matmul time).
Measured (laptop, Ryzen 3 3200U, PERFORMANCE.md): decode GEMV 9.1–10.6× ggml over four runs (output 151669×4096: 492 → 51 ms); prefill 1×4 tile 12.3–13.1×. One whole token's 253 matmuls with the model resident: 0.231 s at three threads against ggml's 2.185 s (9.45×), 0.435 s against 4.145 s at one thread. At three threads bankml's ternary matmuls take less time than ggml's 1-bit ones (0.23–0.25 s against 0.34–0.35 s), although the ternary weights are twice the bytes.
Where it sits. About 2× above the measured memory floor (15–17 GB/s on the laptop); compute-bound per core and at its instruction-throughput limit on Zen+. Three further experiments (repacked quads, a two-row tile, software prefetch) were bit-exact and not reliably faster, and are kept in testing/experiments/ (0.0.5).
Upstream. The same plain-AVX2 approach was prepared as a C drop-in for llama.cpp (upstream/): bit-exact on 200,000 of 200,000 cases and 3.4× ggml's shipped scalar kernel (0.2.0).
Practice. Run-time feature detection; #[target_feature(enable = "avx2,fma,f16c")] functions with documented safety contracts behind safe wrappers; Q8Act2 fields private since 0.0.2 because the kernels trust n; aligned Line storage so no 32-byte activation load straddles a cache line (not measured to matter; kept because it cannot cost).

Design notes

Source and binary. The format and kernels were read from the b11192 tag's source on 2026-09-26. The shipped haswell symbol was read from the disassembly: the generic C, unvectorised (64 imul per block), with GCC's default FP contraction turning sumi += d1·s (twice per block) and sumf += d0·sumi into vfmadd231ss. Decode and prefill both run through it one (row, column) pair at a time. This scalar fallback is where bankml's ternary speed advantage comes from (an earlier source note put it at 5×; the current measurements are under Advantages).
Two shipped builds. vec_dot_ref models the haswell build (FMA). vec_dot_ref_nofma models the baseline libggml-cpu-x64.so (SSE only, no FMA), where each multiply and add rounds separately.
No saturation. In dot2, maddubs(c, q) with c ≤ 3 and q ≥ −128 stays within ±768 per i16 lane, and the sum of the four code planes within ±4·768, so the integer sums are exact.
Q8Act2 fields. 0.0.1 made n, d and sum public, so safe code could send the AVX2 kernels past their buffers; they are private since 0.0.2 and set only by from_q8_0/quantize.
Line alignment. The 64-byte alignment was not measured to matter on the dev box, whose run-to-run swings are the boost clock (both kernels move together); it is kept because it cannot cost.

Limitations

x86_64 AVX2 is the only SIMD path; elsewhere the portable scalar path runs with the same bits. NEON, AVX-512 / VNNI and a Zen 3 row are open (TODO.md).
The oracle uses real weights with synthetic activations; bit-exactness on dumped real activations is open.
The whole-token figures need the 2.31 GB file resident in the page cache. When memory is short, every token re-reads weights from disk and the gates measured 4.5–5.2 s per token for both runtimes (PERFORMANCE.md).
Legacy group-128 Q2_0 files, PrismML's PQ2_0 and Bonsai 2's rotated basis are refused by the guard (gguf.md); there is no kernel for them. The ternary GPU kernel is open (TODO 0.5.0).

bankML/q2_0.rs — the ternary Q2_0_g64 kernel