Repository navigation
Expand file tree
/
Copy pathllama.py
More file actions
5184 lines (4698 loc) · 235 KB
/
Copy pathllama.py
File metadata and controls
5184 lines (4698 loc) · 235 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
from __future__ import annotations
import contextlib
import ctypes
import fnmatch
import json
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 ._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
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.
"""
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
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,
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,
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,
# Adaptive KV streaming (kv-stream fork only, MiB; None = read env LLAMA_ARG_KV_STREAM_STAGE_MIB)
kv_stream_stage_mib: 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.
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.
type_k: KV cache data type for K (default: f16)
type_v: KV cache data type for V (default: f16)
kv_stream_stage_mib: block-streaming KV resident+staging pool in MiB (kv-stream fork; 0 = off)
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()
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.
if not Llama.__backend_initialized:
with suppress_stdout_stderr(disable=verbose):
llama_cpp_lib.llama_backend_init()
# Wheels built with `GGML_BACKEND_DL` ship ggml backends as separate
# dynamic libraries under llama_cpp/lib, for example:
#
# ggml-cpu-x64.dll
# ggml-cpu-haswell.dll
# ggml-cpu-alderlake.dll
# ggml-cuda.dll
#
# With the dynamic backend layout, llama_backend_init() initializes
# the global backend system but does not necessarily register every
# packaged backend. Loading the package lib directory ensures ggml can
# discover CPU variants and optional accelerator backends before model
# loading.
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}")
# Load all dynamic ggml backend plugins from the packaged lib directory.
ggml_backend_load_all_from_path(
ctypes.c_char_p(str(lib_dir).encode("utf-8"))
)
# Print the number of backend registrations to confirm whether the DLL is loaded.
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 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()
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.lazy_mode = lazy_mode
# 本构建基于的 llama.cpp 没有懒加载字段，新版上游才有
if hasattr(self.model_params, "lazy_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.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
# Adaptive KV streaming; 未显式给值时允许用环境变量全局打开
if kv_stream_stage_mib is None:
kv_stream_stage_mib = int(os.environ.get("LLAMA_ARG_KV_STREAM_STAGE_MIB") or 0)
self.context_params.kv_stream_stage_mib = max(0, int(kv_stream_stage_mib))
# 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