feat(kv): activation rotation + packed 4-bit KV cache (#61)

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.
 

include/geist.h

@@ -196,6 +196,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 {

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 \

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,6 +1006,36 @@ 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 (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_qbits = getenv("GEIST_KV_QBITS");
+        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 */
+            }
+        }
+        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");
+        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

src/archs/transformer/arch_state.h

@@ -175,6 +175,23 @@ 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;
+    /* 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_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. */

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 <math.h>
@@ -323,3 +324,118 @@ 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.
+ *
+ * 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,
+                                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];
+                }
+            }
+        }
+    }
+}
+#if defined(__GNUC__) && !defined(__clang__)
+#pragma GCC diagnostic pop
+#endif

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);