Fix cross-process recall: MEMB trailer + raw eval/sample in chat()

Two bugs were blocking memba's main promise (load .memb in a fresh
process → model continues with full recalled context):

1. llama_state_set_data() restores the C-level KV-cache + SSM hidden
   state, but llama-cpp-python's Python wrapper still reports
   n_tokens=0. The next eval() then decodes new tokens at offset 0
   and overwrites the loaded state.

   Fix: extend MEMB format with an optional 12-byte trailer appended
   after the CRC32. It carries the wrapper's n_tokens. The C library
   reads up to CRC and ignores anything past it, so files stay
   backward-compatible with libmemba; only the Python loader uses it.

2. Llama.__call__ / create_chat_completion / generate all re-tokenise
   the prompt on every call and clear the KV-cache when the new
   tokens don't prefix-match input_ids. That destroys any state we
   just loaded.

   Fix: rewrite Session.chat() to use raw tokenize → eval → sample.
   eval() appends tokens to the live state without resetting, and we
   handle stop-token detection ourselves.

Verified end-to-end on Nemotron-3-Nano-4B (hybrid 21x Mamba-2 + 4x
attention) — see experiments/README.md for the full findings log.
diag_session_nemotron.py, mood_batch_poc.py, mood_stream_poc.py and
recall_poc.py now all pass their cross-process tests; Falcon-Mamba
still fails because the trained model itself can't do cross-turn
recall — that was the original misdiagnosis.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

experiments/README.md

@@ -1,32 +1,57 @@
 # experiments/
 
-Throwaway scripts used to probe capabilities of SSM models with memba.
+Throwaway scripts used to probe SSM/hybrid model capabilities with memba.
 Not part of the library API — kept in the repo as reference and reproducible
 evidence for product decisions.
 
-Each script is self-contained and prints what it finds; read the source for
-the test's claim and run it yourself if you want to verify on different
-models or hardware.
+Each script is self-contained. Read the source for what it claims to test
+and run it yourself if you want to verify on different models or hardware.
 
 ## Scripts
 
 | File | What it measures |
 |------|------------------|
-| `recall_poc.py` | Can a memba state, built from N days of git activity, answer "what did I work on last month" in a fresh process? |
-| `mood_poc.py` | Batch sentiment-trajectory test (single prompt with full chat log). Compares Falcon-Mamba vs a Transformer. |
-| `mood_stream_poc.py` | The same trajectory but fed turn-by-turn through `Session.chat()`, then queried cross-process. |
-| `mood_batch_poc.py` | Batch ingest in build process, save, then query in a fresh process. |
-| `diag_saveload.py` | Minimal diagnostic: tell the model one fact, ask it back before save, after save, after cross-process load. |
+| `recall_poc.py` | git log → state → cross-process "what did I work on last month" |
+| `mood_poc.py` | Batch sentiment trajectory in one prompt — Falcon-Mamba vs Gemma |
+| `mood_batch_poc.py` | Batch ingest, save, query in fresh process |
+| `mood_stream_poc.py` | 15 separate observation turns, save, query in fresh process |
+| `diag_saveload.py` | Minimal hamster recall test on Falcon-Mamba |
+| `diag_nemotron.py` | Same hamster test on Nemotron 4B hybrid |
+| `diag_nemotron2.py` | Cross-process recall via raw `generate()` (bypass `create_chat_completion`) |
+| `diag_nemotron3.py` | Same with llama-cpp-python's native `save_state()`/`load_state()` |
+| `diag_session_nemotron.py` | Full hamster recall through the rewritten `Session.chat()` |
 
-## Headline finding (2026-05-16, Falcon-Mamba-7B-Instruct Q4_K_M)
+## Findings log
 
-- **Batch single-prompt analysis** (all input + question in one call): works
-  for both sentiment and recall.
-- **Multi-turn fact recall** (ingest in turn 1, ask in turn 2): fails even
-  in the *same process*. The model does not preserve specific facts in its
-  hidden state across conversation turns.
-- **Save/load roundtrip**: byte-identical, no information loss attributable
-  to memba's file format. The persistence layer works correctly; the
-  trained model just doesn't use the state for cross-turn recall.
+### 2026-05-16 (initial, Falcon-Mamba-7B-Instruct Q4_K_M)
 
-See the script outputs (or rerun) for the raw evidence.
+- Batch single-prompt analysis works for both sentiment and recall.
+- Multi-turn fact recall failed even in-process. *Misdiagnosed at first as
+  a model capability issue.*
+
+### 2026-05-16 (revised, after Nemotron-3-Nano-4B Q4_K_M test)
+
+- Real root cause: `Llama.create_chat_completion()` and `Llama.generate()`
+  retokenise the entire prompt each call and reset KV-cache when the new
+  tokens don't prefix-match `input_ids`. Loaded memba state was being wiped.
+- Secondary issue: `llama_state_set_data()` restores the C-level cache but
+  llama-cpp-python's wrapper still reports `n_tokens=0`, so the next
+  `eval()` writes new tokens at offset 0 and overwrites the loaded state.
+- Two fixes in memba:
+  1. MEMB file now carries a Python-only trailer with `n_tokens` so the
+     wrapper position is restored after `load_state` (12 extra bytes,
+     backward-compatible — the C library ignores anything after the CRC).
+  2. `Session.chat()` rewritten to use raw `tokenize → eval → sample`
+     instead of `Llama.__call__`, avoiding the prefix-matching reset.
+- Results on Nemotron-3-Nano-4B (hybrid: 21 Mamba-2 + 4 attention layers)
+  through the new `Session.chat()`:
+  - `diag_session_nemotron.py` — cross-process hamster recall: ✅
+  - `mood_batch_poc.py` — cross-process mood trajectory: ✅
+  - `mood_stream_poc.py` — 15 streaming turns + cross-process: ✅
+  - `recall_poc.py` — per-project digest from 30-day git log: ✅ accurate,
+    no hallucinated projects.
+
+Falcon-Mamba was rerun through the new `Session.chat()` but still fails on
+the same tests — the trained model genuinely doesn't have the cross-turn
+recall capability that the hybrid Nemotron does. The 4 attention layers
+make the difference.

experiments/diag_nemotron.py

@@ -0,0 +1,95 @@
+"""
+diag_nemotron.py — hamster recall test on Nemotron 3 Nano 4B (hybrid Mamba-Transformer).
+
+The Falcon-Mamba 0/4 hamster failure was the killshot for several product ideas.
+This rerun tests whether the hybrid architecture (21 Mamba-2 layers + 4 attention)
+fixes cross-turn recall.
+
+Three scenarios are measured:
+  1. In-process multi-turn (tell fact, ask next turn)
+  2. Same process: save then ask after save
+  3. Cross process: build (ingest, save, exit), then query (load, ask)
+
+Uses create_chat_completion which applies the GGUF's own chat template.
+"""
+from __future__ import annotations
+import sys, argparse, time
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
+from llama_cpp import Llama
+from memba import core
+
+MODEL = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
+STATE = "/tmp/diag_nemotron.memb"
+
+INGEST = ("I'm going to tell you a fact about my pet. My pet hamster is named "
+          "Bartholomew. He is 4 years old. Reply with just 'noted'.")
+QUERY  = "What is the name of my pet?"
+
+
+def make_llama():
+    return Llama(model_path=MODEL, n_ctx=4096, n_gpu_layers=-1, verbose=False)
+
+
+def chat_continued(m, messages):
+    """Send accumulated chat history, return assistant text + cleaned (strip reasoning)."""
+    out = m.create_chat_completion(
+        messages=messages,
+        max_tokens=200,
+        temperature=0.1,
+    )
+    full = out["choices"][0]["message"]["content"].strip()
+    # Nemotron leaks reasoning — try to extract the final answer if present
+    short = full[-300:] if len(full) > 300 else full
+    return full, short
+
+
+def build():
+    print(f"[build] loading Nemotron 4B…")
+    t0 = time.time()
+    m = make_llama()
+    print(f"[build] loaded in {time.time()-t0:.1f}s")
+
+    messages = [{"role": "user", "content": INGEST}]
+    full, _ = chat_continued(m, messages)
+    print(f"[build] ack (full):\n{full!r}\n")
+    messages.append({"role": "assistant", "content": full})
+
+    # Test 1: in-process recall WITHIN the same chat
+    messages.append({"role": "user", "content": QUERY})
+    full, short = chat_continued(m, messages)
+    print(f"[build] in-process query BEFORE save:\n{full}\n")
+    messages.append({"role": "assistant", "content": full})
+
+    # Save state at this point
+    core.save_state(m, MODEL, STATE)
+    print(f"[build] saved state ({Path(STATE).stat().st_size:,} B)")
+
+    # Test 2: in-process query AFTER save — should still work
+    messages.append({"role": "user", "content": QUERY})
+    full, _ = chat_continued(m, messages)
+    print(f"[build] in-process query AFTER save:\n{full}\n")
+
+
+def query():
+    if not Path(STATE).exists():
+        print("[query] no state — run build first"); return 1
+    print(f"[query] loading model + state ({Path(STATE).stat().st_size:,} B)…")
+    t0 = time.time()
+    m = make_llama()
+    core.load_state(m, MODEL, STATE)
+    print(f"[query] loaded in {time.time()-t0:.1f}s")
+
+    # Cross-process: send a fresh user turn with only the question
+    # The state should already encode the prior conversation
+    out = m.create_chat_completion(
+        messages=[{"role": "user", "content": QUERY}],
+        max_tokens=200,
+        temperature=0.1,
+    )
+    print(f"[query] cross-process answer:\n{out['choices'][0]['message']['content']}")
+
+
+if __name__ == "__main__":
+    cmd = sys.argv[1] if len(sys.argv) > 1 else "build"
+    {"build": build, "query": query}[cmd]()

experiments/diag_nemotron2.py

@@ -0,0 +1,79 @@
+"""
+diag_nemotron2.py — same test, but in cross-process query phase we use
+raw __call__ with reset=False to NOT clear KV-cache before generation.
+
+Hypothesis: create_chat_completion() resets KV-cache on each call, which
+defeats memba's loaded state. Bypassing that should let state work.
+"""
+import sys, time
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
+from llama_cpp import Llama
+from memba import core
+
+MODEL = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
+STATE = "/tmp/diag_nemotron.memb"   # reuse state from previous test
+
+def make_llama():
+    return Llama(model_path=MODEL, n_ctx=4096, n_gpu_layers=-1, verbose=False)
+
+
+def query_raw():
+    """Cross-process query, bypassing create_chat_completion's reset."""
+    if not Path(STATE).exists():
+        print("[query] no state — run diag_nemotron.py build first")
+        return
+    print(f"[query] loading state ({Path(STATE).stat().st_size:,} B)…")
+    t0 = time.time()
+    m = make_llama()
+    print(f"[query] model loaded in {time.time()-t0:.1f}s")
+
+    core.load_state(m, MODEL, STATE)
+    print(f"[query] state loaded in {time.time()-t0:.1f}s total")
+
+    # Continue the conversation: just append a new user turn manually.
+    # Nemotron uses ChatML-like markers based on what we've seen:
+    # <|im_start|>user\n…<|im_end|>\n<|im_start|>assistant\n
+    # If wrong, model will still answer something interpretable.
+    continuation = "<|im_end|>\n<|im_start|>user\nWhat is the name of my pet?<|im_end|>\n<|im_start|>assistant\n"
+
+    # Tokenize continuation WITHOUT BOS — we're mid-sequence
+    tokens = m.tokenize(continuation.encode("utf-8"), add_bos=False, special=True)
+    print(f"[query] continuation tokenized to {len(tokens)} tokens")
+
+    # Use generate() which feeds via eval (does NOT reset KV-cache)
+    print(f"[query] generating tokens (preserving loaded state)…")
+    out_tokens = []
+    eos = m.token_eos()
+    for token in m.generate(tokens, top_k=1, temp=0.0):  # greedy
+        if token == eos:
+            break
+        out_tokens.append(token)
+        if len(out_tokens) >= 200:
+            break
+        # Also stop at <|im_end|>
+        snippet = m.detokenize(out_tokens).decode("utf-8", errors="ignore")
+        if "<|im_end|>" in snippet:
+            break
+
+    text = m.detokenize(out_tokens).decode("utf-8", errors="ignore")
+    print(f"\n[query] answer:\n{text}")
+
+
+def query_reset_default():
+    """Control: same continuation but with default reset=True (clears KV)."""
+    if not Path(STATE).exists():
+        return
+    print(f"\n[control] same continuation, default reset=True (should fail):")
+    m = make_llama()
+    core.load_state(m, MODEL, STATE)
+    out = m(
+        "<|im_end|>\n<|im_start|>user\nWhat is the name of my pet?<|im_end|>\n<|im_start|>assistant\n",
+        max_tokens=120, temperature=0.1, echo=False,  # reset=True is default
+    )
+    print(out["choices"][0]["text"])
+
+
+if __name__ == "__main__":
+    query_raw()
+    # query_reset_default()  # optionally enable for direct comparison

experiments/diag_nemotron3.py

@@ -0,0 +1,91 @@
+"""
+diag_nemotron3.py — test hypothesis that llama-cpp-python's built-in
+save_state/load_state preserves Python-side trackers (n_tokens, input_ids)
+which memba's raw C-level save/load is missing.
+
+If built-in works → memba's MEMB format needs to be extended to include
+those trackers. If built-in also fails → the issue is elsewhere.
+"""
+import sys, time, pickle
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
+from llama_cpp import Llama
+
+MODEL = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
+STATE = "/tmp/diag_nemotron_native.pickle"
+
+INGEST = ("I'm going to tell you a fact about my pet. My pet hamster is named "
+          "Bartholomew. He is 4 years old. Reply with just 'noted'.")
+QUERY  = "What is the name of my pet?"
+
+
+def make_llama():
+    return Llama(model_path=MODEL, n_ctx=4096, n_gpu_layers=-1, verbose=False)
+
+
+def build():
+    m = make_llama()
+    print(f"[build] initial n_tokens={m.n_tokens}")
+
+    msgs = [{"role": "user", "content": INGEST}]
+    out = m.create_chat_completion(messages=msgs, max_tokens=80, temperature=0.1)
+    ack = out["choices"][0]["message"]["content"]
+    msgs.append({"role": "assistant", "content": ack})
+    print(f"[build] after ingest: n_tokens={m.n_tokens}")
+    print(f"[build] ack snippet: {ack[-100:]!r}")
+
+    # Use llama-cpp-python's NATIVE save_state — captures Python trackers too
+    state = m.save_state()
+    with open(STATE, "wb") as f:
+        pickle.dump(state, f)
+    print(f"[build] saved native state to {STATE} ({Path(STATE).stat().st_size:,} B)")
+
+
+def query():
+    if not Path(STATE).exists():
+        print("[query] no state file"); return
+    m = make_llama()
+    print(f"[query] before load: n_tokens={m.n_tokens}")
+
+    with open(STATE, "rb") as f:
+        state = pickle.load(f)
+    m.load_state(state)
+    print(f"[query] after load: n_tokens={m.n_tokens}")
+
+    # Now ask via create_chat_completion. Pass ONLY the new question
+    # (history is already in KV-cache+n_tokens).
+    # The template will format this as if it's turn 1 — see what happens.
+    out = m.create_chat_completion(
+        messages=[{"role": "user", "content": QUERY}],
+        max_tokens=150, temperature=0.1,
+    )
+    print(f"\n[query] answer via chat_completion (resets context):\n{out['choices'][0]['message']['content']}")
+
+    # Alternative: try to continue manually, using raw eval
+    print(f"\n[query] re-loading state for raw continuation test")
+    with open(STATE, "rb") as f:
+        state = pickle.load(f)
+    m.load_state(state)
+    print(f"[query] re-loaded: n_tokens={m.n_tokens}")
+
+    # Append a new user turn via raw tokens, then generate
+    continuation = "<|im_end|>\n<|im_start|>user\nWhat is the name of my pet?<|im_end|>\n<|im_start|>assistant\n"
+    tokens = m.tokenize(continuation.encode(), add_bos=False, special=True)
+
+    out_tokens = []
+    eos = m.token_eos()
+    for tok in m.generate(tokens, top_k=1, temp=0.0):
+        if tok == eos or len(out_tokens) >= 150:
+            break
+        out_tokens.append(tok)
+        snippet = m.detokenize(out_tokens).decode("utf-8", errors="ignore")
+        if "<|im_end|>" in snippet:
+            break
+
+    text = m.detokenize(out_tokens).decode("utf-8", errors="ignore")
+    print(f"\n[query] raw continuation answer:\n{text}")
+
+
+if __name__ == "__main__":
+    cmd = sys.argv[1] if len(sys.argv) > 1 else "build"
+    {"build": build, "query": query}[cmd]()

experiments/diag_session_nemotron.py

@@ -0,0 +1,32 @@
+"""Same hamster test, but through the new Session.chat() (which uses eval+sample)."""
+import sys, argparse
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
+from memba import Session
+
+MODEL = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
+STATE_DIR = "/tmp/diag_session_nemotron"
+SESSION   = "hamster"
+
+
+def build():
+    s = Session(model_path=MODEL, session_id=SESSION, state_dir=STATE_DIR,
+                n_gpu_layers=-1, n_ctx=4096, chat_format="chatml")
+    # Wipe any previous state file so it's a real fresh run
+    p = Path(STATE_DIR) / f"{SESSION}.memb"
+    if p.exists(): p.unlink()
+    print(f"[build] ack: {s.chat('My pet hamster is named Bartholomew. He is 4 years old. Reply noted.', max_tokens=40)!r}")
+    print(f"[build] in-process query: {s.chat('What is the name of my pet?', max_tokens=80)!r}")
+    s.save()
+    print(f"[build] saved")
+
+
+def query():
+    s = Session(model_path=MODEL, session_id=SESSION, state_dir=STATE_DIR,
+                n_gpu_layers=-1, n_ctx=4096, chat_format="chatml")
+    print(f"[query] loaded state, n_tokens={s._llama.n_tokens}")
+    print(f"[query] CROSS-PROCESS query: {s.chat('What is the name of my pet?', max_tokens=80)!r}")
+
+
+if __name__ == "__main__":
+    {"build": build, "query": query}[sys.argv[1]]()

experiments/mood_batch_poc.py

@@ -11,7 +11,7 @@ from pathlib import Path
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
 from memba import Session
 
-MAMBA = "/home/emil/Desktop/Coding/AI/Memba/falcon-mamba-7B-instruct-Q4_K_M.gguf"
+MAMBA = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
 STATE_DIR = "/tmp/mood_batch_test"
 SESSION  = "mood_batch"
 

experiments/mood_stream_poc.py

@@ -14,7 +14,7 @@ from pathlib import Path
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
 from memba import Session
 
-MAMBA = "/home/emil/Desktop/Coding/AI/Memba/falcon-mamba-7B-instruct-Q4_K_M.gguf"
+MAMBA = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
 STATE_DIR = "/tmp/mood_stream_test"
 SESSION  = "mood_stream"
 

experiments/recall_poc.py

@@ -34,7 +34,7 @@ SCAN_ROOTS = [
 ]
 DAYS_BACK    = 30
 STATE_DIR    = Path.home() / ".recall/states"
-SESSION_ID   = "recall_poc"
+SESSION_ID   = "recall_poc_nemotron"
 DEFAULT_Q    = (
     "Summarise what I've been working on over the last month. "
     "Group by project. Mention the main themes per project. "