Smart GPU Management — how the app shares a small GPU

1. The problem this solves

This app runs three heavy AI models on one machine:

A typical laptop has an 8 GB GPU (RTX 4060 class) — often less, sometimes shared with the display. You cannot fit all three models on the GPU at once, and even two at once leaves no room for the activations (the temporary memory a batch of text needs while being processed). The naïve approach — just load everything onto the GPU and hope — ends in CUDA out of memory at the worst moment.

The app's answer is a Smart GPU Manager (backend/legal_retrieval/gpu_manager.py) that acts like a tiny operating system for VRAM: it decides which model sits on the GPU at any moment, evicts the right one when space runs out, and even shrinks the batch size so a burst of work never exceeds the budget.

2. The core idea in one picture

Models live in RAM. They visit the GPU only when they work. When the work is done, they come back to RAM. Re-visiting is fast — a PCIe copy (~0.1–0.3 s), not a slow reload from disk.

        HOST RAM (always holds every model's weights)
   ┌───────────────────────────────────────────────────────┐
   │  BGE-M3 embedder   BGE-Reranker   LLM (via Ollama)    │
   └───────────────────────────────────────────────────────┘
            │  promote (load weights to VRAM)
            ▼
   ┌───────────────────────────────────────────────────────┐
   │                    GPU / VRAM (8 GB)                  │
   │   ▶ only the model whose phase is ACTIVE is here ◀    │
   └───────────────────────────────────────────────────────┘
            ▲
            │  demote (move weights back to RAM)

This is called RAM-first residency: loading a model into host RAM is cheap and happens once; moving it to VRAM and back per phase is a fast copy that leaves the model resident in RAM so the next promotion is near-instant.

3. What actually happens during a chat query

Walk through one RAG search and watch the GPU change hands:

1. Embed the query → the manager promotes the embedder to VRAM, runs encode(), then (usually) demotes it back to RAM.
2. Retrieve + rerank → the manager promotes the reranker, runs predict(), demotes it.
3. LLM generates the answer → the LLM is not wrapped in promote/demote: it runs on CPU by default (PLR_LLM_NUM_GPU=0) so it never steals VRAM from the latency-critical embedding/reranking phases.

So during a single answer, VRAM is handed from the embedder to the reranker and then left alone while the LLM writes text. The GPU is never idle and never overcommitted.

Priority — who wins when space is tight

Each model has a priority (lower number = higher priority):

embedder (0)  >  reranker (1)  >  LLM (2)

When a model needs VRAM and there isn't enough, the manager evicts idle lower-priority models first. It will never evict a higher-priority model to make room for a lower one — the embedder keeps its spot even if the reranker really wants it. Among equal priorities, it evicts cold models before hot ones, then the least-recently-used.

The one exception: the LLM

Ollama owns its own weights, so the manager can't just "demote" it. Instead, when VRAM is truly tight, the manager calls ollama stop on GPU-resident LLMs to free their VRAM. It uses the CLI when available, and falls back to the Ollama HTTP API (keep_alive=0) when the CLI isn't on PATH — which is common on Windows desktop installs, exactly the low-end target this manager is built for. This is off by default in the sense that it only triggers under pressure (PLR_GPU_EVICT_OLLAMA=1).

4. Three residency modes — and "hot" models

After a phase ends, should the model stay on the GPU or go back to RAM? That's the residency mode (residency_mode), and it's the main tuning knob:

What does "hot" mean?

In adaptive mode, a model becomes hot when it is used hot_min_uses (default 2) times within hot_window_seconds (default 120 s) — a sliding window, so old uses expire. A model that keeps getting used (you're searching repeatedly) stays on the GPU and repeat queries skip the promote cost entirely. A model used once ages out and frees VRAM.

Three guards stop "hot" from becoming "greedy":

1. Pressure always wins — if free VRAM drops below keep_resident_free_mb (1 GB), even hot models are demoted.
2. Resident cap — at most max_resident (2) models stay resident at once; the least valuable idle resident is evicted to make room.
3. Never evict higher priority — the cap cannot force out the embedder to keep the reranker.

5. Auto-tuned batch sizes (the OOM safety net)

Even with one model on the GPU, a huge batch can exhaust VRAM during inference. So the manager also auto-tunes the batch size at runtime:

batch = clamp( (free_VRAM − model_footprint − safety_reserve) ÷ MB_per_sample ,
               min=1 , max=configured_batch )

Every encode() / predict() call asks the manager for a safe batch size. Free VRAM is measured live via torch.cuda.mem_get_info(), so if the display or another app is eating VRAM, the batch shrinks — and grows back when space returns. The configured batch size (embedding_batch_size, reranker_batch_size) is now just the cap.

Per-sample estimates (tunable in config.py):

If the model couldn't be promoted at all (VRAM too tight) it runs on CPU, where the VRAM budget is meaningless — the manager keeps the configured batch instead of collapsing to batch 1, so low-VRAM machines that fell back to CPU aren't slowed down twice.

6. Where the code lives

The public API for other code

pythonfrom legal_retrieval.gpu_manager import get_gpu_manager

gpu = get_gpu_manager()
gpu.register("embedder", vram_estimate_mb=1100,
             promote=my_promote_fn, demote=my_demote_fn,
             per_sample_mb=24.0)

with gpu.acquire("embedder"):          # promotes to VRAM for the phase
    batch = gpu.suggest_batch_size("embedder", 32)
    device = gpu.current_device("embedder")   # "cuda" if resident, else "cpu"
    ... run inference ...
# on exit: demoted per residency mode

7. Per-project GPU profiles

Every workspace project can carry its own GPU residency profile (mode, hot window, max resident, …), stored inside <project>/workspace/history.json under the gpu_profile key. Switching folders reloads that project's profile automatically; a project with no saved profile resets to defaults — settings never leak between projects. The GPU Residency card in the UI's Workspace panel is the friendly way to edit it.

8. Configuration reference (all env vars)

HTTP endpoints

9. How to watch it work

Watch the backend log — the manager logs its decisions:

[GPU] registered 'embedder' (priority=0, ~1082MB, per_sample=24.0MB, device=cuda)
[GPU] promoted 'embedder' to VRAM (~1082MB)
[GPU] auto-tuned embedder batch_size 32 -> 8 (free VRAM 1200MB, ~24MB/sample)
[GPU] demoted 'embedder' back to RAM
[GPU] ollama stop lfm2.5:8b (freed VRAM)

Or query the live snapshot:

bashcurl http://localhost:8765/api/gpu/status

which reports each model's state (ram/vram), hot, use_count, and per_sample_mb — everything the manager is doing right now.

10. Design rules (so future changes don't break the balance)

1. RAM-first, always. Never keep an idle model on the GPU when a lower priority model needs the space.
2. Never evict a higher-priority model for a lower one.
3. Pressure beats preference. eager/adaptive/persistent all yield to insufficient VRAM.
4. Slow operations never hold the lock. ollama stop (up to 30 s) runs outside both the manager lock and the FastAPI event loop, so a concurrent chat phase is never blocked.
5. The refcount is set before promotion, so a concurrently-acquired slot can never be evicted mid-promotion.
6. The configured batch size is a cap, never a floor. Auto-tuning only ever shrinks it.