# KAI-NATIVE LLM UPGRADE PLAN

**Purpose:** design an LLM that is *built around KAI's real strengths* — so it inherits KAI's
power-efficiency, speed, accuracy, and tool-use — instead of being a vanilla transformer that
merely reads KAI's memory.

**Status:** design doc only. No code changed. Read-only reconnaissance of the real repo.
Grounded against `C:\KAI\src\`, `C:\KAI\tools\oracle-discord\`, and the Codex (v9.10.163).

**Context that frames this:** we already validated the base recipe with 4 clean experiments —
*pretrain → finetune a small transformer beats every RSHL-in-training hybrid; RSHL belongs at
**inference** as RAG, not baked into the weights.* The Codex masthead confirms the project just
landed on the same conclusion independently (v9.10.163, "INFERENCE-TIME RAG — the correct use of
RSHL", read-only `src/bin/rag_*`). So the frozen conclusion is: **train a clean small LM; wrap it
in KAI's powers at inference.** This doc says which powers, and how.

---

## 0. TL;DR (read this if you read nothing else)

- **The headline asset for efficiency + speed is REAL and already in the repo:** a complete,
  hand-written **BitNet b1.58 (2B-4T) ternary transformer in Rust on Candle**, with a genuine
  **zero-multiply ternary matmul kernel**, KV-cache, and nucleus sampling. Ternary weights
  {−1,0,+1} = the same nature as RSHL. Today it is **inference-only on a frozen Microsoft GGUF**;
  it is not trainable as written.
- **Tool-use is the quick win.** KAI already has a working ~40-tool function-calling system
  (`gemini-live-bridge.mjs` / `native-tools.mjs`) with live dispatchers. Wiring a trained LM's
  generated tool-calls into those existing dispatchers is very doable and is the first thing to build.
- **Accuracy is already being built the right way:** ternary-cosine RAG over the lattice
  (`/api/rshl/query`) + a continuation-predictor for ranking. Inject retrieved memory text into the
  prompt at inference. No weight surgery needed.
- **The "ten functions" (KAI's cognition faculties) mostly belong AROUND the LLM as callable
  tools/engine modules — NOT baked into weights.** This is consistent with the 4 experiments.
- **Recommended build order:** (1) tool-calling on the current model → (2) RAG fusion into the
  prompt → (3) move the base to a *trained* ternary/BitNet model → (4) fuse cognition faculties as
  tools. Tool-calling is a quick win; ternary *training* is the real lift.

---

## 1. KAI's strengths inventory (grounded in real files)

### 1.1 The BitNet / ternary asset — the efficiency + speed answer (the key finding)

A full **LLaMA-style BitNet b1.58 decoder** exists natively in Rust, built on Candle. This is the
single most important thing in the repo for "power-efficiency and speed," and it already matches
RSHL's ternary nature.

| Piece | File | What it is |
|---|---|---|
| Model | `src/cognition/bitnet_llama.rs` | 30-layer decoder, hidden 2560, 20 heads / 5 KV heads (GQA), vocab 128256, RoPE θ=500000, SwiGLU MLP, BitNet sub-norms (`attn_sub_norm`, `ffn_sub_norm`), tied lm_head. Candle tensors. |
| BitLinear | `src/cognition/bitnet_inference.rs` | Ternary linear layer wrapper. |
| Weight loader | `src/cognition/bitnet_brain.rs` | Loads per-tensor `.kai` files (`models/BitNet/Native/`). |
| Ternary math | `src/cognition/ternary_math.rs` | I2_S pack/unpack (2 bits/val, 32-byte / 128-elem blocks), and a **real zero-multiply matmul** (`bitlinear_matmul_buf`) — add/subtract only, rayon-parallel. |
| Generation | `src/cognition/language_warehouse.rs` | `native_decode(prompt, max_tokens)` — encode → BOS → KV-cache decode loop → logit-bias pass → `sample()` (repeat-penalty 1.15, temp 0.7, top-p 0.9, greedy fallback). |
| Tokenizer | (HF `tokenizers` 0.23.1) | `models/BitNet/tokenizer.json`, llama-3 style, vocab 128256. |

**Quantization:** I2_S ternary. `MAP_2_BIT = [-1.0, 0.0, 1.0, 0.0]`, one f32 absmean-style scale per
weight matrix, hard-threshold encode at ±0.5·(w/scale). A genuine popcount/add-subtract kernel
exists and is unit-tested.

**Two honest caveats (verified in source):**
1. **The live decode path does NOT use the zero-multiply kernel.** `BitLinear::new` dequantizes
   each ternary matrix to full **f32** up front and runs a normal Candle matmul → ~10 GB RAM for the
   2B model. The fast packed kernel is built and tested but **not wired into the hot path.** The
   code comments already flag "keep weights PACKED and do a true BitLinear ternary matmul" as the
   unrealized endgame.
2. **It is inference-only.** No `Var`/`VarMap`/`VarBuilder`, no loss, no `.backward()`, no optimizer
   anywhere in the BitNet stack. Weights come from a **frozen** Microsoft `bitnet-b1.58-2B-4T.gguf`,
   extracted offline by Python into `.kai` files. The ternary quantizer is a hard threshold with **no
   straight-through estimator** and no full-precision shadow weights — i.e. none of the machinery
   quant-aware training needs. It is a clean inference engine, by design.

**How BitNet is used today:** (a) native Rust decoder, **default OFF** (`KAI_NATIVE_BRAIN=1` to
mount); (b) offline **distillation teacher** via `bitnet.cpp`'s `llama-cli.exe` — asks BitNet each
prompt, ingests (prompt→answer) pairs into the lattice (data-level distillation, no gradients);
(c) frozen **embedding encoder** for a small separate projection head (`neural_mapper.rs`). Deps:
`candle-core/nn/transformers = 0.10.2`, `tokenizers = 0.23.1`, `memmap2`, `rayon`.

### 1.2 Tool-use — a working ~40-tool function-calling system

| Piece | File | What it is |
|---|---|---|
| Declarations | `tools/oracle-discord/shared/gemini-live-bridge.mjs` (L222–611) | `LIVE_TOOL_DECLARATIONS` — ~40 tools in Gemini function-calling format (`{name, description, parameters}`). |
| Voice dispatch | same file (L1295–1349, 1685+) | `_handleToolCall` normalizes calls; a big `if/else if` on `fn.name` runs each; `sendToolResponse()` returns `{toolResponse:{functionResponses:[…]}}` to the model. |
| Text dispatch | `tools/oracle-discord/shared/native-tools.mjs` (L367) | `executeToolCall(name, args, bot, meta)` — a `switch(toolName)` with a per-bot permission gate `botCanUseTool`. |

Real, working tools include: `search_lattice` (RAG into the engine), `consult_codex`,
`recall_memory`, `search_web`, `ask_google`, `kai_status`, `get_directions`/`find_place`/`geocode`/
`nearby_search`/`get_weather` (live Google Maps/Weather), `calculate`, `simulate_emergence`
(SRHT math), `read_channel_feed`, `consult_oracle`, `request_code_change`, plus Codex search/paging.
Note: there are **two parallel dispatchers** (voice `if/else`, text `switch`) — a duplication to
unify. Tool-call traces are console-logged and there is JSONL transcript infra
(`logs/audit.json`, `kai-transcript.jsonl`) but **no purpose-built tool-call→result SFT log** yet.

### 1.3 Speed / efficiency kernels — sparse ternary VSA

`src/core/sparse_vec.rs` (1633 lines; note: under `src/core/`, not `src/`). 16384-dim sparse ternary
hypervectors {−1,0,+1}, ~4% density, `nz: Vec<u16>` + `vals: Vec<i8>` + cached √nnz norm. The hot
loop `dot()` is a two-pointer merge over sorted index arrays with bounds elision; `cosine()` =
dot/(‖a‖‖b‖). Bitpacked paths (`PackedMask` 2-bit, `DenseMask` pos/neg u64) do popcount cosine with
**runtime AVX-512 VPOPCNTDQ dispatch** (`cosine_avx512`, scalar fallback, parity-tested).
`InvertedIndex` + `KMeansIndex` prefilter for sub-ms queries. (Correction to prior notes: there is
**no `bench_kernel.rs`** and the "179 Gdots/s" figure is not in source — it lives in notes/Codex, not
code. SIMD is AVX-512 only; no AVX2; rayon parallelism is at the caller/batch level, not inside `dot`.)

### 1.4 Accuracy — RAG + predictor self-check

- **RAG:** `POST /api/rshl/query` → `handle_rshl_query` (`src/bridge/oracle_server.rs` L5867). Dense
  path projects floats into 16384-dim ternary ("THE CRUSHER") then `query_vec`; associative path uses
  `NeuralBus::query_associative`. Returns top-n `QueryHit`s; callers join `.text` and inject as
  context. The masthead's new read-only `src/bin/rag_*` (v9.10.163) is the sanctioned inference-time
  RAG entry.
- **Predictor (ranking, not verifier):** `src/core/predictive.rs` + `universe.rs`
  (`predictive_query*`). Score = `0.20·sim + 0.55·predict_match + 0.15·multi_head_consensus −
  0.20·recency`; `predict_match` (the dominant term) is cosine to each cell's stored *continuation*
  vector — so retrieval prefers cells that fit the conversation flow. `ConversationTrace` is a rolling
  hyperdimensional "residual-stream analog"; multi-head consensus checks 4 seeded role projections.
- **Self-check:** `src/core/reasoning.rs` (L1003) has a confidence-based "answer-found" halt
  (`phi_g` + cosine>0.6). **There is no P5a-named grounding/hallucination verifier** that checks a
  generated answer against retrieved evidence — that would be net-new. `confidence` in the codebase
  is `phi_g` emergence-math confidence, not retrieval-grounded verification.

### 1.5 The "ten functions" — what the owner most likely means

**No literal "ten" list exists in the repo** (grepped the Codex + src). Two real candidate sets:

- **The 10-agent fleet** (Codex §14.11, L5861–5883): Leo (voice), KAI (lattice core), Gemini,
  Claudey, X, Groq (social), Analyst, Researcher, Kai Coder (work), Oracle (control). Exactly ten
  rows. **But these are Discord personas routed to external LLMs** (Groq/Gemini/Claude/Ollama) — they
  *front* KAI, they don't make the *mind* powerful.
- **The cognition/engine faculties** (real `.rs` modules under `src/cognition/`, ~90 files). The
  Codex groups them into 8 "faculty" clusters (§14.19, L7223–7241). This is the "zero neural weights"
  special sauce the whole project rests on.

**Judgment:** "the functions that made him super good" is about the *mind*, so it means the
**cognition faculties (interpretation B)**, not the persona bots. The ten that best fit, each a real
module:

| # | Faculty | File | One-liner |
|---|---|---|---|
| 1 | Predictor | `predictor.rs` | Predictive processing; prediction-error drives encoding + curiosity |
| 2 | Hippocampus | `hippocampus.rs` | Pattern separation/completion, memory indexing |
| 3 | Sleep/Dreams | `sleep.rs` | Replay consolidation, synaptic homeostasis, pruning |
| 4 | Amygdala | `amygdala.rs` | Emotional-salience / threat gating that boosts encoding |
| 5 | Drives | `drive/mod.rs` | 6 native drives (curiosity, pain, fatigue, satisfaction, social, pred-error) |
| 6 | Theory of Mind | `theory_of_mind.rs` | Per-person mental models |
| 7 | Figurative language | `figurative.rs` | Hyperbole/idiom/metaphor/sarcasm interpreter |
| 8 | Self-model / metacognition | `self_state_hub.rs` | Biases + accuracy/usefulness/coherence targets |
| 9 | Dopamine/reward | `dopamine.rs` (VTA/NAcc) | Reward-prediction error that scales learning-rate (LTP) |
| 10 | Value / moral anchors | `vmpfc.rs`/`ofc.rs`/`mcc.rs` | Value judgment + conflict monitoring + epistemic anchors |

Soft spots: "curiosity" is a *drive* coupled to the predictor (not its own file); "moral" is spread
across vmPFC/OFC/ACC, no single module. Everything else is a discrete named `.rs`.

---

## 2. Strength → LLM-upgrade mapping

| KAI strength (real asset) | LLM upgrade it delivers | Where it lives | Lift |
|---|---|---|---|
| BitNet ternary transformer + zero-mult kernel (`bitnet_llama.rs`, `ternary_math.rs`) | **Power-efficiency + speed** — ternary weights like RSHL, add/subtract matmul, tiny memory if kept packed | IN the LLM (the base model itself) | **Big** (needs trainable ternary path) |
| ~40-tool function-calling + live dispatchers (`gemini-live-bridge.mjs`, `native-tools.mjs`) | **Capability / agency** — the LLM can act: search, maps, weather, code-change, consult engine | AROUND the LLM (inference-time) | **Quick win** |
| Ternary-cosine RAG (`/api/rshl/query`) + continuation predictor | **Accuracy / grounding** — memory text in the prompt; ranked by conversational fit | AROUND the LLM (inference-time RAG) | **Small** (already being built) |
| Reasoning confidence-halt + (new) grounding self-check | **Accuracy / anti-hallucination** — verify answer vs retrieved evidence | AROUND the LLM (inference-time module) | **Medium** (verifier is net-new) |
| Cognition faculties (predictor, hippocampus, drives, ToM, …) | **Persona, memory dynamics, motivation** — but as *behaviors the LLM invokes*, not weights | AROUND / stays in RSHL engine | Faculties already exist; only need tool wrappers |

**The load-bearing claim (consistent with the 4 experiments):** cognition faculties do NOT belong
baked into LLM weights. Baking RSHL-style structure into training hurt every hybrid we tried. The
faculties are already excellent as **engine modules**; the LLM should **call** them (as tools /
context providers), not imitate them in its parameters.

### 2.1 Where each of the ten faculties should live — honest verdict

| Faculty | Bake IN weights? | AROUND as tool/module? | Stay in RSHL engine? | Rationale |
|---|---|---|---|---|
| Predictor | No | Yes — as a self-check/verifier at inference | Yes (ranking) | Ranking stays in engine; grounding-check is an inference wrapper |
| Hippocampus | No | — | **Yes** | It *is* the memory store the RAG reads |
| Sleep/Dreams | No | — | **Yes** | Offline consolidation; nothing to bake |
| Amygdala | No | Optional — salience tag on retrieved context | Yes | Emotional weighting can bias RAG ranking |
| Drives | No | Optional — expose drive state as prompt context | **Yes** | Sets temperature/mood; feed as a system-context tool |
| Theory of Mind | No | **Yes** — `who_am_i_talking_to` tool returns the PersonProfile | Yes (store) | Per-person facts belong in the prompt, retrieved per-user |
| Figurative | Partly (a good base LM already handles figurative language) | Optional | Yes | Least valuable to special-case; base LM covers most |
| Self-model / metacog | No | **Yes** — `kai_status`/self-state as context | Yes | Already surfaced via `kai_status` tool |
| Dopamine/reward | No | — | **Yes** | Learning-rate/consolidation signal, engine-internal |
| Value / moral anchors | Light — SFT/DPO can instill values | **Yes** — anchors as retrieval constraints | Yes | Some values via alignment finetuning; hard constraints as tools/filters |

Only **values** have a mild case for touching weights (via alignment SFT/DPO), and even that is
better reinforced by anchor-as-constraint at inference. Everything else = engine or tool.

---

## 3. Recommended TARGET ARCHITECTURE

```
                    ┌─────────────────────────────────────────────┐
   user turn ─────▶ │  INFERENCE ORCHESTRATOR (Rust, extends       │
                    │  src/bin/rag_* + language_warehouse loop)    │
                    └───────────────┬─────────────────────────────┘
                                    │  builds the prompt:
                                    │  [system + values/anchors]
                                    │  [RAG context ← /api/rshl/query]     ← ACCURACY
                                    │  [ToM: who am I talking to]          ← faculty-as-tool
                                    │  [tool declarations, ~40]            ← CAPABILITY
                                    │  [drive/self-state context]          ← faculty-as-tool
                                    ▼
                    ┌─────────────────────────────────────────────┐
                    │  BASE LM  =  TERNARY / BitNet model          │  ← EFFICIENCY + SPEED
                    │  (bitnet_llama.rs decode; ternary weights,   │
                    │   packed zero-multiply matmul)               │
                    └───────────────┬─────────────────────────────┘
                                    │  generates text + tool-call JSON
                                    ▼
                    ┌─────────────────────────────────────────────┐
                    │  TOOL PARSER + DISPATCH (unify the two       │  ← CAPABILITY
                    │  existing dispatchers) → execute → feed       │
                    │  results back into the decode loop           │
                    └───────────────┬─────────────────────────────┘
                                    │  final draft
                                    ▼
                    ┌─────────────────────────────────────────────┐
                    │  SELF-CHECK / GROUNDING (predictor-based):   │  ← ACCURACY
                    │  does the answer cohere with retrieved       │
                    │  evidence? confidence-halt from reasoning.rs │
                    └───────────────┬─────────────────────────────┘
                                    ▼
                                 answer
```

**In one line:** a **ternary base LLM** (efficiency/speed) + **tool-calling** (capability) + **RSHL
RAG** (memory/accuracy) + **the engine's cognition faculties as callable tools** (persona/memory
dynamics) = KAI's powers, expressed as an LLM. The base model provides fluent language; everything
that "made KAI good" stays where it already works and is *invoked*, not imitated.

---

## 4. Phased build order (what to build first)

### Phase 1 — Tool-calling on the CURRENT model  ·  quick win
Get any decent small LM (the current base, or even the frozen BitNet decoder) to emit tool calls and
wire them to the **existing** dispatchers.
- Pick one function-calling format (reuse the Gemini `{name, arguments}` shape already in
  `LIVE_TOOL_DECLARATIONS`). Put declarations in the prompt.
- Add a parser that extracts tool-call JSON from generation, calls the existing dispatcher, and feeds
  the result back into the decode loop (multi-turn tool loop).
- **Unify the two dispatchers** (voice `if/else` + text `switch`) behind one table-driven map so the
  LLM path and the bot path share tools.
- **Start harvesting SFT data now:** log every `(prompt, tool-call, result)` as JSONL. This becomes
  the finetuning corpus for Phase 4. (Infra half-exists — `audit.json`, transcript JSONL.)
- *Feasibility: high. This is prompt-format + a parser + a dispatch merge. Days, not weeks.*

### Phase 2 — RAG fusion into the prompt  ·  small, mostly done
- Call `/api/rshl/query` (or the new `src/bin/rag_*`) before generation; inject joined `.text` as
  context. Add a `who_am_i_talking_to` (ToM) context block and a `kai_status`/drive-state block.
- Add the grounding **self-check**: after a draft, score answer-vs-evidence with the predictor /
  reasoning confidence-halt; regenerate or hedge if low. (The verifier is the one net-new piece.)
- *Feasibility: high. RAG endpoint exists and is sanctioned; predictor exists as ranker.*

### Phase 3 — Move the base to a TRAINED ternary / BitNet model  ·  the real lift
This is the headline upgrade and the hard one. Options, cheapest-to-most-work:

- **3a. Ship inference on the existing frozen BitNet first**, but **wire the packed zero-multiply
  kernel into the hot path** so it stops dequantizing 2B weights to ~10 GB f32. Pure speed/RAM win,
  no training. Reuses `bitlinear_matmul_buf` that's already built + tested.
- **3b. Post-training quantize** our validated small transformer to BitNet I2_S (absmean ternary) and
  run it through `bitnet_llama.rs`. Fast; expect some quality loss; measures the ternary tax.
- **3c. Quantization-aware training (QAT).** Retrain the small model with BitLinear +
  straight-through estimator + full-precision shadow weights, same pretrain→finetune recipe. Candle
  supports autodiff + optimizers, so it's *possible*, but the current BitNet code is inference-only —
  QAT means adding `VarBuilder`-backed weights, an STE-wrapped BitLinear, a loss + AdamW loop, and a
  training harness. **Non-trivial; the biggest single piece of work in this plan.**
- **Hardware reality:** the owner is on a laptop (RTX 4050). Training even a *small* BitNet from
  scratch is a cloud job (rented A100/H100 hours), not a laptop job. 3a/3b are laptop-friendly; 3c is
  cloud. Recommend 3a → 3b to de-risk, and only fund 3c once tool-calling + RAG prove the wrapper.

### Phase 4 — Fuse: finetune on KAI's own tool + conversation traces
- Finetune the (ideally ternary) base on the SFT corpus harvested in Phase 1 (tool traces) + KAI's
  transcript JSONL, so tool-calling and persona become native rather than few-shot.
- Optionally light DPO on value/moral anchors.
- Keep all faculties as tools; the LLM just gets better at *when* to call them.

**Ordering rationale:** capability (tools) and accuracy (RAG) are cheap and independent of the base
model, so they ship first and immediately make the current model more useful. The expensive base-model
swap (ternary training) happens after the wrapper is proven, so we never train a ternary model we
can't yet drive.

---

## 5. Honest feasibility per piece

| Piece | Verdict | Why |
|---|---|---|
| Tool-calling wrapper | **Quick win** | Declarations + dispatchers already exist; needs a parser + a tool loop + dispatcher unification. |
| RAG into prompt | **Quick win / mostly done** | `/api/rshl/query` + `src/bin/rag_*` exist and are sanctioned (v9.10.163). |
| Grounding self-check | **Medium** | Predictor + confidence-halt exist as *rankers*; an answer-vs-evidence *verifier* is net-new. |
| Wire packed ternary kernel into hot path (3a) | **Medium, high-value** | Kernel built + tested; just not connected. Removes the ~10 GB f32 dequant. |
| PTQ our small model to BitNet (3b) | **Medium** | I2_S codec exists; expect quality loss to measure. |
| QAT ternary training (3c) | **Big lift** | BitNet stack is inference-only (no autograd/STE/shadow weights). Needs a new training harness + cloud GPU. |
| Faculties-as-tools | **Small each** | Modules exist; each needs a thin tool wrapper + declaration. |
| SFT on tool/convo traces (Phase 4) | **Medium** | Straightforward finetune *once* the trace corpus is harvested (start Phase 1). |

**What is genuinely already there (not vaporware):** the ternary transformer, the ternary kernel,
the tokenizer, the generation loop, the ~40 tools + dispatchers, the RAG endpoint, the predictor, and
~90 cognition modules. **What is missing:** a trainable ternary path, a tool-call parser + unified
dispatch for the LLM, an answer-grounding verifier, and the SFT trace corpus.

---

## 6. Open decisions for the owner

1. **"Ten functions" — confirm the reading.** This plan assumes you mean the **cognition faculties**
   (the mind), not the 10 persona bots. If you meant the fleet, the mapping changes (bots become
   multi-agent tool endpoints the LLM routes to). Which did you mean?
2. **Base model: which small LM do we ternarize?** The transformer we validated in the 4 experiments,
   or start from the existing BitNet-2B and *finetune* it? (Finetuning the existing BitNet skips
   pretraining but inherits Microsoft's model, not "ours.")
3. **How far to push ternary training now (Phase 3)?** Cheap-and-safe (3a wire the kernel + 3b PTQ) is
   laptop-friendly and proves the ternary path. Full QAT (3c) needs cloud GPU budget. Do we fund 3c
   now or after tools+RAG ship?
4. **Unify the two tool dispatchers?** Merging voice `if/else` + text `switch` into one map is the
   right long-term move but touches working fleet code (restart risk). Do it now or keep the LLM on
   its own copy first?
5. **Grounding self-check strictness.** How aggressive should the answer-vs-evidence verifier be —
   regenerate on low confidence, or just hedge/annotate? Trades latency for safety.
6. **Where does the trained LLM run** — as a new `src/bin/` binary alongside `rag_*`, or folded into
   `language_warehouse.rs`'s decode loop? (Recommend a new binary to avoid touching the live engine.)

---

*Doc: `C:\KAI\KAI-NATIVE-LLM-UPGRADE-PLAN.md`. No code changed; no Codex/Cargo version bump (design
only). Every asset above was verified against real source paths/line numbers during reconnaissance.*