for tokens in ([], [3, 1, 3]):

storage = (ctypes.c_int32 * len(tokens))(*tokens)

calls = []

def get_suppressed(vocab, count):

calls.append(vocab)

ctypes.cast(count, ctypes.POINTER(ctypes.c_int32))[0] = len(tokens)

return storage if tokens else None

monkeypatch.setattr(internals.llama_cpp, "llama_vocab_get_suppress_tokens", get_suppressed)

model = SimpleNamespace(vocab=123)

result = internals.LlamaModel.suppress_tokens(model)

if tokens:

storage[0] = 99

assert result == tokens

model.vocab = None

with pytest.raises(RuntimeError, match="vocab is None"):

internals.LlamaModel.suppress_tokens(model)

internals.llama_cpp.llama_backend_init()

monkeypatch.setattr(internals.llama_cpp, "llama_model_get_vocab", lambda _: None)

return SimpleNamespace(model=None, n_vocab=lambda: 32,

lib = internals.llama_cpp

sampling_model.suppress_tokens.return_value = suppressed

logits = np.arange(32, dtype=np.float32) * 0.25

saved = logits.copy()

# Without suppression, doubling this bias incorrectly moves token 29 to first.

biases = [lib.llama_logit_bias(29, 0.3)]

if suppressed:

biases.append(lib.llama_logit_bias(31, 100.0))

original_biases = [(item.token, item.bias) for item in biases]

samplers = [internals.CommonSamplerType.TOP_K]

if probability_filters:

samplers.extend([internals.CommonSamplerType.TOP_P, internals.CommonSamplerType.MIN_P,

internals.CommonSamplerType.TEMPERATURE])

params = internals.LlamaSamplingParams(

temp=temp, top_k=10, top_p=0.8, min_p=0.12, min_keep=1, seed=42,

logit_bias=biases, samplers=samplers,

)

native = internals.LlamaTokenDataArray(n_vocab=32)

native.copy_logits(logits)

with closing(internals.LlamaSampler()) as reference:

reference.add_logit_bias(32, biases + [lib.llama_logit_bias(token, float("-inf")) for token in suppressed])

reference.add_top_k(10)

if probability_filters:

reference.add_top_p(0.8, 1)

reference.add_min_p(0.12, 1)

reference.add_temp(temp)

if temp == 0:

reference.add_greedy()

else:

reference.add_dist(42)

lib.llama_sampler_apply(reference.sampler, ctypes.byref(native.candidates))

expected = native.candidates_data[:native.candidates.size]

token = int(expected["id"][native.candidates.selected])

assert token not in suppressed

context = SimpleNamespace(

get_sampled_token_ith=lambda _: lib.LLAMA_TOKEN_NULL,

get_logits_ith=lambda _: logits.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),

)

with closing(internals.LlamaSamplingContext(params, sampling_model)) as sampler:

for source in (None, logits): # Context outputs and restored snapshots.

sampler.reset()

with closing(sampler.cp()) as clone:

for current in (sampler, clone):

assert current.sample(context, logits=source) == token

actual = current._cur_p.candidates_data[:current._cur_p.candidates.size]

np.testing.assert_array_equal(actual["id"], expected["id"])

# Greedy selection does not populate probabilities.

for field in (("logit", "p") if temp else ("logit",)):

np.testing.assert_allclose(actual[field], expected[field], rtol=2e-6, atol=1e-8)

sampling_model.suppress_tokens.assert_called_once_with()

assert [(item.token, item.bias) for item in params.logit_bias] == original_biases

lib = internals.llama_cpp

logits = np.arange(32, dtype=np.float32)

params = internals.LlamaSamplingParams(temp=0.0, top_k=1,

samplers=[internals.CommonSamplerType.TOP_K])

with closing(internals.LlamaSampler()) as grammar, closing(internals.LlamaSamplingContext(params, sampling_model)) as sampler:

# A native mask stands in for a grammar that rejects the first allowed token.

grammar.add_logit_bias(32, [lib.llama_logit_bias(30, float("-inf"))])

sampler.grammar_sampler = SimpleNamespace(grammar=grammar.sampler, close=lambda: None)

assert sampler.sample(SimpleNamespace(), logits=logits, grammar_first=grammar_first) == 29

context = SimpleNamespace(get_sampled_token_ith=lambda _: 31,

get_logits_ith=Mock(side_effect=AssertionError("CPU fallback is not allowed")))

params = internals.LlamaSamplingParams(samplers=[internals.CommonSamplerType.TOP_K])

with closing(internals.LlamaSamplingContext(params, sampling_model)) as sampler:

assert sampler.sample(context) == 31

assert sampler._cur_p is None

lib = internals.llama_cpp

sampling_model.suppress_tokens.return_value = suppressed

sampling_model.tokenize = lambda *args, **kwargs: [30]

params = internals.LlamaSamplingParams(

temp=1.0, top_k=32, seed=42, reasoning_budget=0,

reasoning_start_in_prompt=True, logit_bias=[lib.llama_logit_bias(30, 0.5)],

samplers=[internals.CommonSamplerType.TOP_K],

)

logits = np.arange(32, dtype=np.float32) * 0.25

# common_reasoning_budget_apply masks other tokens and preserves the forced

# token's original score; common_sampler_sample then applies the bias chain.

expected = np.full(32, -np.inf, dtype=np.float32)

expected[30] = logits[30] + 0.5

expected[suppressed] = -np.inf

candidates = internals.LlamaTokenDataArray(n_vocab=32)

candidates.copy_logits(logits)

with closing(internals.LlamaSamplingContext(params, sampling_model)) as sampler:

chain = sampler.sampler_chain.sampler

# Inspect preprocessing before top-k/dist, including the all-masked

# conflict case for which no valid sampling distribution exists.

for index in range(lib.llama_sampler_chain_n(chain) - 2):

lib.llama_sampler_apply(lib.llama_sampler_chain_get(chain, index), ctypes.byref(candidates.candidates))

np.testing.assert_array_equal(candidates._logit_view, expected)

assert not np.isnan(candidates._logit_view).any()

if 30 not in suppressed:

assert sampler.sample(SimpleNamespace(), logits=logits) == 30

probabilities = [sampler._cur_p.candidates.data[i].p for i in range(32)]

assert np.isfinite(probabilities).all() and sum(probabilities) == pytest.approx(1.0)

model = LlamaEmbedding.__new__(LlamaEmbedding)

with pytest.raises(RuntimeError, match="not initialized"):

model.embed_multimodal("hello")

model._ctx = SimpleNamespace(ctx=123)

try:

with pytest.raises(ValueError, match="mmproj_path"):

model.embed_multimodal("hello")

finally:

import io

from PIL import Image

from llama_cpp.llama_multimodal import MTMDEmbeddingHandler

model_path = os.environ.get(f"LLAMA_TEST_{profile}_EMBEDDING_MODEL")

mmproj_path = os.environ.get(f"LLAMA_TEST_{profile}_EMBEDDING_MMPROJ")

if not model_path or not mmproj_path:

pytest.skip(f"Set LLAMA_TEST_{profile}_EMBEDDING_MODEL/MMPROJ paths")

image = io.BytesIO()

Image.new("RGB", (224, 224), "red").save(image, format="PNG")

handler = MTMDEmbeddingHandler(mmproj_path, use_gpu=False,

image_max_tokens=64 if causal else 70, verbose=False)

try:

with closing(LlamaEmbedding(model_path, embedding_handler=handler, n_ctx=256,

n_batch=256, n_ubatch=256, n_threads=4, verbose=False,

n_gpu_layers=int(os.environ.get("LLAMA_TEST_N_GPU_LAYERS", "0")),

pooling_type=pooling)) as model:

assert handler.mtmd_ctx is None

prefix = ("<|im_start|>system\nRepresent the user's input.<|im_end|>\n"

"<|im_start|>user\n") if causal else "title: none | text: "

suffix = "<|im_end|>\n<|im_start|>assistant\n" if causal else ""

parts = [prefix, image.getvalue(), suffix]

# Call directly after construction; no manual projector initialization.

vector = model.embed_multimodal(parts=parts)

assert len(vector) == width and np.isfinite(vector).all()

assert model._ctx.get_causal_attn() is causal

np.testing.assert_allclose(model.embed_multimodal(parts=parts), vector, atol=1e-5)

if causal:

# Match a direct native submission using the same decode shape.

model.reset()

with handler.tokenize_parts(model, parts) as request:

encoded = request.encode()

for chunk in encoded:

if chunk.tokens:

assert all(isinstance(pos, int) for pos in chunk.positions)

else:

assert all(len(pos) == 4 for pos in chunk.positions)

native = llama_cpp.llama_batch_ext_init(model._ctx.ctx)

assert native

try:

for chunk in encoded:

rows = chunk.tokens if chunk.tokens else chunk.embeddings

for row, pos in zip(rows, chunk.positions):

if chunk.tokens:

idx = llama_cpp.llama_batch_ext_add_token(native, 0, row)

coords = (llama_cpp.llama_pos * 1)(pos)

else:

embd = llama_cpp.llama_embd(row.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), 1, len(row))

idx = llama_cpp.llama_batch_ext_add_embd(native, 0, embd)

coords = (llama_cpp.llama_pos * 4)(*pos)

assert idx >= 0 and llama_cpp.llama_batch_ext_set_pos(native, idx, coords)

assert llama_cpp.llama_batch_ext_set_output_logits(native, idx, True)

assert llama_cpp.llama_process(model._ctx.ctx, llama_cpp.llama_process_type.LLAMA_PROCESS_TYPE_DECODE, native) == 0

raw = np.ctypeslib.as_array(model._ctx.get_embeddings_seq(0), shape=(width,)).copy()

assert np.dot(vector, raw / np.linalg.norm(raw)) > 0.99999

finally:

llama_cpp.llama_batch_ext_free(native)

model.reset()

handler.is_support_vision = False

with pytest.raises(ValueError, match="does not support image"):

model.embed_multimodal(parts=["title: none | text: ", image.getvalue()])

handler.is_support_vision = True

model._ctx.set_embeddings(False)

vector = model.embed_multimodal(prefix + "hello" + suffix)

assert len(vector) == width and np.linalg.norm(vector) == pytest.approx(1, abs=1e-6)

assert model._ctx.get_causal_attn() is causal

assert handler.mtmd_ctx is None

with pytest.raises(RuntimeError, match="closed"):

model.embed_multimodal("hello")

finally:

ctx = internals.LlamaContext.__new__(internals.LlamaContext)

ctx.ctx = 123

ctx.model = SimpleNamespace(

n_vocab=lambda: 100,

rope_type=lambda: internals.llama_cpp.llama_rope_type.LLAMA_ROPE_TYPE_MROPE,

)

ctx.params = object()

ctx.verbose = False

ctx._loras_applied = False

ctx._cvec_applied = False

ctx._lora_refs = ()

monkeypatch.setattr(internals.llama_cpp, "llama_free", Mock())

monkeypatch.setattr(internals.llama_cpp, "llama_n_batch", lambda _: 3)

monkeypatch.setattr(internals.llama_cpp, "llama_n_seq_max", lambda _: 2)

try:

yield ctx

finally:

ctx.ctx = None

ctx._batch_ext_refs = []

lib = internals.llama_cpp

rows, calls, freed = [], [], []

state = SimpleNamespace(fail=False, status=0)

context = runtime_context

monkeypatch.setattr(lib, "llama_batch_ext_init", lambda _: 456)

monkeypatch.setattr(lib, "llama_batch_ext_free", lambda p: freed.append(p))

monkeypatch.setattr(lib, "llama_batch_ext_clear", lambda _: rows.clear())

def add(_, sid, token):

rows.append(dict(token=token, seq_ids=[sid]))

return len(rows) - 1

def payload(_, idx, embd, field):

if state.fail == field:

return False

rows[idx][field] = np.ctypeslib.as_array(embd.data, shape=(embd.n_rows * embd.n_embd,)).copy()

return True

def add_embd(p, sid, embd):

idx = add(p, sid, None)

return idx if payload(p, idx, embd, "embedding") else -2

def sequence(_, idx, sid):

rows[idx]["seq_ids"].append(sid)

return True

def position(_, idx, pos):

n = 1 if rows[idx]["token"] is not None else 4

rows[idx]["position"] = tuple(pos[i] for i in range(n))

return True

def output(_, idx, flag):

rows[idx]["output"] = flag

return True

def process(ctx, kind, native):

calls.append((ctx, kind, native))

return state.status

monkeypatch.setattr(lib, "llama_batch_ext_add_token", add)

monkeypatch.setattr(lib, "llama_batch_ext_add_embd", add_embd)

monkeypatch.setattr(lib, "llama_batch_ext_set_embd_token", lambda *args: payload(*args, "embedding"))

monkeypatch.setattr(lib, "llama_batch_ext_set_embd_state", lambda *args: payload(*args, "state"))

monkeypatch.setattr(lib, "llama_batch_ext_add_seq", sequence)

monkeypatch.setattr(lib, "llama_batch_ext_set_pos", position)

monkeypatch.setattr(lib, "llama_batch_ext_set_output_logits", output)

monkeypatch.setattr(lib, "llama_process", process)

try:

with internals.LlamaBatchExt(context=context) as batch:

yield SimpleNamespace(context=context, batch=batch, rows=rows, calls=calls, freed=freed, state=state)

finally:

from llama_cpp._utils import normalize_embedding

for mode in (-1, False, True, 0, 1, 2, 6):

for values in ([0.0, 0.0], [3.0, -4.0], [1e7, -2e7]):

case = f"mode={mode!r}, values={values}"

with np.errstate(over="raise", invalid="raise"):

result = np.asarray(normalize_embedding(values, mode), dtype=np.float64)

assert np.isfinite(result).all(), case

if mode is False or mode == -1 or not any(values):

np.testing.assert_array_equal(result, values, err_msg=case)

elif mode == 0:

assert np.max(np.abs(result)) == pytest.approx(32760, rel=1e-6), case

else:

assert np.linalg.norm(result, ord=2 if mode is True else mode) == pytest.approx(1, rel=1e-6), case

if any(values):

from llama_cpp.llama_multimodal import MTMDEncodedChunk, MTMDEmbeddingHandler

r = batch_runtime

raw = np.array([3, 4], dtype=np.float32)

monkeypatch.setattr(r.context, "get_causal_attn", lambda: causal)

monkeypatch.setattr(r.context, "set_causal_attn", Mock())

monkeypatch.setattr(r.context, "pooling_type", lambda: llama_cpp.LLAMA_POOLING_TYPE_LAST)

monkeypatch.setattr(r.context, "n_ubatch", lambda: 3)

monkeypatch.setattr(r.context, "set_embeddings", Mock())

monkeypatch.setattr(r.context, "get_embeddings_seq", lambda seq: raw.ctypes.data_as(ctypes.POINTER(ctypes.c_float)))

model = SimpleNamespace(_ctx=r.context, context_params=SimpleNamespace(embeddings=True),

n_ctx=lambda: 3, n_embd_out=lambda: 2, reset=Mock())

chunks = (

MTMDEncodedChunk(0, (1,), None, (0,), 1, False),

MTMDEncodedChunk(1, (), np.ones((1, 2), dtype=np.float32), ((1, 2, 3, 0),), 1, media_non_causal),

)

request = SimpleNamespace(n_tokens=2, n_pos=2, encode=lambda: chunks)

released = []

@contextmanager

def tokenize(*args, **kwargs):

try:

yield request

finally:

released.append(True)

with MTMDEmbeddingHandler(str(tmp_path), verbose=False) as handler:

monkeypatch.setattr(handler, "_init_mtmd_context", Mock())

monkeypatch.setattr(handler, "tokenize", tokenize)

handler.is_support_vision = True

if causal and media_non_causal:

with pytest.raises(NotImplementedError, match="chunk-local"):

handler.embed(model, "prompt")

assert not r.calls

else:

r.state.status = 1

with pytest.raises(RuntimeError, match="decode failed"):

handler.embed(model, "prompt")

r.state.status = 0

np.testing.assert_allclose(handler.embed(model, "prompt"), [0.6, 0.8])

assert model.reset.call_count == len(released) * 2

r.context.set_causal_attn.assert_not_called()

runtime_context.close()

for method, args, native in (

("n_ctx", (), "llama_n_ctx"),

("n_ctx_seq", (), "llama_n_ctx_seq"),

("n_batch", (), "llama_n_batch"),

("n_ubatch", (), "llama_n_ubatch"),

("n_seq_max", (), "llama_n_seq_max"),

("n_rs_seq", (), "llama_n_rs_seq"),

("pooling_type", (), "llama_pooling_type"),

("set_n_threads", (1, 1), "llama_set_n_threads"),

("n_threads", (), "llama_n_threads"),

("n_threads_batch", (), "llama_n_threads_batch"),

("set_causal_attn", (True,), "llama_set_causal_attn"),

("get_causal_attn", (), "llama_get_causal_attn"),

("clear_loras", (), "llama_set_adapters_lora"),

("apply_loras", ([],), "llama_set_adapters_lora"),

("clear_cvec", (), "llama_set_adapter_cvec"),

("apply_cvec", ([], 1, 1, 1), "llama_set_adapter_cvec"),

("graph_reserve", (1, 1, 1), "llama_graph_reserve"),

):

with monkeypatch.context() as patch:

call = Mock(name=native)

patch.setattr(internals.llama_cpp, native, call)

with pytest.raises(RuntimeError, match="closed"):

getattr(runtime_context, method)(*args)

context = runtime_context

events, graph = [], object()

reserve = Mock(side_effect=lambda *args: events.append(args) or graph)

sync = Mock(side_effect=lambda _: events.append('sync'))

monkeypatch.setattr(internals.llama_cpp, 'llama_synchronize', sync)

monkeypatch.setattr(internals.llama_cpp, 'llama_graph_reserve', reserve)

for dimensions, error in (

((0, 1, 1), ValueError), ((1, 0, 1), ValueError),

((1, 1, 0), ValueError), ((-1, 1, 1), ValueError),

((2**32, 1, 1), ValueError), ((1, 3, 1), ValueError),

((1, 1, 2), ValueError), ((2**31 - 1, 2, 1), ValueError),

((True, 1, 1), TypeError), ((1.5, 1, 1), TypeError),

):

with pytest.raises(error):

context.graph_reserve(*dimensions)

reserve.assert_not_called()

sync.assert_not_called()

for failure, message in ((None, 'failed to reserve'),

(RuntimeError('missing symbol'), 'missing symbol')):

reserve.side_effect = failure

reserve.return_value = None

with pytest.raises(RuntimeError, match=message):

context.graph_reserve(1, 1, 1)

assert context.ctx == 123

events.clear()

reserve.side_effect = lambda *args: events.append(args) or graph

# Reservation synchronizes first and permits native sequence rounding.

assert context.graph_reserve(3, 2, 4) is graph

context = runtime_context

native = Mock(return_value=0)

monkeypatch.setattr(internals.llama_cpp, "llama_set_adapters_lora", native)

with pytest.raises(RuntimeError, match="closed or invalid"):

context.apply_loras([(SimpleNamespace(adapter=None), 1.0)])

native.assert_not_called()

original = MockLoraAdapter(1)

reference = weakref.ref(original)

context.apply_loras([(original, 1.0)])

del original

gc.collect()

assert reference() is not None

native.return_value = -1

with pytest.raises(RuntimeError, match="Failed to set LoRA"):

context.apply_loras([(MockLoraAdapter(2), 1.0)])

with pytest.raises(RuntimeError, match="detach failed"):

context.clear_loras()

gc.collect()

assert reference() is not None

native.return_value = 0

context.clear_loras()

gc.collect()

context = runtime_context

events = []

def fail():

events.append("batch")

if failure == "batch":

raise RuntimeError("batch cleanup failed")

context._batch_ext_refs = [SimpleNamespace(close=fail)]

context._checkpoint_caches = [SimpleNamespace(close=lambda: events.append("checkpoint"))]

context._exit_stack = SimpleNamespace(close=lambda: events.append("exit"))

adapter = MockLoraAdapter(1)

reference = weakref.ref(adapter)

context._lora_refs = (adapter,)

del adapter

def free(pointer):

assert pointer == 123 and reference() is not None

events.append("native")

if failure == "native":

raise RuntimeError("native cleanup failed")

monkeypatch.setattr(internals.llama_cpp, "llama_free", free)

with pytest.raises(RuntimeError, match=f"{failure} cleanup failed"):

context.close()

assert events == ["batch", "checkpoint", "native", "exit"]

assert context.ctx is None and context.model is None and context.params is None

gc.collect()

assert reference() is None

context._batch_ext_refs = []

context.close()

r = batch_runtime

r.batch.add_token(1, 0)

source = np.arange(8, dtype=np.float32).reshape(2, 4)

r.batch.add_embeddings(source, positions=[(1, 2, 3, 0), (2, 3, 4, 0)], seq_ids=(0, 1))

source[:] = -1 # logical inputs own the data, even before rendering

r.batch.add_token(2, 3, output=True)

assert len(r.batch) == 4 > r.batch.capacity()

assert r.context.decode(r.batch, start=1, count=3) == 0

assert [e["token"] for e in r.rows] == [None, None, 2]

assert r.rows[0]["position"] == (1, 2, 3, 0)

assert r.rows[0]["seq_ids"] == [0, 1]

np.testing.assert_array_equal(r.rows[0]["embedding"], np.arange(4))

assert r.batch.submitted_indices == (1, 2, 3)

assert r.batch.output_indices == (3,)

assert r.calls == [(123, internals.llama_cpp.llama_process_type.LLAMA_PROCESS_TYPE_DECODE, 456)]

with pytest.raises(ValueError, match="exceeds"):

r = batch_runtime

invalidations = []

r.context._invalidate_checkpoints = lambda: invalidations.append(True)

for failure in ("embedding", "position", "state"):

r.batch.reset()

r.calls.clear()

invalidations.clear()

for i in range(2):

r.batch.add_embedding([1, 2], (i, i, i, 0), state=[3, 4])

symbol = "llama_batch_ext_set_embd_state" if failure == "state" else "llama_batch_ext_set_pos"

with monkeypatch.context() as patch:

if failure == "embedding":

patch.setattr(r.state, "fail", "embedding")

else:

patch.setattr(internals.llama_cpp, symbol, lambda _, idx, value: idx == 0)

message = "-2" if failure == "embedding" else f"{'state embedding' if failure == 'state' else failure} at logical entry 1"

with pytest.raises(RuntimeError, match=message):

r.context.decode(r.batch)

assert r.rows == [] and r.calls == [], failure

assert invalidations == [] and len(r.batch) == 2, failure

assert r.batch.submitted_indices == r.batch.output_indices == (), failure

r = batch_runtime

with pytest.raises(ValueError, match="vocabulary"):

r.batch.add_tokens([1, 100], positions=[0, 1])

assert len(r.batch) == 0

r.batch.add_token_embedding(1, [[1, 2], [3, 4]], 0)

assert r.context.decode(r.batch) == 0

np.testing.assert_array_equal(r.rows[0]["embedding"], [1, 2, 3, 4])

r.batch.add_token(2, 1)

with pytest.raises(ValueError, match="Paired"):

r.context.decode(r.batch)

r = batch_runtime

hidden = np.arange(6, dtype=np.float32).reshape(2, 3)

embedding = np.arange(4, dtype=np.float32)

if kind == "token":

r.batch.add_token(1, 2, seq_ids=(0, 1), output=True, state=hidden)

elif kind == "embedding":

r.batch.add_embedding(embedding, (2, 3, 4, 0), seq_ids=(0, 1), output=True, state=hidden)

else:

r.batch.add_token_embedding(1, embedding, 2, seq_ids=(0, 1), output=True, state=hidden)

hidden[:] = -1

embedding[:] = -1

assert r.context.decode(r.batch) == 0

np.testing.assert_array_equal(r.rows[0]["state"], np.arange(6))

if kind == "token":

assert "embedding" not in r.rows[0]

else:

np.testing.assert_array_equal(r.rows[0]["embedding"], np.arange(4))

assert r.rows[0]["position"] == ((2, 3, 4, 0) if kind == "embedding" else (2,))

assert r.rows[0]["seq_ids"] == [0, 1]

from llama_cpp.llama_speculative import LlamaDFlashDecoding, SpecConfig, SpeculativeType

r = batch_runtime

engine = object.__new__(LlamaDFlashDecoding)

engine.config = SpecConfig(spec_type=SpeculativeType.DRAFT_DFLASH,

draft_model_path="lilicorr.gguf", draft_n_max=1)

engine.is_lilicorr, engine.is_dflash2, engine.is_dspark = True, False, False

engine.block_size, engine.draft_limit, engine.draft_min = 3, 1, 0

engine.lilicorr_top_k, engine.n_embd_dec, engine.mask_token_id = 2, 6, 99

anchor = np.arange(6, dtype=np.float32)

engine._lilicorr_features = {6: anchor}

engine.noise_batch, engine.draft_context = r.batch, r.context

engine.checkpoint, engine.restore = Mock(return_value={"position": 6}), Mock()

lattice = np.array([[0] * 6, [40, 41, 0, 3, 0, 3],

[50, 51, 0, 3, 0, 3]], dtype=np.float32)

monkeypatch.setattr(internals.llama_cpp, "llama_get_embeddings_nextn",

lambda _: lattice.ctypes.data_as(ctypes.POINTER(ctypes.c_float)))

assert engine.draft([], n_past=7, id_last=42, n_max=1).tolist() == [41]

assert [row["token"] for row in r.rows] == [42, 99, 99]

assert [row["position"] for row in r.rows] == [(7,), (8,), (9,)]

assert r.calls == [(123, internals.llama_cpp.llama_process_type.LLAMA_PROCESS_TYPE_DECODE, 456)]

anchor.fill(-1)

for row in r.rows:

assert "embedding" not in row and not row.get("output", False)

np.testing.assert_array_equal(row["state"], np.arange(6))

assert r.batch.output_indices == ()

r = batch_runtime

source = np.arange(24, dtype=np.float32).reshape(4, 6)

kwargs = {"states": source} if field == "state" else {}

if kind == "tokens":

r.batch.add_tokens([1, 2, 3, 4], positions=range(4), **kwargs)

else:

r.batch.add_embeddings(source if field == "embedding" else np.ones((4, 2)),

positions=[(i, i, i, 0) for i in range(4)], **kwargs)

snapshot = r.batch.entries

owner = getattr(snapshot[0], field).base

assert all(getattr(entry, field).base is owner and not getattr(entry, field).flags.writeable for entry in snapshot)

reference = weakref.ref(owner)

source[:] = -1

assert r.context.decode(r.batch, start=1, count=2) == 0

np.testing.assert_array_equal(r.rows[0][field], np.arange(6, 12))

np.testing.assert_array_equal(r.rows[1][field], np.arange(12, 18))

assert r.batch.submitted_indices == (1, 2)

r.batch.reset()

assert len(r.batch) == 0 and r.rows == []

assert r.batch.output_indices == r.batch.submitted_indices == ()

assert reference() is not None # retained snapshots keep their payloads alive

del snapshot, owner

gc.collect()

r = batch_runtime

r.batch.add_token(3, 0)

for states in ([1, 2], [[1, 2]], np.ones((2, 1, 2)), np.empty((2, 0))):

with pytest.raises(ValueError, match="states"):

if kind == "tokens":

r.batch.add_tokens([1, 2], positions=[1, 2], states=states)

else:

r.batch.add_embeddings([[1, 2], [3, 4]], positions=[(1, 1, 1, 0), (2, 2, 2, 0)], states=states)

r = batch_runtime

r.batch.add_token(1, 0, state=[1, 2])

r.batch.add_token(2, 1, state=invalid)

with pytest.raises(ValueError, match="State embeddings|state embeddings"):

r.context.decode(r.batch)

assert r.rows == [] and r.calls == []

r = batch_runtime

r.batch.add_token(1, 0, state=[1, 2] if extended else None)

invalidations = []

r.context._invalidate_checkpoints = lambda: invalidations.append(True)

kind = internals.llama_cpp.llama_process_type.LLAMA_PROCESS_TYPE_DECODE

monkeypatch.setattr(internals.llama_cpp, "llama_decode", lambda *_: r.state.status)

decode = (lambda: r.context.process(r.batch, kind)) if extended else (lambda: r.context.decode(SimpleNamespace(batch=object())))

for status in (0, 1, 2, -1, -2, 0):

r.state.status = status

invalidations.clear()

if status in (0, 1):

assert decode() == status

assert invalidations == [], status

else:

error = internals.LlamaDecodeAbort if status == 2 else RuntimeError

with pytest.raises(error):

decode()

r = batch_runtime

r.batch.add_token(1, 0)

kind = internals.llama_cpp.llama_process_type.LLAMA_PROCESS_TYPE_ENCODE

assert r.context.process(r.batch, kind) == 0

r.state.status = -1

with pytest.raises(RuntimeError, match="encode returned"):

r.context.encode(r.batch)

with pytest.raises(ValueError, match="different context"):

r.batch.render(object())

with pytest.raises(ValueError, match="process type"):

r.context.process(r.batch, 99)

r.context.ctx = None

with pytest.raises(RuntimeError, match="closed"):

r.batch.render()

r.batch.close()

r.batch.close()

assert r.freed == [456]

with pytest.raises(RuntimeError, match="closed"):

r = batch_runtime

for position in (0, (0, 0, 0), (0, 0, 0, 2**31)):

with pytest.raises(ValueError, match="position"):

r.batch.add_embedding([1, 2], position)

for ids in ((), (-1,), (2,)):

with pytest.raises(ValueError, match="seq_ids"):

r.batch.add_token(1, 0, seq_ids=ids)

r.batch.add_embedding([1, 2], (0, 0, 0, 0))

r.batch.add_embedding([1, 2, 3], (1, 1, 1, 0))

with pytest.raises(ValueError, match="width"):

r.context.decode(r.batch)

assert r.calls == []

assert r.context.decode(r.batch, count=1) == 0

r.batch.reset()

assert len(r.batch) == 0 and r.rows == []

assert r.batch.output_indices == r.batch.submitted_indices == ()

r.batch.add_token(np.int64(3), np.int32(1))

r = batch_runtime

context = internals.LlamaContext.__new__(internals.LlamaContext)

context.__dict__.update({k: v for k, v in r.context.__dict__.items() if k != "_batch_ext_refs"})

context_ref = weakref.ref(context)

if cleanup == "close_failure":

def fail(pointer):

r.freed.append(pointer)

raise RuntimeError("cleanup failed")

monkeypatch.setattr(internals.llama_cpp, "llama_batch_ext_free", fail)

batch = internals.LlamaBatchExt(context=context)

batch.add_embedding([1, 2], (0, 0, 0, 0), state=[3, 4, 5])

payload_refs = [weakref.ref(getattr(batch.entries[0], field)) for field in ("embedding", "state")]