fix: reserve memory headroom before prefill - #2251
Open
abendrothj wants to merge 2 commits into
Open
Conversation
golebier
pushed a commit
to golebier/exo
that referenced
this pull request
Aug 20, 2026
…ore#11) Combines PR exo-explore#2251 ("evict before prefill") with oMLX's preflight admission + per-chunk EWMA transient tracking, and adds placement-time KV/activation headroom (exo-explore#1709, exo-explore#2240, exo-explore#2241). The admission ceiling is a port of oMLX's reclaim-based model (min(static, dynamic, metal_cap)) so a model that legitimately fills 80%+ of memory is no longer rejected on every prefill. Phases: - Phase 1: evict_for_prefill_headroom + preflight_or_raise + estimate_prefill_peak_bytes + PrefillMemoryExceededError (cache.py, exceptions.py). - Phase 2: per-chunk EWMA transient guard (PrefillTransientTracker, guard_prefill_chunk_or_raise) wired through prefill()'s progress callback (generate.py, batch_generate.py). - Phase 3: placement reserves weights + KV(context) + activation margin per node (placement_memory.py); EXO_PLACEMENT_CONTEXT_TOKENS lever; num_attention_heads added to the model card. Reclaim-based ceiling (memory_guard.py, port of oMLX ProcessMemoryEnforcer): hard_limit = min(static, dynamic, metal_cap) dynamic = phys_footprint + free + inactive + active*reclaim_ratio Honours iogpu.wired_limit_mb (Metal's max_recommended_working_set_size does not). Tiers (safe/balanced/aggressive) via EXO_MEMORY_GUARD_TIER. Two deployment regressions found and fixed: - Over-aggressive admission ceiling (flat 0.75 fraction rejected a 201 GiB model on a 256 GiB box) -> reclaim-based ceiling. - SIGBUS from a truncated _RusageInfoV4 ctypes struct (80 bytes vs the 296 bytes proc_pid_rusage writes) -> full 36-field struct. Ship default: OFF (opt-in via EXO_ENABLE_PREFILL_GUARD=1 or the UI toggle). Runtime on/off toggle (PUT /v1/memory-guard) in the dashboard Advanced Options, plus a build-version badge in the top-left (GET /v1/version, EXO_APP_VERSION injected by the Swift host). 69 new tests (prefill-headroom, memory-guard, placement-memory). basedpyright 0 errors, ruff clean, 173 passed. Build tag: 1.0.72-memory-headroom-dev4 Adapted from oMLX (jundot/omlx, commit c1a3d44) and upstream PR exo-explore#2251.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Persistent prefix-cache allocations can leave too little room for temporary prefill activations. Existing eviction only runs while adding a completed cache entry, after the dangerous allocation has already happened.
Changes
EXO_PREFILL_MEMORY_THRESHOLD, defaulting to 10 percentage points belowEXO_MEMORY_THRESHOLDTesting
uv run pytest src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py -m "not slow"uv run basedpyright src/exo/worker/engines/mlx/cache.py src/exo/worker/engines/mlx/generator/batch_generate.pyuv run ruff check src/exo/worker/engines/mlx/cache.py src/exo/worker/engines/mlx/generator/batch_generate.py src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.pyFixes #2182
Refiled from #2212, which was closed unintentionally.