llama.py

from __future__ import annotations

import contextlib
import ctypes
import fnmatch
import json
import math
import multiprocessing
import os
import sys
import time
import threading
import typing
import uuid
import warnings

import numpy as np
import numpy.typing as npt

from typing import (
    Any,
    List,
    Literal,
    Optional,
    Union,
    Generator,
    Sequence,
    Iterator,
    Deque,
    Callable,
    Dict,
)
from collections import deque
from pathlib import Path

from .llama_types import *
from .llama_grammar import LlamaGrammar
from .llama_cache import (
    BaseLlamaCache,
    LlamaCache,            # type: ignore
    LlamaDiskCache,        # type: ignore
    LlamaRAMCache,         # type: ignore
    LlamaTrieCache,        # type: ignore
    HybridCheckpointCache, # type: ignore
)
from .llama_tokenizer import BaseLlamaTokenizer, LlamaTokenizer
import llama_cpp.llama_cpp as llama_cpp_lib
import llama_cpp.llama_chat_format as llama_chat_format
import llama_cpp.llama_multimodal as llama_multimodal

from llama_cpp.llama_speculative import (
    LlamaDraftModel,
    LlamaSpecEngine,
    SpecConfig,
    SpeculativeType,
    create_native_spec_engine,
    create_spec_engine,
    _speculative_generation_timing_stats,
    speculative_output_limits,
)

import llama_cpp._internals as internals
from ._rpc import normalize_rpc_servers, register_rpc_devices
from ._internals import (
    LlamaSamplingContext,
    LlamaSamplingParams,
    CommonSamplerType,
    CustomSampler,
)
from ._ggml import (
    ggml_backend_cpu_buffer_type,
    ggml_backend_load_all_from_path,
    ggml_backend_reg_count
)
from ._logger import (
    configure_logging,
    get_verbosity,
    set_verbosity,
    get_log_filters,
    set_log_filters,
    add_log_filters,
    clear_log_filters,
    reset_log_filters,
)
from ._utils import suppress_stdout_stderr, normalize_embedding


def _format_speculative_duration(seconds: float) -> str:
    """Format speculative phase timings without hiding sub-second costs."""
    if seconds >= 1.0:
        return f"{seconds:.3f} s"
    return f"{seconds * 1000.0:.3f} ms"


class AbortCriteria:
    """
    Listen for external interruption signals to trigger a stop condition.
    When an external thread calls `llama.abort()`, a loop interrupt is generated.
    """
    requires_logits = False

    def __init__(self, abort_event: threading.Event):
        self.abort_event = abort_event

    def __call__(self, _input_ids: npt.NDArray[np.intc], _logits: npt.NDArray[np.single]) -> bool:
        # Note: _input_ids and _logits are required by the signature but unused here.
        return self.abort_event.is_set()


@llama_cpp_lib.ggml_abort_callback
def _llama_native_abort_callback(data: ctypes.c_void_p) -> bool:
    """Read an instance abort flag without retaining its Llama object."""
    if not data:
        return False
    return bool(ctypes.cast(data, ctypes.POINTER(ctypes.c_bool)).contents.value)


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

    __backend_initialized = False
    __backend_lock = threading.Lock()

    LLM_FFN_EXPS_REGEX = rb"\.ffn_(up|down|gate|gate_up)_(ch|)exps"

    def __init__(
        self,
        model_path: str,
        mmproj_path: Optional[str] = None,
        *,
        # Model Params
        n_gpu_layers: Union[int, Literal["auto", "all"]] = "auto",
        cpu_moe: bool = False,
        n_cpu_moe: int = 0,
        split_mode: int = llama_cpp_lib.llama_split_mode.LLAMA_SPLIT_MODE_LAYER,
        load_mode: int = llama_cpp_lib.llama_load_mode.LLAMA_LOAD_MODE_AUTO,
        lazy_mode: int = llama_cpp_lib.llama_lazy_mode.LLAMA_LAZY_MODE_AUTO,
        main_gpu: int = 0,
        tensor_split: Optional[List[float]] = None,
        rpc_servers: Optional[Union[str, Sequence[str]]] = None,
        rpc_local_devices: Optional[Sequence[str]] = None,
        kv_overrides: Optional[Dict[str, Union[bool, int, float, str]]] = None,
        use_mmap: bool = False,
        use_direct_io: bool = False,
        use_mlock: bool = False,
        vocab_only: bool = False,
        check_tensors: bool = False,
        use_extra_bufts: bool = True,
        no_host: bool = False,
        no_alloc: bool = False,
        load_mtp: bool = False,
        # Context Params
        seed: int = llama_cpp_lib.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_rs_seq: int = 0,
        n_outputs_max: int = 0,
        n_outputs_max_per_seq: int = 1,
        moe_cache_size: int = 0,
        n_threads: Optional[int] = None,
        n_threads_batch: Optional[int] = None,
        ctx_type: int = llama_cpp_lib.llama_context_type.LLAMA_CONTEXT_TYPE_DEFAULT,
        rope_scaling_type: Optional[
            int
        ] = llama_cpp_lib.llama_rope_scaling_type.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED,
        pooling_type: int = llama_cpp_lib.LLAMA_POOLING_TYPE_UNSPECIFIED,
        attention_type: Optional[int] = llama_cpp_lib.llama_attention_type.LLAMA_ATTENTION_TYPE_UNSPECIFIED,
        flash_attn_type: Optional[int] = llama_cpp_lib.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 = 16,
        checkpoint_interval: int = 4096,
        checkpoint_on_device: bool = False,
        # Sampling Params
        last_n_tokens_size: int = 64,
        # 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,
        speculative: Optional[Union[SpecConfig, LlamaSpecEngine]] = 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,
        # Log
        verbose: bool = True,
        verbosity: Optional[Union[int, str, bool]] = None,
        log_filters: Optional[Sequence[str]] = None,
        log_filters_case_sensitive: bool = True,
        # Extra Params
        chat_template_name: Optional[str] = None,
        chat_handler_kwargs: Dict[str, Any] = {},
        **kwargs,  # type: ignore
    ):
        """Load a llama.cpp model from `model_path`.

        Examples:
            Basic usage

            >>> import llama_cpp
            >>> model = llama_cpp.Llama(
            ...     model_path="path/to/model",
            ... )
            >>> print(model("The quick brown fox jumps ", stop=["."])["choices"][0]["text"])
            the lazy dog

            Loading a chat model

            >>> import llama_cpp
            >>> model = llama_cpp.Llama(
            ...     model_path="path/to/model",
            ...     chat_format="llama-2",
            ... )
            >>> print(model.create_chat_completion(
            ...     messages=[{
            ...         "role": "user",
            ...         "content": "what is the meaning of life?"
            ...     }]
            ... ))

        Args:
            model_path: Path to the model.
            n_gpu_layers: Max number of model layers to store in VRAM (-ngl).
                Accepts an exact integer, "auto", or "all".
                "auto" / -1 lets llama.cpp choose automatically.
                "all" / -2 stores all possible layers in VRAM.
                0 disables model layer offload.
            cpu_moe: Keep all Mixture of Experts (MoE) weights in the CPU
            n_cpu_moe: Keep the MoE expert weights of the first N layers on CPU.
                Useful when VRAM is insufficient for MoE models.
            split_mode: How to split the model across GPUs. See llama_cpp.LLAMA_SPLIT_* for options.
            load_mode: How to load the model. See llama_cpp.LLAMA_LOAD_MODE_* for options.
            main_gpu: main_gpu interpretation depends on split_mode: LLAMA_SPLIT_MODE_NONE: the GPU that is used for the entire model. LLAMA_SPLIT_MODE_ROW: the GPU that is used for small tensors and intermediate results. LLAMA_SPLIT_MODE_LAYER: ignored
            tensor_split: How split tensors should be distributed across GPUs. If None, the model is not split.
            rpc_servers: RPC endpoints as host:port strings, or a comma-separated string.
                RPC devices precede local GPUs in tensor_split and main_gpu indices.
            rpc_local_devices: Local GPU names to include alongside RPC devices, in order.
                None selects the default local GPUs; an empty list selects only RPC devices.
            kv_overrides: Key-value overrides for the model.
            vocab_only: Only load the vocabulary no weights.
            check_tensors: validate model tensor data
            use_extra_bufts: use extra buffer types (used for weight repacking)
            no_host: bypass host buffer allowing extra buffers to be used
            no_alloc: only load metadata and simulate memory allocations
            load_mtp: whether to load MTP layers
            seed: RNG seed, -1 for random
            n_ctx: Text context, 0 = from model
            n_keep: Number of tokens to keep from initial prompt
            n_batch: Prompt processing maximum batch size
            n_ubatch: Physical batch size
            n_seq_max: max number of sequences (i.e. distinct states for recurrent models)
            n_rs_seq: Number of recurrent-state snapshots per sequence for rollback. 0 disables rollback snapshots. Experimental.
            n_outputs_max: Maximum outputs in a physical batch. 0 lets llama.cpp use the effective n_batch.
            n_outputs_max_per_seq: Maximum outputs per sequence. 0 lets llama.cpp use the effective n_outputs_max.
            n_threads: Number of threads to use for generation
            n_threads_batch: Number of threads to use for batch processing
            ctx_type: Context implementation type, such as the MTP context type.
            rope_scaling_type: RoPE scaling type, from `enum llama_rope_scaling_type`. ref: https://github.com/ggml-org/llama.cpp/pull/2054
            pooling_type: Pooling type, from `enum llama_pooling_type`.
            attention_type: attention type to use for embeddings
            flash_attn_type: when to enable Flash Attention
            rope_freq_base: RoPE base frequency, 0 = from model
            rope_freq_scale: RoPE frequency scaling factor, 0 = from model
            yarn_ext_factor: YaRN extrapolation mix factor, negative = from model
            yarn_attn_factor: YaRN magnitude scaling factor
            yarn_beta_fast: YaRN low correction dim
            yarn_beta_slow: YaRN high correction dim
            yarn_orig_ctx: YaRN original context size
            logits_all: Return logits for all tokens, not just the last token. Must be True for completion to return logprobs.
            embeddings: Embedding mode only. if true, extract embeddings (together with logits)
            offload_kqv: Offload K, Q, V to GPU.
            no_perf: Disable performance timing collection.
            op_offload: whether to offload host tensor operations to device
            swa_full: whether to use full-size SWA cache
            kv_unified: use single unified KV buffer for the KV cache of all sequences
            ctx_checkpoints: max number of context checkpoints to create per slot (default: 16)[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)
            checkpoint_interval: Hybrid model checkpoint token intervals, and archiving of text with interval sizes along the way.
            checkpoint_on_device: Store hybrid/recurrent checkpoint tensor payloads in llama_context-owned device buffers via LLAMA_STATE_SEQ_FLAGS_ON_DEVICE.
            last_n_tokens_size: Maximum number of tokens to keep in the last_n_tokens deque.
            numa: numa policy
            chat_format: String specifying the chat format to use when calling create_chat_completion.
            chat_handler: Optional chat handler to use when calling create_chat_completion.
            draft_model: Optional draft model to use for speculative decoding.
                Deprecated; use ``speculative=SpecConfig(...)`` for the
                llama.cpp-compatible stateful speculative-decoding lifecycle.
            speculative: Speculative decoding configuration or an initialized
                speculative engine. Its fields mirror llama.cpp's ``--spec-*``
                arguments, including draft length, probability threshold,
                draft model, cache types and backend sampling.
            tokenizer: Optional tokenizer to override the default tokenizer from llama.cpp.
            moe_cache_size: Device cache size in bytes for MoE experts kept in host memory,
                split among devices like layers. Requires a GPU and MoE model;
                tensor parallelism is unsupported. 0 disables caching. Experimental.
            type_k: KV cache data type for K (default: f16)
            type_v: KV cache data type for V (default: f16)
            spm_infill: Use Suffix/Prefix/Middle pattern for infill (instead of Prefix/Suffix/Middle) as some models prefer this.
            verbose: Backward-compatible boolean switch for native llama.cpp / ggml runtime logs.
                False keeps only error-level native logs; True enables debug-level native logs.
                If `verbosity` is provided, `verbosity` takes precedence over `verbose`.
            verbosity: Fine-grained llama.cpp-style native runtime log verbosity.
                Accepts 0-5, bool, or string aliases.
                Numeric levels:
                    0 = output only
                    1 = error
                    2 = warning
                    3 = info
                    4 = trace
                    5 = debug
                Use `verbosity=3` for llama.cpp-style default info logs.
                `verbose=False` remains equivalent to error-only logging, while
                `verbose=True` remains equivalent to debug logging.
            log_filters: Optional substring filters for native runtime logs.
                If any provided substring appears in a decoded backend log message,
                that message is suppressed. By default, the logger may include built-in
                filters for noisy low-level logs such as CUDA Graph reuse spam messages.
                Pass an empty list to disable all substring filtering for this instance.
            log_filters_case_sensitive: Whether `log_filters` should match case-sensitively.
                Defaults to True for predictable low-level backend log filtering.
        Raises:
            ValueError: If the model path does not exist.

        Returns:
            A Llama instance.
        """
        self.verbose = verbose
        self.verbosity = verbosity
        self._stack = contextlib.ExitStack()
        if moe_cache_size < 0:
            raise ValueError("moe_cache_size must be nonnegative")
        self.rpc_servers = normalize_rpc_servers(rpc_servers)
        if isinstance(rpc_local_devices, (str, bytes)):
            raise TypeError("rpc_local_devices must be a sequence of device names")
        self.rpc_local_devices = list(rpc_local_devices) if rpc_local_devices is not None else None
        self._c_rpc_devices = None
        if self.rpc_local_devices is not None and not self.rpc_servers:
            raise ValueError("rpc_local_devices requires rpc_servers")
        if self.rpc_servers and split_mode == llama_cpp_lib.llama_split_mode.LLAMA_SPLIT_MODE_ROW:
            raise ValueError("RPC devices do not support LLAMA_SPLIT_MODE_ROW")

        configure_logging(
            verbose=verbose,
            verbosity=verbosity,
            log_filters=log_filters,
            log_filters_case_sensitive=log_filters_case_sensitive,
        )

        # llama.cpp / ggml backend initialization is process-global.
        # Run it once before loading any model.
        with Llama.__backend_lock:
            if not Llama.__backend_initialized:
                with suppress_stdout_stderr(disable=verbose):
                    llama_cpp_lib.llama_backend_init()

                    # Dynamic ggml backends must be loaded before model selection.
                    lib_dir = Path(llama_cpp_lib.__file__).resolve().parent / "lib"

                    if not lib_dir.exists():
                        raise FileNotFoundError(f"Llama.__init__: llama_cpp lib directory not found: {lib_dir}")

                    ggml_backend_load_all_from_path(
                        ctypes.c_char_p(str(lib_dir).encode("utf-8"))
                    )

                    if self.verbose:
                        count = ggml_backend_reg_count()
                        print(f"Llama.__init__: Loaded ggml backend registry count: {count}", file=sys.stderr)

                Llama.__backend_initialized = True

            if self.rpc_servers:
                if tensor_split is not None:
                    if not all(math.isfinite(value) and value >= 0 for value in tensor_split):
                        raise ValueError("tensor_split values must be finite and non-negative")
                    if not any(tensor_split):
                        raise ValueError("tensor_split must include a positive value")
                rpc_devices = register_rpc_devices(
                    self.rpc_servers,
                    self.rpc_local_devices,
                    tensor_parallel=(split_mode == llama_cpp_lib.llama_split_mode.LLAMA_SPLIT_MODE_TENSOR),
                    max_devices=llama_cpp_lib.LLAMA_MAX_DEVICES,
                    expected_devices=len(tensor_split) if tensor_split is not None else None,
                )
                self._c_rpc_devices = (ctypes.c_void_p * (len(rpc_devices) + 1))(
                    *rpc_devices, None
                )

        if isinstance(numa, bool):
            self.numa = (
                llama_cpp_lib.GGML_NUMA_STRATEGY_DISTRIBUTE
                if numa
                else llama_cpp_lib.GGML_NUMA_STRATEGY_DISABLED
            )
        else:
            self.numa = numa

        if self.numa != llama_cpp_lib.GGML_NUMA_STRATEGY_DISABLED:
            with suppress_stdout_stderr(disable=verbose):
                llama_cpp_lib.llama_numa_init(self.numa)

        self.model_path = model_path

        if draft_model is not None and speculative is not None:
            raise ValueError("draft_model and speculative cannot be used together")

        if isinstance(speculative, SpecConfig):
            speculative.validate()
            if (
                speculative.spec_type == SpeculativeType.DRAFT_MTP
                and speculative.draft_model_path is None
            ):
                load_mtp = True

        if (use_mmap or use_direct_io or use_mlock) and verbose:
            print(
                "Llama.__init__: WARNING: "
                "Legacy load options (`use_mmap`, `use_direct_io`, `use_mlock`) "
                "are deprecated. Use `load_mode` instead.",
                file=sys.stderr,
            )

        # Model Params
        self.model_params = llama_cpp_lib.llama_model_default_params()
        if self._c_rpc_devices is not None:
            self.model_params.devices = self._c_rpc_devices
        self.model_params.n_gpu_layers = self._parse_n_gpu_layers(n_gpu_layers)
        self.model_params.split_mode = split_mode
        self.model_params.load_mode = load_mode
        self.model_params.lazy_mode = lazy_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_lib.LLAMA_MAX_DEVICES:
                raise ValueError(
                    f"Attempt to split tensors that exceed maximum supported devices. Current LLAMA_MAX_DEVICES={llama_cpp_lib.LLAMA_MAX_DEVICES}"
                )
            # Type conversion and expand the list to the length of LLAMA_MAX_DEVICES
            FloatArray = ctypes.c_float * llama_cpp_lib.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.check_tensors = check_tensors
        self.model_params.use_extra_bufts = use_extra_bufts
        self.model_params.no_host = no_host
        self.model_params.no_alloc = no_alloc
        self.model_params.load_mtp = load_mtp

        # Logic of cpu_moe, n_cpu_moe
        # Reference from llama.cpp/tools/llama-bench/llama-bench.cpp
        self.cpu_moe = cpu_moe
        self.n_cpu_moe = n_cpu_moe
        self._cpu_moe_patterns = None
        self._cpu_moe_tensor_buft_overrides = None

        if self.n_cpu_moe < 0:
            raise ValueError("n_cpu_moe must be >= 0")

        if self.cpu_moe and self.n_cpu_moe != 0 and self.verbose:
            print(
                "Llama.__init__: cpu_moe=True already keeps all MoE expert weights on CPU; "
                "n_cpu_moe is redundant.",
                file=sys.stderr,
            )

        if self.cpu_moe or self.n_cpu_moe > 0:
            cpu_buft = ggml_backend_cpu_buffer_type()

            if self.cpu_moe:
                patterns = [self.LLM_FFN_EXPS_REGEX]
            else:
                patterns = [
                    self._make_cpu_moe_pattern(i)
                    for i in range(self.n_cpu_moe)
                ]

            # keep pattern bytes alive
            self._cpu_moe_patterns = patterns

            TensorBuftOverrideArray = (
                llama_cpp_lib.llama_model_tensor_buft_override
                * (len(patterns) + 1)
            )
            self._cpu_moe_tensor_buft_overrides = TensorBuftOverrideArray()

            for i, pattern in enumerate(self._cpu_moe_patterns):
                self._cpu_moe_tensor_buft_overrides[i].pattern = pattern
                self._cpu_moe_tensor_buft_overrides[i].buft = cpu_buft

            self._cpu_moe_tensor_buft_overrides[len(patterns)].pattern = None
            self._cpu_moe_tensor_buft_overrides[len(patterns)].buft = None

            self.model_params.tensor_buft_overrides = (
                self._cpu_moe_tensor_buft_overrides
            )

        # 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_lib.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_lib.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_lib.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_lib.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_lib.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_lib.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_rs_seq = n_rs_seq
        self.n_outputs_max = n_outputs_max
        self.n_outputs_max_per_seq = n_outputs_max_per_seq
        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_lib.LLAMA_DEFAULT_SEED

        # Context Params
        self.context_params = llama_cpp_lib.llama_context_default_params()
        self.context_params.moe_cache_size = moe_cache_size
        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 = max(1, self.n_seq_max)
        if self.context_params.n_seq_max > llama_cpp_lib.LLAMA_MAX_SEQ:
            raise RuntimeError(f"n_seq_max must be <= {llama_cpp_lib.LLAMA_MAX_SEQ}")

        if (
            isinstance(speculative, SpecConfig)
            and speculative.enabled()
            and self.n_batch > 0
        ):
            required_total, required_per_seq = speculative_output_limits(
                self.n_batch,
                max(1, self.n_seq_max),
                speculative.max_draft_tokens(),
            )
            # common_context_params_to_llama derives recurrent snapshots from
            # the speculative method and draft length.
            if speculative.spec_type.is_draft():
                self.n_rs_seq = max(self.n_rs_seq, speculative.draft_n_max)
            self.n_outputs_max = required_total
            self.n_outputs_max_per_seq = required_per_seq

        self.context_params.n_rs_seq = max(self.n_rs_seq, 0)
        self.context_params.n_outputs_max = max(self.n_outputs_max, 0)
        self.context_params.n_outputs_max_per_seq = max(self.n_outputs_max_per_seq, 0)
        self.context_params.n_threads = self.n_threads
        self.context_params.n_threads_batch = self.n_threads_batch

        self.context_params.ctx_type = ctx_type
        self.context_params.ctx_other = None
        self.context_params.rope_scaling_type = (
            rope_scaling_type
            if rope_scaling_type is not None
            else llama_cpp_lib.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_lib.LLAMA_POOLING_TYPE_UNSPECIFIED
        )
        self.context_params.attention_type = (
            attention_type
            if attention_type is not None
            else llama_cpp_lib.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_lib.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._speculative_verifying = False
        # Set only while generate() is active so eval() can attribute native
        # process() calls to the current request.
        self._active_speculative_phase_stats: Optional[Dict[str, Any]] = None
        self.last_speculative_stats: Dict[str, Any] = {
            "drafted": 0,
            "verified": 0,
            "accepted": 0,
            "begin_calls": 0,
            "draft_calls": 0,
            "process_calls": 0,
            "accept_calls": 0,
            "generated_drafts": 0,
            "accepted_drafts": 0,
            "draft_batch_acceptance_rate": 0.0,
            "accepted_draft_tokens": 0,
            "draft_token_acceptance_rate": 0.0,
            "mean_accepted_length": 0.0,
            "acceptance_rate_per_position": [],
            "begin_seconds": 0.0,
            "draft_seconds": 0.0,
            "target_decode_seconds": 0.0,
            "target_sync_seconds": 0.0,
            "process_seconds": 0.0,
            "accept_seconds": 0.0,
            "checkpoint_captures": 0,
            "checkpoint_restores": 0,
            "checkpoint_verification_reuses": 0,
            "checkpoint_native_captures": 0,
            "checkpoint_native_restores": 0,
            "checkpoint_device_captures": 0,
            "checkpoint_device_restores": 0,
            "checkpoint_native_verification_rollbacks": 0,
            "checkpoint_buffer_bytes": 0,
            "checkpoint_capture_seconds": 0.0,
            "checkpoint_restore_seconds": 0.0,
            "verification_steps": 0,
            "rollbacks": 0,
            "native_rollbacks": 0,
            "checkpoint_rollbacks": 0,
            "acceptance_rate": 0.0,
            "generation_tokens": 0,
            "generation_seconds": 0.0,
            "generation_tokens_per_second": 0.0,
            "time_to_first_token_seconds": 0.0,
            "sustained_tokens": 0,
            "sustained_seconds": 0.0,
            "sustained_tokens_per_second": 0.0,
            # Backward-compatible aliases for the old, ambiguously named fields.
            "decode_tokens": 0,
            "decode_seconds": 0.0,
            "decode_tokens_per_second": 0.0,
        }
        # Stateful speculative decoding requests every native verification
        # output only while checking a draft. It does not require copying every
        # vocabulary row into the long-lived Python ``scores`` matrix.
        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.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)

        # n_ctx=0 resolves n_batch only after model metadata is available, so
        # derive the speculative output limits again at this point.
        if isinstance(speculative, SpecConfig) and speculative.enabled():
            required_total, required_per_seq = speculative_output_limits(
                self.n_batch,
                max(1, self.n_seq_max),
                speculative.max_draft_tokens(),
            )
            self.n_outputs_max = required_total
            self.n_outputs_max_per_seq = required_per_seq
            if speculative.spec_type.is_draft():
                self.n_rs_seq = max(self.n_rs_seq, speculative.draft_n_max)
            self.context_params.n_rs_seq = self.n_rs_seq
            self.context_params.n_outputs_max = required_total
            self.context_params.n_outputs_max_per_seq = required_per_seq

        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}, "
                    f"n_swa: {_n_swa}, swa_full: {self.context_params.swa_full}). "
                    f"Enabling HybridCheckpointCache("
                    f"ctx_checkpoints={ctx_checkpoints}, "
                    f"checkpoint_interval={checkpoint_interval}, "
                    f"on_device={checkpoint_on_device}).",
                    file=sys.stderr,
                )
            self.ctx_checkpoints = ctx_checkpoints
            self.checkpoint_interval = checkpoint_interval
            self.checkpoint_on_device = checkpoint_on_device
            self._hybrid_cache_mgr = HybridCheckpointCache(
                self._ctx.ctx,
                max_checkpoints=self.ctx_checkpoints,
                on_device=self.checkpoint_on_device,
                verbose=self.verbose,
            )
            self._ctx._register_checkpoint_cache(self._hybrid_cache_mgr)
        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,
                )
            )
        )

        if self.verbose:
            print(llama_cpp_lib.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
        if isinstance(speculative, SpecConfig):
            if not speculative.enabled():
                self.speculative = None
            elif speculative.spec_type.is_draft():
                # Draft-family engines depend on the already initialized native
                # target model/context. The engine may also create and own a
                # separate draft model/context, as external MTP does.
                self.speculative: Optional[LlamaSpecEngine] = (
                    create_native_spec_engine(
                        speculative,
                        target_model=self._model,
                        target_context=self._ctx,
                        model_params=self.model_params,
                        context_params=self.context_params,
                        verbose=self.verbose,
                    )
                )
            else:
                # Model-free engines, currently NGram variants, only need token
                # history and can be constructed directly from their config.
                self.speculative = create_spec_engine(speculative)
            self.speculative_config: Optional[SpecConfig] = speculative
        else:
            self.speculative = speculative
            self.speculative_config = None

        self._n_vocab = self.n_vocab()
        self._n_ctx = self.n_ctx()

        # Candidate storage is owned by LlamaSamplingContext. Keeping a second,
        # unused full-vocabulary array here wastes several MiB on large-vocab
        # models (for Qwen3.8: 248320 llama_token_data entries plus token IDs).
        self._candidates = None

        self.n_tokens = 0
        self._last_eval_output_start = 0
        self._last_eval_output_count = 0
        self._restored_logits = None
        self._prefilled_prompt = None
        self._state_needs_speculative_reset = False
        self.input_ids: npt.NDArray[np.intc] = np.ndarray((self._n_ctx,), dtype=np.intc)
        self.scores: npt.NDArray[np.single] = np.ndarray((self._n_ctx if self._logits_all else 1, self._n_vocab), dtype=np.single)

        try:
            self.metadata = self._model.metadata()
            self.model_desc = self._model.model_desc()
            # The total size of all the tensors in the model in bytes
            self.model_size = self._model.model_size()

        except Exception as e:
            self.metadata = {}
            if self.verbose:
                print(f"Failed to load metadata: {e}", file=sys.stderr)
        
        if mmproj_path is not None:
            if self.chat_handler is not None and self.verbose:
                print("Warning: Both `chat_handler` and `mmproj_path` are not null. Chat handler will be overwritten.", flush = True)

            self.chat_handler = llama_multimodal.GenericMTMDChatHandler(
                chat_format = self.metadata.get("tokenizer.chat_template", None),
                mmproj_path = mmproj_path,
                chat_template_name=chat_template_name,
                **chat_handler_kwargs
            )

        if self.verbose:
            print(f"Model desc: {self.model_desc}, "
                  f"Model size: {self.model_size / (1024 * 1024):.2f} MB, "
                  f"Model metadata: {self.metadata}",
                  file=sys.stderr)

        eos_token_id = self.token_eos()
        bos_token_id = self.token_bos()
        eot_token_id = self.token_eot()
        sep_token_id = self.token_sep()
        nl_token_id = self.token_nl()
        pad_token_id = self.token_pad()
        mask_token_id = self.token_mask()

        def _token_text(token_id: int) -> str:
            return self._model.token_get_text(token_id) if token_id != -1 else ""

        bos_token = _token_text(bos_token_id)
        eos_token = _token_text(eos_token_id)

        special_tokens_map = {
            name: text
            for name, token_id in {
                "eot_token": eot_token_id,
                "sep_token": sep_token_id,
                "nl_token": nl_token_id,
                "pad_token": pad_token_id,
                "mask_token": mask_token_id,
            }.items()
            if token_id != -1 and (text := _token_text(token_id))
        }

        stop_token_ids = [
            token_id
            for token_id in (eos_token_id, eot_token_id)
            if token_id != -1
        ]

        if not stop_token_ids:
            stop_token_ids = None

        # 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.
                # Keep this guarded because model metadata may contain malformed or
                # model-specific Jinja templates that still cannot be rendered by this runtime.
                self._chat_handlers[name] = llama_chat_format.Jinja2ChatFormatter(
                    template=template,
                    eos_token=eos_token,
                    bos_token=bos_token,
                    stop_token_ids=stop_token_ids,
                    special_tokens_map=special_tokens_map,
                ).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

        # Create a thread-safe interrupt event
        self._abort_event = threading.Event()
        self._native_abort_flag = ctypes.c_bool(False)
        self._ctx.set_abort_callback(
            _llama_native_abort_callback,
            ctypes.cast(ctypes.pointer(self._native_abort_flag), ctypes.c_void_p),
        )

    def close(self) -> None:
        """Explicitly free the model from memory."""
        resources = []
        for name in ("_sampling_ctx", "speculative", "_candidates",
                     "_hybrid_cache_mgr", "chat_handler", "_stack"):
            resource = getattr(self, name, None)
            if resource is not None and callable(getattr(resource, "close", None)):
                resources.append(resource)
            setattr(self, name, None)

        # Release Python-owned output and history even if a native destructor
        # raises. The public cache is borrowed: detach without clearing other
        # callers' snapshots or closing their disk cache.
        self._restored_logits = None
        self._prefilled_prompt = None
        self.cache = None
        self.n_tokens = 0
        self._last_eval_output_start = 0
        self._last_eval_output_count = 0
        self._state_needs_speculative_reset = False
        self.chat_handler = None
        self.input_ids = None
        self.metadata = None
        self.scores = None
        self.tokenizer_ = None

        # Preserve dependency order and attempt every close on failure.
        try:
            with contextlib.ExitStack() as cleanup:
                for resource in reversed(resources):
                    cleanup.callback(resource.close)
        finally:
            # Native destructors may still need the model parameter storage.
            self.model_params = None
            self.context_params = None
            self._c_tensor_split = None
            self._c_rpc_devices = None
            self._kv_overrides_array = None

    def __del__(self) -> None:
        # __del__ can run after Python has started clearing module globals and
        # disabled imports. Explicit close() still reports cleanup failures,
        # while finalization must not emit an unraisable exception.
        try:
            self.close()
        except Exception:
            pass

    @staticmethod
    def _parse_n_gpu_layers(n_gpu_layers: Union[int, str]) -> int:
        if isinstance(n_gpu_layers, str):
            value = n_gpu_layers.strip().lower()
            if value == "auto":
                return -1
            if value == "all":
                return -2
            try:
                return int(value)
            except ValueError as exc:
                raise ValueError("n_gpu_layers must be an int, 'auto', or 'all'") from exc

        if isinstance(n_gpu_layers, int):
            return n_gpu_layers

        raise TypeError("n_gpu_layers must be an int, 'auto', or 'all'")

    @staticmethod
    def _make_cpu_moe_pattern(i: int) -> bytes:
        return f"blk\\.{i}".encode("utf-8") + Llama.LLM_FFN_EXPS_REGEX

    @property
    def ctx(self) -> llama_cpp_lib.llama_context_p:
        return self._ctx.ctx

    @property
    def model(self) -> llama_cpp_lib.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,
        )

    # Logger API

    def set_verbosity(self, verbosity: Union[int, str, bool, None]) -> None:
        """Set native llama.cpp / ggml runtime log verbosity for this process.

        Levels:
            0 = output only
            1 = error
            2 = warning
            3 = info
            4 = trace
            5 = debug

        Note:
            Native backend logging is process-global because llama.cpp / ggml use
            a global log callback. Changing this affects all Llama instances in
            the current Python process.
        """
        set_verbosity(verbosity)
        self.verbosity = get_verbosity()
        self.verbose = self.verbosity >= 5


    def get_verbosity(self) -> int:
        """Return the current native runtime log verbosity."""
        return get_verbosity()


    def set_log_filters(
        self,
        filters: Sequence[str],
        *,
        case_sensitive: bool = True,
    ) -> None:
        """Replace substring filters for native runtime logs.

        Any backend log message containing one of these substrings will be
        suppressed. Pass an empty list to disable all substring filtering.

        Note:
            Native backend logging is process-global, so this affects all Llama
            instances in the current Python process.
        """
        set_log_filters(filters, case_sensitive=case_sensitive)


    def add_log_filters(self, filters: Sequence[str]) -> None:
        """Append substring filters for native runtime logs."""
        add_log_filters(filters)


    def get_log_filters(self) -> List[str]:
        """Return the current substring filters for native runtime logs."""
        return get_log_filters()


    def clear_log_filters(self) -> None:
        """Clear all substring filters, including default filters."""
        clear_log_filters()


    def reset_log_filters(self) -> None:
        """Restore default substring filters for native runtime logs."""
        reset_log_filters()

    # LoRA / Adapter Management API

    def load_lora(self, name: str, path: str):
        """Loads a LoRA adapter into VRAM without applying it yet."""
        self._model.load_lora(name, path)

    def unload_lora(self, name: str):
        """Actively unloads a specific LoRA to free up VRAM."""
        self._model.unload_lora(name)

    @property
    def loaded_lora_count(self) -> int:
        """Returns the total number of LoRA adapters currently loaded in VRAM."""
        return self._model.loaded_lora_count

    def list_loras(self) -> List[str]:
        """Returns a list of all registered LoRA names."""
        return self._model.list_loras()

    def unload_all_loras(self):
        """Iterates through the registry and forces VRAM release for all loaded LoRAs."""
        self._model.unload_all_loras()

    def tokenize(
        self, text: bytes, add_bos: bool = True, special: bool = False
    ) -> List[int]:
        """Tokenize a string.

        Args:
            text: The utf-8 encoded string to tokenize.
            add_bos: Whether to add a beginning of sequence token.
            special: Whether to tokenize special tokens.

        Raises:
            RuntimeError: If the tokenization failed.

        Returns:
            A list of tokens.
        """
        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.

        Args:
            tokens: The list of tokens to detokenize.
            prev_tokens: The list of previous tokens. Offset mapping will be performed if provided.
            special: Whether to detokenize special tokens.

        Returns:
            The detokenized string.
        """
        return self.tokenizer_.detokenize(
            tokens, prev_tokens=prev_tokens, special=special
        )

    def set_cache(self, cache: Optional[BaseLlamaCache]):
        """Set the cache.

        Args:
            cache: The cache to set.
        """
        self.cache = cache

    def set_seed(self, seed: int):
        """Set the random seed.

        Args:
            seed: The random seed.
        """
        self._seed = seed

    def attach_threadpool(self, threadpool, threadpool_batch=None) -> None:
        """Attach externally owned ggml threadpools to the native context.

        If ``threadpool_batch`` is omitted, the generation pool is used for
        batch processing too. The pools remain owned by the caller.
        """
        self._ctx.attach_threadpool(threadpool, threadpool_batch)

    def detach_threadpool(self) -> None:
        """Detach externally owned ggml threadpools from the native context."""
        self._ctx.detach_threadpool()

    def reset(self):
        """Reset all Python and native model state."""
        # Use a full memory clear rather than sequence removal: recurrent state
        # cannot always be partially truncated, and hybrid memory must clear
        # both its attention KV cache and recurrent state.
        self._ctx.memory_clear(True)

        # Keep the Python-side token cursor in sync with the empty native state.
        self.n_tokens = 0

        self._last_eval_output_start = 0
        self._last_eval_output_count = 0
        self._restored_logits = None
        self._prefilled_prompt = None
        self._state_needs_speculative_reset = False

        if self.speculative is not None:
            self.speculative.clear()

    def _mark_prefilled_prompt(self) -> None:
        """Hand off a freshly decoded MTMD prompt to exactly one generation."""
        # MTMD batches have their own sparse output mapping. Own the last row
        # instead of assuming Python token positions are native output indices.
        self._restored_logits = np.ctypeslib.as_array(
            self._ctx.get_logits_ith(-1), shape=(self._n_vocab,)
        ).copy() if self.n_tokens else None
        if self._restored_logits is not None:
            self.scores[self.n_tokens - 1 if self._logits_all else 0] = self._restored_logits
        self._last_eval_output_start = self.n_tokens - 1
        self._last_eval_output_count = int(self.n_tokens > 0)
        self._prefilled_prompt = tuple(self.input_ids[:self.n_tokens])
        self._state_needs_speculative_reset = False

    def _sample_output(self, sampler, token_index: int) -> int:
        output_index = token_index - self._last_eval_output_start
        if not 0 <= output_index < self._last_eval_output_count:
            raise RuntimeError(
                "Sampling output is unavailable for this token; decode a valid "
                "suffix or regenerate the prompt before sampling"
            )
        logits = getattr(self, "_restored_logits", None)
        if logits is not None:
            return sampler.sample(self._ctx, idx=output_index, logits=logits)
        return sampler.sample(self._ctx, idx=output_index)

    def abort(self) -> None:
        """
        Safely aborts any ongoing text generation.
        Useful for async API environments or UI interruption buttons.
        """
        if self.verbose:
            print(f"Llama.abort: Abort signal received. Terminating generation...", file=sys.stderr)
        native_abort_flag = getattr(self, "_native_abort_flag", None)
        if native_abort_flag is not None:
            native_abort_flag.value = True
        self._abort_event.set()

    def _validate_eval_tokens(
            self,
            tokens: Sequence[int],
    ) -> None:
        """Validate token ids before passing them to llama_decode.

        This mirrors llama.cpp server-side token validation and prevents invalid
        token ids from reaching the native decode path, where they may cause hard
        crashes instead of Python exceptions.
        """
        if not tokens:
            return

        for i, tok in enumerate(tokens):
            if not isinstance(tok, int):
                raise ValueError(
                    f"Llama.eval: invalid token type at index {i}: "
                    f"{type(tok).__name__}"
                )

            if tok < 0:
                raise ValueError(
                    f"Llama.eval: invalid negative token id at index {i}: {tok}"
                )

            if tok >= self._n_vocab:
                raise ValueError(
                    f"Llama.eval: token out of vocab at index {i}: "
                    f"{tok} >= n_vocab({self._n_vocab})"
                )

    def _memory_seq_rm_or_raise(
        self, seq_id: int, p0: int, p1: int, operation: str
    ) -> None:
        """Remove a native memory range or stop before state can diverge."""
        if not self._ctx.memory_seq_rm(seq_id, p0, p1):
            raise RuntimeError(
                f"{operation}: failed to remove sequence {seq_id} "
                f"memory range [{p0}, {p1})"
            )

    def _speculative_start_position(self, token_cursor: int) -> int:
        """Check that drafting and token-based verification use the same position."""
        pos0 = self._ctx.memory_seq_pos_max(0) + 1
        if pos0 != token_cursor:
            raise NotImplementedError(
                "High-level speculative verification requires contiguous text "
                f"positions: next native position={pos0}, token cursor={token_cursor}. "
                "Non-contiguous media positions require a separate position ledger."
            )
        return pos0

    def _limit_speculative_draft_n_max(self, requested: int) -> int:
        """Keep verification within one batch and the available rollback window."""
        limit = min(max(0, int(requested)), max(0, self.n_batch - 1))
        if (
            self.is_hybrid and self.speculative is not None
            and (self._hybrid_cache_mgr is None or self._hybrid_cache_mgr.max_checkpoints <= 0)
        ):
            limit = min(limit, self._ctx.n_rs_seq())
        return limit

    def _decode_eval_batch(
        self, chunk: Sequence[int], initial_batch_size: int
    ) -> int:
        """Decode one batch, keeping speculative verification atomic."""
        current_batch_size = initial_batch_size

        while current_batch_size > 0:
            self._batch.batch.n_tokens = current_batch_size
            phase_stats = self._active_speculative_phase_stats
            decode_started = time.perf_counter()
            try:
                status = self._ctx.decode(self._batch)
            except internals.LlamaDecodeAbort:
                # llama.cpp may have committed an unknown number of ubatches.
                # A full reset is the only backend-independent way to realign
                # native memory with the Python token ledger (including hybrid
                # and recurrent contexts).
                try:
                    self.reset()
                except Exception as reset_exc:
                    raise RuntimeError(
                        "Llama.eval(decode): failed to reset context after "
                        "native decode abort"
                    ) from reset_exc
                raise
            except Exception as exc:
                failed_pos = self.n_tokens
                self.reset()
                min_pos = min(current_batch_size, 128)
                preview = chunk[:min_pos]
                raise RuntimeError(
                    "Llama.eval(decode): Fatal Decode Error at Pos "
                    f"{failed_pos}, Batch size {current_batch_size}, "
                    f"chunk[:{min_pos}]={preview}: {exc}"
                ) from exc
            finally:
                if phase_stats is not None:
                    phase_stats["target_decode_seconds"] += (
                        time.perf_counter() - decode_started
                    )

            if status == 0:
                if self.speculative is not None:
                    # Match llama.cpp server: synchronize target verification
                    # before timing speculative hidden-state processing.
                    sync_started = time.perf_counter()
                    try:
                        self._ctx.synchronize()
                    finally:
                        if phase_stats is not None:
                            phase_stats["target_sync_seconds"] += (
                                time.perf_counter() - sync_started
                            )
                return current_batch_size

            if status == 1:
                if self._speculative_verifying:
                    raise RuntimeError(
                        "Llama.eval: speculative verification batch cannot be "
                        "split after the backend reported no KV slot; increase "
                        "n_batch/n_ctx or reduce the speculative draft length"
                    )
                if current_batch_size == 1:
                    break
                if self.verbose:
                    print(
                        "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
                continue

            raise RuntimeError(
                "Llama.eval(decode): backend returned fatal status "
                f"{status} at position {self.n_tokens}"
            )

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

    def _process_speculative_batch(self) -> None:
        """Process synchronized target outputs and account only engine work."""
        if self.speculative is None:
            return
        phase_stats = self._active_speculative_phase_stats
        process_started = time.perf_counter()
        try:
            self.speculative.process(self._batch.batch, seq_id=0)
        finally:
            if phase_stats is not None:
                phase_stats["process_calls"] += 1
                phase_stats["process_seconds"] += (
                    time.perf_counter() - process_started
                )

    def _recover_interrupted_speculation(
        self,
        *,
        verification_start: int,
        evaluated_tokens: Sequence[int],
        delivered_accepted: int,
        speculative_checkpoint: Any,
        use_native_rollback: bool,
        active_loras: Optional[List[Dict[str, Union[str, float]]]],
        control_vector: Optional[Dict[str, Any]],
    ) -> str:
        """Align target and draft state after speculative verification is interrupted.

        Keep ``id_last`` and accepted draft tokens already delivered by ``yield``;
        discard all uncommitted speculative tokens. Hybrid contexts use native
        rollback when available, otherwise restore checkpoints and replay the
        committed inputs. Transformer contexts truncate both caches directly.

        Args:
            verification_start: Token count before the verification batch.
            evaluated_tokens: ``id_last`` followed by proposed draft tokens.
            delivered_accepted: Accepted draft tokens already delivered by ``yield``.
            speculative_checkpoint: Draft state saved before verification.
            use_native_rollback: Whether native hybrid rollback is available.
            active_loras: LoRA configuration used when replay is required.
            control_vector: Control vector used when replay is required.

        Returns:
            Recovery strategy: ``"native"``, ``"checkpoint"``, or ``"truncate"``.

        Raises:
            RuntimeError: If target and draft state cannot be aligned. The caller
                resets the model on this failure.
        """
        if self.speculative is None:
            raise RuntimeError("Interrupted verification has no speculative engine")

        # memory_seq_rm uses an exclusive boundary. Keep id_last and only the
        # accepted draft tokens already delivered to the caller.
        committed_position = verification_start + 1 + delivered_accepted

        if self.is_hybrid:
            if use_native_rollback:
                # Roll back target and draft to the same committed boundary.
                if not self._ctx.memory_seq_rm(0, committed_position, -1):
                    raise RuntimeError(
                        "Interrupted native recurrent-state rollback failed"
                    )
                self.speculative.rollback_verified(
                    speculative_checkpoint,
                    delivered_accepted,
                    seq_id=0,
                )
                self.n_tokens = committed_position
                self._last_eval_output_count = max(
                    0, committed_position - self._last_eval_output_start
                )
                return "native"

            # Hybrid state requires checkpoint restore plus committed-input replay.
            if self._hybrid_cache_mgr is None:
                raise RuntimeError(
                    "Interrupted hybrid verification has no checkpoint cache"
                )
            best_ckpt = self._hybrid_cache_mgr.find_best_checkpoint(
                self.input_ids[:verification_start].tolist(), 0
            )
            if (
                best_ckpt is None
                or best_ckpt.pos != verification_start
                or not self._hybrid_cache_mgr.restore_checkpoint(best_ckpt, seq_id=0)
            ):
                raise RuntimeError(
                    "Failed to restore interrupted hybrid verification checkpoint"
                )
            self.speculative.restore(speculative_checkpoint, seq_id=0)
            self.n_tokens = verification_start
            accepted_inputs = list(evaluated_tokens[: 1 + delivered_accepted])
            if accepted_inputs:
                self.eval(
                    accepted_inputs,
                    active_loras=active_loras,
                    control_vector=control_vector,
                    copy_logits=False,
                )
            return "checkpoint"

        # Transformer target and draft caches both support direct truncation.
        self._memory_seq_rm_or_raise(
            0,
            committed_position,
            -1,
            "Llama.generate interrupted speculative rollback",
        )
        self.speculative.truncate(committed_position, seq_id=0)
        self.n_tokens = committed_position
        self._last_eval_output_count = max(
            0, committed_position - self._last_eval_output_start
        )
        return "truncate"

    def eval(
            self,
            tokens: Sequence[int],
            active_loras: Optional[List[Dict[str, Union[str, float]]]] = None,
            control_vector: Optional[Dict[str, Any]] = None,
            copy_logits: bool = True,
    ):
        """Evaluate a list of tokens.

        Args:
            tokens: The token ids to evaluate.
            active_loras: Optional LoRA adapters to apply for this evaluation.
                Each item should contain a ``name`` and an optional ``scale``.
            control_vector: Optional control vector configuration to apply during
                this evaluation.
            copy_logits: Whether to copy the final logits into ``self.scores`` when
                ``logits_all`` is disabled. Set to ``False`` for native sampler paths
                that sample directly from the llama context and do not need
                Python-side logits.
        """
        n_eval = len(tokens)
        if n_eval == 0:
            return
        if self._speculative_verifying and n_eval > self.n_batch:
            raise RuntimeError(
                "Llama.eval: speculative verification batch exceeds n_batch "
                f"({n_eval} > {self.n_batch}); reduce the speculative draft length"
            )

        # Validate token ids before any context shifting, batch construction, or
        # native llama_decode call. Invalid ids may otherwise reach the C/C++ backend
        # and cause hard crashes instead of Python exceptions.
        self._validate_eval_tokens(tokens)
        self._restored_logits = None
        self._prefilled_prompt = None
        self._last_eval_output_count = 0

        # 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._memory_seq_rm_or_raise(
                        0,
                        _n_keep,
                        _n_keep + _n_discard,
                        "Llama.eval context shift",
                    )

                    # 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.
        checkpointing_enabled = (
            self.is_hybrid and self._hybrid_cache_mgr is not None
            and self._hybrid_cache_mgr.max_checkpoints > 0
        )
        if checkpointing_enabled:
            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 or self._speculative_verifying:
                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
            )

            # JIT Dynamic LoRAs Weights Mounting

            # Dynamic LoRA Routing
            if active_loras is not None:
                adapters_to_apply = []
                for lora in active_loras:
                    name = lora.get("name")
                    scale = float(lora.get("scale", 1.0))
                    adapter_obj = getattr(self._model, "_lora_registry", {}).get(name)
                    if adapter_obj:
                        adapters_to_apply.append((adapter_obj, scale))
                    elif self.verbose:
                        print(f"Llama.eval: Warning! LoRA '{name}' not found in registry. Skipping.", file=sys.stderr)

                self._ctx.apply_loras(adapters_to_apply)
            else:
                # Crucial Fallback: Wipe the graph clean if no LoRAs are requested.
                # This guarantees zero weight contamination between different users/slots in a multiplexed environment.
                self._ctx.clear_loras()

            # Dynamic Control Vector (CVec) Injection
            if control_vector is not None:
                data = control_vector.get("data", [])
                il_start = control_vector.get("layer_start", 1)
                il_end = control_vector.get("layer_end", self.n_layer())
                n_embd = self.n_embd()

                self._ctx.apply_cvec(data, n_embd, il_start, il_end)
            else:
                # Ensure the control vector is cleared for a clean state
                self._ctx.clear_cvec()

            # Ordinary prefill may retry with a smaller batch. A speculative
            # [id_last, draft...] verification batch must remain atomic.
            current_batch_size = self._decode_eval_batch(chunk, n_chunk)
            if current_batch_size < current_max_batch:
                current_max_batch = current_batch_size

            self._last_eval_output_start = n_past
            self._last_eval_output_count = current_batch_size

            self._process_speculative_batch()

            # 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 checkpointing_enabled:
                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 logits only when Python-side logits are required.
        # Native sampler can sample directly from ctx, so normal generation does not
        # need to copy n_vocab floats into self.scores on every token.
        if not self._logits_all and copy_logits:
            logits_ptr = self._ctx.get_logits_ith(-1)
            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_lib.llama_logit_bias]:
        if not logit_bias:
            return []
        bias_list = []
        for token, bias in logit_bias.items():
            lb = llama_cpp_lib.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
        penalty_last_n: int = 64,     # last n tokens to penalize (0 = disable penalty, -1 = context size)
        repeat_penalty: float = 1.0,  # 1.0 = disabled
        frequency_penalty: float = 0.0,    # 0.0 = disabled
        present_penalty: float = 0.0, # 0.0 = disabled
        # Mirostat
        mirostat_mode: int = 0,       # 0 = disabled, 1 = mirostat, 2 = mirostat 2.0
        mirostat_eta: float = 0.1,    # learning rate
        mirostat_tau: float = 5.0,    # target entropy
        # XTC
        xtc_probability: float = 0.0, # 0.0 = disabled
        xtc_threshold: float = 0.1,   # > 0.5 disables XTC
        # DRY
        dry_multiplier: float = 0.0,  # 0.0 = disabled;      DRY repetition penalty for tokens extending repetition:
        dry_base: float = 1.75,       # 0.0 = disabled;      multiplier * base ^ (length of sequence before token - allowed length)
        dry_allowed_length: int = 2,  # tokens extending repetitions beyond this receive penalty
        dry_penalty_last_n:int = 64,  # how many tokens to scan for repetitions (0 = disable penalty, -1 = context size)
        dry_seq_breakers: list[str] = ["\n", ":", "\"", "*"], # default sequence breakers for DRY
        # Adaptive
        adaptive_target : float = -1.0, # select tokens near this probability (valid range 0.0 to 1.0; negative = disabled)
        adaptive_decay : float = 0.9,   # EMA decay for adaptation; history ≈ 1/(1-decay) tokens (0.0 - 0.99)
        # Config
        ignore_eos: bool = False,
        # Extra
        logit_bias: Optional[Dict[int, float]] = None,  # logit biases to apply
        logits_processor: Optional[LogitsProcessorList] = None,
        grammar: Optional[LlamaGrammar] = None, # optional BNF-like grammar to constrain sampling
        grammar_lazy: bool = False,
        idx: Optional[int] = None,
        seed: Optional[int] = None,
        # Reasoning Budget Params
        reasoning_budget: int = -1,
        reasoning_start: str = "<think>",
        reasoning_end: str = "</think>",
        reasoning_budget_message: Optional[str] = None,
        reasoning_start_in_prompt: bool = False,
        reasoning_start_max_tokens: Optional[int] = 32,
    ):
        """Sample a token from the model.
        Returns:
            The sampled token.
        """
        assert self.n_tokens > 0

        s_ctx = self._sampling_ctx
        is_temp_ctx = False

        if s_ctx is None:
            is_temp_ctx = True
            params = LlamaSamplingParams(
                # Core
                top_k=top_k,
                top_p=top_p,
                min_p=min_p,
                typical_p=typical_p,
                temp=temp,
                top_n_sigma=top_n_sigma,
                min_keep=min_keep,
                seed=seed if seed is not None else self._seed,

                # Dynamic Temp
                dynatemp_range=dynatemp_range,
                dynatemp_exponent=dynatemp_exponent,

                # Penalties
                penalty_last_n=penalty_last_n if penalty_last_n != 0 else self.last_n_tokens_size,
                penalty_repeat=repeat_penalty,
                penalty_freq=frequency_penalty,
                penalty_present=present_penalty,

                # Mirostat
                mirostat=mirostat_mode,
                mirostat_tau=mirostat_tau,
                mirostat_eta=mirostat_eta,

                # XTC
                xtc_probability=xtc_probability,
                xtc_threshold=xtc_threshold,

                # DRY
                dry_multiplier=dry_multiplier,
                dry_base=dry_base,
                dry_allowed_length=dry_allowed_length,
                dry_penalty_last_n=dry_penalty_last_n,
                dry_sequence_breakers=dry_seq_breakers,

                # Adaptive
                adaptive_target=adaptive_target,
                adaptive_decay=adaptive_decay,

                # Misc
                ignore_eos=ignore_eos,
                logit_bias=self._convert_logit_bias(logit_bias),
                grammar=grammar.grammar if grammar else "",
                grammar_root=grammar.root if grammar else "root",
                grammar_triggers=list(grammar.triggers) if grammar else [],
                grammar_lazy=grammar_lazy,

                # Reasoning Budget
                # This generic controller only counts the first visible reasoning
                # block. Use reasoning_budget=-1 to leave it disabled.
                reasoning_budget=reasoning_budget,
                reasoning_start=reasoning_start,
                reasoning_end=reasoning_end,
                reasoning_budget_message=reasoning_budget_message,
                reasoning_start_in_prompt=reasoning_start_in_prompt,
                reasoning_start_max_tokens=reasoning_start_max_tokens,
            )

            # LogitsProcessor Adapter
            if logits_processor:
                def adapter(token_data_array: llama_cpp_lib.llama_token_data_array):
                    if self._logits_all:
                        current_scores = self._scores[self.n_tokens - 1, :]
                    else:
                        current_scores = self._scores[0, :]
                    new_scores = logits_processor(self._input_ids, current_scores)
                    size = token_data_array.size
                    data_ptr = token_data_array.data
                    for i in range(size):
                        tid = data_ptr[i].id
                        if tid < len(new_scores):
                            data_ptr[i].logit = new_scores[tid]

                params.custom_samplers.append(CustomSampler(adapter))
                # When logits_processor is used, CommonSamplerType.CUSTOM is automatically injected into the samplers.
                if CommonSamplerType.CUSTOM not in params.samplers:
                    params.samplers.insert(3, CommonSamplerType.CUSTOM)

            s_ctx = LlamaSamplingContext(params, self._model)

        assert s_ctx is not None

        try:
            token = self._sample_output(s_ctx, self.n_tokens - 1 if idx is None else idx)
        finally:
            if is_temp_ctx:
                s_ctx.close()

        return token

    def generate(
        self,
        tokens: Sequence[int],
        top_k: int = 40,
        top_p: float = 0.95,
        min_p: float = 0.05,
        typical_p: float = 1.0,
        temp: float = 0.80,
        dynatemp_range: float = 0.0,
        dynatemp_exponent: float = 1.0,
        top_n_sigma: float = -1.00,
        min_keep: int = 0,
        penalty_last_n: int = 64,
        repeat_penalty: float = 1.0,
        frequency_penalty: float = 0.0,
        present_penalty: float = 0.0,
        reset: bool = True,
        mirostat_mode: int = 0,
        mirostat_tau: float = 5.0,
        mirostat_eta: float = 0.1,
        xtc_threshold: float = 0.1,
        xtc_probability: float = 0.0,
        dry_multiplier: float = 0.0,
        dry_base: float = 1.75,
        dry_allowed_length: int = 2,
        dry_penalty_last_n:int = 64,
        dry_seq_breakers: list[str] = ["\n", ":", "\"", "*"],
        adaptive_target : float = -1.0,
        adaptive_decay : float = 0.9,
        use_infill: bool = False,
        ignore_eos: bool = False,
        logit_bias: Optional[Dict[int, float]] = None,
        logits_processor: Optional[LogitsProcessorList] = None,
        stopping_criteria: Optional[StoppingCriteriaList] = None,
        grammar: Optional[LlamaGrammar] = None,
        grammar_lazy: bool = False,
        seed: Optional[int] = None,
        active_loras: Optional[List[Dict[str, Union[str, float]]]] = None,
        control_vector: Optional[Dict[str, Any]] = None,
        # Reasoning Budget Params
        reasoning_budget: int = -1,
        reasoning_start: str = "<think>",
        reasoning_end: str = "</think>",
        reasoning_budget_message: Optional[str] = None,
        reasoning_start_in_prompt: bool = False,
        reasoning_start_max_tokens: Optional[int] = 32,
    ) -> Generator[int, Optional[Sequence[int]], None]:
        """Create a generator of tokens from a prompt.

        Examples:
            >>> llama = Llama("models/ggml-7b.bin")
            >>> tokens = llama.tokenize(b"Hello, world!")
            >>> for token in llama.generate(tokens, top_k=40, top_p=0.95, temp=1.0, repeat_penalty=1.0):
            ...     print(llama.detokenize([token]))

        Args:
            tokens: The prompt tokens to evaluate.
            top_k: Limit the next token selection to the K most probable tokens. (<=0 to use vocab size)
            top_p: Nucleus sampling. Limits selection to a cumulative probability of P.
            min_p: Minimum P sampling. Drops tokens with a probability less than min_p relative to the most likely token.
            typical_p: Locally typical sampling. (1.0 = disabled)
            temp: Temperature. Controls randomness. (<=0.0 greedy, 0.0 no probabilities)
            dynatemp_range: Range of dynamic temperature.
            dynatemp_exponent: Exponent of dynamic temperature.
            top_n_sigma: Limit selection to tokens within n * sigma of the max logit. (-1.0 = disabled)
            min_keep: Minimum tokens to keep for sampling.
            penalty_last_n: Last n tokens to penalize (0 = disable penalty, -1 = context size).
            repeat_penalty: General penalty for repeated tokens. (1.0 = disabled)
            frequency_penalty: Penalty based on the absolute frequency of a token in the prompt.
            present_penalty: Flat penalty applied if a token is present anywhere in the context.
            reset: If True, attempts to automatically match the KV cache prefix to avoid re-evaluation. If False, blindly appends tokens to existing context.
            mirostat_mode: Mirostat sampling mode (0 = disabled, 1 = Mirostat, 2 = Mirostat 2.0).
            mirostat_tau: Target cross-entropy (surprisal) for Mirostat.
            mirostat_eta: Learning rate for Mirostat.
            xtc_threshold: Minimum probability threshold for XTC token removal.
            xtc_probability: Chance for token removal in XTC sampling.
            dry_multiplier: DRY (Don't Repeat Yourself) repetition penalty multiplier (0.0 = disabled).
            dry_base: DRY repetition penalty base value.
            dry_allowed_length: DRY maximum allowed sequence length without penalty.
            dry_penalty_last_n: DRY tokens to scan for repetitions (0 = disabled, -1 = context size).
            dry_seq_breakers: Array of sequence breakers for DRY sampling.
            adaptive_target: Adaptive-p target probability (0.0 to 1.0, negative = disabled).
            adaptive_decay: Adaptive-p decay rate (0.0 to 0.99).
            use_infill: Activate specialized fill-in-the-middle sampler.
            ignore_eos: If True, ignore the End-of-Sequence token.
            logit_bias: Dictionary mapping token IDs to their bias values.
            logits_processor: List of custom Python callbacks to modify logits in-place.
            stopping_criteria: List of custom callbacks to halt generation dynamically.
            grammar: Optional BNF-like grammar (GBNF) to constrain sampling syntax.
            grammar_lazy: If True, activates grammar constraints only on specific trigger tokens.
            seed: RNG seed for sampling. Overrides the instance seed.
            reasoning_budget: Token budget for the first visible reasoning block.
                -1 disables the reasoning budget sampler, 0 forces the block to end
                immediately after it starts, and N > 0 allows at most N generated tokens.
            reasoning_start: Token/text sequence that marks the beginning of the first reasoning block.
                Defaults to "<think>". Pass a model-specific value for non-default tags.
            reasoning_end: Token/text sequence that marks the natural and forced end of the reasoning block.
                Defaults to "</think>".
            reasoning_budget_message: Optional message inserted before reasoning_end when the budget is exhausted.
            reasoning_start_in_prompt: Set True when the prompt/template has already inserted reasoning_start,
                so counting starts from the first generated token.
            reasoning_start_max_tokens: Safety window for non-reasoning models. If reasoning_start is not
                generated within this many output tokens, the sampler becomes a no-op. Set None to wait indefinitely.
            active_loras: A list of dictionaries specifying the LoRA adapters to dynamically apply during generation.
                Each dictionary must contain a "name" key (matching a LoRA previously loaded into VRAM via `load_lora()`)
                and an optional "scale" key (float, defaults to 1.0).
                Example: `[{"name": "role_A", "scale": 0.85}, {"name": "role_B", "scale": 0.5}]`.
            control_vector: A dictionary containing Control Vector (CVec) data for representation engineering.
                Must contain a "data" key with a flattened 1D list of floats.
                Optionally accepts "layer_start" (int, defaults to 1) and "layer_end" (int, defaults to the model's total layer count).
                Note: The length of the "data" list MUST be at least `n_embd * layer_end`, with zero-padding for any skipped early layers.

        Yields:
            The generated tokens.

        KeyboardInterrupt during the generation loop cancels the request and
        clears partial state. Completion APIs report finish_reason="abort".
        """
        original_tokens = list(tokens)
        if (
            self.is_hybrid and self.speculative is not None
            and (self._hybrid_cache_mgr is None or self._hybrid_cache_mgr.max_checkpoints <= 0)
            and (self._ctx.n_rs_seq() <= 0 or not self.speculative.can_follow_target_native_rollback())
        ):
            raise RuntimeError(
                "Speculative decoding on this hybrid/recurrent target requires "
                "ctx_checkpoints > 0 or native recurrent-state rollback support"
            )
        prefilled = getattr(self, "_prefilled_prompt", None)
        self._prefilled_prompt = None
        use_prefill = (
            reset and prefilled is not None and tuple(original_tokens) == prefilled
            and len(original_tokens) == self.n_tokens
            and self._last_eval_output_count > 0
        )
        if use_prefill:
            # MTMD already decoded the prompt, including embeddings which
            # cannot be reconstructed by replaying its virtual negative IDs.
            reset = False
            tokens = []
        if getattr(self, "_state_needs_speculative_reset", False):
            if self.speculative is not None and not reset:
                raise RuntimeError(
                    "LlamaState does not restore draft state; start speculative "
                    "generation with reset=True and the full text prompt"
                )
        # The Python MTP engine maintains a second context and pending hidden
        # state. Until speculative checkpoints are persisted alongside the
        # public prompt cache, rebuild both contexts together for a new reset
        # generation instead of reusing only the target KV cache.
        # Check for kv cache prefix match
        if reset and self.speculative is None and self.n_tokens > 0:
            # 1. First, check for a 100% exact match of the entire sequence
            full_match_prefix = self.longest_token_prefix(self._input_ids, tokens, self.verbose)

            # Reuse a valid final output when rollback is disabled or the snapshot
            # owns its logits. MTMD prefill uses the separate handoff above.
            if (
                full_match_prefix == len(tokens)
                and full_match_prefix == self.n_tokens
                and self.is_hybrid
                and (
                    self._hybrid_cache_mgr is None or self._hybrid_cache_mgr.max_checkpoints <= 0
                    or self._restored_logits is not None
                )
                and self._last_eval_output_start <= self.n_tokens - 1
                < self._last_eval_output_start + self._last_eval_output_count
            ):
                reset = False
                longest_prefix = len(tokens)
                tokens = tokens[longest_prefix:] # Empties the tokens array to bypass evaluation
                if self.verbose:
                    print(f"Llama.generate: Hybrid single-turn full match ({longest_prefix} tokens). Bypassing rollback/truncation.", file=sys.stderr)

            # --- STANDARD PATH: Force N-1 re-evaluation ---
            else:
                # By matching against `tokens[:-1]`, we intentionally drop the last token.
                # This forces the engine to re-evaluate the final token to refresh sampling logits.
                longest_prefix = min(full_match_prefix, max(0, len(tokens) - 1))

                if longest_prefix > 0:
                    reset = False

                    # Physically erase trailing "ghost" tokens from the C++ KV cache
                    # to prevent attention misalignment in multi-round chats.
                    if longest_prefix < self.n_tokens:
                        if self.is_hybrid:
                            if self.verbose:
                                print(f"Llama.generate: Hybrid model rollback triggered.", file=sys.stderr)

                            best_ckpt = None
                            if self._hybrid_cache_mgr is not None and self._hybrid_cache_mgr.max_checkpoints > 0:
                                best_ckpt = self._hybrid_cache_mgr.find_best_checkpoint(original_tokens[:-1], 0)
                            if best_ckpt is not None and self._hybrid_cache_mgr.restore_checkpoint(best_ckpt, seq_id=0):
                                actual_prefix = best_ckpt.pos
                            else:
                                # Fallback: No checkpoint found, must fully clear the context to prevent poisoning
                                actual_prefix = 0
                                self._ctx.memory_clear(True)

                            self.n_tokens = actual_prefix
                            self._restored_logits = None
                            self._last_eval_output_count = 0
                            tokens = original_tokens[actual_prefix:]
                            if self.verbose:
                                print(
                                    f"Llama.generate: {actual_prefix} prefix-match hit, "
                                    f"remaining {len(tokens)} prompt tokens to eval",
                                    file=sys.stderr,
                                )
                        else:
                            if self.verbose:
                                print(f"Llama.generate: Truncating KV cache size from {self.n_tokens} to {longest_prefix}", file=sys.stderr)
                            self._memory_seq_rm_or_raise(
                                0,
                                longest_prefix,
                                -1,
                                "Llama.generate prefix truncation",
                            )

                            # Adjust the tokens array and cursor to reuse the matched cache
                            self.n_tokens = longest_prefix
                            tokens = tokens[longest_prefix:]

                            if self.verbose:
                                print(
                                    f"Llama.generate: {longest_prefix} prefix-match hit, "
                                    f"remaining {len(tokens)} prompt tokens to eval",
                                    file=sys.stderr,
                                )
                    else:
                        # The live context already ends at the matched prefix.
                        # Appending the full prompt would evaluate that prefix twice.
                        tokens = original_tokens[longest_prefix:]
        if reset:
            # No prefix matched at all. Completely clear the KV cache to prevent context poisoning.
            self.reset()
            if self.verbose:
                print("Llama.generate: Context reset requested or no prefix match. Cleared KV cache.", file=sys.stderr)

        # Reset mirostat sampling
        params = LlamaSamplingParams(
            # Core Sampling
            top_k=top_k,
            top_p=top_p,
            min_p=min_p,
            typical_p=typical_p,
            temp=temp,
            top_n_sigma=top_n_sigma,
            min_keep=min_keep,

            # Dynamic Temperature
            dynatemp_range=dynatemp_range,
            dynatemp_exponent=dynatemp_exponent,

            # Penalties
            penalty_last_n=penalty_last_n,
            penalty_repeat=repeat_penalty,
            penalty_freq=frequency_penalty,
            penalty_present=present_penalty,

            # Mirostat
            mirostat=mirostat_mode,
            mirostat_tau=mirostat_tau,
            mirostat_eta=mirostat_eta,

            # XTC
            xtc_probability=xtc_probability,
            xtc_threshold=xtc_threshold,

            # DRY (Don't Repeat Yourself)
            dry_multiplier=dry_multiplier,
            dry_base=dry_base,
            dry_allowed_length=dry_allowed_length,
            dry_penalty_last_n=dry_penalty_last_n,
            dry_sequence_breakers=dry_seq_breakers,

            # Adaptive P
            adaptive_target=adaptive_target,
            adaptive_decay=adaptive_decay,

            # Misc
            ignore_eos=ignore_eos,
            logit_bias=self._convert_logit_bias(logit_bias),
            grammar=grammar.grammar if grammar else "",
            grammar_root=grammar.root if grammar else "root",
            grammar_triggers=list(grammar.triggers) if grammar else [],
            grammar_lazy=grammar_lazy,
            seed=seed if seed is not None else self._seed,

            # Reasoning Budget
            # Keeps the core sampler model-agnostic: callers provide the visible
            # reasoning start/end tags, and -1 keeps the controller disabled.
            reasoning_budget=reasoning_budget,
            reasoning_start=reasoning_start,
            reasoning_end=reasoning_end,
            reasoning_budget_message=reasoning_budget_message,
            reasoning_start_in_prompt=reasoning_start_in_prompt,
            reasoning_start_max_tokens=reasoning_start_max_tokens,
        )

        # Register custom python-level logits processors if provided
        if logits_processor:
            def adapter(token_data_array: llama_cpp_lib.llama_token_data_array):
                if self._logits_all:
                    current_scores = self._scores[sample_idx, :]
                else:
                    current_scores = self._scores[0, :]
                new_scores = logits_processor(self._input_ids[:sample_idx + 1], current_scores)

                size = token_data_array.size
                data_ptr = token_data_array.data
                for i in range(size):
                    tid = data_ptr[i].id
                    if tid < len(new_scores):
                        data_ptr[i].logit = new_scores[tid]

            custom_sampler = CustomSampler(adapter)
            params.custom_samplers.append(custom_sampler)

            if CommonSamplerType.CUSTOM not in params.samplers:
                params.samplers.insert(3, CommonSamplerType.CUSTOM)

        # Free previous sampling context to prevent memory leaks
        if getattr(self, "_sampling_ctx", None) is not None:
            self._sampling_ctx.close()
            self._sampling_ctx = None

        self._sampling_ctx = LlamaSamplingContext(params, self._model)

        # Native sampler samples directly from ctx. Python-side logits are only needed
        # for compatibility hooks that explicitly consume self._scores.
        copy_logits = (
            self._logits_all
            or logits_processor is not None
            or (stopping_criteria is not None and getattr(stopping_criteria, "requires_logits", True))
        )

        if not tokens and self.n_tokens > 0 and copy_logits:
            logits = self._restored_logits
            if logits is None:
                logits = np.ctypeslib.as_array(
                    self._ctx.get_logits_ith(self.n_tokens - 1 - self._last_eval_output_start),
                    shape=(self._n_vocab,),
                )
            self.scores[self.n_tokens - 1 if self._logits_all else 0] = logits

        sample_idx = self.n_tokens + len(tokens) - 1
        tokens = list(tokens)

        # llama.cpp calls begin() after the prompt batch has been decoded and fed
        # through common_speculative_process(). Keep the same ordering so model-based
        # engines can validate/capture their prompt-side state first.
        speculative_begun = self.speculative is None

        # Main evaluation and generation loop
        pending_draft_count = 0
        speculative_drafted = 0
        speculative_verified = 0
        speculative_accepted = 0
        speculative_verification_steps = 0
        speculative_rollbacks = 0
        speculative_native_rollbacks = 0
        speculative_checkpoint_rollbacks = 0
        speculative_decode_tokens = 0
        speculative_decode_seconds = 0.0
        speculative_decode_started: Optional[float] = None
        speculative_ttft_seconds = 0.0
        speculative_time_to_last_token_seconds = 0.0
        verification_active = False
        verification_start = 0
        verification_checkpoint = None
        verification_use_native_rollback = False
        verification_evaluated_tokens: List[int] = []
        verification_delivered_accepted = 0
        interrupted_verification_reconciled = True
        speculative_phase_stats: Dict[str, Any] = {
            "begin_calls": 0,
            "draft_calls": 0,
            "process_calls": 0,
            "accept_calls": 0,
            "generated_drafts": 0,
            "accepted_drafts": 0,
            "accepted_tokens": 0,
            "accepted_tokens_per_position": [],
            "begin_seconds": 0.0,
            "draft_seconds": 0.0,
            "target_decode_seconds": 0.0,
            "target_sync_seconds": 0.0,
            "process_seconds": 0.0,
            "accept_seconds": 0.0,
        }
        if self.speculative is not None:
            self._active_speculative_phase_stats = speculative_phase_stats
            self.speculative.reset_checkpoint_stats()

        def speculative_begin(prompt_tokens: Sequence[int]) -> None:
            assert self.speculative is not None
            started = time.perf_counter()
            try:
                self.speculative.begin(prompt_tokens, seq_id=0)
            finally:
                speculative_phase_stats["begin_calls"] += 1
                speculative_phase_stats["begin_seconds"] += (
                    time.perf_counter() - started
                )

        def speculative_draft(
            history: npt.NDArray[np.intc],
            *,
            n_past: int,
            id_last: int,
            n_max: int,
        ) -> npt.NDArray[np.intc]:
            assert self.speculative is not None
            n_max = self._limit_speculative_draft_n_max(n_max)
            if n_max <= 0:
                return np.empty(0, dtype=np.intc)
            pos0 = self._speculative_start_position(n_past)
            started = time.perf_counter()
            try:
                result = self.speculative.draft_at_position(
                    history,
                    pos0=pos0,
                    id_last=id_last,
                    n_max=n_max,
                    seq_id=0,
                )
            finally:
                speculative_phase_stats["draft_calls"] += 1
                speculative_phase_stats["draft_seconds"] += (
                    time.perf_counter() - started
                )
            if len(result) > 0:
                speculative_phase_stats["generated_drafts"] += 1
            return result

        def time_speculative_accept(operation: Callable[[], None]) -> None:
            started = time.perf_counter()
            try:
                operation()
            finally:
                speculative_phase_stats["accept_seconds"] += (
                    time.perf_counter() - started
                )

        try:
            # Match examples/speculative-simple: keep the last prompt token as
            # id_last, process the preceding prompt first, then verify
            # [id_last, draft...] together. This lets speculation cover the very
            # first generated token instead of starting one token late.
            if self.speculative is not None and tokens:
                id_last = int(tokens[-1])
                prompt_prefix = tokens[:-1]
                if prompt_prefix:
                    self.eval(
                        prompt_prefix,
                        active_loras=active_loras,
                        control_vector=control_vector,
                        copy_logits=False,
                    )

                speculative_begin(self.input_ids[: self.n_tokens].tolist())
                speculative_begun = True
                # The prompt prefix is ingested; id_last is intentionally held
                # back for the first verification batch. Time active speculative
                # work from here, excluding time suspended at yield.
                speculative_decode_started = time.perf_counter()

                history_end = self.n_tokens + 1
                self.input_ids[self.n_tokens] = id_last
                room = self._n_ctx - history_end - 1
                initial_draft = np.empty(0, dtype=np.intc)
                if room > 0:
                    initial_draft = speculative_draft(
                        self.input_ids[:history_end],
                        n_past=self.n_tokens,
                        id_last=id_last,
                        n_max=room,
                    )
                tokens = [id_last] + initial_draft.astype(int).tolist()
                pending_draft_count = len(initial_draft)
                speculative_drafted += len(initial_draft)

            while True:
                n_drafted = pending_draft_count
                pending_draft_count = 0
                self._speculative_verifying = n_drafted > 0
                if n_drafted > 0:
                    speculative_verified += n_drafted
                    speculative_verification_steps += 1
                n_accepted = 0
                verification_delivered_accepted = 0
                accept_handled = False
                evaluated_tokens = list(tokens)
                verification_start = self.n_tokens
                speculative_checkpoint = None
                use_native_speculative_rollback = False
                if n_drafted > 0 and self.speculative is not None:
                    speculative_checkpoint = (
                        self.speculative.take_verification_checkpoint(seq_id=0)
                    )
                    if self.is_hybrid:
                        use_native_speculative_rollback = (
                            self._ctx.n_rs_seq() >= n_drafted
                            # The target capacity check above and the engine's
                            # ability to realign its own state are independent.
                            and self.speculative.can_follow_target_native_rollback()
                        )
                        if not use_native_speculative_rollback:
                            if (
                                self._hybrid_cache_mgr is None
                                or self._hybrid_cache_mgr.max_checkpoints <= 0
                            ):
                                raise RuntimeError(
                                    "Speculative decoding on this hybrid/recurrent "
                                    "target requires ctx_checkpoints > 0"
                                )
                            if not self._hybrid_cache_mgr.save_checkpoint(
                                current_pos=verification_start,
                                tokens=self.input_ids[:verification_start].tolist(),
                                seq_id=0,
                            ):
                                raise RuntimeError(
                                    "Failed to checkpoint hybrid target before draft verification"
                                )
                if len(tokens) > 0:
                    # For hybrid models processing a prompt (len > 1), force an N-1 checkpoint
                    # to safely allow 1-token rollbacks (e.g., for seed changes on 100% prompt matches).
                    # ONLY apply this if rollback capabilities are enabled (max_checkpoints > 0).
                    if (
                        self.is_hybrid
                        and self.speculative is None
                        and self._hybrid_cache_mgr is not None
                        and self._hybrid_cache_mgr.max_checkpoints > 0
                        and len(tokens) > 1
                    ):
                        body_tokens = tokens[:-1]
                        last_token = [tokens[-1]]

                        # 1. Evaluate up to N-1 without copying logits.
                        self.eval(
                            body_tokens,
                            active_loras=active_loras,
                            control_vector=control_vector,
                            copy_logits=False,
                        )

                        # 2. Save the N-1 state snapshot
                        current_history = self._input_ids[:self.n_tokens].tolist()
                        self._hybrid_cache_mgr.save_checkpoint(
                            current_pos=self.n_tokens,
                            tokens=current_history,
                            seq_id=0
                        )
                        # 3. Evaluate final token. Copy logits only if Python-side hooks need them.
                        self.eval(
                            last_token,
                            active_loras=active_loras,
                            control_vector=control_vector,
                            copy_logits=copy_logits,
                        )
                    else:
                        # Standard evaluation or single-token generation step
                        self.eval(
                            tokens,
                            active_loras=active_loras,
                            control_vector=control_vector,
                            copy_logits=copy_logits and not self._speculative_verifying,
                        )

                if self.speculative is not None and not speculative_begun:
                    speculative_begin(self.input_ids[: self.n_tokens].tolist())
                    speculative_begun = True
                    speculative_decode_started = time.perf_counter()

                if n_drafted > 0 and self.speculative is not None:
                    # From this point until accept/rollback completes, the target
                    # and speculative contexts contain a verification transaction.
                    # A return, exception, or GeneratorExit at yield must reconcile
                    # the uncommitted suffix before the context can be reused.
                    verification_active = True
                    verification_checkpoint = speculative_checkpoint
                    verification_use_native_rollback = use_native_speculative_rollback
                    verification_evaluated_tokens = evaluated_tokens

                # Sample loop
                while sample_idx < self.n_tokens:
                    if self._abort_event.is_set():
                        return

                    output_idx = sample_idx - self._last_eval_output_start
                    if not 0 <= output_idx < self._last_eval_output_count:
                        raise RuntimeError(
                            "Llama.generate: sampling index is outside the most recent "
                            "decode output batch: "
                            f"token_index={sample_idx}, output_start="
                            f"{self._last_eval_output_start}, output_count="
                            f"{self._last_eval_output_count}"
                        )
                    if self._speculative_verifying and copy_logits and not self._logits_all:
                        # Verification contains several output rows. Refresh the
                        # hook buffer for this sample, not the batch's last row.
                        self.scores[0] = np.ctypeslib.as_array(
                            self._ctx.get_logits_ith(output_idx), shape=(self._n_vocab,),
                        )
                    token = self._sample_output(self._sampling_ctx, sample_idx)
                    self._sampling_ctx.accept(token, False if grammar is None else True)

                    sample_idx += 1

                    if (
                        n_drafted > 0
                        and sample_idx < self.n_tokens
                        and token == self._input_ids[sample_idx]
                    ):
                        n_accepted += 1
                        speculative_accepted += 1

                    if stopping_criteria is not None:
                        if stopping_criteria(
                            self._input_ids[: sample_idx],
                            self._scores[sample_idx - 1 if self._logits_all else 0, :]
                        ):
                            return

                    # Yield the generated token to the caller
                    if self.speculative is not None:
                        now = time.perf_counter()
                        if speculative_decode_started is not None:
                            speculative_decode_seconds += (
                                now - speculative_decode_started
                            )
                        speculative_decode_started = None
                        if speculative_decode_tokens == 0:
                            speculative_ttft_seconds = speculative_decode_seconds
                        speculative_decode_tokens += 1
                        # Stop throughput timing at token delivery. Work performed
                        # after the final yield must not reduce the reported rate.
                        speculative_time_to_last_token_seconds = (
                            speculative_decode_seconds
                        )
                    # Record only accepted draft tokens whose output has actually
                    # crossed the generator boundary. stopping_criteria returns
                    # above this point and therefore must not commit this token.
                    verification_delivered_accepted = n_accepted
                    tokens_or_none = yield token
                    if self.speculative is not None:
                        speculative_decode_started = time.perf_counter()

                    tokens.clear()
                    tokens.append(token)

                    if tokens_or_none is not None:
                        tokens.extend(tokens_or_none)

                    # Rollback Check: A previously evaluated token (e.g. from speculative decoding)
                    # mismatched the newly sampled token. We must rollback the KV cache.
                    if sample_idx < self.n_tokens and token != self._input_ids[sample_idx]:
                        if self.speculative is not None:
                            speculative_rollbacks += 1
                        if self.is_hybrid:
                            if self.speculative is not None:
                                if use_native_speculative_rollback:
                                    speculative_native_rollbacks += 1
                                    if not self._ctx.memory_seq_rm(
                                        0, sample_idx, -1
                                    ):
                                        raise RuntimeError(
                                            "Native recurrent-state speculative rollback failed"
                                        )
                                    time_speculative_accept(
                                        lambda: self.speculative.rollback_verified(
                                            speculative_checkpoint,
                                            n_accepted,
                                            seq_id=0,
                                        )
                                    )
                                    self.n_tokens = sample_idx
                                else:
                                    speculative_checkpoint_rollbacks += 1
                                    assert self._hybrid_cache_mgr is not None
                                    best_ckpt = self._hybrid_cache_mgr.find_best_checkpoint(
                                        self.input_ids[:verification_start].tolist(), 0
                                    )
                                    if (
                                        best_ckpt is None
                                        or best_ckpt.pos != verification_start
                                        or not self._hybrid_cache_mgr.restore_checkpoint(
                                            best_ckpt, seq_id=0
                                        )
                                    ):
                                        raise RuntimeError(
                                            "Failed to restore the exact hybrid checkpoint "
                                            "for speculative rejection"
                                        )
                                    time_speculative_accept(
                                        lambda: self.speculative.restore(
                                            speculative_checkpoint, seq_id=0
                                        )
                                    )
                                    self.n_tokens = verification_start
                                    accepted_inputs = evaluated_tokens[: 1 + n_accepted]
                                    if accepted_inputs:
                                        self._speculative_verifying = False
                                        self.eval(
                                            accepted_inputs,
                                            active_loras=active_loras,
                                            control_vector=control_vector,
                                            copy_logits=False,
                                        )
                                accept_handled = True
                                verification_active = False
                            else:
                                best_ckpt = self._hybrid_cache_mgr.find_best_checkpoint(
                                    self._input_ids[:sample_idx].tolist(), 0
                                )
                                if best_ckpt and self._hybrid_cache_mgr.restore_checkpoint(
                                    best_ckpt, seq_id=0
                                ):
                                    self.n_tokens = best_ckpt.pos
                                else:
                                    self.reset()
                        else:
                            if self.verbose and self.speculative is None:
                                print(f"Llama.generate: Draft token rejected. Truncating context to {sample_idx}.", file=sys.stderr)
                            if self.speculative is not None:
                                speculative_native_rollbacks += 1
                            self._memory_seq_rm_or_raise(
                                0,
                                sample_idx,
                                -1,
                                "Llama.generate speculative rollback",
                            )
                            if self.speculative is not None:
                                time_speculative_accept(
                                    lambda: self.speculative.truncate(
                                        sample_idx, seq_id=0
                                    )
                                )
                            self.n_tokens = sample_idx

                        break

                # llama.cpp-compatible stateful speculative decoding.
                if self.speculative is not None:
                    if n_drafted > 0 and not accept_handled:
                        time_speculative_accept(
                            lambda: self.speculative.accept(n_accepted, seq_id=0)
                        )
                        verification_active = False

                    if n_drafted > 0:
                        speculative_phase_stats["accept_calls"] += 1
                        speculative_phase_stats["accepted_tokens"] += n_accepted
                        if n_accepted > 0:
                            speculative_phase_stats["accepted_drafts"] += 1
                        per_position = speculative_phase_stats[
                            "accepted_tokens_per_position"
                        ]
                        if len(per_position) < n_accepted:
                            per_position.extend([0] * (n_accepted - len(per_position)))
                        for position in range(n_accepted):
                            per_position[position] += 1

                    self.input_ids[self.n_tokens : self.n_tokens + len(tokens)] = tokens
                    history_end = self.n_tokens + len(tokens)
                    # Match server_slot::get_n_draft_max(): id_last is evaluated at
                    # the current target position, and one extra context position is
                    # kept available for shifting/continuation.
                    room = self._n_ctx - history_end - 1
                    if room > 0:
                        draft_tokens = speculative_draft(
                            self.input_ids[:history_end],
                            n_past=self.n_tokens,
                            id_last=int(tokens[-1]),
                            n_max=room,
                        )
                        tokens.extend(draft_tokens.astype(int).tolist())
                        pending_draft_count = len(draft_tokens)
                        speculative_drafted += len(draft_tokens)

                # Deprecated stateless draft-model compatibility path.
                elif self.draft_model is not None:
                    if self.is_hybrid:
                        if self.verbose:
                            print("Llama.generate: Speculative decoding is skipped for Hybrid models.", file=sys.stderr)
                    else:
                        self.input_ids[self.n_tokens : self.n_tokens + len(tokens)] = tokens
                        draft_tokens = self.draft_model(
                            self.input_ids[: self.n_tokens + len(tokens)]
                        )
                        tokens.extend(
                            draft_tokens.astype(int)[
                                : self._n_ctx - self.n_tokens - len(tokens)
                            ]
                        )
        except KeyboardInterrupt:
            # Ctrl+C can arrive while a logits accessor synchronizes queued
            # decode work. Discard partial state instead of attempting rollback.
            verification_active = False
            self.reset()
            self._abort_event.set()
            if self.verbose:
                print(
                    "Llama.generate: KeyboardInterrupt received; generation cancelled, "
                    "context state cleared (finish_reason=abort).",
                    file=sys.stderr,
                )
            return
        except internals.LlamaDecodeAbort:
            # Convert the recoverable native control-flow signal into normal
            # generator termination. _decode_eval_batch() has already cleared
            # the possibly partial native state; marking the event lets the
            # completion layer report finish_reason="abort".
            self._abort_event.set()
            verification_active = False
            return
        except Exception:
            # Failed rollback may have changed target or draft memory already.
            verification_active = False
            self.reset()
            raise
        finally:
            self._speculative_verifying = False
            if verification_active and self.speculative is not None:
                try:
                    rollback_mode = (
                        self._recover_interrupted_speculation(
                            verification_start=verification_start,
                            evaluated_tokens=verification_evaluated_tokens,
                            delivered_accepted=verification_delivered_accepted,
                            speculative_checkpoint=verification_checkpoint,
                            use_native_rollback=verification_use_native_rollback,
                            active_loras=active_loras,
                            control_vector=control_vector,
                        )
                    )
                    if rollback_mode == "checkpoint":
                        speculative_checkpoint_rollbacks += 1
                    else:
                        speculative_native_rollbacks += 1

                    speculative_rollbacks += 1
                except Exception as exc:
                    interrupted_verification_reconciled = False
                    if self.verbose:
                        print(
                            "Llama.generate: failed to reconcile interrupted "
                            f"speculative verification; resetting context: {exc}",
                            file=sys.stderr,
                        )
                    try:
                        self.reset()
                    except Exception as reset_exc:
                        if self.verbose:
                            print(
                                "Llama.generate: failed to reset after interrupted "
                                f"verification cleanup error: {reset_exc}",
                                file=sys.stderr,
                            )
                finally:
                    verification_active = False
            if self._active_speculative_phase_stats is speculative_phase_stats:
                self._active_speculative_phase_stats = None
            # Throughput ends at delivery of the last output token. In
            # particular, do not count work performed after the last yield when
            # the caller closes or resumes the generator.
            speculative_decode_started = None
            acceptance_rate = (
                speculative_accepted / speculative_verified
                if speculative_verified > 0
                else 0.0
            )
            timing_stats = _speculative_generation_timing_stats(
                speculative_decode_tokens,
                speculative_ttft_seconds,
                speculative_time_to_last_token_seconds,
            )
            accept_calls = speculative_phase_stats["accept_calls"]
            accepted_tokens = speculative_phase_stats["accepted_tokens"]
            mean_accepted_length = (
                1.0 + accepted_tokens / accept_calls
                if accept_calls > 0
                else 0.0
            )
            acceptance_rate_per_position = [
                count / accept_calls
                for count in speculative_phase_stats[
                    "accepted_tokens_per_position"
                ]
            ] if accept_calls > 0 else []
            draft_token_acceptance_rate = (
                accepted_tokens / speculative_drafted
                if speculative_drafted > 0
                else 0.0
            )
            generated_drafts = speculative_phase_stats["generated_drafts"]
            accepted_drafts = speculative_phase_stats["accepted_drafts"]
            draft_batch_acceptance_rate = (
                accepted_drafts / generated_drafts
                if generated_drafts > 0
                else 0.0
            )
            checkpoint_stats = (
                self.speculative.checkpoint_stats()
                if self.speculative is not None
                else {}
            )
            self.last_speculative_stats = {
                "drafted": speculative_drafted,
                "verified": speculative_verified,
                "accepted": speculative_accepted,
                "begin_calls": speculative_phase_stats["begin_calls"],
                "draft_calls": speculative_phase_stats["draft_calls"],
                "process_calls": speculative_phase_stats["process_calls"],
                "accept_calls": accept_calls,
                "generated_drafts": generated_drafts,
                "accepted_drafts": accepted_drafts,
                "draft_batch_acceptance_rate": draft_batch_acceptance_rate,
                "accepted_draft_tokens": accepted_tokens,
                "draft_token_acceptance_rate": draft_token_acceptance_rate,
                "mean_accepted_length": mean_accepted_length,
                "acceptance_rate_per_position": acceptance_rate_per_position,
                "begin_seconds": speculative_phase_stats["begin_seconds"],
                "draft_seconds": speculative_phase_stats["draft_seconds"],
                "target_decode_seconds": speculative_phase_stats[
                    "target_decode_seconds"
                ],
                "target_sync_seconds": speculative_phase_stats[
                    "target_sync_seconds"
                ],
                "process_seconds": speculative_phase_stats["process_seconds"],
                "accept_seconds": speculative_phase_stats["accept_seconds"],
                "checkpoint_captures": int(checkpoint_stats.get("captures", 0)),
                "checkpoint_restores": int(checkpoint_stats.get("restores", 0)),
                "checkpoint_verification_reuses": int(
                    checkpoint_stats.get("verification_reuses", 0)
                ),
                "checkpoint_native_captures": int(
                    checkpoint_stats.get("native_captures", 0)
                ),
                "checkpoint_native_restores": int(
                    checkpoint_stats.get("native_restores", 0)
                ),
                "checkpoint_device_captures": int(
                    checkpoint_stats.get("device_captures", 0)
                ),
                "checkpoint_device_restores": int(
                    checkpoint_stats.get("device_restores", 0)
                ),
                "checkpoint_native_verification_rollbacks": int(
                    checkpoint_stats.get("native_verification_rollbacks", 0)
                ),
                "checkpoint_buffer_bytes": int(
                    checkpoint_stats.get("buffer_bytes", 0)
                ),
                "checkpoint_capture_seconds": float(
                    checkpoint_stats.get("capture_seconds", 0.0)
                ),
                "checkpoint_restore_seconds": float(
                    checkpoint_stats.get("restore_seconds", 0.0)
                ),
                "verification_steps": speculative_verification_steps,
                "rollbacks": speculative_rollbacks,
                "native_rollbacks": speculative_native_rollbacks,
                "checkpoint_rollbacks": speculative_checkpoint_rollbacks,
                "acceptance_rate": acceptance_rate,
                **timing_stats,
                # Keep the pre-existing keys as aliases. Their timing now ends at
                # the last output token instead of including generator cleanup.
                "decode_tokens": timing_stats["generation_tokens"],
                "decode_seconds": timing_stats["generation_seconds"],
                "decode_tokens_per_second": timing_stats[
                    "generation_tokens_per_second"
                ],
            }
            if self.verbose and self.speculative is not None and speculative_begun:
                spec_name = (
                    self.speculative_config.spec_type.to_str()
                    if self.speculative_config is not None
                    else type(self.speculative).__name__
                )
                per_position_text = ", ".join(
                    f"{rate:.1%}" for rate in acceptance_rate_per_position
                )
                native_captures = int(checkpoint_stats.get("native_captures", 0))
                device_captures = int(checkpoint_stats.get("device_captures", 0))
                if native_captures and device_captures:
                    checkpoint_mode = (
                        f"mixed (native {native_captures:,}, "
                        f"device {device_captures:,})"
                    )
                elif native_captures:
                    checkpoint_mode = "native-rs"
                elif device_captures:
                    checkpoint_mode = "on-device"
                else:
                    checkpoint_mode = "none"
                stats_lines = [
                    f"Llama.generate: {spec_name} summary",
                    "  Calls       "
                    f"begin {speculative_phase_stats['begin_calls']:,} | "
                    f"draft {speculative_phase_stats['draft_calls']:,} | "
                    f"process {speculative_phase_stats['process_calls']:,} | "
                    f"accept {accept_calls:,}",
                    "  Acceptance  "
                    f"batches {accepted_drafts:,} / {generated_drafts:,} = "
                    f"{draft_batch_acceptance_rate:.1%} | "
                    f"tokens {accepted_tokens:,} / {speculative_drafted:,} = "
                    f"{draft_token_acceptance_rate:.1%} | "
                    f"mean step length {mean_accepted_length:.2f} | "
                    f"by position [{per_position_text}]",
                    "  Phase time  "
                    f"begin {_format_speculative_duration(speculative_phase_stats['begin_seconds'])} | "
                    f"draft {_format_speculative_duration(speculative_phase_stats['draft_seconds'])} | "
                    f"target decode {_format_speculative_duration(speculative_phase_stats['target_decode_seconds'])} | "
                    f"target sync {_format_speculative_duration(speculative_phase_stats['target_sync_seconds'])} | "
                    f"process {_format_speculative_duration(speculative_phase_stats['process_seconds'])} | "
                    f"accept {_format_speculative_duration(speculative_phase_stats['accept_seconds'])}",
                    "  Checkpoint  "
                    f"capture {int(checkpoint_stats.get('captures', 0)):,} in "
                    f"{_format_speculative_duration(float(checkpoint_stats.get('capture_seconds', 0.0)))} | "
                    f"restore {int(checkpoint_stats.get('restores', 0)):,} in "
                    f"{_format_speculative_duration(float(checkpoint_stats.get('restore_seconds', 0.0)))} | "
                    f"reuse {int(checkpoint_stats.get('verification_reuses', 0)):,} | "
                    f"mode {checkpoint_mode}",
                    "  Output      "
                    f"{timing_stats['generation_tokens']:,} tokens / "
                    f"{timing_stats['generation_seconds']:.3f} s = "
                    f"{timing_stats['generation_tokens_per_second']:.2f} tok/s | "
                    f"sustained {timing_stats['sustained_tokens_per_second']:.2f} tok/s | "
                    f"TTFT {timing_stats['time_to_first_token_seconds'] * 1000.0:.2f} ms | "
                    f"rollbacks {speculative_rollbacks:,} "
                    f"(native {speculative_native_rollbacks:,}, "
                    f"checkpoint {speculative_checkpoint_rollbacks:,})",
                ]
                print("\n".join(stats_lines), file=sys.stderr)
            # Preserve ordinary prompt history. Speculative requests reset both
            # contexts on the next request; their target-only final snapshot
            # cannot restore the engine and is never used for verification.
            if (
                self.is_hybrid
                and self.speculative is None
                and self.n_tokens > 0
                and self._hybrid_cache_mgr is not None
                and self._hybrid_cache_mgr.max_checkpoints > 0
                and interrupted_verification_reconciled
            ):
                current_history = self._input_ids[:self.n_tokens].tolist()

                self._hybrid_cache_mgr.save_checkpoint(
                    current_pos=self.n_tokens,
                    tokens=current_history,
                    seq_id=0
                )

    def create_embedding(
        self,
        input: Union[str, List[str]],
        model: Optional[str] = None,
        normalize: Union[bool, int] = False,
        truncate: bool = True,
    ) -> CreateEmbeddingResponse:
        """Create an OpenAI-compatible embedding response.

        Args:
            input: A string or list of strings to embed.
            model: Model name reported in the response.
            normalize: ``False`` disables normalization, ``True`` uses L2
                normalization, and integer values select a llama.cpp
                normalization mode.
            truncate: Truncate inputs to the available context/batch capacity.

        Returns:
            An OpenAI-compatible embedding response.
        """
        model_name: str = model if model is not None else self.model_path

        input = input if isinstance(input, list) else [input]

        # get numeric embeddings
        embeds: Union[List[List[float]], List[List[List[float]]]]
        total_tokens: int
        embeds, total_tokens = self.embed(  # type: ignore
            input,
            normalize=normalize,
            truncate=truncate,
            return_count=True,
        )

        # convert to CreateEmbeddingResponse
        data: List[Embedding] = [
            {
                "object": "embedding",
                "embedding": emb,
                "index": idx,
            }
            for idx, emb in enumerate(embeds)
        ]

        return {
            "object": "list",
            "data": data,
            "model": model_name,
            "usage": {
                "prompt_tokens": total_tokens,
                "total_tokens": total_tokens,
            },
        }

    def embed(
        self,
        input: Union[str, List[str], List[List[int]]],
        normalize: Union[bool, int] = False,
        truncate: bool = True,
        separator: Optional[str] = None,
        return_count: bool = False,
    ):
        """Embed strings or pre-tokenized inputs.

        Args:
            input: A string, a list of strings, or a list of token-id lists.
            normalize: ``False``/``-1`` disables normalization, ``True`` uses
                L2 normalization. Integer modes follow llama.cpp's embedding
                example: 0=max-absolute (scaled to 32760), 1=L1, 2=L2, and
                values greater than 2 use the corresponding p-norm.
            truncate: Truncate inputs that exceed the context/batch capacity.
            separator: Split a single string into multiple inputs.
            return_count: Return ``(embeddings, token_count)``.

        Returns:
            Sequence embeddings, token-level embeddings for pooling type NONE,
            or scalar/vector scores for pooling type RANK. RANK scores are not
            normalized. Token counts reflect inputs after truncation.
        """
        if self.context_params.embeddings is False:
            raise RuntimeError(
                "Llama model must be created with embeddings=True to call this method"
            )

        ctx = self._ctx.ctx
        n_batch = self.n_batch
        # Bidirectional attention must see the entire input in one microbatch.
        if not self._ctx.get_causal_attn():
            n_batch = min(n_batch, self._ctx.n_ubatch())
        n_ctx = self._n_ctx
        n_seq_max = self.context_params.n_seq_max

        pooling_type = self.pooling_type()
        is_rank = pooling_type == llama_cpp_lib.LLAMA_POOLING_TYPE_RANK
        is_none = pooling_type == llama_cpp_lib.LLAMA_POOLING_TYPE_NONE

        # Ranking heads can have a different output width from the hidden size.
        out_dim = (
            llama_cpp_lib.llama_model_n_cls_out(self._model.model)
            if is_rank
            else self.n_embd_out()
        )

        # Preserve the historical bool API while accepting llama.cpp's integer
        # normalization modes used by LlamaEmbedding. Check bool first because
        # bool subclasses int: False means -1 (NONE), True means 2 (EUCLIDEAN).
        if isinstance(normalize, bool):
            normalize_mode = 2 if normalize else -1
        elif isinstance(normalize, int):
            normalize_mode = normalize
        else:
            raise TypeError("normalize must be a bool or int")

        def normalize_vector(vector: Sequence[float]) -> List[float]:
            return list(vector) if is_rank else normalize_embedding(vector, normalize_mode)

        if isinstance(input, str):
            inputs: List[Union[str, List[int]]] = (
                input.split(separator) if separator is not None else [input]
            )
            is_single = separator is None
        else:
            inputs = input
            is_single = False

        # Embedding batches reuse sequence IDs and positions from zero, so old
        # generation memory, output mappings and draft state cannot be retained.
        self.reset()
        try:
            self._batch.reset()

            perf_enabled = not self.context_params.no_perf
            if perf_enabled:
                self._ctx.reset_timings()

            data: List[Any] = []
            seq_sizes: List[int] = []
            total_tokens = 0

            def decode_batch() -> None:
                nonlocal seq_sizes
                if not seq_sizes:
                    return

                if self._ctx.decode(self._batch) != 0:
                    raise RuntimeError("Embedding decode failed: no KV slot available")

                if is_none:
                    # Every token requests an output; rows follow batch order.
                    # Split the flat output stream back into its input sequences.
                    token_index = 0
                    for size in seq_sizes:
                        token_embeddings: List[List[float]] = []
                        for _ in range(size):
                            ptr = llama_cpp_lib.llama_get_embeddings_ith(
                                ctx, token_index
                            )
                            token_embeddings.append(
                                [0.0] * out_dim
                                if ptr is None
                                else normalize_vector(ptr[:out_dim])
                            )
                            token_index += 1
                        data.append(token_embeddings)
                else:
                    # The backend pools each sequence according to its pooling
                    # type. Read by sequence ID, not by the last token's row.
                    for seq_id in range(len(seq_sizes)):
                        ptr = llama_cpp_lib.llama_get_embeddings_seq(ctx, seq_id)
                        if ptr is None:
                            embedding = [0.0] * out_dim
                        else:
                            embedding = list(ptr[:out_dim])

                        if is_rank:
                            data.append(
                                embedding[0] if len(embedding) == 1 else embedding
                            )
                        else:
                            data.append(normalize_vector(embedding))

                # Output pointers are borrowed: copy all vectors before clearing
                # memory and reusing the sequence IDs in the next batch.
                self._batch.reset()
                self._ctx.memory_clear(True)
                seq_sizes = []

            for item in inputs:
                if isinstance(item, str):
                    tokens = self.tokenize(item.encode("utf-8"))
                elif isinstance(item, list) and (
                    not item or isinstance(item[0], int)
                ):
                    tokens = item
                else:
                    raise ValueError("Input item must be str or List[int]")

                max_tokens = min(n_ctx, n_batch)
                if truncate and len(tokens) > max_tokens:
                    tokens = tokens[:max_tokens]

                n_tokens = len(tokens)
                total_tokens += n_tokens

                if n_tokens > n_batch:
                    raise ValueError(
                        f"Requested tokens ({n_tokens}) exceed batch size of {n_batch}"
                    )

                if n_tokens == 0:
                    # Keep result ordering stable when an empty pre-tokenized input
                    # follows sequences that are still waiting to be decoded.
                    decode_batch()
                    data.append(0.0 if is_rank else [])
                    continue

                # Pack whole inputs subject to both token and sequence capacity;
                # splitting an input across clears would break sequence pooling.
                if (
                    self._batch.n_tokens() + n_tokens > n_batch
                    or len(seq_sizes) >= n_seq_max
                ):
                    decode_batch()

                seq_id = len(seq_sizes)
                # The batch logits mask also selects embedding outputs: NONE
                # needs every token, while pooled output needs one per sequence.
                logits_array = (
                    [True] * n_tokens
                    if is_none
                    else [False] * (n_tokens - 1) + [True]
                )
                self._batch.add_sequence(
                    token_array=tokens,
                    pos_array=list(range(n_tokens)),
                    seq_ids=[seq_id],
                    logits_array=logits_array,
                )
                seq_sizes.append(n_tokens)

            decode_batch()

            if self.verbose and perf_enabled:
                self._ctx.print_timings()

            output = data[0] if is_single else data

            if return_count:
                return output, total_tokens
            return output
        finally:
            # A failed decode can leave earlier ubatches committed. Reset both
            # native and Python state even when no embedding result is returned.
            self.reset()
            self._batch.reset()

    def _create_completion(
        self,
        prompt: Union[str, List[int]],
        suffix: Optional[str] = None,
        max_tokens: Optional[int] = 128,
        temperature: float = 0.8,
        top_p: float = 0.95,
        min_p: float = 0.05,
        typical_p: float = 1.0,
        logprobs: Optional[int] = None,
        echo: bool = False,
        stop: Optional[Union[str, List[str]]] = [],
        frequency_penalty: float = 0.0,
        present_penalty: float = 0.0,
        repeat_penalty: float = 1.0,
        penalty_last_n: int = 64,
        top_k: int = 40,
        top_n_sigma: float = -1.00,
        dynatemp_range: float = 0.0,
        dynatemp_exponent: float = 1.0,
        min_keep: int = 0,
        stream: bool = False,
        mirostat_mode: int = 0,
        mirostat_tau: float = 5.0,
        mirostat_eta: float = 0.1,
        xtc_threshold: float = 0.1,
        xtc_probability: float = 0.0,
        dry_multiplier: float = 0.0,
        dry_base: float = 1.75,
        dry_allowed_length: int = 2,
        dry_penalty_last_n:int = 64,
        dry_seq_breakers: list[str] = ["\n", ":", "\"", "*"],
        adaptive_target : float = -1.0,
        adaptive_decay : float = 0.9,
        use_infill: bool = False,
        model: Optional[str] = None,
        stopping_criteria: Optional[StoppingCriteriaList] = None,
        logit_bias: Optional[Dict[int, float]] = None,
        logits_processor: Optional[LogitsProcessorList] = None,
        grammar: Optional[LlamaGrammar] = None,
        grammar_lazy: bool = False,
        seed: Optional[int] = None,
        active_loras: Optional[List[Dict[str, Union[str, float]]]] = None,
        control_vector: Optional[Dict[str, Any]] = None,
        # Reasoning Budget Params
        reasoning_budget: int = -1,
        reasoning_start: str = "<think>",
        reasoning_end: str = "</think>",
        reasoning_budget_message: Optional[str] = None,
        reasoning_start_in_prompt: bool = False,
        reasoning_start_max_tokens: Optional[int] = 32,
        ignore_eos: bool = False,
    ) -> Union[
        Iterator[CreateCompletionResponse], Iterator[CreateCompletionStreamResponse]
    ]:
        assert suffix is None or suffix.__class__ is str
        # Each time a new request is initiated, the previous abort state must be cleared.
        native_abort_flag = getattr(self, "_native_abort_flag", None)
        if native_abort_flag is not None:
            native_abort_flag.value = False
        self._abort_event.clear()

        completion_id: str = f"cmpl-{str(uuid.uuid4())}"
        created: int = int(time.time())
        bos_token_id: int = self._model.token_bos()
        eos_token_id: int = self._model.token_eos()
        sep_token_id: int = self._model.token_sep()
        prefix_token_id: int = self._model.token_fim_pre()
        middle_token_id: int = self._model.token_fim_mid()
        suffix_token_id: int = self._model.token_fim_suf()
        add_space_prefix: bool = (
            self.metadata.get("tokenizer.ggml.add_space_prefix", "true") == "true"
        )
        bos_tokens: List[int] = [bos_token_id]
        eos_tokens: List[int] = [
            sep_token_id if self._model.get_add_sep() else eos_token_id
        ]

        if (
            (isinstance(prompt, list) and suffix is None)
            or not self._model.get_add_bos()
            or bos_tokens[:1] == [-1]
        ):
            bos_tokens = []

        if (isinstance(prompt, list) and suffix is None) or (
            not self._model.get_add_eos() and not self._model.get_add_sep()
        ):
            eos_tokens = []

        suffix_space_prefix: int = 0
        # Tokenizer hack to remove leading space
        if add_space_prefix and suffix_token_id >= 0 and suffix:
            suffix = "☺" + suffix
            suffix_space_prefix = 2

        # If prompt is empty, initialize completion with BOS token to avoid
        # detokenization including a space at the beginning of the completion
        completion_tokens: List[int] = [] if len(prompt) > 0 else [bos_token_id]
        # Add blank space to start of prompt to match OG llama tokenizer
        prefix_tokens: List[int] = (
            [prefix_token_id] if prefix_token_id >= 0 and suffix is not None else []
        ) + (
            (
                self.tokenize(
                    prompt.encode("utf-8"),
                    add_bos=False,
                    special=(prefix_token_id < 0 or suffix is None),
                )
                if prompt != ""
                else []
            )
            if isinstance(prompt, str)
            else prompt
        )
        suffix_tokens: List[int] = (
            (
                [suffix_token_id]
                + (
                    self.tokenize(suffix.encode("utf-8"), add_bos=False, special=False)[
                        suffix_space_prefix:
                    ]
                    if suffix
                    else []
                )
            )
            if suffix_token_id >= 0 and suffix is not None
            else []
        )
        middle_tokens: List[int] = (
            [middle_token_id] if middle_token_id >= 0 and suffix is not None else []
        )
        prompt_tokens: List[int] = (
            bos_tokens
            + (
                (suffix_tokens + prefix_tokens + middle_tokens)
                if self.spm_infill
                else (prefix_tokens + suffix_tokens + middle_tokens)
            )
            + eos_tokens
        )
        text: bytes = b""
        returned_tokens: int = 0
        stop = (
            stop if isinstance(stop, list) else [stop] if isinstance(stop, str) else []
        )
        model_name: str = model if model is not None else self.model_path

        if prompt_tokens[:2] == [self.token_bos()] * 2:
            warnings.warn(
                f'Detected duplicate leading "{self._model.token_get_text(self.token_bos())}" in prompt, this will likely reduce response quality, consider removing it...',
                RuntimeWarning,
            )

        if len(prompt_tokens) >= self._n_ctx:
            raise ValueError(
                f"Requested tokens ({len(prompt_tokens)}) exceed context window of {llama_cpp_lib.llama_n_ctx(self.ctx)}"
            )

        if max_tokens is None or max_tokens <= 0:
            # Unlimited, depending on n_ctx.
            max_tokens = self._n_ctx - len(prompt_tokens)

        # Truncate max_tokens if requested tokens would exceed the context window
        max_tokens = (
            max_tokens
            if max_tokens + len(prompt_tokens) < self._n_ctx
            else (self._n_ctx - len(prompt_tokens))
        )

        if stop != []:
            stop_sequences = [s.encode("utf-8") for s in stop]
        else:
            stop_sequences = []

        if logprobs is not None and self._logits_all is False:
            raise ValueError(
                "logprobs is not supported for models created with logits_all=False"
            )

        if self.cache and getattr(self, "speculative", None) is None:
            live_prefix_length = Llama.longest_token_prefix(
                self._input_ids, prompt_tokens, self.verbose
            )
            try:
                get_matching_state = getattr(self.cache, "get_matching_state", None)
                cached_state = (
                    get_matching_state(prompt_tokens, min_prefix_length=live_prefix_length)
                    if get_matching_state is not None else self.cache[prompt_tokens]
                )
                cached_prefix_length = Llama.longest_token_prefix(
                    cached_state.input_ids[:cached_state.n_tokens], prompt_tokens, self.verbose
                )
                # Loading a full state invalidates partial checkpoints. A hybrid
                # snapshot must therefore be appendable or own its final logits.
                can_reuse_cached_state = not self.is_hybrid or (
                    cached_prefix_length == cached_state.n_tokens
                    and (cached_prefix_length < len(prompt_tokens) or getattr(cached_state, "last_logits", None) is not None)
                )
                if cached_prefix_length > live_prefix_length and can_reuse_cached_state:
                    self.load_state(cached_state)
                    if self.verbose:
                        print("Llama._create_completion: cache hit", file=sys.stderr)
            except KeyError:
                if self.verbose:
                    print("Llama._create_completion: cache miss", file=sys.stderr)

        if stopping_criteria is None:
            stopping_criteria = StoppingCriteriaList([AbortCriteria(self._abort_event)])
        else:
            stopping_criteria.append(AbortCriteria(self._abort_event))

        finish_reason = "length"
        multibyte_fix = 0
        for token in self.generate(
            prompt_tokens,
            top_k=top_k,
            top_n_sigma=top_n_sigma,
            top_p=top_p,
            min_p=min_p,
            typical_p=typical_p,
            temp=temperature,
            dynatemp_range=dynatemp_range,
            dynatemp_exponent=dynatemp_exponent,
            min_keep=min_keep,
            mirostat_mode=mirostat_mode,
            mirostat_tau=mirostat_tau,
            mirostat_eta=mirostat_eta,
            xtc_threshold=xtc_threshold,
            xtc_probability=xtc_probability,
            dry_multiplier=dry_multiplier,
            dry_base=dry_base,
            dry_allowed_length=dry_allowed_length,
            dry_penalty_last_n=dry_penalty_last_n,
            dry_seq_breakers=dry_seq_breakers,
            frequency_penalty=frequency_penalty,
            present_penalty=present_penalty,
            repeat_penalty=repeat_penalty,
            penalty_last_n=penalty_last_n,
            stopping_criteria=stopping_criteria,
            adaptive_target=adaptive_target,
            adaptive_decay=adaptive_decay,
            use_infill=use_infill,
            ignore_eos=ignore_eos,
            logit_bias=logit_bias,
            logits_processor=logits_processor,
            grammar=grammar,
            grammar_lazy=grammar_lazy,
            seed=seed if seed is not None else self._seed,
            active_loras=active_loras,
            control_vector=control_vector,
            reasoning_budget=reasoning_budget,
            reasoning_start=reasoning_start,
            reasoning_end=reasoning_end,
            reasoning_budget_message=reasoning_budget_message,
            reasoning_start_in_prompt=reasoning_start_in_prompt,
            reasoning_start_max_tokens=reasoning_start_max_tokens,
        ):
            if (
                not ignore_eos
                and llama_cpp_lib.llama_token_is_eog(self._model.vocab, token)
            ):
                text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens)
                finish_reason = "stop"
                break

            if self._abort_event.is_set():
                text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens)
                finish_reason = "abort"
                break

            completion_tokens.append(token)

            all_text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens)

            # Contains multi-byte UTF8
            for k, char in enumerate(all_text[-3:]):
                k = 3 - k
                for num, pattern in [(2, 192), (3, 224), (4, 240)]:
                    # Bitwise AND check
                    if num > k and pattern & char == pattern:
                        multibyte_fix = num - k

            # Stop incomplete bytes from passing
            if multibyte_fix > 0:
                multibyte_fix -= 1
                continue

            any_stop = [s for s in stop_sequences if s in all_text]
            if len(any_stop) > 0:
                first_stop = any_stop[0]
                text = all_text[: all_text.index(first_stop)]
                finish_reason = "stop"
                break

            if stream:
                remaining_tokens = completion_tokens[returned_tokens:]
                remaining_text = self.detokenize(
                    remaining_tokens,
                    prev_tokens=prompt_tokens + completion_tokens[:returned_tokens],
                )
                remaining_length = len(remaining_text)

                # We want to avoid yielding any characters from
                # the generated text if they are part of a stop
                # sequence.
                first_stop_position = 0
                for s in stop_sequences:
                    for i in range(min(len(s), remaining_length), 0, -1):
                        if remaining_text.endswith(s[:i]):
                            if i > first_stop_position:
                                first_stop_position = i
                            break

                token_end_position = 0

                if logprobs is not None:
                    # not sure how to handle this branch when dealing
                    # with CJK output, so keep it unchanged
                    for token in remaining_tokens:
                        if token == bos_token_id:
                            continue
                        token_end_position += len(
                            self.detokenize(
                                [token],
                                prev_tokens=prompt_tokens
                                + completion_tokens[:returned_tokens],
                            )
                        )
                        # Check if stop sequence is in the token
                        if token_end_position > (
                            remaining_length - first_stop_position
                        ):
                            break
                        token_str = self.detokenize(
                            [token],
                            prev_tokens=prompt_tokens
                            + completion_tokens[:returned_tokens],
                        ).decode("utf-8", errors="ignore")
                        text_offset = len(prompt) + len(
                            self.detokenize(
                                completion_tokens[:returned_tokens],
                                prev_tokens=prompt_tokens
                                + completion_tokens[:returned_tokens],
                            ).decode("utf-8", errors="ignore")
                        )
                        token_offset = len(prompt_tokens) + returned_tokens
                        if self._logits_all:
                            logits = self._scores[token_offset - 1, :]
                        else:
                            logits = self._scores[0, :]
                        current_logprobs = Llama.logits_to_logprobs(logits).tolist()
                        sorted_logprobs = list(
                            sorted(
                                zip(current_logprobs, range(len(current_logprobs))),
                                reverse=True,
                            )
                        )
                        top_logprob = {
                            self.detokenize([i]).decode(
                                "utf-8", errors="ignore"
                            ): logprob
                            for logprob, i in sorted_logprobs[:logprobs]
                        }
                        top_logprob.update({token_str: current_logprobs[int(token)]})
                        logprobs_or_none = {
                            "tokens": [
                                self.detokenize(
                                    [token],
                                    prev_tokens=prompt_tokens
                                    + completion_tokens[:returned_tokens],
                                ).decode("utf-8", errors="ignore")
                            ],
                            "text_offset": [text_offset],
                            "token_logprobs": [current_logprobs[int(token)]],
                            "top_logprobs": [top_logprob],
                        }
                        returned_tokens += 1
                        yield {
                            "id": completion_id,
                            "object": "text_completion",
                            "created": created,
                            "model": model_name,
                            "choices": [
                                {
                                    "text": self.detokenize(
                                        [token],
                                        prev_tokens=prompt_tokens
                                        + completion_tokens[:returned_tokens],
                                    ).decode("utf-8", errors="ignore"),
                                    "index": 0,
                                    "logprobs": logprobs_or_none,
                                    "finish_reason": None,
                                }
                            ],
                        }
                else:
                    while len(remaining_tokens) > 0:
                        decode_success = False
                        for i in range(1, len(remaining_tokens) + 1):
                            try:
                                bs = self.detokenize(
                                    remaining_tokens[:i],
                                    prev_tokens=prompt_tokens
                                    + completion_tokens[:returned_tokens],
                                )
                                ts = bs.decode("utf-8")
                                decode_success = True
                                break
                            except UnicodeError:
                                pass
                        else:
                            break
                        if not decode_success:
                            # all remaining tokens cannot be decoded to a UTF-8 character
                            break
                        token_end_position += len(bs)
                        if token_end_position > (
                            remaining_length - first_stop_position
                        ):
                            break
                        remaining_tokens = remaining_tokens[i:]
                        returned_tokens += i

                        yield {
                            "id": completion_id,
                            "object": "text_completion",
                            "created": created,
                            "model": model_name,
                            "choices": [
                                {
                                    "text": ts,
                                    "index": 0,
                                    "logprobs": None,
                                    "finish_reason": None,
                                }
                            ],
                        }

            if len(completion_tokens) >= max_tokens:
                text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens)
                finish_reason = "length"
                break

        if stopping_criteria is not None and stopping_criteria(
            self._input_ids, self._scores[-1, :]
        ):
            text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens)
            finish_reason = "stop"

        # If the abort is triggered externally, force the `finish_reason` to be changed to "abort".
        if self._abort_event.is_set():
            text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens)
            finish_reason = "abort"

        if stream:
            remaining_tokens = completion_tokens[returned_tokens:]
            remaining_text = self.detokenize(
                remaining_tokens,
                prev_tokens=prompt_tokens + completion_tokens[:returned_tokens],
            )
            any_stop = [s for s in stop_sequences if s in remaining_text]
            if len(any_stop) > 0:
                end = min(remaining_text.index(stop) for stop in any_stop)
            else:
                end = len(remaining_text)

            token_end_position = 0
            for token in remaining_tokens:
                token_end_position += len(
                    self.detokenize(
                        [token],
                        prev_tokens=prompt_tokens + completion_tokens[:returned_tokens],
                    )
                )

                logprobs_or_none: Optional[CompletionLogprobs] = None
                if logprobs is not None:
                    if token == bos_token_id:
                        continue
                    token_str = self.detokenize([token]).decode(
                        "utf-8", errors="ignore"
                    )
                    text_offset = len(prompt) + len(
                        self.detokenize(
                            completion_tokens[:returned_tokens],
                            prev_tokens=prompt_tokens
                            + completion_tokens[:returned_tokens],
                        )
                    )
                    token_offset = len(prompt_tokens) + returned_tokens - 1
                    if self._logits_all:
                        logits = self._scores[token_offset, :]
                    else:
                        logits = self._scores[0, :]
                    current_logprobs = Llama.logits_to_logprobs(logits).tolist()
                    sorted_logprobs = list(
                        sorted(
                            zip(current_logprobs, range(len(current_logprobs))),
                            reverse=True,
                        )
                    )
                    top_logprob = {
                        self.detokenize([i]).decode("utf-8", errors="ignore"): logprob
                        for logprob, i in sorted_logprobs[:logprobs]
                    }
                    top_logprob.update({token_str: current_logprobs[int(token)]})
                    logprobs_or_none = {
                        "tokens": [
                            self.detokenize([token]).decode("utf-8", errors="ignore")
                        ],
                        "text_offset": [text_offset],
                        "token_logprobs": [current_logprobs[int(token)]],
                        "top_logprobs": [top_logprob],
                    }

                if token_end_position >= end:
                    last_text = self.detokenize([token])
                    if token_end_position == end - 1:
                        break
                    returned_tokens += 1
                    yield {
                        "id": completion_id,
                        "object": "text_completion",
                        "created": created,
                        "model": model_name,
                        "choices": [
                            {
                                "text": last_text[
                                    : len(last_text) - (token_end_position - end)
                                ].decode("utf-8", errors="ignore"),
                                "index": 0,
                                "logprobs": logprobs_or_none,
                                "finish_reason": None,
                            }
                        ],
                    }
                    break
                returned_tokens += 1
                yield {
                    "id": completion_id,
                    "object": "text_completion",
                    "created": created,
                    "model": model_name,
                    "choices": [
                        {
                            "text": self.detokenize([token]).decode(
                                "utf-8", errors="ignore"
                            ),
                            "index": 0,
                            "logprobs": logprobs_or_none,
                            "finish_reason": None,
                        }
                    ],
                }
            yield {
                "id": completion_id,
                "object": "text_completion",
                "created": created,
                "model": model_name,
                "choices": [
                    {
                        "text": "",
                        "index": 0,
                        "logprobs": None,
                        "finish_reason": finish_reason,
                    }
                ],
            }
            if self.cache and finish_reason != "abort" and getattr(self, "speculative", None) is None:
                if self.verbose:
                    print("Llama._create_completion: cache save", file=sys.stderr)
                state = self.save_state()
                self.cache[state.input_ids[:state.n_tokens].tolist()] = state
                if self.verbose:
                    print("Llama._create_completion: cache saved", file=sys.stderr)
            return

        if self.cache and finish_reason != "abort" and getattr(self, "speculative", None) is None:
            if self.verbose:
                print("Llama._create_completion: cache save", file=sys.stderr)
            state = self.save_state()
            self.cache[state.input_ids[:state.n_tokens].tolist()] = state

        text_str = text.decode("utf-8", errors="ignore")

        if echo:
            text_str = prompt + text_str

        if suffix_token_id < 0 and suffix is not None:
            text_str = text_str + suffix

        logprobs_or_none: Optional[CompletionLogprobs] = None
        if logprobs is not None:
            text_offset = 0 if echo else len(prompt)
            token_offset = 0 if echo else len(prompt_tokens[1:])
            text_offsets: List[int] = []
            token_logprobs: List[Optional[float]] = []
            tokens: List[str] = []
            top_logprobs: List[Optional[Dict[str, float]]] = []

            if echo:
                # Remove leading BOS token if exists
                all_tokens = (
                    prompt_tokens[1 if prompt_tokens[0] == self.token_bos() else 0 :]
                    + completion_tokens
                )
            else:
                all_tokens = completion_tokens

            all_token_strs = [
                self.detokenize([token], prev_tokens=all_tokens[:i]).decode(
                    "utf-8", errors="ignore"
                )
                for i, token in enumerate(all_tokens)
            ]
            all_logprobs = Llama.logits_to_logprobs(self._scores)[token_offset:]
            # TODO: may be able to change this loop to use np.take_along_dim
            for idx, (token, token_str, logprobs_token) in enumerate(
                zip(all_tokens, all_token_strs, all_logprobs)
            ):
                if token == bos_token_id:
                    continue
                text_offsets.append(
                    text_offset
                    + len(
                        self.detokenize(all_tokens[:idx]).decode(
                            "utf-8", errors="ignore"
                        )
                    )
                )
                tokens.append(token_str)
                sorted_logprobs = list(
                    sorted(
                        zip(logprobs_token, range(len(logprobs_token))), reverse=True
                    )
                )
                token_logprobs.append(logprobs_token[int(token)])
                top_logprob: Optional[Dict[str, float]] = {
                    self.detokenize([i], prev_tokens=all_tokens[:idx]).decode(
                        "utf-8", errors="ignore"
                    ): logprob
                    for logprob, i in sorted_logprobs[:logprobs]
                }
                top_logprob.update({token_str: logprobs_token[int(token)]})
                top_logprobs.append(top_logprob)
            # Weird idosincracy of the OpenAI API where
            # token_logprobs and top_logprobs are null for
            # the first token.
            if echo and len(all_tokens) > 0:
                token_logprobs[0] = None
                top_logprobs[0] = None
            logprobs_or_none = {
                "tokens": tokens,
                "text_offset": text_offsets,
                "token_logprobs": token_logprobs,
                "top_logprobs": top_logprobs,
            }

        yield {
            "id": completion_id,
            "object": "text_completion",
            "created": created,
            "model": model_name,
            "choices": [
                {
                    "text": text_str,
                    "index": 0,
                    "logprobs": logprobs_or_none,
                    "finish_reason": finish_reason,
                }
            ],
            "usage": {
                "prompt_tokens": len(prompt_tokens),
                "completion_tokens": len(completion_tokens),
                "total_tokens": len(prompt_tokens) + len(completion_tokens),
            },
        }

    def create_completion(
        self,
        prompt: Union[str, List[int]],
        suffix: Optional[str] = None,
        max_tokens: Optional[int] = 16,
        temperature: float = 0.8,
        top_p: float = 0.95,
        min_p: float = 0.05,
        typical_p: float = 1.0,
        logprobs: Optional[int] = None,
        echo: bool = False,
        stop: Optional[Union[str, List[str]]] = [],
        frequency_penalty: float = 0.0,
        present_penalty: float = 0.0,
        presence_penalty: Optional[float] = None,
        repeat_penalty: float = 1.0,
        penalty_last_n: int = 64,
        top_k: int = 40,
        top_n_sigma: float = -1.00,
        stream: bool = False,
        seed: Optional[int] = None,
        mirostat_mode: int = 0,
        mirostat_tau: float = 5.0,
        mirostat_eta: float = 0.1,
        dynatemp_range: float = 0.0,
        dynatemp_exponent: float = 1.0,
        min_keep: int = 0,
        xtc_threshold: float = 0.1,
        xtc_probability: float = 0.0,
        dry_multiplier: float = 0.0,
        dry_base: float = 1.75,
        dry_allowed_length: int = 2,
        dry_penalty_last_n:int = 64,
        dry_seq_breakers: list[str] = ["\n", ":", "\"", "*"],
        adaptive_target : float = -1.0,
        adaptive_decay : float = 0.9,
        use_infill: bool = False,
        model: Optional[str] = None,
        stopping_criteria: Optional[StoppingCriteriaList] = None,
        logit_bias: Optional[Dict[int, float]] = None,
        logits_processor: Optional[LogitsProcessorList] = None,
        grammar: Optional[LlamaGrammar] = None,
        grammar_lazy: bool = False,
        active_loras: Optional[List[Dict[str, Union[str, float]]]] = None,
        control_vector: Optional[Dict[str, Any]] = None,
        # Reasoning Budget Params
        reasoning_budget: int = -1,
        reasoning_start: str = "<think>",
        reasoning_end: str = "</think>",
        reasoning_budget_message: Optional[str] = None,
        reasoning_start_in_prompt: bool = False,
        reasoning_start_max_tokens: Optional[int] = 32,
        ignore_eos: bool = False,
    ) -> Union[CreateCompletionResponse, Iterator[CreateCompletionStreamResponse]]:
        """Generate text from a prompt.

        Args:
prompt: The prompt to generate text from.
            suffix: A suffix to append to the generated text. If None, no suffix is appended.
            max_tokens: The maximum number of tokens to generate. If max_tokens <= 0 or None, the maximum number of tokens to generate is unlimited and depends on n_ctx.
            temperature: The temperature to use for sampling.
            top_p: The top-p value to use for nucleus sampling. Nucleus sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751
            min_p: The min-p value to use for minimum p sampling. Minimum P sampling as described in https://github.com/ggml-org/llama.cpp/pull/3841
            typical_p: The typical-p value to use for sampling. Locally Typical Sampling implementation described in the paper https://arxiv.org/abs/2202.00666.
            logprobs: The number of logprobs to return. If None, no logprobs are returned.
            echo: Whether to echo the prompt.
            stop: A list of strings to stop generation when encountered.
            frequency_penalty: The penalty to apply to tokens based on their frequency in the prompt.
            present_penalty: The penalty to controls whether to apply a penalty to tokens that are already present in the current context, helping to reduce repetition and encourage more diverse generation.
            presence_penalty: Compatibility alias for `present_penalty`. It is used only when `present_penalty` remains at its default value.
            repeat_penalty: The penalty to apply to repeated tokens.
            penalty_last_n: last n tokens to penalize (0 = disable penalty, -1 = context size).
            top_k: The top-k value to use for sampling. Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751
            top_n_sigma: Limit the next token selection to a subset of tokens with pre-softmax logits that are within n * σ less than the max logit (default: -1.00, -1.00 = disabled).
            stream: Whether to stream the results.
            seed: The seed to use for sampling.
            mirostat_mode: The mirostat sampling mode.
            mirostat_tau: The target cross-entropy (or surprise) value you want to achieve for the generated text. A higher value corresponds to more surprising or less predictable text, while a lower value corresponds to less surprising or more predictable text.
            mirostat_eta: The learning rate used to update `mu` based on the error between the target and observed surprisal of the sampled word. A larger learning rate will cause `mu` to be updated more quickly, while a smaller learning rate will result in slower updates.
            dynatemp_range: Range of dynamic temperature.
            dynatemp_exponent: Exponent of dynamic temperature.
            min_keep: Minimum tokens to keep for sampling.
            xtc-probability: Sets the chance for token removal (checked once on sampler start) (default: 0.0). XTC sampler as described in https://github.com/oobabooga/text-generation-webui/pull/6335
            xtc-threshold: Sets a minimum probability threshold for tokens to be removed (default: 0.1). XTC sampler as described in https://github.com/oobabooga/text-generation-webui/pull/6335
            dry_multiplier: Set the DRY (Don't Repeat Yourself) repetition penalty multiplier. Default: `0.0`, which is disabled.
            dry_base`: Set the DRY repetition penalty base value. Default: `1.75`
            dry_allowed_length: Tokens that extend repetition beyond this receive exponentially increasing penalty: multiplier * base ^ (length of repeating sequence before token - allowed length). Default: `2`
            dry_penalty_last_n: How many tokens to scan for repetitions. Default: `64`; `0` disables scanning and `-1` uses the context size.
            dry_seq_breakers: Specify an array of sequence breakers for DRY sampling. Only a JSON array of strings is accepted. Default: `['\n', ':', '"', '*']`
            adaptive-target: Adaptive-p: select tokens near this probability (valid range 0.0 to 1.0; negative = disabled) (default: %.2f) [(more info)](https://github.com/ggml-org/llama.cpp/pull/17927)
            adaptive-decay: Adaptive-p: decay rate for target adaptation over time. lower values are more reactive, higher values are more stable. (valid range 0.0 to 0.99) (default: %.2f)
            use_infill: Determines whether to activate the specialized fill-in-the-middle sampler that consolidates probabilities of tokens sharing common prefixes to ensure the generated text coherently bridges the gap between the prefix and suffix.
            ignore_eos: If True, suppress end-of-generation tokens and continue until another stopping condition is reached.
            model: The name to use for the model in the completion object.
            stopping_criteria: A list of stopping criteria to use.
            logit_bias: A logit bias to use.
            logits_processor: A list of logits processors to use.
            grammar: A grammar to use for constrained sampling.
            grammar_lazy: If True, enables lazy evaluation.
            reasoning_budget: Token budget for the first visible reasoning block.
                -1 disables the sampler, 0 forces an immediate end after reasoning starts,
                and N > 0 allows at most N generated tokens inside the block.
            reasoning_start: Token/text sequence that marks the beginning of the first reasoning block.
            reasoning_end: Token/text sequence that naturally and forcibly ends the reasoning block.
            reasoning_budget_message: Optional message inserted before reasoning_end when the budget is exhausted.
            reasoning_start_in_prompt: Set True when the prompt/template already inserted reasoning_start.
            reasoning_start_max_tokens: Safety window before disabling the sampler for non-reasoning outputs.
            active_loras: A list of dictionaries specifying the LoRA adapters to dynamically apply during generation.
                Each dictionary must contain a "name" key (matching a LoRA previously loaded into VRAM via `load_lora()`)
                and an optional "scale" key (float, defaults to 1.0).
                Example: `[{"name": "role_A", "scale": 0.85}, {"name": "role_B", "scale": 0.5}]`.
            control_vector: A dictionary containing Control Vector (CVec) data for representation engineering.
                Must contain a "data" key with a flattened 1D list of floats.
                Optionally accepts "layer_start" (int, defaults to 1) and "layer_end" (int, defaults to the model's total layer count).
                Note: The length of the "data" list MUST be at least `n_embd * layer_end`, with zero-padding for any skipped early layers.

        Raises:
            ValueError: If the requested tokens exceed the context window.
            RuntimeError: If the prompt fails to tokenize or the model fails to evaluate the prompt.

        Returns:
            Response object containing the generated text.
        """
        if presence_penalty is not None and present_penalty == 0.0:
            present_penalty = presence_penalty

        completion_or_chunks = self._create_completion(
            prompt=prompt,
            suffix=suffix,
            max_tokens=-1 if max_tokens is None else max_tokens,
            temperature=temperature,
            top_p=top_p,
            min_p=min_p,
            typical_p=typical_p,
            logprobs=logprobs,
            echo=echo,
            stop=stop,
            frequency_penalty=frequency_penalty,
            present_penalty=present_penalty,
            repeat_penalty=repeat_penalty,
            penalty_last_n=penalty_last_n,
            top_k=top_k,
            top_n_sigma=top_n_sigma,
            stream=stream,
            seed=seed,
            mirostat_mode=mirostat_mode,
            mirostat_tau=mirostat_tau,
            mirostat_eta=mirostat_eta,
            dynatemp_range=dynatemp_range,
            dynatemp_exponent=dynatemp_exponent,
            min_keep=min_keep,
            xtc_threshold=xtc_threshold,
            xtc_probability=xtc_probability,
            dry_multiplier=dry_multiplier,
            dry_base=dry_base,
            dry_allowed_length=dry_allowed_length,
            dry_penalty_last_n=dry_penalty_last_n,
            dry_seq_breakers=dry_seq_breakers,
            adaptive_target=adaptive_target,
            adaptive_decay=adaptive_decay,
            use_infill=use_infill,
            ignore_eos=ignore_eos,
            model=model,
            stopping_criteria=stopping_criteria,
            logit_bias=logit_bias,
            logits_processor=logits_processor,
            grammar=grammar,
            grammar_lazy=grammar_lazy,
            active_loras=active_loras,
            control_vector=control_vector,
            reasoning_budget=reasoning_budget,
            reasoning_start=reasoning_start,
            reasoning_end=reasoning_end,
            reasoning_budget_message=reasoning_budget_message,
            reasoning_start_in_prompt=reasoning_start_in_prompt,
            reasoning_start_max_tokens=reasoning_start_max_tokens,
        )

        # Keep the native counters request-scoped even when verbose logging is
        # disabled. The wrapper also guarantees that an abandoned streaming
        # iterator prints its partial timings when it is explicitly closed.
        perf_enabled = (
            hasattr(self, "_ctx")
            and not getattr(getattr(self, "context_params", None), "no_perf", False)
        )
        if perf_enabled:
            raw_chunks = completion_or_chunks

            def timed_chunks() -> Iterator[
                Union[CreateCompletionResponse, CreateCompletionStreamResponse]
            ]:
                self._ctx.reset_timings()
                try:
                    yield from raw_chunks
                finally:
                    if self.verbose:
                        self._ctx.print_timings()

            completion_or_chunks = timed_chunks()

        if stream:
            chunks: Iterator[CreateCompletionStreamResponse] = completion_or_chunks
            return chunks
        try:
            completion: Completion = next(completion_or_chunks)  # type: ignore
            return completion
        finally:
            close = getattr(completion_or_chunks, "close", None)
            if close is not None:
                close()

    def __call__(
        self,
        prompt: str,
        suffix: Optional[str] = None,
        max_tokens: Optional[int] = 128,
        temperature: float = 0.8,
        top_p: float = 0.95,
        min_p: float = 0.05,
        typical_p: float = 1.0,
        logprobs: Optional[int] = None,
        echo: bool = False,
        stop: Optional[Union[str, List[str]]] = [],
        frequency_penalty: float = 0.0,
        present_penalty: float = 0.0,
        presence_penalty: Optional[float] = None,
        repeat_penalty: float = 1.0,
        penalty_last_n: int = 64,
        top_k: int = 40,
        top_n_sigma: float = -1.00,
        stream: bool = False,
        seed: Optional[int] = None,
        mirostat_mode: int = 0,
        mirostat_tau: float = 5.0,
        mirostat_eta: float = 0.1,
        dynatemp_range: float = 0.0,
        dynatemp_exponent: float = 1.0,
        min_keep: int = 0,
        xtc_threshold: float = 0.1,
        xtc_probability: float = 0.0,
        dry_multiplier: float = 0.0,
        dry_base: float = 1.75,
        dry_allowed_length: int = 2,
        dry_penalty_last_n:int = 64,
        dry_seq_breakers: list[str] = ["\n", ":", "\"", "*"],
        adaptive_target : float = -1.0,
        adaptive_decay : float = 0.9,
        use_infill: bool = False,
        model: Optional[str] = None,
        stopping_criteria: Optional[StoppingCriteriaList] = None,
        logit_bias: Optional[Dict[int, float]] = None,
        logits_processor: Optional[LogitsProcessorList] = None,
        grammar: Optional[LlamaGrammar] = None,
        grammar_lazy: bool = False,
        active_loras: Optional[List[Dict[str, Union[str, float]]]] = None,
        control_vector: Optional[Dict[str, Any]] = None,
        # Reasoning Budget Params
        reasoning_budget: int = -1,
        reasoning_start: str = "<think>",
        reasoning_end: str = "</think>",
        reasoning_budget_message: Optional[str] = None,
        reasoning_start_in_prompt: bool = False,
        reasoning_start_max_tokens: Optional[int] = 32,
        ignore_eos: bool = False,
    ) -> Union[CreateCompletionResponse, Iterator[CreateCompletionStreamResponse]]:
        """Generate text from a prompt.

        Args:
            prompt: The prompt to generate text from.
            suffix: A suffix to append to the generated text. If None, no suffix is appended.
            max_tokens: The maximum number of tokens to generate. If max_tokens <= 0 or None, the maximum number of tokens to generate is unlimited and depends on n_ctx.
            temperature: The temperature to use for sampling.
            top_p: The top-p value to use for nucleus sampling. Nucleus sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751
            min_p: The min-p value to use for minimum p sampling. Minimum P sampling as described in https://github.com/ggml-org/llama.cpp/pull/3841
            typical_p: The typical-p value to use for sampling. Locally Typical Sampling implementation described in the paper https://arxiv.org/abs/2202.00666.
            logprobs: The number of logprobs to return. If None, no logprobs are returned.
            echo: Whether to echo the prompt.
            stop: A list of strings to stop generation when encountered.
            frequency_penalty: The penalty to apply to tokens based on their frequency in the prompt.
            present_penalty: The penalty to controls whether to apply a penalty to tokens that are already present in the current context, helping to reduce repetition and encourage more diverse generation.
            presence_penalty: Compatibility alias for `present_penalty`. It is used only when `present_penalty` remains at its default value.
            repeat_penalty: The penalty to apply to repeated tokens.
            penalty_last_n: last n tokens to penalize (0 = disable penalty, -1 = context size).
            top_k: The top-k value to use for sampling. Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751
            top_n_sigma: Limit the next token selection to a subset of tokens with pre-softmax logits that are within n * σ less than the max logit (default: -1.00, -1.00 = disabled).
            stream: Whether to stream the results.
            seed: The seed to use for sampling.
            mirostat_mode: The mirostat sampling mode.
            mirostat_tau: The target cross-entropy (or surprise) value you want to achieve for the generated text. A higher value corresponds to more surprising or less predictable text, while a lower value corresponds to less surprising or more predictable text.
            mirostat_eta: The learning rate used to update `mu` based on the error between the target and observed surprisal of the sampled word. A larger learning rate will cause `mu` to be updated more quickly, while a smaller learning rate will result in slower updates.
            dynatemp_range: Range of dynamic temperature.
            dynatemp_exponent: Exponent of dynamic temperature.
            min_keep: Minimum tokens to keep for sampling.
            xtc-probability: Sets the chance for token removal (checked once on sampler start) (default: 0.0). XTC sampler as described in https://github.com/oobabooga/text-generation-webui/pull/6335
            xtc-threshold: Sets a minimum probability threshold for tokens to be removed (default: 0.1). XTC sampler as described in https://github.com/oobabooga/text-generation-webui/pull/6335
            dry_multiplier: Set the DRY (Don't Repeat Yourself) repetition penalty multiplier. Default: `0.0`, which is disabled.
            dry_base`: Set the DRY repetition penalty base value. Default: `1.75`
            dry_allowed_length: Tokens that extend repetition beyond this receive exponentially increasing penalty: multiplier * base ^ (length of repeating sequence before token - allowed length). Default: `2`
            dry_penalty_last_n: How many tokens to scan for repetitions. Default: `64`; `0` disables scanning and `-1` uses the context size.
            dry_seq_breakers: Specify an array of sequence breakers for DRY sampling. Only a JSON array of strings is accepted. Default: `['\n', ':', '"', '*']`
            adaptive-target: Adaptive-p: select tokens near this probability (valid range 0.0 to 1.0; negative = disabled) (default: %.2f) [(more info)](https://github.com/ggml-org/llama.cpp/pull/17927)
            adaptive-decay: Adaptive-p: decay rate for target adaptation over time. lower values are more reactive, higher values are more stable. (valid range 0.0 to 0.99) (default: %.2f)
            use_infill: Determines whether to activate the specialized fill-in-the-middle sampler that consolidates probabilities of tokens sharing common prefixes to ensure the generated text coherently bridges the gap between the prefix and suffix.
            ignore_eos: If True, suppress end-of-generation tokens and continue until another stopping condition is reached.
            model: The name to use for the model in the completion object.
            stopping_criteria: A list of stopping criteria to use.
            logit_bias: A logit bias to use.
            logits_processor: A list of logits processors to use.
            grammar: A grammar to use for constrained sampling.
            grammar_lazy: If True, enables lazy evaluation.
            reasoning_budget: Token budget for the first visible reasoning block.
                -1 disables the sampler, 0 forces an immediate end after reasoning starts,
                and N > 0 allows at most N generated tokens inside the block.
            reasoning_start: Token/text sequence that marks the beginning of the first reasoning block.
            reasoning_end: Token/text sequence that naturally and forcibly ends the reasoning block.
            reasoning_budget_message: Optional message inserted before reasoning_end when the budget is exhausted.
            reasoning_start_in_prompt: Set True when the prompt/template already inserted reasoning_start.
            reasoning_start_max_tokens: Safety window before disabling the sampler for non-reasoning outputs.
            active_loras: A list of dictionaries specifying the LoRA adapters to dynamically apply during generation.
                Each dictionary must contain a "name" key (matching a LoRA previously loaded into VRAM via `load_lora()`)
                and an optional "scale" key (float, defaults to 1.0).
                Example: `[{"name": "role_A", "scale": 0.85}, {"name": "role_B", "scale": 0.5}]`.
            control_vector: A dictionary containing Control Vector (CVec) data for representation engineering.
                Must contain a "data" key with a flattened 1D list of floats.
                Optionally accepts "layer_start" (int, defaults to 1) and "layer_end" (int, defaults to the model's total layer count).
                Note: The length of the "data" list MUST be at least `n_embd * layer_end`, with zero-padding for any skipped early layers.

        Raises:
            ValueError: If the requested tokens exceed the context window.
            RuntimeError: If the prompt fails to tokenize or the model fails to evaluate the prompt.

        Returns:
            Response object containing the generated text.
        """
        return self.create_completion(
            prompt=prompt,
            suffix=suffix,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            min_p=min_p,
            typical_p=typical_p,
            logprobs=logprobs,
            echo=echo,
            stop=stop,
            frequency_penalty=frequency_penalty,
            present_penalty=present_penalty,
            repeat_penalty=repeat_penalty,
            penalty_last_n=penalty_last_n,
            top_k=top_k,
            top_n_sigma=top_n_sigma,
            stream=stream,
            seed=seed,
            mirostat_mode=mirostat_mode,
            mirostat_tau=mirostat_tau,
            mirostat_eta=mirostat_eta,
            dynatemp_range=dynatemp_range,
            dynatemp_exponent=dynatemp_exponent,
            min_keep=min_keep,
            xtc_threshold=xtc_threshold,
            xtc_probability=xtc_probability,
            dry_multiplier=dry_multiplier,
            dry_base=dry_base,
            dry_allowed_length=dry_allowed_length,
            dry_penalty_last_n=dry_penalty_last_n,
            dry_seq_breakers=dry_seq_breakers,
            adaptive_target=adaptive_target,
            adaptive_decay=adaptive_decay,
            use_infill=use_infill,
            ignore_eos=ignore_eos,
            model=model,
            stopping_criteria=stopping_criteria,
            logit_bias=logit_bias,
            logits_processor=logits_processor,
            grammar=grammar,
            grammar_lazy=grammar_lazy,
            active_loras=active_loras,
            control_vector=control_vector,
            reasoning_budget=reasoning_budget,
            reasoning_start=reasoning_start,
            reasoning_end=reasoning_end,
            reasoning_budget_message=reasoning_budget_message,
            reasoning_start_in_prompt=reasoning_start_in_prompt,
            reasoning_start_max_tokens=reasoning_start_max_tokens,
            presence_penalty=presence_penalty,
        )

    def create_chat_completion(
        self,
        messages: List[ChatCompletionRequestMessage],
        functions: Optional[List[ChatCompletionFunction]] = None,
        function_call: Optional[ChatCompletionRequestFunctionCall] = None,
        tools: Optional[List[ChatCompletionTool]] = None,
        tool_choice: Optional[ChatCompletionToolChoiceOption] = None,
        temperature: float = 0.2,
        top_p: float = 0.95,
        top_k: int = 40,
        top_n_sigma: float = -1.00,
        min_p: float = 0.05,
        typical_p: float = 1.0,
        stream: bool = False,
        stop: Optional[Union[str, List[str]]] = [],
        seed: Optional[int] = None,
        response_format: Optional[ChatCompletionRequestResponseFormat] = None,
        max_tokens: Optional[int] = None,
        present_penalty: float = 0.0,
        presence_penalty: Optional[float] = None,
        frequency_penalty: float = 0.0,
        repeat_penalty: float = 1.0,
        penalty_last_n: int = 64,
        mirostat_mode: int = 0,
        mirostat_tau: float = 5.0,
        mirostat_eta: float = 0.1,
        dynatemp_range: float = 0.0,
        dynatemp_exponent: float = 1.0,
        min_keep: int = 0,
        xtc_threshold: float = 0.1,
        xtc_probability: float = 0.0,
        dry_multiplier: float = 0.0,
        dry_base: float = 1.75,
        dry_allowed_length: int = 2,
        dry_penalty_last_n:int = 64,
        dry_seq_breakers: list[str] = ["\n", ":", "\"", "*"],
        adaptive_target : float = -1.0,
        adaptive_decay : float = 0.9,
        use_infill: bool = False,
        model: Optional[str] = None,
        logit_bias: Optional[Dict[int, float]] = None,
        logits_processor: Optional[LogitsProcessorList] = None,
        grammar: Optional[LlamaGrammar] = None,
        grammar_lazy: bool = False,
        active_loras: Optional[List[Dict[str, Union[str, float]]]] = None,
        control_vector: Optional[Dict[str, Any]] = None,
        logprobs: Optional[bool] = None,
        top_logprobs: Optional[int] = None,
        assistant_prefill: bool = False,
        add_generation_prompt: bool = True,
        # Reasoning Budget Params
        reasoning_budget: int = -1,
        reasoning_start: str = "<think>",
        reasoning_end: str = "</think>",
        reasoning_budget_message: Optional[str] = None,
        reasoning_start_in_prompt: bool = False,
        reasoning_start_max_tokens: Optional[int] = 32,
    ) -> Union[
        CreateChatCompletionResponse, Iterator[CreateChatCompletionStreamResponse]
    ]:
        """Generate a chat completion from a list of messages.

        Args:
            messages: A list of messages to generate a response for.
            functions: A list of functions to use for the chat completion.
            function_call: A function call to use for the chat completion.
            tools: A list of tools to use for the chat completion.
            tool_choice: A tool choice to use for the chat completion.
            temperature: The temperature to use for sampling.
            top_p: The top-p value to use for nucleus sampling. Nucleus sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751
            top_k: The top-k value to use for sampling. Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751
            top_n_sigma: Limit the next token selection to a subset of tokens with pre-softmax logits that are within n * σ less than the max logit (default: -1.00, -1.00 = disabled).
            min_p: The min-p value to use for minimum p sampling. Minimum P sampling as described in https://github.com/ggml-org/llama.cpp/pull/3841
            typical_p: The typical-p value to use for sampling. Locally Typical Sampling implementation described in the paper https://arxiv.org/abs/2202.00666.
            stream: Whether to stream the results.
            stop: A list of strings to stop generation when encountered.
            seed: The seed to use for sampling.
            response_format: The response format to use for the chat completion. Use { "type": "json_object" } to contstrain output to only valid json.
            max_tokens: The maximum number of tokens to generate. If max_tokens <= 0 or None, the maximum number of tokens to generate is unlimited and depends on n_ctx.
            frequency_penalty: The penalty to apply to tokens based on their frequency in the prompt.
            present_penalty: The penalty to controls whether to apply a penalty to tokens that are already present in the current context, helping to reduce repetition and encourage more diverse generation.
            presence_penalty: Compatibility alias for `present_penalty`. It is used only when `present_penalty` remains at its default value.
            repeat_penalty: The penalty to apply to repeated tokens.
            penalty_last_n: last n tokens to penalize (0 = disable penalty, -1 = context size).
            mirostat_mode: The mirostat sampling mode.
            mirostat_tau: The mirostat sampling tau parameter.
            mirostat_eta: The mirostat sampling eta parameter.
            dynatemp_range: Range of dynamic temperature.
            dynatemp_exponent: Exponent of dynamic temperature.
            min_keep: Minimum tokens to keep for sampling.
            xtc-probability: Sets the chance for token removal (checked once on sampler start) (default: 0.0). XTC sampler as described in https://github.com/oobabooga/text-generation-webui/pull/6335
            xtc-threshold: Sets a minimum probability threshold for tokens to be removed (default: 0.1).XTC sampler as described in https://github.com/oobabooga/text-generation-webui/pull/6335
            dry_multiplier: Set the DRY (Don't Repeat Yourself) repetition penalty multiplier. Default: `0.0`, which is disabled.
            dry_base`: Set the DRY repetition penalty base value. Default: `1.75`
            dry_allowed_length: Tokens that extend repetition beyond this receive exponentially increasing penalty: multiplier * base ^ (length of repeating sequence before token - allowed length). Default: `2`
            dry_penalty_last_n: How many tokens to scan for repetitions. Default: `64`; `0` disables scanning and `-1` uses the context size.
            dry_seq_breakers: Specify an array of sequence breakers for DRY sampling. Only a JSON array of strings is accepted. Default: `['\n', ':', '"', '*']`
            adaptive-target: Adaptive-p: select tokens near this probability (valid range 0.0 to 1.0; negative = disabled) (default: %.2f) [(more info)](https://github.com/ggml-org/llama.cpp/pull/17927)
            adaptive-decay: Adaptive-p: decay rate for target adaptation over time. lower values are more reactive, higher values are more stable. (valid range 0.0 to 0.99) (default: %.2f)
            use_infill: Determines whether to activate the specialized fill-in-the-middle sampler that consolidates probabilities of tokens sharing common prefixes to ensure the generated text coherently bridges the gap between the prefix and suffix.
            model: The name to use for the model in the completion object.
            logit_bias: A logit bias to use.
            logits_processor: A list of logits processors to use.
            grammar: A grammar to use.
            grammar_lazy: If True, enables lazy evaluation.
            reasoning_budget: Token budget for the first visible reasoning block.
                -1 disables the sampler, 0 forces an immediate end after reasoning starts,
                and N > 0 allows at most N generated tokens inside the block.
            reasoning_start: Token/text sequence that marks the beginning of the first reasoning block.
            reasoning_end: Token/text sequence that naturally and forcibly ends the reasoning block.
            reasoning_budget_message: Optional message inserted before reasoning_end when the budget is exhausted.
            reasoning_start_in_prompt: Set True when the prompt/template already inserted reasoning_start.
            reasoning_start_max_tokens: Safety window before disabling the sampler for non-reasoning outputs.
            active_loras: A list of dictionaries specifying the LoRA adapters to dynamically apply during generation.
                Each dictionary must contain a "name" key (matching a LoRA previously loaded into VRAM via `load_lora()`)
                and an optional "scale" key (float, defaults to 1.0).
                Example: `[{"name": "role_A", "scale": 0.85}, {"name": "role_B", "scale": 0.5}]`.
            control_vector: A dictionary containing Control Vector (CVec) data for representation engineering.
                Must contain a "data" key with a flattened 1D list of floats.
                Optionally accepts "layer_start" (int, defaults to 1) and "layer_end" (int, defaults to the model's total layer count).
                Note: The length of the "data" list MUST be at least `n_embd * layer_end`, with zero-padding for any skipped early layers.

        Returns:
            Generated chat completion or a stream of chat completion chunks.
        """
        if presence_penalty is not None and present_penalty == 0.0:
            present_penalty = presence_penalty

        handler = (
            self.chat_handler
            or self._chat_handlers.get(self.chat_format)
            or llama_chat_format.get_chat_completion_handler(self.chat_format)
        )
        return handler(
            llama=self,
            messages=messages,
            functions=functions,
            function_call=function_call,
            tools=tools,
            tool_choice=tool_choice,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            top_n_sigma=top_n_sigma,
            min_p=min_p,
            typical_p=typical_p,
            logprobs=logprobs,
            top_logprobs=top_logprobs,
            stream=stream,
            stop=stop,
            seed=seed,
            response_format=response_format,
            max_tokens=max_tokens,
            present_penalty=present_penalty,
            frequency_penalty=frequency_penalty,
            repeat_penalty=repeat_penalty,
            penalty_last_n=penalty_last_n,
            mirostat_mode=mirostat_mode,
            mirostat_tau=mirostat_tau,
            mirostat_eta=mirostat_eta,
            dynatemp_range=dynatemp_range,
            dynatemp_exponent=dynatemp_exponent,
            min_keep=min_keep,
            xtc_threshold=xtc_threshold,
            xtc_probability=xtc_probability,
            dry_multiplier=dry_multiplier,
            dry_base=dry_base,
            dry_allowed_length=dry_allowed_length,
            dry_penalty_last_n=dry_penalty_last_n,
            dry_seq_breakers=dry_seq_breakers,
            adaptive_target=adaptive_target,
            adaptive_decay=adaptive_decay,
            use_infill=use_infill,
            model=model,
            logit_bias=logit_bias,
            logits_processor=logits_processor,
            grammar=grammar,
            grammar_lazy=grammar_lazy,
            active_loras=active_loras,
            control_vector=control_vector,
            assistant_prefill=assistant_prefill,
            add_generation_prompt=add_generation_prompt,
            reasoning_budget=reasoning_budget,
            reasoning_start=reasoning_start,
            reasoning_end=reasoning_end,
            reasoning_budget_message=reasoning_budget_message,
            reasoning_start_in_prompt=reasoning_start_in_prompt,
            reasoning_start_max_tokens=reasoning_start_max_tokens,
        )

    def create_chat_completion_openai_v1(
        self,
        *args: Any,
        **kwargs: Any,
    ):
        """Generate a chat completion with return type based on the the OpenAI v1 API.

        OpenAI python package is required to use this method.

        You can install it with `pip install openai`.

        Args:
            *args: Positional arguments to pass to create_chat_completion.
            **kwargs: Keyword arguments to pass to create_chat_completion.

        Returns:
            Generated chat completion or a stream of chat completion chunks.
        """
        try:
            from openai.types.chat import ChatCompletion, ChatCompletionChunk

            stream = kwargs.get("stream", False)  # type: ignore
            assert isinstance(stream, bool)
            if stream:
                return (ChatCompletionChunk(**chunk) for chunk in self.create_chat_completion(*args, **kwargs))  # type: ignore
            else:
                return ChatCompletion(**self.create_chat_completion(*args, **kwargs))  # type: ignore
        except ImportError:
            raise ImportError(
                "To use create_chat_completion_openai_v1, you must install the openai package."
                "You can install it with `pip install openai`."
            )

    def __getstate__(self):
        return dict(
            model_path=self.model_path,
            # Model Params
            n_gpu_layers=self.model_params.n_gpu_layers,
            cpu_moe=self.cpu_moe,
            n_cpu_moe=self.n_cpu_moe,
            split_mode=self.model_params.split_mode,
            load_mode=self.model_params.load_mode,
            lazy_mode=self.model_params.lazy_mode,
            main_gpu=self.model_params.main_gpu,
            tensor_split=self.tensor_split,
            rpc_servers=self.rpc_servers,
            rpc_local_devices=self.rpc_local_devices,
            kv_overrides=self.kv_overrides,
            vocab_only=self.model_params.vocab_only,
            check_tensors=self.model_params.check_tensors,
            use_extra_bufts=self.model_params.use_extra_bufts,
            no_host=self.model_params.no_host,
            no_alloc=self.model_params.no_alloc,
            load_mtp=self.model_params.load_mtp,
            # Context Params
            seed=self._seed,
            n_ctx=self.context_params.n_ctx,
            n_batch=self.context_params.n_batch,
            n_ubatch=self.context_params.n_ubatch,
            n_seq_max=self.context_params.n_seq_max,
            n_rs_seq=self.context_params.n_rs_seq,
            n_outputs_max=self.context_params.n_outputs_max,
            n_outputs_max_per_seq=self.context_params.n_outputs_max_per_seq,
            moe_cache_size=self.context_params.moe_cache_size,
            n_threads=self.context_params.n_threads,
            n_threads_batch=self.context_params.n_threads_batch,
            ctx_type=self.context_params.ctx_type,
            rope_scaling_type=self.context_params.rope_scaling_type,
            pooling_type=self.context_params.pooling_type,
            attention_type=self.context_params.attention_type,
            flash_attn_type=self.context_params.flash_attn_type,
            rope_freq_base=self.context_params.rope_freq_base,
            rope_freq_scale=self.context_params.rope_freq_scale,
            yarn_ext_factor=self.context_params.yarn_ext_factor,
            yarn_attn_factor=self.context_params.yarn_attn_factor,
            yarn_beta_fast=self.context_params.yarn_beta_fast,
            yarn_beta_slow=self.context_params.yarn_beta_slow,
            yarn_orig_ctx=self.context_params.yarn_orig_ctx,
            logits_all=self._logits_all,
            embedding=self.context_params.embeddings,
            offload_kqv=self.context_params.offload_kqv,
            op_offload=self.context_params.op_offload,
            swa_full=self.context_params.swa_full,
            kv_unified= self.context_params.kv_unified,
            # Sampling Params
            no_perf=self.context_params.no_perf,
            last_n_tokens_size=self.last_n_tokens_size,
            # Backend Params
            numa=self.numa,
            # Chat Format Params
            chat_format=self.chat_format,
            chat_handler=self.chat_handler,
            # Speculative Decidng
            draft_model=self.draft_model,
            speculative=self.speculative_config,
            # KV cache quantization
            type_k=self.context_params.type_k,
            type_v=self.context_params.type_v,
            # Misc
            spm_infill=self.spm_infill,
            verbose=self.verbose,
        )

    def __setstate__(self, state):
        self.__init__(**state)

    def _state_compatibility(self) -> Dict[str, Any]:
        """Conservative same-model/context check, not a model content hash."""
        model_stat = os.stat(self.model_path)
        return {
            "model_path": os.path.normcase(os.path.abspath(self.model_path)),
            "model_size": model_stat.st_size,
            "model_mtime_ns": model_stat.st_mtime_ns,
            "n_ctx": self._n_ctx,
            "n_vocab": self._n_vocab,
            "logits_all": self._logits_all,
            **{name: getattr(self.context_params, name) for name in (
                "type_k", "type_v", "n_seq_max", "n_rs_seq", "rope_scaling_type",
                "rope_freq_base", "rope_freq_scale", "attention_type",
            )},
        }

    def save_state(self) -> LlamaState:
        """Own a memory snapshot and last output; sampler/draft state is not saved."""
        if getattr(self, "_speculative_verifying", False):
            raise RuntimeError("Cannot save LlamaState during speculative verification")
        last_logits = None
        if self.n_tokens > 0 and (
            self._last_eval_output_start <= self.n_tokens - 1
            < self._last_eval_output_start + self._last_eval_output_count
        ):
            restored = getattr(self, "_restored_logits", None)
            # LlamaState takes ownership by copying below. Native state export
            # serializes memory without changing the decoded output rows.
            last_logits = (restored if restored is not None else
                np.ctypeslib.as_array(
                    self._ctx.get_logits_ith(self.n_tokens - 1 - self._last_eval_output_start),
                    shape=(self._n_vocab,),
                ))
        if self.verbose:
            print("Llama.save_state: saving llama state", file=sys.stderr)

        # Query the backend for the required buffer size to store the current state.
        state_size = llama_cpp_lib.llama_state_get_size(self._ctx.ctx)
        if self.verbose:
            print(f"Llama.save_state: got state size: {state_size}", file=sys.stderr)

        # Allocate a ctypes uint8 array (buffer) of the required size.
        llama_state = (ctypes.c_uint8 * int(state_size))()
        if self.verbose:
            print("Llama.save_state: allocated state", file=sys.stderr)

        # Copy the raw state data from the internal C context into our Python-managed buffer.
        # Returns the actual number of bytes written (n_bytes).
        n_bytes = llama_cpp_lib.llama_state_get_data(self._ctx.ctx, llama_state, state_size)
        if self.verbose:
            print(f"Llama.save_state: copied llama state: {n_bytes}", file=sys.stderr)

        # Safety check to prevent buffer overflow issues.
        if not 0 < int(n_bytes) <= int(state_size):
            raise RuntimeError("Failed to copy llama state data")

        # Directly read 'n_bytes' from the buffer's memory address to create the Python bytes object.
        # Significantly reducing memory overhead by avoiding an intermediate array allocation.
        llama_state_bytes = ctypes.string_at(ctypes.addressof(llama_state), int(n_bytes))
        del llama_state  # Release the export buffer before copying score arrays.
        if self.verbose:
            print(
                f"Llama.save_state: saving {n_bytes} bytes of llama state",
                file=sys.stderr,
            )

        # Create and return the snapshot object.
        return LlamaState(
            scores=(self.scores[:self.n_tokens] if self._logits_all else
                    last_logits.reshape(1, -1) if last_logits is not None else
                    np.empty((0, self._n_vocab), dtype=np.single)),
            input_ids=self.input_ids[:self.n_tokens],
            n_tokens=self.n_tokens,
            llama_state=llama_state_bytes,
            llama_state_size=n_bytes,
            seed=self._seed,
            last_logits=last_logits,
            compatibility=self._state_compatibility(),
        )

    def load_state(self, state: LlamaState) -> None:
        """Restore memory and optional owned output, starting a new sampler session."""
        if getattr(self, "_speculative_verifying", False):
            raise RuntimeError("Cannot load LlamaState during speculative verification")
        compatibility = getattr(state, "compatibility", None)
        if compatibility is not None and compatibility != self._state_compatibility():
            raise ValueError("LlamaState model/context configuration does not match")
        if not 0 <= state.n_tokens <= self._n_ctx:
            raise ValueError("LlamaState token count exceeds the context")
        if state.llama_state_size <= 0 or state.llama_state_size != len(state.llama_state):
            raise ValueError("LlamaState native byte size is invalid")
        ids = np.asarray(state.input_ids)
        scores = np.asarray(state.scores)
        if ids.ndim != 1 or len(ids) < state.n_tokens or not np.issubdtype(ids.dtype, np.integer):
            raise ValueError("LlamaState input_ids are invalid")
        if scores.ndim != 2 or scores.shape[1] != self._n_vocab:
            raise ValueError("LlamaState scores have an incompatible vocabulary")
        last_logits = getattr(state, "last_logits", None)
        if last_logits is not None:
            if np.shape(last_logits) != (self._n_vocab,) or state.n_tokens == 0:
                raise ValueError("LlamaState last_logits are invalid")
            last_logits = np.array(last_logits, dtype=np.single, copy=True)
        # Allocate and validate before mutating either side. Native reads can
        # partially mutate memory on failure, so failure leaves a cleared model.
        new_ids = np.zeros(self._n_ctx, dtype=np.intc)
        new_ids[:state.n_tokens] = ids[:state.n_tokens]
        if self._logits_all:
            limit = min(state.n_tokens, len(scores))
            new_scores = np.asarray(scores[:limit], dtype=np.single)
        else:
            new_scores = np.asarray(scores[-1:], dtype=np.single)
        # Normal snapshots own separate float32 arrays. Only copy for dtype
        # conversion above or manually constructed aliases of the live buffer.
        if np.may_share_memory(new_scores, self.scores):
            new_scores = new_scores.copy()
        state_size = state.llama_state_size
        # Native input is const and consumed synchronously. Keep immutable bytes
        # alive for the call rather than duplicating the entire memory snapshot.
        state_bytes = bytes(state.llama_state)
        llama_state = ctypes.cast(ctypes.c_char_p(state_bytes), ctypes.POINTER(ctypes.c_uint8))
        try:
            if self.speculative is not None:
                self.speculative.clear()
            if getattr(self, "_sampling_ctx", None) is not None:
                self._sampling_ctx.close()
                self._sampling_ctx = None
            if self._ctx.set_state_data(llama_state, state_size) != state_size:
                raise RuntimeError("Failed to set llama state data")
        except Exception:
            self.reset()
            raise
        self.input_ids = new_ids
        # Reuse the context-sized allocation after successful native restore.
        # Snapshot scores normally need no additional allocation.
        self.scores.fill(0)
        self.scores[:len(new_scores)] = new_scores
        if last_logits is not None:
            self.scores[state.n_tokens - 1 if self._logits_all else 0] = last_logits
        self.n_tokens = state.n_tokens
        self._seed = state.seed
        self._prefilled_prompt = None
        self._restored_logits = last_logits
        self._last_eval_output_start = max(0, self.n_tokens - 1)
        self._last_eval_output_count = int(last_logits is not None)
        self._state_needs_speculative_reset = self.speculative is not None

    def n_ctx(self) -> int:
        """Return the context window size."""
        return self._ctx.n_ctx()

    def n_ctx_train(self) -> int:
        """Return the training context window size."""
        return self._model.n_ctx_train()

    def n_embd(self) -> int:
        """Return the embedding size."""
        return self._model.n_embd()

    def n_embd_inp(self) -> int:
        """Return the input embedding size."""
        return self._model.n_embd_inp()

    def n_embd_out(self) -> int:
        """Return the output embedding size."""
        return self._model.n_embd_out()

    def n_layer(self) -> int:
        """Return the n_layer value."""
        return self._model.n_layer()

    def n_layer_nextn(self) -> int:
        """Return the n_layer_nextn value."""
        return self._model.n_layer_nextn()

    def n_head(self) -> int:
        """Return the head size."""
        return self._model.n_head()

    def n_head_kv(self) -> int:
        """Return the head_kv size."""
        return self._model.n_head_kv()

    def n_swa(self) -> int:
        """Return the swa size."""
        return self._model.n_swa()

    def n_params(self) -> int:
        """Returns the total number of parameters in the model"""
        return self._model.n_params()

    def n_vocab(self) -> int:
        """Return the vocabulary size."""
        return self._model.n_vocab()

    def tokenizer(self) -> LlamaTokenizer:
        """Return the llama tokenizer for this model."""
        return LlamaTokenizer(self)

    def token_bos(self) -> int:
        """Return the beginning-of-sequence token."""
        return self._model.token_bos()

    def token_eos(self) -> int:
        """Return the end-of-sequence token."""
        return self._model.token_eos()

    def token_eot(self) -> int:
        """Return the end-of-turn token."""
        return self._model.token_eot()

    def token_sep(self) -> int:
        """Return the sentence-separator token."""
        return self._model.token_sep()

    def token_nl(self) -> int:
        """Return the next-line token."""
        return self._model.token_nl()

    def token_pad(self) -> int:
        """Return the padding token."""
        return self._model.token_pad()

    def token_mask(self) -> int:
        """Return the mask token."""
        return self._model.token_mask()

    def pooling_type(self) -> str:
        """Return the pooling type."""
        return self._ctx.pooling_type()

    @staticmethod
    def logits_to_logprobs(
        logits: Union[npt.NDArray[np.single], List], axis: int = -1
    ) -> npt.NDArray[np.single]:
        # https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.log_softmax.html
        logits_maxs: np.ndarray = np.amax(logits, axis=axis, keepdims=True)
        if logits_maxs.ndim > 0:
            logits_maxs[~np.isfinite(logits_maxs)] = 0
        elif not np.isfinite(logits_maxs):
            logits_maxs = 0
        subtract_maxs = np.subtract(logits, logits_maxs, dtype=np.single)
        exp = np.exp(subtract_maxs)
        # Suppress warnings about log of zero
        with np.errstate(divide="ignore"):
            summed = np.sum(exp, axis=axis, keepdims=True)
            out = np.log(summed)
        return subtract_maxs - out

    @staticmethod
    def longest_token_prefix(
        current_ids: Union[Sequence[int], npt.NDArray[np.intc]],
        new_tokens: Union[Sequence[int], npt.NDArray[np.intc]],
        verbose: bool = False
    ) -> int:
        """
        Calculates the length of the longest common prefix between two token sequences.

        This implementation uses NumPy for vectorized comparison (SIMD), which offers
        significant performance improvements (up to 2x~100x+ speedup) over standard Python
        loops for long contexts (e.g., RAG or chat history).

        Args:
            current_ids: The existing token sequence (e.g., KV cache).
            new_tokens: The new input token sequence.
            verbose: If True, prints detailed debug information to stderr.

        Returns:
            int: The number of matching tokens from the start.
        """
        # Fast exit for empty sequences to avoid unnecessary processing
        if len(current_ids) == 0 or len(new_tokens) == 0:
            if verbose:
                print(
                    f"Llama.longest_token_prefix [Fast Exit 1]: Empty sequence detected. "
                    f"len(current_ids)={len(current_ids)}, len(new_tokens)={len(new_tokens)}",
                    file=sys.stderr
                )
            return 0

        # Determine the comparison range (limited by the shorter sequence)
        min_len = min(len(current_ids), len(new_tokens))

        # Probe inspection: Use Python to quickly compare the first token
        # If the tokens are different from the beginning, return immediately to avoid any NumPy overhead.
        if current_ids[0] != new_tokens[0]:
            if verbose:
                print(
                    f"Llama.longest_token_prefix [Fast Exit 2]: First token mismatch. "
                    f"current_ids[0]={current_ids[0]} vs new_tokens[0]={new_tokens[0]}",
                    file=sys.stderr
                )
            return 0

        # Accelerating SIMD for Large Data Volumes
        # Only transform necessary slices, avoid processing irrelevant data
        # Use asarray to ensure zero-copy (if the input is already an array)
        current_ids_array = np.asarray(current_ids[:min_len], dtype=np.intc)
        new_tokens_array = np.asarray(new_tokens[:min_len], dtype=np.intc)

        # Perform vectorized element-wise comparison (SIMD instruction set usage)
        # Creates a boolean array where True indicates a match (e.g., [True, True, False, ...])
        matches = (current_ids_array == new_tokens_array)

        # Find the index of the first mismatch efficiently
        # np.argmin returns the index of the minimum value. Since False (0) < True (1),
        # this locates the first False value (mismatch).
        idx = np.argmin(matches)

        # Handle the "Full Match" edge case
        # This means that the match between the two arrays will still result in True in the end.
        if matches[idx]:
            return int(min_len)

        # Otherwise, idx is the position of the first mismatch, which equals the prefix length.
        return int(idx)

    @classmethod
    def from_pretrained(
        cls,
        repo_id: str,
        filename: Optional[str],
        additional_files: Optional[List] = None,
        local_dir: Optional[Union[str, os.PathLike[str]]] = None,
        local_dir_use_symlinks: Union[bool, Literal["auto"]] = "auto",
        cache_dir: Optional[Union[str, os.PathLike[str]]] = None,
        **kwargs: Any,
    ) -> "Llama":
        """Create a Llama model from a pretrained model name or path.
        This method requires the huggingface_hub package.
        You can install it with `pip install --upgrade huggingface_hub`.

        Args:
            repo_id: The model repo id.
            filename: A filename or glob pattern to match the model file in the repo.
            additional_files: A list of filenames or glob patterns to match additional model files in the repo.
            local_dir: The local directory to save the model to.
            local_dir_use_symlinks: Whether to use symlinks when downloading the model.
            **kwargs: Additional keyword arguments to pass to the Llama constructor.

        Returns:
            A Llama model."""
        try:
            from huggingface_hub import hf_hub_download, HfFileSystem
            from huggingface_hub.utils import validate_repo_id
        except ImportError:
            raise ImportError(
                "Llama.from_pretrained requires the huggingface-hub package. "
                "You can install it with `pip install --upgrade huggingface_hub`."
            )

        validate_repo_id(repo_id)

        hffs = HfFileSystem()

        files = [
            file["name"] if isinstance(file, dict) else file
            for file in hffs.ls(repo_id, recursive=True)
        ]

        # split each file into repo_id, subfolder, filename
        file_list: List[str] = []
        for file in files:
            rel_path = Path(file).relative_to(repo_id)
            file_list.append(str(rel_path))

        # find the only/first shard file:
        matching_files = [file for file in file_list if fnmatch.fnmatch(file, filename)]  # type: ignore

        if len(matching_files) == 0:
            raise ValueError(
                f"No file found in {repo_id} that match {filename}\n\n"
                f"Available Files:\n{json.dumps(file_list)}"
            )

        if len(matching_files) > 1:
            raise ValueError(
                f"Multiple files found in {repo_id} matching {filename}\n\n"
                f"Available Files:\n{json.dumps(files)}"
            )

        (matching_file,) = matching_files

        subfolder = str(Path(matching_file).parent)
        filename = Path(matching_file).name

        # download the file
        hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            subfolder=subfolder,
            local_dir=local_dir,
            local_dir_use_symlinks=local_dir_use_symlinks,
            cache_dir=cache_dir,
        )

        if additional_files:
            for additonal_file_name in additional_files:
                # find the additional shard file:
                matching_additional_files = [file for file in file_list if fnmatch.fnmatch(file, additonal_file_name)]

                if len(matching_additional_files) == 0:
                    raise ValueError(
                        f"No file found in {repo_id} that match {additonal_file_name}\n\n"
                        f"Available Files:\n{json.dumps(file_list)}"
                    )

                if len(matching_additional_files) > 1:
                    raise ValueError(
                        f"Multiple files found in {repo_id} matching {additonal_file_name}\n\n"
                        f"Available Files:\n{json.dumps(files)}"
                    )

                (matching_additional_file,) = matching_additional_files

                # download the additional file
                hf_hub_download(
                    repo_id=repo_id,
                    filename=matching_additional_file,
                    subfolder=subfolder,
                    local_dir=local_dir,
                    local_dir_use_symlinks=local_dir_use_symlinks,
                    cache_dir=cache_dir,
                )

        if local_dir is None:
            model_path = hf_hub_download(
                repo_id=repo_id,
                filename=filename,
                subfolder=subfolder,
                local_dir=local_dir,
                local_dir_use_symlinks=local_dir_use_symlinks,
                cache_dir=cache_dir,
                local_files_only=True,
            )
        else:
            model_path = os.path.join(local_dir, filename)

        # loading the first file of a sharded GGUF loads all remaining shard files in the subfolder
        return cls(
            model_path=model_path,
            **kwargs,
        )


class LlamaState:
    """Owned host snapshot with optional last output and compatibility metadata.

    This is not an exact sampler/draft continuation or a device checkpoint.
    Arrays are independent of the source context and remain pickleable.
    """
    def __init__(
        self,
        input_ids: npt.NDArray[np.intc],
        scores: npt.NDArray[np.single],
        n_tokens: int,
        llama_state: bytes,
        llama_state_size: int,
        seed: int,
        *,
        last_logits: Optional[npt.NDArray[np.single]] = None,
        compatibility: Optional[Dict[str, Any]] = None,
    ):
        self.input_ids = input_ids.copy()
        self.scores = scores.copy()
        self.n_tokens = n_tokens
        self.llama_state = bytes(llama_state)
        self.llama_state_size = llama_state_size
        self.seed = seed
        self.last_logits = None if last_logits is None else last_logits.copy()
        self.compatibility = None if compatibility is None else dict(compatibility)

    @property
    def nbytes(self) -> int:
        """Owned payload bytes, excluding Python object/allocator overhead."""
        logits = getattr(self, "last_logits", None)
        return (len(self.llama_state) + self.input_ids.nbytes + self.scores.nbytes
                + (0 if logits is None else logits.nbytes))


LogitsProcessor = Callable[
    [npt.NDArray[np.intc], npt.NDArray[np.single]], npt.NDArray[np.single]
]


class LogitsProcessorList(List[LogitsProcessor]):
    def __call__(
        self, input_ids: npt.NDArray[np.intc], scores: npt.NDArray[np.single]
    ) -> npt.NDArray[np.single]:
        for processor in self:
            scores = processor(input_ids, scores)
        return scores


StoppingCriteria = Callable[[npt.NDArray[np.intc], npt.NDArray[np.single]], bool]


class StoppingCriteriaList(List[StoppingCriteria]):
    @property
    def requires_logits(self) -> bool:
        return any(getattr(criteria, "requires_logits", True) for criteria in self)

    def __call__(
        self, input_ids: npt.NDArray[np.intc], logits: npt.NDArray[np.single]
    ) -> bool:
        return any([stopping_criteria(input_ids, logits) for stopping_criteria in self])


class MinTokensLogitsProcessor(LogitsProcessor):
    def __init__(self, min_tokens: int, token_eos: int):
        self.min_tokens = min_tokens
        self.token_eos = token_eos
        self.prompt_tokens = None

    def __call__(
        self, input_ids: npt.NDArray[np.intc], scores: npt.NDArray[np.single]
    ) -> npt.NDArray[np.single]:
        if self.prompt_tokens is None:
            self.prompt_tokens = len(input_ids)
        if len(input_ids) - self.prompt_tokens < self.min_tokens:
            scores[self.token_eos] = -np.inf
        return scores