From 91665844b7f7dc5805afbf987b36ff72d0551239 Mon Sep 17 00:00:00 2001 From: germar Date: Sun, 5 Jul 2026 09:42:46 +0200 Subject: [PATCH 01/10] feat(kv): opt-in Hadamard (FWHT) rotation for INT8 KV cache (#61) Rotate Q/K/V by an orthonormal Walsh-Hadamard transform before INT8 quantization to suppress activation outliers, per llama.cpp#21038/QuaRot. The transform is self-inverse (one H serves forward Q/K/V and backward output rotation) and dot-preserving (Q and K rotated by the same H leave QK scores intact). Gated behind GEIST_KV_ROT=1; INT8 path only; head_dim a power of two. FWHT is O(n log n), no n*n matmul. Adds fwht.{c,h} + a unit test covering the three properties the trick relies on. --- mk/common.mk | 1 + src/archs/transformer/arch_state.c | 5 ++ src/archs/transformer/arch_state.h | 4 + src/archs/transformer/forward/kv_store.c | 41 +++++++-- src/backends/common/fwht.c | 25 ++++++ src/backends/common/fwht.h | 29 ++++++ tests/test_fwht_unit.c | 110 +++++++++++++++++++++++ 7 files changed, 209 insertions(+), 6 deletions(-) create mode 100644 src/backends/common/fwht.c create mode 100644 src/backends/common/fwht.h create mode 100644 tests/test_fwht_unit.c diff --git a/mk/common.mk b/mk/common.mk index 0c7ff00..e11c60a 100644 --- a/mk/common.mk +++ b/mk/common.mk @@ -165,6 +165,7 @@ LIB_SOURCES := \ src/backends/common/geist_gemm.c \ src/backends/common/gemma4_kernels.c \ src/backends/common/kivi.c \ + src/backends/common/fwht.c \ src/formats/ptqtp/gguf_ptqtp.c \ src/formats/ptqtp/ptqtp_kernel.c \ src/formats/ptqtp/ptqtp_awq.c \ diff --git a/src/archs/transformer/arch_state.c b/src/archs/transformer/arch_state.c index 8a88e26..4a6ce49 100644 --- a/src/archs/transformer/arch_state.c +++ b/src/archs/transformer/arch_state.c @@ -998,6 +998,11 @@ struct transformer_arch_session *transformer_session_alloc(struct transformer_ar const enum geist_kv_mode mode = resolve_kv_mode(opts); sess->kv_kivi_enabled = (mode == GEIST_KV_KIVI); sess->kv_int8_enabled = (mode == GEIST_KV_INT8); + /* Issue #61: opt-in Hadamard rotation, only meaningful on the INT8 path. */ + { + const char *env_rot = getenv("GEIST_KV_ROT"); + sess->kv_rot_enabled = sess->kv_int8_enabled && env_rot != nullptr && env_rot[0] == '1'; + } /* F16 cache: explicit request, or AUTO-resolved FP32 upgraded when the * backend has the fused converting append (env GEIST_KV_F16=0 forces * FP32, =1 requests it under AUTO). Without the slot F16 silently diff --git a/src/archs/transformer/arch_state.h b/src/archs/transformer/arch_state.h index e76b1b9..4b02f0e 100644 --- a/src/archs/transformer/arch_state.h +++ b/src/archs/transformer/arch_state.h @@ -175,6 +175,10 @@ struct transformer_arch_session { * shared across layers (lock-step drain). */ bool kv_int8_enabled; bool kv_kivi_enabled; + /* Experiment (issue #61): FWHT-rotate Q/K/V before INT8 quant to + * suppress activation outliers. Honored only in the INT8 path and + * only when head_dim is a power of two. Env: GEIST_KV_ROT=1. */ + bool kv_rot_enabled; /* F16 KV cache: k_cache[]/v_cache[] hold half floats (2 bytes/elem); * appends convert through the backend's kv_append_f16 slot and * attention reads F16 views. Only set when that slot is non-null. */ diff --git a/src/archs/transformer/forward/kv_store.c b/src/archs/transformer/forward/kv_store.c index 4b9768d..f866b7a 100644 --- a/src/archs/transformer/forward/kv_store.c +++ b/src/archs/transformer/forward/kv_store.c @@ -7,6 +7,7 @@ #include "internal.h" #include +#include "fwht.h" #include "kivi.h" #include @@ -105,14 +106,27 @@ enum geist_status transformer_kv_store_append(struct transformer_layer_forward_c float *v_sca = (float *) v->buffer_map(ctx->v_cache_scale_buf); const size_t row_elems = kv_out; const size_t scales_per_row = st->n_kv_heads; + /* Issue #61: rotate each K/V head row before quantizing. Q is + * rotated symmetrically at attention time (kv_store_attention). */ + const bool rot = st->sess->kv_rot_enabled && fwht_supported(hd) && hd <= 512; + float krot[512]; + float vrot[512]; for (size_t t = 0; t < seq; t++) { const size_t slot = q_position + t; for (size_t h = 0; h < st->n_kv_heads; h++) { - const float *k_row = k_src + t * row_elems + h * hd; - const float *v_row = v_src + t * row_elems + h * hd; - const float k_amax = kv_row_absmax(k_row, hd); - const float v_amax = kv_row_absmax(v_row, hd); - float k_scale = k_amax / 127.0f; + const float *k_row = k_src + t * row_elems + h * hd; + const float *v_row = v_src + t * row_elems + h * hd; + if (rot) { + memcpy(krot, k_row, hd * sizeof(float)); + memcpy(vrot, v_row, hd * sizeof(float)); + fwht_orthonormal(krot, hd); + fwht_orthonormal(vrot, hd); + k_row = krot; + v_row = vrot; + } + const float k_amax = kv_row_absmax(k_row, hd); + const float v_amax = kv_row_absmax(v_row, hd); + float k_scale = k_amax / 127.0f; if (k_scale == 0.0f) { k_scale = 1.0f; } @@ -213,12 +227,22 @@ enum geist_status transformer_kv_store_attention(struct transformer_layer_forwar v->buffer_unmap(ctx->v_residual_buf); v->buffer_unmap(st->sess->scratch_attn); } else if (ctx->kv_int8_enabled) { - const float *qp = (const float *) v->buffer_map(st->sess->scratch_q); + float *qp = (float *) v->buffer_map(st->sess->scratch_q); const int8_t *k_q8p = (const int8_t *) v->buffer_map(ctx->k_cache_q8_buf); const int8_t *v_q8p = (const int8_t *) v->buffer_map(ctx->v_cache_q8_buf); const float *k_scalep = (const float *) v->buffer_map(ctx->k_cache_scale_buf); const float *v_scalep = (const float *) v->buffer_map(ctx->v_cache_scale_buf); float *outp = (float *) v->buffer_map(st->sess->scratch_attn); + /* Issue #61: rotate Q by the same H used on K/V so QK scores are + * unchanged; the kernel then quantizes rotated Q, and we rotate the + * (V-rotated) output back below. H is its own inverse. */ + const bool rot = st->sess->kv_rot_enabled && fwht_supported(ctx->hd) && ctx->hd <= 512; + const size_t n_rows = ctx->seq * st->n_q_heads; + if (rot) { + for (size_t r = 0; r < n_rows; r++) { + fwht_orthonormal(qp + r * ctx->hd, ctx->hd); + } + } /* `scores` scratch is now private per query position inside the kernel * (the loop is parallelized), so no shared arena buffer is needed. */ attention_int8_via_buffers(qp, @@ -234,6 +258,11 @@ enum geist_status transformer_kv_store_attention(struct transformer_layer_forwar ctx->q_position, L->sliding_window, outp); + if (rot) { + for (size_t r = 0; r < n_rows; r++) { + fwht_orthonormal(outp + r * ctx->hd, ctx->hd); + } + } v->buffer_unmap(st->sess->scratch_q); v->buffer_unmap(ctx->k_cache_q8_buf); v->buffer_unmap(ctx->v_cache_q8_buf); diff --git a/src/backends/common/fwht.c b/src/backends/common/fwht.c new file mode 100644 index 0000000..d35a075 --- /dev/null +++ b/src/backends/common/fwht.c @@ -0,0 +1,25 @@ +/* + * fwht.c — orthonormal Fast Walsh–Hadamard Transform. See fwht.h. + */ +#include "fwht.h" + +#include + +void fwht_orthonormal(float *a, size_t n) { + /* Butterfly passes: unnormalized Hadamard (H_n = H_2 ⊗ H_{n/2}). */ + for (size_t len = 1; len < n; len <<= 1) { + for (size_t i = 0; i < n; i += (len << 1)) { + for (size_t j = i; j < i + len; j++) { + const float x = a[j]; + const float y = a[j + len]; + a[j] = x + y; + a[j + len] = x - y; + } + } + } + /* Normalize by 1/sqrt(n) → orthonormal, self-inverse, dot-preserving. */ + const float s = 1.0f / sqrtf((float) n); + for (size_t i = 0; i < n; i++) { + a[i] *= s; + } +} diff --git a/src/backends/common/fwht.h b/src/backends/common/fwht.h new file mode 100644 index 0000000..a2485ab --- /dev/null +++ b/src/backends/common/fwht.h @@ -0,0 +1,29 @@ +/* + * fwht.h — in-place orthonormal Fast Walsh–Hadamard Transform. + * + * Rotation used to suppress activation outliers before INT8 KV-cache + * quantization (issue #61, after llama.cpp#21038 / QuaRot). The Hadamard + * matrix normalized by 1/sqrt(n) is orthonormal, so: + * - it is its own inverse: fwht(fwht(x)) == x (modulo fp rounding); + * - it preserves dot products: (Hx)·(Hy) == x·y. + * The first property lets one transform serve both the forward rotation + * (of Q/K/V) and the backward rotation (of the attention output); the + * second is why rotating Q and K by the same H leaves the QK scores intact. + */ +#ifndef GEIST_FWHT_H +#define GEIST_FWHT_H + +#include +#include + +/* True iff n is a nonzero power of two — the FWHT precondition. Callers + * gate on this and fall back to the unrotated path when it fails. */ +static inline bool fwht_supported(size_t n) { + return n != 0 && (n & (n - 1)) == 0; +} + +/* In-place orthonormal FWHT over a[0..n). n MUST be a power of two + * (see fwht_supported). O(n log n). */ +void fwht_orthonormal(float *a, size_t n); + +#endif /* GEIST_FWHT_H */ diff --git a/tests/test_fwht_unit.c b/tests/test_fwht_unit.c new file mode 100644 index 0000000..e22e129 --- /dev/null +++ b/tests/test_fwht_unit.c @@ -0,0 +1,110 @@ +/* + * test_fwht_unit — verifies the orthonormal FWHT used to rotate attention + * Q/K/V before INT8 KV-cache quantization (issue #61). + * + * Three properties that the rotation trick depends on: + * 1. Self-inverse: H(H(x)) == x. One matrix serves forward (Q/K/V) and + * backward (attention output) rotation. + * 2. Dot-product preserving: (Hx)·(Hy) == x·y. This is why rotating Q and + * K by the same H leaves the QK attention scores unchanged. + * 3. Outlier suppression: a spiky vector's peak magnitude shrinks after + * rotation, so symmetric INT8 quant (scale = amax/127) wastes fewer + * levels on the outlier — the whole reason to rotate before quantizing. + * + * Deterministic — fixed seed, no model needed. + */ +#include "fwht.h" +#include "test_helpers.h" + +#include + +/* Deterministic N(0,1) via Box-Muller over a tiny LCG. */ +static float gauss(uint32_t *seed) { + uint32_t a = (*seed = (*seed) * 1103515245u + 12345u); + uint32_t b = (*seed = (*seed) * 1103515245u + 12345u); + float u1 = ((float) (a & 0xffffff) + 1.0f) / (float) 0x1000000; + float u2 = ((float) (b & 0xffffff)) / (float) 0x1000000; + return sqrtf(-2.0f * logf(u1)) * cosf(6.2831853f * u2); +} + +static float dot(const float *a, const float *b, size_t n) { + double s = 0.0; + for (size_t i = 0; i < n; i++) + s += (double) a[i] * (double) b[i]; + return (float) s; +} + +static float absmax(const float *a, size_t n) { + float m = 0.0f; + for (size_t i = 0; i < n; i++) + m = fabsf(a[i]) > m ? fabsf(a[i]) : m; + return m; +} + +/* Property 1: H(H(x)) reproduces x. */ +static int test_self_inverse(void) { + enum { N = 256 }; + float x[N], y[N]; + uint32_t seed = 0x1234u; + for (size_t i = 0; i < N; i++) + x[i] = y[i] = gauss(&seed); + fwht_orthonormal(y, N); + fwht_orthonormal(y, N); + for (size_t i = 0; i < N; i++) { + if (fabsf(y[i] - x[i]) > 1e-4f) { + printf("self-inverse: y[%zu]=%.6f != x=%.6f\n", i, y[i], x[i]); + return 1; + } + } + return 0; +} + +/* Property 2: rotation preserves the dot product used by QK scores. */ +static int test_dot_preserved(void) { + enum { N = 128 }; + float a[N], b[N], ra[N], rb[N]; + uint32_t seed = 0x9e37u; + for (size_t i = 0; i < N; i++) { + a[i] = ra[i] = gauss(&seed); + b[i] = rb[i] = gauss(&seed); + } + fwht_orthonormal(ra, N); + fwht_orthonormal(rb, N); + const float d0 = dot(a, b, N); + const float d1 = dot(ra, rb, N); + if (fabsf(d0 - d1) > 1e-3f * (1.0f + fabsf(d0))) { + printf("dot not preserved: %.6f vs %.6f\n", d0, d1); + return 1; + } + return 0; +} + +/* Property 3: a single big outlier is spread out, dropping the peak + * magnitude that symmetric INT8 quant has to reach. */ +static int test_outlier_suppressed(void) { + enum { N = 256 }; + float x[N]; + uint32_t seed = 0x5eedu; + for (size_t i = 0; i < N; i++) + x[i] = 0.1f * gauss(&seed); + x[42] = 20.0f; /* the outlier */ + const float peak0 = absmax(x, N); + fwht_orthonormal(x, N); + const float peak1 = absmax(x, N); + /* 20.0 spread over sqrt(256)=16 → peak ~1.25, far below 20. Assert a + * clear drop rather than the exact figure. */ + if (!(peak1 < 0.5f * peak0)) { + printf("outlier not suppressed: peak %.4f -> %.4f\n", peak0, peak1); + return 1; + } + return 0; +} + +int main(void) { + if (fwht_supported(256) && fwht_supported(128) && !fwht_supported(0) && !fwht_supported(96) && + test_self_inverse() == 0 && test_dot_preserved() == 0 && test_outlier_suppressed() == 0) { + printf("PASS: FWHT is self-inverse, dot-preserving, and suppresses outliers\n"); + return GEIST_TEST_PASS; + } + return GEIST_TEST_FAIL; +} From 81afa6872175581d7e99c18d44bb40e3406b6283 Mon Sep 17 00:00:00 2001 From: germar Date: Sun, 5 Jul 2026 09:49:33 +0200 Subject: [PATCH 02/10] test(kv): sweep all four KV modes in bench_kv_quality (#61) Run FP32/INT8/INT8+ROT/KIVI in one invocation (reload per mode so the GEIST_KV_* env resolution stays untouched) and print the fraction of the INT8->KIVI quality gap that the Hadamard rotation recovers. Also fall back to the GGUF-embedded tokenizer so Gemma/BitNet models work, not just external tokenizer.bin ones. --- tests/bench_kv_quality.c | 245 ++++++++++++++++++++++++--------------- 1 file changed, 154 insertions(+), 91 deletions(-) diff --git a/tests/bench_kv_quality.c b/tests/bench_kv_quality.c index a192706..6468850 100644 --- a/tests/bench_kv_quality.c +++ b/tests/bench_kv_quality.c @@ -10,10 +10,13 @@ * matches the actual next token in the text. A monotonic scalar quality * signal sensitive to KV-cache quantization noise. * - * Run with different GEIST_KV_* env vars to compare modes: - * GEIST_KV_INT8=0 ./bench_kv_quality # FP32 KV (reference) - * GEIST_KV_INT8=1 ./bench_kv_quality # INT8 KV - * GEIST_KV_KIVI=1 ./bench_kv_quality # KIVI (future) + * Runs all four KV modes in one pass (FP32 / INT8 / INT8+ROT / KIVI) and + * prints a comparison table plus the fraction of the INT8→KIVI quality gap + * that the issue-#61 Hadamard rotation recovers. The KV mode is resolved + * from GEIST_KV_* env vars at model-load time, so each mode reloads the + * model with those vars set (weights re-mmap, cheap). Override a single + * mode the old way by exporting GEIST_KV_* before the run — it is honored + * as the process default but the sweep sets its own per iteration. * * Uses arch_ops->verify_forward via geist_model_internal_arch_meta — a * test-only accessor that bypasses the public session API to keep the @@ -27,6 +30,7 @@ #include "src/engine/model.h" #include "src/engine/sp_bpe_tokenizer.h" +#include "src/engine/gguf_tokenizer.h" #include "src/archs/transformer/arch.h" /* geist_arch_transformer descriptor */ #include @@ -66,93 +70,69 @@ static const char *DEFAULT_TEXT = "also exposes a PCI Express interface via a flat ribbon connector " "for high-speed peripherals such as NVMe storage."; -int main(int argc, char **argv) { - const char *model_path = argc > 1 ? argv[1] : geist_test_find_gguf(); - GEIST_SKIP_IF(model_path == nullptr, "no GGUF model found — pass path or set GEIST_GGUF_PATH"); +/* One KV mode: which GEIST_KV_* vars to set before the model load that + * resolves it. A nullptr field means unset that var. */ +struct kv_mode_cfg { + const char *label; + const char *int8; /* GEIST_KV_INT8 */ + const char *kivi; /* GEIST_KV_KIVI */ + const char *rot; /* GEIST_KV_ROT */ +}; - const char *text = (argc > 2) ? argv[2] : DEFAULT_TEXT; +static const struct kv_mode_cfg MODES[] = { + {"FP32", "0", nullptr, nullptr}, + {"INT8", "1", nullptr, "0"}, + {"INT8+ROT", "1", nullptr, "1"}, + {"KIVI", nullptr, "1", nullptr}, +}; - struct geist_backend *be = nullptr; - enum geist_status s = geist_backend_create("cpu_neon", nullptr, nullptr, &be); - if (s != GEIST_OK) { - s = geist_backend_create("cpu_scalar", nullptr, nullptr, &be); - } - if (s != GEIST_OK) { - fprintf(stderr, "backend create: %s\n", geist_last_create_error()); - return GEIST_TEST_ERROR; - } +static void set_or_unset(const char *name, const char *val) { + if (val != nullptr) + setenv(name, val, 1); + else + unsetenv(name); +} +/* Load the model under the current GEIST_KV_* env, run top-1 next-token + * prediction over `ids`, and return accuracy. Returns -1.0 on error. The + * caller owns `ids` (tokenized once and reused across modes). */ +static double run_one(const char *model_path, + struct geist_backend *be, + const uint32_t *ids, + size_t n_ids, + size_t *n_eval_out) { struct geist_model *model = nullptr; - s = geist_model_load(model_path, be, &model); + enum geist_status s = geist_model_load(model_path, be, &model); if (s != GEIST_OK) { fprintf(stderr, - "model_load(%s): %s — %s\n", - model_path, + "model_load: %s — %s\n", geist_status_to_string(s), geist_last_create_error()); - geist_backend_destroy(be); - return GEIST_TEST_FAIL; - } - - struct sp_bpe_tokenizer *tok = geist_model_internal_tokenizer(model); - GEIST_SKIP_IF(tok == nullptr, "no tokenizer.bin reachable — set GEIST_TOKENIZER_PATH"); - - uint32_t *ids = nullptr; - size_t n_ids = 0; - if (!sp_bpe_tokenizer_encode(tok, text, &ids, &n_ids)) { - fprintf(stderr, "tokenizer encode failed\n"); - geist_model_destroy(model); - geist_backend_destroy(be); - return GEIST_TEST_FAIL; - } - if (n_ids < 2) { - fprintf(stderr, "text too short (n_ids=%zu)\n", n_ids); - free(ids); - geist_model_destroy(model); - geist_backend_destroy(be); - return GEIST_TEST_FAIL; + return -1.0; } - /* verify_forward runs the WHOLE input through the layer stack in one - * batched pass (using prefill chunks of m_max internally) and writes - * argmax at each position. We don't need a session — go straight to - * arch_ops + arch_meta. */ void *arch_meta = geist_model_internal_arch_meta(model); const struct geist_arch_ops_decoder *ops = &geist_arch_transformer; - if (ops->state_reset == nullptr || ops->verify_forward == nullptr || ops->kv_truncate == nullptr) { - fprintf(stderr, - "arch lacks state_reset / verify_forward / kv_truncate — " - "expected from transformer decoder\n"); - free(ids); + fprintf(stderr, "arch lacks state_reset / verify_forward / kv_truncate\n"); geist_model_destroy(model); - geist_backend_destroy(be); - return GEIST_TEST_FAIL; + return -1.0; } ops->state_reset(arch_meta); - ops->kv_truncate(arch_meta, 0); /* belt-and-braces: prefix_length may be nonzero */ + ops->kv_truncate(arch_meta, 0); - const size_t k = n_ids - 1; /* we predict positions 1..N-1 from inputs 0..N-2 */ + const size_t k = n_ids - 1; geist_token_t *preds = (geist_token_t *) calloc(k, sizeof *preds); if (preds == nullptr) { - fprintf(stderr, "alloc preds (%zu) failed\n", k); - free(ids); geist_model_destroy(model); - geist_backend_destroy(be); - return GEIST_TEST_FAIL; + return -1.0; } - /* Chunk because verify_forward caps at m_max (default 64). Each chunk - * extends the KV cache; predictions are appended in order. - * - * verify_forward does NOT drain KIVI residuals (it's the tentative- - * write path), so we follow each chunk with kv_truncate(kv_len) — - * which is a no-op on kv_len but forces a drain if the residual - * grew past R. This makes the harness exercise the actual 2-bit - * drained-cache attention path for KIVI mode, matching what an - * accept-only decode_step stream would do. For non-KIVI modes the - * truncate is harmless. */ + + /* Chunk because verify_forward caps at m_max (64). kv_truncate after + * each chunk forces a KIVI residual drain so the 2-bit path is + * exercised; harmless for the other modes. */ const size_t M_MAX = 64; size_t kv_len_acc = 0; for (size_t off = 0; off < k; off += M_MAX) { @@ -160,48 +140,131 @@ int main(int argc, char **argv) { enum geist_status vs = ops->verify_forward( arch_meta, chunk, (const geist_token_t *) ids + off, preds + off); if (vs != GEIST_OK) { - fprintf(stderr, - "verify_forward(chunk@%zu, %zu): %s\n", - off, - chunk, - geist_status_to_string(vs)); + fprintf(stderr, "verify_forward(@%zu): %s\n", off, geist_status_to_string(vs)); free(preds); - free(ids); geist_model_destroy(model); - geist_backend_destroy(be); - return GEIST_TEST_FAIL; + return -1.0; } kv_len_acc += chunk; ops->kv_truncate(arch_meta, kv_len_acc); } - /* preds[i] is the model's argmax given prefix ids[0..i]; the target - * for position i+1 in the text is ids[i+1]. */ size_t n_correct = 0; for (size_t i = 0; i < k - 1; i++) { if ((uint32_t) preds[i] == ids[i + 1]) n_correct++; } const size_t n_eval = k - 1; - const double acc = n_eval > 0 ? (double) n_correct / (double) n_eval : 0.0; + free(preds); + geist_model_destroy(model); + if (n_eval_out != nullptr) + *n_eval_out = n_eval; + return n_eval > 0 ? (double) n_correct / (double) n_eval : 0.0; +} + +int main(int argc, char **argv) { + const char *model_path = argc > 1 ? argv[1] : geist_test_find_gguf(); + GEIST_SKIP_IF(model_path == nullptr, "no GGUF model found — pass path or set GEIST_GGUF_PATH"); - const char *mode_kivi = getenv("GEIST_KV_KIVI"); - const char *mode_int8 = getenv("GEIST_KV_INT8"); - const char *mode_label = (mode_kivi != nullptr && mode_kivi[0] == '1') ? "KIVI" - : (mode_int8 != nullptr && mode_int8[0] == '1') ? "INT8" - : (mode_int8 != nullptr && mode_int8[0] == '0') - ? "FP32" - : "FP32 (Apple-default)"; + const char *text = (argc > 2) ? argv[2] : DEFAULT_TEXT; + + struct geist_backend *be = nullptr; + enum geist_status s = geist_backend_create("cpu_neon", nullptr, nullptr, &be); + if (s != GEIST_OK) { + s = geist_backend_create("cpu_scalar", nullptr, nullptr, &be); + } + if (s != GEIST_OK) { + fprintf(stderr, "backend create: %s\n", geist_last_create_error()); + return GEIST_TEST_ERROR; + } + + /* Tokenize once (mode-independent) via a throwaway load. */ + struct geist_model *model = nullptr; + s = geist_model_load(model_path, be, &model); + if (s != GEIST_OK) { + fprintf(stderr, "model_load(%s): %s\n", model_path, geist_status_to_string(s)); + geist_backend_destroy(be); + return GEIST_TEST_FAIL; + } + /* Try the sp_bpe tokenizer (external tokenizer.bin models) first, then + * the GGUF-embedded tokenizer (Gemma/BitNet unigram, Qwen bpe). */ + uint32_t *ids = nullptr; + size_t n_ids = 0; + bool enc_ok = false; + bool had_tokenizer = false; + struct sp_bpe_tokenizer *tok = geist_model_internal_tokenizer(model); + if (tok != nullptr) { + had_tokenizer = true; + enc_ok = sp_bpe_tokenizer_encode(tok, text, &ids, &n_ids); + } else { + struct gguf_tokenizer *gtok = geist_model_internal_gguf_tokenizer(model); + if (gtok != nullptr) { + had_tokenizer = true; + const size_t cap = strlen(text) + 16; + int32_t *tmp = (int32_t *) malloc(cap * sizeof(int32_t)); + if (tmp != nullptr && gguf_tokenizer_encode(gtok, text, tmp, cap, &n_ids)) { + ids = (uint32_t *) malloc(n_ids * sizeof(uint32_t)); + if (ids != nullptr) { + for (size_t i = 0; i < n_ids; i++) + ids[i] = (uint32_t) tmp[i]; + enc_ok = true; + } + } + free(tmp); + } + } + geist_model_destroy(model); /* done with this load; sweep reloads per mode */ + if (!had_tokenizer) { + free(ids); + geist_backend_destroy(be); + GEIST_SKIP_IF(true, "model carries no usable tokenizer"); + } + if (!enc_ok || n_ids < 2) { + fprintf(stderr, "tokenizer encode failed or text too short (n_ids=%zu)\n", n_ids); + free(ids); + geist_backend_destroy(be); + return GEIST_TEST_FAIL; + } printf("model: %s\n", model_path); printf("backend: %s\n", geist_backend_name(be)); - printf("kv_mode: %s\n", mode_label); - printf("n_tokens: %zu (n_eval=%zu)\n", n_ids, n_eval); - printf("top-1 acc: %zu / %zu = %.4f\n", n_correct, n_eval, acc); + printf("n_tokens: %zu (n_eval=%zu)\n\n", n_ids, n_ids - 2); + + const size_t n_modes = sizeof(MODES) / sizeof(MODES[0]); + double acc[sizeof(MODES) / sizeof(MODES[0])]; + printf("%-10s %s\n", "kv_mode", "top-1 acc"); + printf("%-10s %s\n", "-------", "---------"); + for (size_t m = 0; m < n_modes; m++) { + set_or_unset("GEIST_KV_INT8", MODES[m].int8); + set_or_unset("GEIST_KV_KIVI", MODES[m].kivi); + set_or_unset("GEIST_KV_ROT", MODES[m].rot); + acc[m] = run_one(model_path, be, ids, n_ids, nullptr); + if (acc[m] < 0.0) { + free(ids); + geist_backend_destroy(be); + return GEIST_TEST_FAIL; + } + printf("%-10s %.4f\n", MODES[m].label, acc[m]); + } + + /* Issue #61 headline: fraction of the INT8→KIVI gap that ROT recovers. + * MODES order is FP32, INT8, INT8+ROT, KIVI. */ + const double gap = acc[3] - acc[1]; /* KIVI - INT8 */ + if (gap > 1e-6) { + const double recovered = (acc[2] - acc[1]) / gap; + printf("\nINT8→KIVI gap recovered by rotation: %.0f%% (INT8 %.4f → ROT %.4f → KIVI " + "%.4f)\n", + 100.0 * recovered, + acc[1], + acc[2], + acc[3]); + } else { + printf("\nINT8 already within noise of KIVI (gap=%.4f) — rotation headroom is negligible " + "here.\n", + gap); + } - free(preds); free(ids); - geist_model_destroy(model); geist_backend_destroy(be); return GEIST_TEST_PASS; } From 238f2035711b3f6d8002f034352b901bed1dcbe9 Mon Sep 17 00:00:00 2001 From: germar Date: Sun, 5 Jul 2026 09:59:48 +0200 Subject: [PATCH 03/10] feat(kv): N-bit low-bit KV quality-sim + INT4/INT2 bench modes (#61) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reuse the INT8 storage path with a scale = amax/(2^(N-1)-1) grid to simulate a symmetric low-bit cache without packing/kernels — a quality probe for whether rotation rescues low-bit KV. GEIST_KV_INT4=1 (4-bit) or GEIST_KV_QBITS=N. bench_kv_quality sweeps FP32/INT8/INT4/INT2 (+ROT) and KIVI in one run. Finding: rotation's benefit scales inversely with bit-width (nil at 8b, recovers the 4b penalty to ~INT8, recovers most of the 2b collapse) but symmetric-2b+ROT still trails asymmetric KIVI. The usable win is a rotated 4-bit cache: ~INT8 quality at half the bytes, no per-channel bookkeeping. --- src/archs/transformer/arch_state.c | 23 +++++++++ src/archs/transformer/arch_state.h | 7 +++ src/archs/transformer/forward/kv_store.c | 10 +++- tests/bench_kv_quality.c | 61 +++++++++++++++++------- 4 files changed, 81 insertions(+), 20 deletions(-) diff --git a/src/archs/transformer/arch_state.c b/src/archs/transformer/arch_state.c index 4a6ce49..9023601 100644 --- a/src/archs/transformer/arch_state.c +++ b/src/archs/transformer/arch_state.c @@ -998,6 +998,29 @@ struct transformer_arch_session *transformer_session_alloc(struct transformer_ar const enum geist_kv_mode mode = resolve_kv_mode(opts); sess->kv_kivi_enabled = (mode == GEIST_KV_KIVI); sess->kv_int8_enabled = (mode == GEIST_KV_INT8); + /* Issue #61: low-bit quality-sim reuses the INT8 storage path with an + * N-bit quant grid. GEIST_KV_INT4=1 → 4 bits; GEIST_KV_QBITS=N (2..8) + * overrides. Any sim (2..7) forces INT8 storage on regardless of the + * resolved mode. Resolve before the rot flag so rotation sees it. */ + { + int qbits = 0; + const char *env_int4 = getenv("GEIST_KV_INT4"); + const char *env_qbits = getenv("GEIST_KV_QBITS"); + if (env_int4 != nullptr && env_int4[0] == '1') { + qbits = 4; + } + if (env_qbits != nullptr) { + const int q = atoi(env_qbits); + if (q >= 2 && q <= 8) { + qbits = (q == 8) ? 0 : q; /* 8-bit is the native path */ + } + } + sess->kv_sim_qbits = qbits; + if (qbits != 0) { + sess->kv_int8_enabled = true; + sess->kv_kivi_enabled = false; + } + } /* Issue #61: opt-in Hadamard rotation, only meaningful on the INT8 path. */ { const char *env_rot = getenv("GEIST_KV_ROT"); diff --git a/src/archs/transformer/arch_state.h b/src/archs/transformer/arch_state.h index 4b02f0e..e52fd86 100644 --- a/src/archs/transformer/arch_state.h +++ b/src/archs/transformer/arch_state.h @@ -179,6 +179,13 @@ struct transformer_arch_session { * suppress activation outliers. Honored only in the INT8 path and * only when head_dim is a power of two. Env: GEIST_KV_ROT=1. */ bool kv_rot_enabled; + /* Experiment (issue #61): quantize the INT8 K/V cache on an N-bit grid + * (scale = amax / (2^(N-1)-1)) — quality-only simulation of a symmetric + * low-bit cache that reuses the INT8 storage + kernel (no packing, no + * memory win yet). Measures whether rotation rescues low-bit quality. + * 0 = native 8-bit; 2..7 forces the INT8 storage path on. Env: + * GEIST_KV_INT4=1 (→4) or GEIST_KV_QBITS=N. */ + int kv_sim_qbits; /* F16 KV cache: k_cache[]/v_cache[] hold half floats (2 bytes/elem); * appends convert through the backend's kv_append_f16 slot and * attention reads F16 views. Only set when that slot is non-null. */ diff --git a/src/archs/transformer/forward/kv_store.c b/src/archs/transformer/forward/kv_store.c index f866b7a..cc41063 100644 --- a/src/archs/transformer/forward/kv_store.c +++ b/src/archs/transformer/forward/kv_store.c @@ -111,6 +111,12 @@ enum geist_status transformer_kv_store_append(struct transformer_layer_forward_c const bool rot = st->sess->kv_rot_enabled && fwht_supported(hd) && hd <= 512; float krot[512]; float vrot[512]; + /* Issue #61: low-bit quality-sim quantizes on an N-bit grid + * (amax / (2^(N-1)-1)); the int8 container then holds values in + * [-(2^(N-1)-1), +...]. Rounding stays in range (|x| <= amax), so no + * clamp is needed. qbits==0 is the native 8-bit path (denom 127). */ + const int qbits = st->sess->kv_sim_qbits; + const float denom = qbits != 0 ? (float) ((1 << (qbits - 1)) - 1) : 127.0f; for (size_t t = 0; t < seq; t++) { const size_t slot = q_position + t; for (size_t h = 0; h < st->n_kv_heads; h++) { @@ -126,11 +132,11 @@ enum geist_status transformer_kv_store_append(struct transformer_layer_forward_c } const float k_amax = kv_row_absmax(k_row, hd); const float v_amax = kv_row_absmax(v_row, hd); - float k_scale = k_amax / 127.0f; + float k_scale = k_amax / denom; if (k_scale == 0.0f) { k_scale = 1.0f; } - float v_scale = v_amax / 127.0f; + float v_scale = v_amax / denom; if (v_scale == 0.0f) { v_scale = 1.0f; } diff --git a/tests/bench_kv_quality.c b/tests/bench_kv_quality.c index 6468850..d9c9838 100644 --- a/tests/bench_kv_quality.c +++ b/tests/bench_kv_quality.c @@ -74,18 +74,31 @@ static const char *DEFAULT_TEXT = * resolves it. A nullptr field means unset that var. */ struct kv_mode_cfg { const char *label; - const char *int8; /* GEIST_KV_INT8 */ - const char *kivi; /* GEIST_KV_KIVI */ - const char *rot; /* GEIST_KV_ROT */ + const char *int8; /* GEIST_KV_INT8 */ + const char *kivi; /* GEIST_KV_KIVI */ + const char *rot; /* GEIST_KV_ROT */ + const char *int4; /* GEIST_KV_INT4 (4-bit quality-sim on the INT8 path) */ + const char *qbits; /* GEIST_KV_QBITS (N-bit quality-sim; overrides int4) */ }; static const struct kv_mode_cfg MODES[] = { - {"FP32", "0", nullptr, nullptr}, - {"INT8", "1", nullptr, "0"}, - {"INT8+ROT", "1", nullptr, "1"}, - {"KIVI", nullptr, "1", nullptr}, + {"FP32", "0", nullptr, nullptr, nullptr, nullptr}, + {"INT8", "1", nullptr, "0", nullptr, nullptr}, + {"INT8+ROT", "1", nullptr, "1", nullptr, nullptr}, + {"INT4", nullptr, nullptr, "0", "1", nullptr}, + {"INT4+ROT", nullptr, nullptr, "1", "1", nullptr}, + {"INT2", nullptr, nullptr, "0", nullptr, "2"}, + {"INT2+ROT", nullptr, nullptr, "1", nullptr, "2"}, + {"KIVI", nullptr, "1", nullptr, nullptr, nullptr}, }; +static size_t mode_index(const char *label) { + for (size_t i = 0; i < sizeof(MODES) / sizeof(MODES[0]); i++) + if (strcmp(MODES[i].label, label) == 0) + return i; + return 0; /* labels are compile-time constants — unreachable */ +} + static void set_or_unset(const char *name, const char *val) { if (val != nullptr) setenv(name, val, 1); @@ -238,6 +251,8 @@ int main(int argc, char **argv) { set_or_unset("GEIST_KV_INT8", MODES[m].int8); set_or_unset("GEIST_KV_KIVI", MODES[m].kivi); set_or_unset("GEIST_KV_ROT", MODES[m].rot); + set_or_unset("GEIST_KV_INT4", MODES[m].int4); + set_or_unset("GEIST_KV_QBITS", MODES[m].qbits); acc[m] = run_one(model_path, be, ids, n_ids, nullptr); if (acc[m] < 0.0) { free(ids); @@ -247,23 +262,33 @@ int main(int argc, char **argv) { printf("%-10s %.4f\n", MODES[m].label, acc[m]); } - /* Issue #61 headline: fraction of the INT8→KIVI gap that ROT recovers. - * MODES order is FP32, INT8, INT8+ROT, KIVI. */ - const double gap = acc[3] - acc[1]; /* KIVI - INT8 */ + /* Issue #61 headline (option A): does rotation rescue the low-bit INT4 + * cache? Report the fraction of the INT4→INT8 quality drop that ROT + * recovers — INT8 is the lossless ceiling, INT4 the degraded floor. */ + const double a_int8 = acc[mode_index("INT8")]; + const double a_int4 = acc[mode_index("INT4")]; + const double a_int4rot = acc[mode_index("INT4+ROT")]; + const double gap = a_int8 - a_int4; /* low-bit penalty */ + printf("\n"); if (gap > 1e-6) { - const double recovered = (acc[2] - acc[1]) / gap; - printf("\nINT8→KIVI gap recovered by rotation: %.0f%% (INT8 %.4f → ROT %.4f → KIVI " - "%.4f)\n", + const double recovered = (a_int4rot - a_int4) / gap; + printf("INT4→INT8 gap recovered by rotation: %.0f%% (INT4 %.4f → +ROT %.4f → INT8 %.4f)\n", 100.0 * recovered, - acc[1], - acc[2], - acc[3]); + a_int4, + a_int4rot, + a_int8); } else { - printf("\nINT8 already within noise of KIVI (gap=%.4f) — rotation headroom is negligible " - "here.\n", + printf("INT4 already within noise of INT8 (gap=%.4f) — no low-bit penalty to recover.\n", gap); } + /* The lazier-KIVI question: can symmetric 2-bit + rotation rival KIVI's + * asymmetric per-channel 2-bit? */ + printf("2-bit: symmetric INT2 %.4f → +ROT %.4f vs KIVI(asym) %.4f\n", + acc[mode_index("INT2")], + acc[mode_index("INT2+ROT")], + acc[mode_index("KIVI")]); + free(ids); geist_backend_destroy(be); return GEIST_TEST_PASS; From 083c115eec729fed1d2d9f358e9d34ac416c1843 Mon Sep 17 00:00:00 2001 From: germar Date: Sun, 5 Jul 2026 10:20:27 +0200 Subject: [PATCH 04/10] test(kv): add KL-divergence metric to bench_kv_quality (#61) top-1 accuracy saturated and couldn't resolve sub-1% quality differences. Switch to the public session API (prefill_tokens + peek_logits) and score each mode by mean KL(P_fp32 || P_mode) over all positions, teacher-forced in lockstep against an FP32 reference session. Model loads once; each mode is a session. KL over the shared finite support so BitNet's non-finite dead-vocab slots don't poison it. KL cleanly shows what top-1 hid: rotation reduces divergence at low bit- widths (BitNet: INT4 -33%, INT2 -44%; ~nil at INT8), but symmetric 2-bit +ROT still trails asymmetric KIVI. INT4+ROT is near-lossless. --- tests/bench_kv_quality.c | 337 ++++++++++++++++++++++----------------- 1 file changed, 187 insertions(+), 150 deletions(-) diff --git a/tests/bench_kv_quality.c b/tests/bench_kv_quality.c index d9c9838..8b86582 100644 --- a/tests/bench_kv_quality.c +++ b/tests/bench_kv_quality.c @@ -1,41 +1,40 @@ /* - * bench_kv_quality — KV-quant quality A/B via single-pass top-1 prediction. + * bench_kv_quality — KV-quant quality sweep via top-1 accuracy AND KL + * divergence vs an FP32 reference. * - * Tokenizes a fixed paragraph, then calls verify_forward on the WHOLE - * sequence in one batched pass. verify_forward produces the model's - * argmax at each position simultaneously, so we get N-1 predictions for - * the cost of one forward pass. + * Tokenizes a fixed paragraph, then for each KV mode teacher-forces the + * sequence one token at a time through the public session API, reading the + * full next-token distribution at every position with peek_logits. Two + * quality signals per mode: * - * Reports top-1 accuracy: fraction of positions where the model's argmax - * matches the actual next token in the text. A monotonic scalar quality - * signal sensitive to KV-cache quantization noise. + * top-1 acc : fraction of positions where the mode's argmax matches the + * actual next token. Coarse — saturates, sub-1% noise. + * mean KL : mean KL(P_fp32 || P_mode) over all positions, in nats. + * Sensitive — resolves quality differences top-1 cannot. * - * Runs all four KV modes in one pass (FP32 / INT8 / INT8+ROT / KIVI) and - * prints a comparison table plus the fraction of the INT8→KIVI quality gap - * that the issue-#61 Hadamard rotation recovers. The KV mode is resolved - * from GEIST_KV_* env vars at model-load time, so each mode reloads the - * model with those vars set (weights re-mmap, cheap). Override a single - * mode the old way by exporting GEIST_KV_* before the run — it is honored - * as the process default but the sweep sets its own per iteration. + * Modes: FP32 / INT8 / INT8+ROT / INT4 / INT4+ROT / INT2 / INT2+ROT / KIVI. + * INT4/INT2 are the issue-#61 low-bit quality-sims (GEIST_KV_INT4 / + * GEIST_KV_QBITS) that reuse the INT8 storage path with an N-bit grid. The + * KV mode is resolved per session from GEIST_KV_* env vars, so the model + * loads once and each mode just creates a session with those vars set. An + * FP32 reference session runs in lockstep to supply P_fp32 for the KL. * - * Uses arch_ops->verify_forward via geist_model_internal_arch_meta — a - * test-only accessor that bypasses the public session API to keep the - * harness independent of the live spec_step plumbing. SKIPs if no GGUF - * or tokenizer is reachable. + * SKIPs if no GGUF or tokenizer is reachable. */ #define GEIST_INTERNAL_ENGINE_LAYER #define GEIST_INTERNAL_ARCH_LAYER #include "test_helpers.h" +#include "src/engine/gguf_tokenizer.h" #include "src/engine/model.h" #include "src/engine/sp_bpe_tokenizer.h" -#include "src/engine/gguf_tokenizer.h" -#include "src/archs/transformer/arch.h" /* geist_arch_transformer descriptor */ #include #include +#include +#include #include #include #include @@ -70,8 +69,8 @@ static const char *DEFAULT_TEXT = "also exposes a PCI Express interface via a flat ribbon connector " "for high-speed peripherals such as NVMe storage."; -/* One KV mode: which GEIST_KV_* vars to set before the model load that - * resolves it. A nullptr field means unset that var. */ +/* One KV mode: which GEIST_KV_* vars to set before creating its session. + * A nullptr field means unset that var. */ struct kv_mode_cfg { const char *label; const char *int8; /* GEIST_KV_INT8 */ @@ -106,92 +105,142 @@ static void set_or_unset(const char *name, const char *val) { unsetenv(name); } -/* Load the model under the current GEIST_KV_* env, run top-1 next-token - * prediction over `ids`, and return accuracy. Returns -1.0 on error. The - * caller owns `ids` (tokenized once and reused across modes). */ -static double run_one(const char *model_path, - struct geist_backend *be, - const uint32_t *ids, - size_t n_ids, - size_t *n_eval_out) { - struct geist_model *model = nullptr; - enum geist_status s = geist_model_load(model_path, be, &model); - if (s != GEIST_OK) { - fprintf(stderr, - "model_load: %s — %s\n", - geist_status_to_string(s), - geist_last_create_error()); - return -1.0; - } +static void apply_mode_env(const struct kv_mode_cfg *m) { + set_or_unset("GEIST_KV_INT8", m->int8); + set_or_unset("GEIST_KV_KIVI", m->kivi); + set_or_unset("GEIST_KV_ROT", m->rot); + set_or_unset("GEIST_KV_INT4", m->int4); + set_or_unset("GEIST_KV_QBITS", m->qbits); +} - void *arch_meta = geist_model_internal_arch_meta(model); - const struct geist_arch_ops_decoder *ops = &geist_arch_transformer; - if (ops->state_reset == nullptr || ops->verify_forward == nullptr || - ops->kv_truncate == nullptr) { - fprintf(stderr, "arch lacks state_reset / verify_forward / kv_truncate\n"); - geist_model_destroy(model); - return -1.0; +static uint32_t argmax_f32(const float *x, size_t n) { + uint32_t best = 0; + for (size_t i = 1; i < n; i++) + if (x[i] > x[best]) + best = (uint32_t) i; + return best; +} + +/* KL(P_ref || P_mode) in nats over the shared finite support. Some models + * (e.g. BitNet) leave non-finite garbage in unused vocab slots; a slot that + * is non-finite in EITHER distribution is dropped from both before + * normalizing, so the two softmaxes are over the same token set. Streaming + * log-sum-exp — no full-vocab buffers. */ +static double kl_div(const float *lr, const float *lm, size_t n) { + bool have = false; + double maxr = 0.0, maxm = 0.0; + for (size_t i = 0; i < n; i++) { + if (!isfinite(lr[i]) || !isfinite(lm[i])) + continue; + if (!have || lr[i] > maxr) + maxr = lr[i]; + if (!have || lm[i] > maxm) + maxm = lm[i]; + have = true; + } + if (!have) + return 0.0; + double sr = 0.0, sm = 0.0; + for (size_t i = 0; i < n; i++) { + if (!isfinite(lr[i]) || !isfinite(lm[i])) + continue; + sr += exp((double) lr[i] - maxr); + sm += exp((double) lm[i] - maxm); } + const double logZr = maxr + log(sr); + const double logZm = maxm + log(sm); + double kl = 0.0; + for (size_t i = 0; i < n; i++) { + if (!isfinite(lr[i]) || !isfinite(lm[i])) + continue; + const double lp_r = (double) lr[i] - logZr; + kl += exp(lp_r) * (lp_r - ((double) lm[i] - logZm)); + } + return kl < 0.0 ? 0.0 : kl; /* clamp fp noise on identical dists */ +} - ops->state_reset(arch_meta); - ops->kv_truncate(arch_meta, 0); +/* Teacher-force `ids` through both sessions in lockstep; fill top-1 acc and + * mean KL(ref||mode). Returns false on a session error. */ +static bool score_mode(struct geist_session *ref, + struct geist_session *mode, + const uint32_t *ids, + size_t n_ids, + double *acc_out, + double *kl_out) { + if (geist_session_reset(ref) != GEIST_OK || geist_session_reset(mode) != GEIST_OK) + return false; + const geist_token_t t0 = (geist_token_t) ids[0]; + if (geist_session_prefill_tokens(ref, 1, &t0) != GEIST_OK || + geist_session_prefill_tokens(mode, 1, &t0) != GEIST_OK) + return false; - const size_t k = n_ids - 1; - geist_token_t *preds = (geist_token_t *) calloc(k, sizeof *preds); - if (preds == nullptr) { - geist_model_destroy(model); - return -1.0; + double kl_sum = 0.0; + size_t correct = 0, n_eval = 0; + for (size_t i = 1; i < n_ids; i++) { + size_t nr = 0, nm = 0; + const float *lr = geist_session_peek_logits(ref, &nr); + const float *lm = geist_session_peek_logits(mode, &nm); + if (lr == nullptr || lm == nullptr || nr != nm || nr == 0) + return false; + kl_sum += kl_div(lr, lm, nr); + if (argmax_f32(lm, nm) == ids[i]) + correct++; + n_eval++; + const geist_token_t ti = (geist_token_t) ids[i]; + if (geist_session_prefill_tokens(ref, 1, &ti) != GEIST_OK || + geist_session_prefill_tokens(mode, 1, &ti) != GEIST_OK) + return false; } + *acc_out = n_eval > 0 ? (double) correct / (double) n_eval : 0.0; + *kl_out = n_eval > 0 ? kl_sum / (double) n_eval : 0.0; + return true; +} - /* Chunk because verify_forward caps at m_max (64). kv_truncate after - * each chunk forces a KIVI residual drain so the 2-bit path is - * exercised; harmless for the other modes. */ - const size_t M_MAX = 64; - size_t kv_len_acc = 0; - for (size_t off = 0; off < k; off += M_MAX) { - const size_t chunk = (k - off > M_MAX) ? M_MAX : (k - off); - enum geist_status vs = ops->verify_forward( - arch_meta, chunk, (const geist_token_t *) ids + off, preds + off); - if (vs != GEIST_OK) { - fprintf(stderr, "verify_forward(@%zu): %s\n", off, geist_status_to_string(vs)); - free(preds); - geist_model_destroy(model); - return -1.0; - } - kv_len_acc += chunk; - ops->kv_truncate(arch_meta, kv_len_acc); +/* Tokenize via the sp_bpe tokenizer, else the GGUF-embedded one. */ +static bool tokenize(struct geist_model *model, + const char *text, + uint32_t **ids_out, + size_t *n_out, + bool *had_tokenizer) { + *had_tokenizer = false; + struct sp_bpe_tokenizer *tok = geist_model_internal_tokenizer(model); + if (tok != nullptr) { + *had_tokenizer = true; + return sp_bpe_tokenizer_encode(tok, text, ids_out, n_out); } - - size_t n_correct = 0; - for (size_t i = 0; i < k - 1; i++) { - if ((uint32_t) preds[i] == ids[i + 1]) - n_correct++; + struct gguf_tokenizer *gtok = geist_model_internal_gguf_tokenizer(model); + if (gtok == nullptr) + return false; + *had_tokenizer = true; + const size_t cap = strlen(text) + 16; + int32_t *tmp = (int32_t *) malloc(cap * sizeof(int32_t)); + bool ok = false; + if (tmp != nullptr && gguf_tokenizer_encode(gtok, text, tmp, cap, n_out)) { + *ids_out = (uint32_t *) malloc(*n_out * sizeof(uint32_t)); + if (*ids_out != nullptr) { + for (size_t i = 0; i < *n_out; i++) + (*ids_out)[i] = (uint32_t) tmp[i]; + ok = true; + } } - const size_t n_eval = k - 1; - free(preds); - geist_model_destroy(model); - if (n_eval_out != nullptr) - *n_eval_out = n_eval; - return n_eval > 0 ? (double) n_correct / (double) n_eval : 0.0; + free(tmp); + return ok; } int main(int argc, char **argv) { const char *model_path = argc > 1 ? argv[1] : geist_test_find_gguf(); GEIST_SKIP_IF(model_path == nullptr, "no GGUF model found — pass path or set GEIST_GGUF_PATH"); - const char *text = (argc > 2) ? argv[2] : DEFAULT_TEXT; struct geist_backend *be = nullptr; enum geist_status s = geist_backend_create("cpu_neon", nullptr, nullptr, &be); - if (s != GEIST_OK) { + if (s != GEIST_OK) s = geist_backend_create("cpu_scalar", nullptr, nullptr, &be); - } if (s != GEIST_OK) { fprintf(stderr, "backend create: %s\n", geist_last_create_error()); return GEIST_TEST_ERROR; } - /* Tokenize once (mode-independent) via a throwaway load. */ struct geist_model *model = nullptr; s = geist_model_load(model_path, be, &model); if (s != GEIST_OK) { @@ -199,97 +248,85 @@ int main(int argc, char **argv) { geist_backend_destroy(be); return GEIST_TEST_FAIL; } - /* Try the sp_bpe tokenizer (external tokenizer.bin models) first, then - * the GGUF-embedded tokenizer (Gemma/BitNet unigram, Qwen bpe). */ - uint32_t *ids = nullptr; - size_t n_ids = 0; - bool enc_ok = false; - bool had_tokenizer = false; - struct sp_bpe_tokenizer *tok = geist_model_internal_tokenizer(model); - if (tok != nullptr) { - had_tokenizer = true; - enc_ok = sp_bpe_tokenizer_encode(tok, text, &ids, &n_ids); - } else { - struct gguf_tokenizer *gtok = geist_model_internal_gguf_tokenizer(model); - if (gtok != nullptr) { - had_tokenizer = true; - const size_t cap = strlen(text) + 16; - int32_t *tmp = (int32_t *) malloc(cap * sizeof(int32_t)); - if (tmp != nullptr && gguf_tokenizer_encode(gtok, text, tmp, cap, &n_ids)) { - ids = (uint32_t *) malloc(n_ids * sizeof(uint32_t)); - if (ids != nullptr) { - for (size_t i = 0; i < n_ids; i++) - ids[i] = (uint32_t) tmp[i]; - enc_ok = true; - } - } - free(tmp); - } - } - geist_model_destroy(model); /* done with this load; sweep reloads per mode */ + + uint32_t *ids = nullptr; + size_t n_ids = 0; + bool had_tokenizer = false; + const bool enc_ok = tokenize(model, text, &ids, &n_ids, &had_tokenizer); if (!had_tokenizer) { free(ids); + geist_model_destroy(model); geist_backend_destroy(be); GEIST_SKIP_IF(true, "model carries no usable tokenizer"); } if (!enc_ok || n_ids < 2) { fprintf(stderr, "tokenizer encode failed or text too short (n_ids=%zu)\n", n_ids); free(ids); + geist_model_destroy(model); + geist_backend_destroy(be); + return GEIST_TEST_FAIL; + } + + /* FP32 reference session, created once with FP32 env, reused for every + * mode's KL. Greedy opts (all-zero) → AUTO mode, env resolves the rest. */ + struct geist_session_opts opts = {0}; + apply_mode_env(&MODES[mode_index("FP32")]); + struct geist_session *ref = nullptr; + if (geist_session_create(model, be, &opts, &ref) != GEIST_OK) { + fprintf(stderr, "ref session create failed\n"); + free(ids); + geist_model_destroy(model); geist_backend_destroy(be); return GEIST_TEST_FAIL; } printf("model: %s\n", model_path); printf("backend: %s\n", geist_backend_name(be)); - printf("n_tokens: %zu (n_eval=%zu)\n\n", n_ids, n_ids - 2); + printf("n_tokens: %zu (n_eval=%zu)\n\n", n_ids, n_ids - 1); const size_t n_modes = sizeof(MODES) / sizeof(MODES[0]); double acc[sizeof(MODES) / sizeof(MODES[0])]; - printf("%-10s %s\n", "kv_mode", "top-1 acc"); - printf("%-10s %s\n", "-------", "---------"); + double kl[sizeof(MODES) / sizeof(MODES[0])]; + printf("%-10s %-9s %s\n", "kv_mode", "top-1", "mean KL(fp32||·) [nats]"); + printf("%-10s %-9s %s\n", "-------", "-----", "-----------------------"); for (size_t m = 0; m < n_modes; m++) { - set_or_unset("GEIST_KV_INT8", MODES[m].int8); - set_or_unset("GEIST_KV_KIVI", MODES[m].kivi); - set_or_unset("GEIST_KV_ROT", MODES[m].rot); - set_or_unset("GEIST_KV_INT4", MODES[m].int4); - set_or_unset("GEIST_KV_QBITS", MODES[m].qbits); - acc[m] = run_one(model_path, be, ids, n_ids, nullptr); - if (acc[m] < 0.0) { + apply_mode_env(&MODES[m]); + struct geist_session *sess = nullptr; + if (geist_session_create(model, be, &opts, &sess) != GEIST_OK || + !score_mode(ref, sess, ids, n_ids, &acc[m], &kl[m])) { + fprintf(stderr, "mode %s failed\n", MODES[m].label); + geist_session_destroy(sess); + geist_session_destroy(ref); free(ids); + geist_model_destroy(model); geist_backend_destroy(be); return GEIST_TEST_FAIL; } - printf("%-10s %.4f\n", MODES[m].label, acc[m]); + geist_session_destroy(sess); + printf("%-10s %.4f %.5f\n", MODES[m].label, acc[m], kl[m]); } + geist_session_destroy(ref); - /* Issue #61 headline (option A): does rotation rescue the low-bit INT4 - * cache? Report the fraction of the INT4→INT8 quality drop that ROT - * recovers — INT8 is the lossless ceiling, INT4 the degraded floor. */ - const double a_int8 = acc[mode_index("INT8")]; - const double a_int4 = acc[mode_index("INT4")]; - const double a_int4rot = acc[mode_index("INT4+ROT")]; - const double gap = a_int8 - a_int4; /* low-bit penalty */ + /* Issue #61 headlines: rotation's effect where it matters (low bit), + * measured by KL reduction (lower KL = closer to FP32). */ + const double kl_i4 = kl[mode_index("INT4")]; + const double kl_i4r = kl[mode_index("INT4+ROT")]; + const double kl_i2 = kl[mode_index("INT2")]; + const double kl_i2r = kl[mode_index("INT2+ROT")]; + const double kl_kivi = kl[mode_index("KIVI")]; printf("\n"); - if (gap > 1e-6) { - const double recovered = (a_int4rot - a_int4) / gap; - printf("INT4→INT8 gap recovered by rotation: %.0f%% (INT4 %.4f → +ROT %.4f → INT8 %.4f)\n", - 100.0 * recovered, - a_int4, - a_int4rot, - a_int8); - } else { - printf("INT4 already within noise of INT8 (gap=%.4f) — no low-bit penalty to recover.\n", - gap); - } - - /* The lazier-KIVI question: can symmetric 2-bit + rotation rival KIVI's - * asymmetric per-channel 2-bit? */ - printf("2-bit: symmetric INT2 %.4f → +ROT %.4f vs KIVI(asym) %.4f\n", - acc[mode_index("INT2")], - acc[mode_index("INT2+ROT")], - acc[mode_index("KIVI")]); + printf("4-bit: rotation cuts KL %.5f → %.5f (%+.0f%%)\n", + kl_i4, + kl_i4r, + kl_i4 > 0.0 ? 100.0 * (kl_i4r - kl_i4) / kl_i4 : 0.0); + printf("2-bit: rotation cuts KL %.5f → %.5f (%+.0f%%) vs KIVI(asym) %.5f\n", + kl_i2, + kl_i2r, + kl_i2 > 0.0 ? 100.0 * (kl_i2r - kl_i2) / kl_i2 : 0.0, + kl_kivi); free(ids); + geist_model_destroy(model); geist_backend_destroy(be); return GEIST_TEST_PASS; } From 4e69addbcb818f072e5369f6abb980d025b1369c Mon Sep 17 00:00:00 2001 From: germar Date: Sun, 5 Jul 2026 10:28:36 +0200 Subject: [PATCH 05/10] test(kv): _POSIX_C_SOURCE for setenv on glibc (Pi5) (#61) --- tests/bench_kv_quality.c | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/bench_kv_quality.c b/tests/bench_kv_quality.c index 8b86582..8de6f57 100644 --- a/tests/bench_kv_quality.c +++ b/tests/bench_kv_quality.c @@ -24,6 +24,12 @@ #define GEIST_INTERNAL_ENGINE_LAYER #define GEIST_INTERNAL_ARCH_LAYER +/* setenv/unsetenv need a POSIX feature macro on glibc (Pi5); no-op on macOS. + * Must precede any system header. */ +#ifndef _POSIX_C_SOURCE +#define _POSIX_C_SOURCE 200809L +#endif + #include "test_helpers.h" #include "src/engine/gguf_tokenizer.h" From df407179ad0c87a81b54a64ae6fbfc824a3a6630 Mon Sep 17 00:00:00 2001 From: germar Date: Sun, 5 Jul 2026 11:00:43 +0200 Subject: [PATCH 06/10] feat(kv): packed symmetric 4-bit KV cache (GEIST_KV_INT4) (#61) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Real packed INT4: two 4-bit values per byte, halving the K/V data buffers vs INT8, with per-token per-head scale and the optional Hadamard rotation (GEIST_KV_ROT). Rides the INT8 storage path — same buffer slots (allocated half-size), scale buffers, and ctx wiring — with a packing append branch and an unpacking attention kernel (attention_int4_via_buffers). New public mode GEIST_KV_INT4; GEIST_KV_QBITS keeps the N-bit quality-sim for probes. Validation: packed KL is bit-identical to the INT8-container sim (BitNet INT4 0.01275, +ROT 0.00861), so pack/unpack is correct; INT8 + KIVI unregressed. int4_kv.h pack/unpack has a unit test (sign-extension). --- include/geist.h | 42 +++++---- src/archs/transformer/arch_state.c | 30 ++++--- src/archs/transformer/arch_state.h | 8 +- src/archs/transformer/forward/attention.c | 105 ++++++++++++++++++++++ src/archs/transformer/forward/internal.h | 16 ++++ src/archs/transformer/forward/kv_store.c | 82 +++++++++++++++++ src/backends/common/int4_kv.h | 32 +++++++ tests/test_int4_kv_unit.c | 93 +++++++++++++++++++ 8 files changed, 378 insertions(+), 30 deletions(-) create mode 100644 src/backends/common/int4_kv.h create mode 100644 tests/test_int4_kv_unit.c diff --git a/include/geist.h b/include/geist.h index 73dc3e6..409dd02 100644 --- a/include/geist.h +++ b/include/geist.h @@ -47,19 +47,19 @@ enum geist_status { GEIST_OK = 0, /* Generic */ - GEIST_E_OOM, /* allocation failed */ - GEIST_E_INVALID_ARG, /* nullptr where not allowed, bad enum, etc. */ - GEIST_E_INTERNAL, /* programmer error, shouldn't happen */ + GEIST_E_OOM, /* allocation failed */ + GEIST_E_INVALID_ARG, /* nullptr where not allowed, bad enum, etc. */ + GEIST_E_INTERNAL, /* programmer error, shouldn't happen */ /* I/O */ GEIST_E_FILE_NOT_FOUND, - GEIST_E_IO, /* read/write failed */ - GEIST_E_FORMAT, /* corrupt file, wrong magic, etc. */ + GEIST_E_IO, /* read/write failed */ + GEIST_E_FORMAT, /* corrupt file, wrong magic, etc. */ /* Capability */ - GEIST_E_UNSUPPORTED, /* backend cannot run this op/dtype/layout */ - GEIST_E_NOT_FOUND, /* tensor name not in model, etc. */ - GEIST_E_BACKEND, /* backend-specific failure */ + GEIST_E_UNSUPPORTED, /* backend cannot run this op/dtype/layout */ + GEIST_E_NOT_FOUND, /* tensor name not in model, etc. */ + GEIST_E_BACKEND, /* backend-specific failure */ /* Lifecycle */ GEIST_E_INVALID_STATE, /* op called in wrong order */ @@ -91,9 +91,9 @@ enum geist_log_level { /* @stability EXPERIMENTAL — categories and call frequency may evolve. */ typedef void (*geist_log_callback_t)(enum geist_log_level level, - const char *category, - const char *message, - void *user_data); + const char *category, + const char *message, + void *user_data); /* ====================================================================== */ /* Memory / Allocator */ @@ -102,8 +102,8 @@ typedef void (*geist_log_callback_t)(enum geist_log_level level, /* @stability STABLE since 0.1.0 */ struct geist_allocator { void *(*alloc)(void *ctx, size_t bytes, size_t alignment); - void (*free)(void *ctx, void *ptr); - void (*free_all)(void *ctx); /* optional, arena-style; nullptr for malloc-based */ + void (*free)(void *ctx, void *ptr); + void (*free_all)(void *ctx); /* optional, arena-style; nullptr for malloc-based */ void *ctx; }; @@ -132,10 +132,10 @@ struct geist_backend_opts { * Create a backend by name (e.g. "cpu_neon", "cpu_scalar", "auto"). The * special name "auto" picks the best linked backend for the host. Pass * nullptr opts/alloc for defaults. */ -enum geist_status geist_backend_create(const char *name, +enum geist_status geist_backend_create(const char *name, const struct geist_backend_opts *opts, - const struct geist_allocator *alloc, - struct geist_backend **out); + const struct geist_allocator *alloc, + struct geist_backend **out); void geist_backend_destroy(struct geist_backend *be); const char *geist_backend_name(const struct geist_backend *be); @@ -152,9 +152,8 @@ struct geist_model; * Loads a GGUF model file. Architecture is detected from the GGUF * `general.architecture` metadata key; returns GEIST_E_UNSUPPORTED if * no architecture matching this build's compiled set is registered. */ -enum geist_status geist_model_load(const char *path, - struct geist_backend *be, - struct geist_model **out); +enum geist_status +geist_model_load(const char *path, struct geist_backend *be, struct geist_model **out); /* @stability STABLE since 0.2.1 * Load a GGUF that is already in memory — e.g. embedded in the executable, so @@ -196,6 +195,11 @@ enum geist_kv_mode { GEIST_KV_INT8 = 2, GEIST_KV_KIVI = 3, GEIST_KV_F16 = 4, + /* Packed symmetric 4-bit KV cache (2 values/byte, per-token per-head + * scale). Half the INT8 footprint, near-lossless — especially with the + * Hadamard rotation (GEIST_KV_ROT=1, issue #61). No per-channel/group + * bookkeeping (unlike KIVI). Env: GEIST_KV_INT4=1. */ + GEIST_KV_INT4 = 5, }; struct geist_session_opts { diff --git a/src/archs/transformer/arch_state.c b/src/archs/transformer/arch_state.c index 9023601..6629ecd 100644 --- a/src/archs/transformer/arch_state.c +++ b/src/archs/transformer/arch_state.c @@ -306,11 +306,15 @@ alloc_pool_buffer(struct transformer_arch_state *st, size_t bytes, struct geist_ return s; } } else if (st->sess->kv_int8_enabled) { - s = alloc_scratch(be, n_elems * sizeof(int8_t), &st->sess->k_cache_q8[li]); + /* Packed INT4 halves the data buffers (2 values/byte); scales are + * unchanged. hd is a power of two (128/256/512) so n_elems is even. */ + const size_t data_bytes = + st->sess->kv_int4_packed_enabled ? n_elems / 2 : n_elems * sizeof(int8_t); + s = alloc_scratch(be, data_bytes, &st->sess->k_cache_q8[li]); if (s != GEIST_OK) { return s; } - s = alloc_scratch(be, n_elems * sizeof(int8_t), &st->sess->v_cache_q8[li]); + s = alloc_scratch(be, data_bytes, &st->sess->v_cache_q8[li]); if (s != GEIST_OK) { return s; } @@ -892,6 +896,10 @@ void transformer_state_destroy(struct transformer_arch_state *st) { } const char *env_kivi = getenv("GEIST_KV_KIVI"); const char *env_int8 = getenv("GEIST_KV_INT8"); + const char *env_int4 = getenv("GEIST_KV_INT4"); + if (env_int4 != nullptr && env_int4[0] == '1') { + return GEIST_KV_INT4; + } if (env_kivi != nullptr && env_kivi[0] == '1') { return GEIST_KV_KIVI; } @@ -998,18 +1006,20 @@ struct transformer_arch_session *transformer_session_alloc(struct transformer_ar const enum geist_kv_mode mode = resolve_kv_mode(opts); sess->kv_kivi_enabled = (mode == GEIST_KV_KIVI); sess->kv_int8_enabled = (mode == GEIST_KV_INT8); + /* Issue #61: packed 4-bit KV rides the INT8 storage path (buffer alloc + + * ctx wiring), with half-size data buffers holding 2 values/byte. */ + sess->kv_int4_packed_enabled = (mode == GEIST_KV_INT4); + if (sess->kv_int4_packed_enabled) { + sess->kv_int8_enabled = true; + } /* Issue #61: low-bit quality-sim reuses the INT8 storage path with an - * N-bit quant grid. GEIST_KV_INT4=1 → 4 bits; GEIST_KV_QBITS=N (2..8) - * overrides. Any sim (2..7) forces INT8 storage on regardless of the - * resolved mode. Resolve before the rot flag so rotation sees it. */ + * N-bit quant grid (no packing, no memory win). GEIST_KV_QBITS=N (2..8) + * forces INT8 storage on. Resolve before the rot flag so rotation sees + * it. Ignored under the real packed-INT4 mode. */ { int qbits = 0; - const char *env_int4 = getenv("GEIST_KV_INT4"); const char *env_qbits = getenv("GEIST_KV_QBITS"); - if (env_int4 != nullptr && env_int4[0] == '1') { - qbits = 4; - } - if (env_qbits != nullptr) { + if (!sess->kv_int4_packed_enabled && env_qbits != nullptr) { const int q = atoi(env_qbits); if (q >= 2 && q <= 8) { qbits = (q == 8) ? 0 : q; /* 8-bit is the native path */ diff --git a/src/archs/transformer/arch_state.h b/src/archs/transformer/arch_state.h index e52fd86..cf5e3bd 100644 --- a/src/archs/transformer/arch_state.h +++ b/src/archs/transformer/arch_state.h @@ -184,8 +184,14 @@ struct transformer_arch_session { * low-bit cache that reuses the INT8 storage + kernel (no packing, no * memory win yet). Measures whether rotation rescues low-bit quality. * 0 = native 8-bit; 2..7 forces the INT8 storage path on. Env: - * GEIST_KV_INT4=1 (→4) or GEIST_KV_QBITS=N. */ + * GEIST_KV_QBITS=N. */ int kv_sim_qbits; + /* Packed symmetric 4-bit KV cache (issue #61). Rides the INT8 storage + * path (kv_int8_enabled is also set for buffer alloc + ctx wiring) but + * the k/v data buffers are allocated half-size and hold two 4-bit values + * per byte; append packs, attention unpacks. Half the INT8 KV footprint. + * Env: GEIST_KV_INT4=1. */ + bool kv_int4_packed_enabled; /* F16 KV cache: k_cache[]/v_cache[] hold half floats (2 bytes/elem); * appends convert through the backend's kv_append_f16 slot and * attention reads F16 views. Only set when that slot is non-null. */ diff --git a/src/archs/transformer/forward/attention.c b/src/archs/transformer/forward/attention.c index af59c84..0269d5f 100644 --- a/src/archs/transformer/forward/attention.c +++ b/src/archs/transformer/forward/attention.c @@ -12,6 +12,7 @@ #include "internal.h" #include "../arch_state.h" +#include "int4_kv.h" #include "kivi.h" #include @@ -323,3 +324,107 @@ void attention_int8_via_buffers(const float *q, } } } + +/* Packed-INT4 attention. Identical to attention_int8_via_buffers except each + * K/V cache row is unpacked from head_dim/2 bytes into a stack int8 row + * before the (reused) int8 dot / weighted-sum. See internal.h. */ +void attention_int4_via_buffers(const float *q, + size_t n_q, + size_t n_q_heads, + size_t head_dim, + const uint8_t *k_q4, + const float *k_scale, + const uint8_t *v_q4, + const float *v_scale, + size_t n_kv, + size_t n_kv_heads, + size_t q_offset, + size_t sliding_window, + float *out) { + + const size_t kv_group_size = n_q_heads / n_kv_heads; + const size_t packed = head_dim / 2; /* bytes per cache row */ +#if defined(_OPENMP) +#pragma omp parallel for schedule(dynamic) +#endif + for (size_t t = 0; t < n_q; t++) { + const size_t q_pos = q_offset + t; + const size_t s_lo = + (sliding_window > 0 && q_pos + 1 > sliding_window) ? q_pos + 1 - sliding_window : 0; + const size_t s_hi = q_pos < n_kv ? q_pos : n_kv - 1; + float scores[n_kv]; + + for (size_t h = 0; h < n_q_heads; h++) { + const size_t kv_h = h / kv_group_size; + const float *qv = q + (t * n_q_heads + h) * head_dim; + + int8_t q_q8[512]; + float amax = 0.0f; + for (size_t i = 0; i < head_dim; i++) { + float a = fabsf(qv[i]); + if (a > amax) { + amax = a; + } + } + float scale_q = amax / 127.0f; + if (scale_q == 0.0f) { + scale_q = 1.0f; + } + const float inv_q = 1.0f / scale_q; + for (size_t i = 0; i < head_dim; i++) { + q_q8[i] = (int8_t) lrintf(qv[i] * inv_q); + } + + for (size_t s = s_lo; s <= s_hi; s++) { + int8_t k[512]; + int4_unpack_row(k_q4 + (s * n_kv_heads + kv_h) * packed, k, head_dim); + const float ks = k_scale[s * n_kv_heads + kv_h]; + int32_t int_dot = 0; +#if defined(__ARM_NEON) + int32x4_t acc = vdupq_n_s32(0); + size_t i = 0; + for (; i + 16 <= head_dim; i += 16) { + acc = vdotq_s32(acc, vld1q_s8(q_q8 + i), vld1q_s8(k + i)); + } + int_dot = vaddvq_s32(acc); + for (; i < head_dim; i++) { + int_dot += (int32_t) q_q8[i] * (int32_t) k[i]; + } +#else + for (size_t i = 0; i < head_dim; i++) { + int_dot += (int32_t) q_q8[i] * (int32_t) k[i]; + } +#endif + scores[s] = (float) int_dot * scale_q * ks; + } + + float max_score = scores[s_lo]; + for (size_t s = s_lo + 1; s <= s_hi; s++) { + if (scores[s] > max_score) { + max_score = scores[s]; + } + } + double sum_exp = 0.0; + for (size_t s = s_lo; s <= s_hi; s++) { + float e = expf(scores[s] - max_score); + scores[s] = e; + sum_exp += e; + } + const float inv_sum = (float) (1.0 / sum_exp); + + float *outv = out + (t * n_q_heads + h) * head_dim; + for (size_t i = 0; i < head_dim; i++) { + outv[i] = 0.0f; + } + for (size_t s = s_lo; s <= s_hi; s++) { + int8_t vv[512]; + int4_unpack_row(v_q4 + (s * n_kv_heads + kv_h) * packed, vv, head_dim); + const float vs = v_scale[s * n_kv_heads + kv_h]; + const float wvs = scores[s] * inv_sum * vs; + for (size_t i = 0; i < head_dim; i++) { + outv[i] += wvs * (float) vv[i]; + } + } + } + } +} diff --git a/src/archs/transformer/forward/internal.h b/src/archs/transformer/forward/internal.h index a8fc4f2..7a99624 100644 --- a/src/archs/transformer/forward/internal.h +++ b/src/archs/transformer/forward/internal.h @@ -260,6 +260,22 @@ void attention_int8_via_buffers(const float *q, size_t sliding_window, float *out); +/* Packed-INT4 variant (issue #61): k_q4/v_q4 hold two 4-bit values per byte + * (head_dim/2 bytes per row); otherwise identical to the INT8 kernel. */ +void attention_int4_via_buffers(const float *q, + size_t n_q, + size_t n_q_heads, + size_t head_dim, + const uint8_t *k_q4, + const float *k_scale, + const uint8_t *v_q4, + const float *v_scale, + size_t n_kv, + size_t n_kv_heads, + size_t q_offset, + size_t sliding_window, + float *out); + /* forward/layer_attn.c */ [[nodiscard]] enum geist_status transformer_layer_run_attention_block(struct transformer_layer_forward_ctx *ctx); diff --git a/src/archs/transformer/forward/kv_store.c b/src/archs/transformer/forward/kv_store.c index cc41063..255239b 100644 --- a/src/archs/transformer/forward/kv_store.c +++ b/src/archs/transformer/forward/kv_store.c @@ -8,6 +8,7 @@ #include #include "fwht.h" +#include "int4_kv.h" #include "kivi.h" #include @@ -99,6 +100,49 @@ enum geist_status transformer_kv_store_append(struct transformer_layer_forward_c } v->buffer_unmap(ctx->k_residual_buf); v->buffer_unmap(ctx->v_residual_buf); + } else if (st->sess->kv_int4_packed_enabled) { + /* Packed 4-bit: 2 values/byte into the half-size int8 slots. Same + * per-token per-head scale + optional rotation as INT8. denom 7 → + * scale = amax/7, values in [-7,7]. */ + uint8_t *k_dst = (uint8_t *) v->buffer_map(ctx->k_cache_q8_buf); + uint8_t *v_dst = (uint8_t *) v->buffer_map(ctx->v_cache_q8_buf); + float *k_sca = (float *) v->buffer_map(ctx->k_cache_scale_buf); + float *v_sca = (float *) v->buffer_map(ctx->v_cache_scale_buf); + const size_t row_elems = kv_out; + const size_t scales_per_row = st->n_kv_heads; + const bool rot = st->sess->kv_rot_enabled && fwht_supported(hd) && hd <= 512; + float krot[512]; + float vrot[512]; + for (size_t t = 0; t < seq; t++) { + const size_t slot = q_position + t; + for (size_t h = 0; h < st->n_kv_heads; h++) { + const float *k_row = k_src + t * row_elems + h * hd; + const float *v_row = v_src + t * row_elems + h * hd; + if (rot) { + memcpy(krot, k_row, hd * sizeof(float)); + memcpy(vrot, v_row, hd * sizeof(float)); + fwht_orthonormal(krot, hd); + fwht_orthonormal(vrot, hd); + k_row = krot; + v_row = vrot; + } + float k_scale = kv_row_absmax(k_row, hd) / 7.0f; + float v_scale = kv_row_absmax(v_row, hd) / 7.0f; + if (k_scale == 0.0f) + k_scale = 1.0f; + if (v_scale == 0.0f) + v_scale = 1.0f; + const size_t byte_off = (slot * row_elems + h * hd) / 2; + int4_pack_row(k_row, 1.0f / k_scale, k_dst + byte_off, hd); + int4_pack_row(v_row, 1.0f / v_scale, v_dst + byte_off, hd); + k_sca[slot * scales_per_row + h] = k_scale; + v_sca[slot * scales_per_row + h] = v_scale; + } + } + v->buffer_unmap(ctx->k_cache_q8_buf); + v->buffer_unmap(ctx->v_cache_q8_buf); + v->buffer_unmap(ctx->k_cache_scale_buf); + v->buffer_unmap(ctx->v_cache_scale_buf); } else if (ctx->kv_int8_enabled) { int8_t *k_dst = (int8_t *) v->buffer_map(ctx->k_cache_q8_buf); int8_t *v_dst = (int8_t *) v->buffer_map(ctx->v_cache_q8_buf); @@ -232,6 +276,44 @@ enum geist_status transformer_kv_store_attention(struct transformer_layer_forwar v->buffer_unmap(ctx->k_residual_buf); v->buffer_unmap(ctx->v_residual_buf); v->buffer_unmap(st->sess->scratch_attn); + } else if (st->sess->kv_int4_packed_enabled) { + float *qp = (float *) v->buffer_map(st->sess->scratch_q); + const uint8_t *k_q4p = (const uint8_t *) v->buffer_map(ctx->k_cache_q8_buf); + const uint8_t *v_q4p = (const uint8_t *) v->buffer_map(ctx->v_cache_q8_buf); + const float *k_scalep = (const float *) v->buffer_map(ctx->k_cache_scale_buf); + const float *v_scalep = (const float *) v->buffer_map(ctx->v_cache_scale_buf); + float *outp = (float *) v->buffer_map(st->sess->scratch_attn); + const bool rot = st->sess->kv_rot_enabled && fwht_supported(ctx->hd) && ctx->hd <= 512; + const size_t n_rows = ctx->seq * st->n_q_heads; + if (rot) { + for (size_t r = 0; r < n_rows; r++) { + fwht_orthonormal(qp + r * ctx->hd, ctx->hd); + } + } + attention_int4_via_buffers(qp, + ctx->seq, + st->n_q_heads, + ctx->hd, + k_q4p, + k_scalep, + v_q4p, + v_scalep, + kv_len_now, + st->n_kv_heads, + ctx->q_position, + L->sliding_window, + outp); + if (rot) { + for (size_t r = 0; r < n_rows; r++) { + fwht_orthonormal(outp + r * ctx->hd, ctx->hd); + } + } + v->buffer_unmap(st->sess->scratch_q); + v->buffer_unmap(ctx->k_cache_q8_buf); + v->buffer_unmap(ctx->v_cache_q8_buf); + v->buffer_unmap(ctx->k_cache_scale_buf); + v->buffer_unmap(ctx->v_cache_scale_buf); + v->buffer_unmap(st->sess->scratch_attn); } else if (ctx->kv_int8_enabled) { float *qp = (float *) v->buffer_map(st->sess->scratch_q); const int8_t *k_q8p = (const int8_t *) v->buffer_map(ctx->k_cache_q8_buf); diff --git a/src/backends/common/int4_kv.h b/src/backends/common/int4_kv.h new file mode 100644 index 0000000..7fbe72b --- /dev/null +++ b/src/backends/common/int4_kv.h @@ -0,0 +1,32 @@ +/* + * int4_kv.h — symmetric signed 4-bit pack/unpack for the packed-INT4 KV + * cache (issue #61). Two values per byte: low nibble = even index, high + * nibble = odd index. Values quantize to [-7,7], stored as 4-bit two's + * complement. n must be even (head_dim is a power of two). + */ +#ifndef GEIST_INT4_KV_H +#define GEIST_INT4_KV_H + +#include +#include +#include + +/* Quantize `n` floats at scale `inv` (= 1/scale) into n/2 packed bytes. */ +static inline void int4_pack_row(const float *x, float inv, uint8_t *out, size_t n) { + for (size_t i = 0; i < n; i += 2) { + const int lo = (int) lrintf(x[i] * inv); + const int hi = (int) lrintf(x[i + 1] * inv); + out[i >> 1] = (uint8_t) ((lo & 0x0F) | ((hi & 0x0F) << 4)); + } +} + +/* Unpack n/2 bytes into `n` sign-extended int8 values in [-8,7]. */ +static inline void int4_unpack_row(const uint8_t *in, int8_t *out, size_t n) { + for (size_t j = 0; j < n / 2; j++) { + const uint8_t b = in[j]; + out[2 * j] = (int8_t) ((int8_t) (b << 4) >> 4); /* low nibble */ + out[2 * j + 1] = (int8_t) ((int8_t) (b & 0xF0) >> 4); /* high nibble */ + } +} + +#endif /* GEIST_INT4_KV_H */ diff --git a/tests/test_int4_kv_unit.c b/tests/test_int4_kv_unit.c new file mode 100644 index 0000000..ed69bbd --- /dev/null +++ b/tests/test_int4_kv_unit.c @@ -0,0 +1,93 @@ +/* + * test_int4_kv_unit — verifies the symmetric 4-bit pack/unpack used by the + * packed-INT4 KV cache (issue #61). The failure-prone part is nibble + * sign-extension: a wrong shift silently turns -1 (0xF) into +15. + * + * Deterministic — no model needed. + */ +#include "int4_kv.h" +#include "test_helpers.h" + +#include + +/* Scenario 1: every representable level round-trips exactly (inv=1 → q=x). */ +static int test_levels_exact(void) { + /* All 15 symmetric levels [-7,7], padded to an even count. */ + float x[16]; + for (int i = 0; i < 15; i++) + x[i] = (float) (i - 7); + x[15] = 0.0f; + uint8_t packed[8]; + int8_t out[16]; + int4_pack_row(x, 1.0f, packed, 16); + int4_unpack_row(packed, out, 16); + for (size_t i = 0; i < 16; i++) { + if (out[i] != (int8_t) lrintf(x[i])) { + printf("level %zu: got %d want %d (packed sign-extension wrong?)\n", + i, + out[i], + (int) lrintf(x[i])); + return 1; + } + } + return 0; +} + +/* Scenario 2: nibble ordering — low nibble is the even index, high the odd. */ +static int test_nibble_order(void) { + float x[2] = {3.0f, -5.0f}; + uint8_t b; + int4_pack_row(x, 1.0f, &b, 2); + /* low = 3 (0x3), high = -5 (0xB) → byte 0xB3 */ + if (b != 0xB3) { + printf("nibble order: byte=0x%02X want 0xB3\n", b); + return 1; + } + int8_t out[2]; + int4_unpack_row(&b, out, 2); + if (out[0] != 3 || out[1] != -5) { + printf("nibble order unpack: %d,%d want 3,-5\n", out[0], out[1]); + return 1; + } + return 0; +} + +/* Scenario 3: scaled round-trip — dequant error stays within half a step. */ +static int test_scaled_roundtrip(void) { + enum { N = 256 }; + float x[N]; + uint32_t seed = 0x1234u; + float amax = 0.0f; + for (size_t i = 0; i < N; i++) { + seed = seed * 1103515245u + 12345u; + x[i] = ((float) (seed >> 8) / (float) 0xFFFFFF - 0.5f) * 4.0f; /* ~[-2,2] */ + if (fabsf(x[i]) > amax) + amax = fabsf(x[i]); + } + const float scale = amax / 7.0f; + uint8_t packed[N / 2]; + int8_t q[N]; + int4_pack_row(x, 1.0f / scale, packed, N); + int4_unpack_row(packed, q, N); + for (size_t i = 0; i < N; i++) { + const float deq = (float) q[i] * scale; + if (fabsf(deq - x[i]) > 0.5f * scale + 1e-4f) { + printf("scaled[%zu]: x=%.4f deq=%.4f err=%.4f > %.4f\n", + i, + x[i], + deq, + fabsf(deq - x[i]), + 0.5f * scale); + return 1; + } + } + return 0; +} + +int main(void) { + if (test_levels_exact() == 0 && test_nibble_order() == 0 && test_scaled_roundtrip() == 0) { + printf("PASS: int4 pack/unpack round-trips with correct sign-extension\n"); + return GEIST_TEST_PASS; + } + return GEIST_TEST_FAIL; +} From 5de9c890341dfc5d7f9657a059977d506809ef86 Mon Sep 17 00:00:00 2001 From: germar Date: Sun, 5 Jul 2026 11:01:57 +0200 Subject: [PATCH 07/10] fix(kv): zero-init int4 unpack buffers for GCC -Werror on Pi5 (#61) --- src/archs/transformer/forward/attention.c | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/archs/transformer/forward/attention.c b/src/archs/transformer/forward/attention.c index 0269d5f..3f9414a 100644 --- a/src/archs/transformer/forward/attention.c +++ b/src/archs/transformer/forward/attention.c @@ -376,7 +376,7 @@ void attention_int4_via_buffers(const float *q, } for (size_t s = s_lo; s <= s_hi; s++) { - int8_t k[512]; + int8_t k[512] = {0}; /* zero-init: silence GCC maybe-uninitialized on the NEON tail */ int4_unpack_row(k_q4 + (s * n_kv_heads + kv_h) * packed, k, head_dim); const float ks = k_scale[s * n_kv_heads + kv_h]; int32_t int_dot = 0; @@ -417,7 +417,7 @@ void attention_int4_via_buffers(const float *q, outv[i] = 0.0f; } for (size_t s = s_lo; s <= s_hi; s++) { - int8_t vv[512]; + int8_t vv[512] = {0}; int4_unpack_row(v_q4 + (s * n_kv_heads + kv_h) * packed, vv, head_dim); const float vs = v_scale[s * n_kv_heads + kv_h]; const float wvs = scores[s] * inv_sum * vs; From 6f50fbd9ab8857cc9e16c7e7b3389a527a1395bc Mon Sep 17 00:00:00 2001 From: germar Date: Sun, 5 Jul 2026 11:07:36 +0200 Subject: [PATCH 08/10] perf(kv): suppress GCC false positive instead of zero-init int4 buffers (#61) --- src/archs/transformer/forward/attention.c | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/src/archs/transformer/forward/attention.c b/src/archs/transformer/forward/attention.c index 3f9414a..3420ac7 100644 --- a/src/archs/transformer/forward/attention.c +++ b/src/archs/transformer/forward/attention.c @@ -327,7 +327,15 @@ void attention_int8_via_buffers(const float *q, /* Packed-INT4 attention. Identical to attention_int8_via_buffers except each * K/V cache row is unpacked from head_dim/2 bytes into a stack int8 row - * before the (reused) int8 dot / weighted-sum. See internal.h. */ + * before the (reused) int8 dot / weighted-sum. See internal.h. + * + * int4_unpack_row fully writes [0,head_dim); the NEON tail reads only that + * range, so GCC's -Wmaybe-uninitialized on the unpack buffers is a false + * positive — suppressed here rather than paid for with a per-row zero-init. */ +#if defined(__GNUC__) && !defined(__clang__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wmaybe-uninitialized" +#endif void attention_int4_via_buffers(const float *q, size_t n_q, size_t n_q_heads, @@ -376,7 +384,7 @@ void attention_int4_via_buffers(const float *q, } for (size_t s = s_lo; s <= s_hi; s++) { - int8_t k[512] = {0}; /* zero-init: silence GCC maybe-uninitialized on the NEON tail */ + int8_t k[512]; int4_unpack_row(k_q4 + (s * n_kv_heads + kv_h) * packed, k, head_dim); const float ks = k_scale[s * n_kv_heads + kv_h]; int32_t int_dot = 0; @@ -417,7 +425,7 @@ void attention_int4_via_buffers(const float *q, outv[i] = 0.0f; } for (size_t s = s_lo; s <= s_hi; s++) { - int8_t vv[512] = {0}; + int8_t vv[512]; int4_unpack_row(v_q4 + (s * n_kv_heads + kv_h) * packed, vv, head_dim); const float vs = v_scale[s * n_kv_heads + kv_h]; const float wvs = scores[s] * inv_sum * vs; @@ -428,3 +436,6 @@ void attention_int4_via_buffers(const float *q, } } } +#if defined(__GNUC__) && !defined(__clang__) +#pragma GCC diagnostic pop +#endif From 80e65453f3e561b745fe78947fcdd7b4ea8aa64e Mon Sep 17 00:00:00 2001 From: germar Date: Sun, 5 Jul 2026 11:11:57 +0200 Subject: [PATCH 09/10] docs: note packed INT4 KV mode + rotation in ARCHITECTURE (#61) --- docs/ARCHITECTURE.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 150ebd4..c592e81 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -163,7 +163,9 @@ keep the score matrix in registers/L1. A `geist_model` is immutable, shared, read-only weights. A `geist_session` owns the mutable per-conversation state: KV cache, pending logits, sampler config, stats. Multiple sessions can share one model. The KV cache supports quantized -modes (`INT8`, `KIVI`) and prefix pinning (`geist_session_pin_prefix`) to +modes (`INT8`, packed `INT4` — half the INT8 footprint, near-lossless with the +optional Hadamard rotation `GEIST_KV_ROT`; and `KIVI` 2-bit) and prefix pinning +(`geist_session_pin_prefix`) to amortize a constant system prompt across chat turns. Speculative decode drafts via an n-gram lookup over history and verifies in one batched forward. From 92309c583f08cbd73aa87a235b4a26267b10efc9 Mon Sep 17 00:00:00 2001 From: germar Date: Sun, 5 Jul 2026 11:36:55 +0200 Subject: [PATCH 10/10] style(kv): drop unrelated clang-format churn in geist.h (#61) Keep the geist.h diff to just the GEIST_KV_INT4 enum value; revert the whole-file reformat a stray clang-format run introduced. --- include/geist.h | 37 +++++++++++++++++++------------------ 1 file changed, 19 insertions(+), 18 deletions(-) diff --git a/include/geist.h b/include/geist.h index 409dd02..b4cc16b 100644 --- a/include/geist.h +++ b/include/geist.h @@ -47,19 +47,19 @@ enum geist_status { GEIST_OK = 0, /* Generic */ - GEIST_E_OOM, /* allocation failed */ - GEIST_E_INVALID_ARG, /* nullptr where not allowed, bad enum, etc. */ - GEIST_E_INTERNAL, /* programmer error, shouldn't happen */ + GEIST_E_OOM, /* allocation failed */ + GEIST_E_INVALID_ARG, /* nullptr where not allowed, bad enum, etc. */ + GEIST_E_INTERNAL, /* programmer error, shouldn't happen */ /* I/O */ GEIST_E_FILE_NOT_FOUND, - GEIST_E_IO, /* read/write failed */ - GEIST_E_FORMAT, /* corrupt file, wrong magic, etc. */ + GEIST_E_IO, /* read/write failed */ + GEIST_E_FORMAT, /* corrupt file, wrong magic, etc. */ /* Capability */ - GEIST_E_UNSUPPORTED, /* backend cannot run this op/dtype/layout */ - GEIST_E_NOT_FOUND, /* tensor name not in model, etc. */ - GEIST_E_BACKEND, /* backend-specific failure */ + GEIST_E_UNSUPPORTED, /* backend cannot run this op/dtype/layout */ + GEIST_E_NOT_FOUND, /* tensor name not in model, etc. */ + GEIST_E_BACKEND, /* backend-specific failure */ /* Lifecycle */ GEIST_E_INVALID_STATE, /* op called in wrong order */ @@ -91,9 +91,9 @@ enum geist_log_level { /* @stability EXPERIMENTAL — categories and call frequency may evolve. */ typedef void (*geist_log_callback_t)(enum geist_log_level level, - const char *category, - const char *message, - void *user_data); + const char *category, + const char *message, + void *user_data); /* ====================================================================== */ /* Memory / Allocator */ @@ -102,8 +102,8 @@ typedef void (*geist_log_callback_t)(enum geist_log_level level, /* @stability STABLE since 0.1.0 */ struct geist_allocator { void *(*alloc)(void *ctx, size_t bytes, size_t alignment); - void (*free)(void *ctx, void *ptr); - void (*free_all)(void *ctx); /* optional, arena-style; nullptr for malloc-based */ + void (*free)(void *ctx, void *ptr); + void (*free_all)(void *ctx); /* optional, arena-style; nullptr for malloc-based */ void *ctx; }; @@ -132,10 +132,10 @@ struct geist_backend_opts { * Create a backend by name (e.g. "cpu_neon", "cpu_scalar", "auto"). The * special name "auto" picks the best linked backend for the host. Pass * nullptr opts/alloc for defaults. */ -enum geist_status geist_backend_create(const char *name, +enum geist_status geist_backend_create(const char *name, const struct geist_backend_opts *opts, - const struct geist_allocator *alloc, - struct geist_backend **out); + const struct geist_allocator *alloc, + struct geist_backend **out); void geist_backend_destroy(struct geist_backend *be); const char *geist_backend_name(const struct geist_backend *be); @@ -152,8 +152,9 @@ struct geist_model; * Loads a GGUF model file. Architecture is detected from the GGUF * `general.architecture` metadata key; returns GEIST_E_UNSUPPORTED if * no architecture matching this build's compiled set is registered. */ -enum geist_status -geist_model_load(const char *path, struct geist_backend *be, struct geist_model **out); +enum geist_status geist_model_load(const char *path, + struct geist_backend *be, + struct geist_model **out); /* @stability STABLE since 0.2.1 * Load a GGUF that is already in memory — e.g. embedded in the executable, so