Nibmeister/entropy-sampler
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix ABI fragility: compile bindings against the real llama.h instead of hand-copied ctypes structs

Browse Source 
llama_capi.py previously hand-copied struct field order/types from one
specific llama.cpp commit's llama.h into Python ctypes Structures. A
different llama.cpp build could silently reorder or resize those fields and
corrupt memory rather than raising any error.

Replaced with a cffi "API mode" extension (build_capi.py) that #includes the
user's actual llama.h and links against their actual libllama.so. Struct
layout now comes from real compilation - a genuinely incompatible field
fails the build loudly instead of corrupting memory at runtime. Verified
byte-for-byte identical generation output against the prior ctypes
implementation at a fixed seed, plus the fork/peek logit round-trip check
(0.000000 max diff).

Requires a one-time `python build_capi.py` setup step (needs a C compiler,
which building llama.cpp itself already requires).
This commit is contained in:
Joey Grasty
2026-07-11 21:21:29 -05:00
parent e25ee046a3
commit 8c844c08e5
7 changed files with 232 additions and 187 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
.gitignore vendored
View File

@@ -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

35
README.md
View File

@@ -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` (`<repo>/build/bin/libllama.so` ->
`<repo>/include/llama.h`, llama.cpp's normal layout) and for `ggml.h` in the
sibling `<repo>/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 |

167
build_capi.py Normal file
View File

@@ -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. <llama.cpp checkout>/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
# (<repo>/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. <llama.cpp checkout>/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}")

2
entropy_sampler.py
View File

@@ -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."""

6
fork_peek_test.py
View File

@@ -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}")

200
llama_capi.py
View File

@@ -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()

2
requirements.txt
View File

@@ -3,3 +3,5 @@ gguf>=0.16
jinja2>=3.1
fastapi>=0.110
uvicorn>=0.30
cffi>=1.16
setuptools>=68