Give the expert cache's VRAM back without unloading the model (#533)

CMakeLists.txt

@@ -109,6 +109,7 @@ if(STRATA_ENABLE_CUDA)
   endforeach()
   message(STATUS "Strata: CUDA enabled, arch ${CMAKE_CUDA_ARCHITECTURES}")
   set(_strata_gpu_runtime_target CUDA::cudart)
+  set(_strata_gpu_driver_target CUDA::cuda_driver)
   set(_strata_gpu_blas_target CUDA::cublas)
   set(_strata_gpu_include_directories ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
 elseif(STRATA_ENABLE_HIP)
@@ -342,7 +343,7 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
     src/core/expert_cache.cpp src/core/native_head.cpp src/core/native_dense.cpp src/core/verify.cpp src/core/mtp.cpp
     src/core/conversation_snapshot.cpp src/core/conversation_state.cpp src/core/conversation_memory.cpp)
   target_include_directories(strata_engine PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
-  target_link_libraries(strata_engine PUBLIC strata_core strata_kernels strata_kernels_cpu ${_strata_gpu_runtime_target})
+  target_link_libraries(strata_engine PUBLIC strata_core strata_kernels strata_kernels_cpu ${_strata_gpu_runtime_target} ${_strata_gpu_driver_target})
   if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tests/core/file_expert_source_test.cpp")
     enable_testing()
     add_executable(file_expert_source_test tests/core/file_expert_source_test.cpp)

docs/DETAILS.md

@@ -397,11 +397,12 @@ Terminal chat: `.venv/bin/python chat.py`.
 ## Sharing the GPU with other programs (optional)
 
 By default the model stays loaded until you close Strata. On a PC that also games, renders or runs another model
-server, three server options (all off by default; also as keys in `strata-<model>.json`) give the VRAM back:
+server, four server options (all off by default; also as keys in `strata-<model>.json`) give the VRAM back:
 
 | Option | Config key | What it does |
 | --- | --- | --- |
 | `--idle-unload 600` | `"idle_unload_s": 600` | unload the model after 600 s without requests; the next request loads it again |
+| `--cache-release-idle 300` | `"cache_release_idle_s": 300` | give the expert cache's VRAM back after 300 s without requests and keep the model loaded; the next request takes it again in a few seconds (see below) |
 | `--min-free-vram-mib 11000` | `"min_free_vram_mib": 11000` | load an unloaded model only when that much VRAM is free (it waits up to 15 s for memory being given back), else answer **503** "the GPU is in use by another program" instead of starting into what a game left (with several GPUs it checks the first one) |
 | `--before-load "cmd"` | `"before_load": "cmd"` or `["cmd", "arg"]` | a command run before the model is loaded again, e.g. one that unloads another server's model |
 
@@ -413,6 +414,31 @@ the OS file cache, so loading again takes seconds while that RAM is not needed e
 16 GB with Q2_0 in the low-RAM mode: unloading takes ~0.3 s, and a request to an unloaded model answered after
 4.6 s (text) or 14.7 s (a picture, image encoder on the CPU).
 
+**Give back only the expert cache (opt-in, #533):** the expert cache is most of the VRAM a MoE model takes, and it is
+only a copy: every expert is also in RAM. With `--expert-cache-release` in the engine's arguments the engine allocates
+the cache so that it can hand the memory back to the OS and take it again while the model, the context and the
+conversation cache stay loaded. `POST /cache/release` gives it back now and `POST /cache/refill` takes it again
+(`409` while a request is running, `400` when the engine cannot, e.g. without the engine option); a request that
+arrives while it is given back refills it first, so a client needs nothing but the usual requests. The server option
+`--cache-release-idle` adds the engine option itself and gives the cache back after that many seconds without
+requests. The cache keeps its address while it is given back, so the CUDA graphs and the experts' pointers stay valid
+and a refill copies the same experts back into the same slots.
+
+Measured on an RTX 3090 24 GB (Windows 11, driver 616.56) with a Qwen3.8-Flash-Next fine-tune at IQ2_XS, `--expert-cache auto`:
+
+| | |
+| --- | --- |
+| VRAM given back | 14.5 GiB (10,800 experts) in ~50 ms |
+| refill | 1.7-2.5 s from RAM; 0 of 10,800 slots differed from their RAM copy after 3 release/refill cycles (`STRATA_VERIFY_REFILL=1` checks every slot) |
+| a ComfyUI video render with the model loaded | 176 s with the cache in VRAM, 128 s with it given back |
+| decoding after a refill | ~72 tok/s (a 400-token answer) |
+
+Only the expert cache is given back: the dense weights, the KV cache and the engine's buffers stay (the card then
+held about 7 GiB, the desktop included). It works on NVIDIA (CUDA's virtual memory API); on AMD (HIP), with a layer
+split and in the resident RAM mode the engine says so and keeps the cache as before. Refill when the VRAM is free: on
+Windows a refill into a card another program still fills makes the driver move memory to shared system RAM, and
+decoding then ran at ~10-17 tok/s in the same setup until the other program let go.
+
 **Keep what the expert cache learned across restarts (opt-in, engine 0.1.36, #477):** a start fills the GPU's expert
 cache from the shipped profile, and the adaptive tier (`--adapt-every`) then moves in the experts your requests use.
 With `"expert_profile_save": "expert-profile-learned.bin"` in `strata-<model>.json` the engine saves that as a

include/strata/core/expert_cache.hpp

@@ -150,6 +150,19 @@ class ExpertCache {
     bool fill_slot_queued(int32_t slot, const uint8_t* host_blob, std::string& err, int64_t bytes = 0);
     bool sync_queued(std::string& err);
 
+    /// Before open(): back the slot arena with CUDA virtual memory, so release_memory()/restore_memory() can give the
+    /// VRAM back and take it again AT THE SAME DEVICE ADDRESS.  Off by default (plain cudaMalloc, unchanged).
+    void set_releasable(bool on) { releasable_ = on; }
+    /// True when the arena was allocated releasable (the device supports VMM and set_releasable(true) was called).
+    bool can_release() const { return vmm_; }
+    bool released() const { return released_; }
+    /// Unmaps and frees the arena's physical memory; the address range stays reserved, slot_of() and the residency
+    /// table are untouched.  The caller makes sure nothing reads the slots until restore_memory() succeeds.
+    bool release_memory(std::string& err);
+    /// Maps new physical memory at the same address.  The slot CONTENTS are undefined afterwards: the caller refills.
+    /// On failure the cache stays released and err says why (e.g. the VRAM is still in use by another program).
+    bool restore_memory(std::string& err);
+
     /// Reads `slot` back to the host and compares it to `host_blob`, byte for byte.  **THE ONLY THING THAT SAYS
     /// THE CACHE HOLDS THE EXPERT IT CLAIMS TO.**  A slot table that is right about indices and wrong about
     /// bytes produces a plausible token, which is exactly the failure this project has paid for most often.
@@ -177,6 +190,16 @@ class ExpertCache {
     bool per_layer_ = false;
     std::vector<int32_t> layer_next_;   ///< [n_layers] -> that layer's next free slot
     std::vector<uint64_t> off_;         ///< plan v0.3 P6: slot offsets (slots + 1 entries) when sized
+    /// --expert-cache-release: the VMM arena (CUDA only; a plain cudaMalloc when the device has no VMM).
+    bool releasable_ = false;
+    bool vmm_ = false;
+    bool released_ = false;
+    uint64_t vmm_handle_ = 0;
+    uint64_t vmm_bytes_ = 0;
+    int vmm_device_ = 0;
+    bool vmm_alloc(uint64_t bytes, std::string& err);
+    bool vmm_map(std::string& err);
+    void vmm_free();
     int64_t admitted_ = 0;
 };
 

serve/server.py

@@ -430,6 +430,39 @@ def restart(self):
         self.__init__(*self.spawn)
         self.info = {**info, **self.info}
 
+    def cache_command(self, cmd: str, timeout: float = 120.0) -> dict:
+        """Send RELEASE or REFILL to the engine and read its one-line answer.
+
+        Returns {"status": "released", "experts": N, "freed_mib": N, "ms": F} or
+        {"status": "refilled", "experts": N, "ms": F}.
+        Raises ValueError on an ERR line, EngineDied if the engine is gone.
+        """
+        try:
+            self.proc.stdin.write(f"{cmd}\n")
+            self.proc.stdin.flush()
+        except OSError:                                  # the pipe is gone: the engine died (not the client)
+            raise EngineDied(f"the engine stopped unexpectedly (exit code {self.exit_code()})") from None
+        deadline = time.monotonic() + timeout
+        while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise ValueError(f"the engine did not answer {cmd} within {timeout:.0f} s")
+            try:
+                line = self.lines.get(timeout=remaining)
+            except queue.Empty:
+                raise ValueError(f"the engine did not answer {cmd} within {timeout:.0f} s")
+            if line is None:
+                raise EngineDied(f"the engine stopped unexpectedly (exit code {self.exit_code()})")
+            if line.startswith("RELEASED "):
+                f = line.split()
+                return {"status": "released", "experts": int(f[1]), "freed_mib": int(f[2]), "ms": float(f[3])}
+            if line.startswith("REFILLED "):
+                f = line.split()
+                return {"status": "refilled", "experts": int(f[1]), "ms": float(f[2])}
+            if line.startswith("ERR "):
+                raise ValueError(line[4:].strip())
+            # any other line (a heartbeat, a log line) is skipped
+
     def _parse_done(self, line):
         f = line.split()
         self.last = {"generated": int(f[1]), "prompt_tokens": int(f[2]), "prompt_ms": float(f[3]),
@@ -797,6 +830,9 @@ def engine_args(cfg: dict) -> list[str]:
     # opt-in: an auto split runs on the first card alone when it holds every profiled expert and the KV
     if len(gpu_list(cfg)) > 1 and cfg.get("split_skip_if_fits") and "--split-skip-if-fits" not in args:
         args.append("--split-skip-if-fits")
+    # giving the expert cache back when idle needs the engine to allocate it so that it can be released
+    if cfg.get("cache_release_idle_s") and "--expert-cache-release" not in args:
+        args.append("--expert-cache-release")
     return learned_profile_args(cfg, args)
 
 
@@ -981,6 +1017,8 @@ def __init__(self, engine: Engine, tokenizer, template: ChatTemplate, model_name
         # only start it again when this much VRAM is free, and run this command first (e.g. to unload another
         # server's model); the next request after an unload starts the engine again
         self.idle_unload_s = 0
+        self.cache_release_idle_s = 0                   # give the expert cache's VRAM back after this many idle s
+        self.cache_released_at = None                   # when the cache was last given back (None: it is in VRAM)
         self.min_free_vram_mib = 0
         self.before_load = None
         self.reasoning_budget_tokens = 0                 # #123: the config's default thinking budget (0: none)
@@ -1137,6 +1175,31 @@ def unload(self, idle_for: float | None = None) -> str:
         finally:
             self.fifo.release()
 
+    def cache(self, action: str, idle_for: float | None = None) -> dict:
+        """POST /cache/release and /cache/refill: give the expert cache's VRAM back to the OS, or take it again.
+        Returns the engine's answer, or {"status": "busy"} / {"status": "not loaded"} / {"status": "unsupported"}
+        (and "busy" with idle_for when a request ran more recently than that)."""
+        if not hasattr(self.engine, "cache_command"):
+            return {"status": "unsupported"}
+        if not self.fifo.acquire(blocking=False):
+            return {"status": "busy"}
+        try:
+            if not self.engine.alive():
+                return {"status": "not loaded"}
+            with self.status_lock:
+                if self.status.get("busy") or self.status.get("queued"):
+                    return {"status": "busy"}
+            if idle_for is not None and time.time() - (self.last_request_at or self.started_at) < idle_for:
+                return {"status": "busy"}
+            result = self.engine.cache_command("RELEASE" if action == "release" else "REFILL")
+            # the engine refills by itself before the next request, which also moves last_request_at past this
+            self.cache_released_at = time.time() if action == "release" else None
+            print(f"[strata] expert cache {result['status']}{f' after {idle_for:.0f} s idle' if idle_for else ''}: "
+                  f"{result.get('experts', 0)} experts", flush=True)
+            return result
+        finally:
+            self.fifo.release()
+
     def start_idle_unload(self):
         if not self.idle_unload_s or not hasattr(self.engine, "unload"):
             return
@@ -1151,6 +1214,27 @@ def loop():
                     print(f"[strata] idle unload: {e}", flush=True)
         threading.Thread(target=loop, daemon=True).start()
 
+    def start_idle_cache_release(self):
+        if not self.cache_release_idle_s or not hasattr(self.engine, "cache_command"):
+            return
+        print(f"[strata] the expert cache is given back after {self.cache_release_idle_s} s without requests; the "
+              "next request takes it again", flush=True)
+
+        def loop():
+            while True:
+                time.sleep(max(1.0, min(30.0, self.cache_release_idle_s / 4)))
+                at = self.cache_released_at
+                if at is not None and at >= (self.last_request_at or self.started_at):
+                    continue                            # already given back, and no request since
+                try:
+                    self.cache("release", idle_for=self.cache_release_idle_s)
+                except ValueError as e:                 # the engine refused (e.g. a layer split): stop trying
+                    print(f"[strata] idle cache release: {e}; not tried again", flush=True)
+                    return
+                except (EngineDied, EngineStuck) as e:  # tried again at the next turn; the thread keeps running
+                    print(f"[strata] idle cache release: {e}", flush=True)
+        threading.Thread(target=loop, daemon=True).start()
+
     def set_shared(self, defaults) -> dict:
         """The Chat settings every client gets for what it leaves out; {} / None = clients use their own again."""
         self.shared = clean_shared_defaults(defaults)
@@ -2181,6 +2265,17 @@ def do_POST(self):
                     return
                 self._json(409 if r == "busy" else 200, {"status": r})
                 return
+            if path in ("/cache/release", "/cache/refill"):   # give the expert cache's VRAM back, or take it again
+                try:
+                    r = svc.cache("release" if path == "/cache/release" else "refill")
+                except ValueError as e:
+                    self._json(400, {"error": {"type": "invalid_request_error", "message": str(e)}})
+                    return
+                except EngineDied as e:
+                    self._json(503, {"error": {"type": "server_error", "message": str(e)}})
+                    return
+                self._json(409 if r.get("status") == "busy" else 200, r)
+                return
             if path == "/load":                              # load now, e.g. ahead of a request
                 try:
                     svc.load()
@@ -2667,6 +2762,10 @@ def main() -> int:
                     help="unload the model after this many seconds without requests, so other programs (games, other "
                          "model servers) can use the VRAM; the next request loads it again (also \"idle_unload_s\" "
                          "in the config; default: never)")
+    ap.add_argument("--cache-release-idle", type=float, default=None, metavar="SECONDS",
+                    help="give the expert cache's VRAM back after this many seconds without requests and keep the "
+                         "model loaded, so other programs can use that VRAM; the next request takes it again in a "
+                         "few seconds (also \"cache_release_idle_s\" in the config; default: never)")
     ap.add_argument("--min-free-vram-mib", type=int, default=None,
                     help="load an unloaded model only when this much VRAM is free, else answer 503 (also "
                          "\"min_free_vram_mib\" in the config; default: always load)")
@@ -2676,6 +2775,8 @@ def main() -> int:
     cfg = json.loads(Path(a.config).read_text(encoding="utf-8-sig")) if a.config else {}   # Notepad adds a BOM
     if a.gpu is not None:
         cfg["gpu"] = int(a.gpu) if a.gpu.strip().isdigit() else a.gpu
+    if a.cache_release_idle is not None:
+        cfg["cache_release_idle_s"] = a.cache_release_idle
     a.host = a.host or cfg.get("host") or "127.0.0.1"   # issue #26: the run scripts pass no --host, the config can
     try:                                                # before the minutes of loading: is the port free?
         Server((a.host, a.port), BaseHTTPRequestHandler).server_close()
@@ -2766,6 +2867,7 @@ def main() -> int:
               + ("" if svc.api_key or "*" not in svc.cors_origins else
                  " - WARNING: any web page may use the model (no API key)"), flush=True)
     svc.idle_unload_s = a.idle_unload if a.idle_unload is not None else float(cfg.get("idle_unload_s") or 0)
+    svc.cache_release_idle_s = float(cfg.get("cache_release_idle_s") or 0)
     svc.min_free_vram_mib = a.min_free_vram_mib if a.min_free_vram_mib is not None else \
         int(cfg.get("min_free_vram_mib") or 0)
     svc.before_load = a.before_load or cfg.get("before_load") or None
@@ -2803,6 +2905,7 @@ def main() -> int:
         atexit.register(hub.close)                      # the servers Strata started end with it
     httpd = serve(svc, host=a.host, port=a.port)
     svc.start_idle_unload()
+    svc.start_idle_cache_release()
     here = "127.0.0.1" if a.host in ("0.0.0.0", "", "::") else a.host
     print(f"ready: http://{here}:{a.port}/v1  (OpenAI: /v1/chat/completions, Anthropic: /v1/messages, "
           f"context {engine.max_context} tokens{', images on' if vision else ''}"