Studio: cap GGUF context to unified memory on Apple Silicon (#6622)
* Studio: cap GGUF context to unified memory on Apple Silicon * Studio: tighten Apple ctx-cap comments and drop the overstated MLX-sync claim * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: reserve flat MTP fraction and floor sparse-KV ctx in the Apple unified-memory cap The Apple Silicon GGUF context cap mirrored the discrete-GPU auto-fit branch but missed two protections the discrete path already applies: - It passed the full unified-memory budget with budget_frac=1.0 without first reserving the flat MTP fraction the discrete path takes off via _pin_fraction. With an MTP draft whose KV cannot be byte-sized (e.g. Qwen3.6-MTP, #6529), the cap filled the whole budget and left nothing for the draft, so unified memory could still over-commit. Reserve _flat_mtp_reserve up front; this is a no-op when MTP is not engaged. - It required _can_estimate_kv(), so a GGUF with sparse KV metadata skipped the cap entirely and launched at full native context. Mirror the discrete file-size-only fallback and floor the auto context to 4096 when the cache cannot be sized. Adds regression tests for both paths. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: tighten comments in the Apple unified-memory context cap Condense the verbose comment blocks in the Apple budget helper, the no-GPU Metal branch, and the context-fit tests. Comments only, no code change (verified with ast-based comment_tools check); suite still green. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <danielhanchen@gmail.com>
This commit is contained in:
oobabooga
GitHub
parent
c3beb921f3
commit
346d96d7f2
2 changed files with 310 additions and 0 deletions
|
|
@ -803,6 +803,10 @@ _MTP_MIN_SIZE_B = 3.0
|
|||
# of free (see _fit_context_to_vram), plus a byte-accurate MTP draft reserve.
|
||||
_CTX_FIT_VRAM_FRACTION = 0.95
|
||||
|
||||
# Apple unified memory is shared with the OS, so tighter than VRAM. Matches the
|
||||
# 0.85 MLX uses in mlx_inference.py (_configure_memory_limits); not kept in sync.
|
||||
_APPLE_UNIFIED_MEMORY_FRACTION = 0.85
|
||||
|
||||
# Flat MTP reserve, used only when GGUF dims are too sparse for the byte-accurate
|
||||
# reserve (_estimate_mtp_overhead_bytes). Applied to both the fit budget and pin.
|
||||
_MTP_VRAM_RESERVE_FRAC = 0.05
|
||||
|
|
@ -2165,6 +2169,34 @@ class LlamaCppBackend:
|
|||
``_get_gpu_memory`` for callers that only need free VRAM."""
|
||||
return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()]
|
||||
|
||||
@staticmethod
|
||||
def _apple_metal_memory_budget_bytes() -> int:
|
||||
"""Unified-memory budget for GGUF context fitting on Apple Silicon.
|
||||
|
||||
No GPU is enumerated on Metal, so the context would default to native and
|
||||
over-commit unified memory ("Compute error." at decode, #5118/#6529). Use a
|
||||
fraction of MLX's Metal working-set, else total RAM; 0 off Apple Silicon or
|
||||
when unresolvable, so callers skip the cap.
|
||||
"""
|
||||
from utils.hardware import is_apple_silicon
|
||||
|
||||
if not is_apple_silicon():
|
||||
return 0
|
||||
rec_bytes = 0
|
||||
try:
|
||||
import mlx.core as mx
|
||||
if mx.metal.is_available():
|
||||
rec_bytes = int(mx.device_info().get("max_recommended_working_set_size") or 0)
|
||||
except Exception:
|
||||
rec_bytes = 0
|
||||
if rec_bytes <= 0:
|
||||
try:
|
||||
import psutil
|
||||
rec_bytes = int(psutil.virtual_memory().total)
|
||||
except Exception:
|
||||
return 0
|
||||
return int(rec_bytes * _APPLE_UNIFIED_MEMORY_FRACTION)
|
||||
|
||||
@staticmethod
|
||||
def _get_gpu_memory() -> list[tuple[int, int, int]]:
|
||||
"""Query free AND total memory per GPU.
|
||||
|
|
@ -5027,6 +5059,8 @@ class LlamaCppBackend:
|
|||
else 0.0
|
||||
)
|
||||
_pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve
|
||||
# Unified-memory budget (0 off Apple Silicon) for the no-GPU Metal cap below.
|
||||
_apple_budget_mib = self._apple_metal_memory_budget_bytes() // (1024 * 1024)
|
||||
|
||||
def _restore_after_tensor_downgrade():
|
||||
# Tensor mode dropped a quantized KV and stripped the cache
|
||||
|
|
@ -5310,6 +5344,52 @@ class LlamaCppBackend:
|
|||
# so the slider isn't on an unusable native ctx.
|
||||
effective_ctx = min(4096, effective_ctx) if effective_ctx > 0 else 4096
|
||||
|
||||
elif _apple_budget_mib > 0 and effective_ctx > 0:
|
||||
# No GPU on Metal: the branches above are skipped and the context
|
||||
# stays at native, over-committing unified memory (#5118, #6529).
|
||||
# Cap with the same fit math (--fit on stays as a backstop); only
|
||||
# auto context shrinks, explicit is honored.
|
||||
native_ctx_for_cap = self._context_length or effective_ctx
|
||||
# Reserve the flat MTP fraction up front like the discrete
|
||||
# _pin_fraction, so an unsized MTP draft (e.g. Qwen3.6-MTP, #6529)
|
||||
# can't over-commit. No-op when MTP is off; exclusive with the
|
||||
# byte-accurate _mtp_bytes reserve.
|
||||
_apple_fit_budget_mib = int(
|
||||
_apple_budget_mib * max(0.0, 1.0 - _flat_mtp_reserve)
|
||||
)
|
||||
if self._can_estimate_kv():
|
||||
cap = self._fit_context_to_vram(
|
||||
native_ctx_for_cap,
|
||||
_apple_fit_budget_mib,
|
||||
model_size_fit,
|
||||
cache_type_kv,
|
||||
n_parallel = n_parallel,
|
||||
mtp_engaged = _mtp_reserves_gpu,
|
||||
mtp_overhead_fn = mtp_overhead_fn,
|
||||
budget_frac = 1.0,
|
||||
total_mib = None,
|
||||
)
|
||||
_cap_footprint_mib = (
|
||||
model_size_fit
|
||||
+ self._estimate_kv_cache_bytes(
|
||||
cap, cache_type_kv, n_parallel = n_parallel
|
||||
)
|
||||
+ _mtp_bytes(cap)
|
||||
) / (1024 * 1024)
|
||||
# Fit returns the request unchanged when it fits OR weights
|
||||
# exceed budget; only the latter over-commits, so floor to 4096.
|
||||
max_available_ctx = (
|
||||
cap
|
||||
if _cap_footprint_mib <= _apple_fit_budget_mib
|
||||
else min(4096, native_ctx_for_cap)
|
||||
)
|
||||
else:
|
||||
# No KV estimate: mirror the discrete file-size-only fallback
|
||||
# and floor to 4096 rather than launch at native and over-commit.
|
||||
max_available_ctx = min(4096, native_ctx_for_cap)
|
||||
if not explicit_ctx:
|
||||
effective_ctx = max_available_ctx
|
||||
|
||||
# MTP reserve at the final context, for the logs below.
|
||||
_mtp_reserve_bytes = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0
|
||||
if _mtp_will_engage:
|
||||
|
|
|
|||
|
|
@ -68,6 +68,7 @@ _httpx_stub.Client = type(
|
|||
sys.modules.setdefault("httpx", _httpx_stub)
|
||||
|
||||
from core.inference.llama_cpp import (
|
||||
_APPLE_UNIFIED_MEMORY_FRACTION,
|
||||
_CTX_FIT_VRAM_FRACTION,
|
||||
LlamaCppBackend,
|
||||
classify_gpu_offload_lines,
|
||||
|
|
@ -120,6 +121,8 @@ def _drive(
|
|||
kv_per_token_bytes = 325_000,
|
||||
can_estimate_kv = True,
|
||||
extra_args = None,
|
||||
apple_budget_mib = 0,
|
||||
flat_mtp_reserve = 0.0,
|
||||
):
|
||||
"""Drive the post-metadata portion of load_model with stubbed inputs.
|
||||
|
||||
|
|
@ -223,6 +226,32 @@ def _drive(
|
|||
gpu_indices, use_fit = inst._select_gpus(model_size, gpus)
|
||||
if use_fit and not explicit_ctx:
|
||||
effective_ctx = min(FALLBACK_CTX, effective_ctx) if effective_ctx > 0 else FALLBACK_CTX
|
||||
elif apple_budget_mib > 0 and effective_ctx > 0:
|
||||
# Mirrors the Apple unified-memory branch in load_model: flat MTP reserve
|
||||
# off the budget up front (no-op at 0), sparse-KV floors to FALLBACK_CTX,
|
||||
# only auto context shrinks.
|
||||
native_ctx_for_cap = context_length or effective_ctx
|
||||
apple_fit_budget_mib = int(apple_budget_mib * max(0.0, 1.0 - flat_mtp_reserve))
|
||||
if inst._can_estimate_kv():
|
||||
cap = inst._fit_context_to_vram(
|
||||
native_ctx_for_cap,
|
||||
apple_fit_budget_mib,
|
||||
model_size,
|
||||
cache_type_kv,
|
||||
budget_frac = 1.0,
|
||||
)
|
||||
cap_footprint_mib = (model_size + inst._estimate_kv_cache_bytes(cap, cache_type_kv)) / (
|
||||
1024 * 1024
|
||||
)
|
||||
max_available_ctx = (
|
||||
cap
|
||||
if cap_footprint_mib <= apple_fit_budget_mib
|
||||
else min(FALLBACK_CTX, native_ctx_for_cap)
|
||||
)
|
||||
else:
|
||||
max_available_ctx = min(FALLBACK_CTX, native_ctx_for_cap)
|
||||
if not explicit_ctx:
|
||||
effective_ctx = max_available_ctx
|
||||
|
||||
return {
|
||||
"c_arg": effective_ctx if effective_ctx > 0 else 0,
|
||||
|
|
@ -704,3 +733,204 @@ def test_select_gpus_reserves_per_device_overhead():
|
|||
small, gpus, total_by_idx = totals, per_device_overhead_bytes = gib
|
||||
)
|
||||
assert a == [0] and b == [0]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Apple Silicon unified-memory context cap (#5118, #6529): no discrete GPU on
|
||||
# Metal, so the auto context defaulted to native and over-committed unified
|
||||
# memory. The fix budgets and caps the auto context (explicit stays verbatim).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _force_apple(monkeypatch):
|
||||
import platform as _platform
|
||||
monkeypatch.setattr(_platform, "system", lambda: "Darwin")
|
||||
monkeypatch.setattr(_platform, "machine", lambda: "arm64")
|
||||
|
||||
|
||||
def _install_fake_mlx(monkeypatch, working_set_bytes):
|
||||
"""Minimal mlx.core stub exposing metal.is_available() and device_info()."""
|
||||
mlx = _types.ModuleType("mlx")
|
||||
mlx_core = _types.ModuleType("mlx.core")
|
||||
mlx_core.metal = _types.SimpleNamespace(is_available = lambda: True)
|
||||
mlx_core.device_info = lambda: {"max_recommended_working_set_size": working_set_bytes}
|
||||
mlx.core = mlx_core
|
||||
monkeypatch.setitem(sys.modules, "mlx", mlx)
|
||||
monkeypatch.setitem(sys.modules, "mlx.core", mlx_core)
|
||||
|
||||
|
||||
class TestAppleUnifiedMemoryBudget:
|
||||
def test_zero_off_apple_silicon(self, monkeypatch):
|
||||
import platform as _platform
|
||||
|
||||
monkeypatch.setattr(_platform, "system", lambda: "Linux")
|
||||
monkeypatch.setattr(_platform, "machine", lambda: "x86_64")
|
||||
assert LlamaCppBackend._apple_metal_memory_budget_bytes() == 0
|
||||
|
||||
def test_uses_metal_working_set(self, monkeypatch):
|
||||
_force_apple(monkeypatch)
|
||||
ws = 27 * GIB # ~recommended working set on a 36 GB Mac
|
||||
_install_fake_mlx(monkeypatch, ws)
|
||||
assert LlamaCppBackend._apple_metal_memory_budget_bytes() == int(
|
||||
ws * _APPLE_UNIFIED_MEMORY_FRACTION
|
||||
)
|
||||
|
||||
def test_falls_back_to_total_ram_without_mlx(self, monkeypatch):
|
||||
_force_apple(monkeypatch)
|
||||
monkeypatch.setitem(sys.modules, "mlx", None) # import mlx.core -> ImportError
|
||||
fake_psutil = _types.ModuleType("psutil")
|
||||
fake_psutil.virtual_memory = lambda: _types.SimpleNamespace(total = 36 * GIB)
|
||||
monkeypatch.setitem(sys.modules, "psutil", fake_psutil)
|
||||
assert LlamaCppBackend._apple_metal_memory_budget_bytes() == int(
|
||||
36 * GIB * _APPLE_UNIFIED_MEMORY_FRACTION
|
||||
)
|
||||
|
||||
def test_zero_when_no_budget_resolvable(self, monkeypatch):
|
||||
_force_apple(monkeypatch)
|
||||
monkeypatch.setitem(sys.modules, "mlx", None)
|
||||
monkeypatch.setitem(sys.modules, "psutil", None)
|
||||
assert LlamaCppBackend._apple_metal_memory_budget_bytes() == 0
|
||||
|
||||
|
||||
class TestAppleContextCap:
|
||||
"""The real ``_fit_context_to_vram`` against the reporter's M3 Pro case."""
|
||||
|
||||
def test_caps_native_context_into_unified_budget(self):
|
||||
# ~15.7 GB weights at native 262144 (~16 GB KV) -> ~32 GB on a 36 GB M3
|
||||
# Pro (~23 GB budget); the fit must reduce the context to fit.
|
||||
inst = _make_backend(native_ctx = 262144)
|
||||
inst._can_estimate_kv = lambda: True
|
||||
inst._estimate_kv_cache_bytes = (
|
||||
lambda n, *a, **k: 0 if n <= 0 else int(n * 64_000) # ~16 GB @ 262144
|
||||
)
|
||||
model_size_fit = int(15.7 * GIB)
|
||||
budget_mib = int(27 * GIB * _APPLE_UNIFIED_MEMORY_FRACTION) // (1024 * 1024)
|
||||
|
||||
# The native footprint over-commits the budget -- this is the bug.
|
||||
native_footprint_mib = (model_size_fit + inst._estimate_kv_cache_bytes(262144)) // (
|
||||
1024 * 1024
|
||||
)
|
||||
assert native_footprint_mib > budget_mib
|
||||
|
||||
capped = inst._fit_context_to_vram(
|
||||
262144, budget_mib, model_size_fit, None, budget_frac = 1.0
|
||||
)
|
||||
assert capped < 262144
|
||||
capped_footprint_mib = (model_size_fit + inst._estimate_kv_cache_bytes(capped)) // (
|
||||
1024 * 1024
|
||||
)
|
||||
assert capped_footprint_mib <= budget_mib
|
||||
|
||||
|
||||
class TestAppleBranchEndToEnd:
|
||||
"""Drive the Apple elif glue (cap / floor / explicit) via _drive, no GPU."""
|
||||
|
||||
def test_auto_context_capped_below_native(self):
|
||||
plan = _drive(
|
||||
n_ctx = 0,
|
||||
model_gib = 15.7,
|
||||
gpus = [],
|
||||
native_ctx = 262144,
|
||||
kv_per_token_bytes = 64_000,
|
||||
apple_budget_mib = 23_000, # ~22 GB: weights fit, native KV doesn't
|
||||
)
|
||||
assert 0 < plan["c_arg"] < 262144
|
||||
assert plan["use_fit"] is True # --fit on still ships as a backstop
|
||||
assert plan["gpu_indices"] is None # no CUDA device pinning on Metal
|
||||
assert plan["max_available_ctx"] == plan["c_arg"]
|
||||
|
||||
def test_floors_to_fallback_when_weights_exceed_budget(self):
|
||||
# Weights alone exceed budget: ctx can't help, so floor to 4096.
|
||||
plan = _drive(
|
||||
n_ctx = 0,
|
||||
model_gib = 100,
|
||||
gpus = [],
|
||||
native_ctx = 262144,
|
||||
apple_budget_mib = 20_000,
|
||||
)
|
||||
assert plan["c_arg"] == FALLBACK_CTX
|
||||
assert plan["use_fit"] is True
|
||||
assert plan["gpu_indices"] is None
|
||||
|
||||
def test_explicit_context_honored_verbatim(self):
|
||||
# Explicit context is never shrunk, but the UI ceiling still tightens.
|
||||
plan = _drive(
|
||||
n_ctx = 200_000,
|
||||
model_gib = 15.7,
|
||||
gpus = [],
|
||||
native_ctx = 262144,
|
||||
kv_per_token_bytes = 64_000,
|
||||
apple_budget_mib = 23_000,
|
||||
)
|
||||
assert plan["c_arg"] == 200_000 # launch context honored verbatim
|
||||
assert plan["use_fit"] is True
|
||||
# Ceiling reflects the budget so the over-budget warning still fires.
|
||||
assert plan["max_available_ctx"] < 262144
|
||||
|
||||
|
||||
class TestAppleMtpFlatReserve:
|
||||
"""Apple cap reserves the flat MTP fraction up front (like _pin_fraction) so
|
||||
an unsized MTP draft (Qwen3.6-MTP, #6529) can't over-commit."""
|
||||
|
||||
def test_flat_reserve_keeps_draft_within_budget(self):
|
||||
# No reserve -> cap fills the budget, leaving nothing for the ~5% draft.
|
||||
kw = dict(
|
||||
n_ctx = 0,
|
||||
model_gib = 15.7,
|
||||
gpus = [],
|
||||
native_ctx = 262144,
|
||||
kv_per_token_bytes = 64_000,
|
||||
apple_budget_mib = 23_000,
|
||||
)
|
||||
no_reserve = _drive(**kw, flat_mtp_reserve = 0.0)
|
||||
with_reserve = _drive(**kw, flat_mtp_reserve = 0.05)
|
||||
|
||||
def footprint_mib(ctx):
|
||||
return (15.7 * GIB + ctx * 64_000) / (1024 * 1024)
|
||||
|
||||
# No reserve: main footprint + 5% draft exceeds the budget.
|
||||
assert footprint_mib(no_reserve["c_arg"]) + 0.05 * 23_000 > 23_000
|
||||
# With reserve: the cap is smaller and the full footprint fits.
|
||||
assert with_reserve["c_arg"] < no_reserve["c_arg"]
|
||||
assert footprint_mib(with_reserve["c_arg"]) + 0.05 * 23_000 <= 23_000
|
||||
|
||||
def test_no_reserve_is_a_noop_when_mtp_absent(self):
|
||||
# flat_mtp_reserve == 0 (the common, non-MTP case) must not change the cap.
|
||||
kw = dict(
|
||||
n_ctx = 0,
|
||||
model_gib = 15.7,
|
||||
gpus = [],
|
||||
native_ctx = 262144,
|
||||
kv_per_token_bytes = 64_000,
|
||||
apple_budget_mib = 23_000,
|
||||
)
|
||||
assert _drive(**kw, flat_mtp_reserve = 0.0) == _drive(**kw)
|
||||
|
||||
|
||||
class TestAppleNoKvMetadataFloor:
|
||||
"""Sparse KV metadata floors the auto context to FALLBACK_CTX (like the
|
||||
discrete file-size-only fallback) instead of launching at native."""
|
||||
|
||||
def test_sparse_kv_floors_auto_context(self):
|
||||
plan = _drive(
|
||||
n_ctx = 0,
|
||||
model_gib = 15.7,
|
||||
gpus = [],
|
||||
native_ctx = 262144,
|
||||
can_estimate_kv = False,
|
||||
apple_budget_mib = 23_000,
|
||||
)
|
||||
assert plan["c_arg"] == FALLBACK_CTX # not native 262144
|
||||
assert plan["use_fit"] is True
|
||||
assert plan["gpu_indices"] is None
|
||||
|
||||
def test_sparse_kv_still_honors_explicit_context(self):
|
||||
plan = _drive(
|
||||
n_ctx = 100_000,
|
||||
model_gib = 15.7,
|
||||
gpus = [],
|
||||
native_ctx = 262144,
|
||||
can_estimate_kv = False,
|
||||
apple_budget_mib = 23_000,
|
||||
)
|
||||
assert plan["c_arg"] == 100_000 # explicit honored even without KV sizing
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue