unsloth/studio/backend/core/inference
Daniel Han 2a05426adb
Auto-install SSM kernels (causal-conv1d, mamba-ssm) for inference loads (#6535)
* Auto-install SSM kernels (causal-conv1d, mamba-ssm) for inference loads

Mamba/SSM hybrids (Nemotron-H/Nano, Falcon-H1, Granite-4.0-H, ...) lazily import
mamba_ssm / causal_conv1d during from_pretrained, so loading them for chat failed
with 'mamba-ssm is required by the Mamba model but cannot be imported'. The training
worker already wheel-first installs these before a fine-tune; the inference worker
did not. Add utils/ssm_runtime.ensure_ssm_runtime and call it from the inference load
path so the same models load for inference. Training worker is untouched; a drift
test keeps the shared detection and pinned versions in lockstep.

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

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

* ssm_runtime: invalidate import caches, skip MLX, cover LoRA base

- Invalidate importlib finder caches in _is_importable and after a successful
  wheel install, so a kernel installed earlier in this same process is actually
  importable when the modeling code lazy-imports it during from_pretrained.
- Skip the SSM kernel install entirely on the MLX (Apple Silicon) load path:
  these are CUDA/ROCm Torch kernels with no MLX use and no macOS prebuilt wheel,
  so the source build would fail before the MLX backend loads the model.
- For LoRA loads, also run detection over the resolved base model, since an
  adapter id like 'me/my-lora' won't match the SSM heuristics but its SSM base
  (Nemotron-H, ...) is what needs the kernels.

Adds tests for cache invalidation and the MLX-skip / LoRA-base worker wiring.

* Tighten SSM autoinstall comments

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

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

* ssm_runtime: verify wheel imports, HIP-aware source build, build heartbeat

Address review feedback:
- Verify a prebuilt wheel actually imports before trusting it; a CUDA/ABI-mismatched
  wheel now falls back to a source build instead of returning success and failing later
  with the cryptic lazy-import error.
- HIP-aware source build: require hipcc on ROCm, inject clang --gcc-install-dir, and use
  the 1800s timeout, mirroring the training worker (ROCm has no prebuilt wheel).
- Emit a status heartbeat every 60s during the source build so a long (ROCm) build does
  not trip the orchestrator's 300s inactivity timeout.

Tests cover the wheel-not-importable fallback and the missing-hipcc ROCm bail.

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

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

* Make causal-conv1d best-effort and harden the SSM source build

- causal-conv1d is a fast path: models that merely want it (Qwen3-Next, LFM2)
  fall back to torch, so a failed install must not reject an otherwise loadable
  chat model on Windows/CPU/macOS or an ABI without a wheel. Only a true SSM
  model's mamba-ssm requirement stays fatal, matching the training worker which
  treats causal-conv1d as best-effort.
- The source build is reached only when not importable, including a wheel that
  installed but failed to import; add --reinstall/--force-reinstall so it
  replaces the broken install instead of no-opping as already satisfied.
- Add --no-cache to the ROCm uv source build to avoid reusing stale artifacts
  from a partial HIP build, mirroring the training worker.

* Address review: install SSM kernels before transformers, harden import + Windows

Codex:
- Install the SSM kernels before importing transformers. run_inference_process
  imported core.inference.inference (which imports unsloth/transformers) before the
  load, and a sidecar transformers can evaluate its optional-backend gates against
  the import state; installing causal_conv1d/mamba_ssm afterwards left those gates
  unsatisfied and a Nemotron/Falcon/Granite load still failed with "mamba-ssm is
  required". The initial model's kernels are now installed in run_inference_process
  before the ML import, via a shared _ensure_ssm_kernels helper; _handle_load keeps
  calling it (idempotent) for a LoRA's base and for later in-process loads.
- _is_importable now treats any import failure as "not importable", not only
  ImportError. An ABI-incompatible native kernel (undefined symbol after a torch/CUDA
  upgrade) raises OSError/RuntimeError; letting those escape reported
  ssm_runtime_install_failed instead of falling back to reinstall/source build.
- Skip causal-conv1d on Windows (no prebuilt wheel), mirroring the training worker.
  A causal-conv1d-only model (Qwen3-Next/LFM2) no longer drops a chat load into a
  multi-minute untimed source build; it uses the torch fallback. mamba-ssm is still
  attempted for true SSM hybrids.

Tests: test_ssm_runtime.py +5 (broken-kernel exceptions read as not-importable;
causal-conv1d skipped on win32 while mamba-ssm still installs). 36 passed.

* Trim comments to be more succinct

* Run security gates before installing SSM kernels

The SSM kernel auto-install is name-based (model_is_ssm is a substring match, no
config fetch), so a model id merely containing an SSM substring triggered a
native-package install (possibly a slow source build) before the malware and
remote-code consent gates ran. Extract those gates into _run_security_gates and
call it before the kernel install in both the pre-import path of
run_inference_process and in _handle_load, so a blocked or nonexistent model is
refused before any build. The gates are metadata-only and do not import
transformers, so they are safe to run before the pre-import install.

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

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

* Resolve remote LoRA bases before importing transformers

_resolve_base_model only reads a local adapter_config.json, so a remote LoRA
adapter whose own id has no SSM substring but whose base is a Nemotron/Falcon/
Granite model had its base discovered only by ModelConfig in _handle_load, after
transformers was imported and its optional-backend availability snapshotted, so
the SSM kernel install there was too late. Add _remote_lora_base, a metadata-only
adapter_config.json fetch (no huggingface_hub / transformers import), and use it
in the pre-import path so the base is gated and its kernels pre-installed.

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

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

* Gate only loaded roots, tier on the resolved base, read offline LoRA cache

Three follow-ups to the pre-import resolution:

- The security gate reused the SSM target list, which for a local full fine-tune
  includes the config.json-recorded base. That base is never loaded, so scanning
  it could falsely block a safe local checkpoint. Gate only the model plus a
  genuine LoRA base (matching _handle_load's mc.is_lora), separate from the
  broader SSM-install list.

- Tier activation ran on the raw adapter id, so a remote LoRA whose base needs a
  sidecar transformers version imported the default and failed. Resolve the base
  once up front and activate on it.

- _remote_lora_base bailed on offline before checking the hub cache, missing a
  cached adapter's base. Read the cached adapter_config.json when offline or when
  the fetch fails.

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

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

* Keep the pre-import gate transformers-free; harden remote LoRA resolution

The pre-import security gate called security_load_subdirs, which imports
model_config and thus transformers, snapshotting optional-backend availability
before the SSM kernels are installed and defeating the ordering. Add
compute_subdirs to _run_security_gates and pass False in the preflight so it scans
from the root only (transformers-free); _handle_load still runs the authoritative
gate with full subdir scoping after the import.

_remote_lora_base now skips existing local relative paths (is_local_path) so a
checkpoint like outputs/run1 is never treated as a Hub repo, and distinguishes a
definitive 404 (not a LoRA -> None) from transient/offline failures (read the
cache), so a repo that is now a full model no longer resolves a stale cached base.

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

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

* Probe a real model id for SSM kernels; respect HF_ENDPOINT

model_is_ssm is a substring match, so an arbitrary name could false-match and
force a mamba-ssm install that fails the load for a non-SSM model:
- a LoRA adapter id like user/falcon-h1-lora (the SSM-relevant code is the base's);
- a local checkpoint under an SSM-named parent dir, e.g. /runs/falcon-h1/llama-ckpt.

Add ssm_probe_identifier, which resolves the base (or a bare local checkpoint's
basename) and feed that to ensure_ssm_runtime from both the pre-import path and
_handle_load, so detection runs against a real model id, never an adapter id or
parent folders.

_remote_lora_base now honors HF_ENDPOINT so enterprise/mirror deployments resolve
the adapter base instead of always hitting huggingface.co.

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

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

* Tighten comments in the pre-import SSM gate/install path

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <michaelhan2050@gmail.com>
2026-06-22 04:48:29 -07:00
..
__init__.py Studio: make code comments and docstrings more succinct (#6029) 2026-06-08 23:07:28 -07:00
_html_to_md.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
anthropic_compat.py Studio: improve OpenAI- and Anthropic-compatible API spec compliance (#6010) 2026-06-09 17:13:25 +02:00
api_monitor.py Studio: trim serving-log noise and surface llama-server engine stats (#6377) 2026-06-17 05:37:57 -07:00
audio_codecs.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
chat_template_helpers.py Studio: make code comments and docstrings more succinct (#6029) 2026-06-08 23:07:28 -07:00
chat_templates.py Studio: bundle Gemma 4 chat templates (E2B/E4B + larger) and auto-apply to unsloth/gemma-4-*-GGUF (#6245) 2026-06-12 05:49:39 -07:00
defaults.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
external_provider.py Studio: ignore unsupported env proxy during Studio startup (#6102) 2026-06-11 05:13:27 -07:00
inference.py Expose runtime context length for hub models (#6154) 2026-06-11 22:13:53 +03:00
key_exchange.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
llama_cpp.py Studio: simplify the inference backend (#6490) 2026-06-21 20:01:09 -03:00
llama_http.py Studio: serialize non-streaming responses once and pool the proxy client (#6393) 2026-06-17 22:38:02 -07:00
llama_server_args.py studio: deterministic VRAM auto-fit for GGUF (MTP reserve, compute buffer, total-based budget) (#6312) 2026-06-17 03:10:22 -07:00
llama_stats.py Studio: trim serving-log noise and surface llama-server engine stats (#6377) 2026-06-17 05:37:57 -07:00
mcp_client.py Studio: enable stdio MCP servers on a loopback bind (#6295) 2026-06-15 03:02:32 +01:00
mcp_config_import.py studio: show MCP "Import config" on the add-server form (#6030) 2026-06-11 16:17:22 +01:00
mlx_inference.py Expose runtime context length for hub models (#6154) 2026-06-11 22:13:53 +03:00
orchestrator.py Studio: simplify the inference orchestrator and worker (#6439) 2026-06-18 17:23:16 -03:00
pricing.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
providers.py Studio: Add custom provider option to Connections (#6112) 2026-06-12 13:09:35 +02:00
runtime_context.py Expose runtime context length for hub models (#6154) 2026-06-11 22:13:53 +03:00
safetensors_agentic.py Studio: Bypass Permissions (skip confirmation, disable tool sandbox) (#5895) 2026-06-15 04:04:22 -07:00
tensor_fallback.py studio: deterministic VRAM auto-fit for GGUF (MTP reserve, compute buffer, total-based budget) (#6312) 2026-06-17 03:10:22 -07:00
tool_call_parser.py Studio: clean-room compact RAG (knowledge bases, hybrid search, fast indexing) (#5910) 2026-06-09 21:17:04 -07:00
tool_loop_controller.py Studio: clean-room compact RAG (knowledge bases, hybrid search, fast indexing) (#5910) 2026-06-09 21:17:04 -07:00
tools.py Rename chat artifacts copy to canvas (#6298) 2026-06-15 13:38:39 +02:00
worker.py Auto-install SSM kernels (causal-conv1d, mamba-ssm) for inference loads (#6535) 2026-06-22 04:48:29 -07:00