diff --git a/CHANGELOG.md b/CHANGELOG.md index 5b337ad923..6f7dcb6683 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +- fix(example): retain recurrent state for server MTP rollback +- feat: update llama.cpp to ggml-org/llama.cpp@adb55e514 + +## [0.3.34] + +- feat: update llama.cpp to ggml-org/llama.cpp@e3546c794 + +## [0.3.33] + +- feat: update llama.cpp to ggml-org/llama.cpp@78d2f5246 + +## [0.3.32] + +- feat(example): support chained NextN heads for server MTP drafting +- feat: update llama.cpp to ggml-org/llama.cpp@b3fed31b9 +- fix: preserve recurrent/hybrid model state when the full prompt is already cached by @allthatido and @abetlen in #2306 + +## [0.3.31] + +- feat: update llama.cpp to ggml-org/llama.cpp@f449e0553 + ## [0.3.30] - feat: update llama.cpp to ggml-org/llama.cpp@e3a74b299 diff --git a/examples/low_level_api/low_level_api_chat_cpp.py b/examples/low_level_api/low_level_api_chat_cpp.py index 20f7a158ac..43b0b31bb8 100644 --- a/examples/low_level_api/low_level_api_chat_cpp.py +++ b/examples/low_level_api/low_level_api_chat_cpp.py @@ -76,8 +76,14 @@ def __init__(self, params: GptParams) -> None: self.lparams.n_parts = self.params.n_parts self.lparams.seed = self.params.seed self.lparams.memory_f16 = self.params.memory_f16 - self.lparams.use_mlock = self.params.use_mlock - self.lparams.use_mmap = self.params.use_mmap + if self.params.use_mmap and self.params.use_mlock: + self.lparams.load_mode = llama_cpp.LLAMA_LOAD_MODE_MMAP_MLOCK + elif self.params.use_mlock: + self.lparams.load_mode = llama_cpp.LLAMA_LOAD_MODE_MLOCK + elif self.params.use_mmap: + self.lparams.load_mode = llama_cpp.LLAMA_LOAD_MODE_MMAP + else: + self.lparams.load_mode = llama_cpp.LLAMA_LOAD_MODE_NONE self.model = llama_cpp.llama_model_load_from_file( self.params.model.encode("utf8"), self.lparams diff --git a/examples/server/server.py b/examples/server/server.py index 16f8c9f7e5..2eb31aed3c 100644 --- a/examples/server/server.py +++ b/examples/server/server.py @@ -1288,7 +1288,14 @@ def __init__( raise RuntimeError("failed to create MTP draft context") ctx_other = llama_cpp_ext.llama_get_ctx_other(self.ctx) self.is_mem_shared = bool(ctx_other and ctx_other == self.target_ctx) - self.sampled_batch_draft = not self.is_mem_shared + self.n_mtp_layers = max( + 1, + int(llama_cpp.llama_model_n_layer_nextn(self.model)), + ) + self.chain_heads = self.n_mtp_layers > 1 and not self.is_mem_shared + if self.chain_heads: + self.num_pred_tokens = min(self.num_pred_tokens, self.n_mtp_layers) + self.sampled_batch_draft = not self.is_mem_shared and not self.chain_heads self.n_batch = int(llama_cpp.llama_n_batch(self.ctx)) mem = llama_cpp.llama_get_memory(self.ctx) if mem is None: @@ -1451,6 +1458,17 @@ def _try_decode_batch(self) -> bool: return False return True + def _set_nextn_layer_offset(self, offset: int) -> None: + if self.chain_heads: + llama_cpp_ext.llama_set_nextn_layer_offset(self.ctx, offset) + + def _try_decode_batch_for_mtp_head(self, head: int) -> bool: + self._set_nextn_layer_offset(head) + try: + return self._try_decode_batch() + finally: + self._set_nextn_layer_offset(0) + def _decode_batch(self) -> None: n_tokens = int(self.batch.n_tokens) if n_tokens <= 0: @@ -1464,6 +1482,22 @@ def _decode_batch(self) -> None: self.decode_failures_total += 1 raise RuntimeError(f"MTP draft decode failed with code {result}") + def _decode_batch_for_mtp_heads( + self, + start_pos_by_seq: Dict[int, int], + ) -> None: + if not self.chain_heads: + self._decode_batch() + return + try: + for head in range(self.n_mtp_layers): + for seq_id, start_pos in start_pos_by_seq.items(): + llama_cpp.llama_memory_seq_rm(self.mem, seq_id, start_pos, -1) + self._set_nextn_layer_offset(head) + self._decode_batch() + finally: + self._set_nextn_layer_offset(0) + def metric_definitions( self, ) -> List[Tuple[str, str, str, Union[int, float]]]: @@ -1577,7 +1611,8 @@ def _process_rows( target_rows_by_seq: Dict[int, List[int]], aligned_by_seq: Dict[int, bool], ) -> None: - added_pos_by_seq: Dict[int, int] = {} + added_start_pos_by_seq: Dict[int, int] = {} + added_end_pos_by_seq: Dict[int, int] = {} self._clear_batch() for index in range(start, end): if int(batch.n_seq_id[index]) != 1: @@ -1609,15 +1644,85 @@ def _process_rows( self._set_batch_embedding_row(slot, self.pending_h[seq_id]) else: self._set_batch_embedding_row(slot, h_tgt_rows[previous_row]) - added_pos_by_seq[seq_id] = pos + added_start_pos_by_seq.setdefault(seq_id, pos) + added_end_pos_by_seq[seq_id] = pos previous_row_by_seq[seq_id] = index target_rows_by_seq.setdefault(seq_id, []).append(index) if int(self.batch.n_tokens) > 0: - self._decode_batch() - for seq_id, pos in added_pos_by_seq.items(): + self._decode_batch_for_mtp_heads(added_start_pos_by_seq) + for seq_id, pos in added_end_pos_by_seq.items(): self.context_pos[seq_id] = max(self.context_pos[seq_id], pos + 1) + def _draft_chain_heads( + self, + *, + seq_id: int, + first_pos: int, + token: int, + n_predict: int, + ) -> np.ndarray: + if self.context_pos[seq_id] > first_pos: + self.truncate(seq_id, first_pos) + if self.context_pos[seq_id] < first_pos: + self.ready[seq_id] = False + return np.array([], dtype=np.intc) + + drafted: List[int] = [] + chain_tokens = [token] + chain_embeddings = [self.pending_h[seq_id].copy()] + self._reset_sampler(seq_id) + + try: + for head in range(min(n_predict, self.n_mtp_layers)): + llama_cpp.llama_memory_seq_rm(self.mem, seq_id, first_pos, -1) + self._clear_batch() + for offset, (chain_token, embedding) in enumerate( + zip(chain_tokens, chain_embeddings) + ): + slot = int(self.batch.n_tokens) + self._add_batch_token( + token=chain_token, + pos=first_pos + offset, + seq_id=seq_id, + logits=offset == len(chain_tokens) - 1, + ) + self._set_batch_embedding_row(slot, embedding) + + output_index = int(self.batch.n_tokens) - 1 + if not self._try_decode_batch_for_mtp_head(head): + break + self.context_pos[seq_id] = max( + self.context_pos[seq_id], + first_pos + len(chain_tokens), + ) + sampled_token = self._sample_token(output_index, seq_id=seq_id) + if sampled_token is None: + break + drafted.append(sampled_token) + if len(drafted) >= n_predict: + break + h_row = llama_cpp_ext.llama_get_embeddings_nextn_ith( + self.ctx, + output_index, + ) + if not h_row: + break + chain_tokens.append(sampled_token) + chain_embeddings.append( + np.ctypeslib.as_array( + h_row, + shape=(self.n_embd,), + ).copy() + ) + finally: + self._set_nextn_layer_offset(0) + self.truncate(seq_id, first_pos) + + if not drafted: + return np.array([], dtype=np.intc) + return np.asarray(drafted, dtype=np.intc) + def draft( self, input_ids: np.ndarray, @@ -1649,6 +1754,13 @@ def draft( token = int(input_ids[-1]) drafted: List[int] = [] + if self.chain_heads: + return self._draft_chain_heads( + seq_id=seq_id, + first_pos=first_pos, + token=token, + n_predict=n_predict, + ) if not self.is_mem_shared and self.context_pos[seq_id] > first_pos: self.truncate(seq_id, first_pos) if not self.is_mem_shared and self.context_pos[seq_id] < first_pos: @@ -1709,6 +1821,15 @@ def draft_many( /, ) -> List[np.ndarray]: results = [np.array([], dtype=np.intc) for _ in requests] + if self.chain_heads: + for result_index, (input_ids, seq_id, max_tokens) in enumerate(requests): + results[result_index] = self.draft( + input_ids, + seq_id=seq_id, + max_tokens=max_tokens, + ) + return results + active: List["MTPDraftProvider.DraftManyState"] = [] for result_index, (input_ids, seq_id, max_tokens) in enumerate(requests): if ( @@ -10178,6 +10299,7 @@ def __init__( llama_cpp.llama_sampler_chain_add( self._sampler, llama_cpp.llama_sampler_init_penalties( + n_vocab, 64, 1.0, frequency_penalty, @@ -10909,7 +11031,9 @@ def _build_prompt_plan_locked( "multiple videos require MTMD to report frame counts" ) input_text = mtmd_cpp.mtmd_input_text() - input_text.text = prompt.encode("utf-8") + input_text_bytes = prompt.encode("utf-8") + input_text.text = input_text_bytes + input_text.text_len = len(input_text_bytes) input_text.add_special = False input_text.parse_special = True chunks = mtmd_cpp.mtmd_input_chunks_init() @@ -11189,6 +11313,7 @@ def __init__( vocab_only=vocab_only, use_mmap=use_mmap, use_mlock=use_mlock, + load_mtp=draft_model == "draft-mtp", kv_overrides=kv_overrides, ) ) @@ -11237,6 +11362,11 @@ def __init__( "speculative decoding is only supported for attention models" ) n_ctx_train = int(llama_cpp.llama_model_n_ctx_train(llama_model)) + target_n_rs_seq = ( + max(1, draft_model_num_pred_tokens) + if normalized_draft_model == "draft-mtp" + else None + ) context_params = self.build_context_params( n_ctx=n_ctx if n_ctx is not None else n_ctx_train, @@ -11264,7 +11394,7 @@ def __init__( type_k=type_k, type_v=type_v, kv_unified=kv_unified, - n_rs_seq=None, + n_rs_seq=target_n_rs_seq, ctx_type=None, ) ctx = llama_cpp.llama_init_from_model(llama_model, context_params) @@ -11294,6 +11424,15 @@ def __init__( "MTP requires runtime n_batch to fit the pending token plus draft tokens " f"(required {required_mtp_batch}, got {self.n_batch})" ) + if ( + target_n_rs_seq is not None + and self.exact_checkpoints_only + and self.n_rs_seq < target_n_rs_seq + ): + raise RuntimeError( + "MTP requires retained recurrent-state slots for rollback " + f"(required {target_n_rs_seq}, got {self.n_rs_seq})" + ) self.n_ctx_train = n_ctx_train self.n_vocab = int(llama_cpp.llama_vocab_n_tokens(self.vocab)) self.n_embd = int(llama_cpp.llama_model_n_embd(self.llama_model)) @@ -11467,6 +11606,7 @@ def build_model_params( vocab_only: Optional[bool], use_mmap: Optional[bool], use_mlock: Optional[bool], + load_mtp: bool, kv_overrides: Optional[Dict[str, Union[bool, int, float, str]]], ) -> Tuple[Any, Optional[Any], Optional[Any]]: model_params = llama_cpp.llama_model_default_params() @@ -11488,10 +11628,17 @@ def build_model_params( model_params.tensor_split = tensor_split_ref if vocab_only is not None: model_params.vocab_only = vocab_only - if use_mmap is not None: - model_params.use_mmap = use_mmap - if use_mlock is not None: - model_params.use_mlock = use_mlock + model_params.load_mtp = load_mtp + if use_mlock and use_mmap is not False: + model_params.load_mode = llama_cpp.LLAMA_LOAD_MODE_MMAP_MLOCK + elif use_mlock: + model_params.load_mode = llama_cpp.LLAMA_LOAD_MODE_MLOCK + elif use_mmap is not None: + model_params.load_mode = ( + llama_cpp.LLAMA_LOAD_MODE_MMAP + if use_mmap + else llama_cpp.LLAMA_LOAD_MODE_NONE + ) kv_overrides_ref = None if kv_overrides is not None: diff --git a/llama_cpp/__init__.py b/llama_cpp/__init__.py index b72459f653..5a0a40d108 100644 --- a/llama_cpp/__init__.py +++ b/llama_cpp/__init__.py @@ -1,4 +1,4 @@ from .llama_cpp import * from .llama import * -__version__ = "0.3.30" +__version__ = "0.3.34" diff --git a/llama_cpp/_internals.py b/llama_cpp/_internals.py index b0fe94d01f..b45d34b2c4 100644 --- a/llama_cpp/_internals.py +++ b/llama_cpp/_internals.py @@ -276,6 +276,8 @@ def free_ctx(): self.ctx = None self._exit_stack.callback(free_ctx) + # The native context must be freed before its model. + self.model._exit_stack.callback(self.close) def close(self): self._exit_stack.close() @@ -784,12 +786,14 @@ def add_grammar_lazy_patterns( def add_penalties( self, + n_vocab: int, penalty_last_n: int, penalty_repeat: float, penalty_freq: float, penalty_present: float, ): sampler = llama_cpp.llama_sampler_init_penalties( + n_vocab, penalty_last_n, penalty_repeat, penalty_freq, @@ -800,7 +804,6 @@ def add_penalties( def add_dry( self, model: LlamaModel, - n_ctx_train: int, dry_multiplier: float, dry_base: float, dry_allowed_length: int, @@ -814,7 +817,6 @@ def add_dry( sampler = llama_cpp.llama_sampler_init_dry( model.vocab, - n_ctx_train, dry_multiplier, dry_base, dry_allowed_length, diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index 4a09b55ee5..14e2f8500f 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -244,8 +244,15 @@ def __init__( ) # keep a reference to the array so it is not gc'd self.model_params.tensor_split = self._c_tensor_split self.model_params.vocab_only = vocab_only - self.model_params.use_mmap = use_mmap if lora_path is None else False - self.model_params.use_mlock = use_mlock + use_mmap = use_mmap and lora_path is None + if use_mmap and use_mlock: + self.model_params.load_mode = llama_cpp.LLAMA_LOAD_MODE_MMAP_MLOCK + elif use_mlock: + self.model_params.load_mode = llama_cpp.LLAMA_LOAD_MODE_MLOCK + elif use_mmap: + self.model_params.load_mode = llama_cpp.LLAMA_LOAD_MODE_MMAP + else: + self.model_params.load_mode = llama_cpp.LLAMA_LOAD_MODE_NONE # kv_overrides is the original python dict self.kv_overrides = kv_overrides @@ -471,6 +478,8 @@ def free_lora_adapter(): self._candidates = internals.LlamaTokenDataArray(n_vocab=self._n_vocab) self.n_tokens = 0 + # Restored or truncated state must decode before sampling. + self._requires_eval = True self.input_ids: npt.NDArray[np.intc] = np.ndarray((n_ctx,), dtype=np.intc) self.scores: npt.NDArray[np.single] = np.ndarray( (n_ctx if logits_all == True else n_batch, self._n_vocab), dtype=np.single @@ -647,6 +656,7 @@ def set_seed(self, seed: int): def reset(self): """Reset the model state.""" self.n_tokens = 0 + self._requires_eval = True if self._is_recurrent or self._is_hybrid: mem = llama_cpp.llama_get_memory(self._ctx.ctx) @@ -689,6 +699,7 @@ def eval(self, tokens: Sequence[int]): pass # Update n_tokens self.n_tokens += n_tokens + self._requires_eval = False def _init_sampler( self, @@ -733,7 +744,7 @@ def apply_func(token_data_array: llama_cpp.llama_token_data_array_p): sampler.add_custom(apply_func) sampler.add_penalties( - # n_vocab=self._n_vocab, + n_vocab=self._n_vocab, # special_eos_id=self._token_eos, # linefeed_id=self._token_nl, penalty_last_n=self.last_n_tokens_size, @@ -900,41 +911,53 @@ def generate( grammar=grammar, ) + tokens = list(tokens) + # Check for kv cache prefix match if reset and self.n_tokens > 0: longest_prefix = 0 - for a, b in zip(self._input_ids, tokens[:-1]): + for a, b in zip(self._input_ids, tokens): if a == b: longest_prefix += 1 else: break - # Recurrent and hybrid models cannot rewind state; reset if needed - if ( - self._is_recurrent or self._is_hybrid - ) and longest_prefix < self.n_tokens: - longest_prefix = 0 - reset = True + prompt_consumed = longest_prefix == len(tokens) + exact_prompt_cached = self.n_tokens == len(tokens) and prompt_consumed + + # Exact cache hits can sample immediately only when the current + # logits were produced by a live decode, not restored state. + if exact_prompt_cached and not self._requires_eval: + reset = False + tokens = [] + reuse_prefix = 0 if self.verbose: print( - "Llama.generate: recurrent/hybrid model requires full state reset", + "Llama.generate: full prompt already cached, skipping reset", file=sys.stderr, ) - - if longest_prefix > 0: - if self._ctx.kv_cache_seq_rm(-1, longest_prefix, -1): + else: + # If there is no suffix to decode, replay one token to refresh + # logits after truncating to a valid prefix. + reuse_prefix = longest_prefix - 1 if prompt_consumed else longest_prefix + + # Prefix hits can reuse memory because the suffix decode refreshes + # logits before sampling. + if reuse_prefix > 0: + if self._ctx.kv_cache_seq_rm(-1, reuse_prefix, -1): reset = False - tokens = tokens[longest_prefix:] - self.n_tokens = longest_prefix + tokens = tokens[reuse_prefix:] + self.n_tokens = reuse_prefix + self._requires_eval = True if self.verbose: print( - f"Llama.generate: {longest_prefix} prefix-match hit, " + f"Llama.generate: {reuse_prefix} prefix-match hit, " f"remaining {len(tokens)} prompt tokens to eval", file=sys.stderr, ) elif self.verbose: print( - f"Llama.generate: {longest_prefix} prefix-match found " + f"Llama.generate: {reuse_prefix} prefix-match found " f"but partial kv removal not supported, re-evaluating full prompt", file=sys.stderr, ) @@ -948,7 +971,6 @@ def generate( # grammar.reset() sample_idx = self.n_tokens + len(tokens) - 1 - tokens = list(tokens) # Eval and sample while True: @@ -988,6 +1010,7 @@ def generate( if sample_idx < self.n_tokens and token != self._input_ids[sample_idx]: self.n_tokens = sample_idx self._ctx.kv_cache_seq_rm(-1, self.n_tokens, -1) + self._requires_eval = True break if self.draft_model is not None: @@ -2126,8 +2149,16 @@ def __getstate__(self): main_gpu=self.model_params.main_gpu, tensor_split=self.tensor_split, vocab_only=self.model_params.vocab_only, - use_mmap=self.model_params.use_mmap, - use_mlock=self.model_params.use_mlock, + use_mmap=self.model_params.load_mode + in ( + llama_cpp.LLAMA_LOAD_MODE_MMAP, + llama_cpp.LLAMA_LOAD_MODE_MMAP_MLOCK, + ), + use_mlock=self.model_params.load_mode + in ( + llama_cpp.LLAMA_LOAD_MODE_MLOCK, + llama_cpp.LLAMA_LOAD_MODE_MMAP_MLOCK, + ), kv_overrides=self.kv_overrides, # Context Params seed=self._seed, @@ -2217,6 +2248,7 @@ def load_state(self, state: LlamaState) -> None: rest[rest > 0] = 0.0 self.input_ids = state.input_ids.copy() self.n_tokens = state.n_tokens + self._requires_eval = True self._seed = state.seed state_size = state.llama_state_size LLamaStateArrayType = ctypes.c_uint8 * state_size diff --git a/llama_cpp/llama_chat_format.py b/llama_cpp/llama_chat_format.py index 0034bdae98..4f41c2eb75 100644 --- a/llama_cpp/llama_chat_format.py +++ b/llama_cpp/llama_chat_format.py @@ -2936,7 +2936,9 @@ def __call__( # Create input text structure input_text = self._mtmd_cpp.mtmd_input_text() - input_text.text = text.encode("utf-8") + input_text_bytes = text.encode("utf-8") + input_text.text = input_text_bytes + input_text.text_len = len(input_text_bytes) input_text.add_special = True input_text.parse_special = True @@ -3485,7 +3487,9 @@ def raise_exception(message: str): bitmap_cleanup.append(bitmap) input_text = self._mtmd_cpp.mtmd_input_text() - input_text.text = text.encode("utf-8") + input_text_bytes = text.encode("utf-8") + input_text.text = input_text_bytes + input_text.text_len = len(input_text_bytes) input_text.add_special = True input_text.parse_special = True diff --git a/llama_cpp/llama_cpp.py b/llama_cpp/llama_cpp.py index 21f85c81c3..c5387a3df2 100644 --- a/llama_cpp/llama_cpp.py +++ b/llama_cpp/llama_cpp.py @@ -89,7 +89,8 @@ def _warn_deprecated(symbol: str, hint: str) -> None: # GGML_TYPE_MXFP4 = 39, # GGML_TYPE_NVFP4 = 40, # GGML_TYPE_Q1_0 = 41, -# GGML_TYPE_COUNT = 42, +# GGML_TYPE_Q2_0 = 42, +# GGML_TYPE_COUNT = 43, # }; GGML_TYPE_F32 = 0 GGML_TYPE_F16 = 1 @@ -122,7 +123,8 @@ def _warn_deprecated(symbol: str, hint: str) -> None: GGML_TYPE_MXFP4 = 39 GGML_TYPE_NVFP4 = 40 GGML_TYPE_Q1_0 = 41 -GGML_TYPE_COUNT = 42 +GGML_TYPE_Q2_0 = 42 +GGML_TYPE_COUNT = 43 # from ggml-backend.h # typedef bool (*ggml_backend_sched_eval_callback)(struct ggml_tensor * t, bool ask, void * user_data); @@ -411,6 +413,7 @@ def _warn_deprecated(symbol: str, hint: str) -> None: # LLAMA_FTYPE_MOSTLY_MXFP4_MOE = 38, // except 1d tensors # LLAMA_FTYPE_MOSTLY_NVFP4 = 39, // except 1d tensors # LLAMA_FTYPE_MOSTLY_Q1_0 = 40, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q2_0 = 41, // except 1d tensors # # LLAMA_FTYPE_GUESSED = 1024, // not specified in the model file # }; @@ -452,6 +455,7 @@ def _warn_deprecated(symbol: str, hint: str) -> None: LLAMA_FTYPE_MOSTLY_MXFP4_MOE = 38 LLAMA_FTYPE_MOSTLY_NVFP4 = 39 LLAMA_FTYPE_MOSTLY_Q1_0 = 40 +LLAMA_FTYPE_MOSTLY_Q2_0 = 41 LLAMA_FTYPE_GUESSED = 1024 # enum llama_rope_scaling_type { @@ -516,6 +520,22 @@ def _warn_deprecated(symbol: str, hint: str) -> None: LLAMA_SPLIT_MODE_TENSOR = 3 +# enum llama_load_mode { +# LLAMA_LOAD_MODE_AUTO = -1, // auto-detect based on device capabilities +# LLAMA_LOAD_MODE_NONE = 0, // no special loading mode +# LLAMA_LOAD_MODE_MMAP = 1, // memory map the model +# LLAMA_LOAD_MODE_MLOCK = 2, // force system to keep model in RAM rather than swapping or compressing +# LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing +# LLAMA_LOAD_MODE_DIRECT_IO = 4, // use direct I/O if available +# }; +LLAMA_LOAD_MODE_AUTO = -1 +LLAMA_LOAD_MODE_NONE = 0 +LLAMA_LOAD_MODE_MMAP = 1 +LLAMA_LOAD_MODE_MLOCK = 2 +LLAMA_LOAD_MODE_MMAP_MLOCK = 3 +LLAMA_LOAD_MODE_DIRECT_IO = 4 + + # enum llama_context_type { # LLAMA_CONTEXT_TYPE_DEFAULT = 0, # LLAMA_CONTEXT_TYPE_MTP = 1, @@ -785,8 +805,9 @@ class llama_model_imatrix_data(ctypes.Structure): # // NULL-terminated list of buffer types to use for tensors that match a pattern # const struct llama_model_tensor_buft_override * tensor_buft_overrides; -# int32_t n_gpu_layers; // number of layers to store in VRAM +# int32_t n_gpu_layers; // number of layers to store in VRAM, a negative value means all layers # enum llama_split_mode split_mode; // how to split the model across multiple GPUs +# enum llama_load_mode load_mode; // how to load the model # // the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE # int32_t main_gpu; @@ -808,13 +829,11 @@ class llama_model_imatrix_data(ctypes.Structure): # // Keep the booleans together to avoid misalignment during copy-by-value. # bool vocab_only; // only load the vocabulary, no weights -# bool use_mmap; // use mmap if possible -# bool use_direct_io; // use direct io, takes precedence over use_mmap when supported -# bool use_mlock; // force system to keep model in RAM # bool check_tensors; // validate model tensor data # bool use_extra_bufts; // use extra buffer types (used for weight repacking) # bool no_host; // bypass host buffer allowing extra buffers to be used # bool no_alloc; // only load metadata and simulate memory allocations +# bool load_mtp; // whether to load MTP layers # }; class llama_model_params(ctypes.Structure): """Parameters for llama_model @@ -822,21 +841,20 @@ class llama_model_params(ctypes.Structure): Attributes: devices (ctypes.Array[ggml_backend_dev_t]): NULL-terminated list of devices to use for offloading (if NULL, all available devices are used) tensor_buft_overrides (ctypes.Array[llama_model_tensor_buft_override]): NULL-terminated list of buffer types to use for tensors that match a pattern - n_gpu_layers (int): number of layers to store in VRAM + n_gpu_layers (int): number of layers to store in VRAM, a negative value means all layers split_mode (int): how to split the model across multiple GPUs + load_mode (int): how to load the model main_gpu (int): the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE tensor_split (ctypes.Array[ctypes.ctypes.c_float]): proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() progress_callback (llama_progress_callback): called with a progress value between 0.0 and 1.0. Pass NULL to disable. If the provided progress_callback returns true, model loading continues. If it returns false, model loading is immediately aborted. progress_callback_user_data (ctypes.ctypes.c_void_p): context pointer passed to the progress callback kv_overrides (ctypes.Array[llama_model_kv_override]): override key-value pairs of the model meta data vocab_only (bool): only load the vocabulary, no weights - use_mmap (bool): use mmap if possible - use_direct_io (bool): use direct io, takes precedence over use_mmap when supported - use_mlock (bool): force system to keep model in RAM check_tensors (bool): validate model tensor data use_extra_bufts (bool): use extra buffer types (used for weight repacking) no_host (bool): bypass host buffer allowing extra buffers to be used - no_alloc (bool): only load metadata and simulate memory allocations""" + no_alloc (bool): only load metadata and simulate memory allocations + load_mtp (bool): whether to load MTP layers""" if TYPE_CHECKING: devices: CtypesArray[ctypes.c_void_p] # NOTE: unused @@ -845,38 +863,36 @@ class llama_model_params(ctypes.Structure): ] # NOTE: unused n_gpu_layers: int split_mode: int + load_mode: int main_gpu: int tensor_split: CtypesArray[ctypes.c_float] progress_callback: Callable[[float, ctypes.c_void_p], bool] progress_callback_user_data: ctypes.c_void_p kv_overrides: CtypesArray[llama_model_kv_override] vocab_only: bool - use_mmap: bool - use_direct_io: bool - use_mlock: bool check_tensors: bool use_extra_bufts: bool no_host: bool no_alloc: bool + load_mtp: bool _fields_ = [ ("devices", ctypes.c_void_p), # NOTE: unnused ("tensor_buft_overrides", ctypes.c_void_p), # NOTE: unused ("n_gpu_layers", ctypes.c_int32), ("split_mode", ctypes.c_int), + ("load_mode", ctypes.c_int), ("main_gpu", ctypes.c_int32), ("tensor_split", ctypes.POINTER(ctypes.c_float)), ("progress_callback", llama_progress_callback), ("progress_callback_user_data", ctypes.c_void_p), ("kv_overrides", ctypes.POINTER(llama_model_kv_override)), ("vocab_only", ctypes.c_bool), - ("use_mmap", ctypes.c_bool), - ("use_direct_io", ctypes.c_bool), - ("use_mlock", ctypes.c_bool), ("check_tensors", ctypes.c_bool), ("use_extra_bufts", ctypes.c_bool), ("no_host", ctypes.c_bool), ("no_alloc", ctypes.c_bool), + ("load_mtp", ctypes.c_bool), ] @@ -898,14 +914,15 @@ class llama_sampler_seq_config(ctypes.Structure): # // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations # // https://github.com/ggml-org/llama.cpp/pull/7544 # struct llama_context_params { -# uint32_t n_ctx; // text context, 0 = from model -# uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode -# uint32_t n_ubatch; // physical maximum batch size -# uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models) -# uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL] -# uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch) -# int32_t n_threads; // number of threads to use for generation -# int32_t n_threads_batch; // number of threads to use for batch processing +# uint32_t n_ctx; // text context, 0 = from model +# uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode +# uint32_t n_ubatch; // physical maximum batch size +# uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models) +# uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL] +# uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch) +# uint32_t n_outputs_max_per_seq; // max outputs per sequence (0 = n_outputs_max) +# int32_t n_threads; // number of threads to use for generation +# int32_t n_threads_batch; // number of threads to use for batch processing # enum llama_context_type ctx_type; // set the context type (e.g. MTP) # enum llama_rope_scaling_type rope_scaling_type; // RoPE scaling type, from `enum llama_rope_scaling_type` @@ -964,6 +981,7 @@ class llama_context_params(ctypes.Structure): n_seq_max (int): max number of sequences (i.e. distinct states for recurrent models) n_rs_seq (int): number of recurrent-state snapshots per sequence for rollback n_outputs_max (int): max outputs in a ubatch, 0 = n_batch + n_outputs_max_per_seq (int): max outputs per sequence, 0 = n_outputs_max n_threads (int): number of threads to use for generation n_threads_batch (int): number of threads to use for batch processing ctx_type (int): context type, from `enum llama_context_type` @@ -1003,6 +1021,7 @@ class llama_context_params(ctypes.Structure): n_seq_max: int n_rs_seq: int n_outputs_max: int + n_outputs_max_per_seq: int n_threads: int n_threads_batch: int ctx_type: int @@ -1041,6 +1060,7 @@ class llama_context_params(ctypes.Structure): ("n_seq_max", ctypes.c_uint32), ("n_rs_seq", ctypes.c_uint32), ("n_outputs_max", ctypes.c_uint32), + ("n_outputs_max_per_seq", ctypes.c_uint32), ("n_threads", ctypes.c_int32), ("n_threads_batch", ctypes.c_int32), ("ctx_type", ctypes.c_int), @@ -1220,6 +1240,13 @@ class llama_chat_message(ctypes.Structure): llama_adapter_lora_p_ctypes = ctypes.POINTER(ctypes.c_void_p) +# LLAMA_API const char * llama_version(void); +@ctypes_function("llama_version", [], ctypes.c_char_p) +def llama_version() -> bytes: + """Get the llama.cpp version.""" + ... + + # // Helpers for getting default parameters # LLAMA_API struct llama_model_params llama_model_default_params(void); @ctypes_function( @@ -1272,6 +1299,28 @@ def llama_flash_attn_type_name(flash_attn_type: int, /) -> Optional[bytes]: ... +# LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode); +@ctypes_function("llama_load_mode_name", [ctypes.c_int], ctypes.c_char_p) +def llama_load_mode_name(load_mode: int, /) -> Optional[bytes]: + """Get the model load mode name.""" + ... + + +# LLAMA_API enum llama_load_mode llama_load_mode_from_str(const char * str); +@ctypes_function("llama_load_mode_from_str", [ctypes.c_char_p], ctypes.c_int) +def llama_load_mode_from_str(value: bytes, /) -> int: + """Get the model load mode from a string.""" + ... + + +# // Get the model file type (quantization) as a string, e.g. "Q8_0" or "Q4_K - Medium" +# LLAMA_API const char * llama_ftype_name(enum llama_ftype ftype); +@ctypes_function("llama_ftype_name", [ctypes.c_int], ctypes.c_char_p) +def llama_ftype_name(ftype: int, /) -> Optional[bytes]: + '''Get the model file type (quantization) as a string, e.g. "Q8_0" or "Q4_K - Medium"''' + ... + + # // Initialize the llama + ggml backend # // If numa is true, use NUMA optimizations # // Call once at the start of the program @@ -1744,6 +1793,11 @@ def llama_model_n_embd_out(model: llama_model_p, /) -> int: def llama_model_n_layer(model: llama_model_p, /) -> int: ... +# LLAMA_API int32_t llama_model_n_layer_nextn(const struct llama_model * model); +@ctypes_function("llama_model_n_layer_nextn", [llama_model_p_ctypes], ctypes.c_int32) +def llama_model_n_layer_nextn(model: llama_model_p, /) -> int: ... + + # LLAMA_API int32_t llama_model_n_head (const struct llama_model * model); @ctypes_function("llama_model_n_head", [llama_model_p_ctypes], ctypes.c_int32) def llama_model_n_head(model: llama_model_p, /) -> int: ... @@ -1772,7 +1826,8 @@ def llama_model_rope_freq_scale_train(model: llama_model_p, /) -> float: ... # LLAMA_API uint32_t llama_model_n_cls_out(const struct llama_model * model); @ctypes_function("llama_model_n_cls_out", [llama_model_p_ctypes], ctypes.c_uint32) def llama_model_n_cls_out(model: llama_model_p, /) -> int: - """Returns the number of classifier outputs (only valid for classifier models)""" + """Returns the number of classifier outputs (only valid for classifier models) + Undefined behavior for non-classifier models""" ... @@ -1838,7 +1893,7 @@ def llama_model_meta_count(model: llama_model_p, /) -> int: # LLAMA_API const char * llama_model_meta_key_str(enum llama_model_meta_key key); @ctypes_function("llama_model_meta_key_str", [ctypes.c_int], ctypes.c_char_p) def llama_model_meta_key_str(key: int, /) -> Optional[bytes]: - """Get sampling metadata key name. Returns None if the key is invalid.""" + """Get sampling metadata key name. Returns None if the key is invalid""" ... @@ -1905,6 +1960,14 @@ def llama_model_desc( ... +# // Get the model file type (quantization), e.g. LLAMA_FTYPE_MOSTLY_Q8_0 +# LLAMA_API enum llama_ftype llama_model_ftype(const struct llama_model * model); +@ctypes_function("llama_model_ftype", [llama_model_p_ctypes], ctypes.c_int) +def llama_model_ftype(model: llama_model_p, /) -> int: + """Get the model file type (quantization), e.g. LLAMA_FTYPE_MOSTLY_Q8_0""" + ... + + # // Returns the total size of all the tensors in the model in bytes # LLAMA_API uint64_t llama_model_size(const struct llama_model * model); @ctypes_function("llama_model_size", [llama_model_p_ctypes], ctypes.c_uint64) @@ -2480,7 +2543,9 @@ def llama_memory_can_shift(mem: llama_memory_t, /) -> bool: # LLAMA_API size_t llama_state_get_size(struct llama_context * ctx); @ctypes_function("llama_state_get_size", [llama_context_p_ctypes], ctypes.c_size_t) def llama_state_get_size(ctx: llama_context_p, /) -> int: - """Returns the *actual* size in bytes of the state (logits, embedding and memory)""" + """Returns the *actual* size in bytes of the state + (logits, embedding and memory) + Only use when saving the state, not when restoring it, otherwise the size may be too small.""" ... @@ -2839,6 +2904,7 @@ def llama_state_seq_save_file( ) -> int: ... +# // If tokens_out is None, only the token count is reported through n_token_count_out and no state is loaded # LLAMA_API size_t llama_state_seq_load_file( # struct llama_context * ctx, # const char * filepath, @@ -3041,9 +3107,12 @@ def llama_batch_free(batch: llama_batch, /): # struct llama_batch batch); @ctypes_function("llama_encode", [llama_context_p_ctypes, llama_batch], ctypes.c_int32) def llama_encode(ctx: llama_context_p, batch: llama_batch, /) -> int: - """Process a batch of tokens using the encoder. + """Process a batch of tokens. + In contrast to llama_decode() - this call does not use KV cache. + For encode-decoder contexts, processes the batch using the encoder. + Can store the encoder output internally for later use by the decoder's cross-attention layers. 0 - success - < 0 - error""" + < 0 - error. the memory state is restored to the state before this call""" ... @@ -3065,9 +3134,15 @@ def llama_encode(ctx: llama_context_p, batch: llama_batch, /) -> int: @ctypes_function("llama_decode", [llama_context_p_ctypes, llama_batch], ctypes.c_int32) def llama_decode(ctx: llama_context_p, batch: llama_batch, /) -> int: """Process a batch of tokens. + Requires the context to have a memory. + For encode-decoder contexts, processes the batch using the decoder. + Positive return values does not mean a fatal error, but rather a warning. + Upon fatal-error or abort, the ubatches that managed to be been processed will remain in the memory state of the context + To handle this correctly, query the memory state using llama_memory_seq_pos_min() and llama_memory_seq_pos_max() + Upon other return values, the memory state is restored to the state before this call 0 - success 1 - could not find a KV slot for the batch (try reducing the size of the batch or increase the context) - 2 - aborted (processed ubatches will remain in the context's memory) + 2 - aborted (processed ubatches will remain in the context's memory) -1 - invalid input batch < -1 - fatal error (processed ubatches will remain in the context's memory)""" ... @@ -3103,7 +3178,7 @@ def llama_set_n_threads( # LLAMA_API int32_t llama_n_threads(struct llama_context * ctx); @ctypes_function("llama_n_threads", [llama_context_p_ctypes], ctypes.c_int32) def llama_n_threads(ctx: llama_context_p, /) -> int: - """Get the number of threads used for generation of a single token""" + """Get the number of threads used for generation of a single token.""" ... @@ -3111,7 +3186,7 @@ def llama_n_threads(ctx: llama_context_p, /) -> int: # LLAMA_API int32_t llama_n_threads_batch(struct llama_context * ctx); @ctypes_function("llama_n_threads_batch", [llama_context_p_ctypes], ctypes.c_int32) def llama_n_threads_batch(ctx: llama_context_p, /) -> int: - """Get the number of threads used for prompt and batch processing (multiple token)""" + """Get the number of threads used for prompt and batch processing (multiple token).""" ... @@ -3120,7 +3195,8 @@ def llama_n_threads_batch(ctx: llama_context_p, /) -> int: # LLAMA_API void llama_set_embeddings(struct llama_context * ctx, bool embeddings); @ctypes_function("llama_set_embeddings", [llama_context_p_ctypes, ctypes.c_bool], None) def llama_set_embeddings(ctx: llama_context_p, embeddings: bool, /): - """Set whether the context outputs embeddings or not""" + """Set whether the context outputs embeddings or not + TODO: rename to avoid confusion with llama_get_embeddings()""" ... @@ -3269,6 +3345,9 @@ def llama_get_embeddings_seq( # // Get the backend sampled token for the ith token. +# // With multiple outputs, sampler state advances when the token is accepted, +# // not when it is read through this function. +# // When accepting multiple outputs, accept a contiguous prefix in output order. # // Returns LLAMA_TOKEN_NULL if no token was sampled. # LLAMA_API llama_token llama_get_sampled_token_ith(struct llama_context * ctx, int32_t i); @ctypes_function( @@ -3493,6 +3572,20 @@ def llama_vocab_get_add_eos(vocab: llama_vocab_p, /) -> bool: ... def llama_vocab_get_add_sep(vocab: llama_vocab_p, /) -> bool: ... +# // model-specific suppress tokens (gguf key: tokenizer.ggml.suppress_tokens) +# LLAMA_API const llama_token * llama_vocab_get_suppress_tokens(const struct llama_vocab * vocab, int32_t * n_suppress_tokens); +@ctypes_function( + "llama_vocab_get_suppress_tokens", + [llama_vocab_p_ctypes, ctypes.POINTER(ctypes.c_int32)], + ctypes.POINTER(llama_token), +) +def llama_vocab_get_suppress_tokens( + vocab: llama_vocab_p, + n_suppress_tokens: CtypesPointer[ctypes.c_int32], + /, +) -> Optional[CtypesPointer[llama_token]]: ... + + # LLAMA_API llama_token llama_vocab_fim_pre(const struct llama_vocab * vocab); @ctypes_function( "llama_vocab_fim_pre", @@ -4241,7 +4334,7 @@ class llama_sampler(ctypes.Structure): llama_sampler_i_clone = ctypes.CFUNCTYPE(llama_sampler_p_ctypes, llama_sampler_p_ctypes) llama_sampler_i_free = ctypes.CFUNCTYPE(None, llama_sampler_p_ctypes) llama_sampler_i_backend_init = ctypes.CFUNCTYPE( - ctypes.c_bool, llama_sampler_p_ctypes, ctypes.c_void_p + ctypes.c_bool, llama_sampler_p_ctypes, ctypes.c_void_p, ctypes.c_uint32 ) llama_sampler_i_backend_accept = ctypes.CFUNCTYPE( None, @@ -4258,6 +4351,10 @@ class llama_sampler(ctypes.Structure): ctypes.POINTER(llama_sampler_data), ) llama_sampler_i_backend_set_input = ctypes.CFUNCTYPE(None, llama_sampler_p_ctypes) +llama_sampler_i_backend_reset = ctypes.CFUNCTYPE(None, llama_sampler_p_ctypes) +llama_sampler_i_copy_state = ctypes.CFUNCTYPE( + None, llama_sampler_p_ctypes, llama_sampler_p_ctypes +) llama_sampler_i._fields_ = [ ("name", llama_sampler_i_name), @@ -4270,6 +4367,8 @@ class llama_sampler(ctypes.Structure): ("backend_accept", llama_sampler_i_backend_accept), ("backend_apply", llama_sampler_i_backend_apply), ("backend_set_input", llama_sampler_i_backend_set_input), + ("backend_reset", llama_sampler_i_backend_reset), + ("copy_state", llama_sampler_i_copy_state), ] @@ -4346,6 +4445,15 @@ def llama_sampler_reset(smpl: llama_sampler_p, /): ... def llama_sampler_clone(smpl: llama_sampler_p, /) -> llama_sampler_p: ... +# LLAMA_API void llama_sampler_copy (const struct llama_sampler * src, struct llama_sampler * dst); +@ctypes_function( + "llama_sampler_copy", + [llama_sampler_p_ctypes, llama_sampler_p_ctypes], + None, +) +def llama_sampler_copy(src: llama_sampler_p, dst: llama_sampler_p, /): ... + + # // important: do not free if the sampler has been added to a llama_sampler_chain (via llama_sampler_chain_add) # LLAMA_API void llama_sampler_free ( struct llama_sampler * smpl); @ctypes_function( @@ -4648,16 +4756,24 @@ def llama_sampler_init_grammar_lazy_patterns( # /// NOTE: Avoid using on the full vocabulary as searching for repeated tokens can become slow. For example, apply top-k or top-p sampling first. # LLAMA_API struct llama_sampler * llama_sampler_init_penalties( -# int32_t penalty_last_n, // last n tokens to penalize (0 = disable penalty, -1 = context size) -# float penalty_repeat, // 1.0 = disabled -# float penalty_freq, // 0.0 = disabled -# float penalty_present); // 0.0 = disabled +# int32_t n_vocab, +# int32_t penalty_last_n, // last n tokens to penalize (0 = disable penalty) +# float penalty_repeat, // must be > 0.0, 1.0 = disabled +# float penalty_freq, // must be finite, 0.0 = disabled +# float penalty_present); // must be finite, 0.0 = disabled @ctypes_function( "llama_sampler_init_penalties", - [ctypes.c_int32, ctypes.c_float, ctypes.c_float, ctypes.c_float], + [ + ctypes.c_int32, + ctypes.c_int32, + ctypes.c_float, + ctypes.c_float, + ctypes.c_float, + ], llama_sampler_p_ctypes, ) def llama_sampler_init_penalties( + n_vocab: int, penalty_last_n: int, penalty_repeat: float, penalty_freq: float, @@ -4669,18 +4785,16 @@ def llama_sampler_init_penalties( # /// @details DRY sampler, designed by p-e-w, as described in: https://github.com/oobabooga/text-generation-webui/pull/5677, porting Koboldcpp implementation authored by pi6am: https://github.com/LostRuins/koboldcpp/pull/982 # LLAMA_API struct llama_sampler * llama_sampler_init_dry( # const struct llama_vocab * vocab, -# int32_t n_ctx_train, # float dry_multiplier, # float dry_base, # int32_t dry_allowed_length, -# int32_t dry_penalty_last_n, +# int32_t dry_penalty_last_n, // last n tokens to penalize (0 = disable penalty) # const char ** seq_breakers, # size_t num_breakers); @ctypes_function( "llama_sampler_init_dry", [ llama_vocab_p_ctypes, - ctypes.c_int32, ctypes.c_float, ctypes.c_float, ctypes.c_int32, @@ -4692,7 +4806,6 @@ def llama_sampler_init_penalties( ) def llama_sampler_init_dry( vocab: llama_vocab_p, - n_ctx_train: int, dry_multiplier: float, dry_base: float, dry_allowed_length: int, @@ -4754,6 +4867,7 @@ def llama_sampler_get_seed(smpl: llama_sampler_p, /) -> int: ... # /// @details Sample and accept a token from the idx-th output of the last evaluation +# // For multiple outputs from one sampler, call this function in output order without gaps. # LLAMA_API llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_context * ctx, int32_t idx); @ctypes_function( "llama_sampler_sample", diff --git a/llama_cpp/llama_cpp_ext.py b/llama_cpp/llama_cpp_ext.py index 284811086a..a4b424eb63 100644 --- a/llama_cpp/llama_cpp_ext.py +++ b/llama_cpp/llama_cpp_ext.py @@ -62,6 +62,25 @@ def llama_set_embeddings_nextn( ... +# LLAMA_API void llama_set_nextn_layer_offset(struct llama_context * ctx, int32_t offset); +@_ctypes_function_from_names( + ( + "llama_set_nextn_layer_offset", + "_Z28llama_set_nextn_layer_offsetP13llama_contexti", + "?llama_set_nextn_layer_offset@@YAXPEAUllama_context@@H@Z", + ), + [llama_cpp.llama_context_p_ctypes, ctypes.c_int32], + None, +) +def llama_set_nextn_layer_offset( + ctx: llama_cpp.llama_context_p, + offset: Union[ctypes.c_int32, int], + /, +): + """Select which appended NextN block the decoder MTP graph runs.""" + ... + + # LLAMA_API float * llama_get_embeddings_nextn(struct llama_context * ctx); @_ctypes_function_from_names( ( diff --git a/llama_cpp/mtmd_cpp.py b/llama_cpp/mtmd_cpp.py index 78f068aa9a..79f417cbaf 100644 --- a/llama_cpp/mtmd_cpp.py +++ b/llama_cpp/mtmd_cpp.py @@ -5,8 +5,10 @@ from ctypes import ( CFUNCTYPE, c_bool, + c_char, c_char_p, c_int, + c_int32, c_int64, c_uint8, c_uint32, @@ -20,6 +22,7 @@ ) import pathlib from typing import ( + Callable, Union, NewType, Optional, @@ -67,6 +70,9 @@ mtmd_helper_video_p = NewType("mtmd_helper_video_p", int) mtmd_helper_video_p_ctypes = c_void_p +mtmd_helper_gen_audio_p = NewType("mtmd_helper_gen_audio_p", int) +mtmd_helper_gen_audio_p_ctypes = c_void_p + mtmd_image_tokens_p = NewType("mtmd_image_tokens_p", int) mtmd_image_tokens_p_ctypes = c_void_p @@ -83,6 +89,19 @@ MTMD_INPUT_CHUNK_TYPE_TEXT = 0 MTMD_INPUT_CHUNK_TYPE_IMAGE = 1 MTMD_INPUT_CHUNK_TYPE_AUDIO = 2 +MTMD_INPUT_CHUNK_TYPE_COUNT = 3 + +MTMD_GEN_AUDIO_TYPE_NONE = 0 +MTMD_GEN_AUDIO_TYPE_QWEN3TTS = 1 +MTMD_GEN_AUDIO_TYPE_POCKETTTS = 2 + +MTMD_GEN_PROCESS_TYPE_GEN_CODE = 0 +MTMD_GEN_PROCESS_TYPE_GEN_WAV = 1 + +MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM = 0 +MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV = 1 + +mtmd_progress_callback = CFUNCTYPE(c_bool, c_float, c_void_p) # Structures @@ -106,6 +125,8 @@ class mtmd_context_params(Structure): cb_eval: llama_cpp.ggml_backend_sched_eval_callback cb_eval_user_data: c_void_p batch_max_tokens: int + progress_callback: Callable[[float, c_void_p], bool] + progress_callback_user_data: c_void_p _fields_ = [ ("use_gpu", c_bool), @@ -120,14 +141,23 @@ class mtmd_context_params(Structure): ("cb_eval", llama_cpp.ggml_backend_sched_eval_callback), ("cb_eval_user_data", c_void_p), ("batch_max_tokens", c_int), + ("progress_callback", mtmd_progress_callback), + ("progress_callback_user_data", c_void_p), ] class mtmd_input_text(Structure): """Text input passed to `mtmd_tokenize`.""" + if TYPE_CHECKING: + text: Optional[bytes] + text_len: int + add_special: bool + parse_special: bool + _fields_ = [ ("text", c_char_p), + ("text_len", c_size_t), ("add_special", c_bool), ("parse_special", c_bool), ] @@ -161,6 +191,122 @@ class mtmd_caps(Structure): ] +# struct mtmd_gen_audio_info { +# enum mtmd_gen_audio_type type; +# int32_t sample_rate; // in Hz, for example 24000 for qwen3tts +# const char * model_variant; // name of the weight variant, can be None if not applicable +# }; +class mtmd_gen_audio_info(Structure): + if TYPE_CHECKING: + type: int + sample_rate: int + model_variant: Optional[bytes] + + _fields_ = [ + ("type", c_int), + ("sample_rate", c_int32), + ("model_variant", c_char_p), + ] + + +# struct mtmd_gen_inp { +# enum mtmd_gen_process_type type; +# +# // for MTMD_GEN_PROCESS_TYPE_GEN_CODE +# int32_t code0; // the sampled codebook 0 entry from backbone +# float * embd; // the hidden state from backbone, must have n_text_embd elements +# int32_t top_k; +# float top_p; +# uint32_t seed; // UINT32_MAX for random +# float temp; // sampling temperature, or noise scale for flow-matching decoders +# +# // for MTMD_GEN_PROCESS_TYPE_GEN_WAV +# // pass either codes (discrete) or feats (continuous), depending on the pipeline +# int32_t * codes; +# size_t n_codes; +# const float * feats; +# size_t n_feats; +# const char * state_data; +# size_t state_size; +# }; +class mtmd_gen_inp(Structure): + if TYPE_CHECKING: + type: int + code0: int + embd: Optional["_Pointer[c_float]"] + top_k: int + top_p: float + seed: int + temp: float + codes: Optional["_Pointer[c_int32]"] + n_codes: int + feats: Optional["_Pointer[c_float]"] + n_feats: int + state_data: Optional["_Pointer[c_char]"] + state_size: int + + _fields_ = [ + ("type", c_int), + ("code0", c_int32), + ("embd", POINTER(c_float)), + ("top_k", c_int32), + ("top_p", c_float), + ("seed", c_uint32), + ("temp", c_float), + ("codes", POINTER(c_int32)), + ("n_codes", c_size_t), + ("feats", POINTER(c_float)), + ("n_feats", c_size_t), + ("state_data", POINTER(c_char)), + ("state_size", c_size_t), + ] + + +# struct mtmd_gen_out { +# // note: output memory is allocated by the context, valid until next process() call +# +# // for MTMD_GEN_PROCESS_TYPE_GEN_CODE +# const int32_t * codes; +# size_t n_codes; +# const float * feats; // continuous counterpart of codes +# size_t n_feats; +# const float * embd; // the generated hidden state, to be fed back to backbone +# // it must have n_text_embd elements +# bool is_eos; // only set by pipelines having the EOS head inside mmproj +# +# // for MTMD_GEN_PROCESS_TYPE_GEN_WAV +# const float * audio; +# size_t n_samples; +# const char * state_data; +# size_t state_size; +# }; +class mtmd_gen_out(Structure): + if TYPE_CHECKING: + codes: Optional["_Pointer[c_int32]"] + n_codes: int + feats: Optional["_Pointer[c_float]"] + n_feats: int + embd: Optional["_Pointer[c_float]"] + is_eos: bool + audio: Optional["_Pointer[c_float]"] + n_samples: int + state_data: Optional["_Pointer[c_char]"] + state_size: int + + _fields_ = [ + ("codes", POINTER(c_int32)), + ("n_codes", c_size_t), + ("feats", POINTER(c_float)), + ("n_feats", c_size_t), + ("embd", POINTER(c_float)), + ("is_eos", c_bool), + ("audio", POINTER(c_float)), + ("n_samples", c_size_t), + ("state_data", POINTER(c_char)), + ("state_size", c_size_t), + ] + + mtmd_bitmap_lazy_callback = CFUNCTYPE( c_int, c_size_t, @@ -221,6 +367,46 @@ class mtmd_helper_video_init_params(Structure): ] +# struct mtmd_helper_gen_audio_inp { +# llama_seq_id seq_id; +# +# const char * prompt; +# size_t prompt_len; +# +# mtmd_bitmap * speaker_ref; // optional, can be NULL +# const char * lang; // optional, can be NULL +# +# int32_t top_k; +# float top_p; +# uint32_t seed; // UINT32_MAX for random (default: random) +# +# enum mtmd_helper_gen_audio_outtype out_type; +# }; +class mtmd_helper_gen_audio_inp(Structure): + if TYPE_CHECKING: + seq_id: int + prompt: Optional[bytes] + prompt_len: int + speaker_ref: Optional[mtmd_bitmap_p] + lang: Optional[bytes] + top_k: int + top_p: float + seed: int + out_type: int + + _fields_ = [ + ("seq_id", llama_cpp.llama_seq_id), + ("prompt", c_char_p), + ("prompt_len", c_size_t), + ("speaker_ref", mtmd_bitmap_p_ctypes), + ("lang", c_char_p), + ("top_k", c_int32), + ("top_p", c_float), + ("seed", c_uint32), + ("out_type", c_int), + ] + + ################################################ # mtmd.h functions ################################################ @@ -540,6 +726,44 @@ def mtmd_input_chunk_free(chunk: mtmd_input_chunk_p, /): ... +# // save/load an input chunk to/from a buffer (useful for KV save/load) +# // important: only chunk's metadata will be saved, the actual image/audio data will not be saved +# // the loaded chunk will always be a placeholder, cannot be used for mtmd_encode() or mtmd_batch_encode() +# // out_buf can be nullptr (to query expected_out_len) +# // returns 0 on success, non-zero on failure +# MTMD_API int32_t mtmd_input_chunk_save(const mtmd_input_chunk * chunk, char * out_buf, size_t out_len, size_t * expected_out_len); +@ctypes_function( + "mtmd_input_chunk_save", + [mtmd_input_chunk_p_ctypes, POINTER(c_char), c_size_t, POINTER(c_size_t)], + c_int32, +) +def mtmd_input_chunk_save( + chunk: mtmd_input_chunk_p, + out_buf: Optional[CtypesArray[c_char]], + out_len: Union[c_size_t, int], + expected_out_len: "_Pointer[c_size_t]", + /, +) -> int: + """Save an input chunk's metadata to a buffer.""" + ... + + +# // returns nullptr on failure +# MTMD_API mtmd_input_chunk * mtmd_input_chunk_load(const char * buf, size_t len); +@ctypes_function( + "mtmd_input_chunk_load", + [c_char_p, c_size_t], + mtmd_input_chunk_p_ctypes, +) +def mtmd_input_chunk_load( + buf: bytes, + length: Union[c_size_t, int], + /, +) -> Optional[mtmd_input_chunk_p]: + """Load an input chunk placeholder from saved metadata.""" + ... + + # MTMD_API size_t mtmd_image_tokens_get_n_tokens(const mtmd_image_tokens * image_tokens); @ctypes_function( "mtmd_image_tokens_get_n_tokens", [mtmd_image_tokens_p_ctypes], c_size_t @@ -685,6 +909,48 @@ def mtmd_get_cap_from_file(mmproj_fname: bytes, /) -> mtmd_caps: ... +# MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx); +@ctypes_function( + "mtmd_gen_audio_get_info", + [mtmd_context_p_ctypes], + mtmd_gen_audio_info, +) +def mtmd_gen_audio_get_info(ctx: mtmd_context_p, /) -> mtmd_gen_audio_info: + """Get audio generation information for an MTMD context.""" + ... + + +# // defaults tuned for the loaded pipeline, callers override only what they care about +# MTMD_API struct mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx); +@ctypes_function( + "mtmd_gen_inp_default", + [mtmd_context_p_ctypes], + mtmd_gen_inp, +) +def mtmd_gen_inp_default(ctx: mtmd_context_p, /) -> mtmd_gen_inp: + """Get default audio generation input parameters for an MTMD context.""" + ... + + +# // note: this API is stateless, caller must handle state management and audio frame accumulation +# MTMD_API int32_t mtmd_gen_audio_process(mtmd_context * ctx, +# const struct mtmd_gen_inp * inp, +# struct mtmd_gen_out * out); +@ctypes_function( + "mtmd_gen_audio_process", + [mtmd_context_p_ctypes, POINTER(mtmd_gen_inp), POINTER(mtmd_gen_out)], + c_int32, +) +def mtmd_gen_audio_process( + ctx: mtmd_context_p, + inp: "_Pointer[mtmd_gen_inp]", + out: "_Pointer[mtmd_gen_out]", + /, +) -> int: + """Process one audio generation step.""" + ... + + # MTMD_API mtmd_input_chunks * mtmd_test_create_input_chunks(void); @ctypes_function("mtmd_test_create_input_chunks", [], mtmd_input_chunks_p_ctypes) def mtmd_test_create_input_chunks() -> Optional[mtmd_input_chunks_p]: @@ -990,6 +1256,157 @@ def mtmd_helper_video_read_next( ... +# // return true if model can be used for chat +# MTMD_API bool mtmd_helper_model_can_chat(struct llama_context * lctx, struct mtmd_context * mctx); +@ctypes_function( + "mtmd_helper_model_can_chat", + [llama_cpp.llama_context_p_ctypes, mtmd_context_p_ctypes], + c_bool, +) +def mtmd_helper_model_can_chat( + lctx: llama_cpp.llama_context_p, + mctx: mtmd_context_p, + /, +) -> bool: + """Return whether the model can be used for chat.""" + ... + + +# MTMD_API mtmd_helper_gen_audio * mtmd_helper_gen_audio_init( +# struct llama_context * lctx, +# struct mtmd_context * mctx); +@ctypes_function( + "mtmd_helper_gen_audio_init", + [llama_cpp.llama_context_p_ctypes, mtmd_context_p_ctypes], + mtmd_helper_gen_audio_p_ctypes, +) +def mtmd_helper_gen_audio_init( + lctx: llama_cpp.llama_context_p, + mctx: mtmd_context_p, + /, +) -> Optional[mtmd_helper_gen_audio_p]: + """Initialize an audio generation helper context.""" + ... + + +# MTMD_API void mtmd_helper_gen_audio_free(mtmd_helper_gen_audio * ctx); +@ctypes_function( + "mtmd_helper_gen_audio_free", + [mtmd_helper_gen_audio_p_ctypes], + None, +) +def mtmd_helper_gen_audio_free(ctx: mtmd_helper_gen_audio_p, /): ... + + +# MTMD_API void mtmd_helper_gen_audio_reset(mtmd_helper_gen_audio * ctx); +@ctypes_function( + "mtmd_helper_gen_audio_reset", + [mtmd_helper_gen_audio_p_ctypes], + None, +) +def mtmd_helper_gen_audio_reset(ctx: mtmd_helper_gen_audio_p, /): ... + + +# MTMD_API int32_t mtmd_helper_gen_audio_set_input( +# mtmd_helper_gen_audio * ctx, +# const struct mtmd_helper_gen_audio_inp * inp); +@ctypes_function( + "mtmd_helper_gen_audio_set_input", + [mtmd_helper_gen_audio_p_ctypes, POINTER(mtmd_helper_gen_audio_inp)], + c_int32, +) +def mtmd_helper_gen_audio_set_input( + ctx: mtmd_helper_gen_audio_p, + inp: "_Pointer[mtmd_helper_gen_audio_inp]", + /, +) -> int: + """Set the audio generation helper input.""" + ... + + +# // processes at most n_batch prompt tokens per call +# // returns: >0 = number of prompt tokens remaining, 0 = done, <0 = error +# MTMD_API int32_t mtmd_helper_gen_audio_step_prompt( +# mtmd_helper_gen_audio * ctx, +# int32_t n_batch); +@ctypes_function( + "mtmd_helper_gen_audio_step_prompt", + [mtmd_helper_gen_audio_p_ctypes, c_int32], + c_int32, +) +def mtmd_helper_gen_audio_step_prompt( + ctx: mtmd_helper_gen_audio_p, + n_batch: int, + /, +) -> int: + """Process up to n_batch prompt tokens.""" + ... + + +# // generates one frame; must only be called after step_prompt() has returned 0 +# // sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token +# // out_stop (optional) is set on end-of-speech, the caller must then stop the loop +# // h_state_out is valid until next step_gen() or reset() call, None if no frame is generated +# MTMD_API int32_t mtmd_helper_gen_audio_step_gen( +# mtmd_helper_gen_audio * ctx, +# llama_token sampled, +# const float * h_state_in, +# const float ** h_state_out, +# bool * out_stop); +@ctypes_function( + "mtmd_helper_gen_audio_step_gen", + [ + mtmd_helper_gen_audio_p_ctypes, + llama_cpp.llama_token, + POINTER(c_float), + POINTER(POINTER(c_float)), + POINTER(c_bool), + ], + c_int32, +) +def mtmd_helper_gen_audio_step_gen( + ctx: mtmd_helper_gen_audio_p, + sampled: llama_cpp.llama_token, + h_state_in: Optional["_Pointer[c_float]"], + h_state_out: "_Pointer[_Pointer[c_float]]", + out_stop: Optional["_Pointer[c_bool]"], + /, +) -> int: + """Generate one audio frame.""" + ... + + +# // out_data valid until next get_output() or reset() call +# // out_n_samples (optional, can be NULL) receives the number of generated PCM samples +# MTMD_API int32_t mtmd_helper_gen_audio_get_output( +# mtmd_helper_gen_audio * ctx, +# int32_t * out_sample_rate, +# const char ** out_data, +# size_t * out_data_len, +# int64_t * out_n_samples); +@ctypes_function( + "mtmd_helper_gen_audio_get_output", + [ + mtmd_helper_gen_audio_p_ctypes, + POINTER(c_int32), + POINTER(POINTER(c_char)), + POINTER(c_size_t), + POINTER(c_int64), + ], + c_int32, +) +def mtmd_helper_gen_audio_get_output( + ctx: mtmd_helper_gen_audio_p, + out_sample_rate: "_Pointer[c_int32]", + out_data: "_Pointer[_Pointer[c_char]]", + out_data_len: "_Pointer[c_size_t]", + out_n_samples: Optional["_Pointer[c_int64]"], + /, +) -> int: + """Get accumulated PCM or WAV audio output.""" + ... + + # MTMD_API void mtmd_log_set(ggml_log_callback log_callback, void * user_data); @ctypes_function( "mtmd_log_set", diff --git a/tests/test_llama.py b/tests/test_llama.py index 336d6a6122..c1c16e30ca 100644 --- a/tests/test_llama.py +++ b/tests/test_llama.py @@ -1,4 +1,5 @@ import ctypes +import itertools import multiprocessing import numpy as np @@ -64,6 +65,14 @@ def llama_cpp_model_path(): return model_path +@pytest.fixture +def llama_cpp_transformer_model_path(): + repo_id = "ggml-org/models" + filename = "tinyllamas/stories15M-q4_0.gguf" + model_path = hf_hub_download(repo_id, filename) + return model_path + + @pytest.fixture def llama_cpp_embedding_model_path(): repo_id = "CompendiumLabs/bge-small-en-v1.5-gguf" @@ -94,8 +103,14 @@ def test_real_model(llama_cpp_model_path): assert os.path.exists(llama_cpp_model_path) params = llama_cpp.llama_model_default_params() - params.use_mmap = llama_cpp.llama_supports_mmap() - params.use_mlock = llama_cpp.llama_supports_mlock() + if llama_cpp.llama_supports_mmap() and llama_cpp.llama_supports_mlock(): + params.load_mode = llama_cpp.LLAMA_LOAD_MODE_MMAP_MLOCK + elif llama_cpp.llama_supports_mlock(): + params.load_mode = llama_cpp.LLAMA_LOAD_MODE_MLOCK + elif llama_cpp.llama_supports_mmap(): + params.load_mode = llama_cpp.LLAMA_LOAD_MODE_MMAP + else: + params.load_mode = llama_cpp.LLAMA_LOAD_MODE_NONE params.check_tensors = False model = internals.LlamaModel(path_model=llama_cpp_model_path, params=params) @@ -146,6 +161,9 @@ def test_real_model(llama_cpp_model_path): assert len(output) == 4 assert output_text + model.close() + assert context.ctx is None + def test_real_llama(llama_cpp_model_path): model = llama_cpp.Llama( @@ -339,6 +357,285 @@ def test_hybrid_model_prompt_cache_reset(llama_cpp_hybrid_model_path): ) +def _create_test_model(model_path): + return llama_cpp.Llama( + model_path, + n_ctx=64, + n_batch=64, + n_ubatch=64, + n_threads=multiprocessing.cpu_count(), + n_threads_batch=multiprocessing.cpu_count(), + logits_all=False, + verbose=False, + ) + + +def _generate_test_tokens(model, tokens, max_tokens=3): + return list( + itertools.islice( + model.generate( + tokens, + temp=0.0, + ), + max_tokens, + ) + ) + + +MODEL_CACHE_CASES = ( + ("llama_cpp_transformer_model_path", False, False), + ("llama_cpp_recurrent_model_path", True, False), + ("llama_cpp_hybrid_model_path", False, True), +) + +RESTORED_CACHE_CASES = MODEL_CACHE_CASES + + +def _eval_alternate_same_length_prompt(model, tokens, expected_next_token): + replacement_tokens = ( + model.token_eos(), + model.token_nl(), + 0, + 1, + 2, + model.n_vocab() - 1, + ) + + for replacement_token in replacement_tokens: + alternate_tokens = list(tokens) + alternate_tokens[-1] = replacement_token + if alternate_tokens == tokens: + continue + + model.reset() + model.eval(alternate_tokens) + if model.sample(temp=0.0, idx=len(tokens) - 1) != expected_next_token: + return + + raise AssertionError("failed to find an alternate same-length prompt") + + +def _assert_exact_cached_prompt_reuse_matches_fresh( + model_path, + *, + is_recurrent: bool, + is_hybrid: bool, +): + prompt = "The quick brown fox" + fresh = _create_test_model(model_path) + tokens = fresh.tokenize(prompt.encode(), add_bos=True, special=True) + + assert fresh._is_recurrent is is_recurrent + assert fresh._is_hybrid is is_hybrid + + expected_tokens = _generate_test_tokens(fresh, tokens) + + cached = _create_test_model(model_path) + assert cached._is_recurrent is is_recurrent + assert cached._is_hybrid is is_hybrid + + cached.eval(tokens) + assert cached.n_tokens == len(tokens) + assert cached.input_ids[: cached.n_tokens].tolist() == tokens + assert cached.sample(temp=0.0, idx=len(tokens) - 1) == expected_tokens[0] + + reset_calls = 0 + original_reset = cached.reset + + def reset_tracker(): + nonlocal reset_calls + reset_calls += 1 + original_reset() + + cached.reset = reset_tracker + + cached_tokens = _generate_test_tokens(cached, tokens) + assert reset_calls == 0 + assert cached_tokens == expected_tokens + assert cached.n_tokens == len(tokens) + len(cached_tokens) - 1 + + +def _assert_loaded_exact_cached_prompt_reuse_matches_fresh( + model_path, + *, + is_recurrent: bool, + is_hybrid: bool, +): + prompt = "The quick brown fox" + fresh = _create_test_model(model_path) + tokens = fresh.tokenize(prompt.encode(), add_bos=True, special=True) + expected_tokens = _generate_test_tokens(fresh, tokens) + + source = _create_test_model(model_path) + assert source._is_recurrent is is_recurrent + assert source._is_hybrid is is_hybrid + + source.eval(tokens) + state = source.save_state() + + loaded = _create_test_model(model_path) + assert loaded._is_recurrent is is_recurrent + assert loaded._is_hybrid is is_hybrid + + _eval_alternate_same_length_prompt( + loaded, + tokens, + expected_tokens[0], + ) + loaded.load_state(state) + + assert loaded.n_tokens == len(tokens) + assert loaded.input_ids[: loaded.n_tokens].tolist() == tokens + + loaded_tokens = _generate_test_tokens(loaded, tokens) + assert loaded_tokens == expected_tokens + assert loaded.n_tokens == len(tokens) + len(loaded_tokens) - 1 + + +def _assert_ram_cache_exact_prompt_hit_matches_fresh( + model_path, + *, + is_recurrent: bool, + is_hybrid: bool, +): + prompt = "The quick brown fox" + fresh = _create_test_model(model_path) + tokens = fresh.tokenize(prompt.encode(), add_bos=True, special=True) + expected = fresh.create_completion( + tokens, + max_tokens=1, + temperature=0.0, + seed=1337, + ) + + cache = llama_cpp.LlamaRAMCache() + writer = _create_test_model(model_path) + writer.set_cache(cache) + writer.create_completion( + tokens, + max_tokens=1, + temperature=0.0, + seed=1337, + ) + + cached = _create_test_model(model_path) + assert cached._is_recurrent is is_recurrent + assert cached._is_hybrid is is_hybrid + cached.set_cache(cache) + + load_state_calls = 0 + original_load_state = cached.load_state + + def load_state_tracker(state): + nonlocal load_state_calls + load_state_calls += 1 + original_load_state(state) + + cached.load_state = load_state_tracker + + actual = cached.create_completion( + tokens, + max_tokens=1, + temperature=0.0, + seed=1337, + ) + + assert load_state_calls == 1 + assert actual["choices"][0]["text"] == expected["choices"][0]["text"] + assert ( + actual["usage"]["completion_tokens"] == expected["usage"]["completion_tokens"] + ) + + +def _assert_shorter_prompt_prefix_reuse_matches_fresh( + model_path, + *, + is_recurrent: bool, + is_hybrid: bool, +): + prompt = "The quick brown fox" + history = " jumps over the lazy dog" + fresh = _create_test_model(model_path) + tokens = fresh.tokenize(prompt.encode(), add_bos=True, special=True) + history_tokens = fresh.tokenize(history.encode(), add_bos=False, special=True) + expected_tokens = _generate_test_tokens(fresh, tokens) + + cached = _create_test_model(model_path) + assert cached._is_recurrent is is_recurrent + assert cached._is_hybrid is is_hybrid + + cached.eval(tokens + history_tokens) + assert cached.n_tokens > len(tokens) + assert cached.input_ids[: len(tokens)].tolist() == tokens + + cached_tokens = _generate_test_tokens(cached, tokens) + assert cached_tokens == expected_tokens + + +@pytest.mark.parametrize( + ("model_path_fixture", "is_recurrent", "is_hybrid"), MODEL_CACHE_CASES +) +def test_exact_cached_prompt_reuse_matches_fresh( + request, + model_path_fixture, + is_recurrent, + is_hybrid, +): + _assert_exact_cached_prompt_reuse_matches_fresh( + request.getfixturevalue(model_path_fixture), + is_recurrent=is_recurrent, + is_hybrid=is_hybrid, + ) + + +@pytest.mark.parametrize( + ("model_path_fixture", "is_recurrent", "is_hybrid"), RESTORED_CACHE_CASES +) +def test_loaded_exact_cached_prompt_reuse_matches_fresh( + request, + model_path_fixture, + is_recurrent, + is_hybrid, +): + _assert_loaded_exact_cached_prompt_reuse_matches_fresh( + request.getfixturevalue(model_path_fixture), + is_recurrent=is_recurrent, + is_hybrid=is_hybrid, + ) + + +@pytest.mark.parametrize( + ("model_path_fixture", "is_recurrent", "is_hybrid"), RESTORED_CACHE_CASES +) +def test_ram_cache_exact_prompt_hit_matches_fresh( + request, + model_path_fixture, + is_recurrent, + is_hybrid, +): + _assert_ram_cache_exact_prompt_hit_matches_fresh( + request.getfixturevalue(model_path_fixture), + is_recurrent=is_recurrent, + is_hybrid=is_hybrid, + ) + + +@pytest.mark.parametrize( + ("model_path_fixture", "is_recurrent", "is_hybrid"), MODEL_CACHE_CASES +) +def test_shorter_prompt_prefix_reuse_matches_fresh( + request, + model_path_fixture, + is_recurrent, + is_hybrid, +): + _assert_shorter_prompt_prefix_reuse_matches_fresh( + request.getfixturevalue(model_path_fixture), + is_recurrent=is_recurrent, + is_hybrid=is_hybrid, + ) + + def test_real_llama_embeddings(llama_cpp_embedding_model_path): model = llama_cpp.Llama( llama_cpp_embedding_model_path, diff --git a/vendor/llama.cpp b/vendor/llama.cpp index e3a74b2990..adb55e5148 160000 --- a/vendor/llama.cpp +++ b/vendor/llama.cpp @@ -1 +1 @@ -Subproject commit e3a74b299085cd00013804f7fca2e03441b2da20 +Subproject commit adb55e5148dc93bcdca7212a2d1df3ccc422959a