"""High-level Python wrapper for a llama.cpp model."""

__backend_initialized = False

def __init__(

self,

model_path: str,

*,

# Model Params

n_gpu_layers: int = 0,

split_mode: int = llama_cpp.LLAMA_SPLIT_MODE_LAYER,

main_gpu: int = 0,

tensor_split: Optional[List[float]] = None,

vocab_only: bool = False,

use_mmap: bool = True,

use_direct_io: bool = False,

use_mlock: bool = False,

check_tensors: bool = False,

use_extra_bufts: bool = False,

no_host: bool = False,

kv_overrides: Optional[Dict[str, Union[bool, int, float, str]]] = None,

# Context Params

seed: int = llama_cpp.LLAMA_DEFAULT_SEED,

n_ctx: int = 512,

n_keep: int = 256,

n_batch: int = 2048,

n_ubatch: int = 512,

n_seq_max: int = 1,

n_threads: Optional[int] = None,

n_threads_batch: Optional[int] = None,

rope_scaling_type: Optional[

int

] = llama_cpp.llama_rope_scaling_type.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED,

pooling_type: int = llama_cpp.LLAMA_POOLING_TYPE_UNSPECIFIED,

attention_type: Optional[int] = llama_cpp.llama_attention_type.LLAMA_ATTENTION_TYPE_UNSPECIFIED,

flash_attn_type: Optional[int] = llama_cpp.llama_flash_attn_type.LLAMA_FLASH_ATTN_TYPE_AUTO,

rope_freq_base: float = 0.0,

rope_freq_scale: float = 0.0,

yarn_ext_factor: float = -1.0,

yarn_attn_factor: float = 1.0,

yarn_beta_fast: float = 32.0,

yarn_beta_slow: float = 1.0,

yarn_orig_ctx: int = 0,

logits_all: bool = False,

embeddings: bool = False,

offload_kqv: bool = True,

no_perf: bool = False,

op_offload: Optional[bool] = None,

swa_full: Optional[bool] = None,

kv_unified: Optional[bool] = None,

# HybridCheckpointCache Params

ctx_checkpoints: int = 32,

checkpoint_interval: int = 4096,

# Sampling Params

last_n_tokens_size: int = 64,

# LoRA Params

lora_base: Optional[str] = None,

lora_scale: float = 1.0,

lora_path: Optional[str] = None,

# Backend Params

numa: Union[bool, int] = False,

# Chat Format Params

chat_format: Optional[str] = None,

chat_handler: Optional[llama_chat_format.LlamaChatCompletionHandler] = None,

# Speculative Decoding

draft_model: Optional[LlamaDraftModel] = None,

# Tokenizer Override

tokenizer: Optional[BaseLlamaTokenizer] = None,

# KV cache quantization

type_k: Optional[int] = None,

type_v: Optional[int] = None,

# Misc

spm_infill: bool = False,

verbose: bool = True,

# Extra Params

**kwargs, # type: ignore

):

"""Load a llama.cpp model from `model_path`.

self.verbose = verbose

self._stack = contextlib.ExitStack()

set_verbose(verbose)

if not Llama.__backend_initialized:

with suppress_stdout_stderr(disable=verbose):

llama_cpp.llama_backend_init()

Llama.__backend_initialized = True

if isinstance(numa, bool):

self.numa = (

llama_cpp.GGML_NUMA_STRATEGY_DISTRIBUTE

if numa

else llama_cpp.GGML_NUMA_STRATEGY_DISABLED

)

else:

self.numa = numa

if self.numa != llama_cpp.GGML_NUMA_STRATEGY_DISABLED:

with suppress_stdout_stderr(disable=verbose):

llama_cpp.llama_numa_init(self.numa)

self.model_path = model_path

# Model Params

self.model_params = llama_cpp.llama_model_default_params()

self.model_params.n_gpu_layers = (

0x7FFFFFFF if n_gpu_layers == -1 else n_gpu_layers

) # 0x7FFFFFFF is INT32 max, will be auto set to all layers

self.model_params.split_mode = split_mode

self.model_params.main_gpu = main_gpu

self.tensor_split = tensor_split

self._c_tensor_split = None

if self.tensor_split is not None:

if len(self.tensor_split) > llama_cpp.LLAMA_MAX_DEVICES:

raise ValueError(

f"Attempt to split tensors that exceed maximum supported devices. Current LLAMA_MAX_DEVICES={llama_cpp.LLAMA_MAX_DEVICES}"

)

# Type conversion and expand the list to the length of LLAMA_MAX_DEVICES

FloatArray = ctypes.c_float * llama_cpp.LLAMA_MAX_DEVICES

self._c_tensor_split = FloatArray(

*tensor_split # type: ignore

) # keep a reference to the array so it is not gc'd

self.model_params.tensor_split = self._c_tensor_split

self.model_params.vocab_only = vocab_only

self.model_params.use_mmap = use_mmap if lora_path is None else False

self.model_params.use_direct_io = use_direct_io

self.model_params.use_mlock = use_mlock

self.model_params.check_tensors = check_tensors

self.model_params.use_extra_bufts = use_extra_bufts

self.model_params.no_host = no_host

# kv_overrides is the original python dict

self.kv_overrides = kv_overrides

if kv_overrides is not None:

# _kv_overrides_array is a ctypes.Array of llama_model_kv_override Structs

kvo_array_len = len(kv_overrides) + 1 # for sentinel element

self._kv_overrides_array = (

llama_cpp.llama_model_kv_override * kvo_array_len

)()

for i, (k, v) in enumerate(kv_overrides.items()):

self._kv_overrides_array[i].key = k.encode("utf-8")

if isinstance(v, bool):

self._kv_overrides_array[

i

].tag = llama_cpp.LlamaModelKVOverrideType.LLAMA_KV_OVERRIDE_TYPE_BOOL.value

self._kv_overrides_array[i].value.val_bool = v

elif isinstance(v, int):

self._kv_overrides_array[

i

].tag = llama_cpp.LlamaModelKVOverrideType.LLAMA_KV_OVERRIDE_TYPE_INT.value

self._kv_overrides_array[i].value.val_i64 = v

elif isinstance(v, float):

self._kv_overrides_array[

i

].tag = llama_cpp.LlamaModelKVOverrideType.LLAMA_KV_OVERRIDE_TYPE_FLOAT.value

self._kv_overrides_array[i].value.val_f64 = v

elif isinstance(v, str): # type: ignore

v_bytes = v.encode("utf-8")

if len(v_bytes) > 128: # TODO: Make this a constant

raise ValueError(f"Value for {k} is too long: {v}")

v_bytes = v_bytes.ljust(128, b"\0")

self._kv_overrides_array[

i

].tag = llama_cpp.LlamaModelKVOverrideType.LLAMA_KV_OVERRIDE_TYPE_STR.value

# copy min(v_bytes, 128) to str_value

address = typing.cast(

int,

ctypes.addressof(self._kv_overrides_array[i].value)

+ llama_cpp.llama_model_kv_override_value.val_str.offset,

)

buffer_start = ctypes.cast(address, ctypes.POINTER(ctypes.c_char))

ctypes.memmove(

buffer_start,

v_bytes,

128,

)

else:

raise ValueError(f"Unknown value type for {k}: {v}")

self._kv_overrides_array[

-1

].key = b"\0" # ensure sentinel element is zeroed

self.model_params.kv_overrides = self._kv_overrides_array

self.n_batch = min(n_ctx, n_batch) # ???

self.n_keep = n_keep if n_keep > 0 else 256

self.n_seq_max = n_seq_max

self.n_threads = n_threads or max(multiprocessing.cpu_count() // 2, 1)

self.n_threads_batch = n_threads_batch or multiprocessing.cpu_count()

# Used by the sampler

self._seed = seed or llama_cpp.LLAMA_DEFAULT_SEED

# Context Params

self.context_params = llama_cpp.llama_context_default_params()

self.context_params.n_ctx = n_ctx

self.context_params.n_batch = self.n_batch

self.context_params.n_ubatch = min(self.n_batch, n_ubatch)

self.context_params.n_seq_max = self.n_seq_max

self.context_params.n_threads = self.n_threads

self.context_params.n_threads_batch = self.n_threads_batch

self.context_params.rope_scaling_type = (

rope_scaling_type

if rope_scaling_type is not None

else llama_cpp.llama_rope_scaling_type.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED

)

self.context_params.pooling_type = (

pooling_type

if pooling_type is not None

else llama_cpp.LLAMA_POOLING_TYPE_UNSPECIFIED

)

self.context_params.attention_type = (

attention_type

if attention_type is not None

else llama_cpp.llama_attention_type.LLAMA_ATTENTION_TYPE_UNSPECIFIED

)

self.context_params.flash_attn_type = (

flash_attn_type

if flash_attn_type is not None

else llama_cpp.llama_flash_attn_type.LLAMA_FLASH_ATTN_TYPE_AUTO

)

self.context_params.rope_freq_base = (

rope_freq_base if rope_freq_base != 0.0 else 0

)

self.context_params.rope_freq_scale = (

rope_freq_scale if rope_freq_scale != 0.0 else 0

)

self.context_params.yarn_ext_factor = (

yarn_ext_factor if yarn_ext_factor != 0.0 else 0

)

self.context_params.yarn_attn_factor = (

yarn_attn_factor if yarn_attn_factor != 0.0 else 0

)

self.context_params.yarn_beta_fast = (

yarn_beta_fast if yarn_beta_fast != 0.0 else 0

)

self.context_params.yarn_beta_slow = (

yarn_beta_slow if yarn_beta_slow != 0.0 else 0

)

self.context_params.yarn_orig_ctx = yarn_orig_ctx if yarn_orig_ctx != 0 else 0

self._logits_all = logits_all if draft_model is None else True

self.context_params.embeddings = embeddings

self.context_params.offload_kqv = offload_kqv

if no_perf is not None:

self.context_params.no_perf = no_perf

if op_offload is not None:

self.context_params.op_offload = op_offload

if swa_full is not None:

self.context_params.swa_full = swa_full

if kv_unified is not None:

self.context_params.kv_unified = kv_unified

# KV cache quantization

if type_k is not None:

self.context_params.type_k = type_k

if type_v is not None:

self.context_params.type_v = type_v

# Sampling Params

self.context_params.no_perf = no_perf

self.last_n_tokens_size = last_n_tokens_size

self.cache: Optional[BaseLlamaCache] = None

self.lora_base = lora_base

self.lora_scale = lora_scale

self.lora_path = lora_path

self.spm_infill = spm_infill

if not os.path.exists(model_path):

raise ValueError(f"Model path does not exist: {model_path}")

self._model = self._stack.enter_context(

contextlib.closing(

internals.LlamaModel(

path_model=self.model_path,

params=self.model_params,

verbose=self.verbose,

)

)

)

# Check for Encoder-Decoder architecture

self._has_encoder = self._model.has_encoder()

self._has_decoder = self._model.has_decoder()

self._decoder_start_token_id = -1

if self._has_encoder:

try:

self._decoder_start_token_id = self._model.decoder_start_token()

except AttributeError:

# LLAMA_TOKEN_NULL = -1

self._decoder_start_token_id = -1

if self._decoder_start_token_id == -1:

# Fallback to BOS if specific start token is not defined

self._decoder_start_token_id = self.token_bos()

if self.verbose:

print(f"Model is Encoder-Decoder. Decoder start token: {self._decoder_start_token_id}", file=sys.stderr)

# Override tokenizer

self.tokenizer_ = tokenizer or LlamaTokenizer(self)

# Set the default value for the context and correct the batch

if n_ctx == 0:

n_ctx = self._model.n_ctx_train()

self.n_batch = min(n_ctx, n_batch)

self.context_params.n_ctx = self._model.n_ctx_train()

self.context_params.n_batch = self.n_batch

self.context_params.n_ubatch = min(self.n_batch, n_ubatch)

self._ctx = self._stack.enter_context(

contextlib.closing(

internals.LlamaContext(

model=self._model,

params=self.context_params,

verbose=self.verbose,

)

)

)

# Hybrid architecture detection

_is_recurrent = self._model.is_recurrent()

_is_hybrid = self._model.is_hybrid()

_n_swa = self._model.n_swa()

# Sync llama.cpp upstream (#20291): warn swa-full is not supported for non-SWA models.

if _n_swa == 0:

if (self.context_params.swa_full):

self.context_params.swa_full = False

if self.verbose:

print("Llama.__init__: swa_full is not supported by this model, it will be disabled", file=sys.stderr)

# checkpoints are created only if:

# - the model uses SWA and we are not using `swa_full`

# - the model architecture is marked as recurrent or hybrid

self.is_hybrid = _is_recurrent or _is_hybrid or (_n_swa > 0 and not self.context_params.swa_full)

if self.is_hybrid:

if self.verbose:

print(f"Llama.__init__: Hybrid/Recurrent model detected."

f"(is_recurrent: {_is_recurrent}, is_hybrid: {_is_hybrid}, n_swa: {_n_swa}, swa_full: {self.context_params.swa_full}). "

f" Enabling HybridCheckpointCache(ctx_checkpoints={ctx_checkpoints}, checkpoint_interval={checkpoint_interval}).",

file=sys.stderr)

self.ctx_checkpoints = ctx_checkpoints

self.checkpoint_interval = checkpoint_interval

self._hybrid_cache_mgr = HybridCheckpointCache(self._ctx.ctx, max_checkpoints=self.ctx_checkpoints, verbose=self.verbose)

else:

self._hybrid_cache_mgr = None

self._batch = self._stack.enter_context(

contextlib.closing(

internals.LlamaBatch(

n_tokens=self.n_batch,

embd=0,

n_seq_max=self.context_params.n_seq_max,

verbose=self.verbose,

)

)

)

self._lora_adapter: Optional[llama_cpp.llama_adapter_lora_p] = None

if self.lora_path:

self._lora_adapter = llama_cpp.llama_adapter_lora_init(

self._model.model,

self.lora_path.encode("utf-8"),

)

if self._lora_adapter is None:

raise RuntimeError(

f"Failed to initialize LoRA adapter from lora path: {self.lora_path}"

)

def free_lora_adapter():

if self._lora_adapter is None:

return

llama_cpp.llama_adapter_lora_free(self._lora_adapter)

self._lora_adapter = None

self._stack.callback(free_lora_adapter)

# Todo(JamePeng): The current LoRa loading logic is outdated and needs to be refactored.

if llama_cpp.llama_set_adapters_lora(

self._ctx.ctx, self._lora_adapter, self.lora_scale

):

raise RuntimeError(

f"Failed to set LoRA adapter from lora path: {self.lora_path}"

)

if self.verbose:

print(llama_cpp.llama_print_system_info().decode("utf-8"), file=sys.stderr)

self.chat_format = chat_format

self.chat_handler = chat_handler

self._chat_handlers: Dict[

str, llama_chat_format.LlamaChatCompletionHandler

] = {}

self.draft_model = draft_model

self._n_vocab = self.n_vocab()

self._n_ctx = self.n_ctx()

self._token_nl = self.token_nl()

self._token_eos = self.token_eos()

self._candidates = internals.LlamaTokenDataArray(n_vocab=self._n_vocab)

self.n_tokens = 0

self.input_ids: npt.NDArray[np.intc] = np.ndarray((n_ctx,), dtype=np.intc)

self.scores: npt.NDArray[np.single] = np.ndarray((n_ctx if self._logits_all else 1, self._n_vocab), dtype=np.single)

self._mirostat_mu = ctypes.c_float(

2.0 * 5.0

) # TODO: Move this to sampling context

try:

self.metadata = self._model.metadata()

except Exception as e:

self.metadata = {}

if self.verbose:

print(f"Failed to load metadata: {e}", file=sys.stderr)

if self.verbose:

print(f"Model metadata: {self.metadata}", file=sys.stderr)

eos_token_id = self.token_eos()

bos_token_id = self.token_bos()

eos_token = (

self._model.token_get_text(eos_token_id) if eos_token_id != -1 else ""

)

bos_token = (

self._model.token_get_text(bos_token_id) if bos_token_id != -1 else ""

)

# Unfortunately the llama.cpp API does not return metadata arrays, so we can't get template names from tokenizer.chat_templates

template_choices = dict(

(name[10:], template)

for name, template in self.metadata.items()

if name.startswith("tokenizer.chat_template.")

)

if "tokenizer.chat_template" in self.metadata:

template_choices["chat_template.default"] = self.metadata[

"tokenizer.chat_template"

]

if self.verbose and template_choices:

print(

f"Available chat formats from metadata: {', '.join(template_choices.keys())}",

file=sys.stderr,

)

# Iterate through all the chat templates found in the model's metadata

for name, template in template_choices.items():

try:

# Attempt to parse and register the template as a valid chat handler.

# We wrap this in a try-block because some models (like LLaVA) contain

# non-standard Jinja2 tags (e.g., {% generation %}) that cause the

# standard parser to crash.

self._chat_handlers[name] = llama_chat_format.Jinja2ChatFormatter(

template=template,

eos_token=eos_token,

bos_token=bos_token,

stop_token_ids=[eos_token_id],

).to_chat_handler()

except Exception as e:

# If parsing fails (e.g., TemplateSyntaxError), log a warning but do not crash.

# This ensures the model still loads even if one metadata template is broken.

if self.verbose:

print(f"Warning: Failed to parse chat template '{name}': {e}", file=sys.stderr)

pass

if (

self.chat_format is None

and self.chat_handler is None

and "chat_template.default" in template_choices

):

chat_format = llama_chat_format.guess_chat_format_from_gguf_metadata(

self.metadata

)

if chat_format is not None:

self.chat_format = chat_format

if self.verbose:

print(f"Guessed chat format: {chat_format}", file=sys.stderr)

else:

if self.verbose:

print(

f"Using gguf chat template: {template_choices['chat_template.default']}",

file=sys.stderr,

)

print(f"Using chat eos_token: {eos_token}", file=sys.stderr)

print(f"Using chat bos_token: {bos_token}", file=sys.stderr)

self.chat_format = "chat_template.default"

if self.chat_format is None and self.chat_handler is None:

self.chat_format = "llama-2"

if self.verbose:

print(

f"Using fallback chat format: {self.chat_format}", file=sys.stderr

)

self._sampling_ctx: Optional[LlamaSamplingContext] = None

def close(self) -> None:

"""Explicitly free the model from memory."""

if getattr(self, "_sampling_ctx", None) is not None:

self._sampling_ctx.close()

self._sampling_ctx = None

if getattr(self, "_candidates", None) is not None:

self._candidates.close()

self._candidates = None

if getattr(self, "_hybrid_cache_mgr", None) is not None and hasattr(self._hybrid_cache_mgr, "close"):

self._hybrid_cache_mgr.close()

self._hybrid_cache_mgr = None

if hasattr(self, "chat_handler") and hasattr(self.chat_handler, "close"):

self.chat_handler.close()

self.model_params =None

self.context_params = None

self.chat_handler = None

self.input_ids = None

self.metadata = None

self.scores = None

self.tokenizer_ = None

self._c_tensor_split = None

self._kv_overrides_array = None

if getattr(self, "_stack", None) is not None and hasattr(self._stack, "close"):

self._stack.close()

self._stack = None

def __del__(self) -> None:

self.close()

@property

def ctx(self) -> llama_cpp.llama_context_p:

return self._ctx.ctx

@property

def model(self) -> llama_cpp.llama_model_p:

return self._model.model

@property

def _input_ids(self) -> npt.NDArray[np.intc]:

return self.input_ids[: self.n_tokens]

@property

def _scores(self) -> npt.NDArray[np.single]:

if self._logits_all:

return self.scores[: self.n_tokens, :]

else:

return self.scores

@property

def eval_tokens(self) -> Deque[int]:

return deque(self.input_ids[: self.n_tokens].tolist(), maxlen=self._n_ctx)

@property

def eval_logits(self) -> Deque[List[float]]:

return deque(

self.scores[: self.n_tokens, :].tolist(),

maxlen=self._n_ctx if self._logits_all else 1,

)

def tokenize(

self, text: bytes, add_bos: bool = True, special: bool = False

) -> List[int]:

"""Tokenize a string.

return self.tokenizer_.tokenize(text, add_bos, special)

def detokenize(

self,

tokens: List[int],

prev_tokens: Optional[List[int]] = None,

special: bool = False,

) -> bytes:

"""Detokenize a list of tokens.

return self.tokenizer_.detokenize(

tokens, prev_tokens=prev_tokens, special=special

)

def set_cache(self, cache: Optional[BaseLlamaCache]):

"""Set the cache.

self.cache = cache

def set_seed(self, seed: int):

"""Set the random seed.

self._seed = seed

def reset(self):

"""Reset the model state."""

self.n_tokens = 0

def eval(self, tokens: Sequence[int]):

"""Evaluate a list of tokens.

if len(tokens) == 0:

return

n_eval = len(tokens)

if n_eval == 0:

return

# Context Shift: Prevent OOM by discarding older tokens when context limit is reached.

if self.n_tokens + n_eval > self._n_ctx:

# 0. Check if the memory supports shifting

if not self._ctx.memory_can_shift():

raise RuntimeError(

f"Llama.eval: Context Shift is explicitly disabled by the C++ backend "

f"(n_pos_per_embd > 1 or incompatible M-RoPE). "

f"You MUST increase n_ctx (currently {self._n_ctx}) to fit the dialogue."

)

# 1. Calculate the absolute minimum number of tokens we must discard to fit the new chunk.

required_discard = (self.n_tokens + n_eval) - self._n_ctx

# 2. Sanity check: If the incoming chunk itself is larger than the entire context window,

# shifting is physically impossible.

if required_discard > self.n_tokens:

raise RuntimeError(f"Llama.eval: Context shift failed. The incoming chunk ({n_eval} tokens) "

f"is larger than the entire context window ({self._n_ctx}).")

# 3. Determine how many tokens to keep at the beginning (usually the System Prompt).

_n_keep_desired = min(self.n_keep, self.n_tokens)

# Ensure that keeping these tokens doesn't prevent us from discarding the required amount.

max_keep_allowed = max(0, self.n_tokens - required_discard)

_n_keep = min(_n_keep_desired, max_keep_allowed)

# 4. Calculate the final discard count. Default strategy is to discard half of the available

# past tokens to minimize frequent shifting, but it must be at least `required_discard`.

_n_discard = max(required_discard, (self.n_tokens - _n_keep) // 2)

# 5. Execute the shift only if there are tokens to discard.

if _n_discard > 0:

if self.verbose:

model_type = "Hybrid/Recurrent/SWA" if getattr(self, 'is_hybrid', False) else "Transformer"

print(f"Llama.eval: {model_type} context limit reached. Shifting context: "

f"keeping {_n_keep}, discarding {_n_discard} tokens...", file=sys.stderr)

try:

# Remove the specified block of tokens from the physical KV cache

self._ctx.memory_seq_rm(0, _n_keep, _n_keep + _n_discard)

# Shift the positional IDs of all subsequent tokens to the left to close the gap

self._ctx.memory_seq_add(0, _n_keep + _n_discard, self.n_tokens, -_n_discard)

except Exception as e:

# Defense-in-depth: Catch any other recoverable backend errors

raise RuntimeError(f"Llama.eval: Context Shift failed at the C++ level. Error: {str(e)}") from e

# 6. Synchronize the Python-side token tracking array (ledger)

remaining_len = self.n_tokens - (_n_keep + _n_discard)

if remaining_len > 0:

self.input_ids[_n_keep : _n_keep + remaining_len] = self.input_ids[_n_keep + _n_discard : self.n_tokens]

# 7. Update the global token counter

self.n_tokens -= _n_discard

# Adaptive batch downgrade limit initialization

current_max_batch = self.n_batch

last_ckpt_pos = self.n_tokens

# Adaptive Periodic Checkpointing for Hybrid Models

# Following the "no more than three times" principle :)

# when pre-filling very large blocks, dilute the save frequency to minimize I/O blocking.

if self.is_hybrid and self._hybrid_cache_mgr is not None:

dynamic_interval = max(self.checkpoint_interval, n_eval // 3) # Maximum of 3 triggers

# If KV slots are full, `current_batch_size` will be halved.

# A `while` loop allows us to correctly resume from the exact cut-off point.

i = 0

while i < n_eval:

# Chunk the tokens using the adaptive current_max_batch

n_chunk = min(n_eval - i, current_max_batch)

chunk = tokens[i : i + n_chunk]

n_past = self.n_tokens

self._batch.reset()

pos_array = [self.n_tokens + j for j in range(n_chunk)]

# Configure logits extraction:

# If _logits_all is True, calculate for every token.

# Otherwise, only calculate for the very last token in the entire evaluation sequence.

if self._logits_all:

logits_array = [True] * n_chunk

else:

logits_array = [False] * n_chunk

if i + n_chunk == n_eval:

logits_array[-1] = True

self._batch.add_sequence(

token_array=chunk,

pos_array=pos_array,

seq_ids=[0],

logits_array=logits_array

)

# Dynamic Batch Downgrade: Attempt to decode, reduce batch size if KV cache is fragmented

current_batch_size = n_chunk

success = False

while current_batch_size > 0:

# Tell the C++ backend to only process up to `current_batch_size` tokens

self._batch.batch.n_tokens = current_batch_size

try:

status = self._ctx.decode(self._batch)

# 0: Success

if status == 0:

success = True

# If we successfully decoded after a downgrade,

# update current_max_batch to prevent repeated failures in next iterations.

if current_batch_size < current_max_batch:

current_max_batch = current_batch_size

break

# 1: No KV slot available (Recoverable)

elif status == 1:

if current_batch_size == 1:

if self.verbose:

print("Llama.eval: KV slots completely full. "

"Cannot reduce batch size below 1. Aborting...", file=sys.stderr)

break

if self.verbose:

print(f"Llama.eval: KV slots full (Code 1). Halving batch size "

f"from {current_batch_size} to {current_batch_size // 2}...", file=sys.stderr)

current_batch_size //= 2

except Exception as e:

# Catch fatal backend failures (e.g., Code -2, -3)

raise RuntimeError(f"Llama.eval(decode): Fatal Decode Error at Pos {self.n_tokens}, "

f"Batch size {current_batch_size}: {str(e)}") from e

if not success:

raise RuntimeError("Llama.eval(decode): Failed completely even with batch size 1.")

# Save successfully processed tokens into the Python-side ledger

self.input_ids[n_past : n_past + current_batch_size] = chunk[:current_batch_size]

# Extract and save all logits if requested, ensuring we only copy the successfully processed rows

if self._logits_all:

logits_ptr = self._ctx.get_logits()

rows = current_batch_size

cols = self._n_vocab

logits_view = np.ctypeslib.as_array(logits_ptr, shape=(rows * cols,))

self.scores[n_past : n_past + current_batch_size, :].reshape(-1)[:] = logits_view

# Update indices based on actual processed batch size

self.n_tokens += current_batch_size

i += current_batch_size

# Periodic Checkpoint: Save states for hybrid models to avoid massive rollbacks

if self.is_hybrid and self._hybrid_cache_mgr is not None:

current_pos = self.n_tokens

if (current_pos - last_ckpt_pos >= dynamic_interval) and (i < n_eval):

if self.verbose:

print(f"Llama.eval: [Periodic Checkpoint] Saving hybrid state at pos {current_pos} "

f"(checkpoint_interval({dynamic_interval}) reached, last={last_ckpt_pos}).", file=sys.stderr)

success = self._hybrid_cache_mgr.save_checkpoint(

current_pos=current_pos,

tokens=self.input_ids[:current_pos].tolist(),

seq_id=0

)

if success:

last_ckpt_pos = current_pos

else:

if self.verbose:

print(f"Llama.eval: [Periodic Checkpoint] HybridCheckpoint save failed at pos {current_pos}, skipping update", file=sys.stderr)

# Save the final logit if not in _logits_all mode

if not self._logits_all:

logits_ptr = self._ctx.get_logits()

logits_view = np.ctypeslib.as_array(logits_ptr, shape=(self._n_vocab,))

self.scores[0, :] = logits_view

# Helper method: Convert dict logit_bias to List[llama_logit_bias]

def _convert_logit_bias(self, logit_bias: Optional[Dict[int, float]]) -> List[llama_cpp.llama_logit_bias]:

if not logit_bias:

return []

bias_list = []

for token, bias in logit_bias.items():

lb = llama_cpp.llama_logit_bias()

lb.token = token

lb.bias = bias

bias_list.append(lb)

return bias_list

def sample(

self,

# Core

top_k: int = 40, # <= 0 to use vocab size

top_p: float = 0.95, # 1.0 = disabled

min_p: float = 0.05, # 0.0 = disabled

typical_p: float = 1.0, # typical_p, 1.0 = disabled

temp: float = 0.80, # <= 0.0 to sample greedily, 0.0 to not output probabilities

# Dynamic Temp

dynatemp_range: float = 0.0, # 0.0 = disabled

dynatemp_exponent: float = 1.0, # controls how entropy maps to temperature in dynamic temperature sampler

# Common

top_n_sigma: float = -1.00, # -1.0 = disabled

min_keep: int = 0, # 0 = disabled, otherwise samplers should return at least min_keep tokens

# Penalties