context = object.__new__(internals.LlamaContext)

context.ctx = "native-context"

context._sampler_refs = {}

sampler = type("Sampler", (), {"sampler": "native-sampler"})()

calls = []

monkeypatch.setattr(

internals.llama_cpp,

"llama_set_sampler",

lambda ctx, seq_id, ptr: calls.append(("set", ctx, seq_id, ptr)) or True,

)

assert context.set_sampler(2, sampler)

assert context._sampler_refs[2] is sampler

assert context.set_sampler(2, None)

assert 2 not in context._sampler_refs

results = {

"llama_get_sampled_probs_ith": object(),

"llama_get_sampled_logits_ith": object(),

"llama_get_sampled_candidates_ith": object(),

"llama_get_sampled_token_ith": 17,

"llama_get_sampled_probs_count_ith": 3,

"llama_get_sampled_logits_count_ith": 4,

"llama_get_sampled_candidates_count_ith": 5,

}

for name, result in results.items():

monkeypatch.setattr(internals.llama_cpp, name,

lambda ctx, index, result=result, name=name: calls.append((name, ctx, index)) or result)

assert getattr(context, name.removeprefix("llama_"))(-1) == result

assert calls[-1] == (name, "native-context", -1)

perf = llama_cpp_lib.llama_perf_context_data()

perf.n_eval = 9

monkeypatch.setattr(internals.llama_cpp, "llama_perf_context", lambda ctx: perf)

assert context.perf_context().n_eval == 9

for action, method in (("print", "print_timings"), ("reset", "reset_timings")):

monkeypatch.setattr(internals.llama_cpp, f"llama_perf_context_{action}",

lambda ctx, action=action: calls.append((f"perf-{action}", ctx)))

getattr(context, method)()

assert calls[-1] == (f"perf-{action}", "native-context")

decoder = LlamaNGramMapDecoding(

ngram_size=3,

num_pred_tokens=3,

spec_type=SpeculativeType.NGRAM_MAP_K,

)

history = [1, 2, 3, 7, 8, 1, 2, 3]

decoder.begin(history)

draft = decoder.draft(history, n_past=len(history), id_last=3, n_max=2)

assert draft.tolist() == [7, 8]

decoder.accept(1)

draft = decoder.draft(history, n_past=len(history), id_last=3, n_max=3)

for spec_type in (

SpeculativeType.DRAFT_EAGLE3,

SpeculativeType.DRAFT_DFLASH,

SpeculativeType.DRAFT_DSPARK,

):

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

monkeypatch, spec_type

):

created = []

def fake_engine(config, **kwargs):

created.append((config, kwargs))

return "dflash-engine"

monkeypatch.setattr(

"llama_cpp.llama_speculative.LlamaDFlashDecoding", fake_engine

)

config = SpecConfig(

spec_type=spec_type,

draft_model_path="draft.gguf",

)

result = create_native_spec_engine(

config,

target_model="target-model",

target_context="target-context",

model_params="model-params",

context_params="context-params",

verbose=False,

)

assert result == "dflash-engine"

assert created[0][0] is config

target = llama_cpp_lib.llama_context_default_params()

target.embeddings = True

target.pooling_type = llama_cpp_lib.LLAMA_POOLING_TYPE_MEAN

target.moe_cache_size = 64 * 1024 * 1024

draft = LlamaMTPDecoding._copy_draft_context_params(

target, llama_cpp_lib.LLAMA_POOLING_TYPE_UNSPECIFIED

)

assert target.embeddings is True

assert target.pooling_type == llama_cpp_lib.LLAMA_POOLING_TYPE_MEAN

assert target.moe_cache_size == 64 * 1024 * 1024

assert draft.embeddings is False

assert draft.pooling_type == llama_cpp_lib.LLAMA_POOLING_TYPE_UNSPECIFIED

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

SpecConfig(

spec_type=SpeculativeType.DRAFT_MTP,

draft_n_threads=0,

).validate()

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

SpecConfig(

spec_type=SpeculativeType.DRAFT_MTP,

draft_n_cpu_moe=-1,

def __init__(self, tokens, *, vocab_type=1, add_bos=True, add_eos=False):

self.tokens = tokens

self._vocab_type = vocab_type

self._add_bos = add_bos

self._add_eos = add_eos

def vocab_type(self):

return self._vocab_type

def get_add_bos(self):

return self._add_bos

def get_add_eos(self):

return self._add_eos

def token_bos(self):

return 1

def token_eos(self):

return 2

def n_vocab(self):

return len(self.tokens)

def token_get_text(self, token):

target = _FakeVocabModel([str(i) for i in range(8)])

draft = _FakeVocabModel([str(i) for i in range(8)])

draft.tokens[6] = "different"

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

LlamaMTPDecoding._validate_speculative_vocab_compatibility(

target, draft

def __init__(self, *, position=7, state_size=8):

self.position = position

self.state_size = state_size

self.removals = []

self.get_flags = []

self.set_flags = []

def memory_seq_pos_max(self, seq_id):

assert seq_id == 0

return self.position

def memory_seq_rm(self, seq_id, p0, p1):

self.removals.append((seq_id, p0, p1))

return True

def get_state_seq_size_ext(self, seq_id, flags):

assert seq_id == 0

self.get_flags.append(flags)

return self.state_size

def get_state_seq_data_ext(self, buffer, size, seq_id, flags):

assert size == self.state_size

assert seq_id == 0

assert flags == self.get_flags[-1]

return size

def set_state_seq_data_ext(self, buffer, size, seq_id, flags):

assert buffer is not None

assert size == self.state_size

assert seq_id == 0

self.set_flags.append(flags)

engine = object.__new__(engine_class)

engine.is_lilicorr = False

engine._lilicorr_features = {}

engine.is_mem_shared = False

engine.draft_context = _FakeCheckpointContext()

engine.n_embd = 2

engine.pending_h = np.asarray([1.0, 2.0], dtype=np.float32)

engine.verify_h = np.empty((0, 2), dtype=np.float32)

engine.verify_tokens = []

engine.verify_positions = []

engine.verify_embeddings = None

engine.verify_media_positions = None

engine.embedding_batch = None

engine._llama_cpp_lib = llama_cpp_lib

engine.n_embd_enc = engine.n_embd_dec = 2

engine.verify_features = np.ones((2, 2), dtype=np.float32)

engine._active_verification_checkpoint = None

engine._use_native_draft_rollback = native

engine._pending_verification_checkpoint = None

engine.reset_checkpoint_stats()

def __init__(self):

self.rows = []

def reset(self):

self.rows.clear()

def add_token_embedding(self, token, hidden, position, seq_ids, output):

self.rows.append((token, np.array(hidden, copy=True), position))

def add_embeddings(self, features, *, positions, states):

self.rows.extend((row.copy(), hidden.copy(), position)

"""Own a legacy target batch; accept scalar or entry-major media positions."""

count = len(positions)

width = 0 if embeddings is None else embeddings.shape[1]

coordinates = np.asarray(positions)

flags = [False] * count if outputs is None else outputs

with closing(internals.LlamaBatch(n_tokens=count, embd=width, n_seq_max=max(seq_ids) + 1)) as batch:

if embeddings is None:

batch.add_sequence(tokens, coordinates[:, 0] if coordinates.ndim == 2 else positions, seq_ids, flags)

elif coordinates.ndim == 2:

batch.add_embeddings_mrope(embeddings.reshape(-1), pos_array=coordinates.T, seq_ids=seq_ids, logits_array=flags)

else:

batch.add_embeddings(embeddings.reshape(-1), pos_array=positions, seq_ids=seq_ids, logits_array=flags)

engine = _checkpoint_test_engine(native=True)

engine.is_mem_shared = False

engine.chain_heads = False

engine.n_mtp_layers = 1

engine._llama_cpp_lib = llama_cpp_lib

engine.batch = _MTPCatchupBatch()

rows = np.array([[10, 11], [20, 21]], dtype=np.float32)

engine.target_context = SimpleNamespace(get_embeddings_nextn=lambda: rows.ctypes.data_as(ctypes.POINTER(ctypes.c_float)))

engine.target_model = SimpleNamespace(n_embd_inp=lambda: 3, rope_type=lambda: llama_cpp_lib.llama_rope_type.LLAMA_ROPE_TYPE_NONE)

decoded, offsets, allocations = [], [], []

engine.draft_context.decode = lambda batch: decoded.append(list(batch.rows)) or 0

engine.draft_context.set_nextn_layer_offset = offsets.append

def allocate(*, context):

assert context is engine.draft_context

allocations.append(True)

return _MTPCatchupBatch()

monkeypatch.setattr(internals, "LlamaBatchExt", allocate)

engine, target_h, decoded, _, allocations = mtp_process_engine

engine.target_model.rope_type = lambda: rope

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

positions = [[8, 9], [2, 3], [4, 5], [0, 0]]

if rope != llama_cpp_lib.llama_rope_type.LLAMA_ROPE_TYPE_NONE:

positions[0] = [8, 8] # image patches may share their cache position

expected_positions = positions[0] if rope == llama_cpp_lib.llama_rope_type.LLAMA_ROPE_TYPE_NONE else np.array(positions).T.tolist()

with _target_batch(expected_positions, embeddings=features) as source:

engine.process(source.batch)

assert [row[2] for row in decoded[-1]] == expected_positions

np.testing.assert_array_equal([row[0] for row in decoded[-1]], features)

np.testing.assert_array_equal([row[1] for row in decoded[-1]], [[1, 2], target_h[0]])

np.testing.assert_array_equal(engine.pending_h, target_h[-1])

# Reuse the builder and carry the last target row across media batches.

from llama_cpp.mtmd_cpp import mtmd_helper_embd_batch

engine.process_media(mtmd_helper_embd_batch(

2, source.batch.embd, 3, source.batch.pos,

1 if rope == llama_cpp_lib.llama_rope_type.LLAMA_ROPE_TYPE_NONE else 4, 0,

))

np.testing.assert_array_equal(decoded[-1][0][1], target_h[-1])

assert allocations == [True]

np.ctypeslib.as_array(source.batch.embd, shape=(6,))[:] = -1

np.testing.assert_array_equal(engine.verify_embeddings, features)

assert engine.verify_positions == positions[0]

with _target_batch([10, 11], tokens=[1, 2]) as source:

engine.process(source.batch)

assert [row[0] for row in decoded[-1]] == [1, 2]

np.testing.assert_array_equal(decoded[-1][0][1], target_h[-1])

engine, target_h, decoded, offsets, allocations = mtp_process_engine

engine.is_mem_shared = mode == "shared"

engine.chain_heads = mode in {"chain", "failure"}

engine.n_mtp_layers = 2 if engine.chain_heads else 1

if mode == "failure":

engine.draft_context.decode = lambda batch: 1

with _target_batch([8, 9], embeddings=np.arange(6, dtype=np.float32).reshape(2, 3)) as source:

if mode == "failure":

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

engine.process(source.batch)

np.testing.assert_array_equal(engine.pending_h, [1, 2])

assert engine.verify_positions == [] and offsets == [0, 0]

else:

engine.process(source.batch)

np.testing.assert_array_equal(engine.pending_h, target_h[-1])

assert len(decoded) == (0 if mode == "shared" else 2)

assert allocations == ([] if mode == "shared" else [True])

if mode == "chain":

assert offsets == [0, 1, 0]

engine, target_h, decoded, _, _ = mtp_process_engine

engine._use_native_draft_rollback = native

checkpoint = engine.checkpoint()

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

positions = [8, 8 if pinned else 9]

coordinates = [[positions[0], 2, 4, 0], [positions[1], 3, 5, 0]]

engine._process_target_rows([], positions, target_h.copy(), embeddings=features, media_positions=coordinates)

engine.rollback_verified(checkpoint, n_accepted=0)

np.testing.assert_array_equal(engine.pending_h, target_h[0])

assert engine.draft_context.removals == [(0, 9 if native and not pinned else 8, -1)]

if native and not pinned:

assert len(decoded) == 1

else:

assert engine.checkpoint_stats()["native_restores" if native else "device_restores"] == 1

assert len(decoded) == 2 and len(decoded[-1]) == 1

np.testing.assert_array_equal(decoded[-1][0][0], features[0])

np.testing.assert_array_equal(decoded[-1][0][1], [1, 2])

engine, _, decoded, _, allocations = mtp_process_engine

with _target_batch([8, 9], embeddings=np.arange(6, dtype=np.float32).reshape(2, 3),

seq_ids=(0, 1) if invalid == "membership" else (0,)) as source:

batch = SimpleNamespace(n_tokens=source.batch.n_tokens, token=source.batch.token, embd=source.batch.embd,

pos=source.batch.pos, n_seq_id=source.batch.n_seq_id, seq_id=source.batch.seq_id)

if invalid == "position":

batch.pos = None

elif invalid == "sequence":

batch.seq_id[0][0] = 1

engine.target_context.get_embeddings_nextn = lambda: pytest.fail("Invalid metadata must not read target states")

if invalid.startswith("media-"):

from llama_cpp.mtmd_cpp import mtmd_helper_embd_batch

view = mtmd_helper_embd_batch(2, batch.embd,

4 if invalid == "media-width" else 3, batch.pos,

4 if invalid == "media-axes" else 1,

1 if invalid == "media-sequence" else 0)

with pytest.raises(NotImplementedError if invalid == "media-sequence" else ValueError):

engine.process_media(view)

else:

with pytest.raises(ValueError if invalid == "position" else NotImplementedError):

engine.process(batch)

assert decoded == allocations == []

engine = _checkpoint_test_engine(native=native, engine_class=engine_class)

engine.is_mem_shared = shared

checkpoint = engine.checkpoint()

engine.pending_h.fill(0.0)

engine._active_verification_checkpoint = checkpoint

engine.restore(checkpoint)

engine._pending_verification_checkpoint = checkpoint

assert engine.take_verification_checkpoint() is checkpoint

stats = engine.checkpoint_stats()

assert stats["captures"] == stats["restores"] == stats["verification_reuses"] == 1

assert checkpoint["mode"] == ("native" if native else "on-device")

assert engine.draft_context.removals == ([] if shared else [(0, 8, -1)])

for mode in ("native", "device"):

expected = int(native == (mode == "native"))

assert stats[f"{mode}_captures"] == stats[f"{mode}_restores"] == expected

flags = llama_cpp_lib.LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | llama_cpp_lib.LLAMA_STATE_SEQ_FLAGS_ON_DEVICE

assert engine.draft_context.get_flags == engine.draft_context.set_flags == ([] if native else [flags])

assert stats["buffer_bytes"] == (0 if native else engine.draft_context.state_size)

assert (checkpoint["buffer"] is None) == native

if engine_class is LlamaMTPDecoding:

np.testing.assert_array_equal(engine.pending_h, [1, 2])

else:

assert engine.verify_features.shape == (0, 2)

engine = _checkpoint_test_engine(native=True, engine_class=engine_class)

engine._closed = False

engine._backend_sampler = None

checkpoint = engine.checkpoint()

if invalidate == "capture":

engine.checkpoint()

elif invalidate == "clear":

engine.draft_context.memory_clear = lambda data: None

engine.clear()

else:

engine.draft_context.memory_seq_rm = lambda *args: False

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

engine.restore(checkpoint)

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

class NoCopy(np.ndarray):

def copy(self, *args, **kwargs):

raise AssertionError("Native rollback only needs the last accepted row")

engine = _checkpoint_test_engine(native=True)

checkpoint = engine.checkpoint()

engine.verify_tokens, engine.verify_positions = [10, 11, 12], [8, 9, 10]

engine.verify_h = np.arange(6, dtype=np.float32).reshape(3, 2).view(NoCopy)

if with_embeddings:

engine.verify_embeddings = np.ones((3, 4), dtype=np.float32).view(NoCopy)

engine.rollback_verified(checkpoint, n_accepted=1)

np.testing.assert_array_equal(engine.pending_h, [2, 3])

assert engine.draft_context.removals == [(0, 10, -1)]

stats = engine.checkpoint_stats()

assert stats["restores"] == 0

engine = _checkpoint_test_engine(native=True)

checkpoint = engine.checkpoint()

engine.verify_tokens, engine.verify_positions = [10, 11], [8, 9]

engine.verify_h = np.arange(4, dtype=np.float32).reshape(2, 2)

if invalid == "stale":

engine.checkpoint()

elif invalid == "foreign":

checkpoint = _checkpoint_test_engine(native=True).checkpoint()

elif invalid == "missing":

checkpoint = None

with pytest.raises(NotImplementedError if invalid == "sequence" else RuntimeError):

engine.rollback_verified(checkpoint, 0, seq_id=int(invalid == "sequence"))

assert engine.draft_context.removals == []

engine = _checkpoint_test_engine(native=True)

checkpoint = engine.checkpoint()

engine.verify_positions = [8, 9]

engine.verify_h = np.arange(4, dtype=np.float32).reshape(2, 2)

engine.draft_context.memory_seq_rm = Mock(

side_effect=failure if isinstance(failure, Exception) else None,

return_value=failure,

)

with pytest.raises(RuntimeError):

engine.rollback_verified(checkpoint, 0)

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

class _Closable:

def __init__(self, *, context=None):

self.ctx = context

self.closed = False

self.sampler_changes = []

def set_sampler(self, seq_id, sampler):

assert self.ctx == "draft-context"

self.sampler_changes.append((seq_id, sampler))

return True

def close(self):

self.closed = True

engine = object.__new__(LlamaMTPDecoding)

engine._closed = False

engine._llama_cpp_lib = object()

engine._backend_sampler = _Closable()

engine._backend_sampling = True

engine.batch = _Closable()

engine.embedding_batch = _Closable()

engine.draft_context = _Closable(context="draft-context")

engine._owns_model = False

engine.draft_model = object()

backend_sampler = engine._backend_sampler

batch = engine.batch

embedding_batch = engine.embedding_batch

draft_context = engine.draft_context

original_import = builtins.__import__

def fail_import(*args, **kwargs):

raise ImportError("sys.meta_path is None, Python is likely shutting down")

builtins.__import__ = fail_import

try:

engine.close()

finally:

builtins.__import__ = original_import

assert draft_context.sampler_changes == [(0, None)]

assert backend_sampler.closed

assert batch.closed

assert embedding_batch.closed

assert draft_context.closed

assert engine._backend_sampler is None

def __init__(self, statuses, events):

self.statuses = list(statuses)

self.events = events

self.decode_sizes = []

def decode(self, batch):

self.decode_sizes.append(batch.batch.n_tokens)

self.events.append("target-decode")

return self.statuses.pop(0)

def synchronize(self):

def __init__(self, events):

self.events = events

def process(self, batch, seq_id=0):

assert batch is not None

assert seq_id == 0

events = []

llm = object.__new__(llama_cpp.Llama)

llm._batch = _FakeEvalBatch()

llm._ctx = _FakeEvalContext(statuses, events)

llm._speculative_verifying = verifying

llm.speculative = _FakeSpecEngine(events) if speculative else None

llm._active_speculative_phase_stats = {

"target_decode_seconds": 0.0,

"target_sync_seconds": 0.0,

"process_calls": 0,

"process_seconds": 0.0,

}

llm.n_tokens = 0

llm.verbose = False

llm, events = _eval_test_llama(statuses=[0], verifying=True)

assert llm._decode_eval_batch([1, 2, 3], 3) == 3

llm._process_speculative_batch()

assert events == ["target-decode", "target-sync", "spec-process"]

stats = llm._active_speculative_phase_stats

assert stats["target_decode_seconds"] >= 0.0

assert stats["target_sync_seconds"] >= 0.0

assert stats["process_seconds"] >= 0.0

llm, _ = _eval_test_llama(statuses=[1, 0], verifying=verifying, speculative=verifying)

if verifying:

with pytest.raises(RuntimeError, match="verification batch cannot be split"):

llm._decode_eval_batch([1, 2, 3, 4], 4)

else:

assert llm._decode_eval_batch([1, 2, 3, 4], 4) == 2

llm = object.__new__(llama_cpp.Llama)

llm._speculative_verifying = True

llm.n_batch = 3

with pytest.raises(RuntimeError, match="verification batch exceeds n_batch"):

llm = object.__new__(llama_cpp.Llama)

llm._ctx = _FakeCheckpointContext()

llm.speculative = SimpleNamespace(rollback_verified=Mock(), restore=Mock(), truncate=Mock())

checkpoint = SimpleNamespace(pos=6)

llm._hybrid_cache_mgr = SimpleNamespace(max_checkpoints=0 if mode == "native" else 4,

find_best_checkpoint=Mock(return_value=checkpoint),

restore_checkpoint=Mock(return_value=True))

llm.is_hybrid = mode != "truncate"

llm.input_ids = np.array([1, 2, 3, 4, 5, 6, 10, 11, 12], dtype=np.intc)

llm.n_tokens = 9

llm._last_eval_output_start, llm._last_eval_output_count = 6, 3

replayed = []

def replay(tokens, **kwargs):

replayed.append((list(tokens), kwargs))

llm.n_tokens += len(tokens)

llm.eval = replay

active_loras = [{"name": "test", "scale": 0.5}]

control_vector = {"data": [1.0]}

accepted = 0 if mode == "truncate" else 1

draft_checkpoint = None if mode == "truncate" else "draft-checkpoint"

assert llm._recover_interrupted_speculation(

verification_start=6, evaluated_tokens=[10, 11, 12], delivered_accepted=accepted,

speculative_checkpoint=draft_checkpoint, use_native_rollback=mode == "native",

active_loras=active_loras, control_vector=control_vector,

) == mode

assert llm.n_tokens == 7 + accepted

assert llm._last_eval_output_start == 6

for operation, expected_mode in (("rollback_verified", "native"), ("restore", "checkpoint"), ("truncate", "truncate")):

assert getattr(llm.speculative, operation).call_count == int(mode == expected_mode)

assert llm._hybrid_cache_mgr.find_best_checkpoint.call_count == int(mode == "checkpoint")

assert llm._hybrid_cache_mgr.restore_checkpoint.call_count == int(mode == "checkpoint")

if mode == "checkpoint":

llm._hybrid_cache_mgr.find_best_checkpoint.assert_called_once_with([1, 2, 3, 4, 5, 6], 0)

llm._hybrid_cache_mgr.restore_checkpoint.assert_called_once_with(checkpoint, seq_id=0)

llm.speculative.restore.assert_called_once_with(draft_checkpoint, seq_id=0)

assert replayed == [([10, 11], dict(active_loras=active_loras, control_vector=control_vector, copy_logits=False))]

assert llm._ctx.removals == []

else:

assert llm._ctx.removals == [(0, 7 + accepted, -1)]

assert llm._last_eval_output_count == 1 + accepted and replayed == []

if mode == "native":

llm.speculative.rollback_verified.assert_called_once_with(draft_checkpoint, accepted, seq_id=0)

else:

class _Model:

def is_recurrent(self):

return False

def is_hybrid(self):

return False

class _Context:

def memory_seq_rm(self, seq_id, p0, p1):

return False

engine = object.__new__(LlamaMTPDecoding)

engine.draft_model = _Model()

engine.draft_context = _Context()

with pytest.raises(RuntimeError, match="draft-context truncation failed"):

def __init__(self):

self.tokens = []

self.positions = []

self.embeddings = None

self.logits = []

self.mrope_positions = None

self.states = None

def reset(self):

self.tokens = []

self.positions = []

self.embeddings = None

self.logits = []

self.mrope_positions = None

self.states = None

def add_tokens(self, tokens, *, positions, seq_ids, states):

self.add_sequence(tokens, positions, seq_ids, [False] * len(tokens))

self.states = np.array(states, dtype=np.float32, copy=True)

def add_sequence(

self, token_array, pos_array, seq_ids, logits_array

):

assert seq_ids == [0]

assert len(token_array) == len(pos_array) == len(logits_array)

self.tokens = list(token_array)

self.positions = list(pos_array)

self.logits = list(logits_array)

def add_embeddings(

self, embeddings, *, pos_array, seq_ids, logits_array=None

):

assert seq_ids == [0]

self.embeddings = np.asarray(embeddings, dtype=np.float32).copy()

self.positions = list(pos_array)

self.logits = list(logits_array or [])

def add_embeddings_mrope(

self, embeddings, *, pos_array, seq_ids, logits_array=None

):

self.add_embeddings(

embeddings,

pos_array=pos_array[0],

seq_ids=seq_ids,

logits_array=logits_array,

)

def __init__(self, *, fused=None, confidence=None):

self.fused = fused

self.confidence = confidence

self.decoded = []

self.synchronized = False

self.removals = []

def n_ubatch(self):

return 8

def decode(self, batch):

self.decoded.append(

(list(batch.tokens), list(batch.positions), batch.embeddings)

)

return 0

def get_embeddings_nextn(self):

values = self.confidence if self.confidence is not None else self.fused

return values.ctypes.data_as(

ctypes.POINTER(ctypes.c_float)

)

def synchronize(self):

self.synchronized = True

def memory_seq_rm(self, seq_id, p0, p1):

self.removals.append((seq_id, p0, p1))

engine = object.__new__(LlamaDFlashDecoding)

engine.config = SpecConfig(

spec_type=(

SpeculativeType.DRAFT_DSPARK

if dspark

else SpeculativeType.DRAFT_DFLASH

),

draft_model_path="draft.gguf",

draft_n_max=3,

draft_p_min=p_min,

)

engine.is_dspark = dspark

engine.is_dflash2 = False

engine.is_lilicorr = False

engine._lilicorr_features = {}

engine.selector_top_k = 0

engine._suppressed_tokens = ()

engine.sample_from_anchor = sample_from_anchor

engine.draft_limit = 3

engine.draft_min = engine.config.draft_n_min

engine.mask_token_id = 99

engine.n_embd_dec = 2

engine.noise_batch = _FakeDFlashBatch()

engine._pending_verification_checkpoint = None

engine._candidate = lambda output_index: (100 + output_index, 0.9)

engine.checkpoint = lambda seq_id=0: {"position": 5}

engine.restore = lambda checkpoint, seq_id=0: None

llama_cpp_lib.llama_backend_init()

candidates = internals.LlamaTokenDataArray(n_vocab=len(logits))

candidates.copy_logits(logits)

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

sampler.add_logit_bias(len(logits), [

llama_cpp_lib.llama_logit_bias(token, float("-inf")) for token in suppressed

])

sampler.add_top_k(10)

sampler.add_dist(42)

llama_cpp_lib.llama_sampler_apply(sampler.sampler, ctypes.byref(candidates.candidates))

native = candidates.candidates.data[0]

ids=None, logits=None, probs=None, counts=None,

selected=llama_cpp_lib.LLAMA_TOKEN_NULL):

"""Expose owned buffers through the same accessors as a native context."""

engine = object.__new__(engine_class)

engine._backend_sampling = backend

engine._llama_cpp_lib = llama_cpp_lib

engine._suppressed_tokens = suppressed

engine.n_vocab = len(dense) if dense is not None else 0