Fix ABI fragility: compile bindings against the real llama.h instead of hand-copied ctypes structs
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:
200
llama_capi.py
200
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()
|
||||
|
||||
Reference in New Issue
Block a user