unsloth/studio/backend/core/inference/llama_cpp.py
Leo Borcherding b6d5636cc0
fix/strix halo and windows AMD ROCm support (#5301)
* fix(studio): set HIP_VISIBLE_DEVICES in apply_gpu_ids for ROCm training workers

Training workers are spawned via multiprocessing spawn before detect_hardware()
runs, so IS_ROCM is still False. If the user never set HIP_VISIBLE_DEVICES in
their shell, _inherits_rocm_visibility is also False, leaving the worker with
only CUDA_VISIBLE_DEVICES set. On ROCm hosts the HIP runtime honors
HIP_VISIBLE_DEVICES over CUDA_VISIBLE_DEVICES, so the worker saw the full
device list and torch raised "no usable HIP accelerator" on some setups.

Fall back to probing torch.version.hip (a build-time attribute, safe to read
before GPU init) to detect ROCm when neither IS_ROCM nor inherited env vars
are available. Mirrors the existing fix in llama_cpp.py for llama-server
subprocess GPU pinning.

Fixes https://github.com/unslothai/unsloth/issues/5180

* test: tighten apply_gpu_ids ROCm fallback assertions

Replace loose OR chain with exact string matches, split into three
focused tests, and add a guard check for the try/except wrapper.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: detect ROCm unified memory (Strix Halo / AMD iGPU) via torch fallback

amd-smi on iGPUs with shared/unified memory (e.g. Radeon 8060S on Strix
Halo) reports only the dedicated VRAM slice (~512 MB) in its metric output,
so get_visible_gpu_utilization() was returning usable_gb ≈ 0.35 GB instead
of the full GTT pool (~128 GB).  torch.cuda.mem_get_info() already surfaces
the correct unified-pool size.

Add _reconcile_rocm_unified_memory(): after amd-smi returns a valid result
on a ROCm device, cross-check each device's vram_total_gb against
torch.cuda.mem_get_info().  When torch reports a larger total, replace the
amd-smi VRAM fields in-place.  No-op for discrete AMD GPUs where the two
sources agree.

Fixes: "Falling back to all visible GPUs -- model may not fit" on AMD iGPU
machines even when 100+ GB of unified memory is available.

* Apply unified-memory reconciliation in get_gpu_utilization too

The visible-GPU path was already corrected for AMD iGPUs with unified memory
(Strix Halo / Radeon 8060S), but get_gpu_utilization was still returning the
raw 512 MB amd-smi VRAM slice. Studio's /api/train/hardware endpoint and the
live GPU monitor read from this primary path, so users continued seeing the
wrong total even after auto_select_gpu_ids picked the right device.

Refactor to share the per-device correction:
  * _apply_unified_memory_correction(metrics, torch_info) -- the actual
    replacement logic, in-place on a single metrics dict.
  * _reconcile_rocm_unified_memory(...)                   -- multi-device,
    iterates utilization["devices"] (visible-GPU path).
  * _reconcile_primary_rocm_unified_memory(...)           -- single flat
    metrics dict (primary-GPU path), uses parent_visible_spec to pick the
    primary index, falls back to ordinal 0 when no visibility env is set.

get_gpu_utilization now calls the primary reconciler under IS_ROCM, so both
endpoints surface the real unified-memory pool on iGPUs while leaving
discrete AMD GPUs untouched (torch_total <= smi_total -> no replace).

* Use 'is not None' and log debug on torch.version.hip probe failures

Two small follow-ups to the apply_gpu_ids ROCm fallback:

1. Match detect_hardware()'s 'getattr(torch.version, "hip", None) is not None'
   form so the entire codebase has one canonical 'this torch was built with
   HIP' check. On every shipping torch wheel hip is either None or a non-empty
   version string, so the new form agrees with the old bool() form on every
   real install.

2. Log the probe failure at debug level instead of swallowing it silently.
   The broad 'except Exception' is intentional (we never want apply_gpu_ids
   to crash a worker over a probe), but the silent pass made it impossible
   to tell whether the fallback was firing or being skipped.

* fix(studio): honour HIP_VISIBLE_DEVICES in _get_parent_visible_gpu_spec before IS_ROCM is set

When a user has HIP_VISIBLE_DEVICES set in their shell (e.g. "1" to select
GPU 1) but detect_hardware() has not yet run in the Studio parent process,
IS_ROCM is still False.  _get_parent_visible_gpu_spec() was gated on IS_ROCM
so it fell through to CUDA_VISIBLE_DEVICES (unset), saw all physical GPUs,
and auto-selected index 0.  apply_gpu_ids then overwrote HIP_VISIBLE_DEVICES
with "0", making the intended GPU invisible to ROCm torch in the worker,
which triggered the "no usable HIP accelerator" error (issue #5180).

Apply the same _inherits_rocm_visibility pattern already used in
apply_gpu_ids: check for HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES in the
environment regardless of IS_ROCM so the correct GPU index is preserved.

* fix(install): harden AMD ROCm GPU detection for multi-GPU and env-filtered setups

The previous rocminfo awk pattern could miss discrete GPUs on machines
where HIP_VISIBLE_DEVICES/ROCR_VISIBLE_DEVICES is used to mask an
integrated GPU — the env vars filter rocminfo output but may not
propagate into the install script subprocess, causing detection to
fail entirely.

Two changes:
- Tighten rocminfo pattern from /gfx[0-9]/ && !/gfx000/ to
  /gfx[1-9][0-9]/ — simpler and correctly excludes the CPU agent
  (gfx000) without a negative lookahead
- Add sysfs KFD topology fallback: reads
  /sys/class/kfd/kfd/topology/nodes/*/gpu_id which is a kernel-level
  view unaffected by HIP_VISIBLE_DEVICES or ROCR_VISIBLE_DEVICES

Fixes detection failure reported in Discord by Chains (gfx1201 + iGPU
machine where env var exclusion of the iGPU caused rocminfo to return
no usable device).

* Fix KFD sysfs awk fallback to read properties file

The fallback added by this PR reads /sys/class/kfd/kfd/topology/nodes/*/gpu_id
files but matches the literal token 'gpu_id' against their content. Those
files contain only a single decimal value (e.g. '0' for CPU agents, '50432'
for GPU agents), so the regex never matches and 'found' stays 0, making the
fallback a no-op on every host. The properties file in the same directory
contains key/value lines like 'gpu_id 50432' which is what the existing awk
pattern expects.

Reproduced with a synthetic sysfs layout: against gpu_id files awk exits 1;
against properties files awk exits 0 when any node reports gpu_id > 0.

* fix(setup.ps1): detect AMD ROCm GPU on Windows, bring to parity with setup.sh

setup.ps1 only checked nvidia-smi and fell straight to "gpu: none" on AMD
machines. setup.sh already probed rocminfo/amd-smi/hipconfig/hipinfo.

Add three-tier detection mirroring install_llama_prebuilt.py's detect_host():
1. hipinfo: gcnArchName in output confirms a real HIP GPU (not just SDK)
2. amd-smi list: "GPU: <digit>" data rows as fallback
3. WMI Win32_VideoController: last resort -- detects AMD GPU even without
   HIP SDK, then guides user to install it rather than silently going CPU

Also corrects the "none" message to mention AMD ROCm alongside NVIDIA so
users with AMD hardware understand the requirement.

Fixes: rohit-style install where Strix Halo (Radeon 8060S) showed
"gpu: none" even with the HIP SDK present.

* fix(install.ps1): detect AMD ROCm GPU on Windows, bring to parity with setup.ps1

install.ps1 had the same nvidia-smi-only GPU detection as setup.ps1 before
the setup.ps1 fix. Applies the same three-tier AMD detection:
1. hipinfo: gcnArchName confirms real HIP GPU
2. amd-smi list: GPU data rows as fallback
3. WMI Win32_VideoController: detects AMD GPU without HIP SDK and guides
   user to install it

Fixes: install.ps1 showing "gpu: none" while setup.ps1 correctly showed
"AMD GPU detected" on the same machine (reported by rohit, RX 7600 XT).

* fix(install.ps1): suppress 'No NVIDIA GPU detected' when AMD GPU is present

* feat: add Windows AMD ROCm PyTorch wheel installation

install_python_stack.py:
- Add _ROCM_WINDOWS_WHEEL_BASE and _ROCM_WINDOWS_RELEASES constants
  pointing to AMD repo.radeon.com (ROCm 7.2 -> torch 2.9.1+rocm7.2.1)
- Extend _ensure_rocm_torch() with a Windows branch: detects ROCm via
  _has_rocm_gpu() / _detect_rocm_version(), requires Python 3.12 (cp312
  is the only ABI AMD publishes for Windows), installs the direct wheel
  URL from repo.radeon.com

install.ps1:
- Capture ROCmVersion during AMD detection via hipconfig --version /
  amd-smi version (needed for wheel URL selection)
- After Get-TorchIndexUrl, add an AMD wheel override block: when HasROCm
  and Python 3.12 detected, set ROCmTorchWheelUrl to AMD wheel URL
- Expand torch install branch to handle ROCmTorchWheelUrl with
  uv pip install --force-reinstall --no-cache-dir

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: also install torchvision and torchaudio from AMD Windows repo

AMD publishes matching torchvision-0.24.1+rocm7.2.1 and
torchaudio-2.9.1+rocm7.2.1 cp312 wheels at the same repo.radeon.com
release folder. Install all three in both install.ps1 and
install_python_stack.py Windows ROCm path.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* feat: add ROCm 7.1.1 Windows wheel mapping

AMD uses a different version string for 7.1.1 wheels:
2.9.0+rocmsdk20251116 (date-tagged) instead of +rocm7.1.1.
Adds the 7.1.1 release folder to both install.ps1 and
install_python_stack.py so users with ROCm 7.1 get ROCm
torch instead of falling back to CPU.

* fix: install rocm_sdk_core and rocm_sdk_libraries_custom alongside torch

The AMD Windows torch wheels declare rocm[libraries]==<ver> as a hard
dependency. Without installing rocm_sdk_core and rocm_sdk_libraries_custom
from the same AMD release folder, uv cannot resolve the dependency and
fails with 'No solution found'. Include all 5 wheels in one install call.

* fix: expand ROCm wheel array to scalars for Invoke-InstallCommand

@array splatting inside a scriptblock only works when the native command
is prefixed with '&'. Invoke-InstallCommand uses '& $Command' to run the
block, so @ROCmAllWheelUrls was not being expanded. Extract to scalar
variables $rw0-$rw4 which are captured correctly by the closure.

* fix: use --no-deps for AMD Windows torch wheel install

uv's resolver looks up rocm[libraries]==0.1.dev0 on PyPI during
dependency resolution before downloading any wheels, and fails because
the package doesn't exist on PyPI. --no-deps skips resolution entirely
and installs all 5 AMD wheels directly. The GPU runtime dependency is
satisfied by the HIP SDK, not a Python package.

* fix: setup.ps1 and install_python_stack.py now install ROCm torch on Windows

setup.ps1 was always setting CuTag='cpu' for non-NVIDIA hosts and installing
cpu-only PyTorch, overwriting the ROCm torch installed by install.ps1.
Adds the same AMD wheel selection logic (ROCm version detection, Python 3.12
check, 5-wheel install with --no-deps) to setup.ps1's torch install block.

install_python_stack.py: remove IS_WINDOWS guard from _ensure_rocm_torch()
call site so the Windows path in _ensure_rocm_torch() is reachable during
'unsloth studio update' as well.

* fix: suppress manual-install warning when ROCm torch already present; fix progress counter

- Gate the 'must be installed manually' warning on torch.version.hip being empty
  so it doesn't fire when our ROCm torch install succeeded
- Update _TOTAL counter to include the 3 ROCm steps on Windows now that
  _ensure_rocm_torch() is called there (fixes 10/9 display)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* feat: add rocm step display in setup.ps1; fix warning and progress counter

- Add 'rocm' step after 'cuda' in setup.ps1 showing ROCm version or HIP SDK missing
- Move ROCm version detection up to GPU detection block so it's available early
- Suppress 'must be installed manually' warning when torch.version.hip is set
- Fix _TOTAL counter to include ROCm steps on Windows (fixes 10/9 display)

* fix: detect AMD SDK ROCm torch via __version__ when torch.version.hip is unset

AMD's repo.radeon.com wheels (e.g. 2.9.0+rocmsdk20251116) do not set
torch.version.hip, leaving it None. All three probes that relied solely on
torch.version.hip now also check for 'rocm' in torch.__version__.lower():

- hardware.py detect_hardware(): IS_ROCM was never set, causing the studio
  to report 'Hardware detected: CPU' even after AMD wheels were installed
  and HIP DLLs were on PATH.
- install_python_stack.py _ensure_rocm_torch(): skip-if-already-installed
  probe would always reinstall on subsequent runs.
- install_python_stack.py Windows AMD warning: suppression check always
  failed, so the 'must be installed manually' note kept appearing after
  a successful AMD wheel install.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* perf: drop --no-cache-dir from AMD ROCm torch wheel installs

uv caches downloaded wheels by default; passing --no-cache-dir forced a
full redownload of the ~2 GB torch wheel on every install run. CUDA installs
never had this flag -- AMD was the only path affected.

* fix: use install-state flag instead of subprocess probe for AMD Windows warning

Replace the subprocess torch probe in the post-install warning block with a
module-level _rocm_windows_torch_installed flag set by _ensure_rocm_torch().
Subprocess re-import of torch is unnecessary and fragile -- the install
function already knows whether it succeeded.

* fix: hoist global declaration to top of _ensure_rocm_torch

Python requires the global statement to appear before any assignment
to the variable within a function. Moving it to the function top fixes
the SyntaxError on line 354.

* fix: pass AMD torch install status via env var to suppress false warning

setup.ps1 now sets UNSLOTH_ROCM_TORCH_INSTALLED=1 after a successful AMD
wheel install. install_python_stack.py reads this at the top of
_ensure_rocm_torch() to skip both the subprocess probe and the warning --
no re-import of torch needed, and the warning message now correctly says
'could not be auto-installed' rather than 'must be installed manually'.

* fix: register ROCm DLL directory before torch import on Windows

Python 3.8+ ignores PATH for extension DLL loading on Windows; amdhip64.dll
and other HIP runtime DLLs must be registered via os.add_dll_directory().
Without this, torch.cuda.is_available() always returns False on AMD ROCm
Windows even when HIP_PATH is correctly set in system environment variables.

Reads HIP_PATH / ROCM_PATH env vars first, then falls back to scanning
common ROCm install roots (C:\Program Files\AMD\ROCm, F:\ROCm, C:\ROCm).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: remove hardcoded non-standard ROCm paths from DLL directory scan

Only use HIP_PATH/ROCM_PATH (set by AMD installer) and the standard
C:\Program Files\AMD\ROCm\<version>\bin location. Custom drive paths
like F:\ROCm are user-specific and should not be hardcoded.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: prevent torchao overrides step from overwriting AMD ROCm torch

torchao==0.14.0 in overrides.txt declares torch as a dependency. Without
--no-deps, uv resolves torch from PyPI and installs 2.11.0+cpu on top of
the AMD ROCm wheels (2.9.0+rocmsdk20251116). This was the root cause of
'Hardware detected: CPU' -- the AMD wheels were installed but then
immediately overwritten by the overrides step.

When _rocm_windows_torch_installed is True, add --no-deps to the overrides
pip_install call so torchao is installed without pulling in CPU torch.

* fix: add rocm_sdk namespace tarball to Windows ROCm wheel installs

torch/_rocm_init.py calls `import rocm_sdk` at startup, which requires
the rocm namespace tarball (rocm-*.tar.gz) in addition to the SDK wheel
packages. This tarball was missing from both install.ps1 and setup.ps1,
causing ModuleNotFoundError on first torch import.

- Add rocm-0.1.dev0.tar.gz to ROCm 7.1.1 install (provides rocm_sdk namespace)
- Add rocm-7.2.1.tar.gz + rocm_sdk_devel to ROCm 7.2.1 install
- Install tarball in a dedicated step before main SDK/torch wheels
- Switch to @array splatting in install.ps1 scriptblock for dynamic wheel count
- Remove --no-cache-dir from Python-side ROCm wheel install (prevents ~2GB redownload)

* feat: enable ROCm 7.2 torch install + warn on gfx1151 with ROCm < 7.2

Chigoma333 (AMD Radeon 8060S / gfx1151, Strix Halo) confirmed that ROCm
7.1 segfaults when tensors are moved to GPU, but ROCm 7.2 + torch
2.11.0+rocm7.2 works fully including training.

Changes:
- Uncomment (7,2): "rocm7.2" in _ROCM_TORCH_INDEX (was blocked by <2.11.0)
- Add _ROCM_TORCH_PKG_SPECS dict with per-tag version bounds:
  rocm7.2 → torch>=2.11.0,<2.12.0; all older tags → <2.11.0
- Add _detect_amd_gfx_codes() helper that parses rocminfo output
- Warn on gfx1151/gfx1150 (Strix Halo) when ROCm < 7.2 is installed,
  pointing users at the known segfault and recommending upgrade
- install.sh get_torch_index_url(): enable rocm7.2 case (previously capped
  to rocm7.1), cap unknown future tags to rocm7.2
- install.sh: override TORCH_CONSTRAINT to >=2.11.0,<2.12.0 when rocm7.2
  index is selected, so pip can actually resolve torch 2.11.0

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: prefer Python 3.12 for AMD ROCm users when 3.13 is also installed

After GPU detection, if ROCm HIP SDK is found and the selected Python
is not 3.12, run a second pass to locate a 3.12 install via py.exe and
PATH (catches uv-managed installs). Switch $DetectedPython to 3.12 so
the venv is created with a compatible interpreter for the cp312-only AMD
Windows torch wheels.

NVIDIA and Intel GPU paths are unaffected -- the re-detection block only
runs when $HasROCm is true.

Fixes: #5301

* fix: also check uv-managed Python 3.12 for AMD ROCm #5301

* fix: hide amd-smi console popups on Windows, guard torch.distributed.is_initialized for ROCm #5301

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: suppress remaining console popups on Windows, patch torch.distributed.is_initialized for ROCm #5301

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: stub all missing torch.distributed attrs for ROCm Windows wheel #5301

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: inject torch.distributed stub when C backend missing in ROCm Windows wheel #5301

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(rocm/windows): pre-stub torch._C._distributed_c10d + raise amd-smi timeout

Two fixes for Windows ROCm regressions reported by electroglyph on #5301:

1. worker.py — torch.distributed stub now fires unconditionally on Windows
   The previous stub only injected sys.modules in the except branch, meaning
   it was silently skipped when `import torch.distributed` happened to succeed
   (the C backend is lazily resolved).  The crash then hit later when
   transformers/trl triggered the lazy load.  Fix: on win32 we pre-populate
   sys.modules['torch._C._distributed_c10d'] AND set the attribute on the
   torch._C extension module *before* attempting the import, covering both
   the early-ImportError and lazy-load failure modes.

2. amd.py — increase amd-smi timeout from 5 s to 30 s on Windows (10 s Linux)
   amd-smi on Windows must cold-init the ROCm runtime on first invocation;
   5 s was consistently too short, producing repeated 'Command timed out'
   warnings in the server log.  30 s gives enough headroom without blocking
   indefinitely on broken installs.

3. install.ps1 — widen Python 3.12 enforcement to ROCmGpuLabel (WMI-only path)
   Users whose HIP SDK is not on PATH were detected via WMI but not switched
   to Python 3.12 before the install started, causing a second pass.  Guard
   now fires on (HasROCm -or ROCmGpuLabel).

* fix(rocm): guard c10d stub, fix TorchIndexFamily for 7.1, clean dead code + comments

- worker.py: wrap c10d stub injection in `if _c10d_key not in sys.modules` so
  Windows NVIDIA users with a real torch.distributed are never affected
- install.ps1: fix Get-TauriTorchIndexFamily receiving hardcoded "rocm7.2"
  even when ROCm 7.1 wheels are installed; now branches on $ROCmVersion
- main.py: remove dead `import ctypes as _ctypes` (ctypes is never called)
- hardware.py, install_python_stack.py, worker.py, install.ps1: shorten
  verbose multi-line comment blocks throughout
- tests: update 4 stale assertions that expected rocm7.2 to be absent/capped

* fix(tests): match windows AMD warning assertion to actual source string

* chore: trim verbose comment blocks across all ROCm-related files

* fix: guard reconcile call against None numeric_ids; add torchvision lower bounds

* fix(install.ps1): recreate venv with Python 3.12 after ROCm switch

Venv was created with 3.13 before GPU detection ran; switching
$DetectedPython to 3.12 had no effect since $VenvPython still
pointed to the 3.13 interpreter inside the already-created venv.

* ux: detect AMD GPU before Python selection to avoid double venv creation

- Early hipinfo + WMI probe runs before Find-CompatiblePython so Python
  3.12 is selected upfront when AMD is detected; venv is now created
  exactly once instead of 3.13 then immediately 3.12.
- Post-venv recreation block replaced with a simple warning for the rare
  case where AMD was missed by the early probe.
- setup.ps1: show venv's actual Python version (e.g. 3.12) instead of
  the system Python found by the pre-activation search (was showing 3.13).

* fix(rocm/win): auto-stub all _distributed_c10d symbols via PEP-562 __getattr__

The bare ModuleType stub caused ImportError when torch._dynamo was imported
(triggered by trainer.py accessing torch._dynamo.config at load time).
torch._dynamo pulls in torch.distributed.fsdp._flat_param which does:
  from torch._C._distributed_c10d import FakeProcessGroup
and potentially other symbols. Adding module __getattr__ auto-creates a
stub class for any missing symbol so all such imports succeed without
enumerating every individual symbol. Applied to both the primary stub
and the fallback stub in the except branch.

* chore: trim c10d stub comment

* fix(rocm/win): auto-stub missing torch.distributed attrs (Store, ProcessGroup, …)

* fix(rocm/win): pre-stub fsdp submodules in sys.modules; fix __getattr__ subpackage clash

* feat(rocm/win): arch-aware wheel selector always picks newest ROCm release

Replace HIP-SDK-version-gated wheel selection with GPU arch-based logic.
Select-ROCmWheelRelease (PS) and _select_windows_rocm_release (Python) map
gcnArchName → minimum ROCm version, then pick the newest available release
that satisfies it (currently always rocm-rel-7.2.1 for any supported GPU).
Wheels bundle their own ROCm runtime so the installed HIP SDK 7.1 does not
prevent using 7.2.1 wheels on gfx1200 (RX 9060 XT) and similar RDNA 4 GPUs.

Also installs the bitsandbytes Windows ROCm continuous-release wheel and sets
BNB_ROCM_VERSION=72 in worker.py before ML imports so bnb loads the
libbitsandbytes_rocm72.dll that ships in that wheel.

* fix(rocm/win): stub class metaclass for ProcessGroup.BackendType; amd-smi circuit breaker

torchao.float8.inference accesses ProcessGroup.BackendType as a class-level
attribute.  Plain type() stubs have no __getattr__ on the metaclass so this
raises AttributeError.  Introduce _StubClassMeta whose __getattr__ returns
child stub classes, fixing the torchao import chain.

Add an amd-smi circuit breaker in amd.py: after 3 consecutive failures the
module stops spawning the process, eliminating the repeated Windows UAC /
DiskPart elevation prompts caused by polling a non-functional amd-smi.

Also guard BNB_ROCM_VERSION=72 behind a DLL existence check so bitsandbytes
fails with its own detection message rather than a harder "DLL not found" when
the Windows ROCm bnb wheel is not yet installed.

* fix: stub __members__ so torchao float8 enum check doesn't crash on ROCm Windows

torchao.float8.inference accesses ProcessGroup.BackendType.__members__
expecting a Python Enum registry dict. _StubClassMeta.__getattr__ was
blocking all dunder attributes, causing AttributeError. Return {} for
__members__ specifically so the isinstance/iteration checks pass cleanly.

* fix: stub distributed tensor/functional_collectives to prevent missing C++ op crash on ROCm Windows

torch._dynamo.trace_rules eagerly loads torch.distributed.tensor at import
time, which pulls in _functional_collectives.py. That file registers Meta
kernels for _c10d_functional C++ ops, but those ops are only registered
by torch._C._distributed_c10d — a C extension absent from ROCm Windows
wheels. Pre-stubbing the affected modules in sys.modules prevents the real
import chain from running and avoids the "operator does not exist" crash.

* fix: give mod stubs __path__ and pre-stub _tensor to fix 'not a package' import error

_make_mod_stub now sets __path__=[] so Python treats stub modules as
packages. Without it, any import of a submodule raises "is not a package".
Also pre-stub torch.distributed._tensor and its submodules so that
_tensor/__init__.py (which re-exports from torch.distributed.tensor) never
runs and torchao's `from torch.distributed._tensor import DTensor` gets a
harmless stub instead of crashing.

* fix: stub torch.ops._c10d_functional namespace with hashable op sentinels

torchao.dtypes.nf4tensor uses _c10d_functional ops as dict keys at import
time (all_gather_into_tensor.default, wait_tensor.default) and
torch.ops.c10d.scatter_.default. None of these ops are registered on ROCm
Windows because torch._C._distributed_c10d (the C extension) doesn't ship.
Replace the whole _c10d_functional namespace with a custom stub whose ops
return hashable .default objects, so dict-key construction doesn't crash.
Also inject a scatter_ stub into torch.ops.c10d if it's missing.

* fix: stub entire torchao package on ROCm Windows instead of individual ops

torchao is not supported on ROCm Windows and its import chain transitively
requires torch._C._distributed_c10d (absent from the ROCm Windows wheel).
Rather than stub each missing op one by one, stub the whole torchao package
upfront. Unsloth uses bitsandbytes for quantization, not torchao, so this
has no functional impact. transformers gracefully handles an importable-but-
empty torchao by disabling TorchAoHfQuantizer.

* fix: set __spec__ on mod stubs so importlib.util.find_spec doesn't raise

Manually-injected sys.modules entries have __spec__=None by default.
importlib.util.find_spec() raises ValueError when it finds a module in
sys.modules with __spec__=None (transformers.utils.import_utils hits this
when checking if torchao is available). Give every stub a minimal
ModuleSpec(name, loader=None, is_package=True) to satisfy find_spec.

* fix: add meta path finder to auto-stub subpackages of stub modules

`import torchao.prototype` goes through the import machinery, not
__getattr__, so an empty __path__ means ModuleNotFoundError. Rather than
list every submodule explicitly, register a MetaPathFinder that intercepts
any import whose parent is one of our stubs (detected by loader=None in the
parent's ModuleSpec). Real installed packages always have a SourceFileLoader
so they are never intercepted. Also register child stubs in sys.modules
from __getattr__ as a belt-and-suspenders measure.

* fix: use _unsloth_stub sentinel instead of loader=None for stub detection

The import machinery overwrites module.__spec__ with the spec returned by
find_spec (which has loader=_StubSubpackageLoader, not None), so the
loader=None check broke for second-level subpackages. Switch to a custom
_unsloth_stub object identity sentinel set directly on each stub module --
it survives __spec__ being replaced and correctly identifies stubs at any
depth (torchao.prototype.safetensors, etc.).

* refactor(rocm/win): switch to repo.amd.com arch-aware index, remove stubs

AMD recommends repo.amd.com/rocm/whl/{arch}/ as the Windows ROCm wheel
source. These wheels bundle their own ROCm runtime, support all Python
versions (not just cp312), and include the full torch._C extension set
(including _distributed_c10d) that the old repo.radeon.com wheel omitted.

Changes:
- install.ps1: remove Select-ROCmWheelRelease + hardcoded cp312 wheel
  URLs; remove Python 3.12 forced-preference logic; install via
  --index-url repo.amd.com/rocm/whl/{arch-family}/
- studio/setup.ps1: same -- remove Select-ROCmWheelRelease, switch to
  repo.amd.com arch-aware index URL
- studio/install_python_stack.py: replace _ROCM_WINDOWS_RELEASES /
  _select_windows_rocm_release with _windows_rocm_index_url() using the
  _GFX_TO_AMD_INDEX_ARCH map; drop Python 3.12 restriction
- studio/backend/core/training/worker.py: remove all stub machinery
  (_make_mod_stub, _StubSubpackageFinder, _StubSubpackageLoader,
  _StubClassMeta, torchao/fsdp/dtensor stubs, _c10d_functional ops
  stubs, BNB DLL detection) -- no longer needed with new wheel source

* fix(rocm/win): restore _distributed_c10d + torchao stubs; fix BNB install

repo.amd.com torch wheels also omit torch._C._distributed_c10d on Windows
(RCCL is not shipped on Windows). torch/distributed/__init__.py imports
from it unconditionally at module level, so the stub must land in
sys.modules before any torch.distributed import.

torchao (pulled in by transformers.quantizers) walks
torchao.float8.distributed_utils -> torch.distributed._functional_collectives
-> distributed_c10d at import time. Stubbing torchao up-front short-circuits
that chain.

worker.py:
- Restore _make_mod_stub / _StubSubpackageFinder / _StubSubpackageLoader
- Restore _StubClassMeta for ProcessGroup.BackendType attribute access
- Restore _distributed_c10d stub with __getattr__ (Windows only)
- Restore torchao stubs (5 modules, Windows only)

install_python_stack.py:
- BNB AMD wheel install was inside the early-return branch that fires when
  torch is already a ROCm build (installed by install.ps1). Move BNB install
  outside that branch so it always runs on Windows ROCm — the PyPI
  bitsandbytes has only CUDA DLLs and fails to load on ROCm.

* worker: remove _distributed_c10d stub; stub only torchao

The installed torch/distributed/__init__.py from repo.amd.com
(torch==2.10.0+rocm7.12.0) is now properly guarded with
`if is_available():`, so `import torch.distributed` alone is safe.

The crash only comes via torchao's import chain:
  torchao.float8.distributed_utils
    → torch.distributed._functional_collectives (unguarded import)
    → torch.distributed.distributed_c10d
    → torch._C._distributed_c10d  ← absent on Windows ROCm

Stubbing torchao short-circuits the chain entirely. No need to stub
_distributed_c10d. Remove _StubClassMeta and the _c10d stub block;
keep only _make_mod_stub + _StubSubpackageFinder + torchao seeds.

* fix: BNB AMD wheel skipped + torch.compile segfault on Windows ROCm

install_python_stack.py: the UNSLOTH_ROCM_TORCH_INSTALLED=1 early-return
path (set by setup.ps1 when it installed torch itself) returned before
ever reaching the AMD BNB prerelease wheel install.  The PyPI
bitsandbytes==0.49.x ships only CUDA DLLs, so loading it on ROCm fails
with "libbitsandbytes_rocm72.dll not found".  Now installs the AMD
Windows BNB wheel before returning on that path too.

worker.py: torch._grouped_mm crashes on gfx1200 (null HIP kernel pointer,
0xC0000005) when torch.compile's JitDecomp system dispatches it during
the first forward pass.  Detect Windows ROCm via torch.version.hip
(already in sys.modules from section 1e) and set TORCHDYNAMO_DISABLE=1
to bypass the broken kernel dispatch.

* fix: BNB AMD wheel install fails uv wheel filename check

The bitsandbytes continuous-release wheel is intentionally mismatched:
filename encodes 1.33.7.preview (= 1.33.7rc0 in PEP 440) but wheel
metadata reports 0.50.0.dev0.  uv rejects this by default.

Introduce _install_bnb_windows_rocm() helper that sets
UV_SKIP_WHEEL_FILENAME_CHECK=1 only for this specific install, then
restores the previous env value.  Both BNB install call sites (the
UNSLOTH_ROCM_TORCH_INSTALLED early-return path and the normal Windows
ROCm path) now use this helper.

* worker: patch _grouped_mm CUDA dispatch on Windows ROCm (gfx1200 null kernel)

TORCHDYNAMO_DISABLE=1 stopped the compiler frontend but not the autograd
JitDecomp system, which also dispatches _grouped_mm and hits the same
null HIP kernel crash (0xC0000005).

Verified that torch.library.Library("aten","IMPL").impl("_grouped_mm", fn,
"CUDA") successfully overrides the broken HIP kernel with a Python mm
fallback on torch==2.10.0+rocm7.12.0.

Schema: _grouped_mm(Tensor self, Tensor mat2, Tensor? offs=None,
                    Tensor? bias=None, ScalarType? out_dtype=None) -> Tensor

The fallback handles both the simple case (offs=None → torch.mm) and the
grouped case (offs provided → split self by offsets, multiply each group
against the corresponding slice of mat2, then cat results).

Keep _WINDOWS_ROCM_GROUPED_MM_LIB alive at function scope to prevent the
C++ dispatch registration from being freed by GC.

* worker: fix torchao stub — return stub classes not modules for isinstance()

peft/tuners/lora/torchao.py does:
  from torchao.dtypes import AffineQuantizedTensor, LinearActivationQuantizedTensor
  isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor))

The stub __getattr__ was returning stub modules, which isinstance() rejects
with "arg 2 must be a type, a tuple of types, or a union".

Add _StubTypeMeta metaclass whose __instancecheck__ always returns False,
and _make_stub_type() to create stub classes via it. Change _make_mod_stub
__getattr__ to return stub classes instead of stub modules for leaf
attribute access, so isinstance() gets a valid type and returns False.

_StubSubpackageFinder still handles import-style subpackage creation
(those still need module objects in sys.modules); __getattr__ only fires
for from-import or direct attribute access, which are the isinstance paths.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* tests: add coverage for Windows ROCm install paths and worker patches

Add conftest.py to fix pre-existing sys.path issue that prevented
test_rocm_support.py from running at all (install_python_stack.py
imports from backend.utils.wheel_utils which needs studio/ on sys.path).

New test classes cover everything added in this session:
- TestWindowsRocmIndexUrl: arch → AMD pip index URL mapping (gfx120X-all,
  gfx1151, gfx1150, gfx110X-all, unknown → None, trailing slash)
- TestDetectWindowsGfxArch: hipinfo output parsing, missing/timeout/bad
  returncode/no-gcnArchName paths
- TestInstallBnbWindowsRocm: UV_SKIP_WHEEL_FILENAME_CHECK set+restored,
  env restored on exception, no-op when URL missing
- TestRocmTorchInstalledEnvVar: UNSLOTH_ROCM_TORCH_INSTALLED=1 skips
  pip_install, calls _install_bnb_windows_rocm, sets flag
- TestWorkerWindowsRocmPatches: _grouped_mm CUDA dispatch override,
  offs/grouped variant handling, GC-prevention sentinel,
  _StubTypeMeta __instancecheck__, _StubSubpackageFinder registration,
  torchao key submodule pre-stubbing, TORCHDYNAMO_DISABLE guard
- TestRocmTorchPkgSpecs: rocm7.2 torch 2.11.x spec, default <2.11 cap,
  3-tuple shape, _GFX_TO_AMD_INDEX_ARCH RDNA4/3.5/3 coverage

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* tests: fix encoding, IS_WINDOWS patching, and wrong assertion

- Add encoding="utf-8" to all read_text() calls (54 occurrences) so
  tests pass on Windows where the default codec is cp1252 and source
  files contain UTF-8 emoji (e.g. ⚠️ in install_python_stack.py)
- Add @patch.object(stack_mod, "IS_WINDOWS", False) to Linux-path
  TestEnsureRocmTorch tests so they reach the Linux code path when run
  on a Windows machine instead of short-circuiting into the Windows branch
- Fix test_grouped_mm_patch_guarded_by_windows_and_hip_check: the source
  uses getattr(_torch_for_rocm, "version", None) not torch.version, so
  check for '"version"' and '"hip"' substrings instead

137 passed, 2 skipped

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: pin BNB_ROCM_VERSION=72 for torch==2.11.0+rocm7.13.0 compatibility

AMD's pip index now ships torch==2.11.0+rocm7.13.0 (ROCm 7.13).
bitsandbytes auto-detects HIP 7.13 from torch.version.hip and looks for
libbitsandbytes_rocm713.dll, which the AMD Windows prerelease wheel does
not ship (it only ships rocm72.dll), causing a load error at training start.

Fix:
- worker.py section 1f: set BNB_ROCM_VERSION=72 (via setdefault) before
  section 2 ML imports, so bitsandbytes always loads rocm72.dll on Windows ROCm
- install_python_stack.py: set BNB_ROCM_VERSION=72 in _install_bnb_windows_rocm()
  for any post-install imports; update comment to document root cause
- tests: 4 new assertions covering the fix (141 passed, 2 skipped)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: detect BNB ROCm DLL suffix dynamically instead of hardcoding '72'

BNB_ROCM_VERSION was pinned to '72' which works today (AMD wheel ships
rocm72.dll) but would break again if AMD ships a future wheel with a
different DLL suffix (e.g. rocm713.dll).

Add _detect_bnb_rocm_dll_ver() to install_python_stack.py: scans the
installed bitsandbytes package dir for libbitsandbytes_rocm{VER}.dll
using importlib.util.find_spec (no BNB import needed) and returns the
suffix.  '72' remains the fallback when detection fails.

Apply the same detection inline in worker.py section 1f.  Both paths
still respect a pre-set BNB_ROCM_VERSION (caller override wins).

Tests: +8 cases covering detection logic and fallback (147 passed, 2 skipped).

* fix: patch torch.distributed stubs in server process for Windows ROCm

On Windows ROCm, torch.distributed ships without process-group helpers
(is_initialized, is_available, get_rank, get_world_size).  The worker
subprocess already patches these in section 1e, but the main server
process calls _determine_attention_impl_for_gpu_estimate() which calls
unsloth's resolve_attention_implementation() → is_initialized(), causing:

  "Could not resolve attention implementation for '...':
   module 'torch.distributed' has no attribute 'is_initialized'"

Fix: patch the missing attrs onto torch.distributed at the top of
_determine_attention_impl_for_gpu_estimate, matching the same stubs
already applied in worker.py section 1e.  No-ops on Linux/CUDA where
torch.distributed is fully populated.

* fix: gate _grouped_mm dispatch patch on HIP < 7.13

AMD fixed the gfx1200 null HIP kernel in ROCm 7.13 (torch 2.11+).
Users on the new wheel now get the real GPU _grouped_mm kernel for
MoE workloads instead of the Python mm fallback.

Changes:
- worker.py: add _hip_ver_at_least() helper; wrap full _grouped_mm
  patch in `if not _hip_ver_at_least(7, 13):` with else branch that
  logs the skip reason; update section-1f comment to document the fix
- test_rocm_support.py: add 5 tests covering the helper definition,
  the (7, 13) gate expression, the else branch, the skip log message,
  and the AMD-format version string parsing (.split(".")[:2])

Verified: torch==2.11.0+rocm7.13.0 — 3D batch and grouped (offs)
variants both succeed; null crash only present on rocm7.12 and earlier.

* fix: stub is_torchelastic_launched on torch.distributed for Windows ROCm

resolve_attention_implementation calls is_torchelastic_launched() which
does not exist in the incomplete torch.distributed shipped with the
Windows ROCm wheel, causing a warning on every model config load in the
server process. Add it to the stub table alongside the four helpers
already patched in _determine_attention_impl_for_gpu_estimate.

Also adds two tests: one confirming the new stub and one confirming all
five core distributed helpers are covered.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: explicit warnings on AMD ROCm arch/version fallbacks + Fast-Install arg order

setup.ps1:
- Fix Fast-Install argument order: packages before flags, consistent with
  all other Fast-Install calls in the file
  (was: Fast-Install --force-reinstall --index-url $url torch ...)
  (now: Fast-Install torch torchvision torchaudio --force-reinstall --index-url $url)
- Add explicit [WARN] substep when $HasROCm is true but arch mapping fails:
  - GPU arch detected but not in supported wheel list → names the arch and
    lists supported families so user knows exactly what to report
  - HIP SDK present (amd-smi path) but gcnArchName unreadable → instructs
    user to re-install the HIP SDK; previously fell back silently to CPU

install.sh:
- Add [WARN] to stderr before silent CPU fallback when AMD GPU is confirmed
  (rocminfo/amd-smi) but ROCm version cannot be read from any source
  (amd-smi, /opt/rocm/.info/version, hipconfig, dpkg, rpm)
- Add [WARN] to stderr when ROCm version is too old (< 6.0) with upgrade link

install.ps1 and setup.sh: no changes needed (already handle these paths correctly)

* fix: robust gfx arch detection for Strix Halo / HIP-runtime-only installs

Covers users who have the HIP runtime (amd-smi available) but not the
full HIP SDK (no hipinfo), which is common on Strix Halo iGPU systems.
Without this, $ROCmGfxArch stays null and the installer silently falls
back to CPU-only PyTorch despite a working GPU.

Detection waterfall (setup.ps1 + install.ps1):
  1. hipinfo gcnArchName          -- full HIP SDK (existing, unchanged)
  2. amd-smi list gfx pattern     -- newer amd-smi versions embed arch
  3. amd-smi static --asic        -- ROCm 6+ ASIC details with GFX target
  4. UNSLOTH_ROCM_GFX_ARCH env    -- manual override escape hatch
  5. GPU name → arch table        -- best-effort from marketing name:
       890M / Strix Halo  → gfx1151 (RDNA 3.5 iGPU, Strix Halo)
       880M / Strix Point → gfx1150 (RDNA 3.5 iGPU, Strix Point)
       780M / Phoenix     → gfx1103 (RDNA 3 iGPU)
       RX 7900/7800/7700  → gfx1100 (RDNA 3 desktop)
       RX 9070 XT / 9080  → gfx1201 (RDNA 4)
       RX 9070 / 9060 XT  → gfx1200 (RDNA 4)

When arch is inferred from name, a Cyan substep tells the user to set
UNSLOTH_ROCM_GFX_ARCH to skip inference on future installs.
WMI block intentionally does not set $HasROCm (no runtime confirmation).

Tests: 11 new tests in TestStrixHaloGfxArchDetection covering all five
detection levels, WMI safety, and gfx regex in both ps1 files.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: resolve hipinfo/hipconfig via HIP_PATH/ROCM_PATH when not on PATH

AMD HIP SDK sets HIP_PATH on Windows but does not always add the bin
directory to PATH.  Get-Command hipinfo therefore silently fails and
detection falls through to WMI, which cannot provide a gfx arch, leaving
the user with a CPU-only PyTorch install and no warning.

Changes:
- setup.ps1 / install.ps1: before falling through to amd-smi, attempt to
  locate hipinfo.exe and hipconfig.exe under $env:HIP_PATH\bin (then
  $env:ROCM_PATH\bin) when Get-Command returns nothing
- Emit a [WARN] with the resolved path and a one-liner to permanently fix
  PATH via SetEnvironmentVariable
- Emit a [WARN] when HIP_PATH/ROCM_PATH is set but the exe is still not
  found (incomplete SDK install)
- Emit a [WARN] with the first hipinfo output line when hipinfo runs but
  returns a non-zero exit code (e.g. "no ROCm-capable device detected")
- 18 new tests in TestHipSdkEnvPathResolution; total 183 passed, 2 skipped

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* feat: print HIP SDK path and full hipconfig version in terminal on AMD detection

Both install.ps1 and setup.ps1 now emit substeps under the gpu step when
AMD ROCm is detected:

  gpu  AMD ROCm (gfx1200)
       HIP SDK: C:\Program Files\AMD\ROCm\7.1
       hipconfig: 7.1.51803-d3a86bd04

Previously only the gpu label (e.g. "AMD ROCm (gfx1200)") was shown with
no indication of where the SDK was found or which exact build was active.
The full hipconfig build string (e.g. 7.1.51803-d3a86bd04 instead of just
7.1) is now stored in ROCmVersionFull and also used in setup.ps1's
'rocm' step label.

9 new tests in TestHipSdkDetectedSubstep; total 192 passed, 2 skipped

* fix: Strix rocm7.1 segfault bypass + Ubuntu 24.04 HIP gcc-install-dir

Issue 1 (install.sh): gfx1151/gfx1150 + ROCm 7.1 causes a segfault in
torch._grouped_mm (moe_utils.py:167). The Radeon repo now ships cp313
wheels for rocm-rel-7.1, so _amd_gpu_radeon=true silently lands on the
broken combo. When Strix Halo/Point is detected and TORCH_INDEX_URL is
rocm7.1, override to rocm7.2 PyTorch index, update TORCH_CONSTRAINT, and
set _amd_gpu_radeon=false to bypass the Radeon repo entirely. Emits a
clear [WARN] explaining the segfault and linking to the ROCm upgrade docs.

Issue 2 (setup.sh): ROCm 7.x ships clang-20 which on Ubuntu 24.04+ picks
/usr/lib/gcc/x86_64-linux-gnu/14/ (runtime dir, no C++ headers), causing
'cstdlib file not found' and a failed llama.cpp HIP build. Iterate gcc
versions 14→11 to find the first install dir that has both runtime and
/usr/include/c++/<ver> headers, then pass --gcc-install-dir to clang via
CMAKE_HIP_FLAGS. Fix confirmed by h34v3nzc0dex (llama.cpp 417/417 clean).

11 new tests across TestStrixRocm71Override and TestSetupShGccInstallDir;
total 203 passed, 2 skipped

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: BNB_ROCM_VERSION in server process + torch._C._distributed_c10d stubs

Two errors visible in training logs on Windows ROCm:

1. Server process bitsandbytes crash:
   "Configured ROCm binary not found at libbitsandbytes_rocm713.dll"
   The installed BNB wheel ships rocm72.dll (not rocm713.dll). The
   training worker already sets BNB_ROCM_VERSION=72 via DLL detection
   but the server process (main.py) imported bitsandbytes before that
   ran. Fix: add the same DLL-scan + BNB_ROCM_VERSION assignment to
   main.py inside the existing win32 guard, before any downstream
   import can pull in bitsandbytes.

2. torch.distributed import failure:
   "No module named 'torch._C._distributed_c10d'; torch._C is not a package"
   torch._C is a C extension on Windows ROCm — Python cannot do
   submodule imports from it, so torch.distributed fails to import
   before our attribute stubs could ever run. Fix: inject empty
   ModuleType stubs for _distributed_c10d, _distributed_autograd and
   _distributed_rpc into sys.modules inside the win32 guard in
   hardware.py BEFORE importing torch.distributed, so the import
   succeeds and our attribute stubs take effect.

9 new tests in TestServerStartupRocmFixes; total 212 passed, 2 skipped

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(win32): populate distributed c10d stub with dummy symbols

torch.distributed tries to `from torch._C._distributed_c10d import
FakeProcessGroup` (and ProcessGroup, Work, Store, etc.).  The previous
empty ModuleType stub caused an AttributeError on those names.

Populate every stub with a _Dummy class for each known symbol so the
import chain completes silently on Windows ROCm where torch._C is a
compiled extension and its _distributed_c10d submodule doesn't exist.

Adds four new tests in TestServerStartupRocmFixes covering FakeProcessGroup,
ProcessGroup, setattr population, and all three _distributed_* siblings.

* fix(win32): distinguish HIP SDK installed vs GPU not ROCm-accessible

Previously, when hipinfo was found but exited non-zero (e.g. "no
ROCm-capable device detected"), both install.ps1 and setup.ps1 fell
through to the WMI-label-only branch and printed "AMD GPU detected --
HIP SDK not found" -- factually wrong since the SDK binary is present.

Add $HipSdkInstalled flag (set true when hipinfo binary is found,
regardless of exit code). When HipSdkInstalled && !HasROCm:
- Show "AMD GPU detected -- not ROCm-accessible (HIP <ver>)" instead
- Explain this is a driver issue, not an SDK issue, with a link
- Still run hipconfig version capture so version shows in output
- CPU-only hint now says "GPU not ROCm-accessible" not "require HIP SDK"

Also applies to setup.ps1 (same detection block, same branches).

Adds TestHipSdkInstalledButDeviceInaccessible (11 tests).

* fix(win32): scope ROCm workarounds to AMD hosts only

Three Codex-flagged issues where Windows ROCm workarounds incorrectly
applied to Windows CUDA (NVIDIA) machines:

main.py (P1): BNB_ROCM_VERSION was set unconditionally on all win32
hosts. On NVIDIA, bitsandbytes sees BNB_ROCM_VERSION and looks for a
ROCm DLL that doesn't exist, breaking bitsandbytes initialisation.
Fix: gate the block on HIP_PATH/ROCM_PATH being present (ROCm hosts only).

worker.py (P2): torchao stubs were seeded for all win32 runs, shadowing
real torchao on Windows CUDA and silently disabling torchao quantization
for NVIDIA users. Fix: gate on HIP_PATH/ROCM_PATH (win32 ROCm only).

install_python_stack.py (P1): _detect_windows_gfx_arch() only checked
shutil.which("hipinfo"), skipping the HIP_PATH/ROCM_PATH fallback that
the PowerShell installers use. On installs where the HIP SDK bin dir is
not on PATH, _ensure_rocm_torch() returned early without installing
ROCm wheels or bitsandbytes. Fix: mirror the env-var fallback.

* fix(linux): route Strix + ROCm 7.1 to AMD arch-specific index

Instead of falling back to pytorch.org/rocm7.2, the Strix override now
routes to repo.amd.com/rocm/whl/gfx1151/ (or gfx1150/) which serves
torch 2.11.0+rocm7.13.0 -- AMD's build containing the actual _grouped_mm
kernel fix, verified on real gfx1151 hardware by h34v3nzc0dex.

This exercises the real GPU kernel path rather than the rocm7.2 workaround.
UNSLOTH_AMD_ROCM_MIRROR can override the base URL for air-gapped installs.

Also teaches _tauri_torch_index_family to recognise AMD arch-specific URLs
(repo.amd.com/rocm/whl/gfx*) and return the rocm7.13 family label so
_tauri_gpu_branch correctly classifies these installs as rocm.

Suggested by h34v3nzc0dex based on hardware-verified probe results.

* fix(studio/rocm): gate ROCm-only side-effects on active torch runtime

Address five edge cases flagged during PR review:

1. studio/backend/main.py: BNB_ROCM_VERSION was set whenever HIP_PATH or
   ROCM_PATH was present in the environment. A Windows CUDA user who once
   installed the HIP SDK and reverted to a CUDA torch wheel still has those
   env vars set, so bitsandbytes would try to load libbitsandbytes_rocm72.dll
   against a CUDA torch and crash. Now probe torch.version.hip inside the
   env-var guard (worker.py already does this).

2. studio/backend/main.py: os.add_dll_directory returned handles were
   discarded. Per CPython docs, the directory leaves the DLL search list when
   the handle is garbage collected. Retain handles in module-level
   _ROCM_DLL_HANDLES list so they survive process lifetime.

3. studio/install_python_stack.py: _install_bnb_windows_rocm() returned None
   regardless of pip_install_try outcome, and the caller flipped
   _rocm_windows_torch_installed to True unconditionally. On a failed BNB
   install the post-install "manual install may be required" warning was
   suppressed and the user was misled. Helper now returns bool; caller gates
   on it.

4. studio/install_python_stack.py: _detect_windows_gfx_arch returned the raw
   capture group, so mixed-case hipinfo output ("Gfx1151") missed the
   lowercase keys in _GFX_TO_AMD_INDEX_ARCH and silently fell back to CPU
   torch. Lowercase the token.

5. studio/install_python_stack.py: UNSLOTH_ROCM_TORCH_INSTALLED=1 early-
   return trusted the env var even when the venv was wiped between runs.
   Subprocess-probe torch importability first; fall through to the full
   install path if the probe fails.

Tests: 231 passed, 1 skipped in tests/studio/install/test_rocm_support.py
(adds one new test for case 5 fall-through).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(studio/rocm): worker.py parity + don't roll back ROCm torch on bnb failure

Addresses findings from a 10x reviewer pass on the prior fix commit:

1. studio/backend/core/training/worker.py (parity with main.py):
   - Gate the torchao stub block on torch.version.hip / 'rocm' in
     torch.__version__ instead of HIP_PATH / ROCM_PATH env-var presence.
     Same root cause as main.py: HIP SDK env vars stick around on CUDA hosts.
   - Add module-level Windows ROCm DLL registration block. Worker subprocesses
     inherit env vars but not the parent's add_dll_directory handles, so the
     first `import torch` in the worker could fail to find amdhip64.dll when
     HIP_PATH\bin is not on PATH. Mirrors main.py setup. Handles retained at
     module scope via _ROCM_DLL_HANDLES.
   - Promote _WINDOWS_ROCM_GROUPED_MM_LIB to module scope with `global` in
     run_training_process so the torch.library.Library registration survives
     past function return / mid-run garbage collection.
   - Harden _torch_has_hip() to also accept 'rocm' in torch.__version__
     (AMD SDK / Radeon wheels may not set torch.version.hip).

2. studio/install_python_stack.py:
   - Don't roll back ROCm torch when bitsandbytes install fails. The prior
     commit gated _rocm_windows_torch_installed on _install_bnb_windows_rocm()
     returning True; if torch installed successfully but bnb failed, the flag
     stayed False and later install steps could overwrite ROCm torch with the
     generic CPU torch wheel. Set the flag after torch install; surface bnb
     failure as a separate warning instead.
   - _detect_windows_gfx_arch now probes in three tiers: UNSLOTH_ROCM_GFX_ARCH
     env-var override (matches the PowerShell installer), then hipinfo (PATH
     or HIP_PATH\bin), then amd-smi (`static --asic`, `list`). Without the
     amd-smi fallback, runtime-only Radeon installs without hipinfo on PATH
     made `studio update` return early and leave the venv on CPU torch.
   - Linux torch-already-rocm probe in _ensure_rocm_torch now matches the
     Windows probe shape: accepts torch.version.hip OR 'rocm' in
     torch.__version__ to cover AMD SDK / Radeon Linux wheels.

3. studio/backend/utils/hardware/hardware.py:
   - apply_gpu_ids() final-fallback torch probe accepts 'rocm' in
     torch.__version__ in addition to torch.version.hip, matching
     detect_hardware(). AMD SDK wheels could otherwise leak through with
     CUDA-only visibility masks on a spawned ROCm worker.

Tests: 231 passed, 1 skipped in tests/studio/install/test_rocm_support.py
(no test changes needed; the probe shape that prints the hip version (or
'rocm' sentinel) preserves the existing non-empty-string contract).

Not addressed in this commit (deferred or out of scope):
- Tag drift / lemonade checksum (PR 5303 surface, not this PR).
- install.sh rocm7.2.1 URL: small fix, separate.
- install.ps1 / setup.ps1 'Radeon 8060S' marketing-name fallback table.
- Strix Halo + ROCm 7.1 routing asymmetry in Python update path.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(studio/rocm): robustness pass - rocm tag normalisation, Strix routing parity, hardened detection

Robustness pass on top of 76137b2d. Four targeted fixes:

1. install.sh ROCm-tag routing normalisation.
   `rocm7.2.1` would route to https://download.pytorch.org/whl/rocm7.2.1
   which does not exist (PyTorch publishes major.minor URLs only). Same
   for any future patch-level tag. Normalise every rocm{maj.min}* pattern
   to the bare {maj.min} index URL.

2. install.ps1 + studio/setup.ps1 marketing-name fallback.
   The gfx1151 row matched 890M / Strix Halo / HX 37x / HX 38x / AI 9 HX
   but not the actual retail name 'AMD Radeon 8060S Graphics' shipped by
   OEMs (Ryzen AI MAX+ 395). Add '8060S' to the regex.

3. install_python_stack.py Strix + ROCm 7.1 routing parity with install.sh.
   The shell installer reroutes Strix Halo / Point + ROCm 7.1 to
   repo.amd.com/rocm/whl/{gfx}/ (which serves torch 2.11.0+rocm7.13.0
   with the upstream _grouped_mm fix). The Python `studio update` path
   only warned and still installed the broken generic rocm7.1 wheel.
   Mirror the override: detect gfx1151/gfx1150 on ROCm 7.1, route to
   the AMD per-gfx index, honour UNSLOTH_AMD_ROCM_MIRROR override.

4. _detect_windows_gfx_arch amd-smi parsing tightened.
   The amd-smi fallback added in the prior commit used a bare
   `\bgfx[1-9][0-9a-z]{2,3}\b` match against the lowercased stdout,
   which could pick up stray gfx references in warnings / device-name
   strings. Anchor on labelled lines first (Target_Graphics_Version,
   ASIC, Arch, gfx) and fall back to the bare match only when no
   labelled line is present.

Tests: 231 passed, 1 skipped in tests/studio/install/test_rocm_support.py;
sim_5301 23 cases pass (6 new sims for the Strix override + amd-smi parsing).

* fix(studio/rocm): multi-GPU selection, Strix sibling handling, defensive cleanups

Round 4 robustness pass based on 5 parallel Opus reviewers of head 21773215.
Seven items from across regression / edge-case / error-paths / architecture
reviews:

1. studio/backend/main.py BNB gate: aligned with the broad ROCm check used
   everywhere else in this PR (torch.version.hip OR 'rocm' in __version__).
   AMD SDK / Radeon Linux wheels do not always populate torch.version.hip;
   without this, main.py would silently skip BNB_ROCM_VERSION while worker.py
   set it.

2. studio/install_python_stack.py _install_bnb_windows_rocm: init _ok = False
   before the try block. Without this, if pip_install_try itself raises
   (e.g. OSError on uv binary missing), the finally block restored env vars
   correctly but the subsequent `if not _ok:` raised UnboundLocalError,
   masking the original exception.

3. studio/install_python_stack.py _detect_windows_gfx_arch:
   - Rewrote to use re.findall (not re.search) on both hipinfo and amd-smi
     output, dedup tokens preserving order, and select via new
     _pick_visible_index() helper.
   - HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES (first comma entry, integer)
     now picks the right GPU on multi-AMD-GPU hosts. Out-of-range or non-int
     values fall back to the first GPU (matches detect_host behaviour in
     install_llama_prebuilt.py).

4. studio/install_python_stack.py Strix override now consults the runtime
   target before flipping:
   - Previous behaviour intersected gfx_codes with {gfx1151, gfx1150} and
     picked the first Strix arch, ignoring whether HIP_VISIBLE_DEVICES
     selected a non-Strix sibling (e.g. discrete RX 7900 in a mixed APU+dGPU
     box). Could install Strix-specific wheels onto a gfx1100 dGPU.
   - Now resolves the runtime gfx via _pick_visible_index() and only
     overrides when that runtime target is in the Strix set.

5. studio/backend/main.py + studio/backend/core/training/worker.py: ROCm
   version dir scan no longer sorts lexically. Previous sort placed "10.0"
   before "7.0" alphabetically, which would mis-prioritise ROCm 10.x bin
   dirs once AMD ships them. New _ver_key() splits on "." and sorts
   numerically with a string fallback.

6. install.sh Strix override URL: replaced ${var%/} (strips one trailing
   slash) with a while-loop that strips all trailing slashes, matching
   Python's .rstrip("/"). A user setting UNSLOTH_AMD_ROCM_MIRROR with
   "http://corp/whl///" no longer ends up with "http://corp/whl///gfx1151/"
   which strict pip proxies (artifactory, sonatype) 404 on.

7. studio/install_python_stack.py: bumped torch import probe timeout from
   30s to 90s. PyTorch's lazy .so loading can take 60-90s on cold NFS or
   USB-backed venvs. The shorter timeout was producing a false "torch
   missing" classification and reinstalling a working ROCm torch.

Tests: 231 passed, 1 skipped. sim_5301 30 cases pass (added 7 new sims for
multi-GPU detection, Strix sibling handling, and _ok-init regression).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(studio/rocm): worker BNB/grouped_mm broad gate, install.sh Strix visibility, runtime-only ROCm detection

Round-5 robustness pass based on 20 parallel reviewers of head 96b9e465.

1. studio/backend/core/training/worker.py - BNB version pin / dynamo disable
   / _grouped_mm fallback block was still gated on torch.version.hip alone
   despite the torchao stub block above already using the broad check. AMD
   SDK / Radeon Windows wheels (torch.__version__ contains "rocm" but
   torch.version.hip is None) silently skipped the Windows ROCm runtime
   patches. Aligned to the same broad check (8/20 reviewers).

2. studio/backend/core/training/worker.py - _hip_ver_at_least() now also
   parses the ROCm version out of torch.__version__ (e.g. "2.11.0+rocm7.13.0")
   when torch.version.hip is missing, so the kernel-fix gate is correct for
   SDK / Radeon wheels too.

3. studio/backend/core/training/worker.py - _grouped_mm_safe_impl with
   offs=None now picks torch.bmm/matmul for 3-D inputs instead of always
   calling torch.mm. The real _grouped_mm accepts 3-D batched matmul; the
   prior fallback raised "self must be a matrix" on MoE workloads (2/20).

4. studio/backend/main.py - dropped the HIP_PATH / ROCM_PATH env-var gate
   from the BNB block; probe torch directly. Runtime-only Radeon / AMD SDK
   Windows installs do not set those SDK env vars but still ship ROCm torch
   (5/20 reviewers).

5. install.sh - Strix override now collects every gfx token from
   rocminfo / amd-smi (in enumeration order), then indexes by
   HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-
   Strix dGPU host where the user selected the dGPU does NOT get rerouted
   to the Strix per-gfx index. Mirrors the Python update path (5/20 reviewers).

6. install.sh - Strix detection chain now also probes `amd-smi static --asic`,
   matching the PowerShell installer (1/20). Closes the gap on runtime-only
   Strix hosts where `amd-smi list` does not surface a gfx token.

7. studio/install_python_stack.py - _has_rocm_gpu() now has the sysfs KFD
   topology fallback (/sys/class/kfd/kfd/topology/nodes/*/gpu_id), matching
   install.sh. On minimal package-managed installs without rocminfo /
   amd-smi GUI tools, `studio update` can now detect the GPU and repair the
   venv instead of returning early (2/20).

8. studio/install_python_stack.py - _detect_amd_gfx_codes() now falls back
   to `amd-smi list` and `amd-smi static --asic` when rocminfo is missing
   (2/20). Strix routing on runtime-only Radeon hosts now matches what
   install.sh has done for a while.

9. studio/install_python_stack.py - Strix override now applies even when
   has_hip_torch is True. The whole point of the override is to repair an
   existing broken torch.version.hip == "7.1" install; skipping the
   reinstall left users on the known _grouped_mm segfaulting stack (3/20).

Tests: 231 passed, 1 skipped. sim_5301 30 cases pass. sim_cross 12 pass.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(studio/rocm): code review hardening pass

- main.py: numeric DLL sort (string sort picked rocm72 over rocm713);
  add basename() to regex; log warning on detection failure; log info
  when BNB_ROCM_VERSION is set (mirrors worker.py)
- worker.py: explicit len-guard in _hip_ver_at_least() with warning
  logs instead of silent IndexError/ValueError swallow
- hardware.py: isinstance(result, dict) guard before result.get() in
  _smi_query() to prevent AttributeError on non-dict backend returns
- amd.py: round() before int() on parsed GPU IDs; log warning when
  truncation occurs (defensive against malformed amd-smi output)
- setup.sh: quote --gcc-install-dir value in CMAKE_HIP_FLAGS so paths
  with spaces do not break the CMake argument
- install.ps1, setup.ps1: apply colon-split + ToLower() to hipinfo
  gcnArchName match (consistent with each other and with setup.sh)
- install.sh: tighten ROCm tag case patterns to explicit
  rocmX.Y|rocmX.Y.* to avoid unintended prefix matches

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(studio/training): GPU OOM guard to prevent system freeze on VRAM exhaustion

On RDNA 4 (gfx1200/gfx1201) and other ROCm GPUs, exhausting VRAM can
cause a HIP driver hang that freezes the entire system rather than
raising a recoverable Python exception.

Two-part fix:
- set_per_process_memory_fraction(0.90) caps the HIP/CUDA allocator at
  90% of VRAM so PyTorch raises OutOfMemoryError before hitting the
  hardware limit, keeping the driver alive and the system responsive
- top-level exception handler detects OOM errors by type and message
  and surfaces a clear actionable message to the UI (reduce
  max_seq_length, enable gradient_checkpointing, lower batch size)
  instead of the raw CUDA/HIP error string

* fix(studio/rocm): OOM guard ROCm-only + unified memory, multi-GPU arch selection

OOM guard (worker.py):
- Scope to _hw.IS_ROCM only -- NVIDIA CUDA has a graceful OOM path and
  does not need the allocator cap
- Detect unified memory by comparing torch VRAM against psutil system RAM;
  use 0.80 on unified-memory APUs (gfx1151 Strix Halo) where the GPU pool
  is carved from host RAM, 0.90 on discrete cards

Multi-GPU arch selection:
- install.ps1 / setup.ps1: replace -match (first hit only) with
  [regex]::Matches() to collect all gcnArchName entries, then index by
  HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES
- install_python_stack.py: index into full token list before dedup so
  HIP_VISIBLE_DEVICES=2 on [gfx1100, gfx1100, gfx1151] resolves gfx1151
- install.sh: remove awk dedup from gfx token collection for same reason

GCC multiarch (setup.sh):
- Only append -linux-gnu when gcc -print-multiarch does not already return
  the full triple, fixing double-suffix on Ubuntu 24.04

* fix(tests): update ROCm version cap expectations from rocm7.1 to rocm7.2

Daniel's normalisation commit updated the cap from rocm7.1 to rocm7.2
since PyTorch now publishes that index and rocm7.2 ships torch 2.11.0.
Test expectations were stale.

* fix(tests): correct MLX smoke test losses_per_step assertion

logging_steps=1 with max_steps=30 produces 30 loss entries, not 7.
The assertion was stale from a previous config.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(studio/worker): detect unified-memory APU by GPU name not VRAM/RAM ratio

The previous heuristic (VRAM > 50 % of system RAM) false-positived on discrete
cards in low-RAM systems — e.g. RX 9060 XT 16 GB on a 16 GB or 24 GB machine
would trip the unified-memory path and log "unified memory host" when it should
say "discrete".

AMD iGPUs (gfx1150/gfx1151 Strix Halo, Strix Point, etc.) expose names with a
digit+M suffix ("AMD Radeon 890M"), while discrete cards use "RX NNNN [XT|XTX]"
naming.  Matching that suffix is reliable across all current ROCm-capable AMD
consumer GPUs and does not require psutil.

Also includes the device name in the log line to ease future debugging.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(install/setup.ps1): force array on hipinfo gcnArchName parse to fix single-GPU arch truncation

When [regex]::Matches() finds exactly one match, PowerShell's pipeline
unwraps the result to a scalar string.  Indexing a scalar string with [0]
returns the first *character*, so a one-GPU system would parse
gcnArchName "gfx1200" as "g", which is not in the supported arch map
and triggers the CPU-only fallback.

Wrapping with @() forces the result to remain an array regardless of
match count.  On a single-GPU machine the arch is now correctly read as
"gfx1200" (or whatever the full name is) so the ROCm wheel index is
selected.

Reproducer: hipinfo exits 0 and outputs exactly one gcnArchName line.
Without @(), $_hipAllArches = "gfx1200" (String); $_hipAllArches[0] = 'g'.
With @(), $_hipAllArches = @("gfx1200") (Object[]); $_hipAllArches[0] = "gfx1200".

* fix(studio/rocm): classify unified-memory APU via VRAM/RAM ratio, not arch list

Replace the gcnArchName allowlist {gfx1150, gfx1151} with a
psutil-based heuristic: unified APUs expose the entire system RAM
as the HIP pool (ratio ≥ 0.90), discrete cards are well below that.
No arch name required — future APUs classify correctly without code changes.

Also removes the stale import re / \d[Mm]\b device-name regex that
5d84704 left behind, and logs vram/sys GiB for easier on-hardware
verification.

Addresses h34v3nzc0dex review: Radeon 8060S (gfx1151, 128 GiB
unified) now correctly gets 0.80 cap instead of 0.90.

* fix(studio/rocm): revert to gcnArchName for unified-memory APU classification

VRAM/RAM ratio >= 0.90 false-positives on machines where discrete VRAM
equals system RAM (e.g. RX 9060 XT 16 GB + 16 GB system RAM → ratio 1.0,
incorrectly classified as unified → wrong 0.80 cap applied).

gcnArchName is the correct signal: naming-independent, stable within a
product family, and already parsed throughout this PR. Unified set is
{gfx1150, gfx1151} (Strix Point + Strix Halo).

* fix(studio/llama-prebuilt): resolve hipinfo via HIP_PATH/ROCM_PATH on Windows

shutil.which("hipinfo") returns None when the HIP SDK bin dir is not on
PATH -- the HIP SDK installer sets HIP_PATH/ROCM_PATH but does not always
add the bin dir to PATH. This caused has_rocm=False in the prebuilt asset
selector, so AMD ROCm machines got the CPU llama.cpp zip instead of the
HIP one, silently running all chat inference on CPU.

Add _resolve_exe() that falls back to %HIP_PATH%\bin and %ROCM_PATH%\bin
when shutil.which() finds nothing, mirroring the same fallback already
present in setup.ps1.

* fix(studio/llama-prebuilt): pass --has-rocm from setup.ps1 to skip re-detection

The Python prebuilt installer re-detects ROCm independently via
shutil.which("hipinfo"), which fails when hipinfo is not on PATH
(HIP SDK sets HIP_PATH but doesn't always add the bin dir to PATH).
This caused has_rocm=False and downloaded the CPU llama.cpp zip even
on confirmed AMD ROCm machines.

setup.ps1 already performs reliable ROCm detection with its own
HIP_PATH/ROCM_PATH fallback. Add --has-rocm flag to
install_llama_prebuilt.py so setup.ps1 can forward its result directly,
and pass it whenever $HasROCm is true. The Python script then overrides
has_rocm=True in the HostInfo without re-probing.

* fix(studio/llama-prebuilt): add HIP asset to simple-policy Windows path

direct_upstream_release_plan (used by --simple-policy, which setup.ps1
always passes) only checked has_usable_nvidia on Windows and fell
straight to CPU for AMD ROCm machines, ignoring has_rocm entirely.
The --has-rocm override had no effect because the simple-policy code
path never reached resolve_asset_choice where has_rocm was checked.

Add an elif branch for has_rocm that tries the upstream HIP asset
(llama-TAG-bin-win-hip-radeon-x64.zip) before falling through to the
CPU fallback, consistent with the non-simple-policy path.

* fix(studio/setup.ps1): auto-remove mismatched llama.cpp install kind

When an existing llama.cpp install is the wrong kind for the current
GPU (e.g. windows-cpu on an AMD ROCm machine that should have
windows-hip), the prebuilt installer skips on tag match and never
upgrades. Read install_kind from UNSLOTH_PREBUILT_INFO.json before
invoking the installer and remove the directory if the kind doesn't
match, forcing a fresh download of the correct variant.

* fix(studio/setup.ps1): show live PyTorch install output in verbose mode for ROCm

The ROCm torch reinstall (setup.ps1 phase) always silently captured
output, so in --verbose mode the torch downgrade mid-install
(2.11.0+rocm → 2.10.0 → 2.11.0+rocm) looked like the final state was
2.10.0. Match the CPU/CUDA blocks which show live uv output when
$script:UnslothVerbose is set.

* fix(rocm/windows): set ROCBLAS_TENSILE_LIBPATH for bundled rocblas.dll

The llama.cpp ROCm prebuilt bundles rocblas.dll next to the binary but
not the Tensile kernel library files it depends on at runtime
(rocblas/library/TensileLibrary*.dat + *.hsaco).  The bundled DLL
searches for these files relative to its own location by default, i.e.
<binary_dir>/rocblas/library/, which does not exist in the prebuilt
install tree.  This causes a silent crash on the very first GEMM
(prefill) with no output from llama-server, seen by the caller as
WinError 10054 / 10061.  Model load and the single-token warmup pass
because they use simpler code paths that do not trigger rocBLAS GEMM.

Fix: set ROCBLAS_TENSILE_LIBPATH in the subprocess env to
<HIP_PATH>/bin/rocblas/library so the bundled DLL finds the kernel
files from the system ROCm installation.  Uses setdefault so a user-
supplied env var is never overwritten.  No-ops on CUDA and CPU (no
HIP_PATH) and on Linux (win32 branch only).

Reproducer log:
  rocBLAS error: Cannot read .../Release/rocblas/library/TensileLibrary.dat
  rocBLAS error: Could not initialize Tensile host:
  directory_iterator: The system cannot find the path specified.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(install.sh): restore gfx token dedup in Strix multi-GPU awk indexer

536a54df removed the per-source `| awk '!seen[$0]++'` dedup from the
_gfx_all collection step but left the indexer awk as bare NF, so on a
mixed-arch host (e.g. dGPU gfx1100 + Strix iGPU gfx1151) where
rocminfo emits each gfx token twice (Name: field + ISA triple),
HIP_VISIBLE_DEVICES=1 indexed vals[1] = the second gfx1100 occurrence
instead of gfx1151, triggering the Strix routing on the wrong GPU.

Add !seen[$0]++ to the indexer awk so duplicate tokens from the same
GPU collapse to one entry before the HIP_VISIBLE_DEVICES index is
applied -- matching exactly what the Python side does with dict.fromkeys()
in _detect_amd_gfx_codes(). The comment above the block ("skip
duplicates") already documented this as the intended behaviour.

* fix(studio/install): correct _TOTAL progress count on Windows

base_total += 3 fired for all non-macOS platforms including Windows,
but flash-attn (line 1620) and ROCm torch final (line 1705) are both
guarded by 'not IS_WINDOWS and not IS_MACOS', so on Windows with torch
enabled _TOTAL was 13 while only 11 _progress() calls actually execute.

Split into +1 for the ROCm torch check (all non-macOS) and +2 for the
two Linux-only steps, so Windows gets _TOTAL=11 and Linux gets 14.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(install.ps1): enforce torch>=2.11.0 for gfx120X and Strix on Windows

The AMD arch-specific index (repo.amd.com/rocm/whl/gfx120X-all/ and
gfx1151/) publishes torch wheels from 2.7.1 through 2.11.0. Without a
version floor pip can resolve to torch 2.10.0+rocm7.12 on RDNA 4
(gfx120X) or torch 2.10.0+rocm7.1 on Strix (gfx1151/gfx1150), both of
which have a null-pointer crash in torch._C._grouped_mm (TheRock
issues #5284 / #3284). torch 2.11.0+rocm7.13 contains the fix.

Add $ROCmTorchFloor alongside $ROCmIndexUrl: set to torch>=2.11.0 for
the two affected arch families, null for all others. Wire it into the
uv pip install call so the broken wheels are never selected.

* fix(rocm/windows): address Codex nits - deterministic DLL suffix, CUDA llama.cpp kind, HIP_VISIBLE_DEVICES arch indexing

- install_python_stack.py / worker.py: _detect_bnb_rocm_dll_ver() and the
  inline worker probe now collect ALL libbitsandbytes_rocm*.dll suffixes and
  return max() by numeric value instead of stopping at the first glob hit.
  Filesystem glob order is not guaranteed; this ensures '713' always wins
  over '72' when both variants are present in the wheel.

- setup.ps1 (expectedKind): add 'windows-cuda' branch so NVIDIA hosts are
  not treated as 'windows-cpu'. Previously an existing windows-cuda prebuilt
  was always considered a mismatch on non-ROCm machines, forcing an
  unnecessary re-download on every update.

- setup.ps1 (amd-smi gfx arch): collect ALL gfx tokens from amd-smi list
  output in GPU order and honour HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES
  when selecting which arch to use. On mixed-arch AMD systems where the
  visible GPU is not the first enumerated one, this prevents installing an
  incompatible wheel index. Falls back to index 0 (same as before) when the
  visibility var is unset or is a comma-separated list.

- test_rocm_support.py: add test_picks_highest_suffix_when_multiple_dlls to
  cover the multi-DLL case that was previously untested.

* fix(rocm): misleading amd-smi log, BNB spec consistency, torch ceiling for AMD index

amd.py: split 'returncode != 0 or not stdout' into two separate branches.
Previously, exit-0 with empty output logged 'amd-smi returned code 0' (which
reads as success, not a warning) and incorrectly incremented the circuit-breaker
counter. Now: non-zero exit logs the code and counts toward the limit as before;
empty stdout on exit 0 logs at DEBUG level and does not penalise the counter
(amd-smi --json always emits at least [] on exit 0, so this branch is rare and
is not a tool failure).

main.py: replace spec.origin / os.path.dirname() with
spec.submodule_search_locations to match install_python_stack.py and worker.py.
For normal wheel installs both approaches reach the same directory, but using
submodule_search_locations is the canonical way and handles editable bitsandbytes
installs correctly. Also use max() by numeric suffix (same as the other two sites)
instead of a sort-then-break loop.

install.ps1: add <2.12.0 ceiling to the torch constraint for gfx120X (RDNA 4)
and gfx1151/gfx1150 (Strix). AMD actively publishes new versions on their
per-arch index; without a ceiling, a future 2.12.0+rocmX.Y wheel would be
pulled in automatically before being validated on these architectures. The
ceiling matches the existing Linux install_python_stack.py constraint for the
same arches. Bump both when 2.12.x is confirmed working.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(rocm): torch floor in setup.ps1, torchvision pin for Strix, rocmsdk in _hip_ver_at_least

setup.ps1: add \ (mirrors install.ps1) and derive \
from it. Previously the AMD index install called 'Fast-Install torch torchvision
torchaudio --force-reinstall --index-url \' with no version
constraint, so pip could resolve torch 2.10.0+rocm7.12 for gfx1151/gfx1200 --
the exact broken wheel the PR is meant to avoid. Now gfx120X and Strix enforce
'torch>=2.11.0,<2.12.0', matching install.ps1 and the Linux constraint.

install_python_stack.py: pin torchvision and torchaudio in _strix_override_pkgs.
The Strix Linux override uses --index-url (exclusive, no PyPI fallback); bare
unversioned 'torchvision' and 'torchaudio' could resolve a build from AMD's
index targeting a different torch major, causing ABI/version mismatches at
runtime. Now pinned to '>=0.26.0,<0.27.0' and '>=2.11.0,<2.12.0' respectively,
matching _ROCM_TORCH_CONSTRAINT['rocm7.2'].

worker.py: extend _hip_ver_at_least to handle AMD SDK wheel version strings.
The fallback regex r'rocm(\d+)\.(\d+)' cannot match '2.9.0+rocmsdk20251116'
(no rocmX.Y component), so the function always returned False on SDK/Radeon
wheels -- installing the Python _grouped_mm workaround on wheels that already
have the working HIP kernel. Added a second check: if the version string
contains '+rocmsdk', assume >= 7.13 (the rocmsdk format post-dates the
gfx120X null-kernel fix) and skip the fallback.

* fix(rocm): warn on OOB HIP_VISIBLE_DEVICES, bail on empty numeric_ids mask

- setup.ps1: when HIP/ROCR_VISIBLE_DEVICES names an index beyond the
  detected GPU count, emit a yellow warning and fall back to GPU 0
  instead of silently reading allGfxArches[-1] (wrong arch)
- hardware.py _reconcile_primary_rocm_unified_memory: distinguish
  numeric_ids=None (no env var, use torch ordinal 0) from numeric_ids=[]
  (empty mask / HIP_VISIBLE_DEVICES=-1, no GPU visible); bail out early
  in the empty case to avoid querying torch.device(0) incorrectly

* fix(rocm): gate StubSubpackageFinder on win32 ROCm, add gcnArchName fallbacks

- worker.py _StubSubpackageFinder: the meta_path append was running on
  every platform on every call to run_training_process; moved it inside
  the if _is_win32_rocm: block since stubs are only seeded there and the
  finder is a pure accumulation on Linux/Windows CUDA
- worker.py OOM guard: AMD SDK / Radeon wheels may not populate
  gcnArchName, causing Strix Halo to be misclassified as discrete and
  get the 0.90 cap (12.8 GB OS headroom) instead of 0.80 (25.6 GB);
  now tries gcn_arch_name / arch_name / gfx_arch_name variants first,
  then falls back to device-name matching (890M -> Strix Halo,
  880M -> Strix Point) with a debug log when the fallback fires

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(rocm): pin torchvision/torchaudio in setup.ps1, remove -Unique from arch array

- setup.ps1 ROCm torch install: torchvision and torchaudio were passed
  bare alongside pinned torch>=2.11.0,<2.12.0 for gfx1151/gfx1200 arches.
  AMD publishes packages independently so a future torchvision 0.27 (for
  torch 2.12) on the same arch index would cause pip ResolutionImpossible
  or an ABI-incompatible install. Added torchvisionFloorMap and
  torchaudioFloorMap mirroring install_python_stack.py's strix override
  (torchvision>=0.26.0,<0.27.0, torchaudio>=2.11.0,<2.12.0) and derived
  ROCmVisionSpec/ROCmAudioSpec used in all three Fast-Install call sites.

- setup.ps1 amd-smi arch detection: Select-Object -Unique was collapsing
  same-arch multi-GPU arrays (e.g. two gfx1151 APUs -> 1-element array)
  causing HIP_VISIBLE_DEVICES=1 to trigger a false out-of-range warning
  and fall back to GPU 0 even though the correct GPU would have been at
  index 1. Removed -Unique; added comment noting the positional-index
  assumption and its non-contiguous-GPU limitation.

* fix(rocm): add 8060s/8050s to OOM guard device-name fallback, extract classifier helper

Path 3 of the OOM guard device-name fallback only checked for 890m/880m
(gfx1150 Strix Point SKU names). Strix Halo (gfx1151) ships as Radeon 8060S
(Ryzen AI MAX+ 395) and Radeon 8050S (cut-down SKU) -- neither matches, so
the fallback returned is_unified=False and applied the 0.90 fraction instead
of 0.80, leaving ~12.8 GiB OS headroom on a 128 GiB pool instead of ~25.6 GiB.

Fix: add 8060s and 8050s to the name-match set. Also correct the comment that
mislabelled 890M as a Strix Halo name (it is Strix Point).

Refactor: extract the three-path classifier into _rocm_classify_unified_memory()
so it can be unit-tested directly. Add 31 test cases in test_rocm_oom_guard.py
covering all three paths and the regression case (Radeon 8060S Graphics).

Reported-by: h34v3nzc0dex

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(rocm): pass explicit dtype on bf16-unsupported hardware (RDNA2)

dtype=None lets unsloth auto-detect the model dtype. On RDNA2 (gfx103x,
e.g. RX 6600) is_bfloat16_supported() incorrectly returns True, so unsloth
picks bf16 and the first bf16 kernel dispatch triggers:

  LLVM ERROR: Cannot select: intrinsic %llvm.amdgcn.fdot2.bf16.bf16

Replace every dtype=None in load_model() with _auto_dtype which resolves
to None when bf16 is supported (all modern NVIDIA + RDNA3+) and
torch.float16 otherwise. This gives RDNA2 users a working float16
training path without touching NVIDIA behaviour at all.

Fixes: https://github.com/unslothai/unsloth/issues/5337

* fix: reduce log noise for expected non-issues on Windows ROCm

Three log lines fired at warning/error level for conditions that are
completely expected on a Windows HIP SDK-only setup:

amd.py
- amd-smi WinError 2 (FileNotFoundError): downgrade warning -> debug.
  amd-smi ships with Adrenalin, not the HIP SDK; absence is normal.
- 'disabling' message: downgrade warning -> info with clearer text
  'not available (not installed; expected on HIP SDK-only systems);
  GPU VRAM polling disabled'

hardware.py
- torch.distributed.Store missing: downgrade warning -> debug.
  The distributed stub added in this PR intentionally omits Store; the
  attention-impl fallback to eager is expected and non-actionable.

worker.py
- causal-conv1d: add early Windows exit (info) in both
  _ensure_causal_conv1d_fast_path and _causal_conv1d_install hook;
  no cp313/win_amd64 wheel exists, so the install always fails.
- FLA: add early Windows exit (info) in
  _ensure_flash_linear_attention_unconditional; triton dependency has
  no cp313/win_amd64 wheel.
- Defense-in-depth: _install_package_wheel_first non-HIP PyPI failure
  logs info+debug on Windows instead of error; FLA failure logs
  info+debug on Windows instead of warning.

* [AMD] FIx installation of bitsandbytes when it's from .dev and skip rebuilding llama.cpp if we build it manually.

* fix: use force_pip for Windows ROCm bitsandbytes prebuilt wheel install

uv rejects the bnb continuous-release wheel due to filename/metadata
version mismatch (1.33.7.preview vs 0.50.0.dev0). Switch to force_pip=True
(pip bypass) instead of the UV_SKIP_WHEEL_FILENAME_CHECK env var workaround
-- cleaner and consistent with how the Linux path handles it.

BNB_ROCM_VERSION is still set post-install to the detected DLL suffix so
the worker subprocess loads the correct libbitsandbytes_rocm{VER}.dll even
when torch.version.hip reports a newer HIP version than the wheel ships.

* fix: three small correctness fixes found in PR review

- _install_bnb_windows_rocm: use UV_SKIP_WHEEL_FILENAME_CHECK=1 with
  try/finally instead of force_pip=True so the env var is always
  restored and the failing CI test passes
- _determine_attention_impl_for_gpu_estimate: gate torch._C distributed
  stubs on IS_ROCM so Windows CUDA users keep the real extension
- install.ps1 amd-smi fallback: collect all gfx tokens and index by
  HIP_VISIBLE_DEVICES, matching the hipinfo path on multi-GPU hosts

* fix: stub torchao in export subprocess on Windows ROCm

On Windows, the ROCm build of PyTorch ships without the distributed
C extension (torch._C._distributed_c10d). torchao, which is pulled in
transitively by transformers.quantizers at import time, walks into
torch.distributed._functional_collectives -> distributed_c10d and
crashes with:

  No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package

This only affected the export subprocess because the training subprocess
already applied an identical torchao stub (introduced separately to fix
the same root cause). The export subprocess had no such guard and died
during 'Importing Unsloth...' before any model loading could happen.

Fix: apply the same _StubSubpackageFinder / torchao stub pattern to the
export subprocess entry point, gated on Windows ROCm detection, before
any import of transformers or unsloth_zoo.

Root cause tracked in ROCm/TheRock#3284 (libuv / torch.distributed
missing on Windows ROCm builds).

Ref: https://github.com/ROCm/TheRock/issues/3284

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* install.sh, setup.sh: add GPU arch step logging to match PS1 scripts

Both shell scripts were missing the step "gpu" terminal log block that
install.ps1 and setup.ps1 emit. This adds equivalent output: GPU label
with gfx arch (e.g. "AMD ROCm (gfx1151)"), ROCm root path, hipconfig
version, and marketing name substep. Includes the same gfx arch detection
chain (rocminfo → amd-smi list → amd-smi static --asic), UNSLOTH_ROCM_GFX_ARCH
env override, and name-based arch inference table (Strix Halo/Point, RDNA 3/4)
as the PS1 versions. install.sh also replaces bare echo blocks for the AMD
ROCm and CPU-only cases with formatted substep output.

* Fix BNB_ROCM_VERSION gate, ROCm GPU mask preference, APU unified memory and Release build for PR #5301

- main.py: gate BNB_ROCM_VERSION on the rocm bnb DLL or HIP_PATH/ROCM_PATH instead of importing torch on every Windows host
- hardware.py: prefer HIP/ROCR visible-device masks only on ROCm hosts so a stale mask cannot override CUDA_VISIBLE_DEVICES on NVIDIA
- llama_cpp.py: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1 only for unified-memory APUs (gfx1150/gfx1151)
- setup.sh: pass -DCMAKE_BUILD_TYPE=Release for the HIP source build
- add test_amd_apu_unified_memory.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: guard recompile_limit + fix AMD VRAM monitor fallback

trainer.py: torch._dynamo.config.recompile_limit does not exist in
some ROCm torch builds (e.g. pytorch.org/whl/rocm6.2 wheels). Guard
the assignment so training doesn't crash on RDNA2/RDNA3.

hardware.py: when amd-smi/nvidia-smi is unavailable or returns no
usable data (HIP SDK-only Windows, Docker, unexpected JSON format),
the existing fallback used torch.cuda.memory_allocated() which is
process-specific and reads near-zero even with a fully loaded model.
Switch to torch.cuda.mem_get_info() via _torch_get_per_device_info()
which reports system-wide VRAM occupancy so the GPU monitor shows
real usage on all AMD systems without requiring amd-smi.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: Windows VRAM monitor via Performance Counter API

When amd-smi/nvidia-smi is unavailable on Windows, query dedicated GPU
VRAM via Windows Performance Counters (same source as Task Manager).
This gives system-wide cross-process usage, fixing the near-zero reading
caused by torch.cuda.mem_get_info only seeing the Studio server process.

Linux fallback path unchanged (mem_get_info is system-wide on ROCm).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: rename to _rocm_windows_perf_counter_vram_gb, scope to IS_ROCM

Function is AMD ROCm specific — amd-smi absent on Windows when only the
HIP SDK is installed. Scoped to IS_ROCM so NVIDIA Windows path is
untouched (nvidia-smi handles that case).

* fix: AMD VRAM monitor — Linux DRM sysfs + Windows perf counter

Linux: read /sys/class/drm/card*/device/mem_info_vram_used|total for
system-wide GPU memory across all processes. No tools required, always
present on Linux AMD systems.

Windows: Windows Performance Counter API (already added).

Both paths are gated on IS_ROCM and only fire when amd-smi is absent.
torch mem_get_info remains as last resort (process-local).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: AMD GPU monitor — utilization, temperature, and power for Windows and Linux fallback paths

- Windows: GPU utilization via \GPU Engine(*engtype_3D*)\Utilization Percentage perf counter
- Windows: temperature and power via ADL (atiadlxx.dll, ships with Adrenalin)
- Linux: GPU utilization via DRM sysfs gpu_busy_percent
- Linux: temperature via hwmon temp1_input (millidegrees C)
- Linux: power via hwmon power1_average / power1_input (microwatts)

All paths are no-op fallbacks (None) when the source is unavailable.
Mirrors what nvidia-smi provides on the CUDA path.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: remove ADL ctypes — does not support AMD iGPU (Strix Halo)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Erland366 <erland.pg366@gmail.com>
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
2026-05-29 22:29:56 -07:00

5584 lines
244 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
llama-server inference backend for GGUF models.
Manages a llama-server subprocess and proxies chat completions
through its OpenAI-compatible /v1/chat/completions endpoint.
"""
import atexit
import contextlib
import json
import os
import re
import struct
import structlog
from loggers import get_logger
import shutil
import socket
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Generator, Iterable, List, Optional
from urllib.parse import urlparse
import httpx
from core.inference.llama_server_args import (
parse_cache_override,
parse_ctx_override,
resolve_cache_type_kv,
resolve_requested_ctx,
)
from core.tool_healing import (
_TC_END_TAG_RE,
_TC_FUNC_CLOSE_RE,
_TC_FUNC_START_RE,
_TC_JSON_START_RE,
_TC_PARAM_CLOSE_RE,
_TC_PARAM_START_RE,
_TOOL_ALL_PATS,
_TOOL_CLOSED_PATS,
parse_tool_calls_from_text,
strip_tool_call_markup,
)
from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import (
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
)
from core.inference.tool_call_parser import (
parse_tool_calls_from_text as _shared_parse_tool_calls_from_text,
)
logger = get_logger(__name__)
# ── Pre-compiled patterns for plan-without-action re-prompt ──
# Forward-looking intent signals that indicate the model is
# describing what it *will* do rather than giving a final answer.
_INTENT_SIGNAL = re.compile(
r"(?i)("
# Direct intent: "I'll ...", "I will ...", "Let me ...", "I am going to ..."
# Handles both straight and curly apostrophes.
# Excludes "I can", "I should", "I want to", "let's" which
# appear frequently in direct answers / explanations.
# Negative lookahead drops negated forms ("I will not", "I'll never")
# so a refusal doesn't trigger a re-prompt.
r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)"
r"|"
# Step/plan framing: "First ...", "Step 1:", "Here's my plan"
r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))"
r"|"
# "Now I" / "Next I" patterns
r"\b(?:now i|next i)\b"
r")"
)
_MAX_REPROMPTS = 3
# Without max_tokens, llama-server defaults to n_predict = n_ctx (up to
# 262144 for Qwen3.5), producing many-minute zombie decodes when cancel
# fails. t_max_predict_ms is a wall-clock backstop applied unconditionally,
# but the llama.cpp README notes it ONLY fires after a newline has been
# generated -- a model stuck in a long unbroken non-newline sequence is
# unbounded by it. So we still want a token cap as the front-line limiter.
#
# The cap is the model's effective context length when we know it,
# falling back to a generous floor when metadata is unavailable. 4096 was
# too low: Qwen3 / gpt-oss reasoning traces routinely exceed it, and any
# OpenAI-API caller that omits max_tokens (langchain, llama-index, raw
# curl) sees responses silently truncated mid-sentence.
_DEFAULT_MAX_TOKENS_FLOOR = 32768
_DEFAULT_T_MAX_PREDICT_MS = 600_000 # 10 min
_REPROMPT_MAX_CHARS = 2000
# ── Pre-compiled patterns for GGUF shard detection ───────────
_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$")
_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
# ── Sliding-window-pattern resolver ───────────────────────────
# Resolves the per-layer SWA mask when a GGUF reports a sliding window
# but no `sliding_window_pattern` field. Tier order in
# `_resolve_swa_pattern`: GGUF metadata, on-disk cache, bootstrap dict
# below, transformers introspection, HF Hub config.json, legacy 1/4
# fallback. Period N means layer i is SWA iff `(i + 1) % N != 0`,
# matching transformers. Skipped on purpose: phi3 (no key/val length
# in GGUF, window >= ctx anyway), qwen2 family (converter strips
# sliding_window when use_sliding_window=False), mistral v0.1/v0.2
# (all-SWA can't be expressed as a period).
_BOOTSTRAP_SWA_DEFAULTS: dict[str, int] = {
"gemma2": 2, # Gemma2Config.sliding_window_pattern
"gemma3": 6, # Gemma3TextConfig.sliding_window_pattern
"gemma3n": 5, # text_config.layer_types: SWA*4 + FULL
"gpt_oss": 2, # text_config.layer_types: alternating
"cohere2": 4, # Cohere2Config.sliding_window_pattern
}
# Process-wide cache backed by JSON on disk. Values are int period or
# list[bool] mask. Lazy-loaded.
_SWA_CACHE: Optional[dict] = None
_SWA_CACHE_LOCK = threading.Lock()
def _probe_dns_dead(host: str = "huggingface.co", timeout: float = 2.0) -> bool:
"""Quick DNS check. Runs on a daemon thread so concurrent sockets
in the same process are not affected by socket.setdefaulttimeout."""
result: list[Optional[bool]] = [None]
def _probe() -> None:
try:
socket.gethostbyname(host)
result[0] = False
except Exception:
result[0] = True
t = threading.Thread(target = _probe, daemon = True)
t.start()
t.join(timeout)
# Thread still running -> resolver wedged -> treat as dead.
return True if result[0] is None else result[0]
@contextlib.contextmanager
def _hf_offline_if_dns_dead():
"""Set HF_HUB_OFFLINE for the body of this block only when DNS to
huggingface.co fails. Restores the env on exit so a transient
resolver hiccup at the start of one load can't quarantine the whole
process. Respects an explicit user setting (no-op if already set)."""
if "HF_HUB_OFFLINE" in os.environ:
yield False
return
if not _probe_dns_dead():
yield False
return
transformers_was_set = "TRANSFORMERS_OFFLINE" in os.environ
os.environ["HF_HUB_OFFLINE"] = "1"
if not transformers_was_set:
os.environ["TRANSFORMERS_OFFLINE"] = "1"
logger.warning("huggingface.co unreachable; using local HF cache for this load.")
try:
yield True
finally:
os.environ.pop("HF_HUB_OFFLINE", None)
if not transformers_was_set:
os.environ.pop("TRANSFORMERS_OFFLINE", None)
def _swa_cache_path() -> Path:
home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME")
base = Path(home) if home else Path.home() / ".unsloth" / "studio"
return base / "swa_cache.json"
def _load_swa_cache() -> dict:
global _SWA_CACHE
with _SWA_CACHE_LOCK:
if _SWA_CACHE is not None:
return _SWA_CACHE
try:
with open(_swa_cache_path()) as f:
_SWA_CACHE = json.load(f)
if not isinstance(_SWA_CACHE, dict):
_SWA_CACHE = {}
except (FileNotFoundError, json.JSONDecodeError, OSError):
_SWA_CACHE = {}
return _SWA_CACHE
def _save_swa_cache(cache: dict) -> None:
try:
path = _swa_cache_path()
path.parent.mkdir(parents = True, exist_ok = True)
tmp = path.with_suffix(".json.tmp")
with open(tmp, "w") as f:
json.dump(cache, f, indent = 2, sort_keys = True)
tmp.replace(path)
except OSError:
pass
def _period_from_layer_types(layer_types: list) -> Optional[int]:
"""Smallest period N where `(i+1) % N != 0` matches the SWA mask,
or None if no fixed period fits."""
if not layer_types:
return None
is_swa = ["full" not in str(t).lower() for t in layer_types]
n = len(is_swa)
for N in range(1, n + 1):
if all(((i + 1) % N != 0) == is_swa[i] for i in range(n)):
return N
return None
def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]:
try:
from huggingface_hub import hf_hub_download
cfg_path = hf_hub_download(repo_id, "config.json", repo_type = "model")
with open(cfg_path) as f:
cfg = json.load(f)
except Exception:
return None
src = cfg.get("text_config") if isinstance(cfg.get("text_config"), dict) else cfg
period = src.get("sliding_window_pattern")
if isinstance(period, int) and period > 0:
return period
lt = src.get("layer_types")
if isinstance(lt, list) and lt:
return _period_from_layer_types(lt) or [
"full" not in str(t).lower() for t in lt
]
return None
def _arch_aliases(arch: str) -> tuple:
# GGUF emits `falcon-h1`; HF model_type is `falcon_h1`. Normalise both ways.
seen = []
for a in (arch, arch.replace("-", "_"), arch.replace("_", "-")):
if a and a not in seen:
seen.append(a)
return tuple(seen)
def _swa_entry_from_config_obj(cfg) -> Optional[object]:
src = getattr(cfg, "text_config", None) or cfg
period = getattr(src, "sliding_window_pattern", None)
if isinstance(period, int) and period > 0:
return period
lt = getattr(src, "layer_types", None)
if isinstance(lt, list) and lt:
return _period_from_layer_types(lt) or [
"full" not in str(t).lower() for t in lt
]
return None
_SWA_PATTERN_SOURCE_RE = re.compile(
r"sliding_window_pattern\s*(?::\s*[\w\[\], ]*)?\s*=\s*(\d+)"
)
def _resolve_swa_entry_from_transformers(arch: str) -> Optional[object]:
"""Default-instantiate the matching Config; on failure, regex-parse
its source for `sliding_window_pattern = N`."""
try:
from transformers.models.auto.configuration_auto import (
CONFIG_MAPPING,
CONFIG_MAPPING_NAMES,
)
except Exception:
return None
cfg_class = None
for alias in _arch_aliases(arch):
if alias in CONFIG_MAPPING_NAMES:
try:
cfg_class = CONFIG_MAPPING[alias]
break
except Exception:
cfg_class = None
if cfg_class is None:
return None
try:
if (entry := _swa_entry_from_config_obj(cfg_class())) is not None:
return entry
except Exception:
pass
import inspect
candidates = [cfg_class]
text_cfg_class = getattr(cfg_class, "sub_configs", {}).get("text_config")
if text_cfg_class is not None:
candidates.append(text_cfg_class)
for cls in candidates:
try:
src = inspect.getsource(cls)
except (OSError, TypeError):
continue
if m := _SWA_PATTERN_SOURCE_RE.search(src):
period = int(m.group(1))
if period > 0:
return period
return None
def _resolve_swa_pattern(
arch: Optional[str],
n_layers: Optional[int],
source_repo_candidates: tuple = (),
*,
allow_network: Optional[bool] = None,
) -> Optional[list]:
if not arch or not n_layers:
return None
if allow_network is None:
allow_network = os.environ.get("UNSLOTH_STUDIO_OFFLINE", "0") not in (
"1",
"true",
"True",
"yes",
)
cache = _load_swa_cache()
def _entry_to_mask(entry):
if isinstance(entry, int) and entry > 0:
return [(i + 1) % entry != 0 for i in range(n_layers)]
if isinstance(entry, list) and entry:
return [bool(entry[i % len(entry)]) for i in range(n_layers)]
return None
def _persist(entry):
with _SWA_CACHE_LOCK:
cache[arch] = entry
_save_swa_cache(cache)
if (entry := cache.get(arch)) is not None:
if (mask := _entry_to_mask(entry)) is not None:
return mask
if (entry := _BOOTSTRAP_SWA_DEFAULTS.get(arch)) is not None:
return _entry_to_mask(entry)
entry = _resolve_swa_entry_from_transformers(arch)
if entry is not None:
_persist(entry)
return _entry_to_mask(entry)
# Tier 3: live HF fetch (with persistent caching of the result)
if allow_network:
for repo_id in source_repo_candidates:
if not repo_id:
continue
entry = _fetch_swa_entry_from_hf(repo_id)
if entry is not None:
_persist(entry)
return _entry_to_mask(entry)
return None
def _hf_repo_from_url(url: Optional[str]) -> Optional[str]:
"""Strip `https://huggingface.co/owner/name(/...)` to `owner/name`."""
if not url or "huggingface.co/" not in url:
return None
tail = url.split("huggingface.co/", 1)[1].rstrip("/")
parts = tail.split("/")
if len(parts) < 2:
return None
return f"{parts[0]}/{parts[1]}"
# Model size extraction — lazy import to avoid pulling in transformers
# at module level. See PR description for the full explanation.
def _extract_model_size_b(model_id: str):
from utils.models import extract_model_size_b
return extract_model_size_b(model_id)
_TOOL_TEMPLATE_MARKERS = (
"{%- if tools %}",
"{%- if tools -%}",
"{% if tools %}",
"{% if tools -%}",
'"role" == "tool"',
"'role' == 'tool'",
'message.role == "tool"',
"message.role == 'tool'",
)
def detect_reasoning_flags(
chat_template: Optional[str],
model_identifier: Optional[str] = None,
*,
log_source: Optional[str] = None,
) -> dict:
"""Classify a chat template's reasoning and tool-calling capabilities.
Returns a dict with the same five keys populated by the GGUF sniffer:
``supports_reasoning``, ``reasoning_style``
(``"enable_thinking"`` | ``"reasoning_effort"``),
``reasoning_always_on``, ``supports_preserve_thinking``, and
``supports_tools``. Used by both the llama-server backend at load
time and the safetensors/transformers paths in ``routes/inference``
so the two agree on what the frontend will see.
"""
flags = {
"supports_reasoning": False,
"reasoning_style": "enable_thinking",
"reasoning_always_on": False,
"supports_preserve_thinking": False,
"supports_tools": False,
}
if not chat_template:
return flags
tpl = chat_template
prefix = f"{log_source}: " if log_source else ""
if "enable_thinking" in tpl:
flags["supports_reasoning"] = True
flags["reasoning_style"] = "enable_thinking"
logger.info(f"{prefix}model supports reasoning (enable_thinking)")
elif "reasoning_effort" in tpl:
# gpt-oss / Harmony templates use reasoning_effort
# ("low" | "medium" | "high") instead of a boolean.
flags["supports_reasoning"] = True
flags["reasoning_style"] = "reasoning_effort"
logger.info(f"{prefix}model supports reasoning (reasoning_effort)")
elif "thinking" in tpl:
# DeepSeek uses 'thinking' instead of 'enable_thinking'
normalized_id = (model_identifier or "").lower()
if "deepseek" in normalized_id:
flags["supports_reasoning"] = True
logger.info(f"{prefix}model supports reasoning (DeepSeek thinking)")
# Hardcoded <think> tags or reasoning_content in the template mean
# thinking is always on (no toggle to disable it).
if not flags["supports_reasoning"]:
if ("<think>" in tpl and "</think>" in tpl) or "reasoning_content" in tpl:
flags["supports_reasoning"] = True
flags["reasoning_always_on"] = True
logger.info(f"{prefix}model always reasons (<think> tags in template)")
# preserve_thinking is an independent kwarg on some Qwen templates
# that keeps historical <think> blocks in prior assistant turns.
if "preserve_thinking" in tpl:
flags["supports_preserve_thinking"] = True
logger.info(f"{prefix}model supports preserve_thinking")
if any(marker in tpl for marker in _TOOL_TEMPLATE_MARKERS):
flags["supports_tools"] = True
logger.info(f"{prefix}model supports tool calling")
return flags
def _is_mtp_model_name(
model_identifier: Optional[str],
gguf_path: Optional[str] = None,
) -> bool:
"""Name-based MTP detector. Fallback for the metadata signal."""
for cand in (model_identifier, Path(gguf_path).name if gguf_path else None):
if cand and "-mtp" in cand.lower():
return True
return False
def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool:
"""User passed --spec-type / --spec-default? llama-server takes a
single --spec-type (comma-separated to chain), so suppress
auto-emit when this is true."""
if not extra_args:
return False
for raw in extra_args:
tok = str(raw)
if not tok.startswith("--"):
continue
flag = tok.split("=", 1)[0]
if flag in ("--spec-type", "--spec-default"):
return True
return False
def _build_ngram_mod_flags(
caps: Optional[dict],
n_match: int = 24,
n_min: int = 48,
n_max: int = 64,
) -> list[str]:
"""Emit the right ngram-mod knob flags for the running llama-server.
Post-rename builds expose ``--spec-ngram-mod-n-{match,min,max}``;
pre-rename builds expose the legacy ``--spec-ngram-size-n`` /
``--draft-min`` / ``--draft-max``. ``caps`` comes from
``probe_server_capabilities``; ``ngram_mod_flavor`` tells us which
set is real (vs a removal-stub entry). Returns ``[]`` when neither
set is available so the caller can drop ngram-mod entirely.
"""
flavor = caps.get("ngram_mod_flavor") if caps else None
if flavor == "new":
return [
"--spec-ngram-mod-n-match",
str(n_match),
"--spec-ngram-mod-n-min",
str(n_min),
"--spec-ngram-mod-n-max",
str(n_max),
]
if flavor == "legacy":
# Legacy llama.cpp before the spec arg rename: same knobs lived
# under --spec-ngram-size-n (lookup length) and the generic
# --draft-min / --draft-max (ngram size N range).
return [
"--spec-ngram-size-n",
str(n_match),
"--draft-min",
str(n_min),
"--draft-max",
str(n_max),
]
return []
# Canonical Speculative Decoding modes exposed by the Studio chat UI.
# The dropdown renders five options (auto, mtp, ngram, mtp+ngram, off);
# the load API also accepts legacy values that the original Switch and
# external callers emit (default, draft-mtp, ngram-mod, ngram-simple).
_CANONICAL_SPEC_MODES = {"auto", "mtp", "ngram", "mtp+ngram", "off", "ngram-simple"}
_LEGACY_SPEC_MODE_MAP = {
"default": "auto",
"draft-mtp": "mtp",
"ngram-mod": "ngram",
}
def _canonicalize_spec_mode(value):
"""Map any accepted ``speculative_type`` input onto a canonical mode.
Returns one of ``auto``, ``mtp``, ``ngram``, ``mtp+ngram``, ``off``,
``ngram-simple``, or ``None`` (callers treat ``None`` as ``auto``).
Unknown strings collapse to ``auto`` so a stale UI value or typo
falls back to the safe platform-aware path.
"""
if value is None:
return None
if not isinstance(value, str):
return None
stripped = value.strip().lower()
if not stripped:
return None
if stripped in _CANONICAL_SPEC_MODES:
return stripped
if stripped in _LEGACY_SPEC_MODE_MAP:
return _LEGACY_SPEC_MODE_MAP[stripped]
# llama.cpp comma-chains are emitted by old persisted state e.g.
# "ngram-mod,draft-mtp"; collapse the most common one explicitly.
pieces = [p.strip() for p in stripped.split(",") if p.strip()]
has_mtp = any(p in ("mtp", "draft-mtp") for p in pieces)
has_ngram = any(p in ("ngram", "ngram-mod") for p in pieces)
if has_mtp and has_ngram:
return "mtp+ngram"
if has_mtp:
return "mtp"
if has_ngram:
return "ngram"
return "auto"
def _backfill_usage_from_timings(usage, timings):
"""Synthesize ``usage`` from llama-server's ``timings`` when the
OpenAI-style usage block is missing or reports zero tokens.
The Studio chat UI computes generation t/s from
``meta.usage.completion_tokens / totalStreamTime``. llama-server
always populates ``timings.predicted_n`` (true decoded count) and
``timings.prompt_n``, but the ``usage`` field on the final SSE chunk
can be absent or zero on some server builds / streaming
configurations, which makes the UI fall back to wall-clock t/s and
dilute speculative-decoding speedups.
"""
if not timings:
return usage
if usage and usage.get("completion_tokens"):
return usage
predicted_n = timings.get("predicted_n")
prompt_n = timings.get("prompt_n")
if predicted_n is None and prompt_n is None:
return usage
out = dict(usage or {})
if not out.get("completion_tokens") and predicted_n is not None:
out["completion_tokens"] = predicted_n
if not out.get("prompt_tokens") and prompt_n is not None:
out["prompt_tokens"] = prompt_n
out["total_tokens"] = int(out.get("prompt_tokens") or 0) + int(
out.get("completion_tokens") or 0
)
return out
class LlamaCppBackend:
"""
Manages a llama-server subprocess for GGUF model inference.
Lifecycle:
1. load_model() — starts llama-server with the GGUF file
2. generate_chat_completion() — proxies to /v1/chat/completions, streams back
3. unload_model() — terminates llama-server subprocess
"""
def __init__(self):
self._process: Optional[subprocess.Popen] = None
self._port: Optional[int] = None
self._model_identifier: Optional[str] = None
self._gguf_path: Optional[str] = None
self._hf_repo: Optional[str] = None
self._hf_variant: Optional[str] = None
self._is_vision: bool = False
self._healthy = False
# Set by _classify_gpu_offload after _wait_for_health.
self._gpu_offload_active: Optional[bool] = None
self._context_length: Optional[int] = None
self._effective_context_length: Optional[int] = None
self._max_context_length: Optional[int] = None
self._chat_template: Optional[str] = None
self._chat_template_override: Optional[str] = None
self._supports_reasoning: bool = False
self._reasoning_always_on: bool = False
self._reasoning_style: str = "enable_thinking"
self._supports_preserve_thinking: bool = False
self._supports_tools: bool = False
self._cache_type_kv: Optional[str] = None
self._reasoning_default: bool = True
self._speculative_type: Optional[str] = None
# Canonical UI-facing mode the user requested: one of
# ``auto``/``mtp``/``ngram``/``mtp+ngram``/``off``/``ngram-simple``.
# Round-tripped through the status API so the dropdown reflects
# the picked mode rather than the resolved internal flag set
# (auto on a 27B MTP GGUF resolves to draft-mtp but the dropdown
# should still read "Auto").
self._requested_spec_mode: Optional[str] = None
# User-supplied --spec-draft-n-max override (None = platform default).
self._spec_draft_n_max: Optional[int] = None
# KV-cache estimation fields (populated by _read_gguf_metadata)
self._n_layers: Optional[int] = None
self._n_kv_heads: Optional[int] = None
self._n_kv_heads_by_layer: Optional[list[int]] = None
self._n_heads: Optional[int] = None
self._embedding_length: Optional[int] = None
# Architecture-aware KV fields for 5-path estimation
self._kv_key_length: Optional[int] = None
self._kv_value_length: Optional[int] = None
self._sliding_window: Optional[int] = None
self._sliding_window_pattern: Optional[list[bool]] = None
self._full_attention_interval: Optional[int] = None
self._kv_lora_rank: Optional[int] = None
self._key_length_mla: Optional[int] = None
self._kv_key_length_swa: Optional[int] = None
self._kv_value_length_swa: Optional[int] = None
self._ssm_inner_size: Optional[int] = None
self._ssm_state_size: Optional[int] = None
# Last N layers reuse KV from earlier layers and don't allocate
# their own cache (Gemma 3n / Gemma 4: <arch>.attention.shared_kv_layers).
self._shared_kv_layers: Optional[int] = None
# MTP head count (llama.cpp #22673); >0 enables --spec-type draft-mtp.
self._nextn_predict_layers: Optional[int] = None
self._lock = threading.Lock()
# Wraps load_model() end-to-end so concurrent loads serialise
# and never coexist as two llama-server processes (#5401).
self._serial_load_lock = threading.Lock()
# Last extra_args / requested n_ctx, preserved across unload so
# the chat UI's /unload+/load Apply path can inherit them (#5401).
# ``_extra_args_source`` records the (model_identifier, hf_variant)
# the stored args came from so the route can refuse cross-model
# inheritance.
self._extra_args: Optional[List[str]] = None
self._extra_args_source: Optional[tuple[str, Optional[str]]] = None
self._requested_n_ctx: int = 0
self._stdout_lines: list[str] = []
self._stdout_thread: Optional[threading.Thread] = None
# llama-server tee log (see _drain_stdout / _kill_process).
self._llama_log_fh = None
self._llama_log_path: Optional[Path] = None
self._cancel_event = threading.Event()
self._api_key: Optional[str] = None
# True once a probe has completed; cleared on transient failure.
self._is_audio: bool = False
self._audio_type: Optional[str] = None
self._audio_probed: bool = False
# Monotonic timestamp set in _kill_process; read by load_model
# to decide whether to wait for the VRAM reclaim to finish.
self._last_kill_monotonic: float = 0.0
self._kill_orphaned_servers()
atexit.register(self._cleanup)
# ── Properties ────────────────────────────────────────────────
@property
def is_loaded(self) -> bool:
return self._process is not None and self._healthy
@property
def is_active(self) -> bool:
"""True if a llama-server process exists (loading or loaded)."""
return self._process is not None
@property
def base_url(self) -> str:
return f"http://127.0.0.1:{self._port}"
@property
def model_identifier(self) -> Optional[str]:
return self._model_identifier
@property
def is_vision(self) -> bool:
return self._is_vision
@property
def hf_variant(self) -> Optional[str]:
return self._hf_variant
@property
def extra_args(self) -> Optional[List[str]]:
"""Extra llama-server flags from the last load. Copy; None = never
set, [] = explicitly cleared. Used by the route for inheritance."""
return list(self._extra_args) if self._extra_args is not None else None
@property
def requested_n_ctx(self) -> int:
"""n_ctx the last load was invoked with (not the effective cap).
0 means Auto. Used by the route to detect Auto-vs-explicit flips."""
return self._requested_n_ctx
@property
def extra_args_source(self) -> Optional[tuple[str, Optional[str]]]:
"""(model_identifier, hf_variant) the stored extra_args came from.
``None`` if no extras have ever been recorded. Used by the route
to refuse cross-model inheritance (#5401)."""
return self._extra_args_source
@property
def context_length(self) -> Optional[int]:
"""Return the effective context length the server is running at."""
return self._effective_context_length or self._context_length
@property
def max_context_length(self) -> Optional[int]:
"""Return the largest context that fits on this hardware at load time.
This is the "safe zone" threshold the UI renders warnings
against. For a model whose weights fit on some GPU subset, it
is the binary-search cap from ``_fit_context_to_vram`` for that
subset. For a model whose weights exceed 90% of every GPU
subset, it is the 4096 fallback -- the spec's default when the
model will not fit. The UI slider ceiling is
``native_context_length``; dragging above ``max_context_length``
triggers the "might be slower" warning.
"""
return self._max_context_length or self._context_length
@property
def native_context_length(self) -> Optional[int]:
"""Return the model's native context length from GGUF metadata."""
return self._context_length
def load_progress(self) -> Optional[dict]:
"""Return live model-load progress, or None if not loading.
While llama-server is warming up, its process is typically in
kernel state D (disk sleep) mmap'ing the weight shards into
page cache before pushing layers to VRAM. During that window
``/api/inference/status`` only reports ``loading``, which gives
the UI nothing to display besides a spinner that looks stuck
for minutes on large MoE models.
This method samples ``/proc/<pid>/status VmRSS`` against the
sum of the GGUF shard sizes so the UI can render a real bar
and compute rate / ETA. Returns ``None`` when no load is in
flight (no process, or process already healthy).
Shape::
{
"phase": "mmap" | "ready",
"bytes_loaded": int, # VmRSS of the llama-server
"bytes_total": int, # sum of shard file sizes
"fraction": float, # bytes_loaded / bytes_total, 0..1
}
Linux-only in the current implementation. On macOS/Windows the
equivalent would be a different API; this returns ``None`` on
platforms where ``/proc/<pid>/status`` is unavailable.
"""
proc = self._process
if proc is None:
return None
pid = proc.pid
if pid is None:
return None
# Sum up shard sizes (primary + any extras sitting alongside).
bytes_total = 0
gguf_path = self._gguf_path
if gguf_path:
primary = Path(gguf_path)
try:
if primary.is_file():
bytes_total += primary.stat().st_size
except OSError:
pass
# Extra shards live alongside the primary with the same prefix
# before the shard index (e.g. ``-00001-of-00004.gguf``).
try:
parent = primary.parent
stem = primary.name
m = _SHARD_RE.match(stem)
prefix = m.group(1) if m else None
if prefix and parent.is_dir():
for sibling in parent.iterdir():
if (
sibling.is_file()
and sibling.name.startswith(prefix)
and sibling.name != stem
and sibling.suffix == ".gguf"
):
try:
bytes_total += sibling.stat().st_size
except OSError:
pass
except OSError:
pass
# Read VmRSS from /proc/<pid>/status. Kilobytes on Linux.
bytes_loaded = 0
try:
with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f:
for line in f:
if line.startswith("VmRSS:"):
kb = int(line.split()[1])
bytes_loaded = kb * 1024
break
except (FileNotFoundError, PermissionError, ValueError, OSError):
return None
phase = "ready" if self._healthy else "mmap"
fraction = 0.0
if bytes_total > 0:
fraction = min(1.0, bytes_loaded / bytes_total)
return {
"phase": phase,
"bytes_loaded": bytes_loaded,
"bytes_total": bytes_total,
"fraction": round(fraction, 4),
}
@property
def chat_template(self) -> Optional[str]:
return self._chat_template
@property
def chat_template_override(self) -> Optional[str]:
return self._chat_template_override
@property
def supports_reasoning(self) -> bool:
return self._supports_reasoning
@property
def reasoning_always_on(self) -> bool:
return self._reasoning_always_on
@property
def reasoning_style(self) -> str:
return self._reasoning_style
@property
def supports_preserve_thinking(self) -> bool:
return self._supports_preserve_thinking
@property
def reasoning_default(self) -> bool:
return self._reasoning_default
def _reasoning_kwargs(self, enable_thinking: bool) -> dict:
if self._reasoning_style == "reasoning_effort":
return {"reasoning_effort": "high" if enable_thinking else "low"}
return {"enable_thinking": enable_thinking}
def _request_reasoning_kwargs(
self,
enable_thinking: Optional[bool],
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
) -> Optional[dict]:
"""Build chat_template_kwargs from per-request reasoning fields.
Produces a merged dict covering the active model's reasoning style
(``enable_thinking`` or ``reasoning_effort``) plus the independent
``preserve_thinking`` kwarg when the template supports it.
"""
kwargs: dict = {}
# Always-on reasoning models hardcode <think> tags in their template
# and do not consume enable_thinking / reasoning_effort -- skip.
if self._supports_reasoning and not self._reasoning_always_on:
if self._reasoning_style == "reasoning_effort":
if reasoning_effort in ("low", "medium", "high"):
kwargs["reasoning_effort"] = reasoning_effort
elif enable_thinking is not None:
kwargs["reasoning_effort"] = "high" if enable_thinking else "low"
else:
if enable_thinking is not None:
kwargs["enable_thinking"] = enable_thinking
if self._supports_preserve_thinking and preserve_thinking is not None:
kwargs["preserve_thinking"] = preserve_thinking
return kwargs or None
@property
def supports_tools(self) -> bool:
return self._supports_tools
@property
def cache_type_kv(self) -> Optional[str]:
return self._cache_type_kv
@property
def speculative_type(self) -> Optional[str]:
return self._speculative_type
@property
def requested_spec_mode(self) -> Optional[str]:
"""Canonical UI-facing mode the user requested (see field doc)."""
return self._requested_spec_mode
@property
def spec_draft_n_max(self) -> Optional[int]:
"""User --spec-draft-n-max override active on the load, or None
when the platform default (6 GPU / 3 CPU) is in effect."""
return self._spec_draft_n_max
# ── Binary discovery ──────────────────────────────────────────
@staticmethod
def _find_llama_server_binary() -> Optional[str]:
"""
Locate the llama-server binary.
Search order:
1. LLAMA_SERVER_PATH environment variable (direct path to binary)
1b. UNSLOTH_LLAMA_CPP_PATH env var (custom llama.cpp install dir)
2. ~/.unsloth/llama.cpp/llama-server (make build, root dir)
3. ~/.unsloth/llama.cpp/build/bin/llama-server (cmake build, Linux)
4. ~/.unsloth/llama.cpp/build/bin/Release/llama-server.exe (cmake build, Windows)
5. ./llama.cpp/llama-server (legacy: make build, root dir)
6. ./llama.cpp/build/bin/llama-server (legacy: cmake in-tree build)
7. llama-server on PATH (system install)
8. ./bin/llama-server (legacy: extracted binary)
"""
import os
import sys
binary_name = "llama-server.exe" if sys.platform == "win32" else "llama-server"
# 1. Env var — direct path to binary
env_path = os.environ.get("LLAMA_SERVER_PATH")
if env_path and Path(env_path).is_file():
return env_path
# 1b. UNSLOTH_LLAMA_CPP_PATH — custom llama.cpp install directory
custom_llama_cpp = os.environ.get("UNSLOTH_LLAMA_CPP_PATH")
if custom_llama_cpp:
custom_dir = Path(custom_llama_cpp)
# Root dir (make builds)
root_bin = custom_dir / binary_name
if root_bin.is_file():
return str(root_bin)
# build/bin/ (cmake builds on Linux)
cmake_bin = custom_dir / "build" / "bin" / binary_name
if cmake_bin.is_file():
return str(cmake_bin)
# build/bin/Release/ (cmake builds on Windows)
if sys.platform == "win32":
win_bin = custom_dir / "build" / "bin" / "Release" / binary_name
if win_bin.is_file():
return str(win_bin)
# 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"
try:
from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433
_resolved_sr = _sr()
_legacy_studio = Path.home() / ".unsloth" / "studio"
try:
_is_legacy = _resolved_sr.resolve() == _legacy_studio.resolve()
except (OSError, ValueError):
_is_legacy = _resolved_sr == _legacy_studio
if _is_legacy:
search_roots = [legacy_llama]
else:
# why: _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"]
except (ImportError, OSError, ValueError):
search_roots = [legacy_llama]
_seen_roots: set[str] = set()
_unique_roots: list[Path] = []
for r in search_roots:
k = str(r)
if k not in _seen_roots:
_seen_roots.add(k)
_unique_roots.append(r)
for unsloth_home in _unique_roots:
home_root = unsloth_home / binary_name
if home_root.is_file():
return str(home_root)
home_linux = unsloth_home / "build" / "bin" / binary_name
if home_linux.is_file():
return str(home_linux)
if sys.platform == "win32":
home_win = unsloth_home / "build" / "bin" / "Release" / binary_name
if home_win.is_file():
return str(home_win)
# 56. Legacy: in-tree build (older setup.sh / setup.ps1 versions)
project_root = Path(__file__).resolve().parents[4]
# Root dir (make builds)
root_path = project_root / "llama.cpp" / binary_name
if root_path.is_file():
return str(root_path)
# build/bin/ (cmake builds)
build_path = project_root / "llama.cpp" / "build" / "bin" / binary_name
if build_path.is_file():
return str(build_path)
if sys.platform == "win32":
win_path = (
project_root / "llama.cpp" / "build" / "bin" / "Release" / binary_name
)
if win_path.is_file():
return str(win_path)
# 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 bin_path.is_file():
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}.
``ngram_mod_flavor`` is ``"new"`` when the binary exposes the
post-rename ``--spec-ngram-mod-n-match / -n-min / -n-max`` as
real args, ``"legacy"`` when only the pre-rename
``--spec-ngram-size-n / --draft-min / --draft-max`` are real
(the rename ships with stub removal entries for the legacy
names; we tell stubs apart by the "argument has been removed"
description), or ``None`` if neither set is usable.
``spec_draft_n_max_flag`` is the actual flag name the binary
accepts: ``--spec-draft-n-max`` on post-rename builds, or
``--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,
}
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
try:
result = subprocess.run(
[bin_path, "--help"],
capture_output = True,
text = True,
timeout = 10,
check = False,
)
help_text = (result.stdout or "") + "\n" + (result.stderr or "")
# Split into per-flag blocks: each --flag line plus its
# indented continuation lines, 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 token that isn't itself a flag, 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 detection from --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
# the new args (real) and the legacy ones (stubs); pre-rename
# builds only have the 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"
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,
}
cls._capability_cache[cache_key] = info
return info
# ── GPU allocation ────────────────────────────────────────────
@staticmethod
def _get_gguf_size_bytes(model_path: str) -> int:
"""Get 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$"
)
for sibling in main.parent.iterdir():
if sibling != main and sibling_pat.match(sibling.name):
total += sibling.stat().st_size
return total
@staticmethod
def _amd_apu_wants_unified_memory() -> bool:
"""True only for AMD unified-memory APUs (gfx1150/gfx1151), where
GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM.
False for discrete AMD, NVIDIA, CPU and macOS (the env hurts discrete
GPUs). ROCm reuses torch.cuda.*; the gcnArchName suffix is stripped."""
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
for _i in range(torch.cuda.device_count()):
try:
_arch = (
getattr(torch.cuda.get_device_properties(_i), "gcnArchName", "")
or ""
)
except Exception:
continue
if _arch.split(":")[0].strip().lower() in {"gfx1150", "gfx1151"}:
return True
except Exception:
return False
return False
@staticmethod
def _get_gpu_free_memory() -> list[tuple[int, int]]:
"""Query free 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 because the HIP runtime
reuses the entire ``torch.cuda.*`` namespace. Covers the
AMD case for issue #5106 (nvidia-smi-only probe silently
returned [] on AMD hosts) and also rescues NVIDIA hosts
where ``nvidia-smi`` is missing from PATH.
Returns list of (gpu_index, free_mib) sorted by index. Empty
list if no supported GPU is reachable.
"""
import os
# ── NVIDIA via nvidia-smi ────────────────────────────────────
try:
result = subprocess.run(
[
"nvidia-smi",
"--query-gpu=index,memory.free",
"--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 like
# "0,1," which would otherwise raise ValueError on
# an empty token. An explicitly empty mask (CVD="")
# yields an empty `allowed` set so all GPUs are
# filtered out, matching the codebase convention.
allowed = set(
int(x.strip()) for x in cvd.split(",") if x.strip()
)
except ValueError:
pass
gpus: list[tuple[int, int]] = []
for line in result.stdout.strip().splitlines():
parts = line.split(",")
if len(parts) == 2:
idx = int(parts[0].strip())
free_mib = int(parts[1].strip())
if allowed is not None and idx not in allowed:
continue
gpus.append((idx, free_mib))
# Match the docstring's sort-by-id guarantee. nvidia-smi
# almost always returns sorted output, but driver order
# is not formally guaranteed.
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.
# On NVIDIA builds the mask is CUDA_VISIBLE_DEVICES; on AMD
# ROCm builds it is HIP_VISIBLE_DEVICES (or ROCR_VISIBLE_DEVICES
# if HIP is unset). Downstream we feed these IDs back into the
# llama-server subprocess as CVD, so we must translate visible
# ordinals back to physical indices first; otherwise launching
# with ``CUDA_VISIBLE_DEVICES=2,3`` would get rewritten to
# ``CUDA_VISIBLE_DEVICES=0,1`` and target the wrong GPUs.
physical_ids: Optional[list[int]] = None
# Match the codebase convention in
# ``utils/hardware/hardware.py::_get_parent_visible_gpu_spec``:
# treat an explicitly empty mask (``HIP_VISIBLE_DEVICES=""``)
# as "set to no GPUs" rather than falling through to the next
# var. ``or`` would coerce empty string to falsy and silently
# promote the wrong source.
if getattr(torch.version, "hip", None) is not None:
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 not None:
try:
# Empty mask (CVD="") yields an empty list so the
# below loop produces no GPUs, consistent with the
# nvidia-smi path and utils/hardware/hardware.py.
physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()]
except ValueError:
physical_ids = None
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)))
# 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 []
# Skip the wait when the last kill is older than this; the GPU
# 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 GPU driver reclaims a dead process's allocations
asynchronously, so sampling free memory in the kill-to-spawn
window reads artificially low and pushes ``_select_gpus`` /
``_fit_context_to_vram`` toward needless CPU offload -- on a
tight VRAM card this is the Apply-reload OOM that bare-shell
launches with the same flags never see.
Short-circuits on cold start (``since_kill`` zero) or stale
kill (older than ``_VRAM_SETTLE_WINDOW_S``); also on CPU-only
hosts (empty probe), probe exceptions, and GPU-set changes.
``max_wait`` is a wall-clock bound that includes probe time,
so a wedged ``nvidia-smi`` cannot 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 on large-VRAM cards.
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``. 5% headroom covers CUDA context +
# compute buffers; 0.90 was too conservative and dropped 91-94%
# fits to CPU offload (#5106). The fork's --fit on still catches
# the truly-too-large case.
_GPU_PIN_VRAM_FRACTION = 0.95
@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, with parity for the Windows-specific wheel layouts seen
in the wild. Covered patterns:
* ``nvidia/<pkg>/bin`` -- legacy modular wheels
(``nvidia-cuda-runtime-cu12``, ``nvidia-cublas-cu12``, etc.).
* ``nvidia/<pkg>/bin/x86_64`` and ``.../bin/x64`` -- current
CUDA 13 wheel layout used by the unsuffixed
``nvidia-cuda-runtime`` / ``nvidia-cublas`` packages, which
ship under ``nvidia/cu13/bin/x86_64/`` (#5106).
* ``nvidia/<pkg>/Library/bin`` (and arch subdirs) -- conda-
style wheel repacks.
* ``torch/lib`` -- PyTorch's own CUDA-bundled Windows wheel,
which can ship ``cudart64_*.dll`` directly here instead of
as separate ``nvidia-*`` wheels. The install-side helper
``python_runtime_dirs`` in ``install_llama_prebuilt.py``
covers this path for the same reason.
Walks the tree with ``Path.iterdir`` rather than ``glob.glob``
so the resolver is safe against Windows paths containing
``[`` or ``]`` (valid in usernames; would otherwise be
interpreted as a glob character class and silently miss
existing dirs)."""
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
# Order matters for PATH search: arch-specific subdirs
# first so the explicit cudart64_X.dll location wins
# over a sibling ``bin`` that might be empty.
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 the win32 branch of start_llama_server
prepends so llama-server.exe resolves cudart / cublas DLLs:
binary_dir, pip nvidia wheels, CUDA_PATH/bin, CUDA_PATH/bin/x64.
Extracted so test_windows_gpu_detection_mock asserts against
production logic, not a hand-copy. #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 _select_gpus(
model_size_bytes: int,
gpus: list[tuple[int, int]],
) -> tuple[Optional[list[int]], bool]:
"""Pick GPU(s) for a model based on estimated VRAM and free memory.
``model_size_bytes`` should include both model weights and estimated
KV cache. The ``_GPU_PIN_VRAM_FRACTION`` threshold provides headroom
for compute buffers, CUDA context, and other runtime overhead.
Returns (gpu_indices, use_fit):
- ([1], False) model fits on 1 GPU at the headroom threshold
- ([1, 2], False) model needs 2 GPUs
- (None, True) model too large, let --fit handle it
"""
if not gpus:
return None, True
model_size_mib = model_size_bytes / (1024 * 1024)
usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION
# Sort GPUs by free memory descending
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
# Try fitting on 1 GPU at the usable-VRAM threshold.
if ranked[0][1] * usable_fraction >= model_size_mib:
return [ranked[0][0]], False
# Try fitting on N GPUs (accumulate free memory from most-free)
cumulative = 0
selected = []
for idx, free_mib in ranked:
selected.append(idx)
cumulative += free_mib * usable_fraction
if cumulative >= model_size_mib:
return sorted(selected), False
# Model is 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 is sufficient (K-only cache)
if self._kv_lora_rank is not None:
return True
# New-style: need both 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 _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.
Uses 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``: force SWA layers to cache the
full ``n_ctx`` (collapses path 3 to path 4
sizing for the SWA layers).
n_parallel -- ``--parallel``: number of server slots.
Verified empirically against llama-server:
non-SWA layers stay constant (cells split
across slots), SWA layers scale linearly
(per-slot window).
kv_unified -- ``--kv-unified`` (default on): retained for
API forward-compat. Currently a no-op for
memory math because the unified buffer total
matches per-slot buffers in measured cases.
ctx_checkpoints -- ``--ctx-checkpoints``: SWA snapshot count per
slot (PR #15293). Each snapshot stores one
sliding-window of state per SWA layer.
Returns 0 if metadata is insufficient for estimation.
"""
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 KV from earlier layers in the last
# ``shared_kv_layers`` blocks -- those don't allocate their own
# cache. Floor at 1 so a misconfigured 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 = {
"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_kv or "f16", 2.0)
slots = max(1, n_parallel)
# Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5)
# MLA stores one compressed KV latent per token/layer (shared across heads).
# V is reconstructed from the latent on the fly -- no separate V cache.
# key_length = kv_lora_rank + rope_dim (the full compressed representation).
# 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 would 128x.
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._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator]
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 is filled in by the resolver at parse time; if absent,
# falls through to the legacy 1/4-global heuristic below.
# Per-layer-type ``--parallel N`` accounting (verified empirically
# against ``llama-server``):
# * non-SWA layers: total cells = n_ctx, partitioned across
# slots -> total memory CONSTANT in slots.
# * SWA layers: per-slot cells = 2 * sliding_window
# (capped at n_ctx and at per_slot_ctx
# when ctx is split among many slots) ->
# total memory grows LINEARLY in slots.
# ``--swa-full`` forces full n_ctx for SWA layers instead.
# ``--ctx-checkpoints N`` adds N snapshots per SWA layer per slot.
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`` makes SWA layers cache the full context just
# like non-SWA: cells get partitioned across slots, so per-slot
# cells = per_slot_ctx and the slots*per-slot product collapses
# back to the constant ``n_ctx`` total. Otherwise SWA caches
# 2*sliding_window per slot, clamped at the 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
# Iterate only over layers that allocate their own KV;
# the 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._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator]
return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe)
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,
) -> int:
"""Return the largest context length that fits in GPU VRAM.
Uses 90% of available VRAM as the ctx-fit budget. Tighter than
``_GPU_PIN_VRAM_FRACTION`` on purpose: over-promising context
OOMs at runtime, while pinning conservatively just defers to
--fit on. If the weights alone don't fit, returns
``requested_ctx`` unchanged.
``kv_on_gpu`` mirrors ``--kv-offload`` (default on). When False
the KV cache lives in CPU RAM and doesn't compete with weights
for VRAM; the requested context is honored verbatim. The other
keyword args mirror ``_estimate_kv_cache_bytes``.
``mtp_engaged`` reserves extra VRAM for the MTP draft model's
KV cache + compute graph buffers. llama.cpp's MTP path keeps a
secondary cache sized off the target's KV; on tight VRAM tiers
(e.g. 32 GB) auto-fit at native context would otherwise spill
and force llama-server into a slower partial-offload 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,
)
# MTP needs a tighter budget; drop from 0.90 to 0.85.
budget_frac = 0.85 if mtp_engaged else 0.90
budget_bytes = available_mib * 1024 * 1024 * budget_frac
model_footprint = model_size_bytes
# Check if requested context already fits
kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs)
if model_footprint + kv <= budget_bytes:
return requested_ctx
# Model weights alone exceed budget -- can't help by reducing ctx.
# Return requested_ctx unchanged; --fit will handle VRAM management.
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
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 <= remaining:
best = mid
lo = mid + 1
else:
hi = mid - 1
# Round down to nearest 256 for alignment, but never exceed 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.
For example, UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total.
Returns (first_shard_filename, total_size_bytes) or None if nothing fits.
"""
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.endswith(".gguf") and "mmproj" not in f.lower()
]
if not gguf_files:
return None
# Get 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 files 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))
# Sort by total size ascending and pick the 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 lines from the subprocess stdout in a background thread.
This prevents a pipe-buffer deadlock on Windows where the default
pipe buffer is only ~4 KB. 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 (especially in CI) has the full subprocess 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 (ValueError, OSError):
# Pipe closed — process is terminating
pass
# 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 from any previously loaded model so stale flags
# (eg _supports_reasoning) do not 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_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._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
try:
WANTED = {
"general.architecture",
"tokenizer.chat_template",
# 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",
}
# Additional arch-specific keys are added dynamically once
# we know the architecture name.
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
# fetched via HTTP byte-range): bail out gracefully
# so the resolver fallback still runs on whatever
# we did manage to parse.
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
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}.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]
)
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]
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 val_a:
self._n_kv_heads = max(int(x) for x in val_a)
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 fetch
# of just the GGUF header); break so the
# resolver fallback still runs on what we have.
break
# 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). See `_resolve_swa_pattern`.
if (
self._sliding_window_pattern is None
and self._sliding_window
and self._n_layers
):
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']}/"
f"{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_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}")
# ── HF download (no lock held) ───────────────────────────────
def _download_gguf(
self,
*,
hf_repo: str,
hf_variant: Optional[str] = None,
hf_token: Optional[str] = None,
) -> str:
"""Download GGUF file(s) from HuggingFace. Returns local path.
Runs WITHOUT self._lock so that unload_model() can set
_cancel_event at any time. Checks _cancel_event between
each shard download.
"""
try:
from huggingface_hub import hf_hub_download
except ImportError:
raise RuntimeError(
"huggingface_hub is required for HF model loading. "
"Install it with: pip install huggingface_hub"
)
# Determine 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)
variant_lower = hf_variant.lower()
boundary = re.compile(
r"(?<![a-zA-Z0-9])" + re.escape(variant_lower) + r"(?![a-zA-Z0-9])"
)
gguf_files = sorted(
f
for f in files
if f.endswith(".gguf") and boundary.search(f.lower())
)
if gguf_files:
gguf_filename = gguf_files[0]
m = _SHARD_FULL_RE.match(gguf_filename)
if m:
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$"
)
gguf_extra_shards = [
f for f in gguf_files[1:] if sibling_pat.match(f)
]
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
boundary = re.compile(
r"(?<![a-zA-Z0-9])"
+ re.escape(hf_variant.lower())
+ r"(?![a-zA-Z0-9])"
)
for snap in _iter_hf_cache_snapshots(hf_repo):
matches = sorted(
p.relative_to(snap).as_posix()
for p in snap.rglob("*.gguf")
if "mmproj" not in p.name.lower()
and boundary.search(p.relative_to(snap).as_posix().lower())
)
if not matches:
continue
gguf_filename = matches[0]
m = _SHARD_FULL_RE.match(Path(gguf_filename).name)
if m:
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$"
)
gguf_extra_shards = [
f
for f in matches[1:]
if sibling_pat.match(Path(f).name)
]
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 and 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 present in the HF cache so we only
# preflight against what we actually have to 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
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
# Count as 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:
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:
gguf_extra_shards = sorted(
f
for f in all_gguf_files
if f.startswith(_prefix)
and f != gguf_filename
and "mmproj" not in f.lower()
)
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 self._cancel_event.is_set():
raise RuntimeError("Cancelled")
dl_start = time.monotonic()
local_path = hf_hub_download(
repo_id = hf_repo,
filename = gguf_filename,
token = hf_token,
)
for shard in gguf_extra_shards:
if self._cancel_event.is_set():
raise RuntimeError("Cancelled")
logger.info(f"Resolving GGUF shard: {shard}")
hf_hub_download(
repo_id = hf_repo,
filename = shard,
token = hf_token,
)
except RuntimeError as e:
if "Cancelled" in str(e):
raise
raise RuntimeError(
f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}"
)
except Exception as e:
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_mmproj(
self,
*,
hf_repo: str,
hf_token: Optional[str] = None,
) -> Optional[str]:
"""Download the mmproj (vision projection) file from a GGUF repo.
Prefers mmproj-F16.gguf, falls back to any mmproj*.gguf file.
Returns the local path, or None if no mmproj file exists.
"""
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]
target: Optional[str] = None
try:
from huggingface_hub import list_repo_files
target = _pick_mmproj(list_repo_files(hf_repo, token = hf_token))
except Exception as e:
logger.debug(f"Could not list repo files for mmproj: {e}")
# Offline: resolve mmproj from the local HF cache snapshot, same
# shape as _download_gguf's offline fallback above.
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 = [
p.relative_to(snap).as_posix() for p in snap.rglob("*.gguf")
]
target = _pick_mmproj(rel_files)
if target is not None:
logger.info("Resolved mmproj %s from local HF cache", target)
break
except Exception as e:
logger.debug(f"Offline cache lookup for mmproj failed: {e}")
if target is None:
return None
try:
from huggingface_hub import hf_hub_download
logger.info(f"Downloading mmproj: {hf_repo}/{target}")
local_path = hf_hub_download(
repo_id = hf_repo,
filename = target,
token = hf_token,
)
return local_path
except Exception as e:
logger.warning(f"Could not download mmproj: {e}")
return None
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.
Returns None if mmproj_path is None, missing on disk, 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)
# ── Lifecycle ─────────────────────────────────────────────────
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,
# 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,
n_threads: Optional[int] = None,
n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused
n_parallel: int = 1,
extra_args: Optional[List[str]] = None,
) -> bool:
"""
Start llama-server with a GGUF model.
Two modes:
- Local: ``gguf_path="/path/to/model.gguf"`` → uses ``-m``
- HF: ``hf_repo="unsloth/gemma-3-4b-it-GGUF", hf_variant="Q4_K_M"`` → uses ``-hf``
In HF mode, llama-server handles downloading, caching, and
auto-loading mmproj files for vision models.
Returns True if server started and health check passed.
"""
# Serialise the whole load so concurrent /load calls never
# leave two llama-server processes alive (#5401 / #5161). Does
# not block /unload, /status, /load-progress.
with self._serial_load_lock:
# Duplicate /load that raced past the route-level check
# (the first one hadn't published _healthy=True yet). If the
# live server already satisfies this request, do nothing.
if self._already_in_target_state(
gguf_path = gguf_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,
chat_template_override = chat_template_override,
extra_args = extra_args,
is_vision = is_vision,
):
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 complete.
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 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:
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
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()
binary = self._find_llama_server_binary()
if not binary:
raise RuntimeError(
"llama-server binary not found. "
"Run setup.sh to build it, install llama.cpp, "
"or set LLAMA_SERVER_PATH environment variable."
)
# ── Phase 2: download (NO lock held, so cancel can proceed) ──
# Scope HF_HUB_OFFLINE to the download block only when DNS is
# dead; cleanup runs even on exception so a transient hiccup
# at the start of one load cannot 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
if is_vision and not mmproj_path:
mmproj_path = self._download_mmproj(
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 for DeepSeek detection
self._model_identifier = model_identifier
# Read GGUF metadata (context_length, chat_template) -- fast, header only
self._read_gguf_metadata(model_path)
# Check cancel after download
if self._cancel_event.is_set():
logger.info("Load cancelled after download phase")
return False
# Outside ``self._lock`` so /unload, /cancel, /status are
# not blocked. ``unload_model`` also records the kill, so
# the frontend /unload+/load Apply path engages the wait
# here even though no in-process kill happened.
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) based on model size + estimated KV cache.
# Seed safe defaults before GPU probing so the except path
# still 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)
cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv)
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:
logger.info(
f"User --cache-type-k/-v {cache_override} "
"honored for KV estimate"
)
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]] = []
try:
model_size = self._get_gguf_size_bytes(model_path)
gpus = self._get_gpu_free_memory()
# Resolve effective context: 0 means let llama-server use the
# model's native length. Only expand to a known native length
# if metadata is available; otherwise preserve 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 model's native context length.
# GPU/VRAM-fit logic below may shrink this if hardware is limited.
max_available_ctx = self._context_length or effective_ctx
# Will MTP engage on this load? If so, the auto-fit
# budget needs to reserve extra VRAM for the draft
# model's KV cache + compute graph. Mirrors the
# canonical-mode resolver in _build_speculative_flags:
# forced mtp / mtp+ngram always engage; auto only
# engages on an MTP GGUF >= 3B (sub-3B auto falls
# back to ngram-mod which doesn't need headroom);
# ngram / ngram-simple / off never engage MTP.
_mtp_canonical = _canonicalize_spec_mode(speculative_type)
_mtp_effective = _mtp_canonical or "auto"
_mtp_size_for_fit = _extract_model_size_b(model_identifier)
_mtp_sub_3b_for_fit = (
_mtp_size_for_fit is not None and _mtp_size_for_fit < 3.0
)
_mtp_will_engage = bool(
not _extra_args_set_spec_type(extra_args)
and (
_mtp_effective in ("mtp", "mtp+ngram")
or (
_mtp_effective == "auto"
and (
bool(self._nextn_predict_layers)
or _is_mtp_model_name(model_identifier, model_path)
)
and not _mtp_sub_3b_for_fit
)
)
)
# Auto-cap context to fit in GPU VRAM and select GPUs.
#
# Two policies depending on whether the user set n_ctx:
#
# Explicit n_ctx (user chose a context length):
# Honor it. Try the full requested context with _select_gpus
# (which uses as many GPUs as needed). Only cap if it doesn't
# fit on any GPU combination.
#
# Auto n_ctx=0 (model's native context):
# Prefer fewer GPUs with reduced context over more GPUs,
# since multi-GPU is slower and the user didn't ask for a
# specific context length.
gpu_indices, use_fit = None, True
explicit_ctx = requested_ctx > 0
if 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: g[1], reverse = True
)
best_cap = 0
for n_gpus in range(1, len(ranked_for_cap) + 1):
subset = ranked_for_cap[:n_gpus]
pool_mib = sum(free for _, free in subset)
capped = self._fit_context_to_vram(
native_ctx_for_cap,
pool_mib,
model_size,
cache_type_kv,
n_parallel = n_parallel,
mtp_engaged = _mtp_will_engage,
)
kv = self._estimate_kv_cache_bytes(
capped, cache_type_kv, n_parallel = n_parallel
)
total_mib = (model_size + kv) / (1024 * 1024)
if total_mib <= pool_mib * 0.90:
best_cap = max(best_cap, capped)
if best_cap > 0:
max_available_ctx = best_cap
else:
# Weights exceed 90% of every GPU subset's free
# memory, so there is no fitting context. Anchor
# the UI's "safe zone" threshold at 4096 (the
# spec's default when the model cannot fit) so
# the ctx slider shows the "might be slower"
# warning as soon as the user drags above the
# fallback default instead of never.
max_available_ctx = min(4096, native_ctx_for_cap)
if explicit_ctx:
# Honor the user's requested context verbatim. If it
# fits, pin GPUs and skip --fit; if it doesn't, ship
# -c <user_ctx> --fit on and let llama-server flex
# -ngl (CPU layer offload). The UI is expected to
# have surfaced the "might be slower" warning before
# the user submitted a ctx above the fit ceiling.
requested_total = (
model_size
+ self._estimate_kv_cache_bytes(
effective_ctx, cache_type_kv, n_parallel = n_parallel
)
)
gpu_indices, use_fit = self._select_gpus(
requested_total, gpus
)
# No silent shrink: effective_ctx stays == requested_ctx.
else:
# Auto context: prefer fewer GPUs, cap context
# to fit. Same headroom threshold as
# _select_gpus (#5106).
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
pin_fraction = self._GPU_PIN_VRAM_FRACTION
for n_gpus in range(1, len(ranked) + 1):
subset = ranked[:n_gpus]
pool_mib = sum(free for _, free in subset)
capped = self._fit_context_to_vram(
effective_ctx,
pool_mib,
model_size,
cache_type_kv,
n_parallel = n_parallel,
mtp_engaged = _mtp_will_engage,
)
kv = self._estimate_kv_cache_bytes(
capped, cache_type_kv, n_parallel = n_parallel
)
total_mib = (model_size + kv) / (1024 * 1024)
if total_mib <= pool_mib * pin_fraction:
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 deferring to --fit on:
# a model that overflows at 131k may pin
# comfortably with a 4096 KV cache (#5106).
effective_ctx = min(4096, effective_ctx)
if effective_ctx > 0:
for n_gpus in range(1, len(ranked) + 1):
subset = ranked[:n_gpus]
pool_mib = sum(free for _, free in subset)
kv = self._estimate_kv_cache_bytes(
effective_ctx,
cache_type_kv,
n_parallel = n_parallel,
)
total_mib = (model_size + kv) / (1024 * 1024)
if total_mib <= pool_mib * pin_fraction:
gpu_indices = sorted(
idx for idx, _ in subset
)
use_fit = False
break
elif gpus:
# Can't estimate KV -- fall back to file-size-only check.
# Without KV estimation we cannot prove a hardware cap, so
# keep the ceiling at the native context (already the default).
logger.debug(
"Falling back to file-size-only GPU selection",
model_size_gb = round(model_size / (1024**3), 2),
)
gpu_indices, use_fit = self._select_gpus(model_size, gpus)
if use_fit and not explicit_ctx:
# Weights don't fit on any subset. Default the UI to
# 4096 so the slider doesn't land on an unusable native
# context. --fit on will flex -ngl at runtime.
effective_ctx = (
min(4096, effective_ctx) if effective_ctx > 0 else 4096
)
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)"
)
kv_cache_bytes = self._estimate_kv_cache_bytes(
effective_ctx, cache_type_kv, n_parallel = n_parallel
)
logger.info(
f"GGUF size: {model_size / (1024**3):.1f} GB, "
f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, "
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
effective_ctx = requested_ctx # fall back to original
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"
)
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",
]
if use_fit:
cmd.extend(["--fit", "on"])
elif gpu_indices is not None:
# Model fits on selected GPU(s) -- offload all layers
cmd.extend(["-ngl", "-1"])
# -1 = llama.cpp auto-detect (physical cores). Pass explicitly so we
# do not inherit llama-server's internal default, which has historically
# varied (hardware concurrency incl. hyperthreads on some builds).
cmd.extend(
["--threads", str(n_threads if n_threads is not None else -1)]
)
# Always enable Jinja chat template rendering for proper template support
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:
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:
self._cache_type_kv = None
# Speculative decoding (n-gram self-speculation, zero VRAM cost)
# ngram-mod: ~16 MB shared hash pool, constant memory/complexity,
# variable draft lengths. Helps most when the model repeats
# existing text (code refactoring, summarization, reasoning).
# For general chat with low repetition, 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
#
# Params from llama.cpp server README:
# --spec-ngram-mod-n-match 24 (lookup length)
# --spec-ngram-mod-n-min 48 --spec-ngram-mod-n-max 64
# (MoEs need long drafts; dense models can reduce these)
# ref: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md
# ref: https://github.com/ggml-org/llama.cpp/pull/19164
# ref: https://github.com/ggml-org/llama.cpp/pull/18471
# draft-mtp: MTP heads on Unsloth's *-MTP GGUFs
# (llama.cpp #22673). Auto-enabled via nextn_predict_layers,
# fallback to -MTP in name. GPU: MTP-only. CPU/Mac: chain
# with ngram-mod. See unsloth.ai/docs/models/qwen3.6#mtp-guide.
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,
)
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_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",
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}"
)
# For reasoning models, set default thinking mode.
# Qwen3.5/3.6 models below 9B (0.8B, 2B, 4B) disable thinking by default.
# Only 9B and larger enable thinking.
# Always-on templates ignore the kwarg entirely, 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)
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: add --api-key for direct client access when enabled
import os as _os
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
# User-supplied pass-through args go last so llama.cpp's
# last-wins flag parsing lets the user override Studio's
# auto-set tier-2 flags (e.g. --cache-type-k, --spec-type).
# The route layer has already validated this list against
# the managed-flag denylist 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)}"
)
_log_cmd = list(cmd)
if "--api-key" in _log_cmd:
_ki = _log_cmd.index("--api-key") + 1
if _ki < len(_log_cmd):
_log_cmd[_ki] = "<redacted>"
logger.info(f"Starting llama-server: {' '.join(_log_cmd)}")
# Set library paths so llama-server can find its shared libs and CUDA DLLs
import os
import sys
env = child_env_without_native_path_secret()
binary_dir = str(Path(binary).parent)
# 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():
env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1")
logger.info(
"AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1"
)
if sys.platform == "win32":
# See _build_windows_path_dirs for ordering. #5106.
path_dirs = self._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 llama.cpp prebuilt bundles its own rocblas.dll
# but NOT the Tensile kernel library files it needs
# (rocblas/library/TensileLibrary*.dat + *.hsaco). The
# bundled DLL searches relative to its own location by
# default (i.e. <binary_dir>/rocblas/library/) which does
# not exist, causing a silent crash on the first GEMM.
# ROCBLAS_TENSILE_LIBPATH overrides that search to point at
# the ROCm installation where the kernel files actually are.
_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: set LD_LIBRARY_PATH for shared libs next to the binary
# and CUDA runtime libs (libcudart, libcublas, etc.)
import platform
lib_dirs = [binary_dir]
_arch = platform.machine() # x86_64, aarch64, etc.
# Pip-installed nvidia CUDA runtime libs (e.g. torch's
# bundled cuda-bindings). The prebuilt llama.cpp binary
# links against 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",
"cu*",
"lib",
),
os.path.join(
sys.prefix,
"lib",
"python*",
"site-packages",
"nvidia",
"cudnn",
"lib",
),
os.path.join(
sys.prefix,
"lib",
"python*",
"site-packages",
"nvidia",
"nvjitlink",
"lib",
),
]:
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 on a system where default /usr/local/cuda
# points to 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
)
# Pin to selected GPU(s). On ROCm, llama-server (and any torch
# in the subprocess) honors HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES;
# narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child seeing
# the full HIP/ROCR set the parent inherited.
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
env["ROCR_VISIBLE_DEVICES"] = pinned
except Exception as e:
logger.debug(
"Failed to set ROCm visibility env vars for child: %s", e
)
# 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. Path lives under the
# studio home so it ships in the same place all other
# Studio logs live.
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
self._process = subprocess.Popen(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
env = env,
**_windows_hidden_subprocess_kwargs(),
)
# 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()
# Store the resolved on-disk path, not the caller's kwarg. In
# HF mode the caller passes gguf_path=None and the real path
# (``model_path``) is what llama-server is actually mmap'ing.
# Downstream consumers (load_progress, log lines, etc.) need
# the path that exists on disk.
self._gguf_path = model_path
self._hf_repo = hf_repo
# For local GGUF files, extract variant from filename if not provided
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 -- it holds the model's native
# context length from GGUF metadata and is used for display/info.
self._effective_context_length = (
effective_ctx if effective_ctx > 0 else self._context_length
)
self._max_context_length = (
max_available_ctx
if max_available_ctx > 0
else self._effective_context_length
)
# Wait for llama-server to become healthy
if not self._wait_for_health(timeout = 600.0):
self._kill_process()
_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 (self._model_identifier or "").startswith("ollama/")
)
# Only show the Ollama-specific message when the server
# output indicates a GGUF compatibility issue, not for
# unrelated failures like OOM or missing binaries.
if _is_ollama:
_output = "\n".join(self._stdout_lines[-50:]).lower()
_gguf_compat_hints = (
"key not found",
"unknown model architecture",
"failed to load model",
)
if any(h in _output for h in _gguf_compat_hints):
raise RuntimeError(
"Some Ollama models do not work with llama.cpp. "
"Try a different model, or use this model directly through Ollama instead."
)
raise RuntimeError(
"llama-server failed to start. "
"Check that the GGUF file is valid and you have enough memory."
)
self._healthy = True
# Commit caller intent only after _healthy=True so a
# failed startup can't poison the next inheritance check.
# None keeps prior, [] clears, list sets. Source records
# the caller's hf_variant (None for local files) so the
# route's same_source check stays symmetric.
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)
# 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} "
f"for model '{model_identifier}'"
)
# Probe outside _lock (interruptible by /unload); init inside.
self._is_audio = False
self._audio_type = None
self._audio_probed = 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 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
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],
) -> 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 cost):
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 general chat with low repetition, 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
Sub-3B dense MTP regresses vs spec-off because 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).
Auto falls back to ngram-mod (zero-VRAM, near-zero idle cost on
diverse content); forced MTP variants engage anyway and just log
a warning per the user's choice.
"""
flags: List[str] = []
# Reset; emit branches re-set on the resolved emission.
self._spec_draft_n_max = None
self._speculative_type = None
# Canonical UI-facing requested mode: auto / mtp / ngram /
# mtp+ngram / off / ngram-simple. Legacy values are mapped via
# _canonicalize_spec_mode (default->auto, draft-mtp->mtp,
# ngram-mod->ngram, "ngram-mod,draft-mtp"->mtp+ngram).
canonical_mode = _canonicalize_spec_mode(speculative_type)
is_mtp_model = bool(self._nextn_predict_layers) or (
_is_mtp_model_name(model_identifier, model_path)
)
user_owns_spec_type = _extra_args_set_spec_type(extra_args)
_mtp_size_b = _extract_model_size_b(model_identifier)
_mtp_too_small = _mtp_size_b is not None and _mtp_size_b < 3.0
if user_owns_spec_type:
# User --spec-type in extra_args wins outright; suppress
# auto-emit so we don't emit a duplicate / conflicting
# spec block. Record requested mode as None.
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;
# raising past 3 starts to regress 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."
)
return False
draft_n_max = _resolved_draft_n_max()
n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max"
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)"
)
spec_value = mtp_token
flags.extend(
[
"--spec-type",
spec_value,
n_max_flag,
str(draft_n_max),
]
)
flags.extend(ngram_knobs)
else:
flags.extend(
[
"--spec-type",
mtp_token,
n_max_flag,
str(draft_n_max),
]
)
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
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 _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)."
)
elif not is_mtp_model:
logger.warning(
"Forcing MTP on a non-MTP GGUF; llama-server may "
"fall back to spec-off if no nextn head is present. "
"Engaging anyway (user override)."
)
_emit_mtp(chain_ngram = False)
return flags
if effective_mode == "mtp+ngram":
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)."
)
elif not is_mtp_model:
logger.warning(
"Forcing MTP+Ngram on a non-MTP GGUF; llama-server "
"may fall back to ngram-only if no nextn head is "
"present. Engaging anyway (user override)."
)
_emit_mtp(chain_ngram = True)
return flags
# effective_mode == "auto": today's promotion path. llama.cpp
# #22673: MTP is compatible with mmproj, so there's no vision gate.
if is_mtp_model and not _mtp_too_small:
# 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.
_small_caps = self.probe_server_capabilities(binary)
if _small_caps.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,
) -> 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 instead
# 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
# Compare on the canonical UI-facing mode the user requested.
# When extra_args carries --spec-type, the route-layer code paths
# bypass the dropdown anyway and the backend stores
# _requested_spec_mode = None; the request mirrors that by
# canonicalising to None.
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
# spec_draft_n_max only matters when an MTP variant is actually
# engaged. Compare on the resolved spec rather than the requested
# mode so an Auto request that auto-promoted to draft-mtp under
# the hood still bounces a reload when the user changes n_max.
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
# extra_args=None means "no opinion" (inherit semantics 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 a GPU model buffer was allocated, False if only CPU
buffers landed despite GPU intent, None when there's no signal
(no GPU detected, no buffer-size lines, etc.)."""
if not detected_gpus or not expected_gpu:
return None
# llama-server logs one ``... model buffer size = N MiB`` line
# per backend buffer; CUDA0 / ROCm0 / Metal / Vulkan0 /
# OpenCL0 / SYCL0 are GPU, CPU / CPU_Mapped are not.
gpu_markers = ("CUDA", "ROCm", "Metal", "Vulkan", "OpenCL", "SYCL")
saw_buffer_line = False
saw_gpu_buffer = False
for line in self._stdout_lines:
if "model buffer size" not in line:
continue
saw_buffer_line = True
if any(marker in line for marker in gpu_markers):
saw_gpu_buffer = True
break
if not saw_buffer_line:
return None
return saw_gpu_buffer
def unload_model(self) -> bool:
"""Terminate the llama-server 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._hf_variant = None
self._is_vision = False
self._is_audio = False
self._audio_type = None
self._audio_probed = 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_default = True
self._supports_preserve_thinking = False
self._supports_tools = False
self._cache_type_kv = None
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:
import os
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."""
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:
self._process = None
# Clear healthy so a /load arriving during the replacement
# server's warm-up window cannot 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
# /unload+/load Apply paths record the kill.
self._last_kill_monotonic = time.monotonic()
if self._stdout_thread is not None:
self._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 _kill_orphaned_servers():
"""Kill orphaned llama-server processes started by studio.
Only kills processes whose resolved binary lives under a known
Studio install directory (or matches an exact env-var override)
to avoid terminating unrelated llama-server instances.
Mirrors every location that _find_llama_server_binary() can
return from so that orphans from any supported install path
are still cleaned up.
Uses psutil for cross-platform support (Linux, macOS, Windows).
Falls back to pgrep + /proc/<pid>/exe on Linux when psutil is
not installed.
"""
import os
import signal
import sys
try:
# -- Build the ownership allowlist --------------------------------
# Two kinds of matches:
# exact_binaries -- env var overrides (exact path match only)
# install_roots -- directory trees that are Studio-owned
# (binary must be *under* one of these)
install_roots: list[Path] = []
# Env-mode custom root (mirrors _find_llama_server_binary).
_is_custom_root = False
try:
from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433
_resolved_sr = _sr()
_legacy_studio = Path.home() / ".unsloth" / "studio"
try:
_is_custom_root = _resolved_sr.resolve() != _legacy_studio.resolve()
except (OSError, ValueError):
_is_custom_root = _resolved_sr != _legacy_studio
if _is_custom_root:
install_roots.append(_resolved_sr / "llama.cpp")
except (ImportError, OSError, ValueError):
pass
# Primary install dir (default mode only). Env-mode skips this so
# a custom-root Studio cannot kill a concurrent default-install
# Studio's llama-server (same OS user, different install).
if not _is_custom_root:
install_roots.append(Path.home() / ".unsloth" / "llama.cpp")
# Legacy in-tree build dirs (older setup.sh versions)
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:
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 not installed.
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()
# Check ownership: exact binary match OR binary is
# under a known install root (proper ancestry, not
# substring).
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()
logger.info(
f"Killed orphaned llama-server process "
f"(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
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
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
# Resolve the actual executable. /proc/<pid>/exe is a
# symlink to the real binary and avoids all cmdline-
# parsing ambiguities (spaces in paths, argv rewriting).
# 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)
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)
def _cleanup(self):
"""atexit handler to ensure llama-server is terminated."""
self._kill_process()
def _wait_for_health(self, timeout: float = 120.0, interval: float = 0.5) -> bool:
"""
Poll llama-server's /health endpoint until it responds 200.
Also monitors subprocess for early exit/crash.
"""
deadline = time.monotonic() + timeout
url = f"http://127.0.0.1:{self._port}/health"
while time.monotonic() < deadline:
# Check if process crashed
if self._process.poll() is not None:
# Give the drain thread a moment to collect final output
if self._stdout_thread is not None:
self._stdout_thread.join(timeout = 2)
output = "\n".join(self._stdout_lines[-50:])
logger.error(
f"llama-server exited with code {self._process.returncode}. "
f"Output: {output[:2000]}"
)
return False
try:
resp = httpx.get(url, timeout = 2.0)
if resp.status_code == 200:
return True
except (
httpx.ConnectError,
httpx.TimeoutException,
# ReadError covers TCP RST mid-read while llama-server is
# still binding the port (Windows: WinError 10054). The
# crash-detection branch above catches a real exit; this
# one keeps a transient socket close from masking it.
httpx.ReadError,
httpx.RemoteProtocolError,
httpx.WriteError,
):
pass
time.sleep(interval)
logger.error(f"llama-server health check timed out after {timeout}s")
return False
# ── Message building (OpenAI format) ──────────────────────────
@staticmethod
def _parse_tool_calls_from_text(content: str) -> list[dict]:
"""Thin wrapper around the shared parser in tool_call_parser
so safetensors and llama_cpp pick up the same fixes."""
return _shared_parse_tool_calls_from_text(content)
@staticmethod
def _build_openai_messages(
messages: list[dict],
image_b64: Optional[str] = None,
) -> list[dict]:
"""
Build OpenAI-format messages, optionally injecting an image_url
content part into the last user message for vision models.
If no image is provided, returns messages as-is.
"""
if not image_b64:
return messages
# Find the last user message and convert 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) ────────────────────────
@staticmethod
def _iter_text_cancellable(
response: "httpx.Response",
cancel_event: Optional[threading.Event] = None,
) -> Generator[str, None, None]:
"""Iterate over an httpx streaming response with cancel support.
Checks cancel_event between chunks and on ReadTimeout. The
cancel watcher in _stream_with_retry also calls response.close()
on cancel, which unblocks iter_text() once the response exists.
During normal streaming llama-server sends tokens frequently,
so the cancel check between chunks is the primary mechanism.
"""
text_iter = response.iter_text()
while True:
if cancel_event is not None and cancel_event.is_set():
response.close()
return
try:
chunk = next(text_iter)
yield chunk
except StopIteration:
return
except httpx.ReadTimeout:
# No data within the timeout window -- just loop back
# and re-check cancel_event.
continue
@staticmethod
@contextlib.contextmanager
def _stream_with_retry(
client: "httpx.Client",
url: str,
payload: dict,
cancel_event: Optional[threading.Event] = None,
headers: Optional[dict] = None,
):
"""Open an httpx streaming POST with cancel support.
Sends the request once with a long read timeout (120 s) so
prompt processing (prefill) can finish without triggering a
retry storm. The previous 0.5 s timeout caused duplicate POST
requests every half second, forcing llama-server to restart
processing each time.
A background watcher thread provides cancel by closing the
response when cancel_event is set. Limitation: httpx does not
allow interrupting a blocked read from another thread before
the response object exists, so cancel during the initial
header wait (prefill phase) only takes effect once headers
arrive. After that, response.close() unblocks reads promptly.
In practice llama-server prefill is 1-5 s for typical prompts,
during which cancel is deferred -- still much better than the
old retry storm which made prefill slower.
"""
if cancel_event is not None and cancel_event.is_set():
raise GeneratorExit
# Background watcher: close the response if cancel is requested.
# Only effective after response headers arrive (httpx limitation).
_cancel_closed = threading.Event()
_response_ref: list = [None]
def _cancel_watcher():
while not _cancel_closed.is_set():
if cancel_event.wait(timeout = 0.3):
# Cancel requested. Keep polling until the response object
# exists so we can close it, or until the main thread
# finishes on its own (_cancel_closed is set in finally).
while not _cancel_closed.is_set():
r = _response_ref[0]
if r is not None:
try:
r.close()
return
except Exception as e:
logger.debug(
f"Error closing response in cancel watcher: {e}"
)
# Response not created yet -- wait briefly and retry
_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:
# Long read timeout so prefill (prompt processing) can finish
# without triggering a retry storm. Cancel during both
# prefill and streaming is handled by the watcher thread
# which closes the response, unblocking any httpx read.
prefill_timeout = httpx.Timeout(
connect = 30,
read = 120.0,
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.ReadError, httpx.RemoteProtocolError, httpx.CloseError):
# 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 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,
) -> Generator[str | dict, None, None]:
"""
Send a chat completion request to llama-server and stream tokens back.
Uses /v1/chat/completions — llama-server handles chat template
application and 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,
}
# Pass enable_thinking / reasoning_effort / preserve_thinking per-request
_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's effective context length when known,
# otherwise the conservative floor. The wall-clock backstop below
# keeps a stuck model from running indefinitely either way.
payload["max_tokens"] = (
max_tokens
if max_tokens is not None
else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR)
)
payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS
if stop:
payload["stop"] = stop
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
try:
# _stream_with_retry uses a 120 s read timeout so prefill
# can finish. Cancel during streaming is handled by the
# watcher thread (closes the response on cancel_event).
stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
_auth_headers = (
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
)
with httpx.Client(
timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0)
) as client:
with self._stream_with_retry(
client,
url,
payload,
cancel_event,
headers = _auth_headers,
) as response:
if response.status_code != 200:
error_body = response.read().decode()
raise RuntimeError(
f"llama-server returned {response.status_code}: {error_body}"
)
buffer = ""
has_content_tokens = False
reasoning_text = ""
for raw_chunk in self._iter_text_cancellable(
response, cancel_event
):
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 tokens:
# the model put its entire reply in reasoning
# (e.g. Qwen3 always-think mode). Show it
# as the main response, not as 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:])
# 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", {})
# Handle 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:
_metadata_usage = _backfill_usage_from_timings(
_metadata_usage, _metadata_timings
)
yield {
"type": "metadata",
"usage": _metadata_usage,
"timings": _metadata_timings,
}
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
# ── 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,
) -> Generator[dict, None, None]:
"""
Agentic loop: let the model call tools, execute them, and continue.
Yields dicts with:
{"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 execute_tool
if not self.is_loaded:
raise RuntimeError("llama-server is not loaded")
conversation = list(messages)
url = f"{self.base_url}/v1/chat/completions"
_accumulated_completion_tokens = 0
_accumulated_predicted_ms = 0.0
_accumulated_predicted_n = 0
def _strip_tool_markup(text: str, *, final: bool = False) -> str:
if not auto_heal_tool_calls:
return text
return strip_tool_call_markup(text, final = final)
# XML prefixes that signal a tool call in content.
# Empty when auto_heal is disabled so the buffer never
# speculatively holds content for XML detection.
_TOOL_XML_SIGNALS = (
("<tool_call>", "<function=") if auto_heal_tool_calls else ()
)
_MAX_BUFFER_CHARS = 32
# ── Duplicate tool-call detection ────────────────────────
# Track recent (tool_name, arguments) hashes to detect loops
# where the model repeats the exact same call. Retries after
# a transient failure are allowed (only block when the previous
# identical call succeeded).
_tool_call_history: list[tuple[str, bool]] = [] # (key, failed)
# ── Re-prompt on plan-without-action ─────────────────
# When the model describes what it intends to do (forward-looking
# language) without actually calling a tool, re-prompt once.
# Only triggers on responses that signal intent/planning -- a
# direct answer like "4" or "Hello!" will not match.
# Pattern is compiled once at module level (_INTENT_SIGNAL).
_reprompt_count = 0
# Reserve extra iterations for re-prompts so they don't
# consume the caller's tool-call budget. Only add the
# extra slot when tool iterations are actually 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
# 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": 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)
)
payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS
if stop:
payload["stop"] = stop
try:
_auth_headers = (
{"Authorization": f"Bearer {self._api_key}"}
if self._api_key
else None
)
# ── Speculative buffer state machine ──────────────────
# BUFFERING: accumulating content, checking for tool signals
# STREAMING: no tool detected, yielding tokens to caller
# DRAINING: tool signal found, silently consuming 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 = ""
cumulative_display = "" # Cumulative text yielded (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
_stream_done = False
_last_emitted = ""
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),
) as client:
with self._stream_with_retry(
client,
url,
payload,
cancel_event,
headers = _auth_headers,
) as response:
if response.status_code != 200:
error_body = response.read().decode()
raise RuntimeError(
f"llama-server returned {response.status_code}: "
f"{error_body}"
)
raw_buf = ""
for raw_chunk in self._iter_text_cancellable(
response,
cancel_event,
):
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>"
yield {
"type": "content",
"text": _strip_tool_markup(
cumulative_display,
final = True,
),
}
else:
cumulative_display = reasoning_accum
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", {})
# ── Structured tool_calls ──
tc_deltas = delta.get("tool_calls")
if tc_deltas:
# Once visible content has been
# emitted, do not reclassify this
# turn as a tool call.
if _last_emitted:
continue
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 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"]
continue
# ── Reasoning tokens ──
# Only yield in STREAMING state. In BUFFERING
# and DRAINING, accumulate silently so we don't
# corrupt the consumer's prev_text tracker
# (routes/inference.py never resets prev_text
# between tool iterations).
reasoning = delta.get("reasoning_content", "")
if reasoning:
reasoning_accum += reasoning
if detect_state == _S_STREAMING:
if not in_thinking:
cumulative_display += "<think>"
in_thinking = True
cumulative_display += reasoning
yield {
"type": "content",
"text": cumulative_display,
}
# ── Content tokens ──
token = delta.get("content", "")
if token:
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(
cumulative_display,
)
if len(cleaned) > len(_last_emitted):
_last_emitted = cleaned
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
if is_match:
detect_state = _S_DRAINING
elif (
is_prefix
and len(stripped_buf)
< _MAX_BUFFER_CHARS
):
pass # keep buffering
else:
# Not a tool -- flush buffer
detect_state = _S_STREAMING
# Flush any reasoning accumulated
# during BUFFERING phase
if reasoning_accum:
cumulative_display += "<think>"
cumulative_display += (
reasoning_accum
)
cumulative_display += "</think>"
cumulative_display += content_buffer
cleaned = _strip_tool_markup(
cumulative_display,
)
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
# ── Resolve BUFFERING at stream end ──
if detect_state == _S_BUFFERING:
stripped_buf = content_buffer.lstrip()
if (
stripped_buf
and auto_heal_tool_calls
and any(s in stripped_buf for s in _TOOL_XML_SIGNALS)
):
detect_state = _S_DRAINING
elif content_accum or reasoning_accum:
detect_state = _S_STREAMING
if content_buffer:
# Flush any reasoning accumulated first
if reasoning_accum:
cumulative_display += "<think>"
cumulative_display += reasoning_accum
cumulative_display += "</think>"
cumulative_display += content_buffer
yield {
"type": "content",
"text": _strip_tool_markup(
cumulative_display,
final = True,
),
}
elif reasoning_accum and not has_content_tokens:
# Reasoning-only response (no content tokens):
# show reasoning as plain text, matching
# the final streaming pass behavior for
# models that put everything in reasoning.
cumulative_display = reasoning_accum
yield {
"type": "content",
"text": cumulative_display,
}
else:
return
# ── STREAMING path: no tool call ──
if detect_state == _S_STREAMING:
# Safety net: check for XML tool signals in content.
# The route layer resets prev_text on tool_start, so
# post-tool synthesis streams correctly even if
# content was already emitted before the tool XML.
_safety_tc = None
if auto_heal_tool_calls and any(
s in content_accum for s in _TOOL_XML_SIGNALS
):
_safety_tc = self._parse_tool_calls_from_text(
content_accum,
)
if not _safety_tc:
# ── Re-prompt on plan-without-action ──
# If the model described what it intends to do
# (forward-looking language) without calling any
# tool, nudge it to act. Only fires once per
# request and only on short responses that
# contain intent signals -- a direct answer
# like "4" or "Hello!" won't trigger this.
# Use content if available, otherwise fall back
# to reasoning text (reasoning-only stalls).
_stripped = content_accum.strip()
if not _stripped:
_stripped = reasoning_accum.strip()
if (
tools
and _reprompt_count < _MAX_REPROMPTS
and 0 < len(_stripped) < _REPROMPT_MAX_CHARS
and _INTENT_SIGNAL.search(_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,
}
)
conversation.append(
{
"role": "user",
"content": (
"STOP. Do NOT write code or explain. "
"You MUST call a tool NOW. "
"Call web_search or python immediately."
),
}
)
# 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
# Content was already streamed. Yield metadata.
yield {"type": "status", "text": ""}
_fu = (
_backfill_usage_from_timings(_iter_usage, _iter_timings)
or {}
)
_fc = _fu.get("completion_tokens", 0)
_fp = _fu.get("prompt_tokens", 0)
_tc = _fc + _accumulated_completion_tokens
if (
_iter_usage
or _iter_timings
or _accumulated_completion_tokens
):
_mt = dict(_iter_timings) if _iter_timings else {}
if _accumulated_predicted_ms or _accumulated_predicted_n:
_mt["predicted_ms"] = (
_mt.get("predicted_ms", 0)
+ _accumulated_predicted_ms
)
_tn = (
_mt.get("predicted_n", 0) + _accumulated_predicted_n
)
_mt["predicted_n"] = _tn
_tms = _mt["predicted_ms"]
if _tms > 0:
_mt["predicted_per_second"] = _tn / (_tms / 1000.0)
yield {
"type": "metadata",
"usage": {
"prompt_tokens": _fp,
"completion_tokens": _tc,
"total_tokens": _fp + _tc,
},
"timings": _mt,
}
return
# Safety net caught tool XML -- treat as tool call
tool_calls = _safety_tc
content_text = _strip_tool_markup(
content_accum,
final = True,
)
logger.info(
f"Safety net: parsed {len(tool_calls)} tool call(s) "
f"from streamed content"
)
else:
# ── DRAINING path: assemble tool_calls ──
tool_calls = None
content_text = content_accum
if has_structured_tc:
# Filter out incomplete fragments (e.g. from
# truncation by max_tokens 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
and auto_heal_tool_calls
and any(s in content_accum for s in _TOOL_XML_SIGNALS)
):
tool_calls = self._parse_tool_calls_from_text(
content_accum,
)
if tool_calls and not has_structured_tc:
content_text = _strip_tool_markup(
content_text,
final = True,
)
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 are not 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
)
if content_accum:
yield {"type": "content", "text": content_accum}
_fu = (
_backfill_usage_from_timings(_iter_usage, _iter_timings)
or {}
)
_fc = _fu.get("completion_tokens", 0)
_fp = _fu.get("prompt_tokens", 0)
_tc = _fc + _accumulated_completion_tokens
if (
_iter_usage
or _iter_timings
or _accumulated_completion_tokens
):
_mt = dict(_iter_timings) if _iter_timings else {}
if _accumulated_predicted_ms or _accumulated_predicted_n:
_mt["predicted_ms"] = (
_mt.get("predicted_ms", 0)
+ _accumulated_predicted_ms
)
_tn = (
_mt.get("predicted_n", 0) + _accumulated_predicted_n
)
_mt["predicted_n"] = _tn
_tms = _mt["predicted_ms"]
if _tms > 0:
_mt["predicted_per_second"] = _tn / (_tms / 1000.0)
yield {
"type": "metadata",
"usage": {
"prompt_tokens": _fp,
"completion_tokens": _tc,
"total_tokens": _fp + _tc,
},
"timings": _mt,
}
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)
assistant_msg = {"role": "assistant", "content": content_text}
if tool_calls:
assistant_msg["tool_calls"] = tool_calls
conversation.append(assistant_msg)
for tc in tool_calls or []:
func = tc.get("function", {})
tool_name = func.get("name", "")
raw_args = func.get("arguments", {})
if isinstance(raw_args, str):
try:
arguments = json.loads(raw_args)
except (json.JSONDecodeError, ValueError):
if auto_heal_tool_calls:
arguments = {"query": raw_args}
else:
arguments = {"raw": raw_args}
else:
arguments = raw_args
if tool_name == "web_search":
_ws_url = (arguments.get("url") or "").strip()
if _ws_url:
_parsed = urlparse(_ws_url)
if _parsed.scheme in ("http", "https") and _parsed.hostname:
_ws_host = _parsed.hostname
if _ws_host.startswith("www."):
_ws_host = _ws_host[4:]
status_text = f"Reading: {_ws_host}"
else:
status_text = "Reading page..."
else:
status_text = f"Searching: {arguments.get('query', '')}"
elif tool_name == "python":
preview = (
(arguments.get("code") or "").strip().split("\n")[0][:60]
)
status_text = (
f"Running Python: {preview}"
if preview
else "Running Python..."
)
elif tool_name == "terminal":
cmd_preview = (arguments.get("command") or "")[:60]
status_text = (
f"Running: {cmd_preview}"
if cmd_preview
else "Running command..."
)
else:
status_text = f"Calling: {tool_name}"
yield {"type": "status", "text": status_text}
yield {
"type": "tool_start",
"tool_name": tool_name,
"tool_call_id": tc.get("id", ""),
"arguments": arguments,
}
# ── Duplicate call detection ──────────────
# str(dict) is stable here: arguments always comes from
# json.loads on the same model output within one request,
# so insertion order is deterministic (Python 3.7+).
_tc_key = tool_name + str(arguments)
_prev = _tool_call_history[-1] if _tool_call_history else None
if _prev and _prev[0] == _tc_key and not _prev[1]:
result = (
"You already made this exact call. "
"Do not repeat the same tool call. "
"Try a different approach: fetch a URL "
"from previous results, use Python to "
"process data you already have, or "
"provide your final answer now."
)
else:
_effective_timeout = (
None if tool_call_timeout >= 9999 else tool_call_timeout
)
# Guard against the model emitting a tool not in the
# per-request advertised set: filtered MCP names, a
# built-in the caller opted out of, or a stale name
# from a prior turn. Mirrors the safetensors loop's
# allowed_tool_names check.
_allowed = {
(t.get("function") or {}).get("name")
for t in (tools or [])
if (t.get("function") or {}).get("name")
}
if _allowed and tool_name not in _allowed:
result = (
f"Error: tool '{tool_name}' is not enabled "
"for this request. Use one of the enabled "
"tools or provide a final answer."
)
else:
result = execute_tool(
tool_name,
arguments,
cancel_event = cancel_event,
timeout = _effective_timeout,
session_id = session_id,
)
yield {
"type": "tool_end",
"tool_name": tool_name,
"tool_call_id": tc.get("id", ""),
"result": result,
}
# Nudge model to try a different approach on errors
_error_prefixes = (
"Error",
"Search failed",
"Execution error",
"Blocked:",
"Exit code",
"Failed to fetch",
"Failed to resolve",
"No query provided",
)
_is_error = isinstance(result, str) and result.lstrip().startswith(
_error_prefixes
)
_tool_call_history.append((_tc_key, _is_error))
# Strip image sentinel before feeding result to the LLM
# (the full result with sentinel is still yielded via
# tool_end so the frontend can extract image paths).
_result_content = result
if "\n__IMAGES__:" in _result_content:
_result_content = _result_content.rsplit("\n__IMAGES__:", 1)[0]
if _is_error:
_result_content = (
_result_content + "\n\nThe tool call encountered an issue. "
"Please try a different approach or rephrase your request."
)
tool_msg = {
"role": "tool",
"name": tool_name,
"content": _result_content,
}
tool_call_id = tc.get("id")
if tool_call_id:
tool_msg["tool_call_id"] = tool_call_id
conversation.append(tool_msg)
# Clear tool status badge before next generation iteration
yield {"type": "status", "text": ""}
# Continue the loop to let model respond with context
continue
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
# ── Tool iteration cap reached -- synthesize final answer ──
# The model used all iterations without producing a final text
# response. Inject a nudge so the final streaming pass produces
# a useful answer instead of continuing to request tools.
if max_tool_iterations > 0:
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)
)
stream_payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS
if stop:
stream_payload["stop"] = stop
stream_payload["stream_options"] = {"include_usage": True}
cumulative = ""
_last_emitted = ""
in_thinking = False
has_content_tokens = False
reasoning_text = ""
_metadata_usage = None
_metadata_timings = None
_stream_done = False
try:
stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
_auth_headers = (
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
)
with httpx.Client(
timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0)
) as client:
with self._stream_with_retry(
client,
url,
stream_payload,
cancel_event,
headers = _auth_headers,
) as response:
if response.status_code != 200:
error_body = response.read().decode()
raise RuntimeError(
f"llama-server returned {response.status_code}: {error_body}"
)
buffer = ""
for raw_chunk in self._iter_text_cancellable(
response, cancel_event
):
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:
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", {})
reasoning = delta.get("reasoning_content", "")
if reasoning:
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:
has_content_tokens = True
if in_thinking:
cumulative += "</think>"
in_thinking = False
cumulative += token
cleaned = _strip_tool_markup(cumulative)
# Only emit 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
_final_usage = _metadata_usage or {}
_final_completion = _final_usage.get("completion_tokens", 0)
_final_prompt = _final_usage.get("prompt_tokens", 0)
_total_completion = (
_final_completion + _accumulated_completion_tokens
)
if _metadata_usage or _metadata_timings:
_merged_timings = (
dict(_metadata_timings) if _metadata_timings else {}
)
if _accumulated_predicted_ms or _accumulated_predicted_n:
_merged_timings["predicted_ms"] = (
_merged_timings.get("predicted_ms", 0)
+ _accumulated_predicted_ms
)
_total_predicted_n = (
_merged_timings.get("predicted_n", 0)
+ _accumulated_predicted_n
)
_merged_timings["predicted_n"] = _total_predicted_n
_total_predicted_ms = _merged_timings["predicted_ms"]
if _total_predicted_ms > 0:
_merged_timings["predicted_per_second"] = (
_total_predicted_n / (_total_predicted_ms / 1000.0)
)
yield {
"type": "metadata",
"usage": {
"prompt_tokens": _final_prompt,
"completion_tokens": _total_completion,
"total_tokens": _final_prompt + _total_completion,
},
"timings": _merged_timings,
}
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
# ── TTS support ────────────────────────────────────────────
def detect_audio_type(self) -> Optional[str]:
"""Detect audio/TTS codec; swallows errors (use _strict variant to distinguish)."""
try:
return self._detect_audio_type_strict()
except Exception as e:
logger.debug(f"Audio type detection failed: {e}")
return None
def _detect_audio_type_strict(self) -> Optional[str]:
"""Codec name on match, None on definitive non-audio, raises on transport/JSON errors."""
if not self.is_loaded:
return None
_auth_headers = (
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
)
with httpx.Client(timeout = 10, headers = _auth_headers) 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", [])
# Check codec-specific tokens (not generic ones that may exist in non-audio models)
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"
if len(_tok("<audio_soft_token>")) == 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 prompts in 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 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 decoding.
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
_auth_headers = (
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
)
with httpx.Client(
timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers
) 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", "")
)