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>
This commit is contained in:
+28
-1
@@ -25,6 +25,16 @@ MODEL_ID_LEN = 64
|
||||
_HEADER_FMT = "<4sI64sIIQ"
|
||||
_HEADER_SIZE = struct.calcsize(_HEADER_FMT) # == 88
|
||||
|
||||
# Optional Python trailer (appended after the V1 CRC32). The C library reads
|
||||
# only up to data_size+CRC and ignores anything past it, so adding a trailer
|
||||
# keeps the file C-compatible. The trailer carries llama-cpp-python's
|
||||
# wrapper state (n_tokens) which is needed for `Llama.eval()` to append new
|
||||
# tokens at the correct position instead of overwriting from offset 0.
|
||||
_TRAILER_MAGIC = b"MTRL"
|
||||
_TRAILER_FMT = "<4sII" # magic | version | n_tokens
|
||||
_TRAILER_SIZE = struct.calcsize(_TRAILER_FMT) # == 12
|
||||
_TRAILER_VERSION = 1
|
||||
|
||||
|
||||
# ── Helpers ───────────────────────────────────────────────────────
|
||||
|
||||
@@ -140,9 +150,14 @@ def save_state(llama_model: "Llama", model_path: str, file_path: str) -> None:
|
||||
len(data),
|
||||
)
|
||||
|
||||
# Python-only trailer with Llama wrapper position so eval() resumes
|
||||
# at the right offset instead of overwriting from 0 after load_state.
|
||||
n_tokens = getattr(llama_model, "n_tokens", 0)
|
||||
trailer = struct.pack(_TRAILER_FMT, _TRAILER_MAGIC, _TRAILER_VERSION, n_tokens)
|
||||
|
||||
out = Path(file_path)
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
out.write_bytes(header + data + struct.pack("<I", crc))
|
||||
out.write_bytes(header + data + struct.pack("<I", crc) + trailer)
|
||||
|
||||
|
||||
def load_state(llama_model: "Llama", model_path: str, file_path: str) -> None:
|
||||
@@ -187,3 +202,15 @@ def load_state(llama_model: "Llama", model_path: str, file_path: str) -> None:
|
||||
restored = _state_set_data(lib, ctx, buf, len(data))
|
||||
if restored == 0:
|
||||
raise RuntimeError("llama_state_set_data returned 0 — state restore failed")
|
||||
|
||||
# Optional trailer: restore llama-cpp-python wrapper position
|
||||
trailer_off = offset + data_size + 4
|
||||
if len(raw) >= trailer_off + _TRAILER_SIZE:
|
||||
magic, version, n_tokens = struct.unpack_from(_TRAILER_FMT, raw, trailer_off)
|
||||
if magic == _TRAILER_MAGIC and version == _TRAILER_VERSION:
|
||||
# Restore Llama wrapper state so eval() appends at the right
|
||||
# offset instead of overwriting the loaded KV-cache from 0.
|
||||
try:
|
||||
llama_model.n_tokens = n_tokens
|
||||
except AttributeError:
|
||||
pass # wrapper version doesn't expose .n_tokens — eval may misbehave
|
||||
|
||||
Reference in New Issue
Block a user