From 7aeb83747ce0d207979079ed74abb997db604781 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 29 Jun 2026 08:05:07 +0000 Subject: [PATCH] Studio: extend CPU-force detection, advertise the capped ceiling, gate MLA MTP Follow-up review on the CPU-only / NUMA hardening: - Detect --device/-dev none as a CPU-only launch too, not just -ngl 0, and treat the GPU-layers and device controls independently (each flag's last value wins). - Lower max_available_ctx when the CPU cap reduces the launched context, so /status and the UI safe-zone advertise the real window instead of the unlaunched native size. - Skip the MTP reserve in the CPU cap and NUMA footprint when Auto drops embedded MTP for MLA models (GLM-5.2/DeepSeek/Kimi); the launch emits ngram-mod there, so reserving a target-KV copy needlessly shrank the context. GPU paths keep the existing _mtp_will_engage. - Stub structlog/loggers/httpx in the new CPU-default and NUMA test files so a dependency-light run can collect them, matching the existing abort test. The two test files now pass in a pytest-only venv (the import previously failed on structlog via loggers). Unit tests (44) and simulations pass. --- studio/backend/core/inference/llama_cpp.py | 43 ++++++++---- .../backend/tests/test_cpu_only_defaults.py | 69 +++++++++++++++++-- studio/backend/tests/test_numa_interleave.py | 44 ++++++++++++ 3 files changed, 138 insertions(+), 18 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index f21124a392..d2edfcbd45 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1115,21 +1115,23 @@ def _extra_args_n_ubatch( def _extra_args_forces_cpu_offload(extra_args: Optional[Iterable[str]]) -> bool: - """True when extras pin GPU offload to zero (-ngl 0 / --n-gpu-layers 0 / - --gpu-layers 0): the launch runs fully on CPU despite a visible GPU, so the - CPU-only safe defaults apply. Last occurrence wins, as llama-server parses it.""" + """True when extras run the launch fully on CPU despite a visible GPU, so the + CPU-only safe defaults apply: zero GPU layers (-ngl 0 / --n-gpu-layers 0 / + --gpu-layers 0) or no device (--device/-dev none). Each flag's last occurrence + wins (as llama-server parses it); the two controls are independent.""" args = [str(a) for a in extra_args] if extra_args else [] - forced = False + ngl_zero = device_none = False for i, raw in enumerate(args): flag, eq, inline = raw.partition("=") - if flag not in ("-ngl", "--n-gpu-layers", "--gpu-layers"): - continue value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") - try: - forced = int(value) == 0 - except (TypeError, ValueError): - continue - return forced + if flag in ("-ngl", "--n-gpu-layers", "--gpu-layers"): + try: + ngl_zero = int(value) == 0 + except (TypeError, ValueError): + continue + elif flag in ("--device", "-dev"): + device_none = value.strip().lower() == "none" + return ngl_zero or device_none def _build_ngram_mod_flags( @@ -5107,6 +5109,16 @@ class LlamaCppBackend: _mtp_will_engage = bool( _user_mtp_via_extras or _user_draft_via_extras or _auto_studio_mtp ) + # Auto drops embedded MTP for MLA models (GLM-5.2/DeepSeek/Kimi) unless + # forced; mirror that gate so the CPU cap / NUMA footprint don't reserve + # a target-KV copy for a drafter the launch will not start. + _mtp_will_engage_cpu = _mtp_will_engage and not ( + _auto_studio_mtp + and bool(self._nextn_predict_layers) + and self._kv_lora_rank is not None + and not bool(mtp_draft_path) + and not _mla_mtp_auto_enabled() + ) # The duplicated full target-KV copy (ctx_tgt) is an MTP-only # cost: the MTP head runs a second context over the target # model's own KV geometry. The separate-drafter spec modes @@ -5752,8 +5764,8 @@ class LlamaCppBackend: kv_on_gpu = True, # KV lives in the RAM budget we fit # mtp_engaged alone is a no-op once budget_frac is set; # pass the byte-accurate overhead so MTP KV is reserved. - mtp_engaged = _mtp_will_engage, - mtp_overhead_fn = (_mtp_bytes if _mtp_will_engage else None), + mtp_engaged = _mtp_will_engage_cpu, + mtp_overhead_fn = (_mtp_bytes if _mtp_will_engage_cpu else None), budget_frac = _CPU_RAM_BUDGET_FRAC, ) _cpu_cap = max(4096, min(_ctx_ceiling, _fit)) @@ -5766,6 +5778,9 @@ class LlamaCppBackend: _cpu_cap, ) effective_ctx = _cpu_cap + # Advertise the capped window as the ceiling too, so /status and + # the UI safe-zone don't steer back to the unlaunched native size. + max_available_ctx = min(max_available_ctx, _cpu_cap) cmd = [ binary, @@ -6037,7 +6052,7 @@ class LlamaCppBackend: try: # Recompute MTP at the post-cap context; the pre-cap # _mtp_reserve_bytes would overstate a capped million-token load. - _numa_mtp = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0 + _numa_mtp = _mtp_bytes(effective_ctx) if _mtp_will_engage_cpu else 0 _numa_footprint = ( _resident + self._estimate_kv_cache_bytes( diff --git a/studio/backend/tests/test_cpu_only_defaults.py b/studio/backend/tests/test_cpu_only_defaults.py index e03061e65b..5b3ae1500e 100644 --- a/studio/backend/tests/test_cpu_only_defaults.py +++ b/studio/backend/tests/test_cpu_only_defaults.py @@ -32,6 +32,41 @@ except ImportError: _s.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") _s.BoundLogger = type("BoundLogger", (), {}) sys.modules["structlog"] = _s +# Importing core.inference.* runs core/inference/__init__.py (orchestrator + loggers + +# httpx); stub those when absent so a dependency-light run can still collect this file. +try: + import loggers # noqa: F401 +except ImportError: + _loggers_stub = _types.ModuleType("loggers") + _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) + sys.modules["loggers"] = _loggers_stub +try: + import httpx # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "RequestError", + ): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) + _httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) + _httpx_stub.Response = type("Response", (), {}) + _httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__exit__": lambda s, *a: None, + }, + ) + sys.modules["httpx"] = _httpx_stub from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 @@ -126,7 +161,7 @@ def test_cpu_context_fit_accounts_for_mtp(): """mtp_engaged alone is a no-op once budget_frac is set, so the CPU fit must pass the byte-accurate MTP overhead fn to actually reserve MTP KV (PR review fix).""" src = _load_model_src() - assert "mtp_overhead_fn = (_mtp_bytes if _mtp_will_engage else None)" in src + assert "mtp_overhead_fn = (_mtp_bytes if _mtp_will_engage_cpu else None)" in src def test_cpu_context_cap_reuses_fit_helper_against_ram(): @@ -169,7 +204,7 @@ def test_numa_decision_uses_footprint_not_just_weights(): # Footprint = fitted weights (incl. compute buffer) + KV + MTP reserve. assert "_resident = model_size_fit or model_size" in src # MTP is recomputed at the post-cap context, not the stale pre-cap reserve. - assert "_numa_mtp = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0" in src + assert "_numa_mtp = _mtp_bytes(effective_ctx) if _mtp_will_engage_cpu else 0" in src # KV must be sized for the launched --parallel slots, not the n_parallel=1 default. assert "effective_ctx, cache_type_kv, n_parallel = n_parallel" in src @@ -191,7 +226,7 @@ def test_explicit_user_numa_skips_auto_interleave_prefix(): def test_extra_args_forces_cpu_offload_helper(): - """The zero-offload detector: -ngl 0 / --n-gpu-layers 0 / --gpu-layers 0 (last wins).""" + """The CPU-force detector: zero GPU layers or --device none (PR review fixes).""" from core.inference.llama_cpp import _extra_args_forces_cpu_offload as f assert f(["-ngl", "0"]) @@ -202,9 +237,35 @@ def test_extra_args_forces_cpu_offload_helper(): assert not f([]) assert not f(None) assert not f(["--flash-attn", "on"]) - # Last occurrence wins, matching llama-server's own parsing. + # Each flag's last occurrence wins, matching llama-server's own parsing. assert f(["-ngl", "99", "-ngl", "0"]) assert not f(["-ngl", "0", "-ngl", "99"]) + # --device/-dev none also forces CPU, independently of -ngl. + assert f(["--device", "none"]) + assert f(["-dev", "none"]) + assert f(["--device=none"]) + assert not f(["--device", "CUDA0"]) + # The two controls are independent: -ngl 0 stays CPU even with a device named. + assert f(["-ngl", "0", "--device", "CUDA0"]) + assert f(["--device", "none", "-ngl", "99"]) + + +def test_cpu_cap_lowers_advertised_ceiling(): + """When the CPU cap reduces the launched context, max_available_ctx must drop too, + so /status and the UI safe-zone reflect the real window, not native (PR review fix).""" + src = _load_model_src() + assert "max_available_ctx = min(max_available_ctx, _cpu_cap)" in src + + +def test_cpu_fit_skips_mtp_reserve_when_mla_auto_drops(): + """Auto drops embedded MTP for MLA models, so the CPU cap / NUMA footprint must not + reserve a target-KV copy for a drafter that won't launch (PR review fix).""" + src = _load_model_src() + assert "_mtp_will_engage_cpu = _mtp_will_engage and not (" in src + assert "not _mla_mtp_auto_enabled()" in src + # The CPU cap and NUMA recompute use the gated flag, not the raw _mtp_will_engage. + assert "mtp_overhead_fn = (_mtp_bytes if _mtp_will_engage_cpu else None)" in src + assert "_numa_mtp = _mtp_bytes(effective_ctx) if _mtp_will_engage_cpu else 0" in src def test_zero_offload_folds_into_cpu_only(): diff --git a/studio/backend/tests/test_numa_interleave.py b/studio/backend/tests/test_numa_interleave.py index 92377ce9db..179504badf 100644 --- a/studio/backend/tests/test_numa_interleave.py +++ b/studio/backend/tests/test_numa_interleave.py @@ -12,12 +12,56 @@ injected, so these are deterministic on any host (no /sys, no numactl needed). from __future__ import annotations import sys +import types as _types from pathlib import Path _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) +# Importing core.inference.numa runs core/inference/__init__.py (orchestrator + structlog +# + loggers + httpx); stub those when absent so a dependency-light run can collect this. +try: + import structlog # noqa: F401 +except ImportError: + _s = _types.ModuleType("structlog") + _s.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") + _s.BoundLogger = type("BoundLogger", (), {}) + sys.modules["structlog"] = _s +try: + import loggers # noqa: F401 +except ImportError: + _loggers_stub = _types.ModuleType("loggers") + _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) + sys.modules["loggers"] = _loggers_stub +try: + import httpx # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "RequestError", + ): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) + _httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) + _httpx_stub.Response = type("Response", (), {}) + _httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__exit__": lambda s, *a: None, + }, + ) + sys.modules["httpx"] = _httpx_stub + from core.inference.numa import ( # noqa: E402 InterleaveDecision, NumaTopology,