ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix(studio): reject Vulkan diffusion gpu_ids before Phase 1 teardown (#7415)

Browse source 
* fix(studio): reject Vulkan diffusion gpu_ids before Phase 1 teardown

Classify local GGUF paths (and cached HF downloads when available) for
diffusion before _kill_process() so unsupported gpu_ids requests return
400 without tearing down the active model. Fixes #7205.

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

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

* fix(studio): always pre-download HF GGUF before Vulkan diffusion preflight

Reverts the cached-path shortcut so partial split caches still run
_download_gguf before Phase 1 teardown. Header-only classification from
resolve_local_gguf_path() does not prove the variant is complete.

* Fix inaccurate shared-constant comment and cover the local pre-teardown branch for PR #7415

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

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

* Add a regression test for the pre-teardown GGUF download for PR #7415

* Tighten the Vulkan diffusion preflight comments for PR #7415

* Trim the Vulkan diffusion preflight comments for PR #7415

---------

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: danielhanchen <unslothshared@gmail.com>
This commit is contained in:
Souravrajvi0 2026-07-27 11:38:31 +05:30 • committed by GitHub
commit 9eaf5c29a5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 199 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

55
studio/backend/core/inference/llama_cpp.py
View file

@ -127,6 +127,15 @@ LLAMA_SERVER_NOT_FOUND_DETAIL = (
"then try again. (Advanced: set LLAMA_SERVER_PATH to an existing binary.)"
)
# Shared by the route, pre-teardown and post-metadata rejections (#7205).
_VULKAN_DIFFUSION_GPU_IDS_ERROR = (
"GPU selection (gpu_ids) is not supported for a DiffusionGemma "
"GGUF on a Vulkan llama.cpp build: the diffusion runner selects "
"its device by CUDA physical index, which has no defined mapping "
"to ggml Vulkan device ordinals. Omit gpu_ids to use the default "
"device."
)
# llama-server can serve HTTP 200 while running a model entirely on CPU when a
# GPU backend fails to init (#5807 / #5106 / #5830). Classify the startup log so
@ -4710,6 +4719,13 @@ class LlamaCppBackend:
probe._read_gguf_metadata(gguf_path)
return probe._is_diffusion
def _reject_vulkan_diffusion_gpu_ids_before_teardown(
self, gguf_path: str, model_identifier: str
) -> None:
"""Reject Vulkan + gpu_ids for diffusion GGUFs before Phase 1 teardown."""
if self._gguf_path_is_diffusion(gguf_path, model_identifier):
raise ValueError(_VULKAN_DIFFUSION_GPU_IDS_ERROR)
def _read_gguf_metadata(self, gguf_path: str) -> None:
"""Read context_length, architecture params, and chat_template from a GGUF header.
@ -6515,12 +6531,7 @@ class LlamaCppBackend:
f"present. Available Vulkan devices: {sorted(_pf_probed)}."
)
# A remote uncached GGUF may only reveal that it needs the
# single-device diffusion runner after download. On Vulkan, an
# explicit gpu_ids request cannot be mapped from ggml ordinals to
# that runner's CUDA physical index. Download and classify the main
# file before killing the healthy server so this late rejection is
# non-destructive. The Phase 2 call below reuses this cached path.
# Classify before killing the healthy server (#7205); Phase 2 reuses this path.
_preflight_model_path = None
if is_vulkan_backend and gpu_ids and hf_repo:
_resolved_repo = _resolve_repo_id_casing(hf_repo)
@ -6537,14 +6548,17 @@ class LlamaCppBackend:
hf_variant = hf_variant,
hf_token = hf_token,
)
if self._gguf_path_is_diffusion(_preflight_model_path, model_identifier):
raise ValueError(
"GPU selection (gpu_ids) is not supported for a DiffusionGemma "
"GGUF on a Vulkan llama.cpp build: the diffusion runner selects "
"its device by CUDA physical index, which has no defined mapping "
"to ggml Vulkan device ordinals. Omit gpu_ids to use the default "
"device."
)
self._reject_vulkan_diffusion_gpu_ids_before_teardown(
_preflight_model_path,
model_identifier,
)
elif is_vulkan_backend and gpu_ids and gguf_path and not hf_repo:
if not Path(gguf_path).is_file():
raise FileNotFoundError(f"GGUF file not found: {gguf_path}")
self._reject_vulkan_diffusion_gpu_ids_before_teardown(
gguf_path,
model_identifier,
)
# ── Phase 1: kill old process (under lock, fast) ──────────
with self._lock:
@ -6621,18 +6635,9 @@ class LlamaCppBackend:
# Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server;
# serve them with the diffusion runner (same OpenAI-compat interface).
if self._is_diffusion:
# The diffusion runner pins its child by CUDA visibility mask, so a
# ggml Vulkan ordinal cannot be honored (wrong GPU / CPU fallback).
# Route and remote-download preflights reject before teardown; keep
# this as a final defense if classification ever disagrees.
# Final defense: route and pre-teardown preflights reject before Phase 1.
if is_vulkan_backend and gpu_ids:
raise ValueError(
"GPU selection (gpu_ids) is not supported for a DiffusionGemma "
"GGUF on a Vulkan llama.cpp build: the diffusion runner selects "
"its device by CUDA physical index, which has no defined mapping "
"to ggml Vulkan device ordinals. Omit gpu_ids to use the default "
"device."
)
raise ValueError(_VULKAN_DIFFUSION_GPU_IDS_ERROR)
# Not a tensor/layer GGUF: clear any preserved-fallback flag from a
# prior load (this path skips the command builder that clears it).
self._layer_preserves_tensor_intent = False

169
studio/backend/tests/test_gpu_memory_mode.py
View file

@ -22,6 +22,7 @@ and MoE offload itself (``--fit off``). These tests pin:
from __future__ import annotations
import inspect
import struct
import sys
import types as _types
from pathlib import Path
@ -702,6 +703,15 @@ def test_remote_vulkan_diffusion_preflight_runs_before_teardown(monkeypatch):
assert "model_path = _preflight_model_path or self._download_gguf(" in src
def test_local_vulkan_diffusion_preflight_runs_before_teardown():
src = inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model)
local_preflight = src.index(
"self._reject_vulkan_diffusion_gpu_ids_before_teardown(\n gguf_path,"
)
teardown = src.index("# ── Phase 1: kill old process")
assert local_preflight < teardown
def test_remote_vulkan_diffusion_rejection_keeps_active_server(monkeypatch):
backend = LlamaCppBackend()
killed = []
@ -737,6 +747,165 @@ def test_remote_vulkan_diffusion_rejection_keeps_active_server(monkeypatch):
assert killed == []
def test_remote_vulkan_preflight_download_failure_keeps_active_server(monkeypatch, tmp_path):
# A resolvable shard-1 file does not prove the variant is complete, so download
# failures must surface from the pre-teardown _download_gguf, not after the kill.
import hub.utils.gguf as hub_gguf
cached_shard = tmp_path / "model-00001-of-00003.gguf"
cached_shard.write_bytes(b"GGUF")
monkeypatch.setattr(
hub_gguf,
"resolve_local_gguf_path",
lambda _repo, _variant: str(cached_shard),
)
for failure in (
FileNotFoundError("shard 2 of 3 missing"),
OSError("[Errno 28] No space left on device"),
ConnectionError("hub unreachable"),
):
backend = LlamaCppBackend()
order = []
def _download(_failure = failure, **_kwargs):
order.append("download")
raise _failure
monkeypatch.setattr(backend, "_find_llama_server_binary", lambda **_kwargs: "/bin/llama")
monkeypatch.setattr(backend, "_is_vulkan_backend", lambda _binary = None: True)
monkeypatch.setattr(backend, "_get_gpu_memory", lambda _binary = None: [(0, 1024, 2048)])
monkeypatch.setattr(backend, "_download_gguf", _download)
monkeypatch.setattr(backend, "_gguf_path_is_diffusion", lambda *_args: False)
monkeypatch.setattr(backend, "_kill_process", lambda: order.append("kill"))
monkeypatch.setattr(llama_cpp_module, "_resolve_repo_id_casing", lambda repo: repo)
monkeypatch.setattr(
llama_cpp_module,
"_hf_offline_if_dns_dead",
lambda: __import__("contextlib").nullcontext(),
)
with pytest.raises(type(failure)):
backend.load_model(
hf_repo = "owner/model",
hf_variant = "Q4_K_M",
model_identifier = "owner/model",
gpu_ids = [0],
)
assert order == ["download"], failure
def test_local_vulkan_diffusion_rejection_keeps_active_server(monkeypatch, tmp_path):
gguf_path = tmp_path / "diffusion.gguf"
gguf_path.write_bytes(b"GGUF")
backend = LlamaCppBackend()
killed = []
monkeypatch.setattr(backend, "_find_llama_server_binary", lambda **_kwargs: "/bin/llama")
monkeypatch.setattr(backend, "_is_vulkan_backend", lambda _binary = None: True)
monkeypatch.setattr(backend, "_get_gpu_memory", lambda _binary = None: [(0, 1024, 2048)])
monkeypatch.setattr(backend, "_gguf_path_is_diffusion", lambda *_args: True)
monkeypatch.setattr(backend, "_kill_process", lambda: killed.append(True))
with pytest.raises(ValueError, match = "DiffusionGemma"):
backend.load_model(
gguf_path = str(gguf_path),
model_identifier = "local/diffusion",
gpu_ids = [0],
)
assert killed == []
class _ReachedServerStart(Exception):
"""Marks a load getting past the pre-teardown preflight."""
def _write_gguf_header(
path: Path,
architecture: str,
*,
diffusion: bool = False,
) -> str:
"""Smallest GGUF the header probe can classify: arch, plus the canvas marker."""
def _kv_str(key: str, value: str) -> bytes:
kb, vb = key.encode(), value.encode()
return (
struct.pack("<Q", len(kb)) + kb + struct.pack("<I", 8) + struct.pack("<Q", len(vb)) + vb
)
def _kv_u32(key: str, value: int) -> bytes:
kb = key.encode()
return struct.pack("<Q", len(kb)) + kb + struct.pack("<I", 4) + struct.pack("<I", value)
body = _kv_str("general.architecture", architecture)
if diffusion:
body += _kv_u32("diffusion.canvas_length", 256)
path.write_bytes(struct.pack("<IIQQ", 0x46554747, 3, 0, 2 if diffusion else 1) + body)
return str(path)
def _vulkan_pinned_backend(monkeypatch, killed: list) -> LlamaCppBackend:
backend = LlamaCppBackend()
monkeypatch.setattr(backend, "_find_llama_server_binary", lambda **_kwargs: "/bin/llama")
monkeypatch.setattr(backend, "_is_vulkan_backend", lambda _binary = None: True)
monkeypatch.setattr(backend, "_get_gpu_memory", lambda _binary = None: [(0, 1024, 2048)])
monkeypatch.setattr(backend, "_kill_process", lambda: killed.append(True))
return backend
def test_local_vulkan_pre_teardown_reads_the_real_gguf_header(monkeypatch, tmp_path):
# Classify from the header, not from Vulkan + gpu_ids alone: normal GGUFs load.
killed = []
backend = _vulkan_pinned_backend(monkeypatch, killed)
monkeypatch.setattr(
backend,
"_wait_for_vram_settle",
lambda **_kwargs: (_ for _ in ()).throw(_ReachedServerStart()),
)
with pytest.raises(_ReachedServerStart):
backend.load_model(
gguf_path = _write_gguf_header(tmp_path / "chat.gguf", "llama"),
model_identifier = "local/chat",
gpu_ids = [0],
)
assert killed == [True]
def test_local_vulkan_diffusion_header_rejects_before_teardown(monkeypatch, tmp_path):
# Same path, real DiffusionGemma canvas marker: rejected with the server intact.
killed = []
backend = _vulkan_pinned_backend(monkeypatch, killed)
with pytest.raises(ValueError, match = "DiffusionGemma"):
backend.load_model(
gguf_path = _write_gguf_header(tmp_path / "d.gguf", "gemma3", diffusion = True),
model_identifier = "local/diffusion",
gpu_ids = [0],
)
assert killed == []
def test_local_vulkan_missing_gguf_is_reported_before_teardown(monkeypatch, tmp_path):
# The preflight existence check must not cost the live model either.
killed = []
backend = _vulkan_pinned_backend(monkeypatch, killed)
with pytest.raises(FileNotFoundError):
backend.load_model(
gguf_path = str(tmp_path / "absent.gguf"),
model_identifier = "local/missing",
gpu_ids = [0],
)
assert killed == []
def test_start_diffusion_server_resets_tensor_parallel():
# A prior tensor-parallel chat load leaves self._tensor_parallel True (load_model
# phase 1 only kills the process, it skips the unload reset). Diffusion is never