diff --git a/.gitignore b/.gitignore index d5225b5..83c6729 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,13 @@ __pycache__/ *.log .claude/ +# Compiled by build_capi.py against your local llama.cpp - machine/build +# specific, not something to commit. +_llama_cffi.c +_llama_cffi.o +_llama_cffi*.so +build/ + # Generated during test/dev sessions, not part of the tool itself resume.txt Prompt-response.txt diff --git a/README.md b/README.md index c555255..c890d17 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ A prototype implementation of the "Future-Entropy Sampler" technique described in [countbayesie's *Making LLMs Better at Creative Writing Using Entropy*](https://www.countbayesie.com/blog/2026/7/1/making-llms-better-at-creative-writing-using-entropy) -(2026-07-01), built directly against llama.cpp's C API via `ctypes`. +(2026-07-01), built directly against llama.cpp's C API. Instead of sampling straight from `p(w | c)`, this sampler forks the model's KV cache once per top-k candidate token, decodes one token into each fork to see @@ -29,9 +29,10 @@ relying on it for anything beyond experimentation. llama.cpp's `llama_memory_seq_cp`/`llama_memory_seq_rm` plus multi-sequence batched `llama_decode` make the fork/peek/discard cheap: all `top_k` candidate forks are decoded in a single batched call, not `top_k` sequential ones. -`llama-cpp-python` wasn't used - `llama_capi.py` binds the C API directly via -`ctypes`, which means the struct layouts in that file must match whatever -`libllama.so` you point it at (see Known limitations). +`llama-cpp-python` wasn't used - `llama_capi.py` binds the C API directly, +via a small `cffi` extension (`build_capi.py`) compiled against your actual +`llama.h` + `libllama.so`, so the struct layouts always match the exact +build you point it at (see **Setup**). ## Requirements @@ -59,6 +60,24 @@ export ENTROPY_SAMPLER_MODEL=/path/to/your-model.gguf `ENTROPY_SAMPLER_MODEL` is just the default - every entry point also accepts `--model` to override it per run. +Then build the bindings once (requires a C compiler - you already need one to +have built llama.cpp itself): + +``` +python build_capi.py +``` + +This compiles a small extension directly against your `llama.h` and +`libllama.so`, so the struct layouts it uses come from real compilation +against your actual build rather than a hand-copied guess. It looks for +`llama.h` next to `LLAMA_CPP_LIB` (`/build/bin/libllama.so` -> +`/include/llama.h`, llama.cpp's normal layout) and for `ggml.h` in the +sibling `/ggml/include`; set `LLAMA_CPP_INCLUDE` / `LLAMA_CPP_GGML_INCLUDE` +explicitly if your layout differs. Re-run this after rebuilding llama.cpp +against a version that might have changed the structs/functions this project +uses - if something here is now incompatible, this step fails with a compiler +error, not a silent runtime crash. + ## Usage **One-shot CLI**, for testing a single generation with a given configuration: @@ -106,11 +125,6 @@ constants: ## Known limitations -- **ABI-fragile**: `llama_capi.py`'s struct definitions are hand-copied from - one specific llama.cpp commit's `llama.h`. A different llama.cpp version - can silently reorder/resize struct fields and corrupt memory instead of - raising a clean error. Check the struct layouts against your build's - `llama.h` if you see crashes or garbage output. - **Single request at a time**: one `EntropySampler` holds one `llama_context` and one KV cache; `entropy_server.py` serializes requests with a lock. Not built for concurrent multi-user serving. @@ -142,7 +156,8 @@ constants: | File | Purpose | |---|---| -| `llama_capi.py` | `ctypes` bindings over `libllama.so` | +| `build_capi.py` | One-time build step: compiles `llama_capi.py`'s bindings against your `llama.h`/`libllama.so` | +| `llama_capi.py` | Bindings over `libllama.so`, backed by the extension `build_capi.py` compiles | | `entropy_sampler.py` | Core `EntropySampler` class: `prime`/`step`/`generate`/`generate_stream` | | `entropy_cli.py` | One-shot parameterized CLI | | `entropy_server.py` | OpenAI-compatible HTTP server | diff --git a/build_capi.py b/build_capi.py new file mode 100644 index 0000000..ff4ae7a --- /dev/null +++ b/build_capi.py @@ -0,0 +1,167 @@ +"""One-time build step for the llama.cpp bindings. + +Compiles a small C extension (via cffi, "API mode") that #includes your +actual llama.h and links against your actual libllama.so. This is the fix +for the ABI fragility that hand-copied ctypes struct definitions had: rather +than a Python-side guess at field order/padding that can silently produce +wrong offsets against a different llama.cpp build, the real struct layout +here comes from actually compiling against your header - if a field this +project depends on doesn't exist or has a genuinely incompatible type, this +build step fails loudly instead of the sampler corrupting memory at runtime. + +Run this once, and again any time you rebuild llama.cpp against a version +that might have changed the structs/functions this project uses. + +Requires: + LLAMA_CPP_LIB path to libllama.so / libllama.dylib / llama.dll + LLAMA_CPP_INCLUDE directory containing llama.h (optional - guessed as + ../../include relative to LLAMA_CPP_LIB, matching + llama.cpp's normal build-tree layout, if unset) +""" +import os +import sys + +from cffi import FFI + +LIB_PATH = os.environ.get("LLAMA_CPP_LIB") +if not LIB_PATH: + sys.exit( + "Set LLAMA_CPP_LIB to the path of your llama.cpp build's shared library first, " + "e.g.:\n export LLAMA_CPP_LIB=~/llama.cpp/build/bin/libllama.so" + ) +if not os.path.exists(LIB_PATH): + sys.exit(f"LLAMA_CPP_LIB does not exist: {LIB_PATH}") + +INCLUDE_DIR = os.environ.get("LLAMA_CPP_INCLUDE") +if not INCLUDE_DIR: + guess = os.path.normpath(os.path.join(os.path.dirname(LIB_PATH), "..", "..", "include")) + if os.path.exists(os.path.join(guess, "llama.h")): + INCLUDE_DIR = guess + else: + sys.exit( + "Couldn't find llama.h next to LLAMA_CPP_LIB (looked in " + f"{guess}). Set LLAMA_CPP_INCLUDE to the directory containing " + "llama.h, e.g. /include" + ) +elif not os.path.exists(os.path.join(INCLUDE_DIR, "llama.h")): + sys.exit(f"No llama.h found in LLAMA_CPP_INCLUDE={INCLUDE_DIR}") + +# llama.h #includes ggml.h, which lives in a separate include tree +# (/ggml/include) in llama.cpp's normal source layout, not alongside +# llama.h itself. +GGML_INCLUDE_DIR = os.environ.get("LLAMA_CPP_GGML_INCLUDE") +if not GGML_INCLUDE_DIR: + guess = os.path.normpath(os.path.join(INCLUDE_DIR, "..", "ggml", "include")) + if os.path.exists(os.path.join(guess, "ggml.h")): + GGML_INCLUDE_DIR = guess + else: + sys.exit( + "Couldn't find ggml.h (looked in " + f"{guess}). Set LLAMA_CPP_GGML_INCLUDE to the directory " + "containing ggml.h, e.g. /ggml/include" + ) +elif not os.path.exists(os.path.join(GGML_INCLUDE_DIR, "ggml.h")): + sys.exit(f"No ggml.h found in LLAMA_CPP_GGML_INCLUDE={GGML_INCLUDE_DIR}") + +LIB_DIR = os.path.dirname(os.path.abspath(LIB_PATH)) +LIB_BASENAME = os.path.basename(LIB_PATH) +# libllama.so -> llama (the -l link name) +LINK_NAME = LIB_BASENAME +if LINK_NAME.startswith("lib"): + LINK_NAME = LINK_NAME[3:] +LINK_NAME = LINK_NAME.split(".")[0] + +ffi = FFI() + +# Only the fields/functions the sampler actually touches are named below. +# The trailing "...;" in each struct tells cffi to ask the real compiler +# (via the #include in set_source, below) to work out the rest of that +# struct's true layout - fields not listed here are simply invisible to +# Python, they don't need to be declared for the layout to be correct. +ffi.cdef(r""" + typedef int32_t llama_token; + typedef int32_t llama_pos; + typedef int32_t llama_seq_id; + + struct llama_vocab; + struct llama_model; + struct llama_context; + struct llama_memory_i; + typedef struct llama_memory_i * llama_memory_t; + + struct llama_model_params { + int32_t n_gpu_layers; + ...; + }; + + struct llama_context_params { + uint32_t n_ctx; + uint32_t n_batch; + uint32_t n_ubatch; + uint32_t n_seq_max; + int32_t n_threads; + int32_t n_threads_batch; + ...; + }; + + typedef struct llama_batch { + int32_t n_tokens; + llama_token * token; + float * embd; + llama_pos * pos; + int32_t * n_seq_id; + llama_seq_id ** seq_id; + int8_t * logits; + ...; + } llama_batch; + + void llama_backend_init(void); + + struct llama_model_params llama_model_default_params(void); + struct llama_context_params llama_context_default_params(void); + + struct llama_model * llama_model_load_from_file(const char * path_model, + struct llama_model_params params); + void llama_model_free(struct llama_model * model); + const struct llama_vocab * llama_model_get_vocab(const struct llama_model * model); + + struct llama_context * llama_init_from_model(struct llama_model * model, + struct llama_context_params params); + void llama_free(struct llama_context * ctx); + + uint32_t llama_n_ctx(const struct llama_context * ctx); + uint32_t llama_n_ctx_seq(const struct llama_context * ctx); + + int32_t llama_vocab_n_tokens(const struct llama_vocab * vocab); + bool llama_vocab_is_eog(const struct llama_vocab * vocab, llama_token token); + + int32_t llama_tokenize(const struct llama_vocab * vocab, const char * text, int32_t text_len, + llama_token * tokens, int32_t n_tokens_max, + bool add_special, bool parse_special); + int32_t llama_token_to_piece(const struct llama_vocab * vocab, llama_token token, char * buf, + int32_t length, int32_t lstrip, bool special); + + struct llama_batch llama_batch_init(int32_t n_tokens, int32_t embd, int32_t n_seq_max); + void llama_batch_free(struct llama_batch batch); + int32_t llama_decode(struct llama_context * ctx, struct llama_batch batch); + float * llama_get_logits_ith(struct llama_context * ctx, int32_t i); + + llama_memory_t llama_get_memory(const struct llama_context * ctx); + bool llama_memory_seq_rm(llama_memory_t mem, llama_seq_id seq_id, llama_pos p0, llama_pos p1); + void llama_memory_seq_cp(llama_memory_t mem, llama_seq_id seq_id_src, llama_seq_id seq_id_dst, + llama_pos p0, llama_pos p1); +""") + +ffi.set_source( + "_llama_cffi", + '#include "llama.h"', + include_dirs=[INCLUDE_DIR, GGML_INCLUDE_DIR], + library_dirs=[LIB_DIR], + libraries=[LINK_NAME], + extra_link_args=[f"-Wl,-rpath,{LIB_DIR}"], +) + +if __name__ == "__main__": + ffi.compile(verbose=True) + print(f"\nBuilt _llama_cffi against {INCLUDE_DIR}/llama.h " + f"(+ {GGML_INCLUDE_DIR}/ggml.h) + {LIB_PATH}") diff --git a/entropy_sampler.py b/entropy_sampler.py index c80481e..fad7850 100644 --- a/entropy_sampler.py +++ b/entropy_sampler.py @@ -135,7 +135,7 @@ class EntropySampler: def _logits_row(self, i): ptr = C.lib.llama_get_logits_ith(self.ctx, i) - return np.ctypeslib.as_array(ptr, shape=(self.n_vocab,)).copy() + return C.as_float_array(ptr, self.n_vocab) def prime(self, prompt): """Decode the prompt on seq 0. Returns base logits for the first step.""" diff --git a/fork_peek_test.py b/fork_peek_test.py index 8262846..2c40e3e 100644 --- a/fork_peek_test.py +++ b/fork_peek_test.py @@ -88,7 +88,7 @@ def main(): n_ctx_pos = len(prompt_tokens) # next free position on seq 0 logits_ptr = C.lib.llama_get_logits_ith(ctx, -1) - base_logits = np.ctypeslib.as_array(logits_ptr, shape=(n_vocab,)).copy() + base_logits = C.as_float_array(logits_ptr, n_vocab) base_probs = softmax(base_logits) top_idx = np.argsort(-base_probs)[:TOP_K_BASE] @@ -105,7 +105,7 @@ def main(): decode_single_seq_batch(ctx, [w], start_pos=n_ctx_pos, seq_id=fork_seq) fork_logits_ptr = C.lib.llama_get_logits_ith(ctx, -1) - fork_logits = np.ctypeslib.as_array(fork_logits_ptr, shape=(n_vocab,)).copy() + fork_logits = C.as_float_array(fork_logits_ptr, n_vocab) fork_probs = softmax(fork_logits) top_future_idx = np.argsort(-fork_probs)[:TOP_N_FUTURE] @@ -123,7 +123,7 @@ def main(): # same candidate directly on seq 0 must reproduce identical logits decode_single_seq_batch(ctx, [w], start_pos=n_ctx_pos, seq_id=0) check_ptr = C.lib.llama_get_logits_ith(ctx, -1) - check_logits = np.ctypeslib.as_array(check_ptr, shape=(n_vocab,)).copy() + check_logits = C.as_float_array(check_ptr, n_vocab) max_diff = float(np.max(np.abs(check_logits - fork_logits))) print(f"\nseq0-vs-fork logit max abs diff (expect ~0): {max_diff:.6f}") diff --git a/llama_capi.py b/llama_capi.py index 3dd738c..4cae748 100644 --- a/llama_capi.py +++ b/llama_capi.py @@ -1,191 +1,45 @@ -"""Minimal ctypes binding over libllama.so - just the calls the entropy sampler needs. +"""Binding over libllama.so - just the calls the entropy sampler needs. -Built against the llama.cpp checkout at commit d132f22fc (2026-04-09). The struct -layouts below (llama_model_params, llama_context_params, llama_batch) are copied -by hand from that checkout's llama.h and must match the exact libllama.so you point -this at - a different llama.cpp version can silently reorder or resize fields and -corrupt memory rather than raising a clean error. If you hit crashes or garbage -output, check llama.h from your build against the struct definitions here first. +Backed by a cffi extension (see build_capi.py) compiled directly against your +llama.h + libllama.so, so struct layout comes from real compilation instead +of a hand-copied Python guess: a mismatched llama.cpp build fails loudly at +build time rather than silently corrupting memory at run time. Run +`python build_capi.py` once before importing this module. -Not a general-purpose binding - only exposes what fork_peek_test.py / the sampler use. +Not a general-purpose binding - only exposes what fork_peek_test.py / the +sampler use. """ -import ctypes -import os -from ctypes import ( - c_bool, c_char_p, c_float, c_int, c_int8, c_int32, c_size_t, c_uint32, - c_void_p, POINTER, Structure, byref, -) +import numpy as np -LIB_PATH = os.environ.get("LLAMA_CPP_LIB") -if not LIB_PATH: +try: + from _llama_cffi import ffi, lib +except ImportError as e: raise RuntimeError( - "Set LLAMA_CPP_LIB to the path of your llama.cpp build's shared library, " - "e.g.:\n export LLAMA_CPP_LIB=~/llama.cpp/build/bin/libllama.so\n" - "Build it with llama.cpp's CMake build (-DBUILD_SHARED_LIBS=ON) if you " - "don't have it yet." - ) - -lib = ctypes.CDLL(LIB_PATH) - -llama_token = c_int32 -llama_pos = c_int32 -llama_seq_id = c_int32 - - -class llama_model_params(Structure): - _fields_ = [ - ("devices", c_void_p), - ("tensor_buft_overrides", c_void_p), - ("n_gpu_layers", c_int32), - ("split_mode", c_int), - ("main_gpu", c_int32), - ("tensor_split", c_void_p), - ("progress_callback", c_void_p), - ("progress_callback_user_data", c_void_p), - ("kv_overrides", c_void_p), - ("vocab_only", c_bool), - ("use_mmap", c_bool), - ("use_direct_io", c_bool), - ("use_mlock", c_bool), - ("check_tensors", c_bool), - ("use_extra_bufts", c_bool), - ("no_host", c_bool), - ("no_alloc", c_bool), - ] - - -class llama_context_params(Structure): - _fields_ = [ - ("n_ctx", c_uint32), - ("n_batch", c_uint32), - ("n_ubatch", c_uint32), - ("n_seq_max", c_uint32), - ("n_threads", c_int32), - ("n_threads_batch", c_int32), - ("rope_scaling_type", c_int), - ("pooling_type", c_int), - ("attention_type", c_int), - ("flash_attn_type", c_int), - ("rope_freq_base", c_float), - ("rope_freq_scale", c_float), - ("yarn_ext_factor", c_float), - ("yarn_attn_factor", c_float), - ("yarn_beta_fast", c_float), - ("yarn_beta_slow", c_float), - ("yarn_orig_ctx", c_uint32), - ("defrag_thold", c_float), - ("cb_eval", c_void_p), - ("cb_eval_user_data", c_void_p), - ("type_k", c_int), - ("type_v", c_int), - ("abort_callback", c_void_p), - ("abort_callback_data", c_void_p), - ("embeddings", c_bool), - ("offload_kqv", c_bool), - ("no_perf", c_bool), - ("op_offload", c_bool), - ("swa_full", c_bool), - ("kv_unified", c_bool), - ("samplers", c_void_p), - ("n_samplers", c_size_t), - ] - - -class llama_batch(Structure): - _fields_ = [ - ("n_tokens", c_int32), - ("token", POINTER(llama_token)), - ("embd", POINTER(c_float)), - ("pos", POINTER(llama_pos)), - ("n_seq_id", POINTER(c_int32)), - ("seq_id", POINTER(POINTER(llama_seq_id))), - ("logits", POINTER(c_int8)), - ] - - -lib.llama_backend_init.argtypes = [] -lib.llama_backend_init.restype = None - -lib.llama_model_default_params.argtypes = [] -lib.llama_model_default_params.restype = llama_model_params - -lib.llama_context_default_params.argtypes = [] -lib.llama_context_default_params.restype = llama_context_params - -lib.llama_model_load_from_file.argtypes = [c_char_p, llama_model_params] -lib.llama_model_load_from_file.restype = c_void_p - -lib.llama_model_free.argtypes = [c_void_p] -lib.llama_model_free.restype = None - -lib.llama_model_get_vocab.argtypes = [c_void_p] -lib.llama_model_get_vocab.restype = c_void_p - -lib.llama_init_from_model.argtypes = [c_void_p, llama_context_params] -lib.llama_init_from_model.restype = c_void_p - -lib.llama_free.argtypes = [c_void_p] -lib.llama_free.restype = None - -lib.llama_vocab_n_tokens.argtypes = [c_void_p] -lib.llama_vocab_n_tokens.restype = c_int32 - -lib.llama_tokenize.argtypes = [ - c_void_p, c_char_p, c_int32, POINTER(llama_token), c_int32, c_bool, c_bool, -] -lib.llama_tokenize.restype = c_int32 - -lib.llama_token_to_piece.argtypes = [ - c_void_p, llama_token, ctypes.c_char_p, c_int32, c_int32, c_bool, -] -lib.llama_token_to_piece.restype = c_int32 - -lib.llama_batch_init.argtypes = [c_int32, c_int32, c_int32] -lib.llama_batch_init.restype = llama_batch - -lib.llama_batch_free.argtypes = [llama_batch] -lib.llama_batch_free.restype = None - -lib.llama_decode.argtypes = [c_void_p, llama_batch] -lib.llama_decode.restype = c_int32 - -lib.llama_get_logits_ith.argtypes = [c_void_p, c_int32] -lib.llama_get_logits_ith.restype = POINTER(c_float) - -lib.llama_get_memory.argtypes = [c_void_p] -lib.llama_get_memory.restype = c_void_p - -lib.llama_memory_seq_cp.argtypes = [c_void_p, llama_seq_id, llama_seq_id, llama_pos, llama_pos] -lib.llama_memory_seq_cp.restype = None - -lib.llama_memory_seq_rm.argtypes = [c_void_p, llama_seq_id, llama_pos, llama_pos] -lib.llama_memory_seq_rm.restype = c_bool - -lib.llama_vocab_is_eog.argtypes = [c_void_p, llama_token] -lib.llama_vocab_is_eog.restype = c_bool - -lib.llama_n_ctx.argtypes = [c_void_p] -lib.llama_n_ctx.restype = c_uint32 - -lib.llama_n_ctx_seq.argtypes = [c_void_p] -lib.llama_n_ctx_seq.restype = c_uint32 + "Compiled bindings not found. Run `python build_capi.py` first " + "(with LLAMA_CPP_LIB, and LLAMA_CPP_INCLUDE if it can't be guessed, set)." + ) from e def tokenize(vocab, text, add_special=True, parse_special=True, max_tokens=8192): - buf = (llama_token * max_tokens)() + text_bytes = text.encode("utf-8") + buf = ffi.new("int32_t[]", max_tokens) n = lib.llama_tokenize( - vocab, text.encode("utf-8"), len(text.encode("utf-8")), - buf, max_tokens, add_special, parse_special, + vocab, text_bytes, len(text_bytes), buf, max_tokens, add_special, parse_special, ) if n < 0: raise RuntimeError(f"tokenize buffer too small, need {-n}") - return list(buf[:n]) + return [buf[i] for i in range(n)] def token_to_str(vocab, token): - buf = ctypes.create_string_buffer(64) + buf = ffi.new("char[]", 64) n = lib.llama_token_to_piece(vocab, token, buf, 64, 0, True) if n < 0: - buf = ctypes.create_string_buffer(-n) + buf = ffi.new("char[]", -n) n = lib.llama_token_to_piece(vocab, token, buf, -n, 0, True) - return buf.raw[:n].decode("utf-8", errors="replace") + return ffi.buffer(buf, n)[:].decode("utf-8", errors="replace") + + +def as_float_array(ptr, n): + """Copy an n-element float* returned by the library into a numpy array.""" + return np.frombuffer(ffi.buffer(ptr, n * ffi.sizeof("float")), dtype=np.float32).copy() diff --git a/requirements.txt b/requirements.txt index dae10c3..a9a0c5b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,5 @@ gguf>=0.16 jinja2>=3.1 fastapi>=0.110 uvicorn>=0.30 +cffi>=1.16 +setuptools>=68