How LRU Eviction Works for Multi-Model Serving in oMLX

oMLX manages limited GPU memory during multi-model serving by evicting the least-recently-used (LRU) model whenever a new model request exceeds the configured memory budget.

When deploying multiple large language models concurrently on Apple Silicon GPUs, memory constraints become the primary bottleneck. The jundot/omlx project solves this through an EnginePool class that implements LRU (Least Recently Used) eviction to enforce strict max-model-memory limits while keeping frequently accessed models ready for inference. This deterministic approach ensures that idle models are seamlessly unloaded to make room for new workloads without manual intervention.

The Eviction Pipeline in omlx/engine_pool.py

The core logic resides in omlx/engine_pool.py, where the EnginePool coordinates model loading across seven distinct phases:

1. Request Processing: When await pool.get_engine(model_id) is invoked (lines 298-311), the pool checks if the model exists in memory. If found, it refreshes the last_access timestamp to time.time() (lines 30-33).

2. Pre-Load Memory Check: Before allocating space, _ensure_memory_available() calculates whether the requested model fits within the remaining budget, accounting for a 25% KV cache headroom on non-audio models (lines 18-31).

3. Victim Selection: If space is needed, _find_lru_victim() scans all EngineEntry objects, skipping entries where is_pinned is True or where engine.has_active_requests() returns True. It selects the entry with the smallest last_access value (lines 45-72).

4. Memory Reclamation: The selected victim is passed to _unload_engine(), which stops the engine instance, triggers a Metal-level memory barrier via mx.synchronize(); mx.clear_cache(), and decrements self._current_model_memory (lines 73-99).

5. Iterative Eviction: The eviction loop repeats until sufficient memory is freed or no evictable models remain, at which point InsufficientMemoryError is raised (lines 30-38).

6. Model Loading: With space available, _load_engine() instantiates the appropriate backend (BatchedEngine, VLMBatchedEngine, etc.) and initializes last_access to the current time.

7. Process-Wide Enforcement: When configured, process_memory_enforcer triggers the same eviction path if the Metal process exceeds system-wide limits, reusing _find_lru_victim() and _unload_engine() to reclaim space.

Key Safeguards and Configuration Details

LRU Timestamp Tracking

Each loaded model is wrapped in an EngineEntry dataclass that stores last_access as a float. This timestamp updates on every successful retrieval via get_engine(), ensuring that frequently accessed models remain resident even if they were loaded early in the session.

Pinned Models and Active Request Protection

The eviction logic respects two critical protections to prevent service disruption:

• Pinned models: Entries with is_pinned = True are permanently resident and excluded from victim selection entirely. These are typically specified via the pinned_models parameter during model discovery.
• Active requests: The code checks engine.has_active_requests() before unloading. Models currently serving inference are skipped to avoid interrupting ongoing generations, forcing the pool to select the next oldest idle model.

KV Cache Headroom

For transformer-based models (excluding audio), the pool reserves 25% of the model's memory footprint as kv_headroom during availability calculations. This ensures that evicting a model leaves sufficient space for remaining loaded models to allocate their KV caches during text generation, preventing out-of-memory errors during inference.

Implementation Example

Configure an EnginePool with a 4 GB limit and protected models:

from omlx.engine_pool import EnginePool
from omlx.scheduler import SchedulerConfig

# Initialize pool with 4GB max model memory

pool = EnginePool(
    max_model_memory=4 * 1024**3, 
    scheduler_config=SchedulerConfig()
)

# Discover models, pinning specific ones to prevent eviction

pool.discover_models("/path/to/models", pinned_models=["llama-3b"])

# Request triggers automatic LRU eviction if memory is tight

engine = await pool.get_engine("qwen-7b")

Inspect the current LRU order for debugging purposes:

# View timestamps and pin status for all entries

for entry_id, entry in pool._entries.items():
    status = "pinned" if entry.is_pinned else "evictable"
    print(f"{entry_id}: last_access={entry.last_access}, {status}")

Summary

• LRU eviction in oMLX is centralized in the EnginePool class within omlx/engine_pool.py, which maintains a memory-bounded cache of loaded models using precise timestamp tracking.
• The eviction process uses _find_lru_victim() to select the oldest idle model for unloading via _unload_engine(), which explicitly clears Metal caches with mx.synchronize() and mx.clear_cache().
• Safety mechanisms include pinned models (is_pinned), active request detection (has_active_requests()), and 25% KV cache headroom to prevent performance degradation or request interruption.
• Both per-pool limits (max_model_memory) and process-wide limits (process_memory_enforcer) utilize the same LRU infrastructure to maintain GPU memory stability.

Frequently Asked Questions

What happens if no models can be evicted?

If _find_lru_victim() returns None because all loaded models are either pinned or actively serving requests, the pool raises InsufficientMemoryError and refuses to load the new model, preventing out-of-memory crashes.

How does pinning affect memory management?

Pinned models (specified via the pinned_models list in discover_models()) set is_pinned=True on their EngineEntry, causing the LRU scanner to skip them entirely. This ensures critical models remain resident even if they haven't been accessed recently, effectively allowing you to keep a "hot" subset of models permanently in GPU memory.

Can eviction interrupt ongoing inference?

No. Before selecting a victim, the code checks engine.has_active_requests(). Models with active generations are excluded from eviction to prevent request interruption, forcing the pool to find the next oldest idle model or raise an error if no suitable victim exists.

Where is the LRU logic tested?

The eviction logic is validated in tests/test_engine_pool.py (for victim selection and edge cases) and tests/test_process_memory_enforcer.py (for process-wide memory limit scenarios), ensuring deterministic behavior under memory pressure and verifying that pinned models and active requests are properly respected.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}