forked from unslothai/unsloth
-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathllama_cpp.py
More file actions
5653 lines (5142 loc) · 247 KB
/
Copy pathllama_cpp.py
File metadata and controls
5653 lines (5142 loc) · 247 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
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
llama-server inference backend for GGUF models.
Manages a llama-server subprocess and proxies chat completions
through its OpenAI-compatible /v1/chat/completions endpoint.
"""
import atexit
import contextlib
import json
import os
import re
import struct
import structlog
from loggers import get_logger
import shutil
import socket
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Generator, Iterable, List, Optional
from urllib.parse import urlparse
import httpx
from core.tool_healing import (
_TC_END_TAG_RE,
_TC_FUNC_CLOSE_RE,
_TC_FUNC_START_RE,
_TC_JSON_START_RE,
_TC_PARAM_CLOSE_RE,
_TC_PARAM_START_RE,
_TOOL_ALL_PATS,
_TOOL_CLOSED_PATS,
parse_tool_calls_from_text,
strip_tool_call_markup,
)
from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import (
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
)
from core.inference.tool_call_parser import (
parse_tool_calls_from_text as _shared_parse_tool_calls_from_text,
)
logger = get_logger(__name__)
# ── Pre-compiled patterns for plan-without-action re-prompt ──
# Forward-looking intent signals that indicate the model is
# describing what it *will* do rather than giving a final answer.
_INTENT_SIGNAL = re.compile(
r"(?i)("
# Direct intent: "I'll ...", "I will ...", "Let me ...", "I am going to ..."
# Handles both straight and curly apostrophes.
# Excludes "I can", "I should", "I want to", "let's" which
# appear frequently in direct answers / explanations.
# Negative lookahead drops negated forms ("I will not", "I'll never")
# so a refusal doesn't trigger a re-prompt.
r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)"
r"|"
# Step/plan framing: "First ...", "Step 1:", "Here's my plan"
r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))"
r"|"
# "Now I" / "Next I" patterns
r"\b(?:now i|next i)\b"
r")"
)
_MAX_REPROMPTS = 3
# Without max_tokens, llama-server defaults to n_predict = n_ctx (up to
# 262144 for Qwen3.5), producing many-minute zombie decodes when cancel
# fails. t_max_predict_ms is a wall-clock backstop applied unconditionally,
# but the llama.cpp README notes it ONLY fires after a newline has been
# generated -- a model stuck in a long unbroken non-newline sequence is
# unbounded by it. So we still want a token cap as the front-line limiter.
#
# The cap is the model's effective context length when we know it,
# falling back to a generous floor when metadata is unavailable. 4096 was
# too low: Qwen3 / gpt-oss reasoning traces routinely exceed it, and any
# OpenAI-API caller that omits max_tokens (langchain, llama-index, raw
# curl) sees responses silently truncated mid-sentence.
_DEFAULT_MAX_TOKENS_FLOOR = 32768
_DEFAULT_T_MAX_PREDICT_MS = 600_000 # 10 min
_REPROMPT_MAX_CHARS = 2000
# ── Pre-compiled patterns for GGUF shard detection ───────────
_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$")
_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
# ── Sliding-window-pattern resolver ───────────────────────────
# Resolves the per-layer SWA mask when a GGUF reports a sliding window
# but no `sliding_window_pattern` field. Tier order in
# `_resolve_swa_pattern`: GGUF metadata, on-disk cache, bootstrap dict
# below, transformers introspection, HF Hub config.json, legacy 1/4
# fallback. Period N means layer i is SWA iff `(i + 1) % N != 0`,
# matching transformers. Skipped on purpose: phi3 (no key/val length
# in GGUF, window >= ctx anyway), qwen2 family (converter strips
# sliding_window when use_sliding_window=False), mistral v0.1/v0.2
# (all-SWA can't be expressed as a period).
_BOOTSTRAP_SWA_DEFAULTS: dict[str, int] = {
"gemma2": 2, # Gemma2Config.sliding_window_pattern
"gemma3": 6, # Gemma3TextConfig.sliding_window_pattern
"gemma3n": 5, # text_config.layer_types: SWA*4 + FULL
"gpt_oss": 2, # text_config.layer_types: alternating
"cohere2": 4, # Cohere2Config.sliding_window_pattern
}
# Process-wide cache backed by JSON on disk. Values are int period or
# list[bool] mask. Lazy-loaded.
_SWA_CACHE: Optional[dict] = None
_SWA_CACHE_LOCK = threading.Lock()
def _probe_dns_dead(host: str = "huggingface.co", timeout: float = 2.0) -> bool:
"""Quick DNS check. Runs on a daemon thread so concurrent sockets
in the same process are not affected by socket.setdefaulttimeout."""
result: list[Optional[bool]] = [None]
def _probe() -> None:
try:
socket.gethostbyname(host)
result[0] = False
except Exception:
result[0] = True
t = threading.Thread(target = _probe, daemon = True)
t.start()
t.join(timeout)
# Thread still running -> resolver wedged -> treat as dead.
return True if result[0] is None else result[0]
@contextlib.contextmanager
def _hf_offline_if_dns_dead():
"""Set HF_HUB_OFFLINE for the body of this block only when DNS to
huggingface.co fails. Restores the env on exit so a transient
resolver hiccup at the start of one load can't quarantine the whole
process. Respects an explicit user setting (no-op if already set)."""
if "HF_HUB_OFFLINE" in os.environ:
yield False
return
if not _probe_dns_dead():
yield False
return
transformers_was_set = "TRANSFORMERS_OFFLINE" in os.environ
os.environ["HF_HUB_OFFLINE"] = "1"
if not transformers_was_set:
os.environ["TRANSFORMERS_OFFLINE"] = "1"
logger.warning("huggingface.co unreachable; using local HF cache for this load.")
try:
yield True
finally:
os.environ.pop("HF_HUB_OFFLINE", None)
if not transformers_was_set:
os.environ.pop("TRANSFORMERS_OFFLINE", None)
def _swa_cache_path() -> Path:
home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME")
base = Path(home) if home else Path.home() / ".unsloth" / "studio"
return base / "swa_cache.json"
def _load_swa_cache() -> dict:
global _SWA_CACHE
with _SWA_CACHE_LOCK:
if _SWA_CACHE is not None:
return _SWA_CACHE
try:
with open(_swa_cache_path()) as f:
_SWA_CACHE = json.load(f)
if not isinstance(_SWA_CACHE, dict):
_SWA_CACHE = {}
except (FileNotFoundError, json.JSONDecodeError, OSError):
_SWA_CACHE = {}
return _SWA_CACHE
def _save_swa_cache(cache: dict) -> None:
try:
path = _swa_cache_path()
path.parent.mkdir(parents = True, exist_ok = True)
tmp = path.with_suffix(".json.tmp")
with open(tmp, "w") as f:
json.dump(cache, f, indent = 2, sort_keys = True)
tmp.replace(path)
except OSError:
pass
def _period_from_layer_types(layer_types: list) -> Optional[int]:
"""Smallest period N where `(i+1) % N != 0` matches the SWA mask,
or None if no fixed period fits."""
if not layer_types:
return None
is_swa = ["full" not in str(t).lower() for t in layer_types]
n = len(is_swa)
for N in range(1, n + 1):
if all(((i + 1) % N != 0) == is_swa[i] for i in range(n)):
return N
return None
def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]:
try:
from huggingface_hub import hf_hub_download
cfg_path = hf_hub_download(repo_id, "config.json", repo_type = "model")
with open(cfg_path) as f:
cfg = json.load(f)
except Exception:
return None
src = cfg.get("text_config") if isinstance(cfg.get("text_config"), dict) else cfg
period = src.get("sliding_window_pattern")
if isinstance(period, int) and period > 0:
return period
lt = src.get("layer_types")
if isinstance(lt, list) and lt:
return _period_from_layer_types(lt) or [
"full" not in str(t).lower() for t in lt
]
return None
def _arch_aliases(arch: str) -> tuple:
# GGUF emits `falcon-h1`; HF model_type is `falcon_h1`. Normalise both ways.
seen = []
for a in (arch, arch.replace("-", "_"), arch.replace("_", "-")):
if a and a not in seen:
seen.append(a)
return tuple(seen)
def _swa_entry_from_config_obj(cfg) -> Optional[object]:
src = getattr(cfg, "text_config", None) or cfg
period = getattr(src, "sliding_window_pattern", None)
if isinstance(period, int) and period > 0:
return period
lt = getattr(src, "layer_types", None)
if isinstance(lt, list) and lt:
return _period_from_layer_types(lt) or [
"full" not in str(t).lower() for t in lt
]
return None
_SWA_PATTERN_SOURCE_RE = re.compile(
r"sliding_window_pattern\s*(?::\s*[\w\[\], ]*)?\s*=\s*(\d+)"
)
def _resolve_swa_entry_from_transformers(arch: str) -> Optional[object]:
"""Default-instantiate the matching Config; on failure, regex-parse
its source for `sliding_window_pattern = N`."""
try:
from transformers.models.auto.configuration_auto import (
CONFIG_MAPPING,
CONFIG_MAPPING_NAMES,
)
except Exception:
return None
cfg_class = None
for alias in _arch_aliases(arch):
if alias in CONFIG_MAPPING_NAMES:
try:
cfg_class = CONFIG_MAPPING[alias]
break
except Exception:
cfg_class = None
if cfg_class is None:
return None
try:
if (entry := _swa_entry_from_config_obj(cfg_class())) is not None:
return entry
except Exception:
pass
import inspect
candidates = [cfg_class]
text_cfg_class = getattr(cfg_class, "sub_configs", {}).get("text_config")
if text_cfg_class is not None:
candidates.append(text_cfg_class)
for cls in candidates:
try:
src = inspect.getsource(cls)
except (OSError, TypeError):
continue
if m := _SWA_PATTERN_SOURCE_RE.search(src):
period = int(m.group(1))
if period > 0:
return period
return None
def _resolve_swa_pattern(
arch: Optional[str],
n_layers: Optional[int],
source_repo_candidates: tuple = (),
*,
allow_network: Optional[bool] = None,
) -> Optional[list]:
if not arch or not n_layers:
return None
if allow_network is None:
allow_network = os.environ.get("UNSLOTH_STUDIO_OFFLINE", "0") not in (
"1",
"true",
"True",
"yes",
)
cache = _load_swa_cache()
def _entry_to_mask(entry):
if isinstance(entry, int) and entry > 0:
return [(i + 1) % entry != 0 for i in range(n_layers)]
if isinstance(entry, list) and entry:
return [bool(entry[i % len(entry)]) for i in range(n_layers)]
return None
def _persist(entry):
with _SWA_CACHE_LOCK:
cache[arch] = entry
_save_swa_cache(cache)
if (entry := cache.get(arch)) is not None:
if (mask := _entry_to_mask(entry)) is not None:
return mask
if (entry := _BOOTSTRAP_SWA_DEFAULTS.get(arch)) is not None:
return _entry_to_mask(entry)
entry = _resolve_swa_entry_from_transformers(arch)
if entry is not None:
_persist(entry)
return _entry_to_mask(entry)
# Tier 3: live HF fetch (with persistent caching of the result)
if allow_network:
for repo_id in source_repo_candidates:
if not repo_id:
continue
entry = _fetch_swa_entry_from_hf(repo_id)
if entry is not None:
_persist(entry)
return _entry_to_mask(entry)
return None
def _hf_repo_from_url(url: Optional[str]) -> Optional[str]:
"""Strip `https://huggingface.co/owner/name(/...)` to `owner/name`."""
if not url or "huggingface.co/" not in url:
return None
tail = url.split("huggingface.co/", 1)[1].rstrip("/")
parts = tail.split("/")
if len(parts) < 2:
return None
return f"{parts[0]}/{parts[1]}"
# Model size extraction — lazy import to avoid pulling in transformers
# at module level. See PR description for the full explanation.
def _extract_model_size_b(model_id: str):
from utils.models import extract_model_size_b
return extract_model_size_b(model_id)
_TOOL_TEMPLATE_MARKERS = (
"{%- if tools %}",
"{%- if tools -%}",
"{% if tools %}",
"{% if tools -%}",
'"role" == "tool"',
"'role' == 'tool'",
'message.role == "tool"',
"message.role == 'tool'",
)
def detect_reasoning_flags(
chat_template: Optional[str],
model_identifier: Optional[str] = None,
*,
log_source: Optional[str] = None,
) -> dict:
"""Classify a chat template's reasoning and tool-calling capabilities.
Returns a dict with the same five keys populated by the GGUF sniffer:
``supports_reasoning``, ``reasoning_style``
(``"enable_thinking"`` | ``"reasoning_effort"``),
``reasoning_always_on``, ``supports_preserve_thinking``, and
``supports_tools``. Used by both the llama-server backend at load
time and the safetensors/transformers paths in ``routes/inference``
so the two agree on what the frontend will see.
"""
flags = {
"supports_reasoning": False,
"reasoning_style": "enable_thinking",
"reasoning_always_on": False,
"supports_preserve_thinking": False,
"supports_tools": False,
}
if not chat_template:
return flags
tpl = chat_template
prefix = f"{log_source}: " if log_source else ""
if "enable_thinking" in tpl:
flags["supports_reasoning"] = True
flags["reasoning_style"] = "enable_thinking"
logger.info(f"{prefix}model supports reasoning (enable_thinking)")
elif "reasoning_effort" in tpl:
# gpt-oss / Harmony templates use reasoning_effort
# ("low" | "medium" | "high") instead of a boolean.
flags["supports_reasoning"] = True
flags["reasoning_style"] = "reasoning_effort"
logger.info(f"{prefix}model supports reasoning (reasoning_effort)")
elif "thinking" in tpl:
# DeepSeek uses 'thinking' instead of 'enable_thinking'
normalized_id = (model_identifier or "").lower()
if "deepseek" in normalized_id:
flags["supports_reasoning"] = True
logger.info(f"{prefix}model supports reasoning (DeepSeek thinking)")
# Hardcoded <think> tags or reasoning_content in the template mean
# thinking is always on (no toggle to disable it).
if not flags["supports_reasoning"]:
if ("<think>" in tpl and "</think>" in tpl) or "reasoning_content" in tpl:
flags["supports_reasoning"] = True
flags["reasoning_always_on"] = True
logger.info(f"{prefix}model always reasons (<think> tags in template)")
# preserve_thinking is an independent kwarg on some Qwen templates
# that keeps historical <think> blocks in prior assistant turns.
if "preserve_thinking" in tpl:
flags["supports_preserve_thinking"] = True
logger.info(f"{prefix}model supports preserve_thinking")
if any(marker in tpl for marker in _TOOL_TEMPLATE_MARKERS):
flags["supports_tools"] = True
logger.info(f"{prefix}model supports tool calling")
return flags
def _is_mtp_model_name(
model_identifier: Optional[str],
gguf_path: Optional[str] = None,
) -> bool:
"""Name-based MTP detector. Fallback for the metadata signal."""
for cand in (model_identifier, Path(gguf_path).name if gguf_path else None):
if cand and "-mtp" in cand.lower():
return True
return False
def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool:
"""User passed --spec-type / --spec-default? llama-server takes a
single --spec-type (comma-separated to chain), so suppress
auto-emit when this is true."""
if not extra_args:
return False
for raw in extra_args:
tok = str(raw)
if not tok.startswith("--"):
continue
flag = tok.split("=", 1)[0]
if flag in ("--spec-type", "--spec-default"):
return True
return False
def _build_ngram_mod_flags(
caps: Optional[dict],
n_match: int = 24,
n_min: int = 48,
n_max: int = 64,
) -> list[str]:
"""Emit the right ngram-mod knob flags for the running llama-server.
Post-rename builds expose ``--spec-ngram-mod-n-{match,min,max}``;
pre-rename builds expose the legacy ``--spec-ngram-size-n`` /
``--draft-min`` / ``--draft-max``. ``caps`` comes from
``probe_server_capabilities``; ``ngram_mod_flavor`` tells us which
set is real (vs a removal-stub entry). Returns ``[]`` when neither
set is available so the caller can drop ngram-mod entirely.
"""
flavor = caps.get("ngram_mod_flavor") if caps else None
if flavor == "new":
return [
"--spec-ngram-mod-n-match",
str(n_match),
"--spec-ngram-mod-n-min",
str(n_min),
"--spec-ngram-mod-n-max",
str(n_max),
]
if flavor == "legacy":
# Legacy llama.cpp before the spec arg rename: same knobs lived
# under --spec-ngram-size-n (lookup length) and the generic
# --draft-min / --draft-max (ngram size N range).
return [
"--spec-ngram-size-n",
str(n_match),
"--draft-min",
str(n_min),
"--draft-max",
str(n_max),
]
return []
# Canonical Speculative Decoding modes exposed by the Studio chat UI.
# The dropdown renders five options (auto, mtp, ngram, mtp+ngram, off);
# the load API also accepts legacy values that the original Switch and
# external callers emit (default, draft-mtp, ngram-mod, ngram-simple).
_CANONICAL_SPEC_MODES = {"auto", "mtp", "ngram", "mtp+ngram", "off", "ngram-simple"}
_LEGACY_SPEC_MODE_MAP = {
"default": "auto",
"draft-mtp": "mtp",
"ngram-mod": "ngram",
}
def _canonicalize_spec_mode(value):
"""Map any accepted ``speculative_type`` input onto a canonical mode.
Returns one of ``auto``, ``mtp``, ``ngram``, ``mtp+ngram``, ``off``,
``ngram-simple``, or ``None`` (callers treat ``None`` as ``auto``).
Unknown strings collapse to ``auto`` so a stale UI value or typo
falls back to the safe platform-aware path.
"""
if value is None:
return None
if not isinstance(value, str):
return None
stripped = value.strip().lower()
if not stripped:
return None
if stripped in _CANONICAL_SPEC_MODES:
return stripped
if stripped in _LEGACY_SPEC_MODE_MAP:
return _LEGACY_SPEC_MODE_MAP[stripped]
# llama.cpp comma-chains are emitted by old persisted state e.g.
# "ngram-mod,draft-mtp"; collapse the most common one explicitly.
pieces = [p.strip() for p in stripped.split(",") if p.strip()]
has_mtp = any(p in ("mtp", "draft-mtp") for p in pieces)
has_ngram = any(p in ("ngram", "ngram-mod") for p in pieces)
if has_mtp and has_ngram:
return "mtp+ngram"
if has_mtp:
return "mtp"
if has_ngram:
return "ngram"
return "auto"
def _backfill_usage_from_timings(usage, timings):
"""Synthesize ``usage`` from llama-server's ``timings`` when the
OpenAI-style usage block is missing or reports zero tokens.
The Studio chat UI computes generation t/s from
``meta.usage.completion_tokens / totalStreamTime``. llama-server
always populates ``timings.predicted_n`` (true decoded count) and
``timings.prompt_n``, but the ``usage`` field on the final SSE chunk
can be absent or zero on some server builds / streaming
configurations, which makes the UI fall back to wall-clock t/s and
dilute speculative-decoding speedups.
"""
if not timings:
return usage
if usage and usage.get("completion_tokens"):
return usage
predicted_n = timings.get("predicted_n")
prompt_n = timings.get("prompt_n")
if predicted_n is None and prompt_n is None:
return usage
out = dict(usage or {})
if not out.get("completion_tokens") and predicted_n is not None:
out["completion_tokens"] = predicted_n
if not out.get("prompt_tokens") and prompt_n is not None:
out["prompt_tokens"] = prompt_n
out["total_tokens"] = int(out.get("prompt_tokens") or 0) + int(
out.get("completion_tokens") or 0
)
return out
# Probe script run in a short-lived subprocess so the Vulkan instance never
# lives in the long-running backend process. Loads the bundled ggml Vulkan
# backend and prints "<idx>\t<free_bytes>\t<total_bytes>" per device. The
# indices are ggml's own Vulkan device ordinals -- the space
# GGML_VK_VISIBLE_DEVICES expects -- which need not match nvidia-smi order.
_VULKAN_PROBE_SCRIPT = r"""
import ctypes, os, sys
bindir = sys.argv[1]
if sys.platform == "win32":
base_name, vk_name = "ggml-base.dll", "ggml-vulkan.dll"
try:
os.add_dll_directory(bindir)
except Exception:
pass
else:
base_name, vk_name = "libggml-base.so", "libggml-vulkan.so"
try:
ctypes.CDLL(os.path.join(bindir, base_name), mode=ctypes.RTLD_GLOBAL)
lib = ctypes.CDLL(os.path.join(bindir, vk_name), mode=ctypes.RTLD_GLOBAL)
except OSError:
sys.exit(0)
lib.ggml_backend_vk_get_device_count.restype = ctypes.c_int
lib.ggml_backend_vk_get_device_count.argtypes = []
lib.ggml_backend_vk_get_device_memory.restype = None
lib.ggml_backend_vk_get_device_memory.argtypes = [
ctypes.c_int,
ctypes.POINTER(ctypes.c_size_t),
ctypes.POINTER(ctypes.c_size_t),
]
rows = []
for i in range(lib.ggml_backend_vk_get_device_count()):
free, total = ctypes.c_size_t(0), ctypes.c_size_t(0)
lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total))
rows.append("%d\t%d\t%d" % (i, free.value, total.value))
sys.stdout.write("\n".join(rows))
"""
def _vulkan_lib_filename() -> str:
return "ggml-vulkan.dll" if sys.platform == "win32" else "libggml-vulkan.so"
class LlamaCppBackend:
"""
Manages a llama-server subprocess for GGUF model inference.
Lifecycle:
1. load_model() — starts llama-server with the GGUF file
2. generate_chat_completion() — proxies to /v1/chat/completions, streams back
3. unload_model() — terminates llama-server subprocess
"""
def __init__(self):
self._process: Optional[subprocess.Popen] = None
self._port: Optional[int] = None
self._model_identifier: Optional[str] = None
self._gguf_path: Optional[str] = None
self._hf_repo: Optional[str] = None
self._hf_variant: Optional[str] = None
self._is_vision: bool = False
self._healthy = False
# Set by _classify_gpu_offload after _wait_for_health.
self._gpu_offload_active: Optional[bool] = None
self._context_length: Optional[int] = None
self._effective_context_length: Optional[int] = None
self._max_context_length: Optional[int] = None
self._chat_template: Optional[str] = None
self._chat_template_override: Optional[str] = None
self._supports_reasoning: bool = False
self._reasoning_always_on: bool = False
self._reasoning_style: str = "enable_thinking"
self._supports_preserve_thinking: bool = False
self._supports_tools: bool = False
self._cache_type_kv: Optional[str] = None
self._reasoning_default: bool = True
self._speculative_type: Optional[str] = None
# Canonical UI-facing mode the user requested: one of
# ``auto``/``mtp``/``ngram``/``mtp+ngram``/``off``/``ngram-simple``.
# Round-tripped through the status API so the dropdown reflects
# the picked mode rather than the resolved internal flag set
# (auto on a 27B MTP GGUF resolves to draft-mtp but the dropdown
# should still read "Auto").
self._requested_spec_mode: Optional[str] = None
# User-supplied --spec-draft-n-max override (None = platform default).
self._spec_draft_n_max: Optional[int] = None
# KV-cache estimation fields (populated by _read_gguf_metadata)
self._n_layers: Optional[int] = None
self._n_kv_heads: Optional[int] = None
self._n_kv_heads_by_layer: Optional[list[int]] = None
self._n_heads: Optional[int] = None
self._embedding_length: Optional[int] = None
# Architecture-aware KV fields for 5-path estimation
self._kv_key_length: Optional[int] = None
self._kv_value_length: Optional[int] = None
self._sliding_window: Optional[int] = None
self._sliding_window_pattern: Optional[list[bool]] = None
self._full_attention_interval: Optional[int] = None
self._kv_lora_rank: Optional[int] = None
self._key_length_mla: Optional[int] = None
self._kv_key_length_swa: Optional[int] = None
self._kv_value_length_swa: Optional[int] = None
self._ssm_inner_size: Optional[int] = None
self._ssm_state_size: Optional[int] = None
# Last N layers reuse KV from earlier layers and don't allocate
# their own cache (Gemma 3n / Gemma 4: <arch>.attention.shared_kv_layers).
self._shared_kv_layers: Optional[int] = None
# MTP head count (llama.cpp #22673); >0 enables --spec-type draft-mtp.
self._nextn_predict_layers: Optional[int] = None
self._lock = threading.Lock()
# Wraps load_model() end-to-end so concurrent loads serialise
# and never coexist as two llama-server processes (#5401).
self._serial_load_lock = threading.Lock()
# Last extra_args / requested n_ctx, preserved across unload so
# the chat UI's /unload+/load Apply path can inherit them (#5401).
# ``_extra_args_source`` records the (model_identifier, hf_variant)
# the stored args came from so the route can refuse cross-model
# inheritance.
self._extra_args: Optional[List[str]] = None
self._extra_args_source: Optional[tuple[str, Optional[str]]] = None
self._requested_n_ctx: int = 0
self._stdout_lines: list[str] = []
self._stdout_thread: Optional[threading.Thread] = None
# llama-server tee log (see _drain_stdout / _kill_process).
self._llama_log_fh = None
self._llama_log_path: Optional[Path] = None
self._cancel_event = threading.Event()
self._api_key: Optional[str] = None
# True once a probe has completed; cleared on transient failure.
self._is_audio: bool = False
self._audio_type: Optional[str] = None
self._audio_probed: bool = False
# Monotonic timestamp set in _kill_process; read by load_model
# to decide whether to wait for the VRAM reclaim to finish.
self._last_kill_monotonic: float = 0.0
self._kill_orphaned_servers()
atexit.register(self._cleanup)
# ── Properties ────────────────────────────────────────────────
@property
def is_loaded(self) -> bool:
return self._process is not None and self._healthy
@property
def is_active(self) -> bool:
"""True if a llama-server process exists (loading or loaded)."""
return self._process is not None
@property
def base_url(self) -> str:
return f"http://127.0.0.1:{self._port}"
@property
def model_identifier(self) -> Optional[str]:
return self._model_identifier
@property
def is_vision(self) -> bool:
return self._is_vision
@property
def hf_variant(self) -> Optional[str]:
return self._hf_variant
@property
def extra_args(self) -> Optional[List[str]]:
"""Extra llama-server flags from the last load. Copy; None = never
set, [] = explicitly cleared. Used by the route for inheritance."""
return list(self._extra_args) if self._extra_args is not None else None
@property
def requested_n_ctx(self) -> int:
"""n_ctx the last load was invoked with (not the effective cap).
0 means Auto. Used by the route to detect Auto-vs-explicit flips."""
return self._requested_n_ctx
@property
def extra_args_source(self) -> Optional[tuple[str, Optional[str]]]:
"""(model_identifier, hf_variant) the stored extra_args came from.
``None`` if no extras have ever been recorded. Used by the route
to refuse cross-model inheritance (#5401)."""
return self._extra_args_source
@property
def context_length(self) -> Optional[int]:
"""Return the effective context length the server is running at."""
return self._effective_context_length or self._context_length
@property
def max_context_length(self) -> Optional[int]:
"""Return the largest context that fits on this hardware at load time.
This is the "safe zone" threshold the UI renders warnings
against. For a model whose weights fit on some GPU subset, it
is the binary-search cap from ``_fit_context_to_vram`` for that
subset. For a model whose weights exceed 90% of every GPU
subset, it is the 4096 fallback -- the spec's default when the
model will not fit. The UI slider ceiling is
``native_context_length``; dragging above ``max_context_length``
triggers the "might be slower" warning.
"""
return self._max_context_length or self._context_length
@property
def native_context_length(self) -> Optional[int]:
"""Return the model's native context length from GGUF metadata."""
return self._context_length
def load_progress(self) -> Optional[dict]:
"""Return live model-load progress, or None if not loading.
While llama-server is warming up, its process is typically in
kernel state D (disk sleep) mmap'ing the weight shards into
page cache before pushing layers to VRAM. During that window
``/api/inference/status`` only reports ``loading``, which gives
the UI nothing to display besides a spinner that looks stuck
for minutes on large MoE models.
This method samples ``/proc/<pid>/status VmRSS`` against the
sum of the GGUF shard sizes so the UI can render a real bar
and compute rate / ETA. Returns ``None`` when no load is in
flight (no process, or process already healthy).
Shape::
{
"phase": "mmap" | "ready",
"bytes_loaded": int, # VmRSS of the llama-server
"bytes_total": int, # sum of shard file sizes
"fraction": float, # bytes_loaded / bytes_total, 0..1
}
Linux-only in the current implementation. On macOS/Windows the
equivalent would be a different API; this returns ``None`` on
platforms where ``/proc/<pid>/status`` is unavailable.
"""
proc = self._process
if proc is None:
return None
pid = proc.pid
if pid is None:
return None
# Sum up shard sizes (primary + any extras sitting alongside).
bytes_total = 0
gguf_path = self._gguf_path
if gguf_path:
primary = Path(gguf_path)
try:
if primary.is_file():
bytes_total += primary.stat().st_size
except OSError:
pass
# Extra shards live alongside the primary with the same prefix
# before the shard index (e.g. ``-00001-of-00004.gguf``).
try:
parent = primary.parent
stem = primary.name
m = _SHARD_RE.match(stem)
prefix = m.group(1) if m else None
if prefix and parent.is_dir():
for sibling in parent.iterdir():
if (
sibling.is_file()
and sibling.name.startswith(prefix)
and sibling.name != stem
and sibling.suffix == ".gguf"
):
try:
bytes_total += sibling.stat().st_size
except OSError:
pass
except OSError:
pass
# Read VmRSS from /proc/<pid>/status. Kilobytes on Linux.
bytes_loaded = 0
try:
with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f:
for line in f:
if line.startswith("VmRSS:"):
kb = int(line.split()[1])
bytes_loaded = kb * 1024
break
except (FileNotFoundError, PermissionError, ValueError, OSError):
return None
phase = "ready" if self._healthy else "mmap"
fraction = 0.0
if bytes_total > 0:
fraction = min(1.0, bytes_loaded / bytes_total)
return {
"phase": phase,
"bytes_loaded": bytes_loaded,
"bytes_total": bytes_total,
"fraction": round(fraction, 4),
}
@property
def chat_template(self) -> Optional[str]:
return self._chat_template
@property
def chat_template_override(self) -> Optional[str]:
return self._chat_template_override
@property
def supports_reasoning(self) -> bool:
return self._supports_reasoning
@property
def reasoning_always_on(self) -> bool:
return self._reasoning_always_on
@property
def reasoning_style(self) -> str:
return self._reasoning_style
@property
def supports_preserve_thinking(self) -> bool:
return self._supports_preserve_thinking
@property
def reasoning_default(self) -> bool:
return self._reasoning_default
def _reasoning_kwargs(self, enable_thinking: bool) -> dict:
if self._reasoning_style == "reasoning_effort":
return {"reasoning_effort": "high" if enable_thinking else "low"}
return {"enable_thinking": enable_thinking}
def _request_reasoning_kwargs(
self,
enable_thinking: Optional[bool],
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> Optional[dict]:
"""Build chat_template_kwargs from per-request reasoning fields.
Produces a merged dict covering the active model's reasoning style
(``enable_thinking`` or ``reasoning_effort``) plus the independent
``preserve_thinking`` kwarg when the template supports it.
"""
kwargs: dict = {}
# Always-on reasoning models hardcode <think> tags in their template
# and do not consume enable_thinking / reasoning_effort -- skip.
if self._supports_reasoning and not self._reasoning_always_on:
if self._reasoning_style == "reasoning_effort":
if reasoning_effort in ("low", "medium", "high"):
kwargs["reasoning_effort"] = reasoning_effort
elif enable_thinking is not None:
kwargs["reasoning_effort"] = "high" if enable_thinking else "low"
else:
if enable_thinking is not None:
kwargs["enable_thinking"] = enable_thinking
if self._supports_preserve_thinking and preserve_thinking is not None:
kwargs["preserve_thinking"] = preserve_thinking
return kwargs or None
@property
def supports_tools(self) -> bool:
return self._supports_tools
@property
def cache_type_kv(self) -> Optional[str]:
return self._cache_type_kv
@property
def speculative_type(self) -> Optional[str]:
return self._speculative_type
@property
def requested_spec_mode(self) -> Optional[str]:
"""Canonical UI-facing mode the user requested (see field doc)."""
return self._requested_spec_mode
@property
def spec_draft_n_max(self) -> Optional[int]:
"""User --spec-draft-n-max override active on the load, or None
when the platform default (6 GPU / 3 CPU) is in effect."""
return self._spec_draft_n_max
# ── Binary discovery ──────────────────────────────────────────
@staticmethod
def _find_llama_server_binary() -> Optional[str]:
"""
Locate the llama-server binary.
Search order:
1. LLAMA_SERVER_PATH environment variable (direct path to binary)
1b. UNSLOTH_LLAMA_CPP_PATH env var (custom llama.cpp install dir)
2. ~/.unsloth/llama.cpp/llama-server (make build, root dir)
3. ~/.unsloth/llama.cpp/build/bin/llama-server (cmake build, Linux)
4. ~/.unsloth/llama.cpp/build/bin/Release/llama-server.exe (cmake build, Windows)
5. ./llama.cpp/llama-server (legacy: make build, root dir)
6. ./llama.cpp/build/bin/llama-server (legacy: cmake in-tree build)