mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
* studio: tool calling for Llama-3, Mistral, Gemma 4 on safetensors + MLX (#5615)
Adds tool calling for Llama-3, Mistral (pre-v11 + v11+ + [ARGS]), and Gemma 4 to the safetensors / transformers and MLX backends. Parser patched against llama.cpp / vLLM / SGLang per-family parsers and normalises to OpenAI shape. 96 targeted unit tests + cross-OS staging CI (ubuntu / macos-14 / windows) green on the multi-format probe.
* studio: tool-call healing parity between safetensors / MLX and GGUF
After the multi-format parser landed in #5615, the safetensors / MLX
agentic loop and the GGUF loop still differed on healing behaviour.
This commit closes the gaps in both directions so the two backends
react the same way to identical model output.
Changes:
1. core/inference/llama_cpp.py -- the GGUF BUFFERING state machine
now wakes on every emission marker the shared parser knows. Was
("<tool_call>", "<function="); is now the five-tuple imported
from core.inference.tool_call_parser (Qwen / Qwen3.5 / Llama-3
<|python_tag|> / Mistral [TOOL_CALLS] / Gemma 4 <|tool_call>).
Stream cleanup is delegated to the same shared strip_tool_markup
so leaked markup from any family is removed from assistant
content.
2. core/inference/llama_cpp.py -- per-tool canonical heal key. When
a tool arguments field is a bare string and JSON parsing fails,
the GGUF path now heals to {"code": raw_args} for python,
{"command": raw_args} for terminal, and {"query": raw_args} for
everything else. Was hard-coded to {"query": raw_args}, which
silently routed every python / terminal emission through
web_search. Mirrors safetensors_agentic._CANONICAL_HEAL_ARG.
3. core/inference/safetensors_agentic.py -- re-prompt on plan-
without-action. When the model emits a short forward-looking
intent ("I'll search for that", "Let me check", "First, I
will...") and no tool call, the loop nudges the model to act
instead of silently returning a plan-only answer. Up to
_MAX_REPROMPTS=3 (matches GGUF). The intent regex, character
cap, and instruction text are byte-identical to the GGUF path.
The buffer-end fall-through is unified so a buffered intent
emission that never exits the BUFFERING state still triggers
the re-prompt.
4. core/inference/safetensors_agentic.py -- extra iteration slots
for re-prompts. The loop now budgets max_tool_iterations +
_MAX_REPROMPTS + 1 total iterations and tracks the tool-call
count separately, so a stalling model can be nudged 3x without
eating the caller's tool-call budget. Mirrors the _extra slot
reservation in the GGUF path.
Tests (14 new safetensors-side units; 5 GGUF parity pins):
TestLoopRePrompt -- intent-trigger, plain-answer,
no-tools, cap-at-three, budget
preserved, buffer-end intent.
TestLoopCanonicalHealKey -- python / terminal / unknown.
TestGGUFSafetensorsHealingParity -- shared markers used, shared
strip used, canonical heal keys
identical, intent regex matches
same phrases, _MAX_REPROMPTS
equal on both backends.
All 110 targeted tests pass locally; the broader tool / inference /
model-config / sandbox / anthropic / mlx suites stay green.
Why this matters
Without this parity, Llama-3.2 / Mistral / Gemma 4 emissions on Mac
(MLX) and Linux-safetensors stop the agentic loop as soon as the
model says "Let me...", because the GGUF re-prompt logic never
existed on these backends. The two-marker GGUF BUFFERING tuple also
let non-Qwen tool emissions stream out as plain prose when
llama-server's structured channel did not pick them up. Both paths
now drain the same way, heal the same way, and re-prompt the same
way -- so a tool call that works on GGUF works identically on
safetensors / MLX.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: fix tool-call parser bugs from gemini review on #5620
Three high-priority gemini findings on the tool-call parsing additions:
1. unicode_escape on UTF-8 bytes corrupts non-ASCII literals
(e.g. ✨ becomes â\x9c¨). Replace with json.loads on a quoted
string -- preserves emoji / CJK / RTL while still handling
\n \t \uXXXX escapes.
2. Llama-3 sentinel stripping is order-dependent. A leading
`<|eot_id|><|begin_of_text|>` left `<|begin_of_text|>` behind
because the loop had already passed that sentinel. Loop until
no sentinel matches at the start.
3. Mistral v11+ `[TOOL_CALLS] name { json }` regex uses non-greedy
`\{.*?\}` which truncates at the first `}` of a nested JSON
argument, leaking the tail (e.g. `}}`) into user-visible
streamed text. Same problem for the v0.3 array pattern with
nested brackets. Strip those with balanced brace/bracket
scanning via a new `_strip_mistral_closed_calls` helper called
from `strip_tool_markup`.
Also fix the inference routes' parallel `_TOOL_XML_RE`:
- Same nested-JSON truncation in the Mistral patterns; route the
strip through the parser's balanced-scan helper via a thin
`_strip_tool_xml` wrapper that all existing callers now use.
- Llama-3 `<|python_tag|>[^\n<]*` stopped at any `<`, leaking the
tail of any tool call whose argument contained a literal `<`
(queries, code snippets). Relax to `[^\n]*` which keeps the
strip confined to the actual end-of-line.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: tool calling for DeepSeek (R1/V3/V3.1), GLM 4.x, Kimi K2
Adds three more emission-family parsers to tool_call_parser.py so the
shared safetensors / MLX / GGUF agentic loop covers the major open-
weight reasoning families. Patterns ported from llama.cpp
(common/chat-parser.cpp legacy pre-PEG branch), vLLM
(tool_parsers/deepseekv3*, glm4_moe, kimi_k2), and SGLang
(function_call/deepseekv31_detector, glm4_moe_detector, kimik2_detector).
All three references are MIT (llama.cpp) or Apache-2.0 (vLLM, SGLang).
Formats covered:
DeepSeek R1 <|tool▁calls▁begin|><|tool▁call▁begin|>function
<|tool▁sep|>NAME\n```json\n{...}\n```<|tool▁call▁end|>
<|tool▁calls▁end|>
-- args wrapped in a Markdown json fence, ``function``
literal prefix per llama.cpp common_chat_parse_
deepseek_r1 (chat-parser.cpp:801-820)
DeepSeek V3/V3.1
<|tool▁calls▁begin|><|tool▁call▁begin|>NAME
<|tool▁sep|>{json}<|tool▁call▁end|><|tool▁calls▁end|>
-- bare JSON, no code fence, no ``function`` prefix
per llama.cpp common_chat_parse_deepseek_v3_1
(chat-parser.cpp:822-879)
GLM 4.5/4.6/4.7 <tool_call>NAME\n<arg_key>k1</arg_key>
\n<arg_value>v1</arg_value>...</tool_call>
-- strings raw, non-strings JSON-encoded per
chat_template.jinja; multi-call is back-to-back
blocks. Per llama.cpp common_chat_parse_glm_4_5
(chat-parser.cpp:1040-1052)
Kimi K2 <|tool_calls_section_begin|><|tool_call_begin|>
functions.NAME:IDX<|tool_call_argument_begin|>{json}
<|tool_call_end|><|tool_calls_section_end|>
-- bare name recovered by stripping ``functions.``
prefix and ``:IDX`` suffix; full id preserved as
tool_calls[i].id so the roundtrip replays verbatim.
Per llama.cpp common_chat_parse_kimi_k2
(chat-parser.cpp:896-913)
Marker collisions
GLM uses the same ``<tool_call>`` opener as Qwen but with a bare
function name + ``<arg_key>`` body (Qwen has ``\s*{`` after the tag).
The dispatch keeps Qwen first; Qwen's _TC_JSON_START_RE returns no
matches on a GLM emission, so the fall-through to _parse_glm_tool_
calls handles it correctly. Existing Qwen tests confirm zero
regression.
Streaming buffer
TOOL_XML_SIGNALS extended from 5 markers to 12 so the BUFFERING state
machine wakes on every new family's section opener. Added the
DeepSeek alternative markers (ASCII underscores, short ``<|tool▁calls|>``
form) because real checkpoints emit those variants.
Strip patterns
_TOOL_CLOSED_PATS adds DeepSeek envelope (``<|tool▁calls▁begin|>...
<|tool▁calls▁end|>``) and Kimi section (``<|tool_calls_section_begin|>
...<|tool_calls_section_end|>``). _TOOL_ALL_PATS adds the same plus
the unclosed-tail variants so a truncated stream does not leak
markup.
Route gate
_detect_safetensors_features._PARSER_MARKERS grows to include
DeepSeek and Kimi markers plus ``<arg_key>`` (the unique GLM signal).
_TOOL_XML_RE (the route-layer markup-strip regex) gets DeepSeek and
Kimi closed-pair patterns. _TOOL_TEMPLATE_MARKERS in llama_cpp.py
adds ``message['role'] == 'tool'``, ``message['tool_calls']``, and
``tool_calls is defined`` so the classifier recognises DeepSeek's
subscripted-access template style (it has no top-level
``{% if tools %}`` block).
Tests (39 new):
TestParserDeepSeek (7) -- R1 fence, short-form opener, V3.1 bare,
multi-call, with-reasoning, strip,
signal-wakes-streaming
TestParserGLM (6) -- single, mixed types, multi-call,
unclosed-heal, no-Qwen-regression, strip
TestParserKimi (6) -- single, multi-call, dotted-name, unclosed,
strip, signal-wakes-streaming
TestParserCrossFormatRouting (2) -- dispatch routing, signal coverage
TestLoopBasic loop integration (3) -- DeepSeek / GLM / Kimi end-to-end
Capability advertise (3) -- DeepSeek / GLM / Kimi templates flip
supports_tools=True
All 398 targeted tests pass locally (115 safetensors + 27 capability
+ rest of tool / inference / sandbox / model-config suites). Builds
on PR #5620 (parser + healing parity for Llama-3 / Mistral / Gemma 4);
will rebase cleanly onto main once #5620 lands. PR opened as draft -
do not merge until validated against real models for each family.
Sources
- llama.cpp common/chat-parser.cpp lines 801-913, 1040-1052 (MIT)
- vLLM vllm/tool_parsers/deepseekv31_tool_parser.py (Apache-2.0)
- vLLM vllm/tool_parsers/glm4_moe_tool_parser.py (Apache-2.0)
- vLLM vllm/tool_parsers/kimi_k2_tool_parser.py (Apache-2.0)
- SGLang python/sglang/srt/function_call/{deepseekv31,glm4_moe,kimik2}_
detector.py (Apache-2.0)
- Live chat templates: deepseek-ai/DeepSeek-V3.1, zai-org/GLM-4.6,
moonshotai/Kimi-K2-Instruct, unsloth/DeepSeek-V3-0324,
unsloth/GLM-4.5-Air, unsloth/Kimi-K2-Instruct
* studio/routes: make python_tag strip multi-line aware
Earlier revisions of _TOOL_XML_RE in studio.backend.routes.inference
oscillated between two bug shapes:
5615 r"<\|python_tag\|>[^\n<]*" -- stopped at any literal "<"
so code='if x < 10: pass'
leaked '< 10: pass)' to the
user.
5620.1 r"<\|python_tag\|>[^\n]*" -- single-line only; the second
line of
python.call(code="a\nb")
leaked.
The full parser (_parse_llama3_python_tag) already handles both via
balanced-brace scanning, so the parsing path was fine; the LEAK was
in the streaming strip path that runs on every cumulative emission
while content is still arriving.
Switch to r"<\|python_tag\|>(?:[^<]|<(?!\|))*" so the strip consumes:
* any character that is not a "<" (newlines, JSON, code, ...),
* a "<" only when it is NOT followed by "|" (i.e. NOT a Llama-3
sentinel start like <|eot_id|>, <|eom_id|>, <|begin_of_text|>).
This means:
* code='if x < 10' stays inside the strip (5615 fix preserved),
* multi-line code stays inside the strip (5620 round 2),
* the strip terminates at the next Llama-3 sentinel so trailing
assistant content survives.
Tests: TestRoutesPythonTagStrip (8 cases)
pytest test_safetensors_tool_loop.py test_safetensors_capability_advertise.py
-> 118 passed in 1.81s (was 110).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: review follow-ups for DeepSeek / GLM / Kimi tool calling
Four fixes addressing review of the parent commit:
1. GLM <arg_value> coercion: tighten the
json.loads -> ast.literal_eval -> raw cascade to only deserialize
when the body unambiguously looks like a JSON literal (object,
array, JSON-encoded string, true/false/null, or numeric). Strings
like ``True`` / ``None`` (Python literals, not JSON) and arbitrary
prose now stay raw. The bare-numeric / bare-boolean ambiguity with
string args remains an inherent limitation of the template without
schema access -- documented in the new comment. Drops the ast
import entirely (closes Gemini's :1036 suggestion).
2. Kimi K2 bare-counter ids (e.g. ``<|tool_call_begin|>3``) are now
dropped rather than surfaced as a tool literally named "3". Matches
vLLM behaviour; SGLang's schema-infer fallback is out of scope at
the parse site. Real Kimi K2 emissions use ``functions.NAME:IDX``
so this is the exception path.
3. Restore the elaborate ``<|python_tag|>(?:[^<]|<(?!\|))*`` clause in
routes.inference._TOOL_XML_RE -- the simpler ``[^\n<]*`` form
regressed PR #5620's multi-line / literal-``<`` python_tag fix.
Restore ``TestRoutesPythonTagStrip`` (8 tests) adapted to call
``_TOOL_XML_RE.sub`` directly since the ``_strip_tool_xml`` helper
was inlined this PR.
4. Add the spaced and backslash-escaped DeepSeek opener variants
(``<|tool calls begin|>``, ``<|tool\_calls\_begin|>``) to
``TOOL_XML_SIGNALS`` for streaming-gate parity with
``_DEEPSEEK_BEGIN_RE``.
Also updates the llama.cpp / vLLM citations in the parser docstrings:
``common/chat-parser.cpp`` was split into ``common/chat.cpp`` +
``common/chat-peg-parser.cpp`` by llama.cpp PR #18675, and vLLM
moved the tool parsers from ``vllm/entrypoints/openai/tool_parsers/``
to ``vllm/tool_parsers/``. Pin to pre-refactor commit ``51fa458a92d6``
where the cited line numbers still resolve.
New regression tests in ``test_pr5624_regressions.py`` cover the GLM
coercion heuristic shapes, GLM literal-``<`` in arg_value, Kimi K2
dotted name, Kimi K2 bare-counter drop, DeepSeek V3.1 truncated
mid-stream, and routes-layer strip across all three new families.
Tests:
pytest studio/backend/tests/test_safetensors_tool_loop.py
studio/backend/tests/test_safetensors_capability_advertise.py
studio/backend/tests/test_pr5624_regressions.py -q
-> 170 passed in 1.91s
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: tighten verbose comments in tool-call parser sections
Comments were narrating what the code already says. Cut historical
"earlier revisions used X, then Y" narratives down to one-line WHY
notes where the footgun still matters (canonical heal-key parity,
balanced-brace vs non-greedy regex, ``(?:[^<]|<(?!\|))*`` over
``[^\n<]*``/``[^\n]*``). Drop section-header banners.
No behaviour change. Re-ran:
pytest studio/backend/tests/test_safetensors_tool_loop.py \
studio/backend/tests/test_safetensors_capability_advertise.py -q
-> 118 passed.
Regression replay (parser + _coerce_arguments on the 5 #5615 inputs)
-> 21/21.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: GLM 4.7 no-newline emission + Kimi multi-section parity
Two fixes surfaced by triple-confirm verification against the live
HF chat templates and upstream llama.cpp / vLLM / SGLang parsers.
1. GLM 4.7 silent drop
``zai-org/GLM-4.7/chat_template.jinja`` line 65 uses
``{{- '<tool_call>' + tc.name -}}`` which Jinja strips trailing
whitespace from, so the first ``<arg_key>`` follows the function
name with NO ``\n`` between them. Real emissions look like
``<tool_call>get_weather<arg_key>city</arg_key><arg_value>London
</arg_value></tool_call>``. The previous ``_GLM_TC_OPEN_RE`` ended
the name with ``\n`` so GLM-4.7 calls were silently dropped
(parser returned ``[]``).
Fix: relax the name terminator to a lookahead that accepts EITHER
``\n`` OR the next ``<arg_key>``:
_GLM_TC_OPEN_RE = re.compile(
r"<tool_call>\s*([^\n<{][^\n<]*?)\s*(?=\n|<arg_key>)"
)
The first-char restriction ``[^\n<{]`` still excludes Qwen's
``<tool_call>{json}`` form so the Qwen-vs-GLM dispatch remains
mutually exclusive.
2. Kimi multi-section parity with vLLM / SGLang
``vllm/tool_parsers/kimi_k2_tool_parser.py`` and SGLang's
``kimik2_detector.py`` both use ``re.findall`` and so collect every
``<|tool_calls_section_begin|>...<|tool_calls_section_end|>`` block
in a single stream. The previous implementation stopped at the
first ``<|tool_calls_section_end|>``. Kimi K2 doesn't emit
multi-section in practice, but parity is cheap.
Fix: wrap the existing per-call body parser in an outer loop that
advances past each ``<|tool_calls_section_end|>`` and continues to
the next ``<|tool_calls_section_begin|>``. Body parsing extracted
to ``_parse_kimi_section_body`` for clarity. Truncated final
section is still surfaced via the existing in-body balanced-brace
walk.
Verified independently against the live HF templates:
* GLM-4.7 emission constructed from the live template parses to the
expected ``{name, arguments}`` shape.
* GLM-4.5 / 4.6 newline shape continues to parse (the lookahead also
matches ``\n``).
* Qwen ``<tool_call>{json}`` still dispatches to the Qwen path -- the
first-char restriction stops the GLM regex from biting JSON bodies.
* Kimi two-section stream surfaces both calls in order with full ids
preserved.
* Bare-counter Kimi ids still drop.
Tests added in ``test_pr5624_regressions.py``:
* ``test_glm_4_7_no_newlines_between_name_and_arg_key``
* ``test_glm_4_7_no_newlines_multi_call``
* ``test_glm_4_7_does_not_break_qwen_path``
* ``test_kimi_two_sections_in_one_stream_both_parse``
pytest studio/backend/tests/test_safetensors_tool_loop.py
studio/backend/tests/test_safetensors_capability_advertise.py
studio/backend/tests/test_pr5624_regressions.py -q
-> 174 passed in 1.93s
pytest studio/backend/tests/ -q -k 'not gpu and not llama_cpp_integration'
-> 2038 passed, 15 failed (pre-existing CI gaps).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: parser robustness fixes for PR #5620
Three surgical extensions to the multi-format tool-call parser, each
covering a real fine-tune / template emission shape that the current
parser silently drops. No path narrows; all changes widen what is
accepted.
1. `_parse_tool_call_json` now accepts both `arguments` and
`parameters` keys. A Hermes / Qwen `<tool_call>{json}</tool_call>`
wrapper around a Llama-3.2 fine-tune that emits the `parameters`
key was extracting the tool name and silently discarding the
args, producing a working-shaped call with an empty payload. The
bare-JSON and python_tag paths already accepted both keys; this
path now matches them.
2. `_TC_FUNC_START_RE`, `_TC_PARAM_START_RE`, and `_TC_PARAM_CLOSE_RE`
now also match the attribute form
`<function name="..."><param name="...">v</param></function>` used
by MiniCPM-5 and MiniMax-M2. Names land in either capture group,
and `</param>` is accepted as a short close.
3. `_parse_llama3_bare_json` sentinel-strip now consumes the role
label inserted between `<|start_header_id|>` and
`<|end_header_id|>` by Meta's official Llama-3.x chat template.
Without this, every assistant turn re-fed through the template
prefix `<|start_header_id|>assistant<|end_header_id|>\n\n{json}`
parsed to zero calls, so any history-with-tool-call round-trip
in production silently dropped.
Tests in `studio/backend/tests/test_safetensors_tool_loop.py`:
* `TestParserRobustness::test_tool_call_json_accepts_parameters_key`
* `TestParserRobustness::test_function_xml_attribute_form`
* `TestParserRobustness::test_function_xml_attribute_form_multi_param`
* `TestParserRobustness::test_function_xml_legacy_equals_form_still_works`
(regression guard for the existing `<function=name>` syntax)
* `TestParserRobustness::test_llama3_chat_template_round_trip`
* `TestParserRobustness::test_llama3_round_trip_all_roles`
* `TestParserRobustness::test_llama3_round_trip_with_eot_prefix`
`pytest studio/backend/tests/test_safetensors_tool_loop.py
studio/backend/tests/test_safetensors_capability_advertise.py -q`
goes from 118 to 125 passed.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim verbose comments in tool-call parser sections for PR #5624
Pure comment / docstring tightening on top of the GLM 4.7 + Kimi
multi-section fixes. No behavioural change.
* Drop multi-paragraph prelude and post-refactor citation chatter in
the DeepSeek, GLM and Kimi parser docstrings; keep the shape and
upstream-commit pin.
* Collapse ``parse_tool_calls_from_text``'s 9 per-family blocks into
a single ordered loop with one combined comment.
* Tighten the GLM coercion, Kimi bare-counter and ``_TOOL_XML_RE``
comments to one or two lines each.
* Same trim pass on ``_PARSER_MARKERS`` and the regression-test
docstrings.
Tests:
pytest studio/backend/tests/test_safetensors_tool_loop.py
studio/backend/tests/test_safetensors_capability_advertise.py
studio/backend/tests/test_pr5624_regressions.py -q
-> 174 passed in 2.00s
* Fix O(N^2) DeepSeek V3.1 backtracking for PR #5624
Adversarial input ``<|tool▁calls▁begin|><|tool▁call▁begin|>fn<|tool▁sep|>``
followed by a long body that does NOT contain a closing brace caused
the V3 path's ``([^\n<]+?)<|tool▁sep|>`` regex to backtrack
quadratically: at each position the lazy quantifier extends one char
at a time looking for a sep that isn't there, taking ~19s on 50k
chars.
Replace the regex search with ``str.find`` on the sep marker plus a
left-walk to recover the name. ``str.find`` is O(N); the walk stops
on ``\n`` (turn boundary), ``<`` (start of a tag), or ``>`` (end of
an optional ``<|tool▁call▁begin|>`` prefix). Same observable
behaviour as the regex on every canonical input.
Tests:
test_deepseek_v3_1_huge_truncated_body_is_linear (new) -- 50k chars
must parse in < 1s.
pytest studio/backend/tests/test_safetensors_tool_loop.py
studio/backend/tests/test_safetensors_capability_advertise.py
studio/backend/tests/test_pr5624_regressions.py -q
-> 175 passed in 1.97s
pytest studio/backend/tests/ -q -k 'not gpu and not llama_cpp_integration'
-> 2038 passed, 15 pre-existing failures unchanged.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: terminate function-XML body at </function>, not just </tool_call>
`_parse_function_xml` was looking for `</tool_call>` (the Hermes
wrapper) as the body terminator. When a model emits a standalone
`<function=NAME><parameter=K>v</parameter></function>` followed by
explanatory prose (which models routinely do), no `</tool_call>` is
present, so the body extended to end-of-string and the trailing
prose leaked into the LAST parameter value.
Pre-existing on main (the legacy `<function=NAME>` form had this
bug too). Same affects PR #5620's new attribute-form
`<function name="NAME"><param name="K">v</param></function>`
emission used by MiniCPM-5 / MiniMax-M2.
Fix: `_TC_END_TAG_RE` now matches either `</tool_call>` OR
`</function>`. The existing `_TC_FUNC_CLOSE_RE` / `_TC_PARAM_CLOSE_RE`
strips are unchanged. Multi-call inputs still bound each function
at the next `<function=` start, so no over-eager consumption.
New tests:
* `test_function_xml_followed_by_prose` (legacy form + prose)
* `test_function_attribute_xml_followed_by_prose` (attribute form + prose)
Existing `test_code_with_embedded_xml` still passes (a parameter
value containing literal `<a></a>` is preserved because the
embedded close tag is `</a>`, not `</function>`).
`pytest studio/backend/tests/test_safetensors_tool_loop.py
studio/backend/tests/test_safetensors_capability_advertise.py -q`
goes from 125 to 127 passed.
* Studio: tighten Llama-3.2 bare-JSON guard
A fuzz pass on PR #5811 turned up that ``_parse_llama3_bare_json``
accepted ``parameters`` as a string, contradicting the docstring's
"parameters or arguments is a dict" guard. Prose JSON like
``{"name":"foo","parameters":"a sentence"}`` would wrongly fire the
parser, which the agentic loop would then heal into a real
``foo(query="a sentence")`` call.
Same code lives on this branch, so the same fix applies here.
Tightened guard:
- ``parameters`` must be a dict (Llama-3 spec).
- ``arguments`` may be a dict, or a JSON-encoded string that
decodes to a dict (OpenAI shape, e.g.
``"arguments":"{\"q\":\"x\"}"``). Plain non-JSON strings or
JSON-strings of lists / scalars / null no longer pass.
Mirrors the fix landed in PR #5811 commit 615b8608. Adds the same
4 regression tests under TestParserMultiFormat.
Existing test suite stays green: 127 -> 131 passing.
* Studio: skip non-scalar args in python_tag JSON form
The JSON sub-path of ``_parse_llama3_python_tag`` was fabricating
``{"value": args}`` when the model emitted a non-dict / non-string
``arguments`` value (e.g. ``42``, ``[1,2,3]``, ``null``, ``true``).
This silently turned a malformed emission into a real tool call,
which the agentic loop would then execute with arguments the model
never intended.
Tightened: skip the call instead of fabricating. The same
behaviour now matches the bare-JSON guard tightened earlier
(strict-guard merge from PR #5620, inherited via merge here).
Added a regression test covering the four non-scalar shapes.
Pass count on this branch: 158 -> 159.
Sites in ``_parse_tool_call_json`` and ``_consume_mistral_call``
keep the existing looser behaviour for now; both are reached
only after explicit ``<tool_call>`` / ``[TOOL_CALLS]`` markers
so the false-positive surface there is much narrower.
* studio: fix safetensors tool-call parser gaps vs llama.cpp (Mistral CALL_ID / THINK, attribute-form signal)
Three GGUF-parity fixes to the safetensors tool-call parser, each matching
llama.cpp's reference behaviour:
- Mistral Small 3.2 emits [TOOL_CALLS]name[CALL_ID]<id>[ARGS]{json}. The
parser stopped after the name on seeing [CALL_ID] (neither [ARGS] nor {),
dropping the call. Skip an optional [CALL_ID]<id> segment in both the
parse and strip paths. llama.cpp parses this (test-chat.cpp:4785).
- Magistral wraps reasoning in [THINK]...[/THINK]. A [TOOL_CALLS] inside the
reasoning was parsed as a real call, producing a phantom call. Strip a
leading [THINK] block before scanning so only the post-reasoning call
counts (test-chat.cpp:2285); a literal [THINK] inside a later argument is
left intact.
- The standalone MiniCPM-5 / MiniMax-M2 <function name="..."> attribute form
parsed correctly but was absent from TOOL_XML_SIGNALS and the markup strip
patterns, so the streaming safety-net parse was gated off (dropping the
call) and markup leaked into displayed text. Add the signal and broaden
the strip regexes.
Adds regression tests for all three.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: fix GLM and Kimi K2 safetensors tool-call parser gaps vs llama.cpp
Four GGUF-parity fixes for the GLM and Kimi K2 families:
- GLM 4.7 zero-argument inline call <tool_call>name</tool_call> was dropped:
the open-tag lookahead only allowed \n or <arg_key> after the name. Allow
</tool_call> too so a no-arg call parses to empty args (vLLM / SGLang /
llama.cpp all parse it).
- GLM string argument values were stripped, losing significant leading /
trailing whitespace in code / diff arguments. Keep the raw value for the
string fallback and only strip the copy used to probe for a JSON literal,
matching vLLM glm4_moe which never strips string args.
- Kimi K2 calls emitted without the <|tool_calls_section_begin|> wrapper
were dropped. llama.cpp makes the section optional (Kimi can call a tool
straight after reasoning without opening a section); parse a bare
<|tool_call_begin|> when no section is present.
- Kimi K2 malformed / truncated JSON in one call dropped every later call in
the section. Skip the bad call and keep parsing so valid subsequent calls
are recovered (vLLM parity).
Adds regression tests for all four.
* studio: fire safetensors tool calls for the bare-JSON (Llama-3.2) form
The agentic loop's streaming safety-net parse was gated on
has_tool_signal(), which is False for the Llama-3.1 / 3.2 bare-JSON tool
form {"name":..,"parameters":..} (no XML marker). Real tool calls were
therefore dropped: the loop logged "model planned without calling tools",
re-prompted three times, then gave up with zero tool calls, while GGUF's
llama-server parses the same emission natively.
Run parse_tool_calls_from_text() unconditionally in the safety net. The
parser is strict (only fires on a valid tool-call shape) so plain answers
are unaffected. Reproduced on a real unsloth/Llama-3.1-8B-Instruct run:
the model emits {"name":"web_search","parameters":{...}} which now
executes the tool instead of being re-prompted into a no-op.
Adds a loop regression test for the bare-JSON form.
* studio: fire safetensors tool calls for Gemma 4 (native template + stripped parser)
Gemma-4 safetensors fired no tools while its GGUF fired reliably. Three gaps:
- The Studio swaps in the Unsloth "gemma-4" chat template, which does not
render the tools schema (the model's native template does), so the model
never saw the tools. Fall back to the model's native template when the
override template renders identically with and without tools. Same fix
helps any family whose override template drops tools.
- skip_special_tokens strips the <|tool_call> wrapper and <|"|> string
markers, so a streamed Gemma-4 call arrives as a bare call:NAME{k:v, ...}
with unquoted values. Parse that form, keeping commas/braces inside a
code or command value, normalising surrounding quotes, and stripping the
leaked markup from the final answer.
- Without a grammar a small model can loop, repeating one call for the whole
tool budget. Collapse exact-duplicate calls within a turn and force a final
answer after a turn that made no new tool progress (llama-server's lazy
grammar prevents this loop on the GGUF side).
Adds parser tests for the bare/stripped Gemma-4 form.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: complete strict-mode contract and fix parser import paths
Address review findings on the multi-format tool-call parser:
- Honor allow_incomplete=False in the remaining sub-parsers. The Llama-3
<|python_tag|>NAME.call(...) parser, the pre-v11 Mistral [TOOL_CALLS] array
parser, and the Gemma 4 <|tool_call> parser ignored strict mode, so a
truncated call (missing closing paren, ], or <tool_call|>) was still healed
and executed with Auto-Heal disabled. Thread strictness through and reject
the unclosed forms, matching the JSON and function-XML paths.
- Drop the duplicate tool_call_parser import block in llama_cpp.py and the
redundant un-aliased TOOL_XML_SIGNALS; only the _SHARED_TOOL_XML_SIGNALS
alias is used as a value.
- Import _strip_mistral_closed_calls from core.inference.tool_call_parser in
routes/inference.py instead of studio.backend.core... The self-contained
run.py launch mode only puts studio/backend on sys.path, so the absolute
package path raised ModuleNotFoundError on the server-tool strip path.
Add strict-mode regression tests for the truncated Llama-3 dot-call and the
unclosed Mistral array.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: harden DeepSeek/Kimi tool-call parsing and strip
Address review findings on the DeepSeek and Kimi parsers:
- Honor allow_incomplete=False for DeepSeek. An envelope with no closing
<|tool▁calls▁end|> is truncated mid-stream; reject it in strict mode
instead of healing the body out to EOF, matching the strict XML and Mistral
paths.
- Do not skip a following tool call when the current call's end marker is
missing. The DeepSeek V3 and Kimi loops advanced by searching forward for the
next <|tool▁call▁end|> / <|tool_call_end|>, which could land on a later
call's end marker and drop the call in between. Advance by the JSON end; the
loop re-locates the next call marker from there.
- Strip truncated DeepSeek and Kimi section blocks in the route-level display
regex. The patterns required the closing marker; add the end-of-text
alternative so a block truncated by EOS does not leak raw markup to the UI.
Add regression tests for the truncated DeepSeek envelope, and for DeepSeek and
Kimi multi-call recovery when the first call's end marker is missing.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: preserve XML param indentation and alias Mistral array parameters
Two parser-correctness fixes found by auditing against the model chat templates
and the SGLang / vLLM reference parsers:
- Qwen3.5 XML parameter values lost their leading indentation. The chat template
emits <parameter=k>\nVALUE\n</parameter>, but the parameter-start regex ate the
wrapping newline AND the value's first-line indentation with a trailing \s*,
then str.strip() removed the rest. Narrow the trailing class to horizontal
whitespace only and trim exactly one wrapping newline (via _trim_param_value),
preserving indentation in code/diff arguments. Matches SGLang's qwen3_coder
detector. Applies to both _parse_function_xml (tool_call_parser.py) and the XML
path in tool_healing.py.
- Mistral pre-v11 array objects keyed on parameters dropped their payload.
_consume_mistral_call read only the arguments key; alias parameters the same way
the JSON/XML paths and SGLang's base detector do.
Add regression tests for preserved multi-line indentation and the array
parameters alias.
* Studio: DeepSeek strip sync, Gemma nested args, GLM/Kimi strict mode
Parser-correctness fixes found by auditing DeepSeek/GLM/Kimi against vLLM,
SGLang, and the model chat templates:
- DeepSeek: the short <|tool▁calls|> opener (and the space / escaped-underscore
spellings) was parsed but never stripped, so a short-opener envelope leaked raw
markup to the UI. Share one opener alternation between _DEEPSEEK_BEGIN_RE and
the strip patterns (and the route-level display regex) so a signal we parse can
never be left un-stripped.
- Gemma wrapper-less stream: a nested object/array argument (loc:{city:NYC},
labels:[bug,ui]) was kept as a literal string. Parse it recursively when the
bare value is a balanced {} / [], falling back to the raw string for a
truncated value.
- GLM and Kimi ignored allow_incomplete. With Auto-Heal off, a GLM block with no
</tool_call>, a Kimi section with no <|tool_calls_section_end|>, or a Kimi call
with no <|tool_call_end|> are truncated and must be rejected, matching the
strict behavior of the JSON/XML/Mistral/DeepSeek paths and vLLM/SGLang.
Add regression tests for the short-opener strip, the Gemma nested args, and GLM /
Kimi strict-mode rejection.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: tighten tool-call parser comments
Make the comments in the multi-format tool-call parser and its callers succinct:
compress verbose docstrings/blocks to one or two lines, drop ones that restate the
code, and trim the tiny balanced-scanner helpers. Correctness rationale and
upstream provenance (SGLang/llama.cpp parity, the strict-mode / Auto-Heal
contract, whitespace-preservation, and the Unicode / full-width-pipe notes) are
kept in compact form.
Comment-only: no code or behavior change (verified with comment_tools.py check
--strip-docstrings; parser suite green).
* Studio: tighten DeepSeek/GLM/Kimi parser comments
Compress the comments added for the DeepSeek/GLM/Kimi parsers and the Gemma
wrapper-less helpers to one or two lines, keeping the upstream provenance
(llama.cpp 51fa458a92d6), the O(N^2) / strict-mode rationale, and the vLLM parity
notes intact.
Comment-only: no code or behavior change (verified with comment_tools.py check
--strip-docstrings; parser suite green).
* Studio: make DeepSeek R1 / GLM parsing linear and close routes strip gaps
Review follow-up for the DeepSeek/GLM/Kimi parser:
- DeepSeek R1 detection used a greedy ``([^\n]+)\n```json`` regex that backtracks
O(N^2) on a fence-less truncated body; scan with str.find instead (mirrors the
V3 path).
- GLM arg pairs used a lazy-group finditer that rescanned to EOF from each bare
<arg_key> in an unclosed body (O(N^2)); walk pairs with str.find.
- The route display strip (_TOOL_XML_RE) accepted fewer DeepSeek openers than the
parser (missed the space / escaped-underscore spellings) and missed bare
section-less Kimi calls, so a call we parse could leak raw markup to the UI.
Reuse the parser's shared _DEEPSEEK_OPEN_RE_SRC and add a bare-Kimi arm.
Add ReDoS-linearity regressions for the R1 and GLM paths, a positive R1
fenced-json parse test, and routes-strip tests for the space/escaped DeepSeek
openers and the bare Kimi call.
* Studio: fix test_mcp_servers _TOOL_XML_RE reconstruction after _DS_OPEN_SRC reuse
The routes strip fix made _TOOL_XML_RE reference the module-level
_DS_OPEN_SRC variable. test_mcp_servers reconstructs the regex by exec-ing
the extracted compile() source in a namespace that only defined _re, so it
raised NameError. Inject _DS_OPEN_SRC into that namespace, matching the same
fix already applied in test_tool_xml_strip.
* Studio: make Llama-3 .call and Mistral-array healing parsing linear
Two more O(n^2) ReDoS paths in the multi-format parser, both reachable from
the agentic loop on a long truncated body with no length cap:
- _LLAMA3_KV_RE.finditer over a .call(...) body retried at every offset of a
long word run / unterminated quote (40K -> 14s). Replace with a hand-scan
that reuses the same key/number/literal sub-regexes via anchored match and
walks the string body by hand, so an unterminated quote is O(n). Verified
byte-identical to the old regex over 200K fuzzed inputs.
- _parse_mistral_array healing ran _balanced_brace_end from every { in the
body (20K -> 17s). Walk top-level objects, advancing past each balanced
{...}; this also drops the phantom call the old scan emitted from a nested
argument object.
Add adversarial-length linearity regressions plus positive .call kwargs and
unclosed-array recovery coverage.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: strengthen #5624 regression assertions and strip-test harness guards
- test_strip_tool_markup_handles_deepseek_envelope used `A or B` where B was the
preservation property the next line already asserts, masking the real check.
Replace with an explicit assertion that the call name and args are stripped.
- The test_tool_xml_strip source-extraction harness reconstructs _TOOL_XML_RE and
_strip_tool_xml_for_display from routes/inference.py via lazy regexes that could
silently grab a shorter slice. Assert the extracted regex carries the DeepSeek /
bare-Kimi arms and the helper body reached the _TOOL_XML_RE.sub call.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: honor strict mode in safety-net, keep empty Gemma args, strip attribute-form function XML
- safetensors safety-net parser now forwards allow_incomplete=auto_heal_tool_calls,
matching the draining path, so a late incomplete tool call is not healed and
executed when Auto-Heal is off.
- Gemma empty bare value ({k:}) now serialises as "" instead of invalid {"k":},
which previously dropped the whole call.
- Route _TOOL_XML_RE also strips the <function name="..."> attribute form
(MiniCPM-5 / MiniMax-M2) so it no longer leaks to the UI.
* Studio: linearize wrapper-less Gemma nested-arg parsing and correct parser provenance
- _gemma_parse_value/_gemma_parse_mapping/_gemma_parse_array now parse nested
{}/[] in a single forward pass instead of pre-scanning each subtree with a
balanced-brace walk and re-parsing it. Deeply nested wrapper-less Gemma args
were O(n^2); they are now ~linear (and ~40x faster at depth 400).
- Correct the DeepSeek/GLM/Kimi provenance comments: the cited commit
51fa458a92d6 is unrelated, and GLM/Kimi were never standalone
common_chat_parse_* functions (llama.cpp uses common_chat_params_init_glm_4_5
plus a generalized XML parser, PRs #15904 / #16932).
- Add tests: Gemma deep-nesting linearity, nested object/array preservation,
same-turn distinct-call cap, and the native-template tool-render fallback.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: guard Gemma value parser against non-advancement and missing tokenizer
Addresses Gemini review:
- _gemma_parse_value now consumes one character when a stray }/]/, sits where a
value is expected, so _gemma_parse_array can never stall at the same index on
malformed input (a latent infinite loop).
- _render_with_native_template returns None when neither a tokenizer nor a
processor is present instead of raising AttributeError.
- Tests for both.
* Studio: fix attribute-form function-XML literal close tag and zero-arg strict call
Addresses Codex review of the <function name="..."> attribute form in
_parse_function_xml (MiniCPM-5 / MiniMax-M2):
- End the call body at the LAST </function> / </tool_call> within the call's
window, so a literal close tag inside a code/search argument (e.g.
print("</function>")) is preserved instead of truncating the call.
- Accept a closed call with no parameters as a valid zero-argument call in strict
mode (the function close is already required), instead of rejecting it as a
truncated call.
- Tests for both, mirroring the legacy <function=...> coverage.
* Studio: drop scratch review/planning artifacts from the branch
* Studio: fix tool-call parser/loop review findings on the multi-format path
Address the live code-review findings on the safetensors/MLX + GGUF tool path:
- routes: include the attribute form <function name="..."> in the safetensors
capability whitelist so MiniCPM-5 / MiniMax-M2 templates keep the tool pill
(parser already handles the form; the post-filter wrongly suppressed it).
- safetensors loop: build the plan-without-action re-prompt from the active
tools instead of a hardcoded web_search/python string, and gate it on
auto_heal_tool_calls, matching the GGUF loop.
- safetensors loop: hold a leading bare-JSON object ({"name":..,"parameters":..})
during BUFFERING until it closes, then drain it as a tool call instead of
streaming the raw JSON to clients. The DRAINING/STREAMING resolvers still
recover a plain JSON answer, so this can never drop content.
- parser: anchor the Llama-3 <|python_tag|>NAME.call(...) scan to the tag and
chain ; -separated calls, so all semicolon-separated built-ins parse and a
literal <|python_tag|>x.call(...) inside a JSON string argument no longer
fires the wrong tool.
- parser: consume the optional trailing </s> after a named Mistral
[TOOL_CALLS]name{json} call, mirroring the array shape.
- GGUF streaming strip: use the shared parser patterns (which know
[TOOL_CALLS] and <|python_tag|>) so a textual tool call entering DRAINING is
stripped instead of leaking the marker to streaming clients.
- routes: hoist the _strip_mistral_closed_calls import to module level.
Adds regression tests covering each fix; existing parser suite stays green.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: fix DeepSeek/GLM/Gemma tool-call review findings
Address the live code-review findings specific to the DeepSeek / GLM / Kimi
and native-template additions:
- parser: in strict mode (Auto-Heal off) require the per-call
<|tool▁call|end|> terminator for DeepSeek V3 calls instead of executing on
a bare balanced object closed only by the envelope end.
- parser: keep GLM string arguments that begin with a quote verbatim (drop
the leading-quote case from the JSON-decode probe) so a quoted search query
is not decoded down to its inner text.
- parser: reject a GLM call with an unclosed <arg_value> in strict mode, and
under Auto-Heal keep the partial value rather than dropping it to a no-arg
call.
- parser: add a balanced wrapper-less Gemma strip (call:NAME{...}) so a nested
object/array argument is removed whole instead of leaving a trailing brace;
run the balanced Mistral and Gemma strips on the streaming display paths too.
- safetensors loop: buffer a leading wrapper-less Gemma call:NAME{...} so it
drains and executes instead of streaming the raw call text.
- inference: render the native-template fallback on a shallow tokenizer copy
instead of mutating the shared tokenizer outside the generation lock, and
load the native template from base_model for LoRA adapters.
Adds regression tests for each; existing parser suite stays green.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: harden multi-format tool-call detection from review findings
Apply five targeted fixes from the review pass over the multi-format tool
path:
- routes: route display strip delegates to _strip_tool_xml so Mistral
[TOOL_CALLS] blocks with nested JSON are removed from streamed display
text, not just the XML forms.
- tool_call_parser: skip function/parameter starts that fall inside an
already-open parameter block (_inside_open_parameter) so nested example
payloads are not mis-parsed as new calls; extract
strip_llama3_leading_sentinels so the bare-JSON guard is shared.
- safetensors_agentic: probe bare JSON through strip_llama3_leading_sentinels
before the balanced-brace check so a leaked header sentinel does not defeat
the guard.
- tool_healing: allow dotted tool names in the Gemma wrapped start pattern.
- llama_cpp (GGUF): buffer wrapper-less Llama-3.2 {"name":..} calls that carry
no XML signal, drain a complete object silently and hold an incomplete one,
and run the end-of-stream safety net unconditionally so markerless calls are
detected and never leak the raw JSON (including truncated fragments).
Adds regression tests for the GGUF bare-JSON streaming path and the Mistral
display strip.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: stop bare-JSON tool calls leaking at EOF, oversized, and into history
The second review pass flagged that the Llama-3.2 bare-JSON tool-call handling
still leaked raw JSON in several spots; ``strip_tool_markup`` only knows
XML/bracket markup, so the bare-JSON form survived it. Fix them symmetrically
across the safetensors and GGUF loops:
- Safetensors stream-end resolver now routes a held bare-JSON fragment to
DRAINING (mirroring GGUF) so a truncated ``{"name":..`` cut off by the end of
the stream is dropped instead of flushed as assistant content. The 7/10
reviewer finding.
- Both loops now drain (suppress) an oversized still-open bare-JSON call once it
passes ``_MAX_BARE_JSON_BUFFER`` instead of streaming the raw prefix, gated on
a ``"name"`` key so a giant plain JSON answer still streams; a complete
oversized call still executes via the safety net.
- Add a shared ``strip_leading_bare_json_call`` helper and apply it to the
content kept for the assistant turn in both loops, so an executed bare-JSON
call is not replayed as visible text or fed back as next-turn history.
Plain JSON answers without a ``"name"`` key are untouched throughout. Adds
regression tests for the EOF, oversized, and next-turn cases on both backends
plus unit tests for the helper.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: bound the Llama-3 python_tag strip on real control sentinels
The route display strip's <|python_tag|> arm ran to the next <| of any kind.
A tool-call argument carrying a literal <|...|> token (for example <|cite|>
inside a string value) truncated the strip early and leaked the call tail into
the visible response. Narrow the stop condition to the genuine Llama control
sentinels (eot_id, eom_id, python_tag, start/end_header_id, begin_of_text,
finetune_right_pad_id) so embedded markup and JSON are consumed while real
header/turn boundaries still bound the strip.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: harden GLM/Gemma parsing, cap GGUF textual calls, share native-template fallback
GLM 4.x parser walked a body pre-bounded by the first </tool_call>, so a string
argument containing a literal </tool_call> (e.g. code that prints it) was
truncated. Walk arg_key/arg_value pairs against the full content instead, since
each <arg_value> is delimited by its own </arg_value> and the call's real close
is the </tool_call> that precedes the next <arg_key>.
Add a truncated wrapper-less Gemma pattern (call:NAME{... with no closing brace)
to the markup strip so a call cut off mid-arguments does not leak raw into the
visible stream. It runs after the closed form, so a complete call keeps trailing
prose.
Cap and dedup tool calls parsed from the GGUF TEXTUAL fallback at
_MAX_TOOL_CALLS_PER_TURN, mirroring the safetensors loop. Structured
delta.tool_calls are grammar-bounded by llama-server, but text parsed straight
from content is not, so one runaway turn could fan out into dozens of
executions.
Extract the native-chat-template fallback into chat_template_helpers
(render_native_template / render_with_native_template_fallback) so the
transformers and MLX text backends share one implementation. The MLX text path
now applies it too, so an Unsloth override template that drops the tools schema
no longer silently stops MLX from advertising tools. The MLX VLM path renders
via the processor for image tokens and is intentionally left on its own render.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: gate markerless bare JSON on enabled tools and close parser/strip asymmetries
The Llama-3.2 custom_tools bare-JSON form has no marker, so any JSON object with a
name key was read as a tool call. An ordinary JSON answer like
{"name":"Alice","parameters":{"age":30}} was misclassified as a call to a
disabled tool and dropped from the visible response. Gate the markerless form on
the enabled tool names (threaded through parse_tool_calls_from_text and
strip_leading_bare_json_call, supplied by both streaming loops): an object whose
name is not an enabled tool is ordinary content. The marker-based forms keep
their name-agnostic behaviour (an explicit signal is a real call attempt), and
unrestricted mode stays ungated.
Also fix two parser/strip asymmetries the parser already tolerated:
- A literal </function> inside a parameter value (print("</function>")) truncated
both the core and route strips at the first close, leaking the tail. Extend the
strip to the call's real close (last </function> before the next opener),
mirroring the parser, without merging separate calls.
- The single-object Mistral [TOOL_CALLS]{...} shape parsed but _strip_mistral_closed_calls
left it, leaking the raw object into display. Strip the balanced object while
keeping trailing prose, matching the array and name shapes.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio tools: fix strip/parse symmetry and native-template token for DeepSeek/GLM/Kimi
Pass-3 review follow-ups on the multi-format tool parser:
- Bare Kimi call (<|tool_call_begin|>...<|tool_call_end|> with no section
wrapper) is accepted by the parser, so add it to the closed strip patterns
so the streaming (non-final) display strip removes it instead of leaking the
markup mid-generation.
- Route display strip now also runs the wrapper-less Gemma cleanup, so a
Gemma 4 call:NAME{..} no longer leaks into the visible answer.
- MLX model record carries base_model for a LoRA adapter so the native-template
fallback loads the base repo template rather than the adapter's
(often template-less) tokenizer.
- Native-template reload forwards the load-time HF token so a gated/private
model's repo template can still be fetched (transformers and MLX text paths).
- GGUF end-of-stream bare-call heuristic is gated on the enabled tool names so a
truncated ordinary JSON object ({"name":"Alice","age":) streams as the answer
instead of being dropped as a tool call.
Adds regression tests for each case.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio tools: gate GGUF bare-JSON suppression on enabled tools and fix python-tag exponent parsing
Pass-4 review follow-ups on the GGUF tool loop and Llama-3 parser:
- The GGUF bare-JSON suppression sites still keyed off a raw "name" substring,
so an ordinary JSON answer whose name is not an enabled tool was dropped when
it was truncated, oversized, or reached the no-tool DRAINING fallback (the
parser, helper, and safetensors paths were already gated). All three sites now
use the shared enabled-name gate, and a held bare-JSON buffer that turns out not
to be an enabled call is shown as the answer instead of dropped at stream end.
- The Llama-3 python-tag numeric kwarg regex matched only the mantissa, so
scientific notation was truncated to its leading digits (1e-3 parsed as 1) and a
tool executed with the wrong value. The regex now accepts exponent and decimal
forms, and the int/float classification keys off the exponent too.
Adds regression tests for the truncated / oversized disabled-name JSON cases (and
a counterpart that a truncated enabled call still does not leak) plus the
scientific-notation kwargs.
* Studio: drop accidentally committed async worker transcripts
Eight generated reviewer / async-worker transcripts were committed under
studio/backend/async_task_outputs/. They are not imported or referenced by any
code and carry only internal task state, so they should never ship in the repo.
Remove them and gitignore the directory so they cannot be re-added.
* Studio tools: gate safetensors bare-JSON drain, fix nested-name gate and function-XML strip
Pass-4 review follow-ups on the shared parser / safetensors loop:
- The safetensors oversized and end-of-stream bare-JSON drain branches keyed off
a raw "name" substring, so a large or truncated ordinary JSON answer whose name
is not an enabled tool was drained instead of streamed. Both now use the shared
enabled-tool-name gate, matching the GGUF path.
- strip_leading_bare_json_call matched the first "name" anywhere, so a plain JSON
answer with a nested name equal to an enabled tool ({"result":{"name":"web_search"}})
was wrongly suppressed. It now extracts the TOP-LEVEL name only, walking past
nested objects/arrays and keeping the text when a top-level value is truncated.
- The function-XML display strip used a regex negative-lookahead that stopped at a
literal <function=...> opener inside a parameter value and then dropped the rest
of the answer to EOF. A scan-based strip mirrors the parser (ignores openers
inside an open <parameter> via _inside_open_parameter) and closes each call at its
real </function>, so trailing assistant text after such a call survives.
Adds regression tests for each.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: keep tools prompt when native-template probe raises; make helper tests hermetic
Pass-4 review follow-ups on the native-template fallback:
- render_with_native_template_fallback re-renders the live template with tools=None
to detect whether it dropped the schema. A template that requires tools can raise
on that probe; that must not discard the already-valid tools prompt. The probe is
now wrapped so any error returns the original formatted_prompt (transformers would
otherwise fall back to manual formatting and lose the schema; MLX would let the
exception escape).
- The native-template helper tests imported InferenceBackend just to reach the
thin wrapper, which pulls in unsloth and its optional vllm package metadata. They
now call the dependency-light render_native_template helper directly so they pass
in a backend/test environment without vllm. Adds a probe-raises regression test.
* Tool parsing: 3.9 import safety, disabled-Auto-Heal contract, capability gate
Round-2 review follow-ups on the multi-format tool-call parser:
- tool_call_parser: add `from __future__ import annotations`. The module
is dependency-light by design (external llama-server wrappers import it
standalone) and the package targets python >=3.9, where its PEP 604
`int | None` return annotations would raise TypeError on import.
- safetensors + GGUF drain fallback: gate the leading bare-JSON strip on
auto_heal_tool_calls. With Auto-Heal off, a truncated enabled-name
fragment that did not parse now stays visible, matching the XML strip
in the same branch and the disabled-Auto-Heal contract. With Auto-Heal
on it is still suppressed.
- safetensors capability gate: match the bare-JSON `{"name":` template
marker with a whitespace/escape-tolerant regex so a pretty-printed
`{ "name" :` or JSON-escaped `{\"name\":` template is not mis-classified
as tool-less. The parser already accepts that whitespace via
raw_decode, so the gate must too.
Regression tests added for each case.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* GLM tool-call display strip: treat literal close tag in arg value as data
Round-2 review follow-up on the GLM 4.x tool-call format.
The GLM call shape is <tool_call>NAME<arg_key>k</arg_key><arg_value>v
</arg_value>...</tool_call>. The parser was hardened to walk arg_key /
arg_value pairs so a literal </tool_call> inside an argument value (e.g.
print("</tool_call>")) is treated as data and the call's real close is the
</tool_call> that precedes the next <arg_key>. The display strips still used a
non-greedy <tool_call>.*?</tool_call> regex, which stopped at the literal and
leaked the call's tail into visible content and stale history.
Add _strip_glm_calls, a scan that mirrors the parser's close detection, and run
it before the regex arms in every strip pipeline: the core strip_tool_markup,
the route _strip_tool_xml display/history cleanup, and the safetensors + GGUF
streaming strips. Qwen / Hermes <tool_call>{json} has no NAME token after the
opener, so it is left to the regex arms unchanged.
Regression tests cover the literal-close-tag leak (core + route), normal GLM
calls, back-to-back GLM calls, zero-arg GLM, truncated GLM, and untouched Qwen.
* Tool parsing: symmetric "function" bare-JSON alias and route strip parity
Round-3 review follow-ups, all parser/strip symmetry fixes.
- Bare-JSON "function" alias: the markerless parser accepts a call name via
obj.get("name") or obj.get("function"), but the strip/gates only knew "name",
so a {"function":<enabled tool>} call executed while its raw JSON leaked. Teach
_top_level_bare_json_name the alias (with "name" precedence and the same nested
and truncated-name guards), and widen the guards in strip_leading_bare_json_call,
the safetensors and GGUF _looks_like_enabled_bare_json gates, and the route
capability marker regex.
- Route display/history cleanup: strip a tail-only </param> alias close (the
parser accepts <param name="...">...</param>), and run the parser's guarded
function-XML scan (_inside_open_parameter) before _TOOL_XML_RE so a literal
nested <function=...></function> inside an argument value does not truncate the
strip and leak the tail.
Regression tests added for each.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio tools: fix DeepSeek strict recovery, Kimi dotted names, Gemma spaced streaming
Round 3 review fixes for the DeepSeek / GLM / Kimi tool-call parsing path.
- DeepSeek R1 and V3/V3.1 strict parsing (Auto-Heal off): when a call is
truncated (missing closing fence or <tool_call_end> terminator), skip it
and keep scanning for later well-formed calls instead of breaking out and
dropping the rest of the envelope. This matches the Kimi strict parser's
recovery behaviour.
- Kimi dotted tool names: keep the full name after stripping only the
functions. prefix and :idx suffix, e.g. functions.mcp.server-list:0 stays
mcp.server-list. The previous split on "." truncated dotted MCP names to
their last segment. This matches current vLLM
(tool_id.split(":")[0].removeprefix("functions.")) and SGLang
(^(?:functions\.)?(?P<name>[\w.\-]+):(?P<index>\d+)$).
- Gemma wrapper-less call streaming: hold the whitespace-tolerant prefix
(call : NAME) in the streaming suppression buffer, matching the parser's
_GEMMA_BARE_TC_RE, so the spaced spelling split across chunks is buffered
instead of leaking as visible text. Applied to both the safetensors and
llama.cpp streaming paths.
- Remove dead _render_with_native_template method and the now-unused copy
import from inference.py; the live path uses render_with_native_template_fallback.
Adds regression tests for DeepSeek R1/V3 strict recovery, Kimi full dotted
name preservation, and the Gemma spaced-call streaming suppression.
* Studio tools: honor tool budget in GGUF loop and guard function-XML streaming strip
Round 4 review fixes. Both are asymmetric-fix bugs where the final/steady path got a
guard the analogous streaming/loop path did not.
- GGUF tool-call budget: the safetensors loop counts real tool-call turns against
max_tool_iterations (re-prompt stalls excepted), but the GGUF loop only bounded the
turn count by the enlarged range (max_tool_iterations + _MAX_REPROMPTS). Since this
PR raised _MAX_REPROMPTS from 1 to 3, a model that keeps making valid tool calls
could run up to three extra tool rounds (with max_tool_iterations=1, four rounds
instead of one). Add a _tool_iters_done counter that increments only when a tool
actually executed in the turn, and stop once the caller's budget is spent so the
post-loop final-answer nudge fires. A duplicate/disabled no-op turn is a correction
turn (like a plan-without-action re-prompt) and does not consume budget, preserving
the existing "already completed" re-prompt behavior.
- Streaming display strip: the final strip runs the guarded _strip_function_xml_calls
scanner (a literal <function=...> inside a parameter value is data, not a nested
call), but the GGUF and safetensors streaming strips still used only the open-ended
regex arms. When a tool-call argument contained literal function markup, the regex
tail ate everything to end-of-text and dropped the real trailing prose after the
call's true </function>. Run the guarded scanner (and the balanced Mistral strip)
before the regex arms in both streaming paths so streaming and final display agree.
Adds regression tests: GGUF valid tool calls respect max_tool_iterations, and the
streaming strip keeps trailing prose after a function-XML call with a literal marker.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio tools: safetensors tool budget counts only executed turns (GGUF parity)
Follow-up to the GGUF budget fix. The safetensors loop charged max_tool_iterations
per non-re-prompt iteration (iteration + 1 - reprompt_count), so a duplicate/disabled
no-op turn spent a budget slot even though no tool ran. With a small cap this dropped
real work: for max_tool_iterations=2, a model that made a valid call, repeated it (an
internal no-op correction turn), then made a distinct valid call executed only the
first -- the third turn was sent with no tools and the distinct call was ignored.
Track whether a turn actually executed a tool (set on record_result) and count only
those turns against the cap, matching the GGUF loop. A duplicate/disabled no-op is a
correction turn -- like a plan-without-action re-prompt -- and no longer consumes
budget, so the model still gets its "already completed" nudge and another tool-enabled
turn. Adds a regression test for the small-cap duplicate-then-distinct-call flow.
* Studio tools: fix stale Kimi dotted-name regression test
test_pr5624_regressions.py still expected functions.my.tool:0 to resolve to the last
segment (tool). The parser now preserves the full dotted name (my.tool) after removing
only the functions. prefix and :idx suffix, matching current vLLM/SGLang so dotted MCP
names like mcp.server-list survive. Update the assertion, name, and module docstring to
the corrected contract (the raw id is still preserved on the call).
* Studio: render the reasoning block for safetensors and MLX like GGUF
enable_thinking chat templates (Qwen3/Qwen3.5/GLM) prefill an unclosed <think>
into the generation prompt, so the model emits only the closing </think> then
the answer. The safetensors/MLX chat stream emitted that as plain content, so
the reasoning showed inline with no collapsible thinking block, while GGUF
(which surfaces reasoning via reasoning_content) rendered one. This brings
safetensors and MLX to parity.
- _ResponsesReasoningExtractor gains a reasoning_prefilled mode that starts
inside the reasoning block and splits on the first </think>; default False
keeps GGUF and every existing caller byte-identical. It suppresses a stray
re-emitted <think> and holds partial markers back across chunk boundaries.
- _sf_reasoning_prefill_mode gates the mode on reasoning being enabled for the
request, an enable_thinking or enable_thinking_effort style, and the template
actually using the standard <think>/</think> markers. Models with a bespoke
reasoning channel (e.g. gemma's <|think|>/<|channel>) are excluded so their
answer is never swallowed; gpt-oss (Harmony) and thinking-off requests are
excluded too.
- sf_tool_stream and stream_chunks (the latter also serves MLX) feed text
through the extractor, emitting reasoning_content then content deltas, with a
per-turn reset in the tool loop and a flush before each tool_start; only the
visible delta reaches the monitor reply. The two non-streaming drains split
reasoning_content the same way.
- Tests: extractor prefilled mode (streaming and edge cases), the gate matrix
including the gemma-style exclusion, and a route-replay of the tool-loop
reasoning stream.
* Studio: render the reasoning block for safetensors and MLX like GGUF
enable_thinking chat templates (Qwen3/Qwen3.5/GLM) prefill an unclosed <think>
into the generation prompt, so the model emits only the closing </think> then
the answer. The safetensors/MLX chat stream emitted that as plain content, so
the reasoning showed inline with no collapsible thinking block, while GGUF
(which surfaces reasoning via reasoning_content) rendered one. This brings
safetensors and MLX to parity.
- _ResponsesReasoningExtractor gains a reasoning_prefilled mode that starts
inside the reasoning block and splits on the first </think>; default False
keeps GGUF and every existing caller byte-identical. It suppresses a stray
re-emitted <think> and holds partial markers back across chunk boundaries.
- _sf_reasoning_prefill_mode gates the mode on reasoning being enabled for the
request, an enable_thinking or enable_thinking_effort style, and the template
actually using the standard <think>/</think> markers. Models with a bespoke
reasoning channel (e.g. gemma's <|think|>/<|channel>) are excluded so their
answer is never swallowed; gpt-oss (Harmony) and thinking-off requests are
excluded too.
- sf_tool_stream and stream_chunks (the latter also serves MLX) feed text
through the extractor, emitting reasoning_content then content deltas, with a
per-turn reset in the tool loop and a flush before each tool_start; only the
visible delta reaches the monitor reply. The two non-streaming drains split
reasoning_content the same way.
- Tests: extractor prefilled mode (streaming and edge cases), the gate matrix
including the gemma-style exclusion, and a route-replay of the tool-loop
reasoning stream.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: don't force a tool re-prompt on a negated intent (safetensors parity)
The safetensors _INTENT_SIGNAL claimed to mirror GGUF but was missing the
negative lookahead, so a refusal like "I will not search the web for that"
matched the "i will" intent and triggered the plan-without-action re-prompt
(STOP... you MUST call a tool), overriding a valid no-tool answer. GGUF already
excludes not/never. Add the same (?!\s+(?:not|never)\b) lookahead so both
backends agree. Extends the intent parity test with negated refusals.
* studio: parse the outer envelope before DeepSeek/Kimi markers embedded in its args
parse_tool_calls_from_text ran the DeepSeek/Kimi marker pre-pass before the shared
<tool_call>/<function=...> parser. When a Qwen/Hermes call's argument contained
literal Kimi/DeepSeek markup (for example a user asking the model to explain that
syntax), the pre-pass matched the embedded marker and returned it, executing the
wrong tool and dropping the real call. Skip the pre-pass when a <tool_call> or
<function=...> envelope opens before the first DeepSeek/Kimi marker, so the shared
parser takes the outer call; a genuine marker-led call (no leading envelope) still
goes through the pre-pass. Tests for the embedded-marker case and the control.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: trim redundant comments (comment-only, AST-verified)
* Studio: trim redundant comments (comment-only, AST-verified)
* Studio: prevent Gemma tool-parser DoS on stray delimiters
_gemma_parse_value returned the input index unchanged when text[i] was a
stray delimiter (,}]), so the list and mapping caller loops that advance
on the returned index spun forever at 100% CPU on malformed input such as
[},]. Advance past the delimiter so parsing always terminates.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: strip Magistral [THINK] reasoning from final display/history
strip_tool_markup removed [TOOL_CALLS] and <function> markup but left a
leading Magistral [THINK]...[/THINK] block intact, so its bracket-form
reasoning (not the <think> the reasoning channel renders) leaked into the
safetensors display and conversation history while GGUF/llama.cpp routes
it natively. Drop the leading reasoning block at end-of-turn (final=True)
via the existing _strip_mistral_reasoning helper; streaming is untouched.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: keep times in wrapper-less Gemma tool arguments
The wrapper-less Gemma value scanner used _GEMMA_KEY_RE = [\w.\-]+ for keys,
which also matches a digit-leading token, so a comma followed by a time or
ratio inside a value (call:web_search{query:meet at 10:00, 11:00 tomorrow})
was misread as a new 11: key, truncating the query and injecting a bogus
argument. Require keys to start with a letter or underscore, matching the
identifier-start rule the wrapped path already uses (_GEMMA_NEXT_KEY_RE).
Add a regression test.
* Studio: treat markers/close-tags inside tool-call arguments as data
Four parser correctness fixes where a valid argument string was mistaken for
structure:
- DeepSeek: find the envelope-end token outside JSON strings, so a query/code
argument containing the literal token no longer truncates the body and drops
the whole call.
- GLM: locate the real </arg_value> as the one whose next token is <arg_key> /
</tool_call> / end, so a value containing a literal </arg_value> (or
</tool_call>) is kept instead of executing the tool with corrupted arguments.
- Attribute-form <function name="..."> envelopes now count in the embedded-marker
guard, so a DeepSeek/Kimi marker inside a parameter value does not hijack the
outer call and run the wrong tool.
- Wrapper-less Gemma call:NAME{...} is gated on the enabled tool names (parse and
display strip), mirroring the Llama bare-JSON gate, so a disabled/example name in
prose is not stolen as a call and the real answer is preserved.
Add regression tests for each.
* Gate route Gemma wrapperless strip by enabled tools; make Kimi section-end search string-aware
Route-level display stripping now threads the enabled tool-name set into the
Gemma wrapperless-call strip, so prose that mentions a disabled tool
(call:foo{...}) is preserved while active tool calls are still stripped. This
mirrors the parser-level gate already used in tool_call_parser.
The Kimi section-end lookup now searches outside JSON string literals, so a
section-end marker appearing inside an argument string no longer triggers a
false truncation that drops a valid tool call.
* Run DeepSeek/Kimi pre-pass when a closed tool-call example precedes a real block
The marker pre-pass was skipped whenever any <tool_call>/<function> opener
appeared before the first DeepSeek/Kimi marker, even when that opener was a
CLOSED syntax example in prose that ends before the real block. In that case
parse_tool_calls_from_text skipped the DeepSeek/Kimi parsers and the genuine
tool call was dropped while a phantom tool named in the example ran instead.
Only treat a marker as embedded in a leading envelope when removing the closed
outer <tool_call>/<function> envelopes also removes every marker (the marker
actually sat inside one). A marker left standing is a real call, so the pre-pass
runs. The legitimate case of a marker inside a closed outer envelope's arguments
is preserved.
* Honor reasoning_effort none in safetensors prefill; strip Magistral reasoning while streaming
Two safetensors/MLX reasoning fixes surfaced in review:
_sf_reasoning_prefill_mode only checked enable_thinking, so an
enable_thinking_effort (GLM-5.2) request that disables thinking via
reasoning_effort=none (without enable_thinking=False) still began in
prefilled-<think> mode. A plain answer with no </think> was then swallowed
whole into reasoning_content and the visible response came back empty. Thread
reasoning_effort into the predicate and treat none as disabled, mirroring
_request_reasoning_kwargs.
strip_tool_markup_streaming stripped tool markup but not the leading Magistral
[THINK]...[/THINK] bracket block, so the raw chain-of-thought leaked into the
streamed safetensors content instead of the reasoning drawer (GGUF routes it
natively). Apply _strip_mistral_reasoning first, matching the final strip; an
unclosed [THINK] is held from the marker on so nothing flickers.
* Heal truncated outer tool envelopes and keep quoted Gemma args intact
Two follow-ups from review of the marker pre-pass and Gemma parsing:
The leading-envelope guard only removed CLOSED outer <tool_call>/<function>
envelopes before deciding whether a DeepSeek/Kimi marker was embedded, so a
truncated outer call missing its close tag (whose argument embeds a marker) was
treated as a standalone marker and the embedded sample ran instead of the
intended outer call being Auto-Healed. Decide on the last outer opener before the
marker and whether it closed before the marker instead, so a closed syntax
example still runs the pre-pass while a real closed-or-truncated outer call keeps
it.
The wrapper-less Gemma argument scan tracked bracket depth but not quotes, so a
quoted value containing a comma followed by a key-like token (a search query such
as "weather, location: Boston") was split mid-string, truncating the value and
fabricating an extra argument. Track quote state (with escapes) so the top-level
comma boundary is only taken outside quoted spans.
* Span outer envelopes to their real close when locating embedded markers
Locating the DeepSeek/Kimi marker relative to a leading outer envelope used the
FIRST close tag after the opener, so a literal </function> or </tool_call> inside
an argument value (for example python code that contains the text) was mistaken
for the envelope boundary. The marker after it was then treated as a standalone
call and the embedded sample ran instead of the intended outer call.
Match the closed outer envelopes with the shared patterns that already extend to
the real final close (a literal close inside a value is data), and treat a marker
that survives their removal as embedded only when a still-open (truncated) outer
opener precedes it, so Auto-Heal still repairs a truncated outer call. A closed
syntax example before a genuine block still runs the pre-pass.
* Span the tool_call outer envelope to its real close in the marker guard
The leading-envelope check reused the lazy <tool_call>.*?</tool_call> strip
pattern, so a Qwen/Hermes JSON argument containing a literal </tool_call> ended
the span early. A DeepSeek/Kimi sample later in that same string then survived
the closed-envelope removal, and the pre-pass executed the embedded call instead
of the outer <tool_call>. The <function> arm already spanned to its real close;
give <tool_call> the same real-close pattern (with the negative lookahead that
keeps back-to-back calls separate) so a literal close inside a value is data.
* Preserve no-tool Gemma prose and keep later R1 calls when healing a close
Two review follow-ups:
_gemma_strip_gate returned None when no tools were enabled, and None means
strip every markerless call:NAME{...} block, so a no-tool answer that documents
the syntax (or the Anthropic display path, which passes an empty tool list as
None) had that prose deleted. It is a display/history gate, so return the
enabled-name set instead -- an empty set when no tool is enabled, which strips
nothing because every call:NAME{...} is then prose.
The DeepSeek R1 heal path located the close fence with an unbounded forward
search, so when a first call had balanced JSON but omitted its fence the search
landed on a LATER call's terminator and pos advanced past that valid call,
dropping it. Match the close immediately after the JSON (whitespace-skipped) like
the strict path, and advance by just the JSON when it is absent, so a multi-call
turn keeps its later well-formed calls (heal is now a superset of strict).
* Resume wrapper-less Gemma scan past a consumed call's balanced body
The markerless call:NAME{...} scan used finditer, which resumes right after the
opening call: token, so a nested call:OTHER{...} mentioned inside the first
call's own quoted string argument (for example a web_search query that quotes the
Gemma tool syntax) was re-matched and returned as a spurious second tool call,
executing an unintended tool. Walk with a manual cursor that resumes after the
outer call's balanced body (brace matching already skips quoted braces), so a
call's arguments are never rescanned. Genuinely separate back-to-back calls and
disabled/example prose are unaffected.
* Mistral outer call wins over XML literals; align healer signals with its parser
Two follow-ups on the shared-parser ordering after the healing-passthrough
merge:
- A well-formed [TOOL_CALLS] call whose JSON arguments quote tool XML parsed
the literal instead of the outer call (executing the wrong tool). When the
first XML signal sits inside a leading balanced Mistral body it is argument
data, so the Mistral parser now runs first; an XML signal before the trigger
keeps the normal order, so a [TOOL_CALLS] literal inside an XML call's
arguments still stays data.
- passthrough_healing buffered streams on the parser module's broadened signal
list (now including <|python_tag|> and [TOOL_CALLS]) but promotes with
core.tool_healing, which does not parse those forms: a streamed Mistral or
Llama text call was held until finalization and flushed as prose. The healer
keeps its own signal list limited to the formats it can promote, restoring
immediate streaming for the rest.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: Gemma wrapper-less marker literals and quotes, GLM embedded close pair
- The Gemma fallback deferral now keys on an actual wrapped opener
(_GEMMA_TC_RE), not the wrapper literal anywhere in content: a wrapper-less
call whose argument merely mentions <|tool_call> has nothing tool_healing
can parse, and deferring it lost the call entirely (not executed and
stripped from display).
- New _gemma_body_brace_end boundary scanner honors single- and double-quoted
strings like _gemma_parse_stripped_body, shared by parse and strip, so a
quoted brace in a code argument (code:print('}')) no longer truncates the
executed arguments or the strip span.
- _glm_value_close now requires a structural </arg_value> to sit at balanced
quote state: the full pair </arg_value></tool_call> embedded inside a string
literal is data, not an early close. When no candidate balances, the first
token-valid close wins as before.
* Address review: leading envelopes win over rehearsed literals
- New _first_foreign_tool_signal shared by the leading-envelope guards adds
<|python_tag|> to the protected signal set: the spelled-out literal inside a
Mistral call's arguments (a query about Llama built-in tool syntax) executed
the inner literal instead of the outer call.
- New _xml_signal_inside_leading_bare_json guard, sibling of the Mistral one:
a leading bare-JSON call whose string argument quotes tool XML (a code value
citing <function=...>) had the literal promoted by the shared XML pass
before the bare-JSON parser ran.
- Magistral [THINK]...[/THINK] is dropped once at parse entry instead of only
inside the Mistral parser, so a call rehearsed in the think block in a
foreign format can no longer be promoted while the real call after the
block is lost. Parse now agrees with the display strip.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: a disabled leading bare-JSON object keeps its literals as data
When the leading bare-JSON object is ordinary content (name not an enabled
tool), the guard proved the first tool signal sits inside it, so falling
through to the XML/python_tag passes promoted quoted string data as a real
call. Drop the object and parse only the tail: a real call after the object
still parses, nothing inside it can be promoted.
* Address review: apostrophes in raw Gemma values, GLM strict key contract, per-model template token
- Quote openers in the wrapper-less Gemma boundary and body scanners now
require value-start context (after : { [ ( , =): an apostrophe inside an
unquoted value (query:what's the weather) opened quote mode, swallowed the
real closing brace, and lost the whole call on common contraction queries.
Quoted values keep hiding delimiters as before.
- A GLM <arg_key> with no <arg_value> tag now rejects the call in strict
mode, matching the unclosed-value contract, instead of executing the tool
with the argument silently dropped; Auto-Heal keeps the lenient skip.
- The native-template fallback reads the hf_token stored on the model record
instead of the instance-wide last-load token, so a later token-less load
cannot break template fetches for a previously loaded gated model (both
the transformers and MLX backends).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: Mistral literals inside leading JSON, whitespace-tolerant wrapped Gemma opener
- The leading bare-JSON guard now treats the [TOOL_CALLS] trigger as a
foreign signal: the Mistral parser runs before the bare-JSON one, so a
literal quoted inside the leading object's strings was promoted over the
outer call (or over ordinary JSON content).
- tool_healing's wrapped Gemma opener tolerates whitespace around call and
the colon: sampling drift emits call: name{ and call : name{, and
rejecting those lost the call entirely because no fallback re-parses the
wrapped form. Strict mode still requires the closing tag.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: DeepSeek/Kimi markers inside leading JSON and Mistral envelopes stay data
The DeepSeek/Kimi pre-pass runs before the outer-call parsers, and
_marker_inside_leading_envelope only protected XML envelopes: a marker
quoted inside a leading bare-JSON or Mistral call's argument strings was
promoted as a separate no-arg call and the real outer call dropped. The
guard now recognizes those two leading envelopes as well; standalone
DeepSeek/Kimi calls keep parsing.
* Address review: accept dotted Gemma argument keys in the key-quoting scanner
The scanner quoted keys of [alnum_-] only, so a dotted key (user.name:...)
was left unquoted, json.loads failed, and the whole wrapped call was lost
(parse empty, strip wipes the markup). Dots now match the parser's own
key/name charset.
* Address review: a real DeepSeek/Kimi call after a disabled leading JSON object still parses
DeepSeek/Kimi markers are foreign signals for the leading bare-JSON guard
too: a marker literal inside a disabled leading object made the envelope
guard skip the pre-pass for the whole message, so a real DeepSeek/Kimi call
after the object was dropped. Routing the case through the guard's
drop-and-parse-the-tail recursion reaches the real call while the literal
inside the object stays data.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: leading Mistral call owns the turn, dotted keys after bare values
- A LEADING parseable [TOOL_CALLS] call now runs the Mistral parser first
unconditionally: literal XML in trailing prose after the call was promoted
by the earlier shared XML pass, executing the quoted example instead of
the real leading call. XML leading keeps the normal order.
- _GEMMA_NEXT_KEY_RE accepts dots so a dotted key after a bare value
(query:foo,user.name:bob) ends the value at the comma instead of being
swallowed into it, matching the round-earlier key-quoting charset.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: a leading wrapper-less Gemma call owns the turn
A quoted foreign literal inside a leading wrapper-less Gemma call's
argument (a query citing another tool syntax) was promoted by tool_healing
before the Gemma fallback ran, executing the quoted example and dropping
the outer call. New leading guard, sibling of the Mistral and bare-JSON
ones, gated on an enabled name since the form is markerless. Foreign markup
leading keeps the normal order.
* Fix merge resolution: restore both leading-guard test classes intact
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: markup quoted inside a nameless leading JSON answer stays data
The leading bare-JSON guard required a top-level name, so a structured JSON
answer quoting tool markup in its strings (a response_format turn
documenting a tool's syntax) had the literal promoted by the later passes.
A nameless leading object that parses as real JSON now routes through the
same decline-then-parse-the-tail path; non-JSON braced prose keeps the old
behaviour, and a real call after the answer still parses.
* Address review: JSON answers stay data, nested Gemma quotes, earliest envelope, no failure caching
- A whole-content JSON value is a structured answer: the markerless Gemma
scan and its strip no longer promote or strip a quoted example of an
enabled tool's syntax inside it.
- Nested stripped-stream Gemma values now unquote quoted string leaves
recursively, so {loc:{city:"New York"}} hands the tool New York, matching
the top-level coercion.
- The DeepSeek/Kimi pre-pass dispatches by earliest envelope opener, so a
leading real call wins over a trailing example of the sibling format in
either direction.
- A failed native-template fetch is no longer cached as no-template: the
next call retries after the model record's token is fixed or a transient
Hub error clears; only definitive loads are cached.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: closed calls precede the marker pre-pass, truncated Gemma scan stops, quoted nested delimiters
- A closed non-DeepSeek/Kimi call preceding the first DS/Kimi marker owns
the turn: a trailing syntax example, or one quoted inside a wrapped Gemma
argument, was promoted by the pre-pass and dropped the real leading call.
Wrapped Gemma joins the outer-envelope pattern sets.
- An unbalanced wrapper-less Gemma call now stops the scan (mirroring the
strip contract) instead of resuming inside its own argument text, where a
quoted enabled call would be promoted.
- Raw-quoted strings in nested stripped-stream Gemma values hide delimiters,
so {city:"New, York"} is one value instead of a split pair, returned
unquoted like the top-level coercion.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: string-marker literals in wrapper-less args, mid-value quoted phrases
- The wrapper-less deferral guard no longer keys on the <|"|> literal: a
real call whose argument merely mentions the string marker was deferred to
tool_healing, which has no wrapped opener to parse, losing the call. The
wrapped-opener check alone owns the deferral.
- Double quotes now also open at the start of a word, so a quoted phrase
mid-value (query:find "weather, location: Boston", limit:3) hides its
delimiters instead of splitting the value into garbage keys; apostrophes
keep the value-start-only rule so contractions stay prose.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: strict GLM refuses in-quote close fallback, Gemma guard covers preambles
- _glm_value_close gains a strict flag: a truncated value whose only close
candidates sit inside a string literal rejects the call in strict mode
(Auto-Heal keeps the lenient partial), restoring the strict contract the
quote-aware fallback had weakened.
- The leading wrapper-less Gemma guard no longer requires the call to open
the response: a visible preamble before call:NAME{...} is the normal
shape, and the quoted foreign literal inside the argument was promoted
again in that shape. An enabled balanced call beginning before the first
foreign signal owns it.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: contextual GLM quote openers, disabled Gemma examples stay prose, JSON array answers
- The GLM value-close quote tracker uses the same contextual openers as the
Gemma scanners (single quote after punctuation context, double quote also
at word start), so strict mode accepts a normal apostrophe value again
while still rejecting a truncated value whose only close candidates sit
inside a string literal.
- A disabled wrapper-less Gemma call is prose by design, so a tool literal
quoted inside it no longer promotes: the span is dropped for parsing and
the tail parsed, mirroring the nameless-JSON guard.
- Leading JSON ARRAY answers join the leading-JSON envelope guard, so a
marker quoted inside a structured array response stays data.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Align closed-envelope regression test with the document-order contract
The test asserted the pre-round-13 behavior (trailing DeepSeek/Kimi block
wins over a leading closed envelope) while the shipped rule is document
order: the leading closed call owns the turn. Rename the test and assert
the leading call so the suite matches the contract exercised by
test_leading_xml_call_wins_over_trailing_kimi_example.
* Parse a leading Llama-3.2 bare-JSON call before the markerless Gemma scan
The bare-JSON form only ever matches a leading call object, and document
order says that call owns the turn. Running the Gemma wrapper-less scan
first let an enabled call:NAME{...} snippet quoted inside the leading
call's string arguments steal the turn when the JSON was not the whole
content (trailing prose or a second ;-separated call), executing the
quoted tool instead of the real one. Reordering cannot take a leading
Gemma call's turn since that content never starts with an object brace.
* Leading-call ownership: Mistral trigger in Gemma guards, closed bare JSON before markers, depth-aware nested Gemma values
Three parser gaps against the document-order contract:
The wrapperless Gemma leading guards did not count [TOOL_CALLS] as a
foreign signal, so a leading Gemma call quoting a Mistral snippet in its
argument lost the turn to the quoted literal. Both the enabled-call and
disabled-example guards now include the trigger, matching the bare-JSON
guard's local inclusion.
_marker_inside_leading_envelope required the DeepSeek/Kimi marker to sit
inside the first closed bare-JSON or Mistral call. A marker after that
closed call (a trailing example or data in a later ;-chained call's
strings) now also defers to the leading call, the same inside-or-after
rule the closed XML envelope patterns already applied.
The nested Gemma primitive value scan split on every comma, corrupting
arguments like opts:{code:print(1,2),lang:py}. It now applies the same
paren/brace depth, contextual quote openers, and comma-only-before-a-key
mapping rule as the top-level scan.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Gemma leading guard: a closed enabled call preceding the signal owns the turn
The wrapperless Gemma guard only claimed the turn when the first foreign
signal sat inside the first enabled balanced call. When that call closed
before the signal (a second call quoting a Mistral or Kimi literal, or a
trailing prose example), the guard forfeited the turn and the foreign
parser promoted the quoted literal, dropping the real Gemma calls. Apply
the same inside-or-after ownership rule as the closed bare-JSON and
Mistral envelopes, gated on an enabled name so the name-agnostic legacy
path is unchanged.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Marker guard: only an executable leading bare-JSON call owns the turn
The bare-JSON branch of the leading-envelope marker guard claimed the
turn for any NAMED leading object. A disabled-name object is prose by
design (the bare-JSON parser will not execute it), so deferring the
DeepSeek/Kimi pre-pass to it lost the real later call entirely. Gate the
ownership claim on the enabled set (or the name-agnostic None path). A
marker inside the disabled object's own strings stays data, matching the
tail-exclusion contract; a marker after it now falls through so the
pre-pass parses the real call. The Mistral branch stays ungated since
[TOOL_CALLS] parsing is never name-gated.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Gemma scan skips leading JSON answers; GLM heal bounds values at structural tags
Two fixes to the document-order data contracts:
The markerless Gemma scan only exempted whole-content JSON, so a leading
JSON answer followed by prose had an enabled call:NAME{...} snippet
inside its strings promoted to a real executed call and stripped from
the displayed answer. Both the parse and strip scans now start after a
balanced json-valid leading value span, keeping parse and strip
mirrored. Real calls after the answer still parse; mid-prose JSON gets
no exemption.
The GLM heal fallback for a missing closing arg_value tag took the
entire remainder as the value, executing markup-contaminated arguments
like city="NYC</tool_call>" and swallowing trailing prose. The healed
value now stops at the next arg_key or tool_call close and the pair walk
resumes there. EOF-truncated values keep the partial heal, strict mode
still rejects, and closed values holding a literal close tag in quotes
are untouched.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Compress docstrings in the multi-format tool parser to their contract essence
* Condense parser guard comments and test narration to contract essentials
* verify_import_hoist: exempt __future__ imports and same-diff relocations
Two false positives fired on this PR's refactor. A from __future__ import
is a compiler directive whose name never appears as a runtime load, so
HOISTED-IMPORT-UNUSED can never see it used, yet the file requires it for
PEP 604 annotations on Python 3.9. TARGET-CHANGED flagged the deliberate
move of the strip-pattern constants into core.inference.tool_call_parser
as a silent re-point even though the old module-level target was removed
and the new one added in the same diff. Both get narrow exemptions; a
re-point to a pre-existing target is still caught, and the self-test
negative controls all pass unchanged.
* Leading bare-JSON calls own the turn; function calls end at the first balanced close
The XML-signal guard for a leading bare-JSON call required the signal
strictly inside the object, so a trailing XML example stole the turn
from the leading call; it now applies the same inside-or-after rule as
the Mistral guard. Function-XML calls also ended at the LAST close tag,
which let prose after a closed call that mentions a literal close tag
get swallowed into the final parameter value; calls now end at the
first close tag that is not inside an open parameter, and the strip
mirrors the same rule so parse and strip agree.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Attribute-form calls end at the first balanced close; bare-JSON strip requires the call shape
The attribute form parser still kept the last close tag in the call
window, folding prose after a closed call into the final parameter
value. It now takes the first close not inside an open parameter, the
same rule the equals form and the strip already use.
The leading bare-JSON strip deleted any closed object whose top-level
name matched an enabled tool, including plain JSON answers the parser
correctly rejects as non-calls. The strip (and the drain gate that
delegates to it) now requires the parser's exact call shape, so answers
like {"name":"web_search","result":...} stream and display intact.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* False-alarm markers keep the answer; the bare-JSON strip consumes the whole chain
The trailing strip arms dropped everything from a bare marker to EOF,
so a normal answer that mentions [TOOL_CALLS] or another marker
literally was truncated (or fully swallowed when it started with the
literal) after the no-call drain fallback. Those arms now require a
call-shaped lookahead or marker-at-EOF before dropping; truncated real
calls still strip.
Chained bare-JSON turns executed both calls but stripped only the first
object, so the second call's raw JSON replayed into the next assistant
history message alongside the structured tool_calls. The strip now
consumes the entire chained run of call-shaped enabled objects while
non-call answers, disabled names, and trailing prose stay intact.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* DeepSeek and Kimi trailing strip arms require a call-shaped lookahead
Same false-alarm rule as the bare-word markers: a prose answer that
mentions a DeepSeek or Kimi marker literally keeps its tail, while
truncated real envelopes and bare end-of-text fragments still drop.
* Attribute-form containment, parameter-close-decides rule, preamble-tolerant Mistral guard, strict strip shape
Four document-order and containment fixes. A leading attribute-form
call now parses before the shared XML pass, so markup quoted in its
parameter stays data. The open-parameter scan lets the parameter's own
close tag decide, so any number of literal function closes inside one
value stay data, restoring the pre-close-scan behavior for multi-close
arguments. The leading-Mistral guard tolerates a visible preamble, with
the leading-bare-JSON guard running first so a trigger quoted inside a
leading JSON object stays data. The bare-JSON strip requires the
parser's top-level name in every mode, so nested-name JSON answers
survive name-agnostic stripping.
* Keep buffering long wrapper-less Gemma tool names instead of leaking the prefix
The streaming buffer stopped holding a call:NAME prefix at a fixed
32-char cap, so a Gemma wrapper-less call to a tool whose name exceeds
that (OpenAI allows 64 chars, MCP names run longer) streamed its raw
call:longname text as visible content before the end-of-turn parser
executed it. Hold the variable-length prefix while it still matches the
call: shape, bounded like the bare-JSON path and self-terminating into
prose, draining once the opening brace arrives.
* Keep prose that only mentions DeepSeek/Kimi markers in the route display strip
The route-level _TOOL_XML_RE DeepSeek/Kimi arms consumed from an opener up to
the end of text whenever the marker appeared, so an answer that merely refers
to a marker (for example "See <|tool_call_begin|> in the docs") had the rest
of the reply truncated. The parser-level _TOOL_ALL_PATS already gates these
arms with a call-shaped lookahead. Mirror it here so a marker is only stripped
when a real call follows it or it is a bare fragment at end of text.
* Tighten tool-calling parser and backend comments
* Pass trust_remote_code when reloading native tokenizers
The native-template fallback re-fetches a model's native chat template from
its repo when an Unsloth override template drops the tools schema. The
secondary AutoTokenizer.from_pretrained threaded hf_token but not
trust_remote_code, so for a model loaded with trust_remote_code=True whose
tokenizer repo carries custom code the reload raised, was swallowed, and the
request silently kept the tool-dropping prompt for a model that supports tools.
Store the loaded trust_remote_code on each backend's per-model info dict and
source it in render_native_template, so the reload re-uses exactly the consent
granted at load. For a LoRA adapter the reload targets the base model, whose
remote code was gated and loaded under the same stored flag, so re-passing it
executes no unconsented code. Falsy stored flag preserves the prior behaviour.
Adds a regression test that fails without the flag (custom-code reload raises,
returns None) and passes with it (tools-advertising native prompt returned).
* Treat <|python_tag|> as an outer marker envelope
A Llama-3 <|python_tag|> tool call (built-in NAME.call(...) or custom
{json} form) whose argument quotes a complete DeepSeek/Kimi example was
hijacked by the DeepSeek/Kimi marker pre-pass: the embedded example (for
example delete_all) executed instead of the real outer call. python_tag
is Llama-3's tool-call envelope, so a marker quoted inside its arguments
is data, the same as for <tool_call>, <function=...>, bare JSON, Mistral
and wrapper-less Gemma, which the guard already covers.
Add <|python_tag|> to _OUTER_ENVELOPE_OPEN_RE with a call-shaped
lookahead (mirroring the _TOOL_ALL_PATS python_tag arm) so the marker
pre-pass is suppressed when a python_tag call opens before the first
marker, while a bare prose <|python_tag|> mention is left untouched.
* Tighten tool-call parser comments
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: Daniel Han <info@unsloth.ai>
Co-authored-by: danielhanchen <danielhanchen@users.noreply.github.com>
9842 lines
462 KiB
Python
9842 lines
462 KiB
Python
# 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
|
|
from loggers import get_logger
|
|
import shutil
|
|
import signal
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Callable, Collection, Generator, Iterable, List, Mapping, Optional, Union
|
|
|
|
import httpx
|
|
|
|
from core.inference.llama_server_args import (
|
|
_effective_tensor_parallel,
|
|
_tensor_parallel_matches_loaded,
|
|
extra_args_disable_mmproj,
|
|
parse_cache_override,
|
|
parse_cache_override_per_axis,
|
|
parse_ctx_override,
|
|
parse_split_mode_override,
|
|
resolve_requested_ctx,
|
|
strip_shadowing_flags,
|
|
strip_split_mode_only,
|
|
)
|
|
|
|
# Share strip / signal constants with the multi-format parser so BUFFERING also
|
|
# catches Llama-3 / Mistral / Gemma 4 (legacy helper only knew <tool_call> / <function=).
|
|
from core.inference.tool_call_parser import (
|
|
_GEMMA_BARE_TC_PREFIX_RE,
|
|
_GEMMA_BARE_TC_RE,
|
|
_TOOL_ALL_PATS,
|
|
_balanced_brace_end,
|
|
_strip_function_xml_calls,
|
|
_strip_gemma_wrapperless_calls,
|
|
_strip_glm_calls,
|
|
_strip_mistral_closed_calls,
|
|
TOOL_XML_SIGNALS as _SHARED_TOOL_XML_SIGNALS,
|
|
RAG_MAX_SEARCHES_PER_TURN,
|
|
RAG_SEARCH_CAP_NUDGE,
|
|
parse_tool_calls_from_text as _shared_parse_tool_calls_from_text,
|
|
strip_leading_bare_json_call,
|
|
strip_llama3_leading_sentinels,
|
|
strip_tool_markup as _shared_strip_tool_markup,
|
|
)
|
|
from utils.native_path_leases import child_env_without_native_path_secret
|
|
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
|
|
from utils.subprocess_compat import (
|
|
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
|
)
|
|
from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs
|
|
from core.inference.tool_loop_controller import (
|
|
ToolLoopController,
|
|
tool_event_provenance,
|
|
)
|
|
from state.tool_approvals import (
|
|
TOOL_REJECTED_MESSAGE,
|
|
abort_tool_decision,
|
|
begin_tool_decision,
|
|
new_approval_id,
|
|
wait_tool_decision,
|
|
)
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class LlamaServerNotFoundError(RuntimeError):
|
|
"""GGUF model needs the llama.cpp runtime but no llama-server is installed.
|
|
Subclasses RuntimeError so existing handlers still catch it."""
|
|
|
|
|
|
# Shared so the from_identifier preflight and the load-time raise stay in sync.
|
|
LLAMA_SERVER_NOT_FOUND_DETAIL = (
|
|
"This is a GGUF model, but the llama.cpp runtime (llama-server) is not "
|
|
"installed. Run `unsloth studio setup` to download the prebuilt runtime, "
|
|
"then try again. (Advanced: set LLAMA_SERVER_PATH to an existing binary.)"
|
|
)
|
|
|
|
|
|
# llama-server can serve HTTP 200 while running a model entirely on CPU when a
|
|
# GPU backend fails to init (#5807 / #5106 / #5830). Classify the startup log so
|
|
# Studio can warn. Priority: explicit "offloaded N/M layers to GPU" counts
|
|
# (authoritative), then GPU "model buffer size" lines (host-pinned _Host
|
|
# excluded), then the "device_info:" device table (disconfirm only).
|
|
_GPU_OFFLOAD_MARKERS = (
|
|
"CUDA",
|
|
"ROCm",
|
|
"ROCM",
|
|
"HIP",
|
|
"Metal",
|
|
"Vulkan",
|
|
"OpenCL",
|
|
"SYCL",
|
|
"MUSA",
|
|
"CANN",
|
|
)
|
|
_OFFLOADED_LAYERS_RE = re.compile(
|
|
r"offloaded\s+(\d+)\s*/\s*(\d+)\s+layers?\s+to\s+gpu", re.IGNORECASE
|
|
)
|
|
_DEVICE_ROW_RE = re.compile(
|
|
r"-\s*(CUDA|ROCm|ROCM|HIP|Metal|Vulkan|SYCL|OpenCL|MUSA|CANN|CPU)\w*\s*:",
|
|
re.IGNORECASE,
|
|
)
|
|
_GPU_DEVICE_PREFIXES = (
|
|
"cuda",
|
|
"rocm",
|
|
"hip",
|
|
"metal",
|
|
"vulkan",
|
|
"sycl",
|
|
"opencl",
|
|
"musa",
|
|
"cann",
|
|
)
|
|
|
|
|
|
def classify_gpu_offload_lines(lines: "list[str]") -> Optional[bool]:
|
|
"""True if the model landed on a GPU, False if it stayed on CPU despite GPU
|
|
intent, None when the log has no usable signal."""
|
|
# Counted offload is authoritative, keyed on the model with the most layers.
|
|
# A separate MTP/draft model logs its own (much smaller) "offloaded N/M"
|
|
# line, so decide on the largest-M line: a drafter that fits on GPU must not
|
|
# mask a main model running on CPU. N>0 on that model is True, 0 is False.
|
|
max_total = -1
|
|
offloaded_at_max = 0
|
|
for line in lines:
|
|
match = _OFFLOADED_LAYERS_RE.search(line)
|
|
if not match:
|
|
continue
|
|
offloaded, total = int(match.group(1)), int(match.group(2))
|
|
if total > max_total or (total == max_total and offloaded > offloaded_at_max):
|
|
max_total, offloaded_at_max = total, offloaded
|
|
if max_total >= 0:
|
|
return offloaded_at_max > 0
|
|
|
|
# GPU marker on a *model* buffer; _Host buffers are CPU-pinned, not offload.
|
|
# Buffer lines are authoritative: present but none on a GPU means CPU-only,
|
|
# so do not let the device table below override that.
|
|
saw_model_buffer = False
|
|
for line in lines:
|
|
if "model buffer size" not in line:
|
|
continue
|
|
saw_model_buffer = True
|
|
if "_Host" not in line and any(m in line for m in _GPU_OFFLOAD_MARKERS):
|
|
return True
|
|
if saw_model_buffer:
|
|
return False
|
|
|
|
# device_info: lists *available* devices (printed whenever a GPU backend is
|
|
# visible), not where the model loaded, so it can only disconfirm: an
|
|
# all-CPU table means no usable GPU. A visible GPU device is not proof the
|
|
# model used it, so it does not return True. Rows after the header only.
|
|
after_header = False
|
|
saw_device_row = False
|
|
saw_gpu_device = False
|
|
for line in lines:
|
|
if "device_info:" in line:
|
|
after_header = True
|
|
continue
|
|
if not after_header:
|
|
continue
|
|
match = _DEVICE_ROW_RE.search(line)
|
|
if not match:
|
|
continue
|
|
saw_device_row = True
|
|
if match.group(1).lower().startswith(_GPU_DEVICE_PREFIXES):
|
|
saw_gpu_device = True
|
|
if saw_device_row and not saw_gpu_device:
|
|
return False
|
|
return None
|
|
|
|
|
|
def _wsl_system_rocm_lib_dirs() -> "list[str]":
|
|
"""System ROCm lib dir(s) to load before a prebuilt's bundled HIP, on WSL.
|
|
|
|
The bundled bare-metal HIP can't drive WSL's /dev/dxg and segfaults on the
|
|
first GPU call; the system ROCm libs (libamdhip64 + librocdxg) can, while
|
|
the bundle still supplies libggml-hip / librocblas (gfx1151 kernels).
|
|
Mirrors install_llama_prebuilt._wsl_system_rocm_lib_dirs so a prebuilt that
|
|
passed install validation runs the same at serve time. No-op off a ROCDXG
|
|
WSL host (needs /dev/dxg, "microsoft" /proc/version, librocdxg in /opt/rocm).
|
|
"""
|
|
try:
|
|
if not os.path.exists("/dev/dxg"):
|
|
return []
|
|
with open("/proc/version", encoding = "utf-8", errors = "replace") as fh:
|
|
if "microsoft" not in fh.read().lower():
|
|
return []
|
|
except OSError:
|
|
return []
|
|
out: "list[str]" = []
|
|
for d in ("/opt/rocm/lib", "/opt/rocm/lib64"):
|
|
if os.path.exists(os.path.join(d, "librocdxg.so")) or os.path.exists(
|
|
os.path.join(d, "librocdxg.so.1")
|
|
):
|
|
out.append(d)
|
|
return out
|
|
|
|
|
|
# ── Pre-compiled patterns for plan-without-action re-prompt ──
|
|
# Forward-looking intent signals: the model is describing what it *will*
|
|
# do rather than giving a final answer.
|
|
_INTENT_SIGNAL = re.compile(
|
|
r"(?i)("
|
|
# Direct intent ("I'll ...", "Let me ...", straight + curly apostrophes).
|
|
# Excludes "I can"/"I should"/"I want to"/"let's" (common in answers).
|
|
# Negative lookahead drops negated forms ("I will not") 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
|
|
|
|
# Default max_tokens to the effective context when known. The floor is high
|
|
# enough for reasoning-heavy GGUFs and max_tokens-omitting API clients.
|
|
_DEFAULT_MAX_TOKENS_FLOOR = 32768
|
|
_DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min
|
|
|
|
# Only large streamed tool payloads get an early provisional card; render_html
|
|
# is exempt because it needs immediate artifact feedback.
|
|
_PROVISIONAL_ARGS_MIN_CHARS = 256
|
|
_DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0 # 2 min
|
|
_REPROMPT_MAX_CHARS = 2000
|
|
# Cap tool calls from a single TEXTUAL-fallback turn (mirrors the safetensors
|
|
# loop). Structured delta.tool_calls are grammar-bounded by llama-server; text
|
|
# parsed from content is not, so one runaway turn could fan out unbounded.
|
|
_MAX_TOOL_CALLS_PER_TURN = 8
|
|
_FORCED_REPEAT_PLAN_SIGNAL = re.compile(
|
|
r"\b(?:i\s+will|i'll|let\s+me|going\s+to|need\s+to|call|use|run|search|fetch|render)\b",
|
|
re.I,
|
|
)
|
|
_FINAL_ANSWER_SIGNAL = re.compile(
|
|
r"\b(?:final\s+answer|answer\s*:|here\s+is|here's|in\s+summary|result\s*:)\b",
|
|
re.I,
|
|
)
|
|
|
|
|
|
def _is_short_intent_without_action(text: str) -> bool:
|
|
stripped = text.strip()
|
|
return 0 < len(stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(stripped) is not None
|
|
|
|
|
|
def _should_suppress_forced_no_tool_output(text: str) -> bool:
|
|
"""Suppress only repeated forced-turn planning text, not final answers."""
|
|
stripped = text.strip()
|
|
if not stripped or len(stripped) >= _REPROMPT_MAX_CHARS:
|
|
return False
|
|
if _FINAL_ANSWER_SIGNAL.search(stripped):
|
|
return False
|
|
return _FORCED_REPEAT_PLAN_SIGNAL.search(stripped) is not None
|
|
|
|
|
|
# ── Pre-compiled patterns for GGUF shard detection ───────────
|
|
_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$", re.IGNORECASE)
|
|
_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$", re.IGNORECASE)
|
|
|
|
|
|
# ── 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 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 on a daemon thread, so concurrent sockets aren't
|
|
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 -> dead.
|
|
return True if result[0] is None else result[0]
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _hf_offline_if_dns_dead():
|
|
"""Set HF_HUB_OFFLINE for this block only when DNS to huggingface.co fails;
|
|
restores env on exit so a transient hiccup can't quarantine the process.
|
|
No-op if the user already set it."""
|
|
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, else None."""
|
|
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 _swa_entry_from_layer_types(lt) -> Optional[object]:
|
|
"""Period int, or per-layer bool mask, from a transformers ``layer_types`` list."""
|
|
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 _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
|
|
return _swa_entry_from_layer_types(src.get("layer_types"))
|
|
|
|
|
|
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
|
|
return _swa_entry_from_layer_types(getattr(src, "layer_types", 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 (result persistently cached)
|
|
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(/...)` -> `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]}"
|
|
|
|
|
|
# Lazy import to avoid pulling transformers in at module level.
|
|
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'",
|
|
# DeepSeek: no top-level ``{% if tools %}`` block; it gates emission on
|
|
# ``message['role'] == 'tool'`` plus ``message['tool_calls'] is defined``.
|
|
"message['role'] == 'tool'",
|
|
'message["role"] == "tool"',
|
|
"message['tool_calls']",
|
|
'message["tool_calls"]',
|
|
"tool_calls is defined",
|
|
)
|
|
|
|
|
|
# Canonical reasoning_effort levels, weakest -> strongest. Used to read the
|
|
# discrete set a template branches on (e.g. GLM-5.2 uses 'high' | 'max') so we
|
|
# only ever offer levels the template actually understands.
|
|
_REASONING_EFFORT_SCALE = ("minimal", "low", "medium", "high", "max")
|
|
|
|
|
|
def _extract_reasoning_effort_levels(chat_template: str) -> list:
|
|
"""Return the reasoning_effort levels a template references, in canonical
|
|
(weakest -> strongest) order.
|
|
|
|
Looks for the quoted literals (e.g. ``'high'`` / ``"max"``) the template
|
|
compares ``reasoning_effort`` against, so we surface exactly the levels it
|
|
branches on and nothing else.
|
|
"""
|
|
return [
|
|
level
|
|
for level in _REASONING_EFFORT_SCALE
|
|
if f"'{level}'" in chat_template or f'"{level}"' in chat_template
|
|
]
|
|
|
|
|
|
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 the same six keys as the GGUF sniffer: ``supports_reasoning``,
|
|
``reasoning_style`` (``"enable_thinking"`` | ``"reasoning_effort"`` |
|
|
``"enable_thinking_effort"``), ``reasoning_always_on``,
|
|
``reasoning_effort_levels``, ``supports_preserve_thinking``,
|
|
``supports_tools``. A falsy ``chat_template`` yields the all-default dict.
|
|
Used by both the llama-server backend at load time and the
|
|
safetensors/transformers paths in ``routes/inference`` so they agree on
|
|
what the frontend sees.
|
|
"""
|
|
flags = {
|
|
"supports_reasoning": False,
|
|
"reasoning_style": "enable_thinking",
|
|
"reasoning_always_on": False,
|
|
"reasoning_effort_levels": [],
|
|
"supports_preserve_thinking": False,
|
|
"supports_tools": False,
|
|
}
|
|
if not chat_template:
|
|
return flags
|
|
tpl = chat_template
|
|
prefix = f"{log_source}: " if log_source else ""
|
|
|
|
effort_levels = (
|
|
_extract_reasoning_effort_levels(tpl)
|
|
if ("reasoning_effort" in tpl and "enable_thinking" in tpl)
|
|
else []
|
|
)
|
|
if effort_levels:
|
|
# GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort
|
|
# level among a discrete set (e.g. 'high' | 'max'). Distinct from
|
|
# gpt-oss (reasoning_effort only, no on/off gate) and Qwen
|
|
# (enable_thinking only). Disabling is enable_thinking=false; the levels
|
|
# are the quoted effort literals the template actually branches on.
|
|
flags["supports_reasoning"] = True
|
|
flags["reasoning_style"] = "enable_thinking_effort"
|
|
flags["reasoning_effort_levels"] = effort_levels
|
|
logger.info(
|
|
f"{prefix}model supports reasoning "
|
|
f"(enable_thinking + reasoning_effort: {effort_levels})"
|
|
)
|
|
elif "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 use reasoning_effort
|
|
# ("low" | "medium" | "high"), not 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', not '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).
|
|
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: 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
|
|
|
|
|
|
# Gemma 4 ships MTP as a separate drafter (no "-mtp" in the name). Gemma 3n
|
|
# ships no drafter, so it is excluded -- it takes the normal non-MTP path.
|
|
_GEMMA_MTP_FAMILY_RE = re.compile(r"gemma[-_]?4[-_]", re.IGNORECASE)
|
|
|
|
|
|
def _is_gemma_mtp_family(name: Optional[str]) -> bool:
|
|
"""Match Gemma 4 by name."""
|
|
return bool(name) and bool(_GEMMA_MTP_FAMILY_RE.search(name))
|
|
|
|
|
|
def _is_gemma_mtp_name(model_identifier: Optional[str], gguf_path: Optional[str] = None) -> bool:
|
|
"""Match Gemma 4 by id or GGUF filename."""
|
|
return _is_gemma_mtp_family(model_identifier) or _is_gemma_mtp_family(
|
|
Path(gguf_path).name if gguf_path else None
|
|
)
|
|
|
|
|
|
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
|
|
# Recognise Gemma 4 too, so a failed drafter download surfaces a
|
|
# fallback reason instead of silently defaulting.
|
|
if cand and _is_gemma_mtp_family(cand):
|
|
return True
|
|
return False
|
|
|
|
|
|
def _is_companion_gguf_path(path: str) -> bool:
|
|
"""True for a non-main GGUF: vision mmproj or a separate MTP drafter
|
|
(repo-root ``mtp-*.gguf`` or the ``MTP/`` subdir copies, Gemma 4).
|
|
|
|
Mirrors hub.utils.gguf so variant resolution never picks a companion as
|
|
the main model -- e.g. a Gemma ``Q8_0`` request must not resolve to the
|
|
``MTP/...-Q8_0-MTP.gguf`` drafter, which sorts ahead of the real weight.
|
|
"""
|
|
p = path.lower()
|
|
if not p.endswith(".gguf"):
|
|
return False
|
|
if "mmproj" in p:
|
|
return True
|
|
name = p.rsplit("/", 1)[-1]
|
|
return name.startswith("mtp-") or "/mtp/" in f"/{p}"
|
|
|
|
|
|
_BIG_ENDIAN_GGUF_FILENAME_RE = re.compile(r"(^|[-_])be(?:[._-]|$)", re.IGNORECASE)
|
|
_GGUF_KNOWN_QUANT_RE = re.compile(
|
|
r"(UD-)?"
|
|
r"(MXFP[0-9]+(?:_[A-Z0-9]+)*"
|
|
r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?"
|
|
r"|TQ[0-9]+_[0-9]+"
|
|
r"|Q[0-9]+_K_[A-Z]+"
|
|
r"|Q[0-9]+_[0-9]+"
|
|
r"|Q[0-9]+_K"
|
|
r"|BF16|F16|F32)",
|
|
re.IGNORECASE,
|
|
)
|
|
|
|
|
|
def _is_big_endian_gguf_path(path: str, variant_key: str = "") -> bool:
|
|
normalized = path.replace("\\", "/")
|
|
name = normalized.rsplit("/", 1)[-1]
|
|
stem = name.rsplit(".", 1)[0].lower()
|
|
variant_key = variant_key.strip().lower()
|
|
variant_index = stem.find(variant_key) if variant_key else -1
|
|
parent = normalized.rsplit("/", 1)[0].lower() if "/" in normalized else ""
|
|
variant_in_parent_only = (
|
|
bool(parent)
|
|
and variant_index < 0
|
|
and (
|
|
(variant_key and variant_key in parent)
|
|
or (not variant_key and _GGUF_KNOWN_QUANT_RE.search(parent) is not None)
|
|
)
|
|
)
|
|
for match in _BIG_ENDIAN_GGUF_FILENAME_RE.finditer(stem):
|
|
if variant_index >= 0 and variant_index < match.start():
|
|
return True
|
|
tail = stem[match.end() :].lstrip("._-")
|
|
if not tail or _GGUF_KNOWN_QUANT_RE.search(tail) is None:
|
|
return not variant_in_parent_only
|
|
return False
|
|
|
|
|
|
def _gguf_snapshot_files(snapshot: Path) -> list[str]:
|
|
return [
|
|
p.relative_to(snapshot).as_posix()
|
|
for p in snapshot.rglob("*")
|
|
if p.is_file() and p.name.lower().endswith(".gguf")
|
|
]
|
|
|
|
|
|
def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]:
|
|
m = _SHARD_FULL_RE.match(first_shard)
|
|
if not m:
|
|
return []
|
|
prefix = m.group(1)
|
|
total = m.group(3)
|
|
sibling_pat = re.compile(
|
|
r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(total) + r"\.gguf$",
|
|
re.IGNORECASE,
|
|
)
|
|
return sorted(f for f in files if f != first_shard and sibling_pat.match(f))
|
|
|
|
|
|
def _gguf_files_for_variant(files: Iterable[str], variant: str) -> list[str]:
|
|
"""Return main GGUF files matching a requested variant.
|
|
|
|
Prefer exact quant-label matches over loose substring matches so a request
|
|
for ``stories260K`` does not resolve to ``stories260K-be.gguf``.
|
|
"""
|
|
variant_key = variant.strip().lower()
|
|
main_files = [
|
|
f
|
|
for f in files
|
|
if f.lower().endswith(".gguf")
|
|
and not _is_companion_gguf_path(f)
|
|
and not _is_big_endian_gguf_path(f, variant_key)
|
|
]
|
|
if not variant_key:
|
|
return sorted(main_files)
|
|
|
|
try:
|
|
from utils.models.model_config import _extract_quant_label
|
|
except Exception:
|
|
_extract_quant_label = None
|
|
|
|
if _extract_quant_label is not None:
|
|
try:
|
|
exact = sorted(f for f in main_files if _extract_quant_label(f).lower() == variant_key)
|
|
if exact:
|
|
return exact
|
|
except Exception as e:
|
|
logger.warning("Failed to extract GGUF quant labels: %s", e)
|
|
|
|
boundary = re.compile(r"(?<![a-zA-Z0-9])" + re.escape(variant_key) + r"(?![a-zA-Z0-9])")
|
|
return sorted(f for f in main_files if boundary.search(f.lower()))
|
|
|
|
|
|
# Below this many B params, draft-mtp regresses vs spec-off (bench in
|
|
# _build_speculative_flags); auto mode drops MTP under it.
|
|
_MTP_MIN_SIZE_B = 3.0
|
|
|
|
# Cap total GPU occupancy at this fraction of the card. The fit reserves an
|
|
# absolute (1 - frac) * total per GPU when total VRAM is known, else a fraction
|
|
# of free (see _fit_context_to_vram), plus a byte-accurate MTP draft reserve.
|
|
# 3%: the context-linear compute buffer is now modelled (_compute_buffer_ctx_bytes),
|
|
# so this cushion no longer covers it - only fragmentation, the per-device CUDA
|
|
# context on a multi-GPU split, and MoE routing, which measure ~2-3% (Qwen3.5-397B on
|
|
# 3 GPUs under-predicts by 2.7%). Below 3% one fragmentation spike overflows to CPU.
|
|
_CTX_FIT_VRAM_FRACTION = 0.97
|
|
|
|
# Apple unified memory is shared with the OS, so tighter than VRAM. Matches the
|
|
# 0.85 MLX uses in mlx_inference.py (_configure_memory_limits); not kept in sync.
|
|
_APPLE_UNIFIED_MEMORY_FRACTION = 0.85
|
|
|
|
# Flat MTP reserve, used only when GGUF dims are too sparse for the byte-accurate
|
|
# reserve (_estimate_mtp_overhead_bytes). Applied to both the fit budget and pin.
|
|
_MTP_VRAM_RESERVE_FRAC = 0.05
|
|
|
|
|
|
def _kv_bytes_per_elem(cache_type: Optional[str]) -> float:
|
|
"""Bytes per KV-cache element for a llama.cpp cache type (f16 default)."""
|
|
return {
|
|
"f32": 4.0,
|
|
"f16": 2.0,
|
|
"bf16": 2.0,
|
|
"q8_0": 34 / 32,
|
|
"q5_1": 0.75,
|
|
"q5_0": 0.6875,
|
|
"q4_1": 0.625,
|
|
"q4_0": 0.5625,
|
|
"iq4_nl": 0.5625,
|
|
}.get((cache_type or "f16").strip().lower(), 2.0)
|
|
|
|
|
|
def _env_main_cache_type_for_budget(env: Optional[Mapping[str, str]] = None) -> Optional[str]:
|
|
"""Heavier of the inherited LLAMA_ARG_CACHE_TYPE_K/_V env types when it
|
|
exceeds the f16 default, else None. Studio emits --cache-type only for the
|
|
param/extras path, so a heavier env (f32) would otherwise reach the child
|
|
unbudgeted; quantized env types stay over-reserved by f16 (-> None)."""
|
|
e = os.environ if env is None else env
|
|
f16_bpe = _kv_bytes_per_elem("f16")
|
|
heaviest: Optional[str] = None
|
|
heaviest_bpe = f16_bpe
|
|
for var in ("LLAMA_ARG_CACHE_TYPE_K", "LLAMA_ARG_CACHE_TYPE_V"):
|
|
raw = (e.get(var) or "").strip().lower()
|
|
if not raw:
|
|
continue
|
|
bpe = _kv_bytes_per_elem(raw)
|
|
if bpe > heaviest_bpe:
|
|
heaviest, heaviest_bpe = raw, bpe
|
|
return heaviest
|
|
|
|
|
|
def _extra_args_main_cache_type_for_budget(extra_args: Optional[Iterable[str]]) -> Optional[str]:
|
|
"""Heavier (max bytes/elem) of the explicit --cache-type-k/-v extras, or None.
|
|
|
|
Extras are appended last and win per axis, so an asymmetric K=f32,V=f16 must be
|
|
budgeted by its heavier axis. resolve_cache_type_kv returns only the last-wins
|
|
single type, which under-reserves the heavier axis when the lighter one is last."""
|
|
k, v = parse_cache_override_per_axis(extra_args)
|
|
candidates = [c for c in (k, v) if c]
|
|
if not candidates:
|
|
return None
|
|
return max(candidates, key = _kv_bytes_per_elem)
|
|
|
|
|
|
def _auto_mode_drops_mtp(
|
|
req_mode: Optional[str],
|
|
size_b: Optional[float],
|
|
*,
|
|
has_separate_drafter: bool = False,
|
|
) -> bool:
|
|
"""Auto mode drops MTP below _MTP_MIN_SIZE_B for an embedded draft head
|
|
(its per-token cost regresses there); a separate drafter (Gemma) is a tiny
|
|
standalone model that still speeds up below 3B, so it never drops. Forced
|
|
mtp / mtp+ngram engage regardless of size."""
|
|
if has_separate_drafter:
|
|
return False
|
|
return req_mode == "auto" and size_b is not None and size_b < _MTP_MIN_SIZE_B
|
|
|
|
|
|
def _mla_mtp_auto_enabled() -> bool:
|
|
"""Whether Auto may pick embedded MTP for an MLA model (GLM-5.2/DeepSeek/Kimi).
|
|
|
|
Off by default: llama.cpp's MLA/DSA MTP path keeps a duplicated full target-KV
|
|
context and recomputes the sparse-attention indexer every draft step, so it runs
|
|
~2x slower than no speculation (GLM-5.2 bench: 27 vs 45 tok/s, flat across draft
|
|
depth and 96-100% acceptance) -- the opposite of the vLLM/SGLang speedup on the
|
|
same model. Set UNSLOTH_MLA_MTP_ENABLED=1 to let Auto promote MLA MTP again once
|
|
that path is optimized upstream. Forced mtp / mtp+ngram ignore this gate."""
|
|
return os.environ.get("UNSLOTH_MLA_MTP_ENABLED", "0").strip().lower() in (
|
|
"1",
|
|
"true",
|
|
"yes",
|
|
"on",
|
|
)
|
|
|
|
|
|
def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool:
|
|
"""User passed --spec-type / --spec-default? llama-server takes one
|
|
--spec-type (comma-separated to chain), so suppress auto-emit."""
|
|
return _extra_args_set_any_flag(extra_args, {"--spec-type", "--spec-default"})
|
|
|
|
|
|
_GPU_OFFLOAD_OVERRIDE_FLAGS = frozenset({"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"})
|
|
_THREAD_OVERRIDE_FLAGS = frozenset({"-t", "--threads"})
|
|
|
|
|
|
def _extra_arg_flag_name(token: str) -> Optional[str]:
|
|
if not token.startswith("-") or token in {"-", "--"}:
|
|
return None
|
|
if len(token) >= 2 and (token[1].isdigit() or token[1] == "."):
|
|
return None
|
|
return token.split("=", 1)[0]
|
|
|
|
|
|
def _extra_args_set_any_flag(extra_args: Optional[Iterable[str]], flags: Collection[str]) -> bool:
|
|
if not extra_args:
|
|
return False
|
|
for raw in extra_args:
|
|
flag = _extra_arg_flag_name(str(raw))
|
|
if flag in flags:
|
|
return True
|
|
return False
|
|
|
|
|
|
def _effective_spec_type(
|
|
extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
|
|
) -> Optional[str]:
|
|
"""The --spec-type llama-server will use: the last CLI --spec-type (or
|
|
--spec-default, which resolves non-MTP), else LLAMA_ARG_SPEC_TYPE. A CLI flag
|
|
overrides the env (matching llama.cpp), so a stale MTP env can't make the
|
|
budget reserve a drafter the launch won't load. None if neither sets it."""
|
|
args = [str(a) for a in extra_args] if extra_args else []
|
|
cli_present = False
|
|
cli_value: Optional[str] = None
|
|
for i, raw in enumerate(args):
|
|
flag, eq, inline = raw.partition("=")
|
|
if flag == "--spec-default":
|
|
cli_present = True
|
|
cli_value = "default"
|
|
continue
|
|
if flag != "--spec-type":
|
|
continue
|
|
cli_present = True
|
|
cli_value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
|
|
if cli_present:
|
|
return cli_value
|
|
return (os.environ if env is None else env).get("LLAMA_ARG_SPEC_TYPE")
|
|
|
|
|
|
def _extra_args_requests_mtp(
|
|
extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
|
|
) -> bool:
|
|
"""True if the effective --spec-type selects MTP (mtp/draft-mtp), so the
|
|
budget must reserve for it."""
|
|
value = _effective_spec_type(extra_args, env)
|
|
if not value:
|
|
return False
|
|
return any(p.strip().lower() in ("mtp", "draft-mtp") for p in value.split(","))
|
|
|
|
|
|
def _extra_args_requests_separate_draft(
|
|
extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
|
|
) -> bool:
|
|
"""True if the effective --spec-type selects a non-MTP model draft mode
|
|
(draft-simple/draft-eagle3), which loads a separate draft model the budget
|
|
must reserve (draft-mtp -> _extra_args_requests_mtp; ngram-* load no model)."""
|
|
value = _effective_spec_type(extra_args, env)
|
|
if not value:
|
|
return False
|
|
return any(p.strip().lower() in ("draft-simple", "draft-eagle3") for p in value.split(","))
|
|
|
|
|
|
def _extra_args_spec_draft_n_max(extra_args: Optional[Iterable[str]]) -> Optional[int]:
|
|
"""Draft depth from extras (``--spec-draft-n-max`` or legacy ``--draft-max``), else None."""
|
|
if not extra_args:
|
|
return None
|
|
args = [str(a) for a in extra_args]
|
|
found: Optional[int] = None
|
|
for i, raw in enumerate(args):
|
|
flag, eq, inline = raw.partition("=")
|
|
if flag not in ("--spec-draft-n-max", "--draft-max"):
|
|
continue
|
|
value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
|
|
try:
|
|
found = int(value)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
return found
|
|
|
|
|
|
def _extra_args_mtp_draft_path(
|
|
extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
|
|
) -> Optional[str]:
|
|
"""Separate drafter path from extras (local --model-draft/-md or HF
|
|
--spec-draft-hf/-hfd/...), else the LLAMA_ARG_SPEC_DRAFT_MODEL/_HF_REPO env,
|
|
else None. An HF repo isn't a local file, so the budget can't size it (falls
|
|
back to the flat reserve), but recognizing it avoids sizing the wrong one."""
|
|
flags = {
|
|
"--model-draft",
|
|
"--spec-draft-model",
|
|
"-md",
|
|
"--spec-draft-hf",
|
|
"-hfd",
|
|
"-hfrd",
|
|
"--hf-repo-draft",
|
|
}
|
|
args = [str(a) for a in extra_args] if extra_args else []
|
|
found: Optional[str] = None
|
|
for i, raw in enumerate(args):
|
|
flag, eq, inline = raw.partition("=")
|
|
if flag not in flags:
|
|
continue
|
|
value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
|
|
if value and not value.startswith("-"):
|
|
found = value
|
|
if found is not None:
|
|
return found
|
|
e = os.environ if env is None else env
|
|
return e.get("LLAMA_ARG_SPEC_DRAFT_MODEL") or e.get("LLAMA_ARG_SPEC_DRAFT_HF_REPO") or None
|
|
|
|
|
|
def _extra_args_draft_cache_types(
|
|
extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
|
|
) -> tuple[Optional[str], Optional[str]]:
|
|
"""Draft KV cache types (k_type, v_type), each from extras else the
|
|
LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K/_V env, else None (f16). K and V are
|
|
independent: a one-sided override must not apply to both."""
|
|
args = [str(a) for a in extra_args] if extra_args else []
|
|
k_flags = {"--cache-type-k-draft", "--spec-draft-type-k", "-ctkd"}
|
|
v_flags = {"--cache-type-v-draft", "--spec-draft-type-v", "-ctvd"}
|
|
k_type: Optional[str] = None
|
|
v_type: Optional[str] = None
|
|
for i, raw in enumerate(args):
|
|
flag, eq, inline = raw.partition("=")
|
|
if flag not in k_flags and flag not in v_flags:
|
|
continue
|
|
value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
|
|
if not value or value.startswith("-"):
|
|
continue
|
|
if flag in k_flags:
|
|
k_type = value
|
|
else:
|
|
v_type = value
|
|
e = os.environ if env is None else env
|
|
if k_type is None:
|
|
k_type = e.get("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K") or None
|
|
if v_type is None:
|
|
v_type = e.get("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V") or None
|
|
return k_type, v_type
|
|
|
|
|
|
def _extra_args_draft_offloaded_to_cpu(
|
|
extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
|
|
) -> bool:
|
|
"""True if the SEPARATE draft model is on CPU (so the budget must not charge
|
|
its weights+KV): --spec-draft-ngl 0, or --spec-draft-device naming only
|
|
cpu/none, else the LLAMA_ARG_N_GPU_LAYERS_DRAFT env the child honors (the
|
|
device flag has no env). An embedded MTP head follows the main -ngl, so these
|
|
draft-only flags don't move it. Last-wins, so only each flag's final value counts."""
|
|
ngl_flags = {"--spec-draft-ngl", "-ngld", "--gpu-layers-draft", "--n-gpu-layers-draft"}
|
|
dev_flags = {"--spec-draft-device", "-devd", "--device-draft"}
|
|
args = [str(a) for a in extra_args] if extra_args else []
|
|
last_ngl: Optional[str] = None
|
|
last_dev: Optional[str] = None
|
|
for i, raw in enumerate(args):
|
|
flag, eq, inline = raw.partition("=")
|
|
value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
|
|
if flag in ngl_flags:
|
|
last_ngl = value
|
|
elif flag in dev_flags:
|
|
last_dev = value
|
|
if last_ngl is None:
|
|
last_ngl = (os.environ if env is None else env).get("LLAMA_ARG_N_GPU_LAYERS_DRAFT")
|
|
if last_ngl is not None:
|
|
try:
|
|
if int(last_ngl) == 0:
|
|
return True
|
|
except (TypeError, ValueError):
|
|
pass
|
|
if last_dev is not None:
|
|
devs = [d.strip().lower() for d in last_dev.split(",") if d.strip()]
|
|
if devs and all(d in ("cpu", "none") for d in devs):
|
|
return True
|
|
return False
|
|
|
|
|
|
def _extra_args_n_ubatch(
|
|
extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
|
|
) -> Optional[int]:
|
|
"""Physical micro-batch from extras (--ubatch-size/-ub) else the LLAMA_ARG_UBATCH
|
|
env, else None. It sizes the compute-graph buffer, so an override must reach
|
|
the VRAM reserve."""
|
|
args = [str(a) for a in extra_args] if extra_args else []
|
|
found: Optional[int] = None
|
|
for i, raw in enumerate(args):
|
|
flag, eq, inline = raw.partition("=")
|
|
if flag not in ("--ubatch-size", "-ub"):
|
|
continue
|
|
value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
|
|
try:
|
|
found = int(value)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
if found is not None:
|
|
return found
|
|
raw = (os.environ if env is None else env).get("LLAMA_ARG_UBATCH")
|
|
if raw:
|
|
try:
|
|
return int(raw)
|
|
except (TypeError, ValueError):
|
|
pass
|
|
return None
|
|
|
|
|
|
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 legacy ``--spec-ngram-size-n`` /
|
|
``--draft-min`` / ``--draft-max``. ``caps`` comes from
|
|
``probe_server_capabilities``; ``ngram_mod_flavor`` says which set is
|
|
real (vs a removal-stub). Returns ``[]`` when neither 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":
|
|
# Pre-rename llama.cpp: same knobs lived under --spec-ngram-size-n
|
|
# (lookup length) and generic --draft-min / --draft-max (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.
|
|
# Dropdown renders five (auto, mtp, ngram, mtp+ngram, off); the load API
|
|
# also accepts legacy values 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 ``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]
|
|
# Old persisted state emits llama.cpp comma-chains 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 final SSE chunk's ``usage`` can be absent
|
|
or zero on some server builds / streaming configs, making 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
|
|
|
|
|
|
def _is_external_link(path: Path) -> bool:
|
|
"""True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink
|
|
or a Windows directory junction / reparse point. Such a link resolves into
|
|
the user's own llama.cpp checkout, which Studio does not own."""
|
|
try:
|
|
if os.path.islink(path):
|
|
return True
|
|
except OSError:
|
|
return False
|
|
if os.name == "nt":
|
|
try:
|
|
import stat
|
|
attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined]
|
|
return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT)
|
|
except (OSError, AttributeError):
|
|
return False
|
|
return False
|
|
|
|
|
|
class LlamaCppBackend:
|
|
"""Manages a llama-server subprocess for GGUF model inference.
|
|
|
|
Lifecycle:
|
|
1. load_model(): start llama-server with the GGUF file
|
|
2. generate_chat_completion(): proxy to /v1/chat/completions, stream back
|
|
3. unload_model(): terminate the 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
|
|
# Separate MTP drafter launched with the current model; reload-dedup
|
|
# key so a drafter that appears next to the weights forces a reload.
|
|
self._mtp_draft_path: Optional[str] = None
|
|
# Why MTP was disabled on the last load that asked for it (auto on an
|
|
# MTP model, or forced mtp / mtp+ngram), else None. Drives the "update
|
|
# llama.cpp" hint in the UI. "binary_no_mtp" / "binary_outdated" ->
|
|
# a newer prebuilt would help; "runtime_error" -> it may not.
|
|
self._spec_fallback_reason: Optional[str] = None
|
|
self._hf_variant: Optional[str] = None
|
|
self._is_vision: bool = False
|
|
# Block-diffusion model (e.g. DiffusionGemma): served by the diffusion
|
|
# runner, not llama-server. Set from the GGUF architecture at load.
|
|
self._architecture: Optional[str] = None
|
|
self._is_diffusion: bool = False
|
|
self._diffusion_visual_bin: Optional[str] = None
|
|
self._healthy = False
|
|
self._load_rss_hwm = (None, 0) # (pid, peak VmRSS) for load_progress
|
|
self._stats_logger = None # vLLM-style engine-stats poller, set on load
|
|
# 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._reasoning_effort_levels: list = []
|
|
self._supports_preserve_thinking: bool = False
|
|
self._supports_tools: bool = False
|
|
self._cache_type_kv: Optional[str] = None
|
|
# Whether --split-mode tensor was applied on the active load.
|
|
self._tensor_parallel: bool = False
|
|
# Layer load kept multi-GPU only to honor a downgraded tensor request, so a
|
|
# later explicit tensor-off reloads instead of deduping to it (#6659).
|
|
self._layer_preserves_tensor_intent: bool = False
|
|
self._reasoning_default: bool = True
|
|
self._speculative_type: Optional[str] = None
|
|
# Canonical UI-facing mode the user requested
|
|
# (auto/mtp/ngram/mtp+ngram/off/ngram-simple). Round-tripped through the
|
|
# status API so the dropdown reflects the picked mode, not the resolved
|
|
# flag set (auto on a 27B MTP GGUF resolves to draft-mtp but reads "Auto").
|
|
self._requested_spec_mode: Optional[str] = None
|
|
# User --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
|
|
# For the compute-graph buffer estimate; vocab from the tokens array len.
|
|
self._feed_forward_length: Optional[int] = None
|
|
self._vocab_size: 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 earlier layers' KV 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). RLock so MTP-crash
|
|
# recovery can re-acquire it for its nested load_model.
|
|
self._serial_load_lock = threading.RLock()
|
|
# Serialises mid-session respawns so many generations hitting a killed
|
|
# server trigger at most one reload (see _respawn_if_dead).
|
|
self._respawn_lock = threading.Lock()
|
|
# Set by the in-app updater while it swaps prebuilt binaries; load_model()
|
|
# rejects fast so no server starts from a half-swapped binary.
|
|
self._llama_update_in_progress = False
|
|
# 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
|
|
# Raw kwargs of the last healthy load, for the MTP-crash reload. Memory-only
|
|
# (carries hf_token, never logged); single-flight via the lock below.
|
|
self._last_load_kwargs: Optional[dict] = None
|
|
self._mtp_runtime_fallback_lock = threading.Lock()
|
|
self._mtp_runtime_fallback_in_progress = False
|
|
# Background watchdog so an MTP+tensor crash recovers even when no request
|
|
# observes it (direct proxy endpoints, or nothing in flight).
|
|
self._mtp_watchdog_thread: Optional[threading.Thread] = None
|
|
self._mtp_watchdog_stop = threading.Event()
|
|
# True when the launch actually runs MTP+tensor (Studio- or user/env-driven);
|
|
# gates the probe, watchdog, and recovery so pass-through MTP is covered.
|
|
self._mtp_runtime_fallback_active = False
|
|
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
|
|
# Audio INPUT capability (distinct from _is_audio, which is TTS output).
|
|
self._has_audio_input: bool = False
|
|
self._mmproj_has_audio: bool = False # clip.has_audio_encoder, set at load
|
|
# 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
|
|
|
|
_reaped = self._kill_orphaned_servers()
|
|
if _reaped:
|
|
# Reaped VRAM frees lazily; arm the settle wait so the first load
|
|
# waits before ranking GPUs by free memory.
|
|
self._last_kill_monotonic = time.monotonic()
|
|
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 _auth_headers(self) -> "Optional[dict[str, str]]":
|
|
"""Bearer header matching the --api-key direct-stream mode uses, else
|
|
None (so unauthenticated llama-server calls don't get a spurious 401)."""
|
|
return {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
|
|
|
@property
|
|
def model_identifier(self) -> Optional[str]:
|
|
return self._model_identifier
|
|
|
|
@property
|
|
def is_vision(self) -> bool:
|
|
return self._is_vision
|
|
|
|
@property
|
|
def is_diffusion(self) -> bool:
|
|
"""True when the loaded GGUF is a block-diffusion model (DiffusionGemma)."""
|
|
return self._is_diffusion
|
|
|
|
@property
|
|
def hf_variant(self) -> Optional[str]:
|
|
return self._hf_variant
|
|
|
|
@property
|
|
def gguf_path(self) -> Optional[str]:
|
|
return self._gguf_path
|
|
|
|
@property
|
|
def hf_repo(self) -> Optional[str]:
|
|
"""HF repo of the loaded model, or None for local/native file loads."""
|
|
return self._hf_repo
|
|
|
|
@property
|
|
def mtp_draft_path(self) -> Optional[str]:
|
|
return self._mtp_draft_path
|
|
|
|
@property
|
|
def spec_fallback_reason(self) -> Optional[str]:
|
|
"""Why MTP was disabled on the last MTP-requesting load, else None."""
|
|
return self._spec_fallback_reason
|
|
|
|
@property
|
|
def extra_args(self) -> Optional[List[str]]:
|
|
"""Extra llama-server flags from the last load (a 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.
|
|
|
|
The UI's "safe zone" warning threshold: the ``_fit_context_to_vram``
|
|
binary-search cap for the best GPU subset, or the 4096 fallback if the
|
|
weights exceed 90% of every subset. The slider ceiling is
|
|
``native_context_length``; dragging above this triggers the 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
|
|
|
|
@staticmethod
|
|
def _read_rss_bytes(pid: int) -> Optional[int]:
|
|
"""Resident set size of ``pid`` in bytes, from /proc/<pid>/status (Linux).
|
|
0 when the status has no VmRSS line (zombie / kernel thread); None where
|
|
/proc is unavailable (macOS/Windows) or the value is unreadable."""
|
|
try:
|
|
with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f:
|
|
for line in f:
|
|
if line.startswith("VmRSS:"):
|
|
# IndexError guards a "VmRSS:" line with no value column.
|
|
return int(line.split()[1]) * 1024 # kB -> bytes
|
|
except (FileNotFoundError, PermissionError, ValueError, IndexError, OSError):
|
|
return None
|
|
return 0 # readable but no VmRSS line
|
|
|
|
def load_progress(self) -> Optional[dict]:
|
|
"""Return live model-load progress, or None if not loading.
|
|
|
|
During warm-up llama-server mmaps weight shards into page cache before
|
|
pushing layers to VRAM, a window where status only reports ``loading``
|
|
and the UI spinner looks stuck for minutes on large MoEs. Samples
|
|
``/proc/<pid>/status VmRSS`` against the sum of GGUF shard sizes for a
|
|
real progress bar. Returns ``None`` when no load is in flight.
|
|
|
|
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; returns ``None`` 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 shard sizes (primary + any extras 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 share the primary's prefix before the shard index.
|
|
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():
|
|
prefix_lower = prefix.lower()
|
|
for sibling in parent.iterdir():
|
|
if (
|
|
sibling.is_file()
|
|
and sibling.name.lower().startswith(prefix_lower)
|
|
and sibling.name != stem
|
|
and sibling.suffix.lower() == ".gguf"
|
|
):
|
|
try:
|
|
bytes_total += sibling.stat().st_size
|
|
except OSError:
|
|
pass
|
|
except OSError:
|
|
pass
|
|
|
|
# VmRSS of the llama-server; None where /proc is unavailable.
|
|
bytes_loaded = LlamaCppBackend._read_rss_bytes(pid)
|
|
if bytes_loaded is None:
|
|
return None
|
|
|
|
# RSS climbs as weights page in, then drops once -ngl offloads them to
|
|
# VRAM and the mmap pages are freed. Hold a per-process high-water mark
|
|
# so the bar never regresses to ~8% mid-load (#5740).
|
|
hwm_pid, hwm = getattr(self, "_load_rss_hwm", (None, 0))
|
|
hwm = bytes_loaded if hwm_pid != pid else max(hwm, bytes_loaded)
|
|
self._load_rss_hwm = (pid, hwm)
|
|
bytes_loaded = hwm
|
|
|
|
phase = "ready" if self._healthy else "mmap"
|
|
fraction = 0.0
|
|
if bytes_total > 0:
|
|
fraction = min(1.0, bytes_loaded / bytes_total)
|
|
# Once llama-server is healthy the load is complete by definition. With
|
|
# layers offloaded to VRAM (-ngl) the process releases the mmap'd weight
|
|
# pages, so VmRSS sinks back well below the shard total; the raw RSS
|
|
# fraction would then report a partial (~8%) load indefinitely and freeze
|
|
# a fraction-driven progress bar even though the model is ready (#5740).
|
|
if self._healthy:
|
|
if bytes_total > 0:
|
|
bytes_loaded = bytes_total
|
|
fraction = 1.0
|
|
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 reasoning_effort_levels(self) -> list:
|
|
"""Discrete reasoning_effort levels the template offers (e.g. GLM-5.2's
|
|
['high', 'max']). Empty unless reasoning_style == 'enable_thinking_effort'."""
|
|
return self._reasoning_effort_levels
|
|
|
|
@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 == "enable_thinking_effort":
|
|
# GLM-5.2-style: enable_thinking is the on/off gate; when on, leave
|
|
# the template's default effort (max) in place.
|
|
return {"enable_thinking": enable_thinking}
|
|
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.
|
|
|
|
Merges 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 and don't consume
|
|
# enable_thinking / reasoning_effort -- skip.
|
|
if self._supports_reasoning and not self._reasoning_always_on:
|
|
if self._reasoning_style == "enable_thinking_effort":
|
|
# GLM-5.2-style: enable_thinking gates thinking on/off, and the
|
|
# reasoning_effort level (e.g. 'high' | 'max') is only meaningful
|
|
# while thinking is on. Disabling is enable_thinking=false; a raw
|
|
# API caller can also disable via the OpenAI-style
|
|
# reasoning_effort="none" sentinel. We never coerce off into a
|
|
# 'low' effort the way gpt-oss does (those models genuinely
|
|
# cannot disable).
|
|
thinking_off = enable_thinking is False or reasoning_effort == "none"
|
|
if enable_thinking is not None or reasoning_effort == "none":
|
|
kwargs["enable_thinking"] = not thinking_off
|
|
if not thinking_off and reasoning_effort in self._reasoning_effort_levels:
|
|
kwargs["reasoning_effort"] = reasoning_effort
|
|
elif self._reasoning_style == "reasoning_effort":
|
|
if reasoning_effort in ("none", "low", "medium", "high"):
|
|
kwargs["reasoning_effort"] = reasoning_effort
|
|
elif reasoning_effort == "minimal":
|
|
kwargs["reasoning_effort"] = "low"
|
|
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:
|
|
# DiffusionGemma serves via the visual runner, whose live per-step canvas
|
|
# frames are dropped by the agentic tool loop; never route it through tools.
|
|
if self._is_diffusion:
|
|
return False
|
|
return self._supports_tools
|
|
|
|
@property
|
|
def cache_type_kv(self) -> Optional[str]:
|
|
return self._cache_type_kv
|
|
|
|
@property
|
|
def tensor_parallel(self) -> bool:
|
|
"""Whether --split-mode tensor is active on the loaded server."""
|
|
return self._tensor_parallel
|
|
|
|
@property
|
|
def layer_preserves_tensor_intent(self) -> bool:
|
|
"""True when a downgraded tensor request kept this layer load multi-GPU."""
|
|
return self._layer_preserves_tensor_intent
|
|
|
|
@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 _resolved_studio_root_and_is_legacy() -> "tuple[Optional[Path], bool]":
|
|
"""Resolve the Studio install root and classify it as the legacy
|
|
~/.unsloth/studio root vs. a custom (env/venv-inferred) root.
|
|
|
|
Returns (resolved_root, is_legacy). On any import/resolution failure the
|
|
root is treated as legacy and resolved_root is None -- callers must read
|
|
resolved_root only when is_legacy is False. Shared by
|
|
_find_llama_server_binary (discovery) and _kill_orphaned_servers
|
|
(cleanup) so the two never disagree on which root is legacy.
|
|
"""
|
|
try:
|
|
from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433
|
|
|
|
resolved = _sr()
|
|
legacy_studio = Path.home() / ".unsloth" / "studio"
|
|
try:
|
|
is_legacy = resolved.resolve() == legacy_studio.resolve()
|
|
except (OSError, ValueError):
|
|
is_legacy = resolved == legacy_studio
|
|
return (None if is_legacy else resolved), is_legacy
|
|
except (ImportError, OSError, ValueError):
|
|
return None, True
|
|
|
|
@staticmethod
|
|
def _find_llama_server_binary(*, include_denied: bool = False) -> 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)
|
|
7. llama-server on PATH (system install)
|
|
8. ./bin/llama-server (legacy: extracted binary)
|
|
"""
|
|
binary_name = "llama-server.exe" if sys.platform == "win32" else "llama-server"
|
|
|
|
def _file_status(p: Path) -> str:
|
|
# "file", "absent", or "denied" (exists but stays access-denied
|
|
# across a short retry: Windows AV/ACL or an install replace in
|
|
# flight). is_file() raises PermissionError (WinError 5) instead of
|
|
# returning False for the locked case, so never treat it as missing.
|
|
for _ in range(5):
|
|
try:
|
|
return "file" if p.is_file() else "absent"
|
|
except PermissionError:
|
|
time.sleep(0.2)
|
|
except OSError:
|
|
return "absent"
|
|
return "denied"
|
|
|
|
def _is_file(p: Path) -> bool:
|
|
return _file_status(p) == "file"
|
|
|
|
def _layout_candidates(d: Path) -> list:
|
|
# build layouts probed under a llama.cpp dir, highest priority first
|
|
cands = [d / binary_name, d / "build" / "bin" / binary_name]
|
|
if sys.platform == "win32":
|
|
cands.append(d / "build" / "bin" / "Release" / binary_name)
|
|
return cands
|
|
|
|
def _unavailable(p: object) -> None:
|
|
# a pinned or managed binary that exists but is access-denied: report
|
|
# it instead of silently downgrading to a lower-priority llama-server
|
|
logger.warning(
|
|
f"llama-server at {p} exists but is access-denied (antivirus or "
|
|
"an in-flight install); not falling back to another binary, "
|
|
"retry once it is released"
|
|
)
|
|
return None
|
|
|
|
def _scan_pinned(paths: list):
|
|
# first existing candidate wins -> (path, None); a present-but-denied
|
|
# one -> (None, denied_path) so the caller reports it rather than
|
|
# skipping to a lower-priority location. include_denied returns the
|
|
# locked path instead: diffusion asset lookup only needs its dir.
|
|
for p in paths:
|
|
st = _file_status(p)
|
|
if st == "file":
|
|
return str(p), None
|
|
if st == "denied":
|
|
return (str(p), None) if include_denied else (None, p)
|
|
return None, None
|
|
|
|
# 1. Env var: direct path to binary
|
|
env_path = os.environ.get("LLAMA_SERVER_PATH")
|
|
if env_path:
|
|
hit, locked = _scan_pinned([Path(env_path)])
|
|
if locked is not None:
|
|
return _unavailable(locked)
|
|
if hit:
|
|
return hit
|
|
|
|
# 1b. UNSLOTH_LLAMA_CPP_PATH: custom llama.cpp install dir
|
|
custom_llama_cpp = os.environ.get("UNSLOTH_LLAMA_CPP_PATH")
|
|
if custom_llama_cpp:
|
|
hit, locked = _scan_pinned(_layout_candidates(Path(custom_llama_cpp)))
|
|
if locked is not None:
|
|
return _unavailable(locked)
|
|
if hit:
|
|
return hit
|
|
|
|
# 2-4. Match installer layout: env-mode -> $STUDIO_HOME/llama.cpp;
|
|
# default/HOME-redirect -> ~/.unsloth/llama.cpp (sibling of studio).
|
|
legacy_llama = Path.home() / ".unsloth" / "llama.cpp"
|
|
_resolved_sr, _is_legacy = LlamaCppBackend._resolved_studio_root_and_is_legacy()
|
|
if _is_legacy:
|
|
search_roots = [legacy_llama]
|
|
else:
|
|
# _kill_orphaned_servers excludes the legacy root in custom mode;
|
|
# discovery must match so we never spawn a server we then refuse to
|
|
# clean up. UNSLOTH_LLAMA_CPP_PATH (handled earlier) is the explicit
|
|
# way to share a build across roots.
|
|
search_roots = [_resolved_sr / "llama.cpp"]
|
|
for unsloth_home in search_roots:
|
|
hit, locked = _scan_pinned(_layout_candidates(unsloth_home))
|
|
if locked is not None:
|
|
return _unavailable(locked)
|
|
if hit:
|
|
return hit
|
|
|
|
# 5-6. Legacy: in-tree build (older setup.sh / setup.ps1). A fallback,
|
|
# so a denied candidate here just continues (no no-fallback halt).
|
|
project_root = Path(__file__).resolve().parents[4]
|
|
for p in _layout_candidates(project_root / "llama.cpp"):
|
|
if _is_file(p):
|
|
return str(p)
|
|
|
|
# 7. System PATH
|
|
system_path = shutil.which("llama-server")
|
|
if system_path:
|
|
return system_path
|
|
|
|
# 8. Legacy: extracted to bin/
|
|
bin_path = project_root / "bin" / binary_name
|
|
if _is_file(bin_path):
|
|
return str(bin_path)
|
|
|
|
return None
|
|
|
|
# ── llama-server capability probe ─────────────────────────────
|
|
|
|
# Cached on (path, mtime); `unsloth studio update` bumps mtime.
|
|
_capability_cache: dict[tuple[str, int], dict[str, object]] = {}
|
|
|
|
@classmethod
|
|
def probe_server_capabilities(cls, binary: Optional[str] = None) -> dict[str, object]:
|
|
"""Parse `llama-server --help` for feature flags. Returns
|
|
{found, mtp_token, supports_mtp, ngram_mod_flavor,
|
|
supports_ngram_mod, spec_draft_n_max_flag, cache flag support}.
|
|
|
|
``ngram_mod_flavor``: ``"new"`` when the post-rename
|
|
``--spec-ngram-mod-n-match / -n-min / -n-max`` are real args;
|
|
``"legacy"`` when only the pre-rename
|
|
``--spec-ngram-size-n / --draft-min / --draft-max`` are real (the
|
|
rename ships stub removal entries for legacy names, told apart by
|
|
the "argument has been removed" description); ``None`` if neither
|
|
set is usable.
|
|
|
|
``spec_draft_n_max_flag``: the flag the binary accepts --
|
|
``--spec-draft-n-max`` post-rename, ``--draft-max`` on legacy.
|
|
``None`` means n_max cannot be set.
|
|
"""
|
|
bin_path = binary or cls._find_llama_server_binary()
|
|
if not bin_path or not Path(bin_path).is_file():
|
|
return {
|
|
"found": False,
|
|
"mtp_token": None,
|
|
"supports_mtp": False,
|
|
"ngram_mod_flavor": None,
|
|
"supports_ngram_mod": False,
|
|
"spec_draft_n_max_flag": None,
|
|
"supports_kv_unified": False,
|
|
"supports_fit_ctx": False,
|
|
"supports_cache_ram": False,
|
|
"supports_ctx_checkpoints": False,
|
|
"supports_no_cache_prompt": False,
|
|
"supports_metrics": False,
|
|
}
|
|
try:
|
|
mtime = int(Path(bin_path).stat().st_mtime)
|
|
except OSError:
|
|
mtime = 0
|
|
cache_key = (bin_path, mtime)
|
|
cached = cls._capability_cache.get(cache_key)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
mtp_token: Optional[str] = None
|
|
ngram_mod_flavor: Optional[str] = None
|
|
spec_draft_n_max_flag: Optional[str] = None
|
|
supports_kv_unified = False
|
|
supports_fit_ctx = False
|
|
supports_cache_ram = False
|
|
supports_ctx_checkpoints = False
|
|
supports_no_cache_prompt = False
|
|
supports_metrics = False
|
|
try:
|
|
probe_env = cls._llama_server_env_for_binary(bin_path)
|
|
result = subprocess.run(
|
|
[bin_path, "--help"],
|
|
capture_output = True,
|
|
text = True,
|
|
errors = "replace",
|
|
timeout = 10,
|
|
check = False,
|
|
env = probe_env,
|
|
)
|
|
help_text = (result.stdout or "") + "\n" + (result.stderr or "")
|
|
# Split into per-flag blocks (each --flag line + its indented
|
|
# continuation), so the "argument has been removed" description
|
|
# sits with its flag.
|
|
blocks: dict[str, str] = {}
|
|
current_flags: list[str] = []
|
|
current_desc: list[str] = []
|
|
for line in help_text.splitlines():
|
|
stripped = line.strip()
|
|
if stripped.startswith("-") and not line.startswith(" "):
|
|
# New flag line; flush previous.
|
|
if current_flags:
|
|
desc = " ".join(current_desc)
|
|
for f in current_flags:
|
|
blocks[f] = desc
|
|
current_flags = []
|
|
current_desc = [stripped]
|
|
# Extract long-form flag tokens from the DECLARATION
|
|
# prefix only (comma-separated aliases). Stop at the
|
|
# first non-flag token so flag references inside
|
|
# descriptions are ignored.
|
|
for tok in re.split(r"[,\s]+", stripped):
|
|
if tok.startswith("--") and re.match(r"--[A-Za-z][A-Za-z0-9_-]*$", tok):
|
|
current_flags.append(tok)
|
|
elif tok.startswith("-") and len(tok) > 1:
|
|
# short alias like -fa; keep scanning aliases.
|
|
continue
|
|
else:
|
|
# First non-flag token marks end of decl.
|
|
break
|
|
else:
|
|
current_desc.append(stripped)
|
|
if current_flags:
|
|
desc = " ".join(current_desc)
|
|
for f in current_flags:
|
|
blocks[f] = desc
|
|
|
|
def _is_real(flag: str) -> bool:
|
|
"""True if the flag exists AND is not a removal stub."""
|
|
desc = blocks.get(flag)
|
|
if desc is None:
|
|
return False
|
|
return "argument has been removed" not in desc
|
|
|
|
# MTP token from the --spec-type line.
|
|
spec_line = ""
|
|
for line in help_text.splitlines():
|
|
if "--spec-type" in line:
|
|
spec_line = line
|
|
break
|
|
# PR #22673 used draft-mtp; later renamed to mtp.
|
|
if "draft-mtp" in spec_line:
|
|
mtp_token = "draft-mtp"
|
|
elif re.search(r"[|,\[]mtp[|,\]]", spec_line):
|
|
mtp_token = "mtp"
|
|
|
|
# ngram-mod flag flavor. Post-rename builds advertise both new
|
|
# args (real) and legacy ones (stubs); pre-rename builds only
|
|
# have legacy ones as real.
|
|
new_ngram_real = (
|
|
_is_real("--spec-ngram-mod-n-match")
|
|
and _is_real("--spec-ngram-mod-n-min")
|
|
and _is_real("--spec-ngram-mod-n-max")
|
|
)
|
|
legacy_ngram_real = (
|
|
_is_real("--spec-ngram-size-n")
|
|
and _is_real("--draft-max")
|
|
and _is_real("--draft-min")
|
|
)
|
|
if new_ngram_real:
|
|
ngram_mod_flavor = "new"
|
|
elif legacy_ngram_real:
|
|
ngram_mod_flavor = "legacy"
|
|
|
|
# n_max flag: prefer post-rename, fall back to legacy.
|
|
if _is_real("--spec-draft-n-max"):
|
|
spec_draft_n_max_flag = "--spec-draft-n-max"
|
|
elif _is_real("--draft-max"):
|
|
spec_draft_n_max_flag = "--draft-max"
|
|
|
|
supports_kv_unified = _is_real("--kv-unified")
|
|
supports_fit_ctx = _is_real("--fit-ctx")
|
|
supports_cache_ram = _is_real("--cache-ram")
|
|
supports_ctx_checkpoints = _is_real("--ctx-checkpoints")
|
|
supports_no_cache_prompt = _is_real("--no-cache-prompt")
|
|
supports_metrics = _is_real("--metrics")
|
|
except (OSError, subprocess.SubprocessError) as exc:
|
|
logger.debug(f"llama-server --help probe failed: {exc}")
|
|
|
|
info = {
|
|
"found": True,
|
|
"mtp_token": mtp_token,
|
|
"supports_mtp": mtp_token is not None,
|
|
"ngram_mod_flavor": ngram_mod_flavor,
|
|
"supports_ngram_mod": ngram_mod_flavor is not None,
|
|
"spec_draft_n_max_flag": spec_draft_n_max_flag,
|
|
"supports_kv_unified": supports_kv_unified,
|
|
"supports_fit_ctx": supports_fit_ctx,
|
|
"supports_cache_ram": supports_cache_ram,
|
|
"supports_ctx_checkpoints": supports_ctx_checkpoints,
|
|
"supports_no_cache_prompt": supports_no_cache_prompt,
|
|
"supports_metrics": supports_metrics,
|
|
}
|
|
cls._capability_cache[cache_key] = info
|
|
return info
|
|
|
|
# ── GPU allocation ────────────────────────────────────────────
|
|
|
|
@staticmethod
|
|
def _get_gguf_size_bytes(model_path: str) -> int:
|
|
"""Total GGUF size in bytes, including split shards."""
|
|
main = Path(model_path)
|
|
total = main.stat().st_size
|
|
|
|
# Check for split shards (e.g. model-00001-of-00003.gguf)
|
|
m = _SHARD_FULL_RE.match(main.name)
|
|
if m:
|
|
prefix, _, num_total = m.group(1), m.group(2), m.group(3)
|
|
sibling_pat = re.compile(
|
|
r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(num_total) + r"\.gguf$",
|
|
re.IGNORECASE,
|
|
)
|
|
for sibling in main.parent.iterdir():
|
|
if sibling != main and sibling_pat.match(sibling.name):
|
|
total += sibling.stat().st_size
|
|
|
|
return total
|
|
|
|
@staticmethod
|
|
def _resolve_visible_physical_ids() -> Optional[list[int]]:
|
|
"""Physical GPU ids behind the active visibility mask (HIP/ROCR/CUDA on
|
|
ROCm, CUDA otherwise). None when no mask is set; empty list for an empty
|
|
mask. Shared by the APU / datacenter / free-memory probes so they agree
|
|
on the ordinal->physical mapping."""
|
|
try:
|
|
import torch
|
|
is_rocm = getattr(torch.version, "hip", None) is not None
|
|
except Exception:
|
|
is_rocm = False
|
|
if is_rocm:
|
|
hip_v = os.environ.get("HIP_VISIBLE_DEVICES")
|
|
rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES")
|
|
cvd = (
|
|
hip_v
|
|
if hip_v is not None
|
|
else rocr_v
|
|
if rocr_v is not None
|
|
else os.environ.get("CUDA_VISIBLE_DEVICES")
|
|
)
|
|
else:
|
|
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
|
if cvd is None:
|
|
return None
|
|
try:
|
|
return [int(x.strip()) for x in cvd.split(",") if x.strip()]
|
|
except ValueError:
|
|
return None
|
|
|
|
@staticmethod
|
|
def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool:
|
|
"""True only for AMD unified-memory APUs (gfx1150/gfx1151), where
|
|
GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM (it
|
|
hurts discrete GPUs). gpu_indices (PHYSICAL ids) scopes the check to the
|
|
selected GPUs, so a dGPU on a mixed host is not treated as unified-memory;
|
|
None means every visible GPU."""
|
|
try:
|
|
import torch
|
|
|
|
if getattr(torch.version, "hip", None) is None:
|
|
return False
|
|
if not (hasattr(torch, "cuda") and torch.cuda.is_available()):
|
|
return False
|
|
# Map visible ordinal -> physical id via the active ROCm mask (HIP,
|
|
# then ROCR, then CUDA), mirroring _get_gpu_memory's ROCm branch.
|
|
physical_ids = LlamaCppBackend._resolve_visible_physical_ids()
|
|
arch_by_id: dict[int, str] = {}
|
|
for ordinal in range(torch.cuda.device_count()):
|
|
try:
|
|
_arch = (
|
|
getattr(torch.cuda.get_device_properties(ordinal), "gcnArchName", "") or ""
|
|
)
|
|
except Exception:
|
|
continue
|
|
pid = (
|
|
physical_ids[ordinal]
|
|
if physical_ids is not None and ordinal < len(physical_ids)
|
|
else ordinal
|
|
)
|
|
arch_by_id[pid] = _arch.split(":")[0].strip().lower()
|
|
for _i in list(gpu_indices) if gpu_indices is not None else list(arch_by_id):
|
|
if arch_by_id.get(_i) in {"gfx1150", "gfx1151"}:
|
|
return True
|
|
except Exception:
|
|
return False
|
|
return False
|
|
|
|
# Datacenter / professional NVIDIA parts that benefit from the llama.cpp
|
|
# FP32-accum / P2P tunings. Whole-word (\b) so short markers don't match
|
|
# workstation parts as substrings: "a100" must not fire on "RTX A1000".
|
|
_DATACENTER_GPU_RE = re.compile(
|
|
r"\b(?:a100|a30|h100|h200|h800|gh200|b200|b100|b300|gb200|gb300|"
|
|
r"l40s?|l4|rtx pro 6000|rtx 6000 ada)\b"
|
|
)
|
|
|
|
@staticmethod
|
|
def _is_datacenter_gpu(gpu_indices = None) -> bool:
|
|
"""True iff every selected NVIDIA GPU is a datacenter/professional part.
|
|
NVIDIA-only, fails open to False (consumer GeForce, ROCm, CPU and errors
|
|
are left untouched); a mixed DC+consumer selection counts as non-DC.
|
|
|
|
gpu_indices are PHYSICAL ids (see _get_gpu_free_memory), but
|
|
get_device_properties wants mask-relative ordinals, so we rebuild the
|
|
ordinal->physical map from CUDA_VISIBLE_DEVICES and key names by physical
|
|
id. Otherwise a masked host (CUDA_VISIBLE_DEVICES=4,5,6,7, selection [4,5])
|
|
would drop the tuning or probe the wrong GPU."""
|
|
try:
|
|
import torch
|
|
|
|
if getattr(torch.version, "hip", None) is not None:
|
|
return False # ROCm reuses torch.cuda.*; not a CUDA part
|
|
if not (hasattr(torch, "cuda") and torch.cuda.is_available()):
|
|
return False
|
|
count = torch.cuda.device_count()
|
|
|
|
# Mirror _get_gpu_free_memory: map visible ordinal -> physical id via
|
|
# CUDA_VISIBLE_DEVICES; unset/unparsable leaves physical id == ordinal.
|
|
physical_ids = LlamaCppBackend._resolve_visible_physical_ids()
|
|
|
|
pattern = LlamaCppBackend._DATACENTER_GPU_RE
|
|
names_by_id: dict[int, str] = {}
|
|
for ordinal in range(count):
|
|
try:
|
|
name = (torch.cuda.get_device_properties(ordinal).name or "").lower()
|
|
except Exception:
|
|
continue
|
|
pid = (
|
|
physical_ids[ordinal]
|
|
if physical_ids is not None and ordinal < len(physical_ids)
|
|
else ordinal
|
|
)
|
|
names_by_id[pid] = name
|
|
|
|
indices = list(gpu_indices) if gpu_indices else list(names_by_id)
|
|
saw = False
|
|
for _i in indices:
|
|
name = names_by_id.get(_i)
|
|
if name is None:
|
|
continue # not visible -> skip (fail conservative)
|
|
saw = True
|
|
if not pattern.search(name):
|
|
return False
|
|
return saw
|
|
except Exception:
|
|
return False
|
|
|
|
@staticmethod
|
|
def _effective_gpu_count(gpu_indices = None) -> int:
|
|
"""GPUs llama-server will use: len(selection), else the visible CUDA
|
|
device count (None = every visible GPU). 0 on error so multi-GPU tuning
|
|
stays off when the count is unknown."""
|
|
if gpu_indices is not None:
|
|
return len(gpu_indices)
|
|
try:
|
|
import torch
|
|
if hasattr(torch, "cuda") and torch.cuda.is_available():
|
|
return torch.cuda.device_count()
|
|
except Exception:
|
|
return 0
|
|
return 0
|
|
|
|
@staticmethod
|
|
def _apply_datacenter_env(env: dict, gpu_indices = None) -> bool:
|
|
"""Inject DC llama.cpp tuning into env in place via setdefault (user
|
|
values win); return whether the box qualified. Opt out with
|
|
UNSLOTH_DISABLE_DC_TUNING=1; only datacenter NVIDIA parts qualify
|
|
(consumer/ROCm/CPU/error are a no-op). Sets GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F
|
|
for any qualifying GPU (FP32 accum: ~0% cost on B200, real cost on GeForce),
|
|
plus GGML_CUDA_P2P + CUDA_SCALE_LAUNCH_QUEUES=4x for multi-GPU (+33-51% pp
|
|
tensor-split, +8-16% pipeline split on B200)."""
|
|
if os.environ.get("UNSLOTH_DISABLE_DC_TUNING") == "1":
|
|
return False
|
|
if not LlamaCppBackend._is_datacenter_gpu(gpu_indices):
|
|
return False
|
|
env.setdefault("GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F", "1")
|
|
if LlamaCppBackend._effective_gpu_count(gpu_indices) > 1:
|
|
env.setdefault("GGML_CUDA_P2P", "1")
|
|
env.setdefault("CUDA_SCALE_LAUNCH_QUEUES", "4x")
|
|
return True
|
|
|
|
@staticmethod
|
|
def _get_gpu_free_memory() -> list[tuple[int, int]]:
|
|
"""Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by
|
|
index; empty if no supported GPU is reachable. Thin wrapper over
|
|
``_get_gpu_memory`` for callers that only need free VRAM."""
|
|
return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()]
|
|
|
|
@staticmethod
|
|
def _apple_metal_memory_budget_bytes() -> int:
|
|
"""Unified-memory budget for GGUF context fitting on Apple Silicon.
|
|
|
|
No GPU is enumerated on Metal, so the context would default to native and
|
|
over-commit unified memory ("Compute error." at decode, #5118/#6529). Use a
|
|
fraction of MLX's Metal working-set, else total RAM; 0 off Apple Silicon or
|
|
when unresolvable, so callers skip the cap.
|
|
"""
|
|
from utils.hardware import is_apple_silicon
|
|
|
|
if not is_apple_silicon():
|
|
return 0
|
|
rec_bytes = 0
|
|
try:
|
|
import mlx.core as mx
|
|
if mx.metal.is_available():
|
|
rec_bytes = int(mx.device_info().get("max_recommended_working_set_size") or 0)
|
|
except Exception:
|
|
rec_bytes = 0
|
|
if rec_bytes <= 0:
|
|
try:
|
|
import psutil
|
|
rec_bytes = int(psutil.virtual_memory().total)
|
|
except Exception:
|
|
return 0
|
|
return int(rec_bytes * _APPLE_UNIFIED_MEMORY_FRACTION)
|
|
|
|
@staticmethod
|
|
def _get_gpu_memory() -> list[tuple[int, int, int]]:
|
|
"""Query free AND total memory per GPU.
|
|
|
|
Order:
|
|
1. ``nvidia-smi`` (NVIDIA CUDA hosts) -- respects
|
|
``CUDA_VISIBLE_DEVICES``.
|
|
2. ``torch.cuda.mem_get_info`` -- universal fallback that works
|
|
on AMD ROCm too (HIP runtime reuses the ``torch.cuda.*``
|
|
namespace). Covers the AMD case for issue #5106 (nvidia-smi
|
|
probe returned [] on AMD) and NVIDIA hosts missing
|
|
``nvidia-smi`` from PATH.
|
|
|
|
Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no
|
|
supported GPU is reachable. ``total`` lets the fit reserve absolute headroom.
|
|
"""
|
|
# ── NVIDIA via nvidia-smi ────────────────────────────────────
|
|
try:
|
|
result = subprocess.run(
|
|
[
|
|
"nvidia-smi",
|
|
"--query-gpu=index,memory.free,memory.total",
|
|
"--format=csv,noheader,nounits",
|
|
],
|
|
capture_output = True,
|
|
text = True,
|
|
timeout = 10,
|
|
env = child_env_without_native_path_secret(),
|
|
**_windows_hidden_subprocess_kwargs(),
|
|
)
|
|
if result.returncode == 0:
|
|
allowed: Optional[set[int]] = None
|
|
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
|
if cvd is not None:
|
|
try:
|
|
# `if x.strip()` filters trailing-comma masks ("0,1,").
|
|
# Empty mask (CVD="") yields an empty set -> all GPUs
|
|
# filtered out, per codebase convention.
|
|
allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip())
|
|
except ValueError:
|
|
pass
|
|
gpus: list[tuple[int, int, int]] = []
|
|
for line in result.stdout.strip().splitlines():
|
|
parts = [p.strip() for p in line.split(",")]
|
|
if len(parts) < 2:
|
|
continue
|
|
# Index and free required; skip a bad line rather than abandon
|
|
# the probe to the torch fallback.
|
|
try:
|
|
idx = int(parts[0])
|
|
free_mib = int(parts[1])
|
|
except ValueError:
|
|
continue
|
|
# Total parsed separately: a two-column line or a non-integer
|
|
# total ("N/A" on MIG/vGPU) keeps the GPU at total 0 (fit uses
|
|
# the free*frac fallback) instead of dropping it.
|
|
total_mib = 0
|
|
if len(parts) >= 3 and parts[2]:
|
|
try:
|
|
total_mib = int(parts[2])
|
|
except ValueError:
|
|
total_mib = 0
|
|
if allowed is not None and idx not in allowed:
|
|
continue
|
|
gpus.append((idx, free_mib, total_mib))
|
|
# Match the docstring's sort-by-id guarantee (driver order isn't).
|
|
gpus.sort(key = lambda g: g[0])
|
|
if gpus:
|
|
return gpus
|
|
except Exception as e:
|
|
logger.debug(f"nvidia-smi probe failed: {e}")
|
|
|
|
# ── Torch fallback (covers AMD ROCm and missing nvidia-smi) ──
|
|
try:
|
|
import torch
|
|
|
|
if not hasattr(torch, "cuda") or not torch.cuda.is_available():
|
|
return []
|
|
if not hasattr(torch.cuda, "mem_get_info"):
|
|
return []
|
|
# torch.cuda enumerates GPUs RELATIVE to the visibility mask. We
|
|
# feed these IDs back into the subprocess as CVD, so visible ordinals
|
|
# must be translated to physical indices first; otherwise CVD=2,3
|
|
# gets rewritten to 0,1 and targets the wrong GPUs.
|
|
# Match utils/hardware/hardware.py::_get_parent_visible_gpu_spec:
|
|
# treat an empty mask (HIP_VISIBLE_DEVICES="") as "no GPUs" rather
|
|
# than falling through. ``or`` would coerce "" to the wrong source.
|
|
# Empty mask (CVD="") yields an empty list -> no GPUs, consistent
|
|
# with the nvidia-smi path.
|
|
physical_ids = LlamaCppBackend._resolve_visible_physical_ids()
|
|
gpus = []
|
|
for ordinal in range(torch.cuda.device_count()):
|
|
free_bytes, total_bytes = torch.cuda.mem_get_info(ordinal)
|
|
idx = (
|
|
physical_ids[ordinal]
|
|
if physical_ids is not None and ordinal < len(physical_ids)
|
|
else ordinal
|
|
)
|
|
gpus.append((idx, free_bytes // (1024 * 1024), total_bytes // (1024 * 1024)))
|
|
# Match the nvidia-smi path's docstring guarantee of sorted-by-id.
|
|
return sorted(gpus, key = lambda g: g[0])
|
|
except Exception as e:
|
|
logger.debug(f"torch GPU probe failed: {e}")
|
|
return []
|
|
|
|
@staticmethod
|
|
def _available_system_memory_mib() -> Optional[int]:
|
|
"""Available system RAM in MiB (psutil, then /proc/meminfo), or None if
|
|
neither is readable. On a unified-memory APU this, not the ROCm-reported
|
|
VRAM, is the real ceiling: the weights load into shared system RAM."""
|
|
try:
|
|
import psutil
|
|
return int(psutil.virtual_memory().available // (1024 * 1024))
|
|
except Exception:
|
|
pass
|
|
try:
|
|
with open("/proc/meminfo") as f:
|
|
for line in f:
|
|
if line.startswith("MemAvailable:"):
|
|
return int(line.split()[1]) // 1024 # kB -> MiB
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
@staticmethod
|
|
def _apu_ram_shortfall_message(
|
|
model_size_bytes: int,
|
|
avail_mib: Optional[int],
|
|
headroom_mib: int = 2048,
|
|
) -> Optional[str]:
|
|
"""On a unified-memory APU, return a user-facing refusal when the weights
|
|
cannot fit in available system RAM (else None). Weights only: KV/context
|
|
auto-reduce, so counting them too would refuse loads that would succeed.
|
|
None avail (unknown RAM) never refuses."""
|
|
if avail_mib is None:
|
|
return None
|
|
need_mib = model_size_bytes / (1024 * 1024)
|
|
if need_mib <= avail_mib - headroom_mib:
|
|
return None
|
|
return (
|
|
f"This model needs about {need_mib / 1024:.0f} GB but only about "
|
|
f"{avail_mib / 1024:.0f} GB of memory is available. On a unified-memory "
|
|
"APU the weights load into system RAM, so a larger model is stopped by "
|
|
"the OS mid-load. Use a smaller or more quantized GGUF, or free memory "
|
|
"(on WSL, raise the memory limit in .wslconfig)."
|
|
)
|
|
|
|
# Skip the wait when the last kill is older than this; the driver has
|
|
# already reclaimed the prior process's allocations.
|
|
_VRAM_SETTLE_WINDOW_S: float = 15.0
|
|
|
|
@staticmethod
|
|
def _wait_for_vram_settle(
|
|
max_wait: float = 2.0,
|
|
interval: float = 0.25,
|
|
tolerance_mib: int = 256,
|
|
since_kill: float = 0.0,
|
|
) -> None:
|
|
"""Poll ``_get_gpu_free_memory`` until free VRAM stabilises.
|
|
|
|
The driver reclaims a dead process's allocations asynchronously, so
|
|
sampling free memory in the kill-to-spawn window reads artificially low
|
|
and pushes GPU selection toward needless CPU offload (the Apply-reload
|
|
OOM bare-shell launches never see).
|
|
|
|
Short-circuits on cold start, stale kill (older than
|
|
``_VRAM_SETTLE_WINDOW_S``), CPU-only hosts, probe exceptions, and GPU-set
|
|
changes. ``max_wait`` bounds wall-clock time so a wedged ``nvidia-smi``
|
|
can't extend the reload.
|
|
"""
|
|
now = time.monotonic()
|
|
if since_kill <= 0.0:
|
|
return
|
|
if now - since_kill > LlamaCppBackend._VRAM_SETTLE_WINDOW_S:
|
|
return
|
|
deadline = now + max_wait
|
|
|
|
def _probe_or_none():
|
|
if time.monotonic() >= deadline:
|
|
return None
|
|
try:
|
|
return LlamaCppBackend._get_gpu_free_memory()
|
|
except Exception:
|
|
return None
|
|
|
|
prev = _probe_or_none()
|
|
if prev is None or not prev:
|
|
return
|
|
while time.monotonic() < deadline:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
return
|
|
# Clip the nap so a near-zero ``max_wait`` is respected.
|
|
time.sleep(min(interval, remaining))
|
|
curr = _probe_or_none()
|
|
if curr is None or not curr or len(curr) != len(prev):
|
|
return
|
|
prev_map = dict(prev)
|
|
stable = True
|
|
for idx, free in curr:
|
|
if idx not in prev_map:
|
|
stable = False
|
|
break
|
|
prev_free = prev_map[idx]
|
|
# Adaptive: 2% of the larger sample dominates the 256 MiB floor.
|
|
per_gpu_tol = max(tolerance_mib, int(max(free, prev_free) * 0.02))
|
|
if abs(free - prev_free) >= per_gpu_tol:
|
|
stable = False
|
|
break
|
|
if stable:
|
|
return
|
|
prev = curr
|
|
|
|
# Free-VRAM fraction at which Studio pins the GPU directly instead of
|
|
# deferring to ``--fit on``. 3% headroom: the compute buffer is now modelled in
|
|
# the fit, so this only guards fragmentation + multi-GPU per-device CUDA context
|
|
# (~2-3%); kept >= 3% as a floor (0.90 dropped 91-94% fits to CPU offload, #5106).
|
|
_GPU_PIN_VRAM_FRACTION = 0.97
|
|
|
|
# Fallback per-device tensor-mode compute buffer (MiB), used only when GGUF
|
|
# dims are unavailable so _estimate_compute_buffer_bytes (the primary, derived
|
|
# path) returns 0.
|
|
_TENSOR_PARALLEL_BUFFER_RESERVE_MIB = 5120
|
|
|
|
# Fixed per-device overhead on every GPU of a LAYER split (CUDA context +
|
|
# scratch), beyond the conserved slot-scaling buffer. ~0.9 GB/device measured
|
|
# (Qwen3.6-27B, b9625), independent of --parallel; reserved per extra GPU so a
|
|
# tight layer split can't advertise a context that OOMs at load.
|
|
_PIPELINE_PER_DEVICE_OVERHEAD_MIB = 1024
|
|
|
|
# KV cache types llama.cpp accepts in tensor mode. A quantized KV cache
|
|
# aborts a --split-mode tensor load, so it's dropped for the tensor attempt.
|
|
_TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"})
|
|
|
|
# (binary, mtime, model) that aborted on --split-mode tensor this process (#6415
|
|
# geometry limit, e.g. MQA n_head_kv=1). Model-keyed so one model's abort doesn't
|
|
# skip tensor for others; tensor is tried by default, recorded only on a real abort.
|
|
_tensor_split_abort_keys: set[tuple[str, int, str]] = set()
|
|
|
|
@classmethod
|
|
def _tensor_split_cache_key(
|
|
cls, binary: Optional[str], model: Optional[str]
|
|
) -> Optional[tuple[str, int, str]]:
|
|
"""(path, mtime_ns, model) key; ns mtime re-probes a same-second binary swap."""
|
|
if not binary or not model:
|
|
return None
|
|
try:
|
|
mtime = Path(binary).stat().st_mtime_ns
|
|
except OSError:
|
|
mtime = 0
|
|
return (binary, mtime, model)
|
|
|
|
@classmethod
|
|
def _tensor_split_aborts(cls, binary: Optional[str], model: Optional[str]) -> bool:
|
|
"""True if (binary, model) aborted on --split-mode tensor this session."""
|
|
key = cls._tensor_split_cache_key(binary, model)
|
|
return key is not None and key in cls._tensor_split_abort_keys
|
|
|
|
@classmethod
|
|
def _record_tensor_split_abort(cls, binary: Optional[str], model: Optional[str]) -> None:
|
|
"""Remember a (binary, model) that aborts on --split-mode tensor."""
|
|
key = cls._tensor_split_cache_key(binary, model)
|
|
if key is not None:
|
|
cls._tensor_split_abort_keys.add(key)
|
|
|
|
@staticmethod
|
|
def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]:
|
|
"""Return DLL dirs from pip-installed CUDA wheels under
|
|
``<prefix>/Lib/site-packages/`` so llama-server.exe can load
|
|
``cudart64_X.dll`` / ``cublas64_X.dll`` without a system CUDA toolkit.
|
|
Mirrors the Linux ``nvidia/cu*/lib`` LD_LIBRARY_PATH block, covering the
|
|
Windows wheel layouts seen in the wild:
|
|
* ``nvidia/<pkg>/bin`` -- legacy modular wheels.
|
|
* ``nvidia/<pkg>/bin/x86_64`` and ``.../bin/x64`` -- CUDA 13 layout
|
|
for unsuffixed packages (#5106).
|
|
* ``nvidia/<pkg>/Library/bin`` (and arch subdirs) -- conda repacks.
|
|
* ``torch/lib`` -- PyTorch's CUDA-bundled wheel can ship
|
|
``cudart64_*.dll`` here; mirrors install_llama_prebuilt.py.
|
|
|
|
Walks with ``Path.iterdir`` not ``glob.glob`` so it's safe against
|
|
Windows paths containing ``[`` or ``]`` (valid in usernames)."""
|
|
site_packages = Path(prefix) / "Lib" / "site-packages"
|
|
out: list[str] = []
|
|
seen: set[str] = set()
|
|
|
|
def _add(path: Path) -> None:
|
|
if not path.is_dir():
|
|
return
|
|
key = os.path.normcase(os.path.abspath(str(path)))
|
|
if key in seen:
|
|
return
|
|
seen.add(key)
|
|
out.append(str(path))
|
|
|
|
nvidia_root = site_packages / "nvidia"
|
|
if nvidia_root.is_dir():
|
|
for pkg_dir in nvidia_root.iterdir():
|
|
if not pkg_dir.is_dir():
|
|
continue
|
|
# Arch-specific subdirs first so the explicit cudart64_X.dll
|
|
# location wins over an empty sibling ``bin``.
|
|
for sub in (
|
|
pkg_dir / "bin" / "x86_64",
|
|
pkg_dir / "bin" / "x64",
|
|
pkg_dir / "bin",
|
|
pkg_dir / "Library" / "bin" / "x86_64",
|
|
pkg_dir / "Library" / "bin" / "x64",
|
|
pkg_dir / "Library" / "bin",
|
|
):
|
|
_add(sub)
|
|
_add(site_packages / "torch" / "lib")
|
|
return out
|
|
|
|
@staticmethod
|
|
def _build_windows_path_dirs(binary_dir: str, prefix: str, cuda_path: str) -> list[str]:
|
|
"""Ordered PATH entries prepended so llama-server.exe resolves cudart /
|
|
cublas DLLs: binary_dir, pip nvidia wheels, CUDA_PATH/bin, .../bin/x64.
|
|
Extracted so test_windows_gpu_detection_mock tests the real logic. #5106."""
|
|
path_dirs = [binary_dir]
|
|
path_dirs.extend(LlamaCppBackend._windows_pip_nvidia_dll_dirs(prefix))
|
|
if cuda_path:
|
|
cuda_bin = os.path.join(cuda_path, "bin")
|
|
if os.path.isdir(cuda_bin):
|
|
path_dirs.append(cuda_bin)
|
|
cuda_bin_x64 = os.path.join(cuda_path, "bin", "x64")
|
|
if os.path.isdir(cuda_bin_x64):
|
|
path_dirs.append(cuda_bin_x64)
|
|
return path_dirs
|
|
|
|
@staticmethod
|
|
def _llama_server_env_for_binary(binary: str) -> dict[str, str]:
|
|
"""Build a subprocess env that lets llama-server resolve native libs."""
|
|
env = child_env_without_native_path_secret()
|
|
binary_dir = str(Path(binary).parent)
|
|
|
|
if sys.platform == "win32":
|
|
# Ordering: see _build_windows_path_dirs. #5106.
|
|
path_dirs = LlamaCppBackend._build_windows_path_dirs(
|
|
binary_dir,
|
|
sys.prefix,
|
|
os.environ.get("CUDA_PATH", ""),
|
|
)
|
|
existing_path = env.get("PATH", "")
|
|
env["PATH"] = ";".join(path_dirs) + ";" + existing_path
|
|
|
|
# ROCm: the prebuilt bundles rocblas.dll but NOT the Tensile
|
|
# kernel files (rocblas/library/*.dat + *.hsaco); the DLL searches
|
|
# <binary_dir>/rocblas/library/ which doesn't exist.
|
|
_hip_path = os.environ.get("HIP_PATH", os.environ.get("ROCM_PATH", ""))
|
|
if _hip_path:
|
|
_rocblas_lib = os.path.join(_hip_path, "bin", "rocblas", "library")
|
|
if os.path.isdir(_rocblas_lib):
|
|
env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib)
|
|
else:
|
|
# Linux: LD_LIBRARY_PATH for shared libs next to the binary plus
|
|
# CUDA runtime libs (libcudart, libcublas, etc.)
|
|
import platform
|
|
|
|
lib_dirs = []
|
|
# WSL: system HIP before the bundle's (which segfaults on /dev/dxg).
|
|
lib_dirs.extend(_wsl_system_rocm_lib_dirs())
|
|
if lib_dirs:
|
|
env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")
|
|
lib_dirs.append(binary_dir)
|
|
_arch = platform.machine() # x86_64, aarch64, etc.
|
|
|
|
# Pip-installed nvidia CUDA runtime libs. The prebuilt binary links
|
|
# libcudart.so.13 / libcublas.so.13 which live here, not in
|
|
# /usr/local/cuda.
|
|
import glob as _glob
|
|
|
|
for _nv_pattern in [
|
|
os.path.join(sys.prefix, "lib", "python*", "site-packages", "nvidia", _sub, "lib")
|
|
for _sub in ("cu*", "cudnn", "nvjitlink")
|
|
]:
|
|
for _nv_dir in _glob.glob(_nv_pattern):
|
|
if os.path.isdir(_nv_dir):
|
|
lib_dirs.append(_nv_dir)
|
|
|
|
for cuda_lib in [
|
|
"/usr/local/cuda/lib64",
|
|
f"/usr/local/cuda/targets/{_arch}-linux/lib",
|
|
# Fallback CUDA compat paths (e.g. binary built with CUDA 12
|
|
# where default /usr/local/cuda is CUDA 13+).
|
|
"/usr/local/cuda-12/lib64",
|
|
"/usr/local/cuda-12.8/lib64",
|
|
f"/usr/local/cuda-12/targets/{_arch}-linux/lib",
|
|
f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib",
|
|
]:
|
|
if os.path.isdir(cuda_lib):
|
|
lib_dirs.append(cuda_lib)
|
|
existing_ld = env.get("LD_LIBRARY_PATH", "")
|
|
new_ld = ":".join(lib_dirs)
|
|
env["LD_LIBRARY_PATH"] = f"{new_ld}:{existing_ld}" if existing_ld else new_ld
|
|
|
|
return env
|
|
|
|
@staticmethod
|
|
def _select_gpus(
|
|
model_size_bytes: int,
|
|
gpus: list[tuple[int, int]],
|
|
usable_fraction: Optional[float] = None,
|
|
total_by_idx: Optional[dict[int, int]] = None,
|
|
per_device_overhead_bytes: int = 0,
|
|
min_gpus: int = 1,
|
|
) -> tuple[Optional[list[int]], bool]:
|
|
"""Pick GPU(s) for a model from estimated VRAM and free memory.
|
|
|
|
``min_gpus`` (default 1, capped at ``len(gpus)``) keeps a downgraded
|
|
tensor/multi-GPU request spread instead of collapsing to one card.
|
|
|
|
``model_size_bytes`` should include weights and estimated KV cache.
|
|
``usable_fraction`` (default ``_GPU_PIN_VRAM_FRACTION``) provides
|
|
headroom for compute buffers, CUDA context, and other runtime
|
|
overhead; callers lower it when MTP reserves VRAM for a draft model.
|
|
``total_by_idx`` (index -> total MiB) makes the headroom an ABSOLUTE
|
|
``(1 - fraction) * total`` per GPU instead of a fraction of free.
|
|
``per_device_overhead_bytes`` is the fixed layer-split cost per GPU beyond
|
|
the first; a k-GPU pin must hold ``model + (k-1) * overhead`` or it can OOM
|
|
a device after -ngl -1 (no --fit fallback). Single-GPU adds none.
|
|
|
|
Returns (gpu_indices, use_fit):
|
|
- ([1], False) fits on 1 GPU at the headroom threshold
|
|
- ([1, 2], False) needs 2 GPUs
|
|
- (None, True) too large, let --fit handle it
|
|
"""
|
|
if not gpus:
|
|
return None, True
|
|
|
|
min_gpus = max(1, min(min_gpus, len(gpus)))
|
|
model_size_mib = model_size_bytes / (1024 * 1024)
|
|
if usable_fraction is None:
|
|
usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION
|
|
overhead_mib = per_device_overhead_bytes / (1024 * 1024)
|
|
|
|
# Per-GPU usable budget: free - (1-frac)*total when total is known, else
|
|
# the legacy free*frac (also covers a total-0 two-column probe).
|
|
def _usable(idx: int, free_mib: int) -> float:
|
|
t = total_by_idx.get(idx, 0) if total_by_idx else 0
|
|
if t > 0:
|
|
return max(0.0, free_mib - (1.0 - usable_fraction) * t)
|
|
return free_mib * usable_fraction
|
|
|
|
# Rank by usable budget (free - reserve), not raw free: a more-used large
|
|
# card can have less usable room than a less-used small one.
|
|
ranked = sorted(gpus, key = lambda g: _usable(g[0], g[1]), reverse = True)
|
|
|
|
# Cap a downgraded multi-GPU request to the usable count so it doesn't pull
|
|
# in a near-full card to hit min_gpus. No-op for the default min_gpus == 1.
|
|
usable_count = sum(1 for idx, free_mib in ranked if _usable(idx, free_mib) > overhead_mib)
|
|
min_gpus = max(1, min(min_gpus, usable_count or 1))
|
|
|
|
# Try 1 GPU at the usable-VRAM threshold (only when one device is allowed).
|
|
if min_gpus <= 1 and _usable(ranked[0][0], ranked[0][1]) >= model_size_mib:
|
|
return [ranked[0][0]], False
|
|
|
|
# Try N GPUs (most-free first); each past the first adds per-device overhead.
|
|
# Require at least min_gpus devices before accepting a fit.
|
|
cumulative = 0.0
|
|
selected = []
|
|
for idx, free_mib in ranked:
|
|
selected.append(idx)
|
|
cumulative += _usable(idx, free_mib)
|
|
if (
|
|
len(selected) >= min_gpus
|
|
and cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib
|
|
):
|
|
return sorted(selected), False
|
|
|
|
# Too large even for all GPUs; let --fit handle it
|
|
logger.debug(
|
|
"Model does not fit in available GPU memory, falling back to --fit",
|
|
model_size_mib = round(model_size_mib, 2),
|
|
ranked_gpus = ranked,
|
|
)
|
|
return None, True
|
|
|
|
# ── KV cache VRAM estimation ─────────────────────────────────────
|
|
|
|
def _can_estimate_kv(self) -> bool:
|
|
"""True if we have enough GGUF metadata to estimate KV cache size."""
|
|
if self._n_layers is None:
|
|
return False
|
|
# MLA: kv_lora_rank suffices (K-only cache).
|
|
if self._kv_lora_rank is not None:
|
|
return True
|
|
# New-style: need explicit key AND value dimensions.
|
|
if self._kv_key_length is not None and self._kv_value_length is not None:
|
|
return True
|
|
# Legacy: need embedding_length + a head count (scalar or per-layer).
|
|
return self._embedding_length is not None and (
|
|
self._n_kv_heads is not None
|
|
or self._n_heads is not None
|
|
or self._n_kv_heads_by_layer is not None
|
|
)
|
|
|
|
def _kv_heads_for_layer(self, layer_idx: int, fallback: int) -> int:
|
|
if self._n_kv_heads_by_layer is not None and layer_idx < len(self._n_kv_heads_by_layer):
|
|
return self._n_kv_heads_by_layer[layer_idx]
|
|
return fallback
|
|
|
|
def _legacy_head_dim(self) -> int:
|
|
"""Head-dim fallback for GGUFs without explicit key/value dims. Reached
|
|
only via the legacy branch of _can_estimate_kv(), so _embedding_length
|
|
is non-None here."""
|
|
return self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator]
|
|
|
|
def _estimate_kv_cache_bytes(
|
|
self,
|
|
n_ctx: int,
|
|
cache_type_kv: Optional[str] = None,
|
|
*,
|
|
swa_full: bool = False,
|
|
n_parallel: int = 1,
|
|
kv_unified: bool = True,
|
|
ctx_checkpoints: int = 0,
|
|
) -> int:
|
|
"""Estimate KV cache VRAM for a given context length.
|
|
|
|
5-path architecture-aware estimation:
|
|
1. MLA -- compressed KV latent + RoPE, K-only (no separate V)
|
|
2. Hybrid -- only attention layers need KV (Mamba layers don't)
|
|
3. SWA -- sliding-window layers cache min(ctx, window) tokens
|
|
4. GQA -- standard full KV with explicit key/value dimensions
|
|
5. Legacy -- fallback using embed // n_heads
|
|
|
|
Server-flag knobs (mirror llama-server's CLI):
|
|
swa_full -- --swa-full: SWA layers cache full n_ctx (path 3->4).
|
|
n_parallel -- --parallel slots: non-SWA constant, SWA scale linearly.
|
|
kv_unified -- --kv-unified: memory no-op (API forward-compat).
|
|
ctx_checkpoints -- --ctx-checkpoints: N SWA snapshots per slot.
|
|
|
|
Returns 0 if metadata is insufficient.
|
|
"""
|
|
if not self._can_estimate_kv() or n_ctx <= 0:
|
|
return 0
|
|
|
|
n_layers = self._n_layers # type: ignore[assignment]
|
|
# Gemma 3n / Gemma 4 reuse earlier KV in the last ``shared_kv_layers``
|
|
# blocks (no cache). Floor at 1 so a bad GGUF can't zero out KV.
|
|
shared = self._shared_kv_layers or 0
|
|
n_layers_kv = max(1, n_layers - shared)
|
|
n_kv = self._n_kv_heads or self._n_heads or 1 # type: ignore[assignment]
|
|
|
|
# Bytes per element depends on KV cache quantization
|
|
bpe = _kv_bytes_per_elem(cache_type_kv)
|
|
|
|
slots = max(1, n_parallel)
|
|
|
|
# Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5)
|
|
# One compressed KV latent per token/layer (shared across heads); V is
|
|
# reconstructed from it, no separate V cache. key_length = kv_lora_rank
|
|
# + rope_dim. MLA GGUFs set head_count_kv=1; default to 1 if absent to
|
|
# avoid falling back to n_heads (e.g. 128 for DeepSeek-V3) which 128x's.
|
|
if self._kv_lora_rank is not None:
|
|
n_kv_mla = self._n_kv_heads or 1
|
|
rope_dim = self._key_length_mla or 64
|
|
key_len = self._kv_key_length or (self._kv_lora_rank + rope_dim)
|
|
return int(n_layers_kv * n_ctx * n_kv_mla * key_len * bpe)
|
|
|
|
key_len = self._kv_key_length
|
|
val_len = self._kv_value_length
|
|
|
|
# Path 2: Hybrid Mamba/Attention (Qwen3.5-27B, Qwen3.5-35B-A3B)
|
|
# Only 1 in N layers is attention; the rest are Mamba (no KV cache).
|
|
if self._ssm_inner_size is not None and self._full_attention_interval is not None:
|
|
fai = self._full_attention_interval
|
|
n_attn = -(-n_layers // fai) if fai > 0 else n_layers # ceiling division
|
|
if key_len is not None and val_len is not None:
|
|
return int(n_attn * n_ctx * n_kv * (key_len + val_len) * bpe)
|
|
head_dim = self._legacy_head_dim()
|
|
return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe)
|
|
|
|
# Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). Pattern
|
|
# from the resolver; if absent, falls through to the legacy 1/4-global
|
|
# heuristic. --parallel N accounting (verified against llama-server):
|
|
# non-SWA cells = n_ctx split across slots (CONSTANT); SWA per-slot cells
|
|
# = 2*sliding_window (capped at n_ctx/per_slot_ctx) -> LINEAR in slots.
|
|
# --swa-full forces full n_ctx for SWA; --ctx-checkpoints N adds snapshots.
|
|
if (
|
|
self._sliding_window is not None
|
|
and self._sliding_window > 0
|
|
and key_len is not None
|
|
and val_len is not None
|
|
):
|
|
swa = self._sliding_window
|
|
per_slot_ctx = max(1, n_ctx // slots)
|
|
# --swa-full caches full per_slot_ctx (constant n_ctx total); else SWA
|
|
# caches 2*sliding_window per slot, clamped at per-slot ctx.
|
|
swa_cells_per_slot = per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx)
|
|
key_len_swa = self._kv_key_length_swa or key_len
|
|
val_len_swa = self._kv_value_length_swa or val_len
|
|
if self._sliding_window_pattern is not None:
|
|
global_bytes = 0.0 # constant across slots
|
|
swa_bytes_per_slot = 0.0 # multiplied by slots
|
|
checkpoint_extra_per_slot = 0.0
|
|
# Only layers that allocate their own KV; trailing shared layers
|
|
# reuse earlier caches.
|
|
for layer_idx in range(n_layers_kv):
|
|
layer_n_kv = self._kv_heads_for_layer(layer_idx, n_kv)
|
|
is_swa = (
|
|
layer_idx < len(self._sliding_window_pattern)
|
|
and self._sliding_window_pattern[layer_idx]
|
|
)
|
|
if is_swa:
|
|
swa_bytes_per_slot += (
|
|
swa_cells_per_slot * layer_n_kv * (key_len_swa + val_len_swa) * bpe
|
|
)
|
|
if ctx_checkpoints > 0 and not swa_full:
|
|
checkpoint_extra_per_slot += (
|
|
ctx_checkpoints
|
|
* swa
|
|
* layer_n_kv
|
|
* (key_len_swa + val_len_swa)
|
|
* bpe
|
|
)
|
|
else:
|
|
global_bytes += n_ctx * layer_n_kv * (key_len + val_len) * bpe
|
|
return int(global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot))
|
|
n_global = max(1, n_layers_kv // 4)
|
|
n_swa = n_layers_kv - n_global
|
|
kv_per_token = n_kv * (key_len + val_len) * bpe
|
|
kv_per_token_swa = n_kv * (key_len_swa + val_len_swa) * bpe
|
|
global_bytes = n_global * n_ctx * kv_per_token
|
|
swa_bytes_per_slot = n_swa * swa_cells_per_slot * kv_per_token_swa
|
|
checkpoint_extra_per_slot = (
|
|
ctx_checkpoints * n_swa * swa * kv_per_token_swa
|
|
if ctx_checkpoints > 0 and not swa_full
|
|
else 0.0
|
|
)
|
|
return int(global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot))
|
|
|
|
# Path 4: Standard GQA with explicit key/value dimensions
|
|
if key_len is not None and val_len is not None:
|
|
return int(n_layers_kv * n_ctx * n_kv * (key_len + val_len) * bpe)
|
|
|
|
# Path 5: Legacy fallback (old GGUFs without explicit dimensions)
|
|
head_dim = self._legacy_head_dim()
|
|
return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe)
|
|
|
|
def _draft_backend_for(self, drafter_path: str) -> Optional["LlamaCppBackend"]:
|
|
"""Lightweight backend with a drafter GGUF's metadata, to size its own KV
|
|
via _estimate_kv_cache_bytes. Cached per path; None if unreadable."""
|
|
cache = getattr(self, "_draft_backend_cache", None)
|
|
if cache is not None and cache[0] == drafter_path:
|
|
return cache[1]
|
|
db: Optional[LlamaCppBackend] = None
|
|
try:
|
|
db = LlamaCppBackend.__new__(LlamaCppBackend)
|
|
for attr in (
|
|
"_context_length",
|
|
"_n_layers",
|
|
"_n_kv_heads",
|
|
"_n_heads",
|
|
"_embedding_length",
|
|
"_kv_key_length",
|
|
"_kv_value_length",
|
|
"_kv_lora_rank",
|
|
"_sliding_window",
|
|
"_sliding_window_pattern",
|
|
"_ssm_inner_size",
|
|
"_full_attention_interval",
|
|
"_key_length_mla",
|
|
"_n_kv_heads_by_layer",
|
|
"_kv_key_length_swa",
|
|
"_kv_value_length_swa",
|
|
"_shared_kv_layers",
|
|
"_nextn_predict_layers",
|
|
):
|
|
setattr(db, attr, None)
|
|
db._model_identifier = "mtp-draft"
|
|
db._read_gguf_metadata(drafter_path)
|
|
except Exception as e: # unreadable drafter -> caller falls back
|
|
logger.debug(f"Could not read drafter GGUF for MTP budget: {e}")
|
|
db = None
|
|
self._draft_backend_cache = (drafter_path, db)
|
|
return db
|
|
|
|
def _mtp_draft_kv_bytes(
|
|
self,
|
|
n_ctx: int,
|
|
*,
|
|
drafter_path: Optional[str] = None,
|
|
draft_cache_type_k: Optional[str] = None,
|
|
draft_cache_type_v: Optional[str] = None,
|
|
n_parallel: int = 1,
|
|
) -> Optional[int]:
|
|
"""Draft KV cache bytes at n_ctx, sized from GGUF dims (K and V types are
|
|
independent). Separate drafter (Gemma): its own KV via _estimate_kv_cache_bytes
|
|
at the heavier type. Embedded head (Qwen): nextn_predict_layers attention
|
|
layers from the main dims. None when dims are missing (flat fallback)."""
|
|
if n_ctx <= 0:
|
|
return None
|
|
bpe_k = _kv_bytes_per_elem(draft_cache_type_k)
|
|
bpe_v = _kv_bytes_per_elem(draft_cache_type_v)
|
|
if drafter_path:
|
|
db = self._draft_backend_for(drafter_path)
|
|
if db is None or not db._can_estimate_kv():
|
|
return None
|
|
heavier = draft_cache_type_k if bpe_k >= bpe_v else draft_cache_type_v
|
|
# The drafter is served under the same --parallel slot count as the
|
|
# main model, so price its KV per slot too: a sliding-window drafter
|
|
# (Gemma) grows KV with slots and would otherwise be under-reserved.
|
|
kv = db._estimate_kv_cache_bytes(n_ctx, heavier, n_parallel = n_parallel)
|
|
return kv or None
|
|
nextn = self._nextn_predict_layers or 0
|
|
n_kv = self._n_kv_heads or self._n_heads
|
|
k_len = self._kv_key_length
|
|
v_len = self._kv_value_length
|
|
if not (nextn and n_kv and k_len and v_len):
|
|
return None
|
|
# The embedded MTP head is one draft layer, so a quantized draft KV can't
|
|
# amortize its overhead and fits *less* context than f16 (llama.cpp#24102).
|
|
# Floor it at f16: a quantized override is priced as f16, f32 keeps its 4
|
|
# bytes. The separate-drafter branch is multi-layer, so it keeps its type.
|
|
f16_bpe = _kv_bytes_per_elem("f16")
|
|
bpe_k = max(bpe_k, f16_bpe)
|
|
bpe_v = max(bpe_v, f16_bpe)
|
|
return int(nextn * n_kv * (k_len * bpe_k + v_len * bpe_v) * n_ctx)
|
|
|
|
def _estimate_mtp_overhead_bytes(
|
|
self,
|
|
n_ctx: int,
|
|
*,
|
|
spec_draft_n_max: int = 0,
|
|
draft_cache_type_k: Optional[str] = None,
|
|
draft_cache_type_v: Optional[str] = None,
|
|
drafter_path: Optional[str] = None,
|
|
draft_weights_bytes: int = 0,
|
|
n_parallel: int = 1,
|
|
mtp_keeps_target_ctx: bool = True,
|
|
) -> Optional[int]:
|
|
"""MTP draft reserve at ``n_ctx`` = draft KV (grows with ctx) + separate-
|
|
drafter weights + (MTP + MLA only) a duplicated target KV context. The
|
|
verify buffer rides in the ctx-fit headroom (no tuned constant). None when
|
|
the draft KV can't be sized (caller keeps the flat fallback).
|
|
``draft_weights_bytes`` is the drafter file size (0 for embedded).
|
|
``mtp_keeps_target_ctx`` is True for MTP draft modes (which keep the
|
|
duplicated target context) and False for separate-drafter spec modes
|
|
(draft-simple/draft-eagle3), which do not."""
|
|
draft_kv = self._mtp_draft_kv_bytes(
|
|
n_ctx,
|
|
drafter_path = drafter_path,
|
|
draft_cache_type_k = draft_cache_type_k,
|
|
draft_cache_type_v = draft_cache_type_v,
|
|
n_parallel = n_parallel,
|
|
)
|
|
weights = max(0, draft_weights_bytes)
|
|
# MLA models (GLM-5.x, DeepSeek, Kimi-K2) under MTP keep a *second* full copy
|
|
# of the target model's KV context for draft verification -- llama.cpp's
|
|
# `ctx_tgt=yes` -- allocated at f16 regardless of the main cache type. It is
|
|
# ~the main KV again and dwarfs the embedded draft head (GLM-5.2 @ 1M ctx:
|
|
# a ~2 GiB head next to a ~89 GiB target copy), so omitting it lets auto-fit
|
|
# pick a context that fits on paper but OOMs cublasCreate at the first
|
|
# decode. Gated on both MLA (kv_lora_rank present) and the engaged mode
|
|
# actually being MTP: non-MLA MTP (Qwen/Gemma) keeps no such copy, and the
|
|
# separate-drafter spec modes (draft-simple/draft-eagle3) load a small
|
|
# distinct drafter with its own KV -- already counted in draft_kv/weights --
|
|
# rather than duplicating the target, so they must not be charged for it.
|
|
target_ctx_copy = 0
|
|
if mtp_keeps_target_ctx and self._kv_lora_rank is not None:
|
|
target_ctx_copy = self._estimate_kv_cache_bytes(n_ctx, "f16", n_parallel = n_parallel)
|
|
if draft_kv is None:
|
|
# KV unsized (exotic/remote drafter): still reserve known weights + any
|
|
# MLA target copy so a large config can't launch over budget (the small
|
|
# unsized draft KV rides in the cushion). Nothing known -> None, so the
|
|
# caller keeps the flat fallback.
|
|
total = weights + target_ctx_copy
|
|
return total if total > 0 else None
|
|
return draft_kv + weights + target_ctx_copy
|
|
|
|
_DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Studio does not override it
|
|
_COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate
|
|
# Soft VRAM the modeled terms omit; charged to the fit budget on tight tiers (#6682).
|
|
_CUDA_CONTEXT_RESERVE_BYTES = 320 * 1024 * 1024 # CUDA ctx + cuBLAS workspace (~330 MiB)
|
|
_MMPROJ_VRAM_SAFETY = 1.4 # mmproj worst-case buffer vs file size (runtime ~1.3x)
|
|
_MTP_DRAFT_COMPUTE_BYTES = 224 * 1024 * 1024 # MTP draft decode graph beyond its KV
|
|
# The flash-attn KQ mask + attention scratch grow ~linearly with context; the flat
|
|
# _estimate_compute_buffer_bytes term only covers ctx -> 0. The per-token rate
|
|
# depends on the KV cache type: a QUANTIZED cache (q8_0/q5/q4/iq4) needs a
|
|
# context-sized dequant scratch that scales with n_embd, measured at 0.74-2.02 x
|
|
# n_embd across Qwen3.5/3.6 (2B/4B/9B/27B) and Gemma-4 (12B/31B) at q8_0; an
|
|
# f16/bf16/f32 cache skips the dequant and pays only the KQ mask, a flat n_ubatch*2
|
|
# bytes per context token regardless of n_embd (measured 1024 B/tok on Qwen-9B and
|
|
# Gemma-31B alike). So Qwen3.5-4B at 256k is 1.30 GiB at q8_0 vs 0.31 GiB at f16.
|
|
# 2.25 covers the worst quantized case (Qwen3.5-4B, ~2.0x) plus the under-modeled
|
|
# flat base; the mask safety covers the f16 base gap. Without this term, tight tiers
|
|
# at extreme context over-pin and spill to CPU (the 3% cushion is only ~0.25 GiB on
|
|
# an 8 GB card, far below the ~1-2.4 GiB quantized buffer at 256k): e.g. Qwen3.5-4B
|
|
# Q4 at 256k needs ~8.5 GiB on a real 8 GB card (weights 2.4 + KV 4.3 + compute 1.3
|
|
# + CUDA ctx) -> CPU spill; with this reserve the auto context caps to ~210k, fits.
|
|
_CTX_COMPUTE_BYTES_PER_EMBD = 2.25 # quantized KV, regular attention (dequant scratch)
|
|
_CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25 # quantized KV, MLA (compressed attn: measured 0.94x)
|
|
_CTX_COMPUTE_F16_MASK_SAFETY = 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok)
|
|
|
|
def _estimate_compute_buffer_bytes(
|
|
self,
|
|
*,
|
|
n_ubatch: Optional[int] = None,
|
|
n_parallel: int = 1,
|
|
per_device_tensor: bool = False,
|
|
) -> int:
|
|
"""Per-device compute-graph buffer (bytes) from GGUF dims: a vocab-width
|
|
output buffer + activation scratch. Context-independent; scales with
|
|
``--parallel`` (serving slots). Tensor mode materializes it on every device.
|
|
A slight upper bound over measured allocations; 0 when dims are missing."""
|
|
n_vocab = self._vocab_size or 0
|
|
n_embd = self._embedding_length or 0
|
|
if n_vocab <= 0 or n_embd <= 0:
|
|
return 0
|
|
ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH))
|
|
par = max(1, int(n_parallel))
|
|
out_buffer = n_vocab * ub * 4 # f32 output/logits buffer
|
|
act_scratch = 4 * n_embd * ub * 4 # a few resident hidden-width buffers
|
|
if per_device_tensor:
|
|
# Output + comm/staging materialized on every device, every slot.
|
|
compute = 2 * act_scratch + out_buffer * par
|
|
else:
|
|
# Each extra concurrent slot adds one output buffer (chat decode sizes
|
|
# ~one logit row per slot; would under-count embeddings/--logits-all,
|
|
# not run here). Matches measured {1:36,2:492,4:1388,8:3220} MiB.
|
|
compute = act_scratch + out_buffer * max(0, par - 1)
|
|
return int(compute * self._COMPUTE_BUFFER_SAFETY)
|
|
|
|
def _compute_buffer_ctx_bytes(
|
|
self,
|
|
n_ctx: int,
|
|
n_ubatch: Optional[int] = None,
|
|
cache_type_kv: Optional[str] = None,
|
|
) -> int:
|
|
"""Context-linear growth of the per-device compute buffer (bytes), charged
|
|
on top of the flat ``_estimate_compute_buffer_bytes``. The flash-attn KQ
|
|
mask + attention scratch scale ~linearly with context and with the micro-
|
|
batch; the flat term only covers ctx -> 0. A quantized KV cache adds a
|
|
context-sized dequant scratch that scales with n_embd; f16/bf16/f32 pays only
|
|
the KQ mask, a flat n_ubatch*2 bytes per context token. ``cache_type_kv`` None
|
|
-> f16 (llama.cpp's default; an env-set quantized cache is budgeted as f16 on
|
|
the KV side, whose over-reservation absorbs the dequant scratch). Returns 0
|
|
when dims are missing or ``n_ctx`` <= 0."""
|
|
n_embd = self._embedding_length or 0
|
|
if n_embd <= 0 or n_ctx <= 0:
|
|
return 0
|
|
ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH))
|
|
if _kv_bytes_per_elem(cache_type_kv) < 2.0:
|
|
# Quantized cache: the dequant scratch dominates and scales with n_embd.
|
|
# MLA (compressed KV) needs far less of it: measured 0.94 x n_embd on
|
|
# GLM-5.2 and Kimi-K2.7 vs up to 2.02x on regular attention.
|
|
ub_scale = ub / self._DEFAULT_N_UBATCH
|
|
rate = (
|
|
self._CTX_COMPUTE_BYTES_PER_EMBD_MLA
|
|
if self._key_length_mla
|
|
else self._CTX_COMPUTE_BYTES_PER_EMBD
|
|
)
|
|
per_tok = rate * n_embd * ub_scale
|
|
else:
|
|
# f16/bf16/f32: only the KQ mask ([n_kv, n_ubatch] f16), n_embd-independent.
|
|
per_tok = ub * 2 * self._CTX_COMPUTE_F16_MASK_SAFETY
|
|
return int(per_tok * n_ctx)
|
|
|
|
def _slots_that_fit_on_gpu(
|
|
self,
|
|
n_parallel: int,
|
|
effective_ctx: int,
|
|
gpus: list[tuple[int, int]],
|
|
total_by_idx: Optional[dict[int, int]],
|
|
base_footprint_bytes: int,
|
|
cache_type_kv: Optional[str],
|
|
pin_fraction: float,
|
|
per_device_overhead_bytes: int,
|
|
min_gpus: int,
|
|
n_ubatch: Optional[int] = None,
|
|
) -> tuple[Optional[list[int]], bool, int]:
|
|
"""Largest serving-slot count in [1, n_parallel) whose fully-on-GPU footprint fits,
|
|
so Studio keeps the model on GPU (-ngl -1) instead of --fit on, which offloads layers
|
|
to host and collapses decode ~3x (oobabooga #6718). ``base_footprint_bytes`` is the
|
|
slot-independent footprint (weights + soft overhead + MTP + context-linear compute,
|
|
minus the folded compute buffer); each candidate re-adds the slot-sized compute buffer
|
|
and KV, then re-selects GPUs like the explicit-context path. Returns (gpu_indices,
|
|
use_fit=False, slots) for the largest fitting count, else (None, True, n_parallel).
|
|
Only ever reduces; deterministic and unit-testable with synthetic VRAM maps."""
|
|
for slots in range(n_parallel - 1, 0, -1):
|
|
cb = self._estimate_compute_buffer_bytes(
|
|
n_ubatch = n_ubatch, n_parallel = slots, per_device_tensor = False
|
|
)
|
|
if cb <= 0:
|
|
cb = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024
|
|
total = (
|
|
base_footprint_bytes
|
|
+ cb
|
|
+ self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = slots)
|
|
)
|
|
gpu_indices, use_fit = self._select_gpus(
|
|
total,
|
|
gpus,
|
|
usable_fraction = pin_fraction,
|
|
total_by_idx = total_by_idx,
|
|
per_device_overhead_bytes = per_device_overhead_bytes,
|
|
min_gpus = min_gpus,
|
|
)
|
|
if not use_fit:
|
|
return gpu_indices, False, slots
|
|
return None, True, n_parallel
|
|
|
|
def _fit_context_to_vram(
|
|
self,
|
|
requested_ctx: int,
|
|
available_mib: int,
|
|
model_size_bytes: int,
|
|
cache_type_kv: Optional[str] = None,
|
|
min_ctx: int = 4096,
|
|
*,
|
|
swa_full: bool = False,
|
|
n_parallel: int = 1,
|
|
kv_unified: bool = True,
|
|
ctx_checkpoints: int = 0,
|
|
kv_on_gpu: bool = True,
|
|
mtp_engaged: bool = False,
|
|
mtp_overhead_fn: Optional[Callable[[int], int]] = None,
|
|
compute_ctx_bytes_fn: Optional[Callable[[int], int]] = None,
|
|
budget_frac: Optional[float] = None,
|
|
total_mib: Optional[int] = None,
|
|
) -> int:
|
|
"""Return the largest context length that fits in GPU VRAM.
|
|
|
|
Budget caps occupancy at ``_CTX_FIT_VRAM_FRACTION`` of the card: an
|
|
absolute ``free - (1 - frac) * total`` when ``total_mib`` is given, else
|
|
``free * frac``. Weights alone over budget returns ``requested_ctx``.
|
|
|
|
``kv_on_gpu`` mirrors ``--kv-offload`` (default on); when False the KV
|
|
cache lives in CPU RAM and the requested context is honored verbatim.
|
|
Other keyword args mirror ``_estimate_kv_cache_bytes``.
|
|
|
|
``mtp_engaged`` reserves extra VRAM for the MTP draft model's KV cache +
|
|
compute buffers, else tight tiers (e.g. 32 GB) spill to a slower path.
|
|
"""
|
|
if not self._can_estimate_kv():
|
|
logger.debug(
|
|
"Skipping context fit because KV cache metadata is unavailable",
|
|
requested_ctx = requested_ctx,
|
|
available_mib = available_mib,
|
|
)
|
|
return requested_ctx
|
|
|
|
# KV lives off-GPU: no VRAM accounting needed for the cache itself.
|
|
if not kv_on_gpu:
|
|
return requested_ctx
|
|
|
|
kv_kwargs = dict(
|
|
swa_full = swa_full,
|
|
n_parallel = n_parallel,
|
|
kv_unified = kv_unified,
|
|
ctx_checkpoints = ctx_checkpoints,
|
|
)
|
|
|
|
# byte-accurate mtp_overhead_fn supersedes the flat fraction (the fallback
|
|
# when dims can't size the draft KV); callers may override budget_frac.
|
|
if budget_frac is None:
|
|
flat_mtp = mtp_engaged and mtp_overhead_fn is None
|
|
budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if flat_mtp else 0.0)
|
|
# Absolute reserve off total when known, else fraction-of-free; clamp >=0.
|
|
if total_mib is not None and total_mib > 0:
|
|
budget_mib = max(0.0, available_mib - (1.0 - budget_frac) * total_mib)
|
|
else:
|
|
budget_mib = available_mib * budget_frac
|
|
budget_bytes = budget_mib * 1024 * 1024
|
|
model_footprint = model_size_bytes
|
|
|
|
def _mtp_at(ctx: int) -> int:
|
|
return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0
|
|
|
|
def _cc_at(ctx: int) -> int:
|
|
# Context-linear compute-buffer growth (flash-attn KQ mask + scratch);
|
|
# the flat term in model_footprint only covers ctx -> 0.
|
|
return compute_ctx_bytes_fn(ctx) if compute_ctx_bytes_fn is not None else 0
|
|
|
|
# Already fits?
|
|
kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs)
|
|
if model_footprint + kv + _mtp_at(requested_ctx) + _cc_at(requested_ctx) <= budget_bytes:
|
|
return requested_ctx
|
|
|
|
# Weights + compute buffer alone exceed budget -- reducing ctx can't help.
|
|
if model_footprint >= budget_bytes:
|
|
logger.debug(
|
|
"Model footprint exceeds GPU budget before KV cache",
|
|
requested_ctx = requested_ctx,
|
|
available_mib = available_mib,
|
|
model_size_gb = round(model_footprint / (1024**3), 2),
|
|
)
|
|
return requested_ctx
|
|
|
|
# Binary search for max context that fits (KV + MTP draft reserve at that ctx)
|
|
remaining = budget_bytes - model_footprint
|
|
effective_min = min(min_ctx, requested_ctx)
|
|
lo, hi = effective_min, requested_ctx
|
|
best = effective_min
|
|
while lo <= hi:
|
|
mid = (lo + hi) // 2
|
|
kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs)
|
|
if kv + _mtp_at(mid) + _cc_at(mid) <= remaining:
|
|
best = mid
|
|
lo = mid + 1
|
|
else:
|
|
hi = mid - 1
|
|
|
|
# Round down to nearest 256 for alignment, never above requested_ctx
|
|
best = (best // 256) * 256
|
|
best = max(effective_min, best)
|
|
best = min(best, requested_ctx)
|
|
return best
|
|
|
|
# ── Variant fallback ────────────────────────────────────────────
|
|
|
|
@staticmethod
|
|
def _find_smallest_fitting_variant(
|
|
hf_repo: str,
|
|
free_bytes: int,
|
|
hf_token: Optional[str] = None,
|
|
) -> Optional[tuple[str, int]]:
|
|
"""Find the smallest GGUF variant (including all shards) that fits.
|
|
|
|
Groups split shards by variant prefix and sums their sizes (e.g.
|
|
UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total).
|
|
|
|
Returns (first_shard_filename, total_size_bytes) or None.
|
|
"""
|
|
try:
|
|
from huggingface_hub import get_paths_info, list_repo_files
|
|
|
|
files = list_repo_files(hf_repo, token = hf_token)
|
|
gguf_files = [
|
|
f
|
|
for f in files
|
|
if f.lower().endswith(".gguf")
|
|
and not _is_companion_gguf_path(f)
|
|
and not _is_big_endian_gguf_path(f)
|
|
]
|
|
if not gguf_files:
|
|
return None
|
|
|
|
# Sizes for all GGUF files
|
|
path_infos = list(get_paths_info(hf_repo, gguf_files, token = hf_token))
|
|
size_map = {p.path: (p.size or 0) for p in path_infos}
|
|
|
|
# Group by variant: shards share a prefix before -NNNNN-of-NNNNN
|
|
variants: dict[str, list[str]] = {}
|
|
for f in gguf_files:
|
|
m = _SHARD_RE.match(f)
|
|
key = m.group(1) if m else f
|
|
variants.setdefault(key, []).append(f)
|
|
|
|
# Sum shard sizes per variant, track the first shard (for download)
|
|
variant_sizes: list[tuple[str, int, list[str]]] = []
|
|
for key, shard_files in variants.items():
|
|
total = sum(size_map.get(f, 0) for f in shard_files)
|
|
first = sorted(shard_files)[0]
|
|
variant_sizes.append((first, total, shard_files))
|
|
|
|
# Smallest that fits
|
|
variant_sizes.sort(key = lambda x: x[1])
|
|
for first_file, total_size, _ in variant_sizes:
|
|
if total_size > 0 and total_size <= free_bytes:
|
|
return first_file, total_size
|
|
|
|
return None
|
|
except Exception:
|
|
return None
|
|
|
|
# ── Port allocation ───────────────────────────────────────────
|
|
|
|
@staticmethod
|
|
def _find_free_port() -> int:
|
|
"""Find an available TCP port."""
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
s.bind(("127.0.0.1", 0))
|
|
return s.getsockname()[1]
|
|
|
|
# ── Stdout drain (prevents pipe deadlock on Windows) ─────────
|
|
|
|
def _drain_stdout(self):
|
|
"""Read subprocess stdout lines in a background thread.
|
|
|
|
Prevents a pipe-buffer deadlock on Windows (~4 KB buffer): without
|
|
draining, llama-server blocks on writes and never becomes healthy.
|
|
Each line is also teed to ``self._llama_log_fh`` when set, so a
|
|
post-mortem has the full output even if the crash predates the
|
|
drain-thread join in ``_wait_for_health``.
|
|
"""
|
|
try:
|
|
for line in self._process.stdout:
|
|
line = line.rstrip()
|
|
if line:
|
|
self._stdout_lines.append(line)
|
|
logger.debug(f"[llama-server] {line}")
|
|
fh = getattr(self, "_llama_log_fh", None)
|
|
if fh is not None:
|
|
try:
|
|
fh.write(line + "\n")
|
|
fh.flush()
|
|
except (ValueError, OSError):
|
|
# Log file closed under us; tee silently.
|
|
pass
|
|
except Exception:
|
|
# Never let the drain thread die: a full stdout pipe can deadlock
|
|
# llama-server (Windows). Pipe-closed on exit is the common case.
|
|
logger.debug("llama-server stdout drain stopped", exc_info = True)
|
|
|
|
# GGUF KV type sizes for fast skipping
|
|
_GGUF_TYPE_SIZE = {
|
|
0: 1,
|
|
1: 1,
|
|
2: 2,
|
|
3: 2,
|
|
4: 4,
|
|
5: 4,
|
|
6: 4,
|
|
7: 1,
|
|
10: 8,
|
|
11: 8,
|
|
12: 8,
|
|
}
|
|
|
|
@staticmethod
|
|
def _gguf_skip_value(f, vtype: int) -> None:
|
|
"""Skip a GGUF KV value without reading it."""
|
|
sz = LlamaCppBackend._GGUF_TYPE_SIZE.get(vtype)
|
|
if sz is not None:
|
|
f.seek(sz, 1)
|
|
elif vtype == 8: # STRING
|
|
slen = struct.unpack("<Q", f.read(8))[0]
|
|
f.seek(slen, 1)
|
|
elif vtype == 9: # ARRAY
|
|
atype = struct.unpack("<I", f.read(4))[0]
|
|
alen = struct.unpack("<Q", f.read(8))[0]
|
|
elem_sz = LlamaCppBackend._GGUF_TYPE_SIZE.get(atype)
|
|
if elem_sz is not None:
|
|
f.seek(elem_sz * alen, 1)
|
|
elif atype == 8:
|
|
for _ in range(alen):
|
|
slen = struct.unpack("<Q", f.read(8))[0]
|
|
f.seek(slen, 1)
|
|
else:
|
|
for _ in range(alen):
|
|
LlamaCppBackend._gguf_skip_value(f, atype)
|
|
|
|
@staticmethod
|
|
def _gguf_read_array_value(f, atype: int, alen: int) -> Optional[list]:
|
|
if atype == 4: # UINT32
|
|
return [struct.unpack("<I", f.read(4))[0] for _ in range(alen)]
|
|
if atype == 5: # INT32
|
|
return [struct.unpack("<i", f.read(4))[0] for _ in range(alen)]
|
|
if atype == 7: # BOOL
|
|
return [struct.unpack("<?", f.read(1))[0] for _ in range(alen)]
|
|
|
|
for _ in range(alen):
|
|
LlamaCppBackend._gguf_skip_value(f, atype)
|
|
return None
|
|
|
|
def _read_gguf_metadata(self, gguf_path: str) -> None:
|
|
"""Read context_length, architecture params, and chat_template from a GGUF header.
|
|
|
|
Parses only the KV pairs we need (~30ms even for multi-GB files).
|
|
For split GGUFs, metadata is always in shard 1.
|
|
"""
|
|
# Reset metadata so stale flags (e.g. _supports_reasoning) don't
|
|
# carry over when switching models.
|
|
self._context_length = None
|
|
self._chat_template = None
|
|
self._supports_reasoning = False
|
|
self._reasoning_always_on = False
|
|
self._reasoning_style = "enable_thinking"
|
|
self._reasoning_effort_levels = []
|
|
self._reasoning_default = True
|
|
self._supports_preserve_thinking = False
|
|
self._supports_tools = False
|
|
self._n_layers = None
|
|
self._n_kv_heads = None
|
|
self._n_kv_heads_by_layer = None
|
|
self._n_heads = None
|
|
self._embedding_length = None
|
|
self._feed_forward_length = None
|
|
self._vocab_size = None
|
|
self._kv_key_length = None
|
|
self._kv_value_length = None
|
|
self._sliding_window = None
|
|
self._sliding_window_pattern = None
|
|
self._full_attention_interval = None
|
|
self._kv_lora_rank = None
|
|
self._key_length_mla = None
|
|
self._kv_key_length_swa = None
|
|
self._kv_value_length_swa = None
|
|
self._ssm_inner_size = None
|
|
self._ssm_state_size = None
|
|
self._shared_kv_layers = None
|
|
self._nextn_predict_layers = None
|
|
self._architecture = None
|
|
self._is_diffusion = False
|
|
|
|
try:
|
|
canvas_seen = False
|
|
WANTED = {
|
|
"general.architecture",
|
|
"tokenizer.chat_template",
|
|
# Vocab size = tokens array length (no vocab_size key in many GGUFs).
|
|
"tokenizer.ggml.tokens",
|
|
# Block-diffusion marker (DiffusionGemma); routes to the diffusion runner.
|
|
"diffusion.canvas_length",
|
|
# Source-repo hints for the SWA resolver's HF fallback.
|
|
"general.source.huggingface.repository",
|
|
"general.source.url",
|
|
"general.source.repo_url",
|
|
"general.base_model.0.repo_url",
|
|
"general.base_model.0.organization",
|
|
"general.base_model.0.name",
|
|
"general.basename",
|
|
"general.organization",
|
|
"general.size_label",
|
|
"general.finetune",
|
|
}
|
|
# Arch-specific keys added dynamically once we know the arch.
|
|
arch_keys: dict[str, str] = {} # gguf_key -> attribute name
|
|
arch = None
|
|
sliding_window_pattern_period: Optional[int] = None
|
|
general: dict[str, str] = {}
|
|
|
|
with open(gguf_path, "rb") as f:
|
|
magic = struct.unpack("<I", f.read(4))[0]
|
|
if magic != 0x46554747: # b"GGUF" as little-endian u32
|
|
return
|
|
_version = struct.unpack("<I", f.read(4))[0]
|
|
_tensor_count, kv_count = struct.unpack("<QQ", f.read(16))
|
|
|
|
for _ in range(kv_count):
|
|
# Tolerate truncated input (e.g. a partial header from an
|
|
# HTTP byte-range fetch): bail out so the resolver
|
|
# fallback runs on whatever we parsed.
|
|
try:
|
|
key_len_bytes = f.read(8)
|
|
if len(key_len_bytes) < 8:
|
|
break
|
|
key_len = struct.unpack("<Q", key_len_bytes)[0]
|
|
key_bytes = f.read(key_len)
|
|
if len(key_bytes) < key_len:
|
|
break
|
|
key = key_bytes.decode("utf-8")
|
|
vtype_bytes = f.read(4)
|
|
if len(vtype_bytes) < 4:
|
|
break
|
|
vtype = struct.unpack("<I", vtype_bytes)[0]
|
|
except (struct.error, UnicodeDecodeError):
|
|
break
|
|
|
|
try:
|
|
if key in WANTED or key in arch_keys:
|
|
if vtype == 8: # STRING
|
|
slen = struct.unpack("<Q", f.read(8))[0]
|
|
val_s = f.read(slen).decode("utf-8")
|
|
if key.startswith("general.") and key != "general.architecture":
|
|
general[key] = val_s
|
|
if key == "general.architecture":
|
|
arch = val_s
|
|
self._architecture = val_s
|
|
arch_keys = {
|
|
f"{arch}.context_length": "context_length",
|
|
f"{arch}.block_count": "n_layers",
|
|
f"{arch}.attention.head_count_kv": "n_kv_heads",
|
|
f"{arch}.attention.head_count": "n_heads",
|
|
f"{arch}.embedding_length": "embedding_length",
|
|
f"{arch}.feed_forward_length": "feed_forward_length",
|
|
f"{arch}.attention.key_length": "kv_key_length",
|
|
f"{arch}.attention.value_length": "kv_value_length",
|
|
f"{arch}.attention.sliding_window": "sliding_window",
|
|
f"{arch}.attention.sliding_window_pattern": "sliding_window_pattern",
|
|
f"{arch}.full_attention_interval": "full_attention_interval",
|
|
f"{arch}.attention.kv_lora_rank": "kv_lora_rank",
|
|
f"{arch}.attention.key_length_mla": "key_length_mla",
|
|
f"{arch}.attention.key_length_swa": "kv_key_length_swa",
|
|
f"{arch}.attention.value_length_swa": "kv_value_length_swa",
|
|
f"{arch}.attention.shared_kv_layers": "shared_kv_layers",
|
|
f"{arch}.ssm.inner_size": "ssm_inner_size",
|
|
f"{arch}.ssm.state_size": "ssm_state_size",
|
|
f"{arch}.nextn_predict_layers": "nextn_predict_layers",
|
|
}
|
|
elif key == "tokenizer.chat_template":
|
|
self._chat_template = val_s
|
|
elif vtype in (4, 10): # UINT32 or UINT64
|
|
val_i = (
|
|
struct.unpack("<I", f.read(4))[0]
|
|
if vtype == 4
|
|
else struct.unpack("<Q", f.read(8))[0]
|
|
)
|
|
if key == "diffusion.canvas_length":
|
|
canvas_seen = True
|
|
attr = arch_keys.get(key)
|
|
if attr:
|
|
if attr == "sliding_window_pattern":
|
|
sliding_window_pattern_period = val_i
|
|
else:
|
|
setattr(self, f"_{attr}", val_i)
|
|
elif vtype == 9: # ARRAY
|
|
atype = struct.unpack("<I", f.read(4))[0]
|
|
alen = struct.unpack("<Q", f.read(8))[0]
|
|
# Vocab size = token count; keep the length, not the strings.
|
|
if key == "tokenizer.ggml.tokens":
|
|
self._vocab_size = int(alen)
|
|
val_a = self._gguf_read_array_value(f, atype, alen)
|
|
attr = arch_keys.get(key)
|
|
if attr == "n_kv_heads" and val_a is not None:
|
|
self._n_kv_heads_by_layer = [int(x) for x in val_a]
|
|
if self._n_kv_heads is None and self._n_kv_heads_by_layer:
|
|
self._n_kv_heads = max(self._n_kv_heads_by_layer)
|
|
elif attr == "sliding_window_pattern" and val_a is not None:
|
|
self._sliding_window_pattern = [bool(x) for x in val_a]
|
|
sliding_window_pattern_period = None
|
|
else:
|
|
self._gguf_skip_value(f, vtype)
|
|
else:
|
|
self._gguf_skip_value(f, vtype)
|
|
except (struct.error, UnicodeDecodeError):
|
|
# Truncated input (e.g. HTTP byte-range header
|
|
# fetch); break so the resolver fallback runs on
|
|
# what we have.
|
|
break
|
|
|
|
# Decide diffusion routing before the SWA resolver below: it can raise on an arch transformers
|
|
# does not know, which would otherwise drop a DiffusionGemma model to plain llama-server.
|
|
self._is_diffusion = bool(
|
|
(arch and arch.lower().startswith("diffusion")) or canvas_seen
|
|
)
|
|
if self._is_diffusion:
|
|
logger.info(
|
|
f"GGUF metadata: diffusion model detected (architecture={arch}); "
|
|
"will serve via the diffusion runner"
|
|
)
|
|
|
|
# Expand a scalar period straight from the GGUF first.
|
|
if (
|
|
self._sliding_window_pattern is None
|
|
and sliding_window_pattern_period
|
|
and self._n_layers
|
|
):
|
|
self._sliding_window_pattern = [
|
|
(i + 1) % sliding_window_pattern_period != 0 for i in range(self._n_layers)
|
|
]
|
|
|
|
# Otherwise hand off to the resolver (cache / bootstrap / transformers / HF). Diffusion models
|
|
# skip it: they do not use Studio's SWA pattern and the resolver can raise for them.
|
|
if (
|
|
self._sliding_window_pattern is None
|
|
and self._sliding_window
|
|
and self._n_layers
|
|
and not self._is_diffusion
|
|
):
|
|
hf_repo_candidates = (
|
|
general.get("general.source.huggingface.repository"),
|
|
_hf_repo_from_url(general.get("general.source.url")),
|
|
_hf_repo_from_url(general.get("general.source.repo_url")),
|
|
_hf_repo_from_url(general.get("general.base_model.0.repo_url")),
|
|
(
|
|
f"{general['general.base_model.0.organization']}/"
|
|
f"{general['general.base_model.0.name']}".replace(" ", "-")
|
|
if general.get("general.base_model.0.organization")
|
|
and general.get("general.base_model.0.name")
|
|
else None
|
|
),
|
|
(
|
|
f"{general['general.organization']}/{general['general.basename']}".replace(
|
|
" ", "-"
|
|
)
|
|
if general.get("general.organization") and general.get("general.basename")
|
|
else None
|
|
),
|
|
)
|
|
self._sliding_window_pattern = _resolve_swa_pattern(
|
|
arch,
|
|
self._n_layers,
|
|
hf_repo_candidates,
|
|
)
|
|
|
|
if self._context_length:
|
|
logger.info(f"GGUF metadata: context_length={self._context_length}")
|
|
if self._chat_template:
|
|
logger.info(f"GGUF metadata: chat_template={len(self._chat_template)} chars")
|
|
# Detect thinking/reasoning support from chat template.
|
|
flags = detect_reasoning_flags(
|
|
self._chat_template,
|
|
self._model_identifier,
|
|
log_source = "GGUF metadata",
|
|
)
|
|
self._supports_reasoning = flags["supports_reasoning"]
|
|
self._reasoning_style = flags["reasoning_style"]
|
|
self._reasoning_effort_levels = flags.get("reasoning_effort_levels", [])
|
|
self._reasoning_always_on = flags["reasoning_always_on"]
|
|
self._supports_preserve_thinking = flags["supports_preserve_thinking"]
|
|
self._supports_tools = flags["supports_tools"]
|
|
except Exception as e:
|
|
logger.warning(f"Failed to read GGUF metadata: {e}")
|
|
|
|
# ── Diffusion runner (DiffusionGemma) ──
|
|
|
|
def _find_diffusion_assets(self) -> Optional[tuple[list, str, Optional[str]]]:
|
|
"""Resolve how to launch the DiffusionGemma runner: (shim argv prefix,
|
|
visual-server binary, optional extra PYTHONPATH dir for the file override).
|
|
|
|
Shim: UNSLOTH_DG_SHIM (a .py file) first, else the installed
|
|
unsloth_zoo.diffusion_studio.shim. Binary: DG_VISUAL_BIN first, else
|
|
alongside llama-server. Returns None if neither can be found.
|
|
"""
|
|
import importlib.util
|
|
|
|
# Visual-server binary: env override, else next to llama-server or in the
|
|
# install's build/bin (where the prebuilt/installer puts it). .exe on Windows.
|
|
visual_bin = os.environ.get("DG_VISUAL_BIN")
|
|
if not visual_bin:
|
|
name = "llama-diffusion-gemma-visual-server" + (".exe" if os.name == "nt" else "")
|
|
# include_denied: a transiently locked llama-server still pins the
|
|
# install dir so the adjacent visual-server can be found
|
|
base = self._find_llama_server_binary(include_denied = True)
|
|
if base:
|
|
base_dir = Path(base).parent
|
|
for cand in (
|
|
base_dir / name,
|
|
base_dir / "build" / "bin" / name,
|
|
base_dir / "build" / "bin" / "Release" / name,
|
|
):
|
|
if cand.is_file():
|
|
visual_bin = str(cand)
|
|
break
|
|
if not (visual_bin and Path(visual_bin).is_file()):
|
|
return None
|
|
|
|
# Shim: a file override (its dir goes on PYTHONPATH), else the zoo package via -m.
|
|
shim_file = os.environ.get("UNSLOTH_DG_SHIM")
|
|
if shim_file and Path(shim_file).is_file():
|
|
return ([sys.executable, shim_file], visual_bin, str(Path(shim_file).parent))
|
|
|
|
# Find the installed shim without importing the heavy unsloth_zoo package
|
|
# (find_spec on the top-level package does not run its __init__).
|
|
try:
|
|
spec = importlib.util.find_spec("unsloth_zoo")
|
|
except Exception:
|
|
spec = None
|
|
if spec is not None and spec.submodule_search_locations:
|
|
pkg_dir = Path(list(spec.submodule_search_locations)[0])
|
|
if (pkg_dir / "diffusion_studio" / "shim.py").is_file():
|
|
return (
|
|
[sys.executable, "-m", "unsloth_zoo.diffusion_studio.shim"],
|
|
visual_bin,
|
|
None,
|
|
)
|
|
|
|
return None
|
|
|
|
def _start_diffusion_server(
|
|
self,
|
|
*,
|
|
model_path: str,
|
|
gguf_path: Optional[str],
|
|
hf_repo: Optional[str],
|
|
hf_variant: Optional[str],
|
|
model_identifier: str,
|
|
n_ctx: int,
|
|
extra_args: Optional[List[str]],
|
|
) -> bool:
|
|
"""Launch the OpenAI-compat diffusion shim (which drives the on-device
|
|
visual decoder) and wait for health. Presents the same /v1 + /health
|
|
interface as llama-server, so the rest of Studio is unchanged.
|
|
"""
|
|
assets = self._find_diffusion_assets()
|
|
if assets is None:
|
|
raise RuntimeError(
|
|
"DiffusionGemma runner not found. Install unsloth_zoo (which ships "
|
|
"unsloth_zoo.diffusion_studio.shim) or set UNSLOTH_DG_SHIM to a shim "
|
|
"file, and provide the visual-server binary via DG_VISUAL_BIN or next "
|
|
"to llama-server in the install tree."
|
|
)
|
|
shim_cmd, visual_bin, extra_pythonpath = assets
|
|
self._diffusion_visual_bin = visual_bin
|
|
|
|
self._kill_process()
|
|
self._port = self._find_free_port()
|
|
# Auto-size (0): the visual server probes the largest context that fits this GPU's VRAM
|
|
# (capped at the training context). An explicit in-range n_ctx overrides it.
|
|
maxtok = n_ctx if (n_ctx and 0 < n_ctx <= 65536) else 0
|
|
gpu = os.environ.get("DG_GPU", "0")
|
|
|
|
cmd = list(shim_cmd) + [
|
|
"--gguf",
|
|
model_path,
|
|
"--host",
|
|
"127.0.0.1",
|
|
"--port",
|
|
str(self._port),
|
|
"--gpu",
|
|
gpu,
|
|
"--maxtok",
|
|
str(maxtok),
|
|
]
|
|
|
|
env = child_env_without_native_path_secret()
|
|
# `python -m unsloth_zoo.diffusion_studio.shim` imports unsloth_zoo, which
|
|
# refuses to load unless UNSLOTH_IS_PRESENT is set (normally by `import
|
|
# unsloth`). The shim never imports unsloth, so set it here as unsloth does.
|
|
env["UNSLOTH_IS_PRESENT"] = "1"
|
|
env["DG_VISUAL_BIN"] = visual_bin
|
|
env["DG_GPU"] = gpu
|
|
# The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH.
|
|
# (The zoo-package shim is an installed module and needs no PYTHONPATH change.)
|
|
if extra_pythonpath:
|
|
existing = env.get("PYTHONPATH")
|
|
env["PYTHONPATH"] = (
|
|
(extra_pythonpath + os.pathsep + existing) if existing else extra_pythonpath
|
|
)
|
|
|
|
logger.info(f"Starting DiffusionGemma runner: {' '.join(cmd)}")
|
|
self._stdout_lines = []
|
|
self._llama_log_fh = None
|
|
self._llama_log_path = None
|
|
try:
|
|
log_dir = _swa_cache_path().parent / "logs" / "diffusion-server"
|
|
log_dir.mkdir(parents = True, exist_ok = True)
|
|
self._llama_log_path = log_dir / f"diffusion-{int(time.time())}-port-{self._port}.log"
|
|
self._llama_log_fh = open(self._llama_log_path, "w", encoding = "utf-8", buffering = 1)
|
|
logger.info(f"diffusion runner stdout/stderr -> {self._llama_log_path}")
|
|
except OSError as e:
|
|
logger.debug(f"Could not open diffusion runner log file: {e}")
|
|
|
|
# The shim (and its visual server) die with this backend process, so a
|
|
# Studio crash/restart never orphans a GPU process.
|
|
self._process = subprocess.Popen(
|
|
cmd,
|
|
stdout = subprocess.PIPE,
|
|
stderr = subprocess.STDOUT,
|
|
text = True,
|
|
env = env,
|
|
**_windows_hidden_subprocess_kwargs(),
|
|
**_child_popen_kwargs(),
|
|
)
|
|
self._stdout_thread = threading.Thread(
|
|
target = self._drain_stdout, daemon = True, name = "diffusion-stdout"
|
|
)
|
|
self._stdout_thread.start()
|
|
|
|
# Publish state before the health wait (mirrors the llama-server path).
|
|
self._gguf_path = model_path
|
|
self._hf_repo = hf_repo
|
|
self._is_vision = False
|
|
self._is_audio = False # clear any prior TTS/audio model's routing flag
|
|
self._model_identifier = model_identifier
|
|
self._cache_type_kv = None
|
|
self._gpu_offload_active = True
|
|
if hf_variant:
|
|
self._hf_variant = hf_variant
|
|
elif gguf_path:
|
|
try:
|
|
from utils.models.model_config import _extract_quant_label
|
|
self._hf_variant = _extract_quant_label(gguf_path)
|
|
except Exception:
|
|
self._hf_variant = None
|
|
else:
|
|
self._hf_variant = None
|
|
# Provisional until the server reports the budget it resolved (auto-size picks it from VRAM).
|
|
self._effective_context_length = maxtok or self._context_length
|
|
self._max_context_length = self._context_length or maxtok or None
|
|
|
|
healthy = self._wait_for_health(timeout = 600.0)
|
|
if healthy:
|
|
self._healthy = True
|
|
self._gpu_offload_active = True
|
|
if extra_args is not None:
|
|
self._extra_args = list(extra_args)
|
|
self._extra_args_source = (model_identifier, hf_variant)
|
|
# The visual server logs "MAXTOK=<N>" with the context budget it actually resolved
|
|
# (auto-sized to VRAM). Read it back so the UI context bar shows the real budget.
|
|
chosen = maxtok
|
|
try:
|
|
for _ln in reversed(self._stdout_lines):
|
|
_m = re.search(r"MAXTOK=(\d+)", _ln)
|
|
if _m:
|
|
chosen = int(_m.group(1))
|
|
break
|
|
except Exception:
|
|
pass
|
|
if chosen and chosen > 0:
|
|
self._effective_context_length = chosen
|
|
self._max_context_length = chosen
|
|
self._requested_n_ctx = int(n_ctx)
|
|
else:
|
|
self._healthy = False
|
|
logger.error("DiffusionGemma runner failed to become healthy")
|
|
return healthy
|
|
|
|
# ── HF download (no lock held) ───────────────────────────────
|
|
|
|
def _download_gguf(
|
|
self,
|
|
*,
|
|
hf_repo: str,
|
|
hf_variant: Optional[str] = None,
|
|
hf_token: Optional[str] = None,
|
|
force: bool = False,
|
|
allow_smaller_fallback: bool = True,
|
|
cancel_event: Optional[threading.Event] = None,
|
|
) -> str:
|
|
"""Download GGUF file(s) from HuggingFace. Returns local path.
|
|
|
|
Runs WITHOUT self._lock so unload_model() can set _cancel_event at
|
|
any time; checks it between each shard download.
|
|
|
|
``force`` re-fetches even when a (possibly stale) blob is cached.
|
|
``allow_smaller_fallback=False`` raises on low disk instead of silently
|
|
switching to a smaller quant. ``cancel_event`` overrides
|
|
``self._cancel_event`` so an update can use a private event without
|
|
touching the shared one; defaults to the shared event.
|
|
"""
|
|
cancel_event = cancel_event if cancel_event is not None else self._cancel_event
|
|
try:
|
|
import huggingface_hub # noqa: F401 -- presence check only
|
|
except ImportError:
|
|
raise RuntimeError(
|
|
"huggingface_hub is required for HF model loading. "
|
|
"Install it with: pip install huggingface_hub"
|
|
)
|
|
|
|
# Resolve the filename from the variant
|
|
gguf_filename = None
|
|
gguf_extra_shards: list[str] = []
|
|
if hf_variant:
|
|
try:
|
|
from huggingface_hub import list_repo_files
|
|
|
|
files = list_repo_files(hf_repo, token = hf_token)
|
|
gguf_files = _gguf_files_for_variant(files, hf_variant)
|
|
if gguf_files:
|
|
gguf_filename = gguf_files[0]
|
|
gguf_extra_shards = _gguf_extra_shards(gguf_files, gguf_filename)
|
|
except Exception as e:
|
|
logger.warning(f"Could not list repo files: {e}")
|
|
|
|
# Offline: resolve variant -> filename from the local HF cache.
|
|
# The heuristic below assumes filenames echo the repo name, which
|
|
# breaks for e.g. Qwen3.6-27B-MTP-GGUF (no "MTP" in file). Match
|
|
# against the rel path (not just basename) so subdir layouts like
|
|
# ``BF16/foo.gguf`` are findable.
|
|
if not gguf_filename:
|
|
try:
|
|
from utils.models.model_config import _iter_hf_cache_snapshots
|
|
for snap in _iter_hf_cache_snapshots(hf_repo):
|
|
cached_files = _gguf_snapshot_files(snap)
|
|
matches = _gguf_files_for_variant(cached_files, hf_variant)
|
|
if not matches:
|
|
continue
|
|
gguf_filename = matches[0]
|
|
gguf_extra_shards = _gguf_extra_shards(matches, gguf_filename)
|
|
logger.info(
|
|
"Resolved variant %s -> %s from local HF cache",
|
|
hf_variant,
|
|
gguf_filename,
|
|
)
|
|
break
|
|
except Exception as e:
|
|
logger.debug(f"Offline cache lookup for variant failed: {e}")
|
|
|
|
if not gguf_filename:
|
|
repo_name = hf_repo.split("/")[-1].replace("-GGUF", "")
|
|
gguf_filename = f"{repo_name}-{hf_variant}.gguf"
|
|
|
|
# Check disk space; fall back to a smaller variant if needed
|
|
all_gguf_files = [gguf_filename] + gguf_extra_shards
|
|
try:
|
|
from huggingface_hub import get_paths_info, try_to_load_from_cache
|
|
|
|
path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token))
|
|
total_bytes = sum((p.size or 0) for p in path_infos)
|
|
|
|
# Subtract bytes already in the HF cache so we only preflight
|
|
# against what we must download. Without this, re-loading a
|
|
# cached large model (e.g. MiniMax-M2.7-GGUF at 131 GB) fails
|
|
# cold whenever free disk is below the full weight footprint,
|
|
# even though nothing needs downloading.
|
|
already_cached_bytes = 0
|
|
if not force:
|
|
for p in path_infos:
|
|
if not p.size:
|
|
continue
|
|
try:
|
|
cached_path = try_to_load_from_cache(hf_repo, p.path)
|
|
except Exception:
|
|
cached_path = None
|
|
if isinstance(cached_path, str) and os.path.exists(cached_path):
|
|
try:
|
|
on_disk = os.path.getsize(cached_path)
|
|
except OSError:
|
|
on_disk = 0
|
|
# Satisfied only when the full blob is present.
|
|
if on_disk >= p.size:
|
|
already_cached_bytes += p.size
|
|
|
|
total_download_bytes = max(0, total_bytes - already_cached_bytes)
|
|
|
|
if total_download_bytes > 0:
|
|
cache_dir = os.environ.get(
|
|
"HF_HUB_CACHE",
|
|
str(Path.home() / ".cache" / "huggingface" / "hub"),
|
|
)
|
|
Path(cache_dir).mkdir(parents = True, exist_ok = True)
|
|
free_bytes = shutil.disk_usage(cache_dir).free
|
|
|
|
total_gb = total_download_bytes / (1024**3)
|
|
free_gb = free_bytes / (1024**3)
|
|
cached_gb = already_cached_bytes / (1024**3)
|
|
|
|
logger.info(
|
|
f"GGUF download: {total_gb:.1f} GB needed "
|
|
f"({cached_gb:.1f} GB already cached), "
|
|
f"{free_gb:.1f} GB free on disk"
|
|
)
|
|
|
|
if total_download_bytes > free_bytes:
|
|
if not allow_smaller_fallback:
|
|
# Update path: never silently switch to a smaller quant;
|
|
# surface the disk shortfall for the requested variant.
|
|
raise RuntimeError(
|
|
f"Not enough disk space to download {gguf_filename}. "
|
|
f"Only {free_gb:.1f} GB free in {cache_dir}"
|
|
)
|
|
smaller = self._find_smallest_fitting_variant(
|
|
hf_repo,
|
|
free_bytes,
|
|
hf_token,
|
|
)
|
|
if smaller:
|
|
fallback_file, fallback_size = smaller
|
|
logger.info(
|
|
f"Selected variant too large ({total_gb:.1f} GB), "
|
|
f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)"
|
|
)
|
|
gguf_filename = fallback_file
|
|
_m = _SHARD_RE.match(gguf_filename)
|
|
_prefix = _m.group(1) if _m else None
|
|
if _prefix:
|
|
prefix_lower = _prefix.lower()
|
|
gguf_extra_shards = sorted(
|
|
f
|
|
for f in all_gguf_files
|
|
if f.lower().startswith(prefix_lower)
|
|
and f != gguf_filename
|
|
and not _is_companion_gguf_path(f)
|
|
)
|
|
else:
|
|
gguf_extra_shards = []
|
|
else:
|
|
raise RuntimeError(
|
|
f"Not enough disk space to download any variant. "
|
|
f"Only {free_gb:.1f} GB free in {cache_dir}"
|
|
)
|
|
except RuntimeError:
|
|
raise
|
|
except Exception as e:
|
|
logger.warning(f"Could not check disk space: {e}")
|
|
|
|
gguf_label = f"{hf_repo}/{gguf_filename}" + (
|
|
f" (+{len(gguf_extra_shards)} shards)" if gguf_extra_shards else ""
|
|
)
|
|
logger.info(f"Resolving GGUF: {gguf_label}")
|
|
try:
|
|
if cancel_event.is_set():
|
|
raise RuntimeError("Cancelled")
|
|
dl_start = time.monotonic()
|
|
# Xet primary, HTTP fallback on stall; per-file so finished shards stay cached.
|
|
local_path = hf_hub_download_with_xet_fallback(
|
|
hf_repo,
|
|
gguf_filename,
|
|
hf_token,
|
|
cancel_event = cancel_event,
|
|
on_status = lambda m: logger.info(m),
|
|
force_download = force,
|
|
)
|
|
for shard in gguf_extra_shards:
|
|
if cancel_event.is_set():
|
|
raise RuntimeError("Cancelled")
|
|
logger.info(f"Resolving GGUF shard: {shard}")
|
|
hf_hub_download_with_xet_fallback(
|
|
hf_repo,
|
|
shard,
|
|
hf_token,
|
|
cancel_event = cancel_event,
|
|
force_download = force,
|
|
)
|
|
except Exception as e:
|
|
if isinstance(e, RuntimeError) and "Cancelled" in str(e):
|
|
raise
|
|
raise RuntimeError(
|
|
f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}"
|
|
)
|
|
|
|
dl_elapsed = time.monotonic() - dl_start
|
|
if dl_elapsed < 2.0:
|
|
logger.info(f"GGUF resolved from cache: {local_path}")
|
|
else:
|
|
logger.info(f"GGUF downloaded in {dl_elapsed:.1f}s: {local_path}")
|
|
return local_path
|
|
|
|
def _download_companion_gguf(
|
|
self,
|
|
*,
|
|
hf_repo: str,
|
|
hf_token: Optional[str],
|
|
pick: Callable[[list[str]], Optional[str]],
|
|
label: str,
|
|
cancel_event: Optional[threading.Event] = None,
|
|
) -> Optional[str]:
|
|
"""Resolve and fetch a companion GGUF (mmproj / MTP drafter) by name.
|
|
|
|
Tries the live repo file list, then the local HF cache snapshots
|
|
(offline, same fallback as _download_gguf), then hf_hub_download.
|
|
Runs WITHOUT self._lock (like _download_gguf); honors _cancel_event so
|
|
an /unload between the main download and here skips the fetch.
|
|
``cancel_event`` overrides ``self._cancel_event`` (defaults to it).
|
|
"""
|
|
cancel_event = cancel_event if cancel_event is not None else self._cancel_event
|
|
if cancel_event.is_set():
|
|
return None
|
|
|
|
target: Optional[str] = None
|
|
from huggingface_hub import list_repo_files
|
|
|
|
# Retry a transient listing blip; permanent repo/auth errors and offline
|
|
# mode are not retried (offline raises at once -> fall through to cache).
|
|
for attempt in range(3):
|
|
if cancel_event.is_set():
|
|
return None
|
|
try:
|
|
target = pick(list_repo_files(hf_repo, token = hf_token))
|
|
break
|
|
except Exception as e:
|
|
if type(e).__name__ in (
|
|
"RepositoryNotFoundError",
|
|
"GatedRepoError",
|
|
"RevisionNotFoundError",
|
|
"EntryNotFoundError",
|
|
"OfflineModeIsEnabled",
|
|
):
|
|
logger.debug(f"Could not list repo files for {label}: {e}")
|
|
break
|
|
logger.debug(
|
|
f"Could not list repo files for {label} (attempt {attempt + 1}/3): {e}"
|
|
)
|
|
if attempt < 2:
|
|
cancel_event.wait(2**attempt)
|
|
|
|
if target is None:
|
|
try:
|
|
from utils.models.model_config import _iter_hf_cache_snapshots
|
|
for snap in _iter_hf_cache_snapshots(hf_repo):
|
|
rel_files = _gguf_snapshot_files(snap)
|
|
target = pick(rel_files)
|
|
if target is not None:
|
|
logger.info("Resolved %s %s from local HF cache", label, target)
|
|
break
|
|
except Exception as e:
|
|
logger.debug(f"Offline cache lookup for {label} failed: {e}")
|
|
|
|
if target is None or cancel_event.is_set():
|
|
return None
|
|
|
|
try:
|
|
logger.info(f"Downloading {label}: {hf_repo}/{target}")
|
|
# Same policy; companions are best-effort (caller below swallows failures to None).
|
|
return hf_hub_download_with_xet_fallback(
|
|
hf_repo,
|
|
target,
|
|
hf_token,
|
|
cancel_event = cancel_event,
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"Could not download {label}: {e}")
|
|
return None
|
|
|
|
def _download_mmproj(
|
|
self,
|
|
*,
|
|
hf_repo: str,
|
|
hf_token: Optional[str] = None,
|
|
cancel_event: Optional[threading.Event] = None,
|
|
) -> Optional[str]:
|
|
"""Download the mmproj (vision projection) file from a GGUF repo.
|
|
|
|
Prefers mmproj-F16.gguf, else any mmproj*.gguf. Returns the local
|
|
path, or None if none exists. ``cancel_event`` overrides
|
|
``self._cancel_event`` (defaults to it).
|
|
"""
|
|
|
|
def _pick_mmproj(candidates: list[str]) -> Optional[str]:
|
|
mmproj_files = sorted(
|
|
f
|
|
for f in candidates
|
|
if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower()
|
|
)
|
|
if not mmproj_files:
|
|
return None
|
|
for f in mmproj_files:
|
|
if f.lower().endswith("-f16.gguf"):
|
|
return f
|
|
return mmproj_files[0]
|
|
|
|
return self._download_companion_gguf(
|
|
hf_repo = hf_repo,
|
|
hf_token = hf_token,
|
|
pick = _pick_mmproj,
|
|
label = "mmproj",
|
|
cancel_event = cancel_event,
|
|
)
|
|
|
|
def _download_mtp(
|
|
self,
|
|
*,
|
|
hf_repo: str,
|
|
hf_token: Optional[str] = None,
|
|
) -> Optional[str]:
|
|
"""Download the separate MTP drafter (speculative head) from a GGUF repo.
|
|
|
|
Targets the repo-root ``mtp-*.gguf`` companion -- the Q8_0 drafter
|
|
unsloth mirrors there for llama.cpp ``-hf`` auto-discovery (smallest,
|
|
recommended for speculation). Repos that bake the MTP head into the
|
|
main GGUF (e.g. Qwen) ship no such sibling and this returns None. The
|
|
higher-precision copies under ``MTP/`` are for explicit selection and
|
|
are intentionally skipped. Returns the local path, or None.
|
|
"""
|
|
|
|
def _pick_mtp(candidates: list[str]) -> Optional[str]:
|
|
mtp_files = sorted(
|
|
f
|
|
for f in candidates
|
|
if f.lower().endswith(".gguf") and Path(f).name.lower().startswith("mtp-")
|
|
)
|
|
return mtp_files[0] if mtp_files else None
|
|
|
|
return self._download_companion_gguf(
|
|
hf_repo = hf_repo,
|
|
hf_token = hf_token,
|
|
pick = _pick_mtp,
|
|
label = "MTP drafter",
|
|
)
|
|
|
|
def _resolve_launch_mmproj_path(
|
|
self, *, model_path: str, mmproj_path: Optional[str]
|
|
) -> Optional[str]:
|
|
"""Return mmproj_path iff it exists on disk AND matches the model family.
|
|
|
|
None if mmproj_path is None, missing, or family-mismatched.
|
|
"""
|
|
if not mmproj_path:
|
|
return None
|
|
|
|
mmproj = Path(mmproj_path)
|
|
if not mmproj.is_file():
|
|
logger.warning(f"mmproj file not found: {mmproj_path}")
|
|
return None
|
|
|
|
from utils.models.model_config import mmproj_matches_model_family
|
|
|
|
if not mmproj_matches_model_family(model_path, str(mmproj)):
|
|
logger.warning(
|
|
f"mmproj does not match model family: model={Path(model_path).name} "
|
|
f"mmproj={mmproj.name}"
|
|
)
|
|
return None
|
|
|
|
return str(mmproj)
|
|
|
|
def _mmproj_vram_bytes(self, launch_mmproj_path: Optional[str]) -> int:
|
|
"""Return resolved mmproj VRAM bytes, or 0 when absent/unreadable."""
|
|
if not launch_mmproj_path:
|
|
return 0
|
|
try:
|
|
return self._get_gguf_size_bytes(launch_mmproj_path)
|
|
except OSError as e:
|
|
logger.debug(f"Could not size mmproj {launch_mmproj_path}: {e}")
|
|
return 0
|
|
|
|
def _resolve_launch_mtp_path(self, *, mtp_draft_path: Optional[str]) -> Optional[str]:
|
|
"""Return mtp_draft_path iff it exists on disk, else None.
|
|
|
|
No family check needed: the drafter is only ever auto-resolved from
|
|
the same repo as the main GGUF (see _download_mtp).
|
|
"""
|
|
if not mtp_draft_path:
|
|
return None
|
|
if not Path(mtp_draft_path).is_file():
|
|
logger.warning(f"MTP drafter file not found: {mtp_draft_path}")
|
|
return None
|
|
return str(mtp_draft_path)
|
|
|
|
# ── Lifecycle ─────────────────────────────────────────────────
|
|
|
|
# GGUF ``general.architecture`` values for diffusion / image models.
|
|
# llama.cpp has no such architectures, so loading one as a chat model dies
|
|
# with "unknown model architecture: '<arch>'". These match the patched
|
|
# stable-diffusion.cpp / ComfyUI-GGUF enums. Unsloth publishes FLUX and
|
|
# Qwen-Image GGUFs under
|
|
# https://huggingface.co/collections/unsloth/unsloth-diffusion-ggufs.
|
|
# Matched exactly (not a substring) so a chat arch containing "wan"/"sd1"
|
|
# (e.g. "taiwan") isn't misrouted to Images.
|
|
_DIFFUSION_ARCHES = frozenset(
|
|
(
|
|
"qwen_image",
|
|
"flux",
|
|
"sd1",
|
|
"sdxl",
|
|
"sd3",
|
|
"aura",
|
|
"hidream",
|
|
"cosmos",
|
|
"ltxv",
|
|
"hyvid",
|
|
"wan",
|
|
"lumina2",
|
|
)
|
|
)
|
|
|
|
@staticmethod
|
|
def _classify_llama_start_failure(
|
|
output: str,
|
|
gguf_path: Optional[str],
|
|
model_identifier: Optional[str],
|
|
returncode: Optional[int] = None,
|
|
) -> str:
|
|
"""Explain *why* llama-server failed to start, from its output.
|
|
|
|
Several distinct failures otherwise collapse into the same opaque
|
|
"invalid GGUF or out of memory" message. Worst case: a diffusion GGUF
|
|
loaded as a chat model -- valid file, plenty of memory, but llama.cpp
|
|
has no such architecture, so the user is told to free memory that was
|
|
never the problem (#5842). Pick the most specific message we can.
|
|
"""
|
|
lowered = (output or "").lower()
|
|
|
|
# Tensor parallelism (--split-mode tensor) is arch-gated in llama.cpp;
|
|
# unsupported architectures abort the load with this marker. Point the
|
|
# user at the toggle instead of a generic invalid-GGUF/OOM message.
|
|
if "split_mode_tensor not implemented" in lowered:
|
|
return (
|
|
"Tensor parallelism is not supported for this model's "
|
|
"architecture. Turn off Tensor Parallelism in the model "
|
|
"settings and reload."
|
|
)
|
|
|
|
# Detect Ollama source up front so the arch branch can keep the
|
|
# Ollama hint instead of the generic "unsupported arch" message.
|
|
gguf = gguf_path or ""
|
|
is_ollama = (
|
|
".studio_links" in gguf
|
|
or os.sep + "ollama_links" + os.sep in gguf
|
|
or os.sep + ".cache" + os.sep + "ollama" + os.sep in gguf
|
|
or (model_identifier or "").startswith("ollama/")
|
|
)
|
|
|
|
# "unknown model architecture: '<arch>'": diffusion -> Images page,
|
|
# Ollama -> Ollama hint, else a precise "unsupported" message. Exact
|
|
# match so chat archs aren't misrouted.
|
|
arch_match = re.search(r"unknown model architecture:\s*'([^']+)'", lowered)
|
|
if arch_match:
|
|
arch = arch_match.group(1)
|
|
if arch in LlamaCppBackend._DIFFUSION_ARCHES:
|
|
return (
|
|
f"'{arch}' is a diffusion (image-generation) GGUF, which "
|
|
"llama-server cannot run as a chat/completion model. Use "
|
|
"Studio's Images page to generate with local diffusion "
|
|
"GGUFs such as FLUX and Qwen-Image."
|
|
)
|
|
if is_ollama:
|
|
return (
|
|
"Some Ollama models do not work with llama.cpp. Try a "
|
|
"different model, or use this model directly through "
|
|
"Ollama instead."
|
|
)
|
|
return (
|
|
f"llama.cpp does not support this GGUF's model architecture "
|
|
f"('{arch}'). The file is valid, but this model type cannot "
|
|
"be run with llama-server."
|
|
)
|
|
|
|
# Other Ollama compat failures that don't name an arch. Only when
|
|
# the output shows a GGUF compat issue, not OOM / missing binaries.
|
|
if is_ollama:
|
|
gguf_compat_hints = (
|
|
"key not found",
|
|
"unknown model architecture",
|
|
"failed to load model",
|
|
)
|
|
if any(h in lowered for h in gguf_compat_hints):
|
|
return (
|
|
"Some Ollama models do not work with llama.cpp. Try a "
|
|
"different model, or use this model directly through "
|
|
"Ollama instead."
|
|
)
|
|
|
|
# SIGKILL with no diagnostic output is the OOM killer (e.g. a model too
|
|
# large for the WSL VM's RAM cap); name it actionably.
|
|
if returncode == -9:
|
|
return (
|
|
"llama-server was stopped by the operating system (signal 9), "
|
|
"most likely out of memory. Try a smaller or more quantized "
|
|
"GGUF, lower the context length, or free memory (on WSL, raise "
|
|
"the memory limit in .wslconfig)."
|
|
)
|
|
# SIGTERM is also how an unload/cancel or a supervisor stops the server,
|
|
# so report it neutrally rather than blaming memory.
|
|
if returncode == -15:
|
|
return (
|
|
"llama-server was terminated (signal 15) before it became "
|
|
"healthy. If you cancelled or unloaded the model this is "
|
|
"expected; otherwise check the llama-server log for the cause."
|
|
)
|
|
|
|
# A live server that never answered 200 on /health is not a bad GGUF:
|
|
# the load is too large for VRAM/context, or a local proxy/VPN grabbed
|
|
# the loopback probe (#5740).
|
|
if "health check timed out" in lowered:
|
|
return (
|
|
"llama-server started but never became healthy on its local "
|
|
"/health endpoint. Try a smaller context length or a more "
|
|
"quantized GGUF, and if you use a VPN or HTTP proxy make sure "
|
|
"localhost bypasses it (NO_PROXY=127.0.0.1,localhost)."
|
|
)
|
|
|
|
# Fallback: genuinely unknown failure (OOM, missing binary ...).
|
|
return (
|
|
"llama-server failed to start. "
|
|
"Check that the GGUF file is valid and you have enough memory."
|
|
)
|
|
|
|
def _plan_tensor_parallel(
|
|
self,
|
|
gpus: list[tuple[int, int]],
|
|
model_size: int,
|
|
target_ctx: int,
|
|
cache_type_kv: Optional[str] = None,
|
|
n_parallel: int = 1,
|
|
mtp_engaged: bool = False,
|
|
mtp_overhead_fn: Optional[Callable[[int], int]] = None,
|
|
mtp_flat_reserve_bytes: int = 0,
|
|
max_target_ctx: Optional[int] = None,
|
|
total_by_idx: Optional[dict[int, int]] = None,
|
|
n_ubatch: Optional[int] = None,
|
|
soft_overhead_bytes: int = 0,
|
|
) -> tuple[int, int, list[int], Optional[list[int]]]:
|
|
"""Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed.
|
|
|
|
``gpus`` is a list of ``(gpu_index, free_mib)``; ``model_size`` is the
|
|
weight size in bytes; ``target_ctx`` is the context to fit (the explicit
|
|
request, or the model's native length for auto). ``max_target_ctx`` is
|
|
the native/hardware ceiling used only for the UI bound (defaults to
|
|
``target_ctx``). Returns
|
|
``(effective_ctx, max_available_ctx, gpu_indices, tensor_split)``.
|
|
|
|
Policy (assumes >= 2 GPUs; the caller drops the toggle below that):
|
|
- Cap context to the KV that fits the pooled VRAM after the weights, one
|
|
per-device flat compute-graph buffer (``_estimate_compute_buffer_bytes``,
|
|
deterministic from dims; flat fallback when dims are unavailable), and the
|
|
per-device context-linear compute growth (``_compute_buffer_ctx_bytes``,
|
|
replicated on every device in tensor mode, so summed over the split).
|
|
llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only
|
|
cap, honored even for an explicit ``-c``. It is more accurate than the
|
|
0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused.
|
|
- ``tensor_split`` is None (llama.cpp's even default, safe for every arch
|
|
incl. Gemma 3n which GGML_ASSERTs on a weighted split) when an even
|
|
share fits the smallest GPU; otherwise it is weighted by usable budget
|
|
so the roomier GPU absorbs more weight and the smallest keeps room for KV.
|
|
``total_by_idx`` enables the total-based occupancy cap; ``n_ubatch`` sizes
|
|
the compute buffer. ``soft_overhead_bytes`` is the CUDA-context / mmproj /
|
|
MTP-draft-graph reserve the layer path folds into ``model_size_fit``;
|
|
charged against the pooled budget so tensor mode reserves the same overhead.
|
|
"""
|
|
|
|
# Per-GPU usable budget: free - (1-frac)*total, else (unknown total, e.g. a
|
|
# two-column probe) the legacy free*frac. Mirrors _select_gpus and
|
|
# _gpu_usable so the 5% cushion is kept on every path, not dropped here.
|
|
def _usable(idx: int, free_mib: int) -> float:
|
|
t = total_by_idx.get(idx, 0) if total_by_idx else 0
|
|
if t > 0:
|
|
return max(0.0, free_mib - (1.0 - _CTX_FIT_VRAM_FRACTION) * t)
|
|
return max(0.0, free_mib * _CTX_FIT_VRAM_FRACTION)
|
|
|
|
# Drop GPUs whose usable budget can't hold the per-device compute-graph
|
|
# buffer; they'd OOM in tensor mode. Admitting on raw free would let a
|
|
# partly-used big card in with no budget left. Defense-in-depth (load_model
|
|
# gates too). Derived per-device reserve; flat fallback.
|
|
_reserve_bytes = self._estimate_compute_buffer_bytes(
|
|
n_ubatch = n_ubatch, n_parallel = n_parallel, per_device_tensor = True
|
|
)
|
|
reserve_mib = (
|
|
_reserve_bytes // (1024 * 1024)
|
|
if _reserve_bytes > 0
|
|
else self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
|
|
)
|
|
usable_gpus = [g for g in gpus if _usable(g[0], g[1]) >= reserve_mib]
|
|
gpu_indices = sorted(idx for idx, _ in usable_gpus)
|
|
if len(gpu_indices) < 2:
|
|
# Tensor parallelism is meaningless on <2 GPUs (the caller drops the
|
|
# toggle before this); be defensive and never emit a split here.
|
|
return (
|
|
target_ctx if target_ctx > 0 else 4096,
|
|
target_ctx if target_ctx > 0 else 4096,
|
|
gpu_indices,
|
|
None,
|
|
)
|
|
free_by_idx = {idx: free for idx, free in usable_gpus}
|
|
usable_by_idx = {idx: _usable(idx, free_by_idx[idx]) for idx in gpu_indices}
|
|
pool_mib = sum(usable_by_idx.values())
|
|
# MTP reserve: byte-accurate per-ctx inside _fit_ctx (mtp_overhead_fn) plus
|
|
# a flat cushion that the byte fn can't size -- 2 GiB when dims are wholly
|
|
# unavailable (no fn), or mtp_flat_reserve_bytes when the fn is weights-only
|
|
# because the draft KV couldn't be sized (_mtp_kv_unsized). Without this the
|
|
# binary search spends the unsized-KV cushion on main context and OOMs.
|
|
flat_mtp_bytes = max(0, mtp_flat_reserve_bytes)
|
|
if mtp_engaged and mtp_overhead_fn is None:
|
|
flat_mtp_bytes = max(flat_mtp_bytes, 2 * 1024**3)
|
|
# soft_overhead_bytes is the CUDA-context / mmproj / MTP-draft-graph reserve
|
|
# the layer path folds into model_size_fit. Tensor mode has no --fit valve, so
|
|
# an unreserved overshoot OOMs at startup rather than offloading; charge it here
|
|
# too. Once (pooled), mirroring the layer path -- the per-device CUDA context is
|
|
# a known slight under-charge, left for real multi-GPU data.
|
|
kv_budget_b = (
|
|
(pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024
|
|
- model_size
|
|
- flat_mtp_bytes
|
|
- max(0, soft_overhead_bytes)
|
|
)
|
|
|
|
def _mtp_at(ctx: int) -> int:
|
|
return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0
|
|
|
|
# Context-linear compute buffer, summed over the split. Tensor mode
|
|
# replicates the compute graph on EVERY device (measured: the per-device
|
|
# buffer grows a flat n_ubatch*2 bytes/token, ~1024 B/tok on Qwen3.5-9B at
|
|
# f16, independent of n_embd), so the growth is n_dev x the per-device
|
|
# term. cache_type_kv here is always non-quantized (tensor forces f16), so
|
|
# _compute_buffer_ctx_bytes returns the light KQ-mask term, not the heavy
|
|
# quantized dequant scratch. The flat reserve_mib above only covers ctx->0;
|
|
# without this the fit over-pins and OOMs at high context on a tight pool
|
|
# (0.5-4 GiB unreserved at 262k-1M across 2-4 GPUs), the tensor-mode analog
|
|
# of the layer-split compute bug.
|
|
n_dev = len(gpu_indices)
|
|
|
|
def _cc_ctx(ctx: int) -> int:
|
|
return n_dev * self._compute_buffer_ctx_bytes(ctx, n_ubatch, cache_type_kv)
|
|
|
|
def _fit_ctx(ctx: int) -> int:
|
|
# Largest context whose KV (+ MTP draft reserve + context-linear
|
|
# compute) fits the pooled budget. Floors small, but never raises an
|
|
# explicit ctx above asked.
|
|
if self._can_estimate_kv() and ctx > 0:
|
|
ctx_floor = min(2048, ctx)
|
|
if kv_budget_b <= 0:
|
|
# Weights + buffers exceed the pool -> floor; the load then
|
|
# falls back to layer split.
|
|
return ctx_floor
|
|
if mtp_overhead_fn is not None:
|
|
# kv(ctx)+mtp(ctx)+compute(ctx) is not single-linear, so binary search.
|
|
def _consumer(c: int) -> int:
|
|
return (
|
|
self._estimate_kv_cache_bytes(c, cache_type_kv, n_parallel = n_parallel)
|
|
+ _mtp_at(c)
|
|
+ _cc_ctx(c)
|
|
)
|
|
|
|
if _consumer(ctx) <= kv_budget_b:
|
|
return ctx
|
|
lo, hi, best = ctx_floor, ctx, ctx_floor
|
|
while lo <= hi:
|
|
mid = (lo + hi) // 2
|
|
if _consumer(mid) <= kv_budget_b:
|
|
best = mid
|
|
lo = mid + 1
|
|
else:
|
|
hi = mid - 1
|
|
return best
|
|
kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel)
|
|
total_at = kv_at + _cc_ctx(ctx) # both ~linear through the origin
|
|
if total_at <= kv_budget_b:
|
|
return ctx
|
|
return max(ctx_floor, int(ctx * kv_budget_b / total_at))
|
|
# KV size unknown -> can't prove a safe cap; floor.
|
|
return min(4096, ctx) if ctx > 0 else 4096
|
|
|
|
# max_available_ctx is the hardware ceiling for the UI bound, sized from
|
|
# the native context independent of an explicit small -c (which only
|
|
# caps effective_ctx).
|
|
max_ctx_target = max_target_ctx if (max_target_ctx and max_target_ctx > 0) else target_ctx
|
|
max_available_ctx = _fit_ctx(max_ctx_target)
|
|
effective_ctx = min(_fit_ctx(target_ctx), max_available_ctx)
|
|
|
|
min_usable_mib = min(usable_by_idx.values())
|
|
kv_bytes = (
|
|
self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = n_parallel)
|
|
if (self._can_estimate_kv() and effective_ctx > 0)
|
|
else 0
|
|
)
|
|
# The MTP reserve also has to fit the even split (mirror the pooled budget):
|
|
# byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above.
|
|
mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes
|
|
# Context-linear compute is replicated per device; charge the whole split so
|
|
# the weighted ratio reflects it (mirrors kv_budget_b's per-device reserve).
|
|
cc_bytes = _cc_ctx(effective_ctx) if effective_ctx > 0 else 0
|
|
even_share_mib = (
|
|
(model_size + kv_bytes + mtp_bytes + cc_bytes) / len(gpu_indices) / (1024 * 1024)
|
|
)
|
|
tensor_split: Optional[list[int]] = None
|
|
if even_share_mib > (min_usable_mib - reserve_mib):
|
|
# Each device also holds its replicated share of the context-linear
|
|
# compute (cc_bytes/n_dev) on top of the flat reserve. The even-share
|
|
# gate above charges cc_bytes; the split weights must subtract it too, or
|
|
# the smaller card is weighted above its real usable budget and OOMs (the
|
|
# per-device analog of the layer path's per-GPU overhead in _select_gpus).
|
|
cc_per_dev_mib = (cc_bytes // len(gpu_indices)) // (1024 * 1024) if cc_bytes else 0
|
|
adj = [
|
|
max(0, int(usable_by_idx[i] - reserve_mib - cc_per_dev_mib)) for i in gpu_indices
|
|
]
|
|
if sum(adj) > 0:
|
|
tensor_split = adj
|
|
return effective_ctx, max_available_ctx, gpu_indices, tensor_split
|
|
|
|
@staticmethod
|
|
def _is_projector_incompatibility(output: str) -> bool:
|
|
"""True when llama-server aborted because it cannot load the model's
|
|
vision/audio projector (mmproj), typically an installed llama.cpp
|
|
that predates the projector format. Conservative: only matches
|
|
projector-format errors so unrelated failures (OOM, bad GGUF, port
|
|
bind, ...) keep their own handling, and a bare 'clip'/'mmproj'
|
|
mention in a normal startup log does not match.
|
|
"""
|
|
text = (output or "").lower()
|
|
if any(
|
|
m in text
|
|
for m in (
|
|
"unknown projector type",
|
|
"unsupported projector",
|
|
"unsupported mmproj",
|
|
)
|
|
):
|
|
return True
|
|
# Builds that phrase it via clip.cpp without the exact words above.
|
|
return (
|
|
"clip" in text
|
|
and "projector" in text
|
|
and ("unknown" in text or "unsupported" in text or "not supported" in text)
|
|
)
|
|
|
|
@staticmethod
|
|
def _output_has_nonprojector_diagnostic(output: str) -> bool:
|
|
"""True when the output already names a concrete non-projector cause (out
|
|
of memory, an unsupported architecture, a tensor-parallel limit). A hard
|
|
crash carrying such a marker must surface that error, not be silently
|
|
retried text-only as if the vision projector were at fault; a bare crash
|
|
with no marker still gets the text-only retry.
|
|
"""
|
|
text = (output or "").lower()
|
|
return any(
|
|
m in text
|
|
for m in (
|
|
"out of memory",
|
|
"failed to allocate",
|
|
"unknown model architecture",
|
|
"split_mode_tensor not implemented",
|
|
)
|
|
)
|
|
|
|
@staticmethod
|
|
def _is_tensor_split_assert(output: str) -> bool:
|
|
"""True only for the #6415 split-axis warmup assert (GGML_BACKEND_SPLIT_AXIS_*),
|
|
not any ggml assert/abort, so an unrelated invariant isn't cached. stderr is
|
|
merged into output."""
|
|
text = (output or "").lower()
|
|
if "ggml_assert" not in text and "ggml_abort" not in text:
|
|
return False
|
|
# the split-axis enum token, unique to this assert (not the source file).
|
|
return "split_axis" in text
|
|
|
|
@staticmethod
|
|
def _is_signal_crash(returncode: Optional[int]) -> bool:
|
|
"""True only on a hard fault (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS or a
|
|
Windows 0xC0000000+ status), not SIGKILL/SIGTERM/SIGINT (OOM killer /
|
|
unload) nor a clean exit or still-running (None) process.
|
|
"""
|
|
if returncode is None:
|
|
return False
|
|
if returncode >= 0xC0000000: # Windows access violation / illegal instruction
|
|
return True
|
|
return -returncode in (4, 6, 7, 8, 11) # SIGILL SIGABRT SIGBUS SIGFPE SIGSEGV
|
|
|
|
@staticmethod
|
|
def _is_abort_exit(returncode: Optional[int]) -> bool:
|
|
"""Windows CRT abort() exit code (3) from GGML_ASSERT on MSVC -- not a POSIX
|
|
signal or 0xC0000000+ NTSTATUS."""
|
|
return returncode == 3
|
|
|
|
@classmethod
|
|
def _should_record_tensor_split_abort(cls, returncode: Optional[int], output: str) -> bool:
|
|
"""The #6415 split-axis abort: the marker plus a hard crash (POSIX signal or
|
|
Windows abort exit). Marker required so a generic crash isn't cached."""
|
|
return cls._is_tensor_split_assert(output) and (
|
|
cls._is_signal_crash(returncode) or cls._is_abort_exit(returncode)
|
|
)
|
|
|
|
@staticmethod
|
|
def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]:
|
|
"""Return cmd with flash attention forced off, or None when its effective
|
|
(last-wins) value is already off/absent so there is nothing to retry. FA
|
|
kernels hard-crash at startup on some ROCm builds; disabling FA keeps
|
|
vision and MTP, the least destructive rung. A bare --flash-attn/-fa reads
|
|
as on, so it counts toward the effective value and is neutralised too;
|
|
every form is flipped in place (length preserved for downstream slices)."""
|
|
out = list(cmd)
|
|
|
|
def explicit(i):
|
|
nxt = out[i + 1] if i + 1 < len(out) else None
|
|
return nxt if nxt in ("on", "auto", "off") else None
|
|
|
|
effective = None
|
|
for i, tok in enumerate(out):
|
|
if tok.startswith(("--flash-attn=", "-fa=")):
|
|
effective = tok.partition("=")[2]
|
|
elif tok in ("--flash-attn", "-fa"):
|
|
effective = explicit(i) or "on"
|
|
if effective not in ("on", "auto"):
|
|
return None
|
|
for i, tok in enumerate(out):
|
|
if tok.startswith(("--flash-attn=", "-fa=")):
|
|
flag, _, value = tok.partition("=")
|
|
if value in ("on", "auto"):
|
|
out[i] = f"{flag}=off"
|
|
elif tok in ("--flash-attn", "-fa"):
|
|
if explicit(i) in ("on", "auto"):
|
|
out[i + 1] = "off"
|
|
elif explicit(i) is None: # bare flag (reads as on) -> explicit off
|
|
out[i] = f"{tok}=off"
|
|
return out
|
|
|
|
@staticmethod
|
|
def _strip_mmproj_args(cmd: list[str]) -> list[str]:
|
|
"""Return cmd without the '--mmproj <path>' pair (text-only retry).
|
|
Every other flag is preserved; a no-op when --mmproj is absent.
|
|
"""
|
|
out: list[str] = []
|
|
skip_value = False
|
|
for tok in cmd:
|
|
if skip_value:
|
|
skip_value = False
|
|
continue
|
|
if tok == "--mmproj":
|
|
skip_value = True
|
|
continue
|
|
out.append(tok)
|
|
return out
|
|
|
|
@staticmethod
|
|
def _redacted_cmd_for_log(cmd: "list[str]") -> "list[str]":
|
|
"""Copy of cmd with the value after --api-key replaced by <redacted>."""
|
|
out = list(cmd)
|
|
if "--api-key" in out:
|
|
ki = out.index("--api-key") + 1
|
|
if ki < len(out):
|
|
out[ki] = "<redacted>"
|
|
return out
|
|
|
|
def _start_llama_process(self, cmd: list[str], env: dict) -> None:
|
|
"""Spawn llama-server from cmd and start draining its output.
|
|
|
|
Caller holds self._lock. Resets the stdout buffer, opens a fresh
|
|
per-attempt tee log, launches the process, and starts the drain
|
|
thread. Used for the initial start and the text-only mmproj retry.
|
|
"""
|
|
# Defensive kill: if a concurrent load slipped past Phase 1
|
|
# (because its `self._process` was None at the time) and already
|
|
# stored a Popen handle here, drop that orphan before we overwrite
|
|
# the reference. See issue #5161.
|
|
self._kill_process()
|
|
|
|
self._stdout_lines = []
|
|
# Tee llama-server output to a dedicated log file so a post-mortem
|
|
# in CI (or after a remote-debug session) has the full subprocess
|
|
# trail even when the parent only stored the last 50 lines.
|
|
self._llama_log_fh = None
|
|
try:
|
|
log_dir = _swa_cache_path().parent / "logs" / "llama-server"
|
|
log_dir.mkdir(parents = True, exist_ok = True)
|
|
self._llama_log_path = log_dir / f"llama-{int(time.time())}-port-{self._port}.log"
|
|
self._llama_log_fh = open(
|
|
self._llama_log_path,
|
|
"w",
|
|
encoding = "utf-8",
|
|
buffering = 1,
|
|
)
|
|
logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}")
|
|
except OSError as e:
|
|
# Best-effort; never block the load on logging.
|
|
logger.debug(f"Could not open llama-server log file: {e}")
|
|
self._llama_log_path = None
|
|
|
|
# Log the argv per attempt (the text-only mmproj retry re-enters here
|
|
# with --mmproj stripped), redacting the API key.
|
|
logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}")
|
|
|
|
self._process = subprocess.Popen(
|
|
cmd,
|
|
stdout = subprocess.PIPE,
|
|
stderr = subprocess.STDOUT,
|
|
text = True,
|
|
env = env,
|
|
**_windows_hidden_subprocess_kwargs(),
|
|
**_child_popen_kwargs(),
|
|
)
|
|
# Cross-session backstop: record the PID so a later startup can reap this
|
|
# server if parent-death cleanup did not run (macOS / best-effort failure).
|
|
self._record_server_pid(self._process.pid)
|
|
|
|
# Start background thread to drain stdout and prevent pipe deadlock
|
|
self._stdout_thread = threading.Thread(
|
|
target = self._drain_stdout, daemon = True, name = "llama-stdout"
|
|
)
|
|
self._stdout_thread.start()
|
|
|
|
def load_model(
|
|
self,
|
|
*,
|
|
# Local mode: pass a path to a .gguf file
|
|
gguf_path: Optional[str] = None,
|
|
# Vision projection (mmproj) for local vision models
|
|
mmproj_path: Optional[str] = None,
|
|
# Separate MTP drafter for local Gemma loads (HF loads auto-resolve it)
|
|
mtp_draft_path: Optional[str] = None,
|
|
# HF mode: let llama-server download via -hf "repo:quant"
|
|
hf_repo: Optional[str] = None,
|
|
hf_variant: Optional[str] = None,
|
|
hf_token: Optional[str] = None,
|
|
# Common
|
|
model_identifier: str,
|
|
is_vision: bool = False,
|
|
n_ctx: int = 4096,
|
|
chat_template_override: Optional[str] = None,
|
|
cache_type_kv: Optional[str] = None,
|
|
speculative_type: Optional[str] = None,
|
|
spec_draft_n_max: Optional[int] = None,
|
|
tensor_parallel: bool = False,
|
|
n_threads: Optional[int] = None,
|
|
n_gpu_layers: Optional[int] = None, # caller compat, unused
|
|
n_parallel: int = 1,
|
|
extra_args: Optional[List[str]] = None,
|
|
# Route-level tensor->layer fallback retry: keep the layer split multi-GPU.
|
|
preserve_multi_gpu_on_layer: bool = False,
|
|
) -> bool:
|
|
"""Start llama-server with a GGUF model.
|
|
|
|
Two modes:
|
|
- Local: ``gguf_path="/path/to/model.gguf"`` → uses ``-m``
|
|
- HF: ``hf_repo="...-GGUF", hf_variant="Q4_K_M"`` → uses ``-hf``
|
|
|
|
Returns True if the server started and the health check passed.
|
|
"""
|
|
# Raw load inputs so the runtime MTP-crash reload can replay this model
|
|
# without MTP. Committed to _last_load_kwargs only on a healthy load.
|
|
_pending_load_kwargs = {
|
|
"gguf_path": gguf_path,
|
|
"mmproj_path": mmproj_path,
|
|
"mtp_draft_path": mtp_draft_path,
|
|
"hf_repo": hf_repo,
|
|
"hf_variant": hf_variant,
|
|
"hf_token": hf_token,
|
|
"model_identifier": model_identifier,
|
|
"is_vision": is_vision,
|
|
"n_ctx": n_ctx,
|
|
"chat_template_override": chat_template_override,
|
|
"cache_type_kv": cache_type_kv,
|
|
"speculative_type": speculative_type,
|
|
"spec_draft_n_max": spec_draft_n_max,
|
|
"tensor_parallel": tensor_parallel,
|
|
"n_threads": n_threads,
|
|
"n_gpu_layers": n_gpu_layers,
|
|
"n_parallel": n_parallel,
|
|
"extra_args": list(extra_args) if extra_args is not None else None,
|
|
# Replayed by _respawn_if_dead so a downgraded model stays multi-GPU.
|
|
"preserve_multi_gpu_on_layer": preserve_multi_gpu_on_layer,
|
|
}
|
|
# Serialise the whole load so concurrent /load calls never leave two
|
|
# llama-server processes alive (#5401 / #5161). Doesn't block /unload.
|
|
with self._serial_load_lock:
|
|
# In-app update swapping binaries: refuse fast (set under this lock,
|
|
# so any in-flight load has drained) instead of using a half-swapped one.
|
|
if getattr(self, "_llama_update_in_progress", False):
|
|
raise RuntimeError("llama.cpp is updating; try again in a moment.")
|
|
# Duplicate /load that raced past the route check: do nothing if the
|
|
# live server already satisfies this request.
|
|
if self._already_in_target_state(
|
|
gguf_path = gguf_path,
|
|
mtp_draft_path = mtp_draft_path,
|
|
model_identifier = model_identifier,
|
|
hf_variant = hf_variant,
|
|
n_ctx = n_ctx,
|
|
cache_type_kv = cache_type_kv,
|
|
speculative_type = speculative_type,
|
|
spec_draft_n_max = spec_draft_n_max,
|
|
tensor_parallel = tensor_parallel,
|
|
chat_template_override = chat_template_override,
|
|
extra_args = extra_args,
|
|
is_vision = is_vision,
|
|
preserve_multi_gpu_on_layer = preserve_multi_gpu_on_layer,
|
|
):
|
|
logger.info(
|
|
f"load_model: backend already in target state for "
|
|
f"'{model_identifier}', skipping reload"
|
|
)
|
|
# Retry probe only if a prior attempt didn't finish.
|
|
if not self._audio_probed:
|
|
try:
|
|
detected = self._detect_audio_type_strict()
|
|
self._audio_probed = True
|
|
except Exception as exc:
|
|
logger.debug("Fast-path audio probe failed: %s", exc)
|
|
detected = None
|
|
if not self._apply_detected_audio(detected):
|
|
return False
|
|
if not self._healthy:
|
|
return False
|
|
return True
|
|
|
|
self._cancel_event.clear()
|
|
|
|
# ── Phase 1: kill old process (under lock, fast) ──────────
|
|
with self._lock:
|
|
self._kill_process()
|
|
|
|
# Resolve llama-server now but defer a not-found error: a block-diffusion
|
|
# GGUF uses the diffusion runner, and its arch is only known after the header.
|
|
binary = self._find_llama_server_binary()
|
|
|
|
# ── Phase 2: download (NO lock held, so cancel can proceed) ──
|
|
# mtp_draft_path arrives set for local Gemma loads (detected
|
|
# sibling); for -hf loads it's None here and resolved just below.
|
|
# Scope HF_HUB_OFFLINE to the download block only when DNS is
|
|
# dead; cleanup runs even on exception so a transient hiccup
|
|
# can't quarantine future loads.
|
|
if hf_repo:
|
|
with _hf_offline_if_dns_dead():
|
|
model_path = self._download_gguf(
|
|
hf_repo = hf_repo,
|
|
hf_variant = hf_variant,
|
|
hf_token = hf_token,
|
|
)
|
|
# Auto-download mmproj for vision models unless opted out.
|
|
if is_vision and not mmproj_path and not extra_args_disable_mmproj(extra_args):
|
|
mmproj_path = self._download_mmproj(
|
|
hf_repo = hf_repo,
|
|
hf_token = hf_token,
|
|
)
|
|
# Auto-download the separate MTP drafter (e.g. Gemma) when
|
|
# the requested spec mode can use it. Repos with the head
|
|
# baked into the main GGUF (Qwen) have no mtp- sibling and
|
|
# this no-ops, so the size gate stays out of it: a separate
|
|
# drafter speeds up even sub-3B (Gemma E2B), and the resolver
|
|
# below decides the final emission. Skipped only when the
|
|
# user disabled MTP or drives --spec-type manually.
|
|
_spec_canon = _canonicalize_spec_mode(speculative_type) or "auto"
|
|
if (
|
|
not mtp_draft_path
|
|
and _spec_canon in ("auto", "mtp", "mtp+ngram")
|
|
and not _extra_args_set_spec_type(extra_args)
|
|
):
|
|
mtp_draft_path = self._download_mtp(
|
|
hf_repo = hf_repo,
|
|
hf_token = hf_token,
|
|
)
|
|
elif gguf_path:
|
|
if not Path(gguf_path).is_file():
|
|
raise FileNotFoundError(f"GGUF file not found: {gguf_path}")
|
|
model_path = gguf_path
|
|
else:
|
|
raise ValueError("Either gguf_path or hf_repo must be provided")
|
|
|
|
# Set identifier early so _read_gguf_metadata can use it (DeepSeek).
|
|
self._model_identifier = model_identifier
|
|
|
|
# Read GGUF metadata (context_length, chat_template); header-only.
|
|
self._read_gguf_metadata(model_path)
|
|
|
|
if self._cancel_event.is_set():
|
|
logger.info("Load cancelled after download phase")
|
|
return False
|
|
|
|
# Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server;
|
|
# serve them with the diffusion runner (same OpenAI-compat interface).
|
|
if self._is_diffusion:
|
|
# Not a tensor/layer GGUF: clear any preserved-fallback flag from a
|
|
# prior load (this path skips the command builder that clears it).
|
|
self._layer_preserves_tensor_intent = False
|
|
with self._lock:
|
|
if self._cancel_event.is_set():
|
|
logger.info("Load cancelled before diffusion server start")
|
|
return False
|
|
return self._start_diffusion_server(
|
|
model_path = model_path,
|
|
gguf_path = gguf_path,
|
|
hf_repo = hf_repo,
|
|
hf_variant = hf_variant,
|
|
model_identifier = model_identifier,
|
|
n_ctx = n_ctx,
|
|
extra_args = extra_args,
|
|
)
|
|
|
|
if not binary:
|
|
# distinguish a transiently locked binary (antivirus / in-flight
|
|
# install) from a missing one so the user retries, not reinstalls
|
|
locked = self._find_llama_server_binary(include_denied = True)
|
|
if locked:
|
|
raise RuntimeError(
|
|
f"llama-server at {locked} is temporarily unavailable "
|
|
"(access-denied; antivirus or an in-flight install). "
|
|
"Retry the load once it is released."
|
|
)
|
|
# Reached only after the diffusion early-return above, so this is a
|
|
# genuine llama-server-backed GGUF with no runtime. Raise the typed
|
|
# error so /load returns the actionable 400 (not a generic 500), the
|
|
# same message remote validation already shows.
|
|
raise LlamaServerNotFoundError(LLAMA_SERVER_NOT_FOUND_DETAIL)
|
|
|
|
# Outside ``self._lock`` so /unload, /cancel, /status aren't
|
|
# blocked. ``unload_model`` also records the kill, so the
|
|
# frontend /unload+/load Apply path engages the wait here even
|
|
# without an in-process kill.
|
|
self._wait_for_vram_settle(since_kill = self._last_kill_monotonic)
|
|
|
|
# ── Phase 3: start llama-server (under lock) ──────────────
|
|
with self._lock:
|
|
# Re-check cancel inside lock
|
|
if self._cancel_event.is_set():
|
|
logger.info("Load cancelled before server start")
|
|
return False
|
|
|
|
self._port = self._find_free_port()
|
|
|
|
# Select GPU(s) from model size + estimated KV cache. Seed
|
|
# safe defaults before probing so the except path has valid
|
|
# state to publish.
|
|
ctx_override = parse_ctx_override(extra_args)
|
|
requested_ctx = resolve_requested_ctx(extra_args, n_ctx)
|
|
cache_override = parse_cache_override(extra_args)
|
|
# Budget the heavier of asymmetric --cache-type-k/-v extras (they
|
|
# win per axis at launch, appended last); resolve_cache_type_kv only
|
|
# returns the last-wins type, which under-reserves the heavier axis.
|
|
# The user's extras still set the real (possibly asymmetric) child
|
|
# cache, so this only affects the reserve, not the emitted command.
|
|
_extras_cache = _extra_args_main_cache_type_for_budget(extra_args)
|
|
cache_type_kv = _extras_cache if _extras_cache is not None else cache_type_kv
|
|
_cache_type_from_env = False
|
|
if cache_type_kv is None:
|
|
# Param/extras set nothing, so the child inherits
|
|
# LLAMA_ARG_CACHE_TYPE_K/_V. Adopt a heavier env type (f32) for
|
|
# the reserve only; the launch does NOT re-emit it (that would
|
|
# rewrite an asymmetric K=f32,V=f16 env into symmetric flags),
|
|
# so _cache_type_from_env keeps it out of the emitted flags.
|
|
cache_type_kv = _env_main_cache_type_for_budget()
|
|
_cache_type_from_env = cache_type_kv is not None
|
|
# A user --split-mode in extras last-wins-overrides the toggle, and
|
|
# an inherited tensor LLAMA_ARG_SPLIT_MODE flips it on (the child
|
|
# would run tensor unbudgeted otherwise). The duplicate-load matchers
|
|
# use the same helper so a healthy env-driven tensor server matches.
|
|
split_mode_override = parse_split_mode_override(extra_args)
|
|
tensor_parallel = _effective_tensor_parallel(extra_args, tensor_parallel)
|
|
# Tensor mode aborts on a quantized KV cache, so drop it for the
|
|
# tensor attempt (and strip any inherited/explicit --cache-type
|
|
# that would re-impose it when appended last). Layer split does
|
|
# support it, so remember the dropped type and the original extras
|
|
# to restore (verbatim, incl. an asymmetric K/V) if we later fall
|
|
# back to layer split below.
|
|
_tensor_dropped_cache_type_kv: Optional[str] = None
|
|
_tensor_dropped_extra_args: Optional[list] = None
|
|
# Tensor mode rejects any quantized axis. cache_type_kv is the
|
|
# heavier-by-bytes budget type, which can mask a quantized axis (an
|
|
# f16 budget hides a paired q4_0), so also test each explicit
|
|
# --cache-type-k/-v extra, not just the budget type.
|
|
_ck_extra, _cv_extra = parse_cache_override_per_axis(extra_args)
|
|
_cache_non_tensor_safe = any(
|
|
c and c.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES
|
|
for c in (cache_type_kv, _ck_extra, _cv_extra)
|
|
)
|
|
if tensor_parallel and _cache_non_tensor_safe:
|
|
logger.info(
|
|
"Tensor parallelism requires a non-quantized KV cache; "
|
|
"ignoring cache type %s for the tensor attempt.",
|
|
cache_type_kv,
|
|
)
|
|
_tensor_dropped_cache_type_kv = cache_type_kv
|
|
cache_type_kv = None
|
|
if extra_args:
|
|
# Keep the originals so a layer downgrade restores the real
|
|
# (possibly asymmetric) --cache-type-k/-v the layer path
|
|
# supports, not just the scalar heavier type.
|
|
_tensor_dropped_extra_args = list(extra_args)
|
|
extra_args = strip_shadowing_flags(
|
|
extra_args,
|
|
strip_context = False,
|
|
strip_cache = True,
|
|
strip_spec = False,
|
|
strip_template = False,
|
|
strip_split_mode = False,
|
|
)
|
|
# The launch keeps an inherited tensor-safe env cache type (the
|
|
# env cleanup only pops quantized ones), so re-adopt a heavier
|
|
# env type (f32) for the budget here too -- mirrors the initial
|
|
# adoption, which was skipped because the param/extras set the
|
|
# (now-dropped) quantized type. Else the child allocates f32 KV
|
|
# against an f16 budget.
|
|
_env_tensor_cache = _env_main_cache_type_for_budget()
|
|
if _env_tensor_cache is not None:
|
|
cache_type_kv = _env_tensor_cache
|
|
_cache_type_from_env = True
|
|
if ctx_override is not None and ctx_override > 0:
|
|
logger.info(f"User --ctx-size {ctx_override} honored; skipping auto-reduce")
|
|
if cache_override is not None:
|
|
_ck, _cv = parse_cache_override_per_axis(extra_args)
|
|
logger.info(
|
|
f"User --cache-type-k/-v (k={_ck}, v={_cv}) honored; "
|
|
"KV estimate budgets the heavier axis"
|
|
)
|
|
if split_mode_override is not None:
|
|
logger.info(
|
|
f"User --split-mode {split_mode_override} honored; "
|
|
"reconciled into tensor_parallel state"
|
|
)
|
|
effective_ctx = requested_ctx if requested_ctx > 0 else (self._context_length or 0)
|
|
max_available_ctx = self._context_length or effective_ctx
|
|
gpus: list[tuple[int, int]] = []
|
|
# Keep fit-budget and launch-flag mmproj resolution in sync.
|
|
launch_mmproj_path = None
|
|
if not extra_args_disable_mmproj(extra_args):
|
|
launch_mmproj_path = self._resolve_launch_mmproj_path(
|
|
model_path = model_path,
|
|
mmproj_path = mmproj_path,
|
|
)
|
|
# Need both a resolved mmproj AND the config vision flag; a stray
|
|
# mmproj passing the family-name heuristic must not flip a non-VLM
|
|
# GGUF into vision mode.
|
|
effective_is_vision = bool(launch_mmproj_path) and bool(is_vision)
|
|
if is_vision and not effective_is_vision:
|
|
logger.warning(
|
|
"Vision-capable GGUF loaded without a usable mmproj; "
|
|
"image input will be disabled for this session"
|
|
)
|
|
model_size = None # set in the fit try; used by the APU RAM guard
|
|
# Layer-fallback min GPUs; raised below on a tensor downgrade. Bound
|
|
# before the try so the --fit-on except path still has it (no UnboundLocal).
|
|
_layer_min_gpus = 1
|
|
try:
|
|
gguf_size = self._get_gguf_size_bytes(model_path)
|
|
# Include GPU-loaded mmproj in the fit budget (#5825).
|
|
mmproj_size = (
|
|
self._mmproj_vram_bytes(launch_mmproj_path) if effective_is_vision else 0
|
|
)
|
|
model_size = gguf_size + mmproj_size
|
|
# 2-tuple gpus for existing logic + a total map for the absolute
|
|
# per-GPU headroom (correct when the GPU is already partly used).
|
|
_gpu_mem = self._get_gpu_memory()
|
|
gpus = [(idx, free) for idx, free, _t in _gpu_mem]
|
|
total_by_idx = {idx: total for idx, _f, total in _gpu_mem}
|
|
|
|
def _gpu_usable(g, frac = _CTX_FIT_VRAM_FRACTION):
|
|
# Per-GPU usable budget for ranking: free - (1-frac)*total.
|
|
# Callers pass the ACTIVE fraction so the ranking matches the
|
|
# budget the fit then tests (else mixed totals mis-order).
|
|
idx, free = g
|
|
t = total_by_idx.get(idx, 0)
|
|
if t > 0:
|
|
return free - (1.0 - frac) * t
|
|
return free * frac
|
|
|
|
def _pool_budget_mib(subset, frac):
|
|
# Sum each GPU's own usable budget. Pooling free and total
|
|
# separately would let an unknown-total GPU (MIG/vGPU/N/A)
|
|
# add full free with no cushion among known-total GPUs.
|
|
return sum(max(0.0, _gpu_usable(g, frac)) for g in subset)
|
|
|
|
# Resolve effective context: 0 means let llama-server use
|
|
# the model's native length. Only expand to a known native
|
|
# length if metadata exists; else keep 0 as a sentinel.
|
|
if requested_ctx > 0:
|
|
effective_ctx = requested_ctx
|
|
elif self._context_length is not None:
|
|
effective_ctx = self._context_length
|
|
else:
|
|
effective_ctx = 0
|
|
original_ctx = effective_ctx
|
|
# Default UI ceiling to the native context length;
|
|
# GPU/VRAM-fit logic below may shrink it on limited HW.
|
|
max_available_ctx = self._context_length or effective_ctx
|
|
|
|
# Will MTP engage? If so, auto-fit reserves draft-model VRAM.
|
|
# Mirrors _build_speculative_flags: forced mtp/mtp+ngram always
|
|
# engage; auto only on an MTP model >= 3B; ngram/off never. A
|
|
# separate drafter (Gemma) counts as an MTP model.
|
|
_mtp_canonical = _canonicalize_spec_mode(speculative_type)
|
|
_mtp_effective = _mtp_canonical or "auto"
|
|
_mtp_size_for_fit = _extract_model_size_b(model_identifier)
|
|
# Sub-3B drops MTP only for an embedded head; a separate
|
|
# drafter (Gemma) engages and needs its VRAM reserved.
|
|
_mtp_sub_3b_for_fit = (
|
|
_mtp_size_for_fit is not None
|
|
and _mtp_size_for_fit < _MTP_MIN_SIZE_B
|
|
and not bool(mtp_draft_path)
|
|
)
|
|
# LLAMA_ARG_SPEC_TYPE only reaches the child when neither extras
|
|
# nor Studio emit a spec flag (mode "off", no user --spec-type),
|
|
# since _build_speculative_flags emits one for every other mode.
|
|
# Consult the env for the reserve only then, else a stale MTP env
|
|
# would over-reserve.
|
|
_spec_env: Mapping[str, str] = (
|
|
os.environ
|
|
if (not _extra_args_set_spec_type(extra_args) and _mtp_canonical == "off")
|
|
else {}
|
|
)
|
|
# Extras can run MTP even when Studio suppresses its own emission.
|
|
_user_mtp_via_extras = _extra_args_requests_mtp(extra_args, env = _spec_env)
|
|
# A non-MTP model-based draft mode (draft-simple/draft-eagle3) in
|
|
# extras also loads a separate draft model that needs reserving;
|
|
# engage only when extras actually name a drafter for it.
|
|
_user_draft_via_extras = _extra_args_requests_separate_draft(
|
|
extra_args, env = _spec_env
|
|
) and bool(_extra_args_mtp_draft_path(extra_args))
|
|
# Mirror _build_speculative_flags: reserve only for MTP the launch
|
|
# resolver will actually emit (needs a head/drafter and a binary
|
|
# that supports --spec-type mtp).
|
|
_mtp_model_for_fit = bool(
|
|
self._nextn_predict_layers
|
|
or _is_mtp_model_name(model_identifier, model_path)
|
|
or bool(mtp_draft_path)
|
|
) and not (
|
|
# Drafterless Gemma falls back to ngram-mod; reserve no
|
|
# drafter VRAM for it (mirrors the launch resolver).
|
|
_is_gemma_mtp_name(model_identifier, model_path)
|
|
and not mtp_draft_path
|
|
and not self._nextn_predict_layers
|
|
)
|
|
_mtp_binary_ok = True
|
|
_mtp_probe_raised = False
|
|
if not _user_mtp_via_extras:
|
|
try:
|
|
_mtp_binary_ok = bool(
|
|
(self.probe_server_capabilities(binary) or {}).get("mtp_token")
|
|
)
|
|
except Exception:
|
|
_mtp_binary_ok = False
|
|
_mtp_probe_raised = True
|
|
_auto_studio_mtp = (
|
|
not _extra_args_set_spec_type(extra_args)
|
|
and _mtp_model_for_fit
|
|
and (
|
|
_mtp_effective in ("mtp", "mtp+ngram")
|
|
or (_mtp_effective == "auto" and not _mtp_sub_3b_for_fit)
|
|
)
|
|
and (
|
|
_mtp_binary_ok
|
|
# Reserve on a raised (uncached) probe too: it re-probes in
|
|
# _build_speculative_flags and may still engage MTP (embedded
|
|
# head or separate drafter -- _mtp_model_for_fit covers both).
|
|
or _mtp_probe_raised
|
|
)
|
|
)
|
|
_mtp_will_engage = bool(
|
|
_user_mtp_via_extras or _user_draft_via_extras or _auto_studio_mtp
|
|
)
|
|
# The duplicated full target-KV copy (ctx_tgt) is an MTP-only
|
|
# cost: the MTP head runs a second context over the target
|
|
# model's own KV geometry. The separate-drafter spec modes
|
|
# (draft-simple/draft-eagle3, reached via _user_draft_via_extras)
|
|
# load a small distinct drafter with its own KV and keep no such
|
|
# copy, so only charge it when the engaged mode is truly MTP.
|
|
_engaged_is_mtp = bool(_user_mtp_via_extras or _auto_studio_mtp)
|
|
|
|
# Effective draft depth: extras win (last-wins at launch), else
|
|
# the field, else the platform default (2 GPU / 3 CPU).
|
|
_extra_n_max = _extra_args_spec_draft_n_max(extra_args)
|
|
_mtp_eff_n_max = _extra_n_max if _extra_n_max is not None else spec_draft_n_max
|
|
if _mtp_eff_n_max is None:
|
|
_mtp_eff_n_max = 2 if gpus else 3
|
|
# Separate-drafter weights live on GPU (an embedded head is
|
|
# already in model_size). Size the drafter the launch loads, by
|
|
# precedence: extras --model-draft (last-wins), else Studio's
|
|
# emitted mtp_draft_path, else the env drafter. Sizing the wrong
|
|
# one would under-reserve and OOM.
|
|
_cli_draft_for_budget = _extra_args_mtp_draft_path(extra_args, env = {})
|
|
_studio_draft_for_budget = (
|
|
mtp_draft_path
|
|
if (
|
|
_mtp_will_engage
|
|
and mtp_draft_path
|
|
and not _extra_args_set_spec_type(extra_args)
|
|
)
|
|
else None
|
|
)
|
|
_env_draft_for_budget = _extra_args_mtp_draft_path([], env = os.environ)
|
|
_mtp_draft_for_budget = (
|
|
_cli_draft_for_budget or _studio_draft_for_budget or _env_draft_for_budget
|
|
)
|
|
# Drafter offloaded to CPU keeps its weights+KV off the GPU, so
|
|
# drop it from the budget (an embedded head stays in the model).
|
|
# Consult the env too: the child honors LLAMA_ARG_N_GPU_LAYERS_DRAFT.
|
|
_draft_on_cpu = _extra_args_draft_offloaded_to_cpu(extra_args, env = os.environ)
|
|
if _draft_on_cpu:
|
|
_mtp_draft_for_budget = None
|
|
_mtp_draft_weights = 0
|
|
if _mtp_draft_for_budget:
|
|
try:
|
|
_mtp_draft_weights = self._get_gguf_size_bytes(_mtp_draft_for_budget)
|
|
except Exception:
|
|
_mtp_draft_weights = 0
|
|
# Draft K/V types (f16 by default; independent extras overrides).
|
|
_mtp_draft_ck, _mtp_draft_cv = _extra_args_draft_cache_types(extra_args)
|
|
|
|
# Byte-accurate reserve when dims allow, else None -> flat fallback.
|
|
mtp_overhead_fn: Optional[Callable[[int], int]] = None
|
|
# True when the byte reserve is the drafter weights ONLY because
|
|
# its KV couldn't be sized; the flat fraction must then stay on
|
|
# as the cushion for that unsized draft KV (it is not covered by
|
|
# the weights-only mtp_overhead_fn).
|
|
_mtp_kv_unsized = False
|
|
if _mtp_will_engage:
|
|
_probe_ctx = self._context_length or (
|
|
effective_ctx if effective_ctx > 0 else 4096
|
|
)
|
|
_draft_kv_probe = self._mtp_draft_kv_bytes(
|
|
_probe_ctx,
|
|
drafter_path = _mtp_draft_for_budget,
|
|
draft_cache_type_k = _mtp_draft_ck,
|
|
draft_cache_type_v = _mtp_draft_cv,
|
|
n_parallel = n_parallel,
|
|
)
|
|
if (
|
|
self._estimate_mtp_overhead_bytes(
|
|
_probe_ctx,
|
|
spec_draft_n_max = _mtp_eff_n_max,
|
|
draft_cache_type_k = _mtp_draft_ck,
|
|
draft_cache_type_v = _mtp_draft_cv,
|
|
drafter_path = _mtp_draft_for_budget,
|
|
draft_weights_bytes = _mtp_draft_weights,
|
|
n_parallel = n_parallel,
|
|
mtp_keeps_target_ctx = _engaged_is_mtp,
|
|
)
|
|
is not None
|
|
):
|
|
# Reserve is weights-only when the draft KV is unsizable.
|
|
_mtp_kv_unsized = _draft_kv_probe is None
|
|
|
|
# Closure binding this load's draft params; ctx varies.
|
|
def mtp_overhead_fn(
|
|
ctx: int,
|
|
_n: int = _mtp_eff_n_max,
|
|
_ck: Optional[str] = _mtp_draft_ck,
|
|
_cv: Optional[str] = _mtp_draft_cv,
|
|
_dp: Optional[str] = _mtp_draft_for_budget,
|
|
_w: int = _mtp_draft_weights,
|
|
_np: int = n_parallel,
|
|
_mtp: bool = _engaged_is_mtp,
|
|
) -> int:
|
|
v = self._estimate_mtp_overhead_bytes(
|
|
ctx,
|
|
spec_draft_n_max = _n,
|
|
draft_cache_type_k = _ck,
|
|
draft_cache_type_v = _cv,
|
|
drafter_path = _dp,
|
|
draft_weights_bytes = _w,
|
|
n_parallel = _np,
|
|
mtp_keeps_target_ctx = _mtp,
|
|
)
|
|
return v if v is not None else 0
|
|
|
|
def _mtp_bytes(ctx: int) -> int:
|
|
return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0
|
|
|
|
# Effective micro-batch (a user --ubatch override scales the
|
|
# compute buffer); None -> the 512 default in the estimate.
|
|
_effective_ubatch = _extra_args_n_ubatch(extra_args)
|
|
|
|
def _cc_bytes(ctx: int, n_gpus: int = 1) -> int:
|
|
# Context-linear compute-buffer growth (flash-attn KQ mask +
|
|
# attention scratch); the flat _compute_buffer_pipeline folded
|
|
# into model_size_fit only covers ctx -> 0. Charged per
|
|
# candidate context so the fit can't over-pin and spill. The
|
|
# rate depends on the KV cache type (quantized adds a dequant
|
|
# scratch), so pass it through. In a layer split this buffer is
|
|
# replicated on EVERY device (measured ~equal per GPU), so scale
|
|
# by the device count; a large model at high context otherwise
|
|
# under-reserves ~(n-1)x it (e.g. Qwen3.5-397B on 3 GPUs).
|
|
return max(1, n_gpus) * self._compute_buffer_ctx_bytes(
|
|
ctx, _effective_ubatch, cache_type_kv
|
|
)
|
|
|
|
# Layer-split compute buffer (one lump; tensor mode reserves it
|
|
# per device in _plan_tensor_parallel). Context-independent, so
|
|
# fold it into the model footprint for the branches below. Falls
|
|
# back to the flat reserve when dims are missing (returns 0), a
|
|
# safe upper bound since the tensor buffer >= the layer one.
|
|
_compute_buffer_pipeline = self._estimate_compute_buffer_bytes(
|
|
n_ubatch = _effective_ubatch,
|
|
n_parallel = n_parallel,
|
|
per_device_tensor = False,
|
|
)
|
|
if _compute_buffer_pipeline <= 0:
|
|
_compute_buffer_pipeline = (
|
|
self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024
|
|
)
|
|
|
|
# Layer split adds a fixed per-device overhead on every GPU. The
|
|
# folded buffer covers one device; reserve the extra devices'
|
|
# share so a k-GPU split can't pin a context that OOMs a device
|
|
# (k=1 adds nothing).
|
|
_pipeline_overhead_bytes = self._PIPELINE_PER_DEVICE_OVERHEAD_MIB * 1024 * 1024
|
|
|
|
# Auto-cap context to fit VRAM and select GPUs. Explicit n_ctx:
|
|
# honor it, cap only if it fits no combination. Auto (native):
|
|
# prefer fewer GPUs with reduced context (multi-GPU is slower).
|
|
gpu_indices, use_fit = None, True
|
|
# Per-GPU weight proportions for tensor mode (None = even).
|
|
tp_tensor_split: Optional[list[int]] = None
|
|
explicit_ctx = requested_ctx > 0
|
|
# Flat MTP reserve fraction: used only as the fallback when the
|
|
# byte-accurate mtp_overhead_fn can't size the draft KV (dims
|
|
# unavailable, or _mtp_kv_unsized = weights-only). A separate
|
|
# drafter on CPU uses no GPU (no reserve); an embedded head is on
|
|
# GPU regardless of draft-offload flags (keep its reserve).
|
|
_flat_mtp_engages = _mtp_will_engage and (
|
|
mtp_overhead_fn is None or _mtp_kv_unsized
|
|
)
|
|
_draft_cpu_no_embedded = _draft_on_cpu and not self._nextn_predict_layers
|
|
# MTP reserves GPU VRAM unless its only drafter is a separate
|
|
# CPU-offloaded one (an embedded head stays on GPU). The tensor
|
|
# path reserves like the layer path; gate both on this.
|
|
_mtp_reserves_gpu = _mtp_will_engage and not _draft_cpu_no_embedded
|
|
_flat_mtp_reserve = (
|
|
_MTP_VRAM_RESERVE_FRAC
|
|
if (_flat_mtp_engages and not _draft_cpu_no_embedded)
|
|
else 0.0
|
|
)
|
|
_pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve
|
|
|
|
# Charge the soft overhead _CTX_FIT_VRAM_FRACTION under-covers on tight
|
|
# tiers, gated so plain dense loads (#5106) only pay the CUDA-ctx base.
|
|
# CUDA/cuBLAS context is discrete-GPU only (not Metal); the mmproj and
|
|
# MTP draft-graph buffers exist on every backend.
|
|
_soft_overhead = self._CUDA_CONTEXT_RESERVE_BYTES if gpus else 0
|
|
if effective_is_vision and mmproj_size > 0:
|
|
_soft_overhead += int(mmproj_size * (self._MMPROJ_VRAM_SAFETY - 1.0))
|
|
if _mtp_reserves_gpu:
|
|
_soft_overhead += self._MTP_DRAFT_COMPUTE_BYTES
|
|
model_size_fit = model_size + _compute_buffer_pipeline + _soft_overhead
|
|
|
|
def _subset_model_size(n_gpus: int) -> int:
|
|
return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes
|
|
|
|
# Unified-memory budget (0 off Apple Silicon) for the no-GPU Metal cap below.
|
|
_apple_budget_mib = self._apple_metal_memory_budget_bytes() // (1024 * 1024)
|
|
|
|
def _restore_after_tensor_downgrade():
|
|
# Restore the quantized KV + extras tensor dropped (layer
|
|
# split supports them), minus --split-mode.
|
|
nonlocal cache_type_kv, _cache_type_from_env, extra_args
|
|
if _tensor_dropped_cache_type_kv is not None:
|
|
cache_type_kv = _tensor_dropped_cache_type_kv
|
|
_cache_type_from_env = False
|
|
extra_args = strip_split_mode_only(
|
|
_tensor_dropped_extra_args
|
|
if _tensor_dropped_extra_args is not None
|
|
else extra_args
|
|
)
|
|
|
|
# The route fallback retry is tensor-off; keep it multi-GPU.
|
|
if preserve_multi_gpu_on_layer:
|
|
_layer_min_gpus = max(_layer_min_gpus, len(gpus))
|
|
|
|
if tensor_parallel and self._tensor_split_aborts(binary, model_identifier):
|
|
# Aborted on tensor for this model this session (#6415); skip
|
|
# tensor upfront, layer split serves it.
|
|
logger.info(
|
|
"Tensor parallelism skipped: this llama.cpp build aborted "
|
|
"on --split-mode tensor for this model earlier this "
|
|
"session; using layer split across %d GPU(s).",
|
|
len(gpus),
|
|
)
|
|
tensor_parallel = False
|
|
# Keep the multi-GPU request (gated on it, not the cache).
|
|
_layer_min_gpus = max(_layer_min_gpus, len(gpus))
|
|
_restore_after_tensor_downgrade()
|
|
|
|
# Tensor mode replicates a compute buffer on every GPU, so drop
|
|
# GPUs below that reserve from the set up front (gpu_indices
|
|
# becomes the CUDA_VISIBLE_DEVICES mask, fully excluding them).
|
|
tp_gpus = gpus
|
|
if tensor_parallel:
|
|
# Deterministic per-device compute buffer (replicated on
|
|
# every device in tensor mode); flat fallback when dims
|
|
# are unavailable. _plan_tensor_parallel uses the same.
|
|
_tp_reserve_bytes = self._estimate_compute_buffer_bytes(
|
|
n_ubatch = _effective_ubatch,
|
|
n_parallel = n_parallel,
|
|
per_device_tensor = True,
|
|
)
|
|
reserve_mib = (
|
|
_tp_reserve_bytes // (1024 * 1024)
|
|
if _tp_reserve_bytes > 0
|
|
else self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
|
|
)
|
|
# Admit by usable budget (free - (1-frac)*total), not raw
|
|
# free: a partly-used big card can clear the reserve on raw
|
|
# free yet have no budget left.
|
|
tp_gpus = [g for g in gpus if _gpu_usable(g) >= reserve_mib]
|
|
|
|
if tensor_parallel and len(tp_gpus) < 2:
|
|
# Tensor parallelism needs >= 2 usable GPUs. On a single
|
|
# GPU --split-mode tensor is a no-op; with 0 GPUs (CPU-only
|
|
# or probe failed) it must not reach llama-server; and a
|
|
# GPU below the buffer reserve can't participate. Drop the
|
|
# flag and fall through to normal layer/CPU allocation.
|
|
logger.info(
|
|
"Tensor parallelism requested but only %d of %d GPU(s) "
|
|
"have enough free VRAM for the compute buffer; "
|
|
"ignoring (needs >= 2).",
|
|
len(tp_gpus),
|
|
len(gpus),
|
|
)
|
|
tensor_parallel = False
|
|
# GPUs below tensor's compute-buffer reserve can still do layer
|
|
# split, so keep multi-GPU (mirrors the budget/geometry drops);
|
|
# _select_gpus caps unusable cards.
|
|
if len(gpus) >= 2:
|
|
_layer_min_gpus = max(_layer_min_gpus, len(gpus))
|
|
# Layer split supports a quantized KV the tensor attempt
|
|
# dropped; restore the original cache type + extras (minus
|
|
# --split-mode) so the layer launch re-emits them.
|
|
_restore_after_tensor_downgrade()
|
|
|
|
if tensor_parallel and tp_gpus:
|
|
# Pooled usable budget (after each device's compute buffer)
|
|
# must hold the non-shrinkable footprint: weights + the MTP
|
|
# reserve. The planner can shrink ctx/KV, not these.
|
|
_tp_weight_budget_mib = (
|
|
sum(_gpu_usable(g) for g in tp_gpus) - len(tp_gpus) * reserve_mib
|
|
)
|
|
_tp_flat_mtp = 2 * 1024**3 # flat reserve when dims unavailable
|
|
if not _mtp_reserves_gpu:
|
|
# No MTP, or its only drafter is CPU-offloaded (no GPU).
|
|
_tp_mtp_floor = 0
|
|
elif mtp_overhead_fn is not None and not _mtp_kv_unsized:
|
|
_tp_mtp_floor = _mtp_bytes(
|
|
min(2048, effective_ctx) if effective_ctx > 0 else 2048
|
|
)
|
|
else:
|
|
# Dims unavailable / weights-only: tensor mode has no
|
|
# --fit valve, so keep the flat reserve as the unsized-KV
|
|
# cushion, never below the known byte reserve.
|
|
_tp_mtp_floor = max(
|
|
_tp_flat_mtp,
|
|
_mtp_bytes(min(2048, effective_ctx) if effective_ctx > 0 else 2048),
|
|
)
|
|
_tp_required_mib = (model_size + _tp_mtp_floor + _soft_overhead) / (
|
|
1024 * 1024
|
|
)
|
|
if _tp_weight_budget_mib <= _tp_required_mib:
|
|
logger.info(
|
|
"Tensor parallelism requested but the pooled VRAM "
|
|
"budget cannot hold the weights, MTP reserve, and "
|
|
"per-device compute buffers; falling back to layer split."
|
|
)
|
|
tensor_parallel = False
|
|
# Weights needed >1 card, so keep multi-GPU across the
|
|
# usable tensor GPUs.
|
|
if len(tp_gpus) >= 2:
|
|
_layer_min_gpus = max(_layer_min_gpus, len(tp_gpus))
|
|
# Restore the dropped quantized KV + cache extras (minus
|
|
# --split-mode); layer split supports them.
|
|
_restore_after_tensor_downgrade()
|
|
|
|
if tensor_parallel and tp_gpus:
|
|
# Tensor-parallel allocation; see _plan_tensor_parallel.
|
|
target_ctx = (
|
|
effective_ctx
|
|
if explicit_ctx
|
|
else (self._context_length or effective_ctx)
|
|
)
|
|
# When the draft KV couldn't be sized (weights-only reserve),
|
|
# the planner's mtp_overhead_fn is non-None but covers only
|
|
# weights, so pass the flat cushion for the unsized KV (else
|
|
# the binary search spends it on context).
|
|
_tp_unsized_mtp_reserve = (
|
|
2 * 1024**3 if (_mtp_reserves_gpu and _mtp_kv_unsized) else 0
|
|
)
|
|
(
|
|
effective_ctx,
|
|
max_available_ctx,
|
|
gpu_indices,
|
|
tp_tensor_split,
|
|
) = self._plan_tensor_parallel(
|
|
tp_gpus,
|
|
model_size,
|
|
target_ctx,
|
|
cache_type_kv = cache_type_kv,
|
|
n_parallel = n_parallel,
|
|
mtp_engaged = _mtp_reserves_gpu,
|
|
mtp_overhead_fn = mtp_overhead_fn,
|
|
mtp_flat_reserve_bytes = _tp_unsized_mtp_reserve,
|
|
# Report the UI ceiling from native ctx, not the
|
|
# explicit small request.
|
|
max_target_ctx = self._context_length or target_ctx,
|
|
total_by_idx = total_by_idx,
|
|
n_ubatch = _effective_ubatch,
|
|
soft_overhead_bytes = _soft_overhead,
|
|
)
|
|
use_fit = False
|
|
elif gpus and self._can_estimate_kv() and effective_ctx > 0:
|
|
# Compute the largest hardware-aware cap from the model's
|
|
# native context across all usable GPU subsets (for UI
|
|
# bounds), independent of the currently requested context.
|
|
native_ctx_for_cap = self._context_length or effective_ctx
|
|
if native_ctx_for_cap > 0:
|
|
ranked_for_cap = sorted(
|
|
gpus,
|
|
key = lambda g: _gpu_usable(
|
|
g, _CTX_FIT_VRAM_FRACTION - _flat_mtp_reserve
|
|
),
|
|
reverse = True,
|
|
)
|
|
best_cap = 0
|
|
_cap_fraction = _CTX_FIT_VRAM_FRACTION - _flat_mtp_reserve
|
|
for n_gpus in range(1, len(ranked_for_cap) + 1):
|
|
subset = ranked_for_cap[:n_gpus]
|
|
# Per-GPU-consistent pool budget (fixes mixed
|
|
# known/unknown totals); pass it as an absolute
|
|
# budget so the fit and the check below agree.
|
|
pool_budget = _pool_budget_mib(subset, _cap_fraction)
|
|
_ms = _subset_model_size(n_gpus)
|
|
# Compute buffer is replicated per device in a layer
|
|
# split, so scale the context term by the subset size.
|
|
_cc_sub = lambda c, n = n_gpus: _cc_bytes(c, n)
|
|
capped = self._fit_context_to_vram(
|
|
native_ctx_for_cap,
|
|
pool_budget,
|
|
_ms,
|
|
cache_type_kv,
|
|
n_parallel = n_parallel,
|
|
mtp_engaged = _mtp_reserves_gpu,
|
|
mtp_overhead_fn = mtp_overhead_fn,
|
|
compute_ctx_bytes_fn = _cc_sub,
|
|
budget_frac = 1.0,
|
|
total_mib = None,
|
|
)
|
|
kv = self._estimate_kv_cache_bytes(
|
|
capped, cache_type_kv, n_parallel = n_parallel
|
|
)
|
|
footprint_mib = (
|
|
_ms + kv + _mtp_bytes(capped) + _cc_sub(capped)
|
|
) / (1024 * 1024)
|
|
if footprint_mib <= pool_budget:
|
|
best_cap = max(best_cap, capped)
|
|
if best_cap > 0:
|
|
max_available_ctx = best_cap
|
|
else:
|
|
# Weights exceed 90% of every GPU subset, so no
|
|
# context fits. Anchor the UI "safe zone" at 4096
|
|
# so the slider warns above the fallback.
|
|
max_available_ctx = min(4096, native_ctx_for_cap)
|
|
|
|
if explicit_ctx:
|
|
# Honor the requested context verbatim. If it fits,
|
|
# pin GPUs and skip --fit; else ship -c <ctx> --fit
|
|
# on and let llama-server flex -ngl (CPU offload).
|
|
requested_total = (
|
|
model_size_fit
|
|
+ self._estimate_kv_cache_bytes(
|
|
effective_ctx, cache_type_kv, n_parallel = n_parallel
|
|
)
|
|
+ _mtp_bytes(effective_ctx)
|
|
+ _cc_bytes(effective_ctx)
|
|
)
|
|
# The compute buffer is replicated on every device in a
|
|
# layer split; fold it into the per-device reserve so a
|
|
# multi-GPU pin sizes each card for its own copy.
|
|
gpu_indices, use_fit = self._select_gpus(
|
|
requested_total,
|
|
gpus,
|
|
usable_fraction = _pin_fraction,
|
|
total_by_idx = total_by_idx,
|
|
per_device_overhead_bytes = _pipeline_overhead_bytes
|
|
+ _cc_bytes(effective_ctx),
|
|
min_gpus = _layer_min_gpus,
|
|
)
|
|
# No silent shrink: effective_ctx stays == requested_ctx.
|
|
else:
|
|
# Auto context: prefer fewer GPUs, cap to fit. Same
|
|
# headroom threshold as _select_gpus (#5106). Rank by the
|
|
# active pin fraction so the order matches the fit budget.
|
|
pin_fraction = _pin_fraction
|
|
ranked = sorted(
|
|
gpus, key = lambda g: _gpu_usable(g, pin_fraction), reverse = True
|
|
)
|
|
# Skips _select_gpus, so apply its cap: count only cards
|
|
# whose usable VRAM clears the per-device layer overhead.
|
|
_pipeline_overhead_mib = _pipeline_overhead_bytes / (1024 * 1024)
|
|
_auto_min_gpus = max(
|
|
1,
|
|
min(
|
|
_layer_min_gpus,
|
|
sum(
|
|
1
|
|
for g in ranked
|
|
if _gpu_usable(g, pin_fraction) > _pipeline_overhead_mib
|
|
)
|
|
or 1,
|
|
),
|
|
)
|
|
for n_gpus in range(_auto_min_gpus, len(ranked) + 1):
|
|
subset = ranked[:n_gpus]
|
|
pool_budget = _pool_budget_mib(subset, pin_fraction)
|
|
_ms = _subset_model_size(n_gpus)
|
|
# Compute buffer is replicated per device in a layer
|
|
# split, so scale the context term by the subset size.
|
|
_cc_sub = lambda c, n = n_gpus: _cc_bytes(c, n)
|
|
capped = self._fit_context_to_vram(
|
|
effective_ctx,
|
|
pool_budget,
|
|
_ms,
|
|
cache_type_kv,
|
|
n_parallel = n_parallel,
|
|
mtp_engaged = _mtp_reserves_gpu,
|
|
mtp_overhead_fn = mtp_overhead_fn,
|
|
compute_ctx_bytes_fn = _cc_sub,
|
|
budget_frac = 1.0,
|
|
total_mib = None,
|
|
)
|
|
kv = self._estimate_kv_cache_bytes(
|
|
capped, cache_type_kv, n_parallel = n_parallel
|
|
)
|
|
footprint_mib = (
|
|
_ms + kv + _mtp_bytes(capped) + _cc_sub(capped)
|
|
) / (1024 * 1024)
|
|
if footprint_mib <= pool_budget:
|
|
effective_ctx = capped
|
|
gpu_indices = sorted(idx for idx, _ in subset)
|
|
use_fit = False
|
|
break
|
|
else:
|
|
# Native ctx doesn't fit. Drop to 4096 and
|
|
# re-check before --fit on: a model overflowing
|
|
# at 131k may pin fine with a 4096 KV (#5106).
|
|
effective_ctx = min(4096, effective_ctx)
|
|
if effective_ctx > 0:
|
|
for n_gpus in range(_auto_min_gpus, len(ranked) + 1):
|
|
subset = ranked[:n_gpus]
|
|
kv = self._estimate_kv_cache_bytes(
|
|
effective_ctx,
|
|
cache_type_kv,
|
|
n_parallel = n_parallel,
|
|
)
|
|
footprint_mib = (
|
|
_subset_model_size(n_gpus)
|
|
+ kv
|
|
+ _mtp_bytes(effective_ctx)
|
|
+ _cc_bytes(effective_ctx, n_gpus)
|
|
) / (1024 * 1024)
|
|
if footprint_mib <= _pool_budget_mib(subset, pin_fraction):
|
|
gpu_indices = sorted(idx for idx, _ in subset)
|
|
use_fit = False
|
|
break
|
|
|
|
elif gpus:
|
|
# Can't estimate KV -- file-size-only check; keep the
|
|
# ceiling at native context (already the default).
|
|
logger.debug(
|
|
"Falling back to file-size-only GPU selection",
|
|
model_size_gb = round(model_size / (1024**3), 2),
|
|
)
|
|
# Add the byte-accurate MTP reserve here too when it is
|
|
# available; otherwise _pin_fraction carries the flat
|
|
# fallback (the two are mutually exclusive by design).
|
|
_fs_total = model_size_fit + _mtp_bytes(
|
|
self._context_length or effective_ctx or 4096
|
|
)
|
|
gpu_indices, use_fit = self._select_gpus(
|
|
_fs_total,
|
|
gpus,
|
|
usable_fraction = _pin_fraction,
|
|
total_by_idx = total_by_idx,
|
|
per_device_overhead_bytes = _pipeline_overhead_bytes,
|
|
min_gpus = _layer_min_gpus,
|
|
)
|
|
if use_fit and not explicit_ctx:
|
|
# Weights don't fit on any subset; default UI to 4096
|
|
# so the slider isn't on an unusable native ctx.
|
|
effective_ctx = min(4096, effective_ctx) if effective_ctx > 0 else 4096
|
|
|
|
elif _apple_budget_mib > 0 and effective_ctx > 0:
|
|
# No GPU on Metal: the branches above are skipped and the context
|
|
# stays at native, over-committing unified memory (#5118, #6529).
|
|
# Cap with the same fit math (--fit on stays as a backstop); only
|
|
# auto context shrinks, explicit is honored.
|
|
native_ctx_for_cap = self._context_length or effective_ctx
|
|
# Reserve the flat MTP fraction up front like the discrete
|
|
# _pin_fraction, so an unsized MTP draft (e.g. Qwen3.6-MTP, #6529)
|
|
# can't over-commit. No-op when MTP is off; exclusive with the
|
|
# byte-accurate _mtp_bytes reserve.
|
|
_apple_fit_budget_mib = int(
|
|
_apple_budget_mib * max(0.0, 1.0 - _flat_mtp_reserve)
|
|
)
|
|
if self._can_estimate_kv():
|
|
cap = self._fit_context_to_vram(
|
|
native_ctx_for_cap,
|
|
_apple_fit_budget_mib,
|
|
model_size_fit,
|
|
cache_type_kv,
|
|
n_parallel = n_parallel,
|
|
mtp_engaged = _mtp_reserves_gpu,
|
|
mtp_overhead_fn = mtp_overhead_fn,
|
|
compute_ctx_bytes_fn = _cc_bytes,
|
|
budget_frac = 1.0,
|
|
total_mib = None,
|
|
)
|
|
_cap_footprint_mib = (
|
|
model_size_fit
|
|
+ self._estimate_kv_cache_bytes(
|
|
cap, cache_type_kv, n_parallel = n_parallel
|
|
)
|
|
+ _mtp_bytes(cap)
|
|
+ _cc_bytes(cap)
|
|
) / (1024 * 1024)
|
|
# Fit returns the request unchanged when it fits OR weights
|
|
# exceed budget; only the latter over-commits, so floor to 4096.
|
|
max_available_ctx = (
|
|
cap
|
|
if _cap_footprint_mib <= _apple_fit_budget_mib
|
|
else min(4096, native_ctx_for_cap)
|
|
)
|
|
else:
|
|
# No KV estimate: mirror the discrete file-size-only fallback
|
|
# and floor to 4096 rather than launch at native and over-commit.
|
|
max_available_ctx = min(4096, native_ctx_for_cap)
|
|
if not explicit_ctx:
|
|
effective_ctx = max_available_ctx
|
|
|
|
# Prefer fewer serving slots on GPU over --fit on offload: when the extra
|
|
# --parallel slots push the footprint past the pin budget, llama-server
|
|
# offloads layers to host and decode collapses ~3x (#6718). Retry the fit
|
|
# at fewer slots, keeping the largest count that stays fully on GPU and the
|
|
# chosen context. Skips tensor mode / Metal / KV-inestimable paths.
|
|
if (
|
|
use_fit
|
|
and n_parallel > 1
|
|
and gpus
|
|
and self._can_estimate_kv()
|
|
and effective_ctx > 0
|
|
):
|
|
# Slot-independent footprint (folded compute buffer swapped out so the
|
|
# helper re-adds a slot-sized one per candidate).
|
|
_base_footprint = (
|
|
model_size_fit
|
|
- _compute_buffer_pipeline
|
|
+ _mtp_bytes(effective_ctx)
|
|
+ _cc_bytes(effective_ctx)
|
|
)
|
|
_gi_slots, _uf_slots, _slots = self._slots_that_fit_on_gpu(
|
|
n_parallel,
|
|
effective_ctx,
|
|
gpus,
|
|
total_by_idx,
|
|
_base_footprint,
|
|
cache_type_kv,
|
|
_pin_fraction,
|
|
_pipeline_overhead_bytes + _cc_bytes(effective_ctx),
|
|
_layer_min_gpus,
|
|
_effective_ubatch,
|
|
)
|
|
if not _uf_slots:
|
|
logger.info(
|
|
"Serving slots reduced %d -> %d to keep the model on GPU "
|
|
"(avoid --fit offload) at context %d.",
|
|
n_parallel,
|
|
_slots,
|
|
effective_ctx,
|
|
)
|
|
gpu_indices, use_fit, n_parallel = _gi_slots, False, _slots
|
|
|
|
# MTP reserve at the final context, for the logs below.
|
|
_mtp_reserve_bytes = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0
|
|
if _mtp_will_engage:
|
|
_mtp_note = (
|
|
f"MTP reserve: {_mtp_reserve_bytes / (1024**3):.2f} GB "
|
|
f"(draft KV @ {effective_ctx} + verify n_max={_mtp_eff_n_max}"
|
|
+ (", flat-frac fallback" if mtp_overhead_fn is None else "")
|
|
+ "), "
|
|
)
|
|
else:
|
|
_mtp_note = ""
|
|
|
|
if effective_ctx < original_ctx:
|
|
kv_est = self._estimate_kv_cache_bytes(
|
|
effective_ctx, cache_type_kv, n_parallel = n_parallel
|
|
)
|
|
logger.info(
|
|
f"Context auto-reduced: {original_ctx} -> {effective_ctx} "
|
|
f"(model: {model_size / (1024**3):.1f} GB, "
|
|
f"est. KV cache: {kv_est / (1024**3):.1f} GB, "
|
|
f"{_mtp_note}".rstrip(", ")
|
|
+ ")"
|
|
)
|
|
|
|
kv_cache_bytes = self._estimate_kv_cache_bytes(
|
|
effective_ctx, cache_type_kv, n_parallel = n_parallel
|
|
)
|
|
mmproj_note = (
|
|
f"mmproj: {mmproj_size / (1024**3):.1f} GB, " if mmproj_size else ""
|
|
)
|
|
logger.info(
|
|
f"GGUF size: {gguf_size / (1024**3):.1f} GB, "
|
|
f"{mmproj_note}"
|
|
f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, "
|
|
f"{_mtp_note}"
|
|
f"context: {effective_ctx}, "
|
|
f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}"
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"GPU selection failed ({e}), using --fit on")
|
|
gpu_indices, use_fit = None, True
|
|
tp_tensor_split = None
|
|
effective_ctx = requested_ctx # fall back to original
|
|
|
|
# Unified-memory APUs load weights into system RAM (under WSL the VM
|
|
# cap, not the ROCm-reported VRAM, is the real ceiling); refuse an
|
|
# oversize load the OS would otherwise kill mid-flight. Base model
|
|
# only: an optional MTP drafter is dropped by the MTP-drop fallback.
|
|
if model_size is not None and self._amd_apu_wants_unified_memory(gpu_indices):
|
|
_ram_msg = self._apu_ram_shortfall_message(
|
|
model_size, self._available_system_memory_mib()
|
|
)
|
|
if _ram_msg:
|
|
raise RuntimeError(_ram_msg)
|
|
|
|
# Audio input straight from the mmproj (clip.has_audio_encoder),
|
|
# independent of token names.
|
|
self._mmproj_has_audio = False
|
|
if launch_mmproj_path:
|
|
try:
|
|
from utils.models.gguf_metadata import (
|
|
read_mmproj_audio_capability,
|
|
)
|
|
self._mmproj_has_audio = bool(
|
|
read_mmproj_audio_capability(launch_mmproj_path)
|
|
)
|
|
except Exception as e:
|
|
logger.debug(f"mmproj audio-capability read failed: {e}")
|
|
|
|
cmd = [
|
|
binary,
|
|
"-m",
|
|
model_path,
|
|
"--port",
|
|
str(self._port),
|
|
"-c",
|
|
str(effective_ctx) if effective_ctx > 0 else "0",
|
|
"--parallel",
|
|
str(n_parallel),
|
|
"--flash-attn",
|
|
"on", # Force flash attention for speed
|
|
# Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length".
|
|
"--no-context-shift",
|
|
]
|
|
|
|
# Report a clean public model id (matching GET /v1/models) rather
|
|
# than the raw -m path in llama-server's own /v1/models and the
|
|
# "model" field of its chat/completions responses.
|
|
from core.inference.model_ids import public_model_id
|
|
|
|
_alias = public_model_id(self._model_identifier or model_path)
|
|
if _alias:
|
|
cmd.extend(["--alias", _alias])
|
|
|
|
fully_gpu_offloaded = False
|
|
if use_fit:
|
|
cmd.extend(["--fit", "on"])
|
|
elif gpu_indices is not None:
|
|
# Fits on selected GPU(s) -- force all layers on GPU. --fit off is
|
|
# required: without it llama.cpp's default --fit on second-guesses
|
|
# and offloads ~1 GB at --parallel 4 even though the model fits.
|
|
cmd.extend(["-ngl", "-1", "--fit", "off"])
|
|
fully_gpu_offloaded = True
|
|
|
|
server_caps = self.probe_server_capabilities(binary)
|
|
# Expose Prometheus /metrics for the engine-stats logger, only
|
|
# when the binary advertises it (older/custom binaries may not).
|
|
if server_caps.get("supports_metrics"):
|
|
cmd.append("--metrics")
|
|
cmd.extend(
|
|
self._ctx_integrity_flags(
|
|
n_parallel,
|
|
use_fit,
|
|
requested_ctx,
|
|
effective_ctx,
|
|
server_caps,
|
|
)
|
|
)
|
|
offload_overridden = _extra_args_set_any_flag(
|
|
extra_args, _GPU_OFFLOAD_OVERRIDE_FLAGS
|
|
)
|
|
threads_overridden = _extra_args_set_any_flag(extra_args, _THREAD_OVERRIDE_FLAGS)
|
|
full_offload_tuning_active = fully_gpu_offloaded and not offload_overridden
|
|
|
|
# Thread count: an unset --threads makes llama.cpp pick physical
|
|
# cores (common_cpu_get_num_math), but an explicit --threads -1
|
|
# resolves to hardware_concurrency() (every hyperthread), which
|
|
# contends on the memory bus and slows CPU / hybrid decode. So
|
|
# omit the flag when unset and only pin it for an explicit
|
|
# override or the Windows full-offload OpenMP cap. Pass-through
|
|
# thread flags in extra_args still win (appended last). #5692
|
|
if (
|
|
sys.platform == "win32"
|
|
and full_offload_tuning_active
|
|
and not threads_overridden
|
|
):
|
|
cmd.extend(["--threads", "2"])
|
|
elif n_threads is not None and n_threads > 0:
|
|
cmd.extend(["--threads", str(n_threads)])
|
|
|
|
# Enable Jinja chat template rendering
|
|
cmd.extend(["--jinja"])
|
|
|
|
# KV cache data type
|
|
_valid_cache_types = {
|
|
"f16",
|
|
"bf16",
|
|
"q8_0",
|
|
"q4_0",
|
|
"q4_1",
|
|
"q5_0",
|
|
"q5_1",
|
|
"iq4_nl",
|
|
"f32",
|
|
}
|
|
if (
|
|
cache_type_kv
|
|
and cache_type_kv in _valid_cache_types
|
|
and not _cache_type_from_env
|
|
):
|
|
cmd.extend(
|
|
[
|
|
"--cache-type-k",
|
|
cache_type_kv,
|
|
"--cache-type-v",
|
|
cache_type_kv,
|
|
]
|
|
)
|
|
self._cache_type_kv = cache_type_kv
|
|
logger.info(f"KV cache type: {cache_type_kv}")
|
|
else:
|
|
# An env-only type is left inherited (untouched) so an
|
|
# asymmetric K/V env reaches the child as set.
|
|
self._cache_type_kv = None
|
|
|
|
# Tensor parallelism: split the model across GPUs by tensor
|
|
# rather than by layer. Multi-GPU only -- a no-op on a single
|
|
# GPU. Default (layer split) is left implicit by omitting the
|
|
# flag. See llama.cpp --split-mode.
|
|
if tensor_parallel:
|
|
cmd.extend(["--split-mode", "tensor"])
|
|
if tp_tensor_split and len(tp_tensor_split) > 1:
|
|
cmd.extend(
|
|
[
|
|
"--tensor-split",
|
|
",".join(str(int(x)) for x in tp_tensor_split),
|
|
]
|
|
)
|
|
self._tensor_parallel = True
|
|
self._layer_preserves_tensor_intent = False
|
|
logger.info(
|
|
"Tensor parallelism: --split-mode tensor, --tensor-split %s",
|
|
tp_tensor_split,
|
|
)
|
|
else:
|
|
self._tensor_parallel = False
|
|
# > 1 only when a tensor request was downgraded but kept multi-GPU.
|
|
self._layer_preserves_tensor_intent = _layer_min_gpus > 1
|
|
|
|
# Speculative decoding. See _build_speculative_flags for the
|
|
# mode resolution, benchmarks, and llama.cpp references.
|
|
launch_mtp_draft_path = self._resolve_launch_mtp_path(
|
|
mtp_draft_path = mtp_draft_path,
|
|
)
|
|
spec_flags = self._build_speculative_flags(
|
|
speculative_type = speculative_type,
|
|
spec_draft_n_max = spec_draft_n_max,
|
|
extra_args = extra_args,
|
|
model_identifier = model_identifier,
|
|
model_path = model_path,
|
|
gpus = bool(gpus),
|
|
binary = binary,
|
|
mtp_draft_path = launch_mtp_draft_path,
|
|
)
|
|
# Remember where the spec block sits so a drafter-load failure
|
|
# can be retried with these flags swapped out (see below).
|
|
_spec_start = len(cmd)
|
|
cmd.extend(spec_flags)
|
|
|
|
# Apply custom chat template override if provided.
|
|
self._chat_template_override = chat_template_override
|
|
if chat_template_override:
|
|
import tempfile
|
|
|
|
flags = detect_reasoning_flags(
|
|
chat_template_override,
|
|
self._model_identifier,
|
|
log_source = "GGUF chat template override",
|
|
)
|
|
self._supports_reasoning = flags["supports_reasoning"]
|
|
self._reasoning_style = flags["reasoning_style"]
|
|
self._reasoning_effort_levels = flags.get("reasoning_effort_levels", [])
|
|
self._reasoning_always_on = flags["reasoning_always_on"]
|
|
self._supports_preserve_thinking = flags["supports_preserve_thinking"]
|
|
self._supports_tools = flags["supports_tools"]
|
|
|
|
self._chat_template_file = tempfile.NamedTemporaryFile(
|
|
mode = "w",
|
|
encoding = "utf-8",
|
|
suffix = ".jinja",
|
|
delete = False,
|
|
prefix = "unsloth_chat_template_",
|
|
)
|
|
self._chat_template_file.write(chat_template_override)
|
|
self._chat_template_file.close()
|
|
cmd.extend(["--chat-template-file", self._chat_template_file.name])
|
|
logger.info(f"Using custom chat template file: {self._chat_template_file.name}")
|
|
|
|
# Default thinking mode for reasoning models. Qwen3.5/3.6 below
|
|
# 9B disable thinking by default; 9B+ enable it. Always-on
|
|
# templates ignore the kwarg, so skip.
|
|
if self._supports_reasoning and not self._reasoning_always_on:
|
|
thinking_default = True
|
|
mid = (model_identifier or "").lower()
|
|
if "qwen3.5" in mid or "qwen3.6" in mid:
|
|
size_val = _extract_model_size_b(mid)
|
|
if size_val is not None and size_val < 9:
|
|
thinking_default = False
|
|
self._reasoning_default = thinking_default
|
|
reasoning_kw = self._reasoning_kwargs(thinking_default)
|
|
# preserve_thinking is an independent kwarg. Default it OFF
|
|
# at launch so direct OpenAI-compatible callers that omit the
|
|
# field match the UI's default-off behavior (the bundled
|
|
# gemma-4 template also defaults it false; the frontend sends
|
|
# preserve_thinking per request once toggled on).
|
|
if self._supports_preserve_thinking:
|
|
reasoning_kw["preserve_thinking"] = False
|
|
cmd.extend(
|
|
[
|
|
"--chat-template-kwargs",
|
|
json.dumps(reasoning_kw),
|
|
]
|
|
)
|
|
logger.info(f"Reasoning model: {reasoning_kw} by default")
|
|
|
|
if launch_mmproj_path and effective_is_vision:
|
|
cmd.extend(["--mmproj", launch_mmproj_path])
|
|
logger.info(f"Using mmproj for vision: {launch_mmproj_path}")
|
|
|
|
# Option C: --api-key for direct client access when enabled
|
|
import secrets as _secrets
|
|
|
|
if os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1":
|
|
self._api_key = _secrets.token_urlsafe(32)
|
|
cmd.extend(["--api-key", self._api_key])
|
|
logger.info("llama-server started with --api-key for direct streaming")
|
|
else:
|
|
self._api_key = None
|
|
|
|
# Windows + full offload: disable KV checkpoints (WDDM/PCI-E
|
|
# overhead). CPU/partial offload keeps prompt caching. #5692.
|
|
if sys.platform == "win32" and full_offload_tuning_active:
|
|
unsupported_cache_flags: list[str] = []
|
|
if server_caps.get("supports_cache_ram"):
|
|
cmd.extend(["--cache-ram", "0"])
|
|
else:
|
|
unsupported_cache_flags.append("--cache-ram")
|
|
if server_caps.get("supports_ctx_checkpoints"):
|
|
cmd.extend(["--ctx-checkpoints", "0"])
|
|
else:
|
|
unsupported_cache_flags.append("--ctx-checkpoints")
|
|
if server_caps.get("supports_no_cache_prompt"):
|
|
cmd.append("--no-cache-prompt")
|
|
else:
|
|
unsupported_cache_flags.append("--no-cache-prompt")
|
|
if unsupported_cache_flags:
|
|
logger.info(
|
|
"Skipping unsupported Windows cache flags for llama-server: %s",
|
|
", ".join(unsupported_cache_flags),
|
|
)
|
|
|
|
# User pass-through args go last so llama.cpp's last-wins parsing
|
|
# lets the user override Studio's auto-set flags. Already
|
|
# validated by the route via validate_extra_args().
|
|
if extra_args:
|
|
cmd.extend(str(a) for a in extra_args)
|
|
logger.info(f"Appending user extra args to llama-server: {list(extra_args)}")
|
|
|
|
logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}")
|
|
|
|
# Library paths so llama-server finds its shared libs and CUDA DLLs.
|
|
env = self._llama_server_env_for_binary(binary)
|
|
# Omitting --threads relies on llama.cpp's physical-core default, so
|
|
# drop an inherited LLAMA_ARG_THREADS that would otherwise feed the
|
|
# arg handler and silently force hardware_concurrency(). #5692
|
|
if "--threads" not in cmd:
|
|
env.pop("LLAMA_ARG_THREADS", None)
|
|
|
|
# Reconcile the inherited LLAMA_ARG_* env with Studio's final
|
|
# decision: stripping CLI extras on a tensor->layer downgrade
|
|
# can't remove env vars, so the child could run a mode/KV Studio
|
|
# didn't budget.
|
|
if not tensor_parallel:
|
|
# Layer split: clear a non-layer inherited split mode (and any
|
|
# paired tensor-split) so the child can't override the layer plan.
|
|
_inherited_sm = (env.get("LLAMA_ARG_SPLIT_MODE") or "").strip().lower()
|
|
if _inherited_sm and _inherited_sm != "layer":
|
|
env.pop("LLAMA_ARG_SPLIT_MODE", None)
|
|
env.pop("LLAMA_ARG_TENSOR_SPLIT", None)
|
|
else:
|
|
# Studio owns the tensor split: it emits --tensor-split when it
|
|
# picks an uneven one (CLI wins) and nothing when an even split
|
|
# is safe. Clear any inherited LLAMA_ARG_TENSOR_SPLIT so the even
|
|
# case can't be overridden by a stale env (the layer branch above
|
|
# clears it too).
|
|
env.pop("LLAMA_ARG_TENSOR_SPLIT", None)
|
|
# Tensor split aborts on a quantized KV; clear an inherited
|
|
# quantized cache type so the child uses the tensor-safe default.
|
|
for _ct_var in ("LLAMA_ARG_CACHE_TYPE_K", "LLAMA_ARG_CACHE_TYPE_V"):
|
|
_ct_raw = (env.get(_ct_var) or "").strip().lower()
|
|
if _ct_raw and _ct_raw not in self._TENSOR_PARALLEL_KV_TYPES:
|
|
env.pop(_ct_var, None)
|
|
|
|
# Windows + full offload: PASSIVE OMP + 2 threads stop
|
|
# spin-wait burning CPU. CPU/partial offload keeps default
|
|
# OMP parallelism. #5692.
|
|
if sys.platform == "win32" and full_offload_tuning_active:
|
|
env.setdefault("OMP_WAIT_POLICY", "PASSIVE")
|
|
if not threads_overridden:
|
|
env.setdefault("OMP_NUM_THREADS", "2")
|
|
|
|
# AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use
|
|
# shared system RAM. setdefault so a user value wins.
|
|
if self._amd_apu_wants_unified_memory(gpu_indices):
|
|
env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1")
|
|
logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1")
|
|
|
|
# DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU).
|
|
# See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1.
|
|
if self._apply_datacenter_env(env, gpu_indices):
|
|
multi_gpu = self._effective_gpu_count(gpu_indices) > 1
|
|
logger.info(
|
|
f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})"
|
|
)
|
|
|
|
# Pin to selected GPU(s). On ROCm, narrowing only
|
|
# CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full
|
|
# set, so set HIP_VISIBLE_DEVICES too.
|
|
if gpu_indices is not None:
|
|
pinned = ",".join(str(i) for i in gpu_indices)
|
|
env["CUDA_VISIBLE_DEVICES"] = pinned
|
|
try:
|
|
import torch as _torch
|
|
if getattr(_torch.version, "hip", None) is not None:
|
|
env["HIP_VISIBLE_DEVICES"] = pinned
|
|
# Do NOT also set ROCR_VISIBLE_DEVICES to the same
|
|
# value. ROCR_VISIBLE_DEVICES filters at the HSA/ROCr
|
|
# layer and HIP_VISIBLE_DEVICES at the HIP layer, so
|
|
# setting both with the same physical indices applies
|
|
# the mask twice: ROCR reduces the visible set and
|
|
# re-indexes it from 0, then HIP indexes into the
|
|
# already-reduced set. A single non-zero pin (e.g.
|
|
# "1") then points out of range at the HIP layer, HIP
|
|
# enumerates 0 devices, and llama.cpp falls back to
|
|
# CPU ("ggml_cuda_init: no ROCm-capable device is
|
|
# detected"). The HIP mask alone narrows correctly;
|
|
# clear any inherited ROCR mask so it can't double up.
|
|
env.pop("ROCR_VISIBLE_DEVICES", None)
|
|
except Exception as e:
|
|
logger.debug("Failed to set ROCm visibility env vars for child: %s", e)
|
|
|
|
# Captured before any text-only fallback strips it from cmd.
|
|
launched_with_mmproj = "--mmproj" in cmd
|
|
|
|
# One-shot --fit off retry: recent llama.cpp runs a "fitting
|
|
# params to device memory" step by default (--fit defaults to
|
|
# 'on') even when -ngl is explicit. That step has aborted on
|
|
# some ROCm hosts (ggml-cuda.cu ROCm error during worst-case
|
|
# estimation, e.g. MTP + mmproj models on gfx1151). When
|
|
# Studio's own VRAM math already placed the model
|
|
# (use_fit=False), the step is redundant second-guessing --
|
|
# retry once with --fit off before declaring the load failed.
|
|
# Never retry when fit was requested (use_fit) or the caller
|
|
# passed an explicit fit flag via extra args.
|
|
# Argv actually launched (post --fit off / MTP); text-only retry strips this.
|
|
_last_spawn_cmd = list(cmd)
|
|
|
|
def _spawn_and_wait(run_cmd, *, label = ""):
|
|
"""Start llama-server with run_cmd and wait for health.
|
|
|
|
Retries once with --fit off when the first attempt
|
|
crashes during startup and run_cmd is eligible (see
|
|
_fit_off_retry_eligible).
|
|
"""
|
|
nonlocal _last_spawn_cmd
|
|
_fit_retry_allowed = self._fit_off_retry_eligible(run_cmd, use_fit)
|
|
for _spawn_attempt in (0, 1):
|
|
# Defensive kill: drop an orphan Popen a concurrent load may
|
|
# have stored before we overwrite the reference (#5161).
|
|
# Also reaps the crashed first attempt on the retry pass.
|
|
self._kill_process()
|
|
|
|
self._stdout_lines = []
|
|
# Tee llama-server output to a dedicated log file so a
|
|
# post-mortem has the full trail even when the parent only
|
|
# kept the last 50 lines. Path is under the studio home.
|
|
# ``label`` (MTP fallback) and the attempt index (--fit
|
|
# off retry) keep a respawn within the same epoch second
|
|
# from truncating the crash log a retry warning just
|
|
# pointed the user at.
|
|
self._llama_log_fh = None
|
|
try:
|
|
log_dir = _swa_cache_path().parent / "logs" / "llama-server"
|
|
log_dir.mkdir(parents = True, exist_ok = True)
|
|
self._llama_log_path = log_dir / (
|
|
f"llama-{int(time.time())}{label}-port-{self._port}"
|
|
f"-try{_spawn_attempt}.log"
|
|
)
|
|
self._llama_log_fh = open(
|
|
self._llama_log_path,
|
|
"w",
|
|
encoding = "utf-8",
|
|
buffering = 1,
|
|
)
|
|
logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}")
|
|
except OSError as e:
|
|
# Best-effort; never block the load on logging.
|
|
logger.debug(f"Could not open llama-server log file: {e}")
|
|
self._llama_log_path = None
|
|
_last_spawn_cmd = list(run_cmd)
|
|
self._process = subprocess.Popen(
|
|
run_cmd,
|
|
stdout = subprocess.PIPE,
|
|
stderr = subprocess.STDOUT,
|
|
text = True,
|
|
env = env,
|
|
**_windows_hidden_subprocess_kwargs(),
|
|
**_child_popen_kwargs(),
|
|
)
|
|
self._record_server_pid(self._process.pid)
|
|
|
|
# Background thread to drain stdout (prevents pipe deadlock)
|
|
self._stdout_thread = threading.Thread(
|
|
target = self._drain_stdout, daemon = True, name = "llama-stdout"
|
|
)
|
|
self._stdout_thread.start()
|
|
if self._wait_for_health(timeout = 600.0):
|
|
return True
|
|
_startup_crashed = (
|
|
self._process.poll() is not None and self._process.returncode != 0
|
|
)
|
|
# A split-axis abort (#6415) is fit-independent: skip the
|
|
# --fit off retry and let the caller latch it.
|
|
_split_axis_crash = self._is_tensor_split_assert(
|
|
"\n".join(self._stdout_lines[-50:])
|
|
)
|
|
if (
|
|
_spawn_attempt == 0
|
|
and fully_gpu_offloaded
|
|
and _startup_crashed
|
|
and not _split_axis_crash
|
|
):
|
|
# We forced --fit off because Studio's (conservative) VRAM
|
|
# math placed the model fully on GPU. A startup crash here
|
|
# means that estimate was optimistic, so fall back to --fit
|
|
# on and let llama.cpp offload rather than fail the load.
|
|
logger.warning(
|
|
"llama-server crashed during startup (exit code %s) "
|
|
"with forced --fit off; the fit estimate was optimistic, "
|
|
"retrying once with --fit on so it can offload. "
|
|
"Crash log: %s",
|
|
self._process.returncode,
|
|
self._llama_log_path,
|
|
)
|
|
# Flip Studio's own --fit off (added first, before any
|
|
# user extra args) to on; a user's later --fit still wins
|
|
# by last-arg. Defensive: if absent, the default is already
|
|
# --fit on, so leave it.
|
|
_run = list(run_cmd)
|
|
if "--fit" in _run:
|
|
_run[_run.index("--fit") + 1] = "on"
|
|
run_cmd = _run
|
|
continue
|
|
if (
|
|
_spawn_attempt == 0
|
|
and _fit_retry_allowed
|
|
and _startup_crashed
|
|
and not _split_axis_crash
|
|
):
|
|
logger.warning(
|
|
"llama-server crashed during startup (exit code %s) "
|
|
"with the default memory-fit step enabled; Studio "
|
|
"already verified the model fits, retrying once "
|
|
"with --fit off. Crash log: %s",
|
|
self._process.returncode,
|
|
self._llama_log_path,
|
|
)
|
|
run_cmd = [*run_cmd, "--fit", "off"]
|
|
continue
|
|
return False
|
|
|
|
# Store the resolved on-disk path, not the caller's kwarg: in
|
|
# HF mode gguf_path is None and ``model_path`` is what
|
|
# llama-server mmap's, which downstream consumers need. Must be
|
|
# set BEFORE the spawn: load_progress() reads _gguf_path for
|
|
# the mmap progress total while the health wait runs.
|
|
self._gguf_path = model_path
|
|
self._hf_repo = hf_repo
|
|
self._mtp_draft_path = launch_mtp_draft_path
|
|
# For local GGUF files, extract variant from filename if absent
|
|
if hf_variant:
|
|
self._hf_variant = hf_variant
|
|
elif gguf_path:
|
|
try:
|
|
from utils.models.model_config import _extract_quant_label
|
|
self._hf_variant = _extract_quant_label(gguf_path)
|
|
except Exception:
|
|
self._hf_variant = None
|
|
else:
|
|
self._hf_variant = None
|
|
self._is_vision = effective_is_vision
|
|
self._model_identifier = model_identifier
|
|
|
|
# Store the effective (possibly capped) context separately; do
|
|
# NOT overwrite _context_length (the native length for display).
|
|
self._effective_context_length = (
|
|
effective_ctx if effective_ctx > 0 else self._context_length
|
|
)
|
|
self._reconcile_effective_ctx_with_server()
|
|
self._max_context_length = (
|
|
max_available_ctx if max_available_ctx > 0 else self._effective_context_length
|
|
)
|
|
|
|
healthy = _spawn_and_wait(cmd)
|
|
# #6415 split-mode tensor warmup abort. Latch it on THIS first spawn:
|
|
# the flash-attn-off retry below can't run tensor (needs flash_attn),
|
|
# so its output drops the marker and recording later would miss it,
|
|
# looping every load. Record and raise to the route's layer fallback,
|
|
# skipping the futile flash-attn/MTP retries.
|
|
if not healthy and self._tensor_parallel and not self._cancel_event.is_set():
|
|
_ts_out = "\n".join(self._stdout_lines[-50:])
|
|
_ts_rc = self._process.poll() if self._process is not None else None
|
|
if self._should_record_tensor_split_abort(_ts_rc, _ts_out):
|
|
LlamaCppBackend._record_tensor_split_abort(binary, model_identifier)
|
|
self._kill_process()
|
|
raise RuntimeError(
|
|
"llama-server aborted on --split-mode tensor "
|
|
"(split-axis geometry); retrying with layer split."
|
|
)
|
|
# Flash-attention kernels hard-crash at startup on some ROCm/GPU
|
|
# builds (frequently inside the vision tower). Disabling FA keeps
|
|
# both vision and MTP, so retry that way before dropping either.
|
|
# Only on a hard fault with FA on; a cancel/unload stops respawn.
|
|
if not healthy and not self._cancel_event.is_set():
|
|
_fa_rc = self._process.poll() if self._process is not None else None
|
|
_fa_cmd = (
|
|
self._with_flash_attn_off(_last_spawn_cmd)
|
|
if self._is_signal_crash(_fa_rc)
|
|
else None
|
|
)
|
|
if _fa_cmd is not None:
|
|
logger.warning(
|
|
"llama-server hard-crashed at startup (exit %s) with "
|
|
"flash attention on; retrying once with --flash-attn "
|
|
"off (keeps vision and MTP).",
|
|
_fa_rc,
|
|
)
|
|
self._kill_process()
|
|
cmd = _fa_cmd
|
|
healthy = _spawn_and_wait(_fa_cmd, label = "-noflash")
|
|
|
|
# MTP from Studio's spec flags or the user's (extra_args
|
|
# --spec-type / LLAMA_ARG_SPEC_TYPE). The env reaches the child
|
|
# only when neither emits a spec flag, so consult it only then.
|
|
_launch_spec_env: Mapping[str, str] = (
|
|
os.environ
|
|
if (not _extra_args_set_spec_type(extra_args) and not spec_flags)
|
|
else {}
|
|
)
|
|
_spec_requested_mtp = any(
|
|
"mtp" in str(t).lower() for t in spec_flags
|
|
) or _extra_args_requests_mtp(extra_args, env = _launch_spec_env)
|
|
# Is the launched server actually running MTP+tensor? Gates the
|
|
# probe/watchdog/recovery; cleared if the MTP-drop fallback wins.
|
|
_mtp_active_for_launched_server = bool(
|
|
self._tensor_parallel and _spec_requested_mtp
|
|
)
|
|
# MTP can pass /health then crash the flash-attn kernel on the
|
|
# first decode under tensor; probe one generation so the fallback
|
|
# catches that too. Tensor-only, so ordinary MTP stays probe-free.
|
|
if (
|
|
healthy
|
|
and self._tensor_parallel
|
|
and _spec_requested_mtp
|
|
and not self._cancel_event.is_set()
|
|
and not self._probe_mtp_decode()
|
|
):
|
|
# A first-decode hard fault is usually the FA kernel: retry
|
|
# FA-off (keeps MTP) before dropping speculative decoding below.
|
|
_probe_rc = self._process.poll() if self._process is not None else None
|
|
_fa_cmd = (
|
|
self._with_flash_attn_off(_last_spawn_cmd)
|
|
if self._is_signal_crash(_probe_rc)
|
|
else None
|
|
)
|
|
healthy = False
|
|
if _fa_cmd is not None:
|
|
logger.warning(
|
|
"MTP first-decode hard-crashed (exit %s) with flash "
|
|
"attention on; retrying with --flash-attn off.",
|
|
_probe_rc,
|
|
)
|
|
self._kill_process()
|
|
cmd = _fa_cmd
|
|
healthy = (
|
|
_spawn_and_wait(_fa_cmd, label = "-noflash-mtp")
|
|
and self._probe_mtp_decode()
|
|
)
|
|
if not healthy:
|
|
logger.warning(
|
|
"MTP speculative decoding crashed on the first decode "
|
|
"under tensor parallelism; retrying without it."
|
|
)
|
|
# Any MTP request can abort the server: a separate drafter
|
|
# (Gemma) on a binary that predates its arch, or an embedded
|
|
# head (Qwen) the binary cannot build. Retry once with the
|
|
# spec slice replaced by --spec-default so the main model still
|
|
# loads. Gate on the spec block (not the drafter path, which
|
|
# off/ngram local loads also carry) and keep
|
|
# _requested_spec_mode so a duplicate /load doesn't thrash. The
|
|
# cancel check stops an /unload-killed attempt respawning. A
|
|
# decode-probe failure above also routes here.
|
|
if not healthy and _spec_requested_mtp and not self._cancel_event.is_set():
|
|
# Blame the binary only when the output shows MTP itself
|
|
# failing (unknown arch / draft or context build); an
|
|
# unrelated crash (e.g. OOM) gets a neutral message.
|
|
_lo = "\n".join(self._stdout_lines).lower()
|
|
# Only an unknown architecture proves the prebuilt predates
|
|
# this MTP model (an update fixes it). The memory/context
|
|
# build failures are generic (VRAM / ctx pressure), where an
|
|
# update may not help, so classify those as runtime_error.
|
|
_arch_unsupported = "unknown model architecture" in _lo
|
|
if (
|
|
_arch_unsupported
|
|
or "failed to measure draft model memory" in _lo
|
|
or "failed to measure mtp context memory" in _lo
|
|
or "failed to create llama_context" in _lo
|
|
):
|
|
_retry_reason = (
|
|
"the prebuilt may predate it; retrying without "
|
|
"speculative decoding -- run `unsloth studio "
|
|
"update` for MTP"
|
|
)
|
|
self._spec_fallback_reason = (
|
|
"binary_outdated" if _arch_unsupported else "runtime_error"
|
|
)
|
|
else:
|
|
_retry_reason = (
|
|
"retrying without speculative decoding in case MTP is the cause"
|
|
)
|
|
self._spec_fallback_reason = "runtime_error"
|
|
_drafter = (
|
|
Path(launch_mtp_draft_path).name
|
|
if launch_mtp_draft_path
|
|
else "embedded head"
|
|
)
|
|
logger.warning(
|
|
"llama-server failed to start with MTP (%s); %s.",
|
|
_drafter,
|
|
_retry_reason,
|
|
)
|
|
self._kill_process()
|
|
fallback_cmd = (
|
|
cmd[:_spec_start]
|
|
+ ["--spec-default"]
|
|
+ cmd[_spec_start + len(spec_flags) :]
|
|
)
|
|
# User/env MTP survives in the tail; llama.cpp takes the last
|
|
# spec flag, so a trailing --spec-default overrides it too.
|
|
if _extra_args_requests_mtp(extra_args, env = _launch_spec_env):
|
|
fallback_cmd.append("--spec-default")
|
|
healthy = _spawn_and_wait(fallback_cmd, label = "-retry")
|
|
if healthy:
|
|
self._speculative_type = "default"
|
|
_mtp_active_for_launched_server = False
|
|
|
|
# A too-old llama.cpp can reject a model's --mmproj projector
|
|
# (format message or a bare SIGSEGV); retry once text-only.
|
|
if not healthy:
|
|
out = "\n".join(self._stdout_lines[-50:])
|
|
# Read the crash code before _kill_process() clears _process.
|
|
_crash_rc = self._process.poll() if self._process is not None else None
|
|
self._kill_process()
|
|
# The #6415 split-axis abort is latched earlier (first spawn).
|
|
# Skip if a cancel/unload is pending (mirrors the MTP guard).
|
|
if (
|
|
launched_with_mmproj
|
|
and not self._cancel_event.is_set()
|
|
and (
|
|
self._is_projector_incompatibility(out)
|
|
or (
|
|
self._is_signal_crash(_crash_rc)
|
|
and not self._output_has_nonprojector_diagnostic(out)
|
|
)
|
|
)
|
|
):
|
|
logger.warning(
|
|
"llama-server could not load this model's vision "
|
|
"projector (--mmproj). The installed llama.cpp build is "
|
|
"likely too old for it. Loading text-only for this "
|
|
"session; run 'unsloth studio update' to enable vision."
|
|
)
|
|
cmd = self._strip_mmproj_args(_last_spawn_cmd)
|
|
self._is_vision = False
|
|
self._mmproj_has_audio = False
|
|
self._start_llama_process(cmd, env)
|
|
if not self._wait_for_health(timeout = 600.0):
|
|
# Read the exit code before _kill_process() clears it, so
|
|
# an OS-killed text-only retry still gets the OOM message.
|
|
_retry_rc = self._process.poll() if self._process is not None else None
|
|
self._kill_process()
|
|
raise RuntimeError(
|
|
"Vision projector incompatible with this llama.cpp "
|
|
"build, and the text-only retry also failed: "
|
|
+ self._classify_llama_start_failure(
|
|
"\n".join(self._stdout_lines[-50:]),
|
|
gguf_path,
|
|
self._model_identifier,
|
|
_retry_rc,
|
|
)
|
|
)
|
|
else:
|
|
raise RuntimeError(
|
|
self._classify_llama_start_failure(
|
|
out,
|
|
gguf_path,
|
|
self._model_identifier,
|
|
_crash_rc,
|
|
)
|
|
)
|
|
|
|
self._healthy = True
|
|
|
|
# Commit caller intent only after _healthy=True so a failed start
|
|
# can't poison the next inheritance check. None keeps prior, []
|
|
# clears, list sets. Source records hf_variant for the route's
|
|
# same_source check.
|
|
if extra_args is not None:
|
|
self._extra_args = list(extra_args)
|
|
self._extra_args_source = (model_identifier, hf_variant)
|
|
self._requested_n_ctx = int(n_ctx)
|
|
# Commit the known-good snapshot + whether MTP+tensor is live, then
|
|
# watch this load for a mid-generation crash.
|
|
self._last_load_kwargs = _pending_load_kwargs
|
|
self._mtp_runtime_fallback_active = _mtp_active_for_launched_server
|
|
self._start_mtp_crash_watchdog()
|
|
|
|
# Catch silent CPU fallback when GPU was intended (#5106).
|
|
self._gpu_offload_active = self._classify_gpu_offload(
|
|
gpu_indices is not None or use_fit, gpus or []
|
|
)
|
|
if self._gpu_offload_active is False:
|
|
logger.warning(
|
|
"llama-server appears to have loaded the model entirely "
|
|
"on CPU even though Studio detected at least one GPU. "
|
|
"This usually means the prebuilt binary's GPU backend "
|
|
"failed to load -- on Windows, cudart64_X.dll / "
|
|
"cublas64_X.dll could not be resolved. Reinstall the "
|
|
"Studio llama.cpp prebuilt or install a matching CUDA "
|
|
"toolkit (issue unslothai/unsloth#5106).",
|
|
)
|
|
|
|
logger.info(
|
|
f"llama-server ready on port {self._port} for model '{model_identifier}'"
|
|
)
|
|
# Poll llama-server /metrics -> vLLM-style engine_stats logs
|
|
# (only when the binary exposes /metrics).
|
|
if server_caps.get("supports_metrics"):
|
|
try:
|
|
from core.inference.llama_stats import maybe_start_stats_logger
|
|
if self._stats_logger is not None:
|
|
self._stats_logger.stop()
|
|
self._stats_logger = maybe_start_stats_logger(self.base_url, logger)
|
|
except Exception as e:
|
|
logger.debug(f"engine-stats logger not started: {e}")
|
|
else:
|
|
self._stats_logger = None
|
|
|
|
# Probe outside _lock (interruptible by /unload); init inside.
|
|
self._is_audio = False
|
|
self._audio_type = None
|
|
self._audio_probed = False
|
|
self._has_audio_input = False
|
|
try:
|
|
detected = self._detect_audio_type_strict()
|
|
self._audio_probed = True
|
|
except Exception as exc:
|
|
logger.debug("Audio probe failed: %s", exc)
|
|
detected = None
|
|
if not self._apply_detected_audio(detected):
|
|
return False
|
|
|
|
if not self._healthy:
|
|
return False
|
|
return True
|
|
|
|
def _build_speculative_flags(
|
|
self,
|
|
*,
|
|
speculative_type: Optional[str],
|
|
spec_draft_n_max: Optional[int],
|
|
extra_args: Optional[List[str]],
|
|
model_identifier: str,
|
|
model_path: Optional[str],
|
|
gpus: bool,
|
|
binary: Optional[str],
|
|
mtp_draft_path: Optional[str] = None,
|
|
) -> List[str]:
|
|
"""Return the llama-server flag list for the requested spec mode.
|
|
|
|
Side effects: sets ``self._speculative_type`` (resolved internal
|
|
emit), ``self._requested_spec_mode`` (canonical UI mode for the
|
|
status round-trip), and ``self._spec_draft_n_max`` (user override
|
|
only; None when the platform default applies).
|
|
|
|
Speculative decoding (n-gram self-speculation, zero VRAM):
|
|
ngram-mod uses a ~16 MB shared hash pool, constant memory /
|
|
complexity, variable draft lengths. Helps most when the model
|
|
repeats existing text (code refactor, summarisation, reasoning);
|
|
for low-repetition chat, overhead is ~5 ms.
|
|
|
|
Benchmarks from upstream llama.cpp speculative-decoding PRs:
|
|
Scenario | Without | With | Speedup
|
|
gpt-oss-120b code refactor | 181 t/s | 446 t/s | 2.5x
|
|
Qwen3-235B offloaded | 12 t/s | 21 t/s | 1.8x
|
|
gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x
|
|
Refs: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md
|
|
https://github.com/ggml-org/llama.cpp/pull/19164
|
|
https://github.com/ggml-org/llama.cpp/pull/18471
|
|
MTP guide: unsloth.ai/docs/models/qwen3.6#mtp-guide
|
|
|
|
Sub-3B dense MTP regresses vs spec-off when the head is baked into the
|
|
main GGUF (Qwen): the draft head's per-token cost exceeds the
|
|
acceptance savings at this scale. Q4_K_XL clean bench (each prompt once
|
|
after an unrelated warmup) on B200 + x86 CPU:
|
|
0.8B GPU: draft-mtp n=2 = 0.58x vs OFF; ngram-only = 1.10x
|
|
2B GPU: draft-mtp n=2 = 0.82x vs OFF; OFF or ngram = 1.00x
|
|
0.8B CPU: chained n=2 = 0.86x vs OFF; ngram-only = 1.19x
|
|
2B CPU: chained n=2 = 0.83x vs OFF; ngram-only = 1.01x
|
|
4B+ GPU/CPU: spec on is a net win (1.08x-1.46x).
|
|
A separate drafter (Gemma's root mtp-*.gguf) is a different, cheaper
|
|
mechanism that wins even below 3B, so it is exempt from the sub-3B drop
|
|
(``mtp_draft_path`` set -> not too small). B200 Q4_K_XL bench, draft-mtp
|
|
n=2 vs OFF: gemma-4-E2B (2B) = 1.21x, accept ~0.65 (vs ngram = 1.00x);
|
|
gemma-4-E4B (4B) and 12B engage as usual.
|
|
Auto falls back to ngram-mod (zero-VRAM, near-zero idle cost on
|
|
diverse content) for an embedded sub-3B head; forced MTP on a model
|
|
with no head/drafter defaults back (mtp -> spec-default, mtp+ngram ->
|
|
ngram-mod) since llama-server aborts otherwise; a drafter the binary
|
|
cannot build (older prebuilt, or a CUDA kernel limit) aborts the spawn
|
|
and the load retries once without speculative decoding.
|
|
"""
|
|
flags: List[str] = []
|
|
# Reset; emit branches re-set on the resolved emission.
|
|
self._spec_draft_n_max = None
|
|
self._speculative_type = None
|
|
self._spec_fallback_reason = None
|
|
|
|
# Canonical UI-facing requested mode (legacy values mapped via
|
|
# _canonicalize_spec_mode).
|
|
canonical_mode = _canonicalize_spec_mode(speculative_type)
|
|
# MTP signals: head baked into the main GGUF (Qwen, via metadata or
|
|
# name), or a separate drafter resolved from the repo (Gemma).
|
|
is_mtp_model = (
|
|
bool(self._nextn_predict_layers)
|
|
or _is_mtp_model_name(model_identifier, model_path)
|
|
or bool(mtp_draft_path)
|
|
)
|
|
user_owns_spec_type = _extra_args_set_spec_type(extra_args)
|
|
_mtp_size_b = _extract_model_size_b(model_identifier)
|
|
# The sub-3B regression is an embedded-head cost; a separate drafter
|
|
# (Gemma) is a cheap standalone model that wins below 3B, so exempt it.
|
|
_mtp_too_small = (
|
|
_mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B and not bool(mtp_draft_path)
|
|
)
|
|
# Drafterless Gemma (name-only MTP, no embedded head): emitting MTP
|
|
# would abort llama-server, so every mode below falls back instead.
|
|
_mtp_drafter_missing = (
|
|
_is_gemma_mtp_name(model_identifier, model_path)
|
|
and not mtp_draft_path
|
|
and not self._nextn_predict_layers
|
|
)
|
|
# Embedded MTP head on an MLA model (GLM-5.2/DeepSeek/Kimi, detected by
|
|
# kv_lora_rank): llama.cpp's MLA/DSA MTP path is ~2x slower than no spec,
|
|
# so Auto drops it (override via the Settings dropdown / forced mtp, or
|
|
# UNSLOTH_MLA_MTP_ENABLED=1). Separate drafters (Gemma, mtp_draft_path) and
|
|
# non-MLA embedded heads (Qwen, no kv_lora_rank) are unaffected.
|
|
_auto_mla_embedded_mtp = (
|
|
bool(self._nextn_predict_layers)
|
|
and self._kv_lora_rank is not None
|
|
and not bool(mtp_draft_path)
|
|
and not _mla_mtp_auto_enabled()
|
|
)
|
|
|
|
if user_owns_spec_type:
|
|
# User --spec-type wins outright; suppress auto-emit to avoid a
|
|
# duplicate spec block.
|
|
self._requested_spec_mode = None
|
|
return flags
|
|
|
|
effective_mode = canonical_mode or "auto"
|
|
self._requested_spec_mode = effective_mode
|
|
|
|
def _resolved_draft_n_max() -> int:
|
|
# User override wins; else platform default (the B200 / x86
|
|
# clean-sweep sweet spot from PR #5582 is n=2 GPU, n=3 CPU;
|
|
# past 3 regresses on essay-style low-acceptance prompts).
|
|
if spec_draft_n_max is not None:
|
|
n = int(spec_draft_n_max)
|
|
self._spec_draft_n_max = n
|
|
return n
|
|
return 2 if gpus else 3
|
|
|
|
def _emit_mtp(*, chain_ngram: bool) -> bool:
|
|
"""Append --spec-type mtp[/draft-mtp][,ngram-mod] + n-max."""
|
|
caps = self.probe_server_capabilities(binary)
|
|
mtp_token = caps.get("mtp_token") if caps else None
|
|
if not mtp_token:
|
|
logger.warning(
|
|
"Requested MTP speculative decoding but "
|
|
"llama-server lacks --spec-type mtp/draft-mtp; "
|
|
"run `unsloth studio update`. Loading without "
|
|
"speculative decoding."
|
|
)
|
|
# Override an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI wins
|
|
# over env) so the child matches the binary-capability gate and
|
|
# the no-MTP budget, like the sibling no-head/non-MTP fallbacks.
|
|
flags.append("--spec-default")
|
|
self._speculative_type = "default"
|
|
self._spec_fallback_reason = "binary_no_mtp"
|
|
return False
|
|
draft_n_max = _resolved_draft_n_max()
|
|
n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max"
|
|
# Separate-file drafter (Gemma): point llama-server at it. Baked-in
|
|
# heads (Qwen) pass no path -- llama-server reads them from the
|
|
# main GGUF.
|
|
if mtp_draft_path:
|
|
flags.extend(["--model-draft", mtp_draft_path])
|
|
logger.info(f"Using separate MTP drafter: {mtp_draft_path}")
|
|
spec_value = mtp_token
|
|
ngram_knobs: list[str] = []
|
|
if chain_ngram:
|
|
ngram_knobs = _build_ngram_mod_flags(caps)
|
|
if ngram_knobs:
|
|
spec_value = f"ngram-mod,{mtp_token}"
|
|
else:
|
|
logger.warning(
|
|
"llama-server lacks ngram-mod tuning "
|
|
"flags; loading MTP only (no ngram chain)"
|
|
)
|
|
flags.extend(["--spec-type", spec_value, n_max_flag, str(draft_n_max)])
|
|
flags.extend(ngram_knobs)
|
|
self._speculative_type = "draft-mtp"
|
|
chain_label = "chained ngram-mod" if chain_ngram else "MTP-only"
|
|
logger.info(f"Spec decoding: {mtp_token} ({chain_label})")
|
|
return True
|
|
|
|
def _emit_ngram_mod() -> bool:
|
|
"""Append --spec-type ngram-mod + flag-set knobs."""
|
|
ngram_caps = self.probe_server_capabilities(binary)
|
|
ngram_knobs = _build_ngram_mod_flags(ngram_caps)
|
|
flags.extend(["--spec-type", "ngram-mod"])
|
|
if not ngram_knobs:
|
|
logger.warning(
|
|
"llama-server lacks ngram-mod tuning "
|
|
"flags; loading without --spec-ngram-mod-* knobs"
|
|
)
|
|
flags.extend(ngram_knobs)
|
|
self._speculative_type = "ngram-mod"
|
|
logger.info("Spec decoding: ngram-mod")
|
|
return True
|
|
|
|
def _fallback_drafter_not_found() -> None:
|
|
"""Drafterless Gemma: use ngram-mod (or spec-default) and record why."""
|
|
logger.warning(
|
|
"Model %s is MTP-capable but no drafter or head was found; "
|
|
"falling back. Check network or run `unsloth studio update`.",
|
|
model_identifier,
|
|
)
|
|
if self.probe_server_capabilities(binary).get("supports_ngram_mod"):
|
|
_emit_ngram_mod()
|
|
else:
|
|
flags.append("--spec-default")
|
|
self._speculative_type = "default"
|
|
self._spec_fallback_reason = "drafter_not_found"
|
|
|
|
if effective_mode == "off":
|
|
return flags # nothing to emit
|
|
if effective_mode == "ngram-simple":
|
|
flags.extend(["--spec-type", "ngram-simple"])
|
|
self._speculative_type = "ngram-simple"
|
|
return flags
|
|
if effective_mode == "ngram":
|
|
_emit_ngram_mod()
|
|
return flags
|
|
if effective_mode == "mtp":
|
|
if not is_mtp_model:
|
|
# No head and no drafter: llama-server aborts on draft-mtp
|
|
# instead of no-op'ing, so default back.
|
|
logger.warning(
|
|
"MTP requested but this GGUF has no MTP head or drafter; "
|
|
"loading without speculative decoding."
|
|
)
|
|
flags.append("--spec-default")
|
|
self._speculative_type = "default"
|
|
return flags
|
|
if _mtp_drafter_missing:
|
|
# Drafterless: draft-mtp would abort llama-server, so fall back.
|
|
_fallback_drafter_not_found()
|
|
return flags
|
|
if _mtp_too_small:
|
|
logger.warning(
|
|
f"Forcing MTP on a {_mtp_size_b:.1f}B model; "
|
|
"the bench shows draft-mtp regresses below 3B. "
|
|
"Engaging anyway (user override)."
|
|
)
|
|
_emit_mtp(chain_ngram = False)
|
|
return flags
|
|
if effective_mode == "mtp+ngram":
|
|
if not is_mtp_model:
|
|
# No head/drafter: keep the ngram half (needs no head),
|
|
# drop the draft-mtp chain that would abort the server.
|
|
logger.warning(
|
|
"MTP+Ngram requested but this GGUF has no MTP head or "
|
|
"drafter; loading ngram-mod only."
|
|
)
|
|
_emit_ngram_mod()
|
|
return flags
|
|
if _mtp_drafter_missing:
|
|
# No head/drafter: keep ngram-mod, drop the draft-mtp chain.
|
|
_fallback_drafter_not_found()
|
|
return flags
|
|
if _mtp_too_small:
|
|
logger.warning(
|
|
f"Forcing MTP+Ngram on a {_mtp_size_b:.1f}B model; "
|
|
"the bench shows the chain regresses below 3B. "
|
|
"Engaging anyway (user override)."
|
|
)
|
|
_emit_mtp(chain_ngram = True)
|
|
return flags
|
|
|
|
# effective_mode == "auto": the promotion path. llama.cpp #22673:
|
|
# MTP is compatible with mmproj, so there's no vision gate.
|
|
if _auto_mla_embedded_mtp:
|
|
# MLA embedded-MTP (GLM-5.2 et al.): the MTP path regresses vs spec-off
|
|
# on llama.cpp today, so Auto drops it and falls back to ngram-mod (or
|
|
# spec-off if unsupported), mirroring the sub-3B branch. Forced mtp /
|
|
# mtp+ngram (handled above) still engage; UNSLOTH_MLA_MTP_ENABLED=1
|
|
# re-enables this promotion once upstream optimizes the path.
|
|
self._spec_fallback_reason = "mla_mtp_disabled"
|
|
_mla_caps = self.probe_server_capabilities(binary)
|
|
if _mla_caps.get("supports_ngram_mod"):
|
|
logger.info(
|
|
"Auto: MLA embedded-MTP model detected; llama.cpp's MLA/DSA "
|
|
"MTP path is slower than no speculation, so using ngram-mod "
|
|
"instead. Override via the Studio Speculative Decoding "
|
|
"dropdown or UNSLOTH_MLA_MTP_ENABLED=1."
|
|
)
|
|
_emit_ngram_mod()
|
|
else:
|
|
logger.info(
|
|
"Auto: MLA embedded-MTP model detected; disabling speculative "
|
|
"decoding (this llama-server does not advertise ngram-mod). "
|
|
"Override via the dropdown or UNSLOTH_MLA_MTP_ENABLED=1."
|
|
)
|
|
# spec-off: emit nothing, mirroring the sub-3B no-ngram path.
|
|
elif is_mtp_model and not _mtp_too_small:
|
|
if _mtp_drafter_missing:
|
|
# Name-only MTP, drafter did not resolve (download failed/absent).
|
|
_fallback_drafter_not_found()
|
|
else:
|
|
# GPU: MTP-only. CPU/Mac: chain ngram-mod + MTP.
|
|
_emit_mtp(chain_ngram = not gpus)
|
|
elif is_mtp_model and _mtp_too_small:
|
|
# Sub-3B fallback: drop the MTP draft head, keep ngram-mod when
|
|
# the binary supports it.
|
|
if _mtp_drafter_missing:
|
|
_fallback_drafter_not_found()
|
|
elif self.probe_server_capabilities(binary).get("supports_ngram_mod"):
|
|
logger.info(
|
|
f"MTP GGUF detected but model size {_mtp_size_b:.1f}B "
|
|
"is below the 3B speedup threshold; using ngram-mod "
|
|
"only (zero-VRAM, no draft head). Override via "
|
|
"--spec-type or the Studio Speculative Decoding "
|
|
"dropdown."
|
|
)
|
|
_emit_ngram_mod()
|
|
else:
|
|
logger.info(
|
|
f"MTP GGUF detected but model size {_mtp_size_b:.1f}B "
|
|
"is below the 3B speedup threshold and the bundled "
|
|
"llama-server does not advertise ngram-mod; "
|
|
"auto-disabling speculative decoding."
|
|
)
|
|
else:
|
|
# Non-MTP model: let llama-server choose its default strategy.
|
|
flags.append("--spec-default")
|
|
self._speculative_type = "default"
|
|
return flags
|
|
|
|
def _already_in_target_state(
|
|
self,
|
|
*,
|
|
model_identifier: str,
|
|
hf_variant: Optional[str],
|
|
n_ctx: int,
|
|
cache_type_kv: Optional[str],
|
|
speculative_type: Optional[str],
|
|
chat_template_override: Optional[str],
|
|
extra_args: Optional[List[str]],
|
|
is_vision: bool,
|
|
gguf_path: Optional[str] = None,
|
|
spec_draft_n_max: Optional[int] = None,
|
|
tensor_parallel: bool = False,
|
|
mtp_draft_path: Optional[str] = None,
|
|
preserve_multi_gpu_on_layer: bool = False,
|
|
) -> bool:
|
|
"""True iff the live server already satisfies these load kwargs.
|
|
|
|
Mirrors ``routes/inference.py:_request_matches_loaded_settings`` but
|
|
compares raw kwargs so ``load_model`` can short-circuit a duplicate
|
|
/load that raced past the route-level check (#5401).
|
|
"""
|
|
if not self.is_loaded:
|
|
return False
|
|
if (self._model_identifier or "").lower() != (model_identifier or "").lower():
|
|
return False
|
|
# Direct-file loads pass hf_variant=None while the backend stores an
|
|
# extracted filename label; compare paths to keep the guard symmetric.
|
|
if gguf_path is not None and self._gguf_path:
|
|
try:
|
|
if Path(self._gguf_path).resolve() != Path(gguf_path).resolve():
|
|
return False
|
|
except OSError:
|
|
return False
|
|
elif (self._hf_variant or "").lower() != (hf_variant or "").lower():
|
|
return False
|
|
if self._requested_n_ctx != int(n_ctx):
|
|
return False
|
|
|
|
def _norm(value):
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, str):
|
|
stripped = value.strip().lower()
|
|
return stripped or None
|
|
return value
|
|
|
|
if _norm(self._cache_type_kv) != _norm(cache_type_kv):
|
|
return False
|
|
|
|
# Reconcile a user --split-mode in extras AND an inherited tensor
|
|
# LLAMA_ARG_SPLIT_MODE env, but only against a server that actually
|
|
# launched tensor: if load_model downgraded to layer split it scrubbed
|
|
# the child env, so the env must not force an endless reload of a healthy
|
|
# server. An identical request would downgrade the same way.
|
|
if not _tensor_parallel_matches_loaded(extra_args, tensor_parallel, self._tensor_parallel):
|
|
return False
|
|
# Preserved tensor->layer fallback + an EXPLICIT tensor drop: reload so
|
|
# placement re-selects instead of keeping the all-GPU mask (mirrors the route,
|
|
# #6659). preserve_multi_gpu_on_layer carries the route's carry-forward decision
|
|
# (True for an implicit same-settings reload), so those still dedupe -- the HF
|
|
# auto-pick / local-dir flows skip the route guard and only reach here.
|
|
if (
|
|
self._layer_preserves_tensor_intent
|
|
and not _effective_tensor_parallel(extra_args, tensor_parallel)
|
|
and not preserve_multi_gpu_on_layer
|
|
):
|
|
return False
|
|
|
|
# Compare on the canonical requested mode. With --spec-type in
|
|
# extra_args the backend stores None; mirror that here.
|
|
if _extra_args_set_spec_type(extra_args):
|
|
req_mode = None
|
|
else:
|
|
req_mode = _canonicalize_spec_mode(speculative_type) or "auto"
|
|
backend_mode = self._requested_spec_mode
|
|
if req_mode != backend_mode:
|
|
return False
|
|
|
|
# Prior HF load fell back with drafter_not_found; a same-settings reload
|
|
# must retry the download in load_model, not dedupe to the stale fallback
|
|
# (HF loads resolve the drafter there, so gguf_path is None here).
|
|
if (
|
|
self._spec_fallback_reason == "drafter_not_found"
|
|
and gguf_path is None
|
|
and req_mode in ("auto", "mtp", "mtp+ngram")
|
|
):
|
|
return False
|
|
|
|
# spec_draft_n_max only matters when an MTP variant is engaged. Compare
|
|
# on the resolved spec so an Auto request promoted to draft-mtp still
|
|
# bounces a reload when n_max changes.
|
|
if (
|
|
self._speculative_type == "draft-mtp"
|
|
and spec_draft_n_max is not None
|
|
and int(spec_draft_n_max) != (self._spec_draft_n_max or 0)
|
|
):
|
|
return False
|
|
|
|
if (self._chat_template_override or None) != (chat_template_override or None):
|
|
return False
|
|
|
|
# A drafter appearing/disappearing next to a local GGUF changes the
|
|
# launch command (--model-draft) when the mode can use it; without
|
|
# this, adding mtp-*.gguf after a load is deduped away and MTP can't
|
|
# engage short of an unload. HF loads resolve the drafter inside
|
|
# load_model (gguf_path is None here), so only local paths compare;
|
|
# the route-level probe covers HF cache repos. No sub-3B gate: both
|
|
# sides come from the same config detection, so a sub-3B mismatch
|
|
# only happens when a drafter genuinely appeared (one benign reload,
|
|
# then the stored path converges).
|
|
if (
|
|
gguf_path is not None
|
|
and req_mode in ("auto", "mtp", "mtp+ngram")
|
|
and (mtp_draft_path or None) != (self._mtp_draft_path or None)
|
|
):
|
|
return False
|
|
|
|
# extra_args=None means "no opinion" (inherit handled at the route
|
|
# layer); only an explicit list forces equality.
|
|
if extra_args is not None:
|
|
current = list(self._extra_args) if self._extra_args is not None else []
|
|
if list(extra_args) != current:
|
|
return False
|
|
return True
|
|
|
|
def _classify_gpu_offload(
|
|
self, expected_gpu: bool, detected_gpus: list[tuple[int, int]]
|
|
) -> Optional[bool]:
|
|
"""True if the model landed on a GPU, False if only CPU buffers landed
|
|
despite GPU intent, None when there's no signal. Delegates to the shared
|
|
classifier so it tracks current llama.cpp logs (offloaded-layer counts /
|
|
device_info), not just the older "model buffer size" lines."""
|
|
if not detected_gpus or not expected_gpu:
|
|
return None
|
|
return classify_gpu_offload_lines(self._stdout_lines)
|
|
|
|
def load_cancelled(self) -> bool:
|
|
"""True if a load was cancelled (e.g. via unload/_cancel_event) and not
|
|
yet consumed by the next load_model. Lets the tensor->layer fallback
|
|
avoid restarting a load the user just cancelled."""
|
|
return self._cancel_event.is_set()
|
|
|
|
def unload_model(self) -> bool:
|
|
"""Terminate the subprocess and cancel any in-flight download."""
|
|
self._cancel_event.set()
|
|
with self._lock:
|
|
self._kill_process()
|
|
logger.info(f"Unloaded GGUF model: {self._model_identifier}")
|
|
self._model_identifier = None
|
|
self._gguf_path = None
|
|
self._hf_repo = None
|
|
self._mtp_draft_path = None
|
|
self._spec_fallback_reason = None
|
|
self._last_load_kwargs = None
|
|
self._mtp_runtime_fallback_active = False
|
|
self._hf_variant = None
|
|
self._is_vision = False
|
|
self._is_audio = False
|
|
self._audio_type = None
|
|
self._audio_probed = False
|
|
self._has_audio_input = False
|
|
self._mmproj_has_audio = False
|
|
self._port = None
|
|
self._healthy = False
|
|
self._context_length = None
|
|
self._effective_context_length = None
|
|
self._max_context_length = None
|
|
self._chat_template = None
|
|
self._chat_template_override = None
|
|
self._supports_reasoning = False
|
|
self._reasoning_always_on = False
|
|
self._reasoning_style = "enable_thinking"
|
|
self._reasoning_effort_levels = []
|
|
self._reasoning_default = True
|
|
self._supports_preserve_thinking = False
|
|
self._supports_tools = False
|
|
self._cache_type_kv = None
|
|
self._tensor_parallel = False
|
|
self._layer_preserves_tensor_intent = False
|
|
self._speculative_type = None
|
|
self._requested_spec_mode = None
|
|
self._spec_draft_n_max = None
|
|
self._n_layers = None
|
|
self._n_kv_heads = None
|
|
self._n_kv_heads_by_layer = None
|
|
self._n_heads = None
|
|
self._embedding_length = None
|
|
self._kv_key_length = None
|
|
self._kv_value_length = None
|
|
self._sliding_window = None
|
|
self._sliding_window_pattern = None
|
|
self._full_attention_interval = None
|
|
self._kv_lora_rank = None
|
|
self._key_length_mla = None
|
|
self._kv_key_length_swa = None
|
|
self._kv_value_length_swa = None
|
|
self._ssm_inner_size = None
|
|
self._ssm_state_size = None
|
|
self._shared_kv_layers = None
|
|
self._nextn_predict_layers = None
|
|
# Clean up temp chat template file.
|
|
if hasattr(self, "_chat_template_file") and self._chat_template_file:
|
|
try:
|
|
os.unlink(self._chat_template_file.name)
|
|
except Exception:
|
|
pass
|
|
self._chat_template_file = None
|
|
# Free audio codec GPU memory.
|
|
if LlamaCppBackend._codec_mgr is not None:
|
|
LlamaCppBackend._codec_mgr.unload()
|
|
LlamaCppBackend._codec_mgr = None
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
return True
|
|
|
|
def _kill_process(self):
|
|
"""Terminate the subprocess if running."""
|
|
# Stop the watchdog before a deliberate kill so a planned reload/unload
|
|
# isn't seen as a crash; a real crash never routes through here.
|
|
self._stop_mtp_crash_watchdog()
|
|
if self._process is None:
|
|
return
|
|
try:
|
|
self._process.terminate()
|
|
self._process.wait(timeout = 5)
|
|
except subprocess.TimeoutExpired:
|
|
logger.warning("llama-server did not exit on SIGTERM, sending SIGKILL")
|
|
self._process.kill()
|
|
self._process.wait(timeout = 5)
|
|
except Exception as e:
|
|
logger.warning(f"Error killing llama-server process: {e}")
|
|
finally:
|
|
# getattr: teardown must tolerate a partially-built backend (failed
|
|
# __init__ or a __new__-built instance), as with _llama_log_fh below.
|
|
if getattr(self, "_stats_logger", None) is not None:
|
|
self._stats_logger.stop()
|
|
self._stats_logger = None
|
|
self._process = None
|
|
self._clear_server_pid()
|
|
# Clear healthy so a /load during the replacement's warm-up can't
|
|
# short-circuit against the previous server's health (#5401).
|
|
self._healthy = False
|
|
# Drives _wait_for_vram_settle in the next load_model; set in finally
|
|
# so both in-process and frontend Apply paths record the kill.
|
|
self._last_kill_monotonic = time.monotonic()
|
|
stdout_thread = getattr(self, "_stdout_thread", None)
|
|
if stdout_thread is not None:
|
|
stdout_thread.join(timeout = 2)
|
|
self._stdout_thread = None
|
|
fh = getattr(self, "_llama_log_fh", None)
|
|
if fh is not None:
|
|
try:
|
|
fh.close()
|
|
except Exception:
|
|
pass
|
|
self._llama_log_fh = None
|
|
|
|
@staticmethod
|
|
def _server_pidfile_path() -> Optional[Path]:
|
|
"""Pidfile recording the live llama-server PID, under the active studio root
|
|
(per-root, so concurrent Studios with distinct UNSLOTH_STUDIO_HOME stay
|
|
isolated, mirroring the reaper's custom-root isolation)."""
|
|
try:
|
|
from utils.paths.storage_roots import studio_root # noqa: WPS433
|
|
return studio_root() / "llama-server.pid"
|
|
except Exception:
|
|
return None
|
|
|
|
@classmethod
|
|
def _record_server_pid(cls, pid: int) -> None:
|
|
"""Best-effort record of the spawned llama-server PID for orphan reaping.
|
|
|
|
Stores ``pid:starttime`` so a later startup can reject a PID that has
|
|
since been recycled to a different process (see ``_pid_start_identity``).
|
|
A bare ``pid`` (no identity) is still accepted on read for compatibility.
|
|
"""
|
|
path = cls._server_pidfile_path()
|
|
if path is None:
|
|
return
|
|
try:
|
|
path.parent.mkdir(parents = True, exist_ok = True)
|
|
path.write_text(f"{pid}:{cls._pid_start_identity(pid)}")
|
|
except Exception as e:
|
|
logger.debug(f"Could not write llama-server pidfile: {e}")
|
|
|
|
@classmethod
|
|
def _clear_server_pid(cls) -> None:
|
|
"""Best-effort removal of the llama-server pidfile."""
|
|
path = cls._server_pidfile_path()
|
|
if path is None:
|
|
return
|
|
try:
|
|
path.unlink(missing_ok = True)
|
|
except Exception as e:
|
|
logger.debug(f"Could not remove llama-server pidfile: {e}")
|
|
|
|
@staticmethod
|
|
def _pid_is_llama_server(pid: int) -> bool:
|
|
"""True only if pid is a live process whose binary is a llama-server. Guards
|
|
against PID reuse before killing a recorded orphan; returns False on any
|
|
uncertainty so an unrelated process is never killed."""
|
|
try:
|
|
import psutil
|
|
try:
|
|
proc = psutil.Process(pid)
|
|
if (proc.name() or "").lower().startswith("llama-server"):
|
|
return True
|
|
return Path(proc.exe() or "").name.lower().startswith("llama-server")
|
|
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
|
|
return False
|
|
except ImportError:
|
|
pass
|
|
if sys.platform != "linux":
|
|
return False
|
|
try:
|
|
if Path(os.readlink(f"/proc/{pid}/exe")).name.lower().startswith("llama-server"):
|
|
return True
|
|
except OSError:
|
|
pass
|
|
try:
|
|
with open(f"/proc/{pid}/cmdline", "rb") as fh:
|
|
tokens = fh.read().split(b"\x00")
|
|
first = tokens[0].decode("utf-8", "replace") if tokens else ""
|
|
return Path(first).name.lower().startswith("llama-server")
|
|
except OSError:
|
|
return False
|
|
|
|
@staticmethod
|
|
def _pid_start_identity(pid: int) -> str:
|
|
"""Stable per-PID identity (process start time) guarding against PID reuse.
|
|
|
|
Returns a token string, or "" when it cannot be determined (the caller
|
|
then falls back to the llama-server name check only)."""
|
|
try:
|
|
import psutil
|
|
try:
|
|
return str(psutil.Process(pid).create_time())
|
|
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
|
|
return ""
|
|
except ImportError:
|
|
pass
|
|
if sys.platform == "linux":
|
|
try:
|
|
with open(f"/proc/{pid}/stat", "rb") as fh:
|
|
data = fh.read()
|
|
# field 22 (starttime), counted from after the ")" that closes comm.
|
|
return data[data.rfind(b")") + 2 :].split()[19].decode()
|
|
except (OSError, IndexError):
|
|
return ""
|
|
return ""
|
|
|
|
@staticmethod
|
|
def _pid_parent_is_alive(pid: int) -> bool:
|
|
"""True if the recorded server's parent is still running, i.e. the server is
|
|
NOT orphaned. Lets the cross-session reap kill only a true orphan (parent
|
|
gone) and never a live server owned by a running Studio, regardless of which
|
|
process performs the sweep. Biased toward "alive" on uncertainty so a live
|
|
server is never mistakenly reaped."""
|
|
try:
|
|
import psutil
|
|
|
|
try:
|
|
ppid = psutil.Process(pid).ppid()
|
|
except psutil.NoSuchProcess:
|
|
return False # the recorded server itself is gone
|
|
except psutil.Error:
|
|
return True # cannot tell -- never risk killing a live server
|
|
if ppid <= 1:
|
|
return False # reparented to init -> orphan
|
|
return psutil.pid_exists(ppid)
|
|
except ImportError:
|
|
pass
|
|
if sys.platform == "linux":
|
|
try:
|
|
with open(f"/proc/{pid}/stat", "rb") as fh:
|
|
data = fh.read()
|
|
ppid = int(data[data.rfind(b")") + 2 :].split()[1])
|
|
except (OSError, IndexError, ValueError):
|
|
return False
|
|
if ppid <= 1:
|
|
return False
|
|
return Path(f"/proc/{ppid}").exists()
|
|
return False
|
|
|
|
@staticmethod
|
|
def _unlink_pidfile(path: Path) -> None:
|
|
"""Best-effort removal of a resolved pidfile path."""
|
|
try:
|
|
path.unlink(missing_ok = True)
|
|
except Exception:
|
|
pass
|
|
|
|
@classmethod
|
|
def _reap_recorded_pid(cls) -> int:
|
|
"""Kill the exact llama-server PID recorded at spawn, but only when it is a
|
|
genuine orphan -- its parent (the Studio that spawned it) is gone. This is
|
|
the cross-session backstop the parent-death reaper (Job Object /
|
|
PR_SET_PDEATHSIG) cannot cover: an orphan left by an already-dead Studio
|
|
(macOS, a best-effort failure, or a pre-existing orphan). Path-independent,
|
|
so it also catches an orphan the install-root match would miss.
|
|
|
|
A live server whose parent is still running is never reaped, so constructing
|
|
a second backend in-process (the helper / advisor paths each build a
|
|
LlamaCppBackend) cannot kill the active chat server. A recorded PID that has
|
|
been recycled to a different process is rejected by the start-time identity
|
|
and the llama-server name check, so unrelated user processes are never
|
|
touched. SIGKILL falls back to SIGTERM on Windows, where os.kill maps it to
|
|
TerminateProcess and SIGKILL is undefined."""
|
|
path = cls._server_pidfile_path()
|
|
if path is None or not path.exists():
|
|
return 0
|
|
|
|
pid = -1
|
|
identity = ""
|
|
try:
|
|
pid_str, _, identity = path.read_text().strip().partition(":")
|
|
pid = int(pid_str)
|
|
except Exception:
|
|
pid = -1
|
|
|
|
if pid <= 0:
|
|
cls._unlink_pidfile(path) # garbage record
|
|
return 0
|
|
if pid == os.getpid():
|
|
return 0 # never our own pid; leave the record alone
|
|
|
|
if cls._pid_parent_is_alive(pid):
|
|
# Live server with a running parent -> not an orphan; keep the record so
|
|
# a later startup can still reap it if that parent later dies abnormally.
|
|
return 0
|
|
|
|
# Parent is gone: candidate orphan. Reject a PID recycled to something else.
|
|
if identity and cls._pid_start_identity(pid) != identity:
|
|
cls._unlink_pidfile(path)
|
|
return 0
|
|
|
|
killed = 0
|
|
if cls._pid_is_llama_server(pid):
|
|
try:
|
|
os.kill(pid, getattr(signal, "SIGKILL", signal.SIGTERM))
|
|
killed = 1
|
|
logger.info(f"Killed orphaned llama-server from pidfile (pid={pid})")
|
|
except (ProcessLookupError, PermissionError):
|
|
pass
|
|
except Exception as e:
|
|
logger.debug(f"Could not kill recorded llama-server pid {pid}: {e}")
|
|
cls._unlink_pidfile(path)
|
|
return killed
|
|
|
|
@staticmethod
|
|
def _kill_orphaned_servers() -> int:
|
|
"""Kill orphaned llama-server processes started by studio.
|
|
|
|
Only kills processes whose resolved binary lives under a known
|
|
Studio install dir (or matches an exact env-var override), to avoid
|
|
terminating unrelated llama-server instances. Mirrors every location
|
|
_find_llama_server_binary() can return, so orphans from any
|
|
supported install path are cleaned up.
|
|
|
|
Uses psutil for cross-platform support (Linux, macOS, Windows);
|
|
falls back to pgrep + /proc/<pid>/exe on Linux when psutil is
|
|
absent.
|
|
|
|
Returns the count of processes killed; callers arm the VRAM-settle
|
|
wait on a positive count.
|
|
"""
|
|
# Cross-session backstop first: reap the exact PID we recorded at spawn,
|
|
# but only if it is a true orphan whose parent is gone (so a helper backend
|
|
# built while a chat server is live can never kill it). The root-gated
|
|
# enumeration below stays as a fallback.
|
|
killed = LlamaCppBackend._reap_recorded_pid()
|
|
try:
|
|
# -- Build the ownership allowlist --------------------------------
|
|
# exact_binaries -- env var overrides (exact path match).
|
|
# install_roots -- Studio-owned dir trees (binary must be under one).
|
|
install_roots: list[Path] = []
|
|
|
|
# Env-mode custom root (mirrors _find_llama_server_binary).
|
|
_resolved_sr, _is_legacy = LlamaCppBackend._resolved_studio_root_and_is_legacy()
|
|
_is_custom_root = not _is_legacy
|
|
if _is_custom_root:
|
|
install_roots.append(_resolved_sr / "llama.cpp")
|
|
|
|
# Primary install dir (default mode only). Env-mode skips this so a
|
|
# custom-root Studio can't kill a default-install Studio's server.
|
|
if not _is_custom_root:
|
|
install_roots.append(Path.home() / ".unsloth" / "llama.cpp")
|
|
|
|
# Legacy in-tree build dirs (older setup.sh)
|
|
project_root = Path(__file__).resolve().parents[4]
|
|
install_roots.append(project_root / "llama.cpp")
|
|
|
|
# Legacy: extracted binary
|
|
install_roots.append(project_root / "bin")
|
|
|
|
# UNSLOTH_LLAMA_CPP_PATH env var (custom install dir)
|
|
custom_dir = os.environ.get("UNSLOTH_LLAMA_CPP_PATH")
|
|
if custom_dir:
|
|
install_roots.append(Path(custom_dir))
|
|
|
|
# LLAMA_SERVER_PATH env var (exact binary path)
|
|
exact_binaries: list[Path] = []
|
|
env_binary = os.environ.get("LLAMA_SERVER_PATH")
|
|
if env_binary:
|
|
try:
|
|
exact_binaries.append(Path(env_binary).resolve())
|
|
except OSError:
|
|
pass
|
|
|
|
# Resolve all roots so is_relative_to works reliably.
|
|
resolved_roots: list[Path] = []
|
|
for root in install_roots:
|
|
try:
|
|
# A --with-llama-cpp-dir local link (symlink/junction)
|
|
# resolves into the user's own checkout. Adding it would let
|
|
# us treat the user's externally-launched llama-server as our
|
|
# orphan and kill it, so leave such roots out of the
|
|
# allowlist (we forgo orphan-reaping for local-link installs).
|
|
if _is_external_link(root):
|
|
continue
|
|
resolved_roots.append(root.resolve())
|
|
except OSError:
|
|
pass
|
|
|
|
my_pid = os.getpid()
|
|
|
|
# -- Enumerate processes -------------------------------------------
|
|
# Prefer psutil (cross-platform); fall back to pgrep + /proc on
|
|
# Linux when psutil is absent.
|
|
try:
|
|
import psutil
|
|
has_psutil = True
|
|
except ImportError:
|
|
has_psutil = False
|
|
|
|
if has_psutil:
|
|
for proc in psutil.process_iter(["pid", "name", "exe"]):
|
|
try:
|
|
if proc.info["pid"] == my_pid:
|
|
continue
|
|
|
|
name = proc.info.get("name") or ""
|
|
if not name.lower().startswith("llama-server"):
|
|
continue
|
|
|
|
exe = proc.info.get("exe")
|
|
if not exe:
|
|
continue
|
|
|
|
exe_path = Path(exe).resolve()
|
|
|
|
# Ownership: exact match OR binary under a known root.
|
|
is_ours = exe_path in exact_binaries or any(
|
|
exe_path.is_relative_to(root) for root in resolved_roots
|
|
)
|
|
if not is_ours:
|
|
continue
|
|
|
|
proc.kill()
|
|
killed += 1
|
|
logger.info(
|
|
f"Killed orphaned llama-server process (pid={proc.info['pid']})"
|
|
)
|
|
except (
|
|
psutil.NoSuchProcess,
|
|
psutil.AccessDenied,
|
|
psutil.ZombieProcess,
|
|
):
|
|
pass
|
|
else:
|
|
# -- Fallback: pgrep + /proc/<pid>/exe (Linux only) -----------
|
|
if sys.platform != "linux":
|
|
return killed
|
|
result = subprocess.run(
|
|
["pgrep", "-a", "-f", "llama-server"],
|
|
capture_output = True,
|
|
text = True,
|
|
timeout = 5,
|
|
env = child_env_without_native_path_secret(),
|
|
)
|
|
if result.returncode != 0:
|
|
return killed
|
|
|
|
for line in result.stdout.strip().splitlines():
|
|
parts = line.strip().split(None, 1)
|
|
if len(parts) < 2:
|
|
continue
|
|
pid = int(parts[0])
|
|
if pid == my_pid:
|
|
continue
|
|
|
|
# /proc/<pid>/exe symlinks the real binary, avoiding
|
|
# cmdline-parsing ambiguities; fall back to the first
|
|
# cmdline token when /proc is unavailable.
|
|
proc_exe = Path(f"/proc/{pid}/exe")
|
|
try:
|
|
binary = proc_exe.resolve(strict = True)
|
|
except (OSError, ValueError):
|
|
cmdline = parts[1]
|
|
token = cmdline.split()[0] if cmdline.strip() else ""
|
|
if not token:
|
|
continue
|
|
binary = Path(token).resolve(strict = False)
|
|
|
|
owned = binary in exact_binaries or any(
|
|
binary.is_relative_to(root) for root in resolved_roots
|
|
)
|
|
if not owned:
|
|
continue
|
|
|
|
try:
|
|
os.kill(pid, signal.SIGKILL)
|
|
killed += 1
|
|
logger.info(f"Killed orphaned llama-server process (pid={pid})")
|
|
except ProcessLookupError:
|
|
pass
|
|
except PermissionError:
|
|
pass
|
|
except Exception:
|
|
logger.warning("Error during orphan server cleanup", exc_info = True)
|
|
return killed
|
|
|
|
def _cleanup(self):
|
|
"""atexit handler to ensure llama-server is terminated."""
|
|
self._kill_process()
|
|
|
|
@staticmethod
|
|
def _fit_off_retry_eligible(cmd: "list[str]", use_fit: bool) -> bool:
|
|
"""Whether a llama-server startup crash may be retried with --fit off.
|
|
|
|
Only when Studio's own VRAM math placed the model (use_fit=False)
|
|
and nothing on the command line set the fit mode explicitly
|
|
(-fit / --fit, space- or equals-form). --fit-ctx / --fit-target /
|
|
-fitc / -fitt tune the fit step but do not select the mode, so
|
|
they do not block the retry.
|
|
"""
|
|
if use_fit:
|
|
return False
|
|
for a in cmd:
|
|
if a in ("-fit", "--fit") or a.startswith(("-fit=", "--fit=")):
|
|
return False
|
|
return True
|
|
|
|
def _probe_mtp_decode(self, timeout: float = 60.0) -> bool:
|
|
"""One tiny /completion to confirm MTP survives the first decode.
|
|
|
|
MTP-draft can pass /health yet crash the flash-attn kernel only once
|
|
tokens generate (e.g. under --split-mode tensor). False on any error so
|
|
the caller can drop MTP and retry.
|
|
"""
|
|
url = f"{self.base_url}/completion"
|
|
payload = {"prompt": "Hi", "n_predict": 4, "temperature": 0.0, "stream": False}
|
|
try:
|
|
resp = httpx.post(
|
|
url,
|
|
json = payload,
|
|
timeout = timeout,
|
|
headers = self._auth_headers,
|
|
trust_env = False,
|
|
)
|
|
except Exception as e:
|
|
logger.debug(f"MTP decode probe failed: {e}")
|
|
return False
|
|
if resp.status_code != 200:
|
|
logger.debug(f"MTP decode probe returned HTTP {resp.status_code}")
|
|
return False
|
|
# A crash can drop the connection or kill the process right after a reply.
|
|
if self._process is not None and self._process.poll() is not None:
|
|
return False
|
|
return True
|
|
|
|
def _maybe_recover_from_mtp_crash(self, exc: Optional[BaseException] = None) -> bool:
|
|
"""Schedule one background reload without MTP after a mid-generation death.
|
|
|
|
MTP+tensor can crash the flash-attn kernel on a later request, after
|
|
load_model returned, past the load-time fallback and decode probe. Not a
|
|
persistent ban: a fresh load re-tries MTP. Returns True if scheduled.
|
|
"""
|
|
# Cheap async-safe gate: only our live MTP+tensor launch, not cancelled,
|
|
# with a snapshot to replay.
|
|
if self._cancel_event.is_set():
|
|
return False
|
|
if not self._mtp_runtime_fallback_active:
|
|
return False
|
|
if not self._last_load_kwargs or self._process is None:
|
|
return False
|
|
# Single-flight: the first failure claims the reload.
|
|
with self._mtp_runtime_fallback_lock:
|
|
if self._mtp_runtime_fallback_in_progress:
|
|
return False
|
|
self._mtp_runtime_fallback_in_progress = True
|
|
snapshot = dict(self._last_load_kwargs)
|
|
proc = self._process
|
|
|
|
def _recover():
|
|
try:
|
|
# Confirm the process really exited (the error can arrive a beat
|
|
# early) so a transient stream error can't disable MTP.
|
|
deadline = time.monotonic() + 5.0
|
|
while proc.poll() is None and time.monotonic() < deadline:
|
|
time.sleep(0.1)
|
|
if proc.poll() is None:
|
|
logger.debug("Generation error but llama-server is alive; keeping MTP.")
|
|
return
|
|
logger.warning(
|
|
"llama-server exited mid-generation with MTP under tensor "
|
|
"parallelism (%s); reloading without speculative decoding.",
|
|
type(exc).__name__ if exc is not None else "server exited",
|
|
)
|
|
# Re-check under the load lock (RLock allows the nested
|
|
# load_model) so a newer load isn't clobbered by this stale replay.
|
|
requested_mode = snapshot.get("speculative_type")
|
|
with self._serial_load_lock:
|
|
if self._cancel_event.is_set():
|
|
logger.info("MTP-crash reload skipped: load was cancelled/unloaded.")
|
|
return
|
|
if self._process is not proc:
|
|
logger.info("MTP-crash reload skipped: a newer load is already active.")
|
|
return
|
|
if self._last_load_kwargs != snapshot:
|
|
logger.info("MTP-crash reload skipped: load settings changed.")
|
|
return
|
|
snapshot["speculative_type"] = "off"
|
|
# Drop user/env MTP too: append a last-wins --spec-default.
|
|
_ea = list(snapshot.get("extra_args") or [])
|
|
if _extra_args_requests_mtp(_ea, env = os.environ):
|
|
_ea.append("--spec-default")
|
|
snapshot["extra_args"] = _ea
|
|
self.load_model(**snapshot)
|
|
# Restore the requested mode + reason load_model("off") cleared,
|
|
# so /status shows the user's mode + note (like the startup fallback).
|
|
self._requested_spec_mode = _canonicalize_spec_mode(requested_mode)
|
|
self._spec_fallback_reason = "runtime_error"
|
|
logger.info("Reloaded without MTP after the tensor-parallel crash.")
|
|
except Exception as e:
|
|
logger.error(f"Reload without MTP failed: {e}")
|
|
finally:
|
|
with self._mtp_runtime_fallback_lock:
|
|
self._mtp_runtime_fallback_in_progress = False
|
|
|
|
threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start()
|
|
return True
|
|
|
|
def _start_mtp_crash_watchdog(self) -> None:
|
|
"""Background poll that recovers on an MTP+tensor crash even when no
|
|
request observes it (direct proxy endpoints, or nothing in flight).
|
|
|
|
Armed only for a live MTP+tensor launch; the no-MTP reload disarms it, so
|
|
it can't loop.
|
|
"""
|
|
if not self._mtp_runtime_fallback_active:
|
|
return
|
|
proc = self._process
|
|
if proc is None:
|
|
return
|
|
# Replace any prior watchdog (loads are serialised, so at most one).
|
|
self._stop_mtp_crash_watchdog()
|
|
stop = threading.Event()
|
|
self._mtp_watchdog_stop = stop
|
|
|
|
def _watch():
|
|
# Exit on stop or process death. _kill_process sets stop before
|
|
# terminating, so re-check it: only a real crash (stop unset) recovers.
|
|
while not stop.wait(1.0):
|
|
if proc.poll() is not None:
|
|
if not stop.is_set():
|
|
self._maybe_recover_from_mtp_crash()
|
|
return
|
|
|
|
t = threading.Thread(target = _watch, daemon = True, name = "mtp-crash-watchdog")
|
|
self._mtp_watchdog_thread = t
|
|
t.start()
|
|
|
|
def _stop_mtp_crash_watchdog(self) -> None:
|
|
"""Signal the crash watchdog to exit; called before any deliberate kill."""
|
|
stop = getattr(self, "_mtp_watchdog_stop", None)
|
|
if stop is not None:
|
|
stop.set()
|
|
self._mtp_watchdog_thread = None
|
|
|
|
def _wait_for_health(
|
|
self,
|
|
timeout: float = 120.0,
|
|
interval: float = 0.5,
|
|
) -> bool:
|
|
"""Poll llama-server's /health until 200; also detect early exit/crash."""
|
|
deadline = time.monotonic() + timeout
|
|
url = f"{self.base_url}/health"
|
|
|
|
while time.monotonic() < deadline:
|
|
# Process crashed?
|
|
if self._process.poll() is not None:
|
|
# Let the drain thread collect final output.
|
|
if self._stdout_thread is not None:
|
|
self._stdout_thread.join(timeout = 2)
|
|
output = "\n".join(self._stdout_lines[-50:])
|
|
# Keep the TAIL: crash details (abort reason, ROCm/CUDA error
|
|
# text) print last, after the long startup banner. Head
|
|
# truncation has cut off exactly the diagnostic line before.
|
|
_log_hint = (
|
|
f" Full log: {self._llama_log_path}"
|
|
if getattr(self, "_llama_log_path", None)
|
|
else ""
|
|
)
|
|
logger.error(
|
|
f"llama-server exited with code {self._process.returncode}. "
|
|
f"Output (tail): {output[-2000:]}{_log_hint}"
|
|
)
|
|
return False
|
|
|
|
try:
|
|
# trust_env=False: skip ambient HTTP(S)_PROXY, which if it 503s
|
|
# for 127.0.0.1 loops the probe until timeout and hangs load.
|
|
resp = httpx.get(url, timeout = 2.0, trust_env = False)
|
|
if resp.status_code == 200:
|
|
return True
|
|
except (
|
|
httpx.ConnectError,
|
|
httpx.TimeoutException,
|
|
# ReadError covers TCP RST mid-read while still binding the port
|
|
# (Windows: WinError 10054); the crash branch catches real exits.
|
|
httpx.ReadError,
|
|
httpx.RemoteProtocolError,
|
|
httpx.WriteError,
|
|
):
|
|
pass
|
|
|
|
time.sleep(interval)
|
|
|
|
# Leave a marker so _classify_llama_start_failure tells a live but
|
|
# never-healthy load (too large, or a proxy hijacking the loopback
|
|
# probe) apart from a bad GGUF (#5740).
|
|
self._stdout_lines.append(f"llama-server health check timed out after {timeout}s")
|
|
logger.error(f"llama-server health check timed out after {timeout}s")
|
|
return False
|
|
|
|
@staticmethod
|
|
def _ctx_integrity_flags(
|
|
n_parallel: int, use_fit: bool, requested_ctx: int, effective_ctx: int, caps: dict
|
|
) -> list[str]:
|
|
"""Flags that keep the per-request window equal to the advertised ctx.
|
|
|
|
Explicit ``--parallel`` disables llama-server's auto-slots
|
|
``--kv-unified`` default, silently splitting ``-c`` into per-slot
|
|
windows of ``-c / N``; restore the shared pool so one request can use
|
|
the full context. With ``--fit on``, ``--fit-ctx`` floors the fit step
|
|
at an explicitly requested ctx (default floor is 4096) so it offloads
|
|
or fails instead of silently shrinking the window.
|
|
"""
|
|
flags: list[str] = []
|
|
if n_parallel > 1 and caps.get("supports_kv_unified"):
|
|
flags.append("--kv-unified")
|
|
if use_fit and requested_ctx > 0 and effective_ctx > 0 and caps.get("supports_fit_ctx"):
|
|
flags.extend(["--fit-ctx", str(effective_ctx)])
|
|
return flags
|
|
|
|
def _query_server_n_ctx(self) -> Optional[int]:
|
|
"""Per-slot context llama-server actually allocated, from ``/props``.
|
|
|
|
The memory-fit step or ``--parallel`` slot split can leave this below
|
|
the requested ``-c``; requests are validated against this value.
|
|
"""
|
|
url = f"{self.base_url}/props"
|
|
try:
|
|
resp = httpx.get(url, timeout = 5.0, trust_env = False)
|
|
if resp.status_code != 200:
|
|
return None
|
|
settings = resp.json().get("default_generation_settings") or {}
|
|
n_ctx = settings.get("n_ctx")
|
|
return int(n_ctx) if n_ctx else None
|
|
except Exception:
|
|
return None
|
|
|
|
def _reconcile_effective_ctx_with_server(self) -> None:
|
|
"""Adopt the server's real ``n_ctx`` when it is below Studio's value.
|
|
|
|
Keeps ``context_length`` (load response, status route, passthrough
|
|
``max_tokens`` ceiling) honest; clients sized to the requested value
|
|
would otherwise hit ``exceed_context_size_error`` 400s early.
|
|
"""
|
|
actual_n_ctx = self._query_server_n_ctx()
|
|
if not actual_n_ctx or actual_n_ctx <= 0:
|
|
return
|
|
if self._effective_context_length and actual_n_ctx < self._effective_context_length:
|
|
logger.warning(
|
|
"llama-server allocated a smaller per-request context than "
|
|
f"requested ({self._effective_context_length} -> {actual_n_ctx}; "
|
|
"memory fit or --parallel slot split); clients must treat "
|
|
f"{actual_n_ctx} as the real context window."
|
|
)
|
|
self._effective_context_length = actual_n_ctx
|
|
elif not self._effective_context_length:
|
|
self._effective_context_length = actual_n_ctx
|
|
|
|
# ── Message building (OpenAI format) ──────────────────────────
|
|
|
|
@staticmethod
|
|
def _parse_tool_calls_from_text(
|
|
content: str,
|
|
*,
|
|
allow_incomplete: bool = True,
|
|
enabled_tool_names: Optional[set] = None,
|
|
) -> list[dict]:
|
|
"""Wrapper around the shared parser; ``enabled_tool_names`` gates the markerless bare-JSON form."""
|
|
return _shared_parse_tool_calls_from_text(
|
|
content,
|
|
allow_incomplete = allow_incomplete,
|
|
enabled_tool_names = enabled_tool_names,
|
|
)
|
|
|
|
@staticmethod
|
|
def _build_openai_messages(messages: list[dict], image_b64: Optional[str] = None) -> list[dict]:
|
|
"""Build OpenAI-format messages, optionally injecting an image_url part
|
|
into the last user message for vision models. As-is if no image."""
|
|
if not image_b64:
|
|
return messages
|
|
|
|
# Convert the last user message to multimodal content parts
|
|
result = [msg.copy() for msg in messages]
|
|
last_user_idx = None
|
|
for i, msg in enumerate(result):
|
|
if msg["role"] == "user":
|
|
last_user_idx = i
|
|
|
|
if last_user_idx is not None:
|
|
text_content = result[last_user_idx].get("content", "")
|
|
result[last_user_idx]["content"] = [
|
|
{"type": "text", "text": text_content},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/png;base64,{image_b64}",
|
|
},
|
|
},
|
|
]
|
|
|
|
return result
|
|
|
|
# ── Generation (proxy to llama-server) ────────────────────────
|
|
|
|
@contextlib.contextmanager
|
|
def _open_stream(self, url: str, payload: dict, cancel_event):
|
|
"""Open a streaming POST to llama-server, retrying through prefill, and
|
|
yield ``(response, first_token_deadline)`` once a 200 lands. Owns the
|
|
httpx.Client + auth headers for the stream's lifetime; raises
|
|
RuntimeError on a non-200. Shared scaffold for the streaming consumers,
|
|
which differ only in how they parse the SSE body."""
|
|
stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
|
|
with httpx.Client(
|
|
timeout = stream_timeout,
|
|
limits = httpx.Limits(max_keepalive_connections = 0),
|
|
trust_env = False,
|
|
) as client:
|
|
first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
|
|
with self._stream_with_retry(
|
|
client,
|
|
url,
|
|
payload,
|
|
cancel_event,
|
|
headers = self._auth_headers,
|
|
first_token_deadline = first_token_deadline,
|
|
) as response:
|
|
if response.status_code != 200:
|
|
error_body = response.read().decode()
|
|
raise RuntimeError(
|
|
f"llama-server returned {response.status_code}: {error_body}"
|
|
)
|
|
yield response, first_token_deadline
|
|
|
|
@staticmethod
|
|
def _iter_text_cancellable(
|
|
response: "httpx.Response",
|
|
cancel_event: Optional[threading.Event] = None,
|
|
stall_timeout_s: float = _DEFAULT_STREAM_STALL_TIMEOUT_S,
|
|
first_token_deadline: Optional[float] = None,
|
|
post_first_chunk_read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S,
|
|
) -> Generator[str, None, None]:
|
|
"""Iterate a stream while polling cancel and stall timeouts."""
|
|
text_iter = response.iter_text()
|
|
if first_token_deadline is None:
|
|
first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
|
|
last_chunk_at: Optional[float] = None
|
|
while True:
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
response.close()
|
|
return
|
|
try:
|
|
if last_chunk_at is None:
|
|
remaining_s = first_token_deadline - time.monotonic()
|
|
if remaining_s <= 0:
|
|
raise httpx.ReadTimeout("The model did not produce a first token in time.")
|
|
LlamaCppBackend._set_stream_read_timeout(response, remaining_s)
|
|
chunk = next(text_iter)
|
|
if chunk:
|
|
if last_chunk_at is None and post_first_chunk_read_timeout_s is not None:
|
|
LlamaCppBackend._set_stream_read_timeout(
|
|
response,
|
|
post_first_chunk_read_timeout_s,
|
|
)
|
|
last_chunk_at = time.monotonic()
|
|
yield chunk
|
|
except StopIteration:
|
|
return
|
|
except httpx.ReadTimeout:
|
|
now = time.monotonic()
|
|
if last_chunk_at is None:
|
|
if now >= first_token_deadline:
|
|
raise
|
|
elif now - last_chunk_at >= stall_timeout_s:
|
|
raise httpx.ReadTimeout("The model stopped producing tokens mid-response.")
|
|
continue
|
|
|
|
@staticmethod
|
|
def _set_stream_read_timeout(response: "httpx.Response", read_timeout_s: float) -> None:
|
|
"""Lower only post-header stream reads; keep prefill timeout long."""
|
|
try:
|
|
timeout_ext = response.request.extensions.get("timeout")
|
|
if isinstance(timeout_ext, dict):
|
|
timeout_ext["read"] = read_timeout_s
|
|
except Exception:
|
|
logger.debug("Could not lower response read timeout", exc_info = True)
|
|
|
|
@staticmethod
|
|
def _shutdown_active_httpx_sockets(client: "httpx.Client") -> None:
|
|
"""Best-effort interrupt for a sync httpx request blocked before headers."""
|
|
try:
|
|
pool = getattr(getattr(client, "_transport", None), "_pool", None)
|
|
connections = list(getattr(pool, "_connections", []) or [])
|
|
for connection in connections:
|
|
inner = getattr(connection, "_connection", None)
|
|
stream = getattr(inner, "_network_stream", None)
|
|
sock = getattr(stream, "_sock", None)
|
|
if sock is None:
|
|
continue
|
|
try:
|
|
sock.shutdown(socket.SHUT_RDWR)
|
|
except OSError:
|
|
pass
|
|
try:
|
|
sock.close()
|
|
except OSError:
|
|
pass
|
|
except Exception:
|
|
logger.debug("Could not shutdown active httpx socket", exc_info = True)
|
|
try:
|
|
client.close()
|
|
except Exception:
|
|
logger.debug("Could not close httpx client", exc_info = True)
|
|
|
|
@staticmethod
|
|
@contextlib.contextmanager
|
|
def _stream_with_retry(
|
|
client: "httpx.Client",
|
|
url: str,
|
|
payload: dict,
|
|
cancel_event: Optional[threading.Event] = None,
|
|
headers: Optional[dict] = None,
|
|
first_token_deadline: Optional[float] = None,
|
|
):
|
|
"""Open one streaming POST and let cancel interrupt prefill or reads."""
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
raise GeneratorExit
|
|
|
|
_cancel_closed = threading.Event()
|
|
_response_ref: list = [None]
|
|
|
|
def _cancel_watcher():
|
|
while not _cancel_closed.is_set():
|
|
if cancel_event.wait(timeout = 0.3):
|
|
while not _cancel_closed.is_set():
|
|
r = _response_ref[0]
|
|
try:
|
|
if r is not None:
|
|
r.close()
|
|
else:
|
|
LlamaCppBackend._shutdown_active_httpx_sockets(client)
|
|
return
|
|
except Exception as e:
|
|
logger.debug(f"Error closing request in cancel watcher: {e}")
|
|
_cancel_closed.wait(timeout = 0.1)
|
|
return
|
|
|
|
watcher = None
|
|
if cancel_event is not None:
|
|
watcher = threading.Thread(target = _cancel_watcher, daemon = True, name = "prefill-cancel")
|
|
watcher.start()
|
|
|
|
try:
|
|
if first_token_deadline is None:
|
|
first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
|
|
prefill_read_timeout = max(0.1, first_token_deadline - time.monotonic())
|
|
prefill_timeout = httpx.Timeout(
|
|
connect = 30,
|
|
read = prefill_read_timeout,
|
|
write = 10,
|
|
pool = 10,
|
|
)
|
|
with client.stream(
|
|
"POST",
|
|
url,
|
|
json = payload,
|
|
timeout = prefill_timeout,
|
|
headers = headers,
|
|
) as response:
|
|
_response_ref[0] = response
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
raise GeneratorExit
|
|
yield response
|
|
return
|
|
except (httpx.RequestError, RuntimeError):
|
|
# Response was closed by the cancel watcher
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
raise GeneratorExit
|
|
raise
|
|
finally:
|
|
_cancel_closed.set()
|
|
|
|
def _respawn_if_dead(self) -> bool:
|
|
"""Relaunch the llama-server if its process has exited.
|
|
|
|
A loaded chat model can be SIGKILL'd mid-session (usually GPU/RAM pressure
|
|
from a training run on the same box), leaving a defunct process while
|
|
``is_loaded`` still reads True. Replay the last ``load_model`` call to
|
|
recover, returning True once healthy. Serialised on ``_respawn_lock`` so
|
|
many generations hitting the dead server trigger at most one reload.
|
|
"""
|
|
with self._respawn_lock:
|
|
proc = self._process
|
|
if proc is None:
|
|
return False
|
|
if proc.poll() is None:
|
|
# Process is alive: either a concurrent caller already respawned
|
|
# it (healthy), or this connection error wasn't a dead server.
|
|
return self._healthy
|
|
kwargs = self._last_load_kwargs
|
|
if not kwargs:
|
|
return False
|
|
logger.warning(
|
|
f"llama-server for '{self._model_identifier}' exited "
|
|
f"(code {proc.returncode}); respawning to recover the session"
|
|
)
|
|
with self._lock:
|
|
self._healthy = False
|
|
try:
|
|
return bool(self.load_model(**kwargs))
|
|
except Exception as exc:
|
|
logger.error(f"Failed to respawn llama-server: {exc}")
|
|
return False
|
|
|
|
def generate_chat_completion(
|
|
self,
|
|
messages: list[dict],
|
|
image_b64: Optional[str] = None,
|
|
temperature: float = 0.6,
|
|
top_p: float = 0.95,
|
|
top_k: int = 20,
|
|
min_p: float = 0.01,
|
|
max_tokens: Optional[int] = None,
|
|
repetition_penalty: float = 1.0,
|
|
presence_penalty: float = 0.0,
|
|
stop: Optional[list[str]] = None,
|
|
cancel_event: Optional[threading.Event] = None,
|
|
enable_thinking: Optional[bool] = None,
|
|
reasoning_effort: Optional[str] = None,
|
|
preserve_thinking: Optional[bool] = None,
|
|
seed: Optional[int] = None,
|
|
_allow_respawn_retry: bool = True,
|
|
) -> Generator[Union[str, dict], None, None]:
|
|
"""
|
|
Send a chat completion to llama-server and stream tokens back.
|
|
|
|
Uses /v1/chat/completions -- llama-server applies the chat template
|
|
and handles vision (multimodal image_url parts) natively.
|
|
|
|
Yields cumulative text (matching InferenceBackend's convention).
|
|
"""
|
|
if not self.is_loaded:
|
|
raise RuntimeError("llama-server is not loaded")
|
|
|
|
openai_messages = self._build_openai_messages(messages, image_b64)
|
|
|
|
payload = {
|
|
"messages": openai_messages,
|
|
"stream": True,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"top_k": top_k if top_k >= 0 else 0,
|
|
"min_p": min_p,
|
|
"repeat_penalty": repetition_penalty,
|
|
"presence_penalty": presence_penalty,
|
|
}
|
|
# Per-request enable_thinking / reasoning_effort / preserve_thinking
|
|
_reasoning_kw = self._request_reasoning_kwargs(
|
|
enable_thinking, reasoning_effort, preserve_thinking
|
|
)
|
|
if _reasoning_kw is not None:
|
|
payload["chat_template_kwargs"] = _reasoning_kw
|
|
# Default cap to the model context when known.
|
|
payload["max_tokens"] = (
|
|
max_tokens
|
|
if max_tokens is not None
|
|
else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR)
|
|
)
|
|
if stop:
|
|
payload["stop"] = stop
|
|
if seed is not None:
|
|
payload["seed"] = seed
|
|
payload["stream_options"] = {"include_usage": True}
|
|
|
|
url = f"{self.base_url}/v1/chat/completions"
|
|
cumulative = ""
|
|
in_thinking = False
|
|
_stream_done = False
|
|
_metadata_usage = None
|
|
_metadata_timings = None
|
|
_metadata_finish_reason = None
|
|
|
|
try:
|
|
with self._open_stream(url, payload, cancel_event) as (
|
|
response,
|
|
first_token_deadline,
|
|
):
|
|
buffer = ""
|
|
has_content_tokens = False
|
|
reasoning_text = ""
|
|
for raw_chunk in self._iter_text_cancellable(
|
|
response,
|
|
cancel_event,
|
|
first_token_deadline = first_token_deadline,
|
|
):
|
|
buffer += raw_chunk
|
|
while "\n" in buffer:
|
|
line, buffer = buffer.split("\n", 1)
|
|
line = line.strip()
|
|
|
|
if not line:
|
|
continue
|
|
if line == "data: [DONE]":
|
|
if in_thinking:
|
|
if has_content_tokens:
|
|
# Real thinking + content: close the tag
|
|
cumulative += "</think>"
|
|
yield cumulative
|
|
else:
|
|
# Only reasoning_content, no content:
|
|
# model put its whole reply in reasoning
|
|
# (e.g. Qwen3 always-think). Show it as
|
|
# the main response, not a thinking block.
|
|
cumulative = reasoning_text
|
|
yield cumulative
|
|
_stream_done = True
|
|
break # exit inner while
|
|
if not line.startswith("data: "):
|
|
continue
|
|
|
|
try:
|
|
data = json.loads(line[6:])
|
|
# Diffusion frame (per-step canvas) from the shim: forward untouched so
|
|
# the frontend renders it in place. No assistant text, so it never enters
|
|
# the cumulative content.
|
|
if data.get("type") == "diffusion_frame":
|
|
yield data
|
|
continue
|
|
# Capture server timings/usage from final chunks.
|
|
_chunk_timings = data.get("timings")
|
|
if _chunk_timings:
|
|
_metadata_timings = _chunk_timings
|
|
_chunk_usage = data.get("usage")
|
|
if _chunk_usage:
|
|
_metadata_usage = _chunk_usage
|
|
choices = data.get("choices", [])
|
|
if choices:
|
|
delta = choices[0].get("delta", {})
|
|
_fr = choices[0].get("finish_reason")
|
|
if _fr:
|
|
_metadata_finish_reason = _fr
|
|
|
|
# Reasoning/thinking tokens: llama-server
|
|
# sends these as "reasoning_content"; wrap
|
|
# in <think> tags for the frontend parser.
|
|
reasoning = delta.get("reasoning_content", "")
|
|
if reasoning:
|
|
reasoning_text += reasoning
|
|
if not in_thinking:
|
|
cumulative += "<think>"
|
|
in_thinking = True
|
|
cumulative += reasoning
|
|
yield cumulative
|
|
|
|
token = delta.get("content", "")
|
|
if token:
|
|
has_content_tokens = True
|
|
if in_thinking:
|
|
cumulative += "</think>"
|
|
in_thinking = False
|
|
cumulative += token
|
|
yield cumulative
|
|
except json.JSONDecodeError:
|
|
logger.debug(f"Skipping malformed SSE line: {line[:100]}")
|
|
if _stream_done:
|
|
break # exit outer for
|
|
if _metadata_usage or _metadata_timings or _metadata_finish_reason:
|
|
_metadata_usage = _backfill_usage_from_timings(
|
|
_metadata_usage, _metadata_timings
|
|
)
|
|
yield {
|
|
"type": "metadata",
|
|
# Never None: a finish-only metadata event (no usage,
|
|
# no timings) would otherwise crash consumers that do
|
|
# usage.get(...) on the non-streaming paths.
|
|
"usage": _metadata_usage or {},
|
|
"timings": _metadata_timings,
|
|
"finish_reason": _metadata_finish_reason,
|
|
}
|
|
|
|
except httpx.ConnectError as e:
|
|
# Server already down. If this was an MTP+tensor crash, recover by
|
|
# reloading without MTP (scheduled in the background) and fail this
|
|
# request. Otherwise the server was likely SIGKILL'd by GPU pressure
|
|
# from a concurrent training run: respawn the same config and retry the
|
|
# generation once (bounded by the private flag, no duplicate output).
|
|
if self._maybe_recover_from_mtp_crash(e):
|
|
raise RuntimeError("Lost connection to llama-server")
|
|
if _allow_respawn_retry and not cumulative and self._respawn_if_dead():
|
|
logger.warning(
|
|
"llama-server was unreachable; respawned it and retrying the generation"
|
|
)
|
|
yield from self.generate_chat_completion(
|
|
messages,
|
|
image_b64 = image_b64,
|
|
temperature = temperature,
|
|
top_p = top_p,
|
|
top_k = top_k,
|
|
min_p = min_p,
|
|
max_tokens = max_tokens,
|
|
repetition_penalty = repetition_penalty,
|
|
presence_penalty = presence_penalty,
|
|
stop = stop,
|
|
cancel_event = cancel_event,
|
|
enable_thinking = enable_thinking,
|
|
reasoning_effort = reasoning_effort,
|
|
preserve_thinking = preserve_thinking,
|
|
seed = seed,
|
|
_allow_respawn_retry = False,
|
|
)
|
|
return
|
|
raise RuntimeError("Lost connection to llama-server")
|
|
except Exception as e:
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
return
|
|
# Died mid-generation: recover MTP, re-raise unchanged for this request.
|
|
self._maybe_recover_from_mtp_crash(e)
|
|
raise
|
|
|
|
# ── Tool-calling agentic loop ──────────────────────────────
|
|
|
|
def generate_chat_completion_with_tools(
|
|
self,
|
|
messages: list[dict],
|
|
tools: list[dict],
|
|
temperature: float = 0.6,
|
|
top_p: float = 0.95,
|
|
top_k: int = 20,
|
|
min_p: float = 0.01,
|
|
max_tokens: Optional[int] = None,
|
|
repetition_penalty: float = 1.0,
|
|
presence_penalty: float = 0.0,
|
|
stop: Optional[list[str]] = None,
|
|
cancel_event: Optional[threading.Event] = None,
|
|
enable_thinking: Optional[bool] = None,
|
|
reasoning_effort: Optional[str] = None,
|
|
preserve_thinking: Optional[bool] = None,
|
|
max_tool_iterations: int = 25,
|
|
auto_heal_tool_calls: bool = True,
|
|
tool_call_timeout: int = 300,
|
|
session_id: Optional[str] = None,
|
|
rag_scope: Optional[dict] = None,
|
|
seed: Optional[int] = None,
|
|
disable_parallel_tool_use: bool = False,
|
|
confirm_tool_calls: bool = False,
|
|
bypass_permissions: bool = False,
|
|
) -> Generator[dict, None, None]:
|
|
"""
|
|
Agentic loop: let the model call tools, execute them, and continue.
|
|
|
|
Yields dicts:
|
|
{"type": "status", "text": "Searching: ..."/"Reading: ..."} -- tool status updates
|
|
{"type": "content", "text": "token"} -- streamed content tokens (cumulative)
|
|
{"type": "reasoning", "text": "token"} -- streamed reasoning tokens (cumulative)
|
|
"""
|
|
from core.inference.tools import build_rag_autoinject, execute_tool
|
|
|
|
if not self.is_loaded:
|
|
raise RuntimeError("llama-server is not loaded")
|
|
|
|
conversation = list(messages)
|
|
|
|
# Forced first-pass RAG so a doc question doesn't lose to web_search. Emits
|
|
# the same tool card + citations a real call would.
|
|
_auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope)
|
|
if _auto:
|
|
for _ev in _auto["events"]:
|
|
yield _ev
|
|
conversation.extend(_auto["messages"])
|
|
|
|
url = f"{self.base_url}/v1/chat/completions"
|
|
_accumulated_completion_tokens = 0
|
|
_accumulated_predicted_ms = 0.0
|
|
_accumulated_predicted_n = 0
|
|
# GGUF buffers reasoning; emit server-side timing before answer text.
|
|
_reasoning_started_at: Optional[float] = None
|
|
_reasoning_summary_emitted = False
|
|
|
|
def _reasoning_summary_event(started_at: float) -> dict:
|
|
return {
|
|
"type": "reasoning_summary",
|
|
"duration_ms": round((time.monotonic() - started_at) * 1000.0),
|
|
}
|
|
|
|
# Enabled-name gate for the markerless Gemma strip (disabled/example
|
|
# names stay visible). Set per iteration; None = pre-loop name-agnostic.
|
|
_enabled_tool_names = None
|
|
|
|
def _strip_tool_markup(
|
|
text: str,
|
|
*,
|
|
final: bool = False,
|
|
force: bool = False,
|
|
) -> str:
|
|
if not (auto_heal_tool_calls or force):
|
|
return text
|
|
return _shared_strip_tool_markup(
|
|
text, final = final, enabled_tool_names = _enabled_tool_names
|
|
)
|
|
|
|
def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str:
|
|
if not (auto_heal_tool_calls or force):
|
|
return text
|
|
# Shared parser patterns (not the legacy tool_healing set) so textual
|
|
# Mistral/python_tag calls entering DRAINING never leak. Balanced strips
|
|
# first (nested JSON removed whole); no final trim so length compares hold.
|
|
text = _strip_mistral_closed_calls(text)
|
|
text = _strip_gemma_wrapperless_calls(text, _enabled_tool_names)
|
|
# Parser-accurate scans close at each call's REAL terminator before
|
|
# the regex arms: literal markup inside a value is data.
|
|
text = _strip_function_xml_calls(text, final = True)
|
|
text = _strip_glm_calls(text, final = True)
|
|
for pat in _TOOL_ALL_PATS:
|
|
text = pat.sub("", text)
|
|
return text
|
|
|
|
def _build_metadata_event(usage, timings, finish_reason):
|
|
"""Final usage+timings metadata event for the given pass, merging its
|
|
usage/timings with the running cross-iteration accumulators. None when
|
|
there is nothing to report."""
|
|
_fu = _backfill_usage_from_timings(usage, timings) or {}
|
|
_fp = _fu.get("prompt_tokens", 0)
|
|
_tc = _fu.get("completion_tokens", 0) + _accumulated_completion_tokens
|
|
if not (usage or timings or _accumulated_completion_tokens or finish_reason):
|
|
return None
|
|
_mt = dict(timings) if timings else {}
|
|
if _accumulated_predicted_ms or _accumulated_predicted_n:
|
|
_mt["predicted_ms"] = _mt.get("predicted_ms", 0) + _accumulated_predicted_ms
|
|
_mt["predicted_n"] = _mt.get("predicted_n", 0) + _accumulated_predicted_n
|
|
if _mt["predicted_ms"] > 0:
|
|
_mt["predicted_per_second"] = _mt["predicted_n"] / (
|
|
_mt["predicted_ms"] / 1000.0
|
|
)
|
|
_usage = {
|
|
"prompt_tokens": _fp,
|
|
"completion_tokens": _tc,
|
|
"total_tokens": _fp + _tc,
|
|
}
|
|
# Preserve KV-cache hit details (cached_tokens) so the tool path
|
|
# reports them like the standard non-tool path does, not always 0.
|
|
if _fu.get("prompt_tokens_details"):
|
|
_usage["prompt_tokens_details"] = _fu["prompt_tokens_details"]
|
|
return {
|
|
"type": "metadata",
|
|
"usage": _usage,
|
|
"timings": _mt,
|
|
"finish_reason": finish_reason,
|
|
}
|
|
|
|
def _flush_reasoning_and_buffer():
|
|
"""Append buffered reasoning (as a <think> block) then the held
|
|
content_buffer to the cumulative display text."""
|
|
nonlocal cumulative_display
|
|
if reasoning_accum:
|
|
cumulative_display += "<think>" + reasoning_accum + "</think>"
|
|
cumulative_display += content_buffer
|
|
|
|
def _looks_like_enabled_bare_json(text: str, enabled_tool_names: set) -> bool:
|
|
"""True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False."""
|
|
probe = strip_llama3_leading_sentinels(text.lstrip())
|
|
if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)):
|
|
return False
|
|
return strip_leading_bare_json_call(probe, enabled_tool_names) != probe
|
|
|
|
tool_controller = ToolLoopController(
|
|
tools = tools,
|
|
auto_heal_tool_calls = auto_heal_tool_calls,
|
|
)
|
|
|
|
def _tool_succeeded(tool_name: str) -> bool:
|
|
key_prefix = f"{tool_name}:"
|
|
return any(
|
|
record.executed and not record.is_error and record.key.startswith(key_prefix)
|
|
for record in tool_controller.history
|
|
)
|
|
|
|
_MAX_BUFFER_CHARS = 32
|
|
# Hold a leading ``{`` well past the 32-char XML cap until it balances (mirrors safetensors).
|
|
_MAX_BARE_JSON_BUFFER = 16384
|
|
_append_budget_exhausted_nudge = True
|
|
# RAG: cap knowledge-base searches per assistant turn. The controller is
|
|
# tool-agnostic, so this gate stays in the loop.
|
|
_kb_search_count = 0
|
|
|
|
# ── Re-prompt on plan-without-action ─────────────────
|
|
# When the model describes what it intends to do (forward-looking
|
|
# language) without calling a tool, re-prompt once. Only triggers on
|
|
# responses signaling intent/planning -- a direct answer like "4" or
|
|
# "Hello!" won't match. Pattern compiled at module level
|
|
# (_INTENT_SIGNAL).
|
|
_reprompt_count = 0
|
|
# Gates ``max_tool_iterations`` on real tool turns (not the enlarged range) so reserved
|
|
# re-prompt slots don't extend the budget. Mirrors the safetensors guard.
|
|
_tool_iters_done = 0
|
|
_forced_tool_call_pending = False
|
|
|
|
# Reserve extra iterations for re-prompts so they don't consume the
|
|
# caller's tool-call budget; only when tool iterations are allowed.
|
|
_extra = _MAX_REPROMPTS if max_tool_iterations > 0 else 0
|
|
for iteration in range(max_tool_iterations + _extra):
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
return
|
|
# Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget.
|
|
_turn_executed_real_tool = False
|
|
|
|
active_tools = tool_controller.active_tools()
|
|
if not active_tools:
|
|
_append_budget_exhausted_nudge = False
|
|
break
|
|
# Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call.
|
|
_enabled_tool_names = {
|
|
(tool.get("function") or {}).get("name")
|
|
for tool in active_tools
|
|
if (tool.get("function") or {}).get("name")
|
|
}
|
|
# Shared signal tuple so GGUF BUFFERING wakes on every format the parser knows (like safetensors).
|
|
_tool_xml_signals = _SHARED_TOOL_XML_SIGNALS
|
|
|
|
# Build payload -- stream: True so we detect tool signals
|
|
# in the first 1-2 chunks without a non-streaming penalty.
|
|
payload = {
|
|
"messages": conversation,
|
|
"stream": True,
|
|
"stream_options": {"include_usage": True},
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"top_k": top_k if top_k >= 0 else 0,
|
|
"min_p": min_p,
|
|
"repeat_penalty": repetition_penalty,
|
|
"presence_penalty": presence_penalty,
|
|
"tools": active_tools,
|
|
"tool_choice": "auto",
|
|
}
|
|
_reasoning_kw = self._request_reasoning_kwargs(
|
|
enable_thinking, reasoning_effort, preserve_thinking
|
|
)
|
|
if _reasoning_kw is not None:
|
|
payload["chat_template_kwargs"] = _reasoning_kw
|
|
payload["max_tokens"] = (
|
|
max_tokens
|
|
if max_tokens is not None
|
|
else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR)
|
|
)
|
|
if stop:
|
|
payload["stop"] = stop
|
|
if seed is not None:
|
|
payload["seed"] = seed
|
|
|
|
try:
|
|
# ── Speculative buffer state machine ──────────────────
|
|
# BUFFERING: accumulate content, check for tool signals
|
|
# STREAMING: no tool detected, yield tokens to caller
|
|
# DRAINING: tool signal found, silently consume rest
|
|
_S_BUFFERING = 0
|
|
_S_STREAMING = 1
|
|
_S_DRAINING = 2
|
|
|
|
detect_state = _S_BUFFERING
|
|
content_buffer = "" # Raw content held during BUFFERING
|
|
content_accum = "" # All content tokens (for tool parsing)
|
|
reasoning_accum = ""
|
|
# Time each reasoning pass so final answers can replace tool timing.
|
|
_reasoning_started_at = None
|
|
_reasoning_summary_emitted = False
|
|
cumulative_display = "" # Cumulative yielded text (with <think>)
|
|
in_thinking = False
|
|
has_content_tokens = False
|
|
tool_calls_acc = {} # Structured delta.tool_calls fragments
|
|
has_structured_tc = False
|
|
_iter_usage = None
|
|
_iter_timings = None
|
|
_iter_finish_reason = None
|
|
_stream_done = False
|
|
_last_emitted = ""
|
|
# Provisional tool_start cards already shown, keyed by tool_call_id.
|
|
provisional_started_tool_calls: dict[str, str] = {}
|
|
resolved_provisional_tool_call_ids: set[str] = set()
|
|
_suppress_visible_output = _forced_tool_call_pending
|
|
|
|
with self._open_stream(url, payload, cancel_event) as (
|
|
response,
|
|
first_token_deadline,
|
|
):
|
|
raw_buf = ""
|
|
for raw_chunk in self._iter_text_cancellable(
|
|
response,
|
|
cancel_event,
|
|
first_token_deadline = first_token_deadline,
|
|
):
|
|
raw_buf += raw_chunk
|
|
while "\n" in raw_buf:
|
|
line, raw_buf = raw_buf.split("\n", 1)
|
|
line = line.strip()
|
|
|
|
if not line:
|
|
continue
|
|
if line == "data: [DONE]":
|
|
# Flush thinking state for STREAMING
|
|
if detect_state == _S_STREAMING and in_thinking:
|
|
if has_content_tokens:
|
|
cumulative_display += "</think>"
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": _strip_tool_markup(
|
|
cumulative_display,
|
|
final = True,
|
|
),
|
|
}
|
|
else:
|
|
cumulative_display = reasoning_accum
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": cumulative_display,
|
|
}
|
|
_stream_done = True
|
|
break # exit inner while
|
|
if not line.startswith("data: "):
|
|
continue
|
|
|
|
try:
|
|
chunk_data = json.loads(line[6:])
|
|
_ct = chunk_data.get("timings")
|
|
if _ct:
|
|
_iter_timings = _ct
|
|
_cu = chunk_data.get("usage")
|
|
if _cu:
|
|
_iter_usage = _cu
|
|
|
|
choices = chunk_data.get("choices", [])
|
|
if not choices:
|
|
continue
|
|
|
|
delta = choices[0].get("delta", {})
|
|
_fr = choices[0].get("finish_reason")
|
|
if _fr:
|
|
_iter_finish_reason = _fr
|
|
|
|
# ── Structured tool_calls ──
|
|
tc_deltas = delta.get("tool_calls")
|
|
if tc_deltas:
|
|
# Preserve any visible preface before draining
|
|
# the structured tool call.
|
|
has_structured_tc = True
|
|
detect_state = _S_DRAINING
|
|
for tc_d in tc_deltas:
|
|
idx = tc_d.get("index", 0)
|
|
if idx not in tool_calls_acc:
|
|
tool_calls_acc[idx] = {
|
|
"id": tc_d.get("id", f"call_{idx}"),
|
|
"type": "function",
|
|
"function": {
|
|
"name": "",
|
|
"arguments": "",
|
|
},
|
|
}
|
|
elif tc_d.get("id"):
|
|
# Update ID if a real one
|
|
# arrives on a later delta.
|
|
tool_calls_acc[idx]["id"] = tc_d["id"]
|
|
func = tc_d.get("function", {})
|
|
if func.get("name"):
|
|
tool_calls_acc[idx]["function"]["name"] += func["name"]
|
|
if func.get("arguments"):
|
|
tool_calls_acc[idx]["function"]["arguments"] += func[
|
|
"arguments"
|
|
]
|
|
current_name = tool_calls_acc[idx]["function"].get(
|
|
"name", ""
|
|
)
|
|
fallback_id = f"call_{idx}"
|
|
current_id = tool_calls_acc[idx].get("id", fallback_id)
|
|
already_started = (
|
|
current_id in provisional_started_tool_calls
|
|
)
|
|
# Empty/synthetic ids cannot reconcile with real starts.
|
|
has_real_id = bool(current_id) and current_id != fallback_id
|
|
# Show one early card per eligible streamed tool call.
|
|
_is_completed_one_shot = (
|
|
current_name == "render_html"
|
|
and _tool_succeeded("render_html")
|
|
)
|
|
# render_html is one-shot.
|
|
_one_shot_already_provisional = (
|
|
current_name == "render_html"
|
|
and "render_html"
|
|
in provisional_started_tool_calls.values()
|
|
)
|
|
# Later parallel cards only reconcile when parallel use is enabled.
|
|
_confirm_gated = (
|
|
confirm_tool_calls and not bypass_permissions
|
|
)
|
|
# Keep small-argument tools on the normal path.
|
|
_args_len = len(
|
|
tool_calls_acc[idx]["function"].get("arguments", "")
|
|
)
|
|
_payload_is_large = (
|
|
current_name == "render_html"
|
|
or _args_len >= _PROVISIONAL_ARGS_MIN_CHARS
|
|
)
|
|
if (
|
|
current_name
|
|
and (idx == 0 or not disable_parallel_tool_use)
|
|
and has_real_id
|
|
and not already_started
|
|
and not _is_completed_one_shot
|
|
and not _one_shot_already_provisional
|
|
and not _confirm_gated
|
|
and _payload_is_large
|
|
and any(
|
|
(tool.get("function") or {}).get("name")
|
|
== current_name
|
|
for tool in active_tools
|
|
)
|
|
):
|
|
provisional_started_tool_calls[current_id] = (
|
|
current_name
|
|
)
|
|
yield {
|
|
"type": "tool_start",
|
|
"tool_name": current_name,
|
|
"tool_call_id": current_id,
|
|
"arguments": {},
|
|
"provenance": tool_event_provenance(
|
|
provisional = True,
|
|
),
|
|
}
|
|
continue
|
|
|
|
# ── Reasoning tokens ──
|
|
# Yield only in STREAMING. In BUFFERING and
|
|
# DRAINING, accumulate silently so we don't
|
|
# corrupt the consumer's prev_text tracker
|
|
# (routes/inference.py never resets it
|
|
# between tool iterations).
|
|
reasoning = delta.get("reasoning_content", "")
|
|
if reasoning:
|
|
if _reasoning_started_at is None:
|
|
_reasoning_started_at = time.monotonic()
|
|
reasoning_accum += reasoning
|
|
if detect_state == _S_STREAMING:
|
|
if not in_thinking:
|
|
cumulative_display += "<think>"
|
|
in_thinking = True
|
|
cumulative_display += reasoning
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": cumulative_display,
|
|
}
|
|
|
|
# ── Content tokens ──
|
|
token = delta.get("content", "")
|
|
if token:
|
|
# First answer token ends reasoning.
|
|
if (
|
|
_reasoning_started_at is not None
|
|
and not _reasoning_summary_emitted
|
|
):
|
|
_reasoning_summary_emitted = True
|
|
yield _reasoning_summary_event(_reasoning_started_at)
|
|
has_content_tokens = True
|
|
content_accum += token
|
|
|
|
if detect_state == _S_DRAINING:
|
|
pass # accumulate silently
|
|
|
|
elif detect_state == _S_STREAMING:
|
|
if in_thinking:
|
|
cumulative_display += "</think>"
|
|
in_thinking = False
|
|
cumulative_display += token
|
|
cleaned = _strip_tool_markup_streaming(cumulative_display)
|
|
if len(cleaned) > len(_last_emitted):
|
|
_last_emitted = cleaned
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": cleaned,
|
|
}
|
|
|
|
elif detect_state == _S_BUFFERING:
|
|
content_buffer += token
|
|
stripped_buf = content_buffer.lstrip()
|
|
if not stripped_buf:
|
|
continue
|
|
|
|
# Check tool signal prefixes.
|
|
is_prefix = False
|
|
is_match = False
|
|
for sig in _tool_xml_signals:
|
|
if stripped_buf.startswith(sig):
|
|
is_match = True
|
|
break
|
|
if sig.startswith(stripped_buf):
|
|
is_prefix = True
|
|
break
|
|
|
|
# Signal-less call shapes (mirror the safetensors
|
|
# loop): Llama-3.2 bare {"name":..} and Gemma
|
|
# call:NAME{...} would otherwise stream raw.
|
|
_hold_buffer = False
|
|
# Whole buffer is the call (no visible prefix) -- drain silently.
|
|
_drain_silently = False
|
|
if not is_match and not is_prefix:
|
|
_bare = strip_llama3_leading_sentinels(stripped_buf)
|
|
if _bare.startswith("{"):
|
|
if _balanced_brace_end(_bare, 0) is None:
|
|
if len(stripped_buf) < _MAX_BARE_JSON_BUFFER:
|
|
_hold_buffer = True
|
|
elif _looks_like_enabled_bare_json(
|
|
_bare, _enabled_tool_names
|
|
):
|
|
# Oversized still-open enabled call: drain
|
|
# rather than leak; a giant ordinary JSON
|
|
# answer still streams.
|
|
_drain_silently = True
|
|
elif self._parse_tool_calls_from_text(
|
|
content_buffer,
|
|
allow_incomplete = auto_heal_tool_calls,
|
|
enabled_tool_names = _enabled_tool_names,
|
|
):
|
|
_drain_silently = True
|
|
elif (
|
|
"call:".startswith(stripped_buf)
|
|
or _GEMMA_BARE_TC_PREFIX_RE.match(stripped_buf)
|
|
is not None
|
|
or _GEMMA_BARE_TC_RE.match(stripped_buf) is not None
|
|
):
|
|
# Whitespace-tolerant like the parser.
|
|
if _GEMMA_BARE_TC_RE.match(stripped_buf):
|
|
_drain_silently = True
|
|
elif len(stripped_buf) < _MAX_BUFFER_CHARS:
|
|
_hold_buffer = True
|
|
|
|
if _drain_silently:
|
|
# No visible prefix -- the buffered text IS
|
|
# the call; drain without yielding it.
|
|
detect_state = _S_DRAINING
|
|
elif is_match:
|
|
# Tool signal -- flush any visible
|
|
# prefix before DRAINING so the
|
|
# route sends it before tool_start.
|
|
_flush_reasoning_and_buffer()
|
|
cleaned = _strip_tool_markup_streaming(
|
|
cumulative_display,
|
|
force = True,
|
|
)
|
|
if len(cleaned) > len(_last_emitted):
|
|
_last_emitted = cleaned
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": cleaned,
|
|
}
|
|
detect_state = _S_DRAINING
|
|
elif _hold_buffer or (
|
|
is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS
|
|
):
|
|
pass # keep buffering
|
|
else:
|
|
# Not a tool -- flush buffer
|
|
detect_state = _S_STREAMING
|
|
# Flush reasoning accumulated
|
|
# during BUFFERING.
|
|
_flush_reasoning_and_buffer()
|
|
cleaned = _strip_tool_markup(
|
|
cumulative_display,
|
|
)
|
|
if len(cleaned) > len(_last_emitted):
|
|
_last_emitted = cleaned
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": cleaned,
|
|
}
|
|
|
|
except json.JSONDecodeError:
|
|
logger.debug(f"Skipping malformed SSE line: {line[:100]}")
|
|
if _stream_done:
|
|
break # exit outer for
|
|
|
|
# ── Resolve BUFFERING at stream end ──
|
|
if detect_state == _S_BUFFERING:
|
|
stripped_buf = content_buffer.lstrip()
|
|
# A held bare-JSON fragment has no XML signal; route it to DRAINING (the signal-only
|
|
# gate below would flush the raw JSON to the user).
|
|
_bare_eos = strip_llama3_leading_sentinels(stripped_buf)
|
|
# Gate on enabled names so an ordinary JSON answer isn't routed to DRAINING and dropped.
|
|
_is_bare_tc = bool(active_tools) and _looks_like_enabled_bare_json(
|
|
_bare_eos, _enabled_tool_names
|
|
)
|
|
if stripped_buf and any(s in stripped_buf for s in _tool_xml_signals):
|
|
detect_state = _S_DRAINING
|
|
elif _is_bare_tc:
|
|
detect_state = _S_DRAINING
|
|
elif content_accum or reasoning_accum:
|
|
detect_state = _S_STREAMING
|
|
if content_buffer:
|
|
# Flush reasoning first.
|
|
_flush_reasoning_and_buffer()
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": _strip_tool_markup(
|
|
cumulative_display,
|
|
final = True,
|
|
),
|
|
}
|
|
elif reasoning_accum and not has_content_tokens:
|
|
# Reasoning-only reply: show it as plain text.
|
|
if _reasoning_started_at is not None and not _reasoning_summary_emitted:
|
|
_reasoning_summary_emitted = True
|
|
yield _reasoning_summary_event(_reasoning_started_at)
|
|
cumulative_display = reasoning_accum
|
|
if not _suppress_visible_output:
|
|
yield {
|
|
"type": "content",
|
|
"text": cumulative_display,
|
|
}
|
|
else:
|
|
# Held buffer was no tool signal and no enabled bare-JSON call: a leading ``{`` is an
|
|
# ordinary JSON answer and must be shown; any other partial-markup prefix is dropped.
|
|
_held = strip_llama3_leading_sentinels(content_buffer.lstrip())
|
|
if _held.startswith("{") and not _suppress_visible_output:
|
|
yield {"type": "content", "text": _held}
|
|
return
|
|
|
|
# ── STREAMING path: no tool call ──
|
|
if detect_state == _S_STREAMING:
|
|
# Safety net: re-parse the full content for tool calls. The
|
|
# route layer resets prev_text on tool_start, so post-tool
|
|
# synthesis streams correctly even if content was emitted
|
|
# before the tool XML.
|
|
# Unconditional (not gated on _tool_xml_signals): bare-JSON and Gemma wrapper-less
|
|
# calls carry no XML signal, so a signal gate would let them slip past.
|
|
_safety_tc = self._parse_tool_calls_from_text(
|
|
content_accum,
|
|
allow_incomplete = auto_heal_tool_calls,
|
|
enabled_tool_names = _enabled_tool_names,
|
|
)
|
|
if not _safety_tc:
|
|
# ── Re-prompt on plan-without-action ──
|
|
# If the model described its intent (forward-looking
|
|
# language) without calling a tool, nudge it to act.
|
|
# Fires at most once per request, only on short
|
|
# responses with intent signals -- "4" or "Hello!"
|
|
# won't trigger it. Use content if available, else
|
|
# fall back to reasoning text (reasoning-only stalls).
|
|
_stripped = content_accum.strip()
|
|
if not _stripped:
|
|
_stripped = reasoning_accum.strip()
|
|
_render_html_already_done_intent = _tool_succeeded(
|
|
"render_html"
|
|
) and re.search(
|
|
r"(?i)\brender[_\s-]?html\b",
|
|
_stripped,
|
|
)
|
|
if (
|
|
auto_heal_tool_calls
|
|
and active_tools
|
|
and not _render_html_already_done_intent
|
|
and _reprompt_count < _MAX_REPROMPTS
|
|
and _is_short_intent_without_action(_stripped)
|
|
):
|
|
_reprompt_count += 1
|
|
logger.info(
|
|
f"Re-prompt {_reprompt_count}/{_MAX_REPROMPTS}: "
|
|
f"model responded without calling tools "
|
|
f"({len(_stripped)} chars)"
|
|
)
|
|
conversation.append(
|
|
{
|
|
"role": "assistant",
|
|
"content": _stripped,
|
|
}
|
|
)
|
|
available_tool_names = [
|
|
(tool.get("function") or {}).get("name")
|
|
for tool in active_tools
|
|
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
|
|
]
|
|
available_tool_names = [name for name in available_tool_names if name]
|
|
tool_hint = " or ".join(available_tool_names) or "an available tool"
|
|
_forced_tool_call_pending = True
|
|
conversation.append(
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
"You have access to enabled tools. If a tool is needed to satisfy "
|
|
"the user's request or complete the action you described, call "
|
|
f"{tool_hint} now. If no tool is needed, provide the final answer "
|
|
"and follow the user's requested format."
|
|
),
|
|
}
|
|
)
|
|
# Accumulate tokens and timing from this iteration.
|
|
_fu_r = _backfill_usage_from_timings(_iter_usage, _iter_timings) or {}
|
|
_accumulated_completion_tokens += _fu_r.get("completion_tokens", 0)
|
|
_it_r = _iter_timings or {}
|
|
_accumulated_predicted_ms += _it_r.get("predicted_ms", 0)
|
|
_accumulated_predicted_n += _it_r.get("predicted_n", 0)
|
|
yield {"type": "status", "text": ""}
|
|
continue
|
|
|
|
if _forced_tool_call_pending:
|
|
_forced_tool_call_pending = False
|
|
if not _should_suppress_forced_no_tool_output(_stripped):
|
|
if cumulative_display:
|
|
forced_visible_text = _strip_tool_markup(
|
|
cumulative_display,
|
|
final = True,
|
|
)
|
|
elif content_accum:
|
|
forced_visible_text = _strip_tool_markup(
|
|
content_accum,
|
|
final = True,
|
|
)
|
|
else:
|
|
forced_visible_text = reasoning_accum
|
|
if forced_visible_text:
|
|
yield {
|
|
"type": "content",
|
|
"text": forced_visible_text,
|
|
}
|
|
|
|
# Content was already streamed. Yield metadata.
|
|
yield {"type": "status", "text": ""}
|
|
_meta = _build_metadata_event(
|
|
_iter_usage, _iter_timings, _iter_finish_reason
|
|
)
|
|
if _meta is not None:
|
|
yield _meta
|
|
return
|
|
|
|
# Safety net caught tool XML -- treat as tool call.
|
|
tool_calls = _safety_tc
|
|
content_text = _strip_tool_markup(
|
|
content_accum,
|
|
final = True,
|
|
force = True,
|
|
)
|
|
logger.info(
|
|
f"Safety net: parsed {len(tool_calls)} tool call(s) from streamed content"
|
|
)
|
|
else:
|
|
# ── DRAINING path: assemble tool_calls ──
|
|
tool_calls = None
|
|
content_text = content_accum
|
|
if has_structured_tc:
|
|
# Drop incomplete fragments (e.g. from max_tokens
|
|
# truncation or disconnect).
|
|
tool_calls = [
|
|
tool_calls_acc[i]
|
|
for i in sorted(tool_calls_acc)
|
|
if (tool_calls_acc[i].get("function", {}).get("name", "").strip())
|
|
] or None
|
|
if not tool_calls:
|
|
# Unconditional re-parse: we only reach DRAINING when the buffer looked like a
|
|
# call, and bare-JSON / Gemma wrapper-less calls carry no XML signal to gate on.
|
|
tool_calls = self._parse_tool_calls_from_text(
|
|
content_accum,
|
|
allow_incomplete = auto_heal_tool_calls,
|
|
enabled_tool_names = _enabled_tool_names,
|
|
)
|
|
if tool_calls and not has_structured_tc:
|
|
content_text = _strip_tool_markup(
|
|
content_text,
|
|
final = True,
|
|
force = True,
|
|
)
|
|
# ``_strip_tool_markup`` only knows XML; also drop a leading bare-JSON call so the
|
|
# executed call isn't replayed as text or next-turn history.
|
|
content_text = strip_leading_bare_json_call(
|
|
content_text, _enabled_tool_names
|
|
)
|
|
if tool_calls:
|
|
logger.info(
|
|
f"Parsed {len(tool_calls)} tool call(s) from "
|
|
f"{'structured delta' if has_structured_tc else 'content text'}"
|
|
)
|
|
if not tool_calls:
|
|
# DRAINING but no tool calls (false positive). Merge
|
|
# accumulated metrics from prior tool iterations so
|
|
# they aren't silently dropped.
|
|
yield {"type": "status", "text": ""}
|
|
if content_accum:
|
|
# Strip leaked tool-call XML before yielding.
|
|
content_accum = _strip_tool_markup(content_accum, final = True)
|
|
# A truncated bare-JSON call has no XML markup to strip and didn't parse. With
|
|
# Auto-Heal on, drop a leading ENABLED-tool fragment (ordinary JSON answers untouched);
|
|
# off keeps it visible per the strict contract.
|
|
if content_accum and active_tools and auto_heal_tool_calls:
|
|
content_accum = strip_leading_bare_json_call(
|
|
content_accum, _enabled_tool_names
|
|
)
|
|
if content_accum:
|
|
yield {"type": "content", "text": content_accum}
|
|
_meta = _build_metadata_event(
|
|
_iter_usage, _iter_timings, _iter_finish_reason
|
|
)
|
|
if _meta is not None:
|
|
yield _meta
|
|
return
|
|
|
|
# ── Execute tool calls ──
|
|
_accumulated_completion_tokens += (
|
|
_backfill_usage_from_timings(_iter_usage, _iter_timings) or {}
|
|
).get("completion_tokens", 0)
|
|
_it = _iter_timings or {}
|
|
_accumulated_predicted_ms += _it.get("predicted_ms", 0)
|
|
_accumulated_predicted_n += _it.get("predicted_n", 0)
|
|
|
|
# Collapse exact-duplicate calls and cap the count for the TEXTUAL
|
|
# fallback (mirrors the safetensors loop; see _MAX_TOOL_CALLS_PER_TURN).
|
|
if tool_calls and not has_structured_tc and len(tool_calls) > 1:
|
|
_seen_keys: set = set()
|
|
_deduped: list = []
|
|
for _tc in tool_calls:
|
|
_fn = _tc.get("function", {}) or {}
|
|
_key = (_fn.get("name", ""), str(_fn.get("arguments", "")))
|
|
if _key in _seen_keys:
|
|
continue
|
|
_seen_keys.add(_key)
|
|
_deduped.append(_tc)
|
|
if len(_deduped) >= _MAX_TOOL_CALLS_PER_TURN:
|
|
break
|
|
if len(_deduped) != len(tool_calls):
|
|
logger.info(
|
|
"GGUF textual fallback: collapsed %d repeated tool call(s) "
|
|
"in one turn to %d",
|
|
len(tool_calls),
|
|
len(_deduped),
|
|
)
|
|
tool_calls = _deduped
|
|
|
|
# disable_parallel_tool_use: execute only the first tool call
|
|
# this turn. Truncate before building assistant_msg so the
|
|
# conversation stays consistent and extra calls are never executed.
|
|
if disable_parallel_tool_use and tool_calls and len(tool_calls) > 1:
|
|
tool_calls = tool_calls[:1]
|
|
|
|
assistant_msg: dict = {"role": "assistant", "content": content_text}
|
|
assistant_appended = False
|
|
|
|
for tc in tool_calls or []:
|
|
func = tc.get("function", {})
|
|
tool_name = func.get("name", "")
|
|
provisional_match = tc.get("id") in provisional_started_tool_calls
|
|
decision = tool_controller.prepare_call(
|
|
tc,
|
|
forced = _forced_tool_call_pending,
|
|
provisional = provisional_match,
|
|
)
|
|
|
|
if not decision.should_execute:
|
|
if content_text and not assistant_appended:
|
|
conversation.append(assistant_msg)
|
|
assistant_appended = True
|
|
if provisional_match:
|
|
# A provisional tool card is already on screen for this
|
|
# id; close it so it never dangles when the controller
|
|
# turns the call into an internal no-op (duplicate /
|
|
# disabled / render_html_repeat).
|
|
resolved_provisional_tool_call_ids.add(decision.tool_call_id)
|
|
yield {
|
|
"type": "tool_end",
|
|
"tool_name": decision.tool_name,
|
|
"tool_call_id": decision.tool_call_id,
|
|
"result": "",
|
|
"provenance": decision.provenance,
|
|
}
|
|
completion = tool_controller.record_noop(decision)
|
|
conversation.append(completion.model_message())
|
|
if _forced_tool_call_pending:
|
|
_forced_tool_call_pending = False
|
|
logger.info(
|
|
"Suppressed local GGUF tool call as internal no-op: "
|
|
f"action={decision.action} tool={decision.tool_name}"
|
|
)
|
|
break
|
|
|
|
if not assistant_appended:
|
|
assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()]
|
|
conversation.append(assistant_msg)
|
|
assistant_appended = True
|
|
else:
|
|
assistant_msg.setdefault("tool_calls", []).append(
|
|
decision.as_assistant_tool_call()
|
|
)
|
|
|
|
# Bypass wins over the confirm gate at the loop level too,
|
|
# so a direct internal caller with both flags never prompts.
|
|
needs_confirm = bool(confirm_tool_calls) and not bypass_permissions
|
|
approval_id = new_approval_id() if needs_confirm else ""
|
|
decision_slot = (
|
|
begin_tool_decision(session_id, approval_id) if needs_confirm else None
|
|
)
|
|
start_event = decision.tool_start_event()
|
|
start_event["approval_id"] = approval_id
|
|
start_event["awaiting_confirmation"] = needs_confirm
|
|
|
|
try:
|
|
yield {"type": "status", "text": decision.status_text}
|
|
yield start_event
|
|
|
|
if (
|
|
decision_slot is not None
|
|
and wait_tool_decision(
|
|
decision_slot,
|
|
approval_id,
|
|
cancel_event = cancel_event,
|
|
)
|
|
== "deny"
|
|
):
|
|
decision_slot = None
|
|
resolved_provisional_tool_call_ids.add(decision.tool_call_id)
|
|
yield {
|
|
"type": "tool_end",
|
|
"tool_name": decision.tool_name,
|
|
"tool_call_id": decision.tool_call_id,
|
|
"result": TOOL_REJECTED_MESSAGE,
|
|
"provenance": decision.provenance,
|
|
}
|
|
denied_message = {
|
|
"role": "tool",
|
|
"name": decision.tool_name,
|
|
"content": TOOL_REJECTED_MESSAGE,
|
|
}
|
|
if decision.tool_call_id:
|
|
denied_message["tool_call_id"] = decision.tool_call_id
|
|
conversation.append(denied_message)
|
|
if _forced_tool_call_pending:
|
|
_forced_tool_call_pending = False
|
|
continue
|
|
decision_slot = None
|
|
finally:
|
|
if decision_slot is not None:
|
|
abort_tool_decision(decision_slot, approval_id)
|
|
|
|
_effective_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout
|
|
# RAG: cap paraphrased KB re-searches that slip past the dup guard.
|
|
if (
|
|
decision.tool_name == "search_knowledge_base"
|
|
and _kb_search_count >= RAG_MAX_SEARCHES_PER_TURN
|
|
):
|
|
result = RAG_SEARCH_CAP_NUDGE
|
|
else:
|
|
result = execute_tool(
|
|
decision.tool_name,
|
|
decision.arguments,
|
|
cancel_event = cancel_event,
|
|
timeout = _effective_timeout,
|
|
session_id = session_id,
|
|
rag_scope = rag_scope,
|
|
disable_sandbox = bypass_permissions,
|
|
)
|
|
if decision.tool_name == "search_knowledge_base":
|
|
_kb_search_count += 1
|
|
completion = tool_controller.record_result(decision, result)
|
|
resolved_provisional_tool_call_ids.add(decision.tool_call_id)
|
|
# A tool ran this turn, so it counts against the caller's budget.
|
|
_turn_executed_real_tool = True
|
|
yield completion.tool_end_event()
|
|
conversation.append(completion.tool_message())
|
|
|
|
if _forced_tool_call_pending:
|
|
_forced_tool_call_pending = False
|
|
|
|
# Close provisional cards not resolved by execution/no-op handling.
|
|
for _pid, _pname in provisional_started_tool_calls.items():
|
|
if _pid not in resolved_provisional_tool_call_ids:
|
|
resolved_provisional_tool_call_ids.add(_pid)
|
|
yield {
|
|
"type": "tool_end",
|
|
"tool_name": _pname,
|
|
"tool_call_id": _pid,
|
|
"result": "",
|
|
"provenance": tool_event_provenance(provisional = True),
|
|
}
|
|
|
|
# Clear tool status badge before next generation/final pass.
|
|
yield {"type": "status", "text": ""}
|
|
if tool_controller.force_final_answer or not tool_controller.active_tools():
|
|
_append_budget_exhausted_nudge = False
|
|
break
|
|
# Count only real tool turns against the cap so reserved re-prompt slots can't become
|
|
# extra tool rounds; a no-op correction turn doesn't consume budget (GGUF parity).
|
|
if _turn_executed_real_tool:
|
|
_tool_iters_done += 1
|
|
if _tool_iters_done >= max_tool_iterations:
|
|
break
|
|
continue
|
|
|
|
except httpx.ConnectError:
|
|
# Mark unresolved provisional cards as failed before raising.
|
|
for _pid, _pname in provisional_started_tool_calls.items():
|
|
if _pid not in resolved_provisional_tool_call_ids:
|
|
resolved_provisional_tool_call_ids.add(_pid)
|
|
yield {
|
|
"type": "tool_end",
|
|
"tool_name": _pname,
|
|
"tool_call_id": _pid,
|
|
"result": "Error: lost connection to llama-server before the tool call completed.",
|
|
"provenance": tool_event_provenance(provisional = True),
|
|
}
|
|
raise RuntimeError("Lost connection to llama-server")
|
|
except Exception as e:
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
return
|
|
# Same cleanup for other mid-iteration failures.
|
|
for _pid, _pname in provisional_started_tool_calls.items():
|
|
if _pid not in resolved_provisional_tool_call_ids:
|
|
resolved_provisional_tool_call_ids.add(_pid)
|
|
yield {
|
|
"type": "tool_end",
|
|
"tool_name": _pname,
|
|
"tool_call_id": _pid,
|
|
"result": "Error: the tool call was interrupted before it completed.",
|
|
"provenance": tool_event_provenance(provisional = True),
|
|
}
|
|
raise
|
|
|
|
# ── Tool iteration cap reached -- synthesize final answer ──
|
|
# The model used all iterations without a final text response. Nudge
|
|
# the final streaming pass to produce a useful answer instead of
|
|
# continuing to request tools.
|
|
if max_tool_iterations > 0 and _append_budget_exhausted_nudge:
|
|
conversation.append(
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
"You have used all available tool calls. Based on "
|
|
"everything you have found so far, provide your final "
|
|
"answer now. Do not call any more tools."
|
|
),
|
|
}
|
|
)
|
|
|
|
# Clear status.
|
|
yield {"type": "status", "text": ""}
|
|
|
|
# Final streaming pass with the full conversation context.
|
|
stream_payload = {
|
|
"messages": conversation,
|
|
"stream": True,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"top_k": top_k if top_k >= 0 else 0,
|
|
"min_p": min_p,
|
|
"repeat_penalty": repetition_penalty,
|
|
"presence_penalty": presence_penalty,
|
|
}
|
|
_reasoning_kw = self._request_reasoning_kwargs(
|
|
enable_thinking, reasoning_effort, preserve_thinking
|
|
)
|
|
if _reasoning_kw is not None:
|
|
stream_payload["chat_template_kwargs"] = _reasoning_kw
|
|
stream_payload["max_tokens"] = (
|
|
max_tokens
|
|
if max_tokens is not None
|
|
else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR)
|
|
)
|
|
if stop:
|
|
stream_payload["stop"] = stop
|
|
if seed is not None:
|
|
stream_payload["seed"] = seed
|
|
stream_payload["stream_options"] = {"include_usage": True}
|
|
|
|
cumulative = ""
|
|
_last_emitted = ""
|
|
in_thinking = False
|
|
has_content_tokens = False
|
|
reasoning_text = ""
|
|
_final_reasoning_started_at: Optional[float] = None
|
|
_final_reasoning_summary_emitted = False
|
|
_metadata_usage = None
|
|
_metadata_timings = None
|
|
_metadata_finish_reason = None
|
|
_stream_done = False
|
|
|
|
try:
|
|
with self._open_stream(url, stream_payload, cancel_event) as (
|
|
response,
|
|
first_token_deadline,
|
|
):
|
|
buffer = ""
|
|
for raw_chunk in self._iter_text_cancellable(
|
|
response,
|
|
cancel_event,
|
|
first_token_deadline = first_token_deadline,
|
|
):
|
|
buffer += raw_chunk
|
|
while "\n" in buffer:
|
|
line, buffer = buffer.split("\n", 1)
|
|
line = line.strip()
|
|
|
|
if not line:
|
|
continue
|
|
if line == "data: [DONE]":
|
|
if in_thinking:
|
|
if (
|
|
_final_reasoning_started_at is not None
|
|
and not _final_reasoning_summary_emitted
|
|
):
|
|
_final_reasoning_summary_emitted = True
|
|
yield _reasoning_summary_event(_final_reasoning_started_at)
|
|
if has_content_tokens:
|
|
cumulative += "</think>"
|
|
yield {
|
|
"type": "content",
|
|
"text": _strip_tool_markup(cumulative, final = True),
|
|
}
|
|
else:
|
|
cumulative = reasoning_text
|
|
yield {"type": "content", "text": cumulative}
|
|
_stream_done = True
|
|
break # exit inner while
|
|
if not line.startswith("data: "):
|
|
continue
|
|
|
|
try:
|
|
chunk_data = json.loads(line[6:])
|
|
# Capture server timings/usage from final chunks.
|
|
_chunk_timings = chunk_data.get("timings")
|
|
if _chunk_timings:
|
|
_metadata_timings = _chunk_timings
|
|
_chunk_usage = chunk_data.get("usage")
|
|
if _chunk_usage:
|
|
_metadata_usage = _chunk_usage
|
|
choices = chunk_data.get("choices", [])
|
|
if choices:
|
|
delta = choices[0].get("delta", {})
|
|
_fr = choices[0].get("finish_reason")
|
|
if _fr:
|
|
_metadata_finish_reason = _fr
|
|
|
|
reasoning = delta.get("reasoning_content", "")
|
|
if reasoning:
|
|
if _final_reasoning_started_at is None:
|
|
_final_reasoning_started_at = time.monotonic()
|
|
reasoning_text += reasoning
|
|
if not in_thinking:
|
|
cumulative += "<think>"
|
|
in_thinking = True
|
|
cumulative += reasoning
|
|
yield {"type": "content", "text": cumulative}
|
|
|
|
token = delta.get("content", "")
|
|
if token:
|
|
if (
|
|
_final_reasoning_started_at is not None
|
|
and not _final_reasoning_summary_emitted
|
|
):
|
|
_final_reasoning_summary_emitted = True
|
|
yield _reasoning_summary_event(_final_reasoning_started_at)
|
|
has_content_tokens = True
|
|
if in_thinking:
|
|
cumulative += "</think>"
|
|
in_thinking = False
|
|
cumulative += token
|
|
cleaned = _strip_tool_markup(cumulative)
|
|
# Emit only when cleaned text grows (monotonic).
|
|
if len(cleaned) > len(_last_emitted):
|
|
_last_emitted = cleaned
|
|
yield {"type": "content", "text": cleaned}
|
|
except json.JSONDecodeError:
|
|
logger.debug(f"Skipping malformed SSE line: {line[:100]}")
|
|
if _stream_done:
|
|
break # exit outer for
|
|
_meta = _build_metadata_event(
|
|
_metadata_usage, _metadata_timings, _metadata_finish_reason
|
|
)
|
|
if _meta is not None:
|
|
yield _meta
|
|
|
|
except httpx.ConnectError:
|
|
raise RuntimeError("Lost connection to llama-server")
|
|
except Exception as e:
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
return
|
|
raise
|
|
|
|
# ── Prompt token counting ──────────────────────────────────
|
|
|
|
def count_chat_tokens(
|
|
self,
|
|
messages,
|
|
system = None,
|
|
tools = None,
|
|
strict: bool = False,
|
|
) -> int:
|
|
"""Count prompt tokens for a chat request via llama-server.
|
|
|
|
Non-strict callers keep the historical best-effort behavior and receive
|
|
0 when a count cannot be determined. Strict callers (public count_tokens
|
|
endpoints) get an exception instead of a successful-looking zero when
|
|
tokenizer/template calls fail or a multimodal prompt would fall back to a
|
|
text-only approximation.
|
|
"""
|
|
if not self.is_loaded:
|
|
if strict:
|
|
raise RuntimeError("llama-server is not loaded")
|
|
return 0
|
|
|
|
def _has_non_text_content(content) -> bool:
|
|
if isinstance(content, list):
|
|
for block in content:
|
|
if isinstance(block, str):
|
|
continue
|
|
if not isinstance(block, dict):
|
|
return True
|
|
if isinstance(block.get("text"), str):
|
|
continue
|
|
return True
|
|
return False
|
|
|
|
def _has_non_text_prompt_parts() -> bool:
|
|
if _has_non_text_content(system):
|
|
return True
|
|
for msg in messages or []:
|
|
if isinstance(msg, dict) and _has_non_text_content(msg.get("content", "")):
|
|
return True
|
|
return False
|
|
|
|
def _block_text(content) -> str:
|
|
if isinstance(content, str):
|
|
return content
|
|
if isinstance(content, list):
|
|
parts = []
|
|
for block in content:
|
|
if isinstance(block, dict):
|
|
if isinstance(block.get("text"), str):
|
|
parts.append(block["text"])
|
|
elif isinstance(block, str):
|
|
parts.append(block)
|
|
return "".join(parts)
|
|
return ""
|
|
|
|
# Normalize system into a leading message / plain text.
|
|
system_text = ""
|
|
if isinstance(system, str):
|
|
system_text = system
|
|
elif isinstance(system, list):
|
|
system_text = _block_text(system)
|
|
|
|
try:
|
|
with httpx.Client(timeout = 10, headers = self._auth_headers, trust_env = False) as client:
|
|
|
|
def _tokenize(text: str) -> int:
|
|
r = client.post(
|
|
f"{self.base_url}/tokenize",
|
|
json = {"content": text, "add_special": True},
|
|
)
|
|
if r.status_code != 200:
|
|
if strict:
|
|
raise RuntimeError("llama-server tokenizer failed")
|
|
return 0
|
|
tokens = r.json().get("tokens", [])
|
|
if not isinstance(tokens, list):
|
|
if strict:
|
|
raise RuntimeError("llama-server tokenizer returned invalid tokens")
|
|
return 0
|
|
return len(tokens)
|
|
|
|
# 1. Try /apply-template to render the real chat prompt.
|
|
template_messages = list(messages) if messages else []
|
|
if system_text:
|
|
template_messages = [
|
|
{"role": "system", "content": system_text}
|
|
] + template_messages
|
|
apply_template_failed = False
|
|
try:
|
|
# llama-server's /apply-template renders tool declarations
|
|
# into the prompt when ``tools`` is supplied, so pass them
|
|
# through, otherwise tool-schema tokens go uncounted.
|
|
template_body = {"messages": template_messages}
|
|
if tools:
|
|
template_body["tools"] = tools
|
|
resp = client.post(
|
|
f"{self.base_url}/apply-template",
|
|
json = template_body,
|
|
)
|
|
if resp.status_code == 200:
|
|
prompt = resp.json().get("prompt", "")
|
|
if isinstance(prompt, str):
|
|
return _tokenize(prompt)
|
|
apply_template_failed = True
|
|
except Exception:
|
|
apply_template_failed = True
|
|
|
|
if strict and apply_template_failed and _has_non_text_prompt_parts():
|
|
raise RuntimeError(
|
|
"cannot fall back to text-only token counting for multimodal messages"
|
|
)
|
|
|
|
# 2. Fallback: concatenate plain text and tokenize. Append a
|
|
# serialized form of the tools so they still contribute to the
|
|
# count when /apply-template is unavailable.
|
|
parts = []
|
|
if system_text:
|
|
parts.append(system_text)
|
|
for msg in messages or []:
|
|
if isinstance(msg, dict):
|
|
parts.append(_block_text(msg.get("content", "")))
|
|
if tools:
|
|
try:
|
|
parts.append(json.dumps(tools, ensure_ascii = False))
|
|
except Exception:
|
|
pass
|
|
return _tokenize("\n".join(p for p in parts if p))
|
|
except Exception:
|
|
if strict:
|
|
raise
|
|
return 0
|
|
|
|
# ── TTS support ────────────────────────────────────────────
|
|
|
|
def detect_audio_type(self) -> Optional[str]:
|
|
"""Detect audio/TTS codec; swallows errors (use _strict to distinguish)."""
|
|
try:
|
|
return self._detect_audio_type_strict()
|
|
except Exception as e:
|
|
logger.debug(f"Audio type detection failed: {e}")
|
|
return None
|
|
|
|
def _apply_detected_audio(self, detected: Optional[str]) -> bool:
|
|
"""Apply a probed audio codec under self._lock. Returns True to continue
|
|
the load (codec inited OK, or nothing to init), False to abort (server
|
|
unhealthy or codec init failed). Shared by the fast-path retry and the
|
|
main load path."""
|
|
if detected in ("snac", "bicodec", "dac"):
|
|
with self._lock:
|
|
if not self._healthy:
|
|
return False
|
|
try:
|
|
self.init_audio_codec(detected)
|
|
self._is_audio = True
|
|
self._audio_type = detected
|
|
except Exception as exc:
|
|
# Surface as HTTP 500 (matches pre-PR contract).
|
|
logger.warning("Failed to init audio codec '%s': %s", detected, exc)
|
|
self._audio_probed = False
|
|
return False
|
|
elif detected:
|
|
# csm / whisper / audio_vlm: track type but keep _is_audio False --
|
|
# GGUF TTS routing only fires for snac/bicodec/dac.
|
|
with self._lock:
|
|
if not self._healthy:
|
|
return False
|
|
self._audio_type = detected
|
|
# Audio input = token probe (audio_vlm/whisper) OR mmproj encoder.
|
|
from utils.models.model_config import is_audio_input_type
|
|
|
|
self._has_audio_input = bool(is_audio_input_type(self._audio_type)) or bool(
|
|
self._mmproj_has_audio
|
|
)
|
|
return True
|
|
|
|
def _detect_audio_type_strict(self) -> Optional[str]:
|
|
"""Codec name on match, None on non-audio, raises on transport/JSON errors."""
|
|
if not self.is_loaded:
|
|
return None
|
|
with httpx.Client(timeout = 10, headers = self._auth_headers, trust_env = False) as client:
|
|
|
|
def _detok(tid: int) -> str:
|
|
# Non-200 means "marker not in vocab" -- keep probing.
|
|
# Transport / JSON errors still raise.
|
|
r = client.post(f"{self.base_url}/detokenize", json = {"tokens": [tid]})
|
|
if r.status_code != 200:
|
|
return ""
|
|
return r.json().get("content", "")
|
|
|
|
def _tok(text: str) -> list[int]:
|
|
r = client.post(
|
|
f"{self.base_url}/tokenize",
|
|
json = {"content": text, "add_special": False},
|
|
)
|
|
if r.status_code != 200:
|
|
return []
|
|
return r.json().get("tokens", [])
|
|
|
|
# Codec-specific tokens (not generic ones that non-audio models may have)
|
|
if "<custom_token_" in _detok(128258) and "<custom_token_" in _detok(128259):
|
|
return "snac"
|
|
if len(_tok("<|AUDIO|>")) == 1 and len(_tok("<|audio_eos|>")) == 1:
|
|
return "csm"
|
|
if len(_tok("<|startoftranscript|>")) == 1:
|
|
return "whisper"
|
|
# Gemma 3n: <audio_soft_token>; Gemma 4: <|audio|> (not csm's <|AUDIO|>).
|
|
if len(_tok("<audio_soft_token>")) == 1 or len(_tok("<|audio|>")) == 1:
|
|
return "audio_vlm"
|
|
if len(_tok("<|bicodec_semantic_0|>")) == 1 and len(_tok("<|bicodec_global_0|>")) == 1:
|
|
return "bicodec"
|
|
if len(_tok("<|c1_0|>")) == 1 and len(_tok("<|c2_0|>")) == 1:
|
|
return "dac"
|
|
return None
|
|
|
|
# Prompt format per codec: (template, stop_tokens, needs_token_ids).
|
|
# Matches InferenceBackend._generate_snac/bicodec/dac.
|
|
_TTS_PROMPTS = {
|
|
"snac": (
|
|
"<custom_token_3>{text}<|eot_id|><custom_token_4>",
|
|
["<custom_token_2>"],
|
|
True,
|
|
),
|
|
"bicodec": (
|
|
"<|task_tts|><|start_content|>{text}<|end_content|><|start_global_token|>",
|
|
["<|im_end|>", "</s>"],
|
|
False,
|
|
),
|
|
"dac": (
|
|
"<|im_start|>\n<|text_start|>{text}<|text_end|>\n<|audio_start|><|global_features_start|>\n",
|
|
["<|im_end|>", "<|audio_end|>"],
|
|
False,
|
|
),
|
|
}
|
|
|
|
_codec_mgr = None # Shared AudioCodecManager instance
|
|
|
|
def init_audio_codec(self, audio_type: str) -> None:
|
|
"""Load the audio codec at model load time (mirrors the non-GGUF path)."""
|
|
import torch
|
|
from core.inference.audio_codecs import AudioCodecManager
|
|
|
|
if LlamaCppBackend._codec_mgr is None:
|
|
LlamaCppBackend._codec_mgr = AudioCodecManager()
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
model_repo_path = None
|
|
|
|
# BiCodec needs a repo with BiCodec/ weights -- download canonical SparkTTS
|
|
if audio_type == "bicodec":
|
|
from huggingface_hub import snapshot_download
|
|
import os
|
|
|
|
repo_path = snapshot_download("unsloth/Spark-TTS-0.5B", local_dir = "Spark-TTS-0.5B")
|
|
model_repo_path = os.path.abspath(repo_path)
|
|
|
|
LlamaCppBackend._codec_mgr.load_codec(audio_type, device, model_repo_path = model_repo_path)
|
|
logger.info(f"Loaded audio codec for GGUF TTS: {audio_type}")
|
|
|
|
def generate_audio_response(
|
|
self,
|
|
text: str,
|
|
audio_type: str,
|
|
temperature: float = 0.6,
|
|
top_p: float = 0.95,
|
|
top_k: int = 50,
|
|
min_p: float = 0.0,
|
|
max_new_tokens: int = 2048,
|
|
repetition_penalty: float = 1.1,
|
|
) -> tuple:
|
|
"""
|
|
Generate TTS audio via llama-server /completion + codec decode.
|
|
Returns (wav_bytes, sample_rate).
|
|
"""
|
|
if audio_type not in self._TTS_PROMPTS:
|
|
raise RuntimeError(f"GGUF TTS does not support '{audio_type}' codec.")
|
|
|
|
tpl, stop, need_ids = self._TTS_PROMPTS[audio_type]
|
|
|
|
payload: dict = {
|
|
"prompt": tpl.format(text = text),
|
|
"stream": False,
|
|
"n_predict": max_new_tokens,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"top_k": top_k if top_k >= 0 else 0,
|
|
"min_p": min_p,
|
|
"repeat_penalty": repetition_penalty,
|
|
}
|
|
if stop:
|
|
payload["stop"] = stop
|
|
if need_ids:
|
|
payload["n_probs"] = 1
|
|
|
|
with httpx.Client(
|
|
timeout = httpx.Timeout(300, connect = 10),
|
|
headers = self._auth_headers,
|
|
trust_env = False,
|
|
) as client:
|
|
resp = client.post(f"{self.base_url}/completion", json = payload)
|
|
if resp.status_code != 200:
|
|
raise RuntimeError(f"llama-server returned {resp.status_code}: {resp.text}")
|
|
|
|
data = resp.json()
|
|
token_ids = (
|
|
[p["id"] for p in data.get("completion_probabilities", []) if "id" in p]
|
|
if need_ids
|
|
else None
|
|
)
|
|
|
|
import torch
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
return LlamaCppBackend._codec_mgr.decode(
|
|
audio_type, device, token_ids = token_ids, text = data.get("content", "")
|
|
)
|