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.
This commit is contained in:
Daniel Han
parent
2ef8fb6696
commit
7aeb83747c
3 changed files with 138 additions and 18 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue